Merge branch 'master' into postrange_hcopts

tinygrad/codegen/__init__.py

@@ -18,7 +18,7 @@ from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_in
 from tinygrad.codegen.late.linearize import block_create, pm_blockend_merge, block_merge, pm_finalize, BlockContext
 from tinygrad.codegen.opt import pm_get_optimization, pm_do_optimize
 from tinygrad.codegen.opt.swizzler import view_left, view_right, fix_kernel_ops
-from tinygrad.codegen.opt.postrange import pm_postrange_opt
+from tinygrad.codegen.opt.postrange import pm_postrange_opt_early, pm_postrange_opt, pm_postrange_opt_merge
 from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen
 
 @dataclass
@@ -57,18 +57,24 @@ def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVEC
   ret.extend(rewrites_for_views)
 
   # this is kernel.py
-  if not _RANGEIFY: ret.append(RewriteStep(pm_get_optimization, ctx=lambda _: opts, name="get optimization"))
+  ret.append(RewriteStep(pm_get_optimization, ctx=lambda _: opts, name="get optimization"))
+
   if not _POSTOPT and not _RANGEIFY: ret.append(RewriteStep(pm_do_optimize, ctx=lambda _: opts, name="optimize ast"))
 
   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))
 
-  if _POSTOPT or _RANGEIFY: ret.append(RewriteStep(pm_postrange_opt, ctx=lambda _: opts, name="post optimize ast"))
-
-  # ** expander (expand_rewrite) **
+  # symbolic before post opt
   ret.append(RewriteStep(sym+migrate_indexing, name="initial symbolic"))
 
-  # expand
+  if _POSTOPT or _RANGEIFY:
+    ret.append(RewriteStep(pm_postrange_opt_merge, ctx=lambda _: ({}, opts), name="early range merge"))
+    ret.append(RewriteStep(sym, name="mid symbolic"))
+    ret.append(RewriteStep(pm_postrange_opt_early, ctx=lambda _: ({}, opts), name="early post opt ast"))
+    ret.append(RewriteStep(sym, name="mid symbolic"))
+    ret.append(RewriteStep(pm_postrange_opt, ctx=lambda _: opts, name="post optimize ast"))
+
+  # ** expander (expand_rewrite) **
   ret.append(RewriteStep(sym+pm_pre_expander+expander, name="expander"))
 
   # add locals

tinygrad/codegen/late/expander.py

@@ -134,15 +134,15 @@ def fix_store_unroll(x:UOp):
   return UOp(Ops.CONTRACT, dtypes.void, (x.replace(src=x.src[:2]+tuple(store_range)),), tuple(flatten(x.arg for x in store_expand)), tag=1)
 
 def fix_group_for_reduce(x:UOp):
-  reduce_gfr, reduce_r = partition(x.src[1:], lambda u: u.op is Ops.RANGE and u.arg[1] == AxisType.GROUP_REDUCE)
+  reduce_gfr, reduce_r = partition(x.src[1:], lambda u: u.op is Ops.RANGE and u.arg[-1] == AxisType.GROUP_REDUCE)
   if len(reduce_gfr) == 0: return None
 
   # NOTE: if there's other locals here, we need them in the buffer too
-  upstream_locals = [u for u in x.toposort() if u.op is Ops.RANGE and u.arg[1] == AxisType.LOCAL]
+  upstream_locals = [u for u in x.toposort() if u.op is Ops.RANGE and u.arg[-1] == AxisType.LOCAL]
 
   # do only the non grouped reduces early
   ret = x.replace(src=(x.src[0],)+tuple(reduce_r))
-  reduce_loop = [x.replace(arg=(x.arg[0]+100, AxisType.REDUCE)) for x in reduce_gfr]
+  reduce_loop = [x.replace(arg=(*x.arg[0:-1], 0, AxisType.REDUCE)) for x in reduce_gfr]
   buf = ret.bufferize(*upstream_locals, *reduce_gfr, arg=(AddrSpace.LOCAL, reduce_gfr[0].arg[0])).index(*upstream_locals, *reduce_loop)
 
   # gate with an if on the store + do the final reduce
@@ -152,8 +152,8 @@ def fix_group_for_reduce(x:UOp):
 pm_pre_expander = PatternMatcher([
   # rewrite UPCAST/UNROLL range to something to be expanded
   (UPat(Ops.RANGE, name="r"),
-   lambda r: UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(s:=r.vmax+1), tuple(range(s))),), ((r.arg[0],s),)) \
-    if r.arg[1] in {AxisType.UNROLL, AxisType.UPCAST} else None),
+   lambda r: UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(s:=r.vmax+1), tuple(range(s))),), ((r.arg[0:-1],s),)) \
+    if r.arg[-1] in {AxisType.UNROLL, AxisType.UPCAST} else None),
   # fix REDUCEs with UNROLLs
   (UPat(Ops.REDUCE, name="x"), fix_reduce_unroll),
   (UPat(Ops.STORE, name="x"), fix_store_unroll),

tinygrad/codegen/lowerer.py

@@ -83,5 +83,5 @@ pm_lowerer = PatternMatcher([
 
   # 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),
+   lambda ctx,x: x.replace(tag=1, arg=tuple([(ctx.idxs[a].arg[0:-1], sz) for a,sz in x.arg])) if x.tag is None else None),
 ])

tinygrad/codegen/opt/__init__.py

@@ -28,7 +28,8 @@ def get_optimized_ast(ast:UOp, renderer:Renderer) -> UOp:
       kb = Kernel(ast, opts=renderer)
       rawbufs = bufs_from_lin(kb, allocate=False)
       k = beam_search(kb, rawbufs, BEAM.value, bool(getenv("BEAM_ESTIMATE", 1)))
-  return ast.replace(arg=KernelInfo(opts_to_apply=tuple(k.applied_opts)))
+  # NOTE: this does simplify_ones/simplify_merge_adjacent for you
+  return Kernel(ast, opts=renderer).get_optimized_ast().replace(arg=KernelInfo(opts_to_apply=tuple(k.applied_opts)))
 
 pm_get_optimization = PatternMatcher([
   (UPat(Ops.SINK, name="ast"), lambda ctx,ast: get_optimized_ast(ast, ctx) if ast.arg is None and ast.src[0].st is not None else None),

tinygrad/codegen/opt/postrange.py

@@ -1,7 +1,168 @@
+import math
 from dataclasses import replace
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo
-from tinygrad.helpers import colored
-from tinygrad.codegen.opt.kernel import axis_colors
+from tinygrad.dtype import dtypes
+from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, _substitute
+from tinygrad.uop.symbolic import symbolic
+from tinygrad.helpers import colored, USE_TC, DEBUG
+from tinygrad.codegen.opt.kernel import axis_colors, AxisType
+from tinygrad.renderer import Renderer
+from tinygrad.codegen.opt.tc import TensorCore
+
+def count_divmod(x:UOp): return len([u for u in x.toposort() if u.op in {Ops.IDIV, Ops.MOD}])
+
+def flatten_range_in_terminators(r:UOp):
+  off = 2 if r.op is Ops.STORE else 1
+  rngs = r.src[off:]
+  if not len(rngs): return None
+  new_rngs = [x for x in UOp.sink(*rngs).toposort() if x.op is Ops.RANGE]
+  return r.replace(src=r.src[:off]+tuple(new_rngs))
+
+pm_flatten_range = PatternMatcher([
+  # flatten ranges
+  (UPat((Ops.REDUCE, Ops.STORE), name="r"), flatten_range_in_terminators),
+])
+
+# NOTE: this one is better than the one in kernel.py
+def simplify_merge_adjacent(ast:UOp):
+  # get all ranges (sorted)
+  rng = sorted([u for u in ast.parents if u.op is Ops.RANGE], key=lambda x: x.arg[0:-1])
+  terminators = [u for u in ast.parents if u.op in {Ops.REDUCE, Ops.STORE}]
+  termination = {}
+  for t in terminators:
+    for u in t.src[1 if t.op is Ops.REDUCE else 2:]: termination[u] = t
+
+  replaces = {}
+  i = 0
+  while i < len(rng)-1:
+    r0, r1 = rng[i], rng[i+1]
+    # same axistype and same termination
+    if r0.arg[1] == r1.arg[1] and termination[r0] == termination[r1]:
+      s0, s1 = r0.src[0], r1.src[0]
+      new_range = r0.replace(src=(s0*s1,)).simplify()
+      # this checks the legality of a merge
+      oidx = ast.simplify()
+      nidx = graph_rewrite(oidx, _substitute+symbolic+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1}, name=f"check_merge_{i}_{i+1}")
+      # it simplifies
+      if count_divmod(nidx) <= count_divmod(oidx):
+        # it is correct
+        midx = graph_rewrite(nidx, _substitute+symbolic+pm_flatten_range, ctx={new_range:r0*s1+r1}, name=f"correct_merge_{i}_{i+1}")
+        if oidx is midx:
+          termination[new_range] = termination[r0]
+          replaces[r0] = new_range//s1
+          replaces[r1] = new_range%s1
+          rng[i] = new_range
+          del rng[i+1]
+          continue
+    i += 1
+  return ast.substitute(replaces, name="simplify_merge_adjacent")
+
+pm_postrange_opt_merge = pm_flatten_range+PatternMatcher([
+  (UPat(Ops.SINK, name="ast"), simplify_merge_adjacent),
+])
+
+def apply_tensor_cores(ctx:tuple[dict, Renderer], in0:UOp, in1:UOp, r_range:UOp, reduceop:UOp):
+  if not USE_TC: return None
+  # tensor cores have three ranges. X, Y, and REDUCE
+  in0_ranges = sorted([u for u in in0.ranges if u not in in1.ranges], key=lambda x: x.arg[0])
+  in1_ranges = sorted([u for u in in1.ranges if u not in in0.ranges], key=lambda x: x.arg[0])
+  if not len(in0_ranges) or not len(in1_ranges): return None
+  in0_range, in1_range = in0_ranges[0], in1_ranges[0]
+  if DEBUG >= 2: print('TC', in0_range.arg, in1_range.arg, r_range.arg)
+
+  # confirm the dtype and size is good
+  tc_opts: list[TensorCore] = []
+  for tc in ctx[1].tensor_cores:
+    if reduceop.dtype == tc.dtype_out and in0.dtype == tc.dtype_in and in1.dtype == tc.dtype_in:
+      if all(i <= j for i,j in zip(tc.dims, [in0_range.vmax+1, in1_range.vmax+1, r_range.vmax+1])):
+        tc_opts.append(tc)
+  if len(tc_opts) == 0: return None
+  tc = tc_opts[0]
+
+  # create the new ranges as speced by the tensor core
+  old_range = [in0_range, in1_range, r_range]
+  new_range = [r.replace(src=(r.src[0]//tc.dims[i],), arg=r.arg[0:-1]+(0, r.arg[-1])) for i,r in enumerate(old_range)]
+  new_range_args = [list(x.arg[0:-1]) for x in new_range]
+  new_reduce_range = new_range[2]
+  red_ranges = []
+
+  # place the warp at -99
+  warp_range = -99
+
+  ne: list[UOp] = []
+  for o in tc.opts:
+    axis = 1-int(o[1])
+    if o[0] == "u":
+      new_range_args[axis][-1] += 1
+      lrange = UOp.range(dtypes.int, 2, *new_range_args[axis], AxisType.UPCAST)
+    else:
+      lrange = UOp.range(dtypes.int, 2, warp_range, AxisType.LOCAL)
+      warp_range += 1
+    ne.append(lrange)
+    new_range[axis] = (2 * new_range[axis]) + lrange
+  for _, amt in tc.get_reduce_axes():
+    new_range_args[2][-1] += 1
+    lrange = UOp.range(dtypes.int, amt, *new_range_args[2], AxisType.UNROLL)
+    ne.append(lrange)
+    red_ranges.append(lrange)
+    new_range[2] = (amt * new_range[2]) + lrange
+  tne = [x.replace(tag=1) for x in ne]
+
+  # replace ranges in other parts of the graph
+  for x,y in zip(old_range, new_range): ctx[0][x] = y
+
+  # apply the swizzled ranges to the srcs
+  srcs = [s.substitute(dict(zip(old_range, new_range))).substitute(dict(zip(ne, tne))) for s in (in0, in1)]
+  srcs = [x.substitute(dict(zip(tne, [ne[i] for i in p]))) for x,p in zip(srcs, tc.permutes_for_shape_str(tc.base_shape_str()))]
+
+  ned = dict(zip(tc.base_shape_str(), ne))
+  tc_reduce_axes = tuple([ned[f"r{i}"].arg[0:-1] for i in range(len(tc.get_reduce_axes()))])
+  base_upcast_axes = tuple([(ned[s].arg[0:-1], 2) for s in tc.base_upcast_axes()])
+  tc_upcast_axes = tuple([base_upcast_axes[:int(math.log2(tc.elements_per_thread[i]))] for i in range(3)])
+
+  # construct the op
+  # TODO: remove tc_upcast_axes from the arg
+  wmma_arg = (str(tc), tc.dims, tc.dtype_in, tc.dtype_out, ctx[1].device, tc.threads, tc_upcast_axes, tc_reduce_axes)
+  wmma = UOp(Ops.WMMA, dtype=tc.dtype_out.vec(tc.elements_per_thread[2]), src=(
+    UOp(Ops.CONTRACT, dtype=srcs[0].dtype.vec(tc.elements_per_thread[0]), src=(srcs[0],), arg=tc_upcast_axes[0]),
+    UOp(Ops.CONTRACT, dtype=srcs[1].dtype.vec(tc.elements_per_thread[1]), src=(srcs[1],), arg=tc_upcast_axes[1]),
+    UOp.const(tc.dtype_out.vec(tc.elements_per_thread[2]), 0.0)), arg=wmma_arg)
+  tc_uop = UOp(Ops.UNROLL, tc.dtype_out, (wmma,), arg=tc_upcast_axes[2])
+  ret = tc_uop.reduce(new_reduce_range, arg=Ops.ADD)
+  # confirm the UNROLLs aren't actually used, these need to be broadcast MUL
+  assert all(u not in red_ranges for u in ret.toposort()), "UNROLLs in TC"
+  return ret
+
+def early_sink(ctx:tuple[dict, Renderer], s:UOp):
+  s = s.substitute(ctx[0])
+  # global_stores_are_global
+  if ctx[1].has_local:
+    rngs = UOp.sink(*s.src[0].src[2:]).parents
+    s = s.substitute({u:u.replace(arg=u.arg[0:-1]+(AxisType.GLOBAL,)) for u in rngs if u.op is Ops.RANGE and u.arg[-1] is AxisType.LOOP})
+  return s
+
+pm_postrange_opt_early = PatternMatcher([
+  # TODO: this is optional (and can have internal options) and we need a way to express that
+  ((UPat.var("in0")*UPat.var("in1")).reduce(UPat(Ops.RANGE, name="r_range"), name="reduceop", arg=Ops.ADD), apply_tensor_cores),
+  (UPat(Ops.SINK, name="s"), early_sink),
+])
+
+# *** late (BEAM goes here) ***
+
+axis_typemap = { # (is_reduce, is_local)
+  (False,False): AxisType.UPCAST, (False,True): AxisType.LOCAL,
+  (True, False): AxisType.UNROLL, (True, True): AxisType.GROUP_REDUCE}
+
+def split_range(r:UOp):
+  if r.arg[-1] not in {AxisType.LOOP, AxisType.GLOBAL, AxisType.REDUCE}: return None
+  if r.tag is not None: return None
+  # any divisor is an option
+  is_local = False if r.arg[-1] is AxisType.REDUCE else False
+  N = 4
+  rd = r.src[0].divides(N)
+  if rd is None: return None
+  sr = r.replace(src=(rd,), arg=r.arg[0:-1]+(0, r.arg[-1]), tag=1)
+  er = UOp(Ops.RANGE, dtypes.int, src=(UOp.const(dtypes.int, N),), arg=r.arg[0:-1]+(1, axis_typemap[(r.arg[-1] is AxisType.REDUCE, is_local)]))
+  return sr*N+er
 
 def rename_sink(s:UOp):
   if s.arg is not None and s.arg.name != "test": return None
@@ -13,6 +174,11 @@ def rename_sink(s:UOp):
   name = "k" + colored('_', 'BLACK').join(['']+[colored(x.src[0].render(), axis_colors[x.arg[-1]]) for x in rngs])
   return s.replace(arg=KernelInfo(name=name) if s.arg is None else replace(s.arg, name=name))
 
-pm_postrange_opt = PatternMatcher([
+pm_postrange_opt = pm_flatten_range+PatternMatcher([
+  # TODO: this is optional (and can have internal options) and we need a way to express that
+  (UPat(Ops.RANGE, name="r"), split_range),
+  # remove axes with 1
+  (UPat(Ops.RANGE, name="r"), lambda r: r.const_like(0) if r.vmax == 0 else None),
+  # run this last
   (UPat(Ops.SINK, name="s"), rename_sink),
 ])