we name the warp warp now

* vmap on full model

* vmap gemm

* reduce sums on end

* outer reduce

* only if there's ranges

* put those rules in symbolic

* ranges

* do opt later

* add zero range

extra/sqtt/roc.py

@@ -1,4 +1,4 @@
-import ctypes, pathlib, argparse, pickle, re, functools, dataclasses, itertools
+import ctypes, pathlib, argparse, pickle, re, functools, dataclasses, itertools, threading
 from tinygrad.helpers import temp, unwrap, DEBUG
 from tinygrad.device import ProfileEvent, ProfileDeviceEvent, ProfileProgramEvent
 from tinygrad.runtime.ops_amd import ProfileSQTTEvent, ProfilePMCEvent
@@ -94,14 +94,14 @@ def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
   ROCParseCtx = _ROCParseCtx(dev_events, sqtt_events, prog_events)
 
   @rocprof.rocprof_trace_decoder_se_data_callback_t
-  def copy_cb(buf, buf_size, data_ptr):
+  def copy_cb(buf, buf_size, _):
     if (prof_info:=ROCParseCtx.next_sqtt()) is None: return 0
     buf[0] = ctypes.cast(prof_info, ctypes.POINTER(ctypes.c_ubyte))
     buf_size[0] = len(prof_info)
     return len(prof_info)
 
   @rocprof.rocprof_trace_decoder_trace_callback_t
-  def trace_cb(record_type, events_ptr, n, data_ptr):
+  def trace_cb(record_type, events_ptr, n, _):
     match record_type:
       case rocprof.ROCPROFILER_THREAD_TRACE_DECODER_RECORD_OCCUPANCY:
         for ev in (rocprof.rocprofiler_thread_trace_decoder_occupancy_t * n).from_address(events_ptr): ROCParseCtx.on_occupancy_ev(ev)
@@ -112,7 +112,7 @@ def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
     return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
 
   @rocprof.rocprof_trace_decoder_isa_callback_t
-  def isa_cb(instr_ptr, mem_size_ptr, size_ptr, pc, data_ptr):
+  def isa_cb(instr_ptr, mem_size_ptr, size_ptr, pc, _):
     instr, mem_size_ptr[0] = ROCParseCtx.disasms[(unwrap(ROCParseCtx.active_kern), pc.address)]
 
     # this is the number of bytes to next instruction, set to 0 for end_pgm
@@ -126,9 +126,11 @@ def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
 
     return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
 
-  try:
-    rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
-  except AttributeError as e: raise RuntimeError("Failed to find rocprof-trace-decoder. Run sudo ./extra/sqtt/install_sqtt_decoder.py to install") from e
+  def worker():
+    try: rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
+    except AttributeError as e: raise RuntimeError("Failed to find rocprof-trace-decoder. Run sudo ./extra/sqtt/install_sqtt_decoder.py to install") from e
+  (t:=threading.Thread(target=worker, daemon=True)).start()
+  t.join()
   return ROCParseCtx
 
 if __name__ == "__main__":

test/test_outerworld.py

@@ -1,5 +1,6 @@
 import unittest
-from tinygrad import Tensor, UOp
+import numpy as np
+from tinygrad import Tensor, UOp, nn
 from tinygrad.uop.ops import AxisType, Ops
 
 class TestOuterworldReduce(unittest.TestCase):
@@ -77,7 +78,8 @@ class TestOuterScan(unittest.TestCase):
     # 3 matmuls with SCAN
     i = UOp.range(3, -100, AxisType.OUTER)
     out = Tensor.empty(3, 1, 10)
-    comp = Tensor(i.eq(0).where(vec.uop, out[(i-1).maximum(0)].uop)) @ mats[i]
+    phi = Tensor(i.eq(0).where(vec.uop, out[(i-1).maximum(0)].uop))
+    comp = phi @ mats[i]
     store = out[i].uop.store(comp.uop).end(i)
     out = Tensor(out.uop.after(store))
     out.realize()
@@ -144,5 +146,85 @@ class TestOuterworld(unittest.TestCase):
     out = out.reshape(1, 3).expand(a, 3).contiguous().realize()
     self.assertListEqual([[0,4,8],[4,8,12],[8,12,16]], out.tolist())
 
+class TestVmap(unittest.TestCase):
+  def test_vmap_inner(self, axis_type=AxisType.LOOP, fuse=False, grad=False):
+    x = Tensor.ones(1, 10).contiguous().requires_grad_()
+    mats = Tensor.ones(3, 10, 10).contiguous().requires_grad_()
+
+    ref = x @ mats
+    if fuse: ref = ref * 2
+
+    # vmap across axis 0
+    a = UOp.range(3, -1, axis_type)
+    out = x @ mats[a]
+    out = out.reshape(1, 10).pad(((a,(3-a)-1), None))
+    out = Tensor(out.uop.reduce(a, arg=Ops.ADD))
+    if fuse: out = out * 2
+    if grad:
+      out.mean().backward()
+      np.testing.assert_allclose(mats.grad.numpy(), (2./30) if fuse else (1./30))
+    out.realize()
+
+    # TODO: testing allclose
+    assert Tensor.allclose(ref, out, atol=1e-6), f"{ref.numpy()=}, {out.numpy()=}"
+  def test_vmap_inner_fuse(self): self.test_vmap_inner(fuse=True)
+  def test_vmap_outer(self): self.test_vmap_inner(AxisType.OUTER)
+  def test_vmap_outer_fuse(self): self.test_vmap_inner(AxisType.OUTER, fuse=True)
+
+  def test_vmap_inner_grad(self): self.test_vmap_inner(grad=True)
+  def test_vmap_inner_fuse_grad(self): self.test_vmap_inner(fuse=True, grad=True)
+  def test_vmap_outer_grad(self): self.test_vmap_inner(AxisType.OUTER, grad=True)
+
+  def test_vmap_convs(self):
+    layers = [
+      nn.Conv2d(1, 8, 3), Tensor.relu,
+      nn.Conv2d(8, 8, 3), Tensor.relu]
+    img = Tensor.randn(4, 1, 16, 16).realize(*nn.state.get_parameters(layers))
+    a = UOp.range(4, -1, AxisType.OUTER)
+    out = img[a:a+1].sequential(layers)
+    out = out.pad(((a,(4-a)-1), None, None, None))
+    out = Tensor(out.uop.reduce(a, arg=Ops.ADD))
+    out.realize()
+    np.testing.assert_allclose(out.numpy(), img.sequential(layers).numpy(), atol=1e-6)
+
+  def test_vmap_gemm(self):
+    layers = [
+      nn.Linear(16, 16, bias=False), Tensor.relu,
+      nn.Linear(16, 16, bias=False), Tensor.relu]
+    img = Tensor.randn(4, 16).realize(*nn.state.get_parameters(layers))
+    a = UOp.range(4, -1, AxisType.OUTER)
+    out = img[a:a+1].sequential(layers)
+    out = out.pad(((a,(4-a)-1), None))
+    out = Tensor(out.uop.reduce(a, arg=Ops.ADD))
+    out.realize()
+    np.testing.assert_allclose(out.numpy(), img.sequential(layers).numpy(), atol=1e-6)
+
+  @unittest.skip("this is broken, we need to lower the outer reduce in the outer graph")
+  def test_vmap_gemm_grad(self):
+    layers = [
+      nn.Linear(16, 16, bias=False), Tensor.relu,
+      nn.Linear(16, 16, bias=False), Tensor.relu]
+    layer_tensors = nn.state.get_parameters(layers)
+    img = Tensor.randn(4, 16).realize(*layer_tensors)
+    for l in layer_tensors: l.requires_grad_()
+    a = UOp.range(4, -1, AxisType.OUTER)
+    out = img[a:a+1].sequential(layers)
+    out = out.pad(((a,(4-a)-1), None))
+    out = Tensor(out.uop.reduce(a, arg=Ops.ADD))
+    out.mean().backward()
+    grads = [l.grad for l in layer_tensors]
+    out.realize(*grads)
+    out_grads = [x.numpy() for x in grads]
+
+    # compute reference grads
+    for l in layer_tensors: l.grad = None
+    img.sequential(layers).mean().backward()
+    grads = [l.grad for l in layer_tensors]
+    out.realize(*grads)
+    ref_grads = [x.numpy() for x in grads]
+
+    # compare
+    for o,r in zip(out_grads, ref_grads): np.testing.assert_allclose(o, r, atol=1e-6)
+
 if __name__ == '__main__':
   unittest.main()

tinygrad/codegen/late/expander.py

@@ -134,7 +134,7 @@ def fix_group_for_reduce(x:UOp):
 
   # 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]+"_gfr", AxisType.REDUCE)) for x in reduce_gfr]
   buf = ret.bufferize(*upstream_locals, *reduce_gfr, arg=BufferizeOpts(reduce_gfr[0].arg[0], AddrSpace.LOCAL)).index(*upstream_locals, *reduce_loop)
 
   # do the final reduce (if/barrier are added in gpudims step)

tinygrad/codegen/opt/postrange.py

@@ -18,6 +18,7 @@ class Scheduler:
     self.ast, self.ren = ast, ren
     self.dont_use_locals = self.ast.arg.dont_use_locals if self.ast.arg is not None else False
     self.applied_opts = list(self.ast.arg.applied_opts) if self.ast.arg is not None else []
+    self.opt_range = itertools.count()
 
   @property
   def rngs(self):
@@ -29,8 +30,6 @@ class Scheduler:
   def full_shape(self): return [ssimplify(x.src[0]) for x in self.rngs]
   @property
   def axis_types(self): return [x.arg[-1] for x in self.rngs]
-  @property
-  def maxarg(self): return max([x.arg[0] for x in self.rngs], default=0)
 
   # strings like ['g0', 'g1', 'l0', 'l1', 'l2', 'l3', 'l4', 'l5', 'R0', 'r0', 'r1', 'r2', 'u0', 'u1', 'u2']
   def shape_str(self) -> list[str]:
@@ -95,10 +94,10 @@ class Scheduler:
   def shift_to(self, rng:UOp, amount:int, new_type:AxisType, top:bool=False, input_new_rng=None):
     if (old_sz:=rng.src[0].divides(amount)) is None:
       raise KernelOptError(f"{amount} can't divide {rng.src[0]} in {self.colored_shape()}")
-    new_rng = UOp.range(amount, self.maxarg+1, new_type) if input_new_rng is None else input_new_rng
+    new_rng = UOp.range(amount, f"o{next(self.opt_range)}", new_type) if input_new_rng is None else input_new_rng
     replaced_rng = rng.replace(src=(UOp.const(dtypes.int, old_sz),))
     sub_axis = (new_rng * old_sz + replaced_rng) if top else (replaced_rng * amount + new_rng)
-    self.ast = self.ast.substitute({rng:sub_axis}, name=f"shift {rng.arg[:-1]} {amount} {str(new_type).split('.')[1].lower()}")
+    self.ast = self.ast.substitute({rng:sub_axis}, name=f"shift {rng.arg[0]} by {amount} {str(new_type).split('.')[1].lower()}")
     return replaced_rng, new_rng
 
   def ranges_of(self, *axis_type:AxisType) -> list[UOp]: return [r for r in self.rngs if r.arg[-1] in axis_type]
@@ -231,9 +230,9 @@ class Scheduler:
       for tc in tensor_cores:
         if tc.dtype_in == in0.dtype.scalar() and tc.dtype_in == in1.dtype.scalar() and tc.dtype_out == reduceop.dtype.scalar():
           # 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])
-          red_ranges = sorted(reduceop.src[1:], key=lambda x: -x.arg[0])
+          in0_ranges = sorted([u for u in in0.ranges if u not in in1.ranges], key=lambda x: x.arg[0], reverse=True)
+          in1_ranges = sorted([u for u in in1.ranges if u not in in0.ranges], key=lambda x: x.arg[0], reverse=True)
+          red_ranges = sorted(reduceop.src[1:], key=lambda x: x.arg[0], reverse=True)
           if DEBUG >= 3:
             print(f"TC({axis}): {[(x.arg[0],x.vmax+1) for x in in0_ranges]}",
                               f"{[(x.arg[0],x.vmax+1) for x in in1_ranges]} {[(x.arg[0],x.vmax+1) for x in red_ranges]}")
@@ -260,7 +259,8 @@ 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)
+          # the $ puts it before any numbered ranges
+          warp = UOp.range(tc.threads, '$warp', AxisType.WARP)
           ne: list[UOp] = []
           for opt in tc.opts:
             if opt[0] == "l":

tinygrad/codegen/simplify.py

@@ -47,7 +47,7 @@ def do_substitute(ctx, x: UOp):
   subs = {}
   for k,v in ctx.items():
     if v is not None:
-      subs[k] = k.replace(src=(k.src[0]//v,), arg=k.arg[0:-1]+(0,k.arg[-1]))*v + k.replace(src=(v,), arg=k.arg[0:-1]+(1,k.arg[-1]))
+      subs[k] = k.replace(src=(k.src[0]//v,), arg=(k.arg[0]+"_0", k.arg[-1]))*v + k.replace(src=(v,), arg=(k.arg[0]+"_1", k.arg[-1]))
   if not len(subs): return None
   ret = x.substitute(subs).simplify()
   ctx.clear()
@@ -152,7 +152,7 @@ def cut_store_range(ctx, store:UOp, r:UOp):
   if r.src[0].op is not Ops.CONST or ctx!="CPU": return None
   if not (cuts:=[c.src[1].arg for c in store.get_consumer_map()[r] if c.op is Ops.CMPLT and r is c.src[0] and c.src[1].op is Ops.CONST]): return None
   cuts = sorted(dedup([0] + cuts + [r.src[0].arg]))
-  ranges = [UOp.range((end-start), *(r.arg[0:-1]+(i,r.arg[-1]))) for i,(start,end) in enumerate(zip(cuts[:-1], cuts[1:]))]
+  ranges = [UOp.range((end-start), r.arg[0]+f"_{i}", r.arg[-1]) for i,(start,end) in enumerate(zip(cuts[:-1], cuts[1:]))]
 
   return UOp.group(*[store.substitute({r: new_r+start}).end(new_r) for new_r, start in zip(ranges, cuts[:-1])])
 

tinygrad/engine/schedule.py

@@ -90,7 +90,7 @@ def create_schedule_with_vars(sched_sink:UOp) -> tuple[list[ScheduleItem], dict[
       if rk.op is Ops.END: schedule.append(rk)
     else:
       raise RuntimeError(f"can't schedule {k.op}")
-    for x in children[k]:
+    for x in children[rk]:
       in_degree[x] -= 1
       if in_degree[x] == 0: queues[_heuristic(x)].append(x)
 

tinygrad/gradient.py

@@ -3,14 +3,15 @@ import math, dataclasses
 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, all_metadata
 from tinygrad.helpers import argsort
 
-def reduce_gradient(ctx:UOp, ret:UOp):
+def reduce_gradient(ctx:UOp, ret:UOp, op:Ops):
   def broadcast_to_input(x): return x.reshape(x.shape+(1,)*(len(ret.src[0].shape)-len(x.shape))).expand(ret.src[0].shape)
-  if ret.arg[0] == Ops.ADD: return (broadcast_to_input(ctx),)
-  if ret.arg[0] == Ops.MAX:
+  if op == Ops.ADD: return (broadcast_to_input(ctx),)
+  if op == Ops.MAX:
+    assert ret.op is Ops.REDUCE_AXIS, "only works on REDUCE_AXIS"
     mask = ret.src[0].eq(broadcast_to_input(ret)).cast(ctx.dtype)
     count = mask.r(Ops.ADD, ret.arg[1])
     return ((mask/broadcast_to_input(count)) * broadcast_to_input(ctx),)
-  if ret.arg[0] == Ops.MUL: return (broadcast_to_input(ctx * ret) / ret.src[0],)
+  if op == Ops.MUL: return (broadcast_to_input(ctx * ret) / ret.src[0],)
 
 # ctx is grad_output
 pm_gradient = PatternMatcher([
@@ -28,7 +29,8 @@ pm_gradient = PatternMatcher([
     ((x>y).where(ctx, (x.eq(y)).where(ctx * 0.5, 0)), (x<y).where(ctx, (x.eq(y)).where(ctx * 0.5, 0)))),
   (UPat(Ops.MUL, name="ret"), lambda ctx, ret: (ret.src[1]*ctx, ret.src[0]*ctx)),
   (UPat(Ops.WHERE, name="ret"), lambda ctx, ret: (None, ret.src[0].where(ctx, ctx.const_like(0)), ret.src[0].where(ctx.const_like(0), ctx))),
-  (UPat(Ops.REDUCE_AXIS, name="ret"), reduce_gradient),
+  (UPat(Ops.REDUCE_AXIS, name="ret"), lambda ctx, ret: reduce_gradient(ctx, ret, ret.arg[0])),
+  (UPat(Ops.REDUCE, name="ret"), lambda ctx, ret: reduce_gradient(ctx, ret, ret.arg) + (None,)*(len(ret.src)-1)),
   (UPat(Ops.CONTIGUOUS), lambda ctx: (ctx,)),
   (UPat(Ops.CONTIGUOUS_BACKWARD), lambda ctx: (ctx.contiguous(),)),
   (UPat(Ops.RESHAPE, name="ret"), lambda ctx, ret: (ctx.reshape(ret.src[0].shape), None)),
@@ -68,4 +70,8 @@ def compute_gradient(root:UOp, root_grad:UOp, targets:set[UOp]) -> dict[UOp, UOp
         # we add the backward metadata to everything new in the graph
         for bw_uop in v.toposort(lambda x: x not in (t0, *t0.src, grads[t0])):
           all_metadata[bw_uop] = all_metadata.get(bw_uop, ())+backward_metadata
+  # end any ranges on grads with a reduce sum
+  for k,v in grads.items():
+    if len(v.ranges):
+      grads[k] = v.reduce(*v.ranges, arg=Ops.ADD)
   return grads

tinygrad/runtime/autogen/am/smu_v14_0_2.py

@@ -476,7 +476,7 @@ PP_GRTAVFS_FW_SEP_FUSE_FREQUENCY_TO_COUNT_SCALER_4 = PP_GRTAVFS_FW_SEP_FUSE_e.de
 PP_GRTAVFS_FW_SEP_FUSE_COUNT = PP_GRTAVFS_FW_SEP_FUSE_e.define('PP_GRTAVFS_FW_SEP_FUSE_COUNT', 19)
 
 class SviTelemetryScale_t(Struct): pass
-int8_t = ctypes.c_char
+int8_t = ctypes.c_byte
 SviTelemetryScale_t._fields_ = [
   ('Offset', int8_t),
   ('Padding', uint8_t),

tinygrad/runtime/autogen/mesa.py

@@ -89,7 +89,7 @@ NIR_CMAT_C_SIGNED = nir_cmat_signed.define('NIR_CMAT_C_SIGNED', 4)
 NIR_CMAT_RESULT_SIGNED = nir_cmat_signed.define('NIR_CMAT_RESULT_SIGNED', 8)
 
 class nir_const_value(ctypes.Union): pass
-int8_t = ctypes.c_char
+int8_t = ctypes.c_byte
 uint8_t = ctypes.c_ubyte
 int16_t = ctypes.c_int16
 uint16_t = ctypes.c_uint16
@@ -3723,7 +3723,7 @@ struct__IO_FILE._fields_ = [
   ('_flags2', ctypes.c_int32),
   ('_old_offset', ctypes.c_int64),
   ('_cur_column', ctypes.c_uint16),
-  ('_vtable_offset', ctypes.c_char),
+  ('_vtable_offset', ctypes.c_byte),
   ('_shortbuf', (ctypes.c_char * 1)),
   ('_lock', ctypes.POINTER(_IO_lock_t)),
   ('_offset', ctypes.c_int64),

tinygrad/runtime/ops_amd.py

@@ -8,6 +8,7 @@ from tinygrad.runtime.support.hcq import MMIOInterface, BumpAllocator, hcq_filte
 from tinygrad.uop.ops import sint
 from tinygrad.device import Compiled, DMAFdRef, BufferSpec, CompilerPairT
 from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, suppress_finalizing, lo32, hi32, colored, prod, ContextVar
+from tinygrad.helpers import VIZ
 from tinygrad.renderer.cstyle import AMDRenderer
 from tinygrad.renderer.llvmir import AMDLLVMRenderer
 from tinygrad.runtime.autogen import kfd, hsa, pci, sqtt
@@ -19,7 +20,7 @@ from tinygrad.runtime.support.amd import AMDReg, AMDIP, import_module, import_so
 from tinygrad.runtime.support.system import System, PCIIfaceBase, PCIAllocationMeta, PCIDevice, USBPCIDevice, MAP_FIXED, MAP_NORESERVE
 if getenv("IOCTL"): import extra.hip_gpu_driver.hip_ioctl  # noqa: F401 # pylint: disable=unused-import
 
-SQTT, SQTT_ITRACE_SE_MASK, PMC = ContextVar("SQTT", 0), ContextVar("SQTT_ITRACE_SE_MASK", 0b11), ContextVar("PMC", 0)
+SQTT, SQTT_ITRACE_SE_MASK, PMC = ContextVar("SQTT", VIZ.value>=2), ContextVar("SQTT_ITRACE_SE_MASK", 0b11), ContextVar("PMC", 0)
 EVENT_INDEX_PARTIAL_FLUSH = 4 # based on a comment in nvd.h
 WAIT_REG_MEM_FUNCTION_EQ  = 3 # ==
 WAIT_REG_MEM_FUNCTION_NEQ = 4 # !=

tinygrad/runtime/support/autogen.py

@@ -103,7 +103,7 @@ def gen(dll, files, args=[], prolog=[], rules=[], epilog=[], recsym=False, use_e
     suggested_name = anon_names.get(f"{loc_file(loc(decl:=clang.clang_getTypeDeclaration(t)))}:{loc_line(loc(decl))}", suggested_name)
     nonlocal lines, types, anoncnt, objc
     tmap = {clang.CXType_Void:"None", clang.CXType_Char_U:"ctypes.c_ubyte", clang.CXType_UChar:"ctypes.c_ubyte", clang.CXType_Char_S:"ctypes.c_char",
-            clang.CXType_SChar:"ctypes.c_char",
+            clang.CXType_SChar:"ctypes.c_byte",
             **{getattr(clang, f'CXType_{k}'):f"ctypes.c_{k.lower()}" for k in ["Bool", "WChar", "Float", "Double", "LongDouble"]},
             **{getattr(clang, f'CXType_{k}'):f"ctypes.c_{'u' if 'U' in k else ''}int{sz}" for sz,k in
                [(16, "UShort"), (16, "Short"), (32, "UInt"), (32, "Int"), (64, "ULong"), (64, "Long"), (64, "ULongLong"), (64, "LongLong")]}}

tinygrad/schedule/indexing.py

@@ -26,6 +26,8 @@ pm_generate_realize_map = PatternMatcher([
   (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/STORE
   (UPat({Ops.COPY, Ops.BUFFER_VIEW, Ops.CONTIGUOUS, Ops.STORE}, name="tr"), realize),
+  # always realize REDUCE on outer ranges
+  (UPat(Ops.REDUCE, name="r"), lambda ctx,r: realize(ctx, r) if any(tr.arg[-1] == AxisType.OUTER for tr in r.src[1:]) else None),
   # 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)

tinygrad/schedule/rangeify.py

@@ -46,7 +46,7 @@ def split_reduceop(reduce:UOp, x:UOp):
   # get expanded by rangeifying the UOp x
   indexed = x.index(*[UOp.range(s, i) if resolve(s>1) else UOp.const(dtypes.index, 0) for i,s in enumerate(x.shape)])
   range_nums = [y.arg[0] for y in indexed.substitute({x.base:UOp(Ops.NOOP)}, extra_pm=pm_mops).ranges]
-  is_expanded = [i not in range_nums for i in range(len(x.shape))]
+  is_expanded = [str(i) not in range_nums for i in range(len(x.shape))]
 
   if not (split_candidates:=[(i,d) for i in reduce.arg[1] for d in range(min(256,2**getenv("REDUCEOP_SPLIT_SIZE",22)//prod(reduce.shape)),8-1,-1)
                              if x.shape[i]%d==0 and not is_expanded[i]]): return None
@@ -289,7 +289,7 @@ def limit_bufs(ctx:IndexingContext, root:UOp):
     for s in root.src:
       if s.op in GroupOp.Elementwise:
         # Insert bufferize: all AxisType.REDUCE before bufferize are AxisType.LOOP
-        orig_ranges, end_ranges = s.ranges, [x.replace(arg=(next(ctx.range_idx), AxisType.LOOP)) if x.op is Ops.RANGE else x for x in s.ranges]
+        orig_ranges, end_ranges = s.ranges, [x.replace(arg=(str(next(ctx.range_idx)), AxisType.LOOP)) if x.op is Ops.RANGE else x for x in s.ranges]
         s = s.substitute(dict(zip(orig_ranges, end_ranges))).bufferize(*end_ranges, arg=BufferizeOpts(device=s.device)).index(*orig_ranges)
       srcs.append(s)
     return root.replace(src=tuple(srcs))
@@ -325,6 +325,18 @@ def bufferize_to_store(ctx:itertools.count|None, x:UOp, idx:UOp, allow_locals=Tr
     for m in mops[::-1]: ret = ret._mop(*m)
     return ret
 
+  # lower outerworld reduce here
+  if x.src[0].op is Ops.REDUCE and len(x.src[0].src) == 2 and x.src[0].src[1].arg[-1] == AxisType.OUTER:
+    assert sdtype.addrspace == AddrSpace.GLOBAL
+    outer_range = x.src[0].src[1]
+    buf = UOp.new_buffer(x.arg.device, size, x.dtype)
+    # NOTE: this has the same number as the outer range, we need string ranges!
+    zero_range = outer_range.replace(src=(UOp.const(dtypes.index, size),), arg=outer_range.arg[:-1]+(AxisType.LOOP,))
+    buf = buf.after(buf.index(zero_range).store(0).end(zero_range))
+    bufi = buf.index(idx, dtype=sdtype)
+    do_store = bufi.store(bufi.load() + x.src[0].src[0], tag=x.tag).end(*rngs).end(outer_range)
+    return buf.after(do_store)
+
   # NOTE: the DEFINE_LOCAL needs to be disambiguated here
   if sdtype.addrspace == AddrSpace.GLOBAL:
     buf = UOp.new_buffer(x.arg.device, size, x.dtype)
@@ -400,7 +412,7 @@ def renumber_range(ctx:LocalAddBufferContext, r:UOp):
   if r.arg[-1] == AxisType.OUTER:
     # for outer range, we replace with a bound variable
     return UOp.variable("range_"+range_str(r), r.vmin, r.vmax).bind(r.replace(tag=None))
-  ret = r.replace(arg=(ctx.range,)+r.arg[1:], tag=None)
+  ret = r.replace(arg=(str(ctx.range),)+r.arg[1:], tag=None)
   ctx.range += 1
   return ret
 
@@ -472,6 +484,7 @@ pm_add_range_tags = PatternMatcher([
 ])
 
 def split_store(ctx:list[UOp], x:UOp) -> UOp|None:
+  # if we have any outer ranges open here, we don't split
   if len([r for r in x.ranges if r.arg[-1] != AxisType.OUTER]): return None
 
   # ends of outer range don't go in kernels
@@ -543,7 +556,7 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
   # convert movement ops to ranges
   tsink, rctx = run_rangeify(tsink, DEBUG_RANGEIFY)
 
-  tsink = graph_rewrite(tsink, symbolic+pm_reduce_simplify+pm_const_buffer_folding, name="symbolic+reduce_collapse")  # this does const folding
+  tsink = graph_rewrite(tsink, symbolic+pm_reduce_simplify+pm_const_buffer_folding, name="symbolic+reduce_collapse")
   tsink = graph_rewrite(tsink, pm_remove_bufferize, bottom_up=True, name="remove bufferize with cost function")
   tsink = graph_rewrite(tsink, symbolic+pm_reduce_simplify+pm_const_buffer_folding, name="symbolic+reduce_collapse pt 2")
   tsink = graph_rewrite(tsink, pm_limit_bufs, ctx=rctx, name="limit buffers")

tinygrad/uop/ops.py

@@ -49,8 +49,8 @@ def ssimplify(uop:sint): return uop.ssimplify() if isinstance(uop, UOp) else uop
 def sym_infer(uop: UOp|int, var_vals: dict[str, int]) -> int: return uop.sym_infer(var_vals) if isinstance(uop, UOp) else uop
 
 def range_str(u:UOp, color=False) -> str:
-  ret = '_'.join([str(x) if x >= 0 else "m"+str(-x) for x in u.arg[0:-1]])
-  return colored(ret, axis_colors[u.arg[-1]]) if color else ret
+  assert len(u.arg) == 2
+  return colored(u.arg[0], axis_colors[u.arg[-1]]) if color else u.arg[0]
 
 def multirange_str(rngs:Iterable[UOp], color=False, pad=None) -> str:
   ret = ','.join([range_str(x, color=color) for x in sorted(rngs, key=lambda x: x.arg)])
@@ -86,11 +86,12 @@ class UOpMetaClass(type):
     if _buffer is not None:
       assert op is Ops.BUFFER, f"trying to set Buffer {_buffer} for {op}"
       buffers[created] = _buffer
+    if op is Ops.RANGE: assert isinstance(arg[0], str)
     if SPEC > 1:
       from tinygrad.uop.spec import full_spec, test_pyrender
       if SPEC > 2: test_pyrender(created)
-      with Context(IGNORE_OOB=1): ret = full_spec.rewrite(created)
-      if cast(bool|None, ret) is not True: raise RuntimeError(f"SPEC ISSUE {ret}: {created}")
+      with Context(IGNORE_OOB=1): fret = cast(bool|None, full_spec.rewrite(created))
+      if fret is not True: raise RuntimeError(f"SPEC ISSUE {fret}: {created}")
     return created
 
 # some uops map to other stuff
@@ -425,8 +426,10 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     if shape is not None: ret = ret.reshape((1,)*len(shape)).expand(shape)
     return ret
   @staticmethod
-  def range(end:sint, axis_id, axis_type=AxisType.LOOP, *arg, dtype=dtypes.index, src=(), **kwargs):
-    return UOp(Ops.RANGE, dtype=dtype, src=(sint_to_uop(end, dtype),)+src, arg=(axis_id, axis_type)+arg, **kwargs)
+  def range(end:sint, axis_id:str|int, axis_type=AxisType.LOOP, dtype=dtypes.index, src=(), **kwargs):
+    assert isinstance(axis_type, AxisType), f"{axis_type} must be an AxisType"
+    str_axis_id = ("m"+str(-axis_id)) if isinstance(axis_id, int) and axis_id < 0 else str(axis_id)
+    return UOp(Ops.RANGE, dtype=dtype, src=(sint_to_uop(end, dtype),)+src, arg=(str_axis_id, axis_type), **kwargs)
   @staticmethod
   def special(end:sint, name:str, dtype=dtypes.index): return UOp(Ops.SPECIAL, dtype=dtype, src=(sint_to_uop(end, dtype),), arg=name)
   def r(self, op:Ops, axis:tuple[int, ...]):
@@ -583,7 +586,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
   def new_buffer(device:str|tuple[str, ...], size:int, dtype:DType, num=None):
     return UOp(Ops.BUFFER, dtype, (UOp.unique(num), UOp(Ops.DEVICE, arg=device)), size)
   @property
-  def device(self) -> str|tuple[str, ...]: return cast(str|tuple[str, ...], unwrap(self._device))
+  def device(self) -> str|tuple[str, ...]: return unwrap(self._device)
   @recursive_property
   def _device(self) -> str|tuple[str, ...]|None:
     if self.op is Ops.DEVICE: return self.arg
@@ -1164,12 +1167,12 @@ class RewriteContext:
 
   def cached_pm_rewrite(self, x:UOp):
     if (ret:=self.pm_cache.get(x,SENTINEL)) is not SENTINEL: return ret
-    ret = self.pm_cache[x] = cast(PatternMatcher, self.pm).rewrite(x, self.ctx)
+    ret = self.pm_cache[x] = unwrap(self.pm).rewrite(x, self.ctx)
     return ret
 
   def cached_bpm_rewrite(self, x:UOp):
     if (ret:=self.bpm_cache.get(x,SENTINEL)) is not SENTINEL: return ret
-    ret = self.bpm_cache[x] = cast(PatternMatcher, self.bpm).rewrite(x, self.ctx)
+    ret = self.bpm_cache[x] = unwrap(self.bpm).rewrite(x, self.ctx)
     return ret
 
   def unified_rewrite(self, root:UOp) -> UOp:
@@ -1350,7 +1353,7 @@ pm_pyrender_extra = PatternMatcher([
   (UPat(Ops.REDUCE_AXIS, name="r"), lambda ctx,r: f"{ctx[r.src[0]]}.r({r.arg[0]}, {r.arg[1]})"),
   # NOTE: range has srcs sometimes after control flow
   (UPat(Ops.RANGE, src=(UPat(Ops.CONST, name="c"),), allow_any_len=True, name="x"), lambda ctx,x,c:
-    "UOp.range("+', '.join([str(c.arg)] + [str(y) for y in x.arg])+
+    "UOp.range("+', '.join([str(c.arg)] + [repr(y) for y in x.arg])+
       (f', src={srcs(ctx, x.src[1:])}' if len(x.src) > 1 else '')+(', dtype='+str(x.dtype) if x.dtype is not dtypes.index else '')+")"),
   # TODO: index shouldn't mismatch dtype
   (UPat(Ops.INDEX, src=(UPat(), UPat()), allow_any_len=True, name="x"), lambda ctx,x:

tinygrad/uop/spec.py

@@ -37,8 +37,7 @@ shared_spec = PatternMatcher([
 
   # RANGE can be in the big graph now
   (UPat(Ops.RANGE, src=(UPat.var("x"),), allow_any_len=True, name="rng"), lambda rng,x:
-    rng.dtype == x.dtype and isinstance(rng.arg, tuple) and len(rng.arg) >= 2 and \
-      all(isinstance(ra, int) for ra in rng.arg[0:-1]) and isinstance(rng.arg[-1], AxisType)),
+    rng.dtype == x.dtype and isinstance(rng.arg, tuple) and len(rng.arg) == 2 and isinstance(rng.arg[0], str) and isinstance(rng.arg[-1], AxisType)),
   (UPat(Ops.INDEX, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(y.dtype == dtypes.index for y in x.src[1:]) or None),
 
   # RANGE/SPECIAL define loops, END closes them

tinygrad/uop/symbolic.py

@@ -2,7 +2,7 @@
 import math, operator, struct, functools
 from collections import defaultdict
 from tinygrad.uop.ops import Ops, PatternMatcher, UPat, UOp, GroupOp, exec_alu
-from tinygrad.dtype import ConstType, dtypes, PtrDType, AddrSpace, can_safe_cast, Invalid
+from tinygrad.dtype import ConstType, dtypes, PtrDType, can_safe_cast, Invalid
 from tinygrad.helpers import partition, all_same, prod, flatten, get_single_element, cdiv, cmod, CORRECT_DIVMOD_FOLDING, unwrap
 from tinygrad.uop.decompositions import xpow
 
@@ -407,14 +407,10 @@ def uop_given_valid(valid:UOp, uop:UOp, try_simplex=True) -> UOp:
     expr, is_upper, c = res
     bounds[expr][int(is_upper)] = c
 
-  # don't simplify any other gates, can lead to OOB, we substitute them back later
-  uop = uop.substitute((load_subs:={u: UOp(Ops.NOOP, dtype=u.dtype, arg=u) for u in uop.toposort() if u.op is Ops.INDEX}))
-
   # simplify uop given that valid is True
   all_candidates = []
   for i,(expr,v) in enumerate(bounds.items()):
     v0, v1 = (expr.vmin if v[0] is None else v[0], expr.vmax if v[1] is None else v[1])
-    expr = expr.substitute(load_subs)  # make sure expr appears in same form in the uop
     # try checking the whole clause
     all_candidates.append((expr, UOp.variable(f"fake{i}", v0, v1, expr.dtype)))
 
@@ -438,8 +434,6 @@ def uop_given_valid(valid:UOp, uop:UOp, try_simplex=True) -> UOp:
   # try all the valids together (but only the whole expressions)
   if (s_uop:=uop.substitute(sub_dict:=dict(all_candidates))) is not uop:
     uop = s_uop.simplify().substitute({newX:X for X,newX in sub_dict.items()}).simplify(full_symbolic=False)
-  # put the loads back in
-  uop = uop.substitute({v:k for k,v in load_subs.items()})
   return uop
 
 def _valid_priority(v: UOp, valids:list[UOp]):
@@ -456,7 +450,7 @@ def simplify_valid(valid:UOp) -> UOp|None:
     if ret[-1] is not stmt: something_changed = True
   return UOp.prod(*ret) if something_changed else None
 
-# ******** phase 3 is the complete symbolic, and deals with very complex things like loop rewriting and threefry transform ********
+# ******** phase 3 is the complete symbolic ********
 
 def reduce_mul_chain(r:UOp):
   if r.arg not in {Ops.ADD, Ops.MAX}: return None
@@ -485,6 +479,8 @@ def where_on_load(c1, buf, x):
   # aditionally we can drop the clause on the where if it already exists in the load
   remaining_clause = UOp.const(dtypes.bool, True).prod(*[c for c in c1.split_uop(Ops.AND) if c not in removed])
   return remaining_clause.where(buf.index(x.get_idx().valid(functools.reduce(operator.and_, moved_clauses, c2))), 0)
+
+# where after gated load becomes alt value, TODO: this is sort of duplicated with rules in devectorizer
 pm_move_where_on_load = PatternMatcher([
   (UPat.var("c1").where(UPat.var("buf").index(UPat.var("x")), 0), where_on_load),
   (UPat.var("c1").where(0, UPat.var("buf").index(UPat.var("x"))), lambda c1,buf,x: where_on_load(c1.logical_not(),buf,x)),
@@ -500,9 +496,6 @@ pm_simplify_valid = PatternMatcher([
 # this is symbolic 2.0
 REMOVE_FROM_SINK_LIKE = {Ops.UNROLL, Ops.NOOP, Ops.VECTORIZE, Ops.SINK}
 sym = symbolic+pm_simplify_valid+PatternMatcher([
-  # LOAD/STORE -> NOOP
-  (UPat.var('x').store(UPat.var('x').load(), allow_any_len=True), lambda x: None if x.dtype.addrspace != AddrSpace.REG else x.src[0].src[0]),
-  (UPat(Ops.LOAD, src=(UPat.cvar('c'))), lambda c: c),
   # VECTORIZE/GEP
   (UPat(Ops.VECTORIZE, src=UPat(Ops.GEP, src=(UPat.var("x"),)), name="vec"), lambda vec,x: x.gep(tuple(y.arg[0] for y in vec.src))),
   # reorder ALU/VECTORIZE
@@ -531,7 +524,6 @@ sym = symbolic+pm_simplify_valid+PatternMatcher([
   # fold gated LOAD/STORE
   (UPat((Ops.LOAD, Ops.STORE), src=(UPat().index(UPat.const(dtypes.index, Invalid)).or_casted(),), allow_any_len=True, name="x"),
     lambda x: UOp(Ops.NOOP) if x.op is Ops.STORE else x.const_like(0)), # invalid store does nothing. invalid load produces 0
-  # # Where after gated load becomes alt value, TODO: this is sort of duplicated with rules in devectorizer
   ((UPat.var("x") * UPat.var("x")).reciprocal(), lambda x: x.reciprocal()*x.reciprocal()),  # 1/(x^c) -> (1/x)^c
   ((UPat.var("x") * UPat.var("x") * UPat.var("x")).reciprocal(), lambda x: x.reciprocal()*x.reciprocal()*x.reciprocal()),
   ((UPat.var("x") * UPat.cvar("c")).reciprocal(), lambda x,c: x.reciprocal()*c.reciprocal()), # 1/(x*c) -> (1/c)*(1/x)

tinygrad/viz/index.html

@@ -271,7 +271,7 @@
   }
   #device-list > div {
     min-height: 32px;
-    width: 132px;
+    width: 134px;
     overflow-x: auto;
     overflow-y: hidden;
     white-space: nowrap;

tinygrad/viz/js/index.js

@@ -36,7 +36,7 @@ const updateProgress = ({ start, err }) => {
   d3.select("#custom").html("");
   if (err) {
     displaySelection("#custom");
-    d3.select("#custom").append(() => d3.create("div").classed("raw-text", true).call(s => s.append(() => codeBlock(err, "txt"))).node());
+    d3.select("#custom").append("div").classed("raw-text", true).call(s => s.append(() => codeBlock(err, "txt"))).node();
   }
 }