weakint issue for symbolic

extra/runbook_digitalocean_mi350x.md

@@ -1,235 +0,0 @@
-# Runbook: Llama 3 8B Training on DigitalOcean MI350X
-
-## Machine Specs
-- 8x MI350X GPUs (gfx950, device ID 75b0), 288GB VRAM each
-- 2TB RAM, 192 CPUs, 2TB disk
-- ROCm 7.14 at `/opt/rocm` (NOT `/opt/rocm-7.1.1` like the submission scripts assume)
-- Python 3.12
-
-## Phase 1: System Setup
-
-### 1.1 Install packages
-```bash
-apt-get update
-apt-get install -y python3-pip python3-venv git tmux rclone clang
-```
-
-### 1.2 Install Python deps
-```bash
-python3 -m pip install --break-system-packages numpy tqdm wandb tiktoken sentencepiece
-```
-
-### 1.3 Install ROCm dev headers
-The base image has ROCm runtime but NOT the HIP dev headers. Need:
-```bash
-apt-get install -y amdrocm-core-dev
-```
-This installs `hip/hip_runtime.h` at `/opt/rocm/core-7.14/include/hip/hip_runtime.h`.
-The symlink `/opt/rocm/include` → `/opt/rocm/core-7.14/include` makes it available at `/opt/rocm/include/hip/hip_runtime.h`.
-
-### 1.4 Configure ROCm comgr
-ROCm 7.14 ships comgr 3.3 at `/opt/rocm/lib/libamd_comgr.so`. tinygrad's DLL loader needs explicit env vars to find it (it searches for `libcomgr.so*` by default, not `libamd_comgr.so*`). Set these in the run command:
-```bash
-export COMGR_PATH=/opt/rocm/lib/libamd_comgr.so
-export COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so
-```
-Also add ROCm libs to ldconfig so comgr's shared library dependencies resolve:
-```bash
-cat > /etc/ld.so.conf.d/rocm.conf << 'EOF'
-/opt/rocm/lib
-/opt/rocm/lib/llvm/lib
-/opt/rocm/lib/rocm_sysdeps/lib
-EOF
-ldconfig
-```
-
-### 1.5 Install geohot tmux config
-```bash
-curl -sL https://raw.githubusercontent.com/geohot/configuration/master/.tmux.conf -o ~/.tmux.conf
-```
-
-### 1.6 Reload amdgpu driver
-tinygrad's HCQ backend needs `/dev/kfd` which is created by the amdgpu kernel driver.
-If the driver was unloaded, reload it:
-```bash
-modprobe amdgpu
-ls /dev/kfd  # should exist
-```
-
-## Phase 2: Clone tinygrad
-```bash
-cd /root
-git clone https://github.com/tinygrad/tinygrad.git
-cd tinygrad
-python3 -m pip install --break-system-packages -e .
-```
-
-## Phase 3: Download C4 Dataset
-
-The C4 data is on the MLCommons Cloudflare R2 bucket in Megatron-LM indexed format.
-
-```bash
-rclone config create mlc-training s3 provider=Cloudflare \
-  access_key_id=76ea42eadb867e854061a1806220ee1e \
-  secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 \
-  endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
-
-mkdir -p /root/datasets/c4-8b
-rclone copy mlc-training:mlcommons-training-wg-public/llama3_1/datasets/c4/llama3_1_8b/ /root/datasets/c4-8b/ -P
-```
-
-Files downloaded (~85GB total, ~6 minutes):
-- `c4-train.en_6_text_document.bin` (79 GB)
-- `c4-train.en_6_text_document.idx` (870 MB)
-- `c4-validation-91205-samples.en_text_document.bin` (159 MB)
-- `c4-validation-91205-samples.en_text_document.idx` (1.8 MB)
-- `LICENSE.txt`, `NOTICE.txt`
-
-### Symlink for the submission script
-The `dev_run.sh` script hardcodes `BASEDIR="/raid/datasets/c4-8b/"`. Symlink:
-```bash
-mkdir -p /raid/datasets
-ln -s /root/datasets/c4-8b /raid/datasets/c4-8b
-```
-
-## Phase 4: wandb Login
-```bash
-wandb login
-```
-Enter API key from https://wandb.ai/authorize
-
-## Phase 5: Run Training
-
-### 5.1 Smoke test (beam search, 2 layers, fake data)
-Always run beam first to validate the pipeline:
-```bash
-cd /root/tinygrad
-COMGR_PATH=/opt/rocm/lib/libamd_comgr.so \
-COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so \
-CC=/opt/rocm/core-7.14/lib/llvm/bin/clang \
-DEV=AMD:HIP \
-ROCM_PATH=/opt/rocm BASEDIR=/root/datasets/c4-8b/ \
-  bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh
-```
-
-### 5.2 Full training run
-```bash
-cd /root/tinygrad
-COMGR_PATH=/opt/rocm/lib/libamd_comgr.so \
-COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so \
-CC=/opt/rocm/core-7.14/lib/llvm/bin/clang \
-DEV=AMD:HIP \
-ROCM_PATH=/opt/rocm BASEDIR=/root/datasets/c4-8b/ \
-WANDB=1 \
-  bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh
-```
-
-## Environment Variable Reference
-
-| Variable | Value | Why |
-|---|---|---|
-| `COMGR_PATH` | `/opt/rocm/lib/libamd_comgr.so` | tinygrad's DLL loader needs explicit path to find comgr 3.3 |
-| `COMGR_3_PATH` | `/opt/rocm/lib/libamd_comgr.so` | comgr 3.x uses a separate `comgr_3` module with its own path var |
-| `CC` | `/opt/rocm/core-7.14/lib/llvm/bin/clang` | System clang doesn't know gfx950; must use ROCm's bundled clang |
-| `DEV` | `AMD:HIP` | Force HIPRenderer (comgr-based) over HIPCCRenderer (hipcc subprocess) |
-| `ROCM_PATH` | `/opt/rocm` | Script defaults to `/opt/rocm-7.1.1` which doesn't exist |
-| `BASEDIR` | `/root/datasets/c4-8b/` | Where C4 dataset was downloaded (script hardcodes `/raid/datasets/c4-8b/`) |
-| `WANDB` | `1` | Enable wandb logging (off by default) |
-
-## Architecture
-
-| Component | Source file |
-|---|---|
-| Model | `examples/mlperf/models/flat_llama.py` — FlatTransformer, FP8 MXFP4 weights, fused QKV, flash attention |
-| Trainer | `examples/mlperf/model_train.py` → `train_llama3()` |
-| Optimizer | `examples/mlperf/optim.py` — GradAccClipAdamW, master weights, FP8 re-quant |
-| LR schedule | `examples/mlperf/lr_schedulers.py` — CosineAnnealingLRWithWarmup |
-| Dataloader | `examples/mlperf/dataloader.py` — Megatron-LM indexed bin format |
-| ASM GEMM | `extra/gemm/cdna_asm_gemm.py` — gfx950 MFMA assembly, MXFP4 |
-| Flash attention | `extra/thunder/amd/fa.py` |
-| Fused kernels | `extra/llama_kernels/` — rmsnorm, silu, quantize, fused_ce |
-| GPU driver | `tinygrad/runtime/ops_amd.py` — HCQ, direct KFD ioctl |
-| Renderer | `tinygrad/renderer/cstyle.py` — HIPRenderer for gfx950 |
-| comgr compiler | `tinygrad/runtime/support/compiler_amd.py` — HIPCompiler using comgr 3.3 |
-
-## Troubleshooting
-
-### `'hip/hip_runtime.h' file not found`
-Install `amdrocm-core-dev`:
-```bash
-apt-get install -y amdrocm-core-dev
-```
-
-### `'gfx950' is not a recognized processor` + LLVM crash
-System clang doesn't know gfx950. Set `CC=/opt/rocm/core-7.14/lib/llvm/bin/clang`.
-
-### `comgr not available: try setting COMGR_PATH?`
-Add ROCm libs to ldconfig and set `COMGR_PATH` and `COMGR_3_PATH`:
-```bash
-# /etc/ld.so.conf.d/rocm.conf should contain /opt/rocm/lib paths
-ldconfig
-```
-
-### `comgr not available: try setting COMGR_3_PATH?`
-comgr 3.x uses a separate module. Set `COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so` too.
-
-### `FileNotFoundError: '/raid/datasets/c4-8b/...'`
-Script hardcodes `BASEDIR`. Either symlink or edit the script:
-```bash
-mkdir -p /raid/datasets && ln -s /root/datasets/c4-8b /raid/datasets/c4-8b
-```
-
-### `No such file or directory: 'clang'`
-Install clang: `apt-get install -y clang` (for CPU compilation).
-For gfx950 HIP compilation, comgr (not clang) is used — ensure the ROCm 7.14 comgr 3.3 is properly loaded via `COMGR_PATH` and `COMGR_3_PATH`.
-
-## Appendix: KVM Virtualization Observations
-
-### Virtualization detection
-```
-$ systemd-detect-virt
-kvm
-$ lspci -nn | grep AMD
-83:00.0 ... Device [1002:75b0]
-```
-CPU flags include `hypervisor`. `dmesg` shows `Hypervisor detected: KVM`.
-
-### PCI device ID
-`lspci -v` shows device ID `0x75b0` and subsystem ID `0x75a0`:
-```
-83:00.0 Processing accelerators: ... Device 75b0
-    Subsystem: ... Device 75a0
-```
-tinygrad's `PCIIface` in `ops_amd.py` and `hive_reset.py` did not list `0x75b0`, so the GPU was not found. Adding `0x75b0` to the device ID list in both files fixes the detection.
-
-### amdgpu driver behavior
-On first boot, amdgpu loaded and bound to all 8 GPUs. On one boot it failed to initialize:
-```
-[  799.780369] amdgpu 0000:83:00.0: Failed to alloc msi vectors
-[  799.781476] amdgpu 0000:83:00.0: sw_init of IP block <vega20_ih> failed -22
-[  799.782724] amdgpu 0000:83:00.0: amdgpu_device_ip_init failed
-[  799.793885] amdgpu 0000:83:00.0: Fatal error during GPU init
-```
-On a subsequent boot, amdgpu initialized successfully (SMU initialized, VRAM ready). After unbinding all 8 GPUs from amdgpu, `rmmod amdgpu` wedged the module (stuck in "Unloading" state in `/proc/modules`), requiring a full VM reboot.
-
-### `/dev/kfd`
-`/dev/kfd` exists when amdgpu is loaded. Opening it returns `OSError: [Errno 22] Invalid argument`.
-
-### VRAM BAR reads all 0xFF
-After amdgpu initializes the GPU and is then unbound, reading the VRAM BAR (via `/sys/bus/pci/devices/0000:83:00.0/resource0`) returns all `0xFF` at all offsets — including the discovery table at `vram_size - 64KB`. tinygrad's `AMDev._run_discovery()` fails with `AssertionError: discovery signatures mismatch`.
-
-A PCI reset (`echo 1 > /sys/bus/pci/devices/0000:83:00.0/reset`) did not change the VRAM contents — still all `0xFF`.
-
-VRAM was also all `0xFF` when read via `/dev/mem` at the BAR physical address (`0xa0000000000`).
-
-### VFIO attempt
-Bound the GPU to `vfio-pci` with `enable_unsafe_noiommu_mode=1`. The GPU bound successfully and `/dev/vfio/noiommu-0` appeared. Running tinygrad with `VFIO=1` still failed with the same `discovery signatures mismatch` — VRAM BAR still reads all `0xFF`.
-
-### No IOMMU in guest
-`dmesg` has no `AMD-Vi` entries. PCI devices have no `iommu_group` symlink.
-
-### No fan control
-No `fan*` or `pwm*` hwmon entries exist. Only `temp*`, `power*`, `freq*` are exposed. GPU temps read 56-63°C, power ~265W per GPU.
-
-### Current status: NOT WORKING
-tinygrad's `PCIIface` finds the GPU (after adding `0x75b0`) but `AMDev._run_discovery()` fails because the VRAM discovery table reads all `0xFF`. This was observed with the GPU unbound from any driver, after PCI reset, and with VFIO bound.

tinygrad/engine/realize.py

@@ -5,7 +5,7 @@ from dataclasses import dataclass, replace, field
 from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansilen, all_int, prod, flatten, Context, getenv, to_tuple
 from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events, wait_cond
 from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, buffers, graph_rewrite
-from tinygrad.device import Device, Buffer, MultiBuffer, ProfileGraphEntry
+from tinygrad.device import Device, Buffer, MultiBuffer
 from tinygrad.renderer import Estimates
 from tinygrad.codegen import to_program
 from tinygrad.codegen.opt.postrange import args_from_ast
@@ -210,30 +210,27 @@ def exec_graph(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
   return t[0]
 
 def exec_hcq(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
-  if (info:=call.arg.aux).inputs is not None:
+  if (inputs:=call.arg.aux.inputs) is not None:
     bufs = [_resolve(ctx.input_uops[i], ctx.input_uops).buffer for i in call.arg.aux.input_idxs]
-    table = call.src[1+info.inputs].buffer
+    table = call.src[1+inputs].buffer
     for j,dev in enumerate(call.arg.aux.device):
       addrs = array.array('Q', [(b.bufs[j] if isinstance(b, MultiBuffer) else b).get_buf(dev).va_addr for b in bufs])
       mv = (table.bufs[j] if isinstance(table, MultiBuffer) else table).ensure_allocated()._buf.cpu_view().view(fmt='Q')
       wait_cond(lambda: mv[0], value=0, timeout_ms=ctx.timeout or getenv("HCQDEV_WAIT_TIMEOUT_MS", 30000), msg=f"{dev} hang detected")
       mv[:len(addrs)] = addrs
 
-  exec_kernel(replace(ctx, update_stats=DEBUG>=3), call, ast)
+  exec_kernel(replace(ctx, update_stats=False), call, ast)
 
-  tms = []
-  for devices,name,estimates,prof in info.kernels:
-    for device in devices:
-      d, tm = cast(Any, Device[device]), None
-      if prof:
-        d.prof_ents[prof[0]] = ProfileGraphEntry(device, name, *prof)
-        if ctx.wait:
-          d.synchronize(timeout=ctx.timeout)
-          st, en = (d.signal(x)._buf.cpu_view().view(fmt='Q')[0] for x in prof)
-          tms.append(tm:=float(en-st)/d.timestamp_divider/1e6)
-      with track_stats(ctx, call.replace(arg=replace(call.arg, name=name, aux=replace(info, estimates=estimates))), d.device, [], ctx.var_vals) as et:
-        et[0] = tm
-  return max(tms) if tms else None
+  tms:list[float|None] = []
+  for e in (aux:=call.arg.aux).prof: cast(Any, Device[e.device]).prof_ents[e.st_id] = e
+  for d in [cast(Any, Device[x]) for x in aux.device]:
+    with track_stats(ctx, call, d.device, [], ctx.var_vals) as et:
+      if ctx.wait:
+        d.synchronize(timeout=ctx.timeout)
+        ts = [d.signal(i)._buf.cpu_view().view(fmt='Q')[0] for e in aux.prof if e.device == d.device for i in (e.st_id, e.en_id)]
+        if ts: et[0] = float(max(ts)-min(ts))/d.timestamp_divider/1e6
+      tms += et
+  return tms[0]
 
 # flatten LINEAR-in-LINEAR: any nested LINEAR child gets inlined into its parent's src
 pm_flatten_linear = PatternMatcher([
@@ -279,7 +276,7 @@ def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:li
   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"): linear = hcq_compile(linear, input_uops, bool(PROFILE or DEBUG >= 2) if profile is None else profile)
+  if getenv("HCQ2"): linear = hcq_compile(linear, input_uops, bool(PROFILE) if profile is None else profile)
   return graph_rewrite(linear, pm_optimize_local_size, name="optimize local size", walk=True)
 
 def link_linear(linear:UOp, cache=True) -> UOp: return hcq_link(linear, cache=cache) if getenv("HCQ2") else linear

tinygrad/runtime/support/hcq2.py

@@ -33,7 +33,7 @@ class HCQInfo:
 
   input_idxs:tuple[int, ...] = () # indexes into input_uops used by this call
   inputs:int|None = None
-  kernels:tuple[tuple[tuple[str, ...], str, Estimates, tuple[int, ...]], ...] = ()
+  prof:tuple[ProfileGraphEntry, ...] = () # st_id/en_id are timestamp signal slots until collect
 
 def all_devices_in(d:Any, c:frozenset[str]) -> bool: return {x.split(":")[0] for x in to_tuple(d)} <= c
 
@@ -199,10 +199,10 @@ def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]], profile:bool) -> li
     signal_tags |= cur_signal_tags
 
   # build fences and finalizers
-  fences, fins, finalizer_signal_tags = _build_finalizers(batch, batch_info, deps_tracker, slots)
+  fences, finalizers, finalizer_signal_tags = _build_finalizers(batch, batch_info, deps_tracker, slots)
   signal_tags |= finalizer_signal_tags
 
-  src, kerns = [], []
+  src, prof = [], []
   for tag, ((call, _), (devices, queue), q) in enumerate(zip(batch, batch_info, call_waits)):
     # first queue use, sync prior device work with the device timeline
     if batch_info.index((devices, queue)) == tag:
@@ -212,18 +212,18 @@ def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]], profile:bool) -> li
     # and make hcq call
     name, info = get_call_name(call, get_call_arg_uops(call)), HCQInfo(devices, estimate_uop(call))
     ts_ids = [next(UOp.unique_num) for _ in range(2)] if profile else []
-    kerns.append((devices, name, info.estimates, tuple(ts_ids)))
+    prof += [ProfileGraphEntry(d, name, *ts_ids) for d in devices if ts_ids]
 
     ts_ins = [UOp(Ops.INS, arg="timestamp", src=(make_signal(devices, s),)) for s in ts_ids]
     q += ts_ins[:1] + [call.replace(arg=replace(call.arg, aux=info))] + ts_ins[1:]
 
     # signal the queue if someone waits for us
     if tag in signal_tags: q += [UOp(Ops.INS, arg="store", src=(make_signal(devices, slots[queue]), UOp.const(tag + 1, dtypes.uint64)))]
-    src.append(make_call(f"submit {name}", make_submit(*q, devs=devices, queue=queue).sink(), info))
+    src.append(make_call(name, make_submit(*q, devs=devices, queue=queue).sink(), info))
 
   # append batch timestamps to finalizers
-  fins = [f.replace(arg=replace(f.arg, aux=replace(a:=f.arg.aux, kernels=tuple(x for x in kerns if set(x[0]) & set(a.device))))) for f in fins]
-  return fences + src + fins
+  finalizers = [f.replace(arg=replace(f.arg, aux=replace(a:=f.arg.aux, prof=tuple(e for e in prof if e.device in a.device)))) for f in finalizers]
+  return fences + src + finalizers
 
 def sched_hcq_batches(l:UOp, profile:bool) -> UOp:
   srcs:list[UOp] = []

tinygrad/schedule/__init__.py

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

tinygrad/schedule/rangeify2.py

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