Merge branch 'master' into simple_rangeify

.github/workflows/test.yml

@@ -521,7 +521,10 @@ jobs:
       - name: Run HCQ2 tests
         run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/test_tiny.py
       - name: Run HCQ2 multi-device tests
-        run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python -m pytest -n=auto test/backend/test_multitensor.py
+        run: |
+          HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/unit/test_multitensor.py \
+            TestMultiTensor.test_simple_add TestMultiTensor.test_shard_reduce \
+            TestMultiTensor.test_backward_sum TestMultiTensor.test_matmul_shard_0_0
       - name: Run HCQ2 JIT tests
         run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/unit/test_jit.py
       - name: Run HCQ2 unit tests

extra/hcq2/ops_amd2.py

@@ -297,8 +297,6 @@ class AMDAllocator(HCQAllocator['AMDDevice']):
 
   def _do_map(self, buf:HCQ2Buffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
 
-  def _do_unmap(self, buf:HCQ2Buffer): self.dev.iface.unmap(buf)
-
 @dataclass
 class AMDQueueDesc:
   ring: Buffer; read_ptr: Buffer; write_ptr: Buffer; doorbell: Buffer; put_value: Buffer  # noqa: E702
@@ -390,24 +388,15 @@ class KFDIface:
     return hcqbuf
 
   def free(self, mem):
-    self._unmap(mem)
-    if mem.va_addr: FileIOInterface.munmap(mem.va_addr, mem.size)
-    kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle)
-
-  def unmap(self, mem):
-    self._unmap(mem)
-    if getattr(mem, '_owns_kfd_handle', False): kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle)
-
-  def _unmap(self, mem):
     gpus = (ctypes.c_int32 * 1)(self.gpu_id)
     stm = kfd.AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(gpus), n_devices=1)
     assert stm.n_success == 1
+    if mem.owner == self.dev:
+      if mem.va_addr: FileIOInterface.munmap(mem.va_addr, mem.size)
+      kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle)
 
   def map(self, mem):
-    if mem.owner is not None and mem.owner._is_cpu():
-      mapped = self.alloc(mem.size, host=True, cpu_addr=mem.va_addr)
-      mapped._owns_kfd_handle = True
-      return mapped
+    if mem.owner is not None and mem.owner._is_cpu(): return self.alloc(mem.size, host=True, cpu_addr=mem.va_addr)
 
     c_gpus = (ctypes.c_int32 * 1)(self.gpu_id)
     stm = kfd.AMDKFD_IOC_MAP_MEMORY_TO_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(c_gpus), n_devices=1)
@@ -479,7 +468,6 @@ class PCIIface(PCIIfaceBase):
 
   def require_profile_mode(self): return True
   def is_wgp_active(self, xcc, se, sa, wgp) -> bool: return True # TODO: account for WGP disablement on some asics.
-  def unmap(self, mem): self.free(mem)
 
   def _compute_props(self):
     self.ip_versions = self.dev_impl.ip_ver

tinygrad/codegen/opt/tc.py

@@ -1,7 +1,6 @@
 import math, functools
 from dataclasses import dataclass
 from tinygrad.dtype import DType, dtypes
-from tinygrad.uop.ops import PatternMatcher, UOp, UPat, Ops
 
 @dataclass(frozen=True)
 class TensorCore: # D = A * B + C, A is (M x K), B is (K x N), C and D are (M x N)
@@ -136,42 +135,6 @@ amd_cdna4 = amd_cdna_1616128 + amd_cdna_161632 + amd_cdna_161616
 
 def get_amd(arch): return {"gfx942": amd_cdna3, "gfx950": amd_cdna4, "gfx1200": amd_rdna4, "gfx1201": amd_rdna4}.get(arch, amd_rdna3)
 
-pm_validate_wmma_rdna3 = PatternMatcher([
-  (UPat(Ops.WMMA, name="x", dtype=dtypes.int32), lambda x: x.replace(
-    src=(x.src[0].bitcast(dtypes.uint32), x.src[1].bitcast(dtypes.uint32), x.src[2]))
-    if x.src[0].dtype == dtypes.int8 and x.src[0].max_numel() == 16 else None),
-  (UPat(Ops.WMMA, name="x", dtype=dtypes.half), lambda x: UOp(Ops.STACK, src=tuple(x.replace(
-      src=(x.src[0], x.src[1], UOp(Ops.STACK, src=tuple(x.src[2].index(UOp.const(j//2, dtypes.int16))
-      if j%2 == 0 else UOp.const(0.0, x.src[2].dtype)
-      for j in range(x.max_numel()*2)))),
-      arg=(*x.arg[:4], None)).index(UOp.const(i*2, dtypes.int16))
-      for i in range(x.max_numel()))) if x.max_numel() == 8 else None),
-  (UPat(Ops.WMMA, name="x"), lambda x: x.replace(
-    src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
-    if x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 16 else None),
-])
-
-pm_validate_wmma_rdna4 = PatternMatcher([
-  (UPat(Ops.WMMA, name="x", dtype=dtypes.bfloat16), lambda x: x.replace(
-    dtype=dtypes.uint16,
-    src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2].bitcast(dtypes.uint16)))
-      .bitcast(dtypes.bfloat16) if x.max_numel() == 8 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 8 else None),
-  (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
-    lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
-    if x.max_numel() == 8 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 8 else None)
-])
-
-pm_validate_wmma_cdna = PatternMatcher([
-  (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
-    lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint32), x.src[1].bitcast(dtypes.uint32), x.src[2]))
-    if x.arg[0][2] == 128 and x.src[0].dtype.itemsize <= 8 else None),
-  (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
-    lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
-    if x.max_numel() == 4 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 4 else None),
-  (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
-    lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint64), x.src[1].bitcast(dtypes.uint64), x.src[2]))
-    if x.max_numel() == 4 and x.src[0].dtype in dtypes.fp8_ocp and x.src[0].max_numel() == 8 else None),
-])
 # ***** Apple Metal *****
 
 metal = [TensorCore(dims=(8,8,8), threads=32, elements_per_thread=(2,2,2), dtype_in=di, dtype_out=do,

tinygrad/engine/realize.py

@@ -140,7 +140,7 @@ class ExecContext:
   cache: bool = True
 
 def _resolve(b:UOp, inputs:tuple[UOp, ...]) -> UOp:
-  if b.op in (Ops.SLICE, Ops.MSELECT, Ops.SHRINK) and b.src[0].op is Ops.PARAM: return b.replace(src=(inputs[b.src[0].arg.slot], *b.src[1:]))
+  if b.op in (Ops.SLICE, Ops.MSELECT) and b.src[0].op is Ops.PARAM: return b.replace(src=(inputs[b.src[0].arg.slot], *b.src[1:]))
   if b.op is Ops.MSTACK: return b.replace(src=tuple(_resolve(x, inputs) for x in b.src))
   return inputs[b.arg.slot] if b.op is Ops.PARAM else b
 def resolve_params(call:UOp, inputs:tuple[UOp, ...]) -> list[UOp]: return [_resolve(b, inputs) for b in get_call_arg_uops(call)]

tinygrad/renderer/llvmir.py

@@ -279,9 +279,43 @@ exit: %packed = phi i32 [%packed_bf8, %do_bf8], [%packed_fp8, %do_fp8]\n  %trunc
       (UPat(Ops.WMMA, name="wmma"), lambda ctx, wmma, rdna4=AMDLLVMRenderer.is_rdna4(target.arch), cdna=self.is_cdna:
         render_wmma_amd(ctx, wmma, cdna, rdna4))
     ])
-    if self.is_cdna: self.extra_matcher += tc.pm_validate_wmma_cdna
-    if target.arch in {"gfx1100", "gfx1151"}: self.extra_matcher += tc.pm_validate_wmma_rdna3
-    if target.arch in {"gfx1200", "gfx1201"}: self.extra_matcher += tc.pm_validate_wmma_rdna4
+    if self.is_cdna:
+      self.extra_matcher += PatternMatcher([
+        (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
+          lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint32), x.src[1].bitcast(dtypes.uint32), x.src[2]))
+          if x.arg[0][2] == 128 and x.src[0].dtype.itemsize <= 8 else None),
+        (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
+          lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
+          if x.max_numel() == 4 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 4 else None),
+        (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
+          lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint64), x.src[1].bitcast(dtypes.uint64), x.src[2]))
+          if x.max_numel() == 4 and x.src[0].dtype in dtypes.fp8_ocp and x.src[0].max_numel() == 8 else None),
+      ])
+    if target.arch in {"gfx1100", "gfx1151"}:
+      self.extra_matcher += PatternMatcher([
+        (UPat(Ops.WMMA, name="x", dtype=dtypes.int32), lambda x: x.replace(
+          src=(x.src[0].bitcast(dtypes.uint32), x.src[1].bitcast(dtypes.uint32), x.src[2]))
+          if x.src[0].dtype == dtypes.int8 and x.src[0].max_numel() == 16 else None),
+        (UPat(Ops.WMMA, name="x", dtype=dtypes.half), lambda x: UOp(Ops.STACK, src=tuple(x.replace(
+          src=(x.src[0], x.src[1], UOp(Ops.STACK, src=tuple(x.src[2].index(UOp.const(j//2, dtypes.int16))
+            if j%2 == 0 else UOp.const(0.0, x.src[2].dtype)
+            for j in range(x.max_numel()*2)))),
+          arg=(*x.arg[:4], None)).index(UOp.const(i*2, dtypes.int16))
+          for i in range(x.max_numel()))) if x.max_numel() == 8 else None),
+        (UPat(Ops.WMMA, name="x"), lambda x: x.replace(
+          src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
+          if x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 16 else None),
+      ])
+    if target.arch in {"gfx1200", "gfx1201"}:
+      self.extra_matcher += PatternMatcher([
+        (UPat(Ops.WMMA, name="x", dtype=dtypes.bfloat16), lambda x: x.replace(
+          dtype=dtypes.uint16,
+          src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2].bitcast(dtypes.uint16)))
+            .bitcast(dtypes.bfloat16) if x.max_numel() == 8 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 8 else None),
+        (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
+          lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
+          if x.max_numel() == 8 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 8 else None)
+      ])
 
   def supported_dtypes(self): return {d for d in super().supported_dtypes()
                                       if (d not in dtypes.fp8_ocp or self.target.arch == "gfx950") and d not in dtypes.fp8_fnuz}

tinygrad/runtime/graph/metal.py

@@ -113,5 +113,5 @@ class MetalGraph(GraphRunner):
   @staticmethod
   def supports_uop(batch_devs, new_call:UOp) -> bool:
     # Metal ICB replay encodes offsets as uint32; reject if any Metal buffer offset exceeds 32-bit range.
-    if any(b.op in {Ops.SLICE, Ops.SHRINK} and b.src[1].val * b.src[0].dtype.itemsize > 0xFFFFFFFF for b in new_call.src[1:]): return False
+    if any(b.op is Ops.SLICE and b.src[1].val * b.src[0].dtype.itemsize > 0xFFFFFFFF for b in new_call.src[1:]): return False
     return GraphRunner.supports_uop(batch_devs, new_call)

tinygrad/runtime/support/hcq2.py

@@ -141,14 +141,14 @@ def _build_wait_cmds(slots:dict[str, int], dep_lanes:list[tuple[tuple, int, int]
 
   # opt2: keep latest dep per (dep device, queue, cur lane)
   latest = {((dep[0][dlane], dep[1]), lane): (dep, dlane) for dep, dlane, lane in sorted(dep_lanes, key=lambda x: x[0][2])}
-  deps:dict[tuple, dict[int, list[int]]] = collections.defaultdict(lambda: collections.defaultdict(list))
-  for (_, lane), (dep, dlane) in latest.items(): deps[dep][lane].append(dlane)
+  deps:dict[tuple, list[int|None]] = collections.defaultdict(lambda: [None]*len(devices))
+  for (_, lane), (dep, dlane) in latest.items(): deps[dep][lane] = dlane
 
   waits = []
-  for (ddevs, dqueue, dtag), by_lane in deps.items():
-    for ls in itertools.zip_longest(*(by_lane[lane] for lane in range(len(devices)))):
-      s = UOp.mstack(*[make_signal(d, tag="sentinel_signal") if dl is None else make_signal(ddevs[dl], slots[dqueue]) for dl, d in zip(ls, devices)])
-      waits.append(UOp(Ops.INS, arg="wait", src=(s, UOp.const(dtag + 1, dtypes.uint64))))
+  for (ddevs, dqueue, dtag), lanes in deps.items():
+    sig = UOp.mstack(*[make_signal(d, tag="sentinel_signal") if dl is None else make_signal(ddevs[dl], slots[dqueue])
+                       for dl, d in zip(lanes, devices)])
+    waits.append(UOp(Ops.INS, arg="wait", src=(sig, UOp.const(dtag + 1, dtypes.uint64))))
   return waits, {dtag for _, _, dtag in deps}
 
 def _build_finalizers(batch:list[tuple[UOp, tuple[str, ...]]], batch_info:list[tuple[tuple[str, ...], str]],
@@ -383,7 +383,7 @@ def resolve_getaddr_slice(bv:UOp, g:UOp) -> UOp:
   return UOp(Ops.GETADDR, src=(base,), arg=g.arg) + UOp.const(bv.src[1].val * itemsize, dtypes.uint64)
 
 pm_early_simplify = PatternMatcher([
-  (UPat(Ops.GETADDR, src=(UPat.any(sl:=UPat((Ops.SLICE, Ops.SHRINK), name="bv"), sl.after(allow_any_len=True)),), name="g"), resolve_getaddr_slice),
+  (UPat(Ops.GETADDR, src=(UPat.any(sl:=UPat(Ops.SLICE, name="bv"), sl.after(allow_any_len=True)),), name="g"), resolve_getaddr_slice),
   (UPat(Ops.INDEX, src=(UPat(Ops.SLICE, name="bv"),), allow_any_len=True, name="x"),
    lambda bv,x: x.replace(src=(bv.src[0], x.src[1] + bv.src[1].cast(x.src[1].dtype), *x.src[2:]))),
 ])
@@ -623,8 +623,6 @@ class HCQAllocator(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):
     if not hasattr(self, '_do_map'): raise NotImplementedError("map failed: no method implemented")
     return self._do_map(buf)
 
-  def _do_unmap(self, mb): self.dev.iface.free(mb)
-
   @suppress_finalizing
   def _free(self, buf:HCQ2Buffer, options:BufferSpec|None=None):
     if options is not None and options.external_ptr is not None: return
@@ -633,6 +631,6 @@ class HCQAllocator(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):
 
   def _unmap(self, mb):
     self.dev.synchronize()
-    self._do_unmap(mb)
+    self.dev.iface.free(mb)
 
   def _offset(self, buf, size:int, offset:int) -> HCQ2Buffer: return buf.offset(offset=offset, size=size)

tinygrad/schedule/__init__.py

@@ -81,8 +81,7 @@ def create_schedule(sched_sink:UOp) -> UOp:
 
 from tinygrad.schedule.memory import memory_plan_rewrite
 from tinygrad.engine.realize import capturing, pm_flatten_linear
-#from tinygrad.schedule.rangeify import get_kernel_graph
-from tinygrad.schedule.rangeify2 import get_kernel_graph
+from tinygrad.schedule.rangeify import get_kernel_graph
 from tinygrad.helpers import CAPTURING
 from tinygrad.uop.ops import PatternMatcher, UPat, ParamArg
 from tinygrad.dtype import AddrSpace

tinygrad/schedule/rangeify.py

@@ -6,11 +6,11 @@ from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, K
 from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, rewrite_group, identity_element
 from tinygrad.uop.symbolic import symbolic, pm_fold_cast_const
 from tinygrad.uop.movement import mop_cleanup
-from tinygrad.helpers import prod, getenv, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY, VIZ, MAX_KERNEL_BUFFERS, SPEC
+from tinygrad.helpers import prod, getenv, dedup, all_int, DEBUG, SPLIT_REDUCEOP, VIZ, MAX_KERNEL_BUFFERS, SPEC
 from tinygrad.helpers import PCONTIG, FLOAT16, OPENPILOT_HACKS, argsort, partition, get_single_element
 from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_simplify
 from tinygrad.codegen.opt import Opt
-from tinygrad.schedule.indexing import run_rangeify, BufferizeOpts, IndexingContext, apply_movement_op
+from tinygrad.schedule.indexing import BufferizeOpts, IndexingContext, apply_movement_op
 from tinygrad.schedule.multi import multi_pm
 from tinygrad.schedule.allreduce import create_allreduce_function
 
@@ -39,7 +39,16 @@ pm_fold_moved_after = PatternMatcher([
 def _mop_index(r:UOp, idx:UOp):
   idxs = idx.src[1:]
   if len(idxs) == len(r.shape):
-    return r.src[0].index(*apply_movement_op(r.op, r.src[0].shape, r.marg, idxs), dtype=idx.dtype, arg=idx.arg)
+    ret = r.src[0].index(*apply_movement_op(r.op, r.src[0].shape, r.marg, idxs), dtype=idx.dtype, arg=idx.arg)
+    if r.op is Ops.PAD:
+      # insert 0 for PAD with where
+      # TODO: does this need simplify
+      a = UOp.const(True)
+      for s in ret.src[1:]:
+        if s.op is Ops.WHERE and s.src[2].op is Ops.CONST and s.src[2].arg == Invalid:
+          a = a & s.src[0]
+      ret = a.where(ret, ret.const_like(0))
+    return ret
   if r.op is Ops.RESHAPE:
     src_prefix = len(r.src[0].shape) - len(r.shape[len(idxs):])
     if src_prefix >= 0 and r.src[0].shape[src_prefix:] == r.shape[len(idxs):]:
@@ -567,9 +576,101 @@ def convert_copy_to_store(ctx, copy:UOp, existing_buf:UOp|None=None):
   # reshape back to input
   return buf.after(buf.store(input_src)).reshape(copy.shape)
 
+def convert_contig_to_store(ctx, copy:UOp):
+  input_src = copy.src[0]
+  # create the output buffer
+  buf = UOp(Ops.BUFFER, src=(shape_to_shape_arg(input_src.max_shape),), arg=ParamArg(next(ctx), copy.dtype, device=copy.device))
+  # reshape back to input
+  view = buf.shrink_to(input_src.shape)
+  return view.after(view.store(input_src))
+
 pm_copy_to_store = PatternMatcher([
   (UPat(name="existing_buf").store(UPat(Ops.COPY, name="copy")), convert_copy_to_store),
   (UPat(Ops.COPY, name="copy"), convert_copy_to_store),
+  (UPat(Ops.CONTIGUOUS, name="copy"), convert_contig_to_store),
+])
+
+# **** simple rangeify ****
+
+from tinygrad.helpers import all_same
+from tinygrad.uop.ops import _broadcast_shape
+
+def expand_broadcast(x:UOp):
+  shapes = [u._shape for u in x.src]
+  if any(s is None for s in shapes) or all_same(shapes): return None
+  shape = _broadcast_shape(*shapes)
+  return x.replace(src=tuple([u.expand(shape) for u in x.src]))
+
+pm_expand_broadcast = PatternMatcher([
+  # expand broadcasts first
+  (UPat(GroupOp.Binary|GroupOp.Ternary|{Ops.STORE}, name="x"), expand_broadcast),
+])
+
+def expand_coeff(sink:UOp) -> dict[UOp,int]:
+  coeff: dict[UOp,int] = {sink: 1}
+  contig: dict[UOp,int] = {}
+  for u in reversed(list(sink.toposort())):
+    c = 1 if u.op is Ops.STORE else coeff.get(u, 0)
+    # symbolic coeffs mark on vmax, an extra CONTIGUOUS is always safe
+    if (c > 1 if isinstance(c, int) else c.vmax > 1) and u.op in (GroupOp.Elementwise | {Ops.REDUCE}) and u.device is not None:
+      contig[u] = c
+      c = 1
+    coeff[u] = c
+    mult = prod(u.shape) // prod(u.src[0].shape) if u.op is Ops.EXPAND else 1
+    for s in u.src: coeff[s] = coeff.get(s, 0) + c * (mult if s is u.src[0] else 1)
+  return contig
+
+def rangeify_on_reduce(ctx, inp:UOp, red:UOp, idx:UOp|None=None):
+  if red.arg[1] == 0: return None
+  if idx is None and len(red.shape) > 0: return None
+  # TODO: is AxisType.REDUCE a real thing?
+  rngs = [UOp.range(s, next(ctx), AxisType.REDUCE) for s in inp.shape[:red.arg[1]]]
+  return inp.index(*rngs, *(idx.src[1:] if idx is not None else ())).reduce(*rngs, arg=(red.arg[0], 0))
+
+def rangeify_on_store(ctx, x:UOp):
+  if x.shape == (): return None
+  rngs = [UOp.range(s, next(ctx)) for s in x.shape]
+  return x.src[0].index(*rngs).store(x.src[1].index(*rngs)).end(*rngs)
+
+def rangeify_on_stage(ctx, x:UOp):
+  if x.src[0].shape == (): return None
+  # size 1 dims don't get ranges, they are reshaped out and back in
+  if all_int(x.shape) and 0 < len(sq := tuple(s for s in x.shape if s != 1)) < len(x.shape):
+    return rangeify_on_stage(ctx, x.src[0].reshape(sq).bufferize(arg=x.arg)).reshape(x.shape)
+  rngs = [UOp.range(s, next(ctx)) for s in x.shape]
+  return x.replace(src=(x.src[0].index(*rngs), *rngs))
+
+def index_on_stack(stack:UOp, idx:UOp):
+  srcs = [s.index(*idx.src[2:]) for s in stack.src]
+  r0 = idx.src[1]
+  ret = srcs[-1]
+  for k in range(len(srcs)-2, -1, -1): ret = r0.eq(k).where(srcs[k], ret)
+  return ret
+
+pm_simple_rangeify = PatternMatcher([
+  # INDEX without src is nothing
+  (UPat(Ops.INDEX, src=(UPat.var('x'),)), lambda x: x),
+  # STAGE on shape () is nothing
+  (UPat(Ops.STAGE, src=(UPat.var('x'),)), lambda x: x if x.shape == () else None),
+  # if INDEX is on STAGE with the same ranges, remove the pair
+  (UPat(Ops.STAGE, allow_any_len=True, name="s").index(allow_any_len=True, name="i"),
+   lambda s,i: s.src[0] if s.src[1:] == i.src[1:] else None),
+  # reshape of a single element shaped value to scalar is an index
+  (UPat(Ops.RESHAPE, name="x"), lambda x: x.src[0].index(0) if x.marg == () and x.src[0].shape == (1,) else None),
+  # handle movement ops on INDEX
+  (UPat(GroupOp.Movement, name="r").index(name="idx", allow_any_len=True), _mop_index),
+  (UPat(Ops.STACK, name="stack").index(name="idx", allow_any_len=True), index_on_stack),
+  # pass index through elementwise
+  (UPat(GroupOp.Elementwise, name="b").index(name="idx", allow_any_len=True),
+   lambda b,idx: b.replace(src=tuple(s.index(*idx.src[1:]) for s in b.src))),
+])
+
+pm_range_creation = PatternMatcher([
+  # reduce/store are what creates ranges
+  (UPat(Ops.REDUCE, src=(UPat.var('inp'),), name="red").index(name="idx", allow_any_len=True), rangeify_on_reduce),
+  (UPat(Ops.REDUCE, src=(UPat.var('inp'),), name="red"), rangeify_on_reduce),
+  (UPat(Ops.STORE, name="x"), rangeify_on_store),
+  (UPat(Ops.STAGE, name="x"), rangeify_on_stage),
 ])
 
 @rewrite_group(new_ctx=False)
@@ -578,15 +679,61 @@ def get_kernel_graph(sink:UOp) -> UOp:
   if OPENPILOT_HACKS: tsink = graph_rewrite(tsink, pm_fold_moved_after, ctx={}, name="fold moved afters")
   tsink = graph_rewrite(tsink, pm_mops+earliest_rewrites, bottom_up=True, name="earliest rewrites")
 
+  # convert movement ops to ranges
+  #tsink, rctx = run_rangeify(tsink, bool(DEBUG_RANGEIFY))
+  tsink = graph_rewrite(tsink, pm_expand_broadcast, bottom_up=True, name="expand broadcast")
+
+  # mark ops that would be recomputed (expand coeff > 1) as CONTIGUOUS, like the realize map in run_rangeify
+  contig = expand_coeff(tsink)
+  subs: dict[UOp, UOp] = {}
+  for u in tsink.toposort():
+    u2 = u.replace(src=tuple(subs.get(s, s) for s in u.src))
+    subs[u] = u2.alu(Ops.STAGE, arg=BufferizeOpts(u2.device)) if u in contig else u2
+  tsink = subs[tsink]
+
+  # add buffers on copy
   tsink = graph_rewrite(tsink, pm_copy_to_store, ctx=itertools.count(0), bottom_up=True, name="convert copy to store")
 
-  # convert movement ops to ranges
-  tsink, rctx = run_rangeify(tsink, bool(DEBUG_RANGEIFY))
+  # simple rangeify
+  tsink = graph_rewrite(tsink, pm_range_creation+pm_simple_rangeify, ctx=itertools.count(0), bottom_up=True, name="simple rangeify")
 
-  tsink = graph_rewrite(tsink,
-                        symbolic+pm_fold_cast_const+pm_reduce_simplify+pm_const_buffer_folding+pm_remove_bufferize+pm_no_indexing_calls,
-                        name="symbolic+reduce_collapse+debuf")
-  tsink = graph_rewrite(tsink, pm_limit_bufs, ctx=rctx, name="limit buffers")
+  # for each index on a stage without children, try to merge the stage into the consumer kernel
+  while 1:
+    indexes: dict[UOp, list[UOp]] = {}
+    consumers: dict[UOp, list[UOp]] = {}
+    for u in tsink.toposort():
+      if u.op is Ops.INDEX and u.src[0].op is Ops.STAGE:
+        indexes.setdefault(u.src[0], []).append(u)
+      for s in u.src: consumers.setdefault(s, []).append(u)
+    # the ranges wrapping u: REDUCE ranges on the path up, plus the enclosing END/STAGE nest
+    def nest_ranges(u:UOp) -> set[UOp]:
+      ret: set[UOp] = set()
+      stack, seen = [u], set()
+      while len(stack):
+        if (x := stack.pop()) in seen: continue
+        seen.add(x)
+        if x.op is Ops.REDUCE: ret.update(*[er.ranges for er in x.ended_ranges])
+        elif x.op in {Ops.END, Ops.STAGE}:
+          ret.update(*[er.ranges for er in x.ended_ranges])
+          continue
+        stack.extend(consumers.get(x, []))
+      return ret
+    subs = {}
+    for k,v in indexes.items():
+      # don't move REDUCE ranges up (real?)
+      if len(v) != 1 or not all(all([r.arg[-1] == AxisType.WEAK for r in s.ranges]) for s in v[0].src[1:]): continue
+      # merging must not add iteration multiplicity around range-bound computation in the stage body:
+      # every range of the enclosing kernel nest must be used by the index, unless the body has no inner loops
+      if not nest_ranges(v[0]) <= set().union(*[s.ranges for s in v[0].src[1:]]) and \
+         any(x.op is Ops.REDUCE for x in k.src[0].toposort(gate=lambda x: x.op is not Ops.STAGE)): continue
+      for old_r, new_r in zip(k.src[1:], v[0].src[1:]):
+        subs[old_r] = new_r
+    if not len(subs): break
+    tsink = tsink.substitute(subs)
+    tsink = graph_rewrite(tsink, pm_simple_rangeify, bottom_up=True, name=f"merge kernels ({len(subs)})")
+
+  tsink = graph_rewrite(tsink, symbolic+pm_reduce_simplify+pm_const_buffer_folding+pm_remove_bufferize+pm_no_indexing_calls, name="symbolic+reduce_collapse+debuf")
+  #tsink = graph_rewrite(tsink, pm_limit_bufs, ctx=rctx, name="limit buffers")
 
   if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Rangeify")
 

tinygrad/schedule/rangeify2.py

@@ -1,238 +0,0 @@
-from dataclasses import dataclass, field, replace
-from typing import cast
-import itertools
-from tinygrad.dtype import dtypes, AddrSpace, Invalid, to_dtype, strong_dtype
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, KernelInfo, ParamArg, shape_to_shape_arg
-from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, rewrite_group, identity_element, remove_all_tags
-from tinygrad.uop.symbolic import symbolic, pm_fold_cast_const
-from tinygrad.uop.movement import mop_cleanup
-from tinygrad.helpers import prod, getenv, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY, VIZ, MAX_KERNEL_BUFFERS, SPEC
-from tinygrad.helpers import PCONTIG, FLOAT16, OPENPILOT_HACKS, argsort, partition, get_single_element, Context
-from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_simplify
-from tinygrad.codegen.opt import Opt
-from tinygrad.schedule.indexing import run_rangeify, BufferizeOpts, IndexingContext, apply_movement_op
-from tinygrad.schedule.multi import multi_pm
-from tinygrad.schedule.allreduce import create_allreduce_function
-
-# *** preparation ***
-
-from tinygrad.helpers import all_same
-from tinygrad.uop.ops import _broadcast_shape
-
-def expand_broadcast(x:UOp):
-  shapes = [u._shape for u in x.src]
-  if any(s is None for s in shapes) or all_same(shapes): return None
-  shape = _broadcast_shape(*shapes)
-  return x.replace(src=tuple([u.expand(shape) for u in x.src]))
-
-pm_expand_broadcast = PatternMatcher([
-  # expand broadcasts first
-  (UPat(GroupOp.Binary|GroupOp.Ternary|{Ops.STORE}, name="x"), expand_broadcast),
-])
-
-def convert_copy_to_store(ctx, copy:UOp, existing_buf:UOp|None=None):
-  input_src = copy.src[0]
-  if not input_src.has_buffer_identity(after_ok=True): input_src = input_src.contiguous()
-  input_src = input_src.flatten()
-  if existing_buf is not None:
-    # if the existing buffer is not a full buffer, we can't use it
-    if not existing_buf.has_buffer_identity(after_ok=True): return None
-    # if there's already a buffer, we just use it
-    return existing_buf.flatten().store(input_src)
-  # create the output buffer
-  buf = UOp(Ops.BUFFER, src=(shape_to_shape_arg(input_src.max_shape),), arg=ParamArg(next(ctx), copy.dtype, device=copy.device))
-  # reshape back to input
-  return buf.after(buf.store(input_src)).reshape(copy.shape)
-
-def convert_contig_to_store(ctx, copy:UOp):
-  input_src = copy.src[0]
-  # create the output buffer
-  buf = UOp(Ops.BUFFER, src=(shape_to_shape_arg(input_src.max_shape),), arg=ParamArg(next(ctx), copy.dtype, device=copy.device))
-  # reshape back to input
-  view = buf.shrink_to(input_src.shape)
-  return view.after(view.store(input_src))
-
-pm_copy_to_store = PatternMatcher([
-  (UPat(name="existing_buf").store(UPat(Ops.COPY, name="copy")), convert_copy_to_store),
-  (UPat(Ops.COPY, name="copy"), convert_copy_to_store),
-  (UPat(Ops.CONTIGUOUS, name="copy"), convert_contig_to_store),
-])
-
-# *** RANGE creation ***
-
-def rangeify_on_reduce(ctx, inp:UOp, red:UOp, idx:UOp|None=None):
-  if red.arg[1] == 0: return None
-  if idx is None and len(red.shape) > 0: return None
-  # TODO: is AxisType.REDUCE a real thing?
-  rngs = [UOp.range(s, next(ctx), AxisType.REDUCE) for s in inp.shape[:red.arg[1]]]
-  return inp.index(*rngs, *(idx.src[1:] if idx is not None else ())).reduce(*rngs, arg=(red.arg[0], 0))
-
-def rangeify_on_store(ctx, x:UOp):
-  if x.shape == (): return None
-  rngs = [UOp.range(s, next(ctx)) for s in x.shape]
-  return x.src[0].index(*rngs).store(x.src[1].index(*rngs)).end(*rngs)
-
-def rangeify_on_stage(ctx, x:UOp):
-  if x.src[0].shape == (): return None
-  # size 1 dims don't get ranges, they are reshaped out and back in
-  if all_int(x.shape) and 0 < len(sq := tuple(s for s in x.shape if s != 1)) < len(x.shape):
-    return rangeify_on_stage(ctx, x.src[0].reshape(sq).bufferize(arg=x.arg)).reshape(x.shape)
-  rngs = [UOp.range(s, next(ctx)) for s in x.shape]
-  return x.replace(src=(x.src[0].index(*rngs), *rngs))
-
-pm_range_creation = PatternMatcher([
-  # reduce/store are what creates ranges
-  (UPat(Ops.REDUCE, src=(UPat.var('inp'),), name="red").index(name="idx", allow_any_len=True), rangeify_on_reduce),
-  (UPat(Ops.REDUCE, src=(UPat.var('inp'),), name="red"), rangeify_on_reduce),
-  (UPat(Ops.STORE, name="x"), rangeify_on_store),
-  (UPat(Ops.STAGE, name="x"), rangeify_on_stage),
-])
-
-# *** RANGE migration ***
-
-# movement op on INDEX as a PatternMatcher
-def _mop_index(r:UOp, idx:UOp):
-  idxs = idx.src[1:]
-  if len(idxs) == len(r.shape):
-    ret = r.src[0].index(*apply_movement_op(r.op, r.src[0].shape, r.marg, idxs), dtype=idx.dtype, arg=idx.arg)
-    if r.op is Ops.PAD:
-      # insert 0 for PAD with where
-      # TODO: does this need simplify to ensure the Invalids are at the base?
-      a = UOp.const(True)
-      for s in ret.src[1:]:
-        if s.op is Ops.WHERE and s.src[2].op is Ops.CONST and s.src[2].arg == Invalid: a = a & s.src[0]
-      ret = a.where(ret, ret.const_like(0))
-    return ret
-  if r.op is Ops.RESHAPE:
-    src_prefix = len(r.src[0].shape) - len(r.shape[len(idxs):])
-    if src_prefix >= 0 and r.src[0].shape[src_prefix:] == r.shape[len(idxs):]:
-      if src_prefix == 0: return r.src[0] if r.src[0].dtype == idx.dtype else None
-      ret = r.src[0].index(*apply_movement_op(r.op, r.src[0].shape[:src_prefix], r.shape[:len(idxs)], idxs), dtype=idx.dtype, arg=idx.arg)
-      return ret if ret.shape == idx.shape else None
-
-# TODO: this should be in _mop_index
-def index_on_stack(stack:UOp, idx:UOp):
-  srcs = [s.index(*idx.src[2:]) for s in stack.src]
-  r0 = idx.src[1]
-  ret = srcs[-1]
-  for k in range(len(srcs)-2, -1, -1): ret = r0.eq(k).where(srcs[k], ret)
-  return ret
-
-def walk_mop(u:UOp):
-  if u.op in GroupOp.Movement or u.op is Ops.INDEX: return u.src[0]
-  assert u.op == Ops.AFTER
-  return u
-
-pm_range_migration = PatternMatcher([
-  # INDEX without src is nothing
-  (UPat(Ops.INDEX, src=(UPat.var('x'),)), lambda x: x),
-  # STAGE on shape () is nothing
-  (UPat(Ops.STAGE, src=(UPat.var('x'),)), lambda x: x if x.shape == () else None),
-  # if INDEX is on STAGE with the same ranges, remove the pair
-  (UPat(Ops.STAGE, allow_any_len=True, name="s").index(allow_any_len=True, name="i"),
-   lambda s,i: s.src[0] if s.src[1:] == i.src[1:] else None),
-  # reshape of a single element shaped value to scalar is an index
-  (UPat(Ops.RESHAPE, name="x"), lambda x: x.src[0].index(0) if x.marg == () and x.src[0].shape == (1,) else None),
-  # handle movement ops on INDEX
-  (UPat(GroupOp.Movement, name="r").index(name="idx", allow_any_len=True), _mop_index),
-  (UPat(Ops.STACK, name="stack").index(name="idx", allow_any_len=True), index_on_stack),
-  # move movement ops and INDEX after AFTER
-  (UPat(GroupOp.Movement|{Ops.INDEX}, name="r").after(name="a", allow_any_len=True),
-   lambda r,a: UOp(r.op, src=(a.replace(src=(r.src[0],)+a.src[1:]),)+r.src[1:], arg=r.arg)),
-  # pass index through elementwise
-  (UPat(GroupOp.Elementwise, name="b").index(name="idx", allow_any_len=True),
-   lambda b,idx: b.replace(src=tuple(s.index(*idx.src[1:]) for s in b.src))),
-  # remove movement ops from SINK. TODO: should be generic
-  (UPat(Ops.SINK, name="s"), lambda s: s.replace(src=tuple(walk_mop(u) for u in s.src))),
-])
-
-# *** split into kernels ***
-
-@dataclass
-class SplitCtx:
-  call_args:list = field(default_factory=list)
-  range_number:int = -1
-
-def _split_graph(ctx:SplitCtx, u:UOp) -> UOp:
-  assert len(u.shape) <= 1, f"rangeify needs to reduce to a single idx, not {u.shape}"
-  ctx.call_args.append(u)
-  return u.param_like(len(ctx.call_args)-1)
-
-def _renumber_range(ctx:SplitCtx, u:UOp) -> UOp:
-  ctx.range_number += 1
-  return u.replace(arg=(ctx.range_number, u.arg[-1]))
-
-pm_split_graph = PatternMatcher([
-  (UPat((Ops.PARAM, Ops.AFTER, Ops.BUFFER), name="u"), _split_graph),
-  (UPat(Ops.RANGE, name="u"), _renumber_range),
-])
-
-def split_store(x:UOp) -> UOp:
-  ret = graph_rewrite(x, pm_split_graph, ctx:=SplitCtx(), name="split kernel", bottom_up=True, walk=True)
-  return ret.sink(arg=KernelInfo()).call(*ctx.call_args)
-
-split_kernels = PatternMatcher([
-  (UPat((Ops.STORE, Ops.END), name="x"), split_store),
-])
-
-# *** main rangeify ***
-
-debug_tag_factor = PatternMatcher([
-  (UPat(GroupOp.All, name="x"), lambda ctx,x: x.rtag(ctx[0][x] if x not in ctx[1] else 'REAL') if x.tag is None else None),
-])
-
-def remove_stage(ctx, x:UOp) -> UOp:
-  buf = UOp.new_buffer(x.arg.device, x.max_numel(), x.dtype, num=next(ctx))
-  return buf.after(buf.reshape(x.shape).index(*x.src[1:]).store(x.src[0]).end(*x.src[1:])).reshape(x.shape)
-pm_remove_stage = PatternMatcher([(UPat(Ops.STAGE, name="x"), remove_stage)])
-
-@rewrite_group(new_ctx=False)
-def get_kernel_graph(sink:UOp) -> UOp:
-  # TODO: multi should just be part of rangeify
-  tsink = graph_rewrite(sink, multi_pm, name="multi_pm")
-
-  # prepare
-  tsink = graph_rewrite(tsink, pm_expand_broadcast, bottom_up=True, name="expand broadcast")
-  tsink = graph_rewrite(tsink, pm_copy_to_store, ctx=itertools.count(0), bottom_up=True, name="convert copy to store")
-
-  # add safe STAGEs to never duplicate compute
-  # we compute the number of times a buffer is consumed. if > 1, we realize
-  realize = {}
-  consumes = {tsink:0}
-  for u in reversed(tsink.toposort()):
-    assert u in consumes, f"{u.op} not in consumes"
-    if (u.op in GroupOp.ALU or u.op is Ops.REDUCE) and consumes[u] > 1 and u.device is not None:
-      # TODO: rename to stage
-      realize[u] = u.rtag(1).bufferize(arg=BufferizeOpts(device=u.device))
-      consumes[u] = 1
-    if u.op is Ops.STORE: consumes[u] = 1
-    if u.op is Ops.EXPAND: consumes[u] *= u.max_numel() // u.src[0].max_numel()
-    for i,s in enumerate(u.src):
-      if s not in consumes: consumes[s] = 0
-      if u.op is not Ops.STORE or i > 0:
-        consumes[s] += consumes[u]
-  if VIZ:
-    with Context(TRACK_MATCH_STATS=0): ctags = graph_rewrite(tsink, debug_tag_factor, ctx=(consumes, realize), bottom_up=True)
-    graph_rewrite(ctags, PatternMatcher([]), name="View Consumes")
-
-  # add stages
-  tsink = graph_rewrite(tsink.substitute(realize), remove_all_tags, name="untag")
-
-  # simple rangeify
-  tsink = graph_rewrite(tsink, pm_range_creation+pm_range_migration, ctx=itertools.count(0), bottom_up=True, name="simple rangeify")
-
-  # TODO: merging and splitting algorithm
-
-  if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Rangeify")
-
-  tsink = graph_rewrite(tsink, pm_remove_stage, ctx=itertools.count(0), bottom_up=True, name="remove stage")
-
-  tsink = graph_rewrite(tsink, split_kernels, bottom_up=True, name="split kernels")
-
-  if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Kernel Graph")
-  if SPEC:
-    # validate the kernel graph
-    from tinygrad.uop.spec import type_verify, spec_kernel_graph
-    type_verify(tsink, spec_kernel_graph, enter_calls=False)
-  return tsink
-

tinygrad/tensor.py

@@ -23,7 +23,6 @@ class AllocCtx:
   bases: set[UOp] = field(default_factory=set)
   assigns: list[UOp] = field(default_factory=list)
   replacements: list[UOp] = field(default_factory=list)
-  views: set[UOp] = field(default_factory=set)
 
 def tag_uop(ctx:AllocCtx, x:UOp):
   if x.tag is not None: return None
@@ -64,37 +63,40 @@ def replace_contig_with_store_after(u:UOp):
 def replace_store_after_with_contig(u:UOp, src:UOp):
   assigned_to = u
   while assigned_to.op in {Ops.BITCAST, Ops.AFTER, Ops.UNSHARD}: assigned_to = assigned_to.src[0].base
-  if assigned_to.op is not Ops.BUFFER: return src.contiguous(tag=u.tag)
+  if assigned_to.op not in {Ops.BUFFER, Ops.SLICE}: return src.contiguous(tag=u.tag)
 
 def _make_buffer_view(src:UOp) -> UOp|None:
-  if (cv := src.contiguous_view()) is None: return None
-  (buf, offset), size = cv, src.max_numel() * src.element_size() // cv[0].element_size()
-  if buf.op is not Ops.BUFFER: return None
-  # NB: make offset a UOp.variable here to do the offset computation in the kernels
-  return buf[offset:offset+size].bitcast(src.dtype)
-
-def contiguous_mops_to_view(ctx:AllocCtx, c:UOp, src:UOp):
-  """MOPS(BUFFER) → SHRINK when movement ops collapse to a contiguous range."""
+  """If movement ops on src collapse to a contiguous range, return SLICE. Otherwise None."""
+  if (offset := src.contiguous_view_offset()) is None: return None
   buf = src.base
-  while buf.op is Ops.BITCAST: buf = buf.src[0].base
-  if buf.op not in {Ops.BUFFER, Ops.UNSHARD}: return None
+  if buf.op is Ops.SLICE:
+    byte_offset = buf.src[1].val * buf.src[0].dtype.itemsize + offset * src.dtype.itemsize
+    buf = buf.src[0]
+    if byte_offset % buf.dtype.itemsize != 0: return None
+    offset = byte_offset // buf.dtype.itemsize
+  return UOp(Ops.SLICE, src.dtype, (buf, UOp.const(offset)), src.numel())
+
+def contiguous_mops_to_view(c:UOp, src:UOp):
+  """MOPS(BUFFER) → SLICE when movement ops collapse to a contiguous range."""
+  buf = src.base
+  if buf.op not in {Ops.BUFFER, Ops.SLICE, Ops.UNSHARD}: return None
+  if src.op is Ops.RESHAPE and src.src[0].op in {Ops.BUFFER, Ops.SLICE} and c.op is not Ops.BITCAST: return None
+  if c.op is not Ops.BITCAST and src.op is Ops.BUFFER: return None
 
   # no symbolic shape
   if not all_int(c.shape): return None
 
   if buf.op is not Ops.UNSHARD and (view := _make_buffer_view(src)) is not None:
-    ctx.views.add(view)
-    view = view.reshape(c.shape)
-    return c.replace(src=(view,)+c.src[1:]) if c.op in {Ops.COPY, Ops.STORE} else view
+    view = (view.replace(dtype=c.dtype, arg=c.numel()) if c.op is Ops.BITCAST else view).reshape(c.shape)
+    return c.replace(src=(view,)) if c.op is Ops.COPY else view
 
-  # for UNSHARD tensors, use multi_pm to resolve per-shard movement ops, then create SHRINK on the resolved result
+  # for UNSHARD tensors, use multi_pm to resolve per-shard movement ops, then create SLICE on the resolved result
   if not isinstance(c.device, str):
     from tinygrad.schedule.multi import multi_pm
     resolved = graph_rewrite(src, multi_pm, name="multi_buffer_view")
     if resolved.op is not Ops.UNSHARD: return None
     if (view := _make_buffer_view(resolved.src[0])) is None: return None
-    ctx.views.add(view)
-    return view.reshape(resolved.src[0].shape).unshard(resolved.arg, resolved.src[1:])
+    return view.reshape(resolved.src[0].shape).unshard(resolved.arg, resolved.src[1:]).contiguous(tag=c.tag)
 
   return None
 
@@ -149,9 +151,8 @@ pm_early_transform_tensor_graph = PatternMatcher([
   # resolve TUPLE+GETTUPLE (for precompiled calls)
   (UPat(Ops.GETTUPLE, src=(UPat(Ops.TUPLE, name="t"),), name="g"), lambda g,t: t.src[g.arg]),
 
-  # fold MOPS+BITCAST over BUFFER into SHRINK when movement ops collapse to contiguous range
-  (UPat((Ops.COPY, Ops.CONTIGUOUS), src=(UPat(GroupOp.Movement|{Ops.BITCAST}, name="src"),), name="c"), contiguous_mops_to_view),
-  (UPat(Ops.STORE, src=(UPat(Ops.BITCAST, name="src"), UPat()), name="c", allow_any_len=True), contiguous_mops_to_view),
+  # fold MOPS+BITCAST over BUFFER/SLICE into SLICE when movement ops collapse to contiguous range
+  (UPat((Ops.BITCAST, Ops.COPY, Ops.CONTIGUOUS), src=(UPat(GroupOp.Movement|{Ops.BUFFER}, name="src"),), name="c"), contiguous_mops_to_view),
 
   # remove contiguous on movement ops before a copy on disk
   (UPat(GroupOp.Movement-{Ops.SHRINK, Ops.RESHAPE}, name="x").f(Ops.CONTIGUOUS).f(Ops.COPY, name="copy"), lambda x,copy:
@@ -200,8 +201,6 @@ def replace_input_buffer(ctx:AllocCtx, b:UOp):
   return UOp.param(len(ctx.replacements)-1, b.dtype, b.shape, b.device,
                    addrspace=b.addrspace if b.addrspace is not None else AddrSpace.GLOBAL)
 
-def replace_input_view(ctx:AllocCtx, b:UOp): return replace_input_buffer(ctx, b) if b in ctx.views else None
-
 pm_finalize_call = PatternMatcher([
   (UPat(Ops.AFTER, name="x"), finalize_after),
   (UPat(Ops.COPY, name="x"), lambda ctx,x: ctx.assigns.append(x) if isinstance(x.device, str) and x.device.startswith(("DISK", "TINYFS")) else None),
@@ -211,9 +210,8 @@ pm_replace_buf = PatternMatcher([
   # replace BUFFER with PARAM for cache key normalization
   (UPat(Ops.BUFFER, src=(UPat(),), name="b"), lambda ctx,b:
    replace_input_buffer(ctx, b) if isinstance(b.arg, ParamArg) and b.addrspace is AddrSpace.GLOBAL else None),
-  # replace SHRINK with PARAM
-  (UPat(Ops.SHRINK, src=(UPat(Ops.BUFFER),), name="b", allow_any_len=True), replace_input_view),
-  (UPat(Ops.BITCAST, src=(UPat.any(UPat(Ops.SHRINK, src=(UPat(Ops.BUFFER),), allow_any_len=True), UPat(Ops.BUFFER)),), name="b"), replace_input_view),
+  # replace SLICE with PARAM. this rewrite is bottom up so BUFFERs we don't need won't be in the input
+  (UPat(Ops.SLICE, src=(UPat(Ops.BUFFER), UPat(Ops.CONST, dtype=dtypes.weakint)), name="b"), replace_input_buffer),
   # strip value from BIND for cache key normalization, so different values hit same cache
   (UPat(Ops.BIND, src=(UPat(Ops.PARAM), UPat(Ops.CONST)), name="b"), replace_input_buffer),
 ])
@@ -231,7 +229,7 @@ def transform_to_call(big_sink:UOp) -> tuple[UOp, dict[UOp, UOp]]:
   big_sink = graph_rewrite(big_sink, add_tags, ctx=ctx, bottom_up=True, name="number the uops")
 
   # here we can break the tensor graph. this is the only place you need to maintain numbered tags
-  big_sink = graph_rewrite(big_sink, pm_early_transform_tensor_graph, ctx=ctx, name="early transform tensor graph")
+  big_sink = graph_rewrite(big_sink, pm_early_transform_tensor_graph, name="early transform tensor graph")
 
   # here we construct the final buffer_map: as-built nodes -> their final storage. values are never keys
   graph_rewrite(big_sink, pm_finalize_call, ctx=ctx, name="finalize call")

tinygrad/uop/ops.py

@@ -568,9 +568,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     in_tuple = self.src[0] if self.op is Ops.FUNCTION else self
     assert in_tuple.op is Ops.TUPLE, f"gettuple requires FUNCTION or TUPLE source, got {self.op}"
     return UOp(Ops.GETTUPLE, src=(self,), arg=idx)
-  def group(*srcs:UOp|None, **kwargs):  # pylint: disable=no-self-argument
+  def group(*srcs:UOp|None):  # pylint: disable=no-self-argument
     if len(srcs) == 1 and isinstance(srcs[0], UOp): return srcs[0]
-    return UOp(Ops.GROUP, src=tuple([x for x in srcs if x is not None]), **kwargs)
+    return UOp(Ops.GROUP, src=tuple([x for x in srcs if x is not None]))
   def index(self, *srcs:UOp|int|None, **kwargs):
     new_srcs: list[UOp] = [UOp.const(x) if isinstance(x, int) else x for x in srcs if x is not None]
     if len(new_srcs) == 1 and new_srcs[0].op is Ops.CONST and self.op is Ops.STACK: return self.src[new_srcs[0].val]
@@ -822,7 +822,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
   unique_num = itertools.count(0)
 
   def getaddr(self, device=None) -> UOp:
-    if self.without_after.op not in {Ops.BUFFER, Ops.SLICE, Ops.SHRINK, Ops.BINARY, Ops.MSTACK, Ops.MSELECT, Ops.PARAM}: return self
+    if self.without_after.op not in {Ops.BUFFER, Ops.SLICE, Ops.BINARY, Ops.MSTACK, Ops.MSELECT, Ops.PARAM}: return self
     return UOp(Ops.GETADDR, src=(self,), arg=device or to_tuple(self.device)[0])
   @staticmethod
   def new_buffer(device:str|tuple[str, ...], size:int, dtype:DType, num=None):
@@ -901,7 +901,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     while len(s.src) and s.op not in {Ops.BUFFER, Ops.PARAM, Ops.STAGE, Ops.MSTACK}: s = s.src[0]
     return s
 
-  def contiguous_view(self) -> tuple[UOp, int]|None:
+  def contiguous_view_offset(self) -> int|None:
+    """If movement ops on a BUFFER collapse to a contiguous range, return `offset` in elements. Otherwise None."""
     from tinygrad.schedule.rangeify import pm_mops
     from tinygrad.uop.symbolic import symbolic
 
@@ -914,10 +915,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
 
     idx = self.flatten().index(UOp.range(self.numel(), 0))
     out = graph_rewrite(idx, pm_mops+symbolic+pm_contiguous_view_offset, ctx=self, name="contiguous_view_offset")
-    if out.op is not Ops.INDEX or not (b:=out.src[0]).tag or (c:=out.src[1]).op is not Ops.CONST or not isinstance(c.val, int): return None
-    return b.rtag(None), c.val
-
-  def contiguous_view_offset(self) -> int|None: return None if (view := self.contiguous_view()) is None else view[1]
+    return out.val if out.op is Ops.CONST and isinstance(out.val, int) else None
 
   def has_buffer_identity(self, after_ok=False):
     """Check if this UOp has a concrete buffer identity in the graph (RESHAPE/UNSHARD -> BUFFER chain)."""
@@ -934,16 +932,18 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
 
   @property
   def buffer(self) -> Buffer|MultiBuffer:
-    if self.op in {Ops.CONTIGUOUS, Ops.CONTIGUOUS_BACKWARD, Ops.RESHAPE, Ops.UNSHARD, Ops.DETACH, Ops.AFTER}: return self.src[0].buffer
+    if self.op in {Ops.CONTIGUOUS, Ops.RESHAPE, Ops.UNSHARD, Ops.DETACH, Ops.AFTER}: return self.src[0].buffer
     # this buffer can process disk tensors and simple movement ops
-    if self is not self.base or self.op is Ops.BITCAST:
-      if (cv := self.contiguous_view()) is None: raise RuntimeError(f"non-contiguous view is not supported for {self.device} buffer")
-      buf, offset = (b:=cv[0]).base.buffer, cv[1]
-      if isinstance(buf, MultiBuffer):
-        mbuf = MultiBuffer.__new__(MultiBuffer)
-        mbuf.bufs = [x.view(prod(self.max_shape), self.dtype, offset*b.dtype.itemsize) for x in buf.bufs]
-        return mbuf
-      return buf.view(prod(self.max_shape), self.dtype, offset*b.dtype.itemsize)
+    if self is not self.base:
+      buf = self.base.buffer
+      assert isinstance(buf, Buffer), "must be a Buffer for movement ops"
+      offset = self.contiguous_view_offset()
+      if offset is None: raise RuntimeError(f"non-contiguous view is not supported for {buf.device} buffer")
+      return buf.view(prod(self.max_shape), self.dtype, offset*self.dtype.itemsize)
+    if self.op is Ops.BITCAST:
+      buf = self.src[0].buffer
+      assert isinstance(buf, Buffer), "must be a Buffer for BITCAST"
+      return buf.view(prod(self.max_shape), self.dtype, 0)
     if self.op is Ops.SLICE:
       if (cret:=buffers.get(self)) is not None: return cret
       buf = self.src[0].buffer
@@ -1775,14 +1775,10 @@ pm_unbind = PatternMatcher([(UPat(Ops.BIND, name="x"), do_unbind)])
 
 # ctx is source UOp for which we are finding a contiguous view for. used in contiguous_view_offset
 pm_contiguous_view_offset = PatternMatcher([
-  # normalize to 1d bitcasts
-  (UPat(Ops.BITCAST, name="b"), lambda b: b.src[0].flatten().bitcast(b.dtype).reshape(b.shape) if len(b.shape) != 1 else None),
-  (UPat(Ops.BITCAST, name="b").index(UPat.cvar("c")), lambda ctx, b, c:
-   b.src[0].flatten().index(UOp.range(ctx.numel() * (osz:=b.element_size())//(isz:=b.src[0].element_size()), 0) + (c * osz//isz)) if b.tag else None),
-  (UPat(Ops.INDEX, src=(UPat.var("b"),)), lambda b: b.rtag().index(0)),
-  (UPat(Ops.INDEX, src=(UPat.var("b"), UPat(Ops.RANGE))), lambda b: b.rtag().index(0)),
-  (UPat(Ops.INDEX, src=(UPat.var("b"), UPat(Ops.RANGE)+UPat.cvar('c'))), lambda ctx, b, c: b.rtag().index(c)),
-  (UPat(Ops.INDEX, src=(UPat.var("b"), UPat.cvar('c'))), lambda ctx, b, c: b.rtag().index(c) if resolve(ctx.numel() == 1, False) else None),
+  (UPat(Ops.INDEX, src=(UPat(),)), lambda: UOp.const(0)),
+  (UPat(Ops.INDEX, src=(UPat(), UPat(Ops.RANGE))), lambda: UOp.const(0)),
+  (UPat(Ops.INDEX, src=(UPat(), UPat(Ops.RANGE)+UPat.cvar('c'))), lambda c: c),
+  (UPat(Ops.INDEX, src=(UPat(), UPat.cvar('c'))), lambda ctx, c: c if resolve(ctx.numel() == 1, False) else None),
 ])
 
 # *** what was symbolic.py ***