Merge branch 'master' into simple_rangeify

docs/developer/am.md

@@ -7,13 +7,8 @@ Make sure that amdgpu module is unloaded and just run tinygrad with `DEV=AMD`!
 
 Optional requirements:
 
-* 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`.
+* System without IOMMU for P2P / SDMA support
+* vfio-pci module for IRQ handling
 
 ## Environment Variables
 
@@ -21,7 +16,6 @@ without address translation: boot with `iommu=pt` and set `VFIO=0`.
 |----------|------------------|-------------|
 | 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

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