bitcast is a mop

* shrink uses offset and shape

* pad does too

* fix

.github/actions/setup-tinygrad/action.yml

@@ -281,7 +281,7 @@ runs:
         cache-name: cache-gpuocelot-build-1
       with:
         path: ${{ github.workspace }}/gpuocelot/ocelot
-        key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.CACHE_VERSION }}
+        key: ${{ runner.os }}-gpuocelot-f463259669c69abce7b3a0567b6c284f348d0f32-rebuild-${{ env.CACHE_VERSION }}
     - name: Cache gpuocelot
       if: inputs.ocelot == 'true' && github.event_name != 'pull_request'
       id: cache-build
@@ -290,14 +290,14 @@ runs:
         cache-name: cache-gpuocelot-build-1
       with:
         path: ${{ github.workspace }}/gpuocelot/ocelot
-        key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.CACHE_VERSION }}
+        key: ${{ runner.os }}-gpuocelot-f463259669c69abce7b3a0567b6c284f348d0f32-rebuild-${{ env.CACHE_VERSION }}
     - name: Clone/compile gpuocelot
       if: inputs.ocelot == 'true' && steps.cache-build-pr.outputs.cache-hit != 'true' && steps.cache-build.outputs.cache-hit != 'true'
       shell: bash
       run: |
-        git clone --recurse-submodules https://github.com/gpuocelot/gpuocelot.git ${{ github.workspace }}/gpuocelot
+        git clone --recurse-submodules https://github.com/tinygrad/gpuocelot.git ${{ github.workspace }}/gpuocelot
         cd ${{ github.workspace }}/gpuocelot/ocelot
-        git checkout b16039dc940dc6bc4ea0a98380495769ff35ed99
+        git checkout f463259669c69abce7b3a0567b6c284f348d0f32
         mkdir build
         cd build
 
@@ -306,10 +306,7 @@ runs:
           sudo xcode-select -s /Applications/Xcode_16.2.app/Contents/Developer
           CMAKE_ARGS="$CMAKE_ARGS -DBoost_INCLUDE_DIR=$(brew --prefix boost)/include -DBoost_LIBRARY_DIR=$(brew --prefix boost)/lib"
         else
-          curl -fL https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvcc/linux-x86_64/cuda_nvcc-linux-x86_64-11.5.119-archive.tar.xz \
-            | sudo tar -xJ -C /usr/ --strip-components=1
-          curl -fL https://developer.download.nvidia.com/compute/cuda/redist/cuda_cudart/linux-x86_64/cuda_cudart-linux-x86_64-11.5.117-archive.tar.xz \
-            | sudo tar -xJ -C /usr/ --strip-components=1
+          CMAKE_ARGS="$CMAKE_ARGS -DLLVM_DIR=$(llvm-config-15 --cmakedir)"
         fi
 
         cmake .. $CMAKE_ARGS

.github/workflows/test.yml

@@ -35,7 +35,7 @@ jobs:
 
   docs:
     name: Docs
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-24.04
     timeout-minutes: 10
     env:
       CHECK_OOB: 0
@@ -89,7 +89,7 @@ jobs:
 
   torchbackend:
     name: Torch Backend Tests
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     timeout-minutes: 15
     steps:
     - name: Checkout Code
@@ -125,7 +125,7 @@ jobs:
 
   torchbackendmore:
     name: Torch Backend Tests More
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     timeout-minutes: 15
     steps:
     - name: Checkout Code
@@ -147,7 +147,7 @@ jobs:
 
   bepython:
     name: Python Backend
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     timeout-minutes: 15
     steps:
     - name: Checkout Code
@@ -215,7 +215,7 @@ jobs:
 
   linter:
     name: Linters
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     timeout-minutes: 10
 
     steps:
@@ -246,7 +246,7 @@ jobs:
 
   nulltest:
     name: Null Tests
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     timeout-minutes: 15
 
     steps:
@@ -277,7 +277,7 @@ jobs:
 
   unittest:
     name: Unit Tests
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     timeout-minutes: 15
 
     steps:
@@ -320,7 +320,7 @@ jobs:
       matrix:
         group: [1, 2]
     name: SPEC=2 (${{ matrix.group }})
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     timeout-minutes: 15
     steps:
     - name: Checkout Code
@@ -336,7 +336,7 @@ jobs:
 
   fuzzing:
     name: Fuzzing
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     timeout-minutes: 10
     steps:
     - name: Checkout Code
@@ -357,7 +357,7 @@ jobs:
 
   testopenclimage:
     name: CL IMAGE Tests
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-24.04
     timeout-minutes: 15
     steps:
       - name: Checkout Code
@@ -377,7 +377,7 @@ jobs:
 
   testgpumisc:
     name: CL Misc tests
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-24.04
     timeout-minutes: 10
     steps:
       - name: Checkout Code
@@ -402,7 +402,7 @@ jobs:
 
   testopenpilot:
     name: openpilot Compile Tests
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-24.04
     timeout-minutes: 15
     steps:
       - name: Checkout Code
@@ -432,7 +432,7 @@ jobs:
 
   testonnxcpu:
     name: ONNX (CPU) Tests
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-24.04
     timeout-minutes: 20
 
     steps:
@@ -460,7 +460,7 @@ jobs:
 
   testopencl:
     name: ONNX (CL)+Optimization Tests
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-24.04
     timeout-minutes: 20
     steps:
       - name: Checkout Code
@@ -519,7 +519,7 @@ jobs:
 
   testmodels:
     name: Models (llvm+cpu+gpu)
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-24.04
     timeout-minutes: 15
     steps:
       - name: Checkout Code
@@ -586,7 +586,7 @@ jobs:
 
   testwebgpu:
     name: Linux (WebGPU)
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-24.04
     timeout-minutes: 20
     steps:
     - name: Checkout Code
@@ -693,7 +693,7 @@ jobs:
         arch: [gfx1100, gfx1201, gfx950]
 
     name: Linux (${{ matrix.backend }} ${{ matrix.arch }})
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-24.04
     timeout-minutes: 15
     env:
       DEV: MOCKKFD+AMD:${{ matrix.backend == 'amdllvm' && 'LLVM' || '' }}:${{ matrix.arch }}
@@ -728,7 +728,7 @@ jobs:
         backend: [ptx, nv]
 
     name: Linux (${{ matrix.backend }})
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-24.04
     timeout-minutes: 20
     env:
       FORWARD_ONLY: 1
@@ -763,7 +763,7 @@ jobs:
         backend: [llvm, cpu, opencl, lvp, x86]
 
     name: Linux (${{ matrix.backend }})
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-24.04
     timeout-minutes: 20
     steps:
       - name: Checkout Code

examples/mlperf/model_train.py

@@ -1500,7 +1500,7 @@ def train_llama3():
   def fake_data(bs, samples):
     import numpy as np
     for _ in range(samples // bs):
-      fake_data_np = np.random.randint(0, model_params["vocab_size"], size=(bs, SEQLEN + 1), dtype=np.int32)
+      fake_data_np = np.random.randint(0, real_vocab_size, size=(bs, SEQLEN + 1), dtype=np.int32)
       yield Tensor(fake_data_np, device="NPY")
 
   def get_train_iter():

examples/mlperf/models/flat_llama.py

@@ -276,9 +276,9 @@ def apply_grad(grad_buf:Tensor, new_grad:UOp):
     grad_buf.uop = grad_buf.uop.after(grad_buf.uop.store(grad_buf.uop + new_grad))
     return
   cur = grad_buf.uop
-  for pad in sorted(pads, key=lambda p: p.marg[0][0] if p.op == Ops.PAD else 0, reverse=True):
+  for pad in sorted(pads, key=lambda p: p.marg[0][1] if p.op == Ops.PAD else 0, reverse=True):
     if pad.op == Ops.PAD:
-      grad_shrink = tuple([(p[0], s+p[0]) for s,p in zip(pad.src[0].shape, pad.marg)])
+      grad_shrink = tuple([(p[1], s+p[1]) for s,p in zip(pad.src[0].shape, pad.marg)])
       buf_slice = cur.shrink(grad_shrink)
       cur = cur.after(buf_slice.store(buf_slice + pad.src[0].cast(cur.dtype)))
     else:

extra/hcq2/hcq2.py

@@ -1,9 +1,9 @@
 from __future__ import annotations
 from typing import cast, Callable, TypeVar, Generic, Any, TYPE_CHECKING
-import struct, functools, time, collections
+import struct, functools, time, collections, importlib, itertools
 from dataclasses import replace
 if TYPE_CHECKING: from tinygrad.engine.realize import ExecContext
-from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, mv_address, round_up, DEBUG, dedup
+from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, mv_address, round_up, DEBUG, dedup, pluralize
 from tinygrad.device import Device, Buffer, BufferSpec, Compiled, LRUAllocator, MultiBuffer
 from tinygrad.uop.ops import Ops, sint, UOp, UPat, PatternMatcher, KernelInfo, graph_rewrite, track_rewrites
 from tinygrad.uop.symbolic import symbolic_simple, symbolic
@@ -12,7 +12,7 @@ from dataclasses import dataclass, field
 from tinygrad.runtime.support.memory import BumpAllocator
 from tinygrad.runtime.support.hcq import MMIOInterface
 from tinygrad.renderer import Renderer, Estimates
-from tinygrad.engine.realize import to_program, track_stats, get_call_arg_uops, resolve_params
+from tinygrad.engine.realize import to_program, track_stats, get_call_arg_uops, resolve_params, pm_flatten_linear
 
 HCQDeviceType = TypeVar('HCQDeviceType', bound='HCQ2Compiled')
 
@@ -22,19 +22,10 @@ class HCQ2Compiled(Compiled):
   """
   timestamp_divider: float = 1000.0  # GPU timestamp counter ticks per microsecond; override per device
 
-  def __init__(self, device:str, allocator:'HCQAllocator', compilers:list[type[Renderer]], runtime,
-               kernargs_size=(16 << 20), can_recover:bool=False, arch=None):
+  def __init__(self, device:str, allocator:'HCQAllocator', compilers:list[type[Renderer]], runtime, can_recover:bool=False, arch=None):
     self.device_id:int = int(device.split(":")[1]) if ":" in device else 0
 
-    from extra.hcq2.graph.hcq import HCQ2Graph
-    super().__init__(device, allocator, compilers, lambda *a, **kw: None, HCQ2Graph, arch=arch)
-
-    self.kernargs_size = kernargs_size
-    self.kernargs_offset_allocator:BumpAllocator = BumpAllocator(kernargs_size, wrap=True)
-
-  @functools.cached_property
-  def kernargs_buf(self) -> Buffer:
-    return Buffer(self.device, self.kernargs_size, dtypes.uint8, options=BufferSpec(cpu_access=True), preallocate=True)
+    super().__init__(device, allocator, compilers, lambda *a, **kw: None, None, arch=arch)
 
   @functools.cached_property
   def timeline_signal(self) -> Buffer:
@@ -50,6 +41,14 @@ class HCQ2Compiled(Compiled):
     buf.as_memoryview(force_zero_copy=True).cast('Q')[0] = 1
     return buf
 
+  @functools.cached_property
+  def pm_bufferize(self) -> PatternMatcher:
+    return PatternMatcher([
+      (UPat(Ops.BUFFER, tag="timeline_signal"), lambda ctx: ctx.timeline_signal),
+      (UPat(Ops.BUFFER, tag="timeline_value"), lambda ctx: ctx.timeline_value),
+      (UPat(Ops.BUFFER, name="b"), lambda ctx, b: Buffer(ctx.device, b.arg, b.dtype, options=BufferSpec(host=True, uncached=True, cpu_access=True))),
+    ])
+
   def synchronize(self, timeout:int|None=None):
     if not hasattr(self, 'iface'): return
     sig = self.timeline_signal._buf.cpu_view().mv.cast('Q')
@@ -126,7 +125,7 @@ class HCQAllocator(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):
   def _copy(self, dst:Buffer, src:Buffer):
     from tinygrad.engine.realize import run_linear
     su = UOp.from_buffer(src)
-    run_linear(UOp(Ops.LINEAR, dtypes.void, (su.copy_to_device(dst.device).call(UOp.from_buffer(dst), su),)), jit=True, update_stats=False)
+    run_linear(UOp(Ops.LINEAR, dtypes.void, (su.copy_to_device(dst.device).call(UOp.from_buffer(dst), su),)), update_stats=False)
 
   def _copyin(self, dest:HCQ2Buffer, src:memoryview):
     s = Buffer(self.dev.device, len(src), dtypes.uint8, options=BufferSpec(host=True), preallocate=True)
@@ -141,33 +140,16 @@ class HCQAllocator(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):
 
   # def _as_buffer(self, buf): return buf.cpu_view().mv
 
-# **************** lower context ****************
-
 def unwrap_after(uop):
   while uop.op is Ops.AFTER: uop = uop.src[0]
   return uop
 
-@dataclass
-class HCQ2DeviceCtx:
-  device:str                       # device name; resolve to instance via Device[device]
-  kernargs_host:UOp                # UOp whose .buffer is dev.kernargs_buf (BUFFER UOp in runtime, PARAM in graph)
-  kernargs_gpu:UOp                 # va_addr const of dev.kernargs_buf
-  kernargs_allocator:BumpAllocator = field(default_factory=lambda: BumpAllocator(2 << 20, wrap=False))
-
-@dataclass
-class HCQ2LowerCtx:
-  name:str
-  inputs:list[Buffer|MultiBuffer] = field(default_factory=list)
-  holds:list[UOp] = field(default_factory=list)
-  dev_ctx:dict[str, HCQ2DeviceCtx] = field(default_factory=dict)
-  addr_table:UOp|None = None
-  next_slot:int = 0
-
 class HCQEncoder:
   def __init__(self): self.blob, self.patches = b'', []
 
   def get_dev_addr(self, uop:UOp) -> UOp:
-    return UOp(Ops.GETADDR, dtypes.uint64, src=(uop,)) if unwrap_after(uop).op in (Ops.BUFFER, Ops.SLICE, Ops.BINARY, Ops.MSTACK, Ops.MSELECT) else uop
+    if unwrap_after(uop).op not in (Ops.BUFFER, Ops.SLICE, Ops.BINARY, Ops.MSTACK, Ops.MSELECT): return uop
+    return UOp(Ops.GETADDR, dtypes.uint64, src=(uop, UOp(Ops.DEVICE, arg=self.dev.device)))
 
   def append(self, *data, dtype=dtypes.uint32):
     for d in data:
@@ -182,302 +164,213 @@ class HCQEncoder:
     buf = UOp.new_buffer(dev, len(self.blob), dtypes.uint8)
     if tag: buf = buf.rtag(tag)
     blob_uop = UOp(Ops.BINARY, dtypes.void, src=(), arg=self.blob)
-    stores = [buf.index(UOp.const(dtypes.int, off)).cast(dt.ptr()).store(val.cast(dt)) for off, val, dt in self.patches]
+    stores = [buf.index(UOp.const(dtypes.int, off), dtype=buf.dtype.ptr()).cast(dt.ptr()).store(val.cast(dt)) for off, val, dt in self.patches]
     return buf.after(buf.store(blob_uop), *stores)
 
-# **************** prepare runtime ****************
+# *****************
+# 1.1. prep runtimes: staging copies
 
-def _devices(x) -> tuple[str, ...]:
-  return tuple(b.device for b in x.bufs) if isinstance(x, MultiBuffer) else (x.device,) if isinstance(x, Buffer) else x if isinstance(x, tuple) else (x,)
+HCQ_DEVS = ("AMD",)
+HCQ_P2P_DEVS = HCQ_DEVS + ("CPU",)
 
-def rebind_program_dev(c:UOp, p:UOp) -> UOp:
-  devs = _devices(c.src[1].buffer)
-  p = p.replace(src=p.src[:1] + (UOp(Ops.DEVICE, arg=devs),) + p.src[2:])
-  return c.replace(src=(Device[devs[0]].pm_lower.rewrite(p),) + c.src[1:])
+def stage_copy(dst:UOp, src:UOp) -> UOp|None:
+  dst_dev, src_dev = dst.device.split(":")[0], src.device.split(":")[0]
+  if not ((src_dev in HCQ_DEVS and dst_dev not in HCQ_P2P_DEVS) or (dst_dev in HCQ_DEVS and src_dev not in HCQ_P2P_DEVS)): return None
+  stage = UOp.new_buffer("CPU", src.buffer.nbytes, dtypes.uint8)
+  return UOp(Ops.LINEAR, dtypes.void, (src.copy_to_device("CPU").call(stage, src), stage.copy_to_device(dst.device).call(dst, stage)))
+pm_insert_copy_staging = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src"))), stage_copy)])
 
-def lower_kernargs(call:UOp, prg:UOp) -> UOp:
+# *****************
+# 1.2. prep runtimes: programs/kernargs
+
+@functools.cache
+def get_pm_prep_program(name:str) -> PatternMatcher|None:
+  try:
+    importlib.import_module(f'tinygrad.runtime.ops_{name.lower()}') # TODO: remove that
+    return importlib.import_module(f'extra.hcq2.ops_{name.lower()}2').pm_prep_program
+  except ImportError: return None
+
+def prep_program(call:UOp, prg:UOp) -> UOp|None:
+  dev = call.src[1].device
+  if (pm:=get_pm_prep_program(dev.split(":")[0])) is None or (lowered:=pm.rewrite(prg)) is None: return None
+  data, image_bytes = lowered
+  buf = UOp.new_buffer(dev, len(image_bytes), dtypes.uint8).rtag("program")
+  blob = UOp(Ops.BINARY, dtypes.void, src=(), arg=image_bytes)
+  return call.replace(src=(prg.replace(src=(buf.after(buf.store(blob)),), arg=(data, prg.arg)),) + call.src[1:])
+
+def prep_kernargs(call:UOp, prg:UOp) -> UOp:
   data, info = prg.arg
-  enc = HCQEncoder()
-  for gi in info.globals: enc.append(call.src[1+gi], dtype=dtypes.uint64)
-  for v in info.vars: enc.append(v, dtype=dtypes.uint32)
+  patches = [(i*dtypes.uint64.itemsize, UOp(Ops.GETADDR, dtypes.uint64, src=(call.src[1+gi], UOp(Ops.DEVICE, arg=call.src[1+gi].device))),
+              dtypes.uint64) for i,gi in enumerate(info.globals)] \
+          + [(len(info.globals)*dtypes.uint64.itemsize + i*dtypes.uint32.itemsize, v, dtypes.uint32) for i,v in enumerate(info.vars)]
+
+  buf = UOp.new_buffer(call.src[1].device, data.kernargs_alloc_size, dtypes.uint8).rtag("kernargs")
+  kernargs = buf.after(*tuple(buf.index(UOp.const(dtypes.int, o), dtype=buf.dtype.ptr()).cast(dt.ptr()).store(val.cast(dt)) for o, val, dt in patches))
 
-  enc.blob += b'\x00' * (data.kernargs_alloc_size - len(enc.blob)) # pad blob
-  kernargs = enc.uop(_devices(call.src[1].buffer), tag="kernargs")
   return call.replace(src=(prg.replace(src=prg.src + (kernargs,), arg=(data, info)),) + call.src[1:])
 
 pm_prep_runtime = PatternMatcher([
   # bind generic PROGRAM device to the call's actual dev(s), then run device-specific lowering
-  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.DEVICE), UPat(), UPat(), UPat(Ops.BINARY)), name="p"),), name="c", allow_any_len=True),
-    rebind_program_dev),
+  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(), UPat(), UPat(), UPat(), UPat(Ops.BINARY)), name="prg"),),
+    name="call", allow_any_len=True), prep_program),
 
   # lower kernargs (PROGRAM.src[0] is now AFTER(BUFFER, COPY) — the lowered program image)
-  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(Ops.AFTER),), name="prg"),), name="call", allow_any_len=True), lower_kernargs),
+  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(Ops.AFTER),), name="prg"),), name="call", allow_any_len=True), prep_kernargs),
 ])
 
-# **************** lower ops ****************
+# *****************
+# 2.1. lowering to hcq ir
 
 def lower_program(call:UOp, prg:UOp) -> UOp:
-  q = UOp(Ops.LINEAR, dtypes.void, (prg,), arg=(_devices(call.src[1].buffer), "COMPUTE"))
-  return UOp(Ops.LINEAR, dtypes.void, (q,), tag=call.tag)
+  q = UOp(Ops.LINEAR, dtypes.void, (prg,), arg=(call.src[1].device, "COMPUTE"))
+  return call.replace(src=(q,) + call.src[1:]).rtag('hcq')
 
-def lower_copy(call:UOp, copy:UOp) -> UOp:
+def lower_copy(call:UOp, copy:UOp) -> UOp|None:
   dst, src = call.src[1], call.src[2]
-  q = UOp(Ops.LINEAR, dtypes.void, (UOp(Ops.COPY, dtypes.void, src=(dst, src), arg=src.buffer.nbytes),), arg=(_devices(dst.buffer), "COPY"))
-  return UOp(Ops.LINEAR, dtypes.void, (q,), tag=call.tag)
+  if (hcq_dev:=next((b.device for b in (dst, src) if b.device.split(":")[0] in HCQ_DEVS), None)) is None: return None
+
+  q = UOp(Ops.LINEAR, dtypes.void, (UOp(Ops.COPY, dtypes.void, src=(dst, src), arg=src.buffer.nbytes),), arg=(hcq_dev, "COPY"))
+  return call.replace(src=(q,) + call.src[1:]).rtag('hcq')
 
 pm_lower_ops = PatternMatcher([
-  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(Ops.AFTER), UPat()), name="prg"),), name="call", allow_any_len=True), lower_program),
+  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(Ops.AFTER), UPat(Ops.AFTER)), name="prg"),), name="call", allow_any_len=True), lower_program),
   (UPat(Ops.CALL, src=(UPat(Ops.COPY, name="copy"),), name="call", allow_any_len=True), lower_copy),
 ])
 
-def split_into_queues(outer:UOp) -> UOp:
-  groups:dict[tuple, list[UOp]] = collections.defaultdict(list)
-  for child in outer.src:
-    wrapper = child.src[0] if child.op is Ops.AFTER else child
-    for q in wrapper.src: groups[q.arg].extend(q.src)
-  return outer.replace(src=tuple(UOp(Ops.LINEAR, dtypes.void, tuple(cmds), arg=k) for k, cmds in groups.items()))
-pm_split_into_queues = PatternMatcher([(UPat(Ops.LINEAR, src=UPat(Ops.LINEAR, src=UPat(Ops.LINEAR)).or_after(), name="outer"), split_into_queues)])
+# *****************
+# 2.2. queue split
 
-def add_signals(q:UOp) -> UOp:
+# def split_into_queues(linear:UOp) -> UOp:
+#   out = []
+#   for k, grp in itertools.groupby(linear.src, lambda c: c.src[0].arg if c.op is Ops.CALL and c.src[0].op is Ops.LINEAR else None):
+#     if k is None: out.extend(grp)
+#     else:
+#       calls = list(grp)
+#       items = tuple(x for c in calls for x in c.src[0].src)
+#       args = tuple(a for c in calls for a in c.src[1:])
+#       out.append(calls[0].replace(src=(UOp(Ops.LINEAR, dtypes.void, items, arg=k),) + args))
+#   return linear.replace(src=tuple(out))
+# pm_split_into_queues = PatternMatcher([(UPat(Ops.LINEAR, name="linear"), split_into_queues)])
+
+# *****************
+# 2.3. barriers / signals / timeline inc
+
+def add_barriers(call:UOp, q:UOp) -> UOp:
+  return call.replace(src=(q.replace(src=(UOp(Ops.BARRIER, dtypes.void), *q.src)),) + call.src[1:])
+pm_add_barriers = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.LINEAR, name="q"),), name="call", allow_any_len=True), add_barriers)])
+
+def add_signals(call:UOp, q:UOp) -> UOp:
   sig = UOp.new_buffer(q.arg[0], 0x100, dtypes.uint8).rtag("timeline_signal")
   tl = UOp.new_buffer(q.arg[0], 1, dtypes.uint64).rtag("timeline_value").index(UOp.const(dtypes.int, 0))
-  return q.replace(src=(sig.wait(tl-1), *q.src, sig.store(tl)), arg=q.arg)
-pm_add_signals = PatternMatcher([(UPat(Ops.LINEAR, src=UPat(Ops.LINEAR), name="outer"),
-  lambda outer: outer.replace(src=tuple(add_signals(q) for q in outer.src)))])
+  return call.replace(src=(q.replace(src=(sig.wait(tl-1), *q.src, sig.store(tl)), arg=q.arg),) + call.src[1:])
+pm_add_signals = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.LINEAR, name="q"),), name="call", allow_any_len=True), add_signals)])
 
-pm_add_barriers = PatternMatcher([(UPat(Ops.LINEAR, src=UPat(Ops.LINEAR), name="outer"),
-  lambda outer: outer.replace(src=tuple(q.replace(src=(UOp(Ops.BARRIER, dtypes.void), *q.src)) for q in outer.src)))])
+# *****************
+# 3.1. encode cmdbufs
 
-def add_timeline_inc(q:UOp) -> UOp:
-  tl = UOp.new_buffer(q.arg[0], 1, dtypes.uint64).rtag("timeline_value")
-  done = tl.after(UOp(Ops.BARRIER, dtypes.void, src=(q,)))
-  return done.index(UOp.const(dtypes.int, 0), dtype=tl.dtype.ptr()).store(tl.index(UOp.const(dtypes.int, 0)) + 1)
-pm_add_timeline_inc = PatternMatcher([(UPat(Ops.LINEAR, src=UPat(Ops.LINEAR), name="outer"),
-  lambda outer: outer.replace(src=tuple(add_timeline_inc(q) for q in outer.src)))])
+@functools.cache
+def get_pm_lower(name:str) -> PatternMatcher|None:
+  try:
+    importlib.import_module(f'tinygrad.runtime.ops_{name.lower()}') # TODO: remove that
+    return importlib.import_module(f'extra.hcq2.ops_{name.lower()}2').pm_lower
+  except ImportError: return None
 
-# **************** build host program ****************
+def encode_cmdbuf(call:UOp, q:UOp) -> UOp|None:
+  if (pm:=get_pm_lower(call.src[1].device.split(":")[0])) is None or (encoded:=pm.rewrite(q)) is None: return None
+  return call.replace(src=(encoded,) + call.src[1:])
+pm_encode_cmdbufs = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.LINEAR, name="q"),), name="call", allow_any_len=True), encode_cmdbuf)])
 
-def calc_kernargs_sizes(ctx:dict[str,int], u:UOp) -> None:
-  if u.tag != "kernargs": return
-  for d in _devices(u.src[1].arg): ctx[d] = ctx.get(d, 0) + round_up(u.arg, 16)
-pm_calc_kernargs_sizes = PatternMatcher([(UPat(Ops.BUFFER, name="u"), calc_kernargs_sizes)])
+# *****************
+# 3.2. add timeline inc
 
-# bufferize
+def add_timeline_inc(call:UOp, s:UOp) -> UOp:
+  tl = UOp.new_buffer(s.device, 1, dtypes.uint64).rtag("timeline_value")
+  return call.replace(src=(tl.after(s).index(UOp.const(dtypes.int, 0), dtype=tl.dtype.ptr()).store(tl.index(UOp.const(dtypes.int, 0)) + 1),) + call.src[1:])
+pm_add_timeline_inc = PatternMatcher([(UPat(Ops.CALL, tag="hcq", src=(UPat(name="s"),), name="call", allow_any_len=True), add_timeline_inc)])
 
-def _maybe_mstack(srcs:tuple[UOp, ...], tag=None) -> UOp: return srcs[0] if len(srcs) == 1 else UOp(Ops.MSTACK, srcs[0].dtype, srcs, tag=tag)
+# *****************
+# 3.3. lift patches to the command buffer (root)
 
-def _lower_stores(host_buf:UOp, buf_node:UOp, stores:tuple[UOp, ...]) -> list[UOp]:
-  def lower(s:UOp) -> UOp:
-    if s.src[1].op is Ops.BINARY: return s.substitute({buf_node: host_buf})
-    idx = s.src[0].src[0]
-    return s.substitute({idx: host_buf.index(UOp.const(dtypes.int, idx.src[1].arg // host_buf.dtype.base.itemsize), dtype=host_buf.dtype.ptr())})
-  return [lower(s) for s in stores]
-
-_program_uop_cache:dict[tuple[bytes,str], tuple[UOp,UOp]] = {}
-def bufferize_program(ctx:HCQ2LowerCtx, target:UOp, buf_node:UOp) -> UOp:
-  blob, addrs = target.src[1].src[1].arg, []
-  for dev in _devices(buf_node.src[1].arg):
-    if (cached:=_program_uop_cache.get((blob, dev))) is None:
-      lib_gpu = Buffer(dev, round_up(len(blob), 0x1000), dtypes.uint8, options=BufferSpec(nolru=True, cpu_access=True), preallocate=True)
-      lib_gpu._buf.cpu_view()[:len(blob)] = memoryview(blob)
-      cached = _program_uop_cache[(blob, dev)] = (UOp.from_buffer(lib_gpu, dev), UOp.const(dtypes.uint64, lib_gpu._buf.va_addr, device=dev))
-    if cached[0] not in ctx.holds: ctx.holds.append(cached[0])
-    addrs.append(cached[1])
-  return _maybe_mstack(tuple(addrs))
-
-def bufferize_kernargs(ctx:HCQ2LowerCtx, target:UOp, buf_node:UOp) -> UOp:
-  hbufs, addrs = [], []
-  for dev in _devices(buf_node.src[1].arg):
-    dctx = ctx.dev_ctx[dev]
-    isz = dctx.kernargs_host.dtype.base.itemsize
-    off = dctx.kernargs_allocator.alloc(buf_node.arg, 16)
-    hbufs.append(UOp(Ops.SLICE, dctx.kernargs_host.dtype,
-                     src=(dctx.kernargs_host, UOp.const(dtypes.weakint, off // isz)), arg=buf_node.arg // isz))
-    addrs.append(dctx.kernargs_gpu + off)
-  return _maybe_mstack(tuple(addrs)).after(*_lower_stores(_maybe_mstack(tuple(hbufs)), buf_node, target.src[1:]))
-
-def bufferize_cmdbuf(ctx:HCQ2LowerCtx, target:UOp, buf_node:UOp) -> UOp:
-  hbufs = tuple(UOp.from_buffer(Buffer("CPU", buf_node.arg // dtypes.uint32.itemsize, dtypes.uint32,
-                                       options=BufferSpec(cpu_access=True, nolru=True), preallocate=True), dev)
-                for dev in _devices(buf_node.src[1].arg))
-  hbuf = _maybe_mstack(hbufs)
-  return hbuf.after(*_lower_stores(hbuf, buf_node, target.src[1:]), tag=buf_node.tag)
-
-def bufferize_binary(ctx:HCQ2LowerCtx, target:UOp, buf_node:UOp) -> UOp|None:
-  if buf_node.tag == "program": return bufferize_program(ctx, target, buf_node)
-  if buf_node.tag == "kernargs": return bufferize_kernargs(ctx, target, buf_node)
-  if buf_node.tag in ("compute", "copy"): return bufferize_cmdbuf(ctx, target, buf_node)
-  return None
-
-# TODO: merge with bufferize_binary
-def resolve_buffer(b:UOp) -> UOp|None:
-  devs = _devices(b.src[1].arg)
-  if b.tag in ("timeline_signal", "timeline_value"):
-    return _maybe_mstack(tuple(UOp.from_buffer(getattr(Device[d], b.tag), d) for d in devs), b.tag)
-  if b.tag == "scratch":
-    return _maybe_mstack(tuple(UOp.from_buffer(Buffer(d, (s:=Device[d].scratch).size, dtypes.uint8, opaque=s, options=BufferSpec(external_ptr=1)), d)
-                               for d in devs), b.tag)
-  if isinstance(b.tag, tuple): # (compute_queue|sdma_queue, ring|write_ptr|doorbell|put_value)
-    return _maybe_mstack(tuple(UOp.from_buffer(getattr(Device[d].compute_queue if b.tag[0] == "compute_queue" else Device[d].sdma_queue(0), b.tag[1]), d)
-                            for d in devs), b.tag)
-  if isinstance(b.device, tuple): return _maybe_mstack(tuple(UOp.from_buffer(buf, buf.device) for buf in b.buffer.bufs))
-  return None
-
-pm_bufferize = PatternMatcher([
-  (UPat(Ops.AFTER, src=(UPat(Ops.BUFFER, name="buf_node"),), allow_any_len=True, name="target"), bufferize_binary),
-  (UPat(Ops.BUFFER, name="b"), resolve_buffer), # TODO: cleanup
+def lift_patches_to_cmdbuf(cmdbuf:UOp) -> UOp|None:
+  if not (patches:=dedup(u for store in cmdbuf.src[1:] for u in store.toposort() if u.op is Ops.AFTER)): return None
+  deps = tuple(d for p in patches for d in p.src[1:])
+  return cmdbuf.replace(src=cmdbuf.src + deps).substitute({p: p.src[0] for p in patches})
+pm_lift_patches_to_cmdbuf = PatternMatcher([
+  (UPat(Ops.AFTER, src=(UPat(Ops.BUFFER, tag={"compute", "copy"}),), allow_any_len=True, name="cmdbuf"), lift_patches_to_cmdbuf),
 ])
 
-def lift_patches_to_cmdbuf(ctx:HCQ2LowerCtx, cmdbuf:UOp) -> UOp|None:
-  if cmdbuf.tag not in ("compute", "copy"): return None
-  patches = dedup(u for store in cmdbuf.src[1:] for u in store.toposort() if u.op is Ops.AFTER)
-  deps = tuple(d for p in patches for d in p.src[1:])
-  return cmdbuf.replace(src=cmdbuf.src+deps, tag=None).substitute({p:p.src[0] for p in patches})
-pm_lift_patches_to_cmdbuf = PatternMatcher([(UPat(Ops.AFTER, name="cmdbuf", allow_any_len=True), lift_patches_to_cmdbuf)])
+# *****************
+# 4. bufferize placeholders: replace placeholders with real buffers.
 
-# resolve patches
+def bufferize_buf(buf:UOp) -> UOp|None:
+  if buf.tag is None: return None
+  d = buf.src[1].arg
+  uops = tuple(UOp.from_buffer((d:=Device[dev]).pm_bufferize.rewrite(buf, ctx=d), dev) for dev in (d if isinstance(d, tuple) else (d,)))
+  return uops[0] if len(uops) == 1 else UOp(Ops.MSTACK, uops[0].dtype, uops)
+pm_bufferize = PatternMatcher([(UPat(Ops.BUFFER, name="buf"), bufferize_buf)])
 
-def fold_const_store(ctx:HCQ2LowerCtx, buf:UOp, off:UOp, val:UOp) -> UOp:
-  bufs = buf.src if buf.op is Ops.MSTACK else (buf,)
-  vals = val.src if val.op is Ops.MSTACK else (val,) * len(bufs)
-  for b, v in zip(bufs, vals):
-    struct.pack_into(f'<{v.dtype.fmt}', b.buffer.ensure_allocated()._buf.cpu_view().mv.cast('B'), off.arg * b.dtype.base.itemsize, v.arg)
+# *****************
+# 5.1. capture buffers reachable from each hcq call as BIND, so resolve_patches doesn't drop their refs
+
+def hold_call_buffers(call:UOp) -> UOp|None:
+  if not (bufs:=tuple(dedup(u for u in call.src[0].toposort() if u.op is Ops.BUFFER))): return None
+  return call.replace(src=call.src + (UOp(Ops.BIND, dtypes.void, src=bufs),))
+pm_hold_call_buffers = PatternMatcher([(UPat(Ops.CALL, tag="hcq", name="call"), hold_call_buffers)])
+
+# *****************
+# 5.2. resolve patches
+
+def fold_blob_store(buf:UOp, blob:UOp) -> UOp:
+  buf.buffer.ensure_allocated()._buf.cpu_view().mv.cast('B')[:len(blob.arg)] = blob.arg
   return UOp(Ops.NOOP)
 
-def fold_blob_store(ctx:HCQ2LowerCtx, buf:UOp, blob:UOp) -> UOp:
-  for b in (buf.src if buf.op is Ops.MSTACK else (buf,)):
-    b.buffer.ensure_allocated()._buf.cpu_view().mv.cast('B')[:len(blob.arg)] = blob.arg
+def fold_const_store(buf:UOp, off:UOp, val:UOp) -> UOp:
+  struct.pack_into(f'<{val.dtype.fmt}', buf.buffer.ensure_allocated()._buf.cpu_view().mv.cast('B'), off.arg * buf.dtype.base.itemsize, val.arg)
   return UOp(Ops.NOOP)
 
-def resolve_getaddr(ctx:HCQ2LowerCtx, m:UOp) -> UOp:
-  srcs = m.src if m.op is Ops.MSTACK else (m,)
-  for s in srcs:
-    if s.op in (Ops.BUFFER, Ops.SLICE) and s not in ctx.holds: ctx.holds.append(s)
-  addrs = [s.arg if s.op is Ops.CONST else s.buffer.get_buf(s.device).va_addr for s in srcs]
-
-  # fast-path: all per-dev VAs equal -> just a const
-  if all(v == addrs[0] for v in addrs): return UOp.const(dtypes.uint64, addrs[0])
-
-  table = _maybe_mstack(tuple(UOp.from_buffer(Buffer("CPU", 1, dtypes.uint64, preallocate=True), "CPU") for _ in range(len(srcs))))
-  vas = _maybe_mstack(tuple(UOp.const(dtypes.uint64, va) for va in addrs))
-  patch = table.index(slot_const:=UOp.const(dtypes.int, 0), dtype=table.dtype.ptr()).store(vas)
-  return table.after(patch).index(slot_const, dtype=table.dtype.ptr()).load(dtype=dtypes.uint64)
-
 pm_resolve_patches = symbolic + PatternMatcher([
-  # resolve getaddrs
-  (UPat(Ops.GETADDR, src=(UPat(Ops.SLICE, name="bv"),)), # getaddr(buffer_view(x)) -> offset+getaddr(x)
-    lambda ctx, bv: UOp(Ops.GETADDR, dtypes.uint64, src=(bv.src[0],)) + UOp.const(dtypes.uint64, bv.src[1].arg * bv.src[0].dtype.itemsize)),
-  (UPat(Ops.GETADDR, src=(UPat((Ops.BUFFER, Ops.MSTACK), name="m"),)), resolve_getaddr), # getaddr(buffer|mstack) -> addr_table load|const
-  (UPat(Ops.GETADDR, src=(UPat.cvar("const"),)), lambda ctx, const: const), # getaddr(const) -> const
+  (UPat(Ops.GETADDR, src=(UPat(Ops.SLICE, name="bv"), UPat(Ops.DEVICE, name="dev"))), # getaddr(slice(x)) -> offset+getaddr(x)
+    lambda bv, dev: UOp(Ops.GETADDR, dtypes.uint64, src=(bv.src[0], dev)) + UOp.const(dtypes.uint64, bv.src[1].arg * bv.src[0].dtype.itemsize)),
+  (UPat(Ops.GETADDR, src=(UPat(Ops.BUFFER, name="buf"), UPat(Ops.DEVICE)), name="g"),
+    lambda buf, g: UOp.const(dtypes.uint64, buf.buffer.get_buf(g.src[1].arg).va_addr)),
+  (UPat(Ops.GETADDR, src=(UPat.cvar("const"), UPat())), lambda const: const),
 
-  # write consts and binaries directly into the buffer (BUFFER or MSTACK of BUFFERs)
-  (UPat((Ops.BUFFER, Ops.SLICE, Ops.MSTACK), name="buf").store(UPat(Ops.BINARY, name="blob")), fold_blob_store),
-  (UPat((Ops.BUFFER, Ops.SLICE, Ops.MSTACK), name="buf").index(UPat.cvar("off")).or_casted()
-    .store(UPat.any(UPat.cvar("val"), UPat(Ops.MSTACK, src=UPat.cvar(), name="val"))), fold_const_store),
+  (UPat(Ops.BUFFER, name="buf").store(UPat(Ops.BINARY, name="blob")), fold_blob_store),
+  (UPat(Ops.BUFFER, name="buf").index(UPat.cvar("off")).or_casted().store(UPat.cvar("val")), fold_const_store),
 ])
 
-def parametrize_host_buffer(ctx:HCQ2LowerCtx, buf:UOp) -> UOp:
-  # register a host buffer as a launcher input and return its placeholder
-  if buf.op is Ops.AFTER:
-    p = parametrize_host_buffer(ctx, buf.src[0])
-    return p.after(*(s.substitute({buf.src[0]: p}) for s in buf.src[1:]))
-  if (b:=buf.buffer) not in ctx.inputs: ctx.inputs.append(b)
-  return UOp.placeholder((b.size,), b.dtype, ctx.inputs.index(b))
+# *****************
+# 6. callify hcq programs
 
-pm_parametrize_host_buffers = PatternMatcher([
-  # resolve buffer views to parametrize only root buffers
-  (UPat(Ops.INDEX, src=(UPat(Ops.SLICE, name="bv"), UPat.var("idx")), name="bi"),
-    lambda bv, idx, bi: bi.replace(src=(bv.src[0], idx * bv.dtype.itemsize // bv.src[0].dtype.itemsize + bv.src[1].arg))),
-
-  # parametrize host buffers
-  (UPat(Ops.AFTER, src=(UPat((Ops.BUFFER, Ops.SLICE, Ops.MSTACK)),), allow_any_len=True, name="buf"), parametrize_host_buffer),
-  (UPat((Ops.BUFFER, Ops.SLICE, Ops.MSTACK), name="buf"), parametrize_host_buffer),
-
-  # remove UNIQUE/DEVICE to dedup CONST
+pm_fixup = PatternMatcher([ # TODO: this should gone?
   (UPat(Ops.CONST, name="c"), lambda c: c.replace(src=()) if len(c.src) else None),
 ])
 
-def hcq_callify(ctx:HCQ2LowerCtx, l:UOp) -> UOp:
-  sink = UOp.sink(*l.src, arg=KernelInfo(name=ctx.name, estimates=Estimates()), tag=1)
-  inputs = [UOp.from_buffer(b, tuple(x.device for x in b.bufs) if isinstance(b, MultiBuffer) else "CPU") for b in ctx.inputs]
-  call = to_program(sink, Device["CPU"].renderer).call(*inputs)
-  return call.replace(src=call.src + (UOp(Ops.BIND, dtypes.void, src=tuple(ctx.holds)),)) if ctx.holds else call
-pm_callify = PatternMatcher([(UPat(Ops.LINEAR, name="l", allow_any_len=True), hcq_callify)])
+def callify_hcq(call:UOp) -> UOp:
+  bufs = dedup([u for u in call.src[0].toposort() if u.op is Ops.BUFFER])
+  subs = {b: UOp.placeholder((b.arg,), b.dtype, i) for i, b in enumerate(bufs)}
+  sink = UOp.sink(call.src[0].substitute(subs), arg=KernelInfo(name="hcq_submit", estimates=Estimates()), tag=1)
+  return to_program(sink, Device["CPU"].renderer).call(*bufs, UOp(Ops.BIND, dtypes.void, src=call.src[1:]))
+pm_callify_hcq = PatternMatcher([(UPat(Ops.CALL, tag="hcq", name="call"), callify_hcq)])
 
-# **************** schedule ****************
+@track_rewrites(lambda _,ret: f"HCQ Schedule {pluralize('Kernel', len(ret.src))}")
+def hcq_schedule(linear:UOp) -> UOp:
+  linear = graph_rewrite(linear, pm_insert_copy_staging + pm_flatten_linear, name="insert copy staging")
+  linear = graph_rewrite(linear, pm_prep_runtime, name="prepare runtime")
 
-@track_rewrites(name=lambda linear,ast,**kw: f"hcq schedule {getattr(ast.arg, 'name', ast.op.name.lower())}")
-def hcq_schedule(linear:UOp, ast:UOp) -> UOp:
-  # runtime preparation: device-specific program, kernargs for each program
-  linear = graph_rewrite(linear, pm_prep_runtime, name="hcq: prepare runtime")
+  linear = graph_rewrite(linear, pm_lower_ops, name="lower ops into hcq ir")
+  # linear = graph_rewrite(linear, pm_split_into_queues, name="split into queues")
+  linear = graph_rewrite(linear, pm_add_barriers, walk=True, name="add barriers")
+  linear = graph_rewrite(linear, pm_add_signals, walk=True, name="add signals")
+  linear = graph_rewrite(linear, pm_encode_cmdbufs, walk=True, name="encode cmdbufs")
+  linear = graph_rewrite(linear, pm_add_timeline_inc, walk=True, name="add timeline inc")
+  linear = graph_rewrite(linear, pm_lift_patches_to_cmdbuf, name="lift patches to cmdbuf", enter_calls=True)
 
-  # lower ops into hcq style per-device operations
-  linear = graph_rewrite(linear, pm_lower_ops, name="hcq: lower ops")
+  # realize starts from here
+  linear = graph_rewrite(linear, pm_bufferize, bottom_up=True, name="bufferize placeholders", enter_calls=True)
+  linear = graph_rewrite(linear, pm_hold_call_buffers, walk=True, name="hold call buffers")
+  linear = graph_rewrite(linear, pm_resolve_patches, bottom_up=False, name="simplify patches", enter_calls=True)
+  linear = graph_rewrite(linear, pm_fixup, bottom_up=False, name="fixup", enter_calls=True)
+  linear = graph_rewrite(linear, pm_callify_hcq, name="parametrize host buffers")
 
-  # split ops into logical queues
-  linear = graph_rewrite(linear, pm_split_into_queues, name="hcq: split into queues")
-
-  # runtime-specific lowering
-  linear = graph_rewrite(linear, pm_add_barriers, walk=True, name="hcq: add barriers")
-  linear = graph_rewrite(linear, pm_add_signals, walk=True, name="hcq: add signals")
-  linear = graph_rewrite(linear, pm_add_timeline_inc, walk=True, name="hcq: add submit")
-
-  # encode cmdbuffers + submits
-  # TODO: remove dev
-  dev = Device["AMD"]
-  return graph_rewrite(linear, dev.pm_lower, walk=True, name="hcq: encode cmdbuf")
-
-@track_rewrites(name=lambda ctx,linear,ast,**kw: f"hcq realize {getattr(ast.arg, 'name', ast.op.name.lower())}")
-def hcq_realize(ctx:HCQ2LowerCtx, linear:UOp, ast:UOp) -> UOp:
-  # allocate lowering structs
-  graph_rewrite(linear, pm_calc_kernargs_sizes, ctx=(sizes:={}), name=None)
-
-  for dev_name, sz in sizes.items():
-    dev = Device[dev_name]
-    off = dev.kernargs_offset_allocator.alloc(sz, 16)
-    ctx.dev_ctx[dev_name] = HCQ2DeviceCtx(dev_name, UOp.from_buffer(dev.kernargs_buf.view(sz, dtypes.uint8, off), dev_name),
-                                          UOp.const(dtypes.uint64, dev.kernargs_buf.get_buf(dev_name).va_addr + off, device=dev_name))
-
-  linear = graph_rewrite(linear, pm_bufferize, ctx=ctx, bottom_up=True, name="realize binaries")
-  linear = graph_rewrite(linear, pm_lift_patches_to_cmdbuf, ctx=ctx, bottom_up=False, name="lift patches to cmdbuf")
-  linear = graph_rewrite(linear, pm_resolve_patches, ctx=ctx, bottom_up=False, name="simplify patches")
-  linear = graph_rewrite(linear, pm_parametrize_host_buffers, ctx=ctx, bottom_up=True, name="parametrize host buffers")
-  return graph_rewrite(linear, pm_callify, ctx=ctx, name="hcq: callify")
-
-def ensure_accessible(ctx:HCQ2LowerCtx, call:UOp, copy:UOp) -> UOp|None:
-  src_buf = call.src[2].buffer # TODO: cleanup
-  dev = call.src[1].buffer.device
-  try: src_buf.get_buf(dev)
-  except Exception:
-    (cpubuf := Buffer("CPU", src_buf.nbytes, dtypes.uint8, preallocate=True)).copyin(src_buf.ensure_allocated().as_memoryview())
-    ctx.holds.append(buf_uop:=UOp.from_buffer(cpubuf, dev))
-    return call.replace(src=call.src[:2] + (buf_uop,) + call.src[3:])
-pm_ensure_bufs_accessible = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.COPY, name="copy"),), name="call", allow_any_len=True), ensure_accessible)])
-
-def hcq_exec(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
-  from tinygrad.engine.realize import run_linear
-
-  if ast.src[1].arg.split(":")[0] != "AMD": return None
-
-  # TODO: this mess should gone
-  resolved_call = call.replace(src=(ast,) + tuple(resolve_params(call, ctx.input_uops)) + tuple(s for s in call.src[1:] if s.op is Ops.BIND))
-  bufs = [cast(Buffer, resolved_call.src[1+gi].buffer) for gi in ast.arg.globals] if ast.op is Ops.PROGRAM \
-    else [cast(Buffer, resolved_call.src[i].buffer) for i in range(1, len(resolved_call.src))]
-  hcq_ctx = HCQ2LowerCtx(name="submit")
-  linear = graph_rewrite(UOp(Ops.LINEAR, dtypes.void, (resolved_call,)), pm_ensure_bufs_accessible, ctx=hcq_ctx)
-
-  linear = hcq_schedule(linear, ast)
-  host_call = hcq_realize(hcq_ctx, linear, ast)
-
-  dev = Device["AMD"]
-  with track_stats(ctx, call, dev.device, bufs, ctx.var_vals) as tm:
-    st = time.perf_counter() if ctx.wait else 0.0
-    run_linear(UOp(Ops.LINEAR, dtypes.void, (host_call,)), var_vals=ctx.var_vals, jit=True, update_stats=DEBUG>=3)
-    if ctx.wait:
-      dev.synchronize()
-      tm[0] = time.perf_counter() - st
-  return tm[0] if tm[0] is not None else 0.0
-
-pm_hcq_exec = PatternMatcher([
-  (UPat(Ops.CALL, src=(UPat({Ops.PROGRAM, Ops.COPY}, name="ast"),), name="call", allow_any_len=True), hcq_exec),
-])
+  return linear

extra/hcq2/ops_amd2.py

@@ -163,7 +163,8 @@ def amd_submit_pm4(cmdbuf, devs):
   ring_idx = ((put + i.cast(put.dtype)) % q.ring.size).cast(dtypes.int)
 
   # copy the cmdbuf into the ring and advance the put/write pointers
-  copy_to_ring = ring.index(ring_idx, dtype=ring.dtype.ptr()).store(cmdbuf.index(i, dtype=dtypes.uint32)).end(i)
+  copy_to_ring = ring.index(ring_idx, dtype=ring.dtype.ptr()).store(
+    cmdbuf.index(i*4, dtype=cmdbuf.dtype.ptr()).cast(dtypes.uint32.ptr()).load()).end(i)
   bump_put_ptr = put_ptr.index(zero, dtype=put_ptr.dtype.ptr()).store(next_put)
   bump_wptr = wptr.index(zero, dtype=wptr.dtype.ptr()).store(next_put)
 
@@ -233,7 +234,8 @@ def amd_submit_sdma(cmdbuf, devs):
   zi = UOp.range(zero_amt_dw, 0, dtype=dtypes.int, src=(cmdbuf,))
   zero_tail = ring.index(tail_off_dw + zi, dtype=ring.dtype.ptr()).store(UOp.const(dtypes.uint32, 0)).end(zi)
   i = UOp.range(UOp.const(dtypes.int, size_dw), 0, dtype=dtypes.int, src=(cmdbuf,))
-  copy_to_ring = ring.index(start_dw + i, dtype=ring.dtype.ptr()).store(cmdbuf.index(i, dtype=dtypes.uint32)).end(i)
+  copy_to_ring = ring.index(start_dw + i, dtype=ring.dtype.ptr()).store(
+    cmdbuf.index(i*4, dtype=cmdbuf.dtype.ptr()).cast(dtypes.uint32.ptr()).load()).end(i)
 
   # advance the put/write pointers past the zeroed tail and the cmdbuf
   next_put_b = put_b + ((zero_amt_dw + size_dw) * 4).cast(put_b.dtype)
@@ -253,8 +255,7 @@ class AMDProgramData:
 _amd_program_cache:dict[tuple[bytes,str], tuple[AMDProgramData,bytes]] = {}
 
 def amd_build_program(prg:UOp) -> UOp:
-  devs = prg.src[1].arg  # tuple[str, ...] from rebind_program_dev
-  dev = Device[devs[0]]
+  dev = Device[prg.src[1].arg] # TODO: rm this
   if (cached:=_amd_program_cache.get(key:=(lib:=prg.src[4].arg, dev.device))) is None:
     image, sections, relocs = elf_loader(lib)
     rodata = next(sh.header.sh_addr for sh in sections if sh.name == ".rodata")
@@ -274,10 +275,11 @@ def amd_build_program(prg:UOp) -> UOp:
       kernargs_alloc_size=desc.kernarg_size + (ctypes.sizeof(hsa.hsa_kernel_dispatch_packet_t) if edp else 0),
       enable_dispatch_ptr=edp,
       enable_private_segment_sgpr=desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER), bytes(image))
-  data, image_bytes = cached
-  buf_uop = UOp.new_buffer(devs, len(image_bytes), dtypes.uint8).rtag("program")
-  blob_uop = UOp(Ops.BINARY, dtypes.void, src=(), arg=image_bytes)
-  return prg.replace(src=(buf_uop.after(buf_uop.store(blob_uop)),), arg=(data, prg.arg))
+  return cached
+
+pm_prep_program = PatternMatcher([
+  (UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.DEVICE, arg="AMD"), UPat(), UPat(), UPat(Ops.BINARY)), name="prg"), amd_build_program),
+])
 
 class AMDAllocator(HCQAllocator['AMDDevice']):
   def __init__(self, dev:AMDDevice):
@@ -373,17 +375,16 @@ def _mock(iface, name=None): return type(name or f"MOCK{iface.__name__}", (iface
 
 def encode_queue(q:UOp) -> UOp|None:
   if not (isinstance(q.arg, tuple) and len(q.arg) == 2 and q.arg[1] in ("COMPUTE", "COPY")): return None
-  devs = q.arg[0]
+  devs = (q.arg[0],) if isinstance(q.arg[0], str) else q.arg[0]  # TODO: make this prettier
   return amd_submit_pm4(amd_lower_pm4(q, devs), devs) if q.arg[1] == "COMPUTE" else amd_submit_sdma(amd_lower_sdma(q, devs), devs)
 
+pm_lower = PatternMatcher([
+  (UPat(Ops.LINEAR, name="q"), encode_queue),
+])
+
 class AMDDevice(HCQ2Compiled):
   timestamp_divider = 100.0  # AMD GPU clock: ticks/us
 
-  pm_lower = PatternMatcher([
-    (UPat(Ops.PROGRAM, src=(UPat(), UPat(), UPat(), UPat(), UPat(Ops.BINARY)), name="prg"), amd_build_program),
-    (UPat(Ops.LINEAR, name="q"), encode_queue),
-  ])
-
   ifaces = [PCIIface]
 
   def is_am(self) -> bool: return isinstance(self.iface, (PCIIface,))
@@ -424,8 +425,7 @@ class AMDDevice(HCQ2Compiled):
     self.sdma_queues:dict = {}
     self.has_sdma_queue = self.sdma_queue(0) is not None
 
-    super().__init__(device, AMDAllocator(self), [HIPRenderer, AMDLLVMRenderer, HIPCCRenderer], None,
-                     kernargs_size=16 << 20, can_recover=self.is_am(), arch=self.arch)
+    super().__init__(device, AMDAllocator(self), [HIPRenderer, AMDLLVMRenderer, HIPCCRenderer], None, can_recover=self.is_am(), arch=self.arch)
 
     # Scratch setup
     self.max_private_segment_size = 0
@@ -530,4 +530,15 @@ class AMDDevice(HCQ2Compiled):
 
   def on_device_hang(self): self.iface.on_device_hang()
 
+  @functools.cached_property
+  def pm_bufferize(self) -> PatternMatcher:
+    return PatternMatcher([
+      (UPat(Ops.BUFFER, tag="scratch"),
+        lambda ctx: Buffer(ctx.device, ctx.scratch.size, dtypes.uint8, opaque=ctx.scratch, options=BufferSpec(external_ptr=1))),
+      (UPat(Ops.BUFFER, tag={("compute_queue", n) for n in ("ring", "write_ptr", "doorbell", "put_value")}, name="b"),
+        lambda ctx, b: getattr(ctx.compute_queue, b.tag[1])),
+      (UPat(Ops.BUFFER, tag={("sdma_queue", n) for n in ("ring", "write_ptr", "doorbell", "put_value")}, name="b"),
+        lambda ctx, b: getattr(ctx.sdma_queue(0), b.tag[1])),
+    ]) + super().pm_bufferize
+
   def device_props(self): return self.iface.props

extra/llama_kernels/cast_amax/__init__.py

@@ -11,13 +11,13 @@ from extra.llama_kernels import FP8_MAX, NUM_WG, THREADS_PER_WG, compile_cpp, al
 _grad_fp8_mailbox:dict[UOp, tuple[UOp, UOp]] = {}
 
 @functools.cache
-def _custom_fused_bwd_w13(grad_xw13:UOp, grad_xw13_fp8:UOp, grad_amax_buf:UOp,
+def _custom_fused_bwd_w13(grad_xw13_fp8:UOp, grad_amax_buf:UOp,
                           xw13:UOp, grad_x2:UOp, amax_state:UOp, grad_amax_state:UOp, dname:str) -> UOp:
   hidden = xw13.shape[2] // 2
   n_elems = xw13.shape[0] * xw13.shape[1] * hidden
   threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
-  mem = n_elems * 2 * 5 + n_elems * 2 + NUM_WG * 4 + 4
-  sink = UOp.sink(grad_xw13.base, grad_xw13_fp8.base, grad_amax_buf.base,
+  mem = n_elems * 2 * 3 + n_elems * 2 + NUM_WG * 4 + 4
+  sink = UOp.sink(grad_xw13_fp8.base, grad_amax_buf.base,
                   xw13.base, grad_x2.base, amax_state.base, grad_amax_state.base, threads, workgroups,
                   arg=KernelInfo(f"fused_silu_mul_bwd_w13_{n_elems}", estimates=Estimates(ops=10*n_elems, mem=mem)))
   src, lib = compile_cpp(pathlib.Path(__file__).parent, "cast_amax_bwd_w13.cpp", n_elems, hidden)
@@ -41,23 +41,23 @@ def _fused_quantize_bwd_w13(gradient:UOp, kernel:UOp):
   _, _, xw13, amax_state, grad_amax_state = kernel.src[1:]
   device = xw13.device
   axis = xw13.axis if isinstance(device, tuple) else None
-  grad_xw13     = alloc_like(xw13.shape, dtypes.bfloat16, device, axis)
   grad_xw13_fp8 = alloc_like(xw13.shape, dtypes.fp8e4m3,  device, axis)
   grad_amax_buf = alloc_local((NUM_WG,), dtypes.float32,  device, axis)
   grad_amax_state_t = Tensor(grad_amax_state, device=device)
   fxn = functools.partial(_custom_fused_bwd_w13, dname=dname_of(device))
-  grad_xw13, grad_xw13_fp8, grad_amax_buf, *_ = Tensor.custom_kernel(
-    grad_xw13, grad_xw13_fp8, grad_amax_buf,
+  grad_xw13_fp8, grad_amax_buf, *_ = Tensor.custom_kernel(
+    grad_xw13_fp8, grad_amax_buf,
     Tensor(xw13, device=device), Tensor(gradient, device=device).cast(dtypes.bfloat16),
     Tensor(amax_state, device=device), grad_amax_state_t, fxn=fxn)
+  grad_xw13_uop = grad_xw13_fp8.uop.cast(dtypes.bfloat16)
   inv_scale = (grad_amax_state_t.float() + 1e-8) / FP8_MAX
   new_grad_amax = scalar_amax(grad_amax_buf)
   store_effect = grad_amax_state_t.uop.store(new_grad_amax.uop)
   assert grad_xw13_fp8.uop.op is Ops.AFTER, f"expected AFTER, got {grad_xw13_fp8.uop.op}"
   grad_xw13_fp8_uop = grad_xw13_fp8.uop.replace(src=grad_xw13_fp8.uop.src + (store_effect,))
   # Stash fp8 companion for cdna_asm_gemm's bwd to attach to grad_a.
-  _grad_fp8_mailbox[grad_xw13.uop] = (grad_xw13_fp8_uop, inv_scale.uop)
-  return (None, None, grad_xw13.uop, None, None)
+  _grad_fp8_mailbox[grad_xw13_uop] = (grad_xw13_fp8_uop, inv_scale.uop)
+  return (None, None, grad_xw13_uop, None, None)
 
 def fused_quantize_fp8_w13(xw13:Tensor, amax_state:Tensor, fp8_dtype, grad_amax_state:Tensor) -> tuple[Tensor, Tensor, Tensor]:
   # NOTE: silu(xw1)*xw3 -> fp8 + amax over fused xw13 layout. Returns (fp8, inv_scale, new_amax)

extra/llama_kernels/cast_amax/cast_amax_bwd_w13.cpp

@@ -21,15 +21,13 @@ constexpr float FP8_MAX = 448.0f;
 static_assert(N_ELEMS % VEC == 0, "N_ELEMS must be divisible by VEC");
 static_assert(HIDDEN % VEC == 0, "HIDDEN must be divisible by VEC");
 
-// fused silu*mul backward, three outputs in a single HBM pass:
-//   1) bf16 grad_xw13      — consumed by downstream bf16 autograd chain
-//   2) fp8  grad_xw13_fp8  — delayed-scale quantize using grad_amax_state (mailbox to matmul bwd)
-//   3) fp32 grad_amax_buf  — per-WG partial |grad_xw13|, reduced into next step's grad_amax_state
+// fused silu*mul backward, two outputs in a single HBM pass:
+//   1) fp8  grad_xw13_fp8  — delayed-scale quantize using grad_amax_state (mailbox to matmul bwd)
+//   2) fp32 grad_amax_buf  — per-WG partial |grad_xw13|, reduced into next step's grad_amax_state
 // grad_amax_state is read for the fp8 scale. The store of new_grad_amax into grad_amax_state's
 // buffer is built in Python as a separate effect and threaded into grad_a via .after(store).
 extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
 fused_silu_mul_bwd_w13(
-    __hip_bfloat16*       __restrict__ grad_xw13_out,        // bf16, 2*N_ELEMS
     __hip_fp8_storage_t*  __restrict__ grad_xw13_fp8_out,    // fp8,  2*N_ELEMS
     float*                __restrict__ grad_amax_buf,        // fp32, NUM_WG per-WG partials
     const __hip_bfloat16* __restrict__ xw13,                 // bf16, 2*N_ELEMS
@@ -62,7 +60,6 @@ fused_silu_mul_bwd_w13(
     const __hip_bfloat16 *x3 = reinterpret_cast<const __hip_bfloat16*>(&x3_raw);
     const __hip_bfloat16 *gv = reinterpret_cast<const __hip_bfloat16*>(&g_raw);
 
-    __hip_bfloat16 out1[VEC], out3[VEC];
     __hip_fp8_storage_t fp8_1[VEC], fp8_3[VEC];
     #pragma unroll
     for (int i = 0; i < VEC; i++) {
@@ -75,15 +72,11 @@ fused_silu_mul_bwd_w13(
       const float gs = fg * scale;
       const float g1 = gs * silu_prime * f3;
       const float g3 = gs * silu;
-      out1[i] = static_cast<__hip_bfloat16>(g1);
-      out3[i] = static_cast<__hip_bfloat16>(g3);
       local_max = fmaxf(local_max, fmaxf(fabsf(g1), fabsf(g3)));
       fp8_1[i] = __hip_cvt_float_to_fp8(fmaxf(-FP8_MAX, fminf(FP8_MAX, g1 * g_scale)), __HIP_SATFINITE, __HIP_E4M3);
       fp8_3[i] = __hip_cvt_float_to_fp8(fmaxf(-FP8_MAX, fminf(FP8_MAX, g3 * g_scale)), __HIP_SATFINITE, __HIP_E4M3);
     }
 
-    *reinterpret_cast<float4*>(&grad_xw13_out[xw1_off]) = *reinterpret_cast<float4*>(out1);
-    *reinterpret_cast<float4*>(&grad_xw13_out[xw3_off]) = *reinterpret_cast<float4*>(out3);
     *reinterpret_cast<uint64_t*>(&grad_xw13_fp8_out[xw1_off]) = *reinterpret_cast<uint64_t*>(fp8_1);
     *reinterpret_cast<uint64_t*>(&grad_xw13_fp8_out[xw3_off]) = *reinterpret_cast<uint64_t*>(fp8_3);
   }

extra/torch_backend/backend.py

@@ -23,7 +23,7 @@ def calculate_storage_offset(x: Tensor) -> int:
   for u in x.uop.toposort():
     if u.op == Ops.SHRINK:
       u_strides = strides_for_shape(u.src[0].shape)
-      for i, (start, _) in enumerate(u.marg): offset += start * u_strides[i]
+      for i, (_, start) in enumerate(u.marg): offset += start * u_strides[i]
   return offset
 def wrap(x: Tensor) -> torch.Tensor:
   x._strides = strides_for_shape(x.shape) # always recalculate

spec/tinyspec.tex

@@ -74,12 +74,13 @@ A \op{Buffer}'s \textbf{addrspace} is \texttt{GLOBAL}, \texttt{LOCAL}, or \textt
 \op{Flip} & $(T,)$ & bools $\mathbf{f}$ & Reverse along flagged axes. \\
 \op{Reshape} & $(T, \mathbf{s'})$ & --- & Reinterpret in row-major order. $\prod s_k = \prod s'_k$. \\
 \op{Expand} & $(T, \mathbf{s'})$ & --- & Broadcast size-1 axes. $s_k \in \{1, s'_k\}$. \\
-\op{Pad} & $(T, \mathbf{b}, \mathbf{e})$ & --- & Pad with $0$s: $b_k$ before, $e_k$ after each axis. \\
-\op{Shrink} & $(T, \mathbf{b}, \mathbf{e})$ & --- & Keep $[b_k, e_k)$ per axis. Inverse of \op{Pad}. \\
+\op{Pad} & $(T, \mathbf{s'}, \mathbf{o})$ & --- & Place $T$ at offset $o_k$ in a zero-filled output of shape $s'_k$. \\
+\op{Shrink} & $(T, \mathbf{s'}, \mathbf{o})$ & --- & Keep $s'_k$ elements starting at offset $o_k$ per axis. Inverse of \op{Pad}. \\
 \op{Index} & $(T, i_0, i_1, \ldots)$ & --- & Index from left. $()$-shaped $i$ removes dim; $(k,)$-shaped makes it $k$. \\
 \op{Stack} & $(T_0, T_1, \ldots)$ & --- & Join along a newly created leading axis. All shapes must match. \\
 \op{Replicated} & $(T,)$ & axes & Mark $T$ as replicated along axes. Collapse axes to $1$. \\
 \op{Slice} & $(T, \mathrm{offset})$ & size, dtype & Zero-copy \textit{size} elems of dtype; offset is elems of $T$ dtype. \\
+\op{Bitcast} & $(T,)$ & dtype & Reinterpret storage as target dtype; preserve total bytes. \\
 \bottomrule
 \end{tabular}
 
@@ -165,8 +166,7 @@ Unary & $(T,)$
   & $\mathrm{trunc}(x)$: round toward zero. \\
 & & \op{Cast}
   & Convert to target dtype (specified in arg). \\
-& & \op{Bitcast}
-  & Reinterpret bits as target dtype. Must be same size. \\[4pt]
+  \\[4pt]
 Binary & $(A, B)$
   & \op{Add}, \op{Mul}, \op{Max}, \op{Mod}, \op{Idiv}
   & $a+b$, $a \cdot b$, $\max(a,b)$, $a \bmod b$, $\lfloor a/b \rfloor$ \\

test/backend/test_custom_kernel.py

@@ -155,7 +155,7 @@ class TestCustomKernel(unittest.TestCase):
     self.assertTrue((ref == tst).all().item())
 
   def test_eye(self):
-    ref = Tensor.eye(1024).contiguous().realize()
+    ref = Tensor.eye(1024).clone().realize()
     tst = Tensor.empty_like(ref)
     tst = tst.custom_kernel(fxn=custom_eye_kernel)[0]
     self.assertTrue((ref == tst).all().item())
@@ -335,7 +335,7 @@ class TestCustomKernel(unittest.TestCase):
     assert all(x == expected for x in result), f"expected all {expected}, got {result}"
 
   def test_custom_kernel_sched(self, use_custom=False):
-    x = Tensor.arange(32).reshape(8, 4).realize()
+    x = Tensor.arange(32).reshape(8, 4).clone().realize()
     y = Tensor.empty_like(x)
     y = Tensor.custom_kernel(y, x, fxn=custom_add_one_kernel)[0]
     if use_custom:
@@ -352,7 +352,7 @@ class TestCustomKernel(unittest.TestCase):
 
   @unittest.expectedFailure
   def test_sliced_buffer_function(self):
-    x = Tensor.arange(32).reshape(8, 4).realize()
+    x = Tensor.arange(32).reshape(8, 4).clone().realize()
     from tinygrad import function
     @function(precompile=True)
     def run(x:Tensor) -> Tensor:

test/backend/test_dtype_alu.py

@@ -111,6 +111,7 @@ def universal_test_cast(a, in_dtype, dtype):
 def universal_test_midcast(a, b, c, op1, op2, d1:DType, d2:DType):
   if not isinstance(op1, tuple): op1 = (op1, op1)
   if not isinstance(op2, tuple): op2 = (op2, op2)
+  if op1[0] == operator.mod and b == 0: return
   # lt and max with nan is undefined in tinygrad
   if op1[0] in (operator.lt, Tensor.maximum) and (math.isnan(a) or math.isnan(b)): return
   if op2[0] in (operator.lt, Tensor.maximum) and math.isnan(c): return

test/backend/test_jit.py

@@ -49,7 +49,7 @@ class TestJit(unittest.TestCase):
       y = (x + 1).contiguous().realize()
       z = x.shrink(((st, st + N),)).contiguous().realize()
       return y, z
-    x = Tensor.arange(2*N).contiguous().realize()
+    x = Tensor.arange(2*N).clone().realize()
     for _ in range(3): y, z = f(x, Variable("a", 0, N).bind(0))
     self.assertEqual(y.shape, (2*N,))
     self.assertEqual(z.shape, (N,))
@@ -92,7 +92,7 @@ class TestJit(unittest.TestCase):
     @TinyJit
     def f(x): return (x[2:5].contiguous() + 1).realize()
     for i in range(5):
-      x = (Tensor.arange(10).float() + i * 10).contiguous().realize()
+      x = (Tensor.arange(10).float() + i * 10).clone().realize()
       np.testing.assert_allclose(f(x).numpy(), x.numpy()[2:5] + 1)
 
   def test_jit_multiple_outputs(self):

test/backend/test_kernel_cache.py

@@ -8,10 +8,10 @@ class TestKernelCache(unittest.TestCase):
     if Device.DEFAULT not in ["CPU"]:
       self.skipTest("No custom kernel cache is implemented")
 
-    unique_const = 0.6765677269
+    const_value = 0.6765677269
     a = Tensor.rand(4,4).realize()
     b = Tensor.rand(4,4).realize()
-    x = a + b + unique_const
+    x = a + b + const_value
     x.realize()
 
     a1 = Tensor.rand(4,4).realize()
@@ -20,7 +20,7 @@ class TestKernelCache(unittest.TestCase):
     Device['CPU'].compiler.compile_cached = None # making it not callable
 
     try:
-      x1 = a1 + b1 + unique_const
+      x1 = a1 + b1 + const_value
       x1.realize() # Same kernel should be from cache.
     finally:
       Device['CPU'].compiler.compile_cached = orig_compile_func

test/backend/test_multitensor.py

@@ -187,7 +187,7 @@ class TestMultiTensor(unittest.TestCase):
     np.testing.assert_allclose(O.numpy(), X.numpy()[0:2]*W.numpy()[0:2] < 2)
 
   def test_shrink_on_shard_axis(self):
-    X = Tensor.arange(4*4).reshape(4,4).realize()
+    X = Tensor.arange(4*4).reshape(4,4).clone().realize()
     X_np = X.numpy()
     X.shard_(devices_2, 0)
     # only shrink on the device that owns the shard, this is enabled by the mselect simplifier
@@ -293,7 +293,7 @@ class TestMultiTensor(unittest.TestCase):
     @TinyJit
     def f(x): return (x+1).contiguous().sum()
     for _ in range(5):
-      tt = Tensor.arange(0, 4).contiguous().realize().shard((d1,d2), 0).realize()
+      tt = Tensor.arange(0, 4).clone().realize().shard((d1,d2), 0).realize()
       out = f(tt)
       assert out.item() == 1+2+3+4
 
@@ -309,7 +309,7 @@ class TestMultiTensor(unittest.TestCase):
     @TinyJit
     def f(x): return (x.shard((d1,d2), 0)+1).contiguous().sum()
     for _ in range(5):
-      tt = Tensor.arange(0, 4).contiguous().realize()
+      tt = Tensor.arange(0, 4).clone().realize()
       out = f(tt)
       assert out.item() == 1+2+3+4
 
@@ -865,7 +865,7 @@ class TestMultiTensor(unittest.TestCase):
   @unittest.skip("RANGEIFY doesn't support multi const folding")
   def test_multi_const_folding(self):
     with Context(TRACK_MATCH_STATS=0):
-      a = Tensor.arange(3).realize()
+      a = Tensor.arange(3).clone().realize()
       zeros = Tensor.zeros(3).realize()
     b = a.to(devices_2)*zeros.to(devices_2)
     sched = b.schedule_linear().src
@@ -904,7 +904,7 @@ class TestShrinkMultiTensorShardedAxis(unittest.TestCase):
 
   # shrink a multitensor on sharded axis
   def test_shrink_bad_args(self):
-    t = Tensor.arange(64).reshape(8, 8).contiguous().realize()
+    t = Tensor.arange(64).reshape(8, 8).clone().realize()
     t.shard_([f"{Device.DEFAULT}:{i}" for i in range(4)], axis=0)
 
     with self.assertRaises(AssertionError):
@@ -927,7 +927,7 @@ class TestShrinkMultiTensorShardedAxis(unittest.TestCase):
   @given(strat.sampled_from([dtypes.float, dtypes.int, dtypes.int64, dtypes.int16]))
   def test_ops(self, dtype):
     if dtype not in Device[Device.DEFAULT].renderer.supported_dtypes(): return
-    t = Tensor.arange(64).reshape(8, 8).contiguous().realize()
+    t = Tensor.arange(64).reshape(8, 8).clone().realize()
     t.shard_([f"{Device.DEFAULT}:{i}" for i in range(4)], axis=0)
     for i in range(4):
       print(f"{i=}")
@@ -971,7 +971,7 @@ class TestShrinkMultiTensorShardedAxis(unittest.TestCase):
       np.testing.assert_allclose(a.flip(-1).numpy(), b.flip(-1).numpy(), rtol=1e-7, atol=1e-3)
 
   def test_add_two_partitions(self):
-    t = Tensor.arange(64).reshape(8, 8).contiguous().realize()
+    t = Tensor.arange(64).reshape(8, 8).clone().realize()
     t.shard_([f"{Device.DEFAULT}:{i}" for i in range(4)], axis=0)
 
     a = t.shrink(((2, 4), None))
@@ -988,7 +988,7 @@ class TestShrinkMultiTensorShardedAxis(unittest.TestCase):
 
   def test_add_different_tensors(self):
     devices = [f"{Device.DEFAULT}:{i}" for i in range(4)]
-    x = Tensor.arange(64).reshape(8, 8).contiguous().realize().shard(devices, axis=0)
+    x = Tensor.arange(64).reshape(8, 8).clone().realize().shard(devices, axis=0)
 
     to_add = []
     for i in range(len(devices)):
@@ -1098,7 +1098,7 @@ class TestBatchNorm(unittest.TestCase):
   @given(strat.sampled_from((False, True)))
   def test_batchnorm(self, is_training):
     devices = [f"{Device.DEFAULT}:{i}" for i in range(4)]
-    x = Tensor.arange(4096).reshape(8, 8, 8, 8).contiguous().realize().shard(devices, axis=0)
+    x = Tensor.arange(4096).reshape(8, 8, 8, 8).clone().realize().shard(devices, axis=0)
 
     with Tensor.train(is_training):
       bns = []
@@ -1184,7 +1184,7 @@ class TestMultiBufferView(unittest.TestCase):
 
   @unittest.skip("flaky on LLVM")
   def test_shrink_non_shard_axis(self):
-    ref = Tensor.arange(8*4*10).reshape(8, 4, 10).contiguous().realize()
+    ref = Tensor.arange(8*4*10).reshape(8, 4, 10).clone().realize()
     a = Tensor.arange(8*4*10).reshape(8, 4, 10).clone().shard(devices_2, axis=1).realize()
     self._check(ref, a, lambda t: t[3])
 
@@ -1296,7 +1296,7 @@ class TestMultiSetitem(unittest.TestCase):
   @needs_second_gpu
   def setUp(self): pass
 
-  def _t(self, axis): return Tensor.arange(16).contiguous().realize().shard(self.device, axis=axis)
+  def _t(self, axis): return Tensor.arange(16).clone().realize().shard(self.device, axis=axis)
 
   def test_setitem_scalar_axis0(self):
     t = self._t(0)

test/backend/test_pickle.py

@@ -67,7 +67,7 @@ class TestPickle(unittest.TestCase):
 
   # NOTE: currently Buffer exists on the uop, not tensor
   def test_pickle_buffer_uop(self):
-    t = Tensor.arange(4).realize()
+    t = Tensor.arange(4).clone().realize()
     a = t.uop
     assert a.is_realized
     self.assertIsNotNone(buffer:=a.base.realized)
@@ -95,7 +95,7 @@ class TestPickle(unittest.TestCase):
     np.testing.assert_equal(vt2.numpy(), 20)
 
   def test_pickle_buffer_view(self):
-    t = Tensor.arange(10, device="CPU").contiguous().realize()
+    t = Tensor.arange(10).clone(device="CPU").realize()
     vt = t[3:5].contiguous().realize()
     assert hasattr(vt.uop.buffer, 'base')
     ref_value = vt.tolist()

test/backend/test_schedule.py

@@ -221,8 +221,8 @@ class TestSchedule(unittest.TestCase):
     np.testing.assert_allclose(out.numpy(), (x.numpy() - x.numpy().max(keepdims=True)).max())
 
   def test_example_matmul_contig(self):
-    x = Tensor.eye(64).contiguous().realize()
-    y = Tensor.eye(64).contiguous().realize()
+    x = Tensor.eye(64).clone().realize()
+    y = Tensor.eye(64).clone().realize()
     z = y.matmul(x).sum()
     z.backward()
     out = x.grad.contiguous()
@@ -826,7 +826,7 @@ class TestSchedule(unittest.TestCase):
     self._test_fusion([(32, 32)], lambda a:a-a.sum(1), 2)
 
   def test_cast_padded_view(self):
-    a = Tensor.arange(4).reshape(1, 4)
+    a = Tensor.arange(4).reshape(1, 4).clone().realize()
     casted_view = a.pad(((0, 1), (0, 0))).cast(dtypes.float)
     casted_view.realize()
     self.assertEqual(casted_view.uop.base.realized.size, 8)
@@ -836,7 +836,7 @@ class TestSchedule(unittest.TestCase):
 
   # NOTE: we only reorder CAST if it's an EXPAND
   def test_cast_after_shrink(self):
-    a = Tensor.arange(4).reshape(1, 4)
+    a = Tensor.arange(4).reshape(1, 4).clone().realize()
     casted_view = a.shrink(((0, 1), (0, 2))).cast(dtypes.float)
     casted_view.realize()
     self.assertEqual(casted_view.uop.base.realized.size, 2)
@@ -991,7 +991,7 @@ class TestSchedule(unittest.TestCase):
 
   def test_assign_non_contiguous_alt(self): self.test_assign_non_contiguous(alt=True)
   def test_assign_non_contiguous(self, alt=False):
-    x = (Tensor.arange(16)-100).reshape(4,4).contiguous().realize()
+    x = (Tensor.arange(16)-100).reshape(4,4).clone().realize()
     xref = x.numpy()
     if alt:
       y = Tensor.randint(2, 4).contiguous().realize()
@@ -1007,7 +1007,7 @@ class TestSchedule(unittest.TestCase):
     np.testing.assert_equal(tst.numpy(), a.numpy())
 
   def test_setitem_sched(self, mop=lambda x:x, expected_kcount=1):
-    a = Tensor.arange(16, device="CPU").reshape(4, 4).contiguous().realize()
+    a = Tensor.arange(16).reshape(4, 4).clone(device="CPU").realize()
     a2 = mop(a)
     expected = (a+a2).tolist()
     a.assign(a+a2)
@@ -1021,7 +1021,7 @@ class TestSchedule(unittest.TestCase):
 
   def test_setitem_const_fused(self):
     # https://github.com/tinygrad/tinygrad/issues/10690
-    a = Tensor.arange(16).contiguous().realize()
+    a = Tensor.arange(16).clone().realize()
     GlobalCounters.reset()
     a[4] = 3
     self.assertEqual(GlobalCounters.kernel_count, 0)
@@ -1278,7 +1278,7 @@ class TestCopyFolding(unittest.TestCase):
     self.assertEqual(x.item(), 2.0)
 
   def test_late_const_copy_folding(self):
-    a = Tensor.arange(3).realize()
+    a = Tensor.arange(3).clone().realize()
     zeros = Tensor.zeros(3, buffer=False).realize()
     b = (a*zeros).to("CPU") + 1
     run_linear(*check_schedule(b, 1, filter_sink=False))
@@ -1353,14 +1353,14 @@ class TestCopyFolding(unittest.TestCase):
     self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])
 
   def test_permute_on_disk(self):
-    with open(temp('dt_arange_4_permute'), "wb") as f: f.write(Tensor.arange(4).realize().uop.base.buffer.as_memoryview())
+    with open(temp('dt_arange_4_permute'), "wb") as f: f.write(Tensor.arange(4).clone().realize().uop.base.buffer.as_memoryview())
     a = Tensor.empty(4, dtype=dtypes.int32, device=f"disk:{temp('dt_arange_4_permute')}")
     b = a.reshape(2, 2).permute(1, 0).to("CPU")
     b.realize()
     self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])
 
   def test_permute_on_disk_contiguous(self):
-    with open(temp('dt_arange_4_permute_contig'), "wb") as f: f.write(Tensor.arange(4).realize().uop.base.buffer.as_memoryview())
+    with open(temp('dt_arange_4_permute_contig'), "wb") as f: f.write(Tensor.arange(4).clone().realize().uop.base.buffer.as_memoryview())
     a = Tensor.empty(4, dtype=dtypes.int32, device=f"disk:{temp('dt_arange_4_permute_contig')}")
     b = a.reshape(2, 2).permute(1, 0).contiguous().to("CPU")
     b.realize()
@@ -1374,7 +1374,7 @@ class TestCopyFolding(unittest.TestCase):
 
   # NOTE: disk permute must come after COPY
   def test_permute_after_shrink_on_disk(self):
-    with open(temp('dt_arange_5_permute'), "wb") as f: f.write(Tensor.arange(5).realize().uop.base.buffer.as_memoryview())
+    with open(temp('dt_arange_5_permute'), "wb") as f: f.write(Tensor.arange(5).clone().realize().uop.base.buffer.as_memoryview())
     a = Tensor.empty(5, dtype=dtypes.int32, device=f"disk:{temp('dt_arange_5_permute')}")
     b = a.shrink(((0, 4),)).reshape(2, 2).permute(1, 0).to("CPU")
     b.realize()

test/backend/test_setitem.py

@@ -32,13 +32,14 @@ class TestSetitem(unittest.TestCase):
     self.assertListEqual(t.tolist(), [0, 1, 11, 3, 11, 5, 6, 7, 8, 9])
 
   def test_setitem_inplace_mul(self):
-    t = Tensor.arange(10).realize()
+    t = Tensor.arange(10).clone().realize()
     t[:3] *= 10
     self.assertListEqual(t.tolist(), [0, 10, 20, 3, 4, 5, 6, 7, 8, 9])
 
+  @unittest.skip("crashed in LLVM CI")
   def test_setitem_fancy_on_unrealized_view(self):
     # fancy indexing setitem on unrealized SHRINK view (triggered infinite loop in graph_rewrite)
-    base = Tensor.arange(20, dtype=dtypes.float).reshape(4, 5)
+    base = Tensor.arange(20, dtype=dtypes.float).reshape(4, 5).clone().realize()
     sub = base[1:3]
     flat = sub.reshape(sub.numel()).contiguous()
     idx = Tensor([0, 3, 7, 9])
@@ -229,7 +230,7 @@ class TestSetitem(unittest.TestCase):
     np.testing.assert_equal(t.numpy(), n)
 
   def test_setitem_swap_rows(self):
-    t = Tensor.arange(6, dtype=dtypes.float).reshape(3, 2).contiguous().realize()
+    t = Tensor.arange(6, dtype=dtypes.float).reshape(3, 2).clone().realize()
     tmp = t[0]
     t[0] = t[1]
     t[2] = tmp
@@ -237,7 +238,7 @@ class TestSetitem(unittest.TestCase):
     np.testing.assert_allclose(t.numpy(), [[2, 3], [2, 3], [2, 3]])
 
     # eager version
-    t = Tensor.arange(6, dtype=dtypes.float).reshape(3, 2).contiguous().realize()
+    t = Tensor.arange(6, dtype=dtypes.float).reshape(3, 2).clone().realize()
     tmp = t[0].realize()
     t[0] = t[1].realize()
     t[2] = tmp.realize()
@@ -269,8 +270,8 @@ class TestSetitem(unittest.TestCase):
   def test_cross_assign_independence(self):
     # when assigning to two tensors using computations from both,
     # both assigns should see the OLD values of both tensors
-    a = Tensor.arange(4, dtype=dtypes.float).contiguous().realize()
-    b = Tensor.arange(4, 8, dtype=dtypes.float).contiguous().realize()
+    a = Tensor.arange(4, dtype=dtypes.float).clone().realize()
+    b = Tensor.arange(4, 8, dtype=dtypes.float).clone().realize()
     new_a = a + b    # [4, 6, 8, 10]
     new_b = a * 2    # [0, 2, 4, 6] -- should use OLD a
     a.assign(new_a)
@@ -283,8 +284,8 @@ class TestSetitem(unittest.TestCase):
       np.testing.assert_allclose(b.numpy(), [8, 12, 16, 20])
 
     # eager version
-    a = Tensor.arange(4, dtype=dtypes.float).contiguous().realize()
-    b = Tensor.arange(4, 8, dtype=dtypes.float).contiguous().realize()
+    a = Tensor.arange(4, dtype=dtypes.float).clone().realize()
+    b = Tensor.arange(4, 8, dtype=dtypes.float).clone().realize()
     new_a = (a + b).realize()
     new_b = (a * 2).realize()
     a.assign(new_a).realize()
@@ -322,8 +323,8 @@ class TestWithGrad(unittest.TestCase):
 
   def test_set_overlapping_backward(self):
     z = Tensor.zeros(6)
-    x = Tensor.ones(4).contiguous()
-    y = Tensor.ones(4).contiguous() * 2
+    x = Tensor.ones(4)
+    y = Tensor.ones(4) * 2
     z[:4] = x
     z[2:] = y
     z.sum().backward()

test/backend/test_subbuffer.py

@@ -36,7 +36,7 @@ class TestSubBuffer(unittest.TestCase):
     assert len(mv) == 5
 
   def test_subbuffer_used(self):
-    t = Tensor.arange(0, 10, dtype=dtypes.uint8).realize()
+    t = Tensor.arange(0, 10, dtype=dtypes.uint8).clone().realize()
     vt = t[2:4].realize()
     out = (vt + 100).tolist()
     assert out == [102, 103]
@@ -44,7 +44,7 @@ class TestSubBuffer(unittest.TestCase):
   @needs_second_gpu
   @unittest.skipIf(Device.DEFAULT not in {"CUDA", "NV", "AMD"} or DEV.interface.startswith("MOCK"), "only NV, AMD, CUDA")
   def test_subbuffer_transfer(self):
-    t = Tensor.arange(0, 10, dtype=dtypes.uint8).realize()
+    t = Tensor.arange(0, 10, dtype=dtypes.uint8).clone().realize()
     vt = t[2:5].contiguous().realize()
     out = vt.to(f"{Device.DEFAULT}:1").realize().tolist()
     assert out == [2, 3, 4]

test/backend/test_tensor.py

@@ -551,7 +551,7 @@ class TestTinygrad(unittest.TestCase):
       Tensor.zeros(2, 2).realize()
 
   def test_shrink(self):
-    t = Tensor.arange(32).contiguous().realize()
+    t = Tensor.arange(32).clone().realize()
     self.assertListEqual(t[16:20].tolist(), [16,17,18,19])
     self.assertListEqual(t.shrink_to(16).tolist(), list(range(16)))
     t = t.reshape(4, 8).contiguous().realize()

test/external/external_uop_gc.py

@@ -29,7 +29,7 @@ def gradient_test():
   z = y.matmul(x).sum()
   z.backward()
 def realized_eye():
-  Tensor.eye(3).realize()
+  Tensor.eye(3).clone().realize()
 def realized_list():
   Tensor([[2.0,0,-2.0]]).realize()
 def kernel_matmul():

test/null/test_const_folding.py

@@ -134,14 +134,14 @@ class TestBitcastConstFolding(unittest.TestCase):
 # folds advance indexing into basic indexing
 class TestIndexingConstFolding(unittest.TestCase):
   def test_scalar_index(self):
-    t = Tensor.arange(16).float().reshape(1,1,4,4).realize()
+    t = Tensor.arange(16).float().reshape(1,1,4,4).clone().realize()
     _check_ast_count(1, t[:,:,Tensor(1),:])
     _check_ast_count(1, t[:,:,Tensor(1)+2,:])
     _check_ast_count(1, t[:,:,Tensor(1),Tensor(0)])
 
   def test_const_tensor_index(self):
     # TODO: these can be 0, implement const tensor folded indexing
-    t = Tensor.arange(16).float().reshape(1,1,4,4).realize()
+    t = Tensor.arange(16).float().reshape(1,1,4,4).clone().realize()
     _check_ast_count(1, t[:,:,Tensor.ones(2,1,dtype=dtypes.int),:])
     _check_ast_count(1, t[:,:,Tensor.ones(1,2,dtype=dtypes.int)+2,:])
     _check_ast_count(1, t[:,:,Tensor.ones(1,1,dtype=dtypes.int),Tensor.zeros(2,1,2,dtype=dtypes.int)])

test/null/test_schedule.py

@@ -456,14 +456,14 @@ class TestSchedule(unittest.TestCase):
       check_schedule(out, 4, [c1.weight, c1.bias, *nn.state.get_parameters(bn)])
 
   def test_fold_conv_batchnorm_optim(self, adam=False):
-    optim, cnt = (nn.optim.Adam, 29) if adam else (nn.optim.SGD, 15)
+    # 2 is too low?
+    optim, cnt = (nn.optim.Adam, 16) if adam else (nn.optim.SGD, 2)
     with Tensor.train():
-      img = Tensor.ones(1,3,4,4).realize()
+      img = Tensor.ones(1,3,4,4)
       c1 = nn.Conv2d(3,32,3)
       bn = nn.BatchNorm2d(32, track_running_stats=False)
       _realize_weights([c1, bn])
       opt = optim(nn.state.get_parameters([c1, bn]))
-      Tensor.realize(*nn.state.get_parameters(opt))
       img_bn = bn(c1(img)).elu().sum()
       opt.zero_grad()
       img_bn.backward()
@@ -477,7 +477,7 @@ class TestSchedule(unittest.TestCase):
       fw = bn(x).contiguous_backward().relu().contiguous()
       fw.sum().backward()
       # TODO: this is too many
-      check_schedule([x.grad, bn.weight.grad, bn.bias.grad, fw], 10, nn.state.get_parameters(bn))
+      check_schedule([x.grad, bn.weight.grad, bn.bias.grad, fw], 9)
 
   def test_fold_conv_relu(self):
     c1 = nn.Conv2d(3,16,3)

test/null/test_tensor.py

@@ -73,10 +73,10 @@ class TestIdxUpcast(unittest.TestCase):
     # Assert the dtype of the INDEX value, This will need be updated if UOp spec changes
     store = next(uop for uop in uops if uop.op is Ops.STORE)
     assert store.op is Ops.STORE
-    idx = self._find_op(store, Ops.SLICE)
-    # PTX and NIR turn Ops.SLICE into pointer arithmetic earlier than cstyle, plus it's already cast to int64
+    idx = self._find_op(store, Ops.INDEX)
+    # PTX and NIR turn Ops.INDEX into pointer arithmetic earlier than cstyle, plus it's already cast to int64
     if not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)):
-      assert idx.op is Ops.SLICE
+      assert idx.op is Ops.INDEX
       idx_val = idx.src[1]
       self.assertIs(idx_val.dtype, dtype)
 

test/null/test_tensor_uop_mixin.py

@@ -3,7 +3,7 @@ from tinygrad import Tensor, dtypes
 from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, graph_rewrite
 
 _strip_unique_pm = PatternMatcher([
-  (UPat(Ops.BUFFER, src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE, name="d")), name="b"), lambda b,d: b.replace(src=(UOp.unique(0), d))),
+  (UPat((Ops.UNIQUE, Ops.LUNIQUE), name="u"), lambda u: u.replace(arg=0) if u.arg != 0 else None),
 ])
 def _strip_unique(u: UOp) -> UOp: return graph_rewrite(u, _strip_unique_pm)
 

test/null/test_uop_graph.py

@@ -538,7 +538,7 @@ class TestUOpGraph(unittest.TestCase):
     uops = to_uops_list([UOp(Ops.STORE, dtypes.void, (glbl0.index(idx), ld1+ld0))])
     ld0 = uops[-2].src[-1]  # -2 to skip SINK
     # the gate and invalid value are deleted from ld1
-    self.assertEqual(ld0, UOp.load(glbl2.slice(idx), dtype=dtypes.int))
+    self.assertEqual(ld0, UOp.load(glbl2.index(idx, ptr=True), dtype=dtypes.int))
 
   def test_fold_gated_load_local(self):
     glbl0 = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0)
@@ -552,9 +552,7 @@ class TestUOpGraph(unittest.TestCase):
 
     ld0 = uops[-2].src[-1]  # -2 to skip SINK
     # the gate and invalid value are deleted from ld1
-    new_barrier = ld0.src[0].src[0].src[1]
-    assert new_barrier.op is Ops.BARRIER
-    self.assertEqual(ld0.src[0], smem.after(new_barrier).slice(lidx+2))
+    self.assertEqual(ld0.src[0], smem.after(barrier).index(lidx+2, ptr=True))
 
   def test_fold_gated_store(self):
     glbl = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0)
@@ -566,7 +564,7 @@ class TestUOpGraph(unittest.TestCase):
     uops = to_uops_list([st0, st1])
     # only the second store happens
     self.assertEqual(len(uops), 6)  # +1 for SINK
-    self.assertEqual(uops[-2], glbl.slice(idx1).store(val))  # -2 to skip SINK
+    self.assertEqual(uops[-2], glbl.index(idx1, ptr=True).store(val))  # -2 to skip SINK
 
   @unittest.skip("this is a uop type error")
   def test_asserts_bad_gate(self):

test/unit/test_assign.py

@@ -281,15 +281,15 @@ class TestAssign(unittest.TestCase):
     np.testing.assert_equal(t.numpy(), [[100, 104, 108, 112], [101, 105, 109, 113], [102, 106, 110, 114], [103, 107, 111, 115]])
 
   def test_assign_contiguous(self):
-    b = Tensor.arange(16).reshape(4,4).contiguous().realize()
-    a = (Tensor.arange(16).reshape(4,4).contiguous().realize() + 1)
+    b = Tensor.arange(16).reshape(4,4).clone().realize()
+    a = (Tensor.arange(16).reshape(4,4).clone().realize() + 1)
     GlobalCounters.reset()
     b.assign(a.contiguous()).realize()
     self.assertEqual(GlobalCounters.kernel_count, 2)
 
   def test_assign_contiguous_permute(self):
-    b = Tensor.arange(16).reshape(4,4).contiguous().realize()
-    a = (Tensor.arange(16).reshape(4,4).contiguous().realize() + 1).permute((1,0))
+    b = Tensor.arange(16).reshape(4,4).clone().realize()
+    a = (Tensor.arange(16).reshape(4,4).clone().realize() + 1).permute((1,0))
     GlobalCounters.reset()
     b.assign(a.contiguous()).realize()
     self.assertEqual(GlobalCounters.kernel_count, 2)
@@ -325,29 +325,29 @@ class TestAssign(unittest.TestCase):
     np.testing.assert_allclose(a.numpy(), np.arange(N*N).reshape((N,N)) + np.arange(N*N).reshape((N,N)).transpose(1,0))
 
   def test_post_permuted_assignment_alt(self):
-    a = Tensor.arange(N*N).reshape(N,N).contiguous().realize()
-    b = Tensor.arange(N*N).reshape(N,N).contiguous().realize()
+    a = Tensor.arange(N*N).reshape(N,N).clone().realize()
+    b = Tensor.arange(N*N).reshape(N,N).clone().realize()
     new_a = (a.T+b).numpy()
     a.assign(a.T+b)
     np.testing.assert_allclose(a.numpy(), new_a)
 
   def test_post_flipped_assignment(self):
-    a = Tensor.arange(N*N).reshape(N,N).contiguous().realize()
-    b = Tensor.arange(N*N).reshape(N,N).contiguous().realize()
+    a = Tensor.arange(N*N).reshape(N,N).clone().realize()
+    b = Tensor.arange(N*N).reshape(N,N).clone().realize()
     new_a = (a.flip(0)+b).numpy()
     a.assign(a.flip(0)+b)
     np.testing.assert_allclose(a.numpy(), new_a)
 
   def test_post_flipped_assignment_axis1(self):
-    a = Tensor.arange(N*N).reshape(N,N).contiguous().realize()
-    b = Tensor.arange(N*N).reshape(N,N).contiguous().realize()
+    a = Tensor.arange(N*N).reshape(N,N).clone().realize()
+    b = Tensor.arange(N*N).reshape(N,N).clone().realize()
     new_a = (a.flip(1)+b).numpy()
     a.assign(a.flip(1)+b)
     np.testing.assert_allclose(a.numpy(), new_a)
 
   def test_post_reshape_assignment_fine(self):
-    a = Tensor.arange(N*N).reshape(N, N).contiguous().realize()
-    b = Tensor.arange(N*N).reshape(N, N).contiguous().realize()
+    a = Tensor.arange(N*N).reshape(N, N).clone().realize()
+    b = Tensor.arange(N*N).reshape(N, N).clone().realize()
     rhs = a.reshape(-1).reshape(N, N)
     new_a = (rhs+b).numpy()
     a.assign(rhs+b)  # self-assign with reshape view is fine
@@ -355,7 +355,7 @@ class TestAssign(unittest.TestCase):
 
   @unittest.skip("multi output not supported anymore")
   def test_simple_assignment_multioutput(self):
-    a = Tensor.arange(32*32).reshape(32, 32).contiguous().realize()
+    a = Tensor.arange(32*32).reshape(32, 32).clone().realize()
     b = Tensor.full((32, ), 1.).contiguous().realize()
     c = Tensor.full((32, ), 2.).contiguous().realize()
     d = Tensor.full((32, ), 3.).contiguous().realize()
@@ -375,15 +375,15 @@ class TestAssign(unittest.TestCase):
   # NOTE: if the assign target is read/write in a single kernel, it should be contiguous
 
   def test_permuted_assignment_correct(self):
-    a = Tensor.arange(4 * 4).reshape(4, 4).contiguous().realize()
-    b = Tensor.arange(4 * 4).reshape(4, 4).contiguous().realize()
+    a = Tensor.arange(4 * 4).reshape(4, 4).clone().realize()
+    b = Tensor.arange(4 * 4).reshape(4, 4).clone().realize()
     a = a.permute(1, 0)
     new_val = a + b
     a.assign(new_val)
     np.testing.assert_equal(a.numpy(), np.arange(4 * 4).reshape(4, 4).transpose(1, 0) + np.arange(4 * 4).reshape(4, 4))
 
   def test_permuted_reduceop_child_dual_use(self):
-    a = Tensor.arange(32*32*32).reshape(32, 32, 32).contiguous().realize()
+    a = Tensor.arange(32*32*32).reshape(32, 32, 32).clone().realize()
     b = Tensor.ones(32, 32, dtype=dtypes.int).contiguous().realize()
     r = a.sum(axis=1)
     b.assign(r + b.permute(1, 0))
@@ -392,7 +392,7 @@ class TestAssign(unittest.TestCase):
 
   @unittest.skip("multi output not supported anymore")
   def test_permuted_reduceop_multioutput_dual_use(self):
-    a = Tensor.arange(32*32*32).reshape(32, 32, 32).contiguous().realize()
+    a = Tensor.arange(32*32*32).reshape(32, 32, 32).clone().realize()
     b = Tensor.full((32, 32), 1.).contiguous().realize()
     c = Tensor.full((32, 32), 2.).contiguous().realize()
 
@@ -405,9 +405,9 @@ class TestAssign(unittest.TestCase):
 
   @unittest.skip("multi output not supported anymore")
   def test_permuted_reduceop_multioutput_dual_use_possible(self):
-    a = Tensor.arange(32*32*32).reshape(32, 32, 32).contiguous().realize()
-    b = Tensor.arange(32 * 32).reshape(32, 32).realize()
-    c = Tensor.arange(32 * 32).reshape(32, 32).realize()
+    a = Tensor.arange(32*32*32).reshape(32, 32, 32).clone().realize()
+    b = Tensor.arange(32 * 32).reshape(32, 32).clone().realize()
+    c = Tensor.arange(32 * 32).reshape(32, 32).clone().realize()
 
     GlobalCounters.reset()
     r = a.sum(axis=1)
@@ -441,7 +441,7 @@ class TestAssign(unittest.TestCase):
     # Forward shift: read index > write index in overlap
     N = 100000
     shift = 1000
-    a = Tensor.arange(N).float().contiguous().realize()
+    a = Tensor.arange(N).float().clone().realize()
     expected = np.arange(N, dtype=np.float32)
     expected[:N-shift] = expected[shift:].copy()
     with Context(NOOPT=1): a[0:N-shift].assign(a[shift:N]).realize()
@@ -451,7 +451,7 @@ class TestAssign(unittest.TestCase):
     # Reverse shift: write index > read index in overlap
     N = 100000
     shift = 1000
-    a = Tensor.arange(N).float().contiguous().realize()
+    a = Tensor.arange(N).float().clone().realize()
     expected = np.arange(N, dtype=np.float32)
     expected[shift:] = expected[:N-shift].copy()
     with Context(NOOPT=1): a[shift:N].assign(a[0:N-shift]).realize()
@@ -459,7 +459,7 @@ class TestAssign(unittest.TestCase):
 
   def test_nonoverlapping_shrink_assignment(self):
     # TODO: non-overlapping shrinks don't actually need contiguous, could be 1 kernel with smarter range analysis
-    a = Tensor.arange(100).float().contiguous().realize()
+    a = Tensor.arange(100).float().clone().realize()
     expected = np.arange(100, dtype=np.float32)
     expected[0:10] = expected[50:60].copy()
     GlobalCounters.reset()

test/unit/test_call.py

@@ -222,7 +222,7 @@ class TestCallSchedule(unittest.TestCase):
     # find the FUNCTION nodes
     c0 = next(u for u in r0.uop.toposort() if u.op is Ops.FUNCTION)
     c1 = next(u for u in r1.uop.toposort() if u.op is Ops.FUNCTION)
-    # the function bodies (src[0]) should have identical keys — local buffer identity must not leak through
+    # the function bodies (src[0]) should have identical keys — unique consts must not leak through
     self.assertEqual(c0.src[0].key, c1.src[0].key)
 
   def test_precompile_symbolic_2d(self):

test/unit/test_function.py

@@ -495,6 +495,25 @@ class TestFunctionTuple(unittest.TestCase):
     Tensor.realize(a.grad)
     np.testing.assert_allclose(a.grad.numpy(), [2., 2., 2., 2.])
 
+  def test_custom_kernel_precompile_multidevice(self):
+    # a custom_kernel output placeholder (invalids) under multi-device @function(precompile=True) must return the
+    # kernel's computed result. read it back through .numpy() so the cross-device gather reads the output buffer
+    devs = ("CPU:0", "CPU:1")
+    def double_kernel(C:UOp, A:UOp) -> UOp:
+      C, A = C.flatten(), A.flatten()
+      i = UOp.range(A.numel(), 0)
+      return C[i].store(A[i] * 2.0).end(i).sink(arg=KernelInfo(name="double_kernel"))
+    def double_grad(d_c:UOp, call:UOp): return (None, (Tensor(d_c) * 2.0).uop)
+
+    @function(precompile=True, precompile_backward=True)
+    def f(a:Tensor):
+      c = Tensor(Tensor.invalids(a.shape[0]//len(devs), a.shape[1], dtype=a.dtype, device=devs).uop.multi(0), device=devs)
+      return Tensor.custom_kernel(c, a, fxn=double_kernel, grad_fxn=double_grad)[0]
+
+    a = Tensor.full((4, 4), 7.0).contiguous().shard(devs, axis=0)
+    Tensor.realize(a)
+    np.testing.assert_allclose(f(a).numpy(), 14.0)
+
   def test_custom_kernel_precompile_further_compute(self):
     def my_kernel(C:UOp, A:UOp) -> UOp:
       i = UOp.range(A.shape[0], 0)

test/unit/test_realize_is_realize.py

@@ -15,10 +15,6 @@ class TestRealizeIsRealized(unittest.TestCase):
     t = Tensor.zeros(10).contiguous().realize()
     assert t.uop.is_realized
 
-  def test_ones(self):
-    t = Tensor.ones(4, 4).realize()
-    assert t.uop.is_realized
-
   def test_bytes(self):
     t = Tensor(b'\x01\x02\x03').realize()
     assert t.uop.is_realized
@@ -55,6 +51,10 @@ class TestRealizeIsRealized(unittest.TestCase):
     t = Tensor(3.14).realize()
     assert not t.uop.is_realized
 
+  def test_ones_not_realized(self):
+    t = Tensor.ones(4, 4).realize()
+    assert not t.uop.is_realized
+
   def test_none_not_realized(self):
     t = Tensor(None).realize()
     assert not t.uop.is_realized

tinygrad/callify.py

@@ -181,6 +181,8 @@ def replace_input_buffer(ctx:AllocCtx, b:UOp):
 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),
+  # remove unique from const. TODO: this is copied in function.py
+  (UPat(Ops.CONST, src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE, name="d")), name="b"), lambda b,d: b.replace(src=(d,))),
 ])
 
 pm_replace_buf = PatternMatcher([

tinygrad/codegen/__init__.py

@@ -115,7 +115,7 @@ pm_linearize_cleanups = PatternMatcher([
   # if statements are not allowed in the graph
   (UPat((Ops.IF, Ops.ENDIF)), lambda: panic(RuntimeError, "if not allowed in graph")),
   # gated STORE becomes IF-STORE-ENDIF. this is the only use of IF-ENDIF
-  (UPat(Ops.STORE, name="u", src=(UPat(), UPat(), UPat(name="gate", dtype=dtypes.bool))),
+  (UPat(Ops.STORE, name="u", src=(UPat(Ops.INDEX).or_casted(), UPat(), UPat(name="gate", dtype=dtypes.bool))),
    lambda u, gate: ((st:=u.replace(src=u.src[0:2])), [mif:=UOp(Ops.IF, src=(gate, u.src[0])), st, UOp(Ops.ENDIF, src=(mif,))]))
 ])
 

tinygrad/codegen/late/devectorizer.py

@@ -283,12 +283,6 @@ pm_render = PatternMatcher([
   (UPat(Ops.GEP, name='gep'), lambda gep: UOp(Ops.STACK, gep.dtype, tuple(gep.src[0].gep(x) for x in gep.arg)) if len(gep.arg) > 1 else None),
   (UPat(Ops.GEP, name='gep'), lambda gep: gep.src[0] if gep.src[0].dtype.vcount == 1 and gep.arg == (0,) else None),
   (UPat(Ops.STACK, src=(UPat(name='x'),)), lambda x: x),
-  # rewrite non-image INDEX to SLICE
-  (UPat(Ops.INDEX, name="x"), lambda x: None if isinstance(x.src[0].dtype, ImageDType) else \
-    UOp(Ops.SLICE, dtype=x.dtype, src=x.src, arg=0 if x.dtype.count == 1 else x.dtype.count)),
-  # rewrite CAST on SLICE to just SLICE
-  (UPat(Ops.SLICE, name="bv").cast(name="x"),
-   lambda bv,x: bv.replace(dtype=x.dtype, arg=0 if x.dtype.count == 1 else x.dtype.count))
 ])
 
 # *** Ops.REDUCE -> Ops.DEFINE_ACC ***

tinygrad/engine/realize.py

@@ -241,14 +241,13 @@ pm_exec = PatternMatcher([
   (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="validate", name="ast"),), name="call", allow_any_len=True), exec_validate),
 ])
 
-if getenv("HCQ2"):
-  from extra.hcq2.hcq2 import pm_hcq_exec
-  pm_exec = pm_hcq_exec + pm_exec
-
 def compile_linear(linear:UOp, beam:int|None=None, validate=False) -> UOp:
   if validate: linear = graph_rewrite(linear, pm_validate, name="validate", walk=True)
   if (beam_val:=BEAM.value if beam is None else beam) >= 1: linear = graph_rewrite(linear, pm_beam, ctx=beam_val, walk=True)
   linear = graph_rewrite(linear, pm_compile, name="precompile kernels", walk=True)
+  if getenv("HCQ2"):
+    from extra.hcq2.hcq2 import hcq_schedule
+    linear = hcq_schedule(linear)
   return graph_rewrite(linear, pm_optimize_local_size, name="optimize local size", walk=True)
 
 def run_linear(linear:UOp, var_vals:dict[str, int]|None=None, input_uops:tuple[UOp, ...]=(), update_stats=True, jit=False, wait=False):

tinygrad/function.py

@@ -10,13 +10,17 @@ def add_to_ctx(ctx, x:UOp):
   ctx[0].append(x)
   return ret
 
+pm_transform_unique_const = PatternMatcher([
+  # transform unique consts to LUNIQUE
+  (UPat(Ops.CONST, src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE)), name="x"),
+   lambda ctx,x: x.replace(src=(UOp(Ops.LUNIQUE, arg=next(ctx[1])), x.src[1]))),
+])
+
 pm_ctx = PatternMatcher([
-  (UPat(Ops.BUFFER, src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE)), name="x"),
-   lambda ctx,x: x.replace(src=(UOp(Ops.LUNIQUE, arg=next(ctx[1])), x.src[1])) if x.src[0].arg > ctx[2] else add_to_ctx(ctx,x)),
-  (UPat(Ops.BIND, name="x"), add_to_ctx),
+  (UPat((Ops.BUFFER, Ops.BIND), name="x"), add_to_ctx),
   (UPat((Ops.AFTER, Ops.CONTIGUOUS), name="x"),
    lambda ctx,x: add_to_ctx(ctx,x) if not x.op_in_backward_slice_with_self(Ops.PARAM) and x.op_in_backward_slice_with_self(Ops.BUFFER) else None),
-])
+])+pm_transform_unique_const
 
 ReturnType = TypeVar('ReturnType')
 class _function(Generic[ReturnType]):
@@ -42,7 +46,6 @@ class _function(Generic[ReturnType]):
     # run it and do surgery later
     with Context(ALLOW_DEVICE_USAGE=getenv("DEVICE_IN_FUNCTION_BUG", 0)):
       _function.depth += 1
-      unique_start = next(UOp.unique_num)
       ret = self.fxn(*args, **kwargs)
       _function.depth -= 1
     if isinstance(ret, Tensor):
@@ -62,7 +65,7 @@ class _function(Generic[ReturnType]):
 
     # the BUFFERs that are left are the implicit inputs
     num_explicit = len(call_uops)
-    uret = graph_rewrite(uret, pm_ctx, (call_uops, itertools.count(0), unique_start), bottom_up=True, name="get_implicit_inputs")
+    uret = graph_rewrite(uret, pm_ctx, (call_uops, itertools.count(0)), bottom_up=True, name="get_implicit_inputs")
     name = getattr(self.fxn, '__qualname__', None) or type(self.fxn).__qualname__
     if not self.allow_implicit:
       implicit_buffers = [x for x in call_uops[num_explicit:] if x.op is Ops.BUFFER]

tinygrad/gradient.py

@@ -1,6 +1,6 @@
 from typing import cast
-import math, dataclasses
-from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, all_metadata
+import math, dataclasses, itertools
+from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, all_metadata, graph_rewrite
 from tinygrad.helpers import argsort
 from tinygrad.dtype import sum_acc_dtype
 
@@ -41,6 +41,9 @@ def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]:
   grad_bodies = [(i, grads[p]) for i in needed if (p:=params.get(i)) is not None and p in grads]
   bwd_body = UOp.maketuple(*(gb for _, gb in grad_bodies)).substitute(fwd_subs, walk=True)
   bwd_body, compact_args = _compact_params(bwd_body, (*args, *grad_args, *fwd_outs))
+  # TODO: is this okay here?
+  from tinygrad.function import pm_transform_unique_const
+  bwd_body = graph_rewrite(bwd_body, pm_transform_unique_const, ctx=(None, itertools.count(0)))
   bwd_call = bwd_body.call(*compact_args, name=(k.arg.name or "")+"_backward", precompile=k.arg.precompile_backward)
   gb_map = {i: idx for idx, (i, _) in enumerate(grad_bodies)}
   return (None,) + tuple(bwd_call.gettuple(gb_map[i]) if i in gb_map else None for i in range(len(args)))
@@ -69,8 +72,8 @@ pm_gradient = PatternMatcher([
   (UPat(Ops.EXPAND, name="ret"), lambda ctx, ret:
     (ctx.cast(sum_acc_dtype(ctx.dtype))._rop(Ops.ADD, tuple(i for i,(s,n) in enumerate(zip(ret.src[0].shape, ret.shape)) if s!=n))
      .cast(ctx.dtype), None)),
-  (UPat(Ops.PAD, name="ret"), lambda ctx, ret: (ctx.shrink(tuple([(p[0], s+p[0]) for s,p in zip(ret.src[0].shape, ret.marg)])), None, None)),
-  (UPat(Ops.SHRINK, name="ret"), lambda ctx, ret: (ctx.pad(tuple([(p[0], s-p[1]) for s,p in zip(ret.src[0].shape, ret.marg)])), None, None)),
+  (UPat(Ops.PAD, name="ret"), lambda ctx, ret: (ctx.shrink(tuple([(p[1], s+p[1]) for s,p in zip(ret.src[0].shape, ret.marg)])), None, None)),
+  (UPat(Ops.SHRINK, name="ret"), lambda ctx, ret: (ctx.pad(tuple([(p[1], s-p[0]-p[1]) for s,p in zip(ret.src[0].shape, ret.marg)])), None, None)),
   (UPat(Ops.PERMUTE, name="ret"), lambda ctx, ret: (ctx.permute(argsort(ret.marg)),)),
   (UPat(Ops.FLIP, name="ret"), lambda ctx, ret: (ctx.flip([i for i,x in enumerate(ret.marg) if x]),)),
   (UPat(Ops.COPY, name="ret"), lambda ctx, ret: (ctx.copy_to_device(ret.src[0].device), None)),

tinygrad/mixin/__init__.py

@@ -7,29 +7,28 @@ from tinygrad.mixin.reduce import ReduceMixin
 from tinygrad.uop import Ops
 from tinygrad.uop.ops import _broadcast_shape, resolve, smax, smin, identity_element
 from tinygrad.device import canonicalize_device
-from tinygrad.dtype import ConstType, DType, DTypeLike, InvalidType, PtrDType, PyConst, dtypes, least_upper_dtype, sum_acc_dtype, to_dtype
+from tinygrad.dtype import ConstType, DType, DTypeLike, Invalid, InvalidType, PtrDType, PyConst, dtypes, least_upper_dtype, sum_acc_dtype, to_dtype
 from tinygrad.helpers import all_int, argfix, ceildiv, flatten, flat_to_grouped, make_tuple, prod, resolve_pool_pads, round_up
 
 if TYPE_CHECKING:
-  from tinygrad.uop.ops import sint, UOp
+  from tinygrad.uop.ops import sint
 
 ReductionStr = Literal["mean", "sum", "none"]
 
 
 class OpMixin(ElementwiseMixin, ReduceMixin):
   @staticmethod
-  def empty(*shape, **kwargs): raise NotImplementedError("creation helpers are only supported on Tensor and UOp")
+  def unique_const(fill_value:ConstType, **kwargs): raise NotImplementedError("creation helpers are only supported on Tensor and UOp")
   @staticmethod
   def const(dtype, b, device=None): raise NotImplementedError("creation helpers are only supported on Tensor and UOp")
 
   @classmethod
-  def full(cls, shape:tuple[sint, ...], fill_value:ConstType|UOp, dtype:DTypeLike|None=None,
-           device:str|tuple[str, ...]|None=None, buffer=True) -> Self:
+  def full(cls, shape:tuple[sint, ...], fill_value:ConstType, **kwargs) -> Self:
     """
     Creates a tensor with the given shape, filled with the given value.
 
     You can pass in `dtype` and `device` keyword arguments to control the data type and device of the tensor.
-    Pass `buffer=False` to get a broadcast const value instead of a materialized buffer.
+    Additionally, all other keyword arguments are passed to the constructor of the tensor.
 
     ```python exec="true" source="above" session="tensor" result="python"
     print(Tensor.full((2, 3), 42).numpy())
@@ -39,24 +38,21 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     ```
     """
     new_shape = argfix(shape)
-    dt = to_dtype(dtype) if dtype is not None else None
-    # build the broadcast const value (deviceless for a buffer, device-placed for a value), then clone into storage iff buffer
-    if isinstance(fill_value, get_args(ConstType)):
-      val = cls.const(dt or dtypes.from_py(fill_value), fill_value, None if buffer else canonicalize_device(device))
-    else:  # symbolic UOp fill: keep the value's own dtype, cast only when one is requested
-      val = cls.const(dt, fill_value)
-      if dt is not None: val = val.cast(dt)
-    val = val.reshape((1,)*len(new_shape)).expand(new_shape)
-    return val.clone(device=device) if buffer else val
+    if not kwargs.pop("buffer", True):
+      dt = to_dtype(kwargs.pop("dtype", None) or dtypes.from_py(fill_value))
+      return cls.const(dt, fill_value, canonicalize_device(kwargs.pop("device", None))).reshape((1,)*len(new_shape)).expand(new_shape)
+    return cls.unique_const(fill_value, **kwargs).reshape((1,)*len(new_shape)).expand(new_shape)
 
   @classmethod
   def invalids(cls, *shape, **kwargs) -> Self:
     """
-    Creates an anonymous uninitialized buffer with the given shape.
+    Creates a tensor with the given shape, filled with Invalid.
 
-    You can pass in `dtype` and `device` keyword arguments to control the data type and device of the tensor.
+    This is an alternative to Tensor.empty when you want an "anonymous" buffer.
+
+    Eventually Tensor.empty will be replaced by this.
     """
-    return cls.empty(argfix(*shape), **kwargs)
+    return cls.full(argfix(*shape), Invalid, **kwargs)
 
   @classmethod
   def zeros(cls, *shape, **kwargs) -> Self:

tinygrad/mixin/movement.py

@@ -178,7 +178,7 @@ class MovementMixin:
   def pad(self, arg:tuple[tuple[sint, sint] | None, ...]) -> Self:
     if self.ndim != len(arg):
       raise ValueError(f"{self.ndim=} != {len(arg)=}")
-    ret = self._mop(Ops.PAD, tuple(x if x is not None else (0, 0) for x in arg))
+    ret = self._mop(Ops.PAD, tuple((s+x[0]+x[1], x[0]) if x is not None else (s, 0) for x, s in zip(arg, self.shape)))
     return self if ret.shape == self.shape else ret
 
   def shrink(self, arg: tuple[tuple[sint, sint] | None, ...]) -> Self:
@@ -200,7 +200,7 @@ class MovementMixin:
     """
     if self.ndim != len(arg):
       raise ValueError(f"{self.ndim=} != {len(arg)=}")
-    ret = self._mop(Ops.SHRINK, arg=[x if x is not None else (0, s) for x, s in zip(arg, self.shape)])
+    ret = self._mop(Ops.SHRINK, arg=[(x[1]-x[0], x[0]) if x is not None else (s, 0) for x, s in zip(arg, self.shape)])
     return self if ret.shape == self.shape else ret
 
   def permute(self, order, *args) -> Self:
@@ -251,7 +251,7 @@ class MovementMixin:
     return self.shrink(tuple([None if ns is None else (0, ns) for ns in argfix(shape, *args)]))
 
   def pad_to(self, shape, *args) -> Self:
-    return self._mop(Ops.PAD, tuple([(0, 0 if ns is None else ns-s) for s,ns in zip(self.shape, argfix(shape, *args), strict=True)]))
+    return self._mop(Ops.PAD, tuple((s if ns is None else ns, 0) for s,ns in zip(self.shape, argfix(shape, *args), strict=True)))
 
   def view(self, shape, *args) -> Self:
     """`.view` is an alias for `.reshape`."""

tinygrad/renderer/__init__.py

@@ -29,8 +29,8 @@ class Estimates:
       def range_gate(x): return x.op is not Ops.RANGE
       for u in uops:
         if u.op in {Ops.LOAD, Ops.STORE}:
-          # if u.src[0] is SLICE, we have to include the buffer since it might be an AFTER
-          dont_count = dont_count.union((UOp.sink(*u.src[0].src[1:]) if u.src[0].op is Ops.SLICE else u.src[0]).toposort(range_gate))
+          # if u.src[0] is INDEX, we have to include the buffer since it might be an AFTER
+          dont_count = dont_count.union((UOp.sink(*u.src[0].src[1:]) if u.src[0].op is Ops.INDEX else u.src[0]).toposort(range_gate))
           # TODO: is this correct? this all needs to be cleaned up
           if len(u.src) > 2: dont_count = dont_count.union(u.src[2].toposort())
         elif u.op is Ops.IF:
@@ -40,7 +40,7 @@ class Estimates:
         buf = u
         while len(buf.src): buf = buf.src[0]
         if buf.op is Ops.PARAM:
-          # u.src[0] is SLICE, cap at buffer size for re-reads (e.g. matmul)
+          # u.src[0] is INDEX, cap at buffer size for re-reads (e.g. matmul)
           accessed = mem.get((buf, u.op), 0) + u.src[0].dtype.base.itemsize * mults
           mem[(buf, u.op)] = smin(accessed, buf.ptrdtype.nbytes()) if buf.ptrdtype.size != -1 else accessed
       if u.op is Ops.RANGE:

tinygrad/renderer/cstyle.py

@@ -43,10 +43,8 @@ base_rewrite = PatternMatcher([
   (UPat(Ops.CONST, (dtypes.int8, dtypes.int16), name="x"), lambda ctx,x: f"({ctx.render_cast(x.dtype, str(x.arg))})"),
   # default const render
   (UPat(Ops.CONST, name="x"), lambda ctx,x: str(x.arg)),
-  # slice is ptr arithmetic
-  (UPat(Ops.SLICE, src=(UPat.var("buf"), UPat.var('idx')), name="x"),
-   lambda ctx,buf,idx,x: ctx.render_cast(x.dtype, f"({ctx[buf]}+{strip_parens(ctx[idx]) if idx.arg == Ops.ADD else ctx[idx]})")),
   # new load/store
+  (UPat.var("buf").index(UPat.var('idx')), lambda ctx,buf,idx: f"({ctx[buf]}+{strip_parens(ctx[idx]) if idx.arg == Ops.ADD else ctx[idx]})"),
   (UPat(Ops.LOAD, src=(UPat.var('bidx'),)), lambda ctx,bidx: f"(*{ctx[bidx]})"),
   (UPat(Ops.LOAD, src=(UPat.var("bidx"), UPat.var("var"), UPat.var("gate"))), lambda ctx,bidx,var,gate: f"({ctx[gate]}?*{ctx[bidx]}:{ctx[var]})"),
   (UPat(Ops.STORE, src=(UPat.var('bidx'), UPat.var("var"))), lambda ctx,bidx,var: f"*{ctx[bidx]} = {ctx[var]};"),
@@ -192,14 +190,14 @@ class CStyleLanguage(Renderer):
       else:
         prefix = {Ops.WMMA: "wmma", Ops.DEFINE_LOCAL: "temp", Ops.CONST: "const",
                   Ops.CAST: "cast", Ops.BITCAST: "cast", Ops.GEP: "gep", Ops.STACK: "cast",
-                  Ops.SLICE: "bidx", Ops.DEFINE_REG: "acc", Ops.LOAD: "val"}.get(u.op, "alu")
+                  Ops.INDEX: "bidx", Ops.DEFINE_REG: "acc", Ops.LOAD: "val"}.get(u.op, "alu")
         r[u] = f"{prefix}{c[prefix]}"
 
       l = cast(str, self.string_rewrite.rewrite(u, ctx=self))
       assert l is not None, f"failed to render {u.op} {u.dtype} {[(x.op,x.dtype) for x in u.src]} {u.arg}"
 
       if u.op in {Ops.ENDIF, Ops.END}: depth -= 1
-      if (u.op is not Ops.CAST or u.dtype.vcount == 1) and (u.op in {Ops.CONST, Ops.GEP, Ops.SLICE, Ops.CUSTOMI} or \
+      if (u.op is not Ops.CAST or u.dtype.vcount == 1) and (u.op in {Ops.CONST, Ops.GEP, Ops.INDEX, Ops.CUSTOMI} or \
         (u.op is Ops.LOAD and u.src[0].ptrdtype.addrspace == AddrSpace.REG) or \
         (u.op is Ops.CAST and isinstance(u.dtype, PtrDType)) or \
         (u.op in {Ops.STACK, *(GroupOp.ALU-{Ops.WHERE}), Ops.CAST, Ops.BITCAST} and child_count[u] == 1 and not getenv("EXPAND_SSA"))):

tinygrad/renderer/llvmir.py

@@ -74,9 +74,8 @@ lop = {**{x:unsigned_lop for x in (dtypes.bool,)+dtypes.uints}, **{x:signed_lop
 
 base_rewrite = PatternMatcher([
   # memory load/store
-  (UPat(Ops.SLICE, name="x"), lambda ctx,x:
-   f"  {ctx[x]}_o = getelementptr inbounds {ldt(x.src[0].dtype.base)}, {ldt(x.src[0].dtype)} {ctx[x.src[0]]}, {ldt(x.src[1].dtype)} {ctx[x.src[1]]}\n"
-   f"  {ctx[x]} = bitcast {ldt(x.src[0].dtype)} {ctx[x]}_o to {ldt(x.dtype)}"),
+  (UPat(Ops.INDEX, name="x"), lambda ctx,x:
+   f"  {ctx[x]} = getelementptr inbounds {ldt(x.dtype.base)}, {ldt(x.src[0].dtype)} {ctx[x.src[0]]}, {ldt(x.src[1].dtype)} {ctx[x.src[1]]}"),
   (UPat(Ops.LOAD, src=(UPat.var("idx"), UPat.var("alt"), UPat.var("mask")), name="x"),
    lambda ctx,x,idx,alt,mask:
    f"  br label {ctx[x]}_entry\n{ctx[x][1:]}_entry:\n"

tinygrad/renderer/wgsl.py

@@ -90,7 +90,7 @@ class WGSLRenderer(CStyleLanguage):
      # (load & mask) | var -> mask = v.src[0].src[1], var = v.src[1]
      f"atomicAnd(&{ctx[b]},{ctx[v.src[0].src[1]]});\n  atomicAdd(&{ctx[b]},{ctx[v.src[1]]});" if is_packed(b.src[0].dtype) \
       else f"{ctx[b]} = {ctx[v]};"),
-    (UPat(Ops.SLICE, src=(UPat.var("b"), UPat.var("idx"))),
+    (UPat(Ops.INDEX, src=(UPat.var("b"), UPat.var("idx"))),
      lambda ctx,b,idx: f"{ctx[b]}[{strip_parens(ctx[idx]) if idx.arg is Ops.ADD else ctx[idx]}]"),
   ]) + base_rewrite
 

tinygrad/runtime/ops_python.py

@@ -101,18 +101,16 @@ class PythonProgram:
           if arg[0] == 'g': values[i] = [idxs[2-int(arg[-1])]] * warp_size
           elif arg[0] == 'l': values[i] = [x[2-int(arg[-1])] for x in warp]
         elif uop is Ops.CONST: values[i] = [arg] * warp_size
-        elif uop is Ops.SLICE:
-          assert len(src_values) == 2, "non-image index must be 2 srcs"
-          ret:list = []
-          for m,o in zip(*src_values): ret.append((m,o))
-          values[i] = ret
         elif uop is Ops.INDEX:
-          assert isinstance(src_dtypes[0], ImageDType), "only image INDEX is supported"
-          ret = []
-          assert len(src_values) == 3, "image index must be 3 srcs"
-          for m,oy,ox in zip(*src_values):
-            if ox < 0 or ox >= src_dtypes[0].shape[1] or oy < 0 or oy >= src_dtypes[0].shape[0]: ret.append((m, None))
-            else: ret.append((m, ox*4 + oy*src_dtypes[0].shape[1]*4))
+          ret:list = []
+          if isinstance(src_dtypes[0], ImageDType):
+            assert len(src_values) == 3, "image index must be 3 srcs"
+            for m,oy,ox in zip(*src_values):
+              if ox < 0 or ox >= src_dtypes[0].shape[1] or oy < 0 or oy >= src_dtypes[0].shape[0]: ret.append((m, None))
+              else: ret.append((m, ox*4 + oy*src_dtypes[0].shape[1]*4))
+          else:
+            assert len(src_values) == 2, "non-image index must be 2 srcs"
+            for m,o in zip(*src_values): ret.append((m,o))
           values[i] = ret
         elif uop is Ops.CAST and isinstance(dtype, PtrDType):
           values[i] = src_values[0]

tinygrad/runtime/ops_webgpu.py

@@ -29,7 +29,7 @@ def write_buffer(device:WGPUDevPtr, buf:WGPUBufPtr, offset:int, src:memoryview|b
   src = bytearray(src)
   webgpu.wgpuQueueWriteBuffer(webgpu.wgpuDeviceGetQueue(device), buf, offset, (ctypes.c_uint8 * len(src)).from_buffer(src), len(src))
 
-def _run(async_fun, cb_info_type, cb_type, status_enum, res_idx:int|None, msg_idx:int|None, *params):
+def _run(async_fun, cb_info_type, cb_type, status_enum:dict|None, res_idx:int|None, msg_idx:int|None, *params):
   result: List[Any] = []
 
   def cb(*params):
@@ -55,7 +55,7 @@ def read_buffer(dev:WGPUDevPtr, buf:WGPUBufPtr) -> memoryview:
   tmp_buffer = webgpu.wgpuDeviceCreateBuffer(dev, webgpu.WGPUBufferDescriptor(size=size,
     usage=webgpu.WGPUBufferUsage_CopyDst | webgpu.WGPUBufferUsage_MapRead, mappedAtCreation=False))
   copy_buffer_to_buffer(dev, buf, 0, tmp_buffer, 0, size)
-  _run(webgpu.wgpuBufferMapAsync2, webgpu.WGPUBufferMapCallbackInfo2, webgpu.WGPUBufferMapCallback2, webgpu.WGPUBufferMapAsyncStatus, None, 0,
+  _run(webgpu.wgpuBufferMapAsync2, webgpu.WGPUBufferMapCallbackInfo2, webgpu.WGPUBufferMapCallback2, webgpu.enum_WGPUBufferMapAsyncStatus, None, 0,
        tmp_buffer, webgpu.WGPUMapMode_Read, 0, size)
   void_ptr = ctypes.cast(webgpu.wgpuBufferGetConstMappedRange(tmp_buffer, 0, size), ctypes.c_void_p)
   buf_copy = bytearray((ctypes.c_uint8 * size).from_address(void_ptr.value))
@@ -141,7 +141,7 @@ class WebGPUProgram:
     compute_desc = webgpu.WGPUComputePipelineDescriptor(layout=pipeline_layout,
       compute=webgpu.WGPUComputeState(module=self.prg, entryPoint=to_wgpu_str(self.name)))
     pipeline_result = _run(webgpu.wgpuDeviceCreateComputePipelineAsync2, webgpu.WGPUCreateComputePipelineAsyncCallbackInfo2,
-    webgpu.WGPUCreateComputePipelineAsyncCallback2, webgpu.WGPUCreatePipelineAsyncStatus, 1, None, self.dev, compute_desc)
+    webgpu.WGPUCreateComputePipelineAsyncCallback2, webgpu.enum_WGPUCreatePipelineAsyncStatus, 1, None, self.dev, compute_desc)
 
     command_encoder = webgpu.wgpuDeviceCreateCommandEncoder(self.dev, webgpu.WGPUCommandEncoderDescriptor())
     comp_pass_desc = webgpu.WGPUComputePassDescriptor()
@@ -196,8 +196,8 @@ class WebGpuDevice(Compiled):
   def __init__(self, device:str):
     # Requesting an adapter
     adapter_res = _run(webgpu.wgpuInstanceRequestAdapterF, webgpu.WGPURequestAdapterCallbackInfo, webgpu.WGPURequestAdapterCallback,
-      webgpu.WGPURequestAdapterStatus, 1, 2, instance, webgpu.WGPURequestAdapterOptions(powerPreference=webgpu.WGPUPowerPreference_HighPerformance,
-      backendType=backend_types.get(os.getenv("WEBGPU_BACKEND", ""), 0)))
+      webgpu.enum_WGPUCreatePipelineAsyncStatus, 1, 2, instance, webgpu.WGPURequestAdapterOptions(
+        powerPreference=webgpu.WGPUPowerPreference_HighPerformance, backendType=backend_types.get(os.getenv("WEBGPU_BACKEND", ""), 0)))
 
     # Get supported features
     supported_features = webgpu.WGPUSupportedFeatures()
@@ -215,11 +215,11 @@ class WebGpuDevice(Compiled):
 
     # Requesting a device
     self.device_res = _run(webgpu.wgpuAdapterRequestDeviceF, webgpu.WGPURequestDeviceCallbackInfo, webgpu.WGPURequestDeviceCallback,
-      webgpu.WGPURequestDeviceStatus, 1, 2, adapter_res, dev_desc)
+      webgpu.enum_WGPURequestDeviceStatus, 1, 2, adapter_res, dev_desc)
 
     program = functools.partial(WebGPUProgram, (self.device_res, webgpu.WGPUFeatureName_TimestampQuery in supported))
     super().__init__(device, WebGpuAllocator(self), [WGSLRenderer], program, arch="shader-f16" * (webgpu.WGPUFeatureName_ShaderF16 in supported))
 
   def synchronize(self):
     _run(webgpu.wgpuQueueOnSubmittedWorkDone2, webgpu.WGPUQueueWorkDoneCallbackInfo2, webgpu.WGPUQueueWorkDoneCallback2,
-    webgpu.WGPUQueueWorkDoneStatus, None, None, webgpu.wgpuDeviceGetQueue(self.device_res))
+    webgpu.enum_WGPUQueueWorkDoneStatus, None, None, webgpu.wgpuDeviceGetQueue(self.device_res))

tinygrad/schedule/indexing.py

@@ -128,15 +128,15 @@ def _apply_reshape(in_shape:tuple[sint,...], out_shape:tuple[sint, ...], urngs:U
 @functools.cache
 def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UOp, ...]) -> tuple[UOp, ...]:
   match op:
-    case Ops.SHRINK:  rngs = tuple(a if ss == 0 else a+ss for a,(ss,_) in zip(rngs, arg))
+    case Ops.SHRINK:  rngs = tuple(a if off == 0 else a+off for a,(_,off) in zip(rngs, arg))
     case Ops.PERMUTE: rngs = tuple(rngs[p] for p in argsort(arg))
     case Ops.FLIP:    rngs = tuple(((s-1)-a) if f else a for a,s,f in zip(rngs, in_shape, arg))
     case Ops.EXPAND:  rngs = tuple(a if in_sh == out_sh else a.const_like(0) for a,in_sh,out_sh in zip(rngs, in_shape, arg))
     case Ops.PAD:
       # NOTE: the .where(r-s, i) is not inside the graph_rewrite so that `convert_pad_to_where_to_keep_behavior_local`
       #       wraps the pad with only the newly added valid
-      rngs = tuple(r if (s == 0 and e == 0) else graph_rewrite((r >= s) & (r < (sh+s)),
-        symbolic+pm_simplify_valid, name="pad").where(r-s, UOp.invalid()) for r,sh,(s,e) in zip(rngs, in_shape, arg))
+      rngs = tuple(r if (sz == sh and off == 0) else graph_rewrite((r >= off) & (r < (sh+off)),
+        symbolic+pm_simplify_valid, name="pad").where(r-off, UOp.invalid()) for r,sh,(sz,off) in zip(rngs, in_shape, arg))
     case Ops.RESHAPE:
       sink = UOp.sink(*rngs).simplify() # NOTE: this applies any commutative flips to the rngs early
       sub_array = {r:UOp.range(r.src[0], i, AxisType.PLACEHOLDER) for i,r in enumerate(sink.ranges)}

tinygrad/schedule/multi.py

@@ -10,7 +10,7 @@ def mstack_early_shrink(ms:UOp, shrink:UOp):
   def apply_shrink(s:UOp, i:int) -> UOp:
     new_arg = [tuple([x.substitute({dvar[0]:dvar[0].const_like(i)}) if isinstance(x, UOp) and
                       (dvar:=[v for v in x.variables() if v.expr=='_device_num']) else x for x in ss]) for ss in shrink.marg]
-    return s.shrink(tuple(new_arg))
+    return s._mop(Ops.SHRINK, tuple(new_arg))
   for i, x in enumerate(ms.src):
     if x.op is Ops.COPY:
       ret.append(apply_shrink(x.src[0], i).copy_to_device(x.device))
@@ -88,22 +88,25 @@ def expand_multi(root:UOp, multi:UOp):
   return multi.src[0].expand(new_shape).multi(multi.axis)
 
 def pad_multi(root:UOp, multi:UOp):
-  assert multi.axis is None or root.marg[multi.axis] == (0,0), f"padding not supported for {root.marg=}"
-  return multi.src[0].pad(root.marg).multi(multi.axis)
+  assert multi.axis is None or root.marg[multi.axis] == (multi.shape[multi.axis], 0), f"padding not supported for {root.marg=}"
+  local_pad = tuple((multi.src[0].shape[multi.axis], 0) if a == multi.axis else s for a,s in enumerate(root.marg))
+  return multi.src[0]._mop(Ops.PAD, local_pad).multi(multi.axis)
 
 def permute_multi(root:UOp, multi:UOp):
   # all permutes supported!
   return multi.src[0].permute(root.marg).multi(root.axis)
 
 def shrink_multi(root:UOp, multi:UOp):
-  assert multi.axis is None or root.marg[multi.axis] == (0, multi.shape[multi.axis]) or root.marg[multi.axis] in multi.bounds, \
+  shard_bounds = tuple((e-s,s) for s,e in multi.bounds) if multi.axis is not None else ()
+  assert multi.axis is None or root.marg[multi.axis] == (multi.shape[multi.axis], 0) or root.marg[multi.axis] in shard_bounds, \
     f"shrinking not supported for {root.marg=}"
-  if multi.axis is not None and root.marg[multi.axis] in multi.bounds and root.marg[multi.axis] != (0, multi.shape[multi.axis]):
+  if multi.axis is not None and root.marg[multi.axis] in shard_bounds and root.marg[multi.axis] != (multi.shape[multi.axis], 0):
     # NOTE: shrink on the shard axis is only allowed when result is a single partition, denoted by the new real
     # we just copy it to all the devices, no real. this will be optimized out later
-    non_shard_shrink = tuple((0, multi.src[0].shape[i]) if i == multi.axis else s for i, s in enumerate(root.marg))
-    return multi.src[0].copy_to_device(multi.device, arg=multi.bounds.index(root.marg[multi.axis])).shrink(non_shard_shrink)
-  return multi.src[0].shrink(tuple((0, multi.src[0].shape[multi.axis]) if a == multi.axis else s for a,s in enumerate(root.marg))).multi(multi.axis)
+    non_shard_shrink = tuple((multi.src[0].shape[i], 0) if i == multi.axis else s for i, s in enumerate(root.marg))
+    return multi.src[0].copy_to_device(multi.device, arg=shard_bounds.index(root.marg[multi.axis]))._mop(Ops.SHRINK, non_shard_shrink)
+  local_shrink = tuple((multi.src[0].shape[multi.axis], 0) if a == multi.axis else s for a,s in enumerate(root.marg))
+  return multi.src[0]._mop(Ops.SHRINK, local_shrink).multi(multi.axis)
 
 def flip_multi(root:UOp, multi:UOp):
   assert multi.axis is None or not root.marg[multi.axis], "flipping not supported on sharded axis"

tinygrad/schedule/rangeify.py

@@ -52,7 +52,7 @@ def found_after(ctx:dict[UOp, UOp], after:UOp, src:UOp):
     if x.op is Ops.PERMUTE: x, after = x.src[0], after.permute(argsort(x.marg))
     elif x.op is Ops.RESHAPE: x, after = x.src[0], after.reshape(x.src[0].shape)
     elif x.op is Ops.WHERE and x.src[2].base.arg == Invalid and x.src[1].op is Ops.PAD:
-      x, after = x.src[1].src[0], after.shrink(tuple((l, s-r) for (l,r),s in zip(x.src[1].marg, x.shape)))
+      x, after = x.src[1].src[0], after.shrink(tuple((o, s+o) for (_,o),s in zip(x.src[1].marg, x.src[1].src[0].shape)))
     else: break
   ctx[x] = after
 

tinygrad/tensor.py

@@ -161,6 +161,11 @@ class Tensor(OpMixin):
   @staticmethod
   def const(dtype:DType, b:ConstType|UOp, device:str|tuple[str, ...]|None=None) -> Tensor:
     return Tensor(b if isinstance(b, UOp) else UOp.const(dtype, b, device))
+  @staticmethod
+  def unique_const(fill_value:ConstType|UOp, **kwargs) -> Tensor:
+    if isinstance(fill_value, UOp): return Tensor(fill_value, **kwargs)
+    dtype, device = kwargs.pop("dtype", None), kwargs.pop("device", None)
+    return Tensor(UOp.unique_const(fill_value, dtype, device), **kwargs)
 
   def is_param_(self, is_param:bool=True) -> Tensor:
     self.is_param = is_param
@@ -593,10 +598,9 @@ class Tensor(OpMixin):
     print(Tensor.full_like(t, 42).numpy())
     ```
     """
-    if isinstance(self.device, tuple):
-      if device is not None: raise RuntimeError("cannot specify `device` on `full_like` of a multi device tensor")
-      return self._multi_like(Tensor.full, fill_value, dtype=dtype or self.dtype)
-    return Tensor.full(self.shape, fill_value, dtype=dtype or self.dtype, device=self.device if device is None else device)
+    if device is None: return super().full_like(fill_value, dtype)
+    if isinstance(self.device, tuple): raise RuntimeError("cannot specify `device` on `full_like` of a multi device tensor")
+    return Tensor.full(self.shape, fill_value, dtype=dtype or self.dtype, device=device)
 
   def rand_like(self, **kwargs) -> Tensor:
     """

tinygrad/uop/ops.py

@@ -267,7 +267,6 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
       case Ops.BINARY: return (len(self.arg),)
       case Ops.BUFFER: return (self.arg,)
       case Ops.SLICE:
-        if self.arg == 0: return ()
         # HACK: SLICE is used inside kernels, so we set the shape to () if it's on an INDEX
         if self.src[0].op is Ops.INDEX: return ()
         return (self.arg,)
@@ -328,13 +327,14 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
           return tuple(ps[i] for i in self.marg)
         case Ops.PAD:
           # TODO: why do i need resolve here?
-          if len(ps) != len(self.marg) or not all(resolve(b>=0) and resolve(e>=0) for b,e in self.marg): raise ValueError(f"invalid pad {self.marg}")
-          return tuple(ssimplify(s+b+e) for s,(b,e) in zip(ps, self.marg))
+          if len(ps) != len(self.marg) or not all(resolve(sz>=0) and resolve(0<=o) and resolve(o+s<=sz) for s,(sz,o) in zip(ps, self.marg)):
+            raise ValueError(f"invalid pad {self.marg} for {ps}")
+          return tuple(ssimplify(sz) for sz,_ in self.marg)
         case Ops.SHRINK:
           # TODO: why do i need resolve here?
-          if len(ps) != len(self.marg) or not all(resolve(0<=b) and resolve(b<=e) and resolve(e<=s) for s,(b,e) in zip(ps, self.marg)):
+          if len(ps) != len(self.marg) or not all(resolve(0<=b) and resolve(sz>=0) and resolve(b+sz<=s) for s,(sz,b) in zip(ps, self.marg)):
             raise ValueError(f"invalid shrink {self.marg} for {ps}")
-          return tuple(ssimplify(e-s) for s,e in self.marg)
+          return tuple(ssimplify(sz) for sz,_ in self.marg)
         case Ops.FLIP:
           if len(ps) != len(self.marg) or not all(isinstance(x, bool) for x in self.marg): raise ValueError(f"bad flip on {ps}, {self.marg}")
           return ps
@@ -460,8 +460,6 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     return UOp(Ops.GROUP, dtypes.void, tuple([x for x in srcs if x is not None]))
   def vectorize(self, *srcs):
     return UOp(Ops.STACK, self.dtype.vec(len(srcs)+1), (self,)+srcs)
-  def slice(self, offset:UOp|int, size:int=0):
-    return UOp(Ops.SLICE, self.dtype, (self, offset if isinstance(offset, UOp) else UOp.const(dtypes.int, offset)), arg=size)
   def index(self, *srcs:UOp|None, ptr=False, **kwargs):
     return UOp(Ops.INDEX, kwargs.pop("dtype", self.dtype if ptr else self.dtype.base), (self,)+tuple([x for x in srcs if x is not None]), **kwargs)
   def __getitem__(self, idx):
@@ -538,6 +536,14 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
       ret = UOp(Ops.CONST, dtype, arg=dtype.const(b), src=(UOp(Ops.DEVICE, arg=device),) if device is not None else ())
     return ret.reshape((1,)*len(shape)).expand(shape) if shape is not None and shape != () and ret.shape != shape else ret
   @staticmethod
+  def unique_const(fill_value:ConstType, dtype:DTypeLike|None=None, device:str|tuple[str, ...]|None=None,  # type: ignore[override]
+                   shape:tuple[sint, ...]|None=None, unique=True):
+    # NOTE: fill_value is ConstType, not ConstLike, so UOps and tuples aren't allowed
+    assert not isinstance(fill_value, (UOp, tuple)), "unique const only works on numbers"
+    ret = UOp.const(to_dtype(dtype) if dtype is not None else dtypes.from_py(fill_value), fill_value, canonicalize_device(device))
+    ret = ret.replace(src=(UOp.unique(None if unique is True else unique),) + ret.src)
+    return ret.reshape((1,)*len(shape)).expand(shape) if shape is not None and ret.shape != shape else ret
+  @staticmethod
   def range(end:sint, axis_id, axis_type=AxisType.LOOP, *arg, dtype=dtypes.weakint, src=(), **kwargs):
     return UOp(Ops.RANGE, dtype=dtype, src=(sint_to_uop(end, dtype),)+src, arg=(axis_id, axis_type)+arg, **kwargs)
   @staticmethod
@@ -613,7 +619,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     if self.op in GroupOp.ALU: return axes[-1] if (axes := dedup([x.axis for x in self.src if x.axis is not None])) else None
     if len(self.src) == 0: return None
     src_axis = self.src[0].axis
-    if self.op is Ops.SHRINK and src_axis is not None and self.marg[src_axis] != (0, self.src[0].shape[src_axis]):
+    if self.op is Ops.SHRINK and src_axis is not None and self.marg[src_axis] != (self.src[0].shape[src_axis], 0):
       return None # SHRINK will remove the sharding if it's on axis
     if self.op is Ops.REDUCE: return None if src_axis is not None and src_axis in self.arg[1] else src_axis
     if self.op is Ops.RESHAPE:
@@ -1077,8 +1083,9 @@ class ProgramInfo:
     for u in sink.toposort():
       if u.op is Ops.DEFINE_VAR: _vars.append(u)
       if u.op is Ops.PARAM: _globals.append(u.arg)
-      if u.op in (Ops.STORE, Ops.LOAD) and (idx:=u.src[0]).op in (Ops.INDEX, Ops.SLICE) and (buf:=idx.src[0]).op is Ops.PARAM:
-        (outs if u.op is Ops.STORE else ins).append(buf.arg)
+      if u.op in (Ops.STORE, Ops.LOAD):
+        if (idx:=u.src[0]).op is Ops.INDEX or (u.src[0].op is Ops.CAST and (idx:=u.src[0].src[0]).op is Ops.INDEX):
+          if (buf:=idx.src[0]).op is Ops.PARAM: (outs if u.op is Ops.STORE else ins).append(buf.arg)
       if u.op is Ops.SPECIAL:
         if u.arg[0] == 'i': local_size = None
         special_size = local_size if u.arg[0] == 'l' else global_size
@@ -1150,13 +1157,15 @@ def get_location() -> tuple[str, int]:
   return frm.f_code.co_filename, frm.f_lineno
 
 class UPat(OpMixin):
-  __slots__ = ("op", "match_dtype", "arg", "name", "src", "is_any")
+  __slots__ = ("op", "match_dtype", "match_tag", "arg", "name", "src", "is_any")
   def __init__(self, op:Ops|tuple[Ops, ...]|set[Ops]|None=None, dtype:DType|tuple[DType, ...]|set[DType]|None=None,
                src:tuple[UPat, ...]|list[UPat]|UPat|None=None, arg:Any=None,
-               name:str|None=None, allow_any_len:bool=False, custom_early_reject:set[Ops]|None=None, location=None, is_any:bool=False):
+               name:str|None=None, allow_any_len:bool=False, custom_early_reject:set[Ops]|None=None, location=None, is_any:bool=False,
+               tag:Any=None):
     assert op is None or isinstance(op, (Ops, tuple, set)), f"op must be Ops or tuple of Ops, not {op!r}"
     self.op: tuple[Ops, ...]|None = (op,) if isinstance(op, Ops) else (tuple(op) if isinstance(op, set) else op)
     self.match_dtype: tuple[DType, ...]|None = (dtype,) if isinstance(dtype, DType) else (tuple(dtype) if isinstance(dtype, set) else dtype)
+    self.match_tag: tuple[Any, ...]|None = (tag,) if isinstance(tag, str) else (tuple(tag) if isinstance(tag, set) else tag)
     self.arg, self.name, self._in_src, self.custom_early_reject = arg, name, src, custom_early_reject
     self.src: Any = None
     self.is_any = is_any
@@ -1186,8 +1195,10 @@ class UPat(OpMixin):
   def _ensure_float(self) -> UPat: return self
 
   def __reduce__(self):
-    return UPat, (self.op, self.match_dtype, self._in_src, self.arg, self.name, not self.strict_length, self.custom_early_reject, self.location)
-  def named(self, name:str): return UPat(self.op, self.match_dtype, self._in_src, self.arg, name, not self.strict_length, self.custom_early_reject)
+    return UPat, (self.op, self.match_dtype, self._in_src, self.arg, self.name, not self.strict_length, self.custom_early_reject, self.location,
+                  self.is_any, self.match_tag)
+  def named(self, name:str):
+    return UPat(self.op, self.match_dtype, self._in_src, self.arg, name, not self.strict_length, self.custom_early_reject, tag=self.match_tag)
 
   @staticmethod
   def any(*src): return UPat(src=src, is_any=True)
@@ -1248,6 +1259,7 @@ class UPat(OpMixin):
        (self.name is not None and store.setdefault(self.name, uop) is not uop) or \
        (self.match_dtype is not None and uop.dtype not in self.match_dtype and uop.dtype.scalar() not in self.match_dtype) or \
        (self.arg is not None and self.arg != uop.arg) or \
+       (self.match_tag is not None and uop.tag not in self.match_tag) or \
        (len(uop.src) < self.required_len) or \
        (self.strict_length and len(uop.src) != self.required_len): return []
     if self.src is None: return [store]

tinygrad/uop/render.py

@@ -79,6 +79,8 @@ def render_marg(ctx,x:UOp):
 sugar = {Ops.SINK, Ops.END, Ops.STORE, Ops.LOAD, Ops.UNIQUE, Ops.SQRT, Ops.INDEX, Ops.REDUCE, Ops.AFTER, Ops.THREEFRY,
          Ops.WHERE, Ops.RECIPROCAL, Ops.EXP2, Ops.LOG2, Ops.SIN, Ops.CONTIGUOUS, Ops.BARRIER, Ops.DETACH}
 pm_pyrender_extra = PatternMatcher([
+  (UPat(Ops.CONST, src=(UPat(Ops.UNIQUE, name="u"), UPat(Ops.DEVICE, name="d")), name="x"),
+   lambda x,u,d: f"UOp.unique_const({x.arg}, dtype={x.dtype}, device={repr(d.arg)}, unique={u.arg})"),
   (UPat(Ops.CONST, src=(UPat(Ops.DEVICE, name="d"),), name="x"), lambda x,d: f"UOp.const({x.dtype}, {x.arg}, device={repr(d.arg)})"),
   (UPat(Ops.CONST, src=(), name="x"), lambda x: f"UOp.const({x.dtype}, {x.arg})"),
   (UPat(Ops.DEFINE_VAR, src=(), name="x"), lambda x:
@@ -104,6 +106,7 @@ pm_pyrender_extra = PatternMatcher([
   # explicit trunc ops: `//` and `%` parse as FLOORDIV/FLOORMOD, so render CDIV/CMOD via .alu()
   (UPat(Ops.CDIV, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.alu(Ops.CDIV, {ctx[x.src[1]]})"),
   (UPat(Ops.CMOD, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.alu(Ops.CMOD, {ctx[x.src[1]]})"),
+  # NOTE: only match CONSTs without UNIQUE (len(src)==1), unique_const needs explicit rendering
   (UPat(set(syms.keys())-{Ops.SUB, Ops.CMPNE, Ops.CDIV, Ops.CMOD}, src=(UPat(Ops.CONST, src=(UPat(Ops.DEVICE),), name="y"), UPat(name="z")),
         name="x"), lambda ctx,x,y,z: strip_binary_parens(x, str(y.arg), ctx[z], lambda a,b: f"({a}{syms[x.op]}{b})") if y.device==z.device else None),
   # NOTE: sub doesn't work cause it's written as add/mul

tinygrad/uop/spec.py

@@ -99,11 +99,11 @@ spec_shared = PatternMatcher([
   (UPat(Ops.INS), lambda: True),
 
   # LOAD(idx) / STORE(idx, val) with gates on the LOAD/STORE
-  (UPat((Ops.INDEX, Ops.SLICE), name="uidx").or_casted().load(), validate_index),
-  (UPat((Ops.INDEX, Ops.SLICE), name="uidx").or_casted().load(UPat.var("alt"), UPat.var("gate", dtype=dtypes.bool), name="load"),
+  (UPat(Ops.INDEX, name="uidx").or_casted().load(), validate_index),
+  (UPat(Ops.INDEX, name="uidx").or_casted().load(UPat.var("alt"), UPat.var("gate", dtype=dtypes.bool), name="load"),
    lambda uidx,gate,alt,load: validate_index(uidx, gate) if alt.dtype == load.dtype else False),
-  (UPat((Ops.INDEX, Ops.SLICE), name="uidx").or_casted().store(UPat()), validate_index),
-  (UPat((Ops.INDEX, Ops.SLICE), name="uidx").or_casted().store(UPat(), UPat.var("gate", dtype=dtypes.bool)), validate_index),
+  (UPat(Ops.INDEX, name="uidx").or_casted().store(UPat()), validate_index),
+  (UPat(Ops.INDEX, name="uidx").or_casted().store(UPat(), UPat.var("gate", dtype=dtypes.bool)), validate_index),
 
   # STORE in tensor graph: store a value into a target
   (UPat(Ops.STORE, dtypes.void, (UPat(name="x"), UPat())), lambda x: True),
@@ -122,8 +122,9 @@ spec_tensor = PatternMatcher([
   (UPat(Ops.UNIQUE, dtypes.void, ()), lambda: True),
   (UPat(Ops.LUNIQUE, dtypes.void, ()), lambda: True),
 
-  # CONST with a DEVICE
+  # CONST with a UNIQUE or DEVICE
   (UPat(Ops.CONST, src=(UPat(Ops.DEVICE),)), lambda: True),
+  (UPat(Ops.CONST, src=(UPat((Ops.UNIQUE, Ops.LUNIQUE)), UPat(Ops.DEVICE))), lambda: True),
 
   # BUFFER
   (UPat(Ops.BUFFER, src=(UPat((Ops.UNIQUE, Ops.LUNIQUE)), UPat(Ops.DEVICE)), name="buf"),
@@ -157,7 +158,8 @@ spec_tensor = PatternMatcher([
 
   # movement ops
   (UPat((Ops.RESHAPE, Ops.EXPAND), src=(UPat(), UPat(dtype=dtypes.weakint))), lambda: True),
-  (UPat((Ops.PAD, Ops.SHRINK), src=(UPat(), UPat(dtype=dtypes.weakint), UPat(dtype=dtypes.weakint))), lambda: True),
+  (UPat((Ops.PAD, Ops.SHRINK), src=(UPat(), UPat(dtype=dtypes.weakint), UPat(dtype=dtypes.weakint)), name="x"),
+   lambda x: x.src[1].dtype.count == x.src[2].dtype.count),
   (UPat((Ops.PERMUTE, Ops.FLIP), name="mv", src=(UPat(),)), lambda mv: isinstance(mv.arg, tuple)),
 
   # REDUCE has arg=(op, axis_tuple), src[1:] are ranges after lowering
@@ -200,10 +202,6 @@ spec_program = PatternMatcher([
   # weakint is not allowed in programs
   (UPat(GroupOp.All, dtypes.weakint), lambda: False),
 
-  # buffer view in program, Image only for ImageDType
-  (UPat(Ops.SLICE), lambda: True),
-  (UPat(Ops.INDEX, name="idx"), lambda idx: isinstance(idx.src[0].dtype, ImageDType)),
-
   # movement ops are not allowed in programs
   (UPat(GroupOp.Movement), lambda: False),
 

tinygrad/uop/upat.py

@@ -26,6 +26,10 @@ def _get_clause(self:UPat, base:UOp, depth=0) -> UOp:
     if len(self.match_dtype) > 1:
       and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.BIND, arg=tuple(self.match_dtype))), arg="({0}.dtype in {1} or {0}.dtype._scalar in {1})"))
     else: and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.BIND, arg=self.match_dtype[0])), arg="({0}.dtype == {1} or {0}.dtype._scalar == {1})"))
+  if self.match_tag is not None:
+    if len(self.match_tag) > 1:
+      and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.BIND, arg=tuple(self.match_tag))), arg="{0}.tag in {1}"))
+    else: and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.BIND, arg=self.match_tag[0])), arg="{0}.tag == {1}"))
   if self.src is not None:
     # single match
     if len(self.src) == 1 and isinstance(self.src[0], tuple):

tinygrad/viz/serve.py

@@ -50,7 +50,7 @@ uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0",
                Ops.RANGE: "#c8a0e0", Ops.BARRIER: "#ff8080", Ops.IF: "#c8b0c0", Ops.SPECIAL: "#c0c0ff",
                Ops.INDEX: "#cef263", Ops.WMMA: "#efefc0", Ops.MULTI: "#f6ccff", Ops.INS: "#eec4ff",
                **{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80",
-               Ops.SLICE: "#a2c148", Ops.BUFFER: "#B0BDFF", Ops.GETADDR: "#9DB1F0", Ops.COPY: "#a040a0", Ops.CUSTOM_FUNCTION: "#bf71b6",
+               Ops.SLICE: "#E5EAFF", Ops.BUFFER: "#B0BDFF", Ops.GETADDR: "#9DB1F0", Ops.COPY: "#a040a0", Ops.CUSTOM_FUNCTION: "#bf71b6",
                Ops.CALL: "#00B7C8", Ops.FUNCTION: "#C07788", Ops.PARAM: "#14686F", Ops.SOURCE: "#c0c0c0", Ops.BINARY: "#404040",
                Ops.LINEAR: "#7DF4FF",
                Ops.ALLREDUCE: "#ff40a0", Ops.MSELECT: "#d040a0", Ops.MSTACK: "#d040a0", Ops.CONTIGUOUS: "#FFC14D",
@@ -117,6 +117,7 @@ def uop_to_json(data:VizData, x:UOp) -> dict[int, dict]:
   for u in (toposort:=x.toposort()):
     # always exclude DEVICE/CONST/UNIQUE
     if u.op in {Ops.DEVICE, Ops.CONST, Ops.UNIQUE, Ops.LUNIQUE} and u is not x: excluded.add(u)
+    if u.op is Ops.CONST and len(u.src) and u.src[0].op in {Ops.UNIQUE, Ops.LUNIQUE}: excluded.remove(u)
     if u.op is Ops.STACK and len(u.src) == 0: excluded.add(u)
     # exclude RESHAPE/EXPAND that only serve to broadcast a CONST
     if u.op in {Ops.RESHAPE, Ops.EXPAND} and len(u.src) >= 1 and u.src[0] in excluded and u is not x: excluded.add(u)