Merge branch 'master' into replace_if_with_range

zero len ranges fail (#12974 )
* zero len ranges fail * fix Python backend * fix llvm * fix ptx * yolo fix nir * this works... * always store... * always store... * Revert "always store..." This reverts commit 0816cf344d.
2025-10-28 23:02:15 +08:00 · 2025-10-28 22:49:55 +08:00 · 2025-10-28 20:58:47 +08:00 · 2025-10-28 19:28:38 +08:00 · 2025-10-28 19:22:59 +08:00 · 2025-10-28 19:15:01 +08:00
107 changed files with 1986 additions and 1777 deletions
@@ -238,6 +238,8 @@ jobs:
      run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_beam.txt
    - name: Run LLaMA-3 8B on 4 GPUs with BEAM
      run: BENCHMARK_LOG=llama3_beam_4gpu NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_four_gpu.txt
+    - name: Run quantized LLaMA3
+      run: BENCHMARK_LOG=llama3_fp8 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --temperature 0 --benchmark --quantize fp8 | tee llama3_fp8.txt
    # - name: Run LLaMA-3 8B on 6 GPUs
    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_six_gpu.txt
    # - name: Run LLaMA-2 70B
@@ -271,6 +273,7 @@ jobs:
          llama3_beam.txt
          llama3_four_gpu.txt
          llama3_six_gpu.txt
+          llama3_fp8.txt
          llama_2_70B.txt
          mixtral.txt
          gpt2_unjitted.txt
@@ -319,9 +322,9 @@ jobs:
    - name: Run 10 CIFAR training steps
      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=270 NV=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
    - name: Run 10 CIFAR training steps w HALF
-      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=310 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=240 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
    - name: Run 10 CIFAR training steps w BF16
-      run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=310 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
+      run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=270 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
    # TODO: too slow
    # - name: Run 10 CIFAR training steps w winograd
    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=350 NV=1 CAPTURE_PROCESS_REPLAY=0 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
@@ -619,18 +622,24 @@ jobs:
        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
    - name: reset process replay
      run: test/external/process_replay/reset.py
-    - name: benchmark openpilot 0.9.9 driving_vision
-      run: BENCHMARK_LOG=openpilot_0_9_9_vision PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
-    - name: benchmark openpilot 0.9.9 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_9_9_policy PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
-    - name: benchmark openpilot 0.9.9 dmonitoring
-      run: BENCHMARK_LOG=openpilot_0_9_9_dmonitoring PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
+    - name: openpilot compile3 0.9.9 driving_vision
+      run: BENCHMARK_LOG=openpilot_0_9_9_vision PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
+    - name: openpilot compile3 0.9.9 driving_policy
+      run: BENCHMARK_LOG=openpilot_0_9_9_policy PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
+    - name: openpilot compile3 0.9.9 dmonitoring
+      run: BENCHMARK_LOG=openpilot_0_9_9_dmonitoring PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
+    - name: openpilot compile3 0.10.0 driving_policy
+      run: BENCHMARK_LOG=openpilot_0_10_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=5 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/driving_policy.onnx
+    - name: openpilot compile3 0.10.0 dmonitoring
+      run: BENCHMARK_LOG=openpilot_0_10_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=13 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/dmonitoring_model.onnx
    - name: openpilot compile3 0.10.1 driving_vision
-      run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=25 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+      # TODO: ASSERT_MIN_STEP_TIME=17
+      run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=25 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
    - name: openpilot compile3 0.10.1 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=7 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/refs/heads/master/selfdrive/modeld/models/driving_policy.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=5 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
    - name: openpilot compile3 0.10.1 dmonitoring
-      run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=12 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/refs/heads/master/selfdrive/modeld/models/dmonitoring_model.onnx
+      # TODO: ASSERT_MIN_STEP_TIME=10
+      run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=13 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
    - name: benchmark MobileNetV2 on DSP
      run: |
        # generate quantized weights
@@ -641,16 +650,6 @@ jobs:
        PYTHONPATH=. CC=clang-19 DSP=1 NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
    - name: Run process replay tests
      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
-    - uses: actions/upload-artifact@v4
-      with:
-        name: Speed (comma)
-        path: |
-          openpilot_compile_0_9_4.txt
-          openpilot_compile_0_9_7.txt
-          openpilot_0_9_4.txt
-          openpilot_0_9_7.txt
-          openpilot_image_0_9_4.txt
-          openpilot_image_0_9_7.txt

  testreddriverbenchmark:
    name: AM Benchmark
@@ -204,7 +204,7 @@ jobs:
        DEBUG=2 EMULATE=CUDA FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
        DEBUG=2 EMULATE=CUDA ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
        DEBUG=2 EMULATE=CUDA_SM75 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
-        DEBUG=2 EMULATE=CUDA ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
+        DEBUG=2 EMULATE=CUDA_SM89 ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
    - name: Test emulated INTEL OpenCL tensor cores
      run: DEBUG=2 EMULATE=INTEL FORWARD_ONLY=1 PYTHON=1 HALF=1 N=64 python3 ./extra/gemm/simple_matmul.py
    - name: Test emulated AMX tensor cores
@@ -264,8 +264,6 @@ jobs:
      run: python -c "from tinygrad import Device; assert Device.DEFAULT == 'CPU', Device.DEFAULT"
    - name: Run unit tests
      run: CPU=1 python -m pytest -n=auto test/unit/ --durations=20
-    - name: Check SPEC=1
-      run: SPEC=1 python3 test/test_tiny.py
    - name: Run targetted tests on NULL backend
      run: NULL=1 python3 -m unittest test.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step test/device/test_null.py
    # TODO: too slow
@@ -294,6 +292,25 @@ jobs:
    - name: Repo line count < 18000 lines
      run: MAX_LINE_COUNT=18000 python sz.py

+  spec:
+    strategy:
+      fail-fast: false
+      matrix:
+        group: [1, 2]
+    name: SPEC=2 (${{ matrix.group }})
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: spec-unit
+        deps: testing_unit
+    - name: Test SPEC=2
+      run: IGNORE_OOB=0 SPEC=2 PYTHONPATH="." pytest --maxfail=10 -n auto --durations=30 --ignore=test/models --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}
+
  fuzzing:
    name: Fuzzing
    runs-on: ubuntu-latest
@@ -351,7 +368,7 @@ jobs:
      - name: Run Kernel Count Test
        run: CL=1 python -m pytest -n=auto test/external/external_test_opt.py
      - name: Run fused optimizer tests
-        run: CL=1 FUSE_OPTIM=1 python -m pytest -n=auto test/models/test_mnist.py
+        run: CL=1 FUSE_OPTIM=1 python -m pytest -n=auto test/models/test_mnist.py test/test_optim.py -k "not muon"
      - name: Upload artifact
        uses: actions/upload-artifact@v4
        with:
@@ -378,7 +395,7 @@ jobs:
      - name: Test openpilot CL compile fp16
        run: FLOAT16=1 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
      - name: Test openpilot CL compile fp32 (test correctness)
-        run: DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/haraschax/filedump/raw/refs/heads/master/driving_vision_fp32.onnx
+        run: DEBUGCL=1 CL=1 IMAGE=2 SELFTEST=1 python examples/openpilot/compile3.py https://github.com/haraschax/filedump/raw/refs/heads/master/driving_vision_fp32.onnx
      - name: Test openpilot LLVM compile fp16
        run: FLOAT16=1 CPU=1 CPU_LLVM=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
      - name: Run process replay tests
@@ -522,11 +539,11 @@ jobs:
        pydeps: "pillow"
        llvm: "true"
    - name: Test LLVM=1 DEVECTORIZE=0
-      run: CPU=1 CPU_LLVM=1 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"
+      run: CPU=1 CPU_LLVM=1 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py
    - name: Test LLVM=1 DEVECTORIZE=0 for model
      run: CPU=1 CPU_LLVM=1 DEVECTORIZE=0 python3 test/models/test_efficientnet.py
    - name: Test CPU=1 DEVECTORIZE=0
-      run: CPU=1 CPU_LLVM=0 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"
+      run: CPU=1 CPU_LLVM=0 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py

  testdsp:
    name: Linux (DSP)
@@ -28,7 +28,7 @@ repos:
        pass_filenames: false
      - id: tests
        name: subset of tests
-        entry: env PYTHONPATH="." python3 -m pytest -n=8 test/test_ops.py test/test_dtype.py test/test_schedule.py test/test_assign.py
+        entry: env OMP_NUM_THREADS=1 PYTHONPATH="." python3 -m pytest -n=8 test/test_ops.py test/test_dtype.py test/test_schedule.py test/test_assign.py
        language: system
        always_run: true
-        pass_filenames: false
+        pass_filenames: false
@@ -520,13 +520,8 @@ generate_mesa() {
  LVP_NIR_OPTIONS=$(./extra/mesa/lvp_nir_options.sh $MESA_SRC)

  fixup $BASE/mesa.py
-  patch_dlopen $BASE/mesa.py tinymesa_cpu "(BASE:=os.getenv('MESA_PATH', f\"/usr{'/local/' if helpers.OSX else '/'}lib\"))+'/libtinymesa_cpu'+(EXT:='.dylib' if helpers.OSX else '.so')" "f'{BASE}/libtinymesa{EXT}'" "brew_path('tinymesa_cpu')" "brew_path('tinymesa')"
+  patch_dlopen $BASE/mesa.py tinymesa_cpu "(BASE:=os.getenv('MESA_PATH', f\"/usr{'/local/' if helpers.OSX else '/'}lib\"))+'/libtinymesa_cpu'+(EXT:='.dylib' if helpers.OSX else '.so')" "f'{BASE}/libtinymesa{EXT}'" "'/opt/homebrew/lib/libtinymesa_cpu.dylib'" "'/opt/homebrew/lib/libtinymesa.dylib'"
  echo "lvp_nir_options = gzip.decompress(base64.b64decode('$LVP_NIR_OPTIONS'))" >> $BASE/mesa.py
-  cat <<EOF | sed -i "/import ctypes.*/r /dev/stdin" $BASE/mesa.py
-def brew_path(nm):
-  try: return f"{subprocess.check_output(['brew', '--prefix', nm]).decode().strip()}/lib/lib{nm}.dylib"
-  except Exception: return 'failed'
-EOF
  sed -i "/in_dll/s/.*/try: &\nexcept (AttributeError, ValueError): pass/" $BASE/mesa.py
  sed -i "s/import ctypes/import ctypes, ctypes.util, os, gzip, base64, subprocess, tinygrad.helpers as helpers/" $BASE/mesa.py
  sed -i "s/ctypes.CDLL('.\+')/(dll := _try_dlopen_tinymesa_cpu())/" $BASE/mesa.py
@@ -1,109 +0,0 @@
-# Kernel Creation
-
-Tinygrad lazily builds up a graph of Tensor operations. The Tensor graph includes a mix of:
-
- Buffer and Assignment Ops: `BUFFER`, `BUFFER_VIEW`, `COPY`, `ASSIGN`
- Movement Ops: `RESHAPE`, `EXPAND`, `PERMUTE`, `PAD`, `SHRINK`, `FLIP`
- Compute Ops: `ADD`, `MUL`, `REDUCE_AXIS`, ...
-
-`Tensor.kernelize` creates the kernels and buffers needed to realize the output Tensor(s).
-
-## Kernelize flow
-
-Let's see how a multiply add Tensor graph becomes a fused elementwise kernel.
-
-```py
-# initialize 3 input buffers on the device
-a = Tensor([1]).realize()
-b = Tensor([2]).realize()
-c = Tensor([3]).realize()
-
-# create the Tensor graph
-mul = a*b
-out = mul+c
-
-print(mul) # <Tensor <UOp METAL (1,) int (<Ops.MUL: 48>, None)> on METAL with grad None>
-print(out) # <Tensor <UOp METAL (1,) int (<Ops.ADD: 52>, None)> on METAL with grad None>
-
-out.kernelize()
-
-print(mul) # <Tensor <UOp METAL (1,) int (<Ops.MUL: 48>, None)> on METAL with grad None>
-print(out) # <Tensor <UOp METAL (1,) int (<Ops.ASSIGN: 66>, None)> on METAL with grad None>
-```
-
-The multiply Tensor stays the same because it is fused. The output Tensor's UOp becomes a new ASSIGN UOp:
-
-```py
-print(out.uop)
-```
-
-The first source is the output BUFFER:
-
-```
-UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
-  UOp(Ops.DEVICE, dtypes.void, arg='METAL', src=()),
-  UOp(Ops.UNIQUE, dtypes.void, arg=6, src=()),))
-```
-
-And the second source is the KERNEL and its 4 buffer edges (output_buffer, a, b, c):
-
-```
-UOp(Ops.KERNEL, dtypes.void, arg=<Kernel 12 SINK(<Ops.STORE: 45>,) (__add__, __mul__)>, src=(
-  UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
-    x1:=UOp(Ops.DEVICE, dtypes.void, arg='METAL', src=()),
-    UOp(Ops.UNIQUE, dtypes.void, arg=6, src=()),)),
-  UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
-     x1,
-    UOp(Ops.UNIQUE, dtypes.void, arg=1, src=()),)),
-  UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
-     x1,
-    UOp(Ops.UNIQUE, dtypes.void, arg=3, src=()),)),
-  UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
-     x1,
-    UOp(Ops.UNIQUE, dtypes.void, arg=5, src=()),)),))
-```
-
-KERNEL describes the compute AST, metadata and memory dependencies.
-
-BUFFER holds a reference to the device memory where the output will be stored.
-
-Once a Tensor is kernelized, all children will LOAD its BUFFER, instead of fusing it:
-
-```py
-child = out+2
-child.kernelize()
-print(child.uop.src[1].arg.ast)
-```
-
-```
-UOp(Ops.SINK, dtypes.void, arg=None, src=(
-  UOp(Ops.STORE, dtypes.void, arg=None, src=(
-    UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1), arg=0, src=()),
-    x2:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1,), strides=(0,), offset=0, mask=None, contiguous=True),)), src=()),
-    UOp(Ops.ADD, dtypes.int, arg=None, src=(
-      UOp(Ops.LOAD, dtypes.int, arg=None, src=(
-        UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1), arg=1, src=()),
-         x2,)),
-      UOp(Ops.CONST, dtypes.int, arg=2, src=(
-         x2,)),)),)),))
-```
-
-`Tensor.realize` will execute the kernels and write outputs to memory:
-
-```py
-Tensor.realize(out)
-print(out)        # <Tensor <UOp METAL (1,) int (<Ops.BUFFER: 23>, <buf real:True device:METAL size:1 dtype:dtypes.int offset:0>)> on METAL with grad None>
-print(out.item()) # 5
-```
-
-<hr />
-
-**Summary**
-
- The large Tensor graph is built from a mix of data, compute and movement Ops.
-
- `Tensor.kernelize` splits the Tensor graph into data (BUFFER), compute (KERNEL) and links dependencies with ASSIGN.
-
- `Tensor.realize` executes KERNELs on device and replaces the Tensor graph with just a BUFFER.
-
- Kernelize can be called multiple times on a Tensor. This allows for incrementally building the kernel fusion layout of a large Tensor graph, without having to call `realize` or `schedule`.
@@ -41,7 +41,7 @@ BEAM                | [#]        | number of beams in kernel beam search
 DEFAULT_FLOAT       | [HALF, ...]| specify the default float dtype (FLOAT32, HALF, BFLOAT16, FLOAT64, ...), default to FLOAT32
 IMAGE               | [1-2]      | enable 2d specific optimizations
 FLOAT16             | [1]        | use float16 for images instead of float32
-VISIBLE_DEVICES     | [list[int]]| restricts the NV/AMD devices that are available. The format is a comma-separated list of identifiers (indexing starts with 0).
+HCQ_VISIBLE_DEVICES | [list[int]]| restricts the HCQ devices that are available. The format is a comma-separated list of identifiers (indexing starts with 0).
 JIT                 | [0-2]      | 0=disabled, 1=[jit enabled](quickstart.md#jit) (default), 2=jit enabled, but graphs are disabled
 VIZ                 | [1]        | 0=disabled, 1=[viz enabled](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/viz)
 ALLOW_TF32          | [1]        | enable TensorFloat-32 tensor cores on Ampere or newer GPUs.
@@ -145,6 +145,41 @@ def NF4Linear(block_size):
      return new_state_dict
  return _NF4Linear

+def quantize_to_fp8(x: Tensor, dtype=dtypes.fp8e4m3):
+  fp8_min = -448.0 if dtype == dtypes.fp8e4m3 else -57344.0
+  fp8_max = 448.0 if dtype == dtypes.fp8e4m3 else 57344.0
+  scale = fp8_max / x.abs().max()
+  x_scl_sat = (x * scale).clamp(fp8_min, fp8_max)
+  return x_scl_sat.cast(dtype), scale.float().reciprocal()
+
+class FP8Linear:
+  def __init__(self, in_features, out_features, bias=True):
+    self.weight = Tensor.empty(out_features, in_features, dtype=dtypes.fp8e4m3)
+    self.bias = Tensor.empty(out_features, dtype=dtypes.float16) if bias else None
+    self.weight_scale = Tensor.empty((), dtype=dtypes.float16)
+
+  def __call__(self, x:Tensor):
+    y = x.dot(self.weight.T.cast(dtypes.float32)) * self.weight_scale
+    if self.bias is not None: y = y + self.bias.cast(y.dtype)
+    return y.cast(x.dtype)
+
+  @staticmethod
+  def quantize(tensors, device, scale_dtype=dtypes.float16, quantize_embeds=False):
+    assert not quantize_embeds
+    new_tensors = {}
+    for name,v in tensors.items():
+      if "feed_forward" in name or "attention.w" in name:
+        assert "weight" in name, name
+        fp8_weight, scale = quantize_to_fp8(v)
+        new_tensors[name] = fp8_weight
+        new_tensors[name.replace('weight', 'weight_scale')] = scale.cast(scale_dtype)
+        if isinstance(device, tuple):
+          new_tensors[name].shard_(device, axis=-1)
+          new_tensors[name.replace('weight', 'weight_scale')].shard_(device, axis=None)
+      else:
+        new_tensors[name] = v
+    return new_tensors
+
 MODEL_PARAMS = {
  "1B": {
    "args": {"dim": 2048, "n_heads": 32, "n_kv_heads": 8, "n_layers": 16, "norm_eps": 1e-5, "rope_theta": 500000, "vocab_size": 128256, "hidden_dim": 8192},
@@ -167,6 +202,7 @@ def build_transformer(model_path: Path, model_size="8B", quantize=None, scale_dt
  # build model
  if quantize == "int8": linear, embedding, quantize_embeds = Int8Linear, Int8Embedding, True
  elif quantize == "nf4": linear, embedding, quantize_embeds = NF4Linear(64), nn.Embedding, False
+  elif quantize == "fp8": linear, embedding, quantize_embeds = FP8Linear, nn.Embedding, False
  else: linear, embedding, quantize_embeds = nn.Linear, nn.Embedding, False
  model = Transformer(**MODEL_PARAMS[model_size]["args"], linear=linear, embedding=embedding, max_context=max_context, jit=True)

@@ -242,7 +278,7 @@ if __name__ == "__main__":
  parser.add_argument("--model", type=Path, help="Model path")
  parser.add_argument("--size", choices=["1B", "8B", "70B", "405B"], default="1B", help="Model size")
  parser.add_argument("--shard", type=int, default=1, help="Shard the model across multiple devices")
-  parser.add_argument("--quantize", choices=["int8", "nf4", "float16"], help="Quantization method")
+  parser.add_argument("--quantize", choices=["int8", "nf4", "float16", "fp8"], help="Quantization method")
  parser.add_argument("--no_api", action="store_true", help="Disable the api and run a cli test interface")
  parser.add_argument("--host", type=str, default="0.0.0.0", help="Web server bind address")
  parser.add_argument("--port", type=int, default=7776, help="Web server port")
@@ -134,7 +134,7 @@ if __name__ == "__main__":
  with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f)

  test_vs_compile(pickle_loaded, inputs, outputs)
-  if not getenv("FLOAT16"):
+  if getenv("SELFTEST"):
    test_vs_onnx(inputs, outputs, onnx_file, 1e-4)

  if getenv("BENCHMARK_LOG", ""):
@@ -0,0 +1,4 @@
+# source extra/cl_android.sh
+export LD_LIBRARY_PATH=/data/data/com.termux/files/usr/lib:/system/vendor/lib64
+export LD_PRELOAD=/system/vendor/lib64/libOpenCL.so
+
@@ -328,8 +328,7 @@ if __name__ == "__main__":
  elif HL == 1: hprg = hl_spec_kernel3()
  else: hprg = hand_spec_kernel3()
  if HL == 3:
-    with Context(BLOCK_REORDER=0):
-      prg = get_program(hprg, Device.default.renderer)
+    prg = get_program(hprg, Device.default.renderer)
  else:
    prg = get_program(hprg, Device.default.renderer)
  print(prg.src)
@@ -5,8 +5,10 @@ from tinygrad.dtype import _to_np_dtype
 from tinygrad.codegen.opt import OptOps
 from tinygrad.engine.realize import lower_schedule

-dtype_in = dtypes.half if getenv("HALF") else dtypes.bfloat16 if getenv("BFLOAT16") else dtypes.float
-acc_dtype = dtypes.half if getenv("ACC_HALF") else dtypes.bfloat16 if getenv("ACC_BFLOAT16") else None
+dtype_in = (dtypes.half if getenv("HALF") else dtypes.bfloat16 if getenv("BFLOAT16") else
+            dtypes.fp8e4m3 if getenv("FP8E4M3") else dtypes.fp8e5m2 if getenv("FP8E5M2") else dtypes.float)
+acc_dtype = (dtypes.half if getenv("ACC_HALF") else dtypes.bfloat16 if getenv("ACC_BFLOAT16") else
+            dtypes.fp8e4m3 if getenv("ACC_FP8E4M3") else dtypes.fp8e5m2 if getenv("ACC_FP8E5M2") else None)
 if getenv("INT"):  dtype_in, acc_dtype = dtypes.int8, dtypes.int32
 if getenv("UINT"): dtype_in, acc_dtype = dtypes.uint8, dtypes.int32

@@ -14,8 +16,10 @@ N = getenv("N", 4096)
 M = getenv("M", N)
 K = getenv("K", N)
 CNT = getenv("CNT", 10)
-ATOL = getenv("ATOL", 1e-4)
-RTOL = getenv("RTOL", 3e-2)
+
+atol, rtol = {dtypes.bfloat16:(1e-3, 1e-2), dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2:(1.0, 5e-1)}.get(dtype_in, (1e-4, 3e-2))
+ATOL, RTOL = getenv("ATOL", atol), getenv("RTOL", rtol)
+
 INT_LOW = getenv("INT_LOW", 0)
 INT_HIGH = getenv("INT_HIGH", 10)

@@ -8,19 +8,22 @@ import torch
 torch.set_num_threads(1)
 from tinygrad.helpers import getenv
 CUDA = getenv("CUDA", 1)
+MPS = getenv("MPS", 0)

-for dtype in [torch.float32, torch.float16]:
+for dtype in [torch.float32, torch.float16, torch.bfloat16]:
  for N in [256, 512, 1024, 2048, 4096]:
    FLOPS = N*N*N*2

    b = torch.rand((N,N), dtype=dtype)
    c = torch.rand((N,N), dtype=dtype)
    if CUDA: b,c = b.cuda(),c.cuda()
+    if MPS: b,c = b.to('mps'),c.to('mps')

    def torch_prog(b, c):
      st = time.perf_counter()
      a = b@c
      if CUDA: torch.cuda.synchronize()
+      if MPS: torch.mps.synchronize()
      return time.perf_counter() - st
    tm = min([torch_prog(b, c) for _ in range(20)])
    print(f"{N*N:10d} {tm*1e6:9.2f} us, would be {FLOPS*1e-9/tm:9.2f} GFLOPS {N:4d}x{N:4d}x{N:4d} matmul in {dtype}")
@@ -1,7 +1,6 @@
 #!/usr/bin/env python3

-from tinygrad.runtime.support.system import System
-import argparse, glob, os, re, time, subprocess, sys
+import argparse, glob, os, time, subprocess, sys

 def scan_devs_based_on_lock(prefix:str, args) -> list[str]:
  target_dev = args.pci_bus if 'pci_bus' in args.__dir__() else ""
@@ -12,7 +11,7 @@ def scan_devs_based_on_lock(prefix:str, args) -> list[str]:
    if os.path.exists(f"/sys/bus/pci/devices/{dev_id}") and dev_id.startswith(target_dev): devs.append(dev_id)
  return devs

-def _do_reset_device(pci_bus): System.pci_reset(pci_bus)
+def _do_reset_device(pci_bus): os.system(f"sudo sh -c 'echo 1 > /sys/bus/pci/devices/{pci_bus}/reset'")
 def _is_module_loaded(name: str) -> bool: return os.path.isdir(f"/sys/module/{name}")

 def cmd_remove_module(args):
@@ -882,6 +882,11 @@ impl<'a> Thread<'a> {
                            let s1 = sign_ext((s1 & 0xffffff) as u64, 24) as i32;
                            (s0 * s1) as u32
                        }
+                        10 => {
+                            let s0 = sign_ext((s0 & 0xffffff) as u64, 24) as i64;
+                            let s1 = sign_ext((s1 & 0xffffff) as u64, 24) as i64;
+                            ((s0 * s1) >> 32) as u32
+                        }
                        17 | 18 | 26 => {
                            let (s0, s1) = (s0 as i32, s1 as i32);
                            (match op {
@@ -930,7 +935,7 @@ impl<'a> Thread<'a> {

            let op = ((instr >> 16) & 0x3ff) as u32;
            match op {
-                764 | 765 | 288 | 289 | 290 | 766 | 767 | 768 | 769 => {
+                764 | 765 | 288 | 289 | 290 | 766 | 767 | 768 | 769 | 770 => {
                    let vdst = (instr & 0xff) as usize;
                    let sdst = ((instr >> 8) & 0x7f) as usize;
                    let f = |i: u32| -> usize { ((instr >> i) & 0x1ff) as usize };
@@ -996,6 +1001,10 @@ impl<'a> Thread<'a> {
                                    let ret = s0.wrapping_sub(s1);
                                    (ret as u32, s1 > s0)
                                }
+                                770 => {
+                                    let ret = s1.wrapping_sub(s0);
+                                    (ret as u32, s0 > s1)
+                                }
                                _ => todo_instr!(instruction)?,
                            };
                            if self.exec.read() {
@@ -1,98 +1,32 @@
 import numpy as np
 import unittest
 import subprocess, struct, math
-from typing import cast
-from tinygrad.runtime.ops_amd import AMDProgram, AMDDevice
-from tinygrad import Tensor, dtypes, Device
-from tinygrad.helpers import diskcache, OSX, getenv
+from tinygrad import Tensor, dtypes, Device, UOp
+from tinygrad.helpers import getenv
+from tinygrad.runtime.support.compiler_amd import amdgpu_disassemble
+from tinygrad.renderer import ProgramSpec
+from tinygrad.engine.realize import CompiledRunner

-@diskcache
-def assemble(code:str) -> bytes:
-  try:
-    LLVM_MC = "llvm-mc" if OSX else "/opt/rocm/llvm/bin/llvm-mc"
-    return subprocess.run([LLVM_MC, "--arch=amdgcn", "--mcpu=gfx1100", "--triple=amdgcn-amd-amdhsa", "-filetype=obj", "-o", "-"],
-                           input=code.encode("utf-8"), stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True).stdout
-  except subprocess.CalledProcessError as e:
-    print("stderr:")
-    print(e.stderr.decode())
-    raise
-
-# copied from extra/rdna
-def get_prg(code:str, v_cnt:int, s_cnt:int):
-  function_name = "test"
-  metadata = f"""
-amdhsa.kernels:
- .args:
-  - .address_space: global
-    .name: buf_0
-    .offset: 0
-    .size: 8
-    .type_name: unsigned int*
-    .value_kind: global_buffer
-  .group_segment_fixed_size: 0
-  .kernarg_segment_align: 8
-  .kernarg_segment_size: 8
-  .language: OpenCL C
-  .language_version:
-  - 1
-  - 2
-  .max_flat_workgroup_size: 256
-  .name: test
-  .private_segment_fixed_size: 0
-  .sgpr_count: {s_cnt}
-  .sgpr_spill_count: 0
-  .symbol: test.kd
-  .uses_dynamic_stack: false
-  .vgpr_count: {v_cnt}
-  .vgpr_spill_count: 0
-  .wavefront_size: 32
-amdhsa.target: amdgcn-amd-amdhsa--gfx1100
-amdhsa.version:
- 1
- 2
-  """
-  boilerplate_start = f"""
-  .rodata
-  .global {function_name}.kd
-  .type {function_name}.kd,STT_OBJECT
-  .align 0x10
-  .amdhsa_kernel {function_name}"""
-  kernel_desc = {
-    '.amdhsa_group_segment_fixed_size': 0, '.amdhsa_private_segment_fixed_size': 0, '.amdhsa_kernarg_size': 0,
-    '.amdhsa_next_free_vgpr': v_cnt,   # this matters!
-    '.amdhsa_reserve_vcc': 0, '.amdhsa_reserve_xnack_mask': 0,
-    '.amdhsa_next_free_sgpr': s_cnt,
-    '.amdhsa_float_round_mode_32': 0, '.amdhsa_float_round_mode_16_64': 0, '.amdhsa_float_denorm_mode_32': 3, '.amdhsa_float_denorm_mode_16_64': 3,
-    '.amdhsa_dx10_clamp': 1, '.amdhsa_ieee_mode': 1, '.amdhsa_fp16_overflow': 0,
-    '.amdhsa_workgroup_processor_mode': 1, '.amdhsa_memory_ordered': 1, '.amdhsa_forward_progress': 0, '.amdhsa_enable_private_segment': 0,
-    '.amdhsa_system_sgpr_workgroup_id_x': 1, '.amdhsa_system_sgpr_workgroup_id_y': 1, '.amdhsa_system_sgpr_workgroup_id_z': 1,
-    '.amdhsa_system_sgpr_workgroup_info': 0, '.amdhsa_system_vgpr_workitem_id': 2, # is amdhsa_system_vgpr_workitem_id real?
-    '.amdhsa_exception_fp_ieee_invalid_op': 0, '.amdhsa_exception_fp_denorm_src': 0,
-    '.amdhsa_exception_fp_ieee_div_zero': 0, '.amdhsa_exception_fp_ieee_overflow': 0, '.amdhsa_exception_fp_ieee_underflow': 0,
-    '.amdhsa_exception_fp_ieee_inexact': 0, '.amdhsa_exception_int_div_zero': 0,
-    '.amdhsa_user_sgpr_dispatch_ptr': 0, '.amdhsa_user_sgpr_queue_ptr': 0, '.amdhsa_user_sgpr_kernarg_segment_ptr': 1,
-    '.amdhsa_user_sgpr_dispatch_id': 0, '.amdhsa_user_sgpr_private_segment_size': 0, '.amdhsa_wavefront_size32': 1, '.amdhsa_uses_dynamic_stack': 0}
-  code_start = f""".end_amdhsa_kernel
-  .text
-  .global {function_name}
-  .type {function_name},@function
-  .p2align 8
-  {function_name}:
-  """
-  ret = ".amdgpu_metadata\n" + metadata + ".end_amdgpu_metadata" + boilerplate_start + "\n" + '\n'.join("%s %d" % x for x in kernel_desc.items()) \
-      + "\n" + code_start + code + f"\n.size {function_name}, .-{function_name}"
-  return AMDProgram(cast(AMDDevice, Device["AMD"]), function_name, assemble(ret))
-
-def get_output(s:str, n_threads:int=1):
-  assert n_threads <= 32
-  code = "\n".join(["s_load_b64 s[0:1], s[0:1], null", "v_lshlrev_b32_e32 v0, 2, v0", s,
-    "s_waitcnt 0",
-    "global_store_b32 v0, v1, s[0:1]",
-    "s_nop 0", "s_sendmsg sendmsg(MSG_DEALLOC_VGPRS)", "s_endpgm"])
-  test = Tensor.zeros((n_threads,), dtype=dtypes.uint32).contiguous().realize().uop.buffer
-  prg = get_prg(code, 32, 32)
-  prg(test._buf, global_size=(1, 1, 1), local_size=(n_threads, 1, 1), wait=True)
-  return test.numpy()
+def get_output(asm:str, n_threads:int=1):
+  input_asm = "\n".join([ln if ln.strip().startswith('asm volatile') else f'asm volatile("{ln.strip().lstrip()}" : "+v"(a), "+v"(b));'
+                         for ln in asm.strip().splitlines() if ln.strip()])
+  src = f"""
+  typedef long unsigned int size_t;
+  extern "C" __attribute__((device, const)) size_t __ockl_get_local_id(unsigned int);
+  extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size(1, {n_threads}))) test(unsigned int* data0_1) {{
+    int l = __ockl_get_local_id(0);
+    unsigned a = 0, b = 0, c = 0;
+    {input_asm}
+    unsigned res;
+    asm volatile("v_mov_b32 %0, %1" : "=v"(res) : "v"(a));
+    *(data0_1+l) = res;
+  }}"""
+  t = Tensor.zeros(n_threads, dtype=dtypes.uint32).contiguous().realize()
+  prg = ProgramSpec("test", src, Device.DEFAULT, UOp.sink(t), global_size=[1, 1, 1], local_size=[n_threads, 1, 1])
+  car = CompiledRunner(prg)
+  if getenv("PRINT_ASM"): amdgpu_disassemble(car.lib)
+  car([t.uop.buffer], {}, wait=True)
+  return t.numpy()

 def f16_to_bits(x:float) -> int: return struct.unpack('<H', struct.pack('<e', x))[0]
 def f32_from_bits(x:int) -> float: return struct.unpack('<f', struct.pack('<I', x))[0]
@@ -105,54 +39,57 @@ class TestHW(unittest.TestCase):

  def test_simple(self):
    out = get_output("""
-    v_mov_b32_e32 v10 42
-    v_mov_b32_e32 v1 v10
-    """, n_threads=2)
+    v_mov_b32_e32 %1 42
+    v_mov_b32_e32 %2 %1
+    """)[0]
    np.testing.assert_equal(out, 42)

  def test_exec_mov(self):
    out = get_output("""
-    v_mov_b32_e32 v10 42
+    v_mov_b32_e32 %1 42
    s_mov_b32_e32 exec_lo 0b10
-    v_mov_b32_e32 v10 10
+    v_mov_b32_e32 %1 10
    s_mov_b32_e32 exec_lo 0b11
-    v_mov_b32_e32 v1 v10
+    v_mov_b32_e32 %2 %1
    """, n_threads=2)
    np.testing.assert_equal(out, [42, 10])

  def test_exec_cmp_vopc(self):
    out = get_output("""
    s_mov_b32 vcc_lo 0 // reset vcc
-    v_mov_b32_e32 v10 42
-    v_mov_b32_e32 v11 10
+    v_mov_b32_e32 %1 42
+    v_mov_b32_e32 %2 10
    s_mov_b32_e32 exec_lo 0b01
-    v_cmp_ne_u32 v10 v11
+    v_cmp_ne_u32 %1 %2
    s_mov_b32_e32 exec_lo 0b11
-    v_mov_b32_e32 v1 vcc_lo
+    v_mov_b32_e32 %2 vcc_lo
    """, n_threads=2)
    np.testing.assert_equal(out, 0b01)

  def test_exec_cmpx_vop3(self):
    out = get_output("""
-    v_mov_b32_e32 v10 42
-    v_mov_b32_e32 v11 10
+    s_mov_b32_e32 exec_lo 0b11
+    v_mov_b32_e32 %1 42
+    v_mov_b32_e32 %2 10
    s_mov_b32_e32 exec_lo 0b01
-    v_cmpx_ne_u32 v10 v11
+    v_cmpx_ne_u32 %1 %2
    s_mov_b32_e32 s10 exec_lo
    s_mov_b32_e32 exec_lo 0b11
-    v_mov_b32_e32 v1 s10
-    """, n_threads=2)
-    np.testing.assert_equal(out, 0b01)
+    v_mov_b32_e32 %2 s10
+    """, n_threads=2)[0]
+    np.testing.assert_equal(out & 0b11, 0b01)

  def test_fmac_vop3_modifier(self):
    init_state = f"""
-    v_mov_b32_e32 v10 {f16_to_bits(4.0)}
-    v_mov_b32_e32 v11 {f16_to_bits(3.0)}
-    v_mov_b32_e32 v1 {f16_to_bits(2.0)}
+    asm volatile("v_mov_b32_e32 %1, {f16_to_bits(4.0)}" : "+v"(a));
+    asm volatile("v_mov_b32_e32 %1, {f16_to_bits(3.0)}" : "+v"(b));
+    asm volatile("v_mov_b32_e32 %1, {f16_to_bits(2.0)}" : "+v"(c));
    """
-    self.assertEqual(get_output(init_state+"\n"+"v_fmac_f16_e64 v1 v11 v10"), f16_to_bits(14.))
-    self.assertEqual(get_output(init_state+"\n"+"v_fmac_f16_e64 v1 -v11 v10"), f16_to_bits(-10.))
-    self.assertEqual(get_output(init_state+"\n"+"v_fmac_f16_e64 v1 -v11 -v10"), f16_to_bits(14.))
+    mov = """asm volatile("v_mov_b32_e32 %1, %2" : "+v"(c), "+v"(a));"""
+    def fmac(a, b, c): return f"""asm volatile("v_fmac_f16_e64 {c}, {a}, {b}" : "+v"(c) : "v"(a), "v"(b));"""+"\n"+mov
+    self.assertEqual(get_output(init_state+"\n"+fmac("%1", "%2", "%3")), f16_to_bits(14.))
+    self.assertEqual(get_output(init_state+"\n"+fmac("%1", "-%2", "%3")), f16_to_bits(-10.))
+    self.assertEqual(get_output(init_state+"\n"+fmac("-%1", "-%2", "%3")), f16_to_bits(14.))

  def test_s_abs_i32(self):
    def s_abs_i32(x, y, dst="s10", scc=0):
@@ -160,7 +97,7 @@ class TestHW(unittest.TestCase):
        self.assertEqual(get_output(f"""
        s_mov_b32_e32 {dst} {x}
        s_abs_i32 {dst} {dst}
-        v_mov_b32_e32 v1 {reg}
+        v_mov_b32_e32 %2 {reg}
        """)[0], val)
    s_abs_i32(0x00000001, 0x00000001, scc=1)
    s_abs_i32(0x7fffffff, 0x7fffffff, scc=1)
@@ -173,8 +110,8 @@ class TestHW(unittest.TestCase):
  def test_v_rcp_f32_neg_vop3(self):
    def v_neg_rcp_f32(x:float, y:float):
      out = get_output(f"""
-      v_mov_b32_e32 v1 {f32_to_bits(x)}
-      v_rcp_f32_e64 v1, -v1
+      v_mov_b32_e32 %2 {f32_to_bits(x)}
+      v_rcp_f32_e64 %2, -%2
      """)[0]
      assert out == f32_to_bits(y), f"{f32_from_bits(out)} != {y} / {out} != {f32_to_bits(y)}"
    v_neg_rcp_f32(math.inf, -0.0)
@@ -186,10 +123,11 @@ class TestHW(unittest.TestCase):

  def test_v_cndmask_b32_neg(self):
    def v_neg(x:int|float, y:float):
+      # always pick -v1
      out = get_output(f"""
-      v_mov_b32_e32 v1 {f32_to_bits(x)}
-      s_mov_b32_e32 s10 1 // always pick -v1
-      v_cndmask_b32 v1, v1, -v1 s10
+      v_mov_b32_e32 %2 {f32_to_bits(x)}
+      s_mov_b32_e32 s10 1
+      v_cndmask_b32 %2, %2, -%2 s10
      """)[0]
      assert out == f32_to_bits(y), f"{f32_from_bits(out)} != {y} / {out} != {f32_to_bits(y)}"
    v_neg(-0.0, 0.0)
@@ -198,5 +136,12 @@ class TestHW(unittest.TestCase):
    v_neg(math.inf, -math.inf)
    v_neg(-math.inf, math.inf)

+  def test_v_subrev_wrap(self):
+    out = get_output("""
+    v_dual_mov_b32 %1, 0xffffffff :: v_dual_mov_b32 %2, 0x0
+    v_subrev_co_u32 %2, vcc_lo, %2, %1
+    """)[0]
+    self.assertEqual(out, 0xffff_ffff)
+
 if __name__ == "__main__":
  unittest.main()
@@ -27,14 +27,18 @@ class _ROCParseCtx:
        self.disasms[prog.base + addr] = info
        self.addr2prg[prog.base + addr] = prog

-  def next_sqtt(self): return next(self.sqtt_evs, None)
+  def next_sqtt(self):
+    x = next(self.sqtt_evs, None)
+    self.active_se = x.se if x is not None else None
+    return x
+
  def find_program(self, addr): return self.addr2prg[addr]

  def on_occupancy_ev(self, ev):
-    if DEBUG >= 4: print("OCC", ev.time, ev.cu, ev.simd, ev.wave_id, ev.start)
+    if DEBUG >= 4: print("OCC", ev.time, self.active_se, ev.cu, ev.simd, ev.wave_id, ev.start)

  def on_wave_ev(self, ev):
-    if DEBUG >= 4: print("WAVE", ev.wave_id, ev.cu, ev.simd, ev.contexts, ev.begin_time, ev.end_time)
+    if DEBUG >= 4: print("WAVE", ev.wave_id, self.active_se, ev.cu, ev.simd, ev.contexts, ev.begin_time, ev.end_time)

    asm = {}
    for j in range(ev.instructions_size):
@@ -46,9 +46,9 @@ __device__ static inline void arrive(int id) {
 #include "memory/memory.cuh"
 #include "shared/shared.cuh"
 #include "register/register.cuh"
+#include "mma/mma.cuh"

 #ifdef KITTENS_HOPPER
-#include "mma/mma.cuh"

 template<int n_reg> __device__ static inline void increase_registers() {
    static_assert(n_reg % 8 == 0, "n_reg must be a multiple of 8");
@@ -93,4 +93,4 @@ __device__ static inline void sync() {
 using warp = group<1>;      // scope used by most pre-Hopper GPUs, and also for most register operations.
 using warpgroup = group<4>; // special scope commonly used by Hopper and later.

-}
+}
@@ -65,8 +65,8 @@ template<typename _T, int _axis=-9999, bool _swizzle_flag=true> struct descripto
 namespace detail {
 template<typename... Args>
 struct descriptor_dict {
-    __host__ descriptor_dict() {}
-    template<typename T> __host__ descriptor_dict(T _, int b, int d, int r, int c) {}
+    __host__ __device__ descriptor_dict() {}
+    template<typename T> __host__ __device__ descriptor_dict(T _, int b, int d, int r, int c) {}
    __host__ __device__ descriptor_dict(const descriptor_dict &other) {}
 #ifdef KITTENS_HOPPER
    template<typename T, int U> __device__ const CUtensorMap* get() const {
@@ -85,8 +85,8 @@ struct descriptor_dict<_T, Args...> {
    using DESC = kittens::tma::descriptor<_T>; // copy or initialize with a default value
    CUtensorMap tma_desc;
    descriptor_dict<Args...> other_descs;
-    __host__ descriptor_dict() {}
-    __host__ descriptor_dict(typename DESC::T::dtype *data, int b, int d, int r, int c): other_descs(data, b, d, r, c) {
+    __host__ __device__ descriptor_dict() {}
+    __host__ __device__ descriptor_dict(typename DESC::T::dtype *data, int b, int d, int r, int c): other_descs(data, b, d, r, c) {
        kittens::detail::tma::create_tensor_map<typename DESC::T, DESC::axis, DESC::swizzle_flag>(&tma_desc, data, b, d, r, c);
    }
    __host__ __device__ inline descriptor_dict(const descriptor_dict &other) :
@@ -135,7 +135,7 @@ struct gl {

    detail::descriptor_dict<TMA_Types...> tma_descs;

-    __host__ inline gl(T *_data,
+    __host__ __device__ inline gl(T *_data,
                        ducks::gl::make_arg_t<b> _batch,
                        ducks::gl::make_arg_t<d> _depth,
                        ducks::gl::make_arg_t<r> _rows,
@@ -425,4 +425,4 @@ __host__ static inline CUtensorMap* allocate_and_create_tensor_map(const typenam

 } // namespace tma
 } // namespace detail
-} // namespace kittens
+} // namespace kittens
@@ -0,0 +1,45 @@
+// https://github.com/HazyResearch/ThunderKittens/blob/main/kernels/matmul/educational/level_04.cu
+#include "kittens.cuh"
+using namespace kittens;
+
+constexpr int g_N = 8192;
+constexpr int BLOCK_SIZE = 32;
+#define NUM_WORKERS (1)
+#define NUM_THREADS (NUM_WORKERS*kittens::WARP_THREADS)
+
+using sub_tile = st_bf<BLOCK_SIZE,BLOCK_SIZE>;
+using tile_gl =  gl<bf16,  1, 1, g_N, g_N, sub_tile>;
+
+__global__ void kernel(bf16 *c_ptr, bf16 *a_ptr, bf16 *b_ptr) {
+  tile_gl g_C{c_ptr, nullptr, nullptr, nullptr, nullptr};
+  tile_gl g_A{a_ptr, nullptr, nullptr, nullptr, nullptr};
+  tile_gl g_B{b_ptr, nullptr, nullptr, nullptr, nullptr};
+
+  extern __shared__ alignment_dummy __shm[];
+  shared_allocator al((int*)&__shm[0]);
+  st_bf<BLOCK_SIZE,BLOCK_SIZE> &As = al.allocate<st_bf<BLOCK_SIZE,BLOCK_SIZE>>();
+  st_bf<BLOCK_SIZE,BLOCK_SIZE> &Bs = al.allocate<st_bf<BLOCK_SIZE,BLOCK_SIZE>>();
+
+  rt_bf<BLOCK_SIZE,BLOCK_SIZE> A_reg;
+  rt_bf<BLOCK_SIZE,BLOCK_SIZE> B_reg;
+  rt_bf<BLOCK_SIZE,BLOCK_SIZE, ducks::rt_layout::col> B_reg_col;
+  rt_fl<BLOCK_SIZE,BLOCK_SIZE> C_accum;
+
+  int col = blockIdx.x;
+  int row = blockIdx.y;
+
+  warp::zero(C_accum);
+  int num_tiles = (g_N + BLOCK_SIZE - 1) / BLOCK_SIZE;
+  for (int tile = 0; tile < num_tiles; ++tile) {
+    warp::load(As, g_A, {0, 0, row, tile});
+    warp::load(Bs, g_B, {0, 0, tile, col});
+    __syncthreads();
+    warp::load(A_reg, As);
+    warp::load(B_reg, Bs);
+    warp::swap_layout(B_reg_col, B_reg);
+    __syncthreads();
+    warp::mma_AB(C_accum, A_reg, B_reg_col, C_accum);
+    __syncthreads();
+  }
+  warp::store(g_C, C_accum, {0, 0, row, col});
+}
@@ -0,0 +1,37 @@
+import pathlib
+from tinygrad import Device, Tensor
+from tinygrad.helpers import Context
+from tinygrad.runtime.support.compiler_cuda import pretty_ptx, NVCCCompiler
+
+if __name__ == "__main__":
+  code = (pathlib.Path(__file__).parent / "matmul.cu").read_text()
+  device = Device["CUDA"]
+  kitten_args = [f"-I{(pathlib.Path(__file__).parent / 'include').as_posix()}", "-std=c++20", "--expt-relaxed-constexpr"]
+  lib = NVCCCompiler(device.compiler.arch, kitten_args).compile(code)
+  kernel_name = lib.decode().split(".globl\t")[1].split("\n")[0]
+  print("kernel name", kernel_name)
+  print(pretty_ptx(lib.decode()))
+
+  prg = device.runtime(kernel_name, lib)
+  prg.smem = 10000
+
+  N = 8192
+  a = Tensor.randn(N, N, device='CUDA', dtype="bfloat16")
+  b = Tensor.randn(N, N, device='CUDA', dtype="bfloat16")
+  c = Tensor.empty(N, N, device='CUDA', dtype="bfloat16")
+  Tensor.realize(a, b, c)
+
+  BLOCK_SIZE = 32
+
+  gsz = (N // BLOCK_SIZE, N // BLOCK_SIZE, 1)
+  for _ in range(5):
+    et = prg(c.uop.buffer.ensure_allocated()._buf, a.uop.buffer._buf, b.uop.buffer._buf,
+             global_size=gsz, local_size=(32,1,1), wait=True)
+    print(f"{N*N*N*2/(et*1e9):2f} GFLOPS")
+
+  for _ in range(5):
+    with Context(DEBUG=2):
+      ref = (a@b).realize()
+
+  ref, c = ref.float(), c.float()
+  print((ref-c).mean().item(), (ref-c).max().item())
@@ -1,75 +0,0 @@
-import torch
-
-#credit to KellerJordan at https://github.com/KellerJordan/Muon/tree/master
-#some changes: classic momentum instead of weighting gradient
-#added ns_steps, ns_params, nesterov as hyperparams
-def zeropower_via_newtonschulz5(G:torch.tensor, steps:int, params:tuple[int, ...]):
-  """
-  Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a
-  quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose
-  of minimizing steps, it turns out to be empirically effective to keep increasing the slope at
-  zero even beyond the point where the iteration no longer converges all the way to one everywhere
-  on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T
-  where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model
-  performance at all relative to UV^T, where USV^T = G is the SVD.
-  """
-  assert G.ndim >= 2 # batched Muon implementation by @scottjmaddox, and put into practice in the record by @YouJiacheng
-
-  a, b, c = params
-  X = G
-  if G.size(-2) > G.size(-1):
-    X = X.mT
-
-  # Ensure spectral norm is at most 1
-  X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7)
-  # Perform the NS iterations
-  for _ in range(steps):
-    A = X @ X.mT
-    B = b * A + c * A @ A # quintic computation strategy adapted from suggestion by @jxbz, @leloykun, and @YouJiacheng
-    X = a * X + B @ X
-
-  if G.size(-2) > G.size(-1):
-    X = X.mT
-
-  return X
-
-def muon_update(grad, momentum, beta=0.95, ns_steps=5, ns_params=(3.4445, -4.7750,  2.0315), nesterov=True):
-  if beta:
-    momentum.mul_(beta).add_(grad)
-    update = grad.add(momentum,alpha=beta) if nesterov else momentum
-  else: update = grad
-  if update.ndim == 4: # for the case of conv filters
-    update = update.view(len(update), -1)
-  update = zeropower_via_newtonschulz5(update, steps=ns_steps, params=ns_params)
-  return update
-
-class SingleDeviceMuon(torch.optim.Optimizer):
-  """
-  Muon variant for usage in non-distributed settings.
-  """
-  def __init__(self, params, lr=0.02, weight_decay=0.0, momentum=0.95, ns_steps=5, ns_params=(3.4445, -4.7750,  2.0315), nesterov=True):
-    defaults = dict(lr=lr, weight_decay=weight_decay, momentum=momentum, ns_steps=ns_steps, ns_params=ns_params, nesterov=nesterov)
-    super().__init__(params, defaults)
-
-  @torch.no_grad()
-  def step(self, closure=None):
-
-    loss = None
-    if closure is not None:
-      with torch.enable_grad():
-        loss = closure()
-
-    for group in self.param_groups:
-      for p in group["params"]:
-        if p.grad is None:
-          p.grad = torch.zeros_like(p)  # Force synchronization
-        state = self.state[p]
-        if len(state) == 0:
-          state["momentum_buffer"] = torch.zeros_like(p)
-        update = muon_update(p.grad, state["momentum_buffer"], beta=group["momentum"], ns_steps=group["ns_steps"],
-                             ns_params=group["ns_params"], nesterov=group["nesterov"])
-        p.mul_(1.0 - group["lr"] * group["weight_decay"])
-
-        p.add_(update.reshape(p.shape), alpha=-group["lr"])
-
-    return loss
@@ -25,8 +25,6 @@ nav:
    - Layout: developer/layout.md
    - Speed: developer/speed.md
    - UOp: developer/uop.md
-    - Grouper:
-      - developer/kernelize.md
    - Runtime:
      - developer/runtime.md
      - HCQ: developer/hcq.md
@@ -9,10 +9,11 @@ with open(directory / 'README.md', encoding='utf-8') as f:

 testing_minimal = [
  "numpy",
-  "torch==2.8.0",
+  "torch==2.9.0",
  "pytest",
  "pytest-xdist",
  "pytest-timeout",
+  "pytest-split",
  "hypothesis",
  "z3-solver",
 ]
@@ -54,6 +54,8 @@ def gen_diff(table_old, table_new):

 def display_diff(diff): return "+"+str(diff) if diff > 0 else str(diff)

+NONCORE_DIRS = {"tinygrad/apps", "tinygrad/nn", "tinygrad/renderer", "tinygrad/runtime", "tinygrad/viz"}
+
 if __name__ == "__main__":
  if len(sys.argv) == 3:
    headers = ["Name", "Lines", "Diff", "Tokens/Line", "Diff"]
@@ -76,9 +78,12 @@ if __name__ == "__main__":
    else:
      print(tabulate([headers] + sorted(table, key=lambda x: -x[1]), headers="firstrow", floatfmt=".1f")+"\n")
      groups = sorted([('/'.join(x[0].rsplit("/", 1)[0].split("/")[0:2]), x[1], x[2]) for x in table])
+      dir_sizes = {}
      for dir_name, group in itertools.groupby(groups, key=lambda x:x[0]):
-        print(f"{dir_name:30s} : {sum([x[1] for x in group]):6d}")
+        dir_sizes[dir_name] = sum([x[1] for x in group])
+        print(f"{dir_name:30s} : {dir_sizes[dir_name]:6d}")
+      print(f"\n core line count: {sum([v for k,v in dir_sizes.items() if k not in NONCORE_DIRS])}")
      total_lines = sum([x[1] for x in table])
-      print(f"\ntotal line count: {total_lines}")
+      print(f"total line count: {total_lines}")
      max_line_count = int(os.getenv("MAX_LINE_COUNT", "-1"))
      assert max_line_count == -1 or total_lines <= max_line_count, f"OVER {max_line_count} LINES"
@@ -1,63 +0,0 @@
-import time, sys, hashlib
-from pathlib import Path
-from tinygrad.nn.onnx import OnnxRunner
-from tinygrad import Tensor, dtypes, TinyJit
-from tinygrad.helpers import IMAGE, GlobalCounters, fetch, colored, getenv, trange
-import numpy as np
-from extra.bench_log import BenchEvent, WallTimeEvent
-
-OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx"
-
-if __name__ == "__main__":
-  run_onnx = OnnxRunner(fetch(OPENPILOT_MODEL))
-
-  Tensor.manual_seed(100)
-  input_shapes = {name: spec.shape for name, spec in run_onnx.graph_inputs.items()}
-  input_types = {name: spec.dtype for name, spec in run_onnx.graph_inputs.items()}
-  new_inputs = {k:Tensor.randn(*shp, dtype=input_types[k]).mul(8).realize() for k,shp in input_shapes.items()}
-  new_inputs_junk = {k:Tensor.randn(*shp, dtype=input_types[k]).mul(8).realize() for k,shp in input_shapes.items()}
-  new_inputs_junk_numpy = {k:v.numpy() for k,v in new_inputs_junk.items()}
-
-  # benchmark
-  for _ in range(5):
-    GlobalCounters.reset()
-    st = time.perf_counter_ns()
-    ret = next(iter(run_onnx(new_inputs_junk).values())).cast(dtypes.float32).numpy()
-    print(f"unjitted: {(time.perf_counter_ns() - st)*1e-6:7.4f} ms")
-
-  # NOTE: the inputs to a JIT must be first level arguments
-  run_onnx_jit = TinyJit(lambda **kwargs: run_onnx(kwargs), prune=True)
-  step_times = []
-  for _ in range(20):
-    GlobalCounters.reset()
-    st = time.perf_counter_ns()
-    with WallTimeEvent(BenchEvent.STEP):
-      # Need to cast non-image inputs from numpy, this is only realistic way to run model
-      inputs = {**{k:v for k,v in new_inputs_junk.items() if 'img' in k},
-                **{k:Tensor(v) for k,v in new_inputs_junk_numpy.items() if 'img' not in k}}
-      ret = next(iter(run_onnx_jit(**inputs).values())).cast(dtypes.float32).numpy()
-      step_times.append(t:=(time.perf_counter_ns() - st)*1e-6)
-    print(f"jitted:  {t:7.4f} ms")
-
-  suffix = ""
-  if IMAGE.value < 2: suffix += f"_image{IMAGE.value}" # image=2 has no suffix for compatibility
-  if getenv("FLOAT16") == 1: suffix += "_float16"
-  path = Path(__file__).parent / "openpilot" / f"{hashlib.md5(OPENPILOT_MODEL.encode()).hexdigest()}{suffix}.npy"
-
-  # validate if we have records
-  tinygrad_out = next(iter(run_onnx_jit(**new_inputs).values())).cast(dtypes.float32).numpy()
-  if getenv("SAVE_OUTPUT"):
-    np.save(path, tinygrad_out)
-    print(f"saved output to {path}!")
-  elif getenv("FUZZ") and path.exists():
-    known_good_out = np.load(path)
-    for _ in trange(1000):
-      ret = next(iter(run_onnx_jit(**new_inputs).values())).cast(dtypes.float32).numpy()
-      np.testing.assert_allclose(known_good_out, ret, atol=1e-2, rtol=1e-2)
-    print(colored("fuzz validated!", "green"))
-  elif path.exists():
-    known_good_out = np.load(path)
-    np.testing.assert_allclose(known_good_out, tinygrad_out, atol=1e-2, rtol=1e-2)
-    print(colored("outputs validated!", "green"))
-  else:
-    print(colored("skipping validation", "yellow"))
@@ -2,9 +2,9 @@ from extra.models.resnet import ResNet50
 from tinygrad import Tensor, nn, Device
 from tinygrad.helpers import Profiling, Timing, getenv
 from tinygrad.uop.ops import Ops
-from tinygrad.codegen import get_rewrites_for_renderer, apply_rewrites
-from tinygrad.codegen.late.control_flow import linearize
-from tinygrad.uop.spec import type_verify
+from tinygrad.codegen import full_rewrite_to_sink
+from tinygrad.codegen.late.linearizer import linearize
+from tinygrad.uop.spec import type_verify, program_spec

 if __name__ == "__main__":
  mdl = ResNet50()
@@ -29,12 +29,11 @@ if __name__ == "__main__":
        asts = list({x.ast.key:x.ast for x in sched if x.ast.op is Ops.SINK}.values())
        if (restrict_kernel := getenv("RESTRICT_KERNEL", -1)) != -1: asts = asts[restrict_kernel:restrict_kernel+1]

-        rewrites = get_rewrites_for_renderer(Device.default.renderer, linearizer=False)
        with Profiling(PROFILE, fn="/tmp/rewrite.prof"):
          with Timing("***** model rewrite in   "):
            rewritten_uops = []
            for u in asts:
-              rewritten_uops.append(apply_rewrites(u, rewrites))
+              rewritten_uops.append(full_rewrite_to_sink(u, ren=Device.default.renderer))

        if LINEARIZE:
          with Timing("***** model linearize in "):
@@ -42,5 +41,5 @@ if __name__ == "__main__":
            for u in rewritten_uops:
              uops_line.append(linearize(u))
          with Timing("***** model verify in    "):
-            for u in uops_line: type_verify(u)
+            for u in uops_line: type_verify(u, program_spec)
          print(sum(len(u) for u in uops_line))
@@ -0,0 +1,38 @@
+from tinygrad import Tensor, nn, Context, GlobalCounters
+
+if __name__ == "__main__":
+  conv = nn.Conv2d(64, 128, 3)
+  img = Tensor.randn((1,64,128,128))
+  with Context(DEBUG=0, BEAM=0):
+    Tensor.realize(img, conv.weight, conv.bias)
+
+  tst = conv(img).permute(0,2,3,1).realize()
+  print(tst.shape)
+
+  print("NEW")
+  img_perm = img.permute(0,2,3,1).contiguous()
+  print(img_perm.shape)
+  pp = img_perm.permute(0,3,1,2)._pool((3,3)).permute(0,2,3,4,5,1)
+
+  def hwio(pp, conv):
+    pp = pp.unsqueeze(-1)
+    weight = conv.weight.permute(2,3,1,0).contiguous()
+    print(pp.shape, weight.shape, (pp*weight).shape)
+    return (pp * weight).sum([-4,-3, -2])
+
+  def ohwi(pp, conv):
+    pp = pp.unsqueeze(-4)
+    weight = conv.weight.permute(0,2,3,1).contiguous()
+    print(pp.shape, weight.shape, (pp*weight).shape)
+    return (pp * weight).sum([-3,-2,-1])
+
+  for f in [hwio, ohwi]:
+    GlobalCounters.reset()
+    print("\n**************", f.__name__, "**************")
+    out = f(pp, conv)
+    out.realize()
+    print(out.shape)
+
+    with Context(DEBUG=0, BEAM=0):
+      err = (tst-out).square()
+      print(err.mean().item(), err.max().item())
@@ -272,6 +272,10 @@ class TestMainOnnxOps(TestOnnxOps):
  def test_qlinearmatmul_2D_int8_float32(self): self._run_qlinearmatmul_test(np.int8, np.float32, 2)
  def test_qlinearmatmul_3D_int8_float32(self): self._run_qlinearmatmul_test(np.int8, np.float32, 3)

+  def test_reduce_l2_half(self):
+    inputs = {"data": np.random.randn(1, 1, 32, 32, 32).astype(np.half)*100}
+    self.helper_test_single_op("ReduceL2", inputs, {}, ["reduced"])
+
 class TestTrainingOnnxOps(TestOnnxOps):
  # NOTE: ORT doesn't actually support training ops on cpu so we test using functions provided by onnx
  DOMAIN = AI_ONNX_PREVIEW_TRAINING_DOMAIN
@@ -487,4 +491,4 @@ class TestContribOnnxOps(TestOnnxOps):
        self.helper_test_single_op("QLinearGlobalAveragePool", inputs, attributes, outputs)

 if __name__ == "__main__":
-  unittest.main()
+  unittest.main()
@@ -1,7 +1,7 @@
 import random
 import z3
 from tinygrad import dtypes
-from tinygrad.uop.spec import uops_to_z3, z3_cdiv
+from tinygrad.uop.validate import uops_to_z3, z3_cdiv
 from tinygrad.uop.ops import UOp
 from tinygrad.uop.decompositions import fast_idiv
 random.seed(42)
@@ -207,7 +207,7 @@ def fuzz_linearizer(lin: Kernel, rtol=1e-2, atol=1e-2, opts_list=None):
        if not FUZZ_ALL_ACTIONS and test_lin.applied_opts: print(f"applied opts: {test_lin.applied_opts}")

        # stop if kernel uops repeat
-        try: tuops = tuplize_uops(get_program(test_lin.get_optimized_ast(), test_lin.opts).uops)
+        try: tuops = tuplize_uops(get_program(test_lin.get_optimized_ast(), test_lin.ren).uops)
        except KeyboardInterrupt: raise
        except BaseException as e:
          print(test_lin.ast)
@@ -224,7 +224,7 @@ def fuzz_linearizer(lin: Kernel, rtol=1e-2, atol=1e-2, opts_list=None):
        (msg, rawbufs, var_vals, ground_truth, state1) = compare_linearizer(test_lin, rawbufs, var_vals, ground_truth, rtol=rtol, atol=atol)
        if state1 is not None and validate_device is not None:
          validate_lin = test_lin.copy()
-          validate_lin.opts = validate_device.renderer
+          validate_lin.ren = validate_device.renderer
          if validate_rawbufs is None:
            validate_rawbufs = [get_fuzz_rawbuf_like(x, copy=True, force_device=validate_device.device) for x in rawbufs]
          (_msg, _, _, _, state2) = compare_linearizer(validate_lin, validate_rawbufs, var_vals, ground_truth, rtol=rtol, atol=atol)
@@ -2,7 +2,7 @@ import random, operator
 import z3
 from tinygrad import Variable, dtypes
 from tinygrad.uop.ops import UOp
-from tinygrad.uop.spec import uops_to_z3
+from tinygrad.uop.validate import uops_to_z3
 from tinygrad.helpers import DEBUG, Context

 seed = random.randint(0, 100)
@@ -13,7 +13,7 @@ try:
  from tinygrad.engine.realize import get_program
  from tinygrad.uop.ops import UOp, Ops, KernelInfo
  from tinygrad.codegen.opt import Opt
-  from tinygrad.helpers import VERSION, Context, ContextVar, colored, db_connection, getenv, tqdm
+  from tinygrad.helpers import VERSION, Context, ContextVar, colored, db_connection, getenv, tqdm, BEAM
  from tinygrad.device import Device
 except ImportError as e:
  print(repr(e))
@@ -51,8 +51,8 @@ def replay_get_rangeify_map(ret:dict[UOp, UOp], big_sink:UOp) -> tuple[str, str,
  return to_str(new_sink), to_str(big_sink.substitute(ret)), (big_sink,)

 def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer|None=None, opts:list[Opt]|None=None) -> tuple[str, str, tuple[Any, ...]]:
-  # NOTE: this always uses the opts_to_apply path
-  sink_arg = ast.arg or KernelInfo(opts_to_apply=p.applied_opts)
+  # the ast.arg is non None if we are inside of search.py
+  sink_arg = ast.arg or KernelInfo(opts_to_apply=tuple(opts) if opts is not None else p.applied_opts if BEAM>=1 else None)
  input_ast = ast.replace(arg=replace(sink_arg, name=p.name))
  # if no renderer was provided, open the device to get it
  if renderer is None: renderer = Device[p.device].renderer
@@ -112,7 +112,7 @@ class TestRealWorld(unittest.TestCase):
        loss.backward()
        optimizer.step()

-      helper_test("train_mnist", lambda: (Tensor.randn(BS, 1, 28, 28),), train, 0.07, 102)
+      helper_test("train_mnist", lambda: (Tensor.randn(BS, 1, 28, 28),), train, 0.07, 103)

  @unittest.skipIf(CI and Device.DEFAULT in {"CPU", "CL"}, "slow")
  def test_forward_cifar(self):
@@ -176,7 +176,7 @@ class TestRealWorld(unittest.TestCase):
      for v in data.values(): v.to_(Device.DEFAULT)

      helper_test("train_bert", lambda: (data["input_ids"], data["segment_ids"], data["input_mask"], data["masked_lm_positions"], \
-          data["masked_lm_ids"], data["masked_lm_weights"], data["next_sentence_labels"]), train, 0.31, 358)
+          data["masked_lm_ids"], data["masked_lm_weights"], data["next_sentence_labels"]), train, 0.31, 427)

 if __name__ == '__main__':
  unittest.main()
@@ -14,6 +14,8 @@ from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
 # TODO: write a clean version of this
 from test.test_linearizer import helper_realized_ast, helper_linearizer_opt

+# NOTE: get_program always passes in Device[Device.DEFAULT].renderer explicitly for process_replay!!!
+
 def helper_tc_ensure_uops_and_opts_count(N: int, M:int, K:int, dtype_in:DType, dtype_out:DType, axis:int=0, tc_select:int=-1, tc_opt:int=0,
                                         ensure_triggered:bool=True):
  a, b = Tensor.rand(M, K, dtype=dtype_in), Tensor.rand(K, N, dtype=dtype_in)
@@ -41,7 +43,7 @@ def helper_tc_allclose(N:int, M:int, K:int, dtype_in:DType, dtype_out:DType, axi
  if dtype_in == dtypes.bfloat16: r = r.float()
  realized_ast, bufs = helper_realized_ast(r)
  opts = [Opt(op=OptOps.TC, axis=axis, arg=(tc_select, tc_opt, use_tensor_cores))]
-  prg = CompiledRunner(replace(get_program(realized_ast, opts=opts), device=Device.DEFAULT))
+  prg = CompiledRunner(replace(get_program(realized_ast, Device[Device.DEFAULT].renderer, opts=opts), device=Device.DEFAULT))
  if use_tensor_cores == 1: assert len([uop for uop in prg.p.uops if uop.op is Ops.WMMA]) > 0, "wmma not triggered"
  assert len([x for x in prg.p.uops[-1].arg.applied_opts if x.op is OptOps.TC]) == 1, "tensor core opt not included"
  prg.exec(bufs)
@@ -68,7 +70,7 @@ class TestTensorCores(unittest.TestCase):
      n, m, k = tc.dims[0], tc.dims[1], 2 if AMX else tc.dims[2]
      a, b = Tensor.rand(m, k, dtype=tc.dtype_in), Tensor.rand(k, n, dtype=tc.dtype_in)
      r = a.matmul(b, dtype=tc.dtype_out)
-      prg = get_program(r.schedule()[-1].ast, opts=[Opt(op=OptOps.TC, axis=0, arg=(-1, 2, 1))])
+      prg = get_program(r.schedule()[-1].ast, Device[Device.DEFAULT].renderer, opts=[Opt(op=OptOps.TC, axis=0, arg=(-1, 2, 1))])
      if Device.DEFAULT == "CPU" and CPU_LLVM:
        assert "0x201000" in prg.src
      elif Device.DEFAULT == "AMD" and AMD_LLVM:
@@ -154,7 +156,7 @@ class TestTensorCores(unittest.TestCase):
    r = x.matmul(y, dtype=tc.dtype_out)
    opts = [Opt(OptOps.UNROLL, 0, 4)]
    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
-    for u in get_program(ast, opts=opts).uops:
+    for u in get_program(ast, Device[Device.DEFAULT].renderer, opts=opts).uops:
      if u.op is Ops.WMMA:
        assert u.src[-1].src[0].op != Ops.STORE

@@ -167,7 +169,7 @@ class TestTensorCores(unittest.TestCase):
    r = x.matmul(y, dtype=tc.dtype_out)
    opts = [Opt(OptOps.UNROLL, 0, 4)]
    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
-    for u in get_program(ast, opts=opts).uops:
+    for u in get_program(ast, Device[Device.DEFAULT].renderer, opts=opts).uops:
      if u.op is Ops.WMMA:
        #assert u.src[-1].dtype == dtypes.float.vec(prod(tc.thread_local_sizes[2]))
        assert u.src[-1].src[0].op != Ops.STORE
@@ -182,7 +184,7 @@ class TestTensorCores(unittest.TestCase):
    r = x.matmul(y, dtype=tc.dtype_out).relu()
    opts = [Opt(OptOps.UNROLL, 0, 4)]
    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
-    for u in get_program(ast, opts=opts).uops:
+    for u in get_program(ast, Device[Device.DEFAULT].renderer, opts=opts).uops:
      if u.op is Ops.WMMA:
        #assert u.src[-1].dtype == dtypes.float.vec(prod(tc.thread_local_sizes[2]))
        assert u.src[-1].src[0].op != Ops.STORE
@@ -67,12 +67,9 @@ class TestBinaryOpsConstFolding(unittest.TestCase):
  def test_tensor_one_mul(self):
    _check_ast_count(0, Tensor.ones(4) * Tensor([1.0, 2, 3, 4]))

-  # TODO: these will be fixed with better folding
-  @unittest.expectedFailure
  def test_bool_tensor_mul_bool(self):
    _check_ast_count(0, Tensor([True, False]) * True)
    _check_ast_count(0, Tensor([True, False]) * False)
-  @unittest.expectedFailure
  def test_bool_mul_bool_tensor(self):
    _check_ast_count(0, True * Tensor([True, False]))
    _check_ast_count(0, False * Tensor([True, False]))
@@ -185,7 +182,7 @@ class TestReduceOpsConstFolding(unittest.TestCase):
    np.testing.assert_equal(Tensor(4).sum().numpy(), 4)

  def test_padded_const_sum(self):
-    _check_ast_count(1, Tensor.ones(4).pad(((1, 1),)).sum())
+    _check_ast_count(0, Tensor.ones(4).pad(((1, 1),)).sum())
    np.testing.assert_equal(Tensor.ones(4).pad(((1, 1),)).sum().numpy(), 4)

    # NOTE: cannot just count the non-padded area because some Ops f do not have f(0) = 0.
@@ -75,7 +75,10 @@ def universal_test_unary(a, dtype, op):
  out: Tensor = op[0](ta)
  tensor_value = out.numpy()
  numpy_value = op[1](ta.numpy())
-  if dtype in dtypes.fp8s: numpy_value = truncate[dtype](numpy_value)
+  if dtype in dtypes.fp8s:
+    # cuda cast f32 inf to f8 MAX, amd cast it to nan(E4M3)/inf(E5M2)
+    if math.isinf(numpy_value): return
+    numpy_value = truncate[dtype](numpy_value)
  if dtype in dtypes.floats:
    atol, rtol = { dtypes.float16:(1e-3, 1e-2), dtypes.bfloat16:(1e-3, 2e-2),
      dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2: (1.0, 5e-1)}.get(dtype, (1e-6, 1e-5))
@@ -51,7 +51,7 @@ class TestFusionOp(unittest.TestCase):
    a = Tensor(val)
    for _ in range(24): a = Tensor.stack(a, a)[0]
    sched = a.schedule()
-    self.assertEqual(len(sched), 1)
+    self.assertEqual(len(sched), 0)
    self.assertLess(time.perf_counter()-st, 2.0)

  def test_recursive_reshape(self):
@@ -52,7 +52,6 @@ class TestImageDType(unittest.TestCase):
    assert isinstance(it.uop.base.realized.dtype, ImageDType)
    np.testing.assert_equal(tst, it.numpy())

-  @unittest.expectedFailure # this isn't supported anymore, CAST to ImageDType stays ImageDType
  def test_image_cast_and_back_collapses(self):
    data = Tensor.randn(9*27*4).realize()
    tst = data.numpy()
@@ -155,6 +155,7 @@ class TestLinearizer(unittest.TestCase):
    assert stores[1].src[1].dtype == dtypes.float
    assert any(x.op is Ops.DEFINE_GLOBAL for x in stores[1].toposort())

+  @unittest.skipIf(Device.DEFAULT=="CPU", "CPU splits the cat so cant upcast")
  def test_zero_fold(self):
    a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize()
    r = Tensor.stack(a, b)
@@ -393,14 +394,15 @@ class TestLinearizer(unittest.TestCase):
    uops = get_program(ast, opts=opt).uops
    local_stores = [u for u in uops if u.op is Ops.STORE and any(x.op is Ops.DEFINE_LOCAL for x in get_recursive(u.src[0]))]
    global_stores = [u for u in uops if u.op is Ops.STORE and any(x.op is Ops.DEFINE_GLOBAL for x in get_recursive(u.src[0]))]
-    barrier = [u for u in uops if u.op is Ops.BARRIER][0]
+    barrier = [u for u in uops if u.op is Ops.BARRIER]
+    assert len(barrier) == 1
    # check that the float4 cast collapses for all stores
    for store in local_stores+global_stores:
      assert store.src[1].dtype.count > 1 # and store.src[2].op is not Ops.VECTORIZE
    # # check the children's vins
    # TODO: src ALU are not the same, should it?
    # assert barrier.src == tuple(local_stores)
-    assert len([u for u in uops if u.op is Ops.IF and u.src[1] == barrier]) == 1
+    #assert len([u for u in uops if u.op is Ops.IF])

  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
@@ -23,7 +23,7 @@ class TestLinearizerFailure(unittest.TestCase):
    c9 = UOp(Ops.DEFINE_GLOBAL, dtypes.uchar.ptr(47040000), arg=2, src=())
    c10 = c9.index((((c3*UOp.const(dtypes.index, 4704000))+c2)+(c6*UOp.const(dtypes.index, 784))).valid(UOp.const(dtypes.bool, True))).load()
    c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(dtypes.index, 6000))+c6)+((c7*UOp.const(dtypes.index, 16))+c8)).alu(Ops.CMPLT, UOp.const(dtypes.index, 59999)).where(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(dtypes.int, -1))).where(UOp.const(dtypes.uchar, 0), c10).reduce(c6, arg=Ops.ADD)
-    c12 = c0.index((((c1*UOp.const(dtypes.index, 7840))+(c2*UOp.const(dtypes.index, 10)))+c3).valid(UOp.const(dtypes.bool, True))).store(c11, c1, c2, c3)
+    c12 = c0.index((((c1*UOp.const(dtypes.index, 7840))+(c2*UOp.const(dtypes.index, 10)))+c3).valid(UOp.const(dtypes.bool, True))).store(c11).end(c1, c2, c3)
    ast = c12.sink(arg=KernelInfo(name='test', axis_types=(), dont_use_locals=False, applied_opts=(Opt(op=OptOps.GROUP, axis=1, arg=16),), opts_to_apply=None))
    _ = get_program(ast, Device["METAL"].renderer)

@@ -16,7 +16,7 @@ class TestLinearizerFailures(unittest.TestCase):
    c7 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(64), arg=2, src=())
    c8 = c7.index(c3).load()
    c9 = ((((c6+(c8*UOp.const(dtypes.float, -1.0)))*(c6+(c8*UOp.const(dtypes.float, -1.0)))).reduce(c5, arg=Ops.ADD)*UOp.const(dtypes.float, 0.000390625))+UOp.const(dtypes.float, 1e-05)).sqrt().reciprocal()
-    c10 = c0.index(c3).store(c9, c1, c2)
+    c10 = c0.index(c3).store(c9).end(c1, c2)
    ast = c10.sink()
    get_program(ast)

@@ -2602,18 +2602,13 @@ class TestOps(unittest.TestCase):
      lambda x: torch.nn.functional.avg_pool2d(x, kernel_size=(111,28)),
      lambda x: Tensor.avg_pool2d(x, kernel_size=(111,28)), rtol=1e-5)

-  @unittest.skipIf(Device.DEFAULT == "AMD" and CI, "remu failure?")
-  def test_avg_pool3d_failure(self):
-    with Context(NOOPT=0):
-      helper_test_op([(1,1,16,16,16)],
-        lambda x: torch.nn.functional.avg_pool3d(x, kernel_size=(8,8,8), stride=5, padding=1, count_include_pad=False),
-        lambda x: Tensor.avg_pool2d(x, kernel_size=(8,8,8), stride=5, padding=1, count_include_pad=False), rtol=1e-5, forward_only=True)
-
-  def test_avg_pool3d_noopt(self):
-    with Context(NOOPT=1):
-      helper_test_op([(1,1,16,16,16)],
-        lambda x: torch.nn.functional.avg_pool3d(x, kernel_size=(8,8,8), stride=5, padding=1, count_include_pad=False),
-        lambda x: Tensor.avg_pool2d(x, kernel_size=(8,8,8), stride=5, padding=1, count_include_pad=False), rtol=1e-5, forward_only=True)
+  def test_avg_pool3d(self):
+    # TODO: AMD_LLVM has larger atol
+    # TODO: PYTHON=1 backward hangs?
+    atol = 1e-2 if AMD_LLVM else 1e-6
+    helper_test_op([(1,1,16,16,16)],
+      lambda x: torch.nn.functional.avg_pool3d(x, kernel_size=(8,8,8), stride=5, padding=1, count_include_pad=False),
+      lambda x: Tensor.avg_pool2d(x, kernel_size=(8,8,8), stride=5, padding=1, count_include_pad=False), atol=atol, rtol=1e-5, forward_only=True)

  def test_interpolate_linear(self):
    for in_sz, out_sz in [((52,),(29,)), ((29,),(52,))]:
@@ -5,7 +5,6 @@ from tinygrad import Tensor, Device, dtypes
 from tinygrad.nn.optim import Adam, SGD, AdamW, Muon
 from tinygrad.helpers import CI
 from tinygrad.device import is_dtype_supported
-from extra.torch_muon import SingleDeviceMuon as TorchMuon

 np.random.seed(1337)
 x_init = np.random.randn(1,4).astype(np.float32)
@@ -58,12 +57,11 @@ class TestOptim(unittest.TestCase):
  def _test_sgd(self, steps, opts, atol, rtol): self._test_optim(SGD, torch.optim.SGD, steps, opts, atol, rtol)
  def _test_adam(self, steps, opts, atol, rtol): self._test_optim(Adam, torch.optim.Adam, steps, opts, atol, rtol)
  def _test_adamw(self, steps, opts, atol, rtol): self._test_optim(AdamW, torch.optim.AdamW, steps, opts, atol, rtol)
-  #TODO: use torch.muon when it comes out
-  def _test_muon(self, steps, opts, atol, rtol): self._test_optim(Muon, TorchMuon, steps, opts, atol, rtol)
+  def _test_muon(self, steps, opts, atol, rtol): self._test_optim(Muon, torch.optim.Muon, steps, opts, atol, rtol)

  def test_multistep_sgd_high_lr_teeny(self): self._test_sgd(2, {'lr': 1.1, 'teeny': True}, 1e-6, 1e-5)
  def test_multistep_adam_high_lr_teeny(self): self._test_adam(2, {'lr': 1.1, 'teeny': True}, 2e-4, 5e-4)
-  def test_multistep_muon_high_lr_teeny(self): self._test_muon(2, {'lr': 1.1, 'teeny': True}, 2e-4, 5e-4)
+  def test_multistep_muon_high_lr_teeny(self): self._test_muon(2, {'lr': 1.1, 'teeny': True}, 1e-2, 5e-4)

  def test_sgd(self): self._test_sgd(1, {'lr': 0.001}, 1e-6, 0)
  def test_sgd_high_lr(self): self._test_sgd(1, {'lr': 10}, 1e-6, 1e-5)
@@ -87,27 +85,34 @@ class TestOptim(unittest.TestCase):
  def test_multistep_sgd_high_lr_nesterov_momentum_wd(self):
    self._test_sgd(10, {'lr': 9, 'momentum': 0.9, 'nesterov': True, 'weight_decay': 0.1}, 1e-5, 3e-4)

-  def test_muon(self): self._test_muon(1, {'lr': 0.001}, 1e-6, 0)
-  def test_muon_high_lr(self): self._test_muon(1, {'lr': 10}, 1e-6, 3e-4)
-  def test_muon_wd(self): self._test_muon(1, {'lr': 0.001, 'weight_decay': 0.01}, 1e-6, 0)
-  def test_muon_high_lr_wd(self): self._test_muon(1, {'lr': 10, 'weight_decay': 0.01}, 1e-6, 5e-4)
+  def test_muon(self): self._test_muon(1, {'lr': 0.001}, 1e-3, 0)
+  # TODO: disabled due to big atol
+  # def test_muon_high_lr(self): self._test_muon(1, {'lr': 10}, 1e-6, 3e-4)
+  def test_muon_wd(self): self._test_muon(1, {'lr': 0.001, 'weight_decay': 0.01}, 1e-3, 3e-4)
+  # TODO: disabled due to big atol
+  # def test_muon_high_lr_wd(self): self._test_muon(1, {'lr': 10, 'weight_decay': 0.01}, 1e-6, 5e-4)

  # NOTE: momentum set to 0.95 by default, nesterov set to True by default
-  def test_multistep_muon_momentum_wd(self): self._test_muon(10, {'lr': 0.001, 'weight_decay': 0.01}, 1e-5, 0)
+  def test_multistep_muon_momentum_wd(self): self._test_muon(10, {'lr': 0.001, 'weight_decay': 0.01}, 3e-3, 0)
  # ns defaults are numerically unstable, but it is tolerable in real training (see nsteps/nparam tests)
-  def test_multistep_muon_high_lr_momentum_wd(self): self._test_muon(10, {'lr': 10, 'weight_decay': 0.01}, 1e-1, 3e-4)
-  def test_multistep_muon_no_nesterov_momentum(self): self._test_muon(10, {'lr': 0.001, 'nesterov': False}, 1e-5, 0)
-  def test_multistep_muon_high_lr_no_nesterov_momentum(self): self._test_muon(10, {'lr': 10, 'nesterov': False}, 0.5e-1, 1e-1)
+  # TODO: disabled due to big atol
+  # def test_multistep_muon_high_lr_momentum_wd(self): self._test_muon(10, {'lr': 10, 'weight_decay': 0.01}, 1e-1, 3e-4)
+  def test_multistep_muon_no_nesterov_momentum(self): self._test_muon(10, {'lr': 0.001, 'nesterov': False}, 1e-3, 0)
+  # TODO: disabled due to big atol
+  # def test_multistep_muon_high_lr_no_nesterov_momentum(self): self._test_muon(10, {'lr': 10, 'nesterov': False}, 5e-2, 1e-1)

-  def test_muon_ns_steps(self): self._test_muon(1, {'lr': 0.001, 'ns_steps': 3}, 1e-6, 0)
-  def test_muon_high_lr_ns_steps(self): self._test_muon(1, {'lr': 10, 'ns_steps': 3}, 1e-5, 3e-4)
-  def test_muon_ns_params(self): self._test_muon(1, {'lr': 0.001,'ns_params': (2.0,-1.5,0.5)}, 1e-6, 0)
-  def test_muon_high_lr_ns_params(self): self._test_muon(1, {'lr': 10,'ns_params': (2.0,-1.5,0.5)}, 1e-5, 3e-4)
+  def test_muon_ns_steps(self): self._test_muon(1, {'lr': 0.001, 'ns_steps': 3}, 1e-4, 0)
+  # TODO: disabled due to big atol
+  # def test_muon_high_lr_ns_steps(self): self._test_muon(1, {'lr': 10, 'ns_steps': 3}, 1e-5, 3e-4)
+  def test_muon_ns_coefficients(self): self._test_muon(1, {'lr': 0.001,'ns_coefficients': (2.0,-1.5,0.5)}, 1e-5, 3e-4)
+  # TODO: disabled due to big atol
+  # def test_muon_high_lr_ns_coefficients(self): self._test_muon(1, {'lr': 10,'ns_coefficients': (2.0,-1.5,0.5)}, 1e-5, 3e-4)

-  def test_muon_momentum_wd_ns_steps_ns_params(self):
-    self._test_muon(10, {'lr': 0.001, 'momentum': 0.90, 'weight_decay': 0.01, 'ns_steps': 3, 'ns_params': (2.0,-1.5,0.5)}, 1e-5, 0)
-  def test_multistep_muon_high_lr_momentum_wd_ns_steps_ns_params(self):
-    self._test_muon(10, {'lr': 10, 'momentum': 0.90, 'weight_decay': 0.01, 'ns_steps': 3, 'ns_params': (2.0,-1.5,0.5)}, 1e-5, 3e-4)
+  def test_muon_momentum_wd_ns_steps_ns_coefficients(self):
+    self._test_muon(10, {'lr': 0.001, 'momentum': 0.90, 'weight_decay': 0.01, 'ns_steps': 3, 'ns_coefficients': (2.0,-1.5,0.5)}, 1e-4, 0)
+  # TODO: disabled due to big atol
+  # def test_multistep_muon_high_lr_momentum_wd_ns_steps_ns_coefficients(self):
+  #   self._test_muon(10, {'lr': 10, 'momentum': 0.90, 'weight_decay': 0.01, 'ns_steps': 3, 'ns_coefficients': (2.0,-1.5,0.5)}, 1e-5, 3e-4)

  def test_adam(self): self._test_adam(1, {'lr': 0.001}, 1e-5, 0)
  def test_adam_high_lr(self): self._test_adam(1, {'lr': 10}, 1e-4, 1e-4)
@@ -1,6 +1,6 @@
 import unittest
 from tinygrad import Tensor, nn, Device
-from tinygrad.helpers import Context, GlobalCounters, CI, getenv
+from tinygrad.helpers import Context, GlobalCounters, CI, getenv, PCONTIG
 from tinygrad.uop.ops import graph_rewrite, PatternMatcher, UPat, Ops
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.nir import NIRRenderer
@@ -42,33 +42,40 @@ elif getenv("BIG") > 0:
 else:
  BS, HEADS, SEQLEN, EMB = 4, 2, 16, 8

+def fa():
+  Tensor.manual_seed(1337)
+  with Context(DEBUG=0): q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize() for _ in range(3)]
+  GlobalCounters.reset()
+  return q.scaled_dot_product_attention(k, v)
+
+def fa_bw():
+  Tensor.manual_seed(1337)
+  with Context(DEBUG=0):
+    q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize().requires_grad_() for _ in range(3)]
+    attn_output = nn.Linear(HEADS*EMB, HEADS*EMB, bias=False)
+    attn_output.weight.requires_grad_().realize()
+    target = Tensor.rand(BS, SEQLEN, HEADS*EMB).contiguous().realize()
+
+  GlobalCounters.reset()
+  attn = q.scaled_dot_product_attention(k, v).contiguous().contiguous_backward()
+  attn = attn.transpose(1, 2).reshape(BS, SEQLEN, -1)
+  out = attn_output(attn)
+  loss = (out - target).square().mean()
+  loss.backward()
+  #ret = [out, Tensor.stack(q.grad, k.grad, v.grad, dim=-1)]
+  #ret = [out, Tensor.stack(q.grad, k.grad, dim=-1), v.grad]
+  ret = [out, q.grad, k.grad, v.grad]
+  Tensor.realize(*ret)
+  return ret
+
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (NIRRenderer, PTXRenderer)), "broken in LVP and PTX")
 class TestPcontig(unittest.TestCase):
  def test_flash_attention_bw(self):
-    def fa_bw():
-      Tensor.manual_seed(1337)
-      with Context(DEBUG=0):
-        q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize().requires_grad_() for _ in range(3)]
-        attn_output = nn.Linear(HEADS*EMB, HEADS*EMB, bias=False)
-        attn_output.weight.requires_grad_().realize()
-        target = Tensor.rand(BS, SEQLEN, HEADS*EMB).contiguous().realize()
-
-      GlobalCounters.reset()
-      attn = q.scaled_dot_product_attention(k, v).contiguous().contiguous_backward()
-      attn = attn.transpose(1, 2).reshape(BS, SEQLEN, -1)
-      out = attn_output(attn)
-      loss = (out - target).square().mean()
-      loss.backward()
-      #ret = [out, Tensor.stack(q.grad, k.grad, v.grad)]
-      ret = [out, q.grad, k.grad, v.grad]
-      Tensor.realize(*ret)
-      return ret
-
-    with Context(PCONTIG=2, DEBUG=2):
+    with Context(PCONTIG=max(2, PCONTIG.value), DEBUG=2):
      grads = fa_bw()
      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")

-    with Context(DEBUG=2):
+    with Context(PCONTIG=0, DEBUG=2):
      cmp_grads = fa_bw()
      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")

@@ -79,17 +86,11 @@ class TestPcontig(unittest.TestCase):
    self.assertLessEqual(mse, 1e-6)

  def test_flash_attention(self):
-    def fa():
-      Tensor.manual_seed(1337)
-      with Context(DEBUG=0): q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize() for _ in range(3)]
-      GlobalCounters.reset()
-      return q.scaled_dot_product_attention(k, v).realize()
-
    with Context(PCONTIG=2, DEBUG=2):
-      ret = fa()
+      ret = fa().realize()
      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
    with Context(DEBUG=2):
-      cmp = fa()
+      cmp = fa().realize()
      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
    with Context(DEBUG=0):
      mse = ((cmp-ret)**2).sum().item()
@@ -370,6 +370,7 @@ class TestSchedule(unittest.TestCase):

  # NOTE: this is causing "LAZYCACHE=1 incorrectly reuses contiguous const" #4562
  # should contiguous dedup?
+  @unittest.skip("we do the exact opposite now")
  def test_dedup_contiguous(self):
    a = Tensor.ones(4).contiguous()
    b = Tensor.ones(4).contiguous()
@@ -446,7 +447,7 @@ class TestSchedule(unittest.TestCase):
  @unittest.skipUnless(is_dtype_supported(dtypes.ulong), "Needs ulong")
  def test_fold_conv_batchnorm_optim(self):
    # this is too high
-    for optim, cnt in [(nn.optim.Adam, 30), (nn.optim.SGD, 13)]:
+    for optim, cnt in [(nn.optim.Adam, 28), (nn.optim.SGD, 8)]:
      with self.subTest(optim=optim.__name__):
        with Tensor.train():
          img = Tensor.ones(1,3,4,4)
@@ -759,7 +760,7 @@ class TestSchedule(unittest.TestCase):

  def test_pow_neg_05_is_rsqrt(self):
    t = Tensor([1.0, 2.0, 3.0]) ** -0.5
-    self.assertEqual(self._alu_from_tensor(t), [Ops.RECIP, Ops.SQRT])
+    self.assertEqual(self._alu_from_tensor(t), [Ops.RECIPROCAL, Ops.SQRT])

  def test_pow_2_has_1_mul(self):
    t = Tensor([1.0, 2.0, 3.0]) ** Tensor(2.0)
@@ -1220,7 +1221,7 @@ class TestSchedule(unittest.TestCase):
      _realize_weights(layer)
      opt = nn.optim.Adam(nn.state.get_parameters(layer), lr=1e-4)
      layer(x).relu().sum().backward()
-      check_schedule(opt.schedule_step(), 16)
+      check_schedule(opt.schedule_step(), 19)

  def test_adam_conv_fuse(self):
    with Tensor.train():
@@ -1230,7 +1231,7 @@ class TestSchedule(unittest.TestCase):
      opt = nn.optim.Adam(nn.state.get_parameters(c1), lr=1e-4)
      opt.zero_grad()
      c1(img).relu().sum().backward()
-      check_schedule(opt.schedule_step(), 16)
+      check_schedule(opt.schedule_step(), 19)

  def test_adam_2convs_fuse(self):
    with Tensor.train():
@@ -1241,7 +1242,7 @@ class TestSchedule(unittest.TestCase):
      opt = nn.optim.Adam(nn.state.get_parameters([c1, c2]), lr=1e-4)
      opt.zero_grad()
      c2(c1(img).relu()).relu().sum().backward()
-      check_schedule(opt.schedule_step(), 18)
+      check_schedule(opt.schedule_step(), 21)

  def test_sgd_conv_fuse(self):
    with Tensor.train():
@@ -1863,7 +1864,7 @@ class TestSchedule(unittest.TestCase):
    yt = Tensor.randn(BS, 10).realize()
    with Context(SPLIT_REDUCEOP=0):
      loss = yt.sparse_categorical_crossentropy(Y_train[samples])
-      run_schedule(check_schedule(loss, 5))
+      run_schedule(check_schedule(loss, 4))
      loss_fused = loss.numpy()
    loss_ref = torch.nn.CrossEntropyLoss()(torch.tensor(yt.numpy()), torch.tensor(Y_train.numpy())[torch.tensor(samples.numpy())])
    np.testing.assert_allclose(loss_fused, loss_ref.numpy(), atol=1e-6, rtol=1e-6)
@@ -2076,6 +2077,11 @@ class TestCopyFolding(unittest.TestCase):
    check_schedule(b, 0, filter_sink=False)
    assert b.item() == 1

+  def test_one_hot_with_copy(self):
+    y = Tensor([1, 2, 3]).to("CPU")
+    x = y.one_hot(10)
+    check_schedule(x, 3, filter_sink=False)
+
  def test_const_copy_multi(self):
    x = Tensor.ones(1, device="CPU").to_(["CPU", "CPU:1"])
    check_schedule(x, 0, filter_sink=False)
@@ -2085,7 +2091,7 @@ class TestCopyFolding(unittest.TestCase):
    a = Tensor.arange(3).realize()
    zeros = Tensor.zeros(3).realize()
    b = (a*zeros).to("CPU")
-    run_schedule(check_schedule(b, 2, filter_sink=False)) # TODO: 0?
+    run_schedule(check_schedule(b, 0, filter_sink=False))
    self.assertListEqual(b.tolist(), [0, 0, 0])
    self.assertEqual(b.device, "CPU")

@@ -839,12 +839,11 @@ class TestTensorMetadata(unittest.TestCase):
    self.assertEqual(y.grad.uop.metadata[0].name, "sigmoid")
    self.assertTrue(y.grad.uop.metadata[0].backward)
    si = Tensor.schedule(out, x.grad, y.grad)[-1]
-    self.assertEqual(len(si.metadata), 4, f"failed with {si.metadata}")
-    self.assertSetEqual(set(m.name for m in si.metadata), {"__mul__", "sigmoid", "relu"})
+    self.assertEqual(len(si.metadata), 3, f"failed with {si.metadata}")
+    self.assertSetEqual(set(m.name for m in si.metadata), {"sigmoid", "relu"})
    bw = [m for m in si.metadata if m.backward]
-    self.assertEqual(len(bw), 2)
-    self.assertEqual(bw[0].name, "__mul__")
-    self.assertEqual(bw[1].name, "sigmoid")
+    self.assertEqual(len(bw), 1)
+    self.assertEqual(bw[0].name, "sigmoid")

 class TestIdxUpcast(unittest.TestCase):
  def _find_op(self, ast: UOp, op: Ops):
@@ -920,5 +919,38 @@ class TestIdxUpcast(unittest.TestCase):
    a = Tensor.empty(2**11, 2**11, 1, dtype=dtypes.int8).permute((2, 0, 1)).expand((2**9+10, -1, -1)).contiguous()
    a.realize()

+class TestTensorUnique(unittest.TestCase):
+  def test_empty_bufs_unique(self):
+    a = Tensor.empty(10, 10).contiguous()
+    b = Tensor.empty(10, 10).contiguous()
+    Tensor.realize(a,b)
+    self.assertIsNot(a.uop.buffer, b.uop.buffer)
+
+  def test_zeros_bufs_unique_sep(self):
+    a = Tensor.zeros(10, 10).contiguous()
+    Tensor.realize(a)
+    b = Tensor.zeros(10, 10).contiguous()
+    Tensor.realize(b)
+    self.assertIsNot(a.uop.buffer, b.uop.buffer)
+
+  def test_zeros_bufs_unique(self):
+    a = Tensor.zeros(10, 10).contiguous()
+    b = Tensor.zeros(10, 10).contiguous()
+    Tensor.realize(a,b)
+    self.assertIsNot(a.uop.buffer, b.uop.buffer)
+
+  def test_eye_bufs_unique(self):
+    a = Tensor.eye(10).contiguous()
+    b = Tensor.eye(10).contiguous()
+    Tensor.realize(a,b)
+    self.assertIsNot(a.uop.buffer, b.uop.buffer)
+
+  def test_times_2_not_unique(self):
+    a = Tensor.zeros(10, 10).contiguous()
+    b = a * 2
+    c = a * 2
+    Tensor.realize(b,c)
+    self.assertIs(b.uop.buffer, c.uop.buffer)
+
 if __name__ == '__main__':
  unittest.main()
@@ -1,12 +1,11 @@
-from typing import List
 import unittest, pytest
 from tinygrad import dtypes, Variable
 from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import DEBUG, Context
-from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, track_rewrites, graph_rewrite, GroupOp, KernelInfo
+from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, track_rewrites, graph_rewrite, GroupOp, AxisType
 from tinygrad.uop.symbolic import sym
-from tinygrad.codegen import full_rewrite, full_rewrite_to_sink
 from tinygrad.codegen.late.expander import expander
+from test.test_uops import to_uops_list

 simple_pm = PatternMatcher([
  (UPat.cvar('x', dtypes.int), lambda x: UOp.const(dtypes.float, 1.0) + UOp.const(dtypes.float, 2.0)),
@@ -15,12 +14,6 @@ simple_pm = PatternMatcher([
  ((UPat.var('x') + UPat.cvar('c1')) + UPat.cvar('c2'), lambda x,c1,c2: x + (c1.arg+c2.arg)),
 ])

-def to_uops_list(u:List[UOp]) -> List[UOp]:
-  # we strip the SINK here for legacy reasons
-  ret = full_rewrite(UOp.sink(*u, arg=KernelInfo(opts_to_apply=())))
-  assert ret[-1].op is Ops.SINK
-  return ret[:-1]
-
 class TestGraphRewriteConst(unittest.TestCase):
  def test_gep_const(self):
    v1 = UOp.const(dtypes.int.vec(3), (0,1,2))
@@ -270,6 +263,7 @@ class TestUOpGraph(unittest.TestCase):
    uops = to_uops_list([out])
    self.assertEqual(len([x for x in uops if x.op is Ops.VECTORIZE]), 0)

+  @unittest.skip("this test isn't valid uops")
  def test_gep_vec_fold(self):
    d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
    d1 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 1)
@@ -459,30 +453,29 @@ class TestUOpGraph(unittest.TestCase):
    idx = d0.index(ridx0)
    ld = idx.load()
    val = (ridx0<50).where(5, ld)
-    st = idx.store(val, ridx0)
+    st = idx.store(val).end(ridx0)
    uops = to_uops_list([st])
    for u in uops:
      assert u.op is not Ops.WHERE
      if u.op is Ops.STORE: assert u.src[1].arg==5

  def test_load_idx_becomes_int(self):
-    # These loads wont overflow int since we know from the gate that the value is bounded
-    r0 = UOp.range(10, 0)
-    d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 0)
-    d1 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 1)
-    l0 = UOp(Ops.LOAD, dtypes.long, (d0.index(UOp.const(dtypes.int, 0)),)).cast(dtypes.index)
-    idx = l0 * 600
-    valid = (l0<-1).ne(True)&(l0<3000)
-    l1 = UOp(Ops.LOAD, dtypes.long, (d1.index(idx.valid(valid)),))
-    uops = to_uops_list([l1])
+    # mnist indexing with split reduceop
+    # Make sure we are not doign math on the loaded index, which would promote it to long
+    c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.uchar.ptr(128000), arg=0, src=())
+    c1 = UOp.range(UOp.const(dtypes.index, 512), 1, AxisType.LOOP)
+    c2 = UOp.range(UOp.const(dtypes.index, 250), 2, AxisType.LOOP)
+    c3 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(512), arg=1, src=())
+    c4 = c3.index(c1).load()
+    c5 = UOp.range(UOp.const(dtypes.index, 240), 0, AxisType.REDUCE)
+    c6 = ((c2*UOp.const(dtypes.index, 240))+c5)
+    c7 = UOp(Ops.DEFINE_GLOBAL, dtypes.uchar.ptr(60000), arg=2, src=())
+    c8 = c7.index(c6).load()
+    c9 = ((c4<0).where((c4+60000), c4)!=c6.cast(dtypes.int)).where(0, c8.cast(dtypes.uint).cast(dtypes.uchar)).reduce(c5, arg=Ops.ADD)
+    c10 = c0.index(((c1*UOp.const(dtypes.index, 250))+c2)).store(c9).end(c1, c2)
+    uops = to_uops_list([c10])
    for u in uops:
-      if u.op is Ops.INDEX: self.assertEqual(u.src[1].dtype, dtypes.int)
-
-    valid = (10*r0<5-l0).ne(True)&(l0<3000)
-    l2 = UOp(Ops.LOAD, dtypes.long, (d1.index(idx.valid(valid)),))
-    uops = to_uops_list([l2])
-    for u in uops:
-      if u.op is Ops.INDEX: self.assertEqual(u.src[1].dtype, dtypes.int)
+      self.assertNotEqual(u.dtype, dtypes.long)

  def test_in_out_of_bounds_access(self):
    with Context(IGNORE_OOB=0):
@@ -518,6 +511,7 @@ class TestUOpGraph(unittest.TestCase):
      st1 = UOp(Ops.STORE, dtypes.void, (glbl0.index(v), v, v<20))
      with self.assertRaises(RuntimeError): to_uops_list([st1])

+  @unittest.skip("if not allowed in graph")
  def test_in_bounds_access_gated_local(self):
    with Context(IGNORE_OOB=0):
      # Define buffers
@@ -579,7 +573,7 @@ class TestUOpGraph(unittest.TestCase):
  def test_in_out_bounds_access_with_mask(self):
    with Context(IGNORE_OOB=0):
      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
-      gidx0 = UOp(Ops.SPECIAL, dtypes.index, (UOp.const(dtypes.index, 42),), "gidx0")
+      gidx0 = UOp.range(42, 0, AxisType.GLOBAL)
      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(gidx0, (5<gidx0)&(gidx0<16)),))
      ld1 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(gidx0, gidx0<16),))
      to_uops_list([ld0, ld1])
@@ -603,7 +597,7 @@ class TestUOpGraph(unittest.TestCase):
    with Context(IGNORE_OOB=0):
      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
      glbl1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(64), (), 0)
-      gidx0 = UOp(Ops.SPECIAL, dtypes.index, (UOp.const(dtypes.index, 42),), "gidx0")
+      gidx0 = UOp.range(42, 0, AxisType.GLOBAL)
      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(gidx0, gidx0<8),)).cast(dtypes.index)
      ld1 = UOp(Ops.LOAD, dtypes.int, (glbl1.index(ld0*2, (ld0>=0)&(ld0<32)),))
      to_uops_list([ld1])
@@ -638,13 +632,13 @@ class TestUOpGraph(unittest.TestCase):
    lidx = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 16),), "lidx0")
    st = UOp(Ops.STORE, dtypes.void, (smem.index(lidx), UOp.load(glbl0.index(lidx), dtype=dtypes.int)))
    barrier = UOp(Ops.BARRIER, dtypes.void, (st, ))
-    ld0 = UOp(Ops.LOAD, dtypes.int, (smem.index(UOp.invalid()), barrier))
-    ld1 = UOp(Ops.LOAD, dtypes.int, (smem.index(lidx+2, UOp.const(dtypes.bool, True)), barrier))
+    ld0 = UOp(Ops.LOAD, dtypes.int, (smem.after(barrier).index(UOp.invalid()),))
+    ld1 = UOp(Ops.LOAD, dtypes.int, (smem.after(barrier).index(lidx+2, UOp.const(dtypes.bool, True)),))
    uops = to_uops_list([UOp(Ops.STORE, dtypes.void, (glbl0.index(lidx), ld1+ld0))])

    ld0 = uops[-1].src[-1]
    # the gate and invalid value are deleted from ld1
-    self.assertEqual(ld0.src[0], smem.index(lidx+2))
+    self.assertEqual(ld0.src[0], smem.after(barrier).index(lidx+2))

  def test_fold_gated_store(self):
    glbl = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
@@ -726,7 +720,7 @@ class TestExpander(unittest.TestCase):
    self.assertTupleEqual(sink.src[0].arg, (0,2,1,3,4,6,5,7))

  def test_contract_no_expand(self):
-    e1 = UOp(Ops.DEFINE_VAR, dtypes.int)
+    e1 = UOp.variable("i", 0, 10, dtype=dtypes.int)
    con = UOp(Ops.CONTRACT, dtypes.int.vec(2), (e1,), ((2,2),))
    sink = expander_rewrite(con)
    assert sink.op is Ops.VECTORIZE and len(sink.src) == 2
@@ -815,54 +809,6 @@ class TestExpander(unittest.TestCase):
    sink = expander_rewrite(sink)
    print(sink)

-class TestIFUOps(unittest.TestCase):
-  def test_create_ifs(self):
-    gbuf = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
-    sbuf = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(size=4, addrspace=AddrSpace.LOCAL), (), "smem")
-    valid = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 10),), "gidx0")<5
-    lidx = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), "lidx0")
-    gate = valid&(lidx.ne(2))
-    idx = UOp.const(dtypes.int, 0)
-    st = UOp(Ops.STORE, dtypes.void, (sbuf.index(idx), UOp.const(dtypes.float, 42)))
-    barrier = UOp(Ops.BARRIER, dtypes.void, (st,))
-    lbuf = UOp(Ops.LOAD, dtypes.float, (sbuf.index(UOp.const(dtypes.int, 0)), barrier))
-    store = UOp(Ops.STORE, dtypes.void, (gbuf.index(UOp.const(dtypes.int, 0), gate), lbuf))
-    sink = UOp(Ops.SINK, dtypes.void, (store,))
-    sink = full_rewrite_to_sink(sink)
-    if_uops = [u for u in sink.toposort() if u.op is Ops.IF]
-    self.assertEqual(len(if_uops), 1)
-    self.assertEqual(if_uops[0].src[0], gate)
-
-  def test_expand_ifs_one_gate(self):
-    gbuf = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
-    sbuf = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(size=16, addrspace=AddrSpace.LOCAL), (), "smem")
-    valid = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), "gidx0")<1
-    lidx = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 16),), "lidx0")
-    gate = valid&(lidx.ne(2))
-    st = UOp(Ops.STORE, dtypes.void, (sbuf, lidx, UOp.const(dtypes.float, 42)))
-    barrier = UOp(Ops.BARRIER, dtypes.void, (st,))
-    lbufs = [UOp(Ops.LOAD, dtypes.float, (sbuf.index(UOp.const(dtypes.int, i)), barrier)) for i in range(4)]
-    stores = [UOp(Ops.STORE, dtypes.void, (gbuf.index(UOp.const(dtypes.int, i), gate), lbufs[i])) for i in range(4)]
-    sink = UOp(Ops.SINK, dtypes.void, tuple(stores))
-    sink = full_rewrite_to_sink(sink)
-    if_uops = [u for u in sink.toposort() if u.op is Ops.IF]
-    self.assertEqual(len(if_uops), 1)
-    self.assertEqual(if_uops[0].src[0], gate)
-
-  # this will be fixed with the merge gated stores bounty
-  @unittest.expectedFailure
-  def test_expand_ifs_dumb(self):
-    buf = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
-    valid = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 10),), "gidx0")<5
-    lidx = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), "lidx0")
-    gate = valid&(lidx.ne(2))
-    stores = [UOp(Ops.STORE, dtypes.void, (buf, UOp.const(dtypes.int, i), UOp.const(dtypes.float, i), gate)) for i in range(4)]
-    sink = UOp(Ops.SINK, dtypes.void, tuple(stores))
-    sink = full_rewrite_to_sink(sink)
-    if_uops = [u for u in sink.toposort() if u.op is Ops.IF]
-    self.assertEqual(len(if_uops), 1)
-    self.assertEqual(if_uops[0].src[0], gate)
-
 class TestUOpTags(unittest.TestCase):
  def test_inc_by_one(self):
    g = UOp.const(dtypes.int, 1) + UOp.const(dtypes.int, 1)
@@ -6,7 +6,7 @@ from tinygrad.helpers import CI, DEBUG, getenv, Timing
 from tinygrad.dtype import dtypes, DType, AddrSpace
 from tinygrad.device import Buffer, Device
 from tinygrad.uop.ops import Ops, UOp, UPat, KernelInfo, exec_alu # noqa F401
-from tinygrad.uop.spec import spec
+from tinygrad.uop.spec import shared_spec
 from tinygrad.renderer import ProgramSpec
 from tinygrad.engine.realize import CompiledRunner, get_program
 from tinygrad.codegen import full_rewrite
@@ -15,10 +15,16 @@ from tinygrad.device import is_dtype_supported
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.renderer.ptx import PTXRenderer

-def to_uops_list(u:list[UOp], opts=None, skip_check=False) -> list[UOp]: return full_rewrite(UOp.sink(*u), opts)
+def to_uops_list(u:list[UOp], ren=None) -> list[UOp]:
+  sink = UOp.group(*u)
+  for r in sink.ranges: sink = sink.end(r)
+  # we strip the SINK here for legacy reasons
+  ret = full_rewrite(sink.sink(arg=KernelInfo(opts_to_apply=())), ren)
+  assert ret[-1].op is Ops.SINK
+  return ret[:-1]

 def _uops_to_prg(uops_list):
-  uops = full_rewrite(ast:=UOp.sink(*uops_list), opts=Device[Device.DEFAULT].renderer)
+  uops = full_rewrite(ast:=UOp.sink(*uops_list), ren=Device[Device.DEFAULT].renderer)
  src = Device[Device.DEFAULT].renderer.render(uops)
  has_local = Device[Device.DEFAULT].renderer.has_local
  return CompiledRunner(ProgramSpec(uops[-1].arg.name if uops[-1].arg is not None else "test", src, Device.DEFAULT, ast, uops=uops,
@@ -109,7 +115,7 @@ class TestFloatUOps(TestUOps):
  def test_log2(self): self._test_uop_fxn(Ops.LOG2, lambda a: math.log2(a) if a > 0 else float('-inf' if a==0 else 'nan'))
  @unittest.skipIf(Device.DEFAULT == "CPU", 'not supported as uop')
  def test_sin(self): self._test_uop_fxn(Ops.SIN, lambda a: math.sin(a))
-  def test_recip(self): self._test_uop_fxn(Ops.RECIP, lambda a: 1/a if a != 0 else float('inf'))
+  def test_recip(self): self._test_uop_fxn(Ops.RECIPROCAL, lambda a: 1/a if a != 0 else float('inf'))
  def test_sqrt(self): self._test_uop_fxn(Ops.SQRT, lambda a: math.sqrt(a) if a >= 0 else float('nan'))

  def test_add(self): self._test_bop_fxn(Ops.ADD, lambda a,b: a+b)
@@ -212,18 +218,18 @@ class TestExecALU(TestUOps):
    self.assertEqual(exec_alu(Ops.IDIV, dtypes.int8, (7, -3)), -2)
    self.assertEqual(exec_alu(Ops.IDIV, dtypes.int8, (-50, 6)), -8)

-    np.testing.assert_allclose(exec_alu(Ops.MUL, dtypes.float32, (7.0, exec_alu(Ops.RECIP, dtypes.float32, (3.0,)))), 2+(1.0/3.0))
-    np.testing.assert_allclose(exec_alu(Ops.MUL, dtypes.float32, (7.0, exec_alu(Ops.RECIP, dtypes.float32, (-3.0,)))), -2-(1.0/3.0))
+    np.testing.assert_allclose(exec_alu(Ops.MUL, dtypes.float32, (7.0, exec_alu(Ops.RECIPROCAL, dtypes.float32, (3.0,)))), 2+(1.0/3.0))
+    np.testing.assert_allclose(exec_alu(Ops.MUL, dtypes.float32, (7.0, exec_alu(Ops.RECIPROCAL, dtypes.float32, (-3.0,)))), -2-(1.0/3.0))

  def test_recip(self):
-    np.testing.assert_allclose(exec_alu(Ops.RECIP, dtypes.float32, (8,)), 1/8)
-    np.testing.assert_allclose(exec_alu(Ops.RECIP, dtypes.float32, (7,)), 1/7)
-    np.testing.assert_allclose(exec_alu(Ops.RECIP, dtypes.float32, (-3,)), 1/-3)
-    np.testing.assert_allclose(exec_alu(Ops.RECIP, dtypes.float32, (-50,)), 1/-50)
+    np.testing.assert_allclose(exec_alu(Ops.RECIPROCAL, dtypes.float32, (8,)), 1/8)
+    np.testing.assert_allclose(exec_alu(Ops.RECIPROCAL, dtypes.float32, (7,)), 1/7)
+    np.testing.assert_allclose(exec_alu(Ops.RECIPROCAL, dtypes.float32, (-3,)), 1/-3)
+    np.testing.assert_allclose(exec_alu(Ops.RECIPROCAL, dtypes.float32, (-50,)), 1/-50)

-    np.testing.assert_allclose(exec_alu(Ops.RECIP, dtypes.float32, ((32+521+3),)), 1/(32+521+3))
-    np.testing.assert_allclose(exec_alu(Ops.RECIP, dtypes.float32, ((34**2),)), 1/(34**2))
-    np.testing.assert_allclose(exec_alu(Ops.RECIP, dtypes.float32, (10,)), 1/10)
+    np.testing.assert_allclose(exec_alu(Ops.RECIPROCAL, dtypes.float32, ((32+521+3),)), 1/(32+521+3))
+    np.testing.assert_allclose(exec_alu(Ops.RECIPROCAL, dtypes.float32, ((34**2),)), 1/(34**2))
+    np.testing.assert_allclose(exec_alu(Ops.RECIPROCAL, dtypes.float32, (10,)), 1/10)

  def test_bool_cmplt(self):
    self.assertEqual(exec_alu(Ops.CMPLT, dtypes.bool, (False, False)), False)
@@ -266,6 +272,7 @@ class TestConstantFolding(unittest.TestCase):
    si = t.schedule()
    assert len(si) == 0

+@unittest.skip("no more if statements")
 class TestGatedStoreRewrite(unittest.TestCase):
  def test_tiny_gate_store(self):
    gmem = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
@@ -302,6 +309,7 @@ class TestGatedStoreRewrite(unittest.TestCase):
    self.assertIs(gated_uops[-1].op, Ops.STORE)

  # scaled down version of TestLinearizerDumb.test_unmerged_ifs
+  @unittest.skip("we don't merge ifs anymore")
  def test_merge_ifs_alt(self):
    gmem0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
    gmem1 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 1)
@@ -331,7 +339,7 @@ class TestLocalAccess(unittest.TestCase):
    smem = uop(uops, Ops.DEFINE_LOCAL, dtypes.float32.ptr(size=16, addrspace=AddrSpace.LOCAL), (), 'smem')
    st = uop(uops, Ops.STORE, dtypes.void, (smem.index(uop(uops, Ops.CONST, dtypes.int32, (), 0)), uop(uops, Ops.CONST, dtypes.float32, (), 42.0)))
    barr = uop(uops, Ops.BARRIER, dtypes.void, (st,))
-    sres = uop(uops, Ops.LOAD, dtypes.float32, (smem.index(uop(uops, Ops.CONST, dtypes.int32, (), 0)), barr))
+    sres = uop(uops, Ops.LOAD, dtypes.float32, (smem.after(barr).index(uop(uops, Ops.CONST, dtypes.int32, (), 0)),))
    self.assertEqual(_test_uops_result(dtypes.float32, uops, sres), 42)

  # NOTE: webgpu specific, since only webgpu performs bitpacking
@@ -341,7 +349,7 @@ class TestLocalAccess(unittest.TestCase):
    smem = uop(uops, Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=16, addrspace=AddrSpace.LOCAL), (), 'smem')
    st = uop(uops, Ops.STORE, dtypes.void, (smem.index(uop(uops, Ops.CONST, dtypes.int32, (), 0)), uop(uops, Ops.CONST, dtypes.uint8, (), 42)))
    barr = uop(uops, Ops.BARRIER, dtypes.void, (st,))
-    sres = uop(uops, Ops.LOAD, dtypes.uint8, (smem.index(uop(uops, Ops.CONST, dtypes.int32, (), 0)), barr))
+    sres = uop(uops, Ops.LOAD, dtypes.uint8, (smem.after(barr).index(uop(uops, Ops.CONST, dtypes.int32, (), 0)),))
    self.assertEqual(_test_uops_result(dtypes.uint8, uops, sres), 42)

  # NOTE: webgpu specific, since only webgpu performs bitpacking
@@ -351,7 +359,7 @@ class TestLocalAccess(unittest.TestCase):
    size = 16
    for dtype in _dtypes:
      temp = UOp(Ops.DEFINE_LOCAL, dtype.ptr(size=size, addrspace=AddrSpace.LOCAL), (), 'smem')
-      uops = to_uops_list([temp], opts=Device[Device.DEFAULT].renderer)
+      uops = to_uops_list([temp], ren=Device[Device.DEFAULT].renderer)
      out = Device[Device.DEFAULT].renderer.render(uops)
      # half is supported in wgsl, so it doesn't have to be packed
      corrected_size = size//(4//dtype.itemsize) if dtype != dtypes.half else size
@@ -378,7 +386,7 @@ class TestAssembly(unittest.TestCase):
    l1 = UOp(Ops.LOAD, dtypes.int, (g1.index(c1),))
    a1 = UOp(Ops.MUL, dtypes.int, (l1, c1))
    a2 = UOp(Ops.MUL, dtypes.int, (l1, c2))
-    uops = to_uops_list([a1,a2], opts=Device[Device.DEFAULT].renderer)
+    uops = to_uops_list([a1,a2], ren=Device[Device.DEFAULT].renderer)
    Device[Device.DEFAULT].renderer.render(uops)
    ops = [x.op for x in uops]
    self.assertIn(Ops.SHL, ops)
@@ -390,7 +398,7 @@ class TestAssembly(unittest.TestCase):
      c = UOp(Ops.CONST, dt, (), 2)
      l = UOp(Ops.LOAD, dt, (g.index(c),))
      a = UOp(Ops.IDIV, dt, (l, c))
-      uops = to_uops_list([a], opts=Device[Device.DEFAULT].renderer)
+      uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
      Device[Device.DEFAULT].renderer.render(uops)
      ops = [x.op for x in uops]
      self.assertIn(Ops.SHR, ops, f"For dtype={dt} divison by power of two did not simplify to shift")
@@ -401,14 +409,14 @@ class TestAssembly(unittest.TestCase):
    c = UOp(Ops.CONST, dtypes.uint, (), 3)
    l = UOp(Ops.LOAD, dtypes.uint, (g.index(c),))
    a = UOp(Ops.IDIV, dtypes.uint, (l, c))
-    uops = to_uops_list([a], opts=Device[Device.DEFAULT].renderer)
+    uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
    Device[Device.DEFAULT].renderer.render(uops)
    ops = [x.op for x in uops]
    self.assertIn(Ops.SHR, ops)
    self.assertNotIn(Ops.IDIV, ops)

    b = UOp(Ops.MOD, dtypes.uint, (l, c))
-    uops = to_uops_list([b], opts=Device[Device.DEFAULT].renderer)
+    uops = to_uops_list([b], ren=Device[Device.DEFAULT].renderer)
    Device[Device.DEFAULT].renderer.render(uops)
    ops = [x.op for x in uops]
    self.assertIn(Ops.SHR, ops)
@@ -421,7 +429,7 @@ class TestAssembly(unittest.TestCase):
    c = UOp(Ops.CONST, dtypes.uint, (), 7)
    l = UOp(Ops.LOAD, dtypes.uint, (g.index(c),))
    a = UOp(Ops.IDIV, dtypes.uint, (l, c))
-    uops = to_uops_list([a], opts=Device[Device.DEFAULT].renderer)
+    uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
    Device[Device.DEFAULT].renderer.render(uops)
    ops = [x.op for x in uops]
    self.assertIn(Ops.SHR, ops)
@@ -429,7 +437,7 @@ class TestAssembly(unittest.TestCase):

  def test_fast_idiv_remove_powers_of_two(self):
    ridx = UOp.range(2**20, 0)
-    uops = to_uops_list([ridx//(7*64)], opts=Device[Device.DEFAULT].renderer)
+    uops = to_uops_list([ridx//(7*64)], ren=Device[Device.DEFAULT].renderer)
    ops = [x.op for x in uops]
    # this requires shifting out the powers of two before doing fast_idiv
    # (((ridx0>>6)*18725)>>17) instead of (int)((((long)(ridx0)*1198373)>>29))
@@ -453,7 +461,7 @@ class TestAssembly(unittest.TestCase):
    c = UOp(Ops.CONST, dtypes.uint, (), 7)
    l = UOp(Ops.LOAD, dtypes.uint, (g.index(c),))
    comp = l.ne(c).ne(True)
-    uops = to_uops_list([comp], opts=Device[Device.DEFAULT].renderer)
+    uops = to_uops_list([comp], ren=Device[Device.DEFAULT].renderer)
    Device[Device.DEFAULT].renderer.render(uops)
    ops = [x.op for x in uops]
    self.assertIn(Ops.CMPEQ, ops)
@@ -512,7 +520,7 @@ class TestUOpStr(unittest.TestCase):
 class TestUPatHelpers(unittest.TestCase):
  def test_location(self):
    self.assertEqual(sym.patterns[-1][0].location[0].replace("\\", "/").split("/")[-1], "symbolic.py")
-    self.assertEqual(spec.patterns[0][0].location[0].replace("\\", "/").split("/")[-1], "spec.py")
+    self.assertEqual(shared_spec.patterns[0][0].location[0].replace("\\", "/").split("/")[-1], "spec.py")
    test_upat = UPat(Ops.CONST, dtypes.bool)
    self.assertEqual(test_upat.location[0].split("/")[-1], __file__.replace("\\", "/").split("/")[-1])
    test_upat_named = test_upat.named("test_name")
@@ -540,11 +548,24 @@ class TestUopsObject(unittest.TestCase):

 class TestUOpRender(unittest.TestCase):
  def test_render_vectorize_same(self):
-    u = UOp(Ops.VECTORIZE, src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0)))
-    self.assertEqual(u.render(), "{0, ...}")
+    u = UOp(Ops.VECTORIZE, dtype=dtypes.int.vec(3), src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0)))
+    self.assertEqual(u.render(simplify=False), "{0, ...}")
  def test_render_vectorize_different(self):
-    u = UOp(Ops.VECTORIZE, src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1), UOp.const(dtypes.int, 2)))
-    self.assertEqual(u.render(), "{0,1,2}")
+    u = UOp(Ops.VECTORIZE, dtype=dtypes.int.vec(3), src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1), UOp.const(dtypes.int, 2)))
+    self.assertEqual(u.render(simplify=False), "{0,1,2}")
+  def test_render_vectorize_same_simplified(self):
+    u = UOp(Ops.VECTORIZE, dtype=dtypes.int.vec(3), src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0)))
+    self.assertEqual(u.render(), "0")
+  def test_render_vectorize_different_simplified(self):
+    u = UOp(Ops.VECTORIZE, dtype=dtypes.int.vec(3), src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1), UOp.const(dtypes.int, 2)))
+    self.assertEqual(u.render(), "(0, 1, 2)")
+
+class TestZeroRange(unittest.TestCase):
+  def test_reduce_variable(self):
+    for i in range(3,-1,-1):
+      v = UOp.variable("i", 0, 5).bind(i)
+      out = Tensor.ones(10, dtype=dtypes.int).contiguous().shrink(((0,v),)).sum()
+      self.assertEqual(out.item(), i)

 if __name__ == '__main__':
  unittest.main(verbosity=2)
@@ -1,11 +1,12 @@
 import unittest, math
 from tinygrad import dtypes
-from tinygrad.helpers import all_same
+from tinygrad.helpers import all_same, Context
 from tinygrad.uop.ops import GroupOp, UOp, Ops, exec_alu, PatternMatcher, TrackedPatternMatcher, UPat
 from tinygrad.codegen import full_rewrite_to_sink
 from hypothesis import given, strategies as strat

 # Helper function to apply the graph rewrite
+@Context(SPEC=0)
 def apply_rewrite(expr):
  return full_rewrite_to_sink(expr.sink()).src[0]

@@ -305,19 +306,19 @@ class TestRecurse(unittest.TestCase):
    graph_rewrite(a, pm, bottom_up=True)

  def test_inf_loop(self):
-    a = UOp.variable('a', 0, 10)
+    a = UOp.const(dtypes.int, 3)
    pm = PatternMatcher([
-      (UPat(Ops.DEFINE_VAR, name="x"), lambda x: x.replace(op=Ops.CONST)),
-      (UPat(Ops.CONST, name="x"), lambda x: x.replace(op=Ops.DEFINE_VAR)),
+      (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)),
+      (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)),
    ])
    with self.assertRaises(RuntimeError):
      graph_rewrite(a, pm)

  def test_inf_loop_bottom_up(self):
-    a = UOp.variable('a', 0, 10)
+    a = UOp.const(dtypes.int, 3)
    pm = PatternMatcher([
-      (UPat(Ops.DEFINE_VAR, name="x"), lambda x: x.replace(op=Ops.CONST)),
-      (UPat(Ops.CONST, name="x"), lambda x: x.replace(op=Ops.DEFINE_VAR)),
+      (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)),
+      (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)),
    ])
    with self.assertRaises(RuntimeError):
      graph_rewrite(a, pm, bottom_up=True)
@@ -3,6 +3,8 @@ import hashlib, random, unittest
 from tinygrad import Tensor, Device, getenv, dtypes
 from tinygrad.device import is_dtype_supported
 from tinygrad.helpers import CI
+from tinygrad.uop.ops import UOp
+from tinygrad.engine.jit import TinyJit

@unittest.skipUnless(is_dtype_supported(dtypes.uint8) and is_dtype_supported(dtypes.uint64), "Device must support uint8 and uint64")
@unittest.skipIf(getenv("MOCKGPU") and Device.DEFAULT == "NV", "crashes in NV CI")
@@ -72,5 +74,52 @@ class TestKeccak(unittest.TestCase):
    data = b"\x00" * 1000
    self.assertEqual(bytes(Tensor(data).keccak("shake_128").tolist()), hashlib.shake_128(data).digest(16))

+  def test_variable_bs(self):
+    data = Tensor([b"abc", b"abc", b"abc"], dtype=dtypes.uint8).repeat(2048, 1)
+
+    bs = UOp.variable("bs", 1, 4096).bind(1)
+    out = data.shrink_to(bs, data.shape[-1]).keccak().shrink_to(1, 32)
+    self.assertEqual(bytes(out[0].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
+
+    bs = UOp.variable("bs", 1, 4096).bind(2)
+    out = data.shrink_to(bs, data.shape[-1]).keccak().shrink_to(2, 32)
+    self.assertEqual(bytes(out[0].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
+    self.assertEqual(bytes(out[1].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
+
+    bs = UOp.variable("bs", 1, 4096).bind(3)
+    data = Tensor([b"abc", b"abc", b"def"], dtype=dtypes.uint8).repeat(2048, 1)
+    out = data.shrink_to(bs, data.shape[-1]).keccak().shrink_to(3, 32)
+    self.assertEqual(bytes(out[0].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
+    self.assertEqual(bytes(out[1].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
+    self.assertEqual(bytes(out[2].tolist()), bytearray.fromhex("8e0d8f672252acb0 ffc5093db8653b18 1513bf9a2097e737 b4f73533dcaf46df"))
+
+  def test_variable_bs_jit(self):
+    def f(data):
+      return data.keccak()
+    jit_f = TinyJit(f)
+
+    data = Tensor([b"abc", b"abc", b"abc"], dtype=dtypes.uint8).repeat(2048, 1)
+
+    # initialize jit
+    for _ in range(3):
+      bs = UOp.variable("bs", 1, 4096).bind(4096)
+      _ = jit_f(data.shrink_to(bs, data.shape[-1]))
+
+    bs = UOp.variable("bs", 1, 4096).bind(1)
+    out = jit_f(data.shrink_to(bs, data.shape[-1])).shrink_to(1, 32)
+    self.assertEqual(bytes(out[0].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
+
+    bs = UOp.variable("bs", 1, 4096).bind(2)
+    out = jit_f(data.shrink_to(bs, data.shape[-1])).shrink_to(2, 32)
+    self.assertEqual(bytes(out[0].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
+    self.assertEqual(bytes(out[1].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
+
+    bs = UOp.variable("bs", 1, 4096).bind(3)
+    data = Tensor([b"abc", b"abc", b"def"], dtype=dtypes.uint8).repeat(2048, 1)
+    out = jit_f(data.shrink_to(bs, data.shape[-1])).shrink_to(3, 32)
+    self.assertEqual(bytes(out[0].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
+    self.assertEqual(bytes(out[1].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
+    self.assertEqual(bytes(out[2].tolist()), bytearray.fromhex("8e0d8f672252acb0 ffc5093db8653b18 1513bf9a2097e737 b4f73533dcaf46df"))
+
 if __name__ == "__main__":
  unittest.main()
@@ -50,7 +50,7 @@ class TestPatternMatcher(unittest.TestCase):
    def fxn(ctx, x):
      ctx.append(True)
      assert len(x.src) == 0
-      return UOp(Ops.CONST, src=(UOp(Ops.CONST),))
+      return x.replace(src=(UOp(Ops.DEVICE, arg="blah"),))
    matcher = PatternMatcher([(UPat(Ops.CONST, src=(), name="x"), fxn)])
    c1 = UOp(Ops.CONST, dtypes.float, arg=1.0)
    # second rewrite shouldn't match anything
@@ -41,13 +41,13 @@ class TestHelpers(unittest.TestCase):
    self.assertTrue(f2.is_increasing())
    self.assertTrue(f3.is_increasing())

-    rng = UOp(Ops.RANGE, dtypes.int, arg=(2, True), src=(UOp(Ops.CONST, dtypes.int, arg=5, src=()),))
+    rng = UOp.range(5, 2)
    self.assertTrue(rng.is_increasing())
    self.assertTrue((rng+2).is_increasing())

 class TestValidIdxSimplification(unittest.TestCase):
  def check(self, load, sidx, svalid):
-    with Context(NOOPT=1):
+    with Context(NOOPT=1, SPEC=0):
      load = full_rewrite_to_sink(load.sink()).src[0]
    idx, valid = load.src[0].src[1], load.src[0].src[2]
    check_uop_against_string(self, idx, sidx)
@@ -213,7 +213,7 @@ class TestValidIdxSimplification(unittest.TestCase):

 class TestImageSimplification(unittest.TestCase):
  def check(self, load, svalid, sidx0, sidx1):
-    with Context(NOOPT=1):
+    with Context(NOOPT=1, SPEC=0):
      load = full_rewrite_to_sink(load.sink()).src[0]
    idx = load.src[0].src[1]
    self.assertEqual(idx.op, Ops.VECTORIZE)
@@ -283,7 +283,8 @@ class TestImageSimplification(unittest.TestCase):

    # empty -> invalid
    load = get_load_image_uop(shape, (gidx0<8) & (gidx0<8).ne(True), idx)
-    load = full_rewrite_to_sink(load.sink()).src[0]
+    with Context(NOOPT=1, SPEC=0):
+      load = full_rewrite_to_sink(load.sink()).src[0]
    self.assertEqual(load.op, Ops.VECTORIZE)
    self.assertEqual(load.dtype.count, 4)

@@ -7,7 +7,7 @@ from tinygrad.codegen import full_rewrite
 from tinygrad.helpers import Context
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer
 from tinygrad.uop.symbolic import sym, commutative
-from tinygrad.uop.spec import uops_to_z3
+from tinygrad.uop.validate import uops_to_z3

 def check_uop_against_string(self, v:UOp, s:str):
  sym_vars = {v.render():v for v in v.toposort() if v.op in (Ops.DEFINE_VAR, Ops.RANGE, Ops.SPECIAL)}
@@ -40,15 +40,14 @@ class TestVminVmaxProperties(unittest.TestCase):
    self.assertEqual(uop.vmin, 0)
    self.assertEqual(uop.vmax, 5)

-    # this can be improved
    uop = x & 15
    self.assertEqual(uop.vmin, 0)
    self.assertEqual(uop.vmax, 15)

-    # this can be improved
+    # TODO: this can be improved
    uop = x & 32
    self.assertEqual(uop.vmin, 0)
-    self.assertEqual(uop.vmax, 20)
+    self.assertEqual(uop.vmax, 20) # shoud be 0

  def test_vmin_vmax_multiplication_with_variable(self):
    # vmin and vmax for multiplication with a variable
@@ -1,5 +1,5 @@
 import unittest
-from tinygrad.helpers import DEBUG
+from tinygrad.helpers import DEBUG, Context
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import UPat, track_rewrites, GroupOp, Ops
 from tinygrad.uop.upat import _get_code, upat_compile
@@ -14,6 +14,7 @@ def do_compile(up):
  if DEBUG >= 2: dis.dis(match)
  return match_code[0]

+@Context(SPEC=0)
 class TestUPatCompile(unittest.TestCase):
  def test_double(self):
    up = UPat.var("x") * UPat.cvar("c0") + UPat.var("x") * UPat.cvar("c1")
@@ -157,11 +157,11 @@ class TestViz(BaseTestViz):
    self.assertEqual(ansistrip(a2["label"]), "CUSTOM\nx\nyzww\nw")

  def test_inf_loop(self):
-    a = UOp.variable('a', 0, 10, dtype=dtypes.int)
-    b = a.replace(op=Ops.CONST)
+    a = UOp.const(dtypes.int, 3)
+    b = UOp.const(dtypes.int, 4)
    pm = PatternMatcher([
-      (UPat(Ops.DEFINE_VAR, name="x"), lambda x: x.replace(op=Ops.CONST)),
-      (UPat(Ops.CONST, name="x"), lambda x: x.replace(op=Ops.DEFINE_VAR)),
+      (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)),
+      (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)),
    ])
    with self.assertRaises(RuntimeError): exec_rewrite(a, [pm])
    graphs = flatten(x["graph"].values() for x in get_viz_details(0, 0))
@@ -1,9 +1,7 @@
-from typing import Any, Callable
-import functools
-from dataclasses import dataclass
-from tinygrad.helpers import QUANTIZE, DEVECTORIZE, TRANSCENDENTAL
+import itertools
+from tinygrad.helpers import QUANTIZE, DEVECTORIZE, TRANSCENDENTAL, SPEC
 from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype
-from tinygrad.uop.spec import type_verify
+from tinygrad.uop.spec import type_verify, program_spec, kernel_spec
 from tinygrad.renderer import Renderer

 # import all pattern matchers here
@@ -14,109 +12,103 @@ from tinygrad.uop.decompositions import get_late_rewrite_patterns
 from tinygrad.codegen.late.expander import migrate_indexing, expander, pm_pre_expander, pm_group_for_reduce
 from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_indexing, devectorize, pm_reduce, \
  ReduceContext, correct_load_store, pm_render
-from tinygrad.codegen.opt.postrange import pm_postrange_opt
-from tinygrad.codegen.simplify import pm_simplify_ranges, pm_reduce_simplify, pm_flatten_range, pm_split_ranges
-from tinygrad.schedule.rangeify import pm_add_buffers, rangeify_codegen
-from tinygrad.codegen.late.control_flow import CFGContext, pm_merge_ends, pm_add_control_flow, linearize
+from tinygrad.codegen.opt.postrange import apply_opts
+from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_split_ranges, pm_load_collapse, pm_split_store
+from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen
+from tinygrad.codegen.late.linearizer import CFGContext, pm_prepare_control_flow, pm_add_control_flow, linearize

-@dataclass
-class RewriteStep:
-  pm: PatternMatcher
-  ctx: Callable[[UOp], Any]|None = None
-  name: str|None = None
-  bottom_up: bool = False
-  def __call__(self, sink:UOp):
-    return graph_rewrite(sink, self.pm, ctx=self.ctx(sink) if self.ctx is not None else None, name=self.name, bottom_up=self.bottom_up)
+def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -> UOp:
+  if ren is None: ren = Renderer()

-def apply_rewrites(sink:UOp, rewrites:list[RewriteStep]): return functools.reduce(lambda x,f: f(x), rewrites, sink)
-
-def get_rewrites_for_renderer(opts:Renderer, optimize:bool=True, linearizer:bool=True) -> list[RewriteStep]:
-  # cache with the values of the context vars
-  return _get_rewrites_for_renderer(opts, optimize, linearizer, QUANTIZE.value, DEVECTORIZE.value, TRANSCENDENTAL.value)
-
-@functools.cache
-def _get_rewrites_for_renderer(opts:Renderer, optimize:bool, linearizer:bool, _QUANTIZE, _DEVECTORIZE, _TRANSCENDENTAL) -> list[RewriteStep]:
-  # ** lowerer **
-  ret: list[RewriteStep] = []
+  if SPEC: type_verify(sink, kernel_spec)

+  # first we optimize
  if optimize:
+    if QUANTIZE and ren.device in {"CPU", "DSP"}: sink = graph_rewrite(sink, pm_quant, name="quantize")

-    # lowerer first
-    if _QUANTIZE and opts.device in {"CPU", "DSP"}: ret.append(RewriteStep(pm_quant, name="quantize"))
+    # TODO: fix expander and remove this
+    sink = graph_rewrite(sink, pm_add_buffers_local, name="add locals early")
+
+    # collapse loads reduce (indexing by a tensor)
+    sink = graph_rewrite(sink, pm_load_collapse, name="load collapse")

    # split ranges
-    ret.append(RewriteStep(pm_split_ranges+pm_flatten_range, ctx=lambda _: {}, name="split ranges"))
+    sink = graph_rewrite(sink, pm_split_ranges+pm_flatten_range, ctx={}, name="split ranges")

    # symbolic (NOTE: this is a requirement for pm_simplify_ranges to be correct)
-    ret.append(RewriteStep(sym+pm_flatten_range, name="initial symbolic"))
+    sink = graph_rewrite(sink, sym+pm_flatten_range, name="initial symbolic")

    # optimize (schedule) the AST
-    ret.append(RewriteStep(pm_simplify_ranges, name="simplify ranges"))
-    ret.append(RewriteStep(pm_reduce_simplify, name="simplify reduces"))
-    ret.append(RewriteStep(pm_postrange_opt, ctx=lambda _: opts, name="post optimize ast"))
+    sink = graph_rewrite(sink, pm_simplify_ranges, name="simplify ranges")
+
+    # split store range (only on CPU for now)
+    sink = graph_rewrite(sink, pm_split_store, ctx=ren.device, name="cut store ranges")
+
+    # do postrange optimization, BEAM or hand_coded_optimizations
+    sink = apply_opts(sink, ren)

  # ** expander (expand_rewrite) **
-  ret.append(RewriteStep(sym+migrate_indexing+pm_move_where_on_load, name="postopt symbolic"))
+  sink = graph_rewrite(sink, sym+migrate_indexing+pm_move_where_on_load, name="postopt symbolic")

  # expand
-  ret.append(RewriteStep(sym+pm_pre_expander+pm_group_for_reduce+expander, name="expander"))
+  sink = graph_rewrite(sink, sym+pm_pre_expander+pm_group_for_reduce+expander, name="expander")

  # add locals
-  ret.append(RewriteStep(pm_add_buffers+rangeify_codegen, name="add local buffers"))
+  sink = graph_rewrite(sink, pm_add_buffers_local+rangeify_codegen, name="add local buffers")

  # ** devectorizer (full_graph_rewrite) **
  # remove reduce
-  ret.append(RewriteStep(pm_reduce+gep_pushing, lambda _: ReduceContext(), name="remove_reduce"))
+  sink = graph_rewrite(sink, pm_reduce+gep_pushing, ctx=ReduceContext(), name="remove_reduce")

  # add gpu dims (late). this works after devectorize, but it's faster here
-  ret.append(RewriteStep(pm_add_gpudims, lambda _: opts, name="add gpudims"))
+  sink = graph_rewrite(sink, pm_add_gpudims, ctx=ren, name="add gpudims")

  # devectorize (TODO: does this need opts?)
-  if _DEVECTORIZE >= 2: pm_devectorize = sym+load_store_folding+load_store_indexing
-  elif _DEVECTORIZE: pm_devectorize = sym+devectorize+load_store_folding+correct_load_store+load_store_indexing
+  if DEVECTORIZE >= 2: pm_devectorize = sym+load_store_folding+load_store_indexing
+  elif DEVECTORIZE: pm_devectorize = sym+devectorize+load_store_folding+correct_load_store+load_store_indexing
  else: pm_devectorize = sym+load_store_folding+correct_load_store+load_store_indexing
-  ret.append(RewriteStep(pm_devectorize, lambda _: opts, name="devectorize"))
-
-  supported_ops = tuple(opts.code_for_op.keys())
-  extra_matcher = opts.extra_matcher if opts.extra_matcher is not None else PatternMatcher([])
+  sink = graph_rewrite(sink, pm_devectorize, ctx=ren, name="devectorize")

  # lower the index dtype to a concrete int
-  ret.append(RewriteStep(pm_lower_index_dtype+load_store_indexing, lambda _: opts.device, name="lower all index dtypes"))
-  ret.append(RewriteStep(symbolic, name="post index symbolic"))
+  sink = graph_rewrite(sink, pm_lower_index_dtype+load_store_indexing, ctx=ren.device, name="lower all index dtypes")
+  sink = graph_rewrite(sink, symbolic, name="post index symbolic")

  # optional pre matcher
-  if opts.pre_matcher is not None: ret.append(RewriteStep(opts.pre_matcher, name="pre_matcher"))
+  if ren.pre_matcher is not None: sink = graph_rewrite(sink, ren.pre_matcher, name="pre_matcher")

  # decompositions
-  pm_decomp = symbolic_simple+get_late_rewrite_patterns(supported_ops, _TRANSCENDENTAL>=2)
-  ret.append(RewriteStep(pm_decomp, lambda _: opts.device, name="decompositions"))
+  supported_ops = tuple(ren.code_for_op.keys())
+  pm_decomp = symbolic_simple+get_late_rewrite_patterns(supported_ops, TRANSCENDENTAL>=2)
+  sink = graph_rewrite(sink, pm_decomp, ctx=ren.device, name="decompositions")

  # final rules for the renderer (without sym)
+  extra_matcher = ren.extra_matcher if ren.extra_matcher is not None else PatternMatcher([])
  pm_final_rewrite = pm_decomp+pm_render+extra_matcher
-  ret.append(RewriteStep(pm_final_rewrite, lambda _: opts.device, name="final rewrite"))
+  sink = graph_rewrite(sink, pm_final_rewrite, ctx=ren.device, name="final rewrite")
+
+  # prepare for control flow
+  sink = graph_rewrite(sink, pm_prepare_control_flow, ctx=itertools.count(10000), name="split ends + add if ranges")

  # this was the linearizer
-  ret.append(RewriteStep(pm_merge_ends, name="merge ends"))
-  ret.append(RewriteStep(pm_add_control_flow, CFGContext, name="add control flow starts", bottom_up=True))
+  sink = graph_rewrite(sink, pm_add_control_flow, ctx=CFGContext(sink), name="add control flow", bottom_up=True)

-  # return the list
-  return ret
+  # return the rewritten sink
+  return sink

-def full_rewrite_to_sink(sink:UOp, opts:Renderer|None=None, optimize:bool=True) -> UOp:
-  return apply_rewrites(sink, get_rewrites_for_renderer(opts if opts is not None else Renderer(), optimize))
-
-def full_rewrite(sink:UOp, opts:Renderer|None=None) -> list[UOp]:
+def full_rewrite(sink:UOp, ren:Renderer|None=None) -> list[UOp]:
  """
  Function to transform the Kernel UOp graph into a linearized program.

  Args:
    sink: The Ops.SINK rooting the Kernel graph.
-    opts: The Renderer (can change how things are processed, fix this).
+    ren: The Renderer (can change how things are processed, fix this).

  Returns:
    Linear program in UOps.
  """

-  lst = linearize(full_rewrite_to_sink(sink, opts, optimize=sink.tag is None))
-  if __debug__: type_verify(lst)
+  full_sink = full_rewrite_to_sink(sink, ren, optimize=sink.tag is None)
+  assert len(full_sink.ranges) == 0, "all ranges must end by the sink"
+  lst = linearize(full_sink)
+  if SPEC: type_verify(lst, program_spec)
  return lst
@@ -1,7 +1,7 @@
 import math, functools, operator
 from tinygrad.uop.ops import UOp, Ops, sint, PatternMatcher, UPat, KernelInfo, ssimplify, AxisType, sint_to_uop
 from tinygrad.helpers import all_int, dedup, get_contraction
-from tinygrad.dtype import dtypes
+from tinygrad.dtype import dtypes, AddrSpace, Invalid
 from tinygrad.renderer import Renderer

 def _group_dims(dims:tuple[sint, ...], max_sizes:tuple[int, ...]):
@@ -79,6 +79,14 @@ def add_gpudims(ctx:Renderer, s:UOp):
  # apply to multiple ranges
  subs = {}
  for r in s_topo:
+    # look for local INDEXes that are not used in the GLOBAL store, then add them as an INVALID
+    if r.op is Ops.STORE and r.src[0].ptrdtype.addrspace == AddrSpace.GLOBAL:
+      idx = r.src[0]
+      missing_locals = [all_ranges[rng] for rng in local_dims if all_ranges[rng] not in idx.ranges]
+      if len(missing_locals):
+        assert len(idx.src) == 2, "index has 2 sources"
+        mask: UOp = functools.reduce(operator.and_, [x.eq(0) for x in missing_locals])
+        subs[idx] = idx.replace(src=(idx.src[0], mask.broadcast(idx.src[1].dtype.count).where(idx.src[1], Invalid)))
    if r.op is not Ops.RANGE: continue
    try:
      ii = (global_dims+local_dims).index(r.arg[0:-1])
@@ -87,15 +95,7 @@ def add_gpudims(ctx:Renderer, s:UOp):
    except ValueError: continue
  return s.substitute(subs)

-def add_barrier_and_if(buf:UOp, e:UOp):
-  # TODO: this is not generic
-  local_ranges = [x for x in e.ended_ranges if x.op is Ops.RANGE and x.arg[-1] == AxisType.GROUP_REDUCE]
-  if len(local_ranges) == 0: return None
-  return buf.after(UOp(Ops.IF, dtype=dtypes.void, src=(functools.reduce(operator.and_, [x.eq(0) for x in local_ranges]), e.barrier())))
-
 pm_add_gpudims = PatternMatcher([
  # add gpudims must be last
  (UPat(Ops.SINK, name="s"), add_gpudims),
-  # add barrier and if
-  (UPat(Ops.AFTER, src=(UPat(Ops.DEFINE_LOCAL, name="buf"), UPat(Ops.END, name="e"))), add_barrier_and_if),
 ])
@@ -1,102 +0,0 @@
-import heapq
-from collections import defaultdict
-from tinygrad.uop.ops import PatternMatcher, UOp, Ops, UPat
-
-def linearize(u:UOp) -> list[UOp]:
-  lst = list(u.toposort())
-  in_this_block = set(lst)
-  local_children: defaultdict[UOp, list[UOp]] = defaultdict(list)
-  in_degree:dict[UOp, int] = {}
-  priorities:dict[UOp, int] = {}
-
-  # get local children and assign priorities
-  # NOTE: this requires the lst be locally toposorted
-  for u in reversed(lst):
-    in_degree[u] = 0
-    for s in u.src:
-      if s in in_this_block:
-        local_children[s].append(u)
-        in_degree[u] += 1
-    # put loads in the beginning of the block and prevent priority inversion. hack for BARRIER grouping too
-    priority = [0] + [priorities[x] for x in local_children[u]]
-    if u.op is Ops.LOAD: priority.append(-1000)
-    if u.op is Ops.BARRIER: priority.append(-1500)
-    # ranges are scheduled as late as possible so anything that can be outside is
-    #if u.op is Ops.RANGE: priority = [2000]
-    # move defines and consts to the top
-    if u.op in {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.DEFINE_VAR, Ops.SPECIAL, Ops.CONST}: priority.append(-2000)
-    priorities[u] = min(priority)
-
-  # number the uops in "ideal" order
-  nkey = {u:i for i,u in enumerate(sorted(lst, key=lambda x: (priorities[x],)+x.tuplize))}
-
-  # then force then to be toposorted in as close to the ideal order as possible
-  heapq.heapify(heap:=[(nkey[u],u) for u in lst if in_degree[u] == 0])
-  newlst = []
-  while heap:
-    newlst.append(u:=heapq.heappop(heap)[1])
-    for v in local_children[u]:
-      in_degree[v] -= 1
-      if in_degree[v] == 0: heapq.heappush(heap, (nkey[v],v))
-
-  assert len(newlst) == len(lst), f"len mismatch {len(newlst)} != {len(lst)}"
-  return newlst
-
-class CFGContext:
-  def __init__(self, sink:UOp):
-    # there are 3 relationships between ranges:
-    # nested, meaning endrange y is a dependency of endrange x and range x is a dependency of endrange y
-    # dependent, meaning endrange y is a dependency of endrange x and range x is not a dependency of endrange y
-    # independent, endrange y is not a dependency of endrange x
-    # everything is nested inside the sink
-    deps: dict[UOp, set[UOp]] = {}
-    nesting: dict[UOp, UOp] = {}
-    for u in sink.toposort():
-      deps[u] = set().union(*(deps[s] for s in u.src))
-      if u.op in (Ops.END, Ops.ENDIF, Ops.SINK):
-        nesting |= {x:u for x in deps[u] if x.op in (Ops.END, Ops.ENDIF) and (u.op is Ops.SINK or u.src[0] in deps[x]) and x not in nesting}
-      if u.op in (Ops.RANGE, Ops.END, Ops.IF, Ops.ENDIF): deps[u] |= {u}
-
-    self.edges: dict[UOp, UOp] = {}
-    siblings: dict[UOp, list[UOp]] = {}
-    for k,vv in nesting.items(): siblings.setdefault(vv, []).append(k)
-    for k,v in siblings.items():
-      # range/if that have dependencies on other siblings need to run after them
-      order = sorted(v, key=lambda x: len(deps[x].intersection(v)))
-      zipped = zip(order, order[1:]) if k.op is Ops.SINK else zip([k.src[0]] + order, order)
-      for x,y in zipped:
-        # TODO: is this check correct?
-        if y.src[0] not in x.backward_slice_with_self:
-          self.edges[y.src[0]] = x
-
-pm_add_control_flow = PatternMatcher([
-  (UPat((Ops.RANGE, Ops.IF), name="x"), lambda ctx,x: x.replace(src=x.src+(y,)) if (y:=ctx.edges.get(x)) is not None else None),
-])
-
-def do_merge_ends(s:UOp):
-  # NOTE: this can fail
-  stacked: dict[UOp, list[UOp]] = {}
-  dangling_ifs = []
-  for x in s.toposort():
-    if x.op in {Ops.END, Ops.ENDIF}:
-      assert x.op is not Ops.END or x.arg == 1, "ends must be single ends for linearizer"
-      stacked.setdefault(x.src[0], []).append(x)
-    if x.op is Ops.IF: dangling_ifs.append(x)
-  dangling_ifs = [x for x in dangling_ifs if x not in stacked]
-  replaces = {}
-  for k,v in stacked.items():
-    if len(v) == 1: continue
-    rep = UOp(v[0].op, src=tuple([k] + [y for x in v for y in x.src[1:]]), arg=x[0].arg)
-    for x in v: replaces[x] = rep
-  if not len(replaces) and not len(dangling_ifs): return None
-  ret = s.substitute(replaces)
-  if len(dangling_ifs):
-    assert len(dangling_ifs) == 1, "we only support 1 dangling if"
-    ret = ret.replace(src=(UOp(Ops.ENDIF, src=(dangling_ifs[0], *ret.src)),))
-  return ret
-
-pm_merge_ends = PatternMatcher([
-  # for renderering and linearizing, all ends must end one loop
-  (UPat(Ops.END, name="e"), lambda e: e.replace(src=e.src[e.arg-1:], arg=1).end(ends=e.src[:e.arg-1]) if e.arg > 1 else None),
-  (UPat(Ops.SINK, name="s"), do_merge_ends),
-])
@@ -50,7 +50,6 @@ def delete_redundant_gates(store:UOp, buf:UOp, idx:UOp, val:UOp, store_gate:UOp,
  # remove the gate from the index
  return UOp.store(buf.index(idx).cast(cast.dtype) if cast is not None else buf.index(idx), val, *store.src[2:])

-def no_load(u:UOp) -> bool: return not any(x.op is Ops.LOAD for x in u.backward_slice_with_self)
 load_store_indexing = PatternMatcher([
  # image load valid idx simplification
  (UPat(Ops.INDEX, src=(UPat.var("buf"), invalid_gate)), lambda buf,x,i,cond: simplify_valid_load(buf, x, cond)),
@@ -61,8 +60,6 @@ load_store_indexing = PatternMatcher([
  # delete_redundant_gates (after expand)
  (UPat(Ops.STORE, src=(UPat.any(stidx:=UPat.var("buf").index(UPat.var("idx"), UPat.var("store_gate")), stidx.cast().named("cast")),
                                  UPat.var("val")), name="store", allow_any_len=True), delete_redundant_gates),
-  # we want to make sure we dont do math on a loaded index since that can cause overflow, this undoes a pattern in reduce_collapse
-  (UPat.var("c")<(UPat.var("x", dtypes.index)+UPat.var("y")), lambda x,y,c: (-x < -(c-y)) if no_load(y) and no_load(c) and not no_load(x) else None),
 ])

 # ***** load/store grouping *****
@@ -112,7 +109,7 @@ def cat_after_store(cat:UOp, data:UOp, sto:UOp):
  for s in cat.src:
    ret.append(s.store(data.gep(tuple(range(offset, offset+s.dtype.count))), *sto.src[2:]))
    offset += s.dtype.count
-  return UOp(Ops.NOOP, src=tuple(ret))
+  return UOp.group(*ret)

 def gep_on_store(gep:UOp, st:UOp, sto:UOp):
  # NOTE: we need to invert the gep here, but it may be an expanding gep
@@ -182,7 +179,7 @@ def split_load_store(ctx:Renderer|None, ls:UOp, idx:UOp):

  # if it wasn't split, we return None. otherwise we CAT them
  if len(ret) <= 1: return None
-  return UOp(Ops.CAT, ls.dtype, tuple(ret)) if ls.op is Ops.LOAD else UOp(Ops.NOOP, src=tuple(ret))
+  return UOp(Ops.CAT, ls.dtype, tuple(ret)) if ls.op is Ops.LOAD else UOp.group(*ret)

 def image_fixup(ls:UOp):
  # normal image load or store, with the CAST from expand_index
@@ -268,10 +265,6 @@ pm_render = PatternMatcher([
    UPat.var("a")), lambda c,idx,l,a: l.replace(src=(l.src[0], a.cast(l.dtype))+l.src[2:]).cast(a.dtype)),
  (UPat.var("c").where(UPat.var("a"), UPat(Ops.LOAD, src=(UPat().index(UPat.var("idx"), UPat.var("c").logical_not()).or_casted(),),
    allow_any_len=True, name="l").or_casted()), lambda c,idx,l,a: l.replace(src=(l.src[0], a.cast(l.dtype))+l.src[2:]).cast(a.dtype)),
-  # gate any stores that aren't gated with if/endif pairs
-  (UPat(Ops.STORE, src=(UPat(src=(UPat(), UPat(), UPat(dtype=dtypes.bool)), name="idx").or_casted(), UPat()), name="store", allow_any_len=True),
-    lambda store,idx: UOp(Ops.ENDIF, src=(uif:=UOp(Ops.IF, src=(idx.src[2],)), UOp(Ops.STORE, src=store.src[:2]+(uif,)+store.src[2:]))) if \
-      len(store.src) <= 2 or store.src[2].op != Ops.IF else None),
 ])

 # *** Ops.REDUCE -> Ops.DEFINE_ACC ***
@@ -295,7 +288,7 @@ def reduce_to_acc(ctx:ReduceContext, red:UOp):
  # if we have a range
  if len(reduce_range) != 0:
    topo = inp.toposort()
-    ended_ranges = flatten([x.src[:x.arg] for x in topo if x.op is Ops.END])
+    ended_ranges = flatten([x.ended_ranges for x in topo if x.op is Ops.END])
    input_ranges = tuple([x for x in topo if x.op is Ops.RANGE and x not in reduce_range and x not in ended_ranges])
    identity = red.const(red.dtype, identity_element(red.arg, red.dtype.scalar()))
    acc = UOp(Ops.DEFINE_REG, red.dtype.ptr(size=1, addrspace=AddrSpace.REG), arg=(ctx.acc_num,))
@@ -305,7 +298,7 @@ def reduce_to_acc(ctx:ReduceContext, red:UOp):
    ctx.acc_num += 1
  ret = functools.reduce(lambda x,y: x.alu(red.arg, y), lst)
  if len(reduce_range) == 0: return ret
-  return acc.after(acc.index(UOp.const(dtypes.int, 0)).store(ret).end(ends=reduce_range[::-1])).index(UOp.const(dtypes.int, 0)).load()
+  return acc.after(acc.index(UOp.const(dtypes.int, 0)).store(ret).end(*reduce_range)).index(UOp.const(dtypes.int, 0)).load()

 pm_reduce = PatternMatcher([
  # REDUCE -> DEFINE_ACC+ASSIGN
@@ -0,0 +1,86 @@
+import heapq
+from collections import defaultdict
+from tinygrad.dtype import dtypes
+from tinygrad.uop.ops import PatternMatcher, UOp, Ops, UPat, AxisType, GroupOp
+
+def linearize(u:UOp) -> list[UOp]:
+  # this is a toposort with priority
+  lst = list(u.toposort())
+  consumers: defaultdict[UOp, list[UOp]] = defaultdict(list)
+  in_degree:dict[UOp, int] = {}
+  priorities:dict[UOp, int] = {}
+
+  # get consumers and assign priorities
+  # NOTE: this requires the lst be locally toposorted
+  for u in reversed(lst):
+    for s in u.src: consumers[s].append(u)
+    in_degree[u] = len(u.src)
+    # put loads in the beginning of the block and prevent priority inversion. hack for BARRIER grouping too
+    priority = [0] + [priorities[x] for x in consumers[u]]
+    if u.op is Ops.LOAD: priority.append(-1000)
+    if u.op is Ops.BARRIER: priority.append(-1500)
+    # ranges are scheduled as late as possible so anything that can be outside is
+    # if u.op is Ops.RANGE: priority = [2000]
+    if u.op is Ops.END: priority = [-1000]
+    # move defines and consts to the top
+    if u.op in {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.DEFINE_VAR, Ops.SPECIAL, Ops.CONST}: priority.append(-2000)
+    priorities[u] = min(priority)
+
+  # number the uops in "ideal" order
+  nkey = {u:i for i,u in enumerate(sorted(lst, key=lambda x: (priorities[x],)+x.tuplize))}
+
+  # then force then to be toposorted in as close to the ideal order as possible
+  heapq.heapify(heap:=[(nkey[u],u) for u in lst if in_degree[u] == 0])
+  newlst = []
+  while heap:
+    newlst.append(u:=heapq.heappop(heap)[1])
+    for v in consumers[u]:
+      in_degree[v] -= 1
+      if in_degree[v] == 0: heapq.heappush(heap, (nkey[v],v))
+  assert len(newlst) == len(lst), f"len mismatch {len(newlst)} != {len(lst)}"
+  return newlst
+
+class CFGContext:
+  def __init__(self, sink:UOp):
+    # there are 3 relationships between ranges:
+    # nested, meaning endrange y is a dependency of endrange x and range x is a dependency of endrange y
+    # dependent, meaning endrange y is a dependency of endrange x and range x is not a dependency of endrange y
+    # independent, endrange y is not a dependency of endrange x
+    # everything is nested inside the sink
+    deps: dict[UOp, dict[UOp, None]] = {}
+    nesting: dict[UOp, UOp] = {}
+    for u in sink.toposort():
+      # get the deps from the src
+      deps[u] = {}
+      for s in u.src: deps[u] |= deps[s]
+
+      if u.op in (Ops.END, Ops.SINK):
+        nesting |= {x:u for x in deps[u] if x.op is Ops.END and (u.op is Ops.SINK or u.src[1] in deps[x]) and x not in nesting}
+      if u.op in (Ops.RANGE, Ops.END): deps[u][u] = None
+
+    self.edges: dict[UOp, UOp] = {}
+    siblings: dict[UOp, list[UOp]] = {}
+    for k,vv in nesting.items(): siblings.setdefault(vv, []).append(k)
+    for k,v in siblings.items():
+      # ranges that have dependencies on other siblings need to be scheduled after them
+      order = sorted(v, key=lambda x: len([u for u in v if u in deps[x]]))
+      zipped = zip(order, order[1:]) if k.op is Ops.SINK else zip([k.src[1]] + order, order)
+      for x,y in zipped: self.edges[y.src[1]] = x
+
+pm_add_control_flow = PatternMatcher([
+  (UPat(Ops.RANGE, name="x"), lambda ctx,x: x.replace(src=x.src+(y,)) if (y:=ctx.edges.get(x)) is not None else None),
+])
+
+def do_split_ends(e:UOp):
+  ret = e.src[0]
+  for r in list(UOp.sink(*e.src[1:]).ranges)[::-1]: ret = ret.end(r)
+  return ret
+
+pm_prepare_control_flow = PatternMatcher([
+  # split the ends
+  (UPat(Ops.END, name="e"), do_split_ends),
+  # add if ranges
+  (UPat(GroupOp.Defines, name="buf").index(UPat.var("idx"), UPat(name="gate", dtype=dtypes.bool)).or_casted("cast").store(UPat.var("val")),
+    lambda ctx,buf,idx,gate,cast,val:
+      buf.after(r:=UOp.range(gate.cast(dtypes.int), next(ctx), AxisType.IF, dtype=dtypes.int)).index(idx, gate).cast(cast.dtype).store(val).end(r)),
+])
@@ -2,7 +2,6 @@
 from __future__ import annotations
 from enum import Enum, auto
 from dataclasses import dataclass
-from tinygrad.uop.ops import AxisType

 class OptOps(Enum):
  TC = auto(); UPCAST = auto(); UNROLL = auto(); LOCAL = auto(); THREAD = auto() # noqa: E702
@@ -16,11 +15,6 @@ class Opt:
  arg: int|tuple|None = None
  def __repr__(self): return f"Opt(op={self.op}, axis={self.axis}, arg={self.arg})"

-axis_letters = {AxisType.GLOBAL: "g", AxisType.THREAD: "t", AxisType.LOCAL: "l", AxisType.WARP: "w", AxisType.LOOP: "L", AxisType.UPCAST: "u",
-                AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.UNROLL: "r"}
-axis_colors = {AxisType.GLOBAL: "blue", AxisType.THREAD: "BLUE", AxisType.LOCAL: "cyan", AxisType.WARP: "CYAN", AxisType.LOOP: "WHITE",
-               AxisType.UPCAST: "yellow", AxisType.GROUP_REDUCE: "RED", AxisType.REDUCE: "red", AxisType.UNROLL: "magenta"}
-
 class KernelOptError(Exception): pass
 def check(cond:bool, msg:str=""):
  if not cond: raise KernelOptError(msg)
@@ -27,15 +27,15 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
  # NOTE: unless TC_OPT is > 0, we only trigger tensor cores if there's only one reduce axis
  if USE_TC > 0 and (len(k.axes_of(AxisType.GROUP_REDUCE, AxisType.REDUCE)) == 1 or (TC_OPT.value >= 1)):
    good_tc_opt = False
+    tk = k.copy()
    try: # check TC first and apply hand-coded opts if successful
-      tk = k.copy()
      rngs = tk.apply_opt(Opt(OptOps.TC, 0, (TC_SELECT.value, TC_OPT.value, USE_TC.value)))
      good_tc_opt = True
    except KernelOptError:
      pass
-    if good_tc_opt:
-      # skip hand-coded TC opts if AMX, upcasting will make kernel slower
-      if rngs is not None and not AMX:
+    # skip hand-coded TC opts if AMX, upcasting will make kernel slower
+    if good_tc_opt and not AMX:
+      if rngs is not None:
        for tc_dim in [1,0]: # attempt to upcast M and N
          szs = [sz for sz in [5,4,3,2] if rngs[tc_dim].src[0].divides(sz) is not None]
          if szs:
@@ -62,8 +62,8 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:

  # should use matvec - TODO: adjust/tune based on the wide vs tall/large vs small mat
  MV_BLOCKSIZE, MV_THREADS_PER_ROW, MV_ROWS_PER_THREAD = getenv("MV_BLOCKSIZE", 4), getenv("MV_THREADS_PER_ROW", 8), getenv("MV_ROWS_PER_THREAD", 4)
-  if k.opts.has_local and getenv("MV",1) != 0 and (MV_BLOCKSIZE > 1 or MV_THREADS_PER_ROW > 1 or MV_ROWS_PER_THREAD > 1) and  \
-    k.reduceop is not None and k.reduceop.arg[0] is Ops.ADD and len(k.full_shape) >= 2 and k.opts.has_shared and \
+  if k.ren.has_local and getenv("MV",1) != 0 and (MV_BLOCKSIZE > 1 or MV_THREADS_PER_ROW > 1 or MV_ROWS_PER_THREAD > 1) and  \
+    k.reduceop is not None and k.reduceop.arg[0] is Ops.ADD and len(k.full_shape) >= 2 and k.ren.has_shared and \
    (mulop:=k.reduceop.src[0]).op is Ops.MUL and mulop.src[0].op is Ops.LOAD and mulop.src[1].op is Ops.LOAD:
    idx0, idx1 = mulop.src[0].src[0].src[1].get_idx(), mulop.src[1].src[0].src[1].get_idx()
    if k.ranges_of(AxisType.REDUCE):
@@ -103,7 +103,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
  for axis in to_upcast[::-1]: k.apply_opt(Opt(OptOps.UPCAST, axis, 0))

  # potentially do more upcasts of non reduce axes based on a heuristic
-  is_dsp = k.opts is not None and k.opts.device == "DSP"
+  is_dsp = k.ren is not None and k.ren.device == "DSP"
  upcasted_axis: set[int] = set()
  while resolve(prod(k.output_shape[i] for i in k.upcastable_dims) >= 1024):
    xb_choices = []
@@ -149,13 +149,12 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:

  # if nothing at all is upcasted and it's easy to, do an upcast
  for splits in [4]:
-    # TODO: somehow this never hits a reduce
    if not k.upcasted and k.upcastable_dims and k.full_shape[k.upcastable_dims[-1]] % splits == 0:
      k.apply_opt(Opt(OptOps.UPCAST, k.upcastable_dims[-1], splits))

  # **** local groups ****

-  if k.opts.has_local:
+  if k.ren.has_local:
    if NOLOCALS:
      k.apply_opt(Opt(OptOps.NOLOCALS))
    else:
@@ -176,13 +175,14 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:

  # **** threading ****

-  if k.opts.has_threads and k.opts.global_max is not None:
+  if k.ren.has_threads and k.ren.global_max is not None:
    for threads in [32,16,12,8,6,5,4,3,2]:
      # Skip if too many threads. Heuristic: use about 128K ops per thread
-      if threads > k.opts.global_max[0] or resolve(prod(k.full_shape) // (128 << 10) < threads): continue
+      if threads > k.ren.global_max[0] or resolve(prod(k.full_shape) // (128 << 10) < threads): continue
      for axis in k.axes_of(AxisType.LOOP):
        if k.full_shape[axis] % threads == 0:
-          k.apply_opt(Opt(OptOps.THREAD, axis, threads))
+          try: k.apply_opt(Opt(OptOps.THREAD, axis, threads))
+          except KernelOptError: pass
          break
      if k.applied_opts and k.applied_opts[-1].op is OptOps.THREAD: break

@@ -2,11 +2,11 @@ from __future__ import annotations
 import math, itertools
 from collections import defaultdict
 from typing import cast, Final
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, GroupOp
+from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, GroupOp, axis_letters, axis_colors
 from tinygrad.device import Buffer
 from tinygrad.dtype import dtypes, ImageDType
-from tinygrad.helpers import colored, BEAM, getenv, DEBUG, to_function_name, NOOPT, argsort, round_up, prod, merge_dicts, get_single_element
-from tinygrad.codegen.opt import axis_colors, Opt, OptOps, KernelOptError, check, axis_letters
+from tinygrad.helpers import colored, BEAM, getenv, DEBUG, to_function_name, NOOPT, argsort, round_up, prod, merge_dicts, get_single_element, flatten
+from tinygrad.codegen.opt import Opt, OptOps, KernelOptError, check
 from tinygrad.codegen.simplify import pm_flatten_range
 from tinygrad.renderer import Renderer

@@ -17,8 +17,8 @@ axis_to_pos = {AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisTy
               AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5}

 class Scheduler:
-  def __init__(self, ast:UOp, opts:Renderer):
-    self.ast, self.opts = ast, opts
+  def __init__(self, ast:UOp, ren:Renderer):
+    self.ast, self.ren = ast, ren
    self.dont_use_locals = self.ast.arg.dont_use_locals if self.ast.arg is not None else False
    self.applied_opts = list(self.ast.arg.applied_opts) if self.ast.arg is not None else []

@@ -46,7 +46,7 @@ class Scheduler:
  def shape_str_to_axis(self, nms:list[str]) -> tuple[int, ...]: return tuple([self.shape_str().index(x) for x in nms])

  def copy(self):
-    ret = Scheduler(self.ast, self.opts)
+    ret = Scheduler(self.ast, self.ren)
    ret.dont_use_locals = self.dont_use_locals
    ret.applied_opts = self.applied_opts[:]
    return ret
@@ -64,11 +64,10 @@ class Scheduler:
    return self.ast.replace(arg=KernelInfo(name=name, applied_opts=tuple(self.applied_opts), dont_use_locals=self.dont_use_locals), tag=1)

  def _globalizable_rngs(self) -> list[UOp]:
-    # all ranges that end before any STOREs
-    return [x for x in self.ast.toposort(lambda x: x.op is not Ops.STORE) if x.op is Ops.RANGE and x not in self.ast.ranges]
+    return flatten([list(UOp.sink(*s.src[1:]).ranges) for s in self.ast.src if s.op is Ops.END])

  def convert_loop_to_global(self):
-    if not self.opts.has_local: return None
+    if not self.ren.has_local: return None

    globalizible_rngs = self._globalizable_rngs()
    rng = [x.replace(arg=x.arg[0:-1]+(AxisType.GLOBAL,)) if x in globalizible_rngs else x for x in self.rngs]
@@ -76,11 +75,11 @@ class Scheduler:
    self.ast = self.ast.substitute(dict(zip(self.rngs, rng)))

  def colors(self) -> list[str]:
-    globalizible_rngs = self._globalizable_rngs()
+    output_rngs = self._globalizable_rngs()
    ret = []
    for x,r in zip(self.axis_types, self.rngs):
      if self.dont_use_locals and x == AxisType.GLOBAL: ret.append("BLUE")
-      elif r not in globalizible_rngs and x == AxisType.LOOP: ret.append("BLACK")
+      elif r not in output_rngs and x == AxisType.LOOP: ret.append("BLACK")
      else: ret.append(axis_colors[x])
    return ret
  def colored_shape(self) -> str: return ' '.join([colored(f'{x.src[0].render():>4s}', color) for x,color in zip(self.rngs, self.colors())])
@@ -122,7 +121,7 @@ class Scheduler:
      return

    if opt.op in {OptOps.LOCAL, OptOps.GROUP, OptOps.GROUPTOP}:
-      check(self.opts.has_local, "locals needed for opt")
+      check(self.ren.has_local, "locals needed for opt")

    rng = self.rngs[real_axis] if (real_axis:=self.real_axis(opt.op, opt.axis)) >= 0 else UOp(Ops.NOOP)

@@ -140,7 +139,7 @@ class Scheduler:
                                        (self.group_for_reduces and opt.op not in {OptOps.NOLOCALS, OptOps.PADTO})):
        upcast_local_sz = prod([self.full_shape[a] for a in self.axes_of(AxisType.UPCAST, AxisType.WARP, AxisType.LOCAL, AxisType.GROUP_REDUCE)])
        smem_sz = amt*upcast_local_sz*self.reduceop.dtype.itemsize
-        check(smem_sz <= self.opts.shared_max, f"exceeds maximum shared memory size: needs {smem_sz}, max {self.opts.shared_max}")
+        check(smem_sz <= self.ren.shared_max, f"exceeds maximum shared memory size: needs {smem_sz}, max {self.ren.shared_max}")
      if self.reduceop is not None and (opt.op in {OptOps.GROUP, OptOps.GROUPTOP}):
        # We currently dont support a group within another rudece, TODO: fix if-contexts
        reduce = [u for u in self.ast.backward_slice if u.op is Ops.REDUCE and rng in merge_dicts([r.ranges for r in u.src[1:]])][0]
@@ -151,14 +150,14 @@ class Scheduler:
        check(amt <= 32, "don't unroll more than 32")
        check(rng.arg[-1] in {AxisType.GROUP_REDUCE, AxisType.REDUCE}, "unroll is for GROUP_REDUCE/REDUCE")
      if opt.op is OptOps.UPCAST:
-        check((self.opts is not None and self.opts.device == "DSP") or amt <= 16, "don't upcast more than 16")
+        check((self.ren is not None and self.ren.device == "DSP") or amt <= 16, "don't upcast more than 16")
        check(rng.arg[-1] in {AxisType.GLOBAL, AxisType.LOCAL, AxisType.LOOP}, f"upcast is for GLOBAL/LOCAL/LOOP, not {rng.arg[-1]}")
      if opt.op is OptOps.LOCAL:
        check(not self.dont_use_locals, "can't use locals")
        check(rng.arg[-1] in {AxisType.GLOBAL, AxisType.LOOP}, "local is for globals")
      if opt.op is OptOps.THREAD:
-        check(self.opts is not None and self.opts.has_threads, "target does not support threads")
-        check(self.opts is not None and self.opts.global_max is not None and amt <= self.opts.global_max[0], "too many threads")
+        check(self.ren is not None and self.ren.has_threads, "target does not support threads")
+        check(self.ren is not None and self.ren.global_max is not None and amt <= self.ren.global_max[0], "too many threads")
        check(all(x is not AxisType.THREAD for x in self.axis_types), "already threaded")
        check(rng in self._globalizable_rngs(), "can't apply range to this dim")
      if opt.op in {OptOps.GROUP, OptOps.GROUPTOP}:
@@ -170,7 +169,7 @@ class Scheduler:
      check(len(self.applied_opts) == 0, "tensor core opts must be first") # TODO: remove the need for this by having warps
      check(opt.axis is not None, "tensor core opts must have an axis")
      check(opt.arg is not None and isinstance(opt.arg, tuple) and len(opt.arg) == 3, "tensor core opts must have valid arg")
-      check(-1 <= (tc_select:=cast(tuple, opt.arg)[0]) < len(self.opts.tensor_cores), "tensor core opts must have valid tc_select")
+      check(-1 <= (tc_select:=cast(tuple, opt.arg)[0]) < len(self.ren.tensor_cores), "tensor core opts must have valid tc_select")
      check(0 <= (tc_opt:=cast(tuple, opt.arg)[1]) <= 2, "tensor core opts must have valid tc_opt")
      check(0 < (use_tensor_cores:=cast(tuple, opt.arg)[2]) <= 2, "use_tensor_cores value is not valid")
      try: ret = self._apply_tc_opt(use_tensor_cores, cast(int, opt.axis), tc_select, tc_opt)
@@ -217,7 +216,7 @@ class Scheduler:
      if mul.op is not Ops.MUL: return None
      in0, in1 = mul.src
      try:
-        tensor_cores = self.opts.tensor_cores if tc_select == -1 else [self.opts.tensor_cores[tc_select]]
+        tensor_cores = self.ren.tensor_cores if tc_select == -1 else [self.ren.tensor_cores[tc_select]]
      except IndexError:
        raise KernelOptError(f"invalid tensor core choice {tc_select}")
      for tc in tensor_cores:
@@ -288,7 +287,7 @@ class Scheduler:
            # TODO: remove tc_upcast_axes from the arg
            # do the reduce_axes always disappear? i think they don't
            # they need to be moved into the WMMA srcs
-            wmma_arg = (str(tc), tc.dims, tc.dtype_in, tc.dtype_out, self.opts.device, tc.threads, tc_upcast_axes, ()) #, tc_reduce_axes)
+            wmma_arg = (str(tc), tc.dims, tc.dtype_in, tc.dtype_out, self.ren.device, tc.threads, tc_upcast_axes, ()) #, tc_reduce_axes)
            wmma = UOp(Ops.WMMA, dtype=tc.dtype_out.vec(tc.elements_per_thread[2]), src=(
              UOp(Ops.CONTRACT, dtype=srcs[0].dtype.vec(tc.elements_per_thread[0]), src=(srcs[0],), arg=tc_upcast_axes[0], tag=1),
              UOp(Ops.CONTRACT, dtype=srcs[1].dtype.vec(tc.elements_per_thread[1]), src=(srcs[1],), arg=tc_upcast_axes[1], tag=1),
@@ -322,15 +321,15 @@ def bufs_from_ast(ast:UOp, dname:str) -> list[Buffer]:
  glbls = sorted([x for x in ast.backward_slice if x.op is Ops.DEFINE_GLOBAL], key=lambda x: x.arg)
  return [Buffer(dname, x.ptrdtype.size, x.dtype.base if not isinstance(x.dtype, ImageDType) else x.dtype) for x in glbls]

-def apply_opts(ctx:Renderer, ast:UOp):
-  if ast.tag is not None: return None
-  k = Scheduler(ast, ctx)
+def apply_opts(ast:UOp, ren:Renderer) -> UOp:
+  if ast.tag is not None: return ast
+  k = Scheduler(ast, ren)
  k.convert_loop_to_global()
  if ast.arg is not None and ast.arg.opts_to_apply is not None:
    for opt in ast.arg.opts_to_apply: k.apply_opt(opt)
  elif BEAM >= 1:
    from tinygrad.codegen.opt.search import beam_search
-    rawbufs = bufs_from_ast(ast, ctx.device)
+    rawbufs = bufs_from_ast(ast, ren.device)
    k = beam_search(k, rawbufs, BEAM.value, bool(getenv("BEAM_ESTIMATE", 1)))
  elif not NOOPT and (ast.arg is None or ast.arg.applied_opts == ()):
    from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
@@ -338,7 +337,3 @@ def apply_opts(ctx:Renderer, ast:UOp):
    if not any(u.op is Ops.AFTER and u.src[0].op is Ops.DEFINE_LOCAL for u in ast.backward_slice):
      k = hand_coded_optimizations(k)
  return k.get_optimized_ast(name_override=ast.arg.name if ast.arg is not None and ast.arg.name != "test" else None)
-
-pm_postrange_opt = PatternMatcher([
-  (UPat(Ops.SINK, name="ast"), apply_opts),
-])
@@ -66,7 +66,7 @@ def _try_compile_linearized_w_idx(x:tuple[int,Scheduler], compiler:Compiler) ->
    signal.alarm(getenv("BEAM_TIMEOUT_SEC", 10))
  ret = None
  try:
-    p = get_program(x[1].copy().get_optimized_ast(name_override="test"), x[1].opts)
+    p = get_program(x[1].copy().get_optimized_ast(name_override="test"), x[1].ren)
    assert p.uops is not None, "uop list wasn't generated?"
    if len(p.uops) >= (uops_max:=getenv("BEAM_UOPS_MAX", 3000)) > 0:
      if getenv("BEAM_LOG_SURPASS_MAX"): print(f"too many uops. {len(p.uops)=}, {uops_max=}")
@@ -119,7 +119,7 @@ def get_kernel_actions(lin:Scheduler, include_0=True, candidates:list[Opt]|None=
 beam_pool, BEAM_DEBUG = None, getenv("BEAM_DEBUG")
 def beam_search(lin:Scheduler, rawbufs:list[Buffer], amt:int, allow_test_size=True, disable_cache=IGNORE_BEAM_CACHE.value):
  global beam_pool
-  key = {"ast": lin.ast.key, "amt": amt, "allow_test_size": allow_test_size, "device": lin.opts.device, "suffix": lin.opts.suffix}
+  key = {"ast": lin.ast.key, "amt": amt, "allow_test_size": allow_test_size, "device": lin.ren.device, "suffix": lin.ren.suffix}
  if not disable_cache and CACHELEVEL >= 1 and (val:=diskcache_get("beam_search", key)) is not None:
    ret = lin.copy()
    for o in val[len(lin.applied_opts):]: ret.apply_opt(o)
@@ -128,7 +128,7 @@ def beam_search(lin:Scheduler, rawbufs:list[Buffer], amt:int, allow_test_size=Tr
  beam: list[tuple[Scheduler, float]] = [(lin, float("inf"))]
  seen_libs = set()

-  default_parallel = multiprocessing.cpu_count() if lin.opts.device in {"CUDA", "AMD", "NV", "METAL", "HIP"} else 0
+  default_parallel = multiprocessing.cpu_count() if lin.ren.device in {"CUDA", "AMD", "NV", "METAL", "HIP"} else 0
  if beam_pool is None and (workers := getenv("PARALLEL", default_parallel)):
    beam_pool = multiprocessing.get_context("spawn").Pool(workers, _init_worker, (), getenv("BEAM_MAX_TASKS_PER_CHILD", 16))
    @atexit.register
@@ -137,14 +137,14 @@ def beam_search(lin:Scheduler, rawbufs:list[Buffer], amt:int, allow_test_size=Tr
  min_progress = getenv("BEAM_MIN_PROGRESS", 0.01)/1e6
  if BEAM_DEBUG:
    print("BEAM_SEARCH:")
-    print('\n'.join(pyrender(lin.ast.replace(arg=None))))
+    print(pyrender(lin.ast.replace(arg=None)))
  if DEBUG >= 2: print(f"   0.00s:                from   1 ->   1 actions {lin.colored_shape()}")

  try:
    rawbufs = _ensure_buffer_alloc(rawbufs)
    var_vals: dict[str, int] = {k.expr:int(k.vmax+k.vmin)//2 for k in lin.ast.variables()}
    exiting, st = False, time.perf_counter()
-    dev = Device[lin.opts.device]
+    dev = Device[lin.ren.device]
    while not exiting:
      acted_lins: list[Scheduler] = flatten([get_kernel_actions(lin, include_0=False).values() for lin,_ in beam])
      timed_lins: list[tuple[Scheduler, float]] = []
@@ -80,6 +80,10 @@ cuda_81616 = [TensorCore(dims=(8,16,16), threads=32, elements_per_thread=(8,4,4)
  swizzle=((('r1', 'r2', 'l2', 'l3', 'l4'), ('u1', 'r3'), ('l0', 'l1', 'u0', 'r0')),
           (('r1', 'r2', 'u0', 'l0', 'l1'), ('r0', 'r3'), ('l2', 'l3', 'l4', 'u1'))))
  for di,do in [(dtypes.half,dtypes.float), (dtypes.bfloat16,dtypes.float), (dtypes.half,dtypes.half)]]
+cuda_81632_f8 = [TensorCore(dims=(8,16,32), threads=32, elements_per_thread=(16,8,4), dtype_in=di, dtype_out=do, opts=cuda_tc_opts,
+  swizzle=((('r2', 'r3', 'l2', 'l3', 'l4'), ('u1', 'r4'), ('l0', 'l1', 'u0', 'r0', 'r1')),
+           (('r2', 'r3', 'u0', 'l0', 'l1'), ('r1', 'r4'), ('l2', 'l3', 'l4', 'u1', 'r0'))))
+  for di,do in [(dtypes.fp8e4m3,dtypes.float),(dtypes.fp8e5m2,dtypes.float)]]
 cuda_8168_f16 = [TensorCore(dims=(8,16,8), threads=32, elements_per_thread=(4,2,4), dtype_in=di, dtype_out=do, opts=cuda_tc_opts,
  swizzle=((('r1', 'r2', 'l2', 'l3', 'l4'), ('r0', 'u1'), ('l0', 'l1', 'u0')),
           (('r1', 'r2', 'u0', 'l0', 'l1'), ('u1', 'r0'), ('l2', 'l3', 'l4'))))
@@ -87,9 +91,10 @@ cuda_8168_f16 = [TensorCore(dims=(8,16,8), threads=32, elements_per_thread=(4,2,
 cuda_8168_tf32 = [TensorCore(dims=(8,16,8), threads=32, elements_per_thread=(4,2,4), dtype_in=dtypes.float, dtype_out=dtypes.float, opts=cuda_tc_opts,
  swizzle=((('r0', 'r1', 'l2', 'l3', 'l4'), ('u1', 'r2'), ('l0', 'l1', 'u0')),
           (('r0', 'r1', 'u0', 'l0', 'l1'), ('u1', 'r2'), ('l2', 'l3', 'l4'))))]
+cuda_sm75: list[TensorCore] = cuda_8168_f16
 cuda_sm80: list[TensorCore] = cuda_81616 + cuda_8168_f16
 if getenv("ALLOW_TF32", 0): cuda_sm80 += cuda_8168_tf32
-cuda_sm75: list[TensorCore] = cuda_8168_f16
+cuda_sm89: list[TensorCore] = cuda_sm80 + cuda_81632_f8

 # ***** AMD *****

@@ -112,6 +117,14 @@ amd_cdna = [TensorCore(dims=(16,16,16), threads=64, elements_per_thread=(4,4,4),
           (('l0', 'l1', 'l2', 'l3', 'r2', 'r3'), ('r0', 'r1'), ('l4', 'l5', 'u0', 'u1'))))
  for di,do in [(dtypes.half,dtypes.float),(dtypes.bfloat16,dtypes.float)]]

+amd_cdna_161632 = [TensorCore(dims=(16,16,32), threads=64, elements_per_thread=(8,8,4), dtype_in=di, dtype_out=do,
+  opts=("l0","l0","l0","l0","u1","u1","l1","l1"),
+  swizzle=((('u0','u1','l4','l5','r3','r4'), ('r0','r1'), ('l0','l1','l2','l3','r2')),
+           (('l0','l1','l2','l3','r3','r4'), ('r0','r1'), ('l4','l5','u0','u1','r2'))))
+  for di,do in [(dtypes.half,dtypes.float),(dtypes.bfloat16,dtypes.float)]]
+
+amd_cdna4 = amd_cdna_161632 + amd_cdna
+
 # ***** Apple Metal *****

 metal = [TensorCore(dims=(8,8,8), threads=32, elements_per_thread=(2,2,2), dtype_in=di, dtype_out=do,
@@ -1,6 +1,6 @@
 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, graph_rewrite, _substitute, range_start, ImageDType
 from tinygrad.uop.symbolic import symbolic_flat
-from tinygrad.helpers import partition
+from tinygrad.helpers import partition, dedup
 from tinygrad.dtype import dtypes

 def flatten_range(r:UOp):
@@ -12,9 +12,7 @@ def flatten_range(r:UOp):

 pm_flatten_range = PatternMatcher([
  # real ranges only
-  (UPat((Ops.REDUCE, Ops.STORE), name="r"), flatten_range),
-  # END is only on RANGES. TODO: this is copied from symbolic
-  (UPat(Ops.END, name="e"), lambda e: UOp.end(*e.src[e.arg:], ends=sorted(UOp.sink(*e.src[:e.arg]).ranges, key=lambda x: x.arg))),
+  (UPat((Ops.REDUCE, Ops.STORE, Ops.END), name="r"), flatten_range),
 ])

 def count_divmod(x:UOp): return len([u for u in x.toposort() if u.op in {Ops.IDIV, Ops.MOD}])
@@ -92,10 +90,7 @@ pm_reduce_collapse = pm_reduce_unparented + PatternMatcher([
  # lift x+y out of reduce on lt
  ((UPat.var("x")+UPat.var("y")).or_casted() < UPat.var("c"), lambda x,y,c: (x < (c.cast(y.dtype)-y)) if no_range(y) and no_range(c) else None),
  # lift x*y out of reduce
-  ((UPat.var("x")*UPat.var("y")) < UPat.var("c"),
-   lambda x,y,c: (x < ((c+y-1) // y)) if no_range(y) and no_range(c) and y.vmin > 0 else None),
-  # lift x+y out of reduce on ne
-  ((UPat.var("x")+UPat.var("y")).or_casted() != UPat.var("c"), lambda x,y,c: (x != (c.cast(y.dtype)-y)) if no_range(y) and no_range(c) else None),
+  ((UPat.var("x")*UPat.var("y")) < UPat.var("c"), lambda x,y,c: (x < ((c+y-1) // y)) if no_range(y) and no_range(c) and y.vmin > 0 else None),
  # fold the range
  ((UPat(Ops.RANGE, name="r") < UPat.var("cut")).where(0, UPat.cvar("val")).reduce(UPat.var("r"), arg=Ops.ADD),
   lambda r,cut,val: (r.src[0]-cut).maximum(0).minimum(r.src[0]).cast(val.dtype) * val),
@@ -106,29 +101,51 @@ pm_reduce_collapse = pm_reduce_unparented + PatternMatcher([
  # REDUCE on ADD
  ((UPat.var("x")+UPat.var("y")).reduce(arg=Ops.ADD, allow_any_len=True, name="r"),
   lambda x,y,r: x.reduce(*r.src[1:], arg=Ops.ADD) + y.reduce(*r.src[1:],arg=Ops.ADD)),
+])+symbolic_flat
+
+pm_reduce_load_collapse = PatternMatcher([
  # MUL casted bool
  ((UPat.var("x") * UPat.var("gate", dtype=dtypes.bool).cast()), lambda x,gate: gate.where(x, 0)),
+  # lift x+y out of reduce on ne
+  ((UPat.var("x")+UPat.var("y")).or_casted() != UPat.var("c"), lambda x,y,c: (x != (c.cast(y.dtype)-y)) if no_range(y) and no_range(c) else None),
  # reduce on gated load becomes can substitute the range and remove the reduce
  ((UPat.var("idx")!=(UPat(Ops.RANGE, name="r").or_casted())).where(0, UPat.var("expr")).reduce(UPat.var("r"), arg=Ops.ADD),
   lambda r,idx,expr: (v:=(idx.cast(r.dtype) >= 0) & (idx.cast(r.dtype) < r.src[0])).where(expr.substitute({r:idx.cast(r.dtype).valid(v)}),0)),
-  # AND on WHERE
-  ((UPat(Ops.DEFINE_VAR, name="x") & UPat.var("y")).where(UPat.cvar("c"), 0).reduce(arg=Ops.ADD, allow_any_len=True, name="r"),
-    lambda x,y,c,r: y.where(c, 0).reduce(*r.src[1:], arg=Ops.ADD)*x.cast(c.dtype)),
 ])+symbolic_flat

-def reduce_collapse(red:UOp):
-  included, not_included = partition(red.backward_slice, lambda x: any(y in x.backward_slice_with_self for y in red.src[1:]))
+def reduce_collapse(red:UOp, pm=pm_reduce_collapse):
+  included = red.src[0].toposort(gate=lambda x: any(y in x.ranges for y in red.src[1:]))
  if any(x.op in {Ops.STORE, Ops.REDUCE} for x in included): return None
  replaces: dict[UOp, UOp] = {}
  for u in included:
    for s in u.src:
-      if s in not_included and s not in replaces and s.op not in {Ops.CONST, Ops.VCONST, Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_VAR}:
-        replaces[s] = UOp(Ops.DEFINE_VAR, dtype=s.dtype, arg=(f'in{len(replaces)}', s.vmin, s.vmax))
+      if s in included or s in replaces or s.op in {Ops.CONST, Ops.VCONST, Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_VAR}: continue
+      replaces[s] = UOp(Ops.DEFINE_VAR, dtype=s.dtype, arg=(f'in{len(replaces)}', s.vmin, s.vmax))
  collapse_fxn = red.substitute(replaces)
-  sink = graph_rewrite(collapse_fxn, pm_reduce_collapse, name="reduce_collapse")
+  sink = graph_rewrite(collapse_fxn, pm, name="reduce_collapse")
  return sink.substitute({v:k for k,v in replaces.items()}) if no_range(sink) else None

-pm_reduce_simplify = pm_reduce_unparented + PatternMatcher([
-  # remove REDUCE without loads (generic arange opt / indexing). TODO: support multi range
-  (UPat(Ops.REDUCE, src=(UPat(), UPat()), name="red"), reduce_collapse),
+def reduce_load_collapse(red:UOp): return reduce_collapse(red, pm=pm_reduce_load_collapse)
+
+# remove REDUCE without loads (generic arange opt / indexing). TODO: support multi range
+pm_reduce_simplify = pm_reduce_unparented + PatternMatcher([(UPat(Ops.REDUCE, src=(UPat(), UPat()), name="red"), reduce_collapse),])
+# remove REDUCE on load, comes from indexing a tensor with another tensor
+def no_load(u:UOp) -> bool: return not any(x.op is Ops.LOAD for x in u.backward_slice_with_self)
+pm_load_collapse = PatternMatcher([
+  (UPat(Ops.REDUCE, src=(UPat(), UPat()), name="red"), reduce_load_collapse),
+  # we want to make sure we dont do math on a loaded index since that can cause overflow, this undoes the rule in pm_reduce_load_collapse
+  ((UPat.var("x", dtypes.index)+UPat.var("y"))<UPat.var("c"), lambda x,y,c: x < c-y if no_load(y) and no_load(c) and not no_load(x) else None),
+])
+
+def cut_store_range(ctx, store:UOp, r:UOp):
+  # only cut ranges on CPU for now
+  if r.src[0].op is not Ops.CONST or ctx!="CPU": return None
+  if not (cuts:=[c.src[1].arg for c in store.get_consumer_map()[r] if c.op is Ops.CMPLT and r is c.src[0] and c.src[1].op is Ops.CONST]): return None
+  cuts = sorted(dedup([0] + cuts + [r.src[0].arg]))
+  ranges = [UOp.range((end-start), *(r.arg[0:-1]+(i,r.arg[-1]))) for i,(start,end) in enumerate(zip(cuts[:-1], cuts[1:]))]
+
+  return UOp.group(*[store.substitute({r: new_r+start}).end(new_r) for new_r, start in zip(ranges, cuts[:-1])])
+
+pm_split_store = pm_flatten_range+PatternMatcher([
+  (UPat(Ops.END, src=(UPat(Ops.STORE, name="store"), UPat.var("r"))), cut_store_range),
 ])
@@ -331,7 +331,9 @@ def is_dtype_supported(dtype:DType, device:str|None=None) -> bool:
    if device in {"CUDA", "NV"}: return not CI and not getenv(f"{device}_PTX") and not getenv("NV_NAK")
    if device in {"CPU"}: return not CI and platform.machine() in {"arm", "arm64", "aarch64", "x86_64", "amd64"} and not getenv("CPU_LVP")
    return device in {"AMD", "PYTHON", "NULL"}
-  if dtype in dtypes.fp8s: return device in {"PYTHON", "NULL"}
+  if dtype in dtypes.fp8s:
+    if device in {"CUDA", "NV"}: return not CI and not getenv(f"{device}_PTX") and not getenv("NV_NAK")
+    return device in {"PYTHON", "NULL"}
  if device == "WEBGPU": return dtype in [dtypes.bool, dtypes.char, dtypes.uchar, dtypes.short,
                                          dtypes.ushort, dtypes.float, dtypes.int32, dtypes.uint32, dtypes.half]
  # for CI GPU and OSX, cl_khr_fp16 isn't supported
@@ -7,13 +7,15 @@ from enum import Enum, auto

 class InvalidTypeMetaClass(type):
  instance:None|InvalidType = None
-  def __call__(cls, *args, **kwargs):
+  def __call__(cls):
    if (ret:=InvalidTypeMetaClass.instance) is not None: return ret
    InvalidTypeMetaClass.instance = ret = super().__call__()
    return ret

 class InvalidType(metaclass=InvalidTypeMetaClass):
  def __eq__(self, other): return self is other
+  def __lt__(self, other): return self is not other
+  def __gt__(self, other): return self is not other
  def __hash__(self): return id(self)
  def __repr__(self): return "Invalid"
  def __reduce__(self): return (InvalidType, ())  # Return the global Invalid instance
@@ -47,7 +49,7 @@ class DType(metaclass=DTypeMetaClass):
  @staticmethod
  def new(priority:int, itemsize:int, name:str, fmt:FmtStr|None): return DType(priority, itemsize, name, fmt, 1, None)
  def __reduce__(self): return type(self), tuple(getattr(self, f.name) for f in fields(self))
-  def __repr__(self): return f"dtypes.{INVERSE_DTYPES_DICT[self.scalar().name]}"+(f".vec({self.count})" if self.count > 1 else "")
+  def __repr__(self): return f"dtypes.{INVERSE_DTYPES_DICT[self.scalar().name]}"+(f".vec({self.count})" if self.count != 1 else "")
  def __lt__(self, o:DType): return (self.priority, self.itemsize, self.name, self.fmt, self.count) < (o.priority, o.itemsize, o.name, o.fmt, o.count)
  @property
  def base(self): return self
@@ -61,7 +63,7 @@ class DType(metaclass=DTypeMetaClass):
  def ptr(self, size=-1, addrspace=AddrSpace.GLOBAL) -> PtrDType:
    return PtrDType(self.priority, self.itemsize, self.name, self.fmt, self.count, None, self, addrspace, 1, size)
  def scalar(self) -> DType: return self._scalar if self._scalar is not None else self
-  def nbytes(self): raise RuntimeError("only ptr types have nbytes")
+  def nbytes(self) -> int: raise RuntimeError("only ptr types have nbytes")
  @property
  def min(self): return dtypes.min(self)
  @property
@@ -82,7 +84,7 @@ class PtrDType(DType):
    if isinstance(self, ImageDType):
      return ImageDType(self.priority, self.itemsize, self.name, self.fmt, self.count, self, self._base, self.addrspace, sz, self.size, self.shape)
    return type(self)(self.priority, self.itemsize, self.name, self.fmt, self.count, self, self._base, self.addrspace, sz, self.size)
-  def ptr(self, size=-1, addrspace=AddrSpace.GLOBAL): raise RuntimeError("can't make a pointer from a pointer")
+  def ptr(self, size=-1, addrspace=AddrSpace.GLOBAL) -> PtrDType: raise RuntimeError("can't make a pointer from a pointer")
  def nbytes(self) -> int:
    if self.size == -1: raise RuntimeError("can't get nbytes of a pointer with unlimited size")
    return self.size*self.itemsize
@@ -26,7 +26,7 @@ def get_program(ast:UOp, renderer:Renderer|None=None, opts:list[Opt]|None=None)
  """

  if getenv("VIZ"): graph_rewrite(ast, PatternMatcher([]), name="View Base AST")
-  if DEBUG >= 5: print('\n'.join(pyrender(ast)))
+  if DEBUG >= 5: print(pyrender(ast))

  # linearize
  if renderer is None: renderer = Device.default.renderer
@@ -38,7 +38,7 @@ def get_program(ast:UOp, renderer:Renderer|None=None, opts:list[Opt]|None=None)
  except RuntimeError as e:
    print("***** LINEARIZE FAILURE *****")
    print(e)
-    print('\n'.join(pyrender(ast)))
+    print(pyrender(ast))
    raise
  assert uops[-1].op is Ops.SINK, "last uop must be sink"

@@ -15,7 +15,7 @@ def reduce_gradient(ctx:UOp, ret:UOp):
 # ctx is grad_output
 pm_gradient = PatternMatcher([
  (UPat(Ops.CAST, name="ret"), lambda ctx, ret: (ctx.cast(ret.src[0].dtype),)),
-  (UPat(Ops.RECIP, name="ret"), lambda ctx, ret: (-ctx * ret * ret,)),
+  (UPat(Ops.RECIPROCAL, name="ret"), lambda ctx, ret: (-ctx * ret * ret,)),
  (UPat(Ops.SIN, name="ret"), lambda ctx, ret: ((math.pi/2 - ret.src[0]).sin() * ctx,)),
  (UPat(Ops.LOG2, name="ret"), lambda ctx, ret: (ctx / (ret.src[0] * math.log(2)),)),
  (UPat(Ops.EXP2, name="ret"), lambda ctx, ret: (ret * ctx * math.log(2),)),
@@ -24,7 +24,7 @@ pm_gradient = PatternMatcher([
  (UPat(Ops.ADD), lambda ctx: (ctx, ctx)),
  (UPat(Ops.POW, name="ret", src=(UPat.var("b"), UPat.var("e"))), lambda ctx, ret, b, e:
    (ctx * (b.eq(0)&e.eq(0)).where(e, e*b.pow(e-1)), ctx * b.eq(0).where((e<0).where(ret.const_like(-math.inf), 0), ret*b.log2()*math.log(2.0)))),
-  (UPat(Ops.MAX, name="ret", src=(UPat.var("x"), UPat.var("y"))), lambda ctx, ret, x, y:
+  (UPat(Ops.MAX, src=(UPat.var("x"), UPat.var("y"))), lambda ctx, x, y:
    ((x>y).where(ctx, (x.eq(y)).where(ctx * 0.5, 0)), (x<y).where(ctx, (x.eq(y)).where(ctx * 0.5, 0)))),
  (UPat(Ops.MUL, name="ret"), lambda ctx, ret: (ret.src[1]*ctx, ret.src[0]*ctx)),
  (UPat(Ops.WHERE, name="ret"), lambda ctx, ret: (None, ret.src[0].where(ctx, ctx.const_like(0)), ret.src[0].where(ctx.const_like(0), ctx))),
@@ -2,7 +2,7 @@ from __future__ import annotations
 import os, functools, platform, time, re, contextlib, operator, hashlib, pickle, sqlite3, tempfile, pathlib, string, ctypes, sys, gzip, getpass
 import urllib.request, subprocess, shutil, math, types, copyreg, inspect, importlib, decimal, itertools
 from dataclasses import dataclass, field
-from typing import ClassVar, Iterable, Any, TypeVar, Callable, Sequence, TypeGuard, Iterator, Generic, Generator, cast
+from typing import ClassVar, Iterable, Any, TypeVar, Callable, Sequence, TypeGuard, Iterator, Generic, Generator, cast, overload

 T = TypeVar("T")
 U = TypeVar("U")
@@ -85,6 +85,9 @@ def word_wrap(x, wrap=80):
  while len(ansistrip(x[:i])) < wrap and i < len(x): i += 1
  return x[:i] + "\n" + word_wrap(x[i:], wrap)
 def pad_bytes(b:bytes, align:int) -> bytes: return b + b'\x00' * ((align - (len(b) % align)) % align)
+def panic(e:Exception|None=None):
+  if e is None: raise RuntimeError("PANIC!")
+  raise e

@functools.cache
 def canonicalize_strides(shape:tuple[T, ...], strides:tuple[T, ...]) -> tuple[T, ...]:
@@ -123,8 +126,13 @@ def polyN(x:T, p:list[float]) -> T: return functools.reduce(lambda acc,c: acc*x+

@functools.cache
 def to_function_name(s:str): return ''.join([c if c in (string.ascii_letters+string.digits+'_') else f'{ord(c):02X}' for c in ansistrip(s)])
+@overload
+def getenv(key:str) -> int: ...
+@overload
+def getenv(key:str, default:T) -> T: ...
@functools.cache
-def getenv(key:str, default=0): return type(default)(os.getenv(key, default))
+def getenv(key:str, default:Any=0): return type(default)(os.getenv(key, default))
+
 def temp(x:str, append_user:bool=False) -> str:
  return (pathlib.Path(tempfile.gettempdir()) / (f"{x}.{getpass.getuser()}" if append_user else x)).as_posix()

@@ -157,7 +165,7 @@ TRANSCENDENTAL, NOLOCALS = ContextVar("TRANSCENDENTAL", 1), ContextVar("NOLOCALS
 SPLIT_REDUCEOP, NO_MEMORY_PLANNER, RING = ContextVar("SPLIT_REDUCEOP", 1), ContextVar("NO_MEMORY_PLANNER", 0), ContextVar("RING", 1)
 PICKLE_BUFFERS, LRU = ContextVar("PICKLE_BUFFERS", 1), ContextVar("LRU", 1)
 CACHELEVEL, IGNORE_BEAM_CACHE, DEVECTORIZE = ContextVar("CACHELEVEL", 2), ContextVar("IGNORE_BEAM_CACHE", 0), ContextVar("DEVECTORIZE", 1)
-DISABLE_COMPILER_CACHE, BLOCK_REORDER = ContextVar("DISABLE_COMPILER_CACHE", 0), ContextVar("BLOCK_REORDER", 1)
+DISABLE_COMPILER_CACHE = ContextVar("DISABLE_COMPILER_CACHE", 0)
 QUANTIZE, VALIDATE_WITH_CPU, DISABLE_FAST_IDIV = ContextVar("QUANTIZE", 0), ContextVar("VALIDATE_WITH_CPU", 0), ContextVar("DISABLE_FAST_IDIV", 0)
 CORRECT_DIVMOD_FOLDING, FUSE_OPTIM = ContextVar("CORRECT_DIVMOD_FOLDING", 0), ContextVar("FUSE_OPTIM", 0)
 ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE = ContextVar("ALLOW_DEVICE_USAGE", 1), ContextVar("MAX_BUFFER_SIZE", 0)
@@ -166,10 +174,11 @@ EMULATE = ContextVar("EMULATE", "")
 CPU_COUNT = ContextVar("CPU_COUNT", max(1, len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1)))
 CPU_LLVM, CPU_LVP, AMD_LLVM = ContextVar("CPU_LLVM", 0), ContextVar("CPU_LVP", 0), ContextVar("AMD_LLVM", 1)
 VIZ = PROFILE = ContextVar("VIZ", 0)
-SPEC = ContextVar("SPEC", 0)
+SPEC = ContextVar("SPEC", 1)
 # TODO: disable by default due to speed
 IGNORE_OOB = ContextVar("IGNORE_OOB", 1)
 PCONTIG = ContextVar("PCONTIG", 0)  # partial contiguous in rangeify
+DEBUG_RANGEIFY = ContextVar("DEBUG_RANGEIFY", 0)

@dataclass(frozen=True)
 class Metadata:
@@ -177,7 +186,6 @@ class Metadata:
  caller: str
  backward: bool = False
  def __hash__(self): return hash(self.name)
-  def __repr__(self): return str(self) + (f" - {self.caller}" if self.caller else "")
  def __str__(self): return self.name + (" bw" if self.backward else "")

 # **************** global state Counters ****************
@@ -5,7 +5,7 @@ from io import BufferedReader
 from tinygrad.nn.state import TensorIO
 from tinygrad.tensor import Tensor, _broadcast_shape, ReductionStr
 from tinygrad.helpers import getenv, DEBUG, all_same, prod, flatten, make_tuple, argsort, is_numpy_ndarray, get_single_element, polyN
-from tinygrad.dtype import DType, ConstType, dtypes, _from_np_dtype, truncate
+from tinygrad.dtype import DType, ConstType, dtypes, _from_np_dtype, truncate, least_upper_dtype
 from tinygrad.device import is_dtype_supported, Device

 # ***** protobuf definitions ******
@@ -670,7 +670,8 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
  def ReduceL1(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
    return ReduceSum(data.abs(), axes, keepdims, noop_with_empty_axes)
  def ReduceL2(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
-    return ReduceSumSquare(data, axes, keepdims, noop_with_empty_axes).sqrt()
+    dtype = dtypes.float if data.dtype in (dtypes.float16, dtypes.bfloat16) else data.dtype
+    return ReduceSum(data.cast(dtype).square(), axes, keepdims, noop_with_empty_axes).sqrt().cast(data.dtype)
  def ReduceLogSum(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
    return ReduceSum(data, axes, keepdims, noop_with_empty_axes).log()
  def ReduceLogSumExp(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
@@ -897,7 +898,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
  def BatchNormalization(X:Tensor, scale:Tensor, B:Tensor, input_mean:Tensor, input_var:Tensor, epsilon:float=1e-05, momentum:float=0.9,
                        training_mode:int=0, spatial=1, is_test=0):
    if training_mode:
-      x_detached = X.detach()
+      x_detached = X.detach().cast(least_upper_dtype(X.dtype, dtypes.float32))
      current_mean = x_detached.mean(axis=(0,2,3))
      y = (x_detached - current_mean.reshape(shape=[1, -1, 1, 1]))
      current_var = (y*y).mean(axis=(0,2,3))
@@ -906,18 +907,20 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
      running_mean = input_mean * momentum + current_mean * (1 - momentum)
      running_var = input_var * momentum + current_var * (1 - momentum)

-      return X.batchnorm(scale, B, current_mean, current_invstd), running_mean, running_var
+      return X.batchnorm(scale, B, current_mean, current_invstd).cast(X.dtype),running_mean.cast(input_mean.dtype),running_var.cast(input_var.dtype)
    return X.batchnorm(scale, B, input_mean, (input_var + epsilon).rsqrt())
-  def GroupNormalization(x:Tensor, scale:Tensor, bias:Tensor, num_groups:int, epsilon:float=1e-05):
-    x = x.reshape(x.shape[0], num_groups, -1).layernorm(eps=epsilon).reshape(x.shape)
+  def GroupNormalization(x:Tensor, scale:Tensor, bias:Tensor, num_groups:int, epsilon:float=1e-05, stash_type:int=1):
+    assert stash_type == 1, "only float32 is supported"
+    x = x.reshape(x.shape[0], num_groups, -1).cast(dtypes.float).layernorm(eps=epsilon).cast(x.dtype).reshape(x.shape)
    return x * scale.reshape(1, -1, *[1] * (x.ndim-2)) + bias.reshape(1, -1, *[1] * (x.ndim-2))
  def InstanceNormalization(x:Tensor, scale:Tensor, bias:Tensor, epsilon:float=1e-05):
    return GroupNormalization(x, scale, bias, num_groups=cast(int, x.shape[1]), epsilon=epsilon)
  def LayerNormalization(x:Tensor, scale:Tensor, bias:Tensor, axis:int=-1, epsilon:float=1e-05, stash_type:int=1):
    assert stash_type == 1, "only float32 is supported"
    axes = tuple(i for i in range(axis if axis >= 0 else x.ndim + axis, x.ndim))
-    mean = x.mean(axis=axes, keepdim=True)
-    return x.layernorm(axes, epsilon).mul(scale).add(bias), mean, (x.sub(mean)).square().mean(axis=axes, keepdim=True).add(epsilon).rsqrt()
+    mean = (x32:=x.cast(dtypes.float)).mean(axis=axes, keepdim=True)
+    inv_std_dev = (x32.sub(mean)).square().mean(axis=axes, keepdim=True).add(epsilon).rsqrt()
+    return (x32.sub(mean)*inv_std_dev).cast(x.dtype).mul(scale).add(bias), mean, inv_std_dev
  def SkipLayerNormalization(x:Tensor, skip:Tensor, gamma:Tensor, beta:Tensor|None=None, bias:Tensor|None=None, epsilon:float=1e-12):
    x = x + skip
    if bias is not None: x = x + bias
@@ -1089,9 +1092,10 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
    return output, present_key, present_value, qk_matmul_return_val
  Attention = {OpSetId(Domain.ONNX, 1): attention_onnx, OpSetId(Domain.MICROSOFT_CONTRIB_OPS, 1): attention_contrib}

-  def RMSNormalization(X:Tensor, scale:Tensor, axis:int=-1, epsilon:float=1e-5):
-    norm = X.square().mean(axis=tuple(range(axis + X.ndim if axis < 0 else axis, X.ndim)), keepdim=True).add(epsilon).rsqrt()
-    return X * norm * scale
+  def RMSNormalization(X:Tensor, scale:Tensor, axis:int=-1, epsilon:float=1e-5, stash_type:int=1):
+    assert stash_type == 1, "only float32 is supported"
+    norm = X.cast(dtypes.float).square().mean(axis=tuple(range(axis + X.ndim if axis < 0 else axis, X.ndim)), keepdim=True).add(epsilon).rsqrt()
+    return X.cast(X.dtype) * norm * scale

  def RotaryEmbedding(X:Tensor, cos_cache:Tensor, sin_cache:Tensor, position_ids:Tensor|None=None, interleaved:int=0, num_heads:int|None=None,
                      rotary_embedding_dim:int=0):
@@ -80,7 +80,7 @@ def SGD(params: list[Tensor], lr=0.001, momentum=0.0, weight_decay=0.0, nesterov
  return LARS(params, lr, momentum, weight_decay, 0, None, nesterov, classic=classic, pre_wd=True, tcoef=0.0, fused=fused)

 # Muon applies the newton schulz algorithm on gradient. also can include momentum, nesterov, and weight decay
-def Muon(params: list[Tensor], lr=0.02, momentum=0.95, weight_decay=0.0, ns_steps=5, ns_params=(3.4445, -4.775, 2.0315),
+def Muon(params: list[Tensor], lr=0.001, momentum=0.95, weight_decay=0.1, ns_steps=5, ns_coefficients=(3.4445, -4.775, 2.0315),
         nesterov=True, fused=FUSE_OPTIM):
  """
  SGD with newton-schulz iteration and post momentum weight decay.
@@ -89,7 +89,7 @@ def Muon(params: list[Tensor], lr=0.02, momentum=0.95, weight_decay=0.0, ns_step
  - Paper: https://arxiv.org/pdf/2502.16982
  """
  assert not fused, "FUSE_OPTIM not allowed for Muon optimizer"
-  return LARS(params, lr, momentum, weight_decay, ns_steps, ns_params, nesterov, classic=False, pre_wd=False, tcoef=0.0, fused=fused)
+  return LARS(params, lr, momentum, weight_decay, ns_steps, ns_coefficients, nesterov, classic=False, pre_wd=False, tcoef=0.0, fused=fused)

 class LARS(Optimizer):
  """
@@ -97,10 +97,10 @@ class LARS(Optimizer):

  - Paper: https://arxiv.org/abs/1708.03888v3
  """
-  def __init__(self, params:list[Tensor], lr=0.001, momentum=0.9, weight_decay=1e-4, ns_steps=0, ns_params=None,
+  def __init__(self, params:list[Tensor], lr=0.001, momentum=0.9, weight_decay=1e-4, ns_steps=0, ns_coefficients=None,
               nesterov=False, classic=True, pre_wd=True, tcoef=0.001, fused=FUSE_OPTIM):
    super().__init__(params, lr, fused)
-    self.momentum, self.wd, self.ns_steps, self.ns_params  = momentum, weight_decay, ns_steps, ns_params
+    self.momentum, self.wd, self.ns_steps, self.ns_coefficients  = momentum, weight_decay, ns_steps, ns_coefficients
    self.nesterov, self.classic, self.pre_wd, self.tcoef = nesterov, classic, pre_wd, tcoef
    self.b = self._new_optim_param() if self.momentum else []

@@ -118,7 +118,7 @@ class LARS(Optimizer):
      if self.momentum:
        self.b[i].assign(self.momentum * self.b[i] + g)  # NOTE: self.b[i] is zero on the first run, no if required
        g = (g + self.momentum * self.b[i]) if self.nesterov else self.b[i]
-      if self.ns_params: g = g.reshape(g.shape[0], -1).newton_schulz(self.ns_steps, self.ns_params).reshape(g.shape)
+      if self.ns_coefficients: g = g.reshape(g.shape[0], -1).newton_schulz(self.ns_steps, self.ns_coefficients).reshape(g.shape)
      # muon does post momentum weight decay
      if not self.pre_wd and self.wd > 0: t = t.detach() * (1.0 - self.wd * self.lr)
      # popular momentum does pre learning rate update
@@ -37,6 +37,8 @@ class Estimates:
          if len(u.src) > 2: dont_count = dont_count.union(u.src[2].toposort())
        elif u.op is Ops.IF:
          dont_count = dont_count.union(u.src[0].toposort())
+        elif u.op is Ops.RANGE:
+          dont_count = dont_count.union(u.src[0].toposort())
    for u in uops:
      if u.op in {Ops.LOAD, Ops.STORE}:
        buf = u
@@ -45,18 +47,18 @@ class Estimates:
          mem[(buf, u.op)] = buf.ptrdtype.size * buf.dtype.itemsize
      if u.op is Ops.RANGE:
        mult_stack.append(mults)
-        mults *= cast(sint, u.src[0].ssimplify())
+        mults = cast(sint, (mults*u.src[0]).ssimplify())
        # SPECIAL are already counted in mults
        mults = mults.substitute({x:x.const_like(0) for x in mults.toposort() if x.op is Ops.SPECIAL}) if isinstance(mults, UOp) else mults
      elif u.op is Ops.END: mults = mult_stack.pop(-1)
-      elif u.op is Ops.SPECIAL: mults *= cast(sint, u.src[0].ssimplify()) # NOTE: we don't push to the mult_stack here, you can't end these
+      elif u.op is Ops.SPECIAL: mults = cast(sint, (mults*u.src[0]).ssimplify()) # NOTE: we don't push to the mult_stack here, you can't end these
      elif u.op is Ops.LOAD and (not isinstance(u.src[0].dtype, PtrDType) or u.src[0].dtype.addrspace != AddrSpace.REG):
        lds += u.dtype.itemsize * mults
      elif u.op is Ops.STORE and (not isinstance(u.src[0].dtype, PtrDType) or u.src[0].dtype.addrspace != AddrSpace.REG):
        lds += u.src[1].dtype.itemsize * mults
      elif u.op in GroupOp.ALU and u not in dont_count: flops += (mults * (2 if u.op is Ops.MULACC else 1)) * u.dtype.count
      elif u.op is Ops.WMMA and u not in dont_count: flops += 2 * prod(u.arg[1]) // u.arg[5] * mults
-    return Estimates(flops, lds, sum(mem.values()))
+    return Estimates(ssimplify(flops), ssimplify(lds), sum(mem.values()))

@dataclass
 class ProgramSpec:
@@ -81,8 +83,12 @@ class ProgramSpec:
      for u in self.uops:
        if u.op is Ops.DEFINE_VAR: self.vars.append(u)
        if u.op is Ops.DEFINE_GLOBAL: self.globals.append(u.arg)
-        if u.op is Ops.STORE: self.outs.extend([x.arg for x in u.src[0].toposort() if x.op is Ops.DEFINE_GLOBAL])
-        if u.op is Ops.LOAD: self.ins.extend([x.arg for x in u.src[0].toposort() if x.op is Ops.DEFINE_GLOBAL])
+        if u.op is Ops.STORE and (u.src[0].op is Ops.INDEX or (u.src[0].op is Ops.CAST and u.src[0].src[0].op is Ops.INDEX)):
+          idx = u.src[0] if u.src[0].op is Ops.INDEX else u.src[0].src[0]
+          if (buf:=idx.src[0]).op is Ops.DEFINE_GLOBAL: self.outs.append(buf.arg)
+        if u.op is Ops.LOAD and (u.src[0].op is Ops.INDEX or (u.src[0].op is Ops.CAST and u.src[0].src[0].op is Ops.INDEX)):
+          idx = u.src[0] if u.src[0].op is Ops.INDEX else u.src[0].src[0]
+          if (buf:=idx.src[0]).op is Ops.DEFINE_GLOBAL: self.ins.append(buf.arg)
        if u.op is Ops.SPECIAL:
          # NOTE: you have to set local_size and global_size to the base [1,1,1] outside this
          if u.arg[0] == 'i': self.local_size = None
@@ -1,8 +1,8 @@
 from typing import Literal, Callable, cast
 import os, math, sys
 from collections import defaultdict, Counter
-from tinygrad.codegen.opt import tc, axis_letters
-from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, range_str
+from tinygrad.codegen.opt import tc
+from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, range_str, axis_letters
 from tinygrad.helpers import strip_parens, getenv, prod, dedup, AMX, CPU_COUNT
 from tinygrad.dtype import ImageDType, dtypes, DType, PtrDType, AddrSpace, truncate
 from tinygrad.renderer import Renderer
@@ -37,7 +37,7 @@ base_rewrite = PatternMatcher([
  (UPat(Ops.CONST, dtype=dtypes.uint32, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.arg)}u"),
  (UPat(Ops.CONST, dtype=dtypes.bool, name="x"), lambda ctx,x: "1" if x.arg else "0"),
  # consts are rendered to larger type and casted
-  (UPat(Ops.CONST, (dtypes.bfloat16, dtypes.half), name="x"), lambda ctx,x: f"({ctx.render_cast(x.dtype, f'{x.arg}f')})"),
+  (UPat(Ops.CONST, (*dtypes.fp8s, dtypes.bfloat16, dtypes.half), name="x"), lambda ctx,x: f"({ctx.render_cast(x.dtype, f'{x.arg}f')})"),
  (UPat(Ops.CONST, (dtypes.uint8, dtypes.uint16), name="x"), lambda ctx,x: f"({ctx.render_cast(x.dtype, f'{x.arg}u')})"),
  (UPat(Ops.CONST, (dtypes.int8, dtypes.int16), name="x"), lambda ctx,x: f"({ctx.render_cast(x.dtype, str(x.arg))})"),
  # default const render
@@ -95,7 +95,7 @@ class CStyleLanguage(Renderer):
  infinity: str = "INFINITY"
  nan: str = "NAN"
  code_for_op: dict = {
-    Ops.SQRT: lambda x,dtype: f"sqrt({x})", Ops.RECIP: lambda x,dtype: f"(1/{x})", Ops.NEG: lambda x,dtype: f"-{x}",
+    Ops.SQRT: lambda x,dtype: f"sqrt({x})", Ops.RECIPROCAL: lambda x,dtype: f"(1/{x})", Ops.NEG: lambda x,dtype: f"-{x}",
    Ops.EXP2: lambda x,dtype: f"exp2({x})", Ops.LOG2: lambda x,dtype: f"log2({x})", Ops.SIN: lambda x,dtype: f"sin({x})",
    Ops.TRUNC: lambda x,dtype: f"trunc({x})",
    Ops.AND: lambda a,b,dtype: f"({a}&{b})", Ops.XOR: lambda a,b,dtype: f"({a}^{b})", Ops.OR: lambda a,b,dtype: f"({a}|{b})",
@@ -143,7 +143,7 @@ class CStyleLanguage(Renderer):
    c: defaultdict[str, int] = defaultdict(int)
    name = "test"
    for u in uops:
-      if u.op is Ops.NOOP: continue
+      if u.op in {Ops.NOOP, Ops.GROUP}: continue
      if u.op is Ops.AFTER:
        r[u] = r[u.src[0]]
        continue
@@ -157,7 +157,8 @@ class CStyleLanguage(Renderer):

      # mark buffers that we store to writable
      if u.op is Ops.STORE:
-        for up in u.src[0].toposort():
+        # NOTE: we gate on RANGE to not follow it back
+        for up in u.src[0].toposort(lambda x: x.op is not Ops.RANGE):
          if up.op is Ops.DEFINE_GLOBAL: bufs[up] = (bufs[up][0], (bufs[up][1][0], True))

      # naming
@@ -208,7 +209,7 @@ class ClangRenderer(CStyleLanguage):
  # language options
  buffer_suffix = " restrict"
  type_map = {dtypes.bool:"_Bool", dtypes.half:"__fp16"}
-  code_for_op = {**({k:v for k,v in CStyleLanguage.code_for_op.items() if k not in [Ops.EXP2, Ops.SIN, Ops.LOG2, Ops.TRUNC, Ops.RECIP]}),
+  code_for_op = {**({k:v for k,v in CStyleLanguage.code_for_op.items() if k not in [Ops.EXP2, Ops.SIN, Ops.LOG2, Ops.TRUNC, Ops.RECIPROCAL]}),
                 Ops.SQRT: lambda x,dtype: f"__builtin_sqrt({x})" if dtype == dtypes.float64 else f"__builtin_sqrtf({x})",
                 Ops.TRUNC: lambda x,dtype: f"__builtin_trunc({x})" if dtype == dtypes.float64 else f"__builtin_truncf({x})",
                 Ops.FDIV: lambda a,b,dtype: f"({a}/{b})"}
@@ -269,7 +270,8 @@ class OpenCLRenderer(CStyleLanguage):
      lambda ctx,buf,idx,var,gate: f"({ctx[gate]}?read_imagef({ctx[buf]}, smp, {ctx[idx]}):{ctx[var]})"),
    (UPat(Ops.LOAD, dtype=dtypes.float.vec(4), src=(UPat.var('buf').index(UPat.var('idx', dtypes.int.vec(2))),)),
      lambda ctx,buf,idx: f"read_imagef({ctx[buf]}, smp, {ctx[idx]})"),
-    (UPat(Ops.STORE, src=(UPat.var('buf').index(UPat.var('idx', dtypes.int.vec(2))), UPat.var("var", dtypes.float.vec(4))), allow_any_len=True),
+    (UPat(Ops.STORE, src=(UPat.var('buf').index(UPat.var('idx', dtypes.int.vec(2)), allow_any_len=True),
+                          UPat.var("var", dtypes.float.vec(4))), allow_any_len=True),
      lambda ctx,buf,idx,var: f"write_imagef({ctx[buf]}, {ctx[idx]}, {ctx[var]});"),
  ]) + base_rewrite

@@ -345,7 +347,8 @@ class CUDARenderer(CStyleLanguage):
  shared_max = 49152

  def __init__(self, arch:str):
-    self.tensor_cores, self.arch = tc.cuda_sm80 if int(arch[3:]) >= 80 else tc.cuda_sm75 if int(arch[3:]) >= 75 else [], arch
+    self.arch = arch
+    self.tensor_cores = tc.cuda_sm89 if int(arch[3:]) >= 89 else tc.cuda_sm80 if int(arch[3:]) >= 80 else tc.cuda_sm75 if int(arch[3:]) >= 75 else []
  def __reduce__(self): return self.__class__, (self.arch,)

  # language options
@@ -363,9 +366,15 @@ class CUDARenderer(CStyleLanguage):
    Ops.LOG2: lambda x,dtype: f"hlog2({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"log2({x})",
    Ops.EXP2: lambda x,dtype: f"hexp2({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"exp2({x})",
    Ops.SQRT: lambda x,dtype: f"hsqrt({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"sqrt({x})",
-    Ops.RECIP: lambda x,dtype: f"hrcp({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"(1/{x})" }
-  type_map = {dtypes.bfloat16: "nv_bfloat16"}
-
+    Ops.RECIPROCAL: lambda x,dtype: f"hrcp({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"(1/{x})" }
+  type_map = {dtypes.bfloat16: "nv_bfloat16", dtypes.fp8e4m3: "__nv_fp8_e4m3", dtypes.fp8e5m2: "__nv_fp8_e5m2"}
+  extra_matcher = PatternMatcher([
+    (UPat(Ops.CAST, dtypes.fp8s, UPat.var("x", dtypes.fp8s), name='y'), lambda x,y: x.cast(dtypes.float).cast(y.dtype) if x.dtype!=y.dtype else None),
+    (UPat(GroupOp.ALU, dtype=dtypes.fp8s, name="x"),
+     lambda x: UOp(x.op, dtypes.float, tuple(vv.cast(dtypes.float) for vv in x.src), x.arg).cast(x.dtype)),
+    (UPat(GroupOp.ALU, dtypes.bool, name="alu", src=(UPat.var("x", dtype=dtypes.fp8s), UPat.var("y", dtype=dtypes.fp8s))),
+     lambda alu,x,y: UOp(alu.op, dtypes.bool, (x.cast(dtypes.float), y.cast(dtypes.float)), alu.arg)),
+  ]) + extra_pm
  def render_vector_prefix(self, dt:DType) -> str:
    vec, scal = self.render_dtype(dt), self.render_dtype(dt.scalar()),
    elems, header = ', '.join(_nms[:dt.count]), ', '.join([f"{scal} {x}" for x in _nms[:dt.count]])
@@ -376,11 +385,12 @@ class CUDARenderer(CStyleLanguage):
    prefix = ["#define INFINITY (__int_as_float(0x7f800000))","#define NAN (__int_as_float(0x7fffffff))"]

    used_dtypes = uops_to_dtypes(uops)
+    if any(dt.scalar() in dtypes.fp8s for dt in used_dtypes): prefix.append("#include <cuda_fp8.h>")
    if any(dt.scalar() == dtypes.half for dt in used_dtypes): prefix.append("#include <cuda_fp16.h>")
    if any(dt.scalar() == dtypes.bfloat16 for dt in used_dtypes): prefix.append("#include <cuda_bf16.h>")
-    prefix += [self.render_vector_prefix(dt) for dt in used_dtypes if dt.count in (4,8) and dt.scalar() in {dtypes.half, dtypes.bfloat16}]
-
-    dt_map_in = { dtypes.float: "tf32", dtypes.half: "f16", dtypes.bfloat16: "bf16" }
+    prefix += [self.render_vector_prefix(dt) for dt in used_dtypes if (dt.count in (4,8) and dt.scalar() in {dtypes.half, dtypes.bfloat16})
+      or (dt.count in (8,16) and dt.scalar() in dtypes.fp8s)]
+    dt_map_in = { dtypes.float: "tf32", dtypes.half: "f16", dtypes.bfloat16: "bf16", dtypes.fp8e4m3: "e4m3", dtypes.fp8e5m2: "e5m2" }
    dt_map_out = { dtypes.float: "f32", dtypes.half: "f16" }
    for name, (N, M, K), dtype_in, dtype_out, _, _, upcast_axes, _ in wmma_args(uops):
      upcast_sizes = [prod(size for _, size in upcast) for upcast in upcast_axes]
@@ -414,7 +424,7 @@ class AMDRenderer(CStyleLanguage):

  @staticmethod
  def get_tensor_cores(arch):
-    return {"gfx942": tc.amd_cdna, "gfx950": tc.amd_cdna, "gfx1200": tc.amd_rdna4, "gfx1201": tc.amd_rdna4}.get(arch.split(":")[0], tc.amd_rdna3)
+    return {"gfx942": tc.amd_cdna, "gfx950": tc.amd_cdna4, "gfx1200": tc.amd_rdna4, "gfx1201": tc.amd_rdna4}.get(arch.split(":")[0], tc.amd_rdna3)
  def __init__(self, arch:str): # gfx942 => MI300, gfx1100 => RX 7900, gfx1201 => RX 9700
    self.arch = arch
    self.tensor_cores = self.get_tensor_cores(arch)
@@ -49,9 +49,11 @@ def render_wmma_amx(ctx, wmma: UOp) -> str:
 def render_wmma_amd(ctx, wmma: UOp, cdna=False) -> str:
  dt_map = {dtypes.half: "f16", dtypes.float: "f32", dtypes.ushort: "bf16.1k" if cdna else "bf16", dtypes.bfloat16: "bf16.1k" if cdna else "bf16"}
  # https://github.com/llvm/llvm-project/blob/main/clang/test/CodeGenOpenCL/builtins-amdgcn-mfma.cl
+  N,M,K = wmma.arg[1]
  if cdna:
+    if K == 32: dt_map.update({dtypes.half: ".f16", dtypes.bfloat16: ".bf16"})
    return f"  {ctx[wmma]} = call {ldt(wmma.dtype)} @llvm.amdgcn.mfma.{dt_map[wmma.src[-1].dtype.scalar()]}" + \
-           f".16x16x16{dt_map[wmma.src[0].dtype.scalar()]}(" + ", ".join([f"{ldt(w.dtype)} {ctx[w]}" for w in wmma.src]) + ", i32 0, i32 0, i32 0)"
+           f".{N}x{M}x{K}{dt_map[wmma.src[0].dtype.scalar()]}(" + ", ".join([f"{ldt(w.dtype)} {ctx[w]}" for w in wmma.src]) + ", i32 0, i32 0, i32 0)"
  # https://github.com/llvm/llvm-project/blob/main/llvm/test/CodeGen/AMDGPU/GlobalISel/llvm.amdgcn.wmma_32.ll
  # example: %wmma0 = call <8 x float> @llvm.amdgcn.wmma.f32.16x16x16.f16(<16 x half> %v99,<16 x half> %v100,<8 x float> %v101)
  return f"  {ctx[wmma]} = call {ldt(wmma.dtype)} @llvm.amdgcn.wmma.{dt_map[wmma.src[-1].dtype.scalar()]}.16x16x16." + \
@@ -104,15 +106,21 @@ base_rewrite = PatternMatcher([
   f"  {ctx[x]} = select {ldt(x.src[0].dtype)} {ctx[x.src[0]]}, {ldt(x.src[1].dtype)} {ctx[x.src[1]]}, {ldt(x.src[2].dtype)} {ctx[x.src[2]]}"),

  # range
-  (UPat(Ops.RANGE, name="x"), lambda ctx,x:
-   f"  br label %loop_entry_{range_str(x)}\nloop_entry_{range_str(x)}:\n"
-   f"  br label %loop_body_{range_str(x)}\nloop_body_{range_str(x)}:\n"
-   f"  {ctx[x]} = phi {ldt(x.dtype)} [ 0, %loop_entry_{range_str(x)} ], [ {ctx[x]}phi, %loop_latch_{range_str(x)} ]"),
-  (UPat(Ops.END, name="x"), lambda ctx,x:
-   f"  br label %loop_latch_{range_str(x.src[0])}\nloop_latch_{range_str(x.src[0])}:\n"
-   f"  {ctx[x.src[0]]}phi = add {ldt(x.src[0].dtype)} {ctx[x.src[0]]}, 1\n"
-   f"  {ctx[x]} = icmp ult {ldt(x.src[0].dtype)} {ctx[x.src[0]]}phi, {ctx[x.src[0].src[0]]}\n"
-   f"  br i1 {ctx[x]}, label %loop_body_{range_str(x.src[0])}, label %loop_exit_{range_str(x.src[0])}\nloop_exit_{range_str(x.src[0])}:"),
+  (UPat(Ops.RANGE, name="r"), lambda ctx,r:
+   f"  br label %loop_entry_{range_str(r)}\n"
+   f"loop_entry_{range_str(r)}:\n"
+   f"  br label %loop_latch_{range_str(r)}\n"
+   f"loop_latch_{range_str(r)}:\n"
+   f"  {ctx[r]} = phi {ldt(r.dtype)} [ 0, %loop_entry_{range_str(r)} ], [ {ctx[r]}phi, %loop_footer_{range_str(r)} ]\n"
+   f"  {ctx[r]}phi = add {ldt(r.dtype)} {ctx[r]}, 1\n"
+   f"  {ctx[r]}cmp = icmp ult {ldt(r.dtype)} {ctx[r]}, {ctx[r.src[0]]}\n"
+   f"  br i1 {ctx[r]}cmp, label %loop_body_{range_str(r)}, label %loop_exit_{range_str(r)}\n"
+   f"loop_body_{range_str(r)}:"),
+  (UPat(Ops.END, src=(UPat(), UPat(Ops.RANGE, name="r"))), lambda r:
+   f"  br label %loop_footer_{range_str(r)}\n"
+   f"loop_footer_{range_str(r)}:\n"
+   f"  br label %loop_latch_{range_str(r)}\n"
+   f"loop_exit_{range_str(r)}:"),

  # if
  (UPat(Ops.IF, name="x"), lambda ctx,x: f"  br i1 {ctx[x.src[0]]}, label %ifbody_{ctx[x][1:]}, label %ifskip_{ctx[x][1:]}\nifbody_{ctx[x][1:]}:"),
@@ -166,7 +174,7 @@ class LLVMRenderer(Renderer):

    name = "test"
    for u in uops:
-      if u.op is Ops.NOOP: continue
+      if u.op in {Ops.NOOP, Ops.GROUP}: continue
      if u.op is Ops.AFTER:
        r[u] = r[u.src[0]]
        continue
@@ -3,7 +3,7 @@ from tinygrad.dtype import AddrSpace, DType, PtrDType, dtypes
 from tinygrad.helpers import DEBUG, OSX, unwrap
 from tinygrad.renderer import Renderer
 from tinygrad.renderer.cstyle import CUDARenderer
-from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat
+from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, range_str
 import tinygrad.runtime.autogen.mesa as mesa
 import base64, ctypes, ctypes.util, struct, functools, inspect

@@ -21,7 +21,7 @@ def glsl_type(t:DType) -> mesa.struct_glsl_type:
 u_aop = { Ops.ADD: "iadd", Ops.MUL: "imul", Ops.IDIV: "udiv", Ops.MOD: "umod", Ops.CMPLT: "ult", Ops.CMPNE: "ine", Ops.CMPEQ: "ieq", Ops.OR: "ior",
          Ops.AND: "iand", Ops.XOR: "ixor", Ops.WHERE: "bcsel", Ops.MAX: "umax"}
 s_aop = {**u_aop, Ops.CMPLT: "ilt", Ops.IDIV: "idiv", Ops.MOD: "irem", Ops.MAX: "imax"}
-f_aop = { Ops.ADD: "fadd", Ops.MUL: "fmul", Ops.CMPLT: "flt", Ops.CMPNE: "fneu", Ops.CMPEQ: "feq", Ops.FDIV: "fdiv", Ops.RECIP: "frcp",
+f_aop = { Ops.ADD: "fadd", Ops.MUL: "fmul", Ops.CMPLT: "flt", Ops.CMPNE: "fneu", Ops.CMPEQ: "feq", Ops.FDIV: "fdiv", Ops.RECIPROCAL: "frcp",
          Ops.MAX: "fmax", Ops.TRUNC: "ftrunc", Ops.SIN: "fsin", Ops.EXP2: "fexp2", Ops.LOG2: "flog2"}
 aop = {**{x:u_aop for x in (dtypes.bool,)+dtypes.uints}, **{x:s_aop for x in dtypes.sints}, **{x:f_aop for x in dtypes.floats}}

@@ -173,7 +173,7 @@ class NIRRenderer(Renderer):
    self.param_idx, ranges = 0, []

    for u in uops:
-      if u.op == Ops.NOOP or u.op == Ops.INDEX: pass
+      if u.op in {Ops.NOOP, Ops.GROUP, Ops.INDEX}: pass
      elif u.op is Ops.AFTER:
        self.r[u] = self.r[u.src[0]]
      elif u.op == Ops.SINK:
@@ -182,13 +182,17 @@ class NIRRenderer(Renderer):
        self.r[u] = nimm(self.b, self.b.shader.contents.info.shared_size, dtypes.long)
        self.b.shader.contents.info.shared_size += u.dtype.nbytes()
      elif u.op == Ops.RANGE:
-        ranges.append(i:=deref_var(self.b, mesa.nir_local_variable_create(self.b.impl, glsl_type(u.dtype), f"idx{u.arg[0]}".encode()).contents))
+        ranges.append(i:=deref_var(self.b, mesa.nir_local_variable_create(self.b.impl, glsl_type(u.dtype), f"idx{range_str(u)}".encode()).contents))
        nstore(self.b, AddrSpace.REG, i, nimm(self.b, 0, u.dtype), u.dtype)
        mesa.nir_push_loop(self.b)
        self.r[u] = nload(self.b, AddrSpace.REG, i, u.dtype)
+        nif(self.b, nalu(self.b, "ilt", self.r[u], self.r[u.src[0]]), lambda: None, lambda: njump(self.b, mesa.nir_jump_break))
      elif u.op == Ops.END:
-        nif(self.b, nalu(self.b, "ilt", x:=nalu(self.b, "iadd", self.r[u.src[0]], nimm(self.b, 1, u.src[0].dtype)), self.r[u.src[0].src[0]]),
-            functools.partial(nstore, self.b, AddrSpace.REG, ranges.pop(), x, u.src[0].dtype), lambda: njump(self.b, mesa.nir_jump_break))
+        r = u.src[1]
+        next_i = nalu(self.b, "iadd", self.r[r], nimm(self.b, 1, r.dtype))
+        # TODO: this nif should be removable ... but TestMultiTensor.test_double_matmul_shard_W_0 segfaults with it gone
+        nif(self.b, nalu(self.b, "ilt", next_i, self.r[r.src[0]]), lambda: None, lambda: njump(self.b, mesa.nir_jump_break))
+        nstore(self.b, AddrSpace.REG, ranges.pop(), next_i, r.dtype),
        mesa.nir_pop_loop(self.b, None)
      else:
        if (d:=self.def_rewrite.rewrite(u, ctx=self)) is None: raise RuntimeError(f"failed to render {u.op} srcs {[x.dtype for x in u.src]}")
@@ -16,7 +16,7 @@ def render_val(x, dtype):
  return str(int(x)) + ("U" if dtypes.is_unsigned(dtype) else "")

 asm_for_op: dict[Ops, Callable] = {
-  Ops.RECIP: lambda d,a,dt,name: f"rcp{'.approx' if dtypes.is_float(dt) else ''}.{name} {d}, {a};",
+  Ops.RECIPROCAL: lambda d,a,dt,name: f"rcp{'.approx' if dtypes.is_float(dt) else ''}.{name} {d}, {a};",
  Ops.EXP2: lambda d,a,dt,name: f"ex2.approx.{name} {d}, {a};", Ops.LOG2: lambda d,a,dt,name: f"lg2.approx.{name} {d}, {a};",
  Ops.SIN: lambda d,a,dt,name: f"sin.approx.{name} {d}, {a};", Ops.SQRT: lambda d,a,dt,name: f"sqrt.approx.{name} {d}, {a};",
  Ops.TRUNC: lambda d,a,dt,name: f"cvt.rzi.{name}.{name} {d}, {a};",
@@ -49,21 +49,23 @@ ptx_matcher = PatternMatcher([
   lambda x: UOp(x.op, dtypes.uint8, x.src[0:1] + ((x.src[1].cast(dtypes.uint8),) if len(x.src) >= 2 else ()) + x.src[2:]).cast(dtypes.bool)),
  (UPat(Ops.STORE, src=(UPat(dtype=dtypes.int64), UPat(dtype=dtypes.bool)), name="x", allow_any_len=True),
   lambda x: UOp(x.op, dtypes.void, x.src[0:1] + (x.src[1].cast(dtypes.uint8),) + x.src[2:])),
+  # indexing on PTX is in uint64, we do the math while it's still in the graph
+  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx")), name="op", allow_any_len=True), lambda buf,idx,op:
+    UOp(Ops.INDEX, dtype=dtypes.int64, src=(buf, buf.cast(dtypes.int64)+idx.cast(dtypes.int64)*buf.dtype.itemsize)+op.src[2:]) \
+      if op.dtype != dtypes.int64 and buf.dtype.addrspace != AddrSpace.REG else None),
  # load/store use pointer arithmetic, and the cast does nothing
-  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx"))),
-   lambda buf,idx: (buf.cast(dtypes.int64) + idx.cast(dtypes.int64)*buf.dtype.itemsize) if buf.dtype.addrspace != AddrSpace.REG else None),
  (UPat(Ops.CAST, name="x"), lambda x: x.src[0] if isinstance(x.dtype, PtrDType) or x.src[0].dtype == dtypes.void else None),
-  # move mask from INDEX to the load/store to enable pointer arithmetic
-  (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx"), UPat.var("gate"))), UPat.var("alt")), allow_any_len=True, name="l"),
-   lambda buf,idx,gate,alt,l: UOp(Ops.LOAD, alt.dtype, (buf.index(idx), alt, gate, *l.src[2:]))),
-  (UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx"), UPat())), UPat.var("val"), UPat.var("gate")), allow_any_len=True),
-   lambda buf,idx,val,gate: UOp.store(buf.index(idx), val, gate)),
  # ptx shr and shl instructions require y to be uint
  (UPat.var("x") << UPat.var("y"), lambda x,y: UOp(Ops.SHL, x.dtype, (x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None),
  (UPat.var("x") >> UPat.var("y"), lambda x,y: UOp(Ops.SHR, x.dtype, (x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None),
 ])

-def mem_type(x: UOp): return 'shared' if any(_x.op is Ops.DEFINE_LOCAL for _x in x.src[0].toposort()) else 'global'
+def mem_type(x:UOp) -> str:
+  match x.op:
+    case Ops.AFTER: return mem_type(x.src[0])
+    case Ops.DEFINE_LOCAL: return 'shared'
+    case Ops.DEFINE_GLOBAL: return 'global'
+    case _: raise RuntimeError(f"{x.op} needs to be memory")

 def render_wmma(ctx: "PTXRenderer", wmma: UOp):
  assert ctx.wmma_r, "registry values for wmma must be populated"
@@ -88,9 +90,6 @@ def modifier(a: DType, b: DType): return '.rzi' if dtypes.is_int(a) and dtypes.i
 string_rewrite = PatternMatcher([
  (UPat.cvar("x", dtypes.bool), lambda ctx, x: f"setp.ne.s16 {ctx.r[x]}, {render_val(x.arg, x.dtype)}, 0;"),
  (UPat.cvar("x"), lambda ctx, x: f"mov.b{ctx.types[x.dtype][1:]} {ctx.r[x]}, {render_val(x.arg, x.dtype)};"),
-  (UPat(Ops.STORE, name="x", src=(UPat.var('bidx'), UPat.var("var")), allow_any_len=True), lambda ctx, x, bidx, var: f"st.{mem_type(bidx)}" + \
-    f"{f'.v{cnt}' if ((cnt:=var.dtype.count)>1) else ''}.{ctx.mem_types[var.dtype.scalar()]} " + \
-    f"[{ctx.r[bidx]}+0], {('{' + ', '.join(ctx.r[var]) + '}') if var.dtype.count > 1 else ctx.r[var]};"),
  (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: f"mov.u32 %{x.arg}, %{'ctaid' if x.arg[0] == 'g' else 'tid'}.{chr(120+int(x.arg[-1]))};"),
  (UPat(Ops.DEFINE_GLOBAL, name="x"), lambda ctx, x: f"ld.param.{ctx.types[dtypes.ulong]} {ctx.r[x]}, [data{x.arg}+0];"),
  (UPat((Ops.CMPLT, Ops.CMPNE, Ops.CMPEQ), name="x", allow_any_len=True, src=(UPat.var("src0"),)),
@@ -103,22 +102,32 @@ string_rewrite = PatternMatcher([
   lambda ctx, x, a: f"setp.ne.b{ctx.types[a.dtype][1:]} {ctx.r[x]}, {ctx.r[a]}, {render_val(0, a.dtype)};"),
  (UPat(Ops.CAST, name="x", src=(UPat.var("a"),)),
   lambda ctx, x, a: f"cvt{modifier(x.dtype, a.dtype)}.{ctx.cast_types[x.dtype]}.{ctx.cast_types[a.dtype]} {ctx.r[x]}, {ctx.r[a]};"),
-  (UPat(Ops.LOAD, name="x", src=(UPat.var('loc'), UPat(name='alt'), UPat(name="gate", op=GroupOp.ALU)), allow_any_len=True),
-    lambda ctx, x, loc, alt, gate: flatten([
+  # store / gated load / load
+  (UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("loc")), allow_any_len=True), UPat.var("var"))),
+   lambda ctx, loc, var, buf: f"st.{mem_type(buf)}" + \
+    f"{f'.v{cnt}' if ((cnt:=var.dtype.count)>1) else ''}.{ctx.mem_types[var.dtype.scalar()]} " + \
+    f"[{ctx.r[loc]}+0], {('{' + ', '.join(ctx.r[var]) + '}') if var.dtype.count > 1 else ctx.r[var]};"),
+  (UPat(Ops.LOAD, name="x", src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("loc"), UPat.var("gate"))), UPat.var("alt")), allow_any_len=True),
+    lambda ctx, x, loc, alt, gate, buf: flatten([
    [f"mov.{ctx.mem_types[x.dtype.scalar()]} {v}, {render_val(0, x.dtype.scalar())};" for v in ctx.r[x]],
-    [f"@{ctx.r[gate]} ld.{mem_type(x)}.v{x.dtype.count}.{ctx.mem_types[x.dtype.scalar()]} {{{', '.join(ctx.r[x])}}}, [{ctx.r[loc]}+0];"]
+    [f"@{ctx.r[gate]} ld.{mem_type(buf)}.v{x.dtype.count}.{ctx.mem_types[x.dtype.scalar()]} {{{', '.join(ctx.r[x])}}}, [{ctx.r[loc]}+0];"]
  ]) if alt.dtype.count > 1 else [
-    f"@{ctx.r[gate]} ld.{mem_type(x)}.{ctx.mem_types[x.dtype.scalar()]} {ctx.r[x]}, [{ctx.r[loc]}+0];",
+    f"@{ctx.r[gate]} ld.{mem_type(buf)}.{ctx.mem_types[x.dtype.scalar()]} {ctx.r[x]}, [{ctx.r[loc]}+0];",
    f"@!{ctx.r[gate]} mov.b{ctx.types[x.dtype.scalar()][1:]} {ctx.r[x]}, {ctx.r[alt]};"]),
-  (UPat(Ops.LOAD, name="x", src=(UPat.var('loc'),), allow_any_len=True),
-   lambda ctx, x, loc: f"ld.{mem_type(x)}.v{x.dtype.count}.{ctx.mem_types[x.dtype.scalar()]} {{{', '.join(ctx.r[x])}}}, [{ctx.r[loc]}+0];" \
-     if x.dtype.count > 1 else f"ld.{mem_type(x)}.{ctx.mem_types[x.dtype]} {ctx.r[x]}, [{ctx.r[loc]}+0];"),
+  (UPat(Ops.LOAD, name="x", src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("loc"))),), allow_any_len=True),
+    lambda ctx, x, loc, buf: f"ld.{mem_type(buf)}.v{x.dtype.count}.{ctx.mem_types[x.dtype.scalar()]} {{{', '.join(ctx.r[x])}}}, [{ctx.r[loc]}+0];" \
+     if x.dtype.count > 1 else f"ld.{mem_type(buf)}.{ctx.mem_types[x.dtype]} {ctx.r[x]}, [{ctx.r[loc]}+0];"),
+  # simple
  (UPat(Ops.DEFINE_REG, src=()), lambda ctx: []),
-  (UPat(Ops.RANGE, name="x"), lambda ctx, x: [f"mov.u32 {ctx.r[x]}, 0;", "LOOP_" + f"{ctx.r[x][1:]}:"]),
-  (UPat(Ops.END, name="x", src=(UPat.var("src0"),), allow_any_len=True), lambda ctx, x, src0: [
-    ctx.code_for_op[Ops.ADD](ctx.r[src0], ctx.r[src0], "1", dtypes.int, ctx.types[dtypes.int]),
-    ctx.code_for_op[Ops.CMPLT](ctx.r[x], ctx.r[x.src[0]], ctx.r[src0.src[0]], dtypes.int, ctx.types[dtypes.int]),
-    f"@{ctx.r[x]} bra LOOP_{ctx.r[src0][1:]};"]),
+  (UPat(Ops.RANGE, name="r"), lambda ctx, r: [
+    f"mov.u32 {ctx.r[r]}, -1;",
+    f"bra END_{ctx.r[r][1:]};",
+    "LOOP_" + f"{ctx.r[r][1:]}:"]),
+  (UPat(Ops.END, name="x", src=(UPat(), UPat(Ops.RANGE, name="r"))), lambda ctx, x, r: [
+    "END_" + f"{ctx.r[r][1:]}:",
+    ctx.code_for_op[Ops.ADD](ctx.r[r], ctx.r[r], "1", dtypes.int, ctx.types[dtypes.int]),
+    ctx.code_for_op[Ops.CMPLT](ctx.r[x], ctx.r[r], ctx.r[r.src[0]], dtypes.int, ctx.types[dtypes.int]),
+    f"@{ctx.r[x]} bra LOOP_{ctx.r[r][1:]};"]),
  (UPat(Ops.DEFINE_LOCAL, name="x"),
   lambda ctx, x: [f".shared .align 16 .b8 local{x.arg}[{x.dtype.size*x.dtype.itemsize}];", f"mov.u64 {ctx.r[x]}, local{x.arg}[0];"]),
  (UPat(Ops.IF, name="x"), lambda ctx, x: f"@!{ctx.r[x.src[0]]} bra IF_{ctx.r[x.src[0]][1:]}_{ctx.uops.index(x)};"),
@@ -178,7 +187,7 @@ class PTXRenderer(Renderer):

    name = "test"
    for u in uops:
-      if u.op is Ops.NOOP: continue
+      if u.op in {Ops.NOOP, Ops.GROUP}: continue
      if u.op is Ops.AFTER:
        self.r[u] = self.r[u.src[0]]
        continue
@@ -207,6 +216,7 @@ class PTXRenderer(Renderer):
            typ = "pred" if u.src[1].dtype == dtypes.bool else ("b"+self.types[u.src[1].dtype][1:])
            kernel.append(f"mov.{typ} {self.r[u.src[0]]}, {self.r[u.src[1]]};")
        continue
+      if u.op is Ops.INDEX: continue  # other index we can skip
      if u.op is Ops.SPECIAL: r[u] = "%" + u.arg
      elif u.op is Ops.DEFINE_VAR: bufs.append((u.arg[0], u.dtype))
      elif u.op is Ops.LOAD:
@@ -7,14 +7,11 @@
 # LONGDOUBLE_SIZE is: 16
 #
 import ctypes, ctypes.util, os, gzip, base64, subprocess, tinygrad.helpers as helpers
-def brew_path(nm):
-  try: return f"{subprocess.check_output(['brew', '--prefix', nm]).decode().strip()}/lib/lib{nm}.dylib"
-  except Exception: return 'failed'
 PATHS_TO_TRY = [
  (BASE:=os.getenv('MESA_PATH', f"/usr{'/local/' if helpers.OSX else '/'}lib"))+'/libtinymesa_cpu'+(EXT:='.dylib' if helpers.OSX else '.so'),
  f'{BASE}/libtinymesa{EXT}',
-  brew_path('tinymesa_cpu'),
-  brew_path('tinymesa'),
+  '/opt/homebrew/lib/libtinymesa_cpu.dylib',
+  '/opt/homebrew/lib/libtinymesa.dylib',
 ]
 def _try_dlopen_tinymesa_cpu():
  library = ctypes.util.find_library("tinymesa_cpu")
@@ -4,7 +4,7 @@ import os, ctypes, struct, hashlib, functools, importlib, mmap, errno, array, co
 assert sys.platform != 'win32'
 from dataclasses import dataclass
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, HWQueue, CLikeArgsState, HCQSignal, HCQProgram, FileIOInterface
-from tinygrad.runtime.support.hcq import MMIOInterface, BumpAllocator
+from tinygrad.runtime.support.hcq import MMIOInterface, BumpAllocator, hcq_filter_visible_devices
 from tinygrad.uop.ops import sint
 from tinygrad.device import Compiled, DMAFdRef, BufferSpec, CompilerPairT
 from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, suppress_finalizing, lo32, hi32, colored
@@ -15,11 +15,11 @@ from tinygrad.runtime.autogen.am import am
 from tinygrad.runtime.support.compiler_amd import HIPCompiler, AMDLLVMCompiler
 from tinygrad.runtime.support.elf import elf_loader
 from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager
-from tinygrad.runtime.support.amd import AMDReg, AMDIP, import_module, import_soc, setup_pci_bars
-from tinygrad.runtime.support.system import System, PCIIfaceBase, PCIAllocationMeta, MAP_FIXED, MAP_NORESERVE
-from tinygrad.runtime.support.usb import ASM24Controller, USBMMIOInterface
+from tinygrad.runtime.support.amd import AMDReg, AMDIP, import_module, import_soc, import_ip_offsets
+from tinygrad.runtime.support.system import System, PCIIfaceBase, PCIAllocationMeta, PCIDevice, USBPCIDevice, MAP_FIXED, MAP_NORESERVE
 if getenv("IOCTL"): import extra.hip_gpu_driver.hip_ioctl  # noqa: F401 # pylint: disable=unused-import

+SQTT = getenv("SQTT", 0)
 EVENT_INDEX_PARTIAL_FLUSH = 4 # based on a comment in nvd.h
 WAIT_REG_MEM_FUNCTION_EQ  = 3 # ==
 WAIT_REG_MEM_FUNCTION_NEQ = 4 # !=
@@ -66,11 +66,15 @@ class AMDComputeQueue(HWQueue):
    if self.dev.xccs > 1:
      self._q[prev_len-1] |= (len(self._q) - prev_len)

-  def wait_reg_mem(self, value, mask=0xffffffff, mem=None, reg_req=None, reg_done=None):
-    wrm_info_dw = self.pm4.WAIT_REG_MEM_MEM_SPACE(int(mem is not None)) | self.pm4.WAIT_REG_MEM_OPERATION(int(mem is None)) \
-                | self.pm4.WAIT_REG_MEM_FUNCTION(WAIT_REG_MEM_FUNCTION_GEQ) | self.pm4.WAIT_REG_MEM_ENGINE(0)
+  def set_grbm_broadcast(self):
+    self.wreg(self.gc.regGRBM_GFX_INDEX, **{f'{f}_broadcast_writes': 1 for f in ['se', 'sh' if self.dev.target[0] == 9 else 'sa', 'instance']})
+  def set_grbm_se(self, se): self.wreg(self.gc.regGRBM_GFX_INDEX, se_index=se, instance_broadcast_writes=1)

-    self.pkt3(self.pm4.PACKET3_WAIT_REG_MEM, wrm_info_dw, *(data64_le(mem) if mem is not None else (reg_req, reg_done)), value, mask, 4)
+  def wait_reg_mem(self, value, mask=0xffffffff, mem=None, reg=None, reg_done=0, op=WAIT_REG_MEM_FUNCTION_GEQ):
+    wrm_info_dw = self.pm4.WAIT_REG_MEM_MEM_SPACE(int(mem is not None)) | self.pm4.WAIT_REG_MEM_OPERATION(int(mem is None and reg_done > 0)) \
+                | self.pm4.WAIT_REG_MEM_FUNCTION(op) | self.pm4.WAIT_REG_MEM_ENGINE(0)
+
+    self.pkt3(self.pm4.PACKET3_WAIT_REG_MEM, wrm_info_dw, *(data64_le(mem) if mem is not None else (reg, reg_done)), value, mask, 4)

  def acquire_mem(self, addr=0x0, sz=(1 << 64)-1, gli=1, glm=1, glk=1, glv=1, gl1=1, gl2=1):
    if self.dev.target >= (10,0,0):
@@ -113,7 +117,7 @@ class AMDComputeQueue(HWQueue):

  def memory_barrier(self):
    pf = '' if self.nbio.version[0] == 2 else '0' if self.nbio.version[:2] != (7, 11) else '1'
-    self.wait_reg_mem(reg_req=getattr(self.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_REQ').addr[0],
+    self.wait_reg_mem(reg=getattr(self.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_REQ').addr[0],
                      reg_done=getattr(self.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_DONE').addr[0], value=0xffffffff)
    self.acquire_mem()
    return self
@@ -124,6 +128,22 @@ class AMDComputeQueue(HWQueue):

  ### SQTT ###

+  def sqtt_setup_exec(self, prg, global_size):
+    self.sqtt_userdata(sqtt.struct_rgp_sqtt_marker_pipeline_bind(
+      _0=sqtt.union_rgp_sqtt_marker_pipeline_bind_0(_0=sqtt.struct_rgp_sqtt_marker_pipeline_bind_0_0(
+        identifier=sqtt.RGP_SQTT_MARKER_IDENTIFIER_BIND_PIPELINE, bind_point=(__BIND_POINT_COMPUTE:=1))),
+      _1=sqtt.union_rgp_sqtt_marker_pipeline_bind_1(api_pso_hash=data64_le(prg.libhash[0]))))
+
+    self.sqtt_userdata(sqtt.struct_rgp_sqtt_marker_event(
+      _0=sqtt.union_rgp_sqtt_marker_event_0(_0=sqtt.struct_rgp_sqtt_marker_event_0_0(has_thread_dims=1)),
+      _2=sqtt.union_rgp_sqtt_marker_event_2(cmd_id=next(prg.dev.sqtt_next_cmd_id))), *global_size)
+
+    for xcc in range(self.dev.xccs):
+      with self.pred_exec(xcc_mask=1 << xcc):
+        for i in range(8 if prg.dev.target >= (11,0,0) else 4):
+          self.wreg(getattr(self.gc, f'regCOMPUTE_STATIC_THREAD_MGMT_SE{i}'),
+            ((prg.dev.sqtt_itrace_se_mask >> ((self.dev.se_cnt // self.dev.xccs) * xcc + i)) & 0b1) if SQTT >= 2 else 0xffffffff)
+
  def sqtt_userdata(self, data, *extra_dwords):
    data_ints = [x[0] for x in struct.iter_unpack('<I', bytes(data))] + list(extra_dwords)
    for i in range(0, len(data_ints), 2):
@@ -134,87 +154,97 @@ class AMDComputeQueue(HWQueue):
    self.wreg(self.gc.regSQ_THREAD_TRACE_CTRL, draw_event_en=1, spi_stall_en=1, sq_stall_en=1, reg_at_hwm=2, hiwater=1, util_timer=1,
      mode=int(tracing), **trace_ctrl)

-  # Magic values from mesa/src/amd/vulkan/radv_sqtt.c:radv_emit_spi_config_cntl and src/amd/common/ac_sqtt.c:ac_sqtt_emit_start
  def sqtt_start(self, buf0s:list[HCQBuffer], se_mask:int):
    self.memory_barrier()
-    self.spi_config(tracing=True)
-    # One buffer for one SE, mesa does it with a single buffer and ac_sqtt_get_data_offset, but this is simpler and should work just as well
-    for se in range(len(buf0s)):
-      self.wreg(self.gc.regGRBM_GFX_INDEX, se_index=se, instance_broadcast_writes=1)
-      buf0_lo, buf0_hi = data64_le(buf0s[se].va_addr >> 12)
-      if self.dev.target >= (12,0,0):
-        self.wreg(self.gc.regSQ_THREAD_TRACE_BUF0_SIZE, size=buf0s[se].size >> 12)
-        self.wreg(self.gc.regSQ_THREAD_TRACE_BUF0_BASE_LO, base_lo=buf0_lo)
-        self.wreg(self.gc.regSQ_THREAD_TRACE_BUF0_BASE_HI, base_hi=buf0_hi)
-      else:
-        self.wreg(self.gc.regSQ_THREAD_TRACE_BUF0_SIZE, base_hi=buf0_hi, size=buf0s[se].size >> 12)
-        self.wreg(self.gc.regSQ_THREAD_TRACE_BUF0_BASE, base_lo=buf0_lo)
-      # NOTE: SQTT can only trace instructions on one simd per se, this selects first simd in first wgp in first sa.
-      # For RGP to display instruction trace it has to see it on first SE. Howerver ACE/MEC/whatever does the dispatching starting with second se,
-      # and on amdgpu/non-AM it also does weird things with dispatch order inside se: around 7 times out of 10 it starts from the last cu, but
-      # sometimes not, especially if the kernel has more than one wavefront which means that kernels with small global size might get unlucky and
-      # be dispatched on something else and not be seen in instruction tracing tab. You can force the wavefronts of a kernel to be dispatched on the
-      # CUs you want to by disabling other CUs via bits in regCOMPUTE_STATIC_THREAD_MGMT_SE<x> and trace even kernels that only have one wavefront.
-      cs_wtype = (1 << 6) if self.dev.target >= (12,0,0) else self.soc.SQ_TT_WTYPE_INCLUDE_CS_BIT
-      self.wreg(self.gc.regSQ_THREAD_TRACE_MASK, wtype_include=cs_wtype, simd_sel=0, wgp_sel=0, sa_sel=0)
-      reg_include = self.soc.SQ_TT_TOKEN_MASK_SQDEC_BIT | self.soc.SQ_TT_TOKEN_MASK_SHDEC_BIT | self.soc.SQ_TT_TOKEN_MASK_GFXUDEC_BIT | \
-                    self.soc.SQ_TT_TOKEN_MASK_COMP_BIT | self.soc.SQ_TT_TOKEN_MASK_CONTEXT_BIT
-      token_exclude = (1 << self.soc.SQ_TT_TOKEN_EXCLUDE_PERF_SHIFT) if self.dev.target < (12,0,0) else 0
+    if self.dev.target[0] == 9:
+      self.set_grbm_broadcast()
+      self.wreg(self.gc.regSQ_THREAD_TRACE_MASK, simd_en=0xf, cu_sel=0, sq_stall_en=1, spi_stall_en=1, reg_stall_en=1, vm_id_mask=0)
+      for se in range(len(buf0s)):
+        mask = (__SQTT_MISC:=1<<0) | (__SQTT_TIME:=1<<1) | (__SQTT_REG:=1<<2) | (__SQTT_WAVE_START:=1<<3) | (__SQTT_WAVE_END:=1<<6) \
+             | (__SQTT_USERDATA:=1<<12) | (__SQTT_REG_CS:=1<<5) | (__SQTT_REG_CS_PRIV:=1<<15)
+        if (se_mask >> se) & 0b1: mask |= (__SQTTINST:=1<<10) | (__SQTT_INST_PC:=1<<11) | (__SQTT_ISSUE:=1<<13)

-      # disable tracing
-      if not (se_mask >> se) & 0b1:
-        # gfx12 doesn't have enums with all fields, so it's hardcoded, but it's the same as gfx11.
-        token_exclude |= (1 << self.soc.SQ_TT_TOKEN_EXCLUDE_VMEMEXEC_SHIFT | 1 << self.soc.SQ_TT_TOKEN_EXCLUDE_ALUEXEC_SHIFT | \
-                         1 << self.soc.SQ_TT_TOKEN_EXCLUDE_VALUINST_SHIFT | 1 << self.soc.SQ_TT_TOKEN_EXCLUDE_IMMEDIATE_SHIFT | \
-                         1 << self.soc.SQ_TT_TOKEN_EXCLUDE_INST_SHIFT) if self.dev.target < (12,0,0) else 0x927
+        with self.pred_exec(xcc_mask=1<<(se // (ses_per_xcc:=(self.dev.se_cnt // self.dev.xccs)))):
+          self.set_grbm_se(se % ses_per_xcc)
+          self.wreg(self.gc.regSQ_THREAD_TRACE_TOKEN_MASK, reg_mask=0xf, token_mask=mask)
+          self.wreg(self.gc.regSQ_THREAD_TRACE_TOKEN_MASK2, inst_mask=0xffffffff)
+          self.wreg(self.gc.regSQ_THREAD_TRACE_BASE, addr=lo32(buf0s[se].va_addr >> 12))
+          self.wreg(self.gc.regSQ_THREAD_TRACE_BASE2, addr_hi=hi32(buf0s[se].va_addr >> 12))
+          self.wreg(self.gc.regSQ_THREAD_TRACE_SIZE, size=buf0s[se].size >> 12)
+          self.wreg(self.gc.regSQ_THREAD_TRACE_CTRL, reset_buffer=1)
+          self.wreg(self.gc.regSQ_THREAD_TRACE_MODE, mask_cs=1, autoflush_en=1, mode=1)
+    else:
+      self.spi_config(tracing=True)
+      # One buffer for one SE, mesa does it with a single buffer and ac_sqtt_get_data_offset, but this is simpler and should work just as well
+      for se in range(len(buf0s)):
+        self.set_grbm_se(se)

-      token_mask = {} if self.dev.target < (12,0,0) else {'exclude_barrier_wait': 1}
-      self.wreg(self.gc.regSQ_THREAD_TRACE_TOKEN_MASK, reg_include=reg_include, token_exclude=token_exclude, bop_events_token_include=1, **token_mask)
-      # Enable SQTT
-      self.sqtt_config(tracing=True)
-    # Restore global broadcasting
-    self.wreg(self.gc.regGRBM_GFX_INDEX, se_broadcast_writes=1, sa_broadcast_writes=1, instance_broadcast_writes=1)
-    self.wreg(self.gc.regCOMPUTE_THREAD_TRACE_ENABLE, 1)
+        buf0_lo, buf0_hi = data64_le(buf0s[se].va_addr >> 12)
+        if self.dev.target >= (12,0,0):
+          self.wreg(self.gc.regSQ_THREAD_TRACE_BUF0_SIZE, size=buf0s[se].size >> 12)
+          self.wreg(self.gc.regSQ_THREAD_TRACE_BUF0_BASE_LO, base_lo=buf0_lo)
+          self.wreg(self.gc.regSQ_THREAD_TRACE_BUF0_BASE_HI, base_hi=buf0_hi)
+        else:
+          self.wreg(self.gc.regSQ_THREAD_TRACE_BUF0_SIZE, base_hi=buf0_hi, size=buf0s[se].size >> 12)
+          self.wreg(self.gc.regSQ_THREAD_TRACE_BUF0_BASE, base_lo=buf0_lo)
+        # NOTE: SQTT can only trace instructions on one simd per se, this selects first simd in first wgp in first sa.
+        # For RGP to display instruction trace it has to see it on first SE. Howerver ACE/MEC/whatever does the dispatching starting with second se,
+        # and on amdgpu/non-AM it also does weird things with dispatch order inside se: around 7 times out of 10 it starts from the last cu, but
+        # sometimes not, especially if the kernel has more than one wavefront which means that kernels with small global size might get unlucky and
+        # be dispatched on something else and not be seen in instruction tracing tab. You can force the wavefronts of a kernel to be dispatched on the
+        # CUs you want to by disabling other CUs via bits in regCOMPUTE_STATIC_THREAD_MGMT_SE<x> and trace even kernels that only have one wavefront.
+        cs_wtype = (1 << 6) if self.dev.target >= (12,0,0) else self.soc.SQ_TT_WTYPE_INCLUDE_CS_BIT
+        self.wreg(self.gc.regSQ_THREAD_TRACE_MASK, wtype_include=cs_wtype, simd_sel=0, wgp_sel=0, sa_sel=0)
+        reg_include = self.soc.SQ_TT_TOKEN_MASK_SQDEC_BIT | self.soc.SQ_TT_TOKEN_MASK_SHDEC_BIT | self.soc.SQ_TT_TOKEN_MASK_GFXUDEC_BIT | \
+                      self.soc.SQ_TT_TOKEN_MASK_COMP_BIT | self.soc.SQ_TT_TOKEN_MASK_CONTEXT_BIT
+        token_exclude = (1 << self.soc.SQ_TT_TOKEN_EXCLUDE_PERF_SHIFT) if self.dev.target < (12,0,0) else 0
+
+        # disable instr tracing
+        if not (se_mask >> se) & 0b1:
+          # gfx12 doesn't have enums with all fields, so it's hardcoded, but it's the same as gfx11.
+          token_exclude |= (1 << self.soc.SQ_TT_TOKEN_EXCLUDE_VMEMEXEC_SHIFT | 1 << self.soc.SQ_TT_TOKEN_EXCLUDE_ALUEXEC_SHIFT | \
+                            1 << self.soc.SQ_TT_TOKEN_EXCLUDE_VALUINST_SHIFT | 1 << self.soc.SQ_TT_TOKEN_EXCLUDE_IMMEDIATE_SHIFT | \
+                            1 << self.soc.SQ_TT_TOKEN_EXCLUDE_INST_SHIFT) if self.dev.target < (12,0,0) else 0x927
+
+        self.wreg(self.gc.regSQ_THREAD_TRACE_TOKEN_MASK, reg_include=reg_include, token_exclude=token_exclude, bop_events_token_include=1,
+                  **({} if self.dev.target < (12,0,0) else {'exclude_barrier_wait': 1}))
+        self.sqtt_config(tracing=True)
+
+    self.set_grbm_broadcast()
+    if self.dev.target[0] > 9: self.wreg(self.gc.regCOMPUTE_THREAD_TRACE_ENABLE, 1)
    self.memory_barrier()
    return self

  # Magic values from src/amd/common/ac_sqtt.c:ac_sqtt_emit_stop and src/amd/common/ac_sqtt.c:ac_sqtt_emit_wait
-  def sqtt_stop(self, ses: int, wptrs: HCQBuffer):
+  def sqtt_stop(self, ses:int, wptrs:HCQBuffer):
    self.memory_barrier()
+    self.set_grbm_broadcast()
+
    # Start shutting everything down
-    self.wreg(self.gc.regCOMPUTE_THREAD_TRACE_ENABLE, 0)
-    self.pkt3(self.pm4.PACKET3_EVENT_WRITE, self.pm4.EVENT_TYPE(self.soc.THREAD_TRACE_FINISH) | self.pm4.EVENT_INDEX(0))
+    if self.dev.target[0] == 9: self.wreg(self.gc.regSQ_THREAD_TRACE_MODE, mask_cs=1, autoflush_en=1, mode=0)
+    else:
+      self.wreg(self.gc.regCOMPUTE_THREAD_TRACE_ENABLE, 0)
+      self.pkt3(self.pm4.PACKET3_EVENT_WRITE, self.pm4.EVENT_TYPE(self.soc.THREAD_TRACE_FINISH) | self.pm4.EVENT_INDEX(0))
+
    # For each SE wait for finish to complete and copy regSQ_THREAD_TRACE_WPTR to know where in the buffer trace data ends
    for se in range(ses):
-      self.wreg(self.gc.regGRBM_GFX_INDEX, se_index=se, instance_broadcast_writes=1)
-      # Wait for FINISH_PENDING==0
-      self.pkt3(self.pm4.PACKET3_WAIT_REG_MEM, self.pm4.WAIT_REG_MEM_FUNCTION(WAIT_REG_MEM_FUNCTION_EQ),
-                self.gc.regSQ_THREAD_TRACE_STATUS.addr[0], 0, 0, self.gc.regSQ_THREAD_TRACE_STATUS.fields_mask('finish_pending'), 4)
-      # Disable SQTT
-      self.sqtt_config(tracing=False)
-      # Wait for BUSY==0
-      self.pkt3(self.pm4.PACKET3_WAIT_REG_MEM, self.pm4.WAIT_REG_MEM_FUNCTION(WAIT_REG_MEM_FUNCTION_EQ),
-                self.gc.regSQ_THREAD_TRACE_STATUS.addr[0], 0, 0, self.gc.regSQ_THREAD_TRACE_STATUS.fields_mask('busy'), 4)
+      self.set_grbm_se(se)
+
+      status_reg = self.gc.regSQ_THREAD_TRACE_STATUS.addr[0] - (self.pm4.PACKET3_SET_UCONFIG_REG_START if self.dev.target[0] == 9 else 0)
+      if self.dev.target >= (10, 0, 0):
+        self.wait_reg_mem(reg=status_reg, mask=self.gc.regSQ_THREAD_TRACE_STATUS.fields_mask('finish_pending'), op=WAIT_REG_MEM_FUNCTION_EQ, value=0)
+        self.sqtt_config(tracing=False)
+      self.wait_reg_mem(reg=status_reg, mask=self.gc.regSQ_THREAD_TRACE_STATUS.fields_mask('busy'), op=WAIT_REG_MEM_FUNCTION_EQ, value=0)
+      self.pkt3(self.pm4.PACKET3_EVENT_WRITE, self.pm4.EVENT_TYPE(self.soc.CS_PARTIAL_FLUSH) | self.pm4.EVENT_INDEX(EVENT_INDEX_PARTIAL_FLUSH))
+
      # Copy WPTR to memory (src_sel = perf, dst_sel = tc_l2, wr_confirm = True)
      self.pkt3(self.pm4.PACKET3_COPY_DATA, 1 << 20 | 2 << 8 | 4, self.gc.regSQ_THREAD_TRACE_WPTR.addr[0], 0, *data64_le(wptrs.va_addr+(se*4)))
-    # Restore global broadcasting
-    self.wreg(self.gc.regGRBM_GFX_INDEX, se_broadcast_writes=1, sa_broadcast_writes=1, instance_broadcast_writes=1)
-    self.spi_config(tracing=False)
+
+    self.set_grbm_broadcast()
+    if self.dev.target[0] > 9: self.spi_config(tracing=False)
    self.memory_barrier()
    return self

-  def sqtt_prg_marker(self, prg:AMDProgram, global_size:tuple[sint, ...]):
-    BIND_POINT_COMPUTE = 1
-
-    self.sqtt_userdata(sqtt.struct_rgp_sqtt_marker_pipeline_bind(
-      _0=sqtt.union_rgp_sqtt_marker_pipeline_bind_0(_0=sqtt.struct_rgp_sqtt_marker_pipeline_bind_0_0(
-        identifier=sqtt.RGP_SQTT_MARKER_IDENTIFIER_BIND_PIPELINE, bind_point=BIND_POINT_COMPUTE)),
-      _1=sqtt.union_rgp_sqtt_marker_pipeline_bind_1(api_pso_hash=data64_le(prg.libhash[0]))))
-
-    self.sqtt_userdata(sqtt.struct_rgp_sqtt_marker_event(
-      _0=sqtt.union_rgp_sqtt_marker_event_0(_0=sqtt.struct_rgp_sqtt_marker_event_0_0(has_thread_dims=1)),
-      _2=sqtt.union_rgp_sqtt_marker_event_2(cmd_id=next(prg.dev.sqtt_next_cmd_id))), *global_size)
-
  def exec(self, prg:AMDProgram, args_state:CLikeArgsState, global_size:tuple[sint, ...], local_size:tuple[sint, ...]):
    self.bind_args_state(args_state)

@@ -238,7 +268,7 @@ class AMDComputeQueue(HWQueue):

    user_regs += [*data64_le(args_state.buf.va_addr)]

-    if prg.dev.sqtt_enabled: self.sqtt_prg_marker(prg, global_size)
+    if prg.dev.sqtt_enabled: self.sqtt_setup_exec(prg, global_size)

    self.wreg(self.gc.regCOMPUTE_PGM_LO, *data64_le(prg.prog_addr >> 8))
    self.wreg(self.gc.regCOMPUTE_PGM_RSRC1, prg.rsrc1, prg.rsrc2)
@@ -254,13 +284,8 @@ class AMDComputeQueue(HWQueue):
    if (10,0,0) <= prg.dev.target < (11,0,0): self.wreg(self.gc.mmCP_COHER_START_DELAY, 0x20)

    self.wreg(self.gc.regCOMPUTE_RESTART_X, 0, 0, 0)
-    self.wreg(self.gc.regCOMPUTE_STATIC_THREAD_MGMT_SE0, 0xFFFFFFFF, 0xFFFFFFFF)
-    self.wreg(self.gc.regCOMPUTE_STATIC_THREAD_MGMT_SE2, 0xFFFFFFFF, 0xFFFFFFFF)
-    if prg.dev.target >= (11,0,0): self.wreg(self.gc.regCOMPUTE_STATIC_THREAD_MGMT_SE4, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF)
-
    self.wreg(self.gc.regCOMPUTE_USER_DATA_0, *user_regs)
    self.wreg(self.gc.regCOMPUTE_RESOURCE_LIMITS, 0)
-
    self.wreg(self.gc.regCOMPUTE_START_X, 0, 0, 0, *local_size, 0, 0)

    gfx10p = {'cs_w32_en': int(prg.wave32)} if prg.dev.target >= (10,0,0) else {}
@@ -321,6 +346,7 @@ class AMDComputeQueue(HWQueue):
 class AMDComputeAQLQueue(AMDComputeQueue):
  def exec(self, prg:AMDProgram, args_state:CLikeArgsState, global_size:tuple[sint, ...], local_size:tuple[sint, ...]):
    self.bind_args_state(args_state)
+    if prg.dev.sqtt_enabled: self.sqtt_setup_exec(prg, global_size)
    self._q.append(pkt:=hsa.hsa_kernel_dispatch_packet_t(header=AQL_HDR | (hsa.HSA_PACKET_TYPE_KERNEL_DISPATCH << hsa.HSA_PACKET_HEADER_TYPE),
      setup=3<<hsa.HSA_KERNEL_DISPATCH_PACKET_SETUP_DIMENSIONS, private_segment_size=prg.private_segment_size,
      group_segment_size=prg.group_segment_size, kernel_object=prg.aql_prog_addr, kernarg_address=args_state.buf.va_addr))
@@ -549,9 +575,7 @@ class KFDIface:
    if KFDIface.kfd is None:
      KFDIface.kfd = FileIOInterface("/dev/kfd", os.O_RDWR)
      gpus = [g for g in FileIOInterface(kfd_topo_path).listdir() if self._is_usable_gpu(FileIOInterface(f"{kfd_topo_path}/{g}/gpu_id"))]
-      gpus = sorted(gpus, key=lambda x: int(x.split('/')[-1]))
-      visible_devices = [int(x) for x in (getenv('VISIBLE_DEVICES', getenv('HIP_VISIBLE_DEVICES', ''))).split(',') if x.strip()]
-      KFDIface.gpus = [gpus[x] for x in visible_devices] if visible_devices else gpus
+      KFDIface.gpus = hcq_filter_visible_devices(sorted(gpus, key=lambda x: int(x.split('/')[-1])))

    if device_id >= len(KFDIface.gpus): raise RuntimeError(f"No device found for {device_id}. Requesting more devices than the system has?")

@@ -561,8 +585,6 @@ class KFDIface:
    id2ip = {am.GC_HWID: am.GC_HWIP, am.SDMA0_HWID: am.SDMA0_HWIP, am.NBIF_HWID: am.NBIF_HWIP}
    ip_hw = [(id2ip[int(hwid)], int(hwid)) for hwid in FileIOInterface(ip_base).listdir() if hwid.isnumeric() and int(hwid) in id2ip]
    self.ip_versions = {ip:tuple(int(FileIOInterface(f'{ip_base}/{hw}/0/{part}').read()) for part in ['major','minor','revision']) for ip,hw in ip_hw}
-    self.ip_offsets = {ip:{int(i):tuple(int(x, 16) for x in FileIOInterface(f'{ip_base}/{hw}/{i}/base_addr').read().splitlines())
-      for i in FileIOInterface(f'{ip_base}/{hw}').listdir()} for ip,hw in ip_hw }
    self.drm_fd = FileIOInterface(f"/dev/dri/renderD{self.props['drm_render_minor']}", os.O_RDWR)

    self.kfd_ver = ((ver_st:=kfd.AMDKFD_IOC_GET_VERSION(KFDIface.kfd)).major_version, ver_st.minor_version)
@@ -673,12 +695,12 @@ class PCIIface(PCIIfaceBase):
  def __init__(self, dev, dev_id):
    super().__init__(dev, dev_id, vendor=0x1002, devices=[0x744c, 0x7480, 0x7550, 0x7590], bars=[0, 2, 5], vram_bar=0,
      va_start=AMMemoryManager.va_allocator.base, va_size=AMMemoryManager.va_allocator.size)
-    self._setup_adev(self.pci_dev.pcibus, self.pci_dev.map_bar(0), self.pci_dev.map_bar(2, fmt='Q'), self.pci_dev.map_bar(5, fmt='I'))
+    self._setup_adev(self.pci_dev)
    self.pci_dev.write_config(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)

-  def _setup_adev(self, name, vram:MMIOInterface, doorbell:MMIOInterface, mmio:MMIOInterface, dma_regions:list[tuple[int, MMIOInterface]]|None=None):
-    self.dev_impl:AMDev = AMDev(name, vram, doorbell, mmio, dma_regions)
-    self.ip_offsets, self.ip_versions = self.dev_impl.regs_offset, self.dev_impl.ip_ver
+  def _setup_adev(self, pci_dev:PCIDevice, dma_regions:list[tuple[int, MMIOInterface]]|None=None):
+    self.dev_impl:AMDev = AMDev(pci_dev, dma_regions)
+    self.ip_versions = self.dev_impl.ip_ver

    gfxver = int(f"{self.dev_impl.ip_ver[am.GC_HWIP][0]:02d}{self.dev_impl.ip_ver[am.GC_HWIP][1]:02d}{self.dev_impl.ip_ver[am.GC_HWIP][2]:02d}")
    array_count = self.dev_impl.gc_info.gc_num_sa_per_se * self.dev_impl.gc_info.gc_num_se
@@ -715,34 +737,29 @@ class PCIIface(PCIIfaceBase):

 class USBIface(PCIIface):
  def __init__(self, dev, dev_id): # pylint: disable=super-init-not-called
-    self.dev = dev
-    self.usb = ASM24Controller()
-    self.bars = setup_pci_bars(self.usb, gpu_bus=4, mem_base=0x10000000, pref_mem_base=(32 << 30))
-
-    self._setup_adev(f"usb:{dev_id}", USBMMIOInterface(self.usb, *self.bars[0], fmt='B'), USBMMIOInterface(self.usb, *self.bars[2], fmt='Q'),
-      USBMMIOInterface(self.usb, *self.bars[5], fmt='I'), dma_regions=[(0x200000, self._dma_view(0xf000, 0x80000))])
-    self.usb._pci_cacheable += [self.bars[2]] # doorbell region is cacheable
+    self.dev, self.pci_dev = dev, USBPCIDevice(f"usb:{dev_id}", bars=[0, 2, 5])
+    self._setup_adev(self.pci_dev, dma_regions=[(0x200000, self.pci_dev.dma_view(0xf000, 0x80000))])
+    self.pci_dev.usb._pci_cacheable += [(self.pci_dev.bar_info[2].addr, self.pci_dev.bar_info[2].size)] # doorbell region is cacheable

    # special regions
    self.copy_bufs = [self._dma_region(ctrl_addr=0xf000, sys_addr=0x200000, size=0x80000)]
    self.sys_buf, self.sys_next_off = self._dma_region(ctrl_addr=0xa000, sys_addr=0x820000, size=0x1000), 0x800

-  def _dma_view(self, ctrl_addr, size): return USBMMIOInterface(self.usb, ctrl_addr, size, fmt='B', pcimem=False)
  def _dma_region(self, ctrl_addr, sys_addr, size):
    region = self.dev_impl.mm.map_range(vaddr:=self.dev_impl.mm.alloc_vaddr(size=size), size, [(sys_addr, size)], system=True, uncached=True)
-    return HCQBuffer(vaddr, size, meta=PCIAllocationMeta(region, has_cpu_mapping=False), view=self._dma_view(ctrl_addr, size), owner=self.dev)
+    return HCQBuffer(vaddr, size, meta=PCIAllocationMeta(region, has_cpu_mapping=False), view=self.pci_dev.dma_view(ctrl_addr, size), owner=self.dev)

  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, **kwargs) -> HCQBuffer:
    if (host or (uncached and cpu_access)) and self.sys_next_off + size < self.sys_buf.size:
      self.sys_next_off += size
      return self.sys_buf.offset(self.sys_next_off - size, size)

-    am_mapping = self.dev_impl.mm.valloc(size:=round_up(size, 4 << 10), uncached=uncached, contiguous=cpu_access)
-    return HCQBuffer(am_mapping.va_addr, size, meta=PCIAllocationMeta(am_mapping, has_cpu_mapping=False),
-      view=USBMMIOInterface(self.usb, self.bars[0][0] + am_mapping.paddrs[0][0], size, fmt='B') if cpu_access else None, owner=self.dev)
+    mapping = self.dev_impl.mm.valloc(size:=round_up(size, 4 << 10), uncached=uncached, contiguous=cpu_access)
+    barview = self.pci_dev.map_bar(bar=0, off=mapping.paddrs[0][0], size=mapping.size) if cpu_access else None
+    return HCQBuffer(mapping.va_addr, size, meta=PCIAllocationMeta(mapping, has_cpu_mapping=False), view=barview, owner=self.dev)

  def create_queue(self, queue_type, ring, gart, rptr, wptr, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0, xcc_id=0):
-    if queue_type == kfd.KFD_IOC_QUEUE_TYPE_COMPUTE: self.usb._pci_cacheable += [(ring.cpu_view().addr, ring.size)]
+    if queue_type == kfd.KFD_IOC_QUEUE_TYPE_COMPUTE: self.pci_dev.usb._pci_cacheable += [(ring.cpu_view().addr, ring.size)]
    return super().create_queue(queue_type, ring, gart, rptr, wptr, eop_buffer, cwsr_buffer, ctl_stack_size, ctx_save_restore_size, xcc_id)

  def sleep(self, timeout): pass
@@ -778,14 +795,15 @@ class AMDDevice(HCQCompiled):
    debug_memory_size = round_up((self.max_cu_id + 1 if self.target >= (10,1,0) else 1) * (self.max_wave_id + 1) * 32, 64)
    if self.target[0] == 10: ctl_stack_size = min(ctl_stack_size, 0x7000)

+    self.ip_off = import_ip_offsets(self.target)
    self.soc = import_soc(self.target)
    self.pm4 = importlib.import_module(f"tinygrad.runtime.autogen.am.pm4_{'nv' if self.target[0] >= 10 else 'soc15'}")
    self.sdma = import_module('sdma', min(self.iface.ip_versions[am.SDMA0_HWIP], (6, 0, 0)))
-    self.gc = AMDIP('gc', self.iface.ip_versions[am.GC_HWIP], self.iface.ip_offsets[am.GC_HWIP])
+    self.gc = AMDIP('gc', self.iface.ip_versions[am.GC_HWIP],
+                    bases={i: tuple(getattr(self.ip_off, f'GC_BASE__INST{i}_SEG{s}', 0) for s in range(6)) for i in range(6)})

-    nbio_name = 'nbio' if self.target[0] < 12 else 'nbif'
-    nbio_pad = (0,) if self.target[0] == 9 else ()
-    self.nbio = AMDIP(nbio_name, self.iface.ip_versions[am.NBIF_HWIP], {i:nbio_pad+x for i,x in self.iface.ip_offsets[am.NBIF_HWIP].items()})
+    self.nbio = AMDIP('nbio' if self.target[0] < 12 else 'nbif', self.iface.ip_versions[am.NBIF_HWIP],
+                      bases={i: tuple(getattr(self.ip_off, f'NBIO_BASE__INST{i}_SEG{s}', 0) for s in range(9)) for i in range(6)})

    self.is_aql = getenv("AMD_AQL", int(self.xccs > 1))
    if self.is_aql:
@@ -812,16 +830,16 @@ class AMDDevice(HCQCompiled):
    self._ensure_has_local_memory(128) # set default scratch size to 128 bytes per thread

    # SQTT is disabled by default because of runtime overhead and big file sizes (~200mb to Tensor.full() two 4096x4096 tensors and matmul them)
-    self.sqtt_enabled = PROFILE and bool(getenv("SQTT", 0))
+    self.sqtt_enabled = PROFILE and SQTT > 0
    if self.sqtt_enabled:
-      if self.target[0] < 11: raise RuntimeError(f'SQ Thread Tracing is not supported on gc:{self.target}')
+      if self.target[0] not in {9, 11, 12}: raise RuntimeError(f'SQ Thread Tracing is not supported on gc:{self.target}')
      if not self.is_am() and (ppfeaturemask:=int(FileIOInterface('/sys/module/amdgpu/parameters/ppfeaturemask', os.O_RDONLY).read(), 16))&0x8000:
        raise RuntimeError("SQTT can't be enabled because of hardware bug, to workaround either use AMD_IFACE=PCI or add "
                           f"ppfeaturemask={(ppfeaturemask&~0x8000):#x} (current {ppfeaturemask=:#x} & ~PP_GFXOFF_MASK) to amdgpu module parameters\n"
                           "For more information read https://github.com/tinygrad/tinygrad/blob/master/extra/sqtt/README.md")
      SQTT_BUFFER_SIZE = getenv("SQTT_BUFFER_SIZE", 256) # in mb, per shader engine
-      self.sqtt_buffers = [self.allocator.alloc(SQTT_BUFFER_SIZE*1024*1024, BufferSpec(nolru=True)) for _ in range(self.se_cnt)]
-      self.sqtt_itrace_se_mask = getenv("SQTT_ITRACE_SE_MASK", 2) # -1 enable all, 0 disable all, >0 bitmask for where to enable instruction tracing
+      self.sqtt_buffers = [self.allocator.alloc(SQTT_BUFFER_SIZE*1024*1024, BufferSpec(nolru=True, uncached=True)) for _ in range(self.se_cnt)]
+      self.sqtt_itrace_se_mask = getenv("SQTT_ITRACE_SE_MASK", -1 if SQTT >= 2 else (1 << 1)) # se bitmask: -1 enable all, 0 disable all
      self.sqtt_next_cmd_id = itertools.count(0)
      cast(AMDComputeQueue, self.hw_compute_queue_t()).sqtt_start(self.sqtt_buffers, self.sqtt_itrace_se_mask).submit(self)

@@ -881,12 +899,13 @@ class AMDDevice(HCQCompiled):
      self.synchronize()
      if DEBUG >= 2: print(f'{self.device}: Saving SQTT in profile...')
      for i,buf0 in enumerate(self.sqtt_buffers):
-        wptr = ((wptrs_buf.cpu_view().view(fmt='I')[i] & 0x1FFFFFFF) - (((buf0.va_addr//32) & 0x1FFFFFFF) if self.target < (12,0,0) else 0)) * 32
+        wptr = ((wptrs_buf.cpu_view().view(fmt='I')[i] & 0x1FFFFFFF) - (((buf0.va_addr//32) & 0x1FFFFFFF) if self.target[0] == 11 else 0)) * 32
        if DEBUG >= 2: print(f'\t{self.device}: SE {i} blob size {wptr:#x}')
        assert wptr >= 0 and wptr <= buf0.size, f"{wptr} > {buf0.size}, should never happen"
        # When sqtt buffer overflows, wptr stops at the last dword
        if wptr >= buf0.size - 32:
          print(colored(f"{self.device}: Warning: SQTT buffer is full (SE {i})! Increase SQTT buffer with SQTT_BUFFER_SIZE=X (in MB)", "yellow"))
        self.allocator._copyout(sqtt_buf:=memoryview(bytearray(wptr)), buf0)
+        if self.target[0] == 9: sqtt_buf = memoryview(struct.pack('<Q', 0x11 | (4 << 13) | (0xf << 16) | (i << 24)) + sqtt_buf)
        Compiled.profile_events += [ProfileSQTTEvent(self.device, i, self.iface.props, bytes(sqtt_buf), bool((self.sqtt_itrace_se_mask >> i) & 0b1))]
    super()._at_profile_finalize()
@@ -5,7 +5,7 @@ from tinygrad.device import Compiled, BufferSpec, LRUAllocator, CompilerPairT
 from tinygrad.renderer.cstyle import CUDARenderer
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.runtime.autogen import cuda
-from tinygrad.runtime.support.compiler_cuda import pretty_ptx, CUDACompiler, PTXCompiler
+from tinygrad.runtime.support.compiler_cuda import pretty_ptx, CUDACompiler, PTXCompiler, NVCCCompiler
 if getenv("IOCTL"): import extra.nv_gpu_driver.nv_ioctl  # noqa: F401  # pylint: disable=unused-import
 if MOCKGPU:=getenv("MOCKGPU"): from test.mockgpu.cuda import cuda # type: ignore # pylint: disable=reimported

@@ -118,7 +118,8 @@ class CUDADevice(Compiled):

    from tinygrad.runtime.graph.cuda import CUDAGraph
    compilers:list[CompilerPairT] = [(functools.partial(CUDARenderer, self.arch), functools.partial(CUDACompiler, self.arch)),
-                                     (functools.partial(PTXRenderer, self.arch), functools.partial(PTXCompiler, self.arch))]
+                                     (functools.partial(PTXRenderer, self.arch), functools.partial(PTXCompiler, self.arch)),
+                                     (functools.partial(CUDARenderer, self.arch), functools.partial(NVCCCompiler, self.arch))]
    super().__init__(device, CUDAAllocator(self), compilers, functools.partial(CUDAProgram, self), None if MOCKGPU else CUDAGraph)

  def synchronize(self):
@@ -4,7 +4,7 @@ assert sys.platform != 'win32'
 from typing import cast, ClassVar
 from dataclasses import dataclass
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, HWQueue, CLikeArgsState, HCQProgram, HCQSignal, BumpAllocator
-from tinygrad.runtime.support.hcq import MMIOInterface, FileIOInterface, MOCKGPU
+from tinygrad.runtime.support.hcq import MMIOInterface, FileIOInterface, MOCKGPU, hcq_filter_visible_devices
 from tinygrad.uop.ops import sint
 from tinygrad.device import BufferSpec, CompilerPairT
 from tinygrad.helpers import getenv, mv_address, round_up, data64, data64_le, prod, OSX, to_mv, hi32, lo32, suppress_finalizing
@@ -321,8 +321,7 @@ class NVKIface:
      with contextlib.suppress(RuntimeError): uvm.mm_initialize(self.fd_uvm_2, uvmFd=self.fd_uvm.fd) # this error is okay, CUDA hits it too

      nv_iowr(NVKIface.fd_ctl, nv_gpu.NV_ESC_CARD_INFO, gpus_info:=(nv_gpu.nv_ioctl_card_info_t*64)())
-      visible_devices = [int(x) for x in (getenv('VISIBLE_DEVICES', getenv('CUDA_VISIBLE_DEVICES', ''))).split(',') if x.strip()]
-      NVKIface.gpus_info = [gpus_info[x] for x in visible_devices] if visible_devices else gpus_info
+      NVKIface.gpus_info = hcq_filter_visible_devices(gpus_info)

    self.dev, self.device_id = dev, device_id
    if self.device_id >= len(NVKIface.gpus_info) or not NVKIface.gpus_info[self.device_id].valid:
@@ -462,9 +461,7 @@ class PCIIface(PCIIfaceBase):
    if not OSX: System.reserve_hugepages(64)

    self.pci_dev.write_config(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)
-    self.dev_impl:NVDev = NVDev(self.pci_dev.pcibus, self.pci_dev.map_bar(0, fmt='I'), self.pci_dev.map_bar(1),
-      self.pci_dev.read_config(pci.PCI_VENDOR_ID, 4), self.pci_dev.read_config(pci.PCI_SUBSYSTEM_VENDOR_ID, 4),
-      self.pci_dev.read_config(pci.PCI_REVISION_ID, 1), self.pci_dev.bar_info)
+    self.dev_impl:NVDev = NVDev(self.pci_dev)
    self.root, self.gpu_instance = 0xc1000000, 0
    self.rm_alloc(0, nv_gpu.NV01_ROOT, nv_gpu.NV0000_ALLOC_PARAMETERS())

@@ -3,7 +3,7 @@
 # works to test the tensor cores, and all the uops in general
 # this is the (living) definition of uops
 from typing import Any, TYPE_CHECKING, cast
-import pickle, base64, itertools, time, struct, sys
+import pickle, base64, itertools, time, struct, sys, functools
 from tinygrad.dtype import DType, dtypes, ImageDType, PtrDType, truncate, float_to_bf16, float_to_fp8, fp8_to_float
 from tinygrad.helpers import all_same, getenv, flatten, get_single_element, EMULATE
 from tinygrad.device import Compiled, Compiler, Allocator
@@ -36,43 +36,54 @@ def _store(m, i, v, dtype: DType):
  if i < 0 or i >= len(m): raise IndexError(f"store out of bounds, size is {len(m)}, access is {i}, value is {v}")
  m[i] = to_storage_scalar(v, dtype)

+# here are the models for the WMMA instruction on the different hardware
+def generic_wmma_helper(inp, warp_size, WARP_THREADS, K, NUM_A, NUM_B, NUM_C, a_elem, b_elem, c_map):
+  for cc, tinp, num in zip(("A", "B", "C"), inp, (NUM_A, NUM_B, NUM_C)):
+    assert len(tinp) == num, f"{cc} must have {num} elements per thread, it has {len(tinp)}"
+    assert len(flatten(tinp)) == num * warp_size, f"WMMA must have {num * warp_size} total elements for {cc} in WMMA"
+  assert warp_size > 0 and warp_size % WARP_THREADS == 0, f"must have multiples of {WARP_THREADS} warp threads"
+  out = [inp[2][elem_idx][:] for elem_idx in range(NUM_C)]
+  for goff in range(0, warp_size, WARP_THREADS):
+    for lane_id in range(WARP_THREADS):
+      for elem_idx in range(NUM_C): # calculate new muls and add to acc
+        (c_i, c_j) = c_map(lane_id, elem_idx)
+        out[elem_idx][goff+lane_id] += sum(a_elem(inp[0], _k, c_j, goff) * b_elem(inp[1], c_i, _k, goff) for _k in range(K))
+  return out
+
 class PythonProgram:
  def __init__(self, name:str, lib:bytes):
-    self.uops: list[tuple[Ops, DType|None, list[int], Any]] = pickle.loads(lib)
+    self.uops: list[tuple[Ops, DType, list[int], Any]] = pickle.loads(lib)
  def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False):
    st = time.perf_counter()
    warp = list(itertools.product(*[range(x) for x in local_size[::-1]]))
    warp_size = len(warp)
+    void_ops = {Ops.END, Ops.BARRIER, Ops.IF, Ops.ENDIF, Ops.SINK, Ops.NOOP, Ops.GROUP, Ops.STORE}
+    loop_ends: dict[int, int] = {srcs[1]:i for i, (uop, _, srcs, _) in enumerate(self.uops) if uop == Ops.END}
    for idxs in itertools.product(*[range(x) for x in global_size[::-1]]):
-      ul: dict[int, Any] = {}
-      dl: dict[int, DType] = {}
+      values: dict[int, Any] = {}
      pbufs: list[memoryview] = list(bufs)
      pvals: list[int] = list(vals)
      i = 0
-      loop_ends: dict[int, int] = {}
      while i < len(self.uops):
-        uop, dtype, idp, arg = self.uops[i]
-        void_ops = {Ops.END, Ops.BARRIER, Ops.IF, Ops.ENDIF, Ops.SINK, Ops.NOOP, Ops.STORE}
-        inp = [ul[v] for v in idp if self.uops[v][0] not in void_ops]
-        dtp = [dl[v] for v in idp if self.uops[v][0] not in void_ops]
-        if getenv("TRACE"): print(i, uop, dtype, arg, inp, dtp)
+        uop, dtype, srcs, arg = self.uops[i]
+        src_values = [values[v] for v in srcs if self.uops[v][0] not in void_ops]
+        src_dtypes = [self.uops[v][1] for v in srcs if self.uops[v][0] not in void_ops]
+        if getenv("TRACE"): print(i, uop, dtype, arg, src_values, src_dtypes)
        if uop is Ops.END:
-          loop_ends[idp[0]] = i
-          i = idp[0]
+          i = srcs[1]
          continue
-        if uop in (Ops.BARRIER, Ops.IF, Ops.ENDIF, Ops.SINK, Ops.NOOP):
+        if uop in (Ops.BARRIER, Ops.IF, Ops.ENDIF, Ops.SINK, Ops.NOOP, Ops.GROUP):
          # in the python emulator, the warp is always in sync
          i += 1
          continue
        assert dtype is not None, f"{uop} is missing a dtype"
-        dl[i] = dtype
        if uop is Ops.STORE:
-          for j,val in enumerate(inp[1] if dtp[1].count > 1 else [inp[1]]):
-            for (m,o,g),v in zip(inp[0], val):
-              if g: _store(m, o+j, v, dtp[1].scalar())
+          for j,val in enumerate(src_values[1] if src_dtypes[1].count > 1 else [src_values[1]]):
+            for (m,o,g),v in zip(src_values[0], val):
+              if g: _store(m, o+j, v, src_dtypes[1].scalar())
          i += 1
          continue
-        if uop is Ops.AFTER: ul[i] = inp[0]
+        if uop is Ops.AFTER: values[i] = src_values[0]
        elif uop in {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG}:
          assert isinstance(dtype, PtrDType), dtype
          storage_fmt = storage_fmt_for_dtype(dtype.base.scalar())
@@ -80,85 +91,73 @@ class PythonProgram:
          if TYPE_CHECKING or sys.version_info < (3, 12): assert storage_fmt != "e"
          if uop is Ops.DEFINE_REG:
            # REGs are per thread
-            ul[i] = [memoryview(bytearray(dtype.size*dtype.itemsize)).cast(storage_fmt) for _ in range(warp_size)]
+            values[i] = [memoryview(bytearray(dtype.size*dtype.itemsize)).cast(storage_fmt) for _ in range(warp_size)]
          else:
            buf = memoryview(bytearray(dtype.size*dtype.itemsize)) if uop is not Ops.DEFINE_GLOBAL else pbufs.pop(0)
-            ul[i] = [buf.cast(storage_fmt)] * warp_size
+            values[i] = [buf.cast(storage_fmt)] * warp_size
        elif uop is Ops.DEFINE_VAR:
-          ul[i] = [pvals.pop(0)] * warp_size
+          values[i] = [pvals.pop(0)] * warp_size
        elif uop is Ops.SPECIAL:
-          if arg[0] == 'g': ul[i] = [idxs[2-int(arg[-1])]] * warp_size
-          elif arg[0] == 'l': ul[i] = [x[2-int(arg[-1])] for x in warp]
-        elif uop is Ops.CONST: ul[i] = [arg] * warp_size
+          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.INDEX:
          ret:list = []
-          if isinstance(dtp[0], ImageDType):
-            for m,ox,oy in zip(inp[0], inp[1][0], inp[1][1]):
-              if ox < 0 or ox >= dtp[0].shape[1] or oy < 0 or oy >= dtp[0].shape[0]: ret.append((m, None))
-              else: ret.append((m, ox*4 + oy*dtp[0].shape[1]*4))
+          if isinstance(src_dtypes[0], ImageDType):
+            for m,ox,oy in zip(src_values[0], src_values[1][0], src_values[1][1]):
+              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:
-            for m,o in zip(inp[0], inp[1]): ret.append((m,o))
-          ul[i] = [(m,o,g) for (m,o),g in zip(ret, inp[2] if len(inp) == 3 else [True]*len(ret))] # set the gate last
+            for m,o in zip(src_values[0], src_values[1]): ret.append((m,o))
+          values[i] = [(m,o,g) for (m,o),g in zip(ret, src_values[2] if len(src_values) == 3 else [True]*len(ret))] # set the gate last
        elif uop is Ops.CAST and isinstance(dtype, PtrDType):
-          ul[i] = inp[0]
+          values[i] = src_values[0]
        elif uop is Ops.RANGE:
-          if i not in ul: ul[i] = [0] * warp_size
+          if i not in values: values[i] = [0] * warp_size
          else:
-            for j in range(len(ul[i])):
-              ul[i][j] += 1
-            if ul[i][0] == inp[0][0]:
-              del ul[i]
-              i = loop_ends[i] + 1
-              continue
-        elif uop is Ops.VECTORIZE: ul[i] = inp
+            for j in range(len(values[i])):
+              values[i][j] += 1
+          if values[i][0] == src_values[0][0]:
+            del values[i]
+            i = loop_ends[i] + 1
+            continue
+        elif uop is Ops.VECTORIZE: values[i] = src_values
        elif uop is Ops.BITCAST:
-          packed = struct.pack(str(warp_size) + storage_fmt_for_dtype(dtp[0].scalar()), *[to_storage_scalar(x, dtp[0].scalar()) for x in inp[0]])
-          ul[i] = list(struct.unpack(str(warp_size) +  storage_fmt_for_dtype(dtype.scalar()), packed))
-          ul[i] = [from_storage_scalar(x, dtype.scalar()) for x in ul[i]]
+          packed = struct.pack(str(warp_size) + storage_fmt_for_dtype(src_dtypes[0].scalar()),
+                               *[to_storage_scalar(x, src_dtypes[0].scalar()) for x in src_values[0]])
+          values[i] = list(struct.unpack(str(warp_size) +  storage_fmt_for_dtype(dtype.scalar()), packed))
+          values[i] = [from_storage_scalar(x, dtype.scalar()) for x in values[i]]
        elif uop is Ops.CAST:
-          ul[i] = [truncate.get(dtype, lambda dt: dt)(dtypes.as_const(x, dtype)) for x in inp[0]]
+          values[i] = [truncate.get(dtype, lambda dt: dt)(dtypes.as_const(x, dtype)) for x in src_values[0]]
        elif uop is Ops.LOAD:
          if dtype.count > 1:
-            ul[i] = [load([inp[i][j] if i != 0 and dtp[i].count > 1 else inp[i] for i in range(len(inp))], j, dtype.scalar()) \
-              for j in range(dtype.count)]
+            values[i] = [load([src_values[i][j] if i != 0 and src_dtypes[i].count > 1 else src_values[i] \
+                               for i in range(len(src_values))], j, dtype.scalar()) for j in range(dtype.count)]
          else:
-            ul[i] = load(inp, 0, dtype)
-        elif uop is Ops.GEP: ul[i] = inp[0][get_single_element(arg)]
+            values[i] = load(src_values, 0, dtype)
+        elif uop is Ops.GEP: values[i] = src_values[0][get_single_element(arg)]
        elif uop is Ops.WMMA:
-          # here are the models for the WMMA instruction on the different hardware
-          def wmma_helper(WARP_THREADS, K, NUM_A, NUM_B, NUM_C, a_elem, b_elem, c_map):
-            for cc, tinp, num in zip(("A", "B", "C"), inp, (NUM_A, NUM_B, NUM_C)):
-              assert len(tinp) == num, f"{cc} must have {num} elements per thread, it has {len(tinp)}"
-              assert len(flatten(tinp)) == num * warp_size, f"WMMA must have {num * warp_size} total elements for {cc} in WMMA"
-            assert warp_size > 0 and warp_size % WARP_THREADS == 0, f"must have multiples of {WARP_THREADS} warp threads"
-            out = [inp[2][elem_idx][:] for elem_idx in range(NUM_C)]
-            for goff in range(0, warp_size, WARP_THREADS):
-              for lane_id in range(WARP_THREADS):
-                for elem_idx in range(NUM_C): # calculate new muls and add to acc
-                  (c_i, c_j) = c_map(lane_id, elem_idx)
-                  out[elem_idx][goff+lane_id] += sum(a_elem(inp[0], _k, c_j, goff) * b_elem(inp[1], c_i, _k, goff) for _k in range(K))
-            return out
-
-          first_src_dtype = self.uops[idp[0]][1]
+          first_src_dtype = self.uops[srcs[0]][1]
          assert isinstance(first_src_dtype, DType) # mypy
          dims, dtype_in, device, threads = arg[1], first_src_dtype.scalar(), arg[4], arg[5]
+          wmma_helper = functools.partial(generic_wmma_helper, src_values, warp_size)
          # TODO: refactor these to a shared TensorCoreLayout in kernel.py
          if device == "METAL":
            # A (2 elements on 32 threads): row major
            def a_b_elem(x, i, j, goff): return x[(i%2)][goff+(i//2)%2+(j%4)*2+(i//4)*8+(j//4)*16]
            # (i, j), C, D (2 elements on 32 threads): row major same as A/B
            def c_map(lane, elem): return (elem + ((lane%2)*2) + ((lane//8)%2)*4, ((lane//2)%4) + (lane//16)*4)
-            ul[i] = wmma_helper(32, 8, 2, 2, 2, a_b_elem, a_b_elem, c_map)
+            values[i] = wmma_helper(32, 8, 2, 2, 2, a_b_elem, a_b_elem, c_map)
          elif device == "AMD" and threads == 64:
-            def a_elem(x, k, row, goff): return x[k%4][goff + (k//4)*16 + row]
-            def b_elem(x, col, k, goff): return a_elem(x, k, col, goff) # pylint: disable=arguments-out-of-order
+            def a_elem(x, k, row, goff): return x[k%(dims[2]//4)][goff + (k//(dims[2]//4))*16 + row]
+            def b_elem(x, col, k, goff): return a_elem(x, k, col, goff)  # pylint: disable=arguments-out-of-order
            def c_map(lane, elem): return (lane%16, (lane//16)*4 + elem)
-            ul[i] = wmma_helper(64, 16, 4, 4, 4, a_elem, b_elem, c_map)
-          elif device == "AMD" and len(inp[0]) == 8: # RDNA4
+            values[i] = wmma_helper(64, dims[2], len(src_values[0]), len(src_values[1]), len(src_values[2]), a_elem, b_elem, c_map)
+          elif device == "AMD" and len(src_values[0]) == 8: # RDNA4
            def a_elem(x, k, row, goff): return x[k - [0, 4, 4, 8][k//4]][goff + row + [0, 16, 0, 16][k//4]]
            def b_elem(x, col, k, goff): return a_elem(x, k, col, goff)
            def c_map(lane, elem): return (lane%16, (lane//16)*8 + elem)
-            ul[i] = wmma_helper(32, 16, 8, 8, 8, a_elem, b_elem, c_map)
+            values[i] = wmma_helper(32, 16, 8, 8, 8, a_elem, b_elem, c_map)
          elif device == "AMD":
            # A (16 elements on 32 threads): col major, lane 16-32 == lane 0-15
            def a_elem(x, k, row, goff):
@@ -167,7 +166,7 @@ class PythonProgram:
            # B (16 elements on 32 threads): row major, lane 16-32 == lane 0-15
            def b_elem(x, col, k, goff): return a_elem(x, k, col, goff)  # pylint: disable=arguments-out-of-order
            def c_map(lane, elem): return (lane%16, lane//16+elem*2) # (i, j), C, D (8 elements on 32 threads): row major
-            ul[i] = wmma_helper(32, 16, 16, 16, 8, a_elem, b_elem, c_map)
+            values[i] = wmma_helper(32, 16, 16, 16, 8, a_elem, b_elem, c_map)
          elif device == "CUDA":
            # (col, row) given (lane, elem) for C & D (4 elements on 32 threads); shared by all tc shapes with M=16 N=8
            def c_map(lane, elem): return (elem%2 + (lane%4)*2, lane//4 + (elem//2)*8)
@@ -175,17 +174,22 @@ class PythonProgram:
            if dims == (8,16,16):
              def a_elem(x, k, row, goff): return x[k%2 + (row//8)*2 + (k//8)*4][goff + (k//2)%4 + (row%8)*4]
              def b_elem(x, col, k, goff): return x[k%2 + (k//8)*2][goff + (k//2)%4 + col*4]
-              ul[i] = wmma_helper(32, 16, 8, 4, 4, a_elem, b_elem, c_map)
+              values[i] = wmma_helper(32, 16, 8, 4, 4, a_elem, b_elem, c_map)
+
+            elif dims == (8,16,32):
+              def a_elem(x, k, row, goff): return x[k%4 + (row//8)*4 + (k//16)*8][goff + (k//4)%4 + (row%8)*4]
+              def b_elem(x, col, k, goff): return x[k%4 + (k//16)*4][goff + (k//4)%4  + col*4]
+              values[i] = wmma_helper(32, 32, 16, 8, 4, a_elem, b_elem, c_map)

            elif dims == (8,16,8) and dtype_in == dtypes.half:
              def a_elem(x, k, row, goff): return x[k%2 + (row//8)*2][goff + k//2 + (row%8)*4]
              def b_elem(x, col, k, goff): return x[k%2][goff + k//2 + col*4]
-              ul[i] = wmma_helper(32, 8, 4, 2, 4, a_elem, b_elem, c_map)
+              values[i] = wmma_helper(32, 8, 4, 2, 4, a_elem, b_elem, c_map)

            elif dims == (8,16,8) and dtype_in == dtypes.float:
              def a_elem(x, k, row, goff): return x[(k//4)*2 + row//8][goff + k%4 + (row%8)*4]
              def b_elem(x, col, k, goff): return x[k//4][goff + k%4 + col*4]
-              ul[i] = wmma_helper(32, 8, 4, 2, 4, a_elem, b_elem, c_map)
+              values[i] = wmma_helper(32, 8, 4, 2, 4, a_elem, b_elem, c_map)

            else: raise NotImplementedError(f"unimplemented tensor core {arg}")
          elif device == "INTEL":
@@ -195,17 +199,17 @@ class PythonProgram:
            def b_elem(x, col, k, goff): return x[k][goff+col]
            # C, D (8 elements on 8 threads)
            def c_map(lane, elem): return (lane, elem)
-            ul[i] = wmma_helper(8, 16, 16, 16, 8, a_elem, b_elem, c_map)
+            values[i] = wmma_helper(8, 16, 16, 16, 8, a_elem, b_elem, c_map)
          elif device == "CPU":
            def elem(x, col, row, _): return x[col+row][0] # k is always 0
-            def c_map(_, elem): return (elem%16, elem//16)
-            ul[i] = wmma_helper(1, 1, 16, 16, 256, elem, elem, c_map)
+            def c_map(lane, elem): return (elem%16, elem//16)
+            values[i] = wmma_helper(1, 1, 16, 16, 256, elem, elem, c_map)
          else: raise NotImplementedError(f"unimplemented tensor core {arg}")
        elif uop in GroupOp.ALU:
-          assert all_same([len(x) for x in inp]), f"{[len(x) for x in inp]} doesn't match on {uop}"
-          assert all_same([dtype] + dtp) or uop in {*GroupOp.Comparison, Ops.WHERE}, f"dtype mismatch on {uop}"
-          ul[i] = [exec_alu(uop, dtype, p) for p in zip(*inp)]
-        assert i in ul, (uop, dtype, idp, arg)
+          assert all_same([len(x) for x in src_values]), f"{[len(x) for x in src_values]} doesn't match on {uop}"
+          assert all_same([dtype] + src_dtypes) or uop in {*GroupOp.Comparison, Ops.WHERE}, f"dtype mismatch on {uop}"
+          values[i] = [exec_alu(uop, dtype, p) for p in zip(*src_values)]
+        assert i in values, (uop, dtype, srcs, arg)
        i += 1
    return time.perf_counter() - st

@@ -216,10 +220,11 @@ class PythonRenderer(Renderer):
    match cast(str, EMULATE.value):
      case "METAL": self.device, self.tensor_cores = "METAL", tc.metal
      case "AMD": self.device, self.tensor_cores = "AMD", tc.amd_rdna3
-      case "AMD_MFMA": self.device, self.tensor_cores = "AMD", tc.amd_cdna
+      case "AMD_MFMA": self.device, self.tensor_cores = "AMD", tc.amd_cdna4
      case "AMD_RDNA4": self.device, self.tensor_cores = "AMD", tc.amd_rdna4
      case "CUDA": self.device, self.tensor_cores = "CUDA", tc.cuda_sm80
      case "CUDA_SM75": self.device, self.tensor_cores = "CUDA", tc.cuda_sm75
+      case "CUDA_SM89": self.device, self.tensor_cores = "CUDA", tc.cuda_sm89
      case "INTEL": self.device, self.suffix, self.tensor_cores = "INTEL", "INTEL", tc.intel
      case "AMX": self.device, self.tensor_cores = "CPU", tc.amx
      case "": pass
@@ -1,4 +1,4 @@
-import socket, uuid, json, asyncio, threading
+import socket, json, asyncio, threading
 from contextlib import asynccontextmanager
 from tinygrad.device import Compiled, Allocator
 from tinygrad.helpers import DEBUG, getenv
@@ -32,9 +32,6 @@ class TinyFSDevice(Compiled):
    self.conn_pools: dict[str, asyncio.Queue] = {}
    self.conn_pools_lock = asyncio.Lock()

-    # current request
-    self.request_id = uuid.UUID(int=0)
-
  def finalize(self):
    self.sfile.close()

@@ -74,9 +71,10 @@ class TinyFSDevice(Compiled):
      await self.conn_pools[loc].put((reader, writer))

 class TinyFSBuffer:
-  def __init__(self, device:TinyFSDevice, size:int, offset=0, copyout_queue=None):
+  def __init__(self, device:TinyFSDevice, size:int, offset=0, copyout_queue=None, hash_buf=None):
    self.device, self.size, self.offset = device, size, offset
    self.copyout_queue = copyout_queue or []
+    self.hash_buf = hash_buf or bytearray()
  def __repr__(self): return f"<TinyFSBuffer size={self.size} offset={self.offset}>"

 class TinyFSAllocator(Allocator[TinyFSDevice]):
@@ -87,40 +85,33 @@ class TinyFSAllocator(Allocator[TinyFSDevice]):
    if DEBUG >= 2: print(f"Copying in {dest.size} bytes to TINYFS:{dest.device.op}")
    self.dev.sfile.write(f"{dest.device.op}_IN {dest.size}\r\n".encode())

-    if dest.device.op == "STORE":
-      self.dev.sfile.flush()
-      self.dev.request_id = uuid.UUID(bytes=self.dev.sfile.read(16))
-      if DEBUG >= 2: print(f"Request ID: {self.dev.request_id}")
-
    self.dev.sfile.write(src)
    self.dev.sfile.flush()

    if dest.device.op == "LOAD":
      locs = self.dev.sfile.readline()
-      locs = json.loads(locs)
-
-      dest.copyout_queue = []
-      for i, loc in enumerate(locs):
-        dest.copyout_queue.append((i, loc, src[i*16:(i+1)*16].tobytes()))
+      dest.copyout_queue = json.loads(locs)
+      dest.hash_buf[:] = src.tobytes()
+    elif dest.device.op == "STORE":
+      expected_hashes = dest.size // Tensor.CHUNK_SIZE
+      dest.hash_buf = bytearray(expected_hashes * 16)
+      self.dev.sfile.readinto(dest.hash_buf)

  def _copyout(self, dest:memoryview, src:TinyFSBuffer):
    if DEBUG >= 2: print(f"Copying out {src.size} bytes from TINYFS:{src.device.op}")
    if src.device.op == "LOAD":
      asyncio.run_coroutine_threadsafe(self._copyout_async(dest, src), src.device.loop).result()
-    else:
-      self.dev.sfile.write(f"{src.device.op}_OUT {src.size} {self.dev.request_id}\r\n".encode())
-      self.dev.sfile.flush()
-      self.dev.sfile.readinto(dest)
+    elif src.device.op == "STORE":
+      dest[:] = src.hash_buf

  async def _copyout_async(self, dest:memoryview, src:TinyFSBuffer):
-    async def _worker(item):
-      i, loc, h = item
+    async def _worker(i, loc):
      async with self.dev.connection(loc) as (reader, writer):
        ptr = i * Tensor.CHUNK_SIZE
        size = min(len(dest[ptr:ptr+Tensor.CHUNK_SIZE]), Tensor.CHUNK_SIZE)

        writer.write(f"CHUNK_OUT {size}\r\n".encode())
-        writer.write(h)
+        writer.write(src.hash_buf[i*16:(i+1)*16])
        await writer.drain()

        chunk = await reader.readexactly(size)
@@ -129,8 +120,8 @@ class TinyFSAllocator(Allocator[TinyFSDevice]):
        view[:] = chunk
        del view

-    workers = [asyncio.create_task(_worker(item)) for item in src.copyout_queue]
+    workers = [asyncio.create_task(_worker(i, loc)) for i, loc in enumerate(src.copyout_queue)]
    await asyncio.gather(*workers)

  def _offset(self, buf:TinyFSBuffer, size:int, offset:int):
-    return TinyFSBuffer(buf.device, size, offset, buf.copyout_queue)
+    return TinyFSBuffer(buf.device, size, offset, buf.copyout_queue, buf.hash_buf)
@@ -5,7 +5,7 @@ from tinygrad.runtime.autogen.am import am
 from tinygrad.runtime.support.hcq import MMIOInterface
 from tinygrad.runtime.support.amd import AMDReg, import_module, import_asic_regs
 from tinygrad.runtime.support.memory import TLSFAllocator, MemoryManager
-from tinygrad.runtime.support.system import System, PCIDevImplBase
+from tinygrad.runtime.support.system import System, PCIDevice, PCIDevImplBase
 from tinygrad.runtime.support.am.ip import AM_SOC, AM_GMC, AM_IH, AM_PSP, AM_SMU, AM_GFX, AM_SDMA

 AM_DEBUG = getenv("AM_DEBUG", 0)
@@ -118,8 +118,10 @@ class AMMemoryManager(MemoryManager):
 class AMDev(PCIDevImplBase):
  Version = 0xA0000006

-  def __init__(self, devfmt, vram:MMIOInterface, doorbell:MMIOInterface, mmio:MMIOInterface, dma_regions:list[tuple[int, MMIOInterface]]|None=None):
-    self.devfmt, self.vram, self.doorbell64, self.mmio, self.dma_regions = devfmt, vram, doorbell, mmio, dma_regions
+  def __init__(self, pci_dev:PCIDevice, dma_regions:list[tuple[int, MMIOInterface]]|None=None):
+    self.pci_dev, self.devfmt, self.dma_regions = pci_dev, pci_dev.pcibus, dma_regions
+    self.vram, self.doorbell64, self.mmio = self.pci_dev.map_bar(0), self.pci_dev.map_bar(2, fmt='Q'), self.pci_dev.map_bar(5, fmt='I')
+
    self.lock_fd = System.flock_acquire(f"am_{self.devfmt}.lock")

    self._run_discovery()
@@ -242,6 +242,7 @@ class AM_GFX(AM_IP):
      cp_mqd_control=self.adev.regCP_MQD_CONTROL.encode(priv_state=1), cp_hqd_vmid=0, cp_hqd_aql_control=int(aql),
      cp_hqd_eop_base_addr_lo=lo32(eop_addr>>8), cp_hqd_eop_base_addr_hi=hi32(eop_addr>>8),
      cp_hqd_eop_control=self.adev.regCP_HQD_EOP_CONTROL.encode(eop_size=(eop_size//4).bit_length()-2))
+    for se in range(8): setattr(mqd_struct, f'compute_static_thread_mgmt_se{se}', 0xffffffff)

    # Copy mqd into memory
    self.adev.vram.view(mqd.paddrs[0][0], ctypes.sizeof(mqd_struct))[:] = memoryview(mqd_struct).cast('B')
@@ -1,9 +1,7 @@
 import functools, importlib, re, urllib
 from collections import defaultdict
 from dataclasses import dataclass
-from tinygrad.helpers import getbits, round_up, fetch
-from tinygrad.runtime.autogen import pci
-from tinygrad.runtime.support.usb import ASM24Controller
+from tinygrad.helpers import getbits, fetch

@dataclass
 class AMDReg:
@@ -37,7 +35,7 @@ def fixup_ip_version(ip:str, version:tuple[int, ...]) -> list[tuple[int, ...]]:
      if version[:len(ver)] == ver: return ovrd_ver
    return version

-  if ip in ['nbio', 'nbif']: version = _apply_ovrd({(3,3): (2,3,0)})
+  if ip in ['nbio', 'nbif']: version = _apply_ovrd({(3,3): (2,3,0), (7,3): (7,2,0)})
  elif ip in ['mp', 'smu']: version = _apply_ovrd({(14,0,3): (14,0,2)})
  elif ip in ['gc']: version = _apply_ovrd({(9,5,0): (9,4,3)})

@@ -49,7 +47,9 @@ def header_download(file, name=None, subdir="defines", url=None) -> str:

 def import_header(path:str, url=None):
  t = re.sub(r'//.*|/\*.*?\*/','', header_download(path, subdir="defines", url=url), flags=re.S)
-  return {k:int(v,0) for k,v in re.findall(r'\b([A-Za-z_]\w*)\s*=\s*(0x[0-9A-Fa-f]+|\d+)', t)}
+  # TODO: refactor when clang2py is replaced
+  return {k:int(v,0) for k,v in re.findall(r'\b([A-Za-z_]\w*)\s*=\s*(0x[0-9A-Fa-f]+|\d+)', t) + \
+                                re.findall(r'^\s*#\s*define\s+([A-Za-z_0-9]\w*)\s+(0x[0-9A-Fa-f]+|\d+)', t, re.M)}

 def import_module(name:str, version:tuple[int, ...], version_prefix:str=""):
  for ver in fixup_ip_version(name, version):
@@ -62,6 +62,8 @@ def import_soc(ip):
  url = "https://raw.githubusercontent.com/ROCm/rocm-systems/cccc350dc620e61ae2554978b62ab3532dc10bd9/projects"
  return type("SOC", (object,), import_header(f"aqlprofile/linux/{({9: 'vega10', 10: 'navi10', 11: 'soc21', 12: 'soc24'}[ip[0]])}_enum.h", url=url))

+def import_ip_offsets(ip): return type("IPOFF", (object,), import_header(f"include/{('sienna_cichlid' if ip[0] > 9 else 'vega20')}_ip_offset.h"))
+
 def import_asic_regs(prefix:str, version:tuple[int, ...], cls=AMDReg) -> dict[str, AMDReg]:
  def _split_name(name): return name[:(pos:=next((i for i,c in enumerate(name) if c.isupper()), len(name)))], name[pos:]
  def _extract_regs(txt):
@@ -89,53 +91,3 @@ def import_asic_regs(prefix:str, version:tuple[int, ...], cls=AMDReg) -> dict[st
    # NOTE: Some registers like regGFX_IMU_FUSESTRAP in gc_11_0_0 are missing base idx, just skip them
    return {reg:cls(name=reg, offset=off, segment=bases[reg], fields=fields[_split_name(reg)[1]]) for reg,off in offsets.items() if reg in bases}
  raise ImportError(f"Failed to load ASIC registers for {prefix.upper()} {'.'.join(map(str, version))}")
-
-def setup_pci_bars(usb:ASM24Controller, gpu_bus:int, mem_base:int, pref_mem_base:int) -> dict[int, tuple[int, int]]:
-  for bus in range(gpu_bus):
-    # All 3 values must be written at the same time.
-    buses = (0 << 0) | ((bus+1) << 8) | ((gpu_bus) << 16)
-    usb.pcie_cfg_req(pci.PCI_PRIMARY_BUS, bus=bus, dev=0, fn=0, value=buses, size=4)
-
-    usb.pcie_cfg_req(pci.PCI_MEMORY_BASE, bus=bus, dev=0, fn=0, value=(mem_base>>16) & 0xffff, size=2)
-    usb.pcie_cfg_req(pci.PCI_MEMORY_LIMIT, bus=bus, dev=0, fn=0, value=0xffff, size=2)
-    usb.pcie_cfg_req(pci.PCI_PREF_MEMORY_BASE, bus=bus, dev=0, fn=0, value=(pref_mem_base>>16) & 0xffff, size=2)
-    usb.pcie_cfg_req(pci.PCI_PREF_MEMORY_LIMIT, bus=bus, dev=0, fn=0, value=0xffff, size=2)
-    usb.pcie_cfg_req(pci.PCI_PREF_BASE_UPPER32,  bus=bus, dev=0, fn=0, value=pref_mem_base >> 32, size=4)
-    usb.pcie_cfg_req(pci.PCI_PREF_LIMIT_UPPER32, bus=bus, dev=0, fn=0, value=0xffffffff, size=4)
-
-    usb.pcie_cfg_req(pci.PCI_COMMAND, bus=bus, dev=0, fn=0, value=pci.PCI_COMMAND_IO | pci.PCI_COMMAND_MEMORY | pci.PCI_COMMAND_MASTER, size=1)
-
-  # resize bar 0
-  cap_ptr = 0x100
-  while cap_ptr:
-    if pci.PCI_EXT_CAP_ID(hdr:=usb.pcie_cfg_req(cap_ptr, bus=gpu_bus, dev=0, fn=0, size=4)) == pci.PCI_EXT_CAP_ID_REBAR:
-      cap = usb.pcie_cfg_req(cap_ptr + 0x04, bus=gpu_bus, dev=0, fn=0, size=4)
-      new_ctrl = (usb.pcie_cfg_req(cap_ptr + 0x08, bus=gpu_bus, dev=0, fn=0, size=4) & ~0x1F00) | ((int(cap >> 4).bit_length() - 1) << 8)
-      usb.pcie_cfg_req(cap_ptr + 0x08, bus=gpu_bus, dev=0, fn=0, value=new_ctrl, size=4)
-
-    cap_ptr = pci.PCI_EXT_CAP_NEXT(hdr)
-
-  mem_space_addr, bar_off, bars = [mem_base, pref_mem_base], 0, {}
-  while bar_off < 24:
-    cfg = usb.pcie_cfg_req(pci.PCI_BASE_ADDRESS_0 + bar_off, bus=gpu_bus, dev=0, fn=0, size=4)
-    bar_mem, bar_64 = bool(cfg & pci.PCI_BASE_ADDRESS_MEM_PREFETCH), cfg & pci.PCI_BASE_ADDRESS_MEM_TYPE_64
-
-    if (cfg & pci.PCI_BASE_ADDRESS_SPACE) == pci.PCI_BASE_ADDRESS_SPACE_MEMORY:
-      usb.pcie_cfg_req(pci.PCI_BASE_ADDRESS_0 + bar_off, bus=gpu_bus, dev=0, fn=0, value=0xffffffff, size=4)
-      lo = (usb.pcie_cfg_req(pci.PCI_BASE_ADDRESS_0 + bar_off, bus=gpu_bus, dev=0, fn=0, size=4) & 0xfffffff0)
-
-      if bar_64: usb.pcie_cfg_req(pci.PCI_BASE_ADDRESS_0 + bar_off + 4, bus=gpu_bus, dev=0, fn=0, value=0xffffffff, size=4)
-      hi = (usb.pcie_cfg_req(pci.PCI_BASE_ADDRESS_0 + bar_off + 4, bus=gpu_bus, dev=0, fn=0, size=4) if bar_64 else 0)
-
-      bar_size = ((~(((hi << 32) | lo) & ~0xf)) + 1) & (0xffffffffffffffff if bar_64 else 0xffffffff)
-
-      usb.pcie_cfg_req(pci.PCI_BASE_ADDRESS_0 + bar_off, bus=gpu_bus, dev=0, fn=0, value=mem_space_addr[bar_mem] & 0xffffffff, size=4)
-      if bar_64: usb.pcie_cfg_req(pci.PCI_BASE_ADDRESS_0 + bar_off + 4, bus=gpu_bus, dev=0, fn=0, value=mem_space_addr[bar_mem] >> 32, size=4)
-
-      bars[bar_off // 4] = (mem_space_addr[bar_mem], bar_size)
-      mem_space_addr[bar_mem] += round_up(bar_size, 2 << 20)
-
-    bar_off += 8 if bar_64 else 4
-
-  usb.pcie_cfg_req(pci.PCI_COMMAND, bus=gpu_bus, dev=0, fn=0, value=pci.PCI_COMMAND_IO | pci.PCI_COMMAND_MEMORY | pci.PCI_COMMAND_MASTER, size=1)
-  return bars
@@ -60,6 +60,19 @@ class NVCompiler(CUDACompiler):
  def __init__(self, arch:str): super().__init__(arch, cache_key="nv")
  def compile(self, src:str) -> bytes: return self._compile_program(src, nvrtc.nvrtcGetCUBIN, nvrtc.nvrtcGetCUBINSize)

+class NVCCCompiler(Compiler):
+  def __init__(self, arch:str, extra_options:list[str]=[]):
+    self.arch, self.extra_options = arch, extra_options
+    super().__init__(f"compile_nvcc_{self.arch}_{hashlib.sha256(' '.join(extra_options).encode()).hexdigest()[:8]}")
+  def compile(self, src:str) -> bytes:
+    with tempfile.NamedTemporaryFile(suffix=".cu") as srcf, tempfile.NamedTemporaryFile(suffix=".ptx") as libf:
+      srcf.write(src.encode())
+      srcf.flush()
+      subprocess.run(["nvcc", f"-arch={self.arch}", "-ptx", "-o", libf.name, srcf.name] + self.extra_options,
+                 check=True)
+      return libf.read()
+  def disassemble(self, lib:bytes): cuda_disassemble(lib, self.arch)
+
 class PTXCompiler(Compiler):
  def __init__(self, arch:str, cache_key="ptx"):
    self.arch = arch
@@ -57,6 +57,9 @@ if MOCKGPU:=getenv("MOCKGPU"): from test.mockgpu.mockgpu import MockFileIOInterf

 # **************** for HCQ Compatible Devices ****************

+def hcq_filter_visible_devices(dev):
+  return [dev[x] for x in ids] if (ids:=[int(x) for x in (getenv('HCQ_VISIBLE_DEVICES', '')).split(',') if x.strip()]) else dev
+
 SignalType = TypeVar('SignalType', bound='HCQSignal')
 HCQDeviceType = TypeVar('HCQDeviceType', bound='HCQCompiled')
 ProgramType = TypeVar('ProgramType', bound='HCQProgram')
@@ -5,7 +5,7 @@ from tinygrad.runtime.autogen.nv import nv
 from tinygrad.helpers import to_mv, lo32, hi32, DEBUG, round_up, round_down, mv_address, fetch, wait_cond
 from tinygrad.runtime.support.system import System
 from tinygrad.runtime.support.elf import elf_loader
-from tinygrad.runtime.autogen import nv_gpu
+from tinygrad.runtime.autogen import nv_gpu, pci

@dataclasses.dataclass(frozen=True)
 class GRBufDesc: size:int; virt:bool; phys:bool; local:bool=False # noqa: E702
@@ -524,9 +524,11 @@ class NV_GSP(NV_IP):
  def rpc_set_gsp_system_info(self):
    def bdf_as_int(s): return 0x000 if s.startswith("usb") else (int(s[5:7],16)<<8) | (int(s[8:10],16)<<3) | int(s[-1],16)

-    data = nv.GspSystemInfo(gpuPhysAddr=self.nvdev.bars[0][0], gpuPhysFbAddr=self.nvdev.bars[1][0], gpuPhysInstAddr=self.nvdev.bars[3][0],
+    pcidev = self.nvdev.pci_dev
+    data = nv.GspSystemInfo(gpuPhysAddr=pcidev.bar_info[0].addr, gpuPhysFbAddr=pcidev.bar_info[1].addr, gpuPhysInstAddr=pcidev.bar_info[3].addr,
      pciConfigMirrorBase=[0x88000, 0x92000][self.nvdev.fmc_boot], pciConfigMirrorSize=0x1000, nvDomainBusDeviceFunc=bdf_as_int(self.nvdev.devfmt),
-      bIsPassthru=1, PCIDeviceID=self.nvdev.venid, PCISubDeviceID=self.nvdev.subvenid, PCIRevisionID=self.nvdev.rev, maxUserVa=0x7ffffffff000)
+      bIsPassthru=1, PCIDeviceID=pcidev.read_config(pci.PCI_VENDOR_ID, 4), PCISubDeviceID=pcidev.read_config(pci.PCI_SUBSYSTEM_VENDOR_ID, 4),
+      PCIRevisionID=pcidev.read_config(pci.PCI_REVISION_ID, 1), maxUserVa=0x7ffffffff000)
    self.cmd_q.send_rpc(nv.NV_VGPU_MSG_FUNCTION_GSP_SET_SYSTEM_INFO, bytes(data))

  def rpc_unloading_guest_driver(self):
@@ -1,10 +1,9 @@
 from __future__ import annotations
 import ctypes, time, functools, re, gzip, struct
 from tinygrad.helpers import getenv, DEBUG, fetch, getbits
-from tinygrad.runtime.support.hcq import MMIOInterface
 from tinygrad.runtime.support.memory import TLSFAllocator, MemoryManager
 from tinygrad.runtime.support.nv.ip import NV_FLCN, NV_FLCN_COT, NV_GSP
-from tinygrad.runtime.support.system import System, PCIDevImplBase
+from tinygrad.runtime.support.system import System, PCIDevice, PCIDevImplBase

 NV_DEBUG = getenv("NV_DEBUG", 0)

@@ -71,8 +70,9 @@ class NVMemoryManager(MemoryManager):
  def on_range_mapped(self): self.dev.NV_VIRTUAL_FUNCTION_PRIV_MMU_INVALIDATE.write((1 << 0) | (1 << 1) | (1 << 6) | (1 << 31))

 class NVDev(PCIDevImplBase):
-  def __init__(self, devfmt:str, mmio:MMIOInterface, vram:MMIOInterface, venid:int, subvenid:int, rev:int, bars:dict):
-    self.devfmt, self.mmio, self.vram, self.venid, self.subvenid, self.rev, self.bars = devfmt, mmio, vram, venid, subvenid, rev, bars
+  def __init__(self, pci_dev:PCIDevice):
+    self.pci_dev, self.devfmt, self.mmio = pci_dev, pci_dev.pcibus, pci_dev.map_bar(0, fmt='I')
+
    self.lock_fd = System.flock_acquire(f"nv_{self.devfmt}.lock")

    self.smi_dev, self.is_booting = False, True
@@ -120,7 +120,7 @@ class NVDev(PCIDevImplBase):
    self.include("src/common/inc/swref/published/turing/tu102/dev_fb.h")
    if self.reg("NV_PFB_PRI_MMU_WPR2_ADDR_HI").read() != 0:
      if DEBUG >= 2: print(f"nv {self.devfmt}: WPR2 is up. Issuing a full reset.", flush=True)
-      System.pci_reset(self.devfmt)
+      self.pci_dev.reset()
      time.sleep(0.5)

    self.include("src/common/inc/swref/published/turing/tu102/dev_vm.h")
@@ -134,6 +134,8 @@ class NVDev(PCIDevImplBase):
    self.pte_t, self.pde_t, self.dual_pde_t = tuple([self.__dict__[name] for name in mmu_pd_names])

    self.vram_size = self.reg("NV_PGC6_AON_SECURE_SCRATCH_GROUP_42").read() << 20
+
+    self.vram, self.mmio = self.pci_dev.map_bar(1), self.pci_dev.map_bar(0, fmt='I')
    self.large_bar = self.vram.nbytes >= self.vram_size

  def _alloc_boot_struct(self, struct:ctypes.Structure) -> tuple[ctypes.Structure, int]:
@@ -1,12 +1,16 @@
 import os, mmap, array, functools, ctypes, select, contextlib, dataclasses, sys, errno, itertools
 from typing import cast, ClassVar
 from tinygrad.helpers import round_up, getenv, OSX, temp, ceildiv
-from tinygrad.runtime.autogen import libc, vfio
-from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface, HCQBuffer
+from tinygrad.runtime.autogen import libc, vfio, pci
+from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface, HCQBuffer, hcq_filter_visible_devices
 from tinygrad.runtime.support.memory import MemoryManager, VirtMapping
+from tinygrad.runtime.support.usb import ASM24Controller, USBMMIOInterface

 MAP_FIXED, MAP_LOCKED, MAP_POPULATE, MAP_NORESERVE = 0x10, 0 if OSX else 0x2000, getattr(mmap, "MAP_POPULATE", 0 if OSX else 0x008000), 0x400

+@dataclasses.dataclass(frozen=True)
+class PCIBarInfo: addr:int; size:int # noqa: E702
+
 class _System:
  @functools.cached_property
  def atomic_lib(self): return ctypes.CDLL(ctypes.util.find_library('atomic')) if sys.platform == "linux" else None
@@ -86,10 +90,6 @@ class _System:
    if data is not None: sysmem_view[:len(data)] = data
    return sysmem_view, [p + i for p, sz in paddrs for i in range(0, sz, 0x1000)][:ceildiv(size, 0x1000)]

-  def pci_reset(self, gpu):
-    if OSX: System.iokit_pci_rpc(__TinyGPURPCReset:=2)
-    else: os.system(f"sudo sh -c 'echo 1 > /sys/bus/pci/devices/{gpu}/reset'")
-
  def pci_scan_bus(self, target_vendor:int, target_devices:list[int]) -> list[str]:
    result = []
    for pcibus in FileIOInterface("/sys/bus/pci/devices").listdir():
@@ -98,6 +98,56 @@ class _System:
      if vendor == target_vendor and device in target_devices: result.append(pcibus)
    return sorted(result)

+  def pci_setup_usb_bars(self, usb:ASM24Controller, gpu_bus:int, mem_base:int, pref_mem_base:int) -> dict[int, PCIBarInfo]:
+    for bus in range(gpu_bus):
+      # All 3 values must be written at the same time.
+      buses = (0 << 0) | ((bus+1) << 8) | ((gpu_bus) << 16)
+      usb.pcie_cfg_req(pci.PCI_PRIMARY_BUS, bus=bus, dev=0, fn=0, value=buses, size=4)
+
+      usb.pcie_cfg_req(pci.PCI_MEMORY_BASE, bus=bus, dev=0, fn=0, value=(mem_base>>16) & 0xffff, size=2)
+      usb.pcie_cfg_req(pci.PCI_MEMORY_LIMIT, bus=bus, dev=0, fn=0, value=0xffff, size=2)
+      usb.pcie_cfg_req(pci.PCI_PREF_MEMORY_BASE, bus=bus, dev=0, fn=0, value=(pref_mem_base>>16) & 0xffff, size=2)
+      usb.pcie_cfg_req(pci.PCI_PREF_MEMORY_LIMIT, bus=bus, dev=0, fn=0, value=0xffff, size=2)
+      usb.pcie_cfg_req(pci.PCI_PREF_BASE_UPPER32,  bus=bus, dev=0, fn=0, value=pref_mem_base >> 32, size=4)
+      usb.pcie_cfg_req(pci.PCI_PREF_LIMIT_UPPER32, bus=bus, dev=0, fn=0, value=0xffffffff, size=4)
+
+      usb.pcie_cfg_req(pci.PCI_COMMAND, bus=bus, dev=0, fn=0, value=pci.PCI_COMMAND_IO | pci.PCI_COMMAND_MEMORY | pci.PCI_COMMAND_MASTER, size=1)
+
+    # resize bar 0
+    cap_ptr = 0x100
+    while cap_ptr:
+      if pci.PCI_EXT_CAP_ID(hdr:=usb.pcie_cfg_req(cap_ptr, bus=gpu_bus, dev=0, fn=0, size=4)) == pci.PCI_EXT_CAP_ID_REBAR:
+        cap = usb.pcie_cfg_req(cap_ptr + 0x04, bus=gpu_bus, dev=0, fn=0, size=4)
+        new_ctrl = (usb.pcie_cfg_req(cap_ptr + 0x08, bus=gpu_bus, dev=0, fn=0, size=4) & ~0x1F00) | ((int(cap >> 4).bit_length() - 1) << 8)
+        usb.pcie_cfg_req(cap_ptr + 0x08, bus=gpu_bus, dev=0, fn=0, value=new_ctrl, size=4)
+
+      cap_ptr = pci.PCI_EXT_CAP_NEXT(hdr)
+
+    mem_space_addr, bar_off, bars = [mem_base, pref_mem_base], 0, {}
+    while bar_off < 24:
+      cfg = usb.pcie_cfg_req(pci.PCI_BASE_ADDRESS_0 + bar_off, bus=gpu_bus, dev=0, fn=0, size=4)
+      bar_mem, bar_64 = bool(cfg & pci.PCI_BASE_ADDRESS_MEM_PREFETCH), cfg & pci.PCI_BASE_ADDRESS_MEM_TYPE_64
+
+      if (cfg & pci.PCI_BASE_ADDRESS_SPACE) == pci.PCI_BASE_ADDRESS_SPACE_MEMORY:
+        usb.pcie_cfg_req(pci.PCI_BASE_ADDRESS_0 + bar_off, bus=gpu_bus, dev=0, fn=0, value=0xffffffff, size=4)
+        lo = (usb.pcie_cfg_req(pci.PCI_BASE_ADDRESS_0 + bar_off, bus=gpu_bus, dev=0, fn=0, size=4) & 0xfffffff0)
+
+        if bar_64: usb.pcie_cfg_req(pci.PCI_BASE_ADDRESS_0 + bar_off + 4, bus=gpu_bus, dev=0, fn=0, value=0xffffffff, size=4)
+        hi = (usb.pcie_cfg_req(pci.PCI_BASE_ADDRESS_0 + bar_off + 4, bus=gpu_bus, dev=0, fn=0, size=4) if bar_64 else 0)
+
+        bar_size = ((~(((hi << 32) | lo) & ~0xf)) + 1) & (0xffffffffffffffff if bar_64 else 0xffffffff)
+
+        usb.pcie_cfg_req(pci.PCI_BASE_ADDRESS_0 + bar_off, bus=gpu_bus, dev=0, fn=0, value=mem_space_addr[bar_mem] & 0xffffffff, size=4)
+        if bar_64: usb.pcie_cfg_req(pci.PCI_BASE_ADDRESS_0 + bar_off + 4, bus=gpu_bus, dev=0, fn=0, value=mem_space_addr[bar_mem] >> 32, size=4)
+
+        bars[bar_off // 4] = PCIBarInfo(mem_space_addr[bar_mem], bar_size)
+        mem_space_addr[bar_mem] += round_up(bar_size, 2 << 20)
+
+      bar_off += 8 if bar_64 else 4
+
+    usb.pcie_cfg_req(pci.PCI_COMMAND, bus=gpu_bus, dev=0, fn=0, value=pci.PCI_COMMAND_IO | pci.PCI_COMMAND_MEMORY | pci.PCI_COMMAND_MASTER, size=1)
+    return bars
+
  def flock_acquire(self, name:str) -> int:
    import fcntl # to support windows

@@ -153,23 +203,33 @@ class PCIDevice:
    self.cfg_fd = FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/config", os.O_RDWR | os.O_SYNC | os.O_CLOEXEC)
    self.bar_fds = {b: FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/resource{b}", os.O_RDWR | os.O_SYNC | os.O_CLOEXEC) for b in bars}

-    bar_info = FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/resource", os.O_RDONLY).read().splitlines()
-    self.bar_info = {j:(int(start,16), int(end,16), int(flgs,16)) for j,(start,end,flgs) in enumerate(l.split() for l in bar_info)}
+    res = FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/resource", os.O_RDONLY).read().splitlines()
+    self.bar_info = {j:PCIBarInfo(int(s,16), int(e,16)-int(s,16)+1) for j,(s,e,_) in enumerate(l.split() for l in res)}

  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 map_bar(self, bar:int, off:int=0, addr:int=0, size:int|None=None, fmt='B') -> MMIOInterface:
-    fd, sz = self.bar_fds[bar], size or (self.bar_info[bar][1] - self.bar_info[bar][0] + 1)
+    fd, sz = self.bar_fds[bar], size or (self.bar_info[bar].size - 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)
    return MMIOInterface(loc, sz, fmt=fmt)
+  def reset(self): os.system(f"sudo sh -c 'echo 1 > /sys/bus/pci/devices/{self.pcibus}/reset'")

 class APLPCIDevice(PCIDevice):
  def __init__(self, pcibus:str, bars:list[int], resize_bars:list[int]|None=None):
    self.pcibus, self.bars = pcibus, {b: System.iokit_pci_memmap(b) for b in bars}
-    self.bar_info = {b:(0, self.bars[b].nbytes-1 if b in self.bars else 0, 0) for b in range(6)} # NOTE: fake bar info for nv.
+    self.bar_info = {b:PCIBarInfo(0, self.bars[b].nbytes-1 if b in self.bars else 0) for b in range(6)} # NOTE: fake bar info for nv.
  def map_bar(self, bar:int, off:int=0, addr:int=0, size:int|None=None, fmt='B') -> MMIOInterface: return self.bars[bar].view(off, size, fmt)
  def read_config(self, offset:int, size:int): return System.iokit_pci_rpc(__TinyGPURPCReadCfg:=0, offset, size)[0]
  def write_config(self, offset:int, value:int, size:int): System.iokit_pci_rpc(__TinyGPURPCWriteCfg:=1, offset, size, value)
+  def reset(self): System.iokit_pci_rpc(__TinyGPURPCReset:=2)
+
+class USBPCIDevice(PCIDevice):
+  def __init__(self, pcibus:str, bars:list[int], resize_bars:list[int]|None=None):
+    self.usb = ASM24Controller()
+    self.pcibus, self.bar_info = pcibus, System.pci_setup_usb_bars(self.usb, gpu_bus=4, mem_base=0x10000000, pref_mem_base=(32 << 30))
+  def map_bar(self, bar, off=0, addr=0, size=None, fmt='B'):
+    return USBMMIOInterface(self.usb, self.bar_info[bar].addr + off, size or self.bar_info[bar].size, fmt)
+  def dma_view(self, ctrl_addr, size): return USBMMIOInterface(self.usb, ctrl_addr, size, fmt='B', pcimem=False)

 class PCIDevImplBase:
  mm: MemoryManager
@@ -183,14 +243,12 @@ class LNXPCIIfaceBase:

  def __init__(self, dev, dev_id, vendor, devices, bars, vram_bar, va_start, va_size):
    if len((cls:=type(self)).gpus) == 0:
-      cls.gpus = System.pci_scan_bus(vendor, devices)
-      visible_devices = [int(x) for x in (getenv('VISIBLE_DEVICES', '')).split(',') if x.strip()]
-      cls.gpus = [cls.gpus[x] for x in visible_devices] if visible_devices else cls.gpus
+      cls.gpus = hcq_filter_visible_devices(System.pci_scan_bus(vendor, devices))

      # Acquire va range to avoid collisions.
      FileIOInterface.anon_mmap(va_start, va_size, 0, mmap.MAP_PRIVATE | mmap.MAP_ANONYMOUS | MAP_NORESERVE | MAP_FIXED, 0)
    self.pci_dev, self.dev, self.vram_bar = PCIDevice(cls.gpus[dev_id], bars=bars, resize_bars=[vram_bar]), dev, vram_bar
-    self.p2p_base_addr = self.pci_dev.bar_info[vram_bar][0]
+    self.p2p_base_addr = self.pci_dev.bar_info[vram_bar].addr

  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, force_devmem=False, **kwargs) -> HCQBuffer:
    # NOTE: logic on macos is different, since bar is small
@@ -141,6 +141,7 @@ def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UO

@profile_matches
 def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
+  if debug: print("**************************")
  rctx = IndexingContext()

  # get ops to realize
@@ -4,8 +4,9 @@ from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, _substitute, ssimplify, KernelInfo
 from tinygrad.uop.ops import track_rewrites, graph_rewrite, identity_element, sint, AxisType, BottomUpGate
 from tinygrad.uop.symbolic import symbolic_flat
-from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, flatten, dedup, all_int, DEBUG, SPLIT_REDUCEOP, Metadata
-from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_unparented
+from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, flatten, dedup, all_int, DEBUG, SPLIT_REDUCEOP, Metadata, DEBUG_RANGEIFY
+from tinygrad.helpers import PCONTIG, partition
+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, ALWAYS_CONTIGUOUS, IndexingContext, apply_movement_op

@@ -155,18 +156,24 @@ def remove_bufferize(src:UOp, buf:UOp, idx:UOp):
    # if we return None, the bufferize is kept

    accessed_buffers: list[UOp] = []
+    indexes: list[UOp] = []
    reduces: list[UOp] = []
    def red_gate(x:UOp):
-      if x.op is Ops.INDEX:
+      if x.op is Ops.BUFFERIZE and x.arg.addrspace == AddrSpace.GLOBAL:
        accessed_buffers.append(x)
        return False
+      if x.op is Ops.BUFFER:
+        accessed_buffers.append(x)
+      if x.op is Ops.INDEX:
+        indexes.append(x)
      if x.op is Ops.REDUCE: reduces.append(x)
      return True
    src.toposort(gate=red_gate)
    del red_gate
+    accessed_buffers = dedup(accessed_buffers)

    # if this is generated from multiple buffers, don't remove this buffer
-    if len(dedup([x.src[0] for x in accessed_buffers])) > 2: return None
+    if len(accessed_buffers) > 2 and not (PCONTIG > 2): return None

    # if any reduces access a buffer, don't remove this buffer
    buffer_in_reduce = False
@@ -176,44 +183,72 @@ def remove_bufferize(src:UOp, buf:UOp, idx:UOp):
      return not buffer_in_reduce
    UOp.sink(*[x.src[0] for x in reduces]).toposort(gate=buf_gate)
    del buf_gate
-    if buffer_in_reduce: return None
+    if buffer_in_reduce:
+      if PCONTIG > 2:
+        out_in_ratio = (prod(buf.shape)+1) / (sum([x.size for x in accessed_buffers])+1)
+        if out_in_ratio < 10: return None
+        # here we have to check the indexes, we might do a partial contig here
+        local_indexes = [x for x in indexes if x.src[0].op is Ops.BUFFERIZE and x.src[0].arg.addrspace == AddrSpace.LOCAL]
+        exclude_ranges = UOp.group(*[UOp.group(*x.src[1:]) for x in local_indexes]).ranges
+        subs = [(k,v) for k,v in zip(buf.src[1:], idx.src[1:]) if k.op is not Ops.CONST]
+        # if it's bufferized or a reduce, it's pcontig
+        is_pcontig, is_subs = partition(subs, lambda x: x[0] in exclude_ranges or any([r.arg[-1] == AxisType.REDUCE for r in x[1].ranges]))
+        if not len(is_subs):
+          return None
+        if len(is_pcontig):
+          ret = src.substitute(dict(is_subs), extra_pm=pm_gate_substitute)
+          return ret.bufferize(*[x[0] for x in is_pcontig], arg=BufferizeOpts(None, AddrSpace.LOCAL)).index(*[x[1] for x in is_pcontig])
+      else:
+        return None

  # if it makes it here, the bufferize is removed
  # this is the ranges replaced
  # NOTE: if buf src is a const, we don't replace it
  return src.substitute({k:v for k,v in zip(buf.src[1:], idx.src[1:]) if k.op is not Ops.CONST}, extra_pm=pm_gate_substitute)

-def pre_bufferize(b:UOp, x:UOp, copy:UOp):
-  nb = b.replace(src=(b.src[0].contiguous(),)+b.src[1:])
-  return copy.replace(src=(x.replace(src=(nb,)+x.src[1:]), copy.src[1]))
+def remove_noop_bufferize(idx,b2):
+  if idx.src[1:] != b2.src[1:] or idx.src[0].op is Ops.BUFFER_VIEW: return None
+  new_tag = (idx.src[0].tag or ()) + (b2.tag or ()) or None
+  return idx.src[0].rtag(new_tag).shrink(tuple((0, s) for s in b2.shape)) if b2.shape else idx.src[0].rtag(new_tag)

-pm_cleanups = pm_mops+PatternMatcher([
+pm_const_buffer_folding = pm_mops+PatternMatcher([
  (UPat(Ops.BUFFERIZE, name="b"), cleanup_dead_axes),
  (UPat(GroupOp.All-{Ops.BUFFERIZE, Ops.BUFFER}, name="x"), lambda x: x.replace(dtype=x.dtype.base) if isinstance(x.dtype, ImageDType) else None),
  (UPat((Ops.BUFFERIZE), name="x"), lambda x: x.replace(dtype=x.dtype.base) if isinstance(x.dtype, ImageDType)
    and (resolve(prod(x.dtype.shape)!=prod(x.shape)) or x.shape[-1]%4!=0) else None),
  # remove noop buffers. if we look at the next index we can remove even more of these
-  # NOTE: this is mostly the same case as below, but if there's no INDEX this gets more
-  (UPat(Ops.INDEX, name="idx").f(Ops.BUFFERIZE, allow_any_len=True, name="b2"),
-   lambda idx,b2: idx.src[0].replace(tag=nt if len(nt:=(idx.src[0].tag or ()) + (b2.tag or ())) else None) if idx.src[1:] == b2.src[1:] \
-       and idx.src[0].op is not Ops.BUFFER_VIEW else None),
-  # remove reindexing with cost function
-  (UPat.var("src").f(Ops.BUFFERIZE, allow_any_len=True, name="buf").f(Ops.INDEX, allow_any_len=True, name="idx"), remove_bufferize),
+  (UPat(Ops.INDEX, name="idx").f(Ops.BUFFERIZE, allow_any_len=True, name="b2"), remove_noop_bufferize),
+  # dont bufferize an arange
+  (UPat.any((r:=UPat(dtype=dtypes.index).cast()).named("src"), r.eq(UPat()).named("src")).f(Ops.BUFFERIZE,
+    allow_any_len=True, name="buf").f(Ops.INDEX, allow_any_len=True, name="idx"), remove_bufferize),
  # no buffers for const
  (UPat(Ops.CONST, name='c').f(Ops.BUFFERIZE, allow_any_len=True, name="b"), lambda c,b: b.const_like(c.arg).rtag(b.tag)),
+  # indexing a const is a const
+  (UPat(Ops.INDEX, src=(UPat(Ops.CONST, name="c"),),), lambda c: c),
  # copy on CONST is CONST
  (UPat(Ops.COPY, src=(UPat.cvar("x"), UPat()), name="copy"), lambda copy,x: copy.const_like(x.arg)),
-  (UPat(Ops.COPY, src=(UPat(GroupOp.All-{Ops.CONTIGUOUS, Ops.COPY}).f(Ops.BUFFERIZE, allow_any_len=True, name="b")
-                       .f(Ops.INDEX, allow_any_len=True, name="x"), UPat()), name="copy"), pre_bufferize),
+  # hack if a noop turned to a const
+  (UPat.cvar("c").f(Ops.NOOP).f(Ops.BUFFERIZE, allow_any_len=True, name="buf"), lambda c,buf: buf.replace(src=(c,)+buf.src[1:])),
  # mstack on CONST is CONST
  (UPat(Ops.MSTACK, src=(UPat.var("s"),), allow_any_len=True).f(Ops.INDEX, allow_any_len=True),
   lambda s: UOp.const(c.dtype, c.arg) if (c:=s.base).op is Ops.CONST else None),
 ])

+def pre_bufferize(b:UOp, x:UOp, copy:UOp):
+  nb = b.replace(src=(b.src[0].contiguous(),)+b.src[1:])
+  return copy.replace(src=(x.replace(src=(nb,)+x.src[1:]), copy.src[1]))
+pm_remove_bufferize = PatternMatcher([
+  # hack so remove_bufferize doesnt remove the buffer before a copy
+  (UPat(Ops.COPY, src=(UPat(GroupOp.All-{Ops.CONTIGUOUS, Ops.COPY}).f(Ops.BUFFERIZE, allow_any_len=True, name="b")
+                      .f(Ops.INDEX, allow_any_len=True, name="x"), UPat()), name="copy"), pre_bufferize),
+  # remove reindexing with cost function
+  (UPat.var("src").f(Ops.BUFFERIZE, allow_any_len=True, name="buf").f(Ops.INDEX, allow_any_len=True, name="idx"), remove_bufferize),
+])
+
 def late_buffer_view(t:UOp, b:UOp):
  if isinstance(b.device, str) and (b.device.startswith("DISK") or b.device.startswith("TINYFS")):
-    rngs = b.src[1:]
-    size = prod(shape := [int(r.vmax+1) for r in rngs])
+    shape = b.shape
+    size = prod(shape)

    # walk up for the INDEX
    x = t
@@ -264,11 +299,11 @@ pm_limit_bufs = PatternMatcher([(UPat(set.union(GroupOp.Binary, GroupOp.Ternary)
 # BUFFERIZE returns the BUFFER ready for INDEXing (doing this will make splitting a lot easier)
 # NOTE: this has been fixed up a bit

-def bufferize_to_store(x:UOp):
+def bufferize_to_store(x:UOp, allow_locals=True):
  rngs = x.src[1:]
-  shape = tuple([int(r.vmax+1) for r in rngs])
+  shape = x.shape
  size = prod(shape)
-  assert size > 0, f"no zero sized buffers {shape}"
+  assert size > 0 and isinstance(size, int), f"no zero sized or symbolic sized buffers {shape}"

  sdtype = x.dtype.ptr(size=size, addrspace=x.arg.addrspace)
  if x.src[0].op is Ops.ASSIGN:
@@ -276,7 +311,7 @@ def bufferize_to_store(x:UOp):
    assert assign_target.op is Ops.INDEX, f"{assign_target.op} is not index"
    # in assign, this is the buffer size, not the bufferize size
    # TODO: assign_mops here
-    do_store = assign_target.replace(dtype=sdtype).store(assign_src).replace(tag=x.tag).end(ends=[x for x in rngs if x.op is Ops.RANGE])
+    do_store = assign_target.replace(dtype=sdtype).store(assign_src, tag=x.tag).end(*[x for x in rngs if x.op is Ops.RANGE])
    ret = assign_target.src[0].after(do_store)
    mops = []
    walk = assign_mops
@@ -289,7 +324,7 @@ def bufferize_to_store(x:UOp):
  # NOTE: the DEFINE_LOCAL needs to be disambiguated here
  if sdtype.addrspace == AddrSpace.GLOBAL:
    buf = UOp.new_buffer(x.arg.device, size, x.dtype)
-    do_store = buf.reshape(shape).index(*rngs, dtype=sdtype).store(x.src[0]).replace(tag=x.tag).end(ends=[x for x in rngs if x.op is Ops.RANGE])
+    do_store = buf.reshape(shape).index(*rngs, dtype=sdtype).store(x.src[0], tag=x.tag).end(*[x for x in rngs if x.op is Ops.RANGE])
    ret = buf.after(do_store).forced_reshape(shape)
    # TODO: is this right? what if it's offset
    if any(r.op is Ops.RANGE and r.src[0].op is not Ops.CONST for r in rngs):
@@ -297,21 +332,26 @@ def bufferize_to_store(x:UOp):
      ret = ret.shrink(tuple([(0,x) for x in sym_shape]))
    return ret.replace(tag=x.tag)

-  # handle locals
-  tag = x.arg.device
-  if tag is None: tag = UOp.unique().arg # TODO: hack
-  buf = UOp(Ops.DEFINE_LOCAL, sdtype, arg=tag)
-  do_store = buf.reshape(shape).index(*rngs, dtype=sdtype).store(x.src[0]).end(ends=[x for x in rngs if x.op is Ops.RANGE])
-  return buf.after(do_store).reshape(shape)
+  if allow_locals:
+    # handle locals
+    tag = x.arg.device
+    if tag is None: tag = UOp.unique().arg # TODO: hack
+    buf = UOp(Ops.DEFINE_LOCAL, sdtype, arg=tag)
+    do_store = buf.reshape(shape).index(*rngs, dtype=sdtype).store(x.src[0]).end(*[x for x in rngs if x.op is Ops.RANGE])
+    return buf.after(do_store.barrier()).reshape(shape)

 pm_add_buffers = pm_mops+to_bufferview+PatternMatcher([
-  (UPat(Ops.BUFFERIZE, name="x"), bufferize_to_store),
+  (UPat(Ops.BUFFERIZE, name="x"), lambda x: bufferize_to_store(x, allow_locals=False)),

  # move RESHAPEs through MSELECT/MSTACK
  (UPat((Ops.MSELECT, Ops.MSTACK), src=UPat(Ops.RESHAPE), name="m"),
   lambda m: m.replace(src=tuple([x.src[0].base for x in m.src]), tag=None).reshape(m.shape).rtag(m.tag)),
 ])

+pm_add_buffers_local = pm_mops+to_bufferview+PatternMatcher([
+  (UPat(Ops.BUFFERIZE, name="x"), bufferize_to_store),
+])
+
 # *****************
 # 5. split into kernels

@@ -344,8 +384,8 @@ def handle_after(ctx:LocalAddBufferContext, after:UOp):
  return buf

 def renumber_range(ctx:LocalAddBufferContext, r:UOp):
-  if r.tag is not None: return None
-  ret = r.replace(arg=(ctx.range,)+r.arg[1:], tag=())
+  if r.tag != (): return None
+  ret = r.replace(arg=(ctx.range,)+r.arg[1:], tag=None)
  ctx.range += 1
  return ret

@@ -403,13 +443,14 @@ pm_remove_tags = PatternMatcher([
  (UPat(GroupOp.All, name="x"), remove_metadata_tags),
 ])

+pm_add_range_tags = PatternMatcher([
+  (UPat(Ops.RANGE, name="x"), lambda x: x.rtag(()))
+])
+
@dataclass(frozen=True)
 class Kernel:
  ast: UOp
  metadata: tuple[Metadata, ...] = ()
-  def __repr__(self):
-    ast_rep = f"SINK{tuple(s.op for s in self.ast.src)}" if self.ast.op is Ops.SINK else repr(self.ast.op)
-    return f"<Kernel {len(list(self.ast.toposort()))} {ast_rep} {self.metadata}>"

 def split_store(ctx:list[UOp], x:UOp) -> UOp|None:
  if len(x.ranges): return None
@@ -477,11 +518,10 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
  tsink = graph_rewrite(tsink, earliest_rewrites+replace_contiguous, ctx={}, name="earliest rewrites")

  # convert movement ops to ranges
-  tsink, rctx = run_rangeify(tsink, getenv("DEBUG_RANGEIFY", 0))
+  tsink, rctx = run_rangeify(tsink, DEBUG_RANGEIFY)

-  # NOTE: sym (vs symbolic_simple) breaks things here because ranges with len 1 aren't handled right
-  tsink = graph_rewrite(tsink, symbolic_flat+pm_reduce_unparented, name="symbolic")  # this supports const folding
-  tsink = graph_rewrite(tsink, pm_cleanups, bottom_up=True, name="remove costly buffers")
+  tsink = graph_rewrite(tsink, symbolic_flat+pm_reduce_simplify+pm_const_buffer_folding, name="symbolic+reduce_collapse")  # this does const folding
+  tsink = graph_rewrite(tsink, pm_remove_bufferize, bottom_up=True, name="remove bufferize with cost function")
  tsink = graph_rewrite(tsink, pm_limit_bufs, ctx=rctx, name="limit buffers")

  # rebuild the sink with all the BUFFERIZEs with tags, this is what's ending up in the tensor graph
@@ -493,7 +533,7 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
  if getenv("VIZ"): graph_rewrite(tsink, PatternMatcher([]), name="View Tagged Rangeify")

  # bufferize -> store
-  tsink = graph_rewrite(tsink, pm_add_buffers, bottom_up=True, name="bufferize to store")
+  tsink = graph_rewrite(tsink, pm_add_buffers+pm_add_range_tags, bottom_up=True, name="bufferize to store")
  tsink = graph_rewrite(tsink, split_kernels, ctx=uop_list, name="split kernels")

  # if a kernel depends on a buffer, and that buffer is later assigned to, make the assign depend on the kernel's assign
@@ -6,12 +6,12 @@ from typing import Callable, ClassVar, Sequence, cast, get_args, Literal, Suppor
 from tinygrad.dtype import DType, DTypeLike, dtypes, ImageDType, ConstType, least_upper_float, least_upper_dtype, sum_acc_dtype, to_dtype, truncate
 from tinygrad.dtype import _from_np_dtype, _to_np_dtype
 from tinygrad.helpers import argfix, make_tuple, flatten, prod, all_int, round_up, merge_dicts, argsort, getenv, all_same, fully_flatten, dedup
-from tinygrad.helpers import IMAGE, WINO, Metadata, TRACEMETA, ceildiv, fetch, polyN, DEBUG, is_numpy_ndarray, FUSE_ATTENTION
+from tinygrad.helpers import IMAGE, WINO, Metadata, TRACEMETA, ceildiv, fetch, polyN, DEBUG, is_numpy_ndarray, FUSE_ATTENTION, SPEC
 from tinygrad.helpers import suppress_finalizing
 from tinygrad.gradient import compute_gradient
 from tinygrad.uop.mathtraits import MathTrait
 from tinygrad.uop.ops import smax, smin, resolve, UOp, Ops, sint, identity_element, all_metadata, _index_to_concrete_int, sint_to_uop, srender
-from tinygrad.uop.spec import tensor_uop_spec, type_verify
+from tinygrad.uop.spec import type_verify, tensor_spec
 from tinygrad.device import Device, Buffer
 from tinygrad.engine.realize import run_schedule
 from tinygrad.engine.memory import memory_planner
@@ -115,7 +115,7 @@ class Tensor(MathTrait):
  training: ClassVar[bool] = False

  def __init__(self, data:ConstType|bytes|list|tuple|UOp|'np.ndarray'|pathlib.Path|None,  # type: ignore [name-defined] # noqa: F821
-               device:str|tuple|list|None=None, dtype:DTypeLike|None=None, requires_grad:bool|None=None):
+               device:str|tuple|list|None=None, dtype:DTypeLike|None=None, requires_grad:bool|None=None, _force_unique:bool=False):
    if device is None and isinstance(data, pathlib.Path): device = f"DISK:{data.resolve()}"  # keep it on the disk if device is None
    _dtype:DType|None = to_dtype(dtype) if dtype is not None else None
    _device:str|tuple[str, ...] = tuple(canonicalize_device(x) for x in device) if isinstance(device, (tuple, list)) else canonicalize_device(device)
@@ -138,8 +138,8 @@ class Tensor(MathTrait):
        # give the bound constant a device
        const = UOp.const(var.dtype, val, _device, ())
        data = data.replace(src=(var.replace(src=const.src), const))  # type: ignore
-    elif data is None: data = UOp.const(_dtype or dtypes.default_float, 0, _device, ())
-    elif isinstance(data, get_args(ConstType)): data = UOp.const(_dtype or dtypes.from_py(data), data, _device, ())
+    elif data is None: data = UOp.const(_dtype or dtypes.default_float, 0, _device, (), unique=_force_unique)
+    elif isinstance(data, get_args(ConstType)): data = UOp.const(_dtype or dtypes.from_py(data), data, _device, (), unique=_force_unique)
    elif isinstance(data, bytes): data = _frompy(data, dtypes.uint8 if _dtype is None else _dtype)
    elif isinstance(data, (list, tuple)):
      if _dtype is None:
@@ -150,7 +150,7 @@ class Tensor(MathTrait):
    elif is_numpy_ndarray(data):
      import numpy as np
      assert isinstance(data, np.ndarray), f"expected np.ndarray, got {data}"
-      if data.shape == (): data = UOp.const(_dtype or _from_np_dtype(data.dtype), data.item(), _device, ())
+      if data.shape == (): data = UOp.const(_dtype or _from_np_dtype(data.dtype), data.item(), _device, (), unique=_force_unique)
      else: data = _fromnp(data.astype(npdtype) if _dtype is not None and (npdtype:=_to_np_dtype(_dtype)) is not None else data)  # type: ignore [name-defined]
    elif isinstance(data, pathlib.Path):
      _dtype = _dtype or dtypes.uint8
@@ -229,7 +229,7 @@ class Tensor(MathTrait):
    big_sink = UOp.sink(*[x.uop for x in (self,)+lst])

    # verify Tensors match the spec
-    if __debug__: type_verify(list(big_sink.toposort()), tensor_uop_spec)
+    if SPEC: type_verify(big_sink, tensor_spec)

    if any(isinstance(x._device, tuple) for x in big_sink.toposort()):
      _apply_map_to_tensors(get_multi_map(big_sink), "Apply Multi Map")
@@ -625,7 +625,7 @@ class Tensor(MathTrait):
    print(Tensor.full((2, 3), False).numpy())
    ```
    """
-    return Tensor(fill_value, **kwargs).reshape((1, )*len(new_shape := argfix(shape))).expand(new_shape)
+    return Tensor(fill_value, _force_unique=True, **kwargs).reshape((1, )*len(new_shape := argfix(shape))).expand(new_shape)

  @staticmethod
  def zeros(*shape, **kwargs) -> Tensor:
@@ -2090,7 +2090,7 @@ class Tensor(MathTrait):

    state = Tensor.zeros(bs, 25, device=self.device, dtype=dtypes.uint64)
    for k in range(int(data.shape[1])):
-      state = state.bitwise_xor(data[:,k].reshape(bs, 25))
+      state = state ^ data.shrink((None, (k, k+1), None)).squeeze(1)
      for i in range(24): # f1600
        # θ step
        p = state.reshape(bs, 5, 5).transpose(2, 1)
@@ -4149,10 +4149,10 @@ class Tensor(MathTrait):
    ```
    """
    assert self.ndim > 1, "NS only works for two or more dims"
+    if self.shape[-2] > self.shape[-1]: return self.transpose(-2, -1).newton_schulz(steps, params, eps).transpose(-2, -1)
    G = self / (self.square().sum(axis=(-2, -1), keepdim=True).sqrt() + eps)
-    if (swap := self.shape[-2] > self.shape[-1]): G = G.transpose(-2, -1)
    for _ in range(steps): G = sum(p * functools.reduce(lambda x, y: (y @ y.transpose(-2, -1)) @ x, [G]*i, G) for i,p in enumerate(params))
-    return G.transpose(-2, -1) if swap else G
+    return G

  def qr(self) -> tuple[Tensor, Tensor]:
    assert self.ndim > 1, f"expected two or more dimensions, got {self.ndim}"
@@ -3,6 +3,7 @@ from enum import auto, IntEnum, Enum
 # wrapper around IntEnum that preserves Enum.__str__ and makes auto() unique across all FastEnum subclasses
 class FastEnum(IntEnum):
  def __str__(self): return Enum.__str__(self)
+  def __repr__(x): return str(x)
  @staticmethod
  def _generate_next_value_(_, __, ___, last_values): return 1 + max([0, *last_values, *[max(c) for c in FastEnum.__subclasses__()]])

@@ -15,6 +16,9 @@ class Ops(FastEnum):
  # AFTER passes src[0] through and promises in the toposort that any consumers of the AFTER run after src[1:]
  AFTER = auto()

+  # GROUP is a NOOP that just merges things together
+  GROUP = auto()
+
  # buffer ops
  COPY = auto(); BUFFER = auto(); BUFFER_VIEW = auto(); MSELECT = auto(); MSTACK = auto() # noqa: E702

@@ -44,7 +48,7 @@ class Ops(FastEnum):
  UNROLL = auto(); CONTRACT = auto(); GEP = auto(); VECTORIZE = auto(); CAT = auto(); PTRCAT = auto() # noqa: E702

  # UnaryOps
-  CAST = auto(); BITCAST = auto(); EXP2 = auto(); LOG2 = auto(); SIN = auto(); SQRT = auto(); RECIP = auto(); NEG = auto(); TRUNC = auto() # noqa: E702
+  CAST = auto(); BITCAST = auto(); EXP2 = auto(); LOG2 = auto(); SIN = auto(); SQRT = auto(); RECIPROCAL = auto(); NEG = auto(); TRUNC = auto() # noqa: E702

  # load/store before math
  LOAD = auto(); STORE = auto() # noqa: E702
@@ -75,7 +79,7 @@ class Ops(FastEnum):
  CUSTOM = auto(); CUSTOMI = auto() # noqa: E702

 class GroupOp:
-  Unary = {Ops.EXP2, Ops.LOG2, Ops.SIN, Ops.SQRT, Ops.RECIP, Ops.NEG, Ops.TRUNC}
+  Unary = {Ops.EXP2, Ops.LOG2, Ops.SIN, Ops.SQRT, Ops.RECIPROCAL, Ops.NEG, Ops.TRUNC}
  Binary = {Ops.ADD, Ops.MUL, Ops.IDIV, Ops.MAX, Ops.MOD, Ops.CMPLT, Ops.CMPNE, Ops.CMPEQ,
            Ops.XOR, Ops.SHL, Ops.SHR, Ops.OR, Ops.AND, Ops.THREEFRY, Ops.SUB, Ops.FDIV, Ops.POW}
  Ternary = {Ops.WHERE, Ops.MULACC}
@@ -104,6 +108,6 @@ class GroupOp:
  Comparison = {Ops.CMPLT, Ops.CMPNE, Ops.CMPEQ}

  # do not preserve f(0) = 0
-  UnsafePad = {Ops.RECIP, Ops.LOG2, Ops.EXP2, Ops.IDIV, Ops.POW}
+  UnsafePad = {Ops.RECIPROCAL, Ops.LOG2, Ops.EXP2, Ops.IDIV, Ops.POW}

  All = set(Ops)
@@ -114,7 +114,8 @@ class MathTrait:
    return self._binop(Ops.IDIV, x, reverse)
  def mod(self:TMT, x:TMT|ConstType, reverse:bool=False): return self._binop(Ops.MOD, x, reverse)
  def sub(self:TMT, x:TMT|ConstType, reverse:bool=False): return self.ufix(x).alu(Ops.ADD, -self) if reverse else self.alu(Ops.ADD, self.ufix(-x))
-  def div(self:TMT, x:TMT|ConstType, reverse:bool=False): return (self.ufix(x)*self.alu(Ops.RECIP)) if reverse else (self*self.ufix(x).alu(Ops.RECIP))
+  def div(self:TMT, x:TMT|ConstType, reverse:bool=False):
+    return (self.ufix(x)*self.alu(Ops.RECIPROCAL)) if reverse else (self*self.ufix(x).alu(Ops.RECIPROCAL))

  def __neg__(self): return self.neg()

@@ -162,7 +163,7 @@ class MathTrait:
    if isinstance(y, type(self)): return self.alu(Ops.WHERE, y.ufix(x), y)
    raise RuntimeError("where needs at least one UOp arg")
  def threefry(self:TMT, seed:TMT): return self.alu(Ops.THREEFRY, seed)
-  def reciprocal(self): return self.alu(Ops.RECIP)
+  def reciprocal(self): return self.alu(Ops.RECIPROCAL)
  def trunc(self): return self.alu(Ops.TRUNC)
  def sqrt(self): return self.alu(Ops.SQRT)
  def sin(self): return self.alu(Ops.SIN)
@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import Any, Callable, cast, TYPE_CHECKING, Type, Sequence
+from typing import Any, Callable, cast, TYPE_CHECKING, Type, Sequence, Iterable
 import sys, time, functools, itertools, math, operator, hashlib, os, types, pickle, pathlib, inspect, weakref, collections
 from dataclasses import dataclass
 from enum import Enum, auto
@@ -8,16 +8,21 @@ from tinygrad.uop.mathtraits import MathTrait
 from tinygrad.dtype import ConstType, ImageDType, dtypes, DType, truncate, PtrDType, least_upper_dtype, Invalid, InvalidType
 from tinygrad.helpers import ContextVar, all_int, prod, getenv, all_same, Context, partition, temp, unwrap, T, argfix, Metadata, flatten, TRACEMETA
 from tinygrad.helpers import PICKLE_BUFFERS, PROFILE, dedup, cdiv, cmod, diskcache_put, to_function_name, cpu_profile, TracingKey, VIZ, SPEC
-from tinygrad.helpers import strip_parens
+from tinygrad.helpers import strip_parens, colored
 if TYPE_CHECKING:
  from tinygrad.device import Buffer, MultiBuffer

 class AxisType(Enum):
  def __repr__(self): return str(self)
  GLOBAL = auto(); WARP = auto(); LOCAL = auto(); LOOP = auto(); GROUP_REDUCE = auto(); REDUCE = auto(); UPCAST = auto(); UNROLL = auto() # noqa: E702
-  THREAD = auto()
+  THREAD = auto(); IF = auto() # noqa: E702
+axis_letters = {AxisType.GLOBAL: "g", AxisType.THREAD: "t", AxisType.LOCAL: "l", AxisType.WARP: "w", AxisType.LOOP: "L", AxisType.UPCAST: "u",
+                AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.UNROLL: "r", AxisType.IF: "I"}
+axis_colors = {AxisType.GLOBAL: "blue", AxisType.THREAD: "BLUE", AxisType.LOCAL: "cyan", AxisType.WARP: "CYAN", AxisType.LOOP: "WHITE",
+               AxisType.UPCAST: "yellow", AxisType.GROUP_REDUCE: "RED", AxisType.REDUCE: "red", AxisType.UNROLL: "magenta",
+               AxisType.IF: "green"}

-range_start = {Ops.BUFFERIZE: 1, Ops.REDUCE: 1, Ops.STORE: 2, Ops.WMMA: 3}
+range_start = {Ops.BUFFERIZE: 1, Ops.REDUCE: 1, Ops.STORE: 2, Ops.WMMA: 3, Ops.END: 1}

 # https://en.wikipedia.org/wiki/Identity_element
 def identity_element(op:Ops, dt:DType) -> ConstType: return dtypes.as_const({Ops.ADD:0, Ops.MUL:1, Ops.MAX:dtypes.min(dt)}[op], dt)
@@ -40,7 +45,16 @@ def srender(x:sint) -> str: return x.render() if isinstance(x, UOp) else str(x)
 def ssimplify(uop:sint): return uop.ssimplify() if isinstance(uop, UOp) else uop
 def sym_infer(uop: UOp|int, var_vals: dict[str, int]) -> int: return uop.sym_infer(var_vals) if isinstance(uop, UOp) else uop

-def range_str(u:UOp) -> str: return '_'.join([str(x) if x >= 0 else "m"+str(-x) for x in u.arg[0:-1]])
+def range_str(u:UOp, color=False) -> str:
+  ret = '_'.join([str(x) if x >= 0 else "m"+str(-x) for x in u.arg[0:-1]])
+  return colored(ret, axis_colors[u.arg[-1]]) if color else ret
+
+def consumer_map_from_toposort(lst:Iterable[UOp]):
+  ret: dict[UOp, dict[UOp, None]] = {}
+  for u in lst:
+    ret[u] = {}
+    for s in u.src: ret[s][u] = None
+  return ret

 # used for UOp and UPat
 def pretty_print(x:Any, rep:Callable, srcfn=lambda x: x.src, cache=None, d=0)->str:
@@ -64,8 +78,9 @@ class UOpMetaClass(type):
    if _buffer is not None:
      assert op is Ops.BUFFER, f"trying to set Buffer {_buffer} for {op}"
      buffers[created] = _buffer
-    if SPEC:
-      from tinygrad.uop.spec import full_spec
+    if SPEC > 1:
+      from tinygrad.uop.spec import full_spec, test_pyrender
+      if SPEC > 2: test_pyrender(created)
      with Context(IGNORE_OOB=1): ret = full_spec.rewrite(created)
      if cast(bool|None, ret) is not True: raise RuntimeError(f"SPEC ISSUE {ret}: {created}")
    return created
@@ -144,12 +159,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
    return ret

  # returns map of UOps to their consumers in the graph rooted by self
-  def get_consumer_map(self) -> dict[UOp, dict[UOp, None]]:
-    ret: dict[UOp, dict[UOp, None]] = {}
-    for u in self.toposort():
-      ret[u] = {}
-      for s in u.src: ret[s][u] = None
-    return ret
+  def get_consumer_map(self) -> dict[UOp, dict[UOp, None]]: return consumer_map_from_toposort(self.toposort())

  def reverse_toposort(self, consumer_map) -> dict[UOp, None]:
    ret: dict[UOp, None] = {}
@@ -250,7 +260,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
          return tuple(1 if i in axis_arg else s for i,s in enumerate(ps))

    # elementwise ops keep the shape the same. all inputs with shape must match
-    if self.op in (GroupOp.Elementwise-{Ops.BITCAST}).union({Ops.COPY, Ops.ASSIGN, Ops.NOOP, Ops.SINK, Ops.ALLREDUCE}):
+    if self.op in (GroupOp.Elementwise-{Ops.BITCAST}).union({Ops.COPY, Ops.ASSIGN, Ops.NOOP, Ops.GROUP, Ops.SINK, Ops.ALLREDUCE}):
      # TODO: remove this hack for 3 op assign
      input_shapes = [x._shape for x in (self.src[:2] if self.op is Ops.ASSIGN else self.src) if x._shape is not None]
      if len(input_shapes) == 0: return None
@@ -268,20 +278,19 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
  @property
  def size(self) -> int: return prod([int(x.vmax) if isinstance(x, UOp) else x for x in self.shape])

+  @functools.cached_property
+  def ended_ranges(self):
+    if self.op in range_start: return self.src[range_start[self.op]:]
+    return ()
+
  # determine what ranges this is in
  @recursive_property
  def _ranges(self) -> dict[UOp, None]:
    ret: dict[UOp, None] = {}
-    if self.op in range_start.keys():
-      for s in self.src[:range_start[self.op]]: ret.update(s.ranges)
-      for s in UOp.sink(*self.src[range_start[self.op]:]).ranges:
+    for s in self.src: ret.update(s.ranges)
+    if (er:=self.ended_ranges):
+      for s in UOp.sink(*er).ranges:
        if s in ret: del ret[s]
-    elif self.op is Ops.END:
-      for s in self.src[self.arg:]: ret.update(s.ranges)
-      for s in UOp.sink(*self.src[:self.arg]).ranges:
-        if s in ret: del ret[s]
-    else:
-      for s in self.src: ret.update(s.ranges)
    return ret

  @property
@@ -289,13 +298,6 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
    if self.op is Ops.RANGE: return {self:None}
    return self._ranges

-  @functools.cached_property
-  def ended_ranges(self):
-    match self.op:
-      case Ops.REDUCE: return self.src[1:]
-      case Ops.END: return self.src[:self.arg]
-      case _: raise RuntimeError(f"{self.op} doesn't end ranges")
-
  # *** uop evaluation ***

  def simplify(self, tracked=False, full_symbolic=True):
@@ -304,6 +306,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
    with Context(TRACK_MATCH_STATS=0 if not tracked else TRACK_MATCH_STATS.value):
      return graph_rewrite(self, symbolic if full_symbolic else commutative, name="simplify")
  def ssimplify(self) -> UOp|ConstType: return ret.arg if (ret:=self.simplify()).op is Ops.CONST else ret
+  def sintify(self) -> sint: return self.arg if self.op is Ops.CONST else self
  def _eval(self, dtype, expected_type:Type[T]) -> T:
    assert self.dtype in dtype, f"eval with wrong dtype {self}"
    vmin, vmax = (simple_self:=self.simplify())._min_max
@@ -333,6 +336,9 @@ class UOp(MathTrait, metaclass=UOpMetaClass):

  def sink(*srcs:UOp|None, **kwargs):  # pylint: disable=no-self-argument
    return UOp(Ops.SINK, dtypes.void, tuple([x for x in srcs if x is not None]), **kwargs)
+  def group(*srcs:UOp|None):  # pylint: disable=no-self-argument
+    if len(srcs) == 1 and isinstance(srcs[0], UOp): return srcs[0]
+    return UOp(Ops.GROUP, dtypes.void, tuple([x for x in srcs if x is not None]))
  def detach(self): return UOp(Ops.DETACH, self.dtype, (self,))
  def index(self, *srcs:UOp|None, **kwargs):
    return UOp(Ops.INDEX, kwargs.pop("dtype", self.dtype), (self,)+tuple([x for x in srcs if x is not None]), **kwargs)
@@ -361,10 +367,10 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
    return UOp(Ops.GEP, self.dtype.scalar().vec(len(i)) if len(i) > 1 else self.dtype.scalar(), (self,), i)
  def load(self, *src:UOp, **kwargs): return UOp(Ops.LOAD, dtype=kwargs.pop("dtype", self.dtype.base), src=(self,)+src, **kwargs)
  def store(self, *src:UOp, **kwargs): return UOp(Ops.STORE, kwargs.pop("dtype", dtypes.void), (self,)+src, **kwargs)
-  def end(self, *src:UOp, ends:Sequence[UOp]):
-    if len(ends) == 0: return self
-    return UOp(Ops.END, src=(*ends, self, *src), arg=len(ends))
-  def after(self, *src:UOp): return UOp(Ops.AFTER, self.dtype, (self,)+src)
+  def end(self, *src:UOp):
+    if len(src) == 0: return self
+    return UOp(Ops.END, src=(self,)+src)
+  def after(self, *src:UOp, **kwargs): return UOp(Ops.AFTER, self.dtype, (self,)+src, **kwargs)
  def assign(self, x:UOp): return UOp(Ops.ASSIGN, self.dtype, (self, x))
  def barrier(self, *src:UOp): return UOp(Ops.BARRIER, src=(self,)+src)
  def alu(self, op, *src:UOp, **kwargs):
@@ -372,18 +378,25 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
    if op in {Ops.CMPLT, Ops.CMPNE, Ops.CMPEQ}: out_dtype = dtypes.bool.vec(out_dtype.count) if out_dtype.count > 1 else dtypes.bool
    return UOp(op, out_dtype, (self,)+src, **kwargs)
  @staticmethod
-  def const(dtype:DType, b:ConstLike, device:str|tuple[str, ...]|None=None, shape:tuple[sint, ...]|None=None, src=None):
+  def const(dtype:DType, b:ConstLike, device:str|tuple[str, ...]|None=None, shape:tuple[sint, ...]|None=None, src=None, unique:bool|int=False):
    if isinstance(b, UOp): return b.unbind()[0] if b.op is Ops.BIND else b
    if isinstance(b, tuple) and all_same(b): b = b[0]  # doesn't have to be a VCONST if they are all the same
+    # NOTE: float('nan') != float('nan'), so we canonicalize here
+    if isinstance(b, float) and math.isnan(b): b = math.nan
    ret = UOp(Ops.VCONST if isinstance(b, tuple) else Ops.CONST, dtype, arg=dtypes.as_const(b, dtype), src=() if src is None else (src,))
-    if device is not None: ret = ret.replace(src=(UOp(Ops.DEVICE, arg=device),))
+    if device is not None:
+      if unique or not isinstance(unique, bool): ret = ret.replace(src=(UOp(Ops.DEVICE, arg=device), UOp.unique(None if unique is True else unique)))
+      else: ret = ret.replace(src=(UOp(Ops.DEVICE, arg=device),))
+    elif unique or not isinstance(unique, bool): raise RuntimeError("unique consts only with DEVICE")
    if shape is not None: ret = ret.reshape((1,)*len(shape)).expand(shape)
    return ret
  @staticmethod
-  def range(end:sint, *arg):
+  def range(end:sint, *arg, dtype=dtypes.index, src=(), **kwargs):
    if len(arg) == 0: raise RuntimeError("range needs an arg")
    if len(arg) == 1: arg = arg+(AxisType.LOOP,)
-    return UOp(Ops.RANGE, dtype=dtypes.index, src=(sint_to_uop(end),), arg=arg)
+    return UOp(Ops.RANGE, dtype=dtype, src=(sint_to_uop(end, dtype),)+src, arg=arg, **kwargs)
+  @staticmethod
+  def special(end:sint, name:str, dtype=dtypes.index): return UOp(Ops.SPECIAL, dtype=dtype, src=(sint_to_uop(end, dtype),), arg=name)
  def r(self, op:Ops, axis:tuple[int, ...]):
    axis = tuple(sorted([x for x in axis if resolve(self.shape[x] != 1)]))
    return UOp(Ops.REDUCE_AXIS, self.dtype, (self,), (op, axis)) if len(axis) else self
@@ -486,7 +499,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
    match self.op:
      case Ops.CONST: return self.arg
      case Ops.VCONST: return self.arg[i]
-      case Ops.VECTORIZE: return cast(sint, self.src[i].ssimplify())
+      case Ops.VECTORIZE: return self.src[i].sintify()
      case _: raise RuntimeError(f"no sgep on {self.op}")

  @functools.cached_property
@@ -508,7 +521,8 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
      if len(arg) == 0: usrcs.append(UOp(Ops.VECTORIZE, dtypes.index.vec(0)))
      elif all(isinstance(x, int) for x in arg): usrcs.append(UOp.const(dtypes.index.vec(len(arg)), arg))
      else: usrcs.append(UOp(Ops.VECTORIZE, dtypes.index.vec(len(arg)), tuple(UOp.const(dtypes.index, x) if isinstance(x, int) else x for x in arg)))
-    ret = UOp(op, self.dtype, (self,)+tuple(usrcs), arg if len(usrcs) == 0 else None)
+    if len(usrcs) == 0: ret = UOp(op, self.dtype, (self,), arg)
+    else: ret = UOp(op, self.dtype, (self,)+UOp.sink(*usrcs).simplify().src)
    # for all movement ops, we check shape property
    if ret.shape == self.shape and same_shape_noop: return self
    return ret
@@ -529,12 +543,13 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
  # TODO: use this in Buffer
  unique_num = itertools.count(0)
  @staticmethod
-  def unique(): return UOp(Ops.UNIQUE, arg=next(UOp.unique_num))
+  def unique(arg:int|None=None): return UOp(Ops.UNIQUE, arg=next(UOp.unique_num) if arg is None else arg)

  # *** uop Buffer stuff ***

  @staticmethod
-  def new_buffer(device:str|tuple[str, ...], size:int, dtype:DType): return UOp(Ops.BUFFER, dtype, (UOp.unique(), UOp(Ops.DEVICE, arg=device)), size)
+  def new_buffer(device:str|tuple[str, ...], size:int, dtype:DType, num=None):
+    return UOp(Ops.BUFFER, dtype, (UOp.unique(num), UOp(Ops.DEVICE, arg=device)), size)
  @property
  def device(self) -> str|tuple[str, ...]: return cast(str|tuple[str, ...], unwrap(self._device))
  @recursive_property
@@ -555,8 +570,8 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
    if self.op is Ops.BUFFER: return self
    if self.op is Ops.MSELECT: return self.src[0].buf_uop.mselect(self.arg)
    if self.op is Ops.MSTACK: return UOp(Ops.MSTACK, self.dtype, src=tuple(x.buf_uop for x in self.src))
-    assert self.op is Ops.AFTER, f"must be AFTER {self.op}"
-    return self.src[0].buf_uop.base
+    assert self.base.op is Ops.AFTER, f"must be AFTER {self.base.op}"
+    return self.base.src[0].buf_uop.base

  def as_buf(self) -> UOp:
    if self.op is Ops.MSELECT: return self.src[0].as_buf().mselect(self.arg)
@@ -661,8 +676,8 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
    return math.prod([*count.elements(), terms[0].const_like(math.gcd(*factors))])  # put the const at the top
  def divide_exact(self, v:UOp) -> UOp|None:
    if self is v: return self.const_like(1)
-    if self.op is Ops.ADD: return None if (s0:=self.src[0].divide_exact(v)) is None or (s1:=self.src[1].divide_exact(v)) is None else s0+s1
    if v.op is Ops.CONST: return self.divides(v.arg)
+    if self.op is Ops.ADD: return None if (s0:=self.src[0].divide_exact(v)) is None or (s1:=self.src[1].divide_exact(v)) is None else s0+s1
    if self.op is Ops.MUL:
      (fac, const), (div_fac, div_const) = self.pop_const(Ops.MUL), v.pop_const(Ops.MUL)
      new_count = collections.Counter(fac.split_uop(Ops.MUL))
@@ -681,7 +696,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
      (s0_vmin, s0_vmax), (s1_vmin, s1_vmax) = self.src[0]._min_max, self.src[1]._min_max
      if self.op is Ops.ADD: return s0_vmin+s1_vmin, s0_vmax+s1_vmax
      if self.op is Ops.SUB: return s0_vmin-s1_vmax, s0_vmax-s1_vmin
-      if self.op is Ops.AND and s1_vmin == s1_vmax and s0_vmin >= 0 and s1_vmin >= 0: return min(0, s0_vmin), min(s0_vmax, s1_vmax)
+      if self.op is Ops.AND and dtypes.is_int(self.dtype) and s1_vmin == s1_vmax >= 0 and s0_vmin >= 0: return min(0, s0_vmin), min(s0_vmax, s1_vmax)
      if self.op is Ops.MUL: return min(vals:=(s0_vmin*s1_vmin, s0_vmin*s1_vmax, s0_vmax*s1_vmin, s0_vmax*s1_vmax)), max(vals)
      # SHL/SHR on consts only
      if self.op is Ops.SHL and s1_vmin == s1_vmax and all_int(t:=(s0_vmin, s0_vmax, s1_vmin)): return t[0] << t[2], t[1] << t[2]
@@ -696,9 +711,8 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
      if self.op is Ops.MAX: return max(s0_vmin, s1_vmin), max(s0_vmax, s1_vmax)
      if self.op is Ops.CMPLT: return (s0_vmax<s1_vmin, s0_vmin<s1_vmax)
      if self.op is Ops.CMPNE: return ((s0_vmax < s1_vmin) or (s1_vmax < s0_vmin), not (s0_vmin == s0_vmax == s1_vmin == s1_vmax))
-      if self.dtype == dtypes.bool:
-        if self.op is Ops.OR: return s0_vmin or s1_vmin, s0_vmax or s1_vmax
-        if self.op is Ops.AND: return s0_vmin and s1_vmin, s0_vmax and s1_vmax
+      if self.op is Ops.OR and self.dtype == dtypes.bool: return s0_vmin or s1_vmin, s0_vmax or s1_vmax
+      if self.op is Ops.AND and self.dtype == dtypes.bool: return s0_vmin and s1_vmin, s0_vmax and s1_vmax
    # float has NAN issue and we use explicit NAN in transcendental
    if self.op is Ops.WHERE and dtypes.is_int(self.dtype): return min(self.src[1].vmin, self.src[2].vmin), max(self.src[1].vmax, self.src[2].vmax)
    # NOTE: returned UOp is assumed to be CONST
@@ -753,7 +767,7 @@ def safe_pow(x, y):

 python_alu: dict[Ops, Callable]  = {
  Ops.LOG2: lambda x: math.log2(x) if x > 0 else -math.inf if x == 0 else math.nan, Ops.EXP2: safe_exp2,
-  Ops.SQRT: lambda x: math.sqrt(x) if x >= 0 else math.nan, Ops.RECIP: lambda x: 1/x if x != 0 else math.copysign(math.inf, x),
+  Ops.SQRT: lambda x: math.sqrt(x) if x >= 0 else math.nan, Ops.RECIPROCAL: lambda x: 1/x if x != 0 else math.copysign(math.inf, x),
  Ops.SIN: lambda x: math.sin(x) if not math.isinf(x) else math.nan, Ops.POW: safe_pow, Ops.TRUNC: math.trunc,
  Ops.NEG: operator.neg, Ops.ADD: operator.add, Ops.SUB: operator.sub, Ops.MUL: operator.mul, Ops.CMPNE: operator.ne, Ops.CMPLT: operator.lt,
  Ops.XOR: operator.xor, Ops.OR: operator.or_, Ops.AND: operator.and_, Ops.SHR: operator.rshift, Ops.SHL: operator.lshift, Ops.MAX: max,
@@ -771,7 +785,7 @@ def exec_alu(op:Ops, dtype:DType, operands, truncate_output=True):
 def print_uops(uops:list[UOp]):
  for i,u in enumerate(uops):
    formatted_srcs = [(uops.index(x) if x.op is not Ops.CONST else f"{x.arg}") if x in uops else "--" for x in u.src]
-    print(f"{i:4d} {str(u.op):20s}: {str(u.dtype):30s} " f"{str(formatted_srcs):32s} {u.arg}")
+    print(f"{i:4d} {str(u.op):20s}: {str(u.dtype):40s} " f"{str(formatted_srcs):32s} {u.arg}")

 # ***** pattern matcher *****

@@ -845,7 +859,8 @@ class UPat(MathTrait):

  # copied from UOp
  def sink(self, *srcs:UPat|None, **kwargs): return UPat(Ops.SINK, dtypes.void, (self,)+tuple([x for x in srcs if x is not None]), **kwargs)
-  def index(self, idx:UPat, valid:UPat|None=None): return UPat(Ops.INDEX, self.dtype, (self,idx,valid) if valid is not None else (self,idx))
+  def index(self, idx:UPat, valid:UPat|None=None, **kwargs):
+    return UPat(Ops.INDEX, self.dtype, (self,idx,valid) if valid is not None else (self,idx), **kwargs)
  def cast(self, dtype=None, **kwargs): return UPat(Ops.CAST, dtype, (self,), **kwargs)
  def bitcast(self, dtype=None): return UPat(Ops.BITCAST, dtype, (self,))
  def gep(self, i:int|None=None, **kwargs): return UPat(Ops.GEP, None, (self,), (i,) if i is not None else None, **kwargs)
@@ -856,6 +871,7 @@ class UPat(MathTrait):
  def fuse(self): return self.alu(Ops.FUSE)
  def broadcast(self, **kwargs): return UPat(Ops.VECTORIZE, self.dtype, src=self, **kwargs)
  def contiguous(self, *args, **kwargs): return UPat(Ops.CONTIGUOUS, dtype=self.dtype, src=(self,)+args, **kwargs)
+  def after(self, *src:UPat, **kwargs): return UPat(Ops.AFTER, self.dtype, (self,)+src, **kwargs)

  def const_like(self, b:ConstLike): return UPat.const(self.dtype, cast(ConstType, b))
  def alu(self, op:Ops, *src:UPat):
@@ -1063,7 +1079,8 @@ if TRACK_MATCH_STATS or PROFILE:
    if not int(os.getenv("VIZ", "0")) and not int(os.getenv("PROFILE", "0")) and not int(os.getenv("SQTT", "0")):
      args = ['--kernels', getenv("VIZ_DATA", "")] if getenv("VIZ_DATA", "") else []
      args += ['--profile', getenv("PROFILE_DATA", "")] if getenv("PROFILE_DATA", "") else []
-      os.execv(sys.executable, [sys.executable] + [pathlib.Path(__file__).resolve().parent.parent / "viz" / "serve.py"] + args)
+      viz_path = pathlib.Path(__file__).resolve().parent.parent / "viz" / "serve.py"
+      os.execv(sys.executable, [sys.executable, viz_path.as_posix()] + args)

 # *** simple graph rewrite engine ***

@@ -1164,7 +1181,7 @@ def graph_rewrite_map(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False, na
    for k,v in input_map.items(): new_map[k] = new_map.get(v,v)
  return new_map

-def sint_to_uop(x:sint) -> UOp: return UOp.const(dtypes.index, x) if isinstance(x, int) else x.cast(dtypes.index)
+def sint_to_uop(x:sint, dtype=dtypes.index) -> UOp: return UOp.const(dtype, x) if isinstance(x, int) else x.cast(dtype)

 def select_dtype(u): return (dtypes.long if u.overflows(dtypes.int32) else dtypes.int).vec(u.dtype.count)
 pm_lower_index_dtype = PatternMatcher([
@@ -1189,8 +1206,8 @@ pm_lower_index_dtype = PatternMatcher([
  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx", dtypes.ints).cast(), UPat.var("valid"))), lambda buf,idx,valid: buf.index(idx, valid)),
  (UPat((Ops.STORE, Ops.LOAD), src=(UPat(), UPat(), UPat().cast(dtypes.index)), allow_any_len=True, name="s"),
    lambda s: s.replace(src=s.src[:2]+tuple(u.src[0] for u in s.src[2:]))),
-  # TODO: this is only triggering if they are all casts, correct?
-  (UPat((Ops.SINK, Ops.NOOP), src=UPat().cast(dtypes.index), name="n"), lambda n: n.replace(src=tuple(s.src[0] for s in n.src))),
+  (UPat((Ops.SINK, Ops.NOOP, Ops.END), name="n"),
+   lambda n: n.replace(src=tuple(s.src[0] if s.op is Ops.CAST and s.dtype == dtypes.index else s for s in n.src))),
 ])
 def _index_to_concrete_int(u:UOp): return graph_rewrite(u.sink(), pm_lower_index_dtype).src[0]

@@ -1215,7 +1232,7 @@ renderer = PatternMatcher([
  (UPat(Ops.BIND, src=UPat(Ops.NOOP), name="x"), lambda x: x.src[0]),
  #(UPat(Ops.BIND, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"{x.src[0].arg}[={x.src[1].arg}]")),
  (UPat(Ops.NEG, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"(-{x.src[0].arg})")),
-  (UPat(Ops.RECIP, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"(1/{x.src[0].arg})")),
+  (UPat(Ops.RECIPROCAL, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"(1/{x.src[0].arg})")),
  (UPat(Ops.MAX, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"max({x.src[0].arg}, {x.src[1].arg})")),
  (UPat(Ops.MULACC, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"({x.src[0].arg}*{x.src[1].arg}+{x.src[2].arg})")),
  (UPat(Ops.WHERE, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"({x.src[1].arg} if {x.src[0].arg} else {x.src[2].arg})")),
@@ -1231,40 +1248,95 @@ renderer_infer = PatternMatcher([
  *renderer.patterns
 ])

-sugar = { Ops.SINK: "sink", Ops.STORE: "store", Ops.LOAD: "load", Ops.SQRT: "sqrt", Ops.INDEX: "index", Ops.REDUCE: "reduce",
-          Ops.WHERE: "where", Ops.RECIP: "reciprocal", Ops.EXP2: "exp2", Ops.LOG2: "log2", Ops.SIN: "sin"}
-pm_pyrender = PatternMatcher([
-  (UPat(Ops.CONST, src=(UPat(Ops.NOOP),), name="x"), lambda x: UOp(Ops.NOOP, arg=f"UOp.const({x.dtype}, {x.arg}, src={x.src[0].arg})")),
-  (UPat(Ops.CONST, name="x"), lambda x: UOp(Ops.NOOP, arg=f"UOp.const({x.dtype}, {x.arg})")),
-  (UPat(Ops.CAST, src=(UPat(Ops.NOOP),), name="x"), lambda x: UOp(Ops.NOOP, arg=f"{x.src[0].arg}.cast({x.dtype})")),
-  (UPat(Ops.BITCAST, src=(UPat(Ops.NOOP),), name="x"), lambda x: UOp(Ops.NOOP, arg=f"{x.src[0].arg}.bitcast({x.dtype})")),
-  (UPat({Ops.MAX, Ops.THREEFRY, Ops.CMPLT, Ops.CMPNE, Ops.POW}, src=UPat(Ops.NOOP), name="x"),
-   lambda x: UOp(Ops.NOOP, arg=f"{x.src[0].arg}.alu({x.op}, {x.src[1].arg})")),
-  (UPat(Ops.RANGE, src=(UPat(Ops.NOOP),), name="x"), lambda x:
-    UOp(Ops.NOOP, arg=f"UOp.range({x.src[0].arg}, {str(x.arg[0])}, {str(x.arg[1])})")),
-  (UPat(set(sugar.keys()), src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP,
-    arg=f"{x.src[0].arg}.{sugar[x.op]}({', '.join([y.arg for y in x.src[1:]] + ([f'arg={str(x.arg)}'] if x.arg is not None else []))})")),
-  (UPat(Ops.REDUCE_AXIS, src=(UPat(Ops.NOOP),), name="x"),
-   lambda x: UOp(Ops.NOOP, arg=f"{x.src[0].arg}.f({x.op}, arg=({', '.join([str(y) for y in x.arg])}))")),
+# *** pyrender ***
+
+def srcs(ctx, src): return f"({ctx[src[0]]},)" if len(src) == 1 else f"({', '.join([ctx[x] for x in src])})"
+def render_marg(ctx,x:UOp):
+  if x.op in {Ops.PERMUTE, Ops.FLIP}: return str(x.marg)
+  pieces = []
+  if x.op in {Ops.RESHAPE, Ops.EXPAND}:
+    pieces = [f"{ctx[a] if isinstance(a, UOp) else str(a)}" for a in x.marg]
+  if x.op in {Ops.PAD, Ops.SHRINK}:
+    pieces = [f"({ctx[a[0]] if isinstance(a[0], UOp) else str(a[0])}, {ctx[a[1]] if isinstance(a[1], UOp) else str(a[1])})" for a in x.marg]
+  return f"({','.join(pieces)})" if len(pieces) != 1 else f"({pieces[0]},)"
+
+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.ASSIGN, Ops.DETACH}
+pm_pyrender_extra = PatternMatcher([
+  (UPat(Ops.CONST, src=(UPat(Ops.DEVICE, name="d"), UPat(Ops.UNIQUE, name="u")), name="x"),
+   lambda x,d,u: f"UOp.const({x.dtype}, {x.arg}, 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, name="x"), lambda x: f"UOp.const({x.dtype}, {x.arg})"),
+  (UPat(Ops.DEFINE_VAR, src=(), name="x"), lambda x:
+    f"UOp.variable(\"{x.arg[0]}\", {x.arg[1]}, {x.arg[2]}{', dtype='+str(x.dtype) if x.dtype is not dtypes.index else ''})"),
+  (UPat((Ops.CAST, Ops.BITCAST), name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({x.dtype})"),
+  (UPat(Ops.SPECIAL, src=(UPat(Ops.CONST),), name="x"), lambda x: f"UOp.special({x.src[0].arg}, {repr(x.arg)}, dtype={x.dtype})"),
+  (UPat(Ops.BUFFER, src=(UPat(Ops.UNIQUE, name="u"), UPat(Ops.DEVICE, name="d")), name="x"), lambda x,u,d:
+    f"UOp.new_buffer({repr(d.arg)}, {x.size}, {x.dtype}, {u.arg})"),
+  (UPat(Ops.COPY, src=(UPat(name="x"), UPat(Ops.DEVICE, name="d"))), lambda ctx,x,d: f"{ctx[x]}.copy_to_device({repr(d.arg)})"),
+  (UPat(Ops.REDUCE_AXIS, name="r"), lambda ctx,r: f"{ctx[r.src[0]]}.r({r.arg[0]}, {r.arg[1]})"),
+  # NOTE: range has srcs sometimes after control flow
+  (UPat(Ops.RANGE, src=(UPat(Ops.CONST, name="c"),), allow_any_len=True, name="x"), lambda ctx,x,c:
+    "UOp.range("+', '.join([str(c.arg)] + [str(y) for y in x.arg])+
+      (f', src={srcs(ctx, x.src[1:])}' if len(x.src) > 1 else '')+(', dtype='+str(x.dtype) if x.dtype is not dtypes.index else '')+")"),
+  # TODO: index shouldn't mismatch dtype
+  (UPat(Ops.INDEX, src=(UPat(), UPat()), name="x"), lambda ctx,x:
+   f"{ctx[x.src[0]]}.index({ctx[x.src[1]]}, dtype={x.dtype})" if x.src[0].dtype != x.dtype else None),
+  # TODO: fix forced_reshape
+  (UPat(Ops.RESHAPE, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.forced_reshape({render_marg(ctx,x)})" if x.src[0].shape == x.shape else None),
+  (UPat(GroupOp.Movement, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({render_marg(ctx,x)})"),
+  # NOTE: CMPNE doesn't work cause there's no __rne__
+  (UPat(set(syms.keys())-{Ops.SUB, Ops.CMPNE}, src=(UPat(Ops.CONST, name="y"), UPat(name="z")), name="x"),
+   lambda ctx,x,y,z: f"({y.arg}{syms[x.op]}{ctx[z]})"),
+  # NOTE: sub doesn't work cause it's written as add/mul
+  (UPat(set(syms.keys())-{Ops.SUB}, src=(UPat(name="y"), UPat(Ops.CONST, name="z")), name="x"), lambda ctx,x,y,z: f"({ctx[y]}{syms[x.op]}{z.arg})"),
+  (UPat(set(syms.keys())-{Ops.SUB}, name="x"), lambda ctx,x: f"({ctx[x.src[0]]}{syms[x.op]}{ctx[x.src[1]]})"),
+  (UPat(sugar, src=(), name="x"), lambda x: f"UOp.{x.op.name.lower()}("+', '.join(([f'arg={repr(x.arg)}'] if x.arg is not None else []))+")"),
+  (UPat(sugar, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}("+', '.join([ctx[y] for y in x.src[1:]] + \
+    ([f'arg={repr(x.arg)}'] if x.arg is not None else []))+")"),
 ])

-@Context(SPEC=0)
-def pyrender(ast:UOp) -> list[str]:
-  cmap = ast.get_consumer_map()
-  to_render = set()
-  for u in ast.toposort():
-    if u.op is Ops.STORE: to_render.add(u.src[1])
-    if len(cmap[u]) == 1 and u.op not in {Ops.DEFINE_GLOBAL, Ops.LOAD} or u.op in {Ops.CONST}: continue
+# NOTE: you can remove pm_pyrender_extra and it'll still be correct
+pm_pyrender = pm_pyrender_extra+PatternMatcher([
+  (UPat(Ops.KERNEL, name="u"), lambda ctx,u: f"UOp(Ops.KERNEL, src={srcs(ctx,u.src)}, arg=Kernel({ctx[u.arg.ast]}(), {u.arg.metadata}))"),
+  (UPat(GroupOp.All, name="u"), lambda ctx,u: f"UOp({u.op}, {u.dtype}, {srcs(ctx,u.src)}"+(f", {repr(u.arg)})" if u.arg is not None else ")")),
+])
+
+def pyrender(ast:UOp) -> str:
+  lst = list(ast.toposort())
+
+  cmap = consumer_map_from_toposort(lst)
+  not_rendered = {Ops.CONST, Ops.VCONST, Ops.DEVICE}
+  always_rendered = {Ops.DEFINE_GLOBAL, Ops.LOAD, Ops.SPECIAL, Ops.RANGE, Ops.CONTIGUOUS, Ops.VECTORIZE,
+                     Ops.BUFFER, Ops.COPY, Ops.KERNEL, Ops.WHERE, Ops.END, Ops.ASSIGN}
+
+  to_render: set[UOp] = {ast}
+  for u in lst:
    if u.op in {Ops.SINK}:
      for s in u.src: to_render.add(s)
+    if u.op is Ops.STORE: to_render.add(u.src[1])
+    if u.op in {Ops.REDUCE, Ops.REDUCE_AXIS}: to_render.add(u.src[0])
+    if u.op in not_rendered: continue
+    # checking the consumers is not enough, you have to make sure it's not used twice by the one consumer
+    if len(cmap[u]) == 1 and len([x for x in list(cmap[u].keys())[0].src if x is u]) == 1 and u.op not in always_rendered: continue
    to_render.add(u)
-  ret: list[str] = []
-  rep: dict[UOp, UOp] = {}
-  for u in ast.toposort():
-    if u not in to_render: continue
-    ret.append(f"c{len(ret)} = {u.substitute(rep).render(simplify=False, pm=pm_pyrender+renderer)}")
-    rep[u] = UOp(Ops.NOOP, arg=f"c{len(ret)-1}")
-  return ret[0:-1] + ["ast ="+ret[-1].split("=", 1)[1]]
+
+  kernels: dict[UOp, tuple[str, str]] = {}
+  r: dict[UOp, str] = {}
+  ret: dict[str, str] = {}
+  for i,u in enumerate(lst):
+    if u.op is Ops.KERNEL:
+      if u.arg.ast not in kernels:
+        kernels[u.arg.ast] = (f"k{len(kernels)}", f"def k{len(kernels)}():\n  " + pyrender(u.arg.ast).replace('\n', '\n  ') + "\n  return ast\n\n")
+      r[u.arg.ast] = kernels[u.arg.ast][0]
+    ren = cast(str, pm_pyrender.rewrite(u, ctx=r))
+    assert isinstance(ren, str)
+    if u.tag is not None: ren += f".rtag({u.tag})"
+    if u not in to_render: r[u] = ren
+    else:
+      r[u] = f"c{i}" if u is not lst[-1] else "ast"
+      ret[r[u]] = ren
+  return ''.join([v[1] for v in kernels.values()]) + '\n'.join([f"{k} = {v}" for k,v in ret.items()])

 # *** what was symbolic.py ***

--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
George HotzandGitHub	bb7aa19d67	Merge branch 'master' into replace_if_with_range	2025-10-28 23:02:15 +08:00
George HotzandGitHub	5e01cc299b	zero len ranges fail (#12974 ) * zero len ranges fail * fix Python backend * fix llvm * fix ptx * yolo fix nir * this works... * always store... * always store... * Revert "always store..." This reverts commit `0816cf344d`.	2025-10-28 22:49:55 +08:00
George HotzandGitHub	e936aa7974	cleanups from if range branch (#12973 )	2025-10-28 20:58:47 +08:00
geohot	117f37ae5f	don't remove the gate	2025-10-28 19:28:38 +08:00
geohot	ba23b097c1	fix image	2025-10-28 19:22:59 +08:00
geohot	3341272771	cleanup patterns	2025-10-28 19:15:01 +08:00
geohot	975f5ccc99	tests pass	2025-10-28 19:12:21 +08:00
qazalandGitHub	901d27b3ba	viz: optional text dims try 2 (#12971 )	2025-10-28 18:54:28 +08:00
geohot	d0de209ad0	don't brick on that	2025-10-28 18:49:00 +08:00
geohot	2d87d89202	replace if with range	2025-10-28 18:30:11 +08:00
geohot	f5a3b33d33	add fun with nhwc convs	2025-10-28 17:12:22 +08:00
George HotzandGitHub	907499b02c	clean up GROUP/SINK (#12969 ) * clean up GROUP/SINK * fix end * range_str color	2025-10-28 16:08:10 +08:00
Sieds LyklesandGitHub	e22c5e7e73	process_replay uses opts argument for KernelInfo.opts_to_apply (#12946 ) * opts_to_apply is opts * skip beamed kernels * simpler change * fix the tensor cores tests for process replay * use opts	2025-10-28 09:00:28 +01:00
George HotzandGitHub	6c9560a846	more syntactic sugar for pyrender (#12968 )	2025-10-28 15:24:33 +08:00
George HotzandGitHub	b0da173f2f	add unique to const, fix longstanding bug (#12965 ) * add unique to const, fix longstanding bug * _force_unique=True * fix tests * fix more tests	2025-10-28 15:11:37 +08:00
Sieds LyklesandGitHub	e110f4632a	split cat (on cpu) (#12864 ) * split ranges but only on cpu * except KernelOptError for threads * use GROUP and END * no more flatten_range needed * remove noop end * always process replay for openpilot * update test * skip test * fix in outs calculation With the new linearizer the toposort is a problem, this matches the spec now * undo that	2025-10-28 07:55:19 +01:00
qazalandGitHub	3b82dee625	viz: match DEBUG=2 for exec item metadata (#12966 ) * viz: match DEBUG=2 for exec item metadata * remove repr from kernel	2025-10-28 14:53:57 +08:00
qazalandGitHub	99589dea81	move viz edge tagging to UOp graph (#12964 )	2025-10-28 12:46:23 +08:00
George HotzandGitHub	bbe0bebbf3	no range tags in kernels (#12962 )	2025-10-28 12:33:48 +08:00
George HotzandGitHub	39c2117dea	cleanup pyrender (#12961 )	2025-10-28 10:47:39 +08:00
George HotzandGitHub	2832954bcb	test with IGNORE_OOB=0 (#12960 )	2025-10-28 10:32:19 +08:00
George HotzandGitHub	7784cec48e	pytest-split on spec (#12959 )	2025-10-28 10:09:01 +08:00
George HotzandGitHub	4d817a289e	simplify spec (#12958 ) * simplify spec * more	2025-10-28 09:52:32 +08:00
George HotzandGitHub	62e62d8760	move verify to spec / cleanup (#12956 ) * move verify to spec / cleanup * lil * more explicit	2025-10-28 08:58:10 +08:00
wozeparrotandGitHub	24884c6768	fix: don't use KITTENS_HOPPER for 4090 (#12954 )	2025-10-27 17:19:53 -07:00
nimlgenandGitHub	372d9e5753	hcq: helper for visible devices (#12950 ) * hcq: helper for visible devices * fix * f	2025-10-28 02:27:56 +08:00
Justin ErenkrantzandGitHub	f2ffe9c8cf	Apply an override for nbio 7.3.0 to 7.2.0. (#12949 )	2025-10-27 11:10:10 -07:00
qazalandGitHub	63484d837e	Revert "viz graph drawing cleanups (#12933 )" (#12947 ) This reverts commit `189582db5e`.	2025-10-28 00:39:37 +08:00
chenyuandGitHub	a79832b01f	control_flow.py -> linearizer.py [pr] (#12948 )	2025-10-27 12:38:13 -04:00
b1tg GitHub chenyu	45e2f916a3	add quantize fp8 in llama3 (#12893 ) * add quantize fp8 in llama3 * don't truncate fp8 alu result * cast to float32 before matmul * --model weights/LLaMA-3/8B-SF-DPO/ --------- Co-authored-by: chenyu <[email protected]>	2025-10-27 10:22:57 -04:00
George HotzandGitHub	25c2da1579	check SPEC=2 in CI (#12945 ) * check SPEC=2 in CI * split SPEC=2 * fast enough	2025-10-27 21:53:57 +08:00
Sieds LyklesandGitHub	072f7c35c5	fix in/outs calculation in ProgramSpec (#12937 ) With the new linearizer the toposort is a problem, this matches the spec now	2025-10-27 12:31:41 +01:00
qazalandGitHub	e93c9bf6a7	viz: extend main code block to full height (#12944 )	2025-10-27 18:43:49 +08:00
George HotzandGitHub	273b1f914d	new pyrender, tested with SPEC=2 (#12934 ) * pyrender always works with SPEC=3 * test pyrender * work * work * work * .sintify * v const * kernelize * pyrender * viz always * optional forced_reshape * cleanups	2025-10-27 18:41:51 +08:00
George HotzandGitHub	701a632907	move VECTORIZE/CONST (#12942 )	2025-10-27 17:37:13 +08:00
nimlgenandGitHub	95748a4518	nv: map vram after resets (#12938 )	2025-10-27 17:17:07 +08:00
George HotzandGitHub	8fb545c475	don't late simplify on marg (#12941 )	2025-10-27 17:07:41 +08:00
George HotzandGitHub	7139e036c5	bugfixes from pyrender (#12940 )	2025-10-27 16:56:53 +08:00
George HotzandGitHub	804133cffd	rename RECIP to RECIPROCAL (#12939 )	2025-10-27 16:53:13 +08:00
nimlgenandGitHub	f4da94af28	system: reset is a method of pcidevice (#12936 )	2025-10-27 16:21:10 +08:00
wozeparrotandGitHub	6b54378eba	working kitten matmul (#12935 )	2025-10-26 23:40:49 -07:00
qazalandGitHub	189582db5e	viz graph drawing cleanups (#12933 ) * viz: make node label dims optional * inplace edge updates * change that	2025-10-27 13:59:32 +08:00
qazalandGitHub	70ba84eb04	viz: generic node label centering (#12925 ) * viz: correct node label centering * matches * overlay * the other way	2025-10-27 12:02:34 +08:00
Sieds LyklesandGitHub	eaeaea2f9c	pyrender Ops.SPECIAL and use correct dtype for Ops.RANGE rendering (#12931 )	2025-10-27 03:21:34 +01:00
nimlgenandGitHub	8c1368cab6	system: class PCIBarInfo (#12930 ) * system: class PCIBarInfo * fix	2025-10-27 03:57:42 +08:00
nimlgenandGitHub	f00009c731	hcq: drivers take pcidev (#12929 ) * hcq: drivers take pcidev * fix nv	2025-10-26 20:43:51 +08:00
ttomsaandGitHub	99a519f068	linearizer cleanup (#12923 ) * cleanup * comments * also this	2025-10-26 18:30:12 +08:00
George HotzandGitHub	c0c24d3a70	cleanup wmma (#12927 ) * cleanup wmma * fix test_ops failures on android	2025-10-26 18:26:47 +08:00
George HotzandGitHub	0a32ab0006	nitpicks from typecheckers (#12926 ) * nitpicks from the typechecker * more	2025-10-26 17:52:55 +08:00
geohot	db5c918215	source extra/cl_android.sh to fix opencl on android	2025-10-26 15:27:51 +08:00
qazalandGitHub	c94e597b3e	viz ui selector cleanups (#12924 )	2025-10-26 14:40:47 +08:00
chenyuandGitHub	94701d4838	clean up divide_exact order [pr] (#12919 ) do the const first since ADD can also call into that	2025-10-25 18:47:57 -04:00
chenyuandGitHub	e18922f111	limit AND const min max to ints [pr] (#12918 )	2025-10-25 16:07:52 -04:00
nimlgenandGitHub	92324172be	amd: refactor usb into usbdevice (#12916 ) * amd: refactor usb into usbdevice * nu * my bad * ops * my bad	2025-10-26 01:00:19 +08:00
qazalandGitHub	3b192f5eac	split viz graph rendering from dag layout (#12914 )	2025-10-25 15:36:44 +08:00
George HotzandGitHub	6415e3e8a7	use Ops.GROUP instead of Ops.NOOP for merging stores (#12912 ) * use Ops.GROUP instead of Ops.NOOP for merging stores * fs noop	2025-10-25 12:26:12 +08:00
George HotzandGitHub	b4f6a2c7a3	add kernel spec (#12911 ) * add kernel spec * fix kernel spec	2025-10-25 11:49:20 +08:00
George HotzandGitHub	8a941d95a4	SPEC=2 is full spec, SPEC=1 is default (#12910 ) * SPEC=1 passes all tests * just use SPEC, not __debug__	2025-10-25 11:10:43 +08:00
wozeparrotandGitHub	456560c1ff	stateless tinyfs copyin (#12908 )	2025-10-24 19:18:38 -07:00
wozeparrotandGitHub	a5b0f57067	clean: cleanup tinyfs copyout (#12907 )	2025-10-24 18:32:55 -07:00
chenyuandGitHub	4b7329001d	clean up test_avg_pool3d (#12905 )	2025-10-24 14:31:36 -04:00
George HotzandGitHub	6b35467f53	stores don't end ranges (#12902 ) * early endrange * bugfixes	2025-10-24 23:05:03 +08:00
nimlgenandGitHub	5b5ba31a86	amd: make sqtt bufs uc (#12898 )	2025-10-24 18:55:14 +08:00
Sieds LyklesandGitHub	e1f8c82938	Onnx Layer/Group/RMS/Batch-Norm ReduceL2 fp32 intermediates for fp16 (#12109 ) * match onnx spec * use least_upper_dtype * promote the square * just cast before the square	2025-10-24 12:26:11 +02:00
George HotzandGitHub	0bde87d8d7	cleanups from flash attention branch (#12897 )	2025-10-24 14:14:56 +08:00
wozeparrotandGitHub	9dac505565	variable bs keccak (#10731 )	2025-10-23 14:10:21 -07:00
chenyuandGitHub	154b4f9f40	test FUSE_OPTIM=1 test/test_optim.py (#12895 )	2025-10-23 15:54:27 -04:00
chenyuandGitHub	6e4ee8deea	small heuristic cleanup [pr] (#12892 )	2025-10-23 10:50:15 -04:00
nimlgenandGitHub	f835566e27	sqtt: correct header (#12891 ) * sqtt: correct header * f	2025-10-23 22:37:17 +08:00
Sieds LyklesandGitHub	c1db62ff7c	move reduce collapse to rangeify (#12845 )	2025-10-23 15:44:17 +02:00
Sieds LyklesandGitHub	04b3e51f1b	remove old reduce collapse rule (#12889 ) * comment this out * remove	2025-10-23 13:51:49 +02:00
qazalandGitHub	cdfb8e31ae	hotfix: correct viz rewrite step counter reset (#12890 )	2025-10-23 19:47:16 +08:00
George HotzandGitHub	6df19a4ac6	lil qol improvements to viz (#12887 )	2025-10-23 18:41:07 +08:00
George HotzandGitHub	ff68a6263b	move locals into codegen (dedup works) (#12885 ) * move locals into codegen (dedup works) * move in optimize	2025-10-23 17:07:39 +08:00
George HotzandGitHub	ddb53d1d48	PCONTIG=3 both saves ram and flops (#12884 ) * PCONTIG=3 both saves ram and flops * group * gate locals * should be correct	2025-10-23 16:37:26 +08:00
qazalandGitHub	2a5c22436e	remove outdated docs (#12881 )	2025-10-23 12:52:36 +08:00
qazalandGitHub	bcc30e5e10	viz: add linearized UOp list view (#12883 ) * viz: add linearized UOp list view * lang	2025-10-23 12:52:14 +08:00
George HotzandGitHub	e85cee0aad	flip Ops.END srcs (#12882 ) * flip Ops.END srcs * backward * late end split	2025-10-23 12:47:50 +08:00
George HotzandGitHub	74b4cfe44b	Ops.GROUP + range check (#12880 ) * simpler * fix that * Ops.GROUP + range check * fix bugs * fix linter * fix test	2025-10-23 12:05:21 +08:00
Sieds LyklesandGitHub	914defd55d	give endrange priority (#12870 ) * uncomment line * try giving endrange priority	2025-10-23 05:19:13 +02:00
qazalandGitHub	2f95c10702	remu new instructions / use volatile in emulator tests (#12862 ) * remu new instructions * start moving to volatile * test_simple works * test_exec_mov works and lid is still here * test_exec_cmp_vopc * clang did s_mov_b32 exec_lo, 1 * don't hardcode v1 * support volatile in tests * hw_test passes * only the volatile version * subrev saturating behavior	2025-10-23 11:13:43 +08:00
George HotzandGitHub	e718254004	simpler end (#12879 ) * simpler * fix that	2025-10-23 10:35:58 +08:00
wozeparrotandGitHub	6e00dec95d	feat: pin openpilot 0.10.1 models (#12878 )	2025-10-22 14:57:54 -07:00
wozeparrotandGitHub	3a9aa05359	feat: extra nvcc options (#12876 )	2025-10-22 13:21:11 -07:00
chenyuandGitHub	f0831c8c30	add 0.10.0 to comma benchmark (#12875 ) * add 0.10.0 to comma benchmark disabled the 0.10.1 ones which are pinned to master. it does not work because benchmark uses the cached old version * that's pinned	2025-10-22 15:18:21 -04:00
nimlgenandGitHub	e7e535cd53	amd: sqtt for gfx9 (#12844 ) * amd: start sqtt for gfx9 * writes something, but sometimes zeroes * HEADER! * w * tiny * mypy	2025-10-23 02:31:07 +08:00
b1tgandGitHub	81108f91ee	amd tc: 16x16x32 (#12874 ) * amd tc: 16x16x32 * test * clean, test amd_cdna4	2025-10-22 13:48:01 -04:00
George HotzandGitHub	bf173c0a37	we don't support multi end yet (#12869 )	2025-10-22 23:43:32 +08:00
nimlgenandGitHub	a7bc0104c2	amd: clean up sqtt_stop (#12872 )	2025-10-22 22:17:03 +08:00
nimlgenandGitHub	b6eb9172ea	amd: fix ip offsets (#12867 )	2025-10-22 20:50:18 +08:00
geohot	174811fc0f	hotfix: slightly looser load spec for AMD bfloat16	2025-10-22 19:55:59 +08:00
George HotzandGitHub	7762b3558b	clean up the spec (#12868 ) * tighten up the spec * move validate into a different file * that moved to validate * after(barr)	2025-10-22 19:50:42 +08:00
George HotzandGitHub	726988fa4b	late ifs try 2 (#12865 ) * late ifs try 2 * fix image * fix that test * panic * ptx fixups * preserve toposort * those pass locally * Revert "those pass locally" This reverts commit `063409f828`. * no ls * make that explicit	2025-10-22 18:49:27 +08:00
George HotzandGitHub	6abe90fb7c	fix linearizer non-determinism (#12866 )	2025-10-22 17:51:35 +08:00
qazalandGitHub	cebc2b5721	cleanup viz profiler metadata ui (#12860 ) * cleanup viz profiler metadata ui * text * select over .args * space	2025-10-22 17:31:12 +08:00
Sieds LyklesandGitHub	8d0256c46b	Move gate to load for loaded index (#12861 ) * change condition * change test to better represent how the uop looks irl	2025-10-22 09:53:07 +02:00
chenyuandGitHub	6d86e962c7	update ASSERT_MIN_STEP_TIME (#12857 ) 0.10.1 driving_policy is good now, still need driving_vision and dmonitoring to be fast	2025-10-21 22:46:07 -04:00
George HotzandGitHub	92778c7a8b	rename opts to ren, add store ranges back (#12856 ) * rename opts to ren * fix docs and bring store back	2025-10-22 09:15:38 +08:00
chenyuandGitHub	c5cee74706	remove BLOCK_REORDER (#12854 ) not used	2025-10-21 19:10:14 -04:00
chenyuandGitHub	0b673eddec	simpler newton_schulz transpose (#12853 )	2025-10-21 17:21:45 -04:00
b1tgandGitHub	60d7e232f2	cuda fp8 (#12782 ) * cuda fp8 * tensor core * tc test * clean * clean pm	2025-10-21 15:05:25 -04:00
Harald SchäferandGitHub	587ccc0e5c	compile3: make selftests opt-in (#12851 )	2025-10-21 11:32:27 -07:00
wozeparrotandGitHub	c3149c618a	feat: nvcc compiler (#12852 )	2025-10-21 11:31:23 -07:00
chenyuandGitHub	8baa61bd67	use torch 2.9 and its Muon in test (#12773 ) * use torch 2.9 and its Muon in test * relax and disable	2025-10-21 13:35:17 -04:00
chenyuandGitHub	f51f9aaa16	muon ns_params -> ns_coefficients (#12850 ) match the official torch one	2025-10-21 12:35:52 -04:00
wozeparrotandGitHub	62e7b8b870	feat: just use compile3 (#12849 )	2025-10-21 07:56:50 -07:00
nimlgenandGitHub	c7336c3e31	amd: sqtt for aql (#12846 )	2025-10-21 22:35:01 +08:00
George HotzandGitHub	8960ac54f3	remove RewriteStep premature optimization (#12840 ) * remove RewriteStep premature optimization * fix ebs * core line count	2025-10-21 21:45:20 +08:00
Sieds LyklesandGitHub	7f798a9630	Cleanup const buffers (#12829 ) * split pm_cleanups * update test_schedule * shrink when we remove bufferize * dont do shrink if shape is empty * update tests * remove 1 from metadata deal with the noop bufferize * only noop on cvar * cleanup * fix if * rename	2025-10-21 14:53:49 +02:00
nimlgenandGitHub	1ad6598963	amd: trace all instructions (#12831 )	2025-10-21 20:52:24 +08:00
sirhcmandGitHub	cdc72556a1	no more brew (#12839 )	2025-10-21 08:12:46 -04:00
George HotzandGitHub	20a232f1c5	bugfixes from multioutput + PCONTIG=3 for fa bw memory fix (#12837 ) * bugfixes from multioutput * PCONTIG=3 fixes fa memory usage * that's base	2025-10-21 19:21:02 +08:00
qazalandGitHub	0435d31f1c	viz: generic back button functionality (#12838 )	2025-10-21 18:52:00 +08:00