Merge branch 'master' into pci_iommu

A misbehaving device doing raw DMA to host memory used to take the whole
system down (fatal machine check from corrupted/poisoned memory). When
the device is managed by a real IOMMU, bind it to vfio-pci and pin every
page exposed to the device (sysmem allocations, mapped CPU buffers) into
the container at iova=paddr, so a wild DMA faults in the IOMMU instead
of corrupting host memory. Falls back to vfio no-iommu (with a warning)
when no IOMMU is available. Also map BARs through the vfio device fd,
since sysfs BAR mappings are revoked under vfio-pci.

docs/developer/am.md

@@ -7,8 +7,13 @@ Make sure that amdgpu module is unloaded and just run tinygrad with `DEV=AMD`!
 
 Optional requirements:
 
-* System without IOMMU for P2P / SDMA support
-* vfio-pci module for IRQ handling
+* vfio-pci module for IRQ handling and IOMMU-protected DMA
+
+When the system IOMMU is enabled (AMD-Vi), the driver must go through vfio so that the GPU's DMA is confined to explicitly
+mapped pages: a device fault then hits an IOMMU page fault (and only kills the GPU session) instead of corrupting host memory
+and taking the whole system down. This is enabled automatically when the device is behind an IOMMU (set `VFIO=0` to opt out,
+e.g. with `iommu=pt`). Note that without an IOMMU (or with `iommu=pt`) DMA is unprotected. P2P between GPUs is only supported
+without address translation: boot with `iommu=pt` and set `VFIO=0`.
 
 ## Environment Variables
 
@@ -16,6 +21,7 @@ Optional requirements:
 |----------|------------------|-------------|
 | AM_RESET | [1] | Performs a full GPU reset (reloading all firmware and IP blocks) |
 | AM_DEBUG | [0-4] | Sets the level of additional debugging information |
+| VFIO | [0, 1] | Force raw PCI access (0) or vfio (1). By default vfio is used automatically when the device is behind an IOMMU, which requires it |
 
 ## AM Driver Details
 

test/external/external_test_pci_iommu.py

@@ -0,0 +1,62 @@
+# Verifies IOMMU containment of misbehaving device DMA: pages revoked from the vfio container but still mapped in the GPU's
+# page tables must fault in the IOMMU (IO_PAGE_FAULT) instead of reaching host memory. Requires an active IOMMU (VFIO type1v2).
+# Run with: DEV=PCI:0+AMD python3 test/external/external_test_pci_iommu.py
+import subprocess, unittest
+from tinygrad import Device
+from tinygrad.device import BufferSpec
+from tinygrad.runtime.support.system import PCIAllocationMeta
+from tinygrad.runtime.support.memory import AddrSpace
+from tinygrad.runtime.support.hcq import HCQBuffer
+
+class TestPCIIOMMU(unittest.TestCase):
+  @classmethod
+  def setUpClass(cls):
+    cls.dev, cls.pci_dev, cls.mm = (d:=Device[Device.DEFAULT]), d.iface.pci_dev, d.iface.dev_impl.mm
+    if not cls.pci_dev.iommu: raise unittest.SkipTest("requires an active IOMMU")
+
+  def bad_buf(self, paddr:int) -> HCQBuffer:
+    bp = self.mm.map_range(va:=self.mm.alloc_vaddr(0x1000), 0x1000, [(paddr, 0x1000)], aspace=AddrSpace.SYS, snooped=True, uncached=True)
+    return HCQBuffer(va, 0x1000, meta=PCIAllocationMeta(bp, has_cpu_mapping=False), owner=self.dev)
+
+  def test_wild_dma_is_contained(self):
+    N, pages = 64, []
+    for i in range(N):
+      view, paddrs = self.pci_dev.alloc_sysmem(0x1000) # legit sysmem page: pinned in the vfio container
+      view[:0x1000] = (b"SENTINEL" + i.to_bytes(2, 'little')) + bytes(0x1000 - 10)
+      self.pci_dev.dma_unmap(paddrs) # revoke it: from now on any device DMA to it must fault in the IOMMU
+      pages.append((view, self.bad_buf(paddrs[0])))
+
+    src = self.dev.allocator._alloc(0x1000, BufferSpec())
+
+    # storm the IOMMU with wild writes (a valid GART entry pointing at a revoked page == misbehaving GPU)
+    q = self.dev.hw_copy_queue_t()
+    for _, bad in pages: q.copy(bad, src, 0x1000)
+    q.signal(self.dev.timeline_signal, tlv:=self.dev.next_timeline()).submit(self.dev)
+    self.dev.timeline_signal.wait(tlv, timeout=10000)
+
+    # and a wild read for good measure
+    self.dev.hw_copy_queue_t().copy(src, pages[0][1], 0x1000).signal(self.dev.timeline_signal, tlv:=self.dev.next_timeline()).submit(self.dev)
+    self.dev.timeline_signal.wait(tlv, timeout=10000)
+
+    # none of the wild DMA may have reached host memory, and there must be no hardware error (MCE)
+    for i, (view, _) in enumerate(pages): self.assertEqual(bytes(view[:10]), b"SENTINEL" + i.to_bytes(2, 'little'))
+    hw_errs = subprocess.run("journalctl -k --no-pager --since '-60s' | grep -ci 'Hardware Error' || true",
+                             shell=True, capture_output=True, text=True).stdout.strip()
+    self.assertIn(hw_errs, ("", "0"), f"unexpected hardware errors in the kernel log: {hw_errs}")
+
+  def test_device_survives_faults(self):
+    view, paddrs = self.pci_dev.alloc_sysmem(0x1000)
+    view[:0x1000] = b"IOMMU-OK!" + bytes(0x1000 - 9)
+    self.pci_dev.dma_unmap(paddrs)
+
+    src = self.dev.allocator._alloc(0x1000, BufferSpec())
+    self.dev.hw_copy_queue_t().copy(self.bad_buf(paddrs[0]), src, 0x1000) \
+                              .signal(self.dev.timeline_signal, tlv:=self.dev.next_timeline()).submit(self.dev)
+    self.dev.timeline_signal.wait(tlv, timeout=10000)
+    self.assertEqual(bytes(view[:9]), b"IOMMU-OK!")
+
+    # device is still usable after the fault
+    self.dev.allocator._copyout(mv:=memoryview(bytearray(4)), src)
+    self.assertEqual(len(mv), 4)
+
+if __name__ == "__main__": unittest.main()

tinygrad/runtime/support/system.py

@@ -31,9 +31,19 @@ class _System:
     try:
       if not FileIOInterface.exists("/sys/module/vfio"): os.system("sudo modprobe vfio-pci disable_idle_d3=1")
 
-      FileIOInterface("/sys/module/vfio/parameters/enable_unsafe_noiommu_mode", os.O_RDWR).write("1")
       vfio_fd = FileIOInterface("/dev/vfio/vfio", os.O_RDWR)
-      vfio.VFIO_CHECK_EXTENSION(vfio_fd, vfio.VFIO_NOIOMMU_IOMMU)
+
+      # IOVA -> refcount for pages pinned into the vfio container. Only pages present here are reachable by the device's DMA,
+      # so a misbehaving device faults in the IOMMU instead of corrupting host memory (which takes the whole system down).
+      self.vfio_dma_pages: dict[int, int] = {}
+
+      try:
+        # Prefer a real IOMMU when one is available. PCIDevice falls back to no-iommu per device when there is none.
+        vfio.VFIO_CHECK_EXTENSION(vfio_fd, vfio.VFIO_TYPE1v2_IOMMU)
+        self.vfio_noiommu = False
+      except OSError:
+        vfio.VFIO_CHECK_EXTENSION(vfio_fd, vfio.VFIO_NOIOMMU_IOMMU)
+        self.vfio_noiommu = True
 
       return vfio_fd
     except OSError: return None
@@ -154,9 +164,12 @@ System = _System()
 # *** PCI Devices
 
 class PCIDevice:
+  iommu:bool = False # True when the device is managed by vfio with a real IOMMU (DMA is confined to dma_map()ed pages)
+
   def __init__(self, devpref:str, pcibus:str):
     self.lock_fd = System.flock_acquire(f"{devpref.lower()}_{pcibus.lower()}.lock")
     self.pcibus, self.irq_poller = pcibus, None
+    self.dma_mapped: dict[int, list[int]] = {}
 
     try: FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/enable", os.O_RDWR)
     except PermissionError: raise PermissionError(f"Cannot access PCI device {pcibus}: run `extra/amdpci/setup_python_cap.sh` or use sudo")
@@ -169,15 +182,30 @@ class PCIDevice:
     for fn in range(1, 8):
       if FileIOInterface.exists(sib:=f"/sys/bus/pci/devices/{self.pcibus[:-1]}{fn}"): FileIOInterface(f"{sib}/remove", os.O_WRONLY).write("1")
 
-    if getenv("VFIO", 0) and (vfio_fd:=System.vfio) is not None:
+    # Devices behind a real IOMMU must go through vfio with type1v2 mappings: programming raw physical addresses faults in the
+    # IOMMU, so PCI access without vfio silently doesn't work (this is the safe failure mode).
+    has_iommu = FileIOInterface.exists(f"/sys/bus/pci/devices/{self.pcibus}/iommu_group")
+    want_vfio = (vfio_num:=getenv("VFIO", -1)) == 1 or (has_iommu and vfio_num != 0)
+    if want_vfio and (vfio_fd:=System.vfio) is not None:
       FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/driver_override", os.O_WRONLY).write("vfio-pci")
       FileIOInterface("/sys/bus/pci/drivers_probe", os.O_WRONLY).write(self.pcibus)
-      iommu_group = FileIOInterface.readlink(f"/sys/bus/pci/devices/{self.pcibus}/iommu_group").split('/')[-1]
 
-      self.vfio_group = FileIOInterface(f"/dev/vfio/noiommu-{iommu_group}", os.O_RDWR)
+      grp_path = f"/sys/bus/pci/devices/{self.pcibus}/iommu_group"
+      if not FileIOInterface.exists(grp_path):
+        # On systems without a real IOMMU vfio-pci refuses to bind: enable unsafe no-iommu mode (unprotected DMA) and retry.
+        if DEBUG >= 1: print(f"pci {self.pcibus}: WARNING: no IOMMU, device DMA is unprotected (a fault can crash the system)")
+        FileIOInterface("/sys/module/vfio/parameters/enable_unsafe_noiommu_mode", os.O_RDWR).write("1")
+        FileIOInterface("/sys/bus/pci/drivers_probe", os.O_WRONLY).write(self.pcibus)
+        System.vfio_noiommu = True
+      iommu_group = FileIOInterface.readlink(grp_path).split('/')[-1]
+
+      vfio_node = iommu_group if FileIOInterface.exists(f"/dev/vfio/{iommu_group}") else f"noiommu-{iommu_group}"
+      self.iommu = not vfio_node.startswith("noiommu-")
+      self.vfio_group = FileIOInterface(f"/dev/vfio/{vfio_node}", os.O_RDWR)
       vfio.VFIO_GROUP_SET_CONTAINER(self.vfio_group, ctypes.c_int(vfio_fd.fd))
 
-      with contextlib.suppress(OSError): vfio.VFIO_SET_IOMMU(vfio_fd, vfio.VFIO_NOIOMMU_IOMMU) # set iommu works only once for the fd.
+      # set iommu works only once for the fd.
+      with contextlib.suppress(OSError): vfio.VFIO_SET_IOMMU(vfio_fd, vfio.VFIO_TYPE1v2_IOMMU if self.iommu else vfio.VFIO_NOIOMMU_IOMMU)
       self.vfio_dev = FileIOInterface(fd=vfio.VFIO_GROUP_GET_DEVICE_FD(self.vfio_group, ctypes.create_string_buffer(self.pcibus.encode())))
 
       self.irq_fd = FileIOInterface.eventfd(0, 0)
@@ -187,7 +215,10 @@ class PCIDevice:
       irqs = vfio.struct_vfio_irq_set(index=vfio.VFIO_PCI_MSI_IRQ_INDEX, flags=vfio.VFIO_IRQ_SET_DATA_EVENTFD|vfio.VFIO_IRQ_SET_ACTION_TRIGGER,
         argsz=ctypes.sizeof(vfio.struct_vfio_irq_set) + ctypes.sizeof(ctypes.c_int), count=1)
       vfio.VFIO_DEVICE_SET_IRQS(self.vfio_dev, (ctypes.c_byte * irqs.argsz).from_buffer(bytearray(bytes(irqs)) + struct.pack('i', self.irq_fd.fd)))
-    else: FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/enable", os.O_RDWR).write("1")
+    else:
+      if has_iommu and vfio_num != 0: raise RuntimeError(f"{pcibus} is behind an active IOMMU: use vfio (VFIO=1) or boot with iommu=pt")
+      if has_iommu and DEBUG >= 1: print(f"pci {pcibus}: WARNING: vfio disabled while an IOMMU is active, device DMA will fault in the IOMMU")
+      FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/enable", os.O_RDWR).write("1")
 
     self.cfg_fd = FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/config", os.O_RDWR | os.O_SYNC | os.O_CLOEXEC)
 
@@ -195,10 +226,53 @@ class PCIDevice:
     assert not contiguous or size <= (2 << 20), "Contiguous allocation is only supported for sizes up to 2MB"
     flags = (libc.MAP_HUGETLB if contiguous and (size:=round_up(size, mmap.PAGESIZE)) > mmap.PAGESIZE else 0) | (MAP_FIXED if vaddr else 0)
     va = FileIOInterface.anon_mmap(vaddr, size, mmap.PROT_READ|mmap.PROT_WRITE, mmap.MAP_SHARED|mmap.MAP_ANONYMOUS|MAP_POPULATE|MAP_LOCKED|flags, 0)
-    sysmem_view, paddrs = MMIOInterface(va, size), [(x, mmap.PAGESIZE) for x in System.system_paddrs(va, size)]
-    return sysmem_view, [p + i for p, sz in paddrs for i in range(0, sz, 0x1000)][:ceildiv(size, 0x1000)]
+    paddrs = [p for x in System.system_paddrs(va, size) for p in range(x, x + mmap.PAGESIZE, 0x1000)][:ceildiv(size, 0x1000)]
+    self.dma_map(va, paddrs)
+    return MMIOInterface(va, size), paddrs
 
-  def reset(self): os.system(f"sudo sh -c 'echo 1 > /sys/bus/pci/devices/{self.pcibus}/reset'")
+  def dma_map(self, vaddr:int, paddrs:list[int]):
+    if not self.iommu: return
+    assert vaddr % mmap.PAGESIZE == 0 and all(p % mmap.PAGESIZE == 0 for p in paddrs), f"unaligned {vaddr=:#x}"
+
+    fresh, va = [], vaddr
+    for p in paddrs:
+      if System.vfio_dma_pages.get(p, 0) == 0: fresh.append((p, va))
+      System.vfio_dma_pages[p] = System.vfio_dma_pages.get(p, 0) + 1
+      va += mmap.PAGESIZE
+
+    # One ioctl per run of contiguous newly-mapped pages.
+    i = 0
+    while i < len(fresh):
+      j = i
+      while j + 1 < len(fresh) and fresh[j+1] == (fresh[j][0] + mmap.PAGESIZE, fresh[j][1] + mmap.PAGESIZE): j += 1
+      dm = vfio.struct_vfio_iommu_type1_dma_map(argsz=ctypes.sizeof(vfio.struct_vfio_iommu_type1_dma_map),
+        flags=vfio.VFIO_DMA_MAP_FLAG_READ|vfio.VFIO_DMA_MAP_FLAG_WRITE, vaddr=fresh[i][1], iova=fresh[i][0], size=(j-i+1)*mmap.PAGESIZE)
+      vfio.VFIO_IOMMU_MAP_DMA(unwrap(System.vfio), dm)
+      i = j + 1
+
+  def dma_unmap(self, paddrs:list[int]):
+    if not self.iommu: return
+
+    stale = []
+    for p in paddrs:
+      if (rc:=System.vfio_dma_pages.get(p, 0)) > 1: System.vfio_dma_pages[p] = rc - 1
+      else:
+        System.vfio_dma_pages.pop(p, None)
+        stale.append(p)
+
+    # One ioctl per run of contiguous newly-unmapped pages.
+    i = 0
+    while i < len(stale):
+      j = i
+      while j + 1 < len(stale) and stale[j+1] == stale[j] + mmap.PAGESIZE: j += 1
+      du = vfio.struct_vfio_iommu_type1_dma_unmap(argsz=ctypes.sizeof(vfio.struct_vfio_iommu_type1_dma_unmap),
+                                                  iova=stale[i], size=(j-i+1)*mmap.PAGESIZE)
+      vfio.VFIO_IOMMU_UNMAP_DMA(unwrap(System.vfio), du)
+      i = j + 1
+
+  def reset(self):
+    if getattr(self, 'vfio_dev', None) is not None: vfio.VFIO_DEVICE_RESET(self.vfio_dev)
+    else: os.system(f"sudo sh -c 'echo 1 > /sys/bus/pci/devices/{self.pcibus}/reset'")
   def read_config(self, offset:int, size:int): return int.from_bytes(self.cfg_fd.read(size, binary=True, offset=offset), byteorder='little')
   def write_config(self, offset:int, value:int, size:int): self.cfg_fd.write(value.to_bytes(size, byteorder='little'), binary=True, offset=offset)
   def write_config_flush(self, offset:int, value:int, size:int):
@@ -207,14 +281,24 @@ class PCIDevice:
 
   @functools.cache
   def bar_fd(self, bar_idx:int) -> FileIOInterface:
+    # With vfio, sysfs BAR mappings are revoked. BARs are mapped via the vfio device fd at the region's offset instead.
+    if self.iommu: return self.vfio_dev
     return FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/resource{bar_idx}", os.O_RDWR | os.O_SYNC | os.O_CLOEXEC)
   @functools.cache
+  def bar_off(self, bar_idx:int) -> int:
+    if not self.iommu: return 0
+    info = vfio.struct_vfio_region_info(argsz=ctypes.sizeof(vfio.struct_vfio_region_info), index=vfio.VFIO_PCI_BAR0_REGION_INDEX + bar_idx)
+    vfio.VFIO_DEVICE_GET_REGION_INFO(self.vfio_dev, info)
+    assert info.flags & vfio.VFIO_REGION_INFO_FLAG_MMAP, f"BAR {bar_idx} is not mmappable"
+    return info.offset
+  @functools.cache
   def bar_info(self, bar_idx:int) -> tuple[int, int]:
     s, e, _ = FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/resource", os.O_RDONLY).read().splitlines()[bar_idx].split()
     return (int(s, 16), int(e, 16) - int(s, 16) + 1)
   def map_bar(self, bar:int, off:int=0, addr:int=0, size:int|None=None, fmt='B') -> MMIOInterface:
     fd, sz = self.bar_fd(bar), size or (self.bar_info(bar)[1] - off)
-    libc.madvise(loc:=fd.mmap(addr, sz, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED | (MAP_FIXED if addr else 0), off), sz, libc.MADV_DONTFORK)
+    libc.madvise(loc:=fd.mmap(addr, sz, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED | (MAP_FIXED if addr else 0), off + self.bar_off(bar)),
+                 sz, libc.MADV_DONTFORK)
     return MMIOInterface(loc, sz, fmt=fmt)
   def resize_bar(self, bar_idx:int):
     rpath = f"/sys/bus/pci/devices/{self.pcibus}/resource{bar_idx}_resize"
@@ -277,9 +361,13 @@ class PCIIfaceBase:
     return HCQBuffer(mapping.va_addr, size, view=barview, meta=PCIAllocationMeta(mapping, cpu_access, hMemory=mapping.paddrs[0][0]), owner=self.dev)
 
   def free(self, b:HCQBuffer):
-    if b.owner != self.dev: self.dev.iface.dev_impl.mm.unmap_range(b.va_addr, round_up(b.size, 0x1000))
+    if b.owner != self.dev:
+      self.dev.iface.dev_impl.mm.unmap_range(b.va_addr, round_up(b.size, 0x1000))
+      if self.pci_dev.iommu and (paddrs:=self.pci_dev.dma_mapped.pop(int(b.va_addr), None)) is not None: self.pci_dev.dma_unmap(paddrs)
     if b.owner == self.dev and b.meta.mapping.aspace is AddrSpace.PHYS: self.dev_impl.mm.vfree(b.meta.mapping)
-    if b.owner == self.dev and self.is_local() and b.meta.has_cpu_mapping: FileIOInterface.munmap(b.va_addr, b.size)
+    if b.owner == self.dev and self.is_local() and b.meta.has_cpu_mapping:
+      if self.pci_dev.iommu and b.meta.mapping.aspace is AddrSpace.SYS: self.pci_dev.dma_unmap([p for p, _ in b.meta.mapping.paddrs])
+      FileIOInterface.munmap(b.va_addr, b.size)
 
   def p2p_paddrs(self, paddrs:list[tuple[int,int]]) -> tuple[list[tuple[int,int]], AddrSpace]:
     return [(p + self.pci_dev.bar_info(self.vram_bar)[0], sz) for p, sz in paddrs], AddrSpace.SYS
@@ -290,8 +378,12 @@ class PCIIfaceBase:
 
       System.lock_memory(int(b.va_addr), b.size)
       paddrs, aspace = [(x, 0x1000) for x in System.system_paddrs(int(b.va_addr), round_up(b.size, 0x1000))], AddrSpace.SYS
+      if self.pci_dev.iommu:
+        self.pci_dev.dma_mapped[int(b.va_addr)] = flat_paddrs = [p for p, _ in paddrs]
+        self.pci_dev.dma_map(int(b.va_addr), flat_paddrs)
       snooped, uncached = True, True
     elif (ifa:=getattr(b.owner, "iface", None)) is not None and isinstance(ifa, PCIIfaceBase):
+      if self.pci_dev.iommu: raise RuntimeError(f"no P2P mappings with an active IOMMU: {b.owner} -> {self.dev} (boot with iommu=pt and VFIO=0)")
       if ifa.is_bar_small(): raise RuntimeError(f"P2P mapping not supported for small bar devices: {b.owner} -> {self.dev}")
 
       snooped, uncached = True, b.meta.mapping.uncached

tinygrad/schedule/rangeify.py

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