something

* gradient: add TRUNC backward

* test: move round quantization gradient to test_ops

.github/workflows/autogen.yml

@@ -33,12 +33,8 @@ jobs:
       uses: ./.github/actions/setup-tinygrad
       with:
         key: 'autogen'
-        opencl: 'true'
         amd: 'true'
-        cuda: 'true'
         llvm: 'true'
-        webgpu: 'true'
-        mesa: 'true'
         pydeps: 'pyyaml mako'
     - name: Install autogen support packages
       run: sudo apt-get install -y --no-install-recommends libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev libdrm-dev
@@ -48,7 +44,7 @@ jobs:
         python3 -c "from tinygrad.runtime.autogen import opencl"
         python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv_580, nv"
         python3 -c "from tinygrad.runtime.autogen import comgr_3, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd, amdgpu_drm"
-        python3 -c "from tinygrad.runtime.autogen.am import am, pm4_soc15, pm4_nv, sdma_4_0_0, sdma_5_0_0, sdma_6_0_0, smu_v13_0_0, smu_v13_0_6, smu_v13_0_12, smu_v14_0_2"
+        python3 -c "from tinygrad.runtime.autogen.am import *"
         python3 -c "from tinygrad.runtime.autogen import libc, kfd, io_uring, ib, pci, vfio"
         python3 -c "from tinygrad.runtime.autogen import llvm"
         python3 -c "from tinygrad.runtime.autogen import webgpu"

.github/workflows/benchmark.yml

@@ -51,34 +51,35 @@ jobs:
     - name: openpilot compile3 0.10.1 driving_vision
       run: FLOAT16=1 DEV=CL IMAGE=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
 
-  testframeworkpytest:
-    name: framework pytest
-    env:
-      CI: ""
-      CAPTURE_PROCESS_REPLAY: "0"
-    runs-on: [self-hosted, framework]
-    timeout-minutes: 10
-    defaults:
-      run:
-        shell: bash -e -o pipefail {0}
-    if: github.repository_owner == 'tinygrad'
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: setup python environment
-      run: |
-        rm -rf /tmp/tinygrad_pytest_ci
-        uv venv /tmp/tinygrad_pytest_ci
-        source /tmp/tinygrad_pytest_ci/bin/activate
-        uv pip install .[testing]
-    - name: setup staging db
-      run: |
-        echo "CACHEDB=/tmp/pytest-db-ci.db" >> $GITHUB_ENV
-        rm -f /tmp/pytest-db-ci*
-    - name: Run pytest -nauto
-      run: |
-        source /tmp/tinygrad_pytest_ci/bin/activate
-        pytest -nauto --durations=20
+  # TODO: reenable when not flaky
+  #testframeworkpytest:
+  #  name: framework pytest
+  #  env:
+  #    CI: ""
+  #    CAPTURE_PROCESS_REPLAY: "0"
+  #  runs-on: [self-hosted, framework]
+  #  timeout-minutes: 10
+  #  defaults:
+  #    run:
+  #      shell: bash -e -o pipefail {0}
+  #  if: github.repository_owner == 'tinygrad'
+  #  steps:
+  #  - name: Checkout Code
+  #    uses: actions/checkout@v6
+  #  - name: setup python environment
+  #    run: |
+  #      rm -rf /tmp/tinygrad_pytest_ci
+  #      uv venv /tmp/tinygrad_pytest_ci
+  #      source /tmp/tinygrad_pytest_ci/bin/activate
+  #      uv pip install .[testing]
+  #  - name: setup staging db
+  #    run: |
+  #      echo "CACHEDB=/tmp/pytest-db-ci.db" >> $GITHUB_ENV
+  #      rm -f /tmp/pytest-db-ci*
+  #  - name: Run pytest -nauto
+  #    run: |
+  #      source /tmp/tinygrad_pytest_ci/bin/activate
+  #      pytest -nauto --durations=20
 
   testmacbenchmark:
     name: Mac Benchmark
@@ -189,7 +190,7 @@ jobs:
         path: |
           onnx_inference_speed.csv
     - 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.11 process_replay.py
+      uses: ./.github/actions/process-replay
 
   testusbgpu:
     name: UsbGPU Benchmark
@@ -324,7 +325,7 @@ jobs:
         path: |
           onnx_inference_speed.csv
     - 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: ./.github/actions/process-replay
 
   testmorenvidiabenchmark:
     name: tinybox green Training Benchmark
@@ -386,7 +387,7 @@ jobs:
       # TODO: remove BERT_LAYERS once scheduler is fast
       run: BENCHMARK_LOG=bert_10steps_6gpu DEV=NV CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
     - 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: ./.github/actions/process-replay
 
   testamdbenchmark:
     name: tinybox red Benchmark
@@ -498,7 +499,7 @@ jobs:
     - name: Run GPT2 w HALF/BEAM
       run: BENCHMARK_LOG=gpt2_half_beam DEV=AMD HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing
     - 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: ./.github/actions/process-replay
 
   testmoreamdbenchmark:
     name: tinybox red Training Benchmark
@@ -555,7 +556,7 @@ jobs:
     #- name: Test full tinyfs load
     #  run: TINYFS_ENDPOINT=10.0.52.11:6767 PYTHONPATH=. python extra/tinyfs/fetch_file.py --hash d734f5e3be9f1e9d863bfaa4fc6c1ef2 --len 175866113 --dest mapping.json --check
     - 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: ./.github/actions/process-replay
 
   testmlperfamdbenchmark:
     name: tinybox red MLPerf Benchmark
@@ -601,7 +602,7 @@ jobs:
       # TODO: remove BERT_LAYERS once scheduler is fast
       run: BENCHMARK_LOG=bert_10steps_6gpu DEV=AMD CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
     - 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: ./.github/actions/process-replay
 
   testqualcommbenchmark:
     name: comma Benchmark
@@ -628,7 +629,7 @@ jobs:
     - name: IR3 openpilot compile3 0.11.0 driving_vision
       run: BENCHMARK_LOG=ir3_openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM:IR3 FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
     - name: openpilot compile3 0.11.0 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_11_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_policy.onnx
+      run: BENCHMARK_LOG=openpilot_0_11_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_policy.onnx
     - name: openpilot compile3 0.11.0 dmonitoring
       run: BENCHMARK_LOG=openpilot_0_11_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/dmonitoring_model.onnx
     - name: DEBUG=2 openpilot compile3 0.10.1 driving_vision
@@ -636,7 +637,7 @@ jobs:
     - name: openpilot compile3 0.10.1 driving_vision
       run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=1 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=4 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=1 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=11 DEV=QCOM FLOAT16=1 IMAGE=1 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
@@ -648,7 +649,7 @@ jobs:
         # benchmark on DSP with NOOPT=1, the devectorizer has issues
         PYTHONPATH=. CC=clang-19 DEV=DSP 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: ./.github/actions/process-replay
 
   testcommausbgpubenchmark:
     name: UsbGPU Benchmark (comma)
@@ -745,7 +746,7 @@ jobs:
         DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6482 AM_RESET=1 DEV=PCI+AMD AMD_AQL=1 python3 test/test_tiny.py
         pkill -f 'extra/remote/serve.py' || true
     - 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: ./.github/actions/process-replay
 
   testgreendriverbenchmark:
     name: NV Benchmark
@@ -808,4 +809,4 @@ jobs:
         DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6483 DEV=NV python3 test/test_tiny.py
         pkill -f 'extra/remote/serve.py' || true
     - 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: ./.github/actions/process-replay

.github/workflows/test.yml

@@ -333,7 +333,7 @@ jobs:
         deps: testing_unit
         python-version: '3.14'
     - name: Test SPEC=2
-      run: SPEC=2 pytest --maxfail=10 -n auto --durations=30 test/unit test/backend test/opt --ignore test/backend/test_custom_kernel.py --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}
+      run: SPEC=2 pytest --maxfail=10 -n auto --durations=30 test/unit test/backend test/opt --ignore test/backend/test_custom_kernel.py --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" -k "not test_conv2d_ceildiv_edge_case" --splits 2 --group ${{ matrix.group }}
 
   fuzzing:
     name: Fuzzing
@@ -417,7 +417,7 @@ jobs:
           llvm: 'true'
       - name: Test openpilot model kernel count and gate usage
         run: |
-          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1486 ALLOWED_GATED_READ_IMAGE=17 FLOAT16=1 DEV=CL IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1486 ALLOWED_GATED_READ_IMAGE=18 FLOAT16=1 DEV=CL IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
       - name: Test openpilot CL compile fp16
         run: FLOAT16=1 DEV=CL IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
       - name: Test openpilot CL compile fp32 (test correctness)

docs/abstractions4.py

@@ -105,7 +105,7 @@ def example_3_custom_uop(a:Tensor, correct):
 def example_5_custom_assembly(a:Tensor, correct):
   # Kernel class copied from amd_asm_matmul
   class Kernel:
-    def __init__(self, arch='gfx1100'): self.instructions, self.labels, self.pos, self.arch = [], {}, 0, arch
+    def __init__(self): self.instructions, self.labels, self.pos = [], {}, 0
     def label(self, name): self.labels[name] = self.pos
     def emit(self, inst, target=None):
       self.instructions.append(inst)

docs/runtime.md

@@ -5,7 +5,7 @@ tinygrad supports various runtimes, enabling your code to scale across a wide ra
 | Runtime | Description | Compiler Options | Requirements |
 |---------|-------------|------------------|--------------|
 | [NV](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_nv.py) | Provides acceleration for NVIDIA GPUs | nvrtc (default)<br>PTX (`DEV=NV:PTX`) | Ampere/Ada/Blackwell series GPUs.<br>You can select an interface via [the `DEV` variable](env_vars.md#dev-variable). See [NV interfaces](#nv-interfaces) for details. |
-| [AMD](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_amd.py) | Provides acceleration for AMD GPUs | LLVM (`DEV=AMD:LLVM`)<br>HIP/COMGR (`DEV=AMD:HIP`) | RDNA2 or newer GPUs.<br>You can select an interface via [the `DEV` variable](env_vars.md#dev-variable). See [AMD interfaces](#amd-interfaces) for details. |
+| [AMD](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_amd.py) | Provides acceleration for AMD GPUs | LLVM (`DEV=AMD:LLVM`)<br>HIP/COMGR (`DEV=AMD:HIP`) | CDNA3, CDNA4, RDNA3 or RDNA4 GPUs.<br>You can select an interface via [the `DEV` variable](env_vars.md#dev-variable). See [AMD interfaces](#amd-interfaces) for details. |
 | [QCOM](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_qcom.py) | Provides acceleration for QCOM GPUs | - | 6xx series GPUs |
 | [METAL](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_metal.py) | Utilizes Metal for acceleration on Apple devices | - | M1+ Macs; Metal 3.0+ for `bfloat` support |
 | [CUDA](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cuda.py) | Utilizes CUDA for acceleration on NVIDIA GPUs | nvrtc (default)<br> PTX (`DEV=CUDA:PTX`) | NVIDIA GPU with CUDA support |

docs/tensor/elementwise.md

@@ -66,8 +66,8 @@ Elementwise ops operate on a per element basis. They don't change the shape of t
 ::: tinygrad.Tensor.sub
 ::: tinygrad.Tensor.mul
 ::: tinygrad.Tensor.div
-::: tinygrad.Tensor.idiv
 ::: tinygrad.Tensor.mod
+::: tinygrad.Tensor.fmod
 ::: tinygrad.Tensor.bitwise_xor
 ::: tinygrad.Tensor.bitwise_and
 ::: tinygrad.Tensor.bitwise_or

docs/tinygpu.md

@@ -4,7 +4,7 @@ TinyGPU app lets you use AMD and NVIDIA GPUs on macOS over USB4/Thunderbolt with
 
 ## Requirements
 
-- macOS (12.1+)
+- macOS (13.0+)
 - USB4/Thunderbolt port
 - A supported GPU (AMD RDNA3+ or NVIDIA Ampere+)
 

examples/llama3.py

@@ -123,7 +123,7 @@ def NF4Linear(block_size):
     def __call__(self, x: Tensor) -> Tensor:
       high_bits = self.weight
       low_bits = (self.weight * 2 ** 4).contiguous()
-      unpacked = Tensor.stack(high_bits, low_bits, dim=-1).idiv(2 ** 4)
+      unpacked = Tensor.stack(high_bits, low_bits, dim=-1).div(2 ** 4, rounding_mode="trunc")
       unscaled = CODE[unpacked].to(x.device).reshape(-1, block_size) * self.scale
       return x.linear(unscaled.reshape(self.out_features, self.in_features).T)
 

examples/mlperf/model_train.py

@@ -1357,6 +1357,7 @@ def train_llama3():
       MLLOGGER.event(key=mllog_constants.OPT_LR_WARMUP_STEPS, value=WARMUP_STEPS)
       MLLOGGER.event(key=mllog_constants.NUM_WARMUP_STEPS, value=WARMUP_STEPS)
       MLLOGGER.event(key=mllog_constants.OPT_LR_DECAY_STEPS, value=MAX_STEPS - WARMUP_STEPS)
+      MLLOGGER.event(key=mllog_constants.OPT_LR_DECAY_SCHEDULE, value="cosine with linear warmup")
       MLLOGGER.event(key=mllog_constants.OPT_GRADIENT_CLIP_NORM, value=1.0)
   else:
     MLLOGGER = None
@@ -1418,7 +1419,10 @@ def train_llama3():
 
   for p in optim.params:
     grad_dtype = dtypes.bfloat16 if p.dtype == FP8_DTYPE else p.dtype
-    p.grad = Tensor.zeros(p.shape, dtype=grad_dtype, device=p.device).contiguous()
+    if isinstance(p.device, tuple) and p.uop.axis is not None:
+      p.grad = Tensor.zeros(p.shape, dtype=grad_dtype, device=p.device[0]).shard_(p.device, axis=p.uop.axis).contiguous()
+    else:
+      p.grad = Tensor.zeros(p.shape, dtype=grad_dtype, device=p.device).contiguous()
   grads = [p.grad for p in optim.params]
 
   scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)
@@ -1433,6 +1437,7 @@ def train_llama3():
     load_state_dict(scheduler, safe_load(fn), realize=False)
 
   fp8_amax = [t for ts in model._fp8_amax.values() for t in ts]
+  fp8_grad_amax = [t for ts in model._fp8_grad_amax.values() for t in ts] if hasattr(model, "_fp8_grad_amax") else []
   fp8_inv_scales = list(model._fp8_inv_scale.values())
 
   from tinygrad.nn.state import get_state_dict
@@ -1444,6 +1449,10 @@ def train_llama3():
       idx = next(j for j, p in enumerate(optim.params) if p is w)
       optim.master_params[idx].assign((optim.master_params[idx] * w._inv_scale.reshape(-1, *([1]*(w.ndim-1)))).contiguous())
 
+  # realize everything here
+  if optim.master_params: Tensor.realize(*optim.master_params)
+  Tensor.realize(*optim.params, *fp8_inv_scales, *fp8_amax, *fp8_grad_amax)
+
   @TinyJit
   def minibatch(tokens:Tensor):
     if is_dp: tokens = tokens.to(None).shard(device, 0)
@@ -1460,7 +1469,7 @@ def train_llama3():
       apply_grad(g, new_g.uop)
 
     loss_cpu = loss.flatten().float().to("CPU")
-    return loss_cpu.realize(*grads, *fp8_amax)
+    return loss_cpu.realize(*grads, *fp8_amax, *fp8_grad_amax)
 
   @TinyJit
   def optim_step():
@@ -1635,7 +1644,6 @@ def train_llama3():
         tqdm.write(f"target achieved after {sequences_seen} sequences")
         if MLLOGGER and RUNMLPERF:
           MLLOGGER.end(key=mllog_constants.EPOCH_STOP, metadata={mllog_constants.SAMPLES_COUNT: sequences_seen})
-          MLLOGGER.event(key=mllog_constants.TRAIN_SAMPLES, value=sequences_seen)
           MLLOGGER.end(key=mllog_constants.RUN_STOP, metadata={mllog_constants.STATUS: mllog_constants.SUCCESS})
         if getenv("CKPT"):
           if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)

examples/mlperf/models/flat_llama.py

@@ -20,55 +20,73 @@ from extra.llama_kernels.rmsnorm import rmsnorm
 from extra.llama_kernels import FP8_MAX, local_abs_max
 
 ASM_GEMM = getenv("ASM_GEMM", 0)
+FUSED_INPUT_QUANTIZE = getenv("FUSED_INPUT_QUANTIZE", 0)
+FUSED_ADD_NORM_MUL_QUANTIZE = getenv("FUSED_ADD_NORM_MUL_QUANTIZE", 0)
+FUSED_SILU_W13 = getenv("FUSED_SILU_W13", 0)
 
 FP8_DTYPE = dtypes.fp8e4m3
 FP8_GRAD_DTYPE = dtypes.fp8e5m2
 
 def quantize_fp8(x:Tensor, amax_state:Tensor|None=None):
-  new_amax = (local_abs_max(x) if isinstance(x.device, tuple) else x.abs().max()).detach()
+  new_amax = (local_abs_max(x) if isinstance(x.device, tuple) else x.abs().max()).detach().cast(dtypes.float32)
   scale = FP8_MAX / ((amax_state if amax_state is not None else new_amax) + 1e-8)
   x_scaled = x * scale
   x_clamped = x_scaled + (x_scaled.detach().clamp(-FP8_MAX, FP8_MAX) - x_scaled.detach())  # STE
   return x_clamped.cast(FP8_DTYPE), scale.float().reciprocal(), new_amax
 
 def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_scale:Tensor|None=None,
-           x_fp8:Tensor|None=None, x_scale:Tensor|None=None, x_new_amax:Tensor|None=None) -> tuple[Tensor,...]:
+           x_fp8:Tensor|None=None, x_scale:Tensor|None=None, x_new_amax:Tensor|None=None,
+           grad_amax_state:Tensor|None=None) -> tuple[Tensor,...]:
   if not fp8:
     if ASM_GEMM:
       from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
       if can_use_asm_gemm(x, w.T): return (asm_gemm(x, w.T),)
     return (x @ w.T,)
   assert w_inv_scale is not None, "fp8 matmul requires w_inv_scale (weights must be stored in fp8 with per-tensor scale)"
-  if x_fp8 is None: x_fp8, x_scale, x_new_amax = quantize_fp8(x, amax_state=amax_x)
+  if x_fp8 is None:
+    if FUSED_INPUT_QUANTIZE and amax_x is not None:
+      from extra.llama_kernels.quantize_fp8_delayed import quantize_fp8_delayed
+      x_fp8, x_scale, x_new_amax, _ = quantize_fp8_delayed(x, amax_x, FP8_DTYPE)
+    else:
+      x_fp8, x_scale, x_new_amax = quantize_fp8(x, amax_state=amax_x)
   if ASM_GEMM:
     from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
-    if can_use_asm_gemm(x_fp8, w.T): return asm_gemm(x_fp8, w.T, x_scale=x_scale, w_scale=w_inv_scale), x_new_amax, x_fp8, w
-  return (x_fp8.dot(w.T, dtype=dtypes.float) * x_scale * w_inv_scale).cast(dtypes.bfloat16), x_new_amax, x_fp8, w
+    if can_use_asm_gemm(x_fp8, w.T):
+      return asm_gemm(x_fp8, w.T, x_scale=x_scale, w_scale=w_inv_scale, grad_amax_state=grad_amax_state), x_new_amax, x_fp8, w
+  return x_fp8.dot(w.T, dtype=dtypes.float) * x_scale * w_inv_scale, x_new_amax, x_fp8, w
 
-def norm_mul_quantize_matmul(x:Tensor, norm:Tensor, amax_x, w_inv_scale, w:Tensor, eps:float):
-  FUSED_NORM_MUL_QUANTIZE = getenv("FUSED_NORM_MUL_QUANTIZE", 0)
-  normed, rrms = rmsnorm(x, eps)
-  if FUSED_NORM_MUL_QUANTIZE:
-    from extra.llama_kernels.fused_mul_quantize_fp8 import fused_mul_quantize_fp8
-    amax_s = amax_x if amax_x is not None else Tensor.full((), 1.0, dtype=dtypes.bfloat16, device=normed.device)
-    x_fp8, x_inv_scale, new_amax = fused_mul_quantize_fp8(normed, norm, amax_s, FP8_DTYPE)
+def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor, grad_amax_state:Tensor):
+  if FUSED_ADD_NORM_MUL_QUANTIZE:
+    from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_rmsnorm_mul_quantize_fp8
+    x_fp8, x_inv_scale, new_amax, x_normed, rrms = fused_rmsnorm_mul_quantize_fp8(x, norm, amax_x, eps, FP8_DTYPE)
+    out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, x_scale=x_inv_scale, x_new_amax=new_amax, grad_amax_state=grad_amax_state)
+    return out, x_normed, rrms, ret
+  x_normed, rrms = rmsnorm(x, eps)
+  out, *ret = matmul(x_normed * norm, w, amax_x=amax_x, w_inv_scale=w_inv_scale, grad_amax_state=grad_amax_state)
+  return out, x_normed, rrms, ret
+
+def add_norm_quantize_matmul(x:Tensor, residual:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor):
+  if FUSED_ADD_NORM_MUL_QUANTIZE:
+    from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_add_rmsnorm_mul_quantize_fp8
+    x_fp8, x_inv_scale, new_amax, h, x_normed, rrms = fused_add_rmsnorm_mul_quantize_fp8(x, residual, norm, amax_x, eps, FP8_DTYPE)
     out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, x_scale=x_inv_scale, x_new_amax=new_amax)
-  else:
-    x = normed * norm
-    out, *ret = matmul(x, w, amax_x=amax_x, w_inv_scale=w_inv_scale)
-  return out, normed, rrms, ret
+    return out, h, x_normed, rrms, ret
+  h = x + residual
+  x_normed, rrms = rmsnorm(h, eps)
+  out, *ret = matmul(x_normed * norm, w, amax_x=amax_x, w_inv_scale=w_inv_scale)
+  return out, h, x_normed, rrms, ret
 
-def silu_w13_matmul(x_w13:Tensor, w2:Tensor, amax_x2, s_2):
-  FUSED_SILU_W13 = getenv("FUSED_SILU_W13", 0)
+def silu_w13_quantize_matmul(x_w13:Tensor, w2:Tensor, s_2:Tensor,
+                             amax_x2:Tensor,
+                             grad_amax_xw13:Tensor, grad_amax_xout:Tensor):
   if FUSED_SILU_W13:
     from extra.llama_kernels.cast_amax import fused_quantize_fp8_w13
-    amax_s = amax_x2 if amax_x2 is not None else Tensor.full((), 1.0, dtype=dtypes.bfloat16, device=x_w13.device)
-    x2_fp8, x2_inv_scale, new_amax_x2 = fused_quantize_fp8_w13(x_w13, amax_s, FP8_DTYPE)
-    out, *ret = matmul(None, w2, w_inv_scale=s_2, x_fp8=x2_fp8, x_scale=x2_inv_scale, x_new_amax=new_amax_x2)
-  else:
-    hidden_dim = x_w13.shape[-1] // 2
-    x_w1, x_w3 = x_w13[..., :hidden_dim], x_w13[..., hidden_dim:]
-    out, *ret = matmul(x_w1.silu() * x_w3, w2, amax_x=amax_x2, w_inv_scale=s_2)
+    x2_fp8, x2_inv_scale, new_amax_x2 = fused_quantize_fp8_w13(x_w13, amax_x2, FP8_DTYPE, grad_amax_state=grad_amax_xw13)
+    out, *ret = matmul(None, w2, w_inv_scale=s_2, x_fp8=x2_fp8, x_scale=x2_inv_scale, x_new_amax=new_amax_x2, grad_amax_state=grad_amax_xout)
+    return out, ret
+  hidden = x_w13.shape[-1] // 2
+  x_w1, x_w3 = x_w13[..., :hidden], x_w13[..., hidden:]
+  out, *ret = matmul(x_w1.silu() * x_w3, w2, amax_x=amax_x2, w_inv_scale=s_2, grad_amax_state=grad_amax_xout)
   return out, ret
 
 class FlatTransformer:
@@ -104,32 +122,33 @@ class FlatTransformer:
     self.output = Tensor.normal(1, vocab_size, dim, mean=0.0, std=0.02, dtype=dtypes.bfloat16)
     self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).contiguous().requires_grad_(False)
 
-    def _amax(): return Tensor.full((), FP8_MAX).contiguous().requires_grad_(False)
+    def _amax(): return Tensor.full((), FP8_MAX, dtype=dtypes.float32).contiguous().requires_grad_(False)
     names = ["xqkv", "xo", "x13", "x2"]
     self._fp8_amax = {name: [_amax() for _ in range(n_layers)] for name in names}
-    # per-weight inv_scale: single (n_layers,) float32 tensor per weight (kernel reads float* pointers)
+    grad_names = ["xqkv", "xo", "xw13", "xout"]
+    self._fp8_grad_amax = {name: [_amax() for _ in range(n_layers)] for name in grad_names}
     w_names = ["wqkv", "wo", "w13", "w2"]
-    self._fp8_inv_scale = {}
-    for wname, inv_scales in zip(w_names, self._init_inv_scales):
-      self._fp8_inv_scale[wname] = inv_scales.float().contiguous().requires_grad_(False)
+    self._fp8_inv_scale = {wname: inv_scales.float().contiguous().requires_grad_(False)
+                           for wname, inv_scales in zip(w_names, self._init_inv_scales)}
     del self._init_inv_scales
 
   def lin_per_layer(self, in_features:int, out_features:int, std:float=0.02):
-    if getenv("ZEROS", 0): w = Tensor.zeros(self.n_layers, out_features, in_features)
+    if getenv("ZEROS"): w = Tensor.zeros(self.n_layers, out_features, in_features)
     else: w = Tensor.normal(self.n_layers, out_features, in_features, mean=0.0, std=std)
-    # per-layer scaled fp8 cast: fill the fp8 range for best precision
     amax = w.abs().flatten(1).max(1).detach()
     scale = FP8_MAX / (amax + 1e-8)
-    self._init_inv_scales.append((amax + 1e-8) / FP8_MAX)  # save for inv_scale init
+    self._init_inv_scales.append((amax + 1e-8) / FP8_MAX)
     return (w * scale.reshape(-1, 1, 1)).clamp(-FP8_MAX, FP8_MAX).cast(FP8_DTYPE)
 
   def attention(self, x:Tensor, freqs_cis:Tensor, attention_norm:Tensor, wqkv:Tensor, wo:Tensor,
-                amax_xqkv=None, amax_xo=None, s_qkv=None, s_o=None):
+                amax_xqkv:Tensor, amax_xo:Tensor, s_qkv:Tensor, s_o:Tensor,
+                grad_amax_xqkv:Tensor, grad_amax_xo:Tensor):
     bsz, seqlen, _ = x.shape
     new_amaxs, saves = [], []
 
-    xqkv, normed, rrms, ret = norm_mul_quantize_matmul(x, attention_norm, amax_xqkv, s_qkv, wqkv, self.norm_eps)
-    saves.extend([normed, rrms])
+    xqkv, x_normed, rrms, ret = norm_quantize_matmul(x, attention_norm, wqkv, s_qkv, self.norm_eps,
+                                                     amax_x=amax_xqkv, grad_amax_state=grad_amax_xqkv)
+    saves.extend([x_normed, rrms])
     new_amaxs.extend(ret[:1])
     saves.extend(ret[1:] + [xqkv])
     xqkv = xqkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
@@ -139,49 +158,52 @@ class FlatTransformer:
 
     xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
     xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16)
-    xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
     if getenv("HK_FLASH_ATTENTION"):
       from extra.thunder.amd.fa import flash_attention
       attn, *save = flash_attention(xq, xk, xv, is_causal=True)
       saves.extend(save)
     else:
-      attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True)
-    attn = attn.transpose(1, 2).reshape(bsz, seqlen, -1)
+      xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
+      attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True).transpose(1, 2)
+    attn = attn.reshape(bsz, seqlen, -1)
 
-    out, *ret = matmul(attn, wo, amax_x=amax_xo, w_inv_scale=s_o)
+    out, *ret = matmul(attn, wo, amax_x=amax_xo, w_inv_scale=s_o, grad_amax_state=grad_amax_xo)
     new_amaxs.extend(ret[:1])
     saves.extend(ret[1:] + [out])
     return (out, *new_amaxs, *saves)
 
-  def feed_forward(self, x:Tensor, ffn_norm:Tensor, w13:Tensor, w2:Tensor,
-                   amax_x13=None, amax_x2=None, s_13=None, s_2=None):
+  def feed_forward(self, x:Tensor, residual:Tensor, ffn_norm:Tensor, w13:Tensor, w2:Tensor,
+                   amax_x13:Tensor, amax_x2:Tensor, s_13:Tensor, s_2:Tensor,
+                   grad_amax_xw13:Tensor, grad_amax_xout:Tensor):
     new_amaxs, saves = [], []
 
-    x_w13, normed, rrms, ret = norm_mul_quantize_matmul(x, ffn_norm, amax_x13, s_13, w13, self.norm_eps)
-    saves.extend([normed, rrms])
+    x_w13, h, x_normed, rrms, ret = add_norm_quantize_matmul(x, residual, ffn_norm, w13, s_13, self.norm_eps,
+                                                             amax_x=amax_x13)
+    saves.extend([x_normed, rrms])
     new_amaxs.extend(ret[:1])
     saves.extend(ret[1:] + [x_w13])
 
-    out, ret = silu_w13_matmul(x_w13, w2, amax_x2, s_2)
+    out, ret = silu_w13_quantize_matmul(x_w13, w2, s_2, amax_x2=amax_x2, grad_amax_xw13=grad_amax_xw13, grad_amax_xout=grad_amax_xout)
     new_amaxs.extend(ret[:1])
     saves.extend(ret[1:] + [out])
-    return (out, *new_amaxs, *saves)
+    return (out, h, *new_amaxs, *saves)
 
   @function(precompile=True, precompile_backward=True)
   def run_layer(self, x:Tensor, freqs_cis:Tensor,
                 attention_norm:Tensor, wqkv:Tensor, wo:Tensor,
                 ffn_norm:Tensor, w13:Tensor, w2:Tensor,
-                amax_xqkv=None, amax_xo=None,
-                amax_x13=None, amax_x2=None,
-                s_qkv=None, s_o=None, s_13=None, s_2=None):
+                amax_xqkv:Tensor, amax_xo:Tensor,
+                amax_x13:Tensor, amax_x2:Tensor,
+                s_qkv:Tensor, s_o:Tensor, s_13:Tensor, s_2:Tensor,
+                grad_amax_xqkv:Tensor, grad_amax_xo:Tensor,
+                grad_amax_xw13:Tensor, grad_amax_xout:Tensor):
     attn, *attn_ret = self.attention(x, freqs_cis, attention_norm, wqkv, wo,
-                                     amax_xqkv=amax_xqkv, amax_xo=amax_xo,
-                                     s_qkv=s_qkv, s_o=s_o)
+                                     amax_xqkv=amax_xqkv, amax_xo=amax_xo, s_qkv=s_qkv, s_o=s_o,
+                                     grad_amax_xqkv=grad_amax_xqkv, grad_amax_xo=grad_amax_xo)
     attn_amaxs, attn_saves = attn_ret[:2], attn_ret[2:]
-    h = x + attn
-    ffn, *ffn_ret = self.feed_forward(h, ffn_norm, w13, w2,
-                                      amax_x13=amax_x13, amax_x2=amax_x2,
-                                      s_13=s_13, s_2=s_2)
+    ffn, h, *ffn_ret = self.feed_forward(x, attn, ffn_norm, w13, w2,
+                                                   amax_x13=amax_x13, amax_x2=amax_x2, s_13=s_13, s_2=s_2,
+                                                   grad_amax_xw13=grad_amax_xw13, grad_amax_xout=grad_amax_xout)
     ffn_amaxs, ffn_saves = ffn_ret[:2], ffn_ret[2:]
     h = h + ffn
     return (h, *attn_amaxs, *ffn_amaxs, *attn_saves, *ffn_saves)
@@ -202,28 +224,31 @@ class FlatTransformer:
       self.tok_embeddings.weight.shard_(device, axis=0).realize()
       self.output.shard_(device, axis=1).realize()
       self.freqs_cis.shard_(device, axis=None).realize()
-      for name in self._fp8_amax:
-        for i in range(len(self._fp8_amax[name])):
-          self._fp8_amax[name][i] = self._fp8_amax[name][i].to(device).contiguous().requires_grad_(False)
+      for amax_dict in (self._fp8_amax, self._fp8_grad_amax):
+        for name in amax_dict:
+          for i in range(len(amax_dict[name])):
+            amax_dict[name][i] = amax_dict[name][i].to(device).contiguous().requires_grad_(False)
       for name in self._fp8_inv_scale:
         self._fp8_inv_scale[name] = self._fp8_inv_scale[name].to(device).contiguous().requires_grad_(False)
 
   def __call__(self, tokens:Tensor):
     h = self.tok_embeddings(tokens)
     freqs_cis = self.freqs_cis.cast(h.dtype)[:, :tokens.shape[1], :, :, :]
-    amaxs, inv_scales = self._fp8_amax, self._fp8_inv_scale
+    a, ga, s = self._fp8_amax, self._fp8_grad_amax, self._fp8_inv_scale
     for i in range(self.n_layers):
       h, *ret = self.run_layer(h, freqs_cis,
                                self.attention_norm[i], self.wqkv[i], self.wo[i],
                                self.ffn_norm[i], self.w13[i], self.w2[i],
-                               amax_xqkv=amaxs["xqkv"][i], amax_xo=amaxs["xo"][i],
-                               amax_x13=amaxs["x13"][i], amax_x2=amaxs["x2"][i],
-                               s_qkv=inv_scales["wqkv"][i], s_o=inv_scales["wo"][i],
-                               s_13=inv_scales["w13"][i], s_2=inv_scales["w2"][i])
+                               amax_xqkv=a["xqkv"][i], amax_xo=a["xo"][i],
+                               amax_x13=a["x13"][i], amax_x2=a["x2"][i],
+                               s_qkv=s["wqkv"][i], s_o=s["wo"][i],
+                               s_13=s["w13"][i], s_2=s["w2"][i],
+                               grad_amax_xqkv=ga["xqkv"][i], grad_amax_xo=ga["xo"][i],
+                               grad_amax_xw13=ga["xw13"][i], grad_amax_xout=ga["xout"][i])
       for name, new_val in zip(["xqkv", "xo", "x13", "x2"], ret[:5]):
-        amaxs[name][i].assign(new_val)
+        a[name][i].assign(new_val)
 
-    logits = matmul(self.norm(h).contiguous().contiguous_backward(), self.output[0], fp8=False)[0].contiguous_backward()
+    logits = matmul(self.norm(h), self.output[0], fp8=False)[0]
     return logits
 
 def _get_pads(uop:UOp) -> list[UOp]:
@@ -232,13 +257,19 @@ def _get_pads(uop:UOp) -> list[UOp]:
 
 def apply_grad(grad_buf:Tensor, new_grad:UOp):
   pads = _get_pads(new_grad)
-  new_grad = new_grad.cast(grad_buf.dtype)
   if len(pads) <= 1:
+    new_grad = new_grad.cast(grad_buf.dtype)
     store = grad_buf.uop.store(grad_buf.uop + new_grad)
     grad_buf.uop = grad_buf.uop.after(store)
     return
   sorted_pads = sorted(pads, key=lambda p: p.marg[0][0] if p.op == Ops.PAD else 0)
-  inners = [Tensor(p.src[0] if p.op == Ops.PAD else p, device=grad_buf.device).cast(grad_buf.dtype) for p in sorted_pads]
+  inners_raw = [Tensor(p.src[0] if p.op == Ops.PAD else p, device=grad_buf.device) for p in sorted_pads]
+  if getenv("FUSED_PAD_GRAD_ACCUM", 0):
+    from extra.llama_kernels.fused_pad_grad_accum import fused_pad_grad_accum, can_fused_pad_grad_accum
+    if can_fused_pad_grad_accum(grad_buf, inners_raw):
+      grad_buf.uop = fused_pad_grad_accum(grad_buf, inners_raw).uop
+      return
+  inners = [t.cast(grad_buf.dtype) for t in inners_raw]
   grad_buf.assign(grad_buf + inners[0].cat(*inners[1:], dim=0))
 
 if __name__ == "__main__":

examples/mlperf/optim.py

@@ -81,7 +81,7 @@ class GradAccClipAdamW(Optimizer):
     if STOCHASTIC_ROUND and t.dtype == dtypes.bfloat16: return stochastic_round_bf16(new_w)
     if t.dtype in dtypes.fp8s:
       from examples.mlperf.models.flat_llama import FP8_MAX
-      amax = new_w.float().abs().flatten(1).max(1).detach()  # per-layer amax for (n_layers, out, in)
+      amax = new_w.float().abs().max(axis=tuple(range(1, new_w.ndim))).detach()  # per-layer amax for (n_layers, out, in)
       scale = FP8_MAX / (amax + 1e-8)
       fp8_w = (new_w * scale.reshape(-1, *([1]*(new_w.ndim-1)))).clamp(-FP8_MAX, FP8_MAX).cast(t.dtype)
       if hasattr(t, '_inv_scale'):

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama405b/implementations/tinybox_8xMI350X/dev_beam.sh

@@ -2,7 +2,6 @@
 
 export PYTHONPATH="."
 export DEV=${DEV:-AMD}
-export EMULATE="AMD_CDNA4"
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
 export DEVICE_IN_FUNCTION_BUG=1
@@ -10,14 +9,22 @@ export DEVICE_IN_FUNCTION_BUG=1
 export DEBUG=${DEBUG:-2}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
 export ALL2ALL=${ALL2ALL:-1}
-export USE_ATOMICS=${USE_ATOMICS:-0}
+export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
+export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
+export FP8=${FP8:-1}
+export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
+export FAST_CE=${FAST_CE:-0}
+export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
+export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
+export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
+export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-1}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=${DP:-1} MP=${MP:-8}
-export BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
+export DP=${DP:-1} MP=${MP:-8} BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
+export GBS=$((BS * GRADIENT_ACC_STEPS))
 
 export MODEL="llama3"
 export BASEDIR="/raid/datasets/c4/"
@@ -30,7 +37,7 @@ export DATA_SEED=${DATA_SEED:-5760}
 export JITBEAM=${JITBEAM:-3}
 export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
 
-export FAKEDATA=1 BENCHMARK=10
+export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
 if [ -z "$FULL_LAYERS" ]; then
   export LLAMA_LAYERS=2
 fi

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_beam.sh

@@ -9,13 +9,18 @@ export DEVICE_IN_FUNCTION_BUG=1
 export DEBUG=${DEBUG:-2}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
 export ALL2ALL=${ALL2ALL:-1}
+export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
 export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
 export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
 export FP8=${FP8:-1}
 export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
-export FAST_CE=${FASE_CE:-1}
+export FAST_CE=${FAST_CE:-1}
+export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
+export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
+export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
+export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-1}
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
 export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_beam_mp.sh

@@ -2,7 +2,6 @@
 
 export PYTHONPATH="."
 export DEV=${DEV:-AMD}
-export EMULATE="AMD_CDNA4"
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
 export DEVICE_IN_FUNCTION_BUG=1
@@ -10,9 +9,17 @@ export DEVICE_IN_FUNCTION_BUG=1
 export DEBUG=${DEBUG:-2}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
 export ALL2ALL=${ALL2ALL:-1}
-export USE_ATOMICS=${USE_ATOMICS:-0}
+export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
+export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
+export FP8=${FP8:-1}
+export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
+export FAST_CE=${FAST_CE:-0}
+export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
+export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
+export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
+export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-1}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
@@ -35,7 +42,7 @@ export DATA_SEED=${DATA_SEED:-5760}
 export JITBEAM=${JITBEAM:-3}
 export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
 
-export FAKEDATA=1 BENCHMARK=10
+export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
 if [ -z "$FULL_LAYERS" ]; then
   export LLAMA_LAYERS=2
 fi

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh

@@ -9,13 +9,18 @@ export DEVICE_IN_FUNCTION_BUG=1
 export DEBUG=${DEBUG:-0}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
 export ALL2ALL=${ALL2ALL:-1}
+export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
 export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
 export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
 export FP8=${FP8:-1}
 export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
-export FAST_CE=${FASE_CE:-1}
+export FAST_CE=${FAST_CE:-1}
+export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
+export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
+export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
+export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-1}
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
 export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/profile.sh

@@ -3,4 +3,4 @@ export BENCHMARK=5
 export EVAL_BS=0
 VIZ=${VIZ:--1} FULL_LAYERS=1 DEBUG=0 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_beam.sh
 SRC="AMD"; [[ $DEV == NULL* ]] && SRC="NULL"
-python -m tinygrad.viz.cli -s "$SRC" --top 20
+python -m tinygrad.viz.cli -s "$SRC" -t

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/run_and_time.sh

@@ -10,6 +10,7 @@ export DEVICE_IN_FUNCTION_BUG=1
 
 export HK_FLASH_ATTENTION=1
 export ALL2ALL=1
+export LATE_ALLREDUCE=0
 export USE_ATOMICS=1
 export ASM_GEMM=1
 export WQKV=1
@@ -17,6 +18,10 @@ export MASTER_WEIGHTS=1
 export FP8=1
 export ALLREDUCE_CAST=1
 export FAST_CE=1
+export FUSED_INPUT_QUANTIZE=1
+export FUSED_ADD_NORM_MUL_QUANTIZE=1
+export FUSED_SILU_W13=1
+export FUSED_PAD_GRAD_ACCUM=1
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
 export DP=8 MP=1 BS=16 EVAL_BS=8 GRADIENT_ACC_STEPS=2

extra/amdpci/am_smi.py

@@ -64,7 +64,7 @@ def get_bar0_size(pcibus):
 
 class AMSMI(AMDev):
   def __init__(self, pcibus, vram_bar:MMIOInterface, doorbell_bar:MMIOInterface, mmio_bar:MMIOInterface):
-    self.pcibus = pcibus
+    self.pcibus, self.devfmt = pcibus, pcibus
     self.vram, self.doorbell64, self.mmio = vram_bar, doorbell_bar, mmio_bar
     self.pci_state = self.read_pci_state()
     if self.pci_state == "D0": self._init_from_d0()
@@ -91,6 +91,7 @@ class SMICtx:
     self.prev_lines_cnt = 0
     self.prev_terminal_width = 0
     self.prev_terminal_height = 0
+    self.prev_metrics = {}
 
     remove_parts = ["Advanced Micro Devices, Inc. [AMD/ATI]", "VGA compatible controller:", "Processing accelerators:"]
     lspci = subprocess.check_output(["lspci"]).decode("utf-8").splitlines()
@@ -235,6 +236,29 @@ class SMICtx:
       case (13,0,12): return self._smuq10_round(metrics.SocketPower), self._smuq10_round(metrics.SocketPowerLimit)
       case _: return metrics.SmuMetrics.AverageSocketPower, metrics.SmuMetrics.dGPU_W_MAX
 
+  def get_throttle_info(self, dev, metrics):
+    match dev.ip_ver[am.MP1_HWIP]:
+      case (13,0,6)|(13,0,12):
+        throttle_fields = [('ProchotResidencyAcc', 'Prochot'), ('PptResidencyAcc', 'PPT'),
+                           ('SocketThmResidencyAcc', 'Socket Thm'), ('VrThmResidencyAcc', 'VR Thm'), ('HbmThmResidencyAcc', 'HBM Thm')]
+        prev = self.prev_metrics.get(dev.pcibus)
+        active = []
+        if prev is not None:
+          acc_delta = metrics.AccumulationCounter - prev.AccumulationCounter
+          if acc_delta > 0:
+            for field, name in throttle_fields:
+              delta = getattr(metrics, field) - getattr(prev, field)
+              if delta > 0 and (pct := min(100, (delta * 100 + acc_delta // 2) // acc_delta)) > 0: active.append((name, pct))
+        return active
+      case _:
+        smu_mod = dev.smu.smu_mod
+        throttler_names = {getattr(smu_mod, a): a[len('THROTTLER_'):-len('_BIT')]
+                           for a in dir(smu_mod) if a.startswith('THROTTLER_') and a.endswith('_BIT')}
+        active = []
+        for i, pct in enumerate(metrics.SmuMetrics.ThrottlingPercentage):
+          if pct > 0: active.append((throttler_names.get(i, f"UNK_{i}"), int(pct)))
+        return active
+
   def get_mem_usage(self, dev):
     usage = 0
     pt_stack = [dev.mm.root_page_table]
@@ -281,6 +305,13 @@ class SMICtx:
                     + [f"MEM Activity {draw_bar(self.get_mem_activity(dev, metrics) / 100, activity_line_width)}"] \
                     + [f"MEM Usage    {draw_bar(mem_used / mem_total, activity_line_width, opt_text=mem_fmt)}"] \
 
+      throttle_info = self.get_throttle_info(dev, metrics)
+      if throttle_info:
+        throttle_text = colored(', '.join(f"{name} {pct}%" for name, pct in throttle_info), "red")
+      else:
+        throttle_text = colored("None", "green")
+      activity_line += [f"Throttle     {throttle_text}" + " " * (activity_line_width + 2)]
+
       temps_data, temps_data_compact = self.get_temps(dev, metrics), self.get_temps(dev, metrics, compact=True)
       temps_table = ["=== Temps (°C) ==="] + [f"{name:<16}: {color_temp(val)}" for name, val in temps_data.items()]
       temps_table_compact = ["Temps (°C):" + '/'.join([f"{color_temp(val)} {name}" for name, val in temps_data_compact.items()])]
@@ -324,6 +355,8 @@ class SMICtx:
 
       dev_content.append(device_line + activity_line + same_line([temps_table, power_table, frequency_table]))
 
+    self.prev_metrics = {dev.pcibus: m for dev, m in dev_metrics.items() if m is not None}
+
     raw_text = 'AM Monitor'.center(terminal_width) + "\n" + "=" * terminal_width + "\n\n"
     for i in range(0, len(dev_content), 2):
       if i + 1 < len(dev_content): raw_text += '\n'.join(same_line([dev_content[i], dev_content[i+1]], split=padding))

extra/amdpci/headers/amdgpu_irq.h

@@ -28,15 +28,7 @@
 // #include "soc15_ih_clientid.h"
 // #include "amdgpu_ih.h"
 
-#define int32_t int
-#define uint32_t unsigned int
-#define int8_t signed char
-#define uint8_t unsigned char
-#define uint16_t unsigned short
-#define int16_t short
-#define uint64_t unsigned long long
 #define bool _Bool
-#define u32 unsigned int
 
 #define AMDGPU_MAX_IRQ_SRC_ID		0x100
 #define AMDGPU_MAX_IRQ_CLIENT_ID	0x100

extra/amdpci/headers/amdgpu_smu.h

@@ -22,15 +22,7 @@
 #ifndef __AMDGPU_SMU_H__
 #define __AMDGPU_SMU_H__
 
-#define int32_t int
-#define uint32_t unsigned int
-#define int8_t signed char
-#define uint8_t unsigned char
-#define uint16_t unsigned short
-#define int16_t short
-#define uint64_t unsigned long long
 #define bool _Bool
-#define u32 unsigned int
 
 #define SMU_THERMAL_MINIMUM_ALERT_TEMP		0
 #define SMU_THERMAL_MAXIMUM_ALERT_TEMP		255

extra/amdpci/headers/amdgpu_ucode.h

@@ -24,15 +24,7 @@
 #define __AMDGPU_UCODE_H__
 
 // #include "amdgpu_socbb.h"
-#define int32_t int
-#define uint32_t unsigned int
-#define int8_t signed char
-#define uint8_t unsigned char
-#define uint16_t unsigned short
-#define int16_t short
-#define uint64_t unsigned long long
 #define bool _Bool
-#define u32 unsigned int
 
 struct common_firmware_header {
 	uint32_t size_bytes; /* size of the entire header+image(s) in bytes */

extra/gemm/amd_asm_matmul.py

@@ -167,7 +167,7 @@ PREFETCH_LOADS = [(V_LDS_A_DATA[4+2*i], V_LDS_A_DATA[4+2*i+1], V_GLOBAL_B_ADDR,
 # =============================================================================
 
 class Kernel:
-  def __init__(self, arch='gfx1100'): self.instructions, self.labels, self.pos, self.arch = [], {}, 0, arch
+  def __init__(self): self.instructions, self.labels, self.pos = [], {}, 0
   def label(self, name): self.labels[name] = self.pos
 
   def emit(self, inst, target=None):
@@ -196,10 +196,10 @@ class Kernel:
 # Kernel builder
 # =============================================================================
 
-def build_kernel(N, arch='gfx1100'):
+def build_kernel(N):
   assert N % 128 == 0, f"N must be a multiple of 128 (tile size), got {N}"
   assert N >= 256, f"N must be >= 256 (prefetch pipeline requires at least 2 K-blocks), got {N}"
-  k = Kernel(arch)
+  k = Kernel()
 
   # ===========================================================================
   # PROLOGUE: Load kernel arguments, compute tile coordinates and addresses
@@ -443,7 +443,7 @@ def test_matmul():
   dev = Device[Device.DEFAULT]
   print(f"Device arch: {dev.renderer.target.arch}")
 
-  insts = build_kernel(N, dev.renderer.target.arch)
+  insts = build_kernel(N)
 
   rng = np.random.default_rng(42)
   a = Tensor(rng.random((N, N), dtype=np.float32) - 0.5)

extra/gemm/cdna_asm_gemm.py

@@ -2628,20 +2628,24 @@ def custom_asm_gemm(C:UOp, A:UOp, B:UOp, dname:str) -> UOp:
 # ** FP8 GEMM custom kernel
 
 @functools.cache
-def custom_hk_fp8_gemm(C:UOp, A:UOp, B:UOp, X_s:UOp, W_s:UOp, dname:str) -> UOp:
-  # A is (batch, M, K), B is (N, K) transposed, X_s is x_scale, W_s is w_scale — kernel multiplies by both
+def custom_hk_fp8_gemm(C:UOp, A:UOp, B:UOp, *args:UOp, dname:str, scale_mode:int=3) -> UOp:
+  # scale_mode: 0=no scale, 1=x only, 2=w only, 3=both
+  n_scales = (1 if scale_mode & 1 else 0) + (1 if scale_mode & 2 else 0)
+  scales, extra = args[:n_scales], args[n_scales:]
   M, K = A.shape[0]*A.shape[1], A.shape[2]
   N, K2 = B.shape[(1 if B.ndim == 3 else 0):]
   assert K == K2, f"{A.shape} {B.shape}"
   block_size = 256
   threads = UOp.special(64 * 8, "lidx0")
   workgroups = UOp.special((M // block_size) * (N // block_size), "gidx0")
-  sink = UOp.sink(C.base, A.base, B.base, X_s.base, W_s.base, threads, workgroups,
+  sink_inputs = (C.base, A.base, B.base) + tuple(s.base for s in scales) + (threads, workgroups)
+  sink = UOp.sink(*sink_inputs,
                   arg=KernelInfo(f"hk_fp8_gemm_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K, mem=(M*K+N*K)*A.dtype.itemsize+M*N*C.dtype.itemsize)))
   kittens_path = pathlib.Path(__file__).parent.parent/"thunder"/"amd"
   src = (kittens_path/"gemm_fp8.cpp").read_text()
   lib = HIPCCCompiler("gfx950", [f"-I{(kittens_path/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-ffast-math",
-                                 "-DHIP_ENABLE_WARP_SYNC_BUILTINS", f"-DGEMM_M={M}", f"-DGEMM_N={N}", f"-DGEMM_K={K}"]).compile_cached(src)
+                                 "-DHIP_ENABLE_WARP_SYNC_BUILTINS", f"-DGEMM_M={M}", f"-DGEMM_N={N}", f"-DGEMM_K={K}",
+                                 f"-DSCALE_MODE={scale_mode}"]).compile_cached(src)
   return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
                                UOp(Ops.BINARY, arg=lib)))
 
@@ -2698,19 +2702,37 @@ def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
 
 def custom_gemm_bw(gradient:UOp, kernel:UOp):
   inputs = kernel.src[1:]
-  # fp8 scaled gemm has 5 inputs (out, a, b, x_scale, w_scale), others have 3 (out, a, b)
-  if len(inputs) == 5:
-    out, a, b, s_x, s_w = inputs
+  if inputs[1].dtype == FP8_DTYPE:
+    grad_amax_state = inputs[5] if len(inputs) == 6 else None
+    out, a, b, s_x, s_w = inputs[:5]
     a_t, b_t, g_t = Tensor(a, device=a.device), Tensor(b, device=a.device), Tensor(gradient, device=a.device)
     s_x_t, s_w_t = Tensor(s_x, device=a.device), Tensor(s_w, device=a.device)
     g_t = g_t[:a.shape[0]]
-    g_fp8, g_scale, _ = quantize_fp8(g_t)
+    from extra.llama_kernels.cast_amax import _grad_fp8_mailbox
+    from extra.llama_kernels.quantize_fp8_delayed import quantize_fp8_delayed
+    gbase = gradient.base if hasattr(gradient, "base") else gradient
+    mailbox_entry = _grad_fp8_mailbox.pop(gbase, None) or _grad_fp8_mailbox.pop(gradient, None)
+    if mailbox_entry is not None:
+      g_fp8_u, inv_scale_u, _new_amax_u, store_effect = mailbox_entry
+      g_fp8 = Tensor(g_fp8_u, device=a.device)[:a.shape[0]]
+      g_scale = Tensor(inv_scale_u, device=a.device)
+    else:
+      assert grad_amax_state is not None, "fp8 matmul bwd needs either a mailbox entry or a grad_amax_state"
+      g_fp8, g_scale, _, store_effect = quantize_fp8_delayed(g_t, Tensor(grad_amax_state, device=a.device))
     # dgrad: uses g_scale * x_scale * w_scale
     grad_a = asm_gemm(g_fp8, b_t, x_scale=g_scale * s_x_t, w_scale=s_w_t)
     # wgrad: no w_scale
-    _one = Tensor(1.0, dtype=dtypes.float, device=a.device)
-    grad_b = asm_gemm(g_fp8.permute(2, 0, 1).reshape(g_t.shape[-1], -1), a_t.reshape(-1, a_t.shape[-1]), x_scale=g_scale * s_x_t, w_scale=_one)
-    return (None, grad_a.uop, grad_b.uop, None, None)
+    g_fp8_2d = g_fp8.reshape(-1, g_fp8.shape[-1])
+    if getenv("FAST_FP8_TRANSPOSE", 0) and g_fp8_2d.shape[0] % 64 == 0 and g_fp8_2d.shape[1] % 64 == 0:
+      from extra.llama_kernels.fp8_transpose import fast_fp8_transpose
+      g_fp8_T = fast_fp8_transpose(g_fp8_2d)
+    else:
+      g_fp8_T = g_fp8.permute(2, 0, 1).reshape(g_t.shape[-1], -1)
+    grad_b = asm_gemm(g_fp8_T, a_t.reshape(-1, a_t.shape[-1]), x_scale=g_scale * s_x_t)
+    # Attach the delayed-amax store effect (if any) to grad_a so realizing grads commits the amax update.
+    ret = (None, grad_a.uop.after(store_effect), grad_b.uop, None, None)
+    if len(inputs) == 6: ret = ret + (None,)
+    return ret
   else:
     out, a, b = inputs
     assert all_same([gradient.device, a.device, b.device, out.device])
@@ -2725,7 +2747,7 @@ def custom_gemm_bw(gradient:UOp, kernel:UOp):
 
 # ** main gemm function
 
-def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=None) -> Tensor:
+def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=None, grad_amax_state:Tensor|None=None) -> Tensor:
   assert can_use_asm_gemm(a, b), f"{counters['todos'][-1]}"
   counters["used"] += 1
   unfold_batch = a.ndim == 3 and isinstance(a.device, tuple) and a.uop.axis == 2 and b.uop.axis == 0
@@ -2759,10 +2781,11 @@ def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=N
   if arch.startswith("gfx950") and getenv("USE_ASM", 1):
     # fp8 gemm computes [email protected], kernel multiplies output by x_scale * w_scale before bf16 store
     if a.dtype == FP8_DTYPE:
-      _one = lambda: Tensor(1.0, dtype=dtypes.float, device=a.device)
-      xs = x_scale if x_scale is not None else _one()
-      ws = w_scale if w_scale is not None else _one()
-      out = Tensor.custom_kernel(out, a, b.T, xs, ws, fxn=functools.partial(custom_hk_fp8_gemm, dname=dname), grad_fxn=custom_gemm_bw)[0]
+      scales = tuple(s for s in (x_scale, w_scale) if s is not None)
+      scale_mode = (1 if x_scale is not None else 0) | (2 if w_scale is not None else 0)
+      extra = [grad_amax_state] if grad_amax_state is not None else []
+      fxn = functools.partial(custom_hk_fp8_gemm, dname=dname, scale_mode=scale_mode)
+      out = Tensor.custom_kernel(out, a, b.T, *scales, *extra, fxn=fxn, grad_fxn=custom_gemm_bw)[0]
     else:
       out = Tensor.custom_kernel(out, a, b, fxn=functools.partial(custom_asm_gemm, dname=dname), grad_fxn=custom_gemm_bw)[0]
   else:

extra/gemm/triton_nv_matmul.py

@@ -4,7 +4,8 @@ import triton.language as tl
 from triton.compiler import AttrsDescriptor, ASTSource, compile as triton_compile
 import numpy as np
 from tinygrad import Tensor, dtypes, Device
-from tinygrad.engine.realize import CompiledRunner
+from tinygrad.engine.realize import get_runtime
+from tinygrad.codegen import to_program
 from tinygrad.uop.ops import Ops, UOp, KernelInfo, ProgramInfo
 from tinygrad.helpers import getenv
 np.set_printoptions(suppress=True)
@@ -92,13 +93,15 @@ if __name__ == "__main__":
   info = ProgramInfo(name="matmul_kernel",
                      global_size=(M//BLOCK_SIZE_M, N//BLOCK_SIZE_N, 1), local_size=(32*compiled.metadata.num_warps, 1, 1))
   sink = UOp.sink(arg=KernelInfo(name="matmul_kernel"))
-  prg_uop = UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT), UOp(Ops.LINEAR), UOp(Ops.SOURCE, arg=src)), arg=info)
-  runner = CompiledRunner(prg_uop, Device.DEFAULT)
+  prg_uop = to_program(UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT), UOp(Ops.LINEAR), UOp(Ops.SOURCE, arg=src)), arg=info),
+                       Device.default.renderer)
+  rt = get_runtime(Device.DEFAULT, prg_uop)
   all_bufs = [x.ensure_allocated() for x in bufs]
-  prg_bufs = [all_bufs[i] for i in runner.p.globals]
+  prg_bufs = [all_bufs[i] for i in info.globals]
+  gsize, lsize = info.launch_dims({})
   tflops = []
   for i in range(5):
-    tm = runner(prg_bufs, {}, wait=True)
+    tm = rt(*[b._buf for b in prg_bufs], global_size=gsize, local_size=lsize, vals=info.vals({}), wait=True)
     tflops.append((2*M*K*N/tm)*1e-12)
   print(f"TFLOPS: {max(tflops):.2f}")
 

extra/llama_kernels/__init__.py

@@ -29,7 +29,23 @@ def shard_shape(shape:tuple, axis:int, ndev:int) -> list:
   s[axis] //= ndev
   return s
 
+def dname_of(device) -> str:
+  if isinstance(device, tuple): return device[0].split(":")[0]
+  return device.split(":")[0] if isinstance(device, str) else device
+
+def alloc_like(shape, dtype, device, axis=None) -> Tensor:
+  if isinstance(device, tuple) and axis is not None:
+    return Tensor(Tensor.invalids(*shard_shape(shape, axis, len(device)), dtype=dtype, device=device).uop.multi(axis), device=device)
+  return Tensor.invalids(*shape, dtype=dtype, device=device)
+
+def alloc_local(shape, dtype, device, axis=None) -> Tensor:
+  if isinstance(device, tuple) and axis is not None:
+    return Tensor(Tensor.invalids(*shape, dtype=dtype, device=device).uop.multi(0), device=device)
+  return Tensor.invalids(*shape, dtype=dtype, device=device)
+
+def compile_hip(src:str, defines:list[str]):
+  return HIPCCCompiler("gfx950", ["-std=c++20", "-ffast-math", *defines]).compile_cached(src)
+
 def compile_cpp(cpp_dir:pathlib.Path, cpp_name:str, n_elems:int, hidden:int):
   src = (cpp_dir/cpp_name).read_text()
-  defines = [f"-DN_ELEMS={n_elems}", f"-DHIDDEN={hidden}", f"-DNUM_WG={NUM_WG}", f"-DTHREADS_PER_WG={THREADS_PER_WG}"]
-  return src, HIPCCCompiler("gfx950", ["-std=c++20", "-ffast-math", *defines]).compile_cached(src)
+  return src, compile_hip(src, [f"-DN_ELEMS={n_elems}", f"-DHIDDEN={hidden}", f"-DNUM_WG={NUM_WG}", f"-DTHREADS_PER_WG={THREADS_PER_WG}"])

extra/llama_kernels/cast_amax/__init__.py

@@ -3,26 +3,34 @@ import functools, pathlib
 from tinygrad import Tensor, dtypes
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.renderer import Estimates
-from extra.llama_kernels import FP8_MAX, NUM_WG, THREADS_PER_WG, compile_cpp, shard_shape, scalar_amax
+from extra.llama_kernels import FP8_MAX, NUM_WG, THREADS_PER_WG, compile_cpp, alloc_like, alloc_local, scalar_amax, dname_of
+
+# module-level mailbox: grad_xw13 UOp -> (grad_xw13_fp8 UOp, inv_scale UOp, new_amax UOp, store_effect)
+# lets cdna_asm_gemm's bwd reuse the fp8 companion produced by the fused silu_mul bwd kernel
+# instead of doing a redundant bf16 -> fp8 quantize.
+_grad_fp8_mailbox:dict = {}
 
 @functools.cache
-def _custom_fused_bwd_w13(grad_xw13:UOp, xw13:UOp, grad_x2:UOp, amax_state:UOp, dname:str) -> UOp:
+def _custom_fused_bwd_w13(grad_xw13:UOp, grad_xw13_fp8:UOp, grad_amax_buf:UOp,
+                          xw13:UOp, grad_x2:UOp, amax_state:UOp, grad_amax_state:UOp, dname:str) -> UOp:
   hidden = xw13.shape[2] // 2
   n_elems = xw13.shape[0] * xw13.shape[1] * hidden
   threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
-  mem = n_elems * 2 * 5
-  sink = UOp.sink(grad_xw13.base, xw13.base, grad_x2.base, amax_state.base, threads, workgroups,
-                  arg=KernelInfo(f"fused_silu_mul_bwd_w13_{n_elems}", estimates=Estimates(ops=8*n_elems, mem=mem)))
+  mem = n_elems * 2 * 5 + n_elems * 2 + NUM_WG * 4 + 4
+  sink = UOp.sink(grad_xw13.base, grad_xw13_fp8.base, grad_amax_buf.base,
+                  xw13.base, grad_x2.base, amax_state.base, grad_amax_state.base, threads, workgroups,
+                  arg=KernelInfo(f"fused_silu_mul_bwd_w13_{n_elems}", estimates=Estimates(ops=10*n_elems, mem=mem)))
   src, lib = compile_cpp(pathlib.Path(__file__).parent, "cast_amax_bwd_w13.cpp", n_elems, hidden)
   return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
                                UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
 
 @functools.cache
-def _custom_fused_cast_amax_w13(fp8_out:UOp, amax_buf:UOp, xw13:UOp, amax_state:UOp, dname:str) -> UOp:
+def _custom_fused_cast_amax_w13(fp8_out:UOp, amax_buf:UOp, xw13:UOp, amax_state:UOp, grad_amax_state:UOp, dname:str) -> UOp:
+  # NOTE: grad_amax_state is plumbed through as an unused fwd input so the bwd kernel can read it via kernel.src
   hidden = xw13.shape[2] // 2
   n_elems = xw13.shape[0] * xw13.shape[1] * hidden
   threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
-  mem = n_elems * 2 * 2 + n_elems + NUM_WG * 2
+  mem = n_elems * 2 * 2 + n_elems + NUM_WG * 4
   sink = UOp.sink(fp8_out.base, amax_buf.base, xw13.base, amax_state.base, threads, workgroups,
                   arg=KernelInfo(f"fused_silu_mul_cast_amax_w13_{n_elems}", estimates=Estimates(ops=5*n_elems, mem=mem)))
   src, lib = compile_cpp(pathlib.Path(__file__).parent, "cast_amax_fwd_w13.cpp", n_elems, hidden)
@@ -30,44 +38,37 @@ def _custom_fused_cast_amax_w13(fp8_out:UOp, amax_buf:UOp, xw13:UOp, amax_state:
                                UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
 
 def _fused_quantize_bwd_w13(gradient:UOp, kernel:UOp):
-  # NOTE: inputs are (fp8_out, amax_buf, xw13, amax_state); grad for xw13 only
-  _, _, xw13, amax_state = kernel.src[1:]
+  _, _, xw13, amax_state, grad_amax_state = kernel.src[1:]
   device = xw13.device
-  if isinstance(device, tuple):
-    axis, ndev = xw13.axis, len(device)
-    assert axis in (0, 1), f"unsupported sharding axis={axis}"
-    grad_xw13 = Tensor(Tensor.invalids(*shard_shape(xw13.shape, axis, ndev), dtype=dtypes.bfloat16,
-                                        device=device).uop.multi(axis), device=device)
-    dname = device[0].split(":")[0]
-  else:
-    grad_xw13 = Tensor.invalids(*xw13.shape, dtype=dtypes.bfloat16, device=device)
-    dname = device.split(":")[0] if isinstance(device, str) else device
-  grad_x2_t = Tensor(gradient, device=device).cast(dtypes.bfloat16)
-  fxn = functools.partial(_custom_fused_bwd_w13, dname=dname)
-  grad_xw13, *_ = Tensor.custom_kernel(grad_xw13, Tensor(xw13, device=device), grad_x2_t,
-                                        Tensor(amax_state, device=device), fxn=fxn)
-  return (None, None, grad_xw13.uop, None)
+  axis = xw13.axis if isinstance(device, tuple) else None
+  grad_xw13     = alloc_like(xw13.shape, dtypes.bfloat16, device, axis)
+  grad_xw13_fp8 = alloc_like(xw13.shape, dtypes.fp8e4m3,  device, axis)
+  grad_amax_buf = alloc_local((NUM_WG,), dtypes.float32,  device, axis)
+  grad_amax_state_t = Tensor(grad_amax_state, device=device)
+  fxn = functools.partial(_custom_fused_bwd_w13, dname=dname_of(device))
+  grad_xw13, grad_xw13_fp8, grad_amax_buf, *_ = Tensor.custom_kernel(
+    grad_xw13, grad_xw13_fp8, grad_amax_buf,
+    Tensor(xw13, device=device), Tensor(gradient, device=device).cast(dtypes.bfloat16),
+    Tensor(amax_state, device=device), grad_amax_state_t, fxn=fxn)
+  inv_scale = (grad_amax_state_t.float() + 1e-8) / FP8_MAX
+  new_grad_amax = scalar_amax(grad_amax_buf)
+  store_effect = grad_amax_state_t.uop.store(new_grad_amax.uop)
+  # Stash fp8 companion + amax store for cdna_asm_gemm's bwd to attach to grad_a.
+  _grad_fp8_mailbox[grad_xw13.uop] = (grad_xw13_fp8.uop, inv_scale.uop, new_grad_amax.uop, store_effect)
+  return (None, None, grad_xw13.uop, None, None)
 
-def fused_quantize_fp8_w13(xw13:Tensor, amax_state:Tensor, fp8_dtype) -> tuple[Tensor, Tensor, Tensor]:
+def fused_quantize_fp8_w13(xw13:Tensor, amax_state:Tensor, fp8_dtype, grad_amax_state:Tensor) -> tuple[Tensor, Tensor, Tensor]:
   # NOTE: silu(xw1)*xw3 -> fp8 + amax over fused xw13 layout. Returns (fp8, inv_scale, new_amax)
+  # grad_amax_state: delayed amax for grad_xw13 fp8 quantization in the backward.
   assert xw13.dtype == dtypes.bfloat16, f"expected bf16, got {xw13.dtype}"
   MBS, SEQ, H2 = xw13.shape
   assert H2 % 2 == 0, f"w13 last-axis must be even, got {H2}"
   HIDDEN = H2 // 2
-  if isinstance(xw13.device, tuple):
-    axis, ndev = xw13.uop.axis, len(xw13.device)
-    assert axis in (0, 1), f"unsupported sharding axis={axis}"
-    fp8_out = Tensor(Tensor.invalids(*shard_shape((MBS, SEQ, HIDDEN), axis, ndev), dtype=fp8_dtype,
-                                     device=xw13.device).uop.multi(axis), device=xw13.device)
-    amax_buf = Tensor(Tensor.invalids(NUM_WG, dtype=dtypes.bfloat16, device=xw13.device).uop.multi(0),
-                      device=xw13.device)
-    dname = xw13.device[0].split(":")[0]
-  else:
-    fp8_out = Tensor.invalids(MBS, SEQ, HIDDEN, dtype=fp8_dtype, device=xw13.device)
-    amax_buf = Tensor.invalids(NUM_WG, dtype=dtypes.bfloat16, device=xw13.device)
-    dname = xw13.device.split(":")[0] if isinstance(xw13.device, str) else xw13.device
-  fxn = functools.partial(_custom_fused_cast_amax_w13, dname=dname)
-  fp8_out, amax_buf, *_ = Tensor.custom_kernel(fp8_out, amax_buf, xw13, amax_state, fxn=fxn,
-                                                grad_fxn=_fused_quantize_bwd_w13)
+  axis = xw13.uop.axis if isinstance(xw13.device, tuple) else None
+  fp8_out  = alloc_like((MBS, SEQ, HIDDEN), fp8_dtype,      xw13.device, axis)
+  amax_buf = alloc_local((NUM_WG,),         dtypes.float32, xw13.device, axis)
+  fxn = functools.partial(_custom_fused_cast_amax_w13, dname=dname_of(xw13.device))
+  fp8_out, amax_buf, *_ = Tensor.custom_kernel(fp8_out, amax_buf, xw13, amax_state, grad_amax_state,
+                                                fxn=fxn, grad_fxn=_fused_quantize_bwd_w13)
   inv_scale = (amax_state.float() + 1e-8) / FP8_MAX
   return fp8_out, inv_scale, scalar_amax(amax_buf)

extra/llama_kernels/cast_amax/cast_amax_bwd_w13.cpp

@@ -1,5 +1,6 @@
 #include <hip/hip_runtime.h>
 #include <hip/hip_bf16.h>
+#include <hip/hip_fp8.h>
 
 #ifndef N_ELEMS
 #define N_ELEMS 234881024
@@ -20,19 +21,32 @@ constexpr float FP8_MAX = 448.0f;
 static_assert(N_ELEMS % VEC == 0, "N_ELEMS must be divisible by VEC");
 static_assert(HIDDEN % VEC == 0, "HIDDEN must be divisible by VEC");
 
+// fused silu*mul backward, three outputs in a single HBM pass:
+//   1) bf16 grad_xw13      — consumed by downstream bf16 autograd chain
+//   2) fp8  grad_xw13_fp8  — delayed-scale quantize using grad_amax_state (mailbox to matmul bwd)
+//   3) fp32 grad_amax_buf  — per-WG partial |grad_xw13|, reduced into next step's grad_amax_state
+// grad_amax_state is read for the fp8 scale. The store of new_grad_amax into grad_amax_state's
+// buffer is built in Python as a separate effect and threaded into grad_a via .after(store).
 extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
 fused_silu_mul_bwd_w13(
-    __hip_bfloat16*       __restrict__ grad_xw13_out,    // bf16, 2*N_ELEMS (interleaved layout)
-    const __hip_bfloat16* __restrict__ xw13,             // bf16, 2*N_ELEMS (interleaved)
-    const __hip_bfloat16* __restrict__ grad_x2,          // bf16, N_ELEMS
-    const __hip_bfloat16* __restrict__ amax_state)       // bf16 scalar
+    __hip_bfloat16*       __restrict__ grad_xw13_out,        // bf16, 2*N_ELEMS
+    __hip_fp8_storage_t*  __restrict__ grad_xw13_fp8_out,    // fp8,  2*N_ELEMS
+    float*                __restrict__ grad_amax_buf,        // fp32, NUM_WG per-WG partials
+    const __hip_bfloat16* __restrict__ xw13,                 // bf16, 2*N_ELEMS
+    const __hip_bfloat16* __restrict__ grad_x2,              // bf16, N_ELEMS
+    const float*          __restrict__ amax_state,           // fp32 scalar (fwd x2 amax)
+    const float*          __restrict__ grad_amax_state)      // fp32 scalar (delayed grad amax)
 {
+  __shared__ float sdata[THREADS_PER_WG];
+
   const int tid = threadIdx.x;
   const int wg  = blockIdx.x;
   const int gid = wg * THREADS_PER_WG + tid;
   const int stride_elems = NUM_WG * THREADS_PER_WG * VEC;
 
   const float scale = FP8_MAX / (static_cast<float>(*amax_state) + 1e-8f);
+  const float g_scale = FP8_MAX / (static_cast<float>(*grad_amax_state) + 1e-8f);
+  float local_max = 0.0f;
 
   for (int base = gid * VEC; base < N_ELEMS; base += stride_elems) {
     const int outer = base / HIDDEN;
@@ -49,6 +63,7 @@ fused_silu_mul_bwd_w13(
     const __hip_bfloat16 *gv = reinterpret_cast<const __hip_bfloat16*>(&g_raw);
 
     __hip_bfloat16 out1[VEC], out3[VEC];
+    __hip_fp8_storage_t fp8_1[VEC], fp8_3[VEC];
     #pragma unroll
     for (int i = 0; i < VEC; i++) {
       const float f1 = static_cast<float>(x1[i]);
@@ -58,11 +73,26 @@ fused_silu_mul_bwd_w13(
       const float silu = f1 * sig;
       const float silu_prime = sig + silu * (1.0f - sig);
       const float gs = fg * scale;
-      out1[i] = static_cast<__hip_bfloat16>(gs * silu_prime * f3);
-      out3[i] = static_cast<__hip_bfloat16>(gs * silu);
+      const float g1 = gs * silu_prime * f3;
+      const float g3 = gs * silu;
+      out1[i] = static_cast<__hip_bfloat16>(g1);
+      out3[i] = static_cast<__hip_bfloat16>(g3);
+      local_max = fmaxf(local_max, fmaxf(fabsf(g1), fabsf(g3)));
+      fp8_1[i] = __hip_cvt_float_to_fp8(fmaxf(-FP8_MAX, fminf(FP8_MAX, g1 * g_scale)), __HIP_SATFINITE, __HIP_E4M3);
+      fp8_3[i] = __hip_cvt_float_to_fp8(fmaxf(-FP8_MAX, fminf(FP8_MAX, g3 * g_scale)), __HIP_SATFINITE, __HIP_E4M3);
     }
 
     *reinterpret_cast<float4*>(&grad_xw13_out[xw1_off]) = *reinterpret_cast<float4*>(out1);
     *reinterpret_cast<float4*>(&grad_xw13_out[xw3_off]) = *reinterpret_cast<float4*>(out3);
+    *reinterpret_cast<uint64_t*>(&grad_xw13_fp8_out[xw1_off]) = *reinterpret_cast<uint64_t*>(fp8_1);
+    *reinterpret_cast<uint64_t*>(&grad_xw13_fp8_out[xw3_off]) = *reinterpret_cast<uint64_t*>(fp8_3);
   }
+
+  sdata[tid] = local_max;
+  __syncthreads();
+  for (int s = THREADS_PER_WG / 2; s > 0; s >>= 1) {
+    if (tid < s) sdata[tid] = fmaxf(sdata[tid], sdata[tid + s]);
+    __syncthreads();
+  }
+  if (tid == 0) grad_amax_buf[wg] = sdata[0];
 }

extra/llama_kernels/cast_amax/cast_amax_fwd_w13.cpp

@@ -24,9 +24,9 @@ static_assert(HIDDEN % VEC == 0, "HIDDEN must be divisible by VEC (so VEC loads
 extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
 fused_silu_mul_cast_amax_w13(
     __hip_fp8_storage_t*  __restrict__ fp8_out,         // fp8, N_ELEMS
-    __hip_bfloat16*       __restrict__ amax_buf,        // bf16, NUM_WG (per-WG amaxes)
+    float*                __restrict__ amax_buf,        // fp32, NUM_WG (per-WG amaxes)
     const __hip_bfloat16* __restrict__ xw13,            // bf16, 2*N_ELEMS
-    const __hip_bfloat16* __restrict__ amax_state)      // bf16 scalar
+    const float*          __restrict__ amax_state)      // fp32 scalar
 {
   __shared__ float sdata[THREADS_PER_WG];
 
@@ -75,5 +75,5 @@ fused_silu_mul_cast_amax_w13(
     __syncthreads();
   }
 
-  if (tid == 0) amax_buf[wg] = static_cast<__hip_bfloat16>(sdata[0]);
+  if (tid == 0) amax_buf[wg] = sdata[0];
 }

extra/llama_kernels/fp8_transpose/__init__.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+import functools, pathlib
+from tinygrad import Tensor, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.renderer import Estimates
+from extra.llama_kernels import THREADS_PER_WG, alloc_like, dname_of, compile_hip
+
+TILE = 64
+
+@functools.cache
+def _custom_fp8_transpose(out:UOp, inp:UOp, dname:str) -> UOp:
+  M, N = inp.shape
+  num_wg = (M // TILE) * (N // TILE)
+  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(num_wg, "gidx0")
+  mem = M * N * 2  # one byte read + one byte write per element
+  sink = UOp.sink(out.base, inp.base, threads, workgroups,
+                  arg=KernelInfo(f"fp8_transpose_{M}_{N}",
+                                 estimates=Estimates(ops=M*N, mem=mem)))
+  src = (pathlib.Path(__file__).parent/"fp8_transpose.cpp").read_text()
+  defines = [f"-DM_DIM={M}", f"-DN_DIM={N}", f"-DTHREADS_PER_WG={THREADS_PER_WG}"]
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))
+
+def fast_fp8_transpose(t:Tensor) -> Tensor:
+  assert t.ndim == 2, f"fast_fp8_transpose needs 2D input, got shape {t.shape}"
+  assert t.dtype in dtypes.fp8s, f"fast_fp8_transpose needs fp8 dtype, got {t.dtype}"
+  M, N = t.shape
+  assert M % TILE == 0 and N % TILE == 0, f"M={M}, N={N} must be multiples of {TILE}"
+
+  device = t.device
+  axis = t.uop.axis if isinstance(device, tuple) else None
+  out_axis = None
+  if axis == 0: out_axis = 1
+  elif axis == 1: out_axis = 0
+  elif axis is not None:
+    raise ValueError(f"fast_fp8_transpose: unsupported axis {axis}")
+
+  out = alloc_like((N, M), t.dtype, device, out_axis)
+  fxn = functools.partial(_custom_fp8_transpose, dname=dname_of(device))
+  out, _ = Tensor.custom_kernel(out, t, fxn=fxn)
+  return out

extra/llama_kernels/fp8_transpose/fp8_transpose.cpp

@@ -0,0 +1,74 @@
+#include <hip/hip_runtime.h>
+
+// LDS-staged 64x64 fp8 transpose.
+//   in : (M_DIM, N_DIM) fp8 contiguous
+//   out: (N_DIM, M_DIM) fp8 contiguous, out[c][r] = in[r][c]
+//
+// One WG processes one 64x64 output tile.  Each thread reads one uint4 (16 fp8) coalesced
+// from input rows, stages into LDS, then writes one uint4 coalesced to the output (whose
+// 16 fp8 come from 16 different input rows via in-LDS gather).
+//
+// LDS layout: lds[64][LDS_STRIDE] with LDS_STRIDE=65 (1 byte pad) to mitigate bank conflicts
+// during the column-direction read of the write phase.
+
+#ifndef M_DIM
+#define M_DIM 16384
+#endif
+#ifndef N_DIM
+#define N_DIM 28672
+#endif
+#ifndef THREADS_PER_WG
+#define THREADS_PER_WG 256
+#endif
+
+constexpr int TILE = 64;
+constexpr int VEC = 16;                    // fp8 per uint4 (128-bit) load/store
+constexpr int LDS_PAD = 1;
+constexpr int LDS_STRIDE = TILE + LDS_PAD; // 65 fp8 per row
+
+static_assert(THREADS_PER_WG * VEC == TILE * TILE, "256 threads * 16 fp8 = 64*64");
+static_assert(M_DIM % TILE == 0, "M_DIM must be a multiple of 64");
+static_assert(N_DIM % TILE == 0, "N_DIM must be a multiple of 64");
+
+constexpr int N_TILES_N = N_DIM / TILE;
+
+struct alignas(16) fp8x16 { uint8_t v[16]; };
+
+extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
+fp8_transpose(uint8_t* __restrict__ out,         // (N_DIM, M_DIM)
+              const uint8_t* __restrict__ in)    // (M_DIM, N_DIM)
+{
+  __shared__ uint8_t lds[TILE * LDS_STRIDE];
+
+  const int tid     = threadIdx.x;
+  const int wg_id   = blockIdx.x;
+  const int tile_r  = wg_id / N_TILES_N;       // tile index along M dim of input
+  const int tile_c  = wg_id % N_TILES_N;       // tile index along N dim of input
+
+  const int a = tid / (TILE / VEC);            // 0..63 (row within tile during read; col within tile during write)
+  const int b = tid % (TILE / VEC);            // 0..3
+  const int b16 = b * VEC;                     // 0,16,32,48
+
+  // ---- Read phase: input rows -> LDS rows
+  {
+    const long long src = (long long)(tile_r * TILE + a) * (long long)N_DIM
+                        + (long long)(tile_c * TILE + b16);
+    fp8x16 v = *reinterpret_cast<const fp8x16*>(&in[src]);
+    *reinterpret_cast<fp8x16*>(&lds[a * LDS_STRIDE + b16]) = v;
+  }
+  __syncthreads();
+
+  // ---- Write phase: LDS columns (gathered) -> output rows
+  // out[(tile_c*TILE + a)][(tile_r*TILE + b16 + i)] = in[(tile_r*TILE + b16 + i)][(tile_c*TILE + a)]
+  //                                                  = lds[b16 + i][a]
+  {
+    fp8x16 v;
+    #pragma unroll
+    for (int i = 0; i < VEC; ++i) {
+      v.v[i] = lds[(b16 + i) * LDS_STRIDE + a];
+    }
+    const long long dst = (long long)(tile_c * TILE + a) * (long long)M_DIM
+                        + (long long)(tile_r * TILE + b16);
+    *reinterpret_cast<fp8x16*>(&out[dst]) = v;
+  }
+}

extra/llama_kernels/fused_ce/__init__.py

@@ -42,8 +42,7 @@ def _fused_ce_loss_bwd(gradient:UOp, kernel:UOp, label_smoothing:float):
   # gradient is the upstream grad w.r.t. per-row loss (shape: (rows,) fp32)
   _, _, lse_u, logits_u, targets_u = kernel.src[1:]
   device = logits_u.device
-  rows_vocab = logits_u.shape  # (rows, VOCAB) after reshape
-  rows, VOCAB = rows_vocab
+  rows, VOCAB = logits_u.shape  # (rows, VOCAB) after reshape
   if isinstance(device, tuple):
     axis = logits_u.axis
     ndev = len(device)
@@ -54,9 +53,8 @@ def _fused_ce_loss_bwd(gradient:UOp, kernel:UOp, label_smoothing:float):
     d_logits = Tensor.invalids(rows, VOCAB, dtype=dtypes.bfloat16, device=device)
     dname = device.split(":")[0] if isinstance(device, str) else device
     rows_per_dev = rows
-  grad_t = Tensor(gradient, device=device).float().reshape(-1)  # (rows,) fp32
   # NOTE: .mean() backward gives same grad per row (1/N), so broadcast is safe; take scalar
-  scale = grad_t[0:1].contiguous()
+  scale = Tensor(gradient, device=device).float().reshape(-1)[0:1].contiguous()
   logits_t = Tensor(logits_u.after(kernel), device=device)
   lse_t = Tensor(lse_u.after(kernel), device=device)
   targets_t = Tensor(targets_u, device=device)

extra/llama_kernels/fused_mul_quantize_fp8/__init__.py

@@ -1,54 +0,0 @@
-from __future__ import annotations
-import functools, pathlib
-from tinygrad import Tensor, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
-from tinygrad.renderer import Estimates
-from extra.llama_kernels import FP8_MAX, NUM_WG, THREADS_PER_WG, compile_cpp, shard_shape, scalar_amax
-
-@functools.cache
-def _custom_mul_quantize_fp8(fp8_out:UOp, amax_buf:UOp, x:UOp, weight:UOp, amax_state:UOp, dname:str) -> UOp:
-  MBS, SEQ, HIDDEN = x.shape
-  n_elems = MBS * SEQ * HIDDEN
-  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
-  mem = n_elems * 2 + HIDDEN * 2 + n_elems + NUM_WG * 2
-  sink = UOp.sink(fp8_out.base, amax_buf.base, x.base, weight.base, amax_state.base, threads, workgroups,
-                  arg=KernelInfo(f"fused_mul_quantize_fp8_{n_elems}_h{HIDDEN}", estimates=Estimates(ops=3*n_elems, mem=mem)))
-  src, lib = compile_cpp(pathlib.Path(__file__).parent, "fused_mul_quantize_fp8.cpp", n_elems, HIDDEN)
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
-                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
-
-def _fused_mul_quantize_fp8_bwd(gradient:UOp, kernel:UOp):
-  # NOTE: inputs are (fp8_out, amax_buf, x, weight, amax_state); grads for x and weight
-  _, _, x_u, weight_u, amax_state_u = kernel.src[1:]
-  device = x_u.device
-  grad_t = Tensor(gradient, device=device).cast(dtypes.bfloat16)
-  x_t, weight_t = Tensor(x_u, device=device), Tensor(weight_u, device=device)
-  scale = FP8_MAX / (Tensor(amax_state_u, device=device).float() + 1e-8)
-  grad_scaled = grad_t.float() * scale
-  # NOTE: grad_x stays bf16 to avoid CSE materializing a (MBS, SEQ, HIDDEN) fp32 intermediate
-  grad_x = (grad_scaled * weight_t.float()).cast(dtypes.bfloat16)
-  grad_weight = (grad_scaled * x_t.float()).sum(axis=(0, 1)).cast(dtypes.bfloat16)
-  return (None, None, grad_x.uop, grad_weight.uop, None)
-
-def fused_mul_quantize_fp8(x:Tensor, weight:Tensor, amax_state:Tensor, fp8_dtype) -> tuple[Tensor, Tensor, Tensor]:
-  # NOTE: (x * weight) -> fp8 + amax, delayed scaling. Returns (fp8, inv_scale, new_amax)
-  assert x.dtype == dtypes.bfloat16 and weight.dtype == dtypes.bfloat16
-  assert x.shape[-1] == weight.shape[-1], f"HIDDEN mismatch: x={x.shape}, weight={weight.shape}"
-  MBS, SEQ, HIDDEN = x.shape
-  if isinstance(x.device, tuple):
-    axis, ndev = x.uop.axis, len(x.device)
-    assert axis in (0, 1), f"unsupported sharding axis={axis}"
-    fp8_out = Tensor(Tensor.invalids(*shard_shape((MBS, SEQ, HIDDEN), axis, ndev), dtype=fp8_dtype,
-                                     device=x.device).uop.multi(axis), device=x.device)
-    amax_buf = Tensor(Tensor.invalids(NUM_WG, dtype=dtypes.bfloat16, device=x.device).uop.multi(0), device=x.device)
-    dname = x.device[0].split(":")[0]
-  else:
-    fp8_out = Tensor.invalids(MBS, SEQ, HIDDEN, dtype=fp8_dtype, device=x.device)
-    amax_buf = Tensor.invalids(NUM_WG, dtype=dtypes.bfloat16, device=x.device)
-    dname = x.device.split(":")[0] if isinstance(x.device, str) else x.device
-  fxn = functools.partial(_custom_mul_quantize_fp8, dname=dname)
-  fp8_out, amax_buf, *_ = Tensor.custom_kernel(fp8_out, amax_buf, x, weight, amax_state, fxn=fxn,
-                                                grad_fxn=_fused_mul_quantize_fp8_bwd)
-  new_amax = scalar_amax(amax_buf)
-  inv_scale = (amax_state.float() + 1e-8) / FP8_MAX
-  return fp8_out, inv_scale, new_amax

extra/llama_kernels/fused_pad_grad_accum/__init__.py

@@ -0,0 +1,55 @@
+from __future__ import annotations
+import functools, pathlib
+from tinygrad import Tensor, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.renderer import Estimates
+from extra.llama_kernels import THREADS_PER_WG, dname_of, compile_hip
+
+ELEMS_PER_THREAD = 8  # vectorized 16-byte load (uint4 = 8 bf16)
+
+def _build_src(n_chunks:int) -> str:
+  template = (pathlib.Path(__file__).parent/"fused_pad_grad_accum.cpp").read_text()
+  params = "".join(f",\n    const __hip_bfloat16* __restrict__ chunk{i}" for i in range(n_chunks))
+  dispatch = "\n    ".join(f"case {i}: chunk_ptr = chunk{i}; break;" for i in range(n_chunks))
+  return (template.replace("__FUSED_PAD_GRAD_ACCUM_PARAMS", params)
+                  .replace("__FUSED_PAD_GRAD_ACCUM_DISPATCH", dispatch))
+
+@functools.cache
+def _custom_fused_pad_grad_accum(grad_buf:UOp, *chunk_uops, dname:str, n_chunks:int, chunk_size:int) -> UOp:
+  total = n_chunks * chunk_size
+  elems_per_block = THREADS_PER_WG * ELEMS_PER_THREAD
+  assert chunk_size % elems_per_block == 0, f"chunk_size {chunk_size} must be multiple of {elems_per_block}"
+  num_wg = total // elems_per_block
+  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(num_wg, "gidx0")
+  mem = total * 2 * 3
+  sink = UOp.sink(grad_buf.base, *(c.base for c in chunk_uops), threads, workgroups,
+                  arg=KernelInfo(f"fused_pad_grad_accum_n{n_chunks}_c{chunk_size}",
+                                 estimates=Estimates(ops=2*total, mem=mem)))
+  src = _build_src(n_chunks)
+  defines = [f"-DCHUNK_SIZE={chunk_size}", f"-DTHREADS_PER_WG={THREADS_PER_WG}", f"-DELEMS_PER_THREAD={ELEMS_PER_THREAD}"]
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))
+
+def can_fused_pad_grad_accum(grad_buf:Tensor, chunks:list[Tensor]) -> bool:
+  if not chunks or grad_buf.dtype != dtypes.bfloat16: return False
+  if any(c.dtype != dtypes.bfloat16 for c in chunks): return False
+  chunk_shape = chunks[0].shape
+  if any(c.shape != chunk_shape for c in chunks): return False
+  chunk_size, total = 1, 1
+  for d in chunk_shape: chunk_size *= d
+  for d in grad_buf.shape: total *= d
+  return total == len(chunks) * chunk_size and chunk_size % (THREADS_PER_WG * ELEMS_PER_THREAD) == 0
+
+def fused_pad_grad_accum(grad_buf:Tensor, chunks:list[Tensor]) -> Tensor:
+  # NOTE: grad_buf += cat(*chunks, dim=0) in one HBM pass (in-place add). Returns new grad_buf Tensor.
+  # Requires uniform chunk shapes and chunk_size % (THREADS_PER_WG*ELEMS_PER_THREAD) == 0.
+  assert chunks and grad_buf.dtype == dtypes.bfloat16
+  for c in chunks: assert c.dtype == dtypes.bfloat16, f"chunk dtype must be bf16, got {c.dtype}"
+  chunk_size, total = 1, 1
+  for d in chunks[0].shape: chunk_size *= d
+  for d in grad_buf.shape: total *= d
+  assert total == len(chunks) * chunk_size, f"grad_buf size {total} != n_chunks {len(chunks)} * chunk_size {chunk_size}"
+  fxn = functools.partial(_custom_fused_pad_grad_accum, dname=dname_of(grad_buf.device),
+                          n_chunks=len(chunks), chunk_size=chunk_size)
+  out, *_ = Tensor.custom_kernel(grad_buf, *chunks, fxn=fxn)
+  return out

extra/llama_kernels/fused_pad_grad_accum/fused_pad_grad_accum.cpp

@@ -0,0 +1,63 @@
+// Fused custom kernel: grad_buf += cat(*chunks, dim=0) in one HBM pass.
+//
+// Template source — chunk parameter list and switch dispatch are filled by codegen
+// in cast_amax.py:_build_fused_pad_grad_accum_src to support arbitrary N.
+//
+// Defines required at compile time:
+//   CHUNK_SIZE       elements per chunk (must be multiple of THREADS_PER_WG * ELEMS_PER_THREAD)
+//   THREADS_PER_WG
+//   ELEMS_PER_THREAD (8 = one uint4 per thread = 16-byte vectorized load)
+//
+// Layout: one block-per-(slice-of-chunk) — blockIdx.x / BLOCKS_PER_CHUNK selects the chunk.
+// All threads in a block read the same chunk → switch is uniform → no warp divergence.
+
+#include <hip/hip_runtime.h>
+#include <hip/hip_bf16.h>
+
+#ifndef THREADS_PER_WG
+#define THREADS_PER_WG 256
+#endif
+#ifndef ELEMS_PER_THREAD
+#define ELEMS_PER_THREAD 8
+#endif
+
+#define ELEMS_PER_BLOCK (THREADS_PER_WG * ELEMS_PER_THREAD)
+#define BLOCKS_PER_CHUNK (CHUNK_SIZE / ELEMS_PER_BLOCK)
+
+extern "C" __attribute__((global))
+__attribute__((amdgpu_flat_work_group_size(1, THREADS_PER_WG)))
+void fused_pad_grad_accum(
+    __hip_bfloat16* __restrict__ grad_buf
+    __FUSED_PAD_GRAD_ACCUM_PARAMS
+) {
+  const int bid = blockIdx.x;
+  const int chunk_idx = bid / BLOCKS_PER_CHUNK;
+  const int block_in_chunk = bid - chunk_idx * BLOCKS_PER_CHUNK;
+  const int tid = threadIdx.x;
+
+  const __hip_bfloat16* chunk_ptr;
+  switch (chunk_idx) {
+    __FUSED_PAD_GRAD_ACCUM_DISPATCH
+    default: chunk_ptr = (const __hip_bfloat16*)0; break;  // unreachable
+  }
+
+  // int64 for global_offset: at 32 chunks × 117M elements = 3.6B, int32 overflows → MEMVIOL.
+  const int local_offset = block_in_chunk * ELEMS_PER_BLOCK + tid * ELEMS_PER_THREAD;
+  const long long global_offset = (long long)chunk_idx * (long long)CHUNK_SIZE + (long long)local_offset;
+
+  // Vectorized 16-byte load (uint4 = 8 bf16). Requires CHUNK_SIZE % 8 == 0 and 16-byte alignment.
+  const uint4 chunk_v = *reinterpret_cast<const uint4*>(&chunk_ptr[local_offset]);
+  const uint4 grad_v  = *reinterpret_cast<const uint4*>(&grad_buf[global_offset]);
+  uint4 out_v;
+
+  const __hip_bfloat16* chunk_bf = reinterpret_cast<const __hip_bfloat16*>(&chunk_v);
+  const __hip_bfloat16* grad_bf  = reinterpret_cast<const __hip_bfloat16*>(&grad_v);
+  __hip_bfloat16*       out_bf   = reinterpret_cast<__hip_bfloat16*>(&out_v);
+
+  #pragma unroll
+  for (int i = 0; i < ELEMS_PER_THREAD; i++) {
+    out_bf[i] = (__hip_bfloat16)((float)grad_bf[i] + (float)chunk_bf[i]);
+  }
+
+  *reinterpret_cast<uint4*>(&grad_buf[global_offset]) = out_v;
+}

extra/llama_kernels/fused_rmsnorm_mul_quantize_fp8/__init__.py

@@ -0,0 +1,153 @@
+from __future__ import annotations
+import functools, pathlib
+from tinygrad import Tensor, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.renderer import Estimates
+from extra.llama_kernels import FP8_MAX, NUM_WG, THREADS_PER_WG, alloc_like, alloc_local, scalar_amax, dname_of, compile_hip
+
+def _src() -> str: return (pathlib.Path(__file__).parent/"fused_rmsnorm_mul_quantize_fp8.cpp").read_text()
+def _src_bwd() -> str: return (pathlib.Path(__file__).parent/"fused_rmsnorm_mul_quantize_fp8_bwd.cpp").read_text()
+
+@functools.cache
+def _custom_fwd(fp8_out:UOp, x_normed_out:UOp, rrms_out:UOp, amax_buf:UOp,
+                x:UOp, weight:UOp, amax_state:UOp, dname:str, eps_val:float) -> UOp:
+  MBS, SEQ, HIDDEN = x.shape
+  n_elems = MBS * SEQ * HIDDEN
+  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
+  mem = n_elems * 2 + n_elems + MBS * SEQ * 4 + n_elems + HIDDEN * 2 + NUM_WG * 4 + 4
+  sink = UOp.sink(fp8_out.base, x_normed_out.base, rrms_out.base, amax_buf.base,
+                  x.base, weight.base, amax_state.base, threads, workgroups,
+                  arg=KernelInfo(f"fused_rmsnorm_mul_quantize_fp8_{n_elems}_h{HIDDEN}_eps{eps_val:.0e}",
+                                 estimates=Estimates(ops=6*n_elems, mem=mem)))
+  defines = [f"-DN_ELEMS={n_elems}", f"-DHIDDEN={HIDDEN}", f"-DNUM_WG={NUM_WG}", f"-DTHREADS_PER_WG={THREADS_PER_WG}",
+             f"-DEPS_LITERAL={eps_val}f"]
+  src = _src()
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))
+
+@functools.cache
+def _custom_fwd_add(fp8_out:UOp, h_out:UOp, x_normed_out:UOp, rrms_out:UOp, amax_buf:UOp,
+                    x:UOp, residual:UOp, weight:UOp, amax_state:UOp, dname:str, eps_val:float) -> UOp:
+  MBS, SEQ, HIDDEN = x.shape
+  n_elems = MBS * SEQ * HIDDEN
+  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
+  mem = n_elems * 2 * 4 + MBS * SEQ * 4 + HIDDEN * 2 + NUM_WG * 4 + 4
+  sink = UOp.sink(fp8_out.base, h_out.base, x_normed_out.base, rrms_out.base, amax_buf.base,
+                  x.base, residual.base, weight.base, amax_state.base, threads, workgroups,
+                  arg=KernelInfo(f"fused_add_rmsnorm_mul_quantize_fp8_{n_elems}_h{HIDDEN}_eps{eps_val:.0e}",
+                                 estimates=Estimates(ops=7*n_elems, mem=mem)))
+  defines = [f"-DN_ELEMS={n_elems}", f"-DHIDDEN={HIDDEN}", f"-DNUM_WG={NUM_WG}", f"-DTHREADS_PER_WG={THREADS_PER_WG}",
+             f"-DEPS_LITERAL={eps_val}f", f"-DHAS_RESIDUAL=1"]
+  src = _src()
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))
+
+@functools.cache
+def _custom_bwd(grad_x:UOp, grad_weight_partial:UOp,
+                grad_fp8:UOp, x_normed:UOp, rrms:UOp, weight:UOp, amax_state:UOp, dname:str) -> UOp:
+  MBS, SEQ, HIDDEN = x_normed.shape
+  n_elems = MBS * SEQ * HIDDEN
+  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
+  mem = n_elems * 2 * 3 + NUM_WG * HIDDEN * 4 + MBS * SEQ * 4 + HIDDEN * 2 + 4
+  sink = UOp.sink(grad_x.base, grad_weight_partial.base,
+                  grad_fp8.base, x_normed.base, rrms.base, weight.base, amax_state.base, threads, workgroups,
+                  arg=KernelInfo(f"fused_rmsnorm_mul_quantize_fp8_bwd_{n_elems}_h{HIDDEN}",
+                                 estimates=Estimates(ops=8*n_elems, mem=mem)))
+  defines = [f"-DN_ELEMS={n_elems}", f"-DHIDDEN={HIDDEN}", f"-DNUM_WG={NUM_WG}", f"-DTHREADS_PER_WG={THREADS_PER_WG}"]
+  src = _src_bwd()
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))
+
+def _bwd_common(fp8_grad_u, h_grad_u, x_u, x_normed_u, rrms_u, weight_u, amax_state_u, kernel:UOp):
+  device = x_u.device
+  MBS, SEQ, HIDDEN = x_normed_u.shape
+  axis = x_normed_u.axis if isinstance(device, tuple) else None
+  grad_x = alloc_like((MBS, SEQ, HIDDEN), dtypes.bfloat16, device, axis)
+  grad_weight_partial = alloc_local((NUM_WG, HIDDEN), dtypes.float32, device, axis)
+  grad_h_from_fp8 = None
+  grad_weight_uop = None
+  if fp8_grad_u is not None:
+    fxn = functools.partial(_custom_bwd, dname=dname_of(device))
+    grad_x_t, grad_weight_partial_t, *_ = Tensor.custom_kernel(
+      grad_x, grad_weight_partial,
+      Tensor(fp8_grad_u, device=device).cast(dtypes.bfloat16),
+      Tensor(x_normed_u.after(kernel), device=device),
+      Tensor(rrms_u.after(kernel), device=device),
+      Tensor(weight_u, device=device),
+      Tensor(amax_state_u, device=device), fxn=fxn)
+    grad_h_from_fp8 = grad_x_t
+    grad_weight_uop = grad_weight_partial_t.sum(axis=0).cast(dtypes.bfloat16).uop
+  if h_grad_u is not None:
+    h_grad_t = Tensor(h_grad_u, device=device).cast(dtypes.bfloat16)
+    grad_total = (grad_h_from_fp8 + h_grad_t) if grad_h_from_fp8 is not None else h_grad_t
+  else:
+    grad_total = grad_h_from_fp8
+  return grad_total.uop, grad_weight_uop
+
+def _fused_bwd(gradient:UOp, kernel:UOp):
+  # NOTE: fwd inputs (fp8_out, x_normed_out, rrms_out, amax_buf, x, weight, amax_state)
+  _, x_normed_u, rrms_u, _, x_u, weight_u, amax_state_u = kernel.src[1:]
+  grad_x, grad_w = _bwd_common(gradient, None, x_u, x_normed_u, rrms_u, weight_u, amax_state_u, kernel)
+  return (None, None, None, None, grad_x, grad_w, None)
+
+def _fused_add_bwd(*args, **kwargs):
+  # Two invocation modes: 1 grad => positional; >1 grads => kwarg `call=`.
+  # Outputs: (fp8_out, h_out, x_normed_out, rrms_out, amax_buf). Both fp8 and h may be consumed
+  # downstream — TUPLE order in gradient.py preserves kernel-output slot order.
+  # Don't dispatch by dtype: matmul's bwd emits fp8 grad as bf16 (no explicit cast), so
+  # dtype-detection collapses both into h_grad and silently drops the rmsnorm-bwd path.
+  if 'call' in kwargs:
+    kernel, all_grads = kwargs['call'], list(args)
+  else:
+    gradient, kernel = args
+    all_grads = [gradient]
+  fp8_grad_u = h_grad_u = None
+  if len(all_grads) >= 2:
+    fp8_grad_u, h_grad_u = all_grads[0], all_grads[1]
+  elif len(all_grads) == 1:
+    g = all_grads[0]
+    if g.dtype == dtypes.bfloat16: h_grad_u = g
+    else: fp8_grad_u = g
+  _, _, x_normed_u, rrms_u, _, x_u, _, weight_u, amax_state_u = kernel.src[1:]
+  grad_h, grad_w = _bwd_common(fp8_grad_u, h_grad_u, x_u, x_normed_u, rrms_u, weight_u, amax_state_u, kernel)
+  return (None, None, None, None, None, grad_h, grad_h, grad_w, None)
+
+def fused_rmsnorm_mul_quantize_fp8(x:Tensor, weight:Tensor, amax_state:Tensor, eps:float, fp8_dtype) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
+  # NOTE: rmsnorm(x) * weight -> fp8 + amax. Returns (fp8, inv_scale, new_amax, x_normed, rrms).
+  # x_normed + rrms are saved for the rmsnorm backward (also recomputed here from x regs).
+  assert x.dtype == dtypes.bfloat16 and weight.dtype == dtypes.bfloat16
+  assert x.shape[-1] == weight.shape[-1], f"HIDDEN mismatch: x={x.shape}, weight={weight.shape}"
+  MBS, SEQ, HIDDEN = x.shape
+  axis = x.uop.axis if isinstance(x.device, tuple) else None
+  if isinstance(x.device, tuple): assert axis in (None, 0, 1), f"unsupported sharding axis={axis}"
+  fp8_out      = alloc_like((MBS, SEQ, HIDDEN), fp8_dtype,       x.device, axis)
+  x_normed_out = alloc_like((MBS, SEQ, HIDDEN), dtypes.bfloat16, x.device, axis)
+  rrms_out     = alloc_like((MBS, SEQ),         dtypes.float32,  x.device, axis)
+  amax_buf     = alloc_local((NUM_WG,),         dtypes.float32,  x.device, axis)
+  fxn = functools.partial(_custom_fwd, dname=dname_of(x.device), eps_val=eps)
+  fp8_out, x_normed_out, rrms_out, amax_buf, *_ = Tensor.custom_kernel(
+    fp8_out, x_normed_out, rrms_out, amax_buf, x, weight, amax_state, fxn=fxn, grad_fxn=_fused_bwd)
+  inv_scale = (amax_state.float() + 1e-8) / FP8_MAX
+  return fp8_out, inv_scale, scalar_amax(amax_buf), x_normed_out, rrms_out
+
+def fused_add_rmsnorm_mul_quantize_fp8(x:Tensor, residual:Tensor, weight:Tensor, amax_state:Tensor,
+                                       eps:float, fp8_dtype) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:
+  # NOTE: h = x + residual; y_normed = rmsnorm(h); fp8 = quantize(y_normed * weight).
+  # Returns (fp8, inv_scale, new_amax, h, x_normed, rrms). h is also written so downstream can
+  # reuse it without recomputing x+residual — eliminates the separate residual-add kernel.
+  assert x.dtype == dtypes.bfloat16 and residual.dtype == dtypes.bfloat16 and weight.dtype == dtypes.bfloat16
+  assert x.shape == residual.shape
+  MBS, SEQ, HIDDEN = x.shape
+  axis = x.uop.axis if isinstance(x.device, tuple) else None
+  if isinstance(x.device, tuple): assert axis in (None, 0, 1), f"unsupported sharding axis={axis}"
+  fp8_out      = alloc_like((MBS, SEQ, HIDDEN), fp8_dtype,       x.device, axis)
+  h_out        = alloc_like((MBS, SEQ, HIDDEN), dtypes.bfloat16, x.device, axis)
+  x_normed_out = alloc_like((MBS, SEQ, HIDDEN), dtypes.bfloat16, x.device, axis)
+  rrms_out     = alloc_like((MBS, SEQ),         dtypes.float32,  x.device, axis)
+  amax_buf     = alloc_local((NUM_WG,),         dtypes.float32,  x.device, axis)
+  fxn = functools.partial(_custom_fwd_add, dname=dname_of(x.device), eps_val=eps)
+  fp8_out, h_out, x_normed_out, rrms_out, amax_buf, *_ = Tensor.custom_kernel(
+    fp8_out, h_out, x_normed_out, rrms_out, amax_buf, x, residual, weight, amax_state,
+    fxn=fxn, grad_fxn=_fused_add_bwd)
+  inv_scale = (amax_state.float() + 1e-8) / FP8_MAX
+  return fp8_out, inv_scale, scalar_amax(amax_buf), h_out, x_normed_out, rrms_out

extra/llama_kernels/fused_rmsnorm_mul_quantize_fp8/fused_rmsnorm_mul_quantize_fp8.cpp

@@ -0,0 +1,155 @@
+#include <hip/hip_runtime.h>
+#include <hip/hip_bf16.h>
+#include <hip/hip_fp8.h>
+
+// Fuses the full pre-matmul preparation for a layer into a single HBM pass:
+//   y = rmsnorm(x) * weight            (reduce-mean-square + rsqrt + per-elem mul)
+//   fp8 = fp8_sat(y * (FP8_MAX / amax_state))
+// Also writes:
+//   rrms[row]        — saved for the rmsnorm backward
+//   amax_buf[wg]     — per-WG |y| partials, reduced later to update amax_state
+//
+// Layout: one WG per row, ROWS_PER_WG rows per WG via grid-stride (ROWS = N_ELEMS / HIDDEN).
+// Each thread handles HIDDEN / THREADS_PER_WG elements per row.
+
+#ifndef N_ELEMS
+#define N_ELEMS 67108864
+#endif
+#ifndef HIDDEN
+#define HIDDEN 4096
+#endif
+#ifndef NUM_WG
+#define NUM_WG 1024
+#endif
+#ifndef THREADS_PER_WG
+#define THREADS_PER_WG 256
+#endif
+#ifndef EPS_LITERAL
+#define EPS_LITERAL 1e-5f
+#endif
+#ifndef HAS_RESIDUAL
+#define HAS_RESIDUAL 0
+#endif
+
+constexpr int VEC = 8;
+constexpr float FP8_MAX = 448.0f;
+
+static_assert(N_ELEMS % HIDDEN == 0, "N_ELEMS must be a multiple of HIDDEN");
+static_assert(HIDDEN % (THREADS_PER_WG * VEC) == 0, "HIDDEN must be divisible by THREADS_PER_WG*VEC");
+
+constexpr int ROWS = N_ELEMS / HIDDEN;
+constexpr int ELEMS_PER_THREAD = HIDDEN / THREADS_PER_WG;  // each thread sees this many elems per row
+constexpr int VECS_PER_THREAD = ELEMS_PER_THREAD / VEC;    // number of 8-wide vec loads
+
+#if HAS_RESIDUAL
+extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
+fused_add_rmsnorm_mul_quantize_fp8(
+    __hip_fp8_storage_t*  __restrict__ fp8_out,         // fp8, ROWS*HIDDEN
+    __hip_bfloat16*       __restrict__ h_out,           // bf16, ROWS*HIDDEN — x + residual (saved for downstream)
+    __hip_bfloat16*       __restrict__ x_normed_out,    // bf16, ROWS*HIDDEN
+    float*                __restrict__ rrms_out,        // fp32, ROWS
+    float*                __restrict__ amax_buf,        // fp32, NUM_WG
+    const __hip_bfloat16* __restrict__ x,               // bf16, ROWS*HIDDEN
+    const __hip_bfloat16* __restrict__ residual,        // bf16, ROWS*HIDDEN — added into x before rmsnorm
+    const __hip_bfloat16* __restrict__ weight,          // bf16, HIDDEN
+    const float*          __restrict__ amax_state)      // fp32 scalar
+{
+#else
+extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
+fused_rmsnorm_mul_quantize_fp8(
+    __hip_fp8_storage_t*  __restrict__ fp8_out,         // fp8, ROWS*HIDDEN
+    __hip_bfloat16*       __restrict__ x_normed_out,    // bf16, ROWS*HIDDEN (saved for rmsnorm bwd)
+    float*                __restrict__ rrms_out,        // fp32, ROWS (fp32 to match rmsnorm_bwd.cpp expectation)
+    float*                __restrict__ amax_buf,        // fp32, NUM_WG per-WG partials
+    const __hip_bfloat16* __restrict__ x,               // bf16, ROWS*HIDDEN
+    const __hip_bfloat16* __restrict__ weight,          // bf16, HIDDEN (per-hidden scale)
+    const float*          __restrict__ amax_state)      // fp32 scalar
+{
+#endif
+  __shared__ float sdata[THREADS_PER_WG];
+
+  const int tid = threadIdx.x;
+  const int wg  = blockIdx.x;
+
+  const float scale = FP8_MAX / (static_cast<float>(*amax_state) + 1e-8f);
+  const float inv_hidden = 1.0f / static_cast<float>(HIDDEN);
+  float local_max = 0.0f;
+
+  // Grid-stride over rows. Each WG processes rows (wg, wg+NUM_WG, wg+2*NUM_WG, ...).
+  for (int row = wg; row < ROWS; row += NUM_WG) {
+    const int row_off = row * HIDDEN;
+
+    // Load row (+ residual if present) into registers.
+    float regs[ELEMS_PER_THREAD];
+    float sum_sq = 0.0f;
+    #pragma unroll
+    for (int v = 0; v < VECS_PER_THREAD; v++) {
+      const int h_base = tid * VEC + v * THREADS_PER_WG * VEC;
+      float4 raw = *reinterpret_cast<const float4*>(&x[row_off + h_base]);
+      const __hip_bfloat16 *xi = reinterpret_cast<const __hip_bfloat16*>(&raw);
+#if HAS_RESIDUAL
+      float4 res_raw = *reinterpret_cast<const float4*>(&residual[row_off + h_base]);
+      const __hip_bfloat16 *ri = reinterpret_cast<const __hip_bfloat16*>(&res_raw);
+      __hip_bfloat16 h_buf[VEC];
+#endif
+      #pragma unroll
+      for (int i = 0; i < VEC; i++) {
+#if HAS_RESIDUAL
+        const float f = static_cast<float>(xi[i]) + static_cast<float>(ri[i]);
+        h_buf[i] = static_cast<__hip_bfloat16>(f);
+#else
+        const float f = static_cast<float>(xi[i]);
+#endif
+        regs[v * VEC + i] = f;
+        sum_sq += f * f;
+      }
+#if HAS_RESIDUAL
+      *reinterpret_cast<float4*>(&h_out[row_off + h_base]) = *reinterpret_cast<float4*>(h_buf);
+#endif
+    }
+
+    // LDS tree-reduce sum_sq across the WG.
+    sdata[tid] = sum_sq;
+    __syncthreads();
+    for (int s = THREADS_PER_WG / 2; s > 0; s >>= 1) {
+      if (tid < s) sdata[tid] = sdata[tid] + sdata[tid + s];
+      __syncthreads();
+    }
+    const float mean_sq = sdata[0] * inv_hidden;
+    const float rrms = 1.0f / sqrtf(mean_sq + EPS_LITERAL);
+
+    if (tid == 0) rrms_out[row] = rrms;
+
+    // Normalize, multiply by weight, quantize. Also write x_normed (for rmsnorm bwd).
+    #pragma unroll
+    for (int v = 0; v < VECS_PER_THREAD; v++) {
+      const int h_base = tid * VEC + v * THREADS_PER_WG * VEC;
+      float4 w_raw = *reinterpret_cast<const float4*>(&weight[h_base]);
+      const __hip_bfloat16 *wi = reinterpret_cast<const __hip_bfloat16*>(&w_raw);
+
+      __hip_fp8_storage_t out[VEC];
+      __hip_bfloat16 xn[VEC];
+      #pragma unroll
+      for (int i = 0; i < VEC; i++) {
+        const float x_normed = regs[v * VEC + i] * rrms;
+        xn[i] = static_cast<__hip_bfloat16>(x_normed);
+        const float y = x_normed * static_cast<float>(wi[i]);
+        local_max = fmaxf(local_max, fabsf(y));
+        const float scaled = fmaxf(-FP8_MAX, fminf(FP8_MAX, y * scale));
+        out[i] = __hip_cvt_float_to_fp8(scaled, __HIP_SATFINITE, __HIP_E4M3);
+      }
+      *reinterpret_cast<uint64_t*>(&fp8_out[row_off + h_base]) = *reinterpret_cast<uint64_t*>(out);
+      *reinterpret_cast<float4*>(&x_normed_out[row_off + h_base]) = *reinterpret_cast<float4*>(xn);
+    }
+    __syncthreads();  // before next row's sum_sq reduce reuses sdata
+  }
+
+  // Final per-WG amax reduce.
+  sdata[tid] = local_max;
+  __syncthreads();
+  for (int s = THREADS_PER_WG / 2; s > 0; s >>= 1) {
+    if (tid < s) sdata[tid] = fmaxf(sdata[tid], sdata[tid + s]);
+    __syncthreads();
+  }
+  if (tid == 0) amax_buf[wg] = sdata[0];
+}

extra/llama_kernels/fused_rmsnorm_mul_quantize_fp8/fused_rmsnorm_mul_quantize_fp8_bwd.cpp

@@ -0,0 +1,147 @@
+#include <hip/hip_runtime.h>
+#include <hip/hip_bf16.h>
+
+// Full backward for fused_rmsnorm_mul_quantize_fp8.cpp. One HBM pass per row produces:
+//   grad_x (bf16)              — gradient w.r.t. pre-rmsnorm x
+//   grad_weight_partial (fp32) — per-WG partial of the weight gradient, reduced later
+//
+// Input (all read):
+//   grad_fp8 (bf16)     — upstream grad w.r.t. fp8_out (bf16-typed gradient value)
+//   x_normed (bf16)     — saved from the fwd kernel, shape (ROWS, HIDDEN)
+//   rrms (fp32)         — saved rrms per row
+//   weight (bf16)       — per-HIDDEN rmsnorm weight
+//   amax_state (bf16)   — delayed amax used to compute the fp8 scale in fwd
+//
+// Chain: y = x_normed * weight; fp8 = sat(y * scale). Through STE: grad_y = grad_fp8 * scale.
+//        grad_x_normed = grad_y * weight.
+//        grad_weight   = sum_rows(grad_y * x_normed).
+//        grad_x        = rrms * (grad_x_normed - x_normed * mean(grad_x_normed * x_normed, last_dim)).
+
+#ifndef N_ELEMS
+#define N_ELEMS 67108864
+#endif
+#ifndef HIDDEN
+#define HIDDEN 4096
+#endif
+#ifndef NUM_WG
+#define NUM_WG 1024
+#endif
+#ifndef THREADS_PER_WG
+#define THREADS_PER_WG 256
+#endif
+
+constexpr int VEC = 8;
+constexpr float FP8_MAX = 448.0f;
+
+static_assert(N_ELEMS % HIDDEN == 0, "N_ELEMS must be a multiple of HIDDEN");
+static_assert(HIDDEN % (THREADS_PER_WG * VEC) == 0, "HIDDEN must be divisible by THREADS_PER_WG*VEC");
+
+constexpr int ROWS = N_ELEMS / HIDDEN;
+constexpr int ELEMS_PER_THREAD = HIDDEN / THREADS_PER_WG;
+constexpr int VECS_PER_THREAD = ELEMS_PER_THREAD / VEC;
+
+extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
+fused_rmsnorm_mul_quantize_fp8_bwd(
+    __hip_bfloat16*       __restrict__ grad_x,              // out: bf16, ROWS*HIDDEN
+    float*                __restrict__ grad_weight_partial, // out: fp32, NUM_WG*HIDDEN
+    const __hip_bfloat16* __restrict__ grad_fp8,            // in:  bf16, ROWS*HIDDEN (grad of fp8_out)
+    const __hip_bfloat16* __restrict__ x_normed,            // in:  bf16, ROWS*HIDDEN
+    const float*          __restrict__ rrms,                // in:  fp32, ROWS
+    const __hip_bfloat16* __restrict__ weight,              // in:  bf16, HIDDEN
+    const float*          __restrict__ amax_state)          // in:  fp32 scalar
+{
+  __shared__ float sdata[THREADS_PER_WG];
+
+  const int tid = threadIdx.x;
+  const int wg  = blockIdx.x;
+
+  const float scale = FP8_MAX / (static_cast<float>(*amax_state) + 1e-8f);
+  const float inv_hidden = 1.0f / static_cast<float>(HIDDEN);
+
+  // Per-thread accumulator for grad_weight (across all rows this WG touches).
+  float gw_accum[ELEMS_PER_THREAD];
+  #pragma unroll
+  for (int i = 0; i < ELEMS_PER_THREAD; i++) gw_accum[i] = 0.0f;
+
+  // Preload weight into registers (same across rows). Use ELEMS_PER_THREAD entries.
+  float w_regs[ELEMS_PER_THREAD];
+  #pragma unroll
+  for (int v = 0; v < VECS_PER_THREAD; v++) {
+    const int h_base = tid * VEC + v * THREADS_PER_WG * VEC;
+    float4 w_raw = *reinterpret_cast<const float4*>(&weight[h_base]);
+    const __hip_bfloat16 *wi = reinterpret_cast<const __hip_bfloat16*>(&w_raw);
+    #pragma unroll
+    for (int i = 0; i < VEC; i++) w_regs[v * VEC + i] = static_cast<float>(wi[i]);
+  }
+
+  for (int row = wg; row < ROWS; row += NUM_WG) {
+    const int row_off = row * HIDDEN;
+    const float rrms_v = rrms[row];
+
+    // Load grad_fp8 and x_normed rows into registers, compute grad_y and grad_x_normed.
+    float g_y_regs[ELEMS_PER_THREAD];
+    float xn_regs[ELEMS_PER_THREAD];
+    float g_xn_regs[ELEMS_PER_THREAD];  // grad_x_normed
+    float local_dot = 0.0f;  // sum(grad_x_normed * x_normed) for mean
+
+    #pragma unroll
+    for (int v = 0; v < VECS_PER_THREAD; v++) {
+      const int h_base = tid * VEC + v * THREADS_PER_WG * VEC;
+      float4 g_raw  = *reinterpret_cast<const float4*>(&grad_fp8[row_off + h_base]);
+      float4 xn_raw = *reinterpret_cast<const float4*>(&x_normed[row_off + h_base]);
+      const __hip_bfloat16 *gi  = reinterpret_cast<const __hip_bfloat16*>(&g_raw);
+      const __hip_bfloat16 *xni = reinterpret_cast<const __hip_bfloat16*>(&xn_raw);
+
+      #pragma unroll
+      for (int i = 0; i < VEC; i++) {
+        const int idx = v * VEC + i;
+        const float g_y = static_cast<float>(gi[i]) * scale;
+        const float xn  = static_cast<float>(xni[i]);
+        g_y_regs[idx] = g_y;
+        xn_regs[idx]  = xn;
+        g_xn_regs[idx] = g_y * w_regs[idx];           // grad_x_normed = grad_y * weight
+        gw_accum[idx] += g_y * xn;                    // grad_weight contrib
+        local_dot += g_xn_regs[idx] * xn;             // for mean
+      }
+    }
+
+    // LDS reduce local_dot to sdata[0].
+    sdata[tid] = local_dot;
+    __syncthreads();
+    for (int s = THREADS_PER_WG / 2; s > 0; s >>= 1) {
+      if (tid < s) sdata[tid] = sdata[tid] + sdata[tid + s];
+      __syncthreads();
+    }
+    const float mean_term = sdata[0] * inv_hidden;
+
+    // Compute grad_x = rrms * (grad_x_normed - x_normed * mean_term) and write.
+    #pragma unroll
+    for (int v = 0; v < VECS_PER_THREAD; v++) {
+      const int h_base = tid * VEC + v * THREADS_PER_WG * VEC;
+      __hip_bfloat16 out[VEC];
+      #pragma unroll
+      for (int i = 0; i < VEC; i++) {
+        const int idx = v * VEC + i;
+        const float dx = rrms_v * (g_xn_regs[idx] - xn_regs[idx] * mean_term);
+        out[i] = static_cast<__hip_bfloat16>(dx);
+      }
+      *reinterpret_cast<float4*>(&grad_x[row_off + h_base]) = *reinterpret_cast<float4*>(out);
+    }
+    __syncthreads();
+  }
+
+  // Write this WG's grad_weight partial to HBM (fp32, NUM_WG x HIDDEN layout).
+  const int gw_row_off = wg * HIDDEN;
+  #pragma unroll
+  for (int v = 0; v < VECS_PER_THREAD; v++) {
+    const int h_base = tid * VEC + v * THREADS_PER_WG * VEC;
+    // Write 8 fp32 values with two float4 stores.
+    float4 out_lo, out_hi;
+    out_lo.x = gw_accum[v * VEC + 0]; out_lo.y = gw_accum[v * VEC + 1];
+    out_lo.z = gw_accum[v * VEC + 2]; out_lo.w = gw_accum[v * VEC + 3];
+    out_hi.x = gw_accum[v * VEC + 4]; out_hi.y = gw_accum[v * VEC + 5];
+    out_hi.z = gw_accum[v * VEC + 6]; out_hi.w = gw_accum[v * VEC + 7];
+    *reinterpret_cast<float4*>(&grad_weight_partial[gw_row_off + h_base + 0]) = out_lo;
+    *reinterpret_cast<float4*>(&grad_weight_partial[gw_row_off + h_base + 4]) = out_hi;
+  }
+}

extra/llama_kernels/quantize_fp8_delayed/__init__.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+import functools, pathlib
+from tinygrad import Tensor, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.renderer import Estimates
+from extra.llama_kernels import FP8_MAX, NUM_WG, THREADS_PER_WG, alloc_like, alloc_local, scalar_amax, dname_of, compile_hip
+
+@functools.cache
+def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_partial:UOp, x:UOp, amax_state:UOp, dname:str) -> UOp:
+  n_elems = 1
+  for d in x.shape: n_elems *= d
+  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
+  mem = n_elems * 2 + n_elems + 4 + NUM_WG * 4
+  sink = UOp.sink(fp8_out.base, amax_partial.base, x.base, amax_state.base, threads, workgroups,
+                  arg=KernelInfo(f"quantize_fp8_with_amax_{n_elems}", estimates=Estimates(ops=3*n_elems, mem=mem)))
+  src = (pathlib.Path(__file__).parent/"quantize_fp8_with_amax.cpp").read_text()
+  defines = [f"-DN_ELEMS={n_elems}", f"-DNUM_WG={NUM_WG}", f"-DTHREADS_PER_WG={THREADS_PER_WG}"]
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))
+
+@functools.cache
+def _custom_quantize_fp8_scalar(fp8_out:UOp, x:UOp, amax_state:UOp, dname:str) -> UOp:
+  n_elems = 1
+  for d in x.shape: n_elems *= d
+  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
+  mem = n_elems * 2 + n_elems
+  sink = UOp.sink(fp8_out.base, x.base, amax_state.base, threads, workgroups,
+                  arg=KernelInfo(f"quantize_fp8_scalar_{n_elems}", estimates=Estimates(ops=2*n_elems, mem=mem)))
+  src = (pathlib.Path(__file__).parent/"quantize_fp8_scalar.cpp").read_text()
+  defines = [f"-DN_ELEMS={n_elems}", f"-DNUM_WG={NUM_WG}", f"-DTHREADS_PER_WG={THREADS_PER_WG}"]
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))
+
+def _quantize_fp8_delayed_bwd(gradient:UOp, kernel:UOp):
+  # NOTE: STE-equivalent backward — grad_x = grad_fp8 * scale, scale = FP8_MAX / amax_state.
+  # `gradient` is bf16 grad w.r.t. fp8 output (asm_gemm bwd already applied x_scale).
+  _, _, x, amax_state = kernel.src[1:]
+  device = x.device
+  scale = FP8_MAX / (Tensor(amax_state, device=device).float() + 1e-8)
+  grad_x = (Tensor(gradient, device=device).float() * scale).cast(dtypes.bfloat16)
+  return (None, None, grad_x.uop, None)
+
+def quantize_fp8_delayed(x:Tensor, amax_state:Tensor, fp8_dtype=dtypes.fp8e4m3) -> tuple[Tensor, Tensor, Tensor, UOp]:
+  # NOTE: one-pass bf16 -> fp8 quantize with delayed scaling. Returns (fp8, inv_scale, new_amax, store_effect).
+  # Fused kernel reads x once and writes fp8 + per-WG |x| partials (then a small reduce produces scalar new_amax).
+  # store_effect writes new_amax into amax_state's buffer — the caller must thread it into a realized
+  # output via `.after(store_effect)`. Calling `amax_state.assign(new_amax)` inside a grad_fxn does
+  # NOT work because .assign mutates only the temp Tensor's .uop, not the original layer-owned buffer.
+  assert x.dtype == dtypes.bfloat16, f"expected bf16, got {x.dtype}"
+  axis = x.uop.axis if isinstance(x.device, tuple) else None
+  fp8_out      = alloc_like(x.shape,  fp8_dtype,      x.device, axis)
+  amax_partial = alloc_local((NUM_WG,), dtypes.float32, x.device, axis)
+  fxn = functools.partial(_custom_quantize_fp8_with_amax, dname=dname_of(x.device))
+  fp8_out, amax_partial, *_ = Tensor.custom_kernel(fp8_out, amax_partial, x, amax_state,
+                                                    fxn=fxn, grad_fxn=_quantize_fp8_delayed_bwd)
+  new_amax = scalar_amax(amax_partial)
+  inv_scale = (amax_state.float() + 1e-8) / FP8_MAX
+  store_effect = amax_state.uop.store(new_amax.uop)
+  return fp8_out, inv_scale, new_amax, store_effect
+
+def quantize_fp8_scalar(x:Tensor, amax_state:Tensor, fp8_dtype=dtypes.fp8e4m3) -> Tensor:
+  # NOTE: pure one-pass bf16 -> fp8 quantize with delayed scalar scale. No amax computation.
+  axis = x.uop.axis if isinstance(x.device, tuple) else None
+  fp8_out = alloc_like(x.shape, fp8_dtype, x.device, axis)
+  fxn = functools.partial(_custom_quantize_fp8_scalar, dname=dname_of(x.device))
+  fp8_out, *_ = Tensor.custom_kernel(fp8_out, x, amax_state, fxn=fxn)
+  return fp8_out

extra/llama_kernels/quantize_fp8_delayed/quantize_fp8_scalar.cpp

@@ -0,0 +1,48 @@
+#include <hip/hip_runtime.h>
+#include <hip/hip_bf16.h>
+#include <hip/hip_fp8.h>
+
+// Pure one-pass bf16 -> fp8 quantize with delayed scalar scale. No amax computation.
+
+#ifndef N_ELEMS
+#define N_ELEMS 67108864
+#endif
+#ifndef NUM_WG
+#define NUM_WG 1024
+#endif
+#ifndef THREADS_PER_WG
+#define THREADS_PER_WG 256
+#endif
+
+constexpr int VEC = 8;
+constexpr float FP8_MAX = 448.0f;
+
+static_assert(N_ELEMS % VEC == 0, "N_ELEMS must be divisible by VEC");
+
+extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
+quantize_fp8_scalar(
+    __hip_fp8_storage_t*  __restrict__ fp8_out,         // fp8, N_ELEMS
+    const __hip_bfloat16* __restrict__ x,               // bf16, N_ELEMS
+    const float*          __restrict__ amax_state)      // fp32 scalar (delayed)
+{
+  const int tid = threadIdx.x;
+  const int wg  = blockIdx.x;
+  const int gid = wg * THREADS_PER_WG + tid;
+  const int stride_elems = NUM_WG * THREADS_PER_WG * VEC;
+
+  const float scale = FP8_MAX / (static_cast<float>(*amax_state) + 1e-8f);
+
+  for (int base = gid * VEC; base < N_ELEMS; base += stride_elems) {
+    float4 x_raw = *reinterpret_cast<const float4*>(&x[base]);
+    const __hip_bfloat16 *xi = reinterpret_cast<const __hip_bfloat16*>(&x_raw);
+
+    __hip_fp8_storage_t out[VEC];
+    #pragma unroll
+    for (int i = 0; i < VEC; i++) {
+      const float v = static_cast<float>(xi[i]);
+      const float scaled = fmaxf(-FP8_MAX, fminf(FP8_MAX, v * scale));
+      out[i] = __hip_cvt_float_to_fp8(scaled, __HIP_SATFINITE, __HIP_E4M3);
+    }
+    *reinterpret_cast<uint64_t*>(&fp8_out[base]) = *reinterpret_cast<uint64_t*>(out);
+  }
+}

extra/llama_kernels/quantize_fp8_delayed/quantize_fp8_with_amax.cpp

@@ -2,12 +2,13 @@
 #include <hip/hip_bf16.h>
 #include <hip/hip_fp8.h>
 
+// One-pass bf16 -> fp8 quantize using a scalar delayed amax state,
+// AND simultaneously computes per-WG |x| max partials for the next step's amax state.
+// Saves one full HBM pass over the grad tensor vs. doing quantize + separate abs().max().
+
 #ifndef N_ELEMS
 #define N_ELEMS 67108864
 #endif
-#ifndef HIDDEN
-#define HIDDEN 4096
-#endif
 #ifndef NUM_WG
 #define NUM_WG 1024
 #endif
@@ -19,15 +20,13 @@ constexpr int VEC = 8;
 constexpr float FP8_MAX = 448.0f;
 
 static_assert(N_ELEMS % VEC == 0, "N_ELEMS must be divisible by VEC");
-static_assert(HIDDEN % VEC == 0, "HIDDEN must be divisible by VEC");
 
 extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
-fused_mul_quantize_fp8(
-    __hip_fp8_storage_t*  __restrict__ fp8_out,         // fp8, N_ELEMS
-    __hip_bfloat16*       __restrict__ amax_buf,        // bf16, NUM_WG
-    const __hip_bfloat16* __restrict__ x,               // bf16, N_ELEMS
-    const __hip_bfloat16* __restrict__ weight,          // bf16, HIDDEN (per-hidden scale)
-    const __hip_bfloat16* __restrict__ amax_state)      // bf16 scalar
+quantize_fp8_with_amax(
+    __hip_fp8_storage_t*  __restrict__ fp8_out,       // out: fp8, N_ELEMS
+    float*                __restrict__ amax_partial,  // out: fp32, NUM_WG per-WG partials
+    const __hip_bfloat16* __restrict__ x,             // in:  bf16, N_ELEMS
+    const float*          __restrict__ amax_state)    // in:  fp32 scalar (delayed)
 {
   __shared__ float sdata[THREADS_PER_WG];
 
@@ -40,32 +39,25 @@ fused_mul_quantize_fp8(
   float local_max = 0.0f;
 
   for (int base = gid * VEC; base < N_ELEMS; base += stride_elems) {
-    const int h = base % HIDDEN;   // 0..HIDDEN-VEC, 8-aligned (since base is 8-aligned and HIDDEN divides VEC)
     float4 x_raw = *reinterpret_cast<const float4*>(&x[base]);
-    float4 w_raw = *reinterpret_cast<const float4*>(&weight[h]);
-
     const __hip_bfloat16 *xi = reinterpret_cast<const __hip_bfloat16*>(&x_raw);
-    const __hip_bfloat16 *wi = reinterpret_cast<const __hip_bfloat16*>(&w_raw);
 
     __hip_fp8_storage_t out[VEC];
     #pragma unroll
     for (int i = 0; i < VEC; i++) {
-      const float val = static_cast<float>(xi[i]) * static_cast<float>(wi[i]);
-      local_max = fmaxf(local_max, fabsf(val));
-      const float scaled = fmaxf(-FP8_MAX, fminf(FP8_MAX, val * scale));
+      const float v = static_cast<float>(xi[i]);
+      local_max = fmaxf(local_max, fabsf(v));
+      const float scaled = fmaxf(-FP8_MAX, fminf(FP8_MAX, v * scale));
       out[i] = __hip_cvt_float_to_fp8(scaled, __HIP_SATFINITE, __HIP_E4M3);
     }
-
     *reinterpret_cast<uint64_t*>(&fp8_out[base]) = *reinterpret_cast<uint64_t*>(out);
   }
 
-  // LDS tree-reduce per-WG amax
   sdata[tid] = local_max;
   __syncthreads();
   for (int s = THREADS_PER_WG / 2; s > 0; s >>= 1) {
     if (tid < s) sdata[tid] = fmaxf(sdata[tid], sdata[tid + s]);
     __syncthreads();
   }
-
-  if (tid == 0) amax_buf[wg] = static_cast<__hip_bfloat16>(sdata[0]);
+  if (tid == 0) amax_partial[wg] = sdata[0];
 }

extra/nv_gpu_driver/fw.h

@@ -0,0 +1,25 @@
+/* adapted from linux/drivers/gpu/drm/nouveau/include/nvfw/fw.h */
+/* SPDX-License-Identifier: MIT */
+#ifndef __NVFW_FW_H__
+#define __NVFW_FW_H__
+typedef unsigned int u32;
+
+struct nvfw_bin_hdr {
+	u32 bin_magic;
+	u32 bin_ver;
+	u32 bin_size;
+	u32 header_offset;
+	u32 data_offset;
+	u32 data_size;
+};
+
+struct nvfw_bl_desc {
+	u32 start_tag;
+	u32 dmem_load_off;
+	u32 code_off;
+	u32 code_size;
+	u32 data_off;
+	u32 data_size;
+};
+
+#endif

extra/nv_gpu_driver/hs.h

@@ -0,0 +1,52 @@
+/* adapted from linux/drivers/gpu/drm/nouveau/include/nvfw/hs.h */
+/* SPDX-License-Identifier: MIT */
+#ifndef __NVFW_HS_H__
+#define __NVFW_HS_H__
+typedef unsigned int u32;
+
+struct nvfw_hs_header {
+	u32 sig_dbg_offset;
+	u32 sig_dbg_size;
+	u32 sig_prod_offset;
+	u32 sig_prod_size;
+	u32 patch_loc;
+	u32 patch_sig;
+	u32 hdr_offset;
+	u32 hdr_size;
+};
+
+struct nvfw_hs_header_v2 {
+	u32 sig_prod_offset;
+	u32 sig_prod_size;
+	u32 patch_loc;
+	u32 patch_sig;
+	u32 meta_data_offset;
+	u32 meta_data_size;
+	u32 num_sig;
+	u32 header_offset;
+	u32 header_size;
+};
+
+struct nvfw_hs_load_header {
+	u32 non_sec_code_off;
+	u32 non_sec_code_size;
+	u32 data_dma_base;
+	u32 data_size;
+	u32 num_apps;
+	u32 apps[];
+};
+
+struct nvfw_hs_load_header_v2 {
+	u32 os_code_offset;
+	u32 os_code_size;
+	u32 os_data_offset;
+	u32 os_data_size;
+	u32 num_apps;
+	struct {
+		u32 offset;
+		u32 size;
+		u32 data_offset;
+		u32 data_size;
+	} app[];
+};
+#endif

extra/sqtt/roc.py

@@ -1,10 +1,13 @@
 #!/usr/bin/env python3
-import ctypes, pathlib, argparse, pickle, dataclasses, threading
+import ctypes, pathlib, argparse, pickle, dataclasses, threading, itertools
+from decimal import Decimal
 from typing import Generator
 from tinygrad.helpers import temp, unwrap, DEBUG
 from tinygrad.runtime.ops_amd import ProfileSQTTEvent
 from tinygrad.runtime.autogen import rocprof
 from tinygrad.renderer.amd.dsl import Inst
+from tinygrad.helpers import ProfileEvent, ProfileRangeEvent, ProfilePointEvent
+from tinygrad.device import ProfileProgramEvent
 from test.amd.disasm import disasm
 
 @dataclasses.dataclass(frozen=True)
@@ -126,6 +129,71 @@ def decode(sqtt_evs:list[ProfileSQTTEvent], disasms:dict[str, dict[int, Inst]])
     raise exc
   return ROCParseCtx
 
+def unpack_occ(viz_data, i:int, j:int, key:tuple[str, int], data:list, p:ProfileProgramEvent, target:str) -> dict:
+  from tinygrad.viz.serve import amd_decode, create_step, row_tuple
+  steps = viz_data.ctxs[i]["steps"]
+  if len(steps[j+1:]) > 0: return {"steps":[{k:v for k,v in s.items() if k != "data"} for s in steps[j+1:]]}
+  base = unwrap(p.base)
+  disasm:dict[int, Inst] = {addr+base:inst for addr,inst in amd_decode(unwrap(p.lib), target).items()}
+  rctx = decode(data, {p.tag:disasm})
+  cu_events:dict[str, list[ProfileEvent]] = {}
+  # ** inst traces
+  wave_insts:dict[str, dict[str, dict]] = {}
+  inst_units:dict[str, itertools.count] = {}
+  for w in rctx.inst_execs.get(key, []):
+    if (u:=w.wave_loc) not in inst_units: inst_units[u] = itertools.count(0)
+    n = next(inst_units[u])
+    if (events:=cu_events.get(w.cu_loc)) is None: cu_events[w.cu_loc] = events = []
+    events.append(ProfileRangeEvent(f"SIMD:{w.simd}", loc:=f"INST WAVE:{w.wave_id} N:{n}", Decimal(w.begin_time), Decimal(w.end_time)))
+    wave_insts.setdefault(w.cu_loc, {})[f"{u} N:{n}"] = {"wave":w, "disasm":disasm, "prg":p, "run_number":n, "loc":loc}
+  # ** occ traces (only WAVESTART/WAVEEND)
+  units:dict[str, itertools.count] = {}
+  wave_start:dict[str, int] = {}
+  for occ in rctx.occ_events.get(key, []):
+    if (u:=occ.wave_loc) not in units: units[u] = itertools.count(0)
+    if u in inst_units: continue
+    if occ.start: wave_start[u] = occ.time
+    else:
+      if (events:=cu_events.get(occ.cu_loc)) is None: cu_events[occ.cu_loc] = events = []
+      events.append(ProfileRangeEvent(f"SIMD:{occ.simd}", f"OCC WAVE:{occ.wave_id} N:{next(units[u])}", Decimal(wave_start.pop(u)),Decimal(occ.time)))
+  # ** split graph by CU
+  for cu in sorted(cu_events, key=row_tuple):
+    steps.append(create_step(f"{cu} {len(cu_events[cu])}", ("/cu-sqtt", i, len(steps)), depth=1,
+                             data=[ProfilePointEvent(unit, "start", unit, ts=Decimal(0)) for unit in units]+cu_events[cu]))
+    for k in sorted(wave_insts.get(cu, []), key=row_tuple):
+      wd = wave_insts[cu][k]
+      steps.append(create_step(k.replace(cu, ""), ("/amd-sqtt-insts", i, len(steps)), loc=wd["loc"], depth=2,
+                               data={"fxn":unpack_insts, "args":(wd,)}))
+  return {"steps":[{k:v for k,v in s.items() if k != "data"} for s in steps[j+1:]]}
+
+def unpack_insts(viz_data, i:int, j:int, data:dict) -> dict:
+  columns = ["PC", "Instruction", "Hits", "Cycles", "Stall", "Type"]
+  inst_columns = ["N", "Clk", "Idle", "Dur", "Stall"]
+  # Idle:     The total time gap between the completion of previous instruction and the beginning of the current instruction.
+  #           The idle time can be caused by:
+  #             * Arbiter loss
+  #             * Source or destination register dependency
+  #             * Instruction cache miss
+  # Stall:    The total number of cycles the hardware pipe couldn't issue an instruction.
+  # Duration: Total latency in cycles, defined as "Stall time + Issue time" for gfx9 or "Stall time + Execute time" for gfx10+.
+  prev_instr = (w:=data["wave"]).begin_time
+  pc_to_inst = data["disasm"]
+  start_pc = None
+  rows:dict[int, dict] = {}
+  for pc, inst in pc_to_inst.items():
+    if start_pc is None: start_pc = pc
+    rows[pc] = {"pc":pc-start_pc, "inst":str(inst), "hit_count":0, "dur":0, "stall":0, "type":"", "hits":{"cols":inst_columns, "rows":[]}}
+  for e in w.unpack_insts():
+    if not (inst:=rows[e.pc]).get("type"): inst["type"] = str(e.typ).split("_")[-1]
+    inst["hit_count"] += 1
+    inst["dur"] += e.dur
+    inst["stall"] += e.stall
+    inst["hits"]["rows"].append((inst["hit_count"]-1, e.time, max(0, e.time-prev_instr), e.dur, e.stall))
+    prev_instr = max(prev_instr, e.time + e.dur)
+  summary = [{"label":"Total Cycles", "value":w.end_time-w.begin_time}, {"label":"SE", "value":w.se}, {"label":"CU", "value":w.cu},
+             {"label":"SIMD", "value":w.simd}, {"label":"Wave ID", "value":w.wave_id}, {"label":"Run number", "value":data["run_number"]}]
+  return {"rows":[tuple(v.values()) for v in rows.values()], "cols":columns, "metadata":[summary], "ref":viz_data.ref_map.get(data["prg"].name)}
+
 def print_data(data:dict) -> None:
   from tabulate import tabulate
   # plaintext

extra/thunder/amd/fa.py

@@ -55,8 +55,6 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
   assert attn_mask is None, "attn_mask not supported"
   assert is_causal, "only causal attention supported"
 
-  xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
-
   B, N, H, D = xq.shape
   H_KV = xk.shape[2]
   assert D == 128, "only D=128 supported"
@@ -81,7 +79,7 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
 
   attn, l_vec = Tensor.custom_kernel(attn, l_vec, xq, xk, xv, fxn=functools.partial(custom_fa_forward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D), grad_fxn=grad)[:2]
 
-  return attn.transpose(1, 2), attn, l_vec
+  return attn, attn, l_vec
 
 @functools.cache
 def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int):

extra/thunder/amd/gemm_fp8.cpp

@@ -93,7 +93,20 @@ constexpr int NUM_WARPS = 8;
 
 using G = kittens::group<NUM_WARPS>;
 
-__global__ __launch_bounds__(512, 2) void hk_fp8_gemm(bf16 *C_ptr, fp8e4m3 *A_ptr, fp8e4m3 *B_ptr, float *x_scale_ptr, float *w_scale_ptr) {
+// scale_mode: 0=no scale, 1=x only, 2=w only, 3=both
+#ifndef SCALE_MODE
+#define SCALE_MODE 3
+#endif
+
+__global__ __launch_bounds__(512, 2) void hk_fp8_gemm(bf16 *C_ptr, fp8e4m3 *A_ptr, fp8e4m3 *B_ptr
+#if SCALE_MODE == 1
+    , float *x_scale_ptr
+#elif SCALE_MODE == 2
+    , float *w_scale_ptr
+#elif SCALE_MODE == 3
+    , float *x_scale_ptr, float *w_scale_ptr
+#endif
+) {
     constexpr int M = GEMM_M, N = GEMM_N, K = GEMM_K;
 
     kittens::gl<fp8e4m3, 1, 1, M, K> A{A_ptr, nullptr, nullptr, nullptr, nullptr};
@@ -333,11 +346,25 @@ __global__ __launch_bounds__(512, 2) void hk_fp8_gemm(bf16 *C_ptr, fp8e4m3 *A_pt
     }
 
     // apply x_scale * w_scale before bf16 store to prevent overflow
+#if SCALE_MODE == 1
+    float scale = *x_scale_ptr;
+    mul(cA, cA, scale);
+    mul(cB, cB, scale);
+    mul(cC, cC, scale);
+    mul(cD, cD, scale);
+#elif SCALE_MODE == 2
+    float scale = *w_scale_ptr;
+    mul(cA, cA, scale);
+    mul(cB, cB, scale);
+    mul(cC, cC, scale);
+    mul(cD, cD, scale);
+#elif SCALE_MODE == 3
     float scale = *x_scale_ptr * *w_scale_ptr;
     mul(cA, cA, scale);
     mul(cB, cB, scale);
     mul(cC, cC, scale);
     mul(cD, cD, scale);
+#endif
 
     store(C, cA, {0, 0, block_row * WARPS_ROW * 2 + warp_m, block_col * WARPS_COL * 2 + warp_n});
     store(C, cB, {0, 0, block_row * WARPS_ROW * 2 + warp_m, block_col * WARPS_COL * 2 + WARPS_COL + warp_n});

extra/usbgpu/tbgpu/installer/TinyGPUDriverExtension.xcodeproj/project.pbxproj

@@ -359,7 +359,7 @@
 					"$(inherited)",
 					"@executable_path/../Frameworks",
 				);
-				MACOSX_DEPLOYMENT_TARGET = 12.1;
+				MACOSX_DEPLOYMENT_TARGET = 13.0;
 				MARKETING_VERSION = 1.0.0;
 				PRODUCT_BUNDLE_IDENTIFIER = org.tinygrad.tinygpu.installer;
 				PRODUCT_NAME = TinyGPU;
@@ -397,7 +397,7 @@
 					"$(inherited)",
 					"@executable_path/../Frameworks",
 				);
-				MACOSX_DEPLOYMENT_TARGET = 12.1;
+				MACOSX_DEPLOYMENT_TARGET = 13.0;
 				MARKETING_VERSION = 1.0.0;
 				PRODUCT_BUNDLE_IDENTIFIER = org.tinygrad.tinygpu.installer;
 				PRODUCT_NAME = TinyGPU;
@@ -446,7 +446,7 @@
 				CLANG_WARN__DUPLICATE_METHOD_MATCH = YES;
 				COPY_PHASE_STRIP = NO;
 				DEBUG_INFORMATION_FORMAT = dwarf;
-				DRIVERKIT_DEPLOYMENT_TARGET = 21.0;
+				DRIVERKIT_DEPLOYMENT_TARGET = 22.0;
 				ENABLE_STRICT_OBJC_MSGSEND = YES;
 				ENABLE_TESTABILITY = YES;
 				GCC_C_LANGUAGE_STANDARD = gnu11;
@@ -506,7 +506,7 @@
 				CODE_SIGN_IDENTITY = "Apple Development";
 				COPY_PHASE_STRIP = NO;
 				DEBUG_INFORMATION_FORMAT = "dwarf-with-dsym";
-				DRIVERKIT_DEPLOYMENT_TARGET = 21.0;
+				DRIVERKIT_DEPLOYMENT_TARGET = 22.0;
 				ENABLE_NS_ASSERTIONS = NO;
 				ENABLE_STRICT_OBJC_MSGSEND = YES;
 				GCC_C_LANGUAGE_STANDARD = gnu11;
@@ -533,7 +533,7 @@
 				CODE_SIGN_STYLE = Automatic;
 				CURRENT_PROJECT_VERSION = 3;
 				DEVELOPMENT_TEAM = 9YG3G8543N;
-				DRIVERKIT_DEPLOYMENT_TARGET = 21.0;
+				DRIVERKIT_DEPLOYMENT_TARGET = 22.0;
 				ENABLE_USER_SCRIPT_SANDBOXING = YES;
 				EXCLUDED_ARCHS = "";
 				FRAMEWORK_SEARCH_PATHS = (
@@ -566,7 +566,7 @@
 				CURRENT_PROJECT_VERSION = 3;
 				DEVELOPMENT_TEAM = "";
 				"DEVELOPMENT_TEAM[sdk=driverkit*]" = 9YG3G8543N;
-				DRIVERKIT_DEPLOYMENT_TARGET = 21.0;
+				DRIVERKIT_DEPLOYMENT_TARGET = 22.0;
 				ENABLE_USER_SCRIPT_SANDBOXING = YES;
 				EXCLUDED_ARCHS = "";
 				FRAMEWORK_SEARCH_PATHS = (

extra/usbgpu/tbgpu/installer/TinyGPUDriverExtension/TinyGPUDriver.cpp

@@ -188,8 +188,8 @@ kern_return_t TinyGPUDriver::CfgWrite(uint32_t off, uint32_t size, uint32_t val)
 kern_return_t TinyGPUDriver::ResetDevice()
 {
 	if (!ivars->pci) return kIOReturnNotReady;
-	ivars->pci->Reset(kIOPCIDeviceResetTypeFunctionReset);
-	return 0;
+	kern_return_t ret = ivars->pci->Reset(kIOPCIDeviceResetTypeFunctionReset);
+	return ret == kIOReturnSuccess ? ret : ivars->pci->Reset(kIOPCIDeviceResetTypeHotReset);
 }
 
 IOPCIDevice* TinyGPUDriver::GetPCI()

pyproject.toml

@@ -251,8 +251,7 @@ select = [
   "F541",
   "F841",
 ]
-"tinygrad/runtime/autogen/**/*.py" = ["E501", "F401", "E722", "E731", "F821", "A006", "A002", "F811"]
-"tinygrad/runtime/autogen/amd/**/*.py" = ["E501"]
+"tinygrad/runtime/autogen/**/*.py" = ["E501", "F401", "E731", "F821", "A006", "A002", "F811", "F822"]
 "test/amd/**/*.py" = ["F403", "F405"]
 
 [tool.ruff.format]

spec/README.md

@@ -0,0 +1 @@
+Run `./render.sh` whenever you update tinyspec.tex to regenerate tinyspec.pdf.

spec/render.sh

@@ -0,0 +1,10 @@
+#!/bin/bash
+set -e
+
+if ! command -v tectonic &>/dev/null; then
+  echo "tectonic not found, installing..."
+  sudo pacman -S --noconfirm tectonic
+fi
+
+tectonic tinyspec.tex
+echo "done: tinyspec.pdf"

spec/tinyspec.tex

@@ -0,0 +1,450 @@
+\documentclass[10pt,letterpaper]{article}
+
+\usepackage[margin=0.75in]{geometry}
+\usepackage{amsmath,amssymb}
+\usepackage{booktabs}
+\usepackage{array}
+\usepackage[dvipsnames]{xcolor}
+\usepackage{enumitem}
+\usepackage{listings}
+\lstset{language=Python, basicstyle=\ttfamily\small, columns=fullflexible, keepspaces=true}
+
+\newcommand{\op}[1]{\textsc{#1}}
+
+\definecolor{movgreen}{HTML}{2E7D32}
+\definecolor{reducered}{HTML}{C62828}
+\definecolor{elwyellow}{HTML}{F9A825}
+\definecolor{callblue}{HTML}{1565C0}
+\definecolor{assignbrown}{HTML}{795548}
+\definecolor{multipurple}{HTML}{7B1FA2}
+\definecolor{markerorange}{HTML}{E65100}
+% AxisType colors (from tinygrad)
+\definecolor{axblue}{HTML}{1565C0}       % GLOBAL
+\definecolor{axcyan}{HTML}{00838F}       % LOCAL
+\definecolor{axbrcyan}{HTML}{00ACC1}     % WARP
+\definecolor{axbrblue}{HTML}{42A5F5}     % THREAD
+\definecolor{axwhite}{HTML}{616161}      % LOOP (gray on white paper)
+\definecolor{axred}{HTML}{C62828}        % REDUCE
+\definecolor{axbrred}{HTML}{E53935}      % GROUP_REDUCE
+\definecolor{axyellow}{HTML}{F9A825}     % UPCAST
+\definecolor{axmagenta}{HTML}{7B1FA2}    % UNROLL
+
+\title{tinygrad: a single dialect from Tensor programs to Command Buffers}
+\author{tinygrad, Corp. \\ \texttt{[email protected]}}
+\date{}
+
+\begin{document}
+\maketitle
+\thispagestyle{empty}
+
+\section*{UOps}
+
+All nodes in the tinygrad graph are \textbf{UOps}. A UOp is a tuple $(\mathrm{op},\;\mathrm{src},\;\mathrm{arg},\;\mathrm{tag})$ where $\mathrm{op}$ is from the set below, $\mathrm{src}$ is a tuple of input UOps, $\mathrm{arg}$ is op-dependent, and $\mathrm{tag}$ is for temporary processing. The full program is a DAG of UOps. Each UOp has five derived properties --- \textbf{dtype}, \textbf{shape}, \textbf{device}, \textbf{min\_max}, and \textbf{axis} --- determined by the rules at the end of this document.
+
+%% ============================================================
+\subsection*{Source Ops \normalfont\small--- leaf nodes}
+
+\begin{tabular}{@{}l p{3.2cm} p{3.0cm} p{6.2cm}@{}}
+\toprule
+\textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
+\midrule
+\op{Buffer} & () & size, dtype, device, addrspace &
+  Shape $(n \cdot \textit{size},)$ if device is $n$-tuple, else $(\textit{size},)$. \\
+\op{BufferView} & (buf,) & size, dtype, offset &
+  Typed access into a buffer. Zero-copy $(\textit{size},)$ slice at offset; inherits addrspace. \\
+\op{Param} & $(\mathbf{s})$ or $(\mathbf{s}, \text{min}, \text{max})$ & slot, dtype, device? &
+  Placeholder with shape $\mathbf{s}$. Substituted in \op{Function}. \\[4pt]
+\op{Const} & () & value, dtype &
+  A scalar constant with shape $(\ )$. \\
+\op{Vconst} & () & values, dtype &
+  A vector constant with shape $(n,)$. \\
+\bottomrule
+\end{tabular}
+
+\smallskip
+A \op{Buffer}'s \textbf{addrspace} is \texttt{GLOBAL}, \texttt{LOCAL}, or \texttt{REG}.
+
+%% ============================================================
+\subsection*{{\color{movgreen}Movement Ops} \normalfont\small--- no arithmetic, shapes are $(k,)$-shaped UOps with dtype \texttt{index} in src}
+
+\begin{tabular}{@{}l l l l@{}}
+\toprule
+\textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
+\midrule
+\op{Permute} & $(T,)$ & axis order $\pi$ & Reorder axes. $\pi = (1,0)$ is transpose. \\
+\op{Flip} & $(T,)$ & bools $\mathbf{f}$ & Reverse along flagged axes. \\
+\op{Reshape} & $(T, \mathbf{s'})$ & --- & Reinterpret in row-major order. $\prod s_k = \prod s'_k$. \\
+\op{Expand} & $(T, \mathbf{s'})$ & --- & Broadcast size-1 axes. $s_k \in \{1, s'_k\}$. \\
+\op{Pad} & $(T, \mathbf{b}, \mathbf{e})$ & --- & Pad with $0$s: $b_k$ before, $e_k$ after each axis. \\
+\op{Shrink} & $(T, \mathbf{b}, \mathbf{e})$ & --- & Keep $[b_k, e_k)$ per axis. Inverse of \op{Pad}. \\
+\op{Index} & $(T, i_0, i_1, \ldots)$ & --- & Index from left. $()$-shaped $i$ removes dim; $(k,)$-shaped makes it $k$. \\
+\op{Stack} & $(T_0, T_1, \ldots)$ & --- & Join along a newly created leading axis. All shapes must match. \\
+\op{Replicated} & $(T,)$ & axes & Mark $T$ as replicated along axes. Collapse axes to $1$. \\
+\bottomrule
+\end{tabular}
+
+%% ============================================================
+\subsection*{{\color{reducered}Reduce Ops} \normalfont\small--- collapse axes to size $1$}
+
+\begin{tabular}{@{}l l l l@{}}
+\toprule
+\textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
+\midrule
+\op{Reduce} & $(T,)$ & op, axes & Reduce $T$ along axes. Op is \op{Add}, \op{Max}, or \op{Mul}. \\
+\bottomrule
+\end{tabular}
+
+%% ============================================================
+\subsection*{{\color{callblue}Call Ops} \normalfont\small--- function abstraction, like the lambda calculus}
+
+\begin{tabular}{@{}l l l l@{}}
+\toprule
+\textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
+\midrule
+\op{Function} & (body, $a_0$, $a_1$, \ldots) & --- & Substitute each \op{Param} $k$ in \op{Tuple} body with $a_k$. Gradient-able. \\
+\op{Call} & (body, $a_0$, $a_1$, \ldots) & --- & Opaque invocation of a compiled kernel or custom function. \\
+\op{Tuple} & $(v_0, v_1, \ldots)$ & --- & Pack values; required as \op{Function} body to return a value. \\
+\op{GetTuple} & $(T,)$ & idx & Extract element at idx from a \op{Tuple}. \\
+\bottomrule
+\end{tabular}
+
+%% ============================================================
+\subsection*{{\color{multipurple}Store Ops} \normalfont\small--- side effects}
+
+\begin{tabular}{@{}l l l l@{}}
+\toprule
+\textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
+\midrule
+\op{Store} & (buf, val, gate?) & --- & Write val into buf. buf.shape $=$ val.shape. \\
+  & & & If gate is present, write only when gate is true. Output is void. \\
+\bottomrule
+\end{tabular}
+
+%% ============================================================
+\subsection*{{\color{assignbrown}Ordering Ops} \normalfont\small--- execution order}
+
+\begin{tabular}{@{}l l l p{6.0cm}@{}}
+\toprule
+\textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
+\midrule
+\op{Range} & $(\text{bound},)$ & type & Iterator from $0$ to bound. \\
+\op{End} & (body, range) & --- & Close a \op{Range} loop. \\
+\op{After} & (buf, deps\ldots) & --- & Passthrough of buf; guarantees deps execute first. \\
+\op{Group} & $(u_0, u_1, \ldots)$ & --- & Void no-op that merges multiple \op{Store}s into one node, unordered. \\
+\op{Sink} & $(s_0, s_1, \ldots)$ & --- & Collect side effects into a single root node. \\
+\op{Linear} & (uops\ldots) & --- & Linearized (toposorted) instruction sequence. \\
+\bottomrule
+\end{tabular}
+
+\smallskip
+Assign is \op{Store} followed by \op{After}: write the value, then return the buffer with an ordering dependency.
+
+%% ============================================================
+\subsection*{{\color{elwyellow}Elementwise Ops} \normalfont\small--- all inputs same shape, output same shape, applied per-element}
+
+\begin{tabular}{@{}l l l l@{}}
+\toprule
+\textbf{Arity} & \textbf{src} & \textbf{Op} & \textbf{Semantics} \\
+\midrule
+Unary & $(T,)$
+  & \op{Recip}
+  & $1/x$ \\
+& & \op{Trunc}
+  & $\mathrm{trunc}(x)$: round toward zero. \\
+& & \op{Cast}
+  & Convert to target dtype (specified in arg). \\
+& & \op{Bitcast}
+  & Reinterpret bits as target dtype. Must be same size. \\[4pt]
+Binary & $(A, B)$
+  & \op{Add}, \op{Mul}, \op{Max}, \op{Mod}, \op{Idiv}
+  & $a+b$, $a \cdot b$, $\max(a,b)$, $a \bmod b$, $\lfloor a/b \rfloor$ \\
+& & \op{CmpLt}, \op{CmpNe}
+  & $[a < b]$, $[a \ne b]$ \\
+& & \op{Xor}, \op{Or}, \op{And}, \op{Shr}, \op{Shl}
+  & $a \oplus b$, $a \mid b$, $a \mathbin{\&} b$, $a \gg b$, $a \ll b$ \\[4pt]
+Ternary & $(P, A, B)$
+  & \op{Where}
+  & $A[\mathbf{i}]$ if $P[\mathbf{i}] \ne 0$, else $B[\mathbf{i}]$ \\
+
+\bottomrule
+\end{tabular}
+
+\medskip
+\textbf{Decomposed elementwise ops} --- defined in terms of the primitives above.
+
+\smallskip
+\begin{tabular}{@{}l l l@{}}
+\toprule
+\textbf{Op} & \textbf{Decomposition} & \textbf{Semantics} \\
+\midrule
+\op{Neg} & \op{Mul}($A$, $-1$) & $-x$ \\
+\op{Sub} & \op{Add}($A$, \op{Neg}($B$)) & $a - b$ \\
+\op{Div} & \op{Mul}($A$, \op{Recip}($B$)) & $a / b$ \\
+\op{CmpGt} & \op{CmpLt}($B$, $A$) & $[a > b]$ \\
+\op{CmpGe} & \op{CmpNe}(\op{CmpLt}($A$, $B$),\, $1$) & $[a \ge b]$ \\
+\op{CmpLe} & \op{CmpNe}(\op{CmpLt}($B$, $A$),\, $1$) & $[a \le b]$ \\
+\op{CmpEq} & \op{CmpNe}(\op{CmpNe}($A$, $B$),\, $1$) & $[a = b]$ \\
+\op{Not} & \op{CmpNe}($A$, $1$) & $\lnot a$ \\[4pt]
+\op{Exp2} & polynomial approx + \op{Mul}, \op{Add} & $2^x$ \\
+\op{Log2} & exponent extract + polynomial approx & $\log_2 x$ \\
+\op{Sin} & argument reduction + polynomial approx & $\sin x$ \\
+\op{Sqrt} & \op{Exp2}($0.5 \cdot$ \op{Log2}($A$)) & $\sqrt{x}$ \\
+\op{Pow} & \op{Exp2}(\op{Log2}($A$) $\cdot\, B$) & $a^b$ \\
+\op{Mulacc} & \op{Add}(\op{Mul}($A$, $B$),\, $C$) & $a \cdot b + c$ \\
+\op{Threefry} & 5 rounds of add-rotate-xor (ARX) & Threefry 2x32 PRNG \\
+\bottomrule
+\end{tabular}
+
+%% ============================================================
+\subsection*{{\color{markerorange}Marker Ops} \normalfont\small--- identity on data}
+
+\begin{tabular}{@{}l l l l@{}}
+\toprule
+\textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
+\midrule
+\op{Contiguous} & $(T,)$ & --- & Force contiguous memory layout. \\
+\op{ContiguousBackward} & $(T,)$ & --- & Force contiguous in backward pass. \\
+\op{Detach} & $(T,)$ & --- & Stops gradient propagation. \\
+\op{Copy} & $(T,)$ & device & Copy to target device. \\
+\bottomrule
+\end{tabular}
+
+%% ============================================================
+\subsection*{Codegen Ops \normalfont\small--- generated code primitives, these do not appear in the main graph}
+
+\begin{tabular}{@{}l l l l@{}}
+\toprule
+\textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
+\midrule
+\op{Load} & (idx,alt?,gate?) & --- & Dereference: read element at index from buffer. \\
+ & & & All loads will be replaced by \op{Store}. \\
+\op{Barrier} & (deps\ldots) & --- & Synchronize threads within a workgroup. \\
+\op{Ins} & \ldots & \ldots & A single machine instruction (e.g.\ AMD ISA). \\
+\op{Special} & (bound,) & name & GPU thread/workgroup index (e.g.\ \texttt{gidx0}, \texttt{lidx1}). \\
+\op{If} & (gate,) & --- & Begin conditional execution block. \\
+\op{Endif} & (if,) & --- & End conditional execution block. \\
+\op{Wmma} & (A, B, acc) & config & Warp matrix multiply-accumulate (tensor cores). \\
+\op{Custom} & (args\ldots) & fmt & Inject custom code string into generated source. \\
+\op{AtomicAdd} & (idx, val) & --- & Atomic read-modify-write: \texttt{buf[idx] += val}. \\[4pt]
+\op{CustomFunction} & (meta\ldots) & name & Opaque device function (e.g.\ HW decode). Via \op{Call}. \\
+\op{Program} & (linear, source, binary) & --- & Compiled kernel: instructions, source, and machine code. \\
+\op{Source} & () & str & Human-readable rendered source code. \\
+\op{Binary} & () & bytes & Compiled machine code. \\
+\bottomrule
+\end{tabular}
+
+\smallskip
+These ops are not part of the core specification and are subject to change.
+
+%% ============================================================
+\subsection*{Derived Properties}
+
+Every UOp has a \textbf{dtype}, \textbf{shape}, \textbf{device}, \textbf{min\_max}, and \textbf{axis}, derived from its op, src, and arg:
+
+\medskip
+\begin{tabular}{@{}l l l l l@{}}
+\toprule
+\textbf{Op} & \textbf{dtype} & \textbf{shape} & \textbf{device} & \textbf{min\_max} \\
+\midrule
+\op{Buffer} & from arg & $(\text{size},)$ from arg & from arg & dtype range \\
+\op{Const} & from arg & $()$ & \textsc{null} & $[v, v]$ \\
+\op{Param} & from arg & from $\mathrm{src}[0]$ & from arg & from src or dtype range \\[3pt]
+Movement ops & $\mathrm{src}[0].\mathrm{dtype}$ & (see op) & $\mathrm{src}[0].\mathrm{device}$ & $\mathrm{src}[0]$ \\
+\op{Reduce} & $\mathrm{src}[0].\mathrm{dtype}$ & collapse axes to $1$ & $\mathrm{src}[0].\mathrm{device}$ & dtype range \\[3pt]
+\op{Cast} & from arg & $\mathrm{src}[0].\mathrm{shape}$ & $\mathrm{src}[0].\mathrm{device}$ & clamped to dtype \\
+\op{Bitcast} & from arg & $\mathrm{src}[0].\mathrm{shape}$ & $\mathrm{src}[0].\mathrm{device}$ & dtype range \\
+\op{Copy} & $\mathrm{src}[0].\mathrm{dtype}$ & $\mathrm{src}[0].\mathrm{shape}$ & from arg & $\mathrm{src}[0]$ \\
+ALU unary & $\mathrm{src}[0].\mathrm{dtype}$ & $\mathrm{src}[0].\mathrm{shape}$ & $\mathrm{src}[0].\mathrm{device}$ & dtype range \\
+\op{Add} & $\mathrm{src}[0].\mathrm{dtype}$ & broadcast & $\mathrm{src}[0].\mathrm{device}$ & $[a+b,\, A+B]$ \\
+\op{Mul} & $\mathrm{src}[0].\mathrm{dtype}$ & broadcast & $\mathrm{src}[0].\mathrm{device}$ & $[\min,\max]$ of products \\
+\op{Max} & $\mathrm{src}[0].\mathrm{dtype}$ & broadcast & $\mathrm{src}[0].\mathrm{device}$ & $[\max(a,b),\, \max(A,B)]$ \\
+Other binary & $\mathrm{src}[0].\mathrm{dtype}$ & broadcast & $\mathrm{src}[0].\mathrm{device}$ & dtype range \\
+\op{CmpLt}, \op{CmpNe} & bool & broadcast & $\mathrm{src}[0].\mathrm{device}$ & from intervals \\
+\op{Where} & $\mathrm{src}[1].\mathrm{dtype}$ & broadcast & $\mathrm{src}[0].\mathrm{device}$ & $[\min(b,c),\, \max(B,C)]$ \\[3pt]
+\op{Function}, \op{Call} & $\mathrm{src}[0].\mathrm{dtype}$ & substitute \op{Param} shapes & $\mathrm{src}[1].\mathrm{device}$ & dtype range \\
+\op{Range} & index & $()$ & \textsc{null} & $[0,\, n{-}1]$ \\
+\op{Index} & $\mathrm{src}[0].\mathrm{dtype}$ & remaining dims & $\mathrm{src}[0].\mathrm{device}$ & $\mathrm{src}[0]$ \\
+\op{Store} & void & $()$ & $\mathrm{src}[0].\mathrm{device}$ & --- \\
+\op{After} & $\mathrm{src}[0].\mathrm{dtype}$ & $\mathrm{src}[0].\mathrm{shape}$ & $\mathrm{src}[0].\mathrm{device}$ & $\mathrm{src}[0]$ \\
+\bottomrule
+\end{tabular}
+
+\smallskip
+$\mathrm{broadcast}$: right-align shapes, element-wise max; each axis must be equal or $1$.
+$[a,A]$, $[b,B]$, $[c,C]$ denote min\_max of $\mathrm{src}[0]$, $\mathrm{src}[1]$, $\mathrm{src}[2]$.
+Default \emph{dtype range}: $[\mathrm{dtype\_min},\, \mathrm{dtype\_max}]$.
+
+\medskip
+\textbf{axis} tracks the multi-device sharding dimension. \op{Buffer} with $n$-tuple device: axis $= 0$ (device dim).
+\op{Reshape} remaps axis to preserve the shard boundary. \op{Permute} follows the permutation.
+\op{Reduce} on the shard axis $\to$ \textsc{null}. \op{Replicated} on the shard axis $\to$ \textsc{null}. \op{Copy} $\to$ \textsc{null}. ALU ops inherit from sources. Default: \textsc{null}.
+
+%% ============================================================
+\subsection*{Kernel Optimizations (OptOps) \normalfont\small--- schedule-level transforms on kernel ranges}
+
+Each kernel's iteration space is a set of \op{Range} axes. Every range has an \textbf{AxisType}:
+
+\medskip
+\begin{tabular}{@{}l l l l l@{}}
+\toprule
+\textbf{AxisType} & \textbf{Letter} & \textbf{Split from} & \textbf{Direction} & \textbf{Semantics} \\
+\midrule
+{\color{axblue}\texttt{GLOBAL}} & \texttt{g} & --- & --- & GPU global workgroup dimension. \\
+{\color{axcyan}\texttt{LOCAL}} & \texttt{l} & g, L & inner & Workgroup local dimension (shared memory). \\
+{\color{axbrcyan}\texttt{WARP}} & \texttt{w} & \multicolumn{2}{l}{(created by \op{TC})} & Warp-level lanes for tensor cores. \\
+{\color{axbrblue}\texttt{THREAD}} & \texttt{t} & g & outer & CPU thread parallelism. \\
+{\color{axwhite}\texttt{LOOP}} & \texttt{L} & --- & --- & Generic sequential loop (initial state). \\
+{\color{axred}\texttt{REDUCE}} & \texttt{R} & --- & --- & Reduction axis. \\
+{\color{axbrred}\texttt{GROUP\_REDUCE}} & \texttt{G} & R & inner/outer & Shared-memory group reduction. \\
+{\color{axyellow}\texttt{UPCAST}} & \texttt{u} & g, l, L & inner & Register-level vectorization. \\
+{\color{axmagenta}\texttt{UNROLL}} & \texttt{r} & R, G & inner & Fully unrolled loop. \\
+\bottomrule
+\end{tabular}
+
+\medskip
+An optimization is a triple $(\mathrm{op},\;\mathrm{axis},\;\mathrm{arg})$:
+
+\smallskip
+\begin{tabular}{@{}l l l p{6.5cm}@{}}
+\toprule
+\textbf{OptOp} & \textbf{axis} & \textbf{arg} & \textbf{Semantics} \\
+\midrule
+\op{Split} & any & (factor $k$, target, top?) &
+  Split axis $n$ by $k$ into $(n/k, k)$ or $(k, n/k)$ if top. New sub-axis gets target AxisType (see table above). \\
+\op{Padto} & any & multiple $m$ &
+  Pad axis to next multiple of $m$ with validity masks. \\[4pt]
+\op{Swap} & axis$_i$ & axis$_j$ &
+  Swap two axes $i \leftrightarrow j$. \\
+\op{Nolocals} & --- & --- &
+  Disable local memory; no workgroup dims emitted. \\
+\op{TC} & reduce idx & (tc, opt, mode) &
+  Apply tensor core \op{Wmma}: split reduce/output axes into \texttt{WARP}, \texttt{UPCAST}, and \texttt{UNROLL} dims. \\
+\bottomrule
+\end{tabular}
+
+\smallskip
+Optimizations compose left-to-right. \op{TC} must be first. The search space is explored by BEAM search or hand-coded heuristics.
+
+%% ============================================================
+\subsection*{Common Ops as Compositions}
+
+All high-level tensor operations decompose into the primitives above.
+
+\begin{lstlisting}
+# gemm: C[M,N] = A[M,K] @ B[K,N]
+def gemm(A, B):
+  M,K = A.shape; _,N = B.shape
+  return (A.reshape(M,K,1) * B.reshape(1,K,N)).sum(1)
+
+# prefix_sum: cumulative sum via repeat+reshape sliding window trick
+def prefix_sum(T):
+  n = T.shape[0]
+  x = T.pad((n-1, 0))                            # (2n-1,)
+  x = x.reshape(1,2*n-1).expand(n+1,2*n-1)       # tile
+  x = x.reshape((n+1)*(2*n-1)).shrink_to(2*n*n)  # trim
+  x = x.reshape(n,2*n).shrink_to(n,n)            # windows
+  return x.sum(-1)                               # reduce
+
+# arange: prefix_sum of all 1s gives [1,2,...,n], subtract 1 for [0,1,...,n-1]
+def arange(n):
+  return prefix_sum(Tensor(1).reshape(1).expand(n)) - 1
+
+# gather: out[i] = T[idx[i]]. one-hot mask along gather axis, then reduce
+def gather(T, idx):
+  K = T.shape[0]
+  pos = arange(K).reshape(K, 1)                     # (K, 1)
+  mask = (pos == idx.reshape(1, -1)).cast(T.dtype)  # (K, D)
+  return (T.reshape(K, 1) * mask).sum(0)            # (D,)
+
+# scatter_add: T[idx[i]] += val[i]
+def scatter_add(T, idx, val):
+  K, D = T.shape[0], idx.shape[0]
+  pos = arange(K).reshape(K, 1)                     # (K, 1)
+  mask = (pos == idx.reshape(1, D)).cast(T.dtype)   # (K, D)
+  return T + (mask * val.reshape(1, D)).sum(1)      # (K,)
+\end{lstlisting}
+
+%% ============================================================
+\subsection*{{\color{multipurple}Multi-Device Collectives} \normalfont\small--- derived from primitives}
+
+Let $D = (d_0, \ldots, d_{n-1})$ be an $n$-tuple device.
+\op{Copy} to an $n$-tuple device reshards with axis $= 0$. \op{Copy} never changes shape.
+
+\begin{lstlisting}
+# T has shape (s,) on a single device.
+
+# broadcast: replicate T to all n devices
+def broadcast(T):
+  return T.reshape(1, s).expand(n, s).copy(D).replicated(0)        # (s,) on D, axis=null
+
+# scatter: split T into n chunks, one per device
+def scatter(T):
+  return T.copy(D)                                                  # (s,) on D, axis=0
+
+# T has shape (n*s,) on D with axis=0, so each device holds (s,) elements.
+
+# gather: collect all shards onto one device
+def gather(T):
+  return T.copy(D[0])                                               # (n*s,) on D[0], axis=null
+
+# reduce: gather + sum
+def reduce(T):
+  return gather(T).reshape(n, s).sum(0)                             # (s,) on D[0], axis=null
+
+# allgather: collect all shards, replicate to all devices
+def allgather(T):
+  return T.reshape(1, n*s).expand(n, n*s).copy(D).replicated(0)    # (n*s,) on D, axis=null
+
+# reduce_scatter: reduce across devices, scatter result
+def reduce_scatter(T):
+  return T.reshape(n, n, s//n).permute(1, 0, 2).copy(D).sum(1).reshape(s)  # (s,) on D, axis=0
+
+# allreduce: reduce_scatter + allgather
+def allreduce(T):
+  return allgather(reduce_scatter(T))                                       # (s,) on D, axis=null
+\end{lstlisting}
+
+%% ============================================================
+\subsection*{{\color{callblue}The \texttt{@function} Decorator} \normalfont\small--- graph capture via tracing}
+
+The \texttt{@function} decorator transforms a Python function on Tensors into a single \op{Function} node.
+
+\begin{lstlisting}
+@function
+def f(a: Tensor, b: Tensor) -> Tensor:
+  return a + b
+\end{lstlisting}
+
+When \texttt{f(x, y)} is called, the decorator:
+
+\begin{enumerate}[leftmargin=1.5em, itemsep=2pt]
+  \item \textbf{Extracts inputs}: walks all arguments to find every Tensor, deduplicates by identity.
+  \item \textbf{Runs the function} lazily (no device execution), building a UOp graph from the result.
+  \item \textbf{Parameterizes}: replaces each input UOp with a \op{Param}$(k)$ placeholder.
+  \item \textbf{Wraps the body} in a \op{Tuple} (even for single returns) and creates\\
+        \op{Function}(\op{Tuple}(body), $x$, $y$).
+  \item \textbf{Returns} the result via \op{GetTuple}$(0)$, or one \op{GetTuple} per element for tuple returns.
+\end{enumerate}
+
+The result is a reusable graph fragment: the body contains only \op{Param} references, not concrete buffers. At schedule time, the \op{Function} is resolved by substituting each \op{Param}$(k)$ back with its corresponding argument $a_k$, or lowered into an opaque \op{Call} if it is to be compiled as a reusable kernel.
+
+%% ============================================================
+\subsection*{Lowering Pipeline \normalfont\small--- from Tensor graph to machine code}
+
+\begin{tabular}{@{}l p{9.7cm}@{}}
+\toprule
+\textbf{Stage} & \textbf{Semantics} \\
+\midrule
+\textbf{Callify} & Transform the Tensor graph into a single stateless function. \\
+\textbf{Rangeify} & Determine the kernel split of the function. Break everything down to shape () \\
+\textbf{Optimize} & Insert local buffers. Swap and split ranges, and determine which axes are parallel and which are serial. \\
+\textbf{Expand} & Expand the parallel ranges into shape. \\
+\textbf{Instruction Selection} & Select target instructions, including WMMA and devectorization. \\
+\textbf{Linearize} & Topologically sort the graph and determine execution order. \\
+\textbf{Register/Memory Plan} & Allocate and reuse \texttt{GLOBAL}, \texttt{LOCAL}, and \texttt{REG} storage for values with non-overlapping lifetimes. \\
+\textbf{Render} & Output the machine code. \\
+\bottomrule
+\end{tabular}
+
+
+\end{document}

test/amd/test_custom_kernel.py

@@ -3,7 +3,7 @@ import functools
 import numpy as np
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
-from tinygrad.engine.realize import run_linear, estimate_uop
+from tinygrad.engine.realize import run_linear, estimate_uop, compile_linear
 from tinygrad.renderer import Estimates
 from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import getenv
@@ -99,13 +99,13 @@ def custom_lds_sync(A:UOp, arch:str) -> UOp:
   sink = UOp.sink(A.base, lds, threads, wg, arg=KernelInfo("custom_lds_sync"))
   return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
 
-def custom_handwritten(A:UOp, arch:str) -> UOp:
+def custom_handwritten(A:UOp) -> UOp:
   A = A.flatten()
   threads = UOp.special(128, "lidx0")
   wg = UOp.special(1, "gidx0")
   lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=512, addrspace=AddrSpace.LOCAL), (), 'lds')  # 128 * 4 bytes
   pipes = {getenv("PIPE", "")} if getenv("PIPE", "") else {"SALU", "VALU", "TRANSCENDENTAL", "WMMA"}
-  k = Kernel(arch)
+  k = Kernel()
   # wrap in loop to filter out icache misses
   LOOP_N, UNROLL_N = 8, 5
   k.emit(r4.s_mov_b32(s[1], LOOP_N))
@@ -145,10 +145,10 @@ def custom_handwritten(A:UOp, arch:str) -> UOp:
   sink = UOp.sink(A.base, threads, wg, lds, arg=KernelInfo("custom_handwritten"))
   return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
 
-def custom_data_deps(A:UOp, arch:str) -> UOp:
+def custom_data_deps(A:UOp) -> UOp:
   A = A.flatten()
   threads = UOp.special(A.numel(), "lidx0")
-  k = Kernel(arch)
+  k = Kernel()
   k.emit(s_load_b64(s[0:1], s[0:1], soffset=NULL))
   k.emit(s_waitcnt_lgkmcnt(sdst=NULL, simm16=0))
   k.emit(v_lshlrev_b32_e32(v[0], 2, v[0]))
@@ -169,7 +169,7 @@ class TestCustomKernel(unittest.TestCase):
     if self.arch != "rdna3": self.skipTest("only rdna3")
     a = Tensor.full((16, 16), 1.).contiguous().realize()
     a = Tensor.custom_kernel(a, fxn=custom_add_one)[0]
-    linear = a.schedule_linear()
+    linear = compile_linear(a.schedule_linear())
     est = estimate_uop(linear.src[-1])
     self.assertEqual(est.ops, a.numel())
     self.assertEqual(est.mem, a.nbytes()*2)
@@ -198,13 +198,13 @@ class TestCustomKernel(unittest.TestCase):
   def test_handwritten(self):
     if self.arch != "rdna4": self.skipTest("only tested on rdna4")
     a = Tensor.empty(1024, dtype=dtypes.int32).contiguous().realize()
-    a = Tensor.custom_kernel(a, fxn=functools.partial(custom_handwritten, arch=self.arch))[0]
+    a = Tensor.custom_kernel(a, fxn=custom_handwritten)[0]
     a.realize()
 
   def test_data_deps(self):
     if self.arch != "rdna3": self.skipTest("only tested on rdna3")
     a = Tensor(np.full(32, 5.0, dtype=np.float32)).realize()
-    a = Tensor.custom_kernel(a, fxn=functools.partial(custom_data_deps, arch=self.arch))[0]
+    a = Tensor.custom_kernel(a, fxn=custom_data_deps)[0]
     a.realize()
     self.assertTrue((a.numpy() == 6.0).all())
 

test/amd/test_mockgpu_invalid.py

@@ -8,17 +8,14 @@ class TestMockGPUInvalidInstruction(unittest.TestCase):
     test_code = '''
 import struct
 from tinygrad import Device, Tensor
-from tinygrad.engine.realize import get_runner
+from tinygrad.engine.realize import compile_linear
 from tinygrad.runtime.ops_amd import AMDProgram
 
 dev = Device["AMD"]
 a = Tensor([1.0]).realize()
 b = a + 1
-si = b.schedule_linear().src[-1]
-runner = get_runner(dev.device, si.src[0])
-
-prg = runner._prg
-lib = bytearray(prg.lib)
+linear = compile_linear(b.schedule_linear())
+lib = bytearray(linear.src[-1].src[0].src[4].arg)
 
 # Find s_endpgm (0xBFB00000) and replace with V_MOVRELD_B32 (op=66) which has no pcode
 # VOP1 encoding: bits[31:25]=0x7E, op=bits[16:9], so op=66 -> 66<<9 = 0x8400

test/amd/test_sqttmap.py

@@ -1,5 +1,5 @@
 # test to compare every packet with the rocprof decoder
-import unittest, pickle, functools
+import unittest, pickle, functools, json
 from typing import Iterator
 from pathlib import Path
 from tinygrad.helpers import DEBUG, getenv, temp, ansistrip, Context
@@ -156,7 +156,7 @@ class TestSQTTMapRDNA4(TestSQTTMapBase):
     row_counts:dict[str, int] = {}
     for e in sqtt_timeline(events[1].blob, lib, target):
       if type(e).__name__ != "ProfileRangeEvent": continue
-      info = e.name.ret or ""
+      info = json.loads(e.name.ret) if e.name.ret else {}
       if e.device.startswith("WAVE"):
         idx = row_counts.get(e.device, 0)
         dispatch_st[f"{e.device}-{idx}"] = int(e.st)

test/backend/test_arange.py

@@ -2,14 +2,14 @@ import unittest
 import numpy as np
 from tinygrad import Tensor, GlobalCounters, dtypes, nn, Device, Variable
 from tinygrad.helpers import Context, getenv, DEV
-from tinygrad.engine.realize import run_linear, estimate_uop
+from tinygrad.engine.realize import run_linear, estimate_uop, compile_linear
 from tinygrad.renderer.ptx import PTXRenderer
 from test.helpers import needs_second_gpu
 
 class TestArange(unittest.TestCase):
   def _get_flops(self, tensor, desired):
     GlobalCounters.reset()
-    linear = tensor.schedule_linear()
+    linear = compile_linear(tensor.schedule_linear())
     self.assertEqual(len(linear.src), 1)
     run_linear(linear)
     np.testing.assert_equal(tensor.numpy(), desired)
@@ -36,7 +36,7 @@ class TestArange(unittest.TestCase):
   def test_tri_complexity(self):
     with Context(NOOPT=1):
       t = Tensor.ones(256, 256).contiguous().realize()
-      linear = t.triu().schedule_linear()
+      linear = compile_linear(t.triu().schedule_linear())
       self.assertLessEqual(estimate_uop(linear.src[-1]).ops, 4 * 256 * 256)
 
 DSET, DDIM = 2048, 32
@@ -229,7 +229,7 @@ class TestIndexing(unittest.TestCase):
     xq = xq.reshape(bs, seqlen, n_heads, head_dim)
     xq_rope, _ = apply_rotary_emb(xq, xq, freqs_cis)
     xq_rope.sum().backward()
-    linear = wq.grad.schedule_linear()
+    linear = compile_linear(wq.grad.schedule_linear())
     assert len(linear.src) == 1, f"expected one kernel for backward, got: {len(linear.src)}"
     bwd_ops = estimate_uop(linear.src[0]).ops
     # bfloat16 on non CDNA4 has ~10x ops overhead because of the software emulation

test/backend/test_custom_kernel.py

@@ -1,5 +1,5 @@
 import unittest
-from tinygrad import Tensor, UOp
+from tinygrad import Tensor, UOp, GlobalCounters
 from tinygrad.dtype import AddrSpace, dtypes
 from tinygrad.uop.ops import KernelInfo, AxisType
 
@@ -308,6 +308,36 @@ class TestCustomKernel(unittest.TestCase):
     expected = (3+2)*2+2
     assert all(x == expected for x in result), f"expected all {expected}, got {result}"
 
+  def test_custom_kernel_sched(self, use_custom=False):
+    x = Tensor.arange(32).reshape(8, 4).realize()
+    y = Tensor.empty_like(x)
+    y = Tensor.custom_kernel(y, x, fxn=custom_add_one_kernel)[0]
+    if use_custom:
+      z = Tensor.empty_like(x)
+      z = Tensor.custom_kernel(y, y.T.T, fxn=custom_add_one_kernel)[0]
+    else: z = y.T.T+1
+    GlobalCounters.reset()
+    z.realize()
+    self.assertEqual(GlobalCounters.kernel_count, 2)
+    self.assertEqual(z.tolist(), x.add(2).tolist())
+
+  @unittest.expectedFailure
+  def test_custom_kernel_sched_copy(self): self.test_custom_kernel_sched(use_custom=True)
+
+  @unittest.expectedFailure
+  def test_sliced_buffer_function(self):
+    x = Tensor.arange(32).reshape(8, 4).realize()
+    from tinygrad import function
+    @function(precompile=True)
+    def run(x:Tensor) -> Tensor:
+      y = Tensor.invalids(*x.shape, dtype=x.dtype)
+      return Tensor.custom_kernel(y, x, fxn=custom_add_one_kernel)[0]
+    GlobalCounters.reset()
+    y = run(x[0]).realize()
+    # it's copying the input and the output
+    self.assertEqual(GlobalCounters.kernel_count, 1)
+    self.assertEqual(y.tolist(), [1, 2, 3, 4])
+
 class TestUOpReduce(unittest.TestCase):
   def test_uop_sum(self):
     a = Tensor([1.0, 2, 3, 4, 5])

test/backend/test_jit.py

@@ -4,6 +4,7 @@ import numpy as np
 
 from hypothesis import given, settings, strategies as strat
 from test.helpers import assert_jit_cache_len, call_is_graph, not_support_multi_device, needs_second_gpu
+from tinygrad import Variable
 from tinygrad.tensor import Tensor
 from tinygrad.engine.jit import TinyJit, JitError, graph_class
 from tinygrad.device import Device
@@ -39,6 +40,19 @@ class TestJit(unittest.TestCase):
     def add(a, b): return (a+b).realize()
     _simple_test(add)
 
+  @unittest.skipUnless(Device.DEFAULT == "CPU", "core_id is a CPU runtimevar")
+  def test_hcq_core_id_runtimevar_merge(self):
+    N = 262144
+    @TinyJit
+    def f(x, st):
+      y = (x + 1).contiguous().realize()
+      z = x.shrink(((st, st + N),)).contiguous().realize()
+      return y, z
+    x = Tensor.arange(2*N).contiguous().realize()
+    for _ in range(3): y, z = f(x, Variable("a", 0, N).bind(0))
+    self.assertEqual(y.shape, (2*N,))
+    self.assertEqual(z.shape, (N,))
+
   def test_jitbeam_triggers_beam(self):
     from unittest.mock import patch
     from tinygrad.helpers import getenv as _getenv

test/backend/test_linearizer.py

@@ -2,10 +2,10 @@ import numpy as np
 import unittest
 
 from tinygrad.codegen.opt import Opt, OptOps
-from tinygrad.uop.ops import UOp, Ops, GroupOp, AxisType
+from tinygrad.uop.ops import UOp, Ops, GroupOp, AxisType, buffers
 from tinygrad.device import Device, Buffer, is_dtype_supported
 from tinygrad.tensor import Tensor, _to_np_dtype
-from tinygrad.engine.realize import run_linear, CompiledRunner
+from tinygrad.engine.realize import run_linear
 from tinygrad.codegen import to_program
 from tinygrad.helpers import Context, flatten, dedup, TC_SELECT, TC_OPT, DEV
 from tinygrad.dtype import DType, dtypes, PtrDType, AddrSpace
@@ -14,7 +14,7 @@ from tinygrad.renderer.cstyle import CUDARenderer
 from test.helpers import replace_opts
 MOCKGPU = DEV.interface.startswith("MOCK")
 
-from tinygrad.uop.ops import print_uops # noqa: F401 # pylint: disable=unused-import
+from tinygrad.uop.render import print_uops # noqa: F401 # pylint: disable=unused-import
 
 class TestLinearizer(unittest.TestCase):
   def test_arg_dedup(self):
@@ -274,7 +274,7 @@ class TestLinearizer(unittest.TestCase):
     sched = [si for si in t.schedule_linear().src if si.src[0].op is Ops.SINK]
     # sum_collapse is a full collapse now
     assert len(sched) == 1
-    assert not any(u.op is Ops.REDUCE_AXIS for u in sched[0].src[0].toposort()), "found reduce in sum collapse"
+    assert not any(u.op is Ops.REDUCE and len(u.arg[1]) > 0 for u in sched[0].src[0].toposort()), "found reduce in sum collapse"
     #lin = Kernel(sched[0].ast)
     #assert not any(u.op is Ops.RANGE for u in lin.linearize().uops), "found loop in sum collapse"
 
@@ -424,30 +424,28 @@ def reset_bufs(bufs:list[Buffer]):
 def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[],
                                apply_tc=False, atol=1e-4, rtol=1e-4, color_sizes=[], wanna_output=[]):
   outbufs = real_bufs[:len(realized_ast.src)]
-  device = real_bufs[0].device
   wanna_output = [np.array(x).flatten() for x in wanna_output]
+  buf_uops = [UOp.new_buffer(b.device, b.size, b.dtype) for b in real_bufs]
+  for u,b in zip(buf_uops, real_bufs): buffers[u] = b
 
-  def get_prg(opts):
+  def run_prg(opts):
     ast = realized_ast if opts is None else replace_opts(realized_ast, list(opts))
-    return CompiledRunner(to_program(ast, renderer=Device[Device.DEFAULT].renderer), device)
+    run_linear(UOp(Ops.LINEAR, src=(ast.call(*buf_uops),)))
 
   def check_opt(opts):
-    prg = get_prg(opts=opts)
     reset_bufs(outbufs)
-    prg.exec(real_bufs)
+    run_prg(opts)
     for x,want in zip(copyout_outputs(outbufs), wanna_output): np.testing.assert_allclose(x, want, atol=atol, rtol=rtol)
 
   # Get baseline if it is not provided, which is not optimized at all.
-  prg = get_prg(opts=())
-  prg.exec(real_bufs)
+  run_prg(opts=())
   if len(wanna_output) == 0: wanna_output = copyout_outputs(outbufs)
   else:
     for buf,want in zip(copyout_outputs(outbufs), wanna_output): np.testing.assert_allclose(buf, want, atol=atol, rtol=rtol)
 
   # Check correctness of handcoded optimiztions.
-  prg = get_prg(opts=None)
   reset_bufs(outbufs)
-  prg.exec(real_bufs)
+  run_prg(opts=None)
   for buf,want in zip(copyout_outputs(outbufs), wanna_output): np.testing.assert_allclose(buf, want, atol=atol, rtol=rtol)
   for x in opts: # Check custom transformations if any.
     check_opt(([Opt(OptOps.TC, 0, (TC_SELECT.value, TC_OPT.value, 1))] if apply_tc else [])+x)

test/backend/test_ops.py

@@ -431,6 +431,9 @@ class TestOps(unittest.TestCase):
     helper_test_op(None, lambda x: x.round(), vals=[[1.499, 1.5, 1.501, 1.0, 2.1, 0.0, -5.0, -2.499, -2.5, -2.501]], forward_only=True)
     helper_test_op(None, lambda x: x.round(), vals=[[2.5, -1.5]], forward_only=True)
 
+  def test_round_quantization_gradient(self):
+    helper_test_op(None, lambda x: x + 0.125 * (x.round() - x), vals=[[-1.2, -0.7, -0.2, 0.2, 0.7, 1.2]])
+
   def test_isinf(self):
     val = [float('-inf'), 0., float('inf'), float('nan'), 1.1]
     helper_test_op(None, torch.isinf, Tensor.isinf, vals=[val], forward_only=True)
@@ -606,10 +609,11 @@ class TestOps(unittest.TestCase):
     helper_test_op(None, lambda x,y: x//y, forward_only=True, vals=[[5, 6, 7],[1, 2, 3]])
     helper_test_op(None, lambda x: x/2, forward_only=True, vals=[[3, 4, 5]])
     helper_test_op(None, lambda x: x//2, forward_only=True, vals=[[3, 4, 5]])
-    helper_test_op(None, functools.partial(torch.div, rounding_mode="trunc"), Tensor.idiv, forward_only=True,
+    helper_test_op(None, functools.partial(torch.div, rounding_mode="trunc"),
+                   functools.partial(Tensor.div, rounding_mode="trunc"), forward_only=True,
                    vals=[[-4, 7, 5, 4, -7, 8], [2, -3, 8, -2, 3, 5]])
     if not COMPILE_ONLY:
-      x = Tensor(2**64 - 1, dtype=dtypes.uint64).idiv(1)
+      x = Tensor(2**64 - 1, dtype=dtypes.uint64).div(1, rounding_mode="trunc")
       np.testing.assert_equal(x.numpy(), 2**64 - 1)
 
   def test_scalar_div(self):
@@ -636,6 +640,17 @@ class TestOps(unittest.TestCase):
         helper_test_op(None, lambda x: 100%x, forward_only=True, vals=[va])
         helper_test_op(None, lambda x: 100.5%x, forward_only=True, vals=[va])
 
+  def test_fmod(self):
+    a = [-4, 7, 5, 4, -7, 8, -9]
+    b = [2, -3, 8, -2, 3, 5, -5]
+    for float_a in [True, False]:
+      for float_b in [True, False]:
+        va = [float(ai) for ai in a] if float_a else a
+        vb = [float(bi) for bi in b] if float_b else b
+        helper_test_op(None, lambda x,y: x.fmod(y), forward_only=True, vals=[va, vb])
+        helper_test_op(None, lambda x: x.fmod(2), forward_only=True, vals=[va])
+        helper_test_op(None, lambda x: x.fmod(3.5), forward_only=True, vals=[va])
+
   def test_mul_naninf(self):
     helper_test_op([(45,65)], lambda x: x*math.inf)
     helper_test_op([(45,65)], lambda x: x*-math.inf)
@@ -867,10 +882,10 @@ class TestOps(unittest.TestCase):
     helper_test_op([], lambda: tor >> 31, lambda: ten >> 31, forward_only=True)
 
   def test_idiv_shift_rewrite_negative(self):
-    a = Tensor(-5).idiv(2).item()
-    b = Tensor(-5).contiguous().idiv(2).item()
+    a = Tensor(-5).div(2, rounding_mode="trunc").item()
+    b = Tensor(-5).contiguous().div(2, rounding_mode="trunc").item()
     self.assertEqual(a, b)
-    self.assertEqual(Tensor(-1).contiguous().idiv(4).item(), 0)  # NOTE this is trunc-div behaviour
+    self.assertEqual(Tensor(-1).contiguous().div(4, rounding_mode="trunc").item(), 0)  # NOTE this is trunc-div behaviour
 
   @unittest.skipIf(DEV.renderer == "NAK", "MUFU.SIN is not accurate enough")
   def test_sin(self):

test/backend/test_profiler.py

@@ -3,7 +3,8 @@ from tinygrad import Device, Tensor, dtypes, TinyJit
 from tinygrad.helpers import CI, DEV, Context, ProfileRangeEvent, cpu_profile, cpu_events, ProfilePointEvent, dedup
 from tinygrad.device import Buffer, BufferSpec, Compiled, ProfileDeviceEvent, ProfileGraphEvent
 from tinygrad.runtime.support.hcq import HCQCompiled
-from tinygrad.engine.realize import get_runner
+from tinygrad.engine.realize import get_runtime
+from tinygrad.codegen import to_program
 
 MOCKGPU = DEV.interface.startswith("MOCK")
 def _dev_base(d):
@@ -46,13 +47,15 @@ class TestProfiler(unittest.TestCase):
     TestProfiler.b = self.a + 1
     si = self.b.schedule_linear().src[-1]
 
-    TestProfiler.runner = get_runner(TestProfiler.d0.device, si.src[0])
+    TestProfiler.prg = to_program(si.src[0], TestProfiler.d0.renderer)
+    TestProfiler.runtime = get_runtime(TestProfiler.d0.device, TestProfiler.prg)
     TestProfiler.b.uop.buffer.allocate()
 
-  def test_profile_kernel_run(self):
-    runner_name = TestProfiler.runner._prg.name
+  def test_profile_kernel_run(self, wait=False):
+    runner_name = TestProfiler.runtime.name
     with helper_collect_profile(TestProfiler.d0) as profile:
-      TestProfiler.runner([TestProfiler.b.uop.buffer, TestProfiler.a.uop.buffer], var_vals={})
+      gs, ls = TestProfiler.prg.arg.launch_dims({})
+      TestProfiler.runtime(TestProfiler.b.uop.buffer._buf, TestProfiler.a.uop.buffer._buf, global_size=gs, local_size=ls, wait=wait)
 
     profile, _ = helper_profile_filter_device(profile, TestProfiler.d0.device)
     kernel_runs = [x for x in profile if isinstance(x, ProfileRangeEvent)]
@@ -60,6 +63,9 @@ class TestProfiler(unittest.TestCase):
     assert kernel_runs[0].name == runner_name, "kernel name is not correct"
     assert _dev_base(kernel_runs[0].device) == kernel_runs[0].device, "kernel should not be on a sub-device"
 
+  def test_profile_kernel_run_wait(self):
+    self.test_profile_kernel_run(wait=True)
+
   def test_profile_copyin(self):
     buf1 = Buffer(Device.DEFAULT, 2, dtypes.float, options=BufferSpec(nolru=True)).ensure_allocated()
 
@@ -70,12 +76,13 @@ class TestProfiler(unittest.TestCase):
     assert len(kernel_runs) == 1, "one kernel run is expected"
 
   def test_profile_multiops(self):
-    runner_name = TestProfiler.runner._prg.name
+    runner_name = TestProfiler.runtime.name
     buf1 = Buffer(Device.DEFAULT, 2, dtypes.float, options=BufferSpec(nolru=True)).ensure_allocated()
 
     with helper_collect_profile(TestProfiler.d0) as profile:
       buf1.copyin(memoryview(bytearray(struct.pack("ff", 0, 1))))
-      TestProfiler.runner([buf1, TestProfiler.a.uop.buffer], var_vals={})
+      gs, ls = TestProfiler.prg.arg.launch_dims({})
+      TestProfiler.runtime(buf1._buf, TestProfiler.a.uop.buffer._buf, global_size=gs, local_size=ls)
       buf1.copyout(memoryview(bytearray(buf1.nbytes)))
 
     evs = [x for x in profile if isinstance(x, ProfileRangeEvent) and x.device.startswith(TestProfiler.d0.device)]

test/backend/test_randomness.py

@@ -307,17 +307,26 @@ class TestRandomness(unittest.TestCase):
     with self.assertRaises(TypeError): Tensor.randint((3, 4), low=0, high=3.5)
     with self.assertRaises(TypeError): Tensor.randint((3, 4), low=1, high=3, dtype="float")
     with self.assertRaises(TypeError): Tensor.randint((3, 4), low=0, high=3, dtype=dtypes.float32)
+    # check low < high
+    with self.assertRaises(ValueError): Tensor.randint((3, 4), low=10, high=5)
+    with self.assertRaises(ValueError): Tensor.randint((3, 4), low=10, high=10)
+    np.testing.assert_array_equal(Tensor.randint(16, low=5, high=6).numpy(), 5)
 
   def test_normal(self):
     self.assertTrue(normal_test(Tensor.normal))
     self.assertTrue(equal_distribution(Tensor.normal, lambda x: torch.nn.init.normal_(torch.empty(x), mean=0, std=1),
                                                       lambda x: np.random.normal(loc=0, scale=1, size=x)))
+    # check std >= 0
+    with self.assertRaises(ValueError): Tensor.normal((3, 4), mean=0, std=-1)
 
   def test_uniform(self):
     self.assertFalse(normal_test(Tensor.uniform))
     self.assertTrue(equal_distribution(Tensor.uniform, lambda x: torch.nn.init.uniform_(torch.empty(x)), lambda x: np.random.uniform(size=x)))
     self.assertTrue(equal_distribution(partial(Tensor.uniform, low=-100, high=100, dtype=dtypes.int32),
                                        numpy_func=lambda x: np.random.randint(low=-100, high=100, size=x)))
+    # check low < high
+    with self.assertRaises(ValueError): Tensor.uniform((3, 4), low=5.0, high=3.0)
+    with self.assertRaises(ValueError): Tensor.uniform((3, 4), low=1.0, high=1.0)
 
   def test_scaled_uniform(self):
     self.assertFalse(normal_test(Tensor.scaled_uniform))
@@ -352,7 +361,7 @@ class TestRandomness(unittest.TestCase):
     _check_with_torch(w=[0.231, 0., 1., 0.5], num_samples=300, replacement=True)
     _check_with_torch(w=[[0.2, 0.8]], num_samples=300, replacement=True)  # 2D but only 1 row
     _check_with_torch(w=[[0.453, 0., 1., 0.81], [0.1, 0.8, 0., 0.1]], num_samples=300, replacement=True)
-    # no-replacement isn't supported, unless taking only one sample
+    # no-replacement
     w = [0.1, 0.9]
     self.assertRaises(AssertionError, lambda: Tensor(w).multinomial(100, replacement=False))
 
@@ -363,6 +372,23 @@ class TestRandomness(unittest.TestCase):
     torch_samples = [torch.tensor(w).multinomial(1, replacement=False).item() for _ in range(1000)]
     self.assertTrue(equal_distribution(lambda *_: Tensor(tiny_samples), lambda _: torch.tensor(torch_samples)))
 
+    w = list(range(32))
+    s1 = Tensor(w).multinomial(5, replacement=False).numpy()
+    self.assertEqual(len(set(s1.tolist())), 5)
+    s2 = Tensor(w).multinomial(5, replacement=False).numpy()
+    self.assertFalse(np.array_equal(s1, s2))
+    full = Tensor(w).multinomial(len(w), replacement=False).numpy()
+    self.assertEqual(sorted(full.tolist()), w)
+
+    w = [0.1, 0.2, 0.3, 0.4]
+    @TinyJit
+    def sample_three(): return Tensor(w).multinomial(3, replacement=False).realize()
+
+    tiny_draws = np.array([sample_three().numpy() for _ in range(1000)])
+    torch_draws = np.array([torch.tensor(w).multinomial(3, replacement=False).numpy() for _ in range(1000)])
+    for pos in range(3):
+      self.assertTrue(equal_distribution(lambda *_: Tensor(tiny_draws[:, pos]), lambda _: torch.tensor(torch_draws[:, pos])))
+
   @unittest.skip("this test is flaky")
   def test_multinomial_counterexample(self):
     tiny_res = Tensor([0.3, 0.6, 0.1]).multinomial(4000, replacement=True)

test/backend/test_renderer_failures.py

@@ -1,9 +1,8 @@
 import unittest
 import numpy as np
-from dataclasses import replace
-from tinygrad.device import Buffer, Device, is_dtype_supported
+from tinygrad.device import Device, is_dtype_supported
 from tinygrad.dtype import dtypes, ConstType
-from tinygrad.engine.realize import CompiledRunner
+from tinygrad.engine.realize import run_linear
 from tinygrad.codegen import to_program
 from tinygrad.helpers import prod
 from tinygrad.renderer.cstyle import CStyleLanguage
@@ -13,17 +12,13 @@ from tinygrad.runtime.ops_python import PythonRenderer
 from tinygrad.uop.ops import UOp, Ops, KernelInfo, python_alu
 from tinygrad.tensor import Tensor, _to_np_dtype
 
-def _test_uop_result(inputs:list[Tensor], prg:UOp, local_size=None):
+def _test_uop_result(inputs:list[Tensor], sink:UOp, local_size=None):
   for x in inputs: x.realize()
-  uops = prg.src[2].src
-  outbufs = [Buffer(Device.DEFAULT, sz:=(1 if local_size is None else prod(local_size)), (dtype:=u.src[1].dtype), \
-      initial_value=np.zeros(sz, dtype=_to_np_dtype(dtype)).data) for u in uops if u.op is Ops.STORE]
-  inbufs = [x.uop.base.buffer for x in inputs]
-  info = prg.arg
-  if local_size is not None: info = replace(info, local_size=tuple(local_size))
-  ei = CompiledRunner(prg.replace(arg=info), Device.DEFAULT)
-  ei.exec(outbufs+inbufs)
-  return [np.frombuffer(x.as_memoryview(), _to_np_dtype(x.dtype)) for x in outbufs]
+  sz = 1 if local_size is None else prod(local_size)
+  outs = [UOp.new_buffer(Device.DEFAULT, sz, u.src[1].dtype) for u in sink.src if u.op is Ops.STORE]
+  for u in outs: u.buffer.allocate().copyin(np.zeros(sz, dtype=_to_np_dtype(u.dtype)).data)
+  run_linear(UOp(Ops.LINEAR, src=(sink.call(*outs, *(x.uop.base for x in inputs)),)))
+  return [u.buffer.numpy() for u in outs]
 
 def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp):
   dtype = alu_src_uops[0].dtype
@@ -33,9 +28,7 @@ def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp):
   ld = b.index(idx)
   alu = ld.alu(alu_op, *alu_src_uops)
   store = UOp.store(a.index(idx), alu)
-  sink = UOp(Ops.SINK, dtypes.void, (store,), arg=KernelInfo())
-  prg = to_program(sink, Device[Device.DEFAULT].renderer)
-  return _test_uop_result([Tensor([input_val])], prg)[0]
+  return _test_uop_result([Tensor([input_val])], UOp(Ops.SINK, dtypes.void, (store,), arg=KernelInfo()))[0]
 
 class TestRendererFailures(unittest.TestCase):
   @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
@@ -44,8 +37,7 @@ class TestRendererFailures(unittest.TestCase):
     gate_alu = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'lidx0')).ne(0)
     gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index(lidx0.valid(gate_alu)), UOp.const(dtypes.int, 1)))
     sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,), arg=KernelInfo())
-    prg = to_program(sink, Device[Device.DEFAULT].renderer)
-    ret = _test_uop_result([], prg, local_size=[4, 1, 1])[0]
+    ret = _test_uop_result([], sink, local_size=[4, 1, 1])[0]
     np.testing.assert_equal(ret, [0, 1, 1, 1])
 
   @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
@@ -55,8 +47,7 @@ class TestRendererFailures(unittest.TestCase):
     gate_alu_1 = (lidx1:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 2),), 'lidx1')).ne(0)
     gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index((lidx0+lidx1*4).valid(gate_alu_0&gate_alu_1)), UOp.const(dtypes.int, 1)))
     sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,), arg=KernelInfo())
-    prg = to_program(sink, Device[Device.DEFAULT].renderer)
-    ret = _test_uop_result([], prg, local_size=[4, 2, 1])[0]
+    ret = _test_uop_result([], sink, local_size=[4, 2, 1])[0]
     np.testing.assert_equal(ret, [0, 0, 0, 0, 0, 1, 1, 1])
 
 @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, CStyleLanguage), "uops are for cstyle")
@@ -102,8 +93,7 @@ class TestPTXFailures(unittest.TestCase):
     if_uop = UOp(Ops.IF, dtypes.void, (gate_alu,))
     gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index(lidx0, if_uop), val))
     sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,), arg=KernelInfo())
-    prg = to_program(sink, Device[Device.DEFAULT].renderer)
-    ret = _test_uop_result([], prg, local_size=[4, 1, 1])[0]
+    ret = _test_uop_result([], sink, local_size=[4, 1, 1])[0]
     np.testing.assert_equal(ret, [0, 1, 1, 1])
 
   @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")

test/backend/test_schedule.py

@@ -957,7 +957,7 @@ class TestSchedule(unittest.TestCase):
 
   def test_div_padded_arange(self):
     x = Tensor.full((2,2), 16)
-    y = x.idiv(Tensor.linspace(2, 8, steps=4, dtype=dtypes.int).reshape(2,2)).pad(((1,1), (1,1)))
+    y = x.div(Tensor.linspace(2, 8, steps=4, dtype=dtypes.int).reshape(2,2), rounding_mode="trunc").pad(((1,1), (1,1)))
     out = y.sum(axis=1)
     run_linear(*check_schedule(out, 1))
     self.assertListEqual(out.tolist(), [0, 12, 4, 0])

test/backend/test_uops.py

@@ -5,18 +5,19 @@ from tinygrad.tensor import Tensor, _to_np_dtype
 from tinygrad.helpers import CI, Context
 from tinygrad.dtype import dtypes, DType, AddrSpace, ConstFloat  # noqa: F401
 from tinygrad.device import Buffer, Device
-from tinygrad.uop.ops import Ops, UOp, KernelInfo, AxisType
+from tinygrad.uop.ops import Ops, UOp, KernelInfo, AxisType, buffers
 from tinygrad.renderer.cstyle import CStyleLanguage
-from tinygrad.engine.realize import CompiledRunner, run_linear
+from tinygrad.engine.realize import run_linear
 from tinygrad.codegen import to_program
 from tinygrad.device import is_dtype_supported
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.renderer.ptx import PTXRenderer
 from test.helpers import to_uops_list
 
-def _uops_to_prg(uops_list):
-  prg = to_program(UOp.sink(*uops_list, arg=KernelInfo()), Device[Device.DEFAULT].renderer)
-  return CompiledRunner(prg, Device.DEFAULT)
+def run_uops(uops_list:list[UOp], bufs:list[Buffer]):
+  buf_uops = [UOp.new_buffer(b.device, b.size, b.dtype) for b in bufs]
+  for u,b in zip(buf_uops, bufs): buffers[u] = b
+  run_linear(UOp(Ops.LINEAR, src=(UOp.sink(*uops_list, arg=KernelInfo()).call(*buf_uops),)))
 
 def uop(uops:list[UOp], op:Ops, dtype:Optional[DType], src:tuple[UOp, ...], arg:Any=None) -> UOp:
   if op is Ops.CONST: uops.append(UOp.const(dtype, arg))
@@ -33,8 +34,7 @@ def _test_single_value(vals, op, dts):
   out = uop(uops, Ops.STORE, dtypes.void, (buf_store.index(uop(uops, Ops.CONST, dtypes.int32, (), 0), ptr=True), alu))
   buf = Buffer(Device.DEFAULT, 1, output_dtype).allocate()
   buf2 = [Buffer(Device.DEFAULT, 1, dtype).allocate().copyin(np.array([a], dtype=_to_np_dtype(dtype)).data) for a,dtype in zip(vals, dts)]
-  prg = _uops_to_prg([out])
-  prg.exec([buf]+buf2)
+  run_uops([out], [buf]+buf2)
   ret = np.empty(1, _to_np_dtype(output_dtype))
   buf.copyout(ret.data)
   return ret[0]
@@ -47,8 +47,7 @@ def _test_single_value_const(vals, op, dts):
   alu = uop(uops, op, output_dtype, loads)
   out = buf_store[UOp.const(dtypes.int32, 0)].store(alu)
   buf = Buffer(Device.DEFAULT, 1, output_dtype).allocate()
-  prg = _uops_to_prg([out])
-  prg.exec([buf])
+  run_uops([out], [buf])
   ret = np.empty(1, _to_np_dtype(output_dtype))
   buf.copyout(ret.data)
   return ret[0]
@@ -59,8 +58,7 @@ def _test_uops_result(output_dtype, uops, res):
   # res = output_fn(uops)
   out = uop(uops, Ops.STORE, dtypes.void, (buf_store.index(uop(uops, Ops.CONST, dtypes.int32, (), 0)), res))
   buf = Buffer(Device.DEFAULT, 1, output_dtype).allocate()
-  prg = _uops_to_prg([out])
-  prg.exec([buf])
+  run_uops([out], [buf])
   ret = np.empty(1, _to_np_dtype(output_dtype))
   buf.copyout(ret.data)
   return ret[0]
@@ -135,11 +133,11 @@ class TestNonFloatUOps(TestUOps):
   @unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, CStyleLanguage)), "only ptx and cstyle use bitshifts")
   def test_shl_int32(self): self._test_bop_fxn(Ops.SHL, lambda a,b: int(a)<<int(b), (dtypes.int32, dtypes.int32), no_b_neg=True)
   def test_div_int32(self):
-    self._test_bop_fxn(Ops.IDIV, lambda a,b: int(a/b), (dtypes.int32, dtypes.int32), no_b_zero=True)
+    self._test_bop_fxn(Ops.CDIV, lambda a,b: int(a/b), (dtypes.int32, dtypes.int32), no_b_zero=True)
   def test_and_int32(self): self._test_bop_fxn(Ops.AND, lambda a,b: int(a)&int(b), (dtypes.int32, dtypes.int32))
   def test_or_int32(self): self._test_bop_fxn(Ops.OR, lambda a,b: int(a)|int(b), (dtypes.int32, dtypes.int32))
   def test_mod_int32(self):
-    self._test_bop_fxn(Ops.MOD,
+    self._test_bop_fxn(Ops.CMOD,
                        lambda a,b: abs(int(a))%abs(int(b))*(1,-1)[a<0], (dtypes.int32, dtypes.int32), no_b_zero=True)
   def test_cmplt_int32(self): self._test_bop_fxn(Ops.CMPLT, lambda a,b: int(a)<int(b), (dtypes.int32, dtypes.int32))
   def test_cmpne_int32(self): self._test_bop_fxn(Ops.CMPNE, lambda a,b: int(a)!=int(b), (dtypes.int32, dtypes.int32))

test/device/test_hcq.py

@@ -6,7 +6,7 @@ from tinygrad.device import Buffer, BufferSpec
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQBuffer
 from tinygrad.runtime.autogen import libc
 from tinygrad.runtime.support.system import PCIIfaceBase
-from tinygrad.engine.realize import get_runner, CompiledRunner
+from tinygrad.engine.realize import get_runtime
 from tinygrad.codegen import to_program
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad import Variable
@@ -22,11 +22,12 @@ class TestHCQ(unittest.TestCase):
     TestHCQ.b = self.a + 1
     si = self.b.schedule_linear().src[-1]
 
-    TestHCQ.runner = get_runner(TestHCQ.d0.device, si.src[0])
+    TestHCQ.prg = to_program(si.src[0], TestHCQ.d0.renderer)
+    TestHCQ.runtime = get_runtime(TestHCQ.d0.device, TestHCQ.prg)
     TestHCQ.b.uop.buffer.allocate()
 
-    TestHCQ.kernargs_ba_ptr = TestHCQ.runner._prg.fill_kernargs([TestHCQ.b.uop.buffer._buf, TestHCQ.a.uop.buffer._buf])
-    TestHCQ.kernargs_ab_ptr = TestHCQ.runner._prg.fill_kernargs([TestHCQ.a.uop.buffer._buf, TestHCQ.b.uop.buffer._buf])
+    TestHCQ.kernargs_ba_ptr = TestHCQ.runtime.fill_kernargs([TestHCQ.b.uop.buffer._buf, TestHCQ.a.uop.buffer._buf])
+    TestHCQ.kernargs_ab_ptr = TestHCQ.runtime.fill_kernargs([TestHCQ.a.uop.buffer._buf, TestHCQ.b.uop.buffer._buf])
 
   def setUp(self):
     TestHCQ.d0.synchronize()
@@ -114,7 +115,7 @@ class TestHCQ(unittest.TestCase):
 
   # Test exec
   def test_exec_one_kernel(self):
-    TestHCQ.d0.hw_compute_queue_t().exec(TestHCQ.runner._prg, TestHCQ.kernargs_ba_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size) \
+    TestHCQ.d0.hw_compute_queue_t().exec(TestHCQ.runtime, TestHCQ.kernargs_ba_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size) \
                                    .signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
 
     TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
@@ -128,8 +129,8 @@ class TestHCQ(unittest.TestCase):
 
     q = TestHCQ.d0.hw_compute_queue_t()
     q.wait(TestHCQ.d0.timeline_signal, virt_val - 1) \
-     .exec(TestHCQ.runner._prg, TestHCQ.kernargs_ba_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size) \
-     .exec(TestHCQ.runner._prg, TestHCQ.kernargs_ab_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size) \
+     .exec(TestHCQ.runtime, TestHCQ.kernargs_ba_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size) \
+     .exec(TestHCQ.runtime, TestHCQ.kernargs_ab_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size) \
      .signal(TestHCQ.d0.timeline_signal, virt_val)
 
     for _ in range(100):
@@ -141,11 +142,11 @@ class TestHCQ(unittest.TestCase):
 
   @unittest.skipIf(Device.DEFAULT in {"CPU"}, "No globals/locals on LLVM/CPU")
   def test_exec_update(self):
-    sint_global = (Variable("sint_global", 0, 0xffffffff, dtypes.uint32),) + tuple(TestHCQ.runner.p.global_size[1:])
-    sint_local = (Variable("sint_local", 0, 0xffffffff, dtypes.uint32),) + tuple(TestHCQ.runner.p.local_size[1:])
+    sint_global = (Variable("sint_global", 0, 0xffffffff, dtypes.uint32),) + tuple(TestHCQ.prg.arg.global_size[1:])
+    sint_local = (Variable("sint_local", 0, 0xffffffff, dtypes.uint32),) + tuple(TestHCQ.prg.arg.local_size[1:])
 
     q = TestHCQ.d0.hw_compute_queue_t()
-    q.exec(TestHCQ.runner._prg, TestHCQ.kernargs_ba_ptr, sint_global, sint_local) \
+    q.exec(TestHCQ.runtime, TestHCQ.kernargs_ba_ptr, sint_global, sint_local) \
      .signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
 
     q.submit(TestHCQ.d0, {sint_global[0].expr: 1, sint_local[0].expr: 1})
@@ -166,17 +167,17 @@ class TestHCQ(unittest.TestCase):
     b = a + 1
     si = b.schedule_linear().src[-1]
 
-    runner = CompiledRunner(to_program(replace_opts(si.src[0], [Opt(op=OptOps.LOCAL, axis=0, arg=3) for _ in range(3)]), TestHCQ.d0.renderer),
-                            Device.DEFAULT)
+    prg = to_program(replace_opts(si.src[0], [Opt(op=OptOps.LOCAL, axis=0, arg=3) for _ in range(3)]), TestHCQ.d0.renderer)
+    runtime = get_runtime(Device.DEFAULT, prg)
 
     zb = Buffer(Device.DEFAULT, 3 * 3 * 3, dtypes.int, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
     zt = Buffer(Device.DEFAULT, 3 * 3 * 3, dtypes.int, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
     ctypes.memset(zb._buf.va_addr, 0, zb.nbytes)
-    kernargs = runner._prg.fill_kernargs([zt._buf, zb._buf])
+    kernargs = runtime.fill_kernargs([zt._buf, zb._buf])
 
     q = TestHCQ.d0.hw_compute_queue_t()
     q.memory_barrier() \
-     .exec(runner._prg, kernargs, (1,1,1), virt_local) \
+     .exec(runtime, kernargs, (1,1,1), virt_local) \
      .signal(TestHCQ.d0.timeline_signal, virt_val)
 
     for x in range(1, 4):
@@ -330,7 +331,7 @@ class TestHCQ(unittest.TestCase):
   def test_speed_exec_time(self):
     sig_st, sig_en = TestHCQ.d0.new_signal(), TestHCQ.d0.new_signal()
     TestHCQ.d0.hw_compute_queue_t().timestamp(sig_st) \
-                                   .exec(TestHCQ.runner._prg, TestHCQ.kernargs_ba_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size) \
+                                   .exec(TestHCQ.runtime, TestHCQ.kernargs_ba_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size) \
                                    .timestamp(sig_en) \
                                    .signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
 
@@ -470,12 +471,13 @@ class TestHCQ(unittest.TestCase):
   def test_memory_barrier(self):
     a = Tensor([0, 1], device=Device.DEFAULT, dtype=dtypes.int8).realize()
     b = a + 1
-    runner = get_runner(TestHCQ.d0.device, b.schedule_linear().src[-1].src[0])
+    prg = to_program(b.schedule_linear().src[-1].src[0], TestHCQ.d0.renderer)
+    runtime = get_runtime(TestHCQ.d0.device, prg)
 
     buf1 = Buffer(Device.DEFAULT, 2, dtypes.int8, options=BufferSpec(nolru=True)).ensure_allocated()
     buf2 = Buffer(Device.DEFAULT, 2, dtypes.int8, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
 
-    kernargs_ptr = runner._prg.fill_kernargs([buf1._buf, buf2._buf])
+    kernargs_ptr = runtime.fill_kernargs([buf1._buf, buf2._buf])
 
     for i in range(255):
       ctypes.memset(buf2._buf.va_addr, i, 2)
@@ -483,7 +485,7 @@ class TestHCQ(unittest.TestCase):
       # Need memory_barrier after direct write to vram
       TestHCQ.d0.hw_compute_queue_t().wait(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value - 1) \
                                      .memory_barrier() \
-                                     .exec(runner._prg, kernargs_ptr, runner.p.global_size, runner.p.local_size) \
+                                     .exec(runtime, kernargs_ptr, prg.arg.global_size, prg.arg.local_size) \
                                      .signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
       TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
       TestHCQ.d0.timeline_value += 1

test/device/test_metal.py

@@ -50,19 +50,6 @@ kernel void r_5(device int* data0, const device int* data1, uint3 gid [[threadgr
       compiled = compiled[:40] # corrupt the compiled program
       MetalProgram(device, "r_5", compiled)
 
-  def test_wait_skips_in_flight(self):
-    device = MetalDevice("metal")
-    compiled = MetalCompiler().compile("""
-#include <metal_stdlib>
-kernel void noop(uint3 gid [[threadgroup_position_in_grid]], uint3 lid [[thread_position_in_threadgroup]]) {}
-""")
-    prg = MetalProgram(device, "noop", compiled)
-    self.assertIsInstance(prg(wait=True), float)
-    self.assertEqual(device.mtl_buffers_in_flight, [])
-    self.assertIsNone(prg(wait=False))
-    self.assertEqual(len(device.mtl_buffers_in_flight), 1)
-    device.synchronize()
-
   def test_free(self):
     size = 2**16
     device = Device['METAL']

test/external/external_benchmark_op_conv.py

@@ -3,7 +3,7 @@ from dataclasses import replace
 from tinygrad import dtypes, Device
 from tinygrad.uop.ops import UOp, AxisType, Ops, KernelInfo
 from tinygrad.codegen.opt import Opt, OptOps # pylint: disable=unused-import
-from tinygrad.engine.realize import CompiledRunner
+from tinygrad.engine.realize import get_runtime
 from tinygrad.codegen import to_program
 from tinygrad.helpers import dedup, getenv
 from tinygrad.device import Buffer
@@ -90,12 +90,13 @@ renderer = Device.default.renderer
 allocator = Device.default.allocator
 
 ps = to_program(ast, renderer)
-cr = CompiledRunner(ps, Device.DEFAULT)
+rt = get_runtime(Device.DEFAULT, ps)
 
 gs = sorted(dedup([u for u in ast.toposort() if u.op is Ops.PARAM]), key=lambda u: u.arg)
 # print(len(gs))
 # print([g.dtype for g in gs])
 bufs = [Buffer(ps.arg.device, g.size, g.dtype if isinstance(g.dtype, ImageDType) else g.dtype._base).ensure_allocated() for g in gs]
 
-t = cr(bufs, wait=True)
+gsize, lsize = ps.arg.launch_dims({})
+t = rt(*[b._buf for b in bufs], global_size=gsize, local_size=lsize, vals=ps.arg.vals({}), wait=True)
 print(f"{t*1e6:.2f} us")

test/external/external_fuzz_beam_timeout_recovery.py

@@ -3,7 +3,7 @@
 Stress test for beam timeout + device recovery on AM devices.
 
 Usage:
-  DEV=AMD python test/external/external_test_beam_timeout_recovery.py
+  DEV=AMD python test/external/external_fuzz_beam_timeout_recovery.py
 """
 from tinygrad import Tensor, Device
 from tinygrad.helpers import Context

test/external/external_test_hcq.py

@@ -2,7 +2,8 @@ import unittest, ctypes, struct, time, array
 from tinygrad import Device, Tensor, dtypes
 from tinygrad.helpers import to_mv, CI
 from tinygrad.device import Buffer, BufferSpec
-from tinygrad.engine.realize import get_runner
+from tinygrad.engine.realize import get_runtime
+from tinygrad.codegen import to_program
 
 def _time_queue(q, d):
   st = time.perf_counter()
@@ -21,13 +22,14 @@ class TestHCQ(unittest.TestCase):
     TestHCQ.a = Tensor([0.,1.], device=Device.DEFAULT).realize()
     TestHCQ.b = self.a + 1
     linear = self.b.schedule_linear()
-    TestHCQ.runner = get_runner(TestHCQ.d0.device, linear.src[-1].src[0])
+    TestHCQ.prg = to_program(linear.src[-1].src[0], TestHCQ.d0.renderer)
+    TestHCQ.runtime = get_runtime(TestHCQ.d0.device, TestHCQ.prg)
     TestHCQ.b.uop.buffer.allocate()
     # wow that's a lot of abstraction layers
     TestHCQ.addr = struct.pack("QQ", TestHCQ.b.uop.buffer._buf, TestHCQ.a.uop.buffer._buf)
     TestHCQ.addr2 = struct.pack("QQ", TestHCQ.a.uop.buffer._buf, TestHCQ.b.uop.buffer._buf)
-    TestHCQ.kernargs_off = TestHCQ.runner._prg.kernargs_offset
-    TestHCQ.kernargs_size = TestHCQ.runner._prg.kernargs_alloc_size
+    TestHCQ.kernargs_off = TestHCQ.runtime.kernargs_offset
+    TestHCQ.kernargs_size = TestHCQ.runtime.kernargs_alloc_size
     ctypes.memmove(TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_off, TestHCQ.addr, len(TestHCQ.addr))
     ctypes.memmove(TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_size+TestHCQ.kernargs_off, TestHCQ.addr2, len(TestHCQ.addr2))
 
@@ -38,8 +40,8 @@ class TestHCQ(unittest.TestCase):
     elif Device.DEFAULT == "NV":
       from tinygrad.runtime.ops_nv import HWQueue, HWQueue
       # nv need to copy constbuffer there as well
-      to_mv(TestHCQ.d0.kernargs_ptr, 0x160).cast('I')[:] = array.array('I', TestHCQ.runner._prg.constbuffer_0)
-      to_mv(TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_size, 0x160).cast('I')[:] = array.array('I', TestHCQ.runner._prg.constbuffer_0)
+      to_mv(TestHCQ.d0.kernargs_ptr, 0x160).cast('I')[:] = array.array('I', TestHCQ.runtime.constbuffer_0)
+      to_mv(TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_size, 0x160).cast('I')[:] = array.array('I', TestHCQ.runtime.constbuffer_0)
       TestHCQ.compute_queue = HWQueue
       TestHCQ.copy_queue = HWQueue
 
@@ -53,11 +55,11 @@ class TestHCQ(unittest.TestCase):
     temp_signal, temp_value = TestHCQ.d0._alloc_signal(value=0), 0
     q = TestHCQ.compute_queue()
     for _ in range(1000):
-      q.exec(TestHCQ.runner._prg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
+      q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
       q.signal(temp_signal, temp_value + 1).wait(temp_signal, temp_value + 1)
       temp_value += 1
 
-      q.exec(TestHCQ.runner._prg, TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_size, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
+      q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_size, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
       q.signal(temp_signal, temp_value + 1).wait(temp_signal, temp_value + 1)
       temp_value += 1
 
@@ -71,10 +73,10 @@ class TestHCQ(unittest.TestCase):
   def test_run_1000_times(self):
     temp_signal = TestHCQ.d0._alloc_signal(value=0)
     q = TestHCQ.compute_queue()
-    q.exec(TestHCQ.runner._prg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
+    q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
     q.signal(temp_signal, 2).wait(temp_signal, 2)
-    q.exec(TestHCQ.runner._prg, TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_size, TestHCQ.runner.p.global_size,
-           TestHCQ.runner.p.local_size)
+    q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_size, TestHCQ.prg.arg.global_size,
+           TestHCQ.prg.arg.local_size)
     for _ in range(1000):
       TestHCQ.d0._set_signal(temp_signal, 1)
       q.submit(TestHCQ.d0)
@@ -87,11 +89,11 @@ class TestHCQ(unittest.TestCase):
   def test_run_to_3(self):
     temp_signal = TestHCQ.d0._alloc_signal(value=0)
     q = TestHCQ.compute_queue()
-    q.exec(TestHCQ.runner._prg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
+    q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
     q.signal(temp_signal, 1).wait(temp_signal, 1)
-    q.exec(TestHCQ.runner._prg, TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_size, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
+    q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_size, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
     q.signal(temp_signal, 2).wait(temp_signal, 2)
-    q.exec(TestHCQ.runner._prg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
+    q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
     q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
     TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
@@ -101,7 +103,7 @@ class TestHCQ(unittest.TestCase):
   def test_update_exec(self):
     q = TestHCQ.compute_queue()
     exec_cmd_idx = len(q)
-    q.exec(TestHCQ.runner._prg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
+    q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
     q.update_exec(exec_cmd_idx, (1,1,1), (1,1,1))
     q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
     TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
@@ -115,10 +117,10 @@ class TestHCQ(unittest.TestCase):
   def test_bind_run(self):
     temp_signal = TestHCQ.d0._alloc_signal(value=0)
     q = TestHCQ.compute_queue()
-    q.exec(TestHCQ.runner._prg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
+    q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
     q.signal(temp_signal, 2).wait(temp_signal, 2)
-    q.exec(TestHCQ.runner._prg, TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_size, TestHCQ.runner.p.global_size,
-           TestHCQ.runner.p.local_size)
+    q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_size, TestHCQ.prg.arg.global_size,
+           TestHCQ.prg.arg.local_size)
     q.bind(TestHCQ.d0)
     for _ in range(1000):
       TestHCQ.d0._set_signal(temp_signal, 1)
@@ -133,7 +135,7 @@ class TestHCQ(unittest.TestCase):
   def test_update_exec_binded(self):
     q = TestHCQ.compute_queue()
     exec_ptr = q.ptr()
-    q.exec(TestHCQ.runner._prg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
+    q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
     q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
     q.bind(TestHCQ.d0)
 
@@ -170,7 +172,7 @@ class TestHCQ(unittest.TestCase):
 
   def test_run_normal(self):
     q = TestHCQ.compute_queue()
-    q.exec(TestHCQ.runner._prg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
+    q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
     q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
     TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
@@ -201,7 +203,7 @@ class TestHCQ(unittest.TestCase):
 
   def test_run_signal(self):
     q = TestHCQ.compute_queue()
-    q.exec(TestHCQ.runner._prg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
+    q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
     q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
     q.submit(TestHCQ.d0)
     TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
@@ -278,7 +280,7 @@ class TestHCQ(unittest.TestCase):
   def test_interleave_compute_and_copy(self):
     q = TestHCQ.compute_queue()
     qc = TestHCQ.copy_queue()
-    q.exec(TestHCQ.runner._prg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)  # b = [1, 2]
+    q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)  # b = [1, 2]
     q.signal(sig:=TestHCQ.d0._alloc_signal(value=0), value=1)
     qc.wait(sig, value=1)
     qc.copy(TestHCQ.a.uop.buffer._buf, TestHCQ.b.uop.buffer._buf, 8)
@@ -315,7 +317,7 @@ class TestHCQ(unittest.TestCase):
     for _ in range(40):
       q = TestHCQ.compute_queue()
       q.wait(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value - 1)
-      q.exec(TestHCQ.runner._prg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
+      q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
       q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
       TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
       TestHCQ.d0.timeline_value += 1

test/external/external_test_onnx_runner.py

@@ -144,7 +144,7 @@ class MetadataOnnxPBParser(OnnxPBParser):
     for fid, wire_type in self._parse_message(self.reader.len):
       match fid:
         case 7: obj["graph"] = self._parse_GraphProto()
-        case 14: obj["metadata_props"].append(self._parse_StringStringEntryProto())
+        case 14: obj["metadata_props"].append(self._parse_proto(self._SIMPLE_PROTOS["StringStringEntryProto"]))
         case _: self.reader.skip_field(wire_type)
     return obj
 

test/external/external_test_speed_llama.py

@@ -5,7 +5,7 @@ from tinygrad.tensor import Tensor
 from tinygrad import Device
 from tinygrad.nn.state import get_state_dict
 from tinygrad.device import Allocator, Compiled
-from tinygrad.engine.realize import method_cache
+from tinygrad.codegen import to_program_cache
 from tinygrad.helpers import Profiling
 
 class FakeProgram:
@@ -31,8 +31,8 @@ class TestLLaMASpeed(unittest.TestCase):
     for v in get_state_dict(model).values(): v.assign(Tensor.empty(*v.shape, dtype=v.dtype))
     print("assigned empty tensors, doing warmup")
 
-    def run_llama(st, empty_method_cache=True):
-      if empty_method_cache: method_cache.clear()
+    def run_llama(st, empty_cache=True):
+      if empty_cache: to_program_cache.clear()
       tms = [time.perf_counter()]
       for i in range(5):
         model(Tensor([[1,2,3,4]]), i).realize()

test/external/external_uop_gc.py

@@ -1,7 +1,6 @@
 import gc
 from tinygrad import Tensor, UOp, Device, nn
 from tinygrad.schedule import schedule_cache
-from tinygrad.engine.realize import method_cache
 from tinygrad.codegen import to_program, to_program_cache
 from tinygrad.schedule.indexing import apply_movement_op, _apply_reshape
 from tinygrad.uop.divandmod import fold_divmod_general
@@ -71,7 +70,6 @@ if __name__ == "__main__":
 
     # these caches will keep uops alive
     schedule_cache.clear()
-    method_cache.clear()
     to_program_cache.clear()
     apply_movement_op.cache_clear()
     _apply_reshape.cache_clear()

test/external/fuzz_symbolic.py

@@ -7,7 +7,7 @@ import z3
 from tinygrad import Variable, dtypes
 from tinygrad.uop.ops import UOp
 from tinygrad.uop.validate import uops_to_z3
-from tinygrad.helpers import DEBUG, Context
+from tinygrad.helpers import DEBUG
 
 seed = int(sys.argv[1]) if len(sys.argv) > 1 else random.randint(0, 100)
 print(f"Seed: {seed}", flush=True)
@@ -56,8 +56,7 @@ if __name__ == "__main__":
     v = [u1,u2,u3]
     expr = random_int_expr(6)
 
-    with Context(CORRECT_DIVMOD_FOLDING=1):
-      simplified_expr = expr.simplify()
+    simplified_expr = expr.simplify()
 
     solver = z3.Solver(ctx=z3.Context())
     solver.set(timeout=5000)  # some expressions take very long verify, but its very unlikely they actually return sat
@@ -74,10 +73,9 @@ if __name__ == "__main__":
       m = solver.model()
       n1, n2, n3 = m[v1], m[v2], m[v3]
       u1_val, u2_val, u3_val = u1.const_like(n1.as_long()), u2.const_like(n2.as_long()), u3.const_like(n3.as_long())
-      with Context(CORRECT_DIVMOD_FOLDING=1):
-        num = expr.simplify().substitute({u1:u1_val, u2:u2_val, u3:u3_val}).ssimplify()
-        rn = expr.substitute({u1:u1_val, u2:u2_val, u3:u3_val}).ssimplify()
-        if num==rn: print("z3 found a mismatch but the expressions are equal!!")
+      num = expr.simplify().substitute({u1:u1_val, u2:u2_val, u3:u3_val}).ssimplify()
+      rn = expr.substitute({u1:u1_val, u2:u2_val, u3:u3_val}).ssimplify()
+      if num==rn: print("z3 found a mismatch but the expressions are equal!!")
       assert False, f"mismatched {expr.render()} at v1={m[v1]}; v2={m[v2]}; v3={m[v3]} = {num} != {rn}\n" +\
             "Reproduce with:\n" +\
             f"v1=Variable(\"{u1.arg[0]}\", {u1.arg[1]}, {u1.arg[2]})\n" +\

test/external/fuzz_symbolic_symbolic_div.py

@@ -2,7 +2,7 @@ import random, sys
 import z3
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.uop.validate import uops_to_z3
-from tinygrad.helpers import DEBUG, Context, colored
+from tinygrad.helpers import DEBUG, colored
 
 seed = int(sys.argv[1]) if len(sys.argv) > 1 else random.randint(0, 100)
 print(f"Seed: {seed}", flush=True)
@@ -16,7 +16,7 @@ def get_random_term(ranges, factors):
 def get_random_expr(ranges, factors):
   num_terms = random.randint(2,4)
   x = UOp.usum(*[get_random_term(ranges, factors) for _ in range(num_terms)])
-  return x.alu(random.choice([Ops.IDIV, Ops.MOD]), x.ufix(random.choice(factors)*random.choice([1, 1, 1, -1])))
+  return x.alu(random.choice([Ops.CDIV, Ops.CMOD]), x.ufix(random.choice(factors)*random.choice([1, 1, 1, -1])))
 
 if __name__ == "__main__":
   skipped = 0
@@ -36,8 +36,7 @@ if __name__ == "__main__":
     variable_names += [f"r{i}" for i in range(num_ranges)]
     expr = get_random_expr(ranges, factors)
 
-    with Context(CORRECT_DIVMOD_FOLDING=1):
-      simplified_expr = expr.simplify()
+    simplified_expr = expr.simplify()
 
     if DEBUG>=1:
       print(expr.render(simplify=False), "  -->  ", simplified_expr.render(simplify=False))

test/helpers.py

@@ -17,11 +17,14 @@ from tinygrad.codegen.late.linearizer import linearize
 slow = unittest.skipUnless(os.getenv("RUN_SLOW"), "slow test, set RUN_SLOW=1 to run")
 from tinygrad.runtime.ops_python import PythonProgram, PythonRenderer, PythonCompiler
 
-def get_uops(sink:UOp, ren:Renderer|None=None) -> list[UOp]:
-  """Extract linearized UOps from a sink. Test helper that only does linearization (no render)."""
+def full_rewrite(sink:UOp, ren:Renderer|None=None) -> UOp:
   if ren is None: ren = Renderer(Target())
   if sink.arg is None: sink = sink.replace(arg=KernelInfo())
-  full_sink = full_rewrite_to_sink(sink, ren, optimize=sink.tag is None)
+  return full_rewrite_to_sink(sink, ren, optimize=sink.tag is None)
+
+def get_uops(sink:UOp, ren:Renderer|None=None) -> list[UOp]:
+  """Extract linearized UOps from a sink. Test helper that only does linearization (no render)."""
+  full_sink = full_rewrite(sink, ren)
   return line_rewrite(linearize(full_sink), pm_linearize_cleanups)
 
 def replace_opts(ast:UOp, opts:list) -> UOp: return ast.replace(arg=replace(ast.arg, opts_to_apply=tuple(opts)))

test/mockgpu/am/amgpu.py

@@ -133,8 +133,7 @@ class MockPSP(MockIPBlock):
 
 class MockSMU(MockIPBlock):
   def __init__(self, gpu, mmio):
-    try: regs = import_asic_regs('mp', (11, 0), cls=functools.partial(AMDReg, bases={0: IP_BASES[am.MP1_HWIP]}))
-    except Exception: regs = {}
+    regs = import_asic_regs('mp', (11, 0, 0), cls=functools.partial(AMDReg, bases={0: IP_BASES[am.MP1_HWIP]}))
     super().__init__(gpu, mmio, regs)
     self._msg_pending = False
     def r(n): return self.reg(f"mmMP1_SMN_C2PMSG_{n}")

test/mockgpu/amd/emu.py

@@ -53,10 +53,11 @@ class _MXCSRContext:
 
 from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
 from tinygrad.dtype import dtypes, AddrSpace
-from tinygrad.device import Buffer, BufferSpec
+from tinygrad.device import Buffer, BufferSpec, Device
 from tinygrad.runtime.autogen import hsa
 from tinygrad.helpers import Context, DEBUG, PROFILE, colored
-from tinygrad.engine.realize import get_runner
+from tinygrad.engine.realize import get_runtime
+from tinygrad.codegen import to_program
 
 from tinygrad.renderer.amd import decode_inst
 from tinygrad.runtime.autogen.amd.rdna3.str_pcode import PCODE as PCODE_RDNA3
@@ -374,7 +375,7 @@ def _mem_store(mem: UOp, addr: UOp, val: UOp, active: UOp, addr_bits: int = 32,
   """Conditional memory store with sub-word support. Returns list of store UOps."""
   adt = dtypes.uint64 if addr_bits == 64 else dtypes.uint32
   word_addr = addr >> UOp.const(adt, 2)
-  idx = mem.index(word_addr.cast(dtypes.int), active)
+  idx = mem.index(word_addr.cast(dtypes.int).valid(active))
   if data_bits == 32: return [idx.store(active.where(_to_u32(val), idx))]
   # Sub-word store: read-modify-write with mask
   byte_pos = addr.cast(dtypes.uint32) & _c(3)
@@ -387,7 +388,7 @@ def _mem_store(mem: UOp, addr: UOp, val: UOp, active: UOp, addr_bits: int = 32,
   is_cross = byte_pos.eq(_c(3))
   cross_word0 = (idx & _c(0x00FFFFFF)) | ((val_u32 & _c(0xFF)) << _c(24))
   store0 = idx.store(active.where(is_cross.where(cross_word0, new_word), idx))
-  next_idx = mem.index((word_addr + UOp.const(adt, 1)).cast(dtypes.int), active & is_cross)
+  next_idx = mem.index((word_addr + UOp.const(adt, 1)).cast(dtypes.int).valid(active & is_cross))
   cross_word1 = (next_idx & _c(0xFFFFFF00)) | ((val_u32 >> _c(8)) & _c(0xFF))
   return [store0, next_idx.store((active & is_cross).where(cross_word1, next_idx))]
 
@@ -397,7 +398,7 @@ def _mem_store_bytes(mem: UOp, addr: UOp, val: UOp, active: UOp, data_bits: int
   val_u32 = val.cast(dtypes.uint32) if val.dtype != dtypes.uint32 else val
   for i in range(data_bits // 8):
     byte_val = (val_u32 >> UOp.const(dtypes.uint32, i * 8)) & UOp.const(dtypes.uint32, 0xFF)
-    stores.append(mem.index((addr + UOp.const(dtypes.uint64, i)).cast(dtypes.int), active).store(byte_val.cast(dtypes.uint8)))
+    stores.append(mem.index((addr + UOp.const(dtypes.uint64, i)).cast(dtypes.int).valid(active)).store(byte_val.cast(dtypes.uint8)))
   return stores
 
 def _collect_data_slices(assigns: list[tuple[str, UOp]], data_prefix: str, pcode_vars: dict | None = None, op_name: str = "") -> dict[int, UOp]:
@@ -515,14 +516,14 @@ class _Ctx:
   # Dynamic register access (takes UOp index instead of int)
   def rsgpr_dyn(self, reg: UOp, valid: UOp | None = None) -> UOp:
     """Read SGPR with dynamic register index."""
-    if valid is not None: return self.sgpr.index(reg.cast(dtypes.int), valid, ptr=True).load()
+    if valid is not None: return self.sgpr.index(reg.cast(dtypes.int).valid(valid), ptr=True).load()
     return self.sgpr.index(reg.cast(dtypes.int), ptr=True).load()
 
   def wsgpr_dyn(self, reg: UOp, val: UOp) -> UOp:
     """Write SGPR with dynamic register index. On RDNA, index 124 = NULL (writes discarded). On CDNA, index 124 = M0 (read/write)."""
     # RDNA: NULL (124) discards writes. CDNA: M0 (124) is writable.
     valid = None if self.wave_size == 64 else reg.ne(_c(124))
-    return self.sgpr.index(reg.cast(dtypes.int), valid).store(val.cast(dtypes.uint32))
+    return self.sgpr.index(reg.cast(dtypes.int).valid(valid) if valid is not None else reg.cast(dtypes.int)).store(val.cast(dtypes.uint32))
 
   def wmask(self, reg: UOp, val: UOp) -> list[UOp]:
     """Write a lane mask (VCC/EXEC). Splits into lo/hi for wave64."""
@@ -539,24 +540,24 @@ class _Ctx:
   def rvgpr_dyn(self, reg: UOp, lane: UOp, valid: UOp | None = None) -> UOp:
     """Read VGPR with dynamic register index."""
     idx = reg.cast(dtypes.int) * _c(self.wave_size, dtypes.int) + lane.cast(dtypes.int)
-    return self.vgpr.index(idx, valid, ptr=True).load() if valid is not None else self.vgpr.index(idx, ptr=True).load()
+    return self.vgpr.index(idx.valid(valid), ptr=True).load() if valid is not None else self.vgpr.index(idx, ptr=True).load()
 
   def wvgpr_dyn(self, reg: UOp, lane: UOp, val: UOp, exec_mask: UOp, after: UOp | None = None) -> UOp:
     """Write VGPR with dynamic register index."""
     buf = self.vgpr.after(after) if after is not None else self.vgpr
     offset = reg.cast(dtypes.int) * _c(self.wave_size, dtypes.int) + lane.cast(dtypes.int)
-    return buf.index(offset, _lane_active(exec_mask, lane)).store(val.cast(dtypes.uint32))
+    return buf.index(offset.valid(_lane_active(exec_mask, lane))).store(val.cast(dtypes.uint32))
 
   def raccvgpr_dyn(self, reg: UOp, lane: UOp, valid: UOp | None = None) -> UOp:
     """Read ACCVGPR with dynamic register index (CDNA only)."""
     idx = reg.cast(dtypes.int) * _c(self.wave_size, dtypes.int) + lane.cast(dtypes.int)
-    return self.accvgpr.index(idx, valid, ptr=True).load() if valid is not None else self.accvgpr.index(idx, ptr=True).load()
+    return self.accvgpr.index(idx.valid(valid), ptr=True).load() if valid is not None else self.accvgpr.index(idx, ptr=True).load()
 
   def waccvgpr_dyn(self, reg: UOp, lane: UOp, val: UOp, exec_mask: UOp, after: UOp | None = None) -> UOp:
     """Write ACCVGPR with dynamic register index (CDNA only)."""
     buf = self.accvgpr.after(after) if after is not None else self.accvgpr
     offset = reg.cast(dtypes.int) * _c(self.wave_size, dtypes.int) + lane.cast(dtypes.int)
-    return buf.index(offset, _lane_active(exec_mask, lane)).store(val.cast(dtypes.uint32))
+    return buf.index(offset.valid(_lane_active(exec_mask, lane))).store(val.cast(dtypes.uint32))
 
   def rsrc_dyn(self, off: UOp, lane: UOp | None, bits: int = 32, literal: UOp | None = None, is_f64: bool = False, do_cast: bool = True) -> UOp:
     """Read source operand with dynamic offset. Handles SGPR/inline constants (<256), VGPR (>=256).
@@ -712,7 +713,7 @@ class _Ctx:
         old = self.vgpr.index(val[0].cast(dtypes.int), ptr=True).load()
         new_val = _set_bits(old, _val_to_bits(val[1]), width, lo_bit).cast(dtypes.uint32)
         active = _lane_active(exec_mask, lane)
-        raw_stores.append(('vgpr_direct', self.vgpr.index(val[0].cast(dtypes.int), active).store(new_val)))
+        raw_stores.append(('vgpr_direct', self.vgpr.index(val[0].cast(dtypes.int).valid(active)).store(new_val)))
         continue
       if 'D0' in dest and '[laneId]' in dest:
         old_vcc = self.rmask(_c(VCC_LO.offset))
@@ -1846,7 +1847,7 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA
     if data_bits < 32:
       # Sub-dword LDS write: read-modify-write within the uint32 slot
       word_addr = (addr >> addr_shift).cast(dtypes.int)
-      idx = mem.index(word_addr, active)
+      idx = mem.index(word_addr.valid(active))
       byte_pos = addr.cast(dtypes.uint32) & _c(3)
       byte_shift = byte_pos * _c(8)
       size_mask = _c(0xFF if data_bits == 8 else 0xFFFF)
@@ -2004,17 +2005,18 @@ def _compile_mubuf(inst: irc.MUBUF, ctx: _Ctx) -> UOp:
       word_addr = (addr + UOp.const(dtypes.uint64, i * 4)) >> UOp.const(dtypes.uint64, 2)
       val = in_bounds.where(mem.index(word_addr.cast(dtypes.int64), ptr=True).load(), _c(0))
       lds_idx = ((lds_addr + _c(i * 4)) >> _c(2)).cast(dtypes.int)
-      stores.append(ctx.lds.index(lds_idx, active).store(active.where(val, ctx.lds.index(lds_idx, active))))
+      lds_slot = ctx.lds.index(lds_idx.valid(active))
+      stores.append(lds_slot.store(active.where(val, lds_slot)))
   elif is_store:
     for i in range(n_dwords):
       word_addr = (addr + UOp.const(dtypes.uint64, i * 4)) >> UOp.const(dtypes.uint64, 2)
-      idx = mem.index(word_addr.cast(dtypes.int64), in_bounds)
+      idx = mem.index(word_addr.cast(dtypes.int64).valid(in_bounds))
       val = (ctx.raccvgpr_dyn if use_acc else ctx.rvgpr_dyn)(vdata + _c(i), lane)
       stores.append(idx.store(in_bounds.where(_to_u32(val), idx)))
   else:
     for i in range(n_dwords):
       word_addr = (addr + UOp.const(dtypes.uint64, i * 4)) >> UOp.const(dtypes.uint64, 2)
-      val = in_bounds.where(mem.index(word_addr.cast(dtypes.int64), in_bounds, ptr=True).load(), _c(0))
+      val = in_bounds.where(mem.index(word_addr.cast(dtypes.int64).valid(in_bounds), ptr=True).load(), _c(0))
       stores.append((ctx.waccvgpr_dyn if use_acc else ctx.wvgpr_dyn)(vdata + _c(i), lane, val, exec_mask))
   return UOp.sink(UOp.group(*stores).end(lane), *ctx.inc_pc())
 
@@ -2045,18 +2047,18 @@ _INST_HANDLERS: dict[type, Callable[..., UOp]] = {
 # PROGRAM DECODE AND COMPILATION
 # ═══════════════════════════════════════════════════════════════════════════════
 
-_canonical_runner_cache: list[tuple[type, int, int, int, object]] = []  # [(inst_type, base, mask, size, runner), ...]
+_canonical_runner_cache: list[tuple[type, int, int, int, tuple[UOp, object]]] = []  # [(inst_type, base, mask, size, (prg, runtime)), ...]
 
 @functools.cache
 def _get_runner(inst_bytes: bytes, arch: str = "rdna3"):
-  """Build and compile instruction to CompiledRunner. Cached by instruction bytes, with canonical dedup."""
+  """Build and compile instruction to (prg, runtime). Cached by instruction bytes, with canonical dedup."""
   inst = decode_inst(inst_bytes, arch)
   inst_size = inst.size()
   inst_int = int.from_bytes(inst_bytes[:inst_size], 'little')
 
   # Check if instruction matches any cached canonical pattern (must also match instruction type to avoid variant conflicts)
-  for inst_type, base, mask, size, runner in _canonical_runner_cache:
-    if type(inst) is inst_type and inst_size == size and (inst_int & mask) == base: return runner
+  for inst_type, base, mask, size, entry in _canonical_runner_cache:
+    if type(inst) is inst_type and inst_size == size and (inst_int & mask) == base: return entry
 
   # Look up handler by type, falling back to base classes for _LIT variants
   handler = _INST_HANDLERS.get(type(inst))
@@ -2075,9 +2077,10 @@ def _get_runner(inst_bytes: bytes, arch: str = "rdna3"):
 
   # NOTE: renderer output is not reproducible because of _MXCSRContext. PROFILE=0 prevents emulator instruction runners from polluting profiling.
   with Context(NOOPT=1, CHECK_OOB=0, TUPLE_ORDER=0, EMULATED_DTYPES="", CAPTURE_PROCESS_REPLAY=0, PROFILE=0):
-    runner = get_runner('CPU', sink)
-  _canonical_runner_cache.append((type(inst), base, mask, size, runner))
-  return runner
+    prg = to_program(sink, Device['CPU'].renderer)
+    runtime = get_runtime('CPU', prg)
+  _canonical_runner_cache.append((type(inst), base, mask, size, (prg, runtime)))
+  return prg, runtime
 
 _BARRIER_OPS = {ir3.SOPPOp.S_BARRIER, irc.SOPPOp.S_BARRIER}
 if hasattr(ir4.SOPPOp, 'S_BARRIER_WAIT'): _BARRIER_OPS.add(ir4.SOPPOp.S_BARRIER_WAIT)
@@ -2208,10 +2211,10 @@ def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int,
   def _ensure_compiled(pc: int) -> tuple[Callable, list[int], bool, Inst]:
     if pc not in program:
       prev_len = len(_canonical_runner_cache)
-      runner, inst = _decode_at(pc, arch)
+      (prg, runtime), inst = _decode_at(pc, arch)
       is_barrier = (isinstance(inst, (ir3.SOPP, ir4.SOPP, irc.SOPP)) and inst.op in _BARRIER_OPS) or \
                    (isinstance(inst, (ir4.SOP1,)) and inst.op in _BARRIER_SOP1_OPS)
-      program[pc] = (runner._prg.fxn, runner.p.globals, is_barrier, inst)
+      program[pc] = (runtime.fxn, prg.arg.globals, is_barrier, inst)
       if DEBUG >= 3:
         msg = f"[emu] PC={pc - lib}: {inst!r}"
         print(colored(msg, 'green') if len(_canonical_runner_cache) > prev_len else msg)

test/mockgpu/amd/pcode.py

@@ -828,28 +828,28 @@ class Parser:
     assert mem is not None, "memory load requires _vmem or _lds"
     adt = dtypes.uint64 if addr.dtype == dtypes.uint64 else dtypes.uint32
     active = self.vars.get('_active')
-    gate = (active,) if active is not None else ()
+    def mindex(idx:UOp, ptr=False): return mem.index(idx.valid(active) if active is not None else idx, ptr=ptr)
     byte_mem = mem.dtype.base == dtypes.uint8
     if byte_mem:
       idx = addr.cast(dtypes.int)
       if dt in (dtypes.uint64, dtypes.int64, dtypes.float64):
         val = _u32(0).cast(dtypes.uint64)
-        for i in range(8): val = val | (mem.index(idx + _const(dtypes.int, i), *gate, ptr=True).load().cast(dtypes.uint64) << _u64(i * 8))
+        for i in range(8): val = val | (mindex(idx + _const(dtypes.int, i), ptr=True).load().cast(dtypes.uint64) << _u64(i * 8))
       elif dt in (dtypes.uint8, dtypes.int8):
-        val = mem.index(idx, *gate, ptr=True).load().cast(dt)
+        val = mindex(idx, ptr=True).load().cast(dt)
       elif dt in (dtypes.uint16, dtypes.int16, dtypes.short):
-        lo = mem.index(idx, *gate, ptr=True).load().cast(dtypes.uint32)
-        hi = mem.index(idx + _const(dtypes.int, 1), *gate, ptr=True).load().cast(dtypes.uint32)
+        lo = mindex(idx, ptr=True).load().cast(dtypes.uint32)
+        hi = mindex(idx + _const(dtypes.int, 1), ptr=True).load().cast(dtypes.uint32)
         val = (lo | (hi << _u32(8))).cast(dt)
       else:
         val = _u32(0)
-        for i in range(4): val = val | (mem.index(idx + _const(dtypes.int, i), *gate, ptr=True).load().cast(dtypes.uint32) << _u32(i * 8))
+        for i in range(4): val = val | (mindex(idx + _const(dtypes.int, i), ptr=True).load().cast(dtypes.uint32) << _u32(i * 8))
     else:
       idx = (addr >> _const(addr.dtype, 2)).cast(dtypes.int)
-      val = mem.index(idx, *gate)
+      val = mindex(idx)
       if dt in (dtypes.uint64, dtypes.int64, dtypes.float64):
         idx2 = ((addr + _const(adt, 4)) >> _const(adt, 2)).cast(dtypes.int)
-        val = val.cast(dtypes.uint64) | (mem.index(idx2, *gate).cast(dtypes.uint64) << _u64(32))
+        val = val.cast(dtypes.uint64) | (mindex(idx2).cast(dtypes.uint64) << _u64(32))
       elif dt in (dtypes.uint8, dtypes.int8): val = (val >> ((addr & _const(adt, 3)).cast(dtypes.uint32) * _u32(8))) & _u32(0xFF)
       elif dt in (dtypes.uint16, dtypes.int16):
         val = (val >> (((addr >> _const(adt, 1)) & _const(adt, 1)).cast(dtypes.uint32) * _u32(16))) & _u32(0xFFFF)
@@ -862,7 +862,7 @@ class Parser:
         idx_native = (addr >> _const(adt, 2)).cast(dtypes.int64)
         idx_hi_native = ((addr + _const(adt, 4)) >> _const(adt, 2)).cast(dtypes.int64)
         safe_idx_hi = is_unaligned.where(idx_hi_native, idx_native)
-        hi = mem.index(safe_idx_hi, *gate)
+        hi = mindex(safe_idx_hi)
         combined = val.cast(dtypes.uint64) | (hi.cast(dtypes.uint64) << UOp.const(dtypes.uint64, 32))
         val = is_unaligned.where((combined >> (byte_off.cast(dtypes.uint64) * UOp.const(dtypes.uint64, 8))).cast(dtypes.uint32), val)
     return _cast_to(val, dt)

test/null/test_const_folding.py

@@ -2,7 +2,7 @@ import unittest, itertools, math
 from tinygrad import Tensor, dtypes, Context
 from tinygrad.dtype import DType, ConstType
 from tinygrad.uop.ops import Ops, UOp
-from tinygrad.codegen import full_rewrite_to_sink
+from test.helpers import full_rewrite
 import numpy as np
 
 def _check_ast_count(desired_count:int, t:Tensor):
@@ -79,9 +79,9 @@ class TestBinaryOpsConstFolding(unittest.TestCase):
   def test_div_tensor_one(self):
     _check_ast_count(0, Tensor([1.0, 2, 3, 4]) / Tensor.ones(4))
 
-  def test_idiv_literal_one(self):
+  def test_floordiv_literal_one(self):
     _check_ast_count(0, Tensor([1, 2, 3, 4]) // 1)
-  def test_idiv_tensor_one(self):
+  def test_floordiv_tensor_one(self):
     _check_ast_count(0, Tensor([1, 2, 3, 4]) // Tensor.ones(4, dtype=dtypes.int32))
 
   def test_pow_literal_zero(self):
@@ -103,7 +103,7 @@ class TestBitcastConstFolding(unittest.TestCase):
     def t(cases: dict[DType, ConstType]):
       for (from_dt, from_v), (to_dt, to_v) in itertools.product(cases.items(), cases.items()):
         if not math.isnan(from_v):
-          r = full_rewrite_to_sink(UOp.const(from_dt, from_v).bitcast(to_dt).sink()).src[0]
+          r = full_rewrite(UOp.const(from_dt, from_v).bitcast(to_dt).sink()).src[0]
           self.assertEqual(r.op, Ops.CONST, msg:=f"{from_dt} -> {to_dt} ({from_v} -> {to_v})")
           self.assertEqual(r.dtype, to_dt, msg)
           np.testing.assert_equal(r.arg, to_v, msg)
@@ -127,7 +127,7 @@ class TestBitcastConstFolding(unittest.TestCase):
 
   def test_vec_bitcast(self):
     with Context(SPEC=0):
-      r = full_rewrite_to_sink(UOp.const(dtypes.int32.vec(3), (-1, -2**31, 75)).bitcast(dtypes.uint32.vec(3)).sink()).src[0]
+      r = full_rewrite(UOp.const(dtypes.int32.vec(3), (-1, -2**31, 75)).bitcast(dtypes.uint32.vec(3)).sink()).src[0]
     self.assertEqual(r.op, Ops.STACK)
     self.assertEqual(r.dtype, dtypes.uint32.vec(3))
     self.assertEqual(tuple(x.arg for x in r.src), (2**32-1, 2**31, 75))

test/null/test_graph_rewrite.py

@@ -2,13 +2,13 @@ import unittest, math
 from tinygrad import dtypes
 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 test.helpers import full_rewrite
 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]
+  return full_rewrite(expr.sink()).src[0]
 
 def evaluate_uop(uop, variables):
   if uop.op == Ops.CONST:
@@ -151,7 +151,7 @@ class TestModuloAndDivisionFolding(unittest.TestCase):
 
 class TestEdgeCasesAndSpecialOperations(unittest.TestCase):
   def test_full_graph_rewrite_transcendental_edge_cases(self):
-    optimized_sink = full_rewrite_to_sink(UOp.const(dtypes.float32, -1.0).log2().sink(UOp.const(dtypes.float32, 0.0).reciprocal()))
+    optimized_sink = full_rewrite(UOp.const(dtypes.float32, -1.0).log2().sink(UOp.const(dtypes.float32, 0.0).reciprocal()))
     optimized_log2_neg, optimized_recip_zero = optimized_sink.src
     self.assertTrue(math.isnan(optimized_log2_neg.arg), f"Expected NaN for log2(-1.0), got {optimized_log2_neg.arg}")
     self.assertTrue(math.isinf(optimized_recip_zero.arg) and optimized_recip_zero.arg > 0,
@@ -160,14 +160,14 @@ class TestEdgeCasesAndSpecialOperations(unittest.TestCase):
   @unittest.skip("broken")
   def test_full_graph_rewrite_modulo_negative_dividend(self):
     x_var_uop = UOp.variable('x', -5, -1)
-    optimized_sink = full_rewrite_to_sink((x_var_uop % 3).sink())
+    optimized_sink = full_rewrite((x_var_uop % 3).sink())
     for x_value in range(-5, 0):
       self.assertEqual(x_value % 3, evaluate_uop(optimized_sink.src[0], {'x': x_value}))
 
   @unittest.skip("broken")
   def test_full_graph_rewrite_division_negative_divisor(self):
     x_var_uop = UOp.variable('x', 1, 5)
-    optimized_sink = full_rewrite_to_sink((x_var_uop // -2).sink())
+    optimized_sink = full_rewrite((x_var_uop // -2).sink())
     for x_value in range(1, 6):
       self.assertEqual(x_value // -2, evaluate_uop(optimized_sink.src[0], {'x': x_value}))
 

test/null/test_helpers.py

@@ -209,6 +209,11 @@ class TestFetch(unittest.TestCase):
           headers={"Range": "bytes=0-100"}).read_bytes()
     assert len(x) == 101, f"{len(x) != 101}"
 
+  def test_fetch_sha(self):
+    self.assertRaises(Exception, fetch, "https://ftp.gnu.org/gnu/gzip/gzip-1.13.tar.gz", allow_caching=False, sha256="a")
+    fetch("https://ftp.gnu.org/gnu/gzip/gzip-1.13.tar.gz", allow_caching=False,
+          sha256="20fc818aeebae87cdbf209d35141ad9d3cf312b35a5e6be61bfcfbf9eddd212a")
+
 class TestFullyFlatten(unittest.TestCase):
   def test_fully_flatten(self):
     self.assertEqual(fully_flatten([[1, 3], [1, 2]]), [1, 3, 1, 2])

test/null/test_schedule.py

@@ -380,7 +380,7 @@ class TestSchedule(unittest.TestCase):
     r1 = (x - r0).sum(axis=0).div(2)
     out = r0 + r1
     linear, _ = check_schedule(out, 2)
-    reduceops = [x for si in linear.src for x in si.src[0].toposort() if x.op in {Ops.REDUCE_AXIS, Ops.REDUCE}]
+    reduceops = [x for si in linear.src for x in si.src[0].toposort() if x.op is Ops.REDUCE]
     assert len(reduceops) == 2
 
   def test_cache_reduce_multiple_children(self):
@@ -391,7 +391,7 @@ class TestSchedule(unittest.TestCase):
     out0 = r0 + y
     out1 = r1 + y
     linear, _ = check_schedule([out0, out1], 3)
-    reduceops = [x for si in linear.src for x in si.src[0].toposort() if x.op in {Ops.REDUCE_AXIS, Ops.REDUCE}]
+    reduceops = [x for si in linear.src for x in si.src[0].toposort() if x.op is Ops.REDUCE]
     self.assertEqual(len(reduceops), 2) # why is RANGEIFY different?
 
   def test_dedup_assign(self):
@@ -1286,6 +1286,47 @@ class TestBufferView(unittest.TestCase):
     a = Tensor.arange(4*2).reshape(4, 2).contiguous().shard(devices, axis=1).realize()
     run_linear(*check_schedule(a.flip(0).contiguous(), 2))
 
+  def test_replicated_reshape_is_buffer_view(self):
+    devices = ("NULL:1", "NULL:2")
+    a = Tensor.arange(24).contiguous().to(devices).realize()
+    run_linear(*check_schedule(a.reshape(4, 6).contiguous(), 0))
+
+  def test_replicated_shrink_is_buffer_view(self):
+    # DP pattern: replicated weight[layer_idx]
+    devices = ("NULL:1", "NULL:2")
+    a = Tensor.arange(8*10).reshape(8, 10).contiguous().to(devices).realize()
+    run_linear(*check_schedule(a[3].contiguous(), 0))
+
+  def test_replicated_chained_mops_is_buffer_view(self):
+    devices = ("NULL:1", "NULL:2")
+    a = Tensor.arange(100).contiguous().to(devices).realize()
+    run_linear(*check_schedule(a.reshape(10, 10).shrink(((2, 7), None)).contiguous(), 0))
+
+  def test_replicated_shard_none_is_buffer_view(self):
+    devices = ("NULL:1", "NULL:2")
+    a = Tensor.arange(24).contiguous().shard(devices, axis=None).realize()
+    run_linear(*check_schedule(a.reshape(4, 6).contiguous(), 0))
+
+  def test_replicated_4_devices_is_buffer_view(self):
+    devices = tuple(f"NULL:{i}" for i in range(4))
+    a = Tensor.arange(8*10).reshape(8, 10).contiguous().to(devices).realize()
+    run_linear(*check_schedule(a[3].contiguous(), 0))
+
+  def test_replicated_expand_not_buffer_view(self):
+    devices = ("NULL:1", "NULL:2")
+    a = Tensor.arange(12).reshape(4, 1, 3).contiguous().to(devices).realize()
+    run_linear(*check_schedule(a.expand(4, 3, 3).contiguous(), 2))
+
+  def test_replicated_permute_not_buffer_view(self):
+    devices = ("NULL:1", "NULL:2")
+    a = Tensor.arange(24).reshape(4, 6).contiguous().to(devices).realize()
+    run_linear(*check_schedule(a.permute(1, 0).contiguous(), 2))
+
+  def test_replicated_flip_not_buffer_view(self):
+    devices = ("NULL:1", "NULL:2")
+    a = Tensor.arange(24).reshape(4, 6).contiguous().to(devices).realize()
+    run_linear(*check_schedule(a.flip(0).contiguous(), 2))
+
 class TestInvalidTensor(unittest.TestCase):
   def test_full_invalid_is_zero_kernels(self):
     from tinygrad.dtype import Invalid

test/null/test_simplify_valid_idx.py

@@ -1,12 +1,18 @@
 import unittest, itertools
 
-from tinygrad.codegen import full_rewrite_to_sink
+from tinygrad.codegen.late.devectorizer import load_store_indexing
 from tinygrad.dtype import dtypes
-from tinygrad.uop.ops import UOp, Ops
-from tinygrad.uop.symbolic import simplify_valid
+from tinygrad.uop.ops import UOp, Ops, graph_rewrite
+from tinygrad.uop.symbolic import simplify_valid, sym, pm_move_where_on_load
 from tinygrad.helpers import Context
+from test.helpers import full_rewrite
 from test.null.test_uop_symbolic import check_uop_against_string
 
+# symbolic-only idx + valid simplification (no late lowering of FLOORDIV/FLOORMOD)
+def simplify_valid_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move_where_on_load, name="simplify_valid_idx")
+# image-aware idx + valid simplification: adds the codegen-layer matcher that drops provably in-bounds gates
+def simplify_image_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move_where_on_load+load_store_indexing, name="simplify_image_idx")
+
 def get_gated_load_uop(valid:UOp, idx:UOp):
   return UOp(Ops.LOAD, dtypes.float, (
     UOp(Ops.PARAM, dtypes.float.ptr(), arg=0).index(idx.valid(valid), ptr=True),
@@ -15,7 +21,7 @@ def get_gated_load_uop(valid:UOp, idx:UOp):
 
 def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UOp]):
   return UOp(Ops.LOAD, dtypes.float.vec(4), (
-    UOp(Ops.PARAM, dtypes.imagef(image_shape), arg=0).index(UOp(Ops.STACK, dtypes.weakint.vec(2), idx).valid(valid), ptr=True),
+    UOp(Ops.PARAM, dtypes.imagef(image_shape), arg=0).index(idx[0].valid(valid), idx[1].valid(valid), ptr=True),
     UOp(Ops.STACK, dtypes.float.vec(4), src=(UOp.const(dtypes.float, 0.0),) * 4)
   ))
 
@@ -47,11 +53,10 @@ class TestHelpers(unittest.TestCase):
 
 class TestValidIdxSimplification(unittest.TestCase):
   def check(self, load, sidx, svalid, extra=()):
-    with Context(NOOPT=1, SPEC=0):
-      load = full_rewrite_to_sink(UOp.sink(load, *extra)).src[0]
-    idx, valid = load.src[0].src[1], load.src[0].src[2]
-    check_uop_against_string(self, idx, sidx)
-    check_uop_against_string(self, valid, svalid)
+    load = simplify_valid_idx(UOp.sink(load, *extra)).src[0]
+    off = load.src[0].src[1]
+    check_uop_against_string(self, off.get_idx(), sidx)
+    check_uop_against_string(self, off.get_valid(), svalid)
 
   def test_cumsum(self):
     gidx0 = Special("gidx0", 5)
@@ -157,7 +162,7 @@ class TestValidIdxSimplification(unittest.TestCase):
     valid = (ridx2<1)&(ridx1<6)
     load = get_gated_load_uop(valid, idx)
     # prevent ridx1 and ridx2 from being shrunk
-    red = UOp(Ops.REDUCE, dtypes.float, (load, ridx1, ridx2), Ops.ADD)
+    red = load.reduce(ridx1, ridx2, arg=Ops.ADD)
     self.check(load,
       "(r0*1568)",
       "((r2<1)&(r1<6))",
@@ -216,18 +221,16 @@ class TestValidIdxSimplification(unittest.TestCase):
 
 class TestImageSimplification(unittest.TestCase):
   def check(self, load, svalid, sidx0, sidx1):
-    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.STACK)
-    self.assertEqual(len(idx.src), 2)
-    idx0, idx1 = idx.src[0], idx.src[1]
+    load = simplify_image_idx(load.sink()).src[0]
+    off = load.src[0]
+    idx0, idx1 = off.src[1].get_idx(), off.src[2].get_idx()
     check_uop_against_string(self, idx0, sidx0)
     check_uop_against_string(self, idx1, sidx1)
+    self.assertEqual(off.src[1].get_valid(), off.src[2].get_valid())
     if svalid is not None:
-      check_uop_against_string(self, load.src[0].src[2], svalid)
+      check_uop_against_string(self, off.src[1].get_valid(), svalid)
     else:
-      self.assertEqual(len(load.src[0].src), 2, "svalid is None but load still has a valid")
+      self.assertEqual(off.src[1].get_valid(), UOp.const(dtypes.bool, True), "svalid is None but valid is not True")
 
   def test_idx_gt_c(self):
     # (idx1 < c+1).ne(True) ? (..., idx1-1+c) : 0 can drop the valid
@@ -287,7 +290,7 @@ class TestImageSimplification(unittest.TestCase):
     # empty -> invalid
     load = get_load_image_uop(shape, (gidx0<8) & (gidx0<8).ne(True), idx)
     with Context(NOOPT=1, SPEC=0):
-      load = full_rewrite_to_sink(load.sink()).src[0]
+      load = full_rewrite(load.sink()).src[0]
     self.assertEqual(load.op, Ops.STACK)
     self.assertEqual(load.dtype.count, 4)
 
@@ -447,12 +450,12 @@ class TestImageSimplification(unittest.TestCase):
     load = get_load_image_uop((32, 1024, 4), valid, (alu0, alu1))
     self.check(load, None, "(lidx1*128+gidx0//2+144)", "(lidx0*2+r0+-3)")
 
-    # TODO: this is the same idx as above, but simplifying idx too early makes it hard to drop the valid
+    # same idx, written without the inline simplification of the inner div/mod
     alu0 = ((gidx0*2+lidx1*512+(lidx0*8192+r0*4096)+-11711)//4%1024)
     alu1 = (lidx0*2+r0+-3)
     valid = ((lidx1<7)&((((lidx0*2+r0)<3)!=1)&((lidx0*2+r0)<35)))
     load = get_load_image_uop((32, 1024, 4), valid, (alu0, alu1))
-    self.check(load, "(lidx1<7)", "((gidx0*2+lidx1*512+(lidx0*8192+r0*4096)+-11711)//4%1024)", "(lidx0*2+r0+-3)")
+    self.check(load, None, "(lidx1*128+gidx0//2+144)", "(lidx0*2+r0+-3)")
 
   def test_simplify8(self):
     # from openpilot compile3, kernel r_4_16_8_16_4_4_3_3n1
@@ -508,27 +511,28 @@ class TestUnfoldableImage(unittest.TestCase):
     with Context(SPEC=0):
       lidx = Special("lidx", 2)
       load = UOp(Ops.LOAD, dtypes.float, (UOp(Ops.PARAM, dtypes.imagef((10, 10, 4)), arg=0).index(lidx, ptr=True), UOp.const(dtypes.float, 0)))
-      res = full_rewrite_to_sink(load.sink()).src[0]
+      res = full_rewrite(load.sink()).src[0]
       self.assertEqual(res.src[0].src[0].dtype, dtypes.float.ptr(400))
 
 class TestDropTrueGate(unittest.TestCase):
   def test_drop_true_gate_on_index(self):
-    # test that INDEX with a constant True gate gets simplified to drop the gate
+    # test that INDEX with a constant True valid gets simplified to drop the valid
     from tinygrad.codegen.late.devectorizer import load_store_indexing
     from tinygrad.uop.ops import graph_rewrite
+    from tinygrad.uop.symbolic import sym
     buf = UOp(Ops.PARAM, dtypes.int.ptr(), arg=0)
     idx = UOp.const(dtypes.weakint, 0)
     true_gate = UOp.const(dtypes.bool, True)
-    index_with_gate = UOp(Ops.INDEX, dtypes.int.ptr(), (buf, idx, true_gate))
+    index_with_gate = UOp(Ops.INDEX, dtypes.int.ptr(), (buf, idx.valid(true_gate)))
     # apply the optimization
-    result = graph_rewrite(index_with_gate, load_store_indexing)
-    # the True gate should be dropped (INDEX should only have 2 sources)
-    self.assertEqual(len(result.src), 2, "True gate should be dropped from INDEX")
+    result = graph_rewrite(index_with_gate, sym+load_store_indexing)
+    # the True valid should be dropped (INDEX should only have 2 sources)
+    self.assertEqual(len(result.src), 2, "True valid should be dropped from INDEX")
 
 class TestRangeShrink(unittest.TestCase):
   def get_ranges(self, sink):
     with Context(NOOPT=1, SPEC=0):
-      result = full_rewrite_to_sink(sink)
+      result = full_rewrite(sink)
     return [u for u in result.toposort() if u.op is Ops.RANGE]
 
   def test_range_shrink_single_guard(self):
@@ -569,7 +573,7 @@ class TestRangeShrink(unittest.TestCase):
     # range used in both a gated load AND directly in the reduce expression -> no shrink
     r = Range(0, 204)
     gated_load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r)
-    red = UOp(Ops.REDUCE, dtypes.float, (r.cast(dtypes.float) + gated_load, r), Ops.ADD)
+    red = (r.cast(dtypes.float) + gated_load).reduce(r, arg=Ops.ADD)
     ranges = self.get_ranges(red.sink())
     self.assertEqual(len(ranges), 1)
     self.assertEqual(ranges[0].src[0].arg, 204)

test/null/test_symbolic_failures.py

@@ -1,16 +1,8 @@
 import unittest
 from tinygrad import Variable
-from tinygrad.helpers import Context
 
 
 class TestFuzzFailure(unittest.TestCase):
-  def setUp(self):
-    self.context = Context(CORRECT_DIVMOD_FOLDING=1)
-    self.context.__enter__()
-
-  def tearDown(self):
-    self.context.__exit__(None, None, None)
-
   def test_fuzz_failure1(self):
     v1=Variable('v1', 0, 8)
     v2=Variable('v2', 0, 2)

test/null/test_tensor_uop_mixin.py

@@ -35,6 +35,19 @@ class TestTensorUOpBinop(unittest.TestCase):
   def test_isclose(self):
     t = _t(4).float()
     self.assertIs(_strip_unique(t.isclose(t).uop), _strip_unique(t.uop.isclose(t.uop)))
+  # __floordiv__/mod/fmod and div(rounding_mode=...) dispatch on dtype in mixin
+  def test_floordiv_int(self):   _check(self, _t(4), lambda x: x // 3)
+  def test_floordiv_float(self): _check(self, _t(4).float() + 1.5, lambda x: x // 2.0)
+  def test_rfloordiv_int(self):  _check(self, _t(4)+1, lambda x: 7 // x)
+  def test_mod_int(self):        _check(self, _t(4), lambda x: x % 3)
+  def test_mod_float(self):      _check(self, _t(4).float() + 1.5, lambda x: x % 2.0)
+  def test_div_trunc_int(self):  _check(self, _t(4), lambda x: x.div(3, rounding_mode="trunc"))
+  def test_div_trunc_float(self):_check(self, _t(4).float() + 1.5, lambda x: x.div(2.0, rounding_mode="trunc"))
+  def test_fmod_int(self):       _check(self, _t(4), lambda x: x.fmod(3))
+  def test_fmod_float(self):     _check(self, _t(4).float() + 1.5, lambda x: x.fmod(2.0))
+  def test_floordiv_bool(self):  _check(self, _t(4).cast(dtypes.bool), lambda x: x // True)
+  def test_mod_bool(self):       _check(self, _t(4).cast(dtypes.bool), lambda x: x % True)
+  def test_fmod_bool(self):      _check(self, _t(4).cast(dtypes.bool), lambda x: x.fmod(True))
 
 class TestTensorUOpGetitem(unittest.TestCase):
   # ---- pure slice patterns ----
@@ -137,6 +150,22 @@ class TestTensorUOpAllclose(unittest.TestCase):
     a, b = _t(4).float(), _t(4).float()
     self.assertIs(_strip_unique(a.allclose(b).uop), _strip_unique(a.uop.allclose(b.uop)))
 
+class TestTensorUOpBitcast(unittest.TestCase):
+  def test_bitcast_same_dtype(self): _check(self, _t(4).float(), lambda x: x.bitcast(dtypes.float32))
+
+class TestTensorUOpRand(unittest.TestCase):
+  def test_random_bits(self):
+    k = UOp.empty((2,), dtype=dtypes.uint32)
+    c = UOp.zeros(2, dtype=dtypes.uint32)
+    for num in (1, 4, 7, 1024):
+      self.assertIs(_strip_unique(Tensor.random_bits(Tensor(k), Tensor(c), num).uop),
+                    _strip_unique(UOp.random_bits(k, c, num)))
+  def test_bits_to_rand_float32(self):
+    bits_uop = UOp.empty((8,), dtype=dtypes.uint32)
+    for shape in ((8,), (2, 4), (5,)):
+      self.assertIs(_strip_unique(Tensor._bits_to_rand(Tensor(bits_uop), shape, dtypes.float32).uop),
+                    _strip_unique(UOp._bits_to_rand(bits_uop, shape, dtypes.float32)))
+
 class TestTensorUOpGather(unittest.TestCase):
   def _check(self, t, dim, idx):
     self.assertIs(_strip_unique(t.gather(dim, idx).uop), _strip_unique(t.uop.gather(dim, idx.uop)))

test/null/test_uop_graph.py

@@ -424,7 +424,7 @@ class TestUOpGraph(unittest.TestCase):
     ld = d0.index(ridx0.valid(ridx0<50))
     w = (ridx0<50).where(ld, 5)
     # prevent ridx0 from being shrunk
-    red = UOp(Ops.REDUCE, dtypes.long, (ridx0.cast(dtypes.long), ridx0), Ops.ADD)
+    red = ridx0.cast(dtypes.long).reduce(ridx0, arg=Ops.ADD)
     uops = to_uops_list([w, red])
     for u in uops:
       assert u.op is not Ops.WHERE
@@ -447,7 +447,7 @@ class TestUOpGraph(unittest.TestCase):
     ld = d0.index(gate_idx).cast(dtypes.float)
     w = (ridx0<50).where(ld, 5.0)
     # prevent ridx0 from being shrunk
-    red = UOp(Ops.REDUCE, dtypes.long, (ridx0.cast(dtypes.long), ridx0), Ops.ADD)
+    red = ridx0.cast(dtypes.long).reduce(ridx0, arg=Ops.ADD)
     uops = to_uops_list([w, red])
     for u in uops:
       assert u.op is not Ops.WHERE
@@ -459,7 +459,7 @@ class TestUOpGraph(unittest.TestCase):
     ld = d0.index(ridx0.valid(ridx0<50))
     w = ((ridx0<50) & (ridx0>30)).where(ld, UOp.const(dtypes.float, 0)).cast(dtypes.half)
     # prevent ridx0 from being shrunk
-    red = UOp(Ops.REDUCE, dtypes.long, (ridx0.cast(dtypes.long), ridx0), Ops.ADD)
+    red = ridx0.cast(dtypes.long).reduce(ridx0, arg=Ops.ADD)
     uops = to_uops_list([w, red])
     for u in uops:
       assert u.op is not Ops.WHERE
@@ -470,7 +470,7 @@ class TestUOpGraph(unittest.TestCase):
     ld = d0.index(ridx0.valid(ridx0<50))
     w = ((ridx0<50) & (ridx0>30)).where(UOp.const(dtypes.float, 0), ld).cast(dtypes.half)
     # prevent ridx0 from being shrunk
-    red = UOp(Ops.REDUCE, dtypes.long, (ridx0.cast(dtypes.long), ridx0), Ops.ADD)
+    red = ridx0.cast(dtypes.long).reduce(ridx0, arg=Ops.ADD)
     uops = to_uops_list([w, red])
     for u in uops:
       assert u.op is not Ops.WHERE
@@ -679,7 +679,7 @@ class TestExpander(unittest.TestCase):
   @unittest.skip("no longer supported")
   def test_reduce_known_axis(self):
     e1 = UOp(Ops.UNROLL, dtypes.int, tuple(UOp.const(dtypes.int, x) for x in range(4)), ((1,4),))
-    sink = UOp(Ops.REDUCE, dtypes.int, (3*e1,e1), Ops.ADD)
+    sink = (3*e1).reduce(e1, arg=Ops.ADD)
     sink = expander_rewrite(sink)
     assert sink.op is Ops.CONST
     self.assertEqual(sink.arg, 3*(0+1+2+3))
@@ -687,7 +687,7 @@ class TestExpander(unittest.TestCase):
   @unittest.skip("no longer supported")
   def test_reduce_const(self):
     e1 = UOp(Ops.UNROLL, dtypes.int, tuple(UOp.const(dtypes.int, x) for x in range(4)), ((1,4),))
-    sink = UOp(Ops.REDUCE, dtypes.int, (UOp.const(dtypes.int, 3), e1), Ops.ADD)
+    sink = UOp.const(dtypes.int, 3).reduce(e1, arg=Ops.ADD)
     sink = expander_rewrite(sink)
     assert sink.op is Ops.CONST
     self.assertEqual(sink.arg, 3*4)
@@ -728,7 +728,7 @@ class TestExpander(unittest.TestCase):
   def test_reduce_different_axis(self):
     e1 = UOp(Ops.UNROLL, dtypes.int, tuple(UOp.const(dtypes.int, x) for x in range(4)), ((1,4),))
     e2 = UOp(Ops.UNROLL, dtypes.int, tuple(UOp.const(dtypes.int, x) for x in range(4)), ((2,4),))
-    sink = UOp(Ops.REDUCE, dtypes.int, (e1,e2), Ops.ADD)
+    sink = e1.reduce(e2, arg=Ops.ADD)
     sink = expander_rewrite(sink)
     print(sink)
 
@@ -756,7 +756,7 @@ class TestLoadStoreFolding(unittest.TestCase):
     buf = UOp(Ops.PARAM, dtypes.float.vec(4).ptr(), (), 0)
     idx = UOp.const(dtypes.int, 0)
     gate = UOp.const(dtypes.bool, True)
-    gated_index = buf.index(idx, gate)
+    gated_index = buf.index(idx.valid(gate))
     gep = gated_index.gep(0)
     alt = UOp.const(dtypes.float, 42.0)
     gated_load = gep.load(alt)

test/null/test_uop_symbolic.py

@@ -3,7 +3,6 @@ import unittest, pickle, functools, math
 import z3
 
 from tinygrad.dtype import dtypes, ConstType, DType, Invalid
-from tinygrad.helpers import Context
 from test.helpers import get_uops
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer
 from tinygrad.uop.symbolic import sym, commutative, pm_simplify_valid, pm_move_where_on_load
@@ -181,8 +180,8 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable(Variable("a", 0, 8)*1, 0, 8, "a")
 
   def test_mul_neg_1(self):
-    self.helper_test_variable((Variable("a", 0, 2)*-1)//3, 0, 0, "0")
-    self.helper_test_variable((Variable("a", 2, 7)*-1)//3, -2, 0, "((a//3)*-1)")
+    self.helper_test_variable((Variable("a", 0, 2)*-1)//3, -1, 0, "((a*-1)//3)")
+    self.helper_test_variable((Variable("a", 2, 7)*-1)//3, -3, -1, "((a*-1)//3)")
 
   def test_mul_2(self):
     self.helper_test_variable(Variable("a", 0, 8)*2, 0, 16, "(a*2)")
@@ -203,8 +202,8 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable(Variable("a", 0, 7) // 20, 0, 0, "0")
 
   def test_div_neg_min_max(self):
-    self.helper_test_variable(Variable("a", 1, 7) // -2, -3, 0, "((a//2)*-1)")
-    self.helper_test_variable(Variable("a", 0, 6) // -2, -3, 0, "((a//2)*-1)")
+    self.helper_test_variable(Variable("a", 1, 7) // -2, -4, -1, "(a//-2)")
+    self.helper_test_variable(Variable("a", 0, 6) // -2, -3, 0, "(a//-2)")
 
   def test_div_mod_zero(self):
     with self.assertRaises(ZeroDivisionError):
@@ -238,14 +237,14 @@ class TestSymbolic(unittest.TestCase):
 
   def test_mod_min_max(self):
     self.helper_test_variable(Variable("x", 0, 10)%Variable("y", 1, 10), 0, 9, "(x%y)")
-    self.helper_test_variable(Variable("x", -10, 0)%Variable("y", 1, 10), -9, 0, "(((x*-1)%y)*-1)")
-    self.helper_test_variable(Variable("x", 0, 10)%Variable("y", -10, -1), 0, 9, "(x%(y*-1))")
-    self.helper_test_variable(Variable("x", -10, 0)%Variable("y", -10, -1), -9, 0, "(((x*-1)%(y*-1))*-1)")
-    self.helper_test_variable(Variable("x", -10, 10)%Variable("y", -10, -1), -9, 9, "(x%(y*-1))")
+    self.helper_test_variable(Variable("x", -10, 0)%Variable("y", 1, 10), 0, 9, "(x%y)")
+    self.helper_test_variable(Variable("x", 0, 10)%Variable("y", -10, -1), -9, 0, "(x%y)")
+    self.helper_test_variable(Variable("x", -10, 0)%Variable("y", -10, -1), -9, 0, "(x%y)")
+    self.helper_test_variable(Variable("x", -10, 10)%Variable("y", -10, -1), -9, 0, "(x%y)")
 
-    # test _min_max directly without the rewrite taking out the sign
+    # test _min_max directly: floor mod with positive divisor is in [0, c-1]; with negative divisor in [c+1, 0]
     self.assertEqual((Variable("x", -10, 0)%Variable("y", -10, -1))._min_max, (-9, 0))
-    self.assertEqual((Variable("x", -10, 0)%Variable("y", 1, 10))._min_max, (-9, 0))
+    self.assertEqual((Variable("x", -10, 0)%Variable("y", 1, 10))._min_max, (0, 9))
 
   def test_range_div_its_symbolic_bound(self):
     a = Variable("a", 1, 10, dtypes.weakint)
@@ -262,12 +261,12 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable(Variable("a", 0, 6) // 2, 0, 3, "(a//2)")
 
     self.helper_test_variable(Variable("x", 0, 10)//Variable("y", 1, 10), 0, 10, "(x//y)")
-    self.helper_test_variable(Variable("x", -10, 0)//Variable("y", 1, 10), -10, 0, "(((x*-1)//y)*-1)")
-    self.helper_test_variable(Variable("x", 0, 10)//Variable("y", -10, -1), -10, 0, "((x//(y*-1))*-1)")
-    self.helper_test_variable(Variable("x", -10, 0)//Variable("y", -10, -1), 0, 10, "((x*-1)//(y*-1))")
+    self.helper_test_variable(Variable("x", -10, 0)//Variable("y", 1, 10), -10, 0, "(x//y)")
+    self.helper_test_variable(Variable("x", 0, 10)//Variable("y", -10, -1), -10, 0, "(x//y)")
+    self.helper_test_variable(Variable("x", -10, 0)//Variable("y", -10, -1), 0, 10, "(x//y)")
 
     self.helper_test_variable(Variable("x", -10, 10)//Variable("y", 1, 10), -10, 10, "(x//y)")
-    self.helper_test_variable(Variable("x", -10, 10)//Variable("y", -10, -1), -10, 10, "((x//(y*-1))*-1)")
+    self.helper_test_variable(Variable("x", -10, 10)//Variable("y", -10, -1), -10, 10, "(x//y)")
 
   def test_mod_factor(self):
     self.helper_test_variable(usum([Variable("a", 0, 7)*100, Variable("b", 0, 3)*50]) % 100, 0, 50, "((b%2)*50)")
@@ -334,12 +333,12 @@ class TestSymbolic(unittest.TestCase):
   def test_mod_mod_wrong_sign(self):
     v1=Variable("v1", 0, 128)
     v3=Variable("v3", 0, 7)
-    self.helper_test_variable((((((v1%2)*2)+((v3+-1)%5))+-2)%5), -3, 4, "(v1%2*2+(v3+-1)%5+-2)")
+    self.helper_test_variable((((((v1%2)*2)+((v3+-1)%5))+-2)%5), 0, 4, "((v3+v1%2*2+-3)%5)")
 
   def test_mod_mod_wrong_sign2(self):
     v2=Variable("v2", 0, 8)
     v3=Variable("v3", 0, 4)
-    self.helper_test_variable((((((v3+3)%7)+(v2+-2))%7)%7), -2, 6, "(((v2+((v3+3)%7))+-2)%7)")
+    self.helper_test_variable((((((v3+3)%7)+(v2+-2))%7)%7), 0, 6, "((v2+v3+1)%7)")
 
   def test_mul_mul(self):
     self.helper_test_variable((Variable("a", 0, 5)*10)*9, 0, 5*10*9, "(a*90)")
@@ -357,21 +356,28 @@ class TestSymbolic(unittest.TestCase):
   def test_div_const_div(self):
     a = Variable("a", 0, 124)
     self.helper_test_variable((a//2+1)//2, 0, 31, "((a+2)//4)")
-    self.helper_test_variable(((-a)//2-1)//2, -31, 0, "(((a+2)//4)*-1)")
-    self.helper_test_variable(((-a)//2+10)//2, -26, 5, "((((a//2)*-1)+10)//2)")
+    self.helper_test_variable(((-a)//2-1)//2, -32, -1, "((a*-1+2)//4+-1)")
+    self.helper_test_variable(((-a)//2+10)//2, -26, 5, "(a*-1//4+5)")
 
   def test_div_const_div_wrong_sign(self):
     a = Variable("a", 0, 124)
-    self.helper_test_variable(((a-10)//2+10)//2, 2, 33, "((((a+-10)//2)+10)//2)")
+    self.helper_test_variable(((a-10)//2+10)//2, 2, 33, "((a+2)//4+2)")
 
   def test_div_const_div_wrong_sign_divisor(self):
     a = Variable("a", 0, 124)
-    self.helper_test_variable(((a+10)//-2+10)//-4, -1, 14, "(((((a//2)*-1)+5)//4)*-1)")
+    self.helper_test_variable(((a+10)//-2+10)//-4, -2, 14, "(((a+10)//-2+10)//-4)")
+
+  def test_nested_div_negative_divisor(self):
+    # (x//c1)//c2 -> x//(c1*c2) only when c2>0
+    a = Variable("a", 0, 124)
+    self.helper_test_variable((a//-2)//-3, 0, 20, "((a//-2)//-3)")
+    self.helper_test_variable((a//2)//-3, -21, 0, "((a//2)//-3)")
+    self.helper_test_variable((a//-2)//3, -21, 0, "(a//-6)")
 
   def test_neg_mod(self):
     a = Variable("a", 0, 124)
-    self.helper_test_variable((-a)%4, -3, 0, "((a%4)*-1)")
-    self.helper_test_variable(a%-4, 0, 3, "(a%4)")
+    self.helper_test_variable((-a)%4, 0, 3, "(a*-1%4)")
+    self.helper_test_variable(a%-4, -3, 0, "(a%-4)")
 
   def test_distribute_mul(self):
     self.helper_test_variable(usum([Variable("a", 0, 3), Variable("b", 0, 5)])*3, 0, 24, "((a*3)+(b*3))")
@@ -387,11 +393,11 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable(Variable("a", 0, 6)%100, 0, 6, "a")
 
   def test_big_mod(self):
-    self.helper_test_variable(Variable("a", -20, 20)%10, -9, 9, "(a%10)")
-    self.helper_test_variable(Variable("a", -20, 0)%10, -9, 0, "(((a*-1)%10)*-1)")
-    self.helper_test_variable(Variable("a", -20, 1)%10, -9, 1, "(a%10)")
+    self.helper_test_variable(Variable("a", -20, 20)%10, 0, 9, "(a%10)")
+    self.helper_test_variable(Variable("a", -20, 0)%10, 0, 9, "(a%10)")
+    self.helper_test_variable(Variable("a", -20, 1)%10, 0, 9, "(a%10)")
     self.helper_test_variable(Variable("a", 0, 20)%10, 0, 9, "(a%10)")
-    self.helper_test_variable(Variable("a", -1, 20)%10, -1, 9, "(a%10)")
+    self.helper_test_variable(Variable("a", -1, 20)%10, 0, 9, "(a%10)")
 
   def test_ge_remove(self):
     self.helper_test_variable(Variable("a", 0, 6) >= 25, 0, 0, "False")
@@ -439,8 +445,8 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable(c & c.logical_not(), False, False, "False")
 
   def test_mod_factor_negative(self):
-    self.helper_test_variable(usum([uconst(-29), Variable("a", 0, 10), Variable("b", 0, 10)*28]) % 28, -27, 27, "(((a+(b*28))+-29)%28)")
-    self.helper_test_variable(usum([uconst(-29), Variable("a", 0, 100), Variable("b", 0, 10)*28]) % 28, -27, 27, "(((a+(b*28))+-29)%28)")
+    self.helper_test_variable(usum([uconst(-29), Variable("a", 0, 10), Variable("b", 0, 10)*28]) % 28, 0, 27, "((a+b*28+-29)%28)")
+    self.helper_test_variable(usum([uconst(-29), Variable("a", 0, 100), Variable("b", 0, 10)*28]) % 28, 0, 27, "((a+b*28+-29)%28)")
 
   def test_sum_combine_num(self):
     self.helper_test_variable(usum([uconst(29), Variable("a", 0, 10), uconst(-23)]), 6, 16, "(a+6)")
@@ -448,22 +454,12 @@ class TestSymbolic(unittest.TestCase):
   def test_sum_num_hoisted_and_factors_cancel_out(self):
     self.helper_test_variable(usum([Variable("a", 0, 1) * -4 + 1, Variable("a", 0, 1) * 4]), 1, 1, "1")
 
-  @unittest.expectedFailure  # only correct for floordiv, not truncdiv
   def test_div_cancel(self):
     self.helper_test_variable(usum([uconst(-40), Variable("a", 0, 10)*2, Variable("b", 0, 10)*40])//40, -1, 9, "(b+-1)")
 
-  def test_div_cancel_correct(self):
-    with Context(CORRECT_DIVMOD_FOLDING=1):
-      self.helper_test_variable(usum([uconst(-40), Variable("a", 0, 10)*2, Variable("b", 0, 10)*40])//40, -1, 9, "(((a+(b*20))+-20)//20)")
-
-  @unittest.expectedFailure  # only correct for floordiv, not truncdiv
   def test_mod_cancel(self):
     self.helper_test_variable(usum([uconst(-40), Variable("a", 0, 10)*2, Variable("b", 0, 10)*40]) % 40, 0, 20, "(a*2)")
 
-  def test_mod_cancel_correct(self):
-    with Context(CORRECT_DIVMOD_FOLDING=1):
-      self.helper_test_variable(usum([uconst(-40), Variable("a", 0, 10)*2, Variable("b", 0, 10)*40]) % 40, -38, 38, "((((a+(b*20))+-20)%20)*2)")
-
   def test_mul_div(self):
     self.helper_test_variable((Variable("a", 0, 10)*4)//4, 0, 10, "a")
 
@@ -475,22 +471,22 @@ class TestSymbolic(unittest.TestCase):
     lidx1 = UOp.variable("lidx1", 0, 1)
     ridx1005 = UOp.variable("ridx1005", 0, 2)
     ridx1006 = UOp.variable("ridx1006", 0, 2)
-    self.helper_test_variable((lidx1+((gidx1*18)+(ridx1005*18)+(lidx0*162))+(gidx0*2)+(ridx1006*2)+-40)//18, -2, 20,
-      "(((((lidx1+(((gidx1*18)+(ridx1005*18))+(lidx0*162)))+(gidx0*2))+(ridx1006*2))+-40)//18)")
+    self.helper_test_variable((lidx1+((gidx1*18)+(ridx1005*18)+(lidx0*162))+(gidx0*2)+(ridx1006*2)+-40)//18, -3, 20,
+      "(gidx1+ridx1005+lidx0*9+(gidx0+ridx1006+7)//9+-3)")
 
   def test_add_div(self):
     # careful about the lower bounds and upper bounds
-    self.helper_test_variable((Variable("a", 0, 5)-2)//4, 0, 0, "0")
-    self.helper_test_variable((Variable("a", 0, 5)-1)//4, 0, 1, "((a+-1)//4)")
+    self.helper_test_variable((Variable("a", 0, 5)-2)//4, -1, 0, "((a+2)//4+-1)")
+    self.helper_test_variable((Variable("a", 0, 5)-1)//4, -1, 1, "((a+3)//4+-1)")
     self.helper_test_variable((Variable("a", 0, 5))//4, 0, 1, "(a//4)")
     self.helper_test_variable((Variable("a", 0, 5)+1)//4, 0, 1, "((a+1)//4)")
     self.helper_test_variable((Variable("a", 0, 5)+2)//4, 0, 1, "((a+2)//4)")
     self.helper_test_variable((Variable("a", 0, 5)+3)//4, 0, 2, "((a+3)//4)")
-    self.helper_test_variable((Variable("a", 0, 5)+4)//4, 1, 2, "((a//4)+1)")
-    self.helper_test_variable((Variable("a", 0, 5)+5)//4, 1, 2, "(((a+1)//4)+1)")
+    self.helper_test_variable((Variable("a", 0, 5)+4)//4, 1, 2, "(a//4+1)")
+    self.helper_test_variable((Variable("a", 0, 5)+5)//4, 1, 2, "((a+1)//4+1)")
 
   def test_div_neg_rem(self):
-    self.helper_test_variable((-Variable("a", 0, 255)+256)//2, 0, 128, "((((a+1)//2)*-1)+128)")
+    self.helper_test_variable((-Variable("a", 0, 255)+256)//2, 0, 128, "(a*-1//2+128)")
 
   def test_mul_div_factor_mul(self):
     self.helper_test_variable((Variable("a", 0, 10)*8)//4, 0, 20, "(a*2)")
@@ -502,7 +498,7 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable((Variable("a", 0, 10)*4)//8, 0, 5, "(a//2)")
 
   def test_mul_div_factor_div_neg(self):
-    self.helper_test_variable((Variable("a", 0, 10)*-4+4)//8, -4, 0, "(((a*-1)+1)//2)")
+    self.helper_test_variable((Variable("a", 0, 10)*-4+4)//8, -5, 0, "((a*-1+1)//2)")
 
   def test_div_symbolic_const_gcd(self):
     a = Variable("a", -10, 10)
@@ -520,8 +516,8 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable((d1*a*d2*b*d1)//(d1*d2),  -1000, 1000, "(a*(b*d1))", test_z3=False)
     self.helper_test_variable((d1*a + b*d1)//(d1), -20, 20, "(a+b)", test_z3=False)
     self.helper_test_variable((d1*a + b*d1 + c*d1)//(d1), -30, 30, "(c+(a+b))", test_z3=False)
-    self.helper_test_variable((3*a*d1 + 9*b*d1)//(3*d1*d2), -40, 40, "(((a+(b*3))//(d2*-1))*-1)", test_z3=False)
-    self.helper_test_variable((3*a*d1 + 9*b*d1+3)//(3*d1*d2), -401, 399, "(((((a*d1)+((b*d1)*3))+1)//((d1*d2)*-1))*-1)", test_z3=False)
+    self.helper_test_variable((3*a*d1 + 9*b*d1)//(3*d1*d2), -40, 40, "((a+b*3)//d2)", test_z3=False)
+    self.helper_test_variable((3*a*d1 + 9*b*d1+3)//(3*d1*d2), -401, 399, "((a*d1+b*d1*3+1)//(d1*d2))", test_z3=False)
 
   def test_symbolic_factor_remainder_div(self):
     a = Variable("a", 0, 10)
@@ -532,7 +528,7 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable((d*a*20+b*d*5+10)//(5*d), 0, 52, "((b+(a*4))+(2//d))")
 
   def test_mod_gcd_factor_neg(self):
-    self.helper_test_variable((Variable("a", 0, 10)*-4+4)%8, -4, 4, "((((a*-1)+1)%2)*4)")
+    self.helper_test_variable((Variable("a", 0, 10)*-4+4)%8, 0, 4, "((a*-1+1)%2*4)")
 
   def test_mod_gcd_fold_neg(self):
     self.helper_test_variable((Variable("a", 0, 10)*-8+20)%4, 0, 0, "0")
@@ -540,22 +536,32 @@ class TestSymbolic(unittest.TestCase):
   def test_sum_div_partial_remove(self):
     self.helper_test_variable(usum([Variable("idx0", 0, 127)*4, Variable("idx2", 0, 3)])//4, 0, 127, "idx0")
 
-  def test_cdiv_const_evaluation(self):
-    self.helper_test_variable((Variable("a", 0, 2)-12)//8, -1, -1, "-1")
-    self.helper_test_variable((-Variable("a", 0, 2))//7, 0, 0, "0")
+  def test_floordiv_const_evaluation(self):
+    self.helper_test_variable((Variable("a", 0, 2)-12)//8, -2, -2, "-2")
+    self.helper_test_variable((-Variable("a", 0, 2))//7, -1, 0, "(a*-1//7)")
 
-  def test_cmod_const_evaluation(self):
-    self.helper_test_variable((Variable("a", 1, 1)*-3)%8, -3, -3, "-3")
-    self.helper_test_variable((-Variable("a", 10, 10))%7, -3, -3, "-3")
+  def test_floormod_const_evaluation(self):
+    self.helper_test_variable((Variable("a", 1, 1)*-3)%8, 5, 5, "5")
+    self.helper_test_variable((-Variable("a", 10, 10))%7, 4, 4, "4")
 
   def test_div_numerator_negative(self):
-    with Context(CORRECT_DIVMOD_FOLDING=1):
-      self.helper_test_variable((Variable("idx", 0, 9)*-10)//11, -8, 0, "(((idx*10)//11)*-1)")
+    self.helper_test_variable((Variable("idx", 0, 9)*-10)//11, -9, 0, "(idx*-1)")
 
   def test_nest_div_negative_factor(self):
     ridx0=Variable("ridx0", 0, 9)
     ridx1=Variable("ridx1", 0, 6)
-    self.helper_test_variable(((((ridx0*-7)+ridx1)+63)//35), 0, 1, "(((ridx0//5)*-1)+1)")
+    self.helper_test_variable(((((ridx0*-7)+ridx1)+63)//35), 0, 1, "((ridx0*-1+4)//5+1)")
+
+  def test_floordiv_factor_nest_negative_numerator(self):
+    # x//c = (x//f)//(c//f) for f|c, any sign of x
+    a = Variable("a", -10, 10)
+    b = Variable("b", 0, 3)
+    self.helper_test_variable((a*4 + b)//12, -4, 3, "(a//3)")
+
+  def test_floordiv_gcd_with_remainder_negative_numerator(self):
+    # factor gcd from numerator, even when x crosses zero, as long as the shifted numerator stays nonneg
+    a = Variable("a", -1, 5)
+    self.helper_test_variable((a*2 + 7)//8, 0, 2, "((a+3)//4)")
 
   def test_div_into_mod(self):
     self.helper_test_variable((Variable("idx", 0, 16)*4)%8//4, 0, 1, "(idx%2)")
@@ -568,11 +574,11 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable(x%12//4*4 + x%4 + x//12*12, 0, 23, "x")
 
   def test_div_neg_cancel(self):
-    self.helper_test_variable((-Variable("idx", 0, 100)+199)//-4 + 50, 1, 26, "((idx//4)+1)")
-    self.helper_test_variable((-Variable("idx", 0, 100)+200)//-4 + 50, 0, 25, "((idx+3)//4)")
-    self.helper_test_variable((-Variable("idx", 0, 100)+201)//-4 + 50, 0, 25, "((idx+2)//4)")
-    self.helper_test_variable((-Variable("idx", 0, 100))//2, -50, 0, "((idx//2)*-1)")
-    self.helper_test_variable(Variable("idx", 0, 100)//-2, -50, 0, "((idx//2)*-1)")
+    self.helper_test_variable((-Variable("idx", 0, 100)+199)//-4 + 50, 0, 25, "((idx*-1+199)//-4+50)")
+    self.helper_test_variable((-Variable("idx", 0, 100)+200)//-4 + 50, 0, 25, "((idx*-1+200)//-4+50)")
+    self.helper_test_variable((-Variable("idx", 0, 100)+201)//-4 + 50, -1, 24, "((idx*-1+201)//-4+50)")
+    self.helper_test_variable((-Variable("idx", 0, 100))//2, -50, 0, "(idx*-1//2)")
+    self.helper_test_variable(Variable("idx", 0, 100)//-2, -50, 0, "(idx//-2)")
 
   def test_sum_div_big_const(self):
     gidx0 = Variable("gidx0", 0, 24)
@@ -647,22 +653,22 @@ class TestSymbolic(unittest.TestCase):
   def test_div_neg_all_range(self):
     gidx = Variable("gidx", 0, 124)
     lidx = Variable("lidx", 0, 7)
-    self.helper_test_variable((-gidx*8-lidx+999)//-4 + 250, 1, 250, "(((gidx*2)+(lidx//4))+1)")
-    self.helper_test_variable((-gidx*8-lidx+1000)//-4 + 250, 0, 250, "((gidx*2)+((lidx+3)//4))")
-    self.helper_test_variable((-gidx*8-lidx+1001)//-4 + 250, 0, 250, "((gidx*2)+((lidx+2)//4))")
-    self.helper_test_variable((-gidx*8-lidx+1002)//-4 + 250, 0, 250, "((gidx*2)+((lidx+1)//4))")
+    self.helper_test_variable((-gidx*8-lidx+999)//-4 + 250, 0, 250, "((gidx*-8+lidx*-1+999)//-4+250)")
+    self.helper_test_variable((-gidx*8-lidx+1000)//-4 + 250, 0, 249, "((gidx*-8+lidx*-1+1000)//-4+250)")
+    self.helper_test_variable((-gidx*8-lidx+1001)//-4 + 250, -1, 249, "((gidx*-8+lidx*-1+1001)//-4+250)")
+    self.helper_test_variable((-gidx*8-lidx+1002)//-4 + 250, -1, 249, "((gidx*-8+lidx*-1+1002)//-4+250)")
 
   def test_div_neg_then_neg(self):
     # taken from arange opts
     lidx0 = Variable("lidx0", 0, 7)
     lidx1 = Variable("lidx1", 0, 7)
     alu2 = -lidx0-lidx1
-    self.helper_test_variable((((alu2+14)//(-32))+4), 4, 4, "4")
-    self.helper_test_variable(-(((alu2+14)//(-32))+4), -4, -4, "-4")
-    self.helper_test_variable((((alu2+134)//(-32))+4), 0, 1, "(((lidx0+lidx1)+25)//32)")
-    self.helper_test_variable((((alu2+142)//(-32))+4), 0, 0, "0")
-    self.helper_test_variable((((alu2+150)//(-32))+4), 0, 0, "0")
-    self.helper_test_variable((((alu2+158)//(-32))+4), 0, 0, "0")
+    self.helper_test_variable((((alu2+14)//(-32))+4), 3, 4, "((lidx0*-1+lidx1*-1+14)//-32+4)")
+    self.helper_test_variable(-(((alu2+14)//(-32))+4), -4, -3, "((lidx0*-1+lidx1*-1+14)//-32*-1+-4)")
+    self.helper_test_variable((((alu2+134)//(-32))+4), -1, 0, "((lidx0*-1+lidx1*-1+134)//-32+4)")
+    self.helper_test_variable((((alu2+142)//(-32))+4), -1, 0, "((lidx0*-1+lidx1*-1+142)//-32+4)")
+    self.helper_test_variable((((alu2+150)//(-32))+4), -1, -1, "-1")
+    self.helper_test_variable((((alu2+158)//(-32))+4), -1, -1, "-1")
 
   def test_div_mod_recombine(self):
     gidx = Variable("gidx", 0, 124)
@@ -696,7 +702,7 @@ class TestSymbolic(unittest.TestCase):
     # negative variable range
     xn = Variable("x", -1000, 1000)
     self.helper_test_variable(xn//3%224*3 + xn%3 + xn//672*672, -1000, 1000, "x")
-    self.helper_test_variable(xn//3%7*3 + xn//21*21, -999, 999, "(x//3*3)")
+    self.helper_test_variable(xn//3%7*3 + xn//21*21, -1002, 999, "(x//3*3)")
     # should NOT simplify: a*c1 != b (3*224 != 600)
     self.helper_test_variable(gidx//3%224*3 + gidx//600*600, 0, 150669, "(gidx//600*600+gidx//3%224*3)")
     # should NOT simplify: c1*c2 != c3 (224*3 != 700)
@@ -709,7 +715,7 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable((30 * b + 1) % 18 + ((30 * b + 1) // 18) * 18, 1, 3001, "((b*30)+1)")
 
   def test_div_partial_quotient(self):
-    # IDIV should extract partial quotients when const_factor > divisor, matching what MOD already does
+    # FLOORDIV should extract partial quotients when const_factor > divisor, matching what FLOORMOD already does
     # (f*x+c)//d -> (f%d*x+c)//d + (f//d)*x  when f >= d
     b = Variable("b", 0, 100)
     self.helper_test_variable((31*b+1)//18, 0, 172, "(((b*13)+1)//18+b)")
@@ -730,8 +736,7 @@ class TestSymbolic(unittest.TestCase):
   def test_div_by_factor_tie_break(self):
     a = Variable("a", 0, 1)
     b = Variable("b", 0, 1)
-    with Context(CORRECT_DIVMOD_FOLDING=1):
-      self.helper_test_variable((a*2+b*3+2)//6, 0, 1, "((a+b+1)//3)")
+    self.helper_test_variable((a*2+b*3+2)//6, 0, 1, "((a+b+1)//3)")
 
   def test_div_mod_recombine_large_coeff(self):
     # recombine must work even when coeff > divisor: both mod and div reduce the coeff the same way
@@ -740,8 +745,22 @@ class TestSymbolic(unittest.TestCase):
     a = Variable("a", 0, 10)
     self.helper_test_variable((25*a+3)%10 + ((25*a+3)//10)*10, 3, 253, "((a*25)+3)")
 
+  def test_div_mod_recombine_negative_div_unsound(self):
+    # ((b//d)%div)*mul + (b//(d*div))*(div*mul) only equals (b//d)*mul when div>0
+    b = Variable("b", -100, 100)
+    self.helper_test_variable(((b//(-3))%(-2)) + (b//6)*(-2), -33, 34, "(b//6*-2+b//-3%-2)")
+
+  def test_mod_recombine_with_outer_mul(self):
+    # ((x//c)%d)*(c*mul) + (x%c)*mul -> (x%(c*d))*mul
+    x = Variable("x", 0, 100)
+    self.helper_test_variable((x%4)*3 + ((x//4)%2)*12, 0, 21, "(x%8*3)")
+    y = Variable("y", 0, 5*7*11-1)
+    self.helper_test_variable((y%11)*5 + ((y//11)%7)*55, 0, 380, "(y%77*5)")
+    # negative mul
+    self.helper_test_variable((x%4)*-2 + ((x//4)%2)*-8, -14, 0, "(x%8*-2)")
+
   def test_mod_nest_by_factor(self):
-    # (a*f+b) % (f*k) = (a%k)*f + b when 0<=b<f — mirrors nest_div_by_factor for MOD
+    # (a*f+b) % (f*k) = (a%k)*f + b when 0<=b<f — mirrors nest_div_by_factor for FLOORMOD
     gidx0 = Variable("gidx0", 0, 15)
     lidx0 = Variable("lidx0", 0, 3)
     # f=4, k=2, c=8: (gidx0*4+lidx0)%8 = (gidx0%2)*4 + lidx0
@@ -755,7 +774,7 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable((a*3+b)%9, 0, 8, "(b+a%3*3)")
 
   def test_mod_nest_by_factor_with_const(self):
-    # nest_by_factor MOD with non-zero constant offset: (a*f+b+const) % (f*k) = (a%k)*f + b + const when 0<=b+const<f
+    # nest_by_factor FLOORMOD with non-zero constant offset: (a*f+b+const) % (f*k) = (a%k)*f + b + const when 0<=b+const<f
     a = Variable("a", 0, 7)
     b = Variable("b", 0, 1)
     # f=4, k=2, const=2: (a*4+b+2)%8 = (a%2)*4 + b + 2
@@ -767,7 +786,7 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable((a*3+b+1)%6, 1, 5, "(b+a%2*3+1)")
 
   def test_div_nest_by_factor_with_const(self):
-    # nest_by_factor IDIV: (160*a + 5*b + 4*c + K) // 60 should pick div=5 (clean) over div=4 (dirty)
+    # nest_by_factor FLOORDIV: (160*a + 5*b + 4*c + K) // 60 should pick div=5 (clean) over div=4 (dirty)
     a = Variable("a", 0, 2)
     b = Variable("b", 0, 31)
     c = Variable("c", 0, 1)
@@ -827,12 +846,26 @@ class TestSymbolic(unittest.TestCase):
     # TODO: simplify the true branch
     self.helper_test_variable((idx<4).where(idx//4, idx.const_like(-1)), -1, 6, "(idx<4).where((idx//4), -1)")
 
-  def test_idiv_lt(self):
+  def test_floordiv_lt(self):
+    # x//d<c <=> x<c*d for d>0
     idx = Variable("idx", 0, 24)
     self.helper_test_variable((idx//4<3), 0, 1, "(idx<12)")
-    self.helper_test_variable(((idx-20)//4<-3), 0, 1, "(idx<5)")
-    self.helper_test_variable(((idx-10)//4<0), 0, 1, "(idx<7)")
-    self.helper_test_variable((idx//-4<-3), 0, 1, "(((idx//4)*-1)<-3)")
+    self.helper_test_variable(((idx-20)//4<-3), 0, 1, "(idx<8)")
+    self.helper_test_variable(((idx-10)//4<0), 0, 1, "(idx<10)")
+    self.helper_test_variable((idx//-4<-3), 0, 1, "((idx//-4)<-3)")
+
+  def test_nested_div_mod_negative_inner_divisor(self):
+    # (x % (k*c)) // c -> (x // c) % k requires k>0; (x % (k*c)) % c -> x % c is unconditional for c>0
+    a = Variable("a", 0, 100)
+    self.helper_test_variable((a % -8) // 2, -4, 0, "(a%-8//2)")
+    self.helper_test_variable((a % -8) % 2, 0, 1, "(a%2)")
+
+  def test_floordiv_lt_negative_c(self):
+    # x//d<c with negative c also reduces to x<c*d for d>0
+    idx = Variable("idx", -20, 20)
+    self.helper_test_variable((idx//4 < 0), 0, 1, "(idx<0)")
+    self.helper_test_variable((idx//4 < -1), 0, 1, "(idx<-4)")
+    self.helper_test_variable((idx//4 < -2), 0, 1, "(idx<-8)")
 
   def test_simplex_lt(self):
     a = Variable("a", 0, 3)
@@ -981,10 +1014,10 @@ class TestSymbolic(unittest.TestCase):
     self.assertIn((a.cast(dtypes.long)*b.cast(dtypes.long)).render(), "(long)((a*b))")
 
   def test_nested_mod_negative_range(self):
-    # (x%(k*c))%c = x%c holds for cmod regardless of signs since sign(x%(k*c)) = sign(x)
+    # (x%(k*c))%c = x%c for positive c
     x = Variable("x", 0, 1575)
-    self.helper_test_variable(((x + (-1064)) % 512) % 4, -3, 3, "((x+-1064)%4)")
-    self.helper_test_variable(((x + (-1064)) % 512) % 128, -127, 127, "((x+-1064)%128)")
+    self.helper_test_variable(((x + (-1064)) % 512) % 4, 0, 3, "((x+-1064)%4)")
+    self.helper_test_variable(((x + (-1064)) % 512) % 128, 0, 127, "((x+-1064)%128)")
 
 class TestSymbolicNumeric(unittest.TestCase):
   def helper_test_numeric(self, f):
@@ -1062,12 +1095,13 @@ class TestSymInfer(unittest.TestCase):
     assert sym_infer(a+b+c, var_vals) == 9
     assert sym_infer(a*b, var_vals) == 6
     assert sym_infer(a*b+c, var_vals) == 10
-  def test_sym_infer_cdiv_cmod(self):
+  def test_sym_infer_floordiv_floormod(self):
     a = Variable("a", -1000, 1)
     b = Variable("b", -1000, 1)
     var_vals = {a.expr: 1, b.expr: -1000}
-    assert sym_infer(a%b, var_vals) == 1
-    assert sym_infer(a//b, var_vals) == 0
+    # floor: 1 % -1000 = -999, 1 // -1000 = -1
+    assert sym_infer(a%b, var_vals) == -999
+    assert sym_infer(a//b, var_vals) == -1
   def test_sym_infer_with_bitcast(self):
     a = Variable("a", 1, 10, dtypes.int)
     expr = ((a.bitcast(dtypes.uint) << UOp.const(dtypes.uint, 1)).bitcast(dtypes.int) + 2)
@@ -1286,7 +1320,8 @@ class TestGatedUopGivenValid(unittest.TestCase):
 
     idx:UOp = (r0 < 3).where((r0 + uconst(-1)) // uconst(3), UOp.invalid())
     idx = graph_rewrite(idx, pm_simplify_valid)
-    self.assertEqual(idx, (r0 < 3).where(uconst(0), UOp.invalid()))
+    # (r0-1)//3 = (r0+2)//3 - 1 (constant offset split)
+    self.assertEqual(idx, (r0 < 3).where((r0 + uconst(2)) // uconst(3) + uconst(-1), UOp.invalid()))
 
   def test_invalid_gate_simplifies_vectorize(self):
     r0 = Variable("r0", 0, 2)
@@ -1295,8 +1330,8 @@ class TestGatedUopGivenValid(unittest.TestCase):
     idx1 = r0 % uconst(3)
     idx:UOp = (r0 < 3).where(UOp(Ops.STACK, dtypes.weakint.vec(2), (idx0, idx1)), UOp.invalid())
     idx = graph_rewrite(idx, pm_simplify_valid)
-    # NOTE: independent simplification: (r0-1)//3 -> 0, r0%3 -> r0 when r0 in [0,2]
-    expected_vec = UOp(Ops.STACK, dtypes.weakint.vec(2), (uconst(0), r0))
+    # independent simplification: (r0-1)//3 -> (r0+2)//3 - 1, and r0%3 -> r0 when r0 in [0,2]
+    expected_vec = UOp(Ops.STACK, dtypes.weakint.vec(2), ((r0 + uconst(2)) // uconst(3) + uconst(-1), r0))
     self.assertEqual(idx, (r0 < 3).where(expected_vec, UOp.invalid()))
 
 class TestRangeSplitting(unittest.TestCase):
@@ -1335,8 +1370,8 @@ class TestBounds(unittest.TestCase):
     alu0 = gidx0 * -1
     assert alu0.vmin == -2559 and alu0.vmax == 0
     assert (alu0+2559).vmin == 0 and (alu0+2559).vmax == 2559
-    assert ((alu0+2559)//-4).vmin == -639 and ((alu0+2559)//-4).vmax == 0
-    assert (((alu0+2559)//-4)*(-1)).vmin == 0 and (((alu0+2559)//-4)*(-1)).vmax == 639
+    assert ((alu0+2559)//-4).vmin == -640 and ((alu0+2559)//-4).vmax == 0
+    assert (((alu0+2559)//-4)*(-1)).vmin == 0 and (((alu0+2559)//-4)*(-1)).vmax == 640
 
 class TestFuzzFailure(unittest.TestCase):
   def test_fuzz_failure1(self):

test/null/test_uop_vmin_vmax.py

@@ -173,17 +173,15 @@ class TestVminVmaxDivMod(unittest.TestCase):
     self.assertEqual(uop.vmax, 10)
 
   def test_vmin_vmax_division_negative(self):
-    # vmin and vmax for division of a variable by a negative constant
-    # always positive
+    # floor division of a variable by a negative constant
     x = UOp.variable('x', 10, 20)
     uop = x // -2
     self.assertEqual(uop.vmin, -10)
     self.assertEqual(uop.vmax, -5)
     uop = x // -3
-    self.assertEqual(uop.vmin, -6)
-    self.assertEqual(uop.vmax, -3)
+    self.assertEqual(uop.vmin, -7)
+    self.assertEqual(uop.vmax, -4)
 
-    # always negative
     x = UOp.variable('x', -20, -10)
     uop = x // -2
     self.assertEqual(uop.vmin, 5)
@@ -192,38 +190,61 @@ class TestVminVmaxDivMod(unittest.TestCase):
     self.assertEqual(uop.vmin, 3)
     self.assertEqual(uop.vmax, 6)
 
+  def test_vmin_vmax_floordiv_floormod(self):
+    x = UOp.variable('x', -7, 7)
+    floordiv = x.alu(Ops.FLOORDIV, x.const_like(3))
+    self.assertEqual(floordiv.vmin, -3)
+    self.assertEqual(floordiv.vmax, 2)
+    floormod = x.alu(Ops.FLOORMOD, x.const_like(3))
+    self.assertEqual(floormod.vmin, 0)
+    self.assertEqual(floormod.vmax, 2)
+    # negative const divisor: floormod range is [c+1, 0]
+    floormod_neg = x.alu(Ops.FLOORMOD, x.const_like(-3))
+    self.assertEqual(floormod_neg.vmin, -2)
+    self.assertEqual(floormod_neg.vmax, 0)
+
     # cross 0
     x = UOp.variable('x', -10, 10)
     uop = x // -2
     self.assertEqual(uop.vmin, -5)
     self.assertEqual(uop.vmax, 5)
     uop = x // -3
-    self.assertEqual(uop.vmin, -3)
+    self.assertEqual(uop.vmin, -4)
     self.assertEqual(uop.vmax, 3)
 
+  def test_vmin_vmax_floordiv_floormod_empty_range(self):
+    # empty numerator range (vmin > vmax, e.g. RANGE with end=0) short-circuits to (0, 0)
+    rng = UOp.range(0, 0)
+    self.assertEqual(rng.vmin, 0)
+    self.assertEqual(rng.vmax, -1)
+    self.assertEqual((rng // 4).vmin, 0)
+    self.assertEqual((rng // 4).vmax, 0)
+    self.assertEqual((rng % 4).vmin, 0)
+    self.assertEqual((rng % 4).vmax, 0)
+
   def test_vmin_vmax_div_symbolic(self):
     x = UOp.variable('x', 1, 10)
     y = UOp.variable('y', 3, 5)
     self.assertEqual((x//y).vmin, 0)
     self.assertEqual((x//y).vmax, 3)
-    self.assertEqual(((-x)//y).vmin, -3)
-    self.assertEqual(((-x)//y).vmax, 0)
-    self.assertEqual((x//(-y)).vmin, -3)
-    self.assertEqual((x//(-y)).vmax, 0)
+    self.assertEqual(((-x)//y).vmin, -4)
+    self.assertEqual(((-x)//y).vmax, -1)
+    self.assertEqual((x//(-y)).vmin, -4)
+    self.assertEqual((x//(-y)).vmax, -1)
     self.assertEqual(((-x)//(-y)).vmin, 0)
     self.assertEqual(((-x)//(-y)).vmax, 3)
 
     self.assertEqual((100//y).vmin, 20)
     self.assertEqual((100//y).vmax, 33)
-    self.assertEqual(((-100)//y).vmin, -33)
+    self.assertEqual(((-100)//y).vmin, -34)
     self.assertEqual(((-100)//y).vmax, -20)
-    self.assertEqual((100//(-y)).vmin, -33)
+    self.assertEqual((100//(-y)).vmin, -34)
     self.assertEqual((100//(-y)).vmax, -20)
     self.assertEqual(((-100)//(-y)).vmin, 20)
     self.assertEqual(((-100)//(-y)).vmax, 33)
 
   def test_vmin_vmax_mod_positive(self):
-    # vmin and vmax for modulo of a variable by a positive constant
+    # floor mod with positive divisor: result in [0, c-1] regardless of dividend sign
     positive = UOp.variable('positive', 10, 20)
     uop = positive % 3
     self.assertEqual(uop.vmin, 0)
@@ -231,20 +252,20 @@ class TestVminVmaxDivMod(unittest.TestCase):
 
     negative = UOp.variable('negative', -20, -10)
     uop = negative % 3
-    self.assertEqual(uop.vmin, -2)
-    self.assertEqual(uop.vmax, 0)
+    self.assertEqual(uop.vmin, 0)
+    self.assertEqual(uop.vmax, 2)
 
     mixed = UOp.variable('mixed', -20, 20)
     uop = mixed % 3
-    self.assertEqual(uop.vmin, -2)
+    self.assertEqual(uop.vmin, 0)
     self.assertEqual(uop.vmax, 2)
 
   def test_vmin_vmax_mod_negative(self):
-    # vmin and vmax for modulo of a variable by a negative constant
+    # floor mod with negative divisor: result in [c+1, 0] regardless of dividend sign
     positive = UOp.variable('positive', 10, 20)
     uop = positive % -3
-    self.assertEqual(uop.vmin, 0)
-    self.assertEqual(uop.vmax, 2)
+    self.assertEqual(uop.vmin, -2)
+    self.assertEqual(uop.vmax, 0)
 
     negative = UOp.variable('negative', -20, -10)
     uop = negative % -3
@@ -254,7 +275,7 @@ class TestVminVmaxDivMod(unittest.TestCase):
     mixed = UOp.variable('mixed', -20, 20)
     uop = mixed % -3
     self.assertEqual(uop.vmin, -2)
-    self.assertEqual(uop.vmax, 2)
+    self.assertEqual(uop.vmax, 0)
 
 class TestVminVmaxVConst(unittest.TestCase):
   def test_vmin_vmax_vconst_single_element(self):
@@ -297,7 +318,7 @@ class TestVminVmaxVConst(unittest.TestCase):
     # vmin and vmax for a vector constant of bool values
     d1 = UOp(Ops.PARAM, dtypes.int.ptr(), (), 1)
     idx = UOp.const(dtypes.int, 0)
-    val = UOp(Ops.LOAD, dtypes.int.vec(2), (d1.index(idx),))
+    val = UOp(Ops.LOAD, dtypes.int.vec(2), (d1.index(idx).cast(dtypes.int.vec(2).ptr()),))
     uop = (val // 32).gep(0)
     self.assertEqual(uop.vmin, -67108864)
     self.assertEqual(uop.vmax, 67108863)

test/null/test_uops.py

@@ -41,10 +41,24 @@ class TestExecALU(unittest.TestCase):
     self.assertEqual(exec_alu(Ops.SQRT, dtypes.float, (0.0,)), 0.0)
 
   def test_div(self):
-    self.assertEqual(exec_alu(Ops.IDIV, dtypes.int8, (8, 2)), 4)
-    self.assertEqual(exec_alu(Ops.IDIV, dtypes.int8, (7, 3)), 2)
-    self.assertEqual(exec_alu(Ops.IDIV, dtypes.int8, (7, -3)), -2)
-    self.assertEqual(exec_alu(Ops.IDIV, dtypes.int8, (-50, 6)), -8)
+    self.assertEqual(exec_alu(Ops.CDIV, dtypes.int8, (8, 2)), 4)
+    self.assertEqual(exec_alu(Ops.CDIV, dtypes.int8, (7, 3)), 2)
+    self.assertEqual(exec_alu(Ops.CDIV, dtypes.int8, (7, -3)), -2)
+    self.assertEqual(exec_alu(Ops.CDIV, dtypes.int8, (-50, 6)), -8)
+
+  def test_floordiv(self):
+    self.assertEqual(exec_alu(Ops.FLOORDIV, dtypes.int8, (8, 2)), 4)
+    self.assertEqual(exec_alu(Ops.FLOORDIV, dtypes.int8, (7, 3)), 2)
+    self.assertEqual(exec_alu(Ops.FLOORDIV, dtypes.int8, (7, -3)), -3)
+    self.assertEqual(exec_alu(Ops.FLOORDIV, dtypes.int8, (-7, 3)), -3)
+    self.assertEqual(exec_alu(Ops.FLOORDIV, dtypes.int8, (-50, 6)), -9)
+
+  def test_floormod(self):
+    self.assertEqual(exec_alu(Ops.FLOORMOD, dtypes.int8, (8, 2)), 0)
+    self.assertEqual(exec_alu(Ops.FLOORMOD, dtypes.int8, (7, 3)), 1)
+    self.assertEqual(exec_alu(Ops.FLOORMOD, dtypes.int8, (7, -3)), -2)
+    self.assertEqual(exec_alu(Ops.FLOORMOD, dtypes.int8, (-7, 3)), 2)
+    self.assertEqual(exec_alu(Ops.FLOORMOD, dtypes.int8, (-50, 6)), 4)
 
     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))
@@ -135,8 +149,8 @@ class TestGatedStoreRewrite(unittest.TestCase):
     gidx0 = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'gidx0')
     idx = gidx0*UOp.const(dtypes.int, 2)
     gate = gidx0<UOp.const(dtypes.int, 1)
-    idx0 = UOp(Ops.INDEX, dtypes.float.ptr(), (gmem0, idx, gate))
-    idx1 = UOp(Ops.INDEX, dtypes.float.ptr(), (gmem1, idx, gate))
+    idx0 = UOp(Ops.INDEX, dtypes.float.ptr(), (gmem0, idx.valid(gate)))
+    idx1 = UOp(Ops.INDEX, dtypes.float.ptr(), (gmem1, idx.valid(gate)))
     val = UOp.const(dtypes.float, 42.0)
     stores = [UOp.store(idx0, val), UOp.store(idx1, val)]
     uops = to_uops_list(stores)
@@ -156,31 +170,55 @@ class TestFastIdiv(unittest.TestCase):
       g = UOp(Ops.PARAM, dt.ptr(), (), 0)
       c = UOp.const(dt, 2)
       l = g.index(c)
-      a = UOp(Ops.IDIV, dt, (l, c))
+      a = UOp(Ops.CDIV, dt, (l, c))
       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")
-      self.assertNotIn(Ops.IDIV, ops, f"For dtype={dt} divison by power of two did not simplify to shift")
+      self.assertNotIn(Ops.CDIV, ops, f"For dtype={dt} divison by power of two did not simplify to shift")
+
+  def test_floormod_power_of_two(self):
+    # FLOORMOD by a power of two lowers to AND (correct floor mod for any sign in two's complement)
+    for dt in (dtypes.int32, dtypes.uint32):
+      g = UOp(Ops.PARAM, dt.ptr(), (), 0)
+      c = UOp.const(dt, 8)
+      a = UOp(Ops.FLOORMOD, dt, (g.index(c), c))
+      uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
+      ops = [x.op for x in uops]
+      self.assertIn(Ops.AND, ops, f"For dtype={dt} FLOORMOD by pow2 did not simplify to AND")
+      self.assertNotIn(Ops.CMOD, ops, f"For dtype={dt} FLOORMOD by pow2 left a MOD")
+      self.assertNotIn(Ops.FLOORMOD, ops, f"For dtype={dt} FLOORMOD survived past late rewrite")
+
+  def test_floordiv_power_of_two_uint(self):
+    # uint FLOORDIV by a power of two lowers to a shift, leaving no IDIV/FLOORDIV in the kernel
+    for dt in (dtypes.uint32, dtypes.uint64):
+      g = UOp(Ops.PARAM, dt.ptr(), (), 0)
+      c = UOp.const(dt, 2)
+      a = UOp(Ops.FLOORDIV, dt, (g.index(c), c))
+      uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
+      ops = [x.op for x in uops]
+      self.assertIn(Ops.SHR, ops, f"For dtype={dt} FLOORDIV by power of two did not simplify to shift")
+      self.assertNotIn(Ops.CDIV, ops, f"For dtype={dt} FLOORDIV by power of two did not simplify to shift")
+      self.assertNotIn(Ops.FLOORDIV, ops, f"For dtype={dt} FLOORDIV survived past late rewrite")
 
   @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU doesn't support long")
   def test_fast_idiv_and_mod(self):
     g = UOp(Ops.PARAM, dtypes.uint32.ptr(), (), 0)
     c = UOp.const(dtypes.uint, 3)
     l = g.index(c)
-    a = UOp(Ops.IDIV, dtypes.uint, (l, c))
+    a = UOp(Ops.CDIV, dtypes.uint, (l, c))
     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)
+    self.assertNotIn(Ops.CDIV, ops)
 
-    b = UOp(Ops.MOD, dtypes.uint, (l, c))
+    b = UOp(Ops.CMOD, dtypes.uint, (l, c))
     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)
-    self.assertNotIn(Ops.MOD, ops)
+    self.assertNotIn(Ops.CMOD, ops)
 
   def test_fast_idiv_remove_powers_of_two(self):
     ridx = UOp.range(2**20, 0)
@@ -196,23 +234,23 @@ class TestFastIdiv(unittest.TestCase):
     g = UOp(Ops.PARAM, dtypes.uint32.ptr(), (), 0)
     c = UOp.const(dtypes.uint, 7)
     l = UOp(Ops.LOAD, dtypes.uint, (g.index(c),))
-    a = UOp(Ops.IDIV, dtypes.uint, (l, c))
+    a = UOp(Ops.CDIV, dtypes.uint, (l, c))
     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)
+    self.assertNotIn(Ops.CDIV, ops)
 
   def test_disable_fast_idiv(self):
     g = UOp(Ops.PARAM, dtypes.uint32.ptr(), (), 0)
     c = UOp.const(dtypes.uint, 3)
     l = g.index(c)
-    a = UOp(Ops.IDIV, dtypes.uint, (l, c))
+    a = UOp(Ops.CDIV, dtypes.uint, (l, c))
     with Context(DISABLE_FAST_IDIV=1):
       uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
     ops = [x.op for x in uops]
     self.assertNotIn(Ops.SHR, ops)
-    self.assertIn(Ops.IDIV, ops)
+    self.assertIn(Ops.CDIV, ops)
 
 class TestUOpMethod(unittest.TestCase):
   @unittest.skip("uops lt no longer ordered")

test/null/test_validate_oob.py

@@ -53,7 +53,7 @@ class TestValidateOOB(unittest.TestCase):
         to_uops_list([buf.index(v.valid(v < 20)).store(0)])  # oob
 
   # ALU ops in index
-  def test_idiv(self):
+  def test_floordiv(self):
     with Context(CHECK_OOB=1, SPEC=2):
       buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
       to_uops_list([buf.index(UOp.range(32, 0, AxisType.GLOBAL) // 2, ptr=True).load(dtype=dtypes.int)])  # 0..15 valid
@@ -154,7 +154,7 @@ class TestValidateOOB(unittest.TestCase):
 
       gate = (gidx<400) & (lidx<8)
 
-      local_store = UOp(Ops.STORE, dtypes.void, (sbuf.index(lidx, lidx<8), UOp.const(dtypes.uint, 1)))
+      local_store = sbuf.index(lidx.valid(lidx<8)).store(UOp.const(dtypes.uint, 1))
 
       barrier = UOp(Ops.BARRIER, dtypes.void, (local_store,))
       if_barrier = UOp(Ops.IF, dtypes.void, (gate, barrier))

test/null/test_viz.py

@@ -1,4 +1,4 @@
-import unittest, decimal, sys, json, contextlib, tempfile, pickle, io, itertools
+import unittest, decimal, sys, json, contextlib, tempfile, pickle, io
 from pathlib import Path
 from dataclasses import dataclass
 from typing import Generator
@@ -11,7 +11,7 @@ from tinygrad.helpers import cpu_profile, ProfilePointEvent, unwrap
 from tinygrad.device import Buffer
 
 from tinygrad.uop.ops import tracked_keys, tracked_ctxs, uop_fields, active_rewrites, active_group, _name_cnt, RewriteTrace
-from tinygrad.viz.serve import load_rewrites, get_full_rewrite, uop_to_json, VizData
+from tinygrad.viz.serve import load_rewrites, get_full_rewrite, uop_to_json, VizData, get_render
 from tinygrad.codegen import to_program_cache
 from tinygrad.codegen import to_program
 
@@ -320,8 +320,7 @@ class TestVizGC(unittest.TestCase):
 
 # VIZ integrates with other parts of tinygrad
 
-from tinygrad import Tensor, Device
-from tinygrad.engine.realize import get_runner
+from tinygrad import Tensor, Device, TinyJit, Variable
 
 class TestVizIntegration(unittest.TestCase):
   # codegen supports rendering of code blocks
@@ -409,6 +408,18 @@ class TestVizIntegration(unittest.TestCase):
     lst = viz.list_items()
     assert len(lst) == 1
 
+  def test_jit(self):
+    with save_viz():
+      @TinyJit
+      def f(a, b, c): return (a+b).contiguous().mul(3), c.add(1).contiguous().assign(a.to(c.device)), b.assign(c.to(b.device))
+      a, b, c = Tensor.empty(16, device="NULL"), Tensor.empty(16, device="NULL"), Tensor.empty(16, device="NULL:1")
+      for _ in range(3): Tensor.realize(*f(a, b, c))
+    out = load_profile(cpu_events)
+    self.assertEqual(["NULL", "NULL Graph", "NULL:SDMA:0", "NULL:1", "NULL:1:SDMA:0"], [k for k in out["layout"] if k.startswith("NULL")])
+    self.assertEqual(len(out["layout"]["NULL"]["events"]), 2*3)
+    self.assertEqual(len(out["layout"]["NULL:SDMA:0"]["events"]), 3)
+    self.assertEqual(len(out["layout"]["NULL Graph"]["events"]), 2)
+
 from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry
 from tinygrad.viz.serve import get_profile
 from tinygrad.viz.cli import decode_profile
@@ -422,9 +433,9 @@ class TestVizProfiler(unittest.TestCase):
       a.to("NULL:1").realize()
     range_events = [e for e in cpu_events if isinstance(e, ProfileRangeEvent)]
     compute_events = [e for e in range_events if e.device == "NULL"]
-    copy_events = [e for e in range_events if e.device.endswith(":COPY")]
+    copy_events = [e for e in range_events if e.device.endswith(":SDMA:0")]
     self.assertGreater(len(compute_events), 0, "expected compute events on base device")
-    self.assertGreater(len(copy_events), 0, "transfer must produce events with ':COPY' device suffix")
+    self.assertGreater(len(copy_events), 0, "transfer must produce events with ':SDMA' device suffix")
 
   def test_node(self):
     prof = [ProfileRangeEvent(device='NV', name='E_2', st=decimal.Decimal(1000), en=decimal.Decimal(1010)),
@@ -465,8 +476,7 @@ class TestVizProfiler(unittest.TestCase):
             ProfileDeviceEvent(device='NV:SDMA:0', tdiff=decimal.Decimal(-1000))]
     j = load_profile(prof)
     event = j['layout']['NV:SDMA:0']['events'][0]
-    gbs = sz/(dur*1e-6)*1e-9
-    self.assertEqual(event['fmt'], f"{gbs:.0f} GB/s\n{sz/1e6:.0f} MB")
+    self.assertEqual(event['fmt'], {"B/s": sz/(dur*1e-6), "B": sz})
 
   def test_graph(self):
     prof = [ProfileDeviceEvent(device='NV', tdiff=decimal.Decimal(-1000)),
@@ -507,8 +517,7 @@ class TestVizProfiler(unittest.TestCase):
 
     j = load_profile(prof)
     sdma_events = j['layout']['NV:1:SDMA:0']['events']
-    gbs = sz/(dur*1e-6)*1e-9
-    self.assertEqual(sdma_events[0]["fmt"], f"{gbs:.0f} GB/s\n{sz/1e6:.0f} MB")
+    self.assertEqual(sdma_events[0]["fmt"], {"B/s": sz/(dur*1e-6), "B": sz})
 
   def test_block_ordering(self):
     prof = [ProfileDeviceEvent(device='NV', tdiff=decimal.Decimal(-1000)),
@@ -572,7 +581,7 @@ class TestVizProfiler(unittest.TestCase):
     profile_ret = load_profile(cpu_events)
     e = profile_ret["layout"]["CUSTOM"]["events"][0]
     self.assertEqual(e["name"], "test_fxn")
-    runtime_trace = json.loads(e["fmt"].replace("TB:", ""))
+    runtime_trace = e["fmt"]["tb"]
     assert any(fxn.__code__.co_filename == f and fxn.__code__.co_firstlineno+1 == l for f,l,*_ in runtime_trace), str(runtime_trace)
 
   # can pack up to 1hr 11 min of trace events
@@ -707,7 +716,6 @@ class TestVizMemoryLayout(unittest.TestCase):
     self.assertEqual(len(programs), len(set(users)), n)
 
 from tinygrad.uop.ops import KernelInfo
-from tinygrad.viz.serve import amdgpu_cfg
 from tinygrad.renderer.amd.dsl import s
 from tinygrad.runtime.autogen.amd.rdna3.ins import (s_add_u32, s_branch, s_cbranch_execz, s_cbranch_scc0, s_cbranch_scc1, s_cmp_eq_i32,
                                                     s_cmp_eq_u64, s_code_end, s_endpgm, s_mov_b32, s_nop)
@@ -723,13 +731,16 @@ class TestCfg(unittest.TestCase):
       gidx = UOp.special(1, "gidx0")
       sink = UOp.sink(out.base, lidx, gidx, arg=KernelInfo(name=name))
       return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="NULL"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
-    with Context(DEV=f"NULL::{self.arch}"):
-      out = Tensor.custom_kernel(Tensor.empty(1), fxn=fxn)[0]
-      runner = get_runner(out.device, out.schedule_linear().src[-1].src[0])
-      return amdgpu_cfg(runner.prg.src[4].arg, self.arch)
+    with save_viz() as viz:
+      with Context(DEV=f"NULL::{self.arch}"):
+        out = Tensor.custom_kernel(Tensor.empty(1), fxn=fxn)[0]
+        _ = to_program(out.schedule_linear().src[-1].src[0], Device[out.device].renderer)
+    codegen_rewrites = next(s for s in viz.list_items() if s["name"] == name)
+    disasm = next(s for s in codegen_rewrites["steps"] if s["name"] == "View Disassembly")
+    return get_render(viz.data, disasm["query"])
 
   def test_simple(self):
-    k = Kernel(arch=self.arch)
+    k = Kernel()
     k.label("entry")
     k.emit(s_branch(), target="bb1")
     k.label("bb1")
@@ -739,7 +750,7 @@ class TestCfg(unittest.TestCase):
     self.assertEqual(len(cfg["blocks"]), 2)
 
   def test_diamond(self):
-    k = Kernel(arch=self.arch)
+    k = Kernel()
     k.label("entry")
     k.emit(s_mov_b32(s[0], 0))
     k.emit(s_mov_b32(s[1], 0))
@@ -773,7 +784,7 @@ class TestCfg(unittest.TestCase):
       assert st.startswith("s_code_end") and st.endswith("x)"), st
 
   def test_loop(self):
-    k = Kernel(arch=self.arch)
+    k = Kernel()
     k.label("entry")
     k.emit(s_mov_b32(s[1], 4))
     k.label("loop")
@@ -785,7 +796,7 @@ class TestCfg(unittest.TestCase):
     self.get_cfg("simple_loop", k)
 
   def test_loop_branch(self):
-    k = Kernel(arch=self.arch)
+    k = Kernel()
     k.label("entry")
     k.emit(s_mov_b32(s[1], 4))
     k.label("loop")
@@ -803,7 +814,7 @@ class TestCfg(unittest.TestCase):
     self.get_cfg("loop_if", k)
 
   def test_loop_break(self):
-    k = Kernel(arch=self.arch)
+    k = Kernel()
     k.label("entry")
     k.emit(s_mov_b32(s[1], 8))
     k.label("loop")
@@ -818,7 +829,7 @@ class TestCfg(unittest.TestCase):
     self.get_cfg("loop_break", k)
 
   def test_switch(self):
-    k = Kernel(arch=self.arch)
+    k = Kernel()
     k.label("entry")
     k.emit(s_cmp_eq_i32(s[0], 0))
     k.emit(s_cbranch_scc1(), target="case0")
@@ -840,7 +851,7 @@ class TestCfg(unittest.TestCase):
     self.get_cfg("switch_case", k)
 
   def test_ping_pong(self):
-    k = Kernel(arch=self.arch)
+    k = Kernel()
     k.label("entry")
     k.emit(s_cmp_eq_i32(s[0], 0))
     k.emit(s_cbranch_scc1(), target="ping")
@@ -859,7 +870,7 @@ class TestCfg(unittest.TestCase):
 
   def test_colored_blocks(self):
     N = 10
-    k = Kernel(arch=self.arch)
+    k = Kernel()
     k.label("entry")
     k.emit(s_branch(), target="init0")
     for i in range(N):
@@ -879,7 +890,7 @@ class TestCfg(unittest.TestCase):
     self.get_cfg("test_colored_blocks", k)
 
   def test_jump_back_to_end(self):
-    k = Kernel(arch=self.arch)
+    k = Kernel()
     k.label("entry")
     k.emit(s_mov_b32(s[1], 2))
     k.emit(s_cbranch_execz(), target="loop")
@@ -900,46 +911,62 @@ def run_cli(*cli_args) -> str:
     main(args)
   return buf.getvalue().strip()
 
+def call_cli(fxn, *cli_args, debug=2) -> str:
+  with save_viz() as viz:
+    fxn()
+  with tempfile.TemporaryDirectory() as tmpdir:
+    (r:=Path(tmpdir)/"rewrites.pkl").write_bytes(pickle.dumps(viz.data.trace))
+    (p:=Path(tmpdir)/"profile.pkl").write_bytes(pickle.dumps(cpu_events))
+    with Context(DEBUG=debug, NO_COLOR=1):
+      stdout = run_cli("--rewrites-path", str(r), "--profile-path", str(p), *cli_args)
+  return stdout
+
 class TestCLI(unittest.TestCase):
-  def test_simple(self):
-    a = Tensor.empty(1, device="NULL")+2.0
-    empty_counter = itertools.count(0)
-    def custom_empty_prg(B:UOp, A:UOp) -> UOp:
-      sink = UOp(Ops.SINK, arg=KernelInfo(name=f"custom_empty_n{next(empty_counter)}"))
-      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=a.device), UOp(Ops.LINEAR, src=(sink,))))
-    def custom_empty_src(B:UOp, A:UOp) -> UOp:
-      sink = UOp(Ops.SINK, arg=KernelInfo(name=f"custom_empty_n{next(empty_counter)}"))
-      src = "void custom_empty_src() { 0; }"
-      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=a.device), UOp(Ops.LINEAR, src=(sink,)), UOp(Ops.SOURCE, arg=src)))
-    b = Tensor.custom_kernel(Tensor.empty_like(a), a, fxn=custom_empty_prg)[0]
-    c = Tensor.custom_kernel(Tensor.empty_like(a), a, fxn=custom_empty_prg)[0]
-    d = Tensor.custom_kernel(Tensor.empty_like(a), a, fxn=custom_empty_src)[0]
-    with save_viz() as viz:
-      b.realize()
+  def test_reconstruct_debug(self):
+    def fxn():
+      Tensor.empty(1, device="NULL").add(2.0).realize()
       profile_marker("marker @ 1")
-      c.realize()
-      d.realize()
-    # save trace to disk for CLI to consume it
-    with tempfile.TemporaryDirectory() as tmpdir:
-      (r:=Path(tmpdir)/"rewrites.pkl").write_bytes(pickle.dumps(viz.data.trace))
-      (p:=Path(tmpdir)/"profile.pkl").write_bytes(pickle.dumps(cpu_events))
-      # reconstruct DEBUG=4 output and see all markers.
-      with Context(DEBUG=4):
-        kernels = run_cli("--rewrites-path", str(r), "--profile-path", str(p), "-s", "NULL")
-      self.assertIn("void custom_empty_n0", kernels)
-      self.assertIn("marker @ 1", kernels)
-      self.assertIn("void custom_empty_n1", kernels)
-      self.assertIn("void custom_empty_src", kernels)
-      self.assertIn("E", kernels)
-      self.assertIn("UOp.const", kernels)
-      # get the top slowest functions across all devices
-      with Context(DEBUG=2):
-        times = run_cli("--rewrites-path", str(r), "--profile-path", str(p), "-s", "ALL", "--top", "-1")
-      self.assertIn("TINY", times)
-      self.assertIn("NULL", times)
-      with Context(DEBUG=3):
-        json_lines = run_cli("--rewrites-path", str(r), "--profile-path", str(p), "-s", "ALL", "--json")
-      for line in json_lines.split("\n"): _ = json.loads(line)
+      Tensor.empty(1, device="NULL").add(3.0).realize()
+    out = call_cli(fxn, "-s", "NULL", debug=4)
+    self.assertIn("void E", out)
+    self.assertIn("marker @ 1", out)
+
+  def test_aggregate(self):
+    N, CNT = 1024, 5
+    def fxn():
+      for _ in range(CNT):
+        (Tensor.empty(N, N, device="NULL")@Tensor.empty(N, N, device="NULL")).realize()
+      for _ in range(CNT):
+        (Tensor.empty(N, N, device="NULL").assign(Tensor.empty(N, N, device="NULL"))).realize()
+    kernels = [json.loads(line) for line in call_cli(fxn, "-s", "NULL", "-t", "--json").splitlines()]
+    self.assertEqual(len(kernels), 2)
+    gemm_summary = [s for s in kernels if s["name"].startswith("r_")][0]
+    copy_summary = [s for s in kernels if s["name"].startswith("E_")][0]
+    self.assertEqual(gemm_summary["count"], CNT)
+    self.assertEqual(copy_summary["count"], CNT)
+
+  def test_flops(self):
+    test_n = [(8, 16), (16, 32), (32, 64)]
+    def fxn():
+      @TinyJit
+      def f(a, b): return (a@a.T), (b@b.T)
+      a = Tensor.empty(64, 64, device="NULL")
+      b = Tensor.empty(64, 64, device="NULL")
+      for i_val, j_val in test_n:
+        i = Variable("i", 1, 64).bind(i_val)
+        j = Variable("j", 1, 64).bind(j_val)
+        Tensor.realize(*f(a[:i], b[:j]))
+    out = [json.loads(line) for line in call_cli(fxn, "-s", "NULL", "--json").splitlines()]
+    self.assertEqual(len(out), 3*2)
+    # flops increases as N gets larger
+    gflops = [row["fmt"]["FLOPS"] for row in out]
+    self.assertGreater(gflops[4], gflops[2])
+    self.assertGreater(gflops[5], gflops[3])
+    # aggregate flops
+    out = [json.loads(line) for line in call_cli(fxn, "-s", "NULL", "-t", "--json").splitlines()]
+    self.assertEqual(len(out), 2)
+    agg_gflops = [row["fmt"]["FLOPS"] for row in out]
+    assert all(min(gflops) < v < max(gflops) for v in agg_gflops), f"{agg_gflops}"
 
 if __name__ == "__main__":
   unittest.main()

test/opt/test_tensor_cores.py

@@ -3,12 +3,12 @@ import unittest
 
 from tinygrad import Device, Tensor, dtypes
 from tinygrad.tensor import _to_np_dtype
-from tinygrad.uop.ops import Ops
+from tinygrad.uop.ops import Ops, UOp, buffers
 from tinygrad.dtype import DType
-from tinygrad.device import is_dtype_supported
+from tinygrad.device import Buffer, is_dtype_supported
 from tinygrad.helpers import DEV, Context
 from test.helpers import slow, replace_opts
-from tinygrad.engine.realize import CompiledRunner
+from tinygrad.engine.realize import run_linear
 from tinygrad.codegen import to_program
 from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
 from tinygrad.codegen.opt.tc import amd_cdna_1616128
@@ -20,6 +20,11 @@ from test.backend.test_linearizer import helper_realized_ast, helper_linearizer_
 
 AMX = "AMX" in DEV.arch
 
+def run_program(prg:UOp, bufs:list[Buffer]):
+  buf_uops = [UOp.new_buffer(b.device, b.size, b.dtype) for b in bufs]
+  for u,b in zip(buf_uops, bufs): buffers[u] = b
+  run_linear(UOp(Ops.LINEAR, src=(prg.call(*buf_uops),)))
+
 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)
@@ -47,11 +52,11 @@ 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))]
-  pu = to_program(replace_opts(realized_ast, opts), Device[Device.DEFAULT].renderer)
+  ast = replace_opts(realized_ast, opts)
+  pu = to_program(ast, Device[Device.DEFAULT].renderer)
   if use_tensor_cores == 1: assert len([uop for uop in pu.src[2].src if uop.op is Ops.WMMA]) > 0, "wmma not triggered"
   assert len([x for x in pu.src[0].arg.applied_opts if x.op is OptOps.TC]) == 1, "tensor core opt not included"
-  prg = CompiledRunner(pu, Device.DEFAULT)
-  prg.exec(bufs)
+  run_program(ast, bufs)
   if dtype_in == dtypes.half: tc_atol, tc_rtol = 1e-2, 1e-3
   elif dtype_in == dtypes.bfloat16: tc_atol, tc_rtol = (1e-1, 2e-2) if dtype_out == dtypes.bfloat16 else (1e-2, 1e-2)
   else: tc_atol, tc_rtol = 5e-3, 1e-4
@@ -145,15 +150,15 @@ class TestTensorCores(unittest.TestCase):
         c = a.conv2d(b, padding=1, dtype=tc.dtype_out)
         realized_ast, real_bufs = helper_realized_ast(c)
 
-        program = to_program(replace_opts(realized_ast, [Opt(OptOps.TC, axis, (-1, 2, 1))]), Device[Device.DEFAULT].renderer)
+        ast = replace_opts(realized_ast, [Opt(OptOps.TC, axis, (-1, 2, 1))])
+        program = to_program(ast, Device[Device.DEFAULT].renderer)
         assert len([uop for uop in tuple(program.src[2].src) if uop.op is Ops.WMMA]) > 0, "tensor core not triggered"
         assert len([x for x in program.src[0].arg.applied_opts if x.op is OptOps.TC]) == 1, "tensor core opt not included"
 
-        prg = CompiledRunner(program, Device.DEFAULT)
         # TODO: support this even if numpy doesn't
         if _to_np_dtype(real_bufs[0].dtype) is None: continue
         real_bufs[0].copyin(np.zeros((real_bufs[0].size, ), dtype=_to_np_dtype(real_bufs[0].dtype)).data) # Zero to check that all values are filled
-        prg.exec(real_bufs)
+        run_program(ast, real_bufs)
         result = np.frombuffer(real_bufs[0].as_memoryview(), _to_np_dtype(real_bufs[0].dtype))
 
         # ensure the results for each choice of axis matches

test/unit/test_assign.py

@@ -903,8 +903,8 @@ class TestAssignToUnrealizedView(unittest.TestCase):
 
   def test_reduce(self):
     a = Tensor([[1,2],[3,4]]).contiguous().realize()
-    r = a.sum(axis=0)  # unrealized REDUCE_AXIS
-    self.assertIs(r.uop.base.op, Ops.REDUCE_AXIS)
+    r = a.sum(axis=0)  # unrealized REDUCE
+    self.assertIs(r.uop.base.op, Ops.REDUCE)
     r[:1].assign(Tensor([99]).realize())
     try:
       self.assertEqual(r.tolist(), [99,6])