fix flaky

.github/workflows/benchmark.yml

@@ -168,7 +168,7 @@ jobs:
       run: BENCHMARK_LOG=cifar_10steps STEPS=10 python3 examples/hlb_cifar10.py
     - name: Run 10 CIFAR training steps w HALF
       env:
-        ASSERT_MIN_STEP_TIME: ${{ matrix.dev == 'NV' && '120' || matrix.dev == 'AMD' && '230' || '3000' }}
+        ASSERT_MIN_STEP_TIME: ${{ matrix.dev == 'NV' && '120' || matrix.dev == 'AMD' && '235' || '3000' }}
       run: BENCHMARK_LOG=cifar_10steps_half STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
     - name: Run full CIFAR training w 1 GPU
       # slow on metal

test/null/test_simplify_valid_idx.py

@@ -1,6 +1,6 @@
 import unittest, itertools
 
-from tinygrad.codegen.late.devectorizer import indexing_simplify
+from tinygrad.codegen.late.coalese import indexing_simplify
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite
 from tinygrad.uop.symbolic import simplify_valid, sym, pm_move_where_on_load
@@ -495,7 +495,7 @@ class TestImageSimplification(unittest.TestCase):
 class TestDropTrueGate(unittest.TestCase):
   def test_drop_true_gate_on_index(self):
     # test that INDEX with a constant True valid gets simplified to drop the valid
-    from tinygrad.codegen.late.devectorizer import indexing_simplify
+    from tinygrad.codegen.late.coalese import indexing_simplify
     from tinygrad.uop.ops import graph_rewrite
     from tinygrad.uop.symbolic import sym
     buf = UOp.param(0, dtypes.int.ptr())

tinygrad/codegen/__init__.py

@@ -16,7 +16,7 @@ from tinygrad.uop.symbolic import sym, symbolic_simple, symbolic, pm_move_where_
 from tinygrad.codegen.decomp.dtype import pm_dtype_decomps
 from tinygrad.codegen.decomp.op import get_late_rewrite_patterns, get_simplifying_rewrite_patterns
 from tinygrad.codegen.decomp.transcendental import get_transcendental_patterns
-from tinygrad.codegen.late.devectorizer import indexing_simplify
+from tinygrad.codegen.late.coalese import indexing_simplify
 from tinygrad.codegen.opt.postrange import apply_opts
 from tinygrad.codegen.late.gater import pm_move_gates_from_index
 from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_split_ranges, pm_load_collapse

tinygrad/codegen/late/coalese.py

@@ -1,11 +1,70 @@
 from typing import Any
-import itertools
+import itertools, functools
 from collections import defaultdict
-from tinygrad.dtype import dtypes, AddrSpace, Invalid, ImageDType
-from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat
-from tinygrad.helpers import getenv, IMAGE
+from tinygrad.dtype import dtypes, AddrSpace, Invalid, ImageDType, DType
+from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat, GroupOp
+from tinygrad.uop.symbolic import uop_given_valid, parse_valid, invalid_gate
+from tinygrad.helpers import getenv, IMAGE, OSX, ceildiv
 from tinygrad.renderer import Renderer
-from tinygrad.codegen.late.devectorizer import image_valid_dims, _drop_valid_stmts, uop_given_valid
+
+# ***** image load valid simplification *****
+
+@functools.cache
+def _drop_valid_stmts(valid:UOp, idx:UOp, height:int, width:int) -> list[UOp]:
+  # can drop valid if idx is out of bound when valid is False
+  drop_stmt = []
+  for i,stmt in enumerate(valid.split_uop(Ops.AND)):
+    if (res:=parse_valid(stmt)) is None: continue
+    X, is_upper_bound, c = res
+
+    # for X0 + X1 + ... >= 1, check if it's out of bound when Xi = 0 for all i
+    if not is_upper_bound and c == 1 and all(u.op in GroupOp.Irreducible and u.vmin == 0 for u in X.split_uop(Ops.ADD)):
+      testidx = functools.reduce(lambda nowidx,u: nowidx.substitute({u:u.const_like(0)}), X.split_uop(Ops.ADD), idx)
+      if testidx.index(0).vmax < 0 or testidx.index(1).vmax < 0:
+        drop_stmt.append(stmt)
+        continue
+
+    # check if idx is out of bound when X is on the wrong side of the bound: X in [c+1, vmax] or [vmin, c-1]
+    lo, hi = (c + 1, X.vmax) if is_upper_bound else (X.vmin, c - 1)
+    if lo <= hi:
+      fake = UOp.variable(f"fake{i}", lo, hi, X.dtype)
+      for coord,b in zip(idx.src, (width, height)):
+        rw = coord.substitute({X:fake}).simplify()
+        if rw.vmin >= b or rw.vmax < 0:
+          drop_stmt.append(stmt)
+          break
+  return drop_stmt
+
+def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None:
+  idx = uop_given_valid(valid, start_idx)
+  return None if idx is start_idx else buf.index(idx.valid(valid), ptr=True)
+
+def simplify_valid_image_load(buf:UOp, idx_y:UOp, idx_x:UOp, valid:UOp) -> UOp|None:
+  if not isinstance(buf.dtype, ImageDType): return None
+  start_idx = idx_x._stack(idx_y)
+  idx = uop_given_valid(valid, start_idx)
+  drop_stmt = _drop_valid_stmts(valid, idx, buf.dtype.shape[0], buf.dtype.shape[1])
+
+  if not drop_stmt and idx is start_idx: return None
+  new_valid = UOp.uprod(*ss) if (ss:=[s for s in valid.split_uop(Ops.AND) if s not in drop_stmt]) else None
+  idx_y, idx_x = idx.index(1), idx.index(0)
+  return buf.index(idx_y.valid(new_valid), idx_x.valid(new_valid), ptr=True) if new_valid is not None else buf.index(idx_y, idx_x, ptr=True)
+
+indexing_simplify = PatternMatcher([
+  # image load valid idx simplification
+  (UPat(Ops.INDEX, src=(UPat.var("buf"), invalid_gate)), lambda buf,x,i,cond: simplify_valid_load(buf, x, cond)),
+  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("valid").where(UPat.var("idx_y"), UPat(arg=Invalid)),
+                                         UPat.var("valid").where(UPat.var("idx_x"), UPat(arg=Invalid)))), simplify_valid_image_load),
+])
+
+# get list of (height, width) that do not require pitch padding
+def image_valid_dims(base:DType, size:int, arch:str) -> list[tuple[int,int]]:
+  if (ALIGN:=next((int(p.split('=')[1]) for p in arch.split(',') if p.startswith("IMAGE_PITCH_ALIGNMENT=")), 0)) == 0: return []
+  MAXW, pxls = 16384, size // 4
+  if base not in (dtypes.half, dtypes.float) or size > 4*MAXW*MAXW: return []
+  # height=1 images just need to abide by alignment requirements in bytes, not pixels!
+  if size % (ALIGN * 4) != 0: return [] if (base.itemsize * size) % (64 if OSX else ALIGN) != 0 or pxls > MAXW else [(1, pxls)]
+  return [(pxls//ALIGN//k, ALIGN*k) for k in range(ceildiv(pxls//ALIGN, MAXW), min(pxls//ALIGN, MAXW//ALIGN)+1) if (pxls//ALIGN)%k == 0]
 
 def transform_to_image(ctx, buf:UOp, x:UOp) -> UOp|None:
   shapes, ren = ctx

tinygrad/codegen/late/devectorizer.py

@@ -1,66 +0,0 @@
-import functools
-from tinygrad.dtype import dtypes, ImageDType, DType, Invalid
-from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, GroupOp
-from tinygrad.uop.symbolic import uop_given_valid, parse_valid, invalid_gate
-from tinygrad.helpers import OSX, ceildiv
-
-# ***** image load valid simplification *****
-
-@functools.cache
-def _drop_valid_stmts(valid:UOp, idx:UOp, height:int, width:int) -> list[UOp]:
-  # can drop valid if idx is out of bound when valid is False
-  drop_stmt = []
-  for i,stmt in enumerate(valid.split_uop(Ops.AND)):
-    if (res:=parse_valid(stmt)) is None: continue
-    X, is_upper_bound, c = res
-
-    # for X0 + X1 + ... >= 1, check if it's out of bound when Xi = 0 for all i
-    if not is_upper_bound and c == 1 and all(u.op in GroupOp.Irreducible and u.vmin == 0 for u in X.split_uop(Ops.ADD)):
-      testidx = functools.reduce(lambda nowidx,u: nowidx.substitute({u:u.const_like(0)}), X.split_uop(Ops.ADD), idx)
-      if testidx.index(0).vmax < 0 or testidx.index(1).vmax < 0:
-        drop_stmt.append(stmt)
-        continue
-
-    # check if idx is out of bound when X is on the wrong side of the bound: X in [c+1, vmax] or [vmin, c-1]
-    lo, hi = (c + 1, X.vmax) if is_upper_bound else (X.vmin, c - 1)
-    if lo <= hi:
-      fake = UOp.variable(f"fake{i}", lo, hi, X.dtype)
-      for coord,b in zip(idx.src, (width, height)):
-        rw = coord.substitute({X:fake}).simplify()
-        if rw.vmin >= b or rw.vmax < 0:
-          drop_stmt.append(stmt)
-          break
-  return drop_stmt
-
-def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None:
-  idx = uop_given_valid(valid, start_idx)
-  return None if idx is start_idx else buf.index(idx.valid(valid), ptr=True)
-
-def simplify_valid_image_load(buf:UOp, idx_y:UOp, idx_x:UOp, valid:UOp) -> UOp|None:
-  if not isinstance(buf.dtype, ImageDType): return None
-  start_idx = idx_x._stack(idx_y)
-  idx = uop_given_valid(valid, start_idx)
-  drop_stmt = _drop_valid_stmts(valid, idx, buf.dtype.shape[0], buf.dtype.shape[1])
-
-  if not drop_stmt and idx is start_idx: return None
-  new_valid = UOp.uprod(*ss) if (ss:=[s for s in valid.split_uop(Ops.AND) if s not in drop_stmt]) else None
-  idx_y, idx_x = idx.index(1), idx.index(0)
-  return buf.index(idx_y.valid(new_valid), idx_x.valid(new_valid), ptr=True) if new_valid is not None else buf.index(idx_y, idx_x, ptr=True)
-
-indexing_simplify = PatternMatcher([
-  # image load valid idx simplification
-  (UPat(Ops.INDEX, src=(UPat.var("buf"), invalid_gate)), lambda buf,x,i,cond: simplify_valid_load(buf, x, cond)),
-  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("valid").where(UPat.var("idx_y"), UPat(arg=Invalid)),
-                                         UPat.var("valid").where(UPat.var("idx_x"), UPat(arg=Invalid)))), simplify_valid_image_load),
-])
-
-# ***** load/store grouping *****
-
-# get list of (height, width) that do not require pitch padding
-def image_valid_dims(base:DType, size:int, arch:str) -> list[tuple[int,int]]:
-  if (ALIGN:=next((int(p.split('=')[1]) for p in arch.split(',') if p.startswith("IMAGE_PITCH_ALIGNMENT=")), 0)) == 0: return []
-  MAXW, pxls = 16384, size // 4
-  if base not in (dtypes.half, dtypes.float) or size > 4*MAXW*MAXW: return []
-  # height=1 images just need to abide by alignment requirements in bytes, not pixels!
-  if size % (ALIGN * 4) != 0: return [] if (base.itemsize * size) % (64 if OSX else ALIGN) != 0 or pxls > MAXW else [(1, pxls)]
-  return [(pxls//ALIGN//k, ALIGN*k) for k in range(ceildiv(pxls//ALIGN, MAXW), min(pxls//ALIGN, MAXW//ALIGN)+1) if (pxls//ALIGN)%k == 0]

tinygrad/codegen/opt/heuristic.py

@@ -3,7 +3,7 @@ from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
 from tinygrad.helpers import getenv, DEBUG, prod, NOLOCALS, TC_OPT, TC_SELECT, USE_TC, IMAGE
 from tinygrad.dtype import PtrDType
 from tinygrad.uop.ops import Ops, resolve, AxisType
-from tinygrad.codegen.late.devectorizer import image_valid_dims
+from tinygrad.codegen.late.coalese import image_valid_dims
 from tinygrad.codegen.opt.postrange import Scheduler
 
 def hand_coded_optimizations(k:Scheduler) -> Scheduler: