Revert "remove the on_stack stuff and any retries"

This reverts commit 49a2b328b9.

examples/train_resnet.py

@@ -32,7 +32,7 @@ if __name__ == "__main__":
 
   lr = 5e-3
   transform = ComposeTransforms([
-    lambda x: [Image.fromarray(xx).resize((64, 64)) for xx in x],
+    lambda x: [Image.fromarray(xx, mode='L').resize((64, 64)) for xx in x],
     lambda x: np.stack([np.asarray(xx) for xx in x], 0),
     lambda x: x / 255.0,
     lambda x: np.tile(np.expand_dims(x, 1), (1, 3, 1, 1)).astype(np.float32),

extra/test_hcopt.py

@@ -0,0 +1,40 @@
+import time
+from extra.optimization.helpers import load_worlds, ast_str_to_ast
+from tinygrad import Device
+from tinygrad.codegen.lowerer import pm_lowerer, get_index
+from tinygrad.uop.ops import graph_rewrite
+from tinygrad.codegen.opt.kernel import Kernel
+from tinygrad.codegen.opt.postrange import Scheduler
+from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
+from tinygrad.helpers import getenv
+
+if __name__ == "__main__":
+  renderer = Device.default.renderer
+  ast_strs = load_worlds()
+  if (n:=getenv("N", -1)) != -1: ast_strs = ast_strs[n:n+1]
+  good = 0
+  for i, ast_str in enumerate(ast_strs):
+    ast = ast_str_to_ast(ast_str)
+
+    st = time.perf_counter()
+    lin = Kernel(ast, renderer)
+    opt1 = hand_coded_optimizations(lin)
+    et_lin = time.perf_counter() - st
+
+    lowered = graph_rewrite(ast, pm_lowerer, ctx=get_index(ast), bottom_up=True)
+    st = time.perf_counter()
+    sch = Scheduler(lowered, renderer)
+    sch.convert_loop_to_global()
+    sch.simplify_merge_adjacent()
+    opt2 = hand_coded_optimizations(sch)
+    et_sch = time.perf_counter() - st
+
+    if opt1 != opt2:
+      print(f"******* {i:6d}")
+      print("Kernel:    ", lin.colored_shape(), "->", lin.apply_opts(opt1).colored_shape())
+      print("Scheduler: ", sch.colored_shape(), "->", sch.apply_opts(opt2).colored_shape())
+      print(opt1)
+      print(opt2)
+    else:
+      good += 1
+      print(f"******* {i:6d} MATCH {good/(i+1)*100:.2f}% -- {et_lin/et_sch:4.2f}x speedup")

tinygrad/codegen/__init__.py

@@ -7,6 +7,7 @@ from tinygrad.uop.spec import type_verify
 from tinygrad.renderer import Renderer
 
 # import all pattern matchers here
+from tinygrad.codegen.lowerer import pm_lowerer, get_index
 from tinygrad.codegen.quantize import pm_quant
 from tinygrad.codegen.gpudims import pm_add_gpudims
 from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing, symbolic
@@ -50,6 +51,7 @@ def _get_rewrites_for_renderer(opts:Renderer, optimize:bool, linearizer:bool, _Q
 
     # lowerer first
     if _QUANTIZE and opts.device in {"CPU", "DSP"}: ret.append(RewriteStep(pm_quant, name="quantize"))
+    ret.append(RewriteStep(pm_lowerer, get_index, name="lowerer", bottom_up=True))
 
     # split ranges
     if _RANGEIFY:

tinygrad/codegen/lowerer.py

@@ -0,0 +1,48 @@
+# the job of the lowerer is to do indexing
+from dataclasses import dataclass
+from tinygrad.uop.ops import KernelInfo, UOp, Ops, PatternMatcher, UPat, sint_to_uop, AxisType, graph_rewrite
+
+# ***** indexing *****
+
+@dataclass
+class IndexContext:
+  axis_types: tuple[AxisType, ...]
+  idxs: list[UOp]
+  start: int = 0
+
+def shape_to_idx(s, axis_types, start=0):
+  return [UOp.range(sint_to_uop(s), start+i, at) for i, (s, at) in enumerate(zip(s, axis_types))]
+
+def get_index(ast:UOp) -> IndexContext:
+  axis_types = ast.arg.axis_types if isinstance(ast.arg, KernelInfo) else ()
+  #if len(ast.full_shape) != len(axis_types) and ast.st is not None:
+  #  axis_types = tuple([AxisType.REDUCE if resolve(s != fs) else AxisType.LOOP for s,fs in zip(ast.shape, ast.full_shape)])
+  return IndexContext(axis_types, [], 0)
+
+# ***** lowering (given index) *****
+
+def subblock(ctx: IndexContext, full_new_idx: list[UOp], src: UOp):
+  lc = IndexContext(ctx.axis_types, full_new_idx, ctx.start+1000)
+  ctx.start = lc.start
+  return graph_rewrite(src, pm_lowerer, lc, name="subblock", bottom_up=True)
+
+def fixup_wmma(ctx:IndexContext, x:UOp):
+  if x.tag is not None: return None
+  new_idxs = shape_to_idx(x.src[0].shape, ctx.axis_types, ctx.start)
+  full_new_idx = list(ctx.idxs)
+  for a in x.arg[-1]: full_new_idx[a] = new_idxs[a]
+
+  srcs = subblock(ctx, full_new_idx, UOp.sink(*x.src)).src
+
+  # NOTE: this assumes these are expanded. which now shouldn't change anything
+  new_x_arg_m2 = tuple([tuple([(full_new_idx[a].arg[0], sz) for a,sz in v]) for v in x.arg[-2]])
+  new_x_arg_m1 = tuple([full_new_idx[a].arg[0] for a in x.arg[-1]])
+  return x.replace(src=srcs, arg=x.arg[:-2]+(new_x_arg_m2, new_x_arg_m1), tag=1)
+
+pm_lowerer = PatternMatcher([
+  (UPat(Ops.WMMA, name="x"), fixup_wmma),
+
+  # axis fixups for WMMA
+  (UPat((Ops.CONTRACT, Ops.UNROLL), name="x"),
+   lambda ctx,x: x.replace(tag=1, arg=tuple([(ctx.idxs[a].arg[0], sz) for a,sz in x.arg])) if x.tag is None else None),
+])

tinygrad/schedule/indexing.py

@@ -25,11 +25,11 @@ def realize_assign(ctx:dict[UOp, None], a:UOp) -> None:
 pm_generate_realize_map = PatternMatcher([
   # always realize SINK src
   (UPat(Ops.SINK, name="s"), lambda ctx,s: ctx.update((x.base, None) for x in s.src if x.base.op not in ALWAYS_CONTIGUOUS)),
-  # always realize COPY/BUFFER_VIEW/CONTIGUOUS
+  # always realize ASSIGN/COPY/BUFFER_VIEW/CONTIGUOUS
   (UPat({Ops.COPY, Ops.BUFFER_VIEW, Ops.CONTIGUOUS}, name="tr"), realize),
   # realize srcs of COPY, MSELECT, MSTACK
   (UPat((Ops.COPY, Ops.MSELECT, Ops.MSTACK), name="rb"), realize_srcs),
-  # realize ASSIGN and input to assign (might be optimized out)
+  # realize input to assign (might be optimized out)
   (UPat(Ops.ASSIGN, name="a"), realize_assign),
 ])
 
@@ -104,6 +104,8 @@ pm_apply_rangeify = PatternMatcher([
 ])
 
 def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
+  tsink_base = UOp.sink(*[x.base for x in tsink.src])
+
   rctx = IndexingContext()
 
   # get ops to realize
@@ -111,7 +113,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
 
   # explicit rangeify
   ending_ranges: dict[UOp, bool] = {}
-  for x in tsink.reverse_toposort(consumer_map:=tsink.get_consumer_map()):
+  for x in tsink_base.reverse_toposort(consumer_map:=tsink_base.get_consumer_map()):
     if x.op in {Ops.DEVICE, Ops.UNIQUE}: continue
     ending_ranges[x] = any(ending_ranges[u] for u in consumer_map[x])