remove noops

tinygrad/codegen/__init__.py

@@ -1,7 +1,7 @@
 from typing import Any, Callable
 import functools
 from dataclasses import dataclass
-from tinygrad.helpers import QUANTIZE, DEVECTORIZE, TRANSCENDENTAL, RANGEIFY
+from tinygrad.helpers import QUANTIZE, DEVECTORIZE, TRANSCENDENTAL, RANGEIFY, getenv
 from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype
 from tinygrad.uop.spec import type_verify
 from tinygrad.renderer import Renderer
@@ -17,7 +17,7 @@ from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_in
   ReduceContext, correct_load_store, pm_render
 from tinygrad.codegen.late.linearize import block_create, pm_blockend_merge, block_merge, pm_finalize, BlockContext
 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, pm_add_local_buffers, pm_pipeline
 from tinygrad.codegen.simplify import pm_simplify_ranges, pm_reduce_simplify, pm_flatten_range, pm_split_ranges
 from tinygrad.schedule.rangeify import pm_add_buffers, rangeify_codegen
 
@@ -77,16 +77,25 @@ def _get_rewrites_for_renderer(opts:Renderer, optimize:bool, linearizer:bool, _Q
   # ** expander (expand_rewrite) **
   ret.append(RewriteStep(sym+migrate_indexing, name="postopt symbolic"))
 
-  # expand
-  ret.append(RewriteStep(sym+pm_pre_expander+pm_group_for_reduce+expander, name="expander"))
+  # locals
+  if getenv("LOCALS") or getenv("PIPELINE"):
+    ret.append(RewriteStep(pm_add_local_buffers, name="add locals"))
 
   # add locals
   ret.append(RewriteStep(pm_add_buffers+rangeify_codegen, name="add local buffers"))
 
+  # expand
+  ret.append(RewriteStep(sym+pm_pre_expander+pm_group_for_reduce+expander, name="expander"))
+
   # ** devectorizer (full_graph_rewrite) **
   # remove reduce
   ret.append(RewriteStep(pm_reduce+gep_pushing, lambda _: ReduceContext(), name="remove_reduce"))
 
+  # pipelining
+  if getenv("PIPELINE"):
+    ret.append(RewriteStep(pm_pipeline, name="pipeline"))
+    ret.append(RewriteStep(sym, name="pipeline sym"))
+
   # add gpu dims (late). this works after devectorize, but it's faster here
   ret.append(RewriteStep(pm_add_gpudims, lambda _: opts, name="add gpudims"))
 

tinygrad/codegen/opt/heuristic.py

@@ -35,7 +35,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
       pass
     if good_tc_opt:
       # skip hand-coded TC opts if AMX, upcasting will make kernel slower
-      if rngs is not None and not AMX:
+      if rngs is not None and not AMX and False:
         for tc_dim in [1,0]: # attempt to upcast M and N
           szs = [sz for sz in [5,4,3,2] if rngs[tc_dim].src[0].divides(sz) is not None]
           if szs:
@@ -43,6 +43,11 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
             rngs[tc_dim] = tk.apply_opt(Opt(OptOps.UPCAST, tk.rngs.index(rngs[tc_dim]), szs[0]))[0]
         if (szs := [sz for sz in [4,2] if rngs[0].src[0].divides(sz) is not None]): # attempt to local N
           tk.apply_opt(Opt(OptOps.LOCAL, tk.rngs.index(rngs[0]), szs[0]))
+      #tk.apply_opt(Opt(OptOps.LOCAL, 0, 2))
+      #tk.apply_opt(Opt(OptOps.LOCAL, 1, 2))
+      #tk.apply_opt(Opt(OptOps.UPCAST, 0, 2))
+      #tk.apply_opt(Opt(OptOps.UPCAST, 1, 2))
+      #tk.apply_opt(Opt(OptOps.UNROLL, 0, 8))
       return tk
 
   # make a copy so it does not mutate the input

tinygrad/codegen/opt/postrange.py

@@ -1,14 +1,15 @@
 from __future__ import annotations
-import math, itertools
+import math, itertools, functools, operator
 from collections import defaultdict
 from typing import cast, Final
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, can_pad, GroupOp
 from tinygrad.device import Buffer
 from tinygrad.dtype import AddrSpace, dtypes, ImageDType
-from tinygrad.helpers import colored, BEAM, getenv, DEBUG, to_function_name, NOOPT, argsort, round_up, prod, merge_dicts, get_single_element
+from tinygrad.helpers import colored, BEAM, getenv, DEBUG, to_function_name, NOOPT, argsort, round_up, prod, merge_dicts, get_single_element, dedup
 from tinygrad.codegen.opt import axis_colors, Opt, OptOps, KernelOptError, check, axis_letters
 from tinygrad.codegen.simplify import pm_flatten_range
 from tinygrad.renderer import Renderer
+from tinygrad.schedule.rangeify import BufferizeOpts
 
 remove_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])
 
@@ -257,12 +258,13 @@ class Scheduler:
           except KernelOptError: continue
 
           # we create the warp as a whole thing, in case some of these ranges are moved/removed later
-          warp = UOp.range(tc.threads, -1, AxisType.WARP)
+          warp_num = -10
           ne: list[UOp] = []
           for opt in tc.opts:
             if opt[0] == "l":
-              axes[int(opt[1])], new_range = self.shift_to(axes[int(opt[1])], 2, AxisType.LOCAL, input_new_rng=warp%2)
-              warp //= 2
+              warp = UOp.range(2, warp_num, AxisType.WARP)
+              axes[int(opt[1])], new_range = self.shift_to(axes[int(opt[1])], 2, AxisType.WARP, input_new_rng=warp)
+              warp_num += 1
             elif opt[0] == "u":
               axes[int(opt[1])], new_range = self.shift_to(axes[int(opt[1])], 2, AxisType.UPCAST)
             else: raise RuntimeError(f"unsupported opt {opt[0]} in tensor cores")
@@ -347,3 +349,135 @@ def apply_opts(ctx:Renderer, ast:UOp):
 pm_postrange_opt = PatternMatcher([
   (UPat(Ops.SINK, name="ast"), apply_opts),
 ])
+
+def add_local_buffer(x:UOp):
+  if x.tag is not None: return None
+  # should UPCAST/UNROLL be here?
+  branges = tuple([r for r in x.ranges if r.arg[-1] in {AxisType.WARP, AxisType.LOCAL, AxisType.UPCAST, AxisType.UNROLL}])[::-1]
+  buf = UOp(Ops.BUFFERIZE, x.dtype, src=(x.replace(tag=1),)+branges, arg=BufferizeOpts(device=None, addrspace=AddrSpace.LOCAL))
+  return UOp(Ops.INDEX, x.dtype, src=(buf,)+branges)
+
+pm_add_local_buffers = PatternMatcher([
+  (UPat(Ops.LOAD, name="x"), add_local_buffer),
+])
+
+def add_pipeline(x:UOp):
+  if x.tag == 1: return None
+  if x.arg[-1] == AxisType.REDUCE:
+    # 3 splits
+    #srcs = (x.const_like(0), x.replace(src=(x.src[0]-2,), tag=1)+1, x.src[0]-1)
+
+    # 4 split
+    rng = x.replace(src=((x.src[0]-2)//2,), tag=1)
+    srcs = (x.const_like(0), rng*2+1, rng*2+2, x.src[0]-1)
+
+    return UOp(Ops.SPLIT, x.dtype, src=srcs, arg=1).simplify()
+    #vec = UOp(Ops.VECTORIZE, x.dtype.vec(3), src=(x.const_like(0), x.replace(src=(x.src[0]-2,), tag=1)+1, x.src[0]-1)).simplify()
+    #return UOp(Ops.UNROLL, x.dtype, src=(vec,), arg=())
+
+def do_split(x:UOp):
+  splits = [x for x in x.src if x.op is Ops.SPLIT]
+  if len(splits) == 0: return None
+  if x.op is Ops.SINK: return x.replace(src=x.src[0].src)
+  #if x.op is Ops.REDUCE:
+    #assert x.src[0].op is Ops.SPLIT
+    #rr = [y for y in x.src[1].toposort() if y.op is Ops.RANGE][0]
+    #return x.replace(src=(functools.reduce(operator.add, x.src[0].src), rr))
+  uu = []
+  for i in range(len(splits[0].src)):
+    new_srcs = []
+    for s in x.src:
+      if s.op is Ops.SPLIT:
+        new_srcs.append(s.src[i])
+      else:
+        new_srcs.append(s)
+    uu.append(UOp(x.op, x.dtype, tuple(new_srcs), x.arg, x.tag))
+  if x.op is Ops.STORE and len(splits) == 2:
+    dls = dedup([x for x in uu[0].toposort() if x.op is Ops.DEFINE_LOCAL])
+    uu2 = []
+    # NOTE: here we have to order the STORES and fix the ranges
+    for i,u in enumerate(uu):
+      subs = {}
+      for dl in dls: subs[dl] = dl.replace(arg=(dl.arg, i%2))
+      uu2.append(u.substitute(subs))
+    # TODO: reorder (is the reorder just a toposort question?)
+    # there's 4 barriers and 4 output stores
+    # load 0
+    # barrier (between load 0 and compute 0)
+    #   load 1 (depends on range)
+    #   compute 0
+    #   barrier (between load 0+2 and load 2)
+    #   load 2 (depends on range)
+    #   compute 1
+    #   barrier (between load 1 and load 1)
+    # load 3
+    # compute 2
+    # barrier (between load 3 and compute 3)
+    # compute 3
+
+    # uncouple based on barriers
+    def do_uncouple(ctx, l:UOp, b:UOp):
+      ctx[1].append(b.src[0])
+      return l.replace(src=l.src[0:1]+(UOp(Ops.NOOP, tag=ctx[0]),))
+    uncouple_barrier = PatternMatcher([
+      (UPat(Ops.LOAD, src=(UPat(), UPat(Ops.BARRIER, name='b')), name='l'), do_uncouple),
+      (UPat(Ops.LOAD, src=(UPat(), UPat(), UPat()), name='x'), lambda x: x.replace(src=x.src[0:2])),
+    ])
+
+    loads = []
+    computes = []
+    for i,u in enumerate(uu2):
+      cc = graph_rewrite(u, uncouple_barrier, ctx=(i,tloads:=[]))
+      computes.append(cc)
+      loads.append(UOp(Ops.NOOP, src=tuple(tloads)))
+
+    # remove the range from here
+    computes[2] = computes[2].replace(src=computes[2].src[0:2])
+
+    # pipelined!
+    ret = computes[3]
+
+    # put computes[2] before compute[3]
+    const_store = [x for x in ret.toposort() if x.op is Ops.STORE and x.src[1].op is Ops.CONST][0]
+    ret = ret.substitute({const_store:const_store.replace(tag=1)})
+    ret = ret.substitute({const_store.replace(tag=1):computes[2].barrier()})
+
+    # put loads[3] before compute[2] (both ports)
+    const_store = [x for x in ret.toposort() if x.op is Ops.STORE and x.src[1].op is Ops.CONST][0]
+    ret = ret.substitute({const_store:const_store.replace(tag=1)})
+    ret = ret.substitute({const_store.replace(tag=1):loads[3], UOp(Ops.NOOP, tag=2):loads[3]})
+
+    # put compute[1] before loads[3]
+    load = [x for x in loads[3].toposort() if x.op is Ops.LOAD][0]
+    ret = ret.substitute({load: load.replace(src=load.src+(computes[1].barrier(),))})
+
+    # put loads[2] before compute[1] (both ports)
+    const_store = [x for x in ret.toposort() if x.op is Ops.STORE and x.src[1].op is Ops.CONST][0]
+    ret = ret.substitute({const_store:const_store.replace(tag=1)})
+    ret = ret.substitute({const_store.replace(tag=1):loads[2], UOp(Ops.NOOP, tag=1):loads[2]})
+
+    # put compute[0] before loads[2]
+    load = [x for x in loads[2].toposort() if x.op is Ops.LOAD][0]
+    ret = ret.substitute({load: load.replace(src=load.src+(computes[0].barrier(),))})
+
+    # put loads[1] before compute[0] (one port)
+    ret = ret.substitute({UOp(Ops.NOOP, tag=0):loads[1]})
+
+    # put loads[0] before loads[1]
+    load = [x for x in loads[1].toposort() if x.op is Ops.LOAD][0]
+    ret = ret.substitute({load: load.replace(src=load.src+(loads[0],))})
+
+    pm_remove_noops = PatternMatcher([
+      (UPat(Ops.NOOP, src=(UPat.var('x'),)), lambda x: x),
+    ])
+    return graph_rewrite(ret, pm_remove_noops, name="remove noops")
+  return UOp(Ops.SPLIT, x.dtype, src=tuple(uu))
+
+pm_pipeline = PatternMatcher([
+  (UPat(Ops.RANGE, name="x"), add_pipeline),
+  # do expansion
+  (UPat(GroupOp.All, name="x", custom_early_reject=set([Ops.SPLIT])), do_split),
+  #(UPat(Ops.STORE, src=(UPat(), UPat(), UPat(Ops.CONST)), name="x"), lambda x: x.replace(src=x.src[:2])),
+  (UPat(Ops.STORE, src=(UPat.var('x'), UPat.var('z'), UPat.var('rr')), name='r'),
+   lambda x,rr,r,z: r.replace(src=(x,z)+tuple([y for y in rr.toposort() if y.op is Ops.RANGE]))),
+])

tinygrad/codegen/opt/search.py

@@ -68,7 +68,7 @@ def _try_compile_linearized_w_idx(x:tuple[int,Scheduler], compiler:Compiler) ->
   try:
     p = get_program(x[1].copy().get_optimized_ast(name_override="test"), x[1].opts)
     assert p.uops is not None, "uop list wasn't generated?"
-    if len(p.uops) >= (uops_max:=getenv("BEAM_UOPS_MAX", 3000)) > 0:
+    if len(p.uops) >= (uops_max:=getenv("BEAM_UOPS_MAX", 6000)) > 0:
       if getenv("BEAM_LOG_SURPASS_MAX"): print(f"too many uops. {len(p.uops)=}, {uops_max=}")
       raise RuntimeError("too many uops")
     st = time.perf_counter()

tinygrad/renderer/cstyle.py

@@ -295,7 +295,7 @@ class MetalRenderer(CStyleLanguage):
   # language options
   kernel_typedef = "kernel void"
   buffer_prefix = "device "
-  smem_prefix = "threadgroup __attribute__((aligned(16))) "
+  smem_prefix = "threadgroup " #__attribute__((aligned(16))) "
   arg_int_prefix = "constant int&"
   barrier = "threadgroup_barrier(mem_flags::mem_threadgroup);"
   float4 = "float4"

tinygrad/schedule/rangeify.py

@@ -364,7 +364,7 @@ pm_rangeify = pm_mops+PatternMatcher([
   (UPat(Ops.INDEX, src=(UPat(GroupOp.Elementwise.union({Ops.REDUCE_AXIS})),), allow_any_len=True, name="idx"), might_end_axis),
 
   # handle size 0
-  (UPat(Ops.INDEX, name="x"), lambda x: x.replace(src=(x.const_like(0),)+x.src[1:]) if x.st is not None and x.size == 0 else None),
+  #(UPat(Ops.INDEX, name="x"), lambda x: x.replace(src=(x.const_like(0),)+x.src[1:]) if x.st is not None and x.size == 0 else None),
 
   # handle assign
   (UPat(Ops.INDEX, src=(UPat(Ops.ASSIGN, name="assign"),), allow_any_len=True, name="x"),

tinygrad/uop/__init__.py

@@ -21,6 +21,7 @@ class Ops(FastEnum):
   # create buffer
   BUFFERIZE = auto()
   SUBSTITUTE = auto()
+  SPLIT = auto()
 
   # ops that adjust the behavior of the scheduler
   CONTIGUOUS = auto(); CONTIGUOUS_BACKWARD = auto(); DETACH = auto(); FUSE = auto() # noqa: E702