don't bufferize 0s

.github/workflows/test.yml

@@ -601,22 +601,11 @@ jobs:
     - name: Setup Environment
       uses: ./.github/actions/setup-tinygrad
       with:
-        key: rangeify-minimal-llvm
+        key: rangeify-minimal
         deps: testing_minimal
-        llvm: "true"
     - name: Test CPU=1 RANGEIFY=1
       # TODO: add more passing tests here
-      # test_symbolic_arange_sym_step is passing now
-      # test_threefry_doesnt_use_long is because there's a contig after the long now
-      run: |
-        CPU=1 RANGEIFY=1 python3 -m pytest -n auto --durations 20 \
-          -k "not test_symbolic_arange_sym_step and not test_threefry_doesnt_use_long" \
-          test/test_tiny.py test/test_rangeify.py test/test_ops.py test/test_tensor_variable.py \
-          test/test_outerworld_range.py test/test_sample.py test/test_randomness.py test/test_tensor_data.py
-    - name: Test CPU=1 RANGEIFY=2
-      run: CPU=1 RANGEIFY=2 python3 -m pytest -n auto test/test_tiny.py test/test_rangeify.py test/test_ops.py --durations 20
-    - name: Test LLVM=1 RANGEIFY=1 (slow tests)
-      run: LLVM=1 RANGEIFY=1 python3 -m pytest -n auto test/models/test_mnist.py --durations 20
+      run: CPU=1 RANGEIFY=1 python3 -m pytest -n auto test/test_tiny.py test/test_rangeify.py test/test_ops.py --durations 20
 
   testdevectorize:
     name: Linux (devectorize)

examples/beautiful_mnist.py

@@ -41,9 +41,6 @@ if __name__ == "__main__":
     if i%10 == 9: test_acc = get_test_acc().item()
     t.set_description(f"loss: {loss.item():6.2f} test_accuracy: {test_acc:5.2f}%")
 
-  GlobalCounters.reset()   # NOTE: this makes it nice for DEBUG=2 timing
-  test_acc = get_test_acc().item()
-
   # verify eval acc
   if target := getenv("TARGET_EVAL_ACC_PCT", 0.0):
     if test_acc >= target and test_acc != 100.0: print(colored(f"{test_acc=} >= {target}", "green"))

ruff.toml

@@ -35,7 +35,6 @@ lint.select = [
 line-length = 150
 
 exclude = [
-  ".git/",
   "docs/",
   "extra/",
   "tinygrad/runtime/autogen",

test/unit/test_viz.py

@@ -5,7 +5,7 @@ from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, TrackedPatternMatch
 from tinygrad.uop.ops import graph_rewrite, track_rewrites, TRACK_MATCH_STATS
 from tinygrad.uop.symbolic import sym
 from tinygrad.dtype import dtypes
-from tinygrad.helpers import PROFILE, colored, ansistrip, flatten, TracingKey, ProfileRangeEvent, Context
+from tinygrad.helpers import PROFILE, colored, ansistrip, flatten, TracingKey, ProfileRangeEvent
 from tinygrad.device import Buffer
 
 @track_rewrites(name=True)
@@ -240,15 +240,6 @@ class TestVizIntegration(BaseTestViz):
     self.assertEqual(lst[0]["name"], "Schedule 1 Kernel n1")
     self.assertEqual(lst[1]["name"], prg.name)
 
-  def test_metadata_tracing(self):
-    with Context(TRACEMETA=2):
-      a = Tensor.empty(1)
-      b = Tensor.empty(1)
-      metadata = (alu:=a+b).uop.metadata
-      alu.kernelize()
-      graph = next(get_details(tracked_ctxs[0][0]))["graph"]
-    self.assertEqual(len([n for n in graph.values() if repr(metadata) in n["label"]]), 1)
-
 from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry
 from tinygrad.viz.serve import get_profile
 

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
 from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp
 from tinygrad.uop.spec import type_verify
 from tinygrad.renderer import Renderer
@@ -18,7 +18,6 @@ from tinygrad.codegen.devectorizer import load_store_folding, load_store_indexin
 from tinygrad.codegen.linearize import block_create, pm_blockend_merge, block_merge, pm_finalize, BlockContext
 from tinygrad.codegen.opt import pm_optimize
 from tinygrad.codegen.opt.swizzler import view_left, view_right, fix_kernel_ops
-from tinygrad.codegen.opt.postrange import pm_postrange_opt
 
 @dataclass
 class RewriteStep:
@@ -45,10 +44,10 @@ rewrites_for_linearizer = [
 
 def get_rewrites_for_renderer(opts:Renderer, linearizer:bool=True) -> list[RewriteStep]:
   # cache with the values of the context vars
-  return _get_rewrites_for_renderer(opts, linearizer, QUANTIZE.value, DEVECTORIZE.value, TRANSCENDENTAL.value, RANGEIFY.value)
+  return _get_rewrites_for_renderer(opts, linearizer, QUANTIZE.value, DEVECTORIZE.value, TRANSCENDENTAL.value)
 
 @functools.cache
-def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVECTORIZE, _TRANSCENDENTAL, _RANGEIFY) -> list[RewriteStep]:
+def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVECTORIZE, _TRANSCENDENTAL) -> list[RewriteStep]:
   # ** lowerer (rewrite_shapetracker_with_index) **
   ret: list[RewriteStep] = []
 
@@ -58,10 +57,6 @@ def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVEC
   # this is kernel.py
   ret.append(RewriteStep(pm_optimize, ctx=lambda _: opts, name="optimize ast"))
 
-  # this is the new optimizer
-  if _RANGEIFY:
-    ret.append(RewriteStep(pm_postrange_opt, ctx=lambda _: opts, name="new 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))
 

tinygrad/codegen/opt/postrange.py

@@ -1,25 +0,0 @@
-from tinygrad.dtype import dtypes
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp
-from tinygrad.helpers import partition, flatten, prod
-
-def unroll_range(r:UOp):
-  # all ranges sub 5 can be UNROLLS
-  if r.src[0].op is Ops.CONST and r.vmax < 5:
-    i = r.arg
-    s = r.vmax+1
-    return UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(s), tuple(range(s))),), ((i,s),), tag=1)
-
-def fix_reduce(r:UOp):
-  reduce_range, reduce_expand = partition(r.src[1:], lambda y: y.op is Ops.RANGE)
-  if len(reduce_expand) == 0: return None
-  assert all(x.op is Ops.UNROLL for x in reduce_expand), f"not all UNROLLS in {reduce_expand}"
-  ret = r.src[0]
-  if len(contract_axis:=flatten(x.arg for x in reduce_expand)):
-    ret = UOp(Ops.CONTRACT, r.dtype.vec(prod(x[1] for x in contract_axis)), (ret,), tuple(contract_axis), tag=1)
-  # REDUCE supports both "horizontal" reduction and range reduction. the horizontal elements are taken in the nearest group
-  return UOp(Ops.REDUCE, r.dtype, (ret,)+tuple(reduce_range), r.arg)
-
-pm_postrange_opt = PatternMatcher([
-  (UPat(Ops.RANGE, name="r"), unroll_range),
-  (UPat(Ops.REDUCE, name="r"), fix_reduce),
-])

tinygrad/helpers.py

@@ -140,7 +140,7 @@ DONT_REALIZE_EXPAND, DONT_GROUP_REDUCES = ContextVar("DONT_REALIZE_EXPAND", 0),
 QUANTIZE, VALIDATE_WITH_CPU, DISABLE_FAST_IDIV = ContextVar("QUANTIZE", 0), ContextVar("VALIDATE_WITH_CPU", 0), ContextVar("DISABLE_FAST_IDIV", 0)
 CORRECT_DIVMOD_FOLDING, FUSE_OPTIM = ContextVar("CORRECT_DIVMOD_FOLDING", 0), ContextVar("FUSE_OPTIM", 0)
 ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, AMD_LLVM = ContextVar("ALLOW_DEVICE_USAGE", 1), ContextVar("MAX_BUFFER_SIZE", 0), ContextVar("AMD_LLVM", 1)
-RANGEIFY = ContextVar("RANGEIFY", 0)
+RANGEIFY, PARTIAL_CONTIG = ContextVar("RANGEIFY", 0), ContextVar("PARTIAL_CONTIG", 0)
 
 @dataclass(frozen=True)
 class Metadata:

tinygrad/renderer/llvmir.py

@@ -45,10 +45,10 @@ def render_wmma_amx(ctx, wmma: UOp) -> str:
       f'  call void asm sideeffect "nop\\0Anop\\0Anop\\0A.word ({0x201000 + (17 << 5) + 1})", "~{{memory}}"() #0; AMX clr',             # clr
       f'  {ctx[wmma]} = load {ldt(wmma.dtype)}, ptr {ctx[wmma]}_amx2, align {wmma.dtype.itemsize}'])
 
-def render_wmma_amd(ctx, wmma: UOp, cdna=False) -> str:
-  dt_map = {dtypes.half: "f16", dtypes.float: "f32", dtypes.ushort: "bf16.1k" if cdna else "bf16", dtypes.bfloat16: "bf16.1k" if cdna else "bf16"}
+def render_wmma_amd(ctx, wmma: UOp, arch: str) -> str:
+  dt_map = {dtypes.half: "f16", dtypes.float: "f32", dtypes.bfloat16: "bf16", dtypes.ushort: "bf16"}
   # https://github.com/llvm/llvm-project/blob/main/clang/test/CodeGenOpenCL/builtins-amdgcn-mfma.cl
-  if cdna:
+  if arch.split(":")[0] in {"gfx942", "gfx950"}:
     return f"  {ctx[wmma]} = call {ldt(wmma.dtype)} @llvm.amdgcn.mfma.{dt_map[wmma.src[-1].dtype.scalar()]}" + \
            f".16x16x16{dt_map[wmma.src[0].dtype.scalar()]}(" + ", ".join([f"{ldt(w.dtype)} {ctx[w]}" for w in wmma.src]) + ", i32 0, i32 0, i32 0)"
   # https://github.com/llvm/llvm-project/blob/main/llvm/test/CodeGen/AMDGPU/GlobalISel/llvm.amdgcn.wmma_32.ll
@@ -221,14 +221,7 @@ class AMDLLVMRenderer(LLVMRenderer):
   def __init__(self, arch:str):
     self.arch = arch
     self.tensor_cores = AMDRenderer.get_tensor_cores(arch)
-    self.is_cdna = arch.split(":")[0] in {"gfx942", "gfx950"}
-    self.string_rewrite += PatternMatcher([(UPat(Ops.WMMA, name="wmma"), lambda ctx, wmma, cdna=self.is_cdna: render_wmma_amd(ctx, wmma, cdna))])
-    if self.is_cdna:
-      self.extra_matcher += PatternMatcher([
-        (UPat(Ops.WMMA, name="x", dtype=dtypes.float.vec(4)),
-          lambda x: UOp(Ops.WMMA, dtypes.float.vec(4), (x.src[0].bitcast(dtypes.uint16.vec(4)), x.src[1].bitcast(dtypes.uint16.vec(4)),
-            x.src[2]), (*x.arg,)) if x.src[0].dtype == dtypes.bfloat16.vec(4) else None)
-      ])
+    self.string_rewrite += PatternMatcher([(UPat(Ops.WMMA, name="wmma"), lambda ctx, wmma, arch=arch: render_wmma_amd(ctx, wmma, arch))])
     if self.arch.split(":")[0] == "gfx1100":
       self.extra_matcher += PatternMatcher([
         (UPat(Ops.WMMA, name="x", dtype=dtypes.half.vec(8)),

tinygrad/runtime/ops_remote.py

@@ -306,10 +306,7 @@ class RemoteHandler:
           case ProgramAlloc():
             lib = dev.compiler.compile_cached(req._h[c.datahash].decode())
             session.programs[(c.name, c.datahash)] = dev.runtime(c.name, lib)
-          case ProgramFree():
-            key = (c.name, c.datahash)
-            # WORKAROUND: should be unconditional once the protocol supports proper exception handling
-            if key in session.programs: del session.programs[key]
+          case ProgramFree(): del session.programs[(c.name, c.datahash)]
           case ProgramExec():
             bufs = [session.buffers[x]._buf for x in c.bufs]
             extra_args = {k:v for k,v in [("global_size", c.global_size), ("local_size", c.local_size)] if v is not None}
@@ -424,24 +421,19 @@ class RemoteConnection:
     conns = RemoteConnection.all.keys()
     datas = {conn: conn.req.serialize() for conn in conns}
     reqs, hashes, hash_datas = sum(len(c.req._q) for c in conns), sum(len(c.req._h) for c in conns), sum(len(data) for data in datas.values())
-    resps = []
     with Timing(f"*** send {reqs:-3d} requests {hashes:-3d} hashes with len {hash_datas/1024:.2f} kB in ", enabled=DEBUG>=3):
       for conn,data in datas.items(): conn.conn.request("POST", "/batch", data)
       for conn in datas.keys():
-        resp = conn.conn.getresponse()
-        body = resp.read()
-        resps.append((conn, resp, body))
-        conn.req = BatchRequest()
-    if take_q: RemoteConnection.q_lock.release()
-    for conn,resp,body in resps:
-      match resp.status:
-        case http.HTTPStatus.OK: pass
-        case http.HTTPStatus.INTERNAL_SERVER_ERROR:
-          exc_wrapper = safe_eval(ast.parse(body.decode(), mode="eval").body)
+        response = conn.conn.getresponse()
+        resp = response.read()
+        conn.req = BatchRequest() # no matter what response, reset conn
+        if response.status == http.HTTPStatus.INTERNAL_SERVER_ERROR:
+          exc_wrapper = safe_eval(ast.parse(resp.decode(), mode="eval").body)
           exc_wrapper.exc.add_note(exc_wrapper.trace)
           raise exc_wrapper.exc
-        case code: raise RuntimeError(f"POST /batch failed with {code}: {body.decode()}")
-      if conn == self: ret = body
+        assert response.status == http.HTTPStatus.OK, f"POST /batch failed: {resp.decode()}"
+        if conn == self: ret = resp
+    if take_q: RemoteConnection.q_lock.release()
     return ret
 
 def parse_hosts(hs:str) -> list[tuple[str, int]]|LazySeq[tuple[str, int]]:

tinygrad/runtime/support/am/amdev.py

@@ -239,7 +239,7 @@ class AMDev(PCIDevImplBase):
       ip_offset = ctypes.addressof(self.bhdr) + ctypes.sizeof(dhdr) + ihdr.die_info[num_die].die_offset
       for _ in range(dhdr.num_ips):
         ip = am.struct_ip_v4.from_address(ip_offset)
-        ba = ((ctypes.c_uint64 if ihdr.base_addr_64_bit else ctypes.c_uint32) * ip.num_base_address).from_address(ip_offset + 8)
+        ba = (ctypes.c_uint32 * ip.num_base_address).from_address(ip_offset + 8)
         for hw_ip in range(1, am.MAX_HWIP):
           if hw_ip in hw_id_map and hw_id_map[hw_ip] == ip.hw_id:
             self.regs_offset[hw_ip][ip.instance_number] = tuple(list(ba))

tinygrad/schedule/rangeify.py

@@ -2,7 +2,7 @@ from typing import Any
 from dataclasses import dataclass, field
 from tinygrad.dtype import dtypes, PtrDType
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, RewriteNotReady, _substitute
-from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, colored, RANGEIFY
+from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, colored, PARTIAL_CONTIG
 from tinygrad.schedule.multi import multi_pm
 
 from tinygrad.schedule.kernelize import Kernel
@@ -10,12 +10,7 @@ from tinygrad.uop.ops import track_rewrites, graph_rewrite_map, graph_rewrite, K
 
 # 0. do some cleanup rewrites, mostly copied from the old stuff
 
-double_reshape = PatternMatcher([
-  # RESHAPE on RESHAPE is the second reshape
-  (UPat(Ops.RESHAPE, src=(UPat(Ops.RESHAPE),), name="x"), lambda x: x.replace(src=(x.src[0].src[0],))),
-])
-
-earliest_rewrites = double_reshape+PatternMatcher([
+earliest_rewrites = PatternMatcher([
   # UOp with size 0 is zero
   (UPat(GroupOp.All-{Ops.SINK}, name="root"), lambda root: root.const_like(0) if root.base.st is not None and root.size == 0 else None),
   # DETACH and CONTIGUOUS_BACKWARD are NOOPs here, so is FUSE
@@ -23,6 +18,8 @@ earliest_rewrites = double_reshape+PatternMatcher([
   # reduce of size 0 is the identity element
   (UPat(Ops.REDUCE_AXIS, name="reduce", src=(UPat.var("x"),)),
    lambda reduce,x: reduce.const_like(identity_element(reduce.arg[0], reduce.dtype)) if x.size == 0 and reduce.size != 0 else None),
+  # RESHAPE on RESHAPE is the second reshape
+  (UPat(Ops.RESHAPE, src=(UPat(Ops.RESHAPE),), name="x"), lambda x: x.replace(src=(x.src[0].src[0],))),
   # non shape changing RESHAPE is NOOP
   (UPat(Ops.RESHAPE, name="x"), lambda x: x.src[0] if x.src[0].shape == x.arg else None),
   # RESHAPE after COPY
@@ -36,8 +33,6 @@ earliest_rewrites = double_reshape+PatternMatcher([
   # assign only to buffer
   (UPat(Ops.ASSIGN, src=(UPat(GroupOp.All-{Ops.BUFFER}, name="target"), UPat(name="x"))),
    lambda x,target: x if target.base.op is not Ops.BUFFER else None),
-  # contiguous/buffer/copy/assign is already contiguous
-  (UPat(Ops.CONTIGUOUS, name="root", src=(UPat((Ops.CONTIGUOUS, Ops.BUFFER, Ops.COPY, Ops.ASSIGN)),)), lambda root: root.src[0]),
 ])
 
 # 1. add contiguous where we have to
@@ -52,9 +47,6 @@ def realize_parents(ctx:dict[UOp, None], rb:UOp) -> None:
   for s in rb.src:
     if s.op not in ALWAYS_CONTIGUOUS: ctx[s] = None
 
-def realize_assign(ctx:dict[UOp, None], a:UOp) -> None:
-  if a.src[1].op not in ALWAYS_CONTIGUOUS: ctx[a.src[1]] = None
-
 do_realize = PatternMatcher([
   # always realize SINK parents
   (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)),
@@ -62,8 +54,6 @@ do_realize = PatternMatcher([
   (UPat({Ops.ASSIGN, Ops.COPY, Ops.BUFFER_VIEW}, name="tr"), realize),
   # realize parents of COPY, MSELECT, MSTACK
   (UPat((Ops.COPY, Ops.MSELECT, Ops.MSTACK), name="rb"), realize_parents),
-  # realize input to assign (might be optimized out)
-  (UPat(Ops.ASSIGN, name="a"), realize_assign),
 ])
 
 add_contiguous = PatternMatcher([
@@ -102,29 +92,35 @@ pm_children = PatternMatcher([
 
 @dataclass
 class RangeifyContext:
-  # block on parent until all children have been seen
+  idx: int = 0
+  regs: int = 0
   seen_children: dict[UOp, dict[int, UOp]] = field(default_factory=dict)
   seen_child: dict[UOp, Any] = field(default_factory=dict)
   progress: int = 0
-
-  # create ranges
-  range_idx: int = 0
+  children: dict[UOp, list[UOp]]|None = None
   def new_range(self, s:sint):
-    ret = UOp.range(dtypes.int, s, self.range_idx)
-    self.range_idx += 1
+    ret = UOp.range(dtypes.int, s, self.idx)
+    self.idx += 1
     return ret
 
-def map_reshape(idx:UOp, r:UOp):
+def collapse_to_1(shp:tuple[sint, ...], idxs:tuple[UOp, ...]) -> UOp:
   acc = 1
   to_sum = []
-  for s,src in list(zip(idx.shape, idx.src[1:]))[::-1]:
+  for s,src in list(zip(shp, idxs))[::-1]:
     to_sum.append(acc*src)
     acc *= s
-  mish = sum(to_sum, start=UOp.const(dtypes.int, 0))
+  return sum(to_sum, start=UOp.const(dtypes.int, 0))
+
+def map_reshape(idx:UOp, r:UOp):
+  mish = collapse_to_1(idx.shape, idx.src[1:])
   ret:list[UOp] = []
   for s in r.src[0].shape[::-1]:
-    ret.append(mish % s) # NOTE: simplify will turn this to CONST
-    mish //= s
+    if resolve(s!=1):
+      # this MOD should limit any ranges outside s
+      ret.append(mish % s)
+      mish //= s
+    else:
+      ret.append(UOp.const(dtypes.int, 0))
   tret = ret[0].sink(*ret[1:]).simplify().src[::-1] if len(ret) else ()
   return r.src[0].index(*tret, dtype=idx.dtype, arg=idx.arg)
 
@@ -138,7 +134,6 @@ def map_pad(idx:UOp, r:UOp):
     if resolve(s > 0): where = where & (ret[i] >= s)
     bigwhere = bigwhere & where
     # this is safe but dumb
-    # TODO (S-Lykles): switch to mixed index/valid
     ret[i] = (ret[i] - s).maximum(0).minimum(r.src[0].shape[i]-1)
   # PAD is with 0
   return bigwhere.simplify().where(r.src[0].index(*ret, dtype=idx.dtype, arg=idx.arg), UOp.const(r.dtype, 0))
@@ -161,24 +156,24 @@ def map_expand(r:UOp, idx:UOp):
 
 pm_mops = PatternMatcher([
   # this is like the definitions of these
-  (UPat(Ops.SHRINK, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"),
+  (UPat(Ops.INDEX, src=(UPat(Ops.SHRINK, name="r"),), allow_any_len=True, name="idx"),
    lambda r,idx: r.src[0].index(*[a+ss if resolve(ss != 0) else a for a,(ss,_) in zip(idx.src[1:], r.arg)], dtype=idx.dtype, arg=idx.arg)),
-  (UPat(Ops.PERMUTE, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"),
+  (UPat(Ops.INDEX, src=(UPat(Ops.PERMUTE, name="r"),), allow_any_len=True, name="idx"),
    lambda r,idx: r.src[0].index(*[idx.src[1+p] for p in argsort(idx.src[0].arg)], dtype=idx.dtype, arg=idx.arg)),
-  (UPat(Ops.FLIP, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"),
+  (UPat(Ops.INDEX, src=(UPat(Ops.FLIP, name="r"),), allow_any_len=True, name="idx"),
    lambda r,idx: r.src[0].index(*[((s-1)-a) if f else a for a,s,f in zip(idx.src[1:], r.shape, r.arg)], dtype=idx.dtype, arg=idx.arg)),
   # expand needs to end ranges
-  (UPat(Ops.EXPAND, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"), map_expand),
+  (UPat(Ops.INDEX, src=(UPat(Ops.EXPAND, name="r"),), allow_any_len=True, name="idx"), map_expand),
   # reshape does a lot of symbolic stuff
-  (UPat(Ops.RESHAPE, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"), map_reshape),
+  (UPat(Ops.INDEX, src=(UPat(Ops.RESHAPE, name="r"),), allow_any_len=True, name="idx"), map_reshape),
   # pad adds min and max
-  (UPat(Ops.PAD, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"), map_pad),
+  (UPat(Ops.INDEX, src=(UPat(Ops.PAD, name="r"),), allow_any_len=True, name="idx"), map_pad),
 ])
 
 def map_partial_contiguous(ctx:RangeifyContext, x:UOp, idx:UOp):
   if x.arg is None: return None  # map_contiguous can handle this
   # NOTE: all partial contiguous can safely be replaced by full contiguous. we should be able to match old functionality like this
-  if not (RANGEIFY > 1): return idx.replace(src=(x.replace(arg=None),)+idx.src[1:])
+  if not PARTIAL_CONTIG: return idx.replace(src=(x.replace(arg=None),)+idx.src[1:])
   ranges = []
   new_ranges = []
   passthrough_idx = []
@@ -197,7 +192,8 @@ def map_contiguous(ctx:RangeifyContext, x:UOp):
   ranges = []
   for s in x.shape:
     ranges.append(ctx.new_range(s) if resolve(s!=1) else UOp.const(dtypes.int, 0))
-  return x.src[0].index(*ranges).bufferize(*[x for x in ranges if x.op is not Ops.CONST], arg=x.device).forced_reshape(x.shape)
+  ret = x.src[0].index(*ranges).bufferize(*[x for x in ranges if x.op is not Ops.CONST], arg=x.device)
+  return ret.forced_reshape(x.shape)
 
 def map_reduce(ctx:RangeifyContext, idx:UOp, red:UOp):
   rngs = list(idx.src[1:])
@@ -236,11 +232,8 @@ def index_child(ctx:RangeifyContext, c:UOp, x:UOp, idx:UOp):
     out_rngs = list(idx.src[1:])
     idx_ranges, end_ranges = ctx.seen_child[c]
     for i,nr in zip(idx_ranges, end_ranges): out_rngs[i] = nr
-  # index based on the shared ranges
-  ret = c.index(*out_rngs)
-  # if all ranges aren't the same between children, we have to bufferize
-  if len(idx_ranges) > 0: ret = ret.bufferize(*end_ranges, arg=x.device).index(*[idx.src[1+i] for i in idx_ranges])
-  return ret
+  if len(idx_ranges) == 0: return c.index(*out_rngs)
+  return c.index(*out_rngs).bufferize(*end_ranges, arg=x.device).index(*[idx.src[1+i] for i in idx_ranges])
 
 def children_gate(ctx:RangeifyContext, idx:UOp, c:UOp):
   if len(ctx.seen_children[c]) != c.arg: raise RuntimeError("all children should have been seen by now")
@@ -248,7 +241,6 @@ def children_gate(ctx:RangeifyContext, idx:UOp, c:UOp):
 
 def might_end_axis(idx:UOp):
   if idx.arg is None: return None
-  # TODO: write a proper cost function here
   if all(x.op not in {Ops.BUFFER, Ops.CONTIGUOUS, Ops.BUFFERIZE} for x in idx.toposort()): return None
   if all(x.op not in {Ops.REDUCE_AXIS} for x in idx.toposort()): return None
   to_end_axis = []
@@ -271,7 +263,7 @@ pm_rangeify = pm_mops+PatternMatcher([
   # if we come across this, remove it. it was a CHILD unused in an INDEX
   (UPat(Ops.CHILD, src=(UPat(Ops.CHILDREN, src=(UPat.var("x"),)),)), lambda x: x),
 
-  # CONST (or DEFINE_VAR) can't have axes. remove srcs when we INDEX it
+  # CONST (or DEFINE_VAR) can't have axes. remove srcs when we idx
   (UPat(Ops.INDEX, src=(UPat((Ops.CONST, Ops.DEFINE_VAR), name="c"),)), lambda c: c.replace(src=())),
 
   # handle arg on any op with weight. old endrange stuff
@@ -286,7 +278,6 @@ pm_rangeify = pm_mops+PatternMatcher([
 # 3.5 cleanups
 
 # you don't know in the first pass if axes are going to die, this happens if there's an EXPAND to the left
-# TODO: figure out how to reenable this
 def cleanup_dead_axes(b:UOp):
   parents = b.src[0].toposort()
   new_rng = []
@@ -311,7 +302,7 @@ def remove_bufferize(b2:UOp, idx2:UOp):
   assert all(x.op is Ops.RANGE for x in b2.src[1:])
   return b2.src[0].substitute(dict(zip(b2.src[1:], idx2.src[1:])))
 
-pm_cleanups = double_reshape+pm_mops+PatternMatcher([
+pm_cleanups = pm_mops+PatternMatcher([
   #(UPat(Ops.BUFFERIZE, name="b"), cleanup_dead_axes),
   # remove noop buffers. if we look at the next index we can remove even more of these
   # NOTE: this is mostly the same case as below, but if there's no INDEX this gets more
@@ -370,9 +361,9 @@ def unbind_kernel(ctx:LocalAddBufferContext, b:UOp):
 
 def handle_assign(ctx:LocalAddBufferContext, assign:UOp):
   buf = assign.as_buf()
+  assert buf not in ctx.map
   # HACK to put the buffer in the MAP instead of MSTACK/MSELECT
   if buf.op in {Ops.MSTACK, Ops.MSELECT}: buf = buf.src[0]
-  assert buf not in ctx.map
   ctx.map[buf] = assign
   return buf
 
@@ -390,9 +381,8 @@ to_define_global = PatternMatcher([
   (UPat(Ops.STORE, name="store").f(Ops.INDEX, allow_any_len=True, name="idx").f(Ops.LOAD),
     lambda store,idx: idx.replace(src=(store.as_buf(),)+idx.src[1:]).load(store)),
 
-  # HACK in case any CONSTs were replaced
-  # this is only needed if you are using symbolic
-  #(UPat(Ops.CONST, name="c"), lambda c: c.replace(src=()) if len(c.src) else None),
+  # HACK
+  (UPat(Ops.CONST, name="c"), lambda c: c.replace(src=()) if len(c.src) else None),
 ])
 
 def split_store(x:UOp):
@@ -406,7 +396,7 @@ def split_store(x:UOp):
 
   # NOTE: the hack for COPY is here
   ret = ret.sink(arg=KernelInfo(name=name)) if ret.src[1].op is not Ops.COPY else ret.src[1]
-  kernel = UOp(Ops.KERNEL, src=tuple(ctx.map.values())+tuple(ctx.vars.keys()), arg=Kernel(ret,()))
+  kernel = UOp(Ops.KERNEL, src=tuple(ctx.map.values())+tuple(ctx.vars.keys()), arg=Kernel(ret, ()))
   return x.as_buf().assign(kernel)
 
 split_kernels = PatternMatcher([

tinygrad/uop/ops.py

@@ -6,7 +6,7 @@ from enum import Enum, auto
 from tinygrad.uop import Ops, GroupOp
 from tinygrad.uop.mathtraits import MathTrait
 from tinygrad.dtype import ConstType, ImageDType, dtypes, DType, truncate, PtrDType
-from tinygrad.helpers import ContextVar, all_int, prod, getenv, all_same, Context, partition, temp, unwrap, T, argfix, Metadata, flatten, TRACEMETA
+from tinygrad.helpers import ContextVar, all_int, prod, getenv, all_same, Context, partition, temp, unwrap, T, argfix, Metadata, flatten
 from tinygrad.helpers import PICKLE_BUFFERS, PROFILE, dedup, cdiv, cmod, diskcache_put, to_function_name, cpu_profile, TracingKey
 if TYPE_CHECKING:
   from tinygrad.shape.shapetracker import ShapeTracker
@@ -795,7 +795,7 @@ def track_uop(u:UOp):
   uop_number[u] = num = next(ucount)
   # KERNEL also has a UOp in the arg
   arg = type(u.arg)(track_uop(u.arg.ast), u.arg.metadata) if u.op is Ops.KERNEL else u.arg
-  uop_fields[num] = (u.op, u.dtype, tuple(track_uop(s) for s in u.src), arg, u.tag)+((u.metadata,) if TRACEMETA>=2 else ())
+  uop_fields[num] = (u.op, u.dtype, tuple(track_uop(s) for s in u.src), arg, u.tag)
   return num
 
 # *** tracking pattern matcher ***

tinygrad/viz/serve.py

@@ -91,9 +91,9 @@ def uop_to_json(x:UOp) -> dict[int, dict]:
 
 @functools.cache
 def _reconstruct(a:int):
-  op, dtype, src, arg, *rest = contexts[2][a]
+  op, dtype, src, arg, tag = contexts[2][a]
   arg = type(arg)(_reconstruct(arg.ast), arg.metadata) if op is Ops.KERNEL else arg
-  return UOp(op, dtype, tuple(_reconstruct(s) for s in src), arg, *rest)
+  return UOp(op, dtype, tuple(_reconstruct(s) for s in src), arg, tag)
 
 def get_details(ctx:TrackedGraphRewrite) -> Generator[GraphRewriteDetails, None, None]:
   yield {"graph":uop_to_json(next_sink:=_reconstruct(ctx.sink)), "uop":str(next_sink), "changed_nodes":None, "diff":None, "upat":None}