Merge branch 'master' into move_gates_to_load_store

better order

.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, fw"
+        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

@@ -628,7 +628,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 +636,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

.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/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

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

examples/mlperf/model_train.py

@@ -1419,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)
@@ -1446,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)

examples/mlperf/models/flat_llama.py

@@ -158,14 +158,14 @@ 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, grad_amax_state=grad_amax_xo)
     new_amaxs.extend(ret[:1])

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

@@ -15,7 +15,7 @@ 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}

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

@@ -15,7 +15,7 @@ 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}

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/gemm/cdna_asm_gemm.py

@@ -2628,21 +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, *extra: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.
-  # extra is unused fwd inputs (e.g. grad_amax_state) plumbed through so the bwd can read them via kernel.src.
+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)))
 
@@ -2699,8 +2702,7 @@ 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) optionally plus grad_amax_state (6 total); plain gemm has 3
-  if len(inputs) >= 5:
+  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)
@@ -2720,8 +2722,7 @@ def custom_gemm_bw(gradient:UOp, kernel:UOp):
     # 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)
+    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)
     # 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,)
@@ -2774,11 +2775,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()
+      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 []
-      out = Tensor.custom_kernel(out, a, b.T, xs, ws, *extra, fxn=functools.partial(custom_hk_fp8_gemm, dname=dname), grad_fxn=custom_gemm_bw)[0]
+      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/llama_kernels/__init__.py

@@ -34,13 +34,12 @@ def dname_of(device) -> str:
   return device.split(":")[0] if isinstance(device, str) else device
 
 def alloc_like(shape, dtype, device, axis=None) -> Tensor:
-  if isinstance(device, tuple):
-    if axis is None: return Tensor(Tensor.invalids(*shape, dtype=dtype, device=device).uop.multi(0), device=device)
+  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) -> Tensor:
-  if isinstance(device, tuple):
+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)
 

extra/llama_kernels/cast_amax/__init__.py

@@ -41,10 +41,9 @@ def _fused_quantize_bwd_w13(gradient:UOp, kernel:UOp):
   _, _, xw13, amax_state, grad_amax_state = kernel.src[1:]
   device = xw13.device
   axis = xw13.axis if isinstance(device, tuple) else None
-  if isinstance(device, tuple): assert axis in (0, 1), f"unsupported sharding axis={axis}"
   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)
+  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(
@@ -66,9 +65,8 @@ def fused_quantize_fp8_w13(xw13:Tensor, amax_state:Tensor, fp8_dtype, grad_amax_
   assert H2 % 2 == 0, f"w13 last-axis must be even, got {H2}"
   HIDDEN = H2 // 2
   axis = xw13.uop.axis if isinstance(xw13.device, tuple) else None
-  if isinstance(xw13.device, tuple): assert axis in (0, 1), f"unsupported sharding axis={axis}"
   fp8_out  = alloc_like((MBS, SEQ, HIDDEN), fp8_dtype,      xw13.device, axis)
-  amax_buf = alloc_local((NUM_WG,),         dtypes.float32, xw13.device)
+  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)

extra/llama_kernels/fused_rmsnorm_mul_quantize_fp8/__init__.py

@@ -63,7 +63,7 @@ def _bwd_common(fp8_grad_u, h_grad_u, x_u, x_normed_u, rrms_u, weight_u, amax_st
   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)
+  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:
@@ -119,11 +119,11 @@ def fused_rmsnorm_mul_quantize_fp8(x:Tensor, weight:Tensor, amax_state:Tensor, e
   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 (0, 1), f"unsupported sharding axis={axis}"
+  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)
+  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)
@@ -139,12 +139,12 @@ def fused_add_rmsnorm_mul_quantize_fp8(x:Tensor, residual:Tensor, weight:Tensor,
   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 (0, 1), f"unsupported sharding axis={axis}"
+  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)
+  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,

extra/llama_kernels/quantize_fp8_delayed/__init__.py

@@ -49,7 +49,7 @@ def quantize_fp8_delayed(x:Tensor, amax_state:Tensor, fp8_dtype=dtypes.fp8e4m3)
   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)
+  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)

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});

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]

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_custom_kernel.py

@@ -324,6 +324,20 @@ class TestCustomKernel(unittest.TestCase):
   @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

@@ -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):

test/backend/test_ops.py

@@ -636,6 +636,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)

test/backend/test_profiler.py

@@ -51,11 +51,11 @@ class TestProfiler(unittest.TestCase):
     TestProfiler.runtime = get_runtime(TestProfiler.d0.device, TestProfiler.prg)
     TestProfiler.b.uop.buffer.allocate()
 
-  def test_profile_kernel_run(self):
+  def test_profile_kernel_run(self, wait=False):
     runner_name = TestProfiler.runtime.name
     with helper_collect_profile(TestProfiler.d0) as profile:
       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)
+      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)]
@@ -63,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()
 

test/backend/test_randomness.py

@@ -361,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))
 
@@ -372,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/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_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/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)
@@ -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

@@ -375,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)
@@ -388,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))]
 
@@ -398,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]:
@@ -516,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."""
@@ -540,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).
@@ -713,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))
@@ -1847,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)
@@ -2005,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())
 

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):
@@ -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_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),
@@ -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)
@@ -216,18 +221,18 @@ 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]
+    load = simplify_image_idx(load.sink()).src[0]
+    off = load.src[0].src[1]
+    idx = off.get_idx()
     self.assertEqual(idx.op, Ops.STACK)
     self.assertEqual(len(idx.src), 2)
     idx0, idx1 = idx.src[0], idx.src[1]
     check_uop_against_string(self, idx0, sidx0)
     check_uop_against_string(self, idx1, sidx1)
     if svalid is not None:
-      check_uop_against_string(self, load.src[0].src[2], svalid)
+      check_uop_against_string(self, off.get_valid(), svalid)
     else:
-      self.assertEqual(len(load.src[0].src), 2, "svalid is None but load still has a valid")
+      self.assertEqual(off.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 +292,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 +452,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,7 +513,7 @@ 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):
@@ -528,7 +533,7 @@ class TestDropTrueGate(unittest.TestCase):
 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):

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_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
@@ -741,7 +746,7 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable((25*a+3)%10 + ((25*a+3)//10)*10, 3, 253, "((a*25)+3)")
 
   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 +760,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 +772,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 +832,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 +1000,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 +1081,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 +1306,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 +1316,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 +1356,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):

test/null/test_uops.py

@@ -46,6 +46,20 @@ class TestExecALU(unittest.TestCase):
     self.assertEqual(exec_alu(Ops.IDIV, dtypes.int8, (7, -3)), -2)
     self.assertEqual(exec_alu(Ops.IDIV, 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))
 
@@ -163,6 +177,30 @@ class TestFastIdiv(unittest.TestCase):
       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")
 
+  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.MOD, 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.IDIV, 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)

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
@@ -320,7 +320,7 @@ class TestVizGC(unittest.TestCase):
 
 # VIZ integrates with other parts of tinygrad
 
-from tinygrad import Tensor, Device
+from tinygrad import Tensor, Device, TinyJit, Variable
 
 class TestVizIntegration(unittest.TestCase):
   # codegen supports rendering of code blocks
@@ -408,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
@@ -421,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)),
@@ -464,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)),
@@ -506,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)),
@@ -571,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
@@ -901,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/unit/test_function.py

@@ -449,8 +449,8 @@ class TestFunctionTuple(unittest.TestCase):
       store_d = D[i].store(A[i] * 3.0)
       return UOp.group(store_c, store_d).end(i).sink(arg=KernelInfo(name="my_kernel"))
 
-    def my_grad(d_combined:UOp, call:UOp):
-      return (None, None, Tensor(d_combined).uop)
+    def my_grad(d_c:UOp, d_d:UOp, call:UOp):
+      return (None, None, (Tensor(d_c) + Tensor(d_d)).uop)
 
     @function(precompile=True, precompile_backward=True)
     def f(a:Tensor):

test/unit/test_gradient.py

@@ -2,6 +2,7 @@ import unittest
 import numpy as np
 from tinygrad import Tensor
 from tinygrad.dtype import dtypes
+from tinygrad.uop.ops import UOp, KernelInfo
 
 class TestTensorGradient(unittest.TestCase):
   def test_example(self):
@@ -76,6 +77,71 @@ class TestTensorGradient(unittest.TestCase):
     x = Tensor.randn(4, 4)
     np.testing.assert_allclose(x.pad(((1,0),(0,0))).gradient(x, gradient=g2)[0].numpy(), np.zeros((4, 4)))
 
+class TestMultiOutputGradient(unittest.TestCase):
+  @staticmethod
+  def addmul_kernel(C:UOp, D:UOp, A:UOp, B:UOp) -> UOp:
+    C, D, A, B = C.flatten(), D.flatten(), A.flatten(), B.flatten()
+    i = UOp.range(C.numel(), 0)
+    store_c = C[i].store(A[i] + B[i])
+    store_d = D[i].store(A[i] * B[i])
+    return UOp.group(store_c, store_d).end(i).sink(arg=KernelInfo(name="addmul")).simplify()
+  @staticmethod
+  def backward_addmul(grad_c, grad_d, call):
+    _c, _d, a, b = call.src[1:]
+    grad_a = (Tensor(grad_c) + Tensor(grad_d) * Tensor(b)).uop
+    grad_b = (Tensor(grad_c) + Tensor(grad_d) * Tensor(a)).uop
+    return (None, None, grad_a, grad_b)
+
+  def test_custom_kernel_multi_output_backward(self):
+    a_np, b_np = np.random.randn(4, 4).astype(np.float32), np.random.randn(4, 4).astype(np.float32)
+    a_ref, b_ref = Tensor(a_np, requires_grad=True), Tensor(b_np, requires_grad=True)
+    ((a_ref + b_ref).sum() + (a_ref * b_ref).sum()).backward()
+
+    a, b = Tensor(a_np, requires_grad=True), Tensor(b_np, requires_grad=True)
+    Tensor.realize(a, b)
+    c, d, _, _ = Tensor.custom_kernel(Tensor.empty(4, 4), Tensor.empty(4, 4), a, b, fxn=self.addmul_kernel, grad_fxn=self.backward_addmul)
+    (c.sum() + d.sum()).backward()
+    np.testing.assert_allclose(a.grad.numpy(), a_ref.grad.numpy(), rtol=1e-5)
+    np.testing.assert_allclose(b.grad.numpy(), b_ref.grad.numpy(), rtol=1e-5)
+
+  def test_custom_kernel_multi_output_backward_interacting(self):
+    a_np, b_np = np.random.randn(4, 4).astype(np.float32), np.random.randn(4, 4).astype(np.float32)
+    a_ref, b_ref = Tensor(a_np, requires_grad=True), Tensor(b_np, requires_grad=True)
+    ((a_ref + b_ref) * (a_ref * b_ref)).sum().backward()
+
+    a, b = Tensor(a_np, requires_grad=True), Tensor(b_np, requires_grad=True)
+    Tensor.realize(a, b)
+    c, d, _, _ = Tensor.custom_kernel(Tensor.empty(4, 4), Tensor.empty(4, 4), a, b, fxn=self.addmul_kernel, grad_fxn=self.backward_addmul)
+    (c * d).sum().backward()
+    np.testing.assert_allclose(a.grad.numpy(), a_ref.grad.numpy(), rtol=1e-5)
+    np.testing.assert_allclose(b.grad.numpy(), b_ref.grad.numpy(), rtol=1e-5)
+
+  def test_custom_kernel_three_output_backward(self):
+    def addmulsub_kernel(C:UOp, D:UOp, E:UOp, A:UOp, B:UOp) -> UOp:
+      C, D, E, A, B = C.flatten(), D.flatten(), E.flatten(), A.flatten(), B.flatten()
+      i = UOp.range(C.numel(), 0)
+      store_c = C[i].store(A[i] + B[i])
+      store_d = D[i].store(A[i] * B[i])
+      store_e = E[i].store(A[i] - B[i])
+      return UOp.group(store_c, store_d, store_e).end(i).sink(arg=KernelInfo(name="addmulsub")).simplify()
+    def backward_addmulsub(grad_c, grad_d, grad_e, call):
+      _c, _d, _e, a, b = call.src[1:]
+      grad_a = (Tensor(grad_c) + Tensor(grad_d) * Tensor(b) + Tensor(grad_e)).uop
+      grad_b = (Tensor(grad_c) + Tensor(grad_d) * Tensor(a) - Tensor(grad_e)).uop
+      return (None, None, None, grad_a, grad_b)
+
+    a_np, b_np = np.random.randn(4, 4).astype(np.float32), np.random.randn(4, 4).astype(np.float32)
+    a_ref, b_ref = Tensor(a_np, requires_grad=True), Tensor(b_np, requires_grad=True)
+    ((a_ref + b_ref).sum() + (a_ref * b_ref).sum() + (a_ref - b_ref).sum()).backward()
+
+    a, b = Tensor(a_np, requires_grad=True), Tensor(b_np, requires_grad=True)
+    Tensor.realize(a, b)
+    c, d, e, _, _ = Tensor.custom_kernel(Tensor.empty(4, 4), Tensor.empty(4, 4), Tensor.empty(4, 4), a, b,
+                                          fxn=addmulsub_kernel, grad_fxn=backward_addmulsub)
+    (c.sum() + d.sum() + e.sum()).backward()
+    np.testing.assert_allclose(a.grad.numpy(), a_ref.grad.numpy(), rtol=1e-5)
+    np.testing.assert_allclose(b.grad.numpy(), b_ref.grad.numpy(), rtol=1e-5)
+
 class TestViewGradient(unittest.TestCase):
   def test_expand(self):
     x = Tensor.randn(5,2)

test/unit/test_hcq_graph.py

@@ -11,7 +11,7 @@ from test.mockgpu.usb import MockUSB
 @unittest.skipUnless(issubclass(type(Device[Device.DEFAULT]), HCQCompiled), "HCQ device required to run")
 class TestHCQUnit(unittest.TestCase):
   @unittest.skipIf(Device.DEFAULT == "CPU", "requires non-CPU HCQ device")
-  def test_supports_exec_item(self):
+  def test_supports_uop(self):
     d0, cpu_dev = Device[Device.DEFAULT], Device["CPU"]
 
     @TinyJit
@@ -20,23 +20,23 @@ class TestHCQUnit(unittest.TestCase):
     inp, inp_cpu = Tensor.randn(10, 10, device=Device.DEFAULT).realize(), Tensor.randn(10, 10, device="CPU").realize()
     for _ in range(5): f(inp, inp_cpu)
 
-    # construct minimal CALL UOps for supports_exec_item (graphs only see PROGRAMs after compile_linear)
+    # construct minimal CALL UOps for supports_uop (graphs only see PROGRAMs after compile_linear)
     gpu_call = UOp(Ops.PROGRAM).call(UOp.new_buffer(Device.DEFAULT, 1, dtypes.float))
     cpu_call = UOp(Ops.PROGRAM).call(UOp.new_buffer("CPU", 1, dtypes.float))
     gpu_devs = [d0]
 
     # local MMIO: GPU works alone and with CPU in batch (cpu_support=True)
-    assert HCQGraph.supports_exec_item(gpu_devs, gpu_call) is True
-    assert HCQGraph.supports_exec_item(gpu_devs, cpu_call) is True
-    assert HCQGraph.supports_exec_item(gpu_devs + [cpu_dev], gpu_call) is True
+    assert HCQGraph.supports_uop(gpu_devs, gpu_call) is True
+    assert HCQGraph.supports_uop(gpu_devs, cpu_call) is True
+    assert HCQGraph.supports_uop(gpu_devs + [cpu_dev], gpu_call) is True
 
     # USB MMIO: GPU-only still works, but CPU batching must be rejected (cpu_support=False)
     orig_view = d0.timeline_signal.base_buf.view
     try:
       d0.timeline_signal.base_buf.view = USBMMIOInterface(MockUSB(bytearray(256)), 0, 16, fmt='B')
-      assert HCQGraph.supports_exec_item(gpu_devs, gpu_call) is True
-      assert HCQGraph.supports_exec_item(gpu_devs, cpu_call) is False
-      assert HCQGraph.supports_exec_item(gpu_devs + [cpu_dev], gpu_call) is False
+      assert HCQGraph.supports_uop(gpu_devs, gpu_call) is True
+      assert HCQGraph.supports_uop(gpu_devs, cpu_call) is False
+      assert HCQGraph.supports_uop(gpu_devs + [cpu_dev], gpu_call) is False
     finally:
       d0.timeline_signal.base_buf.view = orig_view
 

test/unit/test_metal_graph.py

@@ -27,18 +27,18 @@ class TestMetalGraph(unittest.TestCase):
     c.src = (MagicMock(op=Ops.PROGRAM),) + tuple(bufs)
     return c
 
-  def test_supports_exec_item_normal_offset(self):
-    assert self.MetalGraph.supports_exec_item([self.dev], self.call(self.metal_buf(0), self.metal_buf(100), self.metal_buf(0xFFFFFFFF))) is True
+  def test_supports_uop_normal_offset(self):
+    assert self.MetalGraph.supports_uop([self.dev], self.call(self.metal_buf(0), self.metal_buf(100), self.metal_buf(0xFFFFFFFF))) is True
 
-  def test_supports_exec_item_overflow_offset(self):
-    assert self.MetalGraph.supports_exec_item([self.dev], self.call(self.metal_buf(0), self.metal_buf(0x100000000))) is False
+  def test_supports_uop_overflow_offset(self):
+    assert self.MetalGraph.supports_uop([self.dev], self.call(self.metal_buf(0), self.metal_buf(0x100000000))) is False
 
-  def test_supports_exec_item_nonmetal_buf(self):
+  def test_supports_uop_nonmetal_buf(self):
     # non-BUFFER_VIEW ops should not be checked for offset
     buf = MagicMock()
     buf.op = Ops.BUFFER
     buf.device = Device.DEFAULT
-    self.MetalGraph.supports_exec_item([self.dev], self.call(buf))
+    self.MetalGraph.supports_uop([self.dev], self.call(buf))
 
 if __name__ == "__main__":
   unittest.main()

tinygrad/codegen/__init__.py

@@ -2,8 +2,9 @@ from typing import cast
 from dataclasses import replace
 import itertools
 from tinygrad.helpers import DISABLE_FAST_IDIV, DEVECTORIZE, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC
-from tinygrad.helpers import ALLOW_TF32, TracingKey, Context, Target, panic
-from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, ProgramInfo, pyrender
+from tinygrad.helpers import ALLOW_TF32, TracingKey, Context, panic
+from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, ProgramInfo
+from tinygrad.uop.render import pyrender
 from tinygrad.uop.spec import type_verify, program_spec, kernel_spec
 from tinygrad.renderer import Renderer, Estimates
 from tinygrad.dtype import dtypes
@@ -15,20 +16,19 @@ from tinygrad.uop.decompositions import get_late_rewrite_patterns, get_transcend
 from tinygrad.codegen.late.expander import expander, pm_pre_expander, pm_group_for_reduce
 from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_indexing, devectorize, pm_reduce, \
   ReduceContext, correct_load_store, pm_render, pm_add_loads, pm_make_images
+from tinygrad.codegen.late.gater import pm_move_gates_from_index
 from tinygrad.codegen.opt.postrange import apply_opts
 from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_split_ranges, pm_load_collapse
 from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen, pm_mops, pm_syntactic_sugar, pm_store_ranges
 from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow, linearize
 
-def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True, beam:int=0) -> UOp:
-  if ren is None: ren = Renderer(Target())
-
-  if VIZ: graph_rewrite(sink, PatternMatcher([]), name="View Base AST")
-  if DEBUG >= 5: print(pyrender(sink))
-  if SPEC: type_verify(sink, kernel_spec)
+def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
+  if VIZ: graph_rewrite(ast, PatternMatcher([]), name="View Base AST")
+  if DEBUG >= 5: print(pyrender(ast))
+  if SPEC: type_verify(ast, kernel_spec)
 
   # preprocess
-  sink = graph_rewrite(sink, pm_mops+pm_syntactic_sugar+pm_store_ranges, ctx=itertools.count(1000), name="early movement ops", bottom_up=True)
+  sink = graph_rewrite(ast, pm_mops+pm_syntactic_sugar+pm_store_ranges, ctx=itertools.count(1000), name="early movement ops", bottom_up=True)
 
   # first we optimize
   if optimize:
@@ -45,7 +45,7 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True, b
     sink = graph_rewrite(sink, pm_flatten_range+pm_simplify_ranges, ctx={}, name="simplify ranges")
 
     # do postrange optimization, BEAM or hand_coded_optimizations
-    sink = apply_opts(sink, ren, beam=beam)
+    sink = apply_opts(sink, ren, beam=ast.arg.beam)
 
   # ** expander (expand_rewrite) **
   sink = graph_rewrite(sink, sym+pm_move_where_on_load, name="postopt symbolic")
@@ -77,8 +77,13 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True, b
   else: pm_devectorize = sym+load_store_folding+correct_load_store+load_store_indexing
   if DEVECTORIZE >= 0: sink = graph_rewrite(sink, pm_devectorize, ctx=ren, name="devectorize")
 
-  # lower the index dtype to a concrete int
+  # lower the index dtype to a concrete int. this needs to happen while gates are still present
   sink = graph_rewrite(sink, pm_lower_index_dtype+load_store_indexing+gep_pushing, name="lower all index dtypes")
+
+  # move the gates from index onto the loads and stores
+  sink = graph_rewrite(sink, pm_move_gates_from_index, name="move gates from index")
+
+  # a final symbolic
   sink = graph_rewrite(sink, symbolic, name="post index symbolic")
 
   # optional pre matcher
@@ -107,8 +112,8 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True, b
 pm_linearize_cleanups = PatternMatcher([
   # if statements are not allowed in the graph
   (UPat((Ops.IF, Ops.ENDIF)), lambda: panic(RuntimeError, "if not allowed in graph")),
-  # gated INDEX becomes IF-STORE-ENDIF. this is the only use of IF-ENDIF
-  (UPat(Ops.STORE, name="u", src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat(name="gate", dtype=dtypes.bool))).or_casted(), UPat())),
+  # gated STORE becomes IF-STORE-ENDIF. this is the only use of IF-ENDIF
+  (UPat(Ops.STORE, name="u", src=(UPat(Ops.INDEX).or_casted(), UPat(), UPat(name="gate", dtype=dtypes.bool))),
    lambda u, gate: (u, [mif:=UOp(Ops.IF, src=(gate, u.src[0])), u, UOp(Ops.ENDIF, src=(mif,))]))
 ])
 
@@ -170,7 +175,7 @@ def do_to_program(ast:UOp, renderer:Renderer) -> UOp:
   if ast.op is Ops.PROGRAM: prg = ast
   elif ast.op is Ops.SINK:
     assert isinstance(ast.arg, KernelInfo), "requires KernelInfo on arg to to_program"
-    full_sink = full_rewrite_to_sink(ast, renderer, optimize=ast.tag is None, beam=ast.arg.beam)
+    full_sink = full_rewrite_to_sink(ast, renderer, optimize=ast.tag is None)
     prg = UOp(Ops.PROGRAM, src=(full_sink, UOp(Ops.DEVICE, arg=renderer.target.device)), arg=ProgramInfo.from_sink(full_sink))
   else: raise RuntimeError(f"can't call to_program on {ast.op}")
   if not isinstance(prg.arg, ProgramInfo): prg = prg.replace(arg=ProgramInfo.from_sink(prg.src[0]))

tinygrad/codegen/late/devectorizer.py

@@ -37,6 +37,7 @@ def _drop_valid_stmts(valid:UOp, idx:UOp, height:int, width:int) -> list[UOp]:
   return drop_stmt
 
 def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None:
+  start_idx = start_idx.simplify()  # if you don't do this, uop_given_valid may simplify things and this might inf loop
   idx = uop_given_valid(valid, start_idx)
   if not isinstance(buf.dtype, ImageDType): return None if idx is start_idx else buf.index(idx.valid(valid), ptr=True)
 
@@ -116,22 +117,22 @@ def fold_expanded_index(midx:UOp):
   post_cat = UOp(Ops.PTRCAT, buf.ptrdtype.base.ptr(size=buf.ptrdtype.size, addrspace=buf.ptrdtype.addrspace).vec(global_offset), tuple(ret))
   return post_cat.gep(tuple(cast(list[int], idxs)))
 
-def cat_after_store(cat:UOp, data:UOp, sto:UOp):
+def cat_after_store(cat:UOp, data:UOp):
   # TODO: this is written in many places
   offset = 0
   ret: list[UOp] = []
   for s in cat.src:
-    ret.append(s.store(data.gep(tuple(range(offset, offset+s.dtype.count))), *sto.src[2:]))
+    ret.append(s.store(data.gep(tuple(range(offset, offset+s.dtype.count)))))
     offset += s.dtype.count
   return UOp.group(*ret)
 
-def gep_on_store(gep:UOp, st:UOp, sto:UOp):
+def gep_on_store(gep:UOp, st:UOp):
   # NOTE: we need to invert the gep here, but it may be an expanding gep
   # fake argsort. TODO: handle duplicates
   a = {}
   for i,x in enumerate(gep.arg): a[x] = i
   new_arg = tuple(x[1] for x in sorted(a.items()))
-  return gep.src[0].store(st.gep(new_arg), *sto.src[2:])
+  return gep.src[0].store(st.gep(new_arg))
 
 load_store_folding = PatternMatcher([
   (UPat(Ops.INDEX, src=(UPat(Ops.STACK, src=UPat(GroupOp.Defines).or_after(name="buf")), UPat.var("vec"))), expand_index),
@@ -140,12 +141,12 @@ load_store_folding = PatternMatcher([
   (UPat(Ops.LOAD, src=(UPat(Ops.GEP, name="gep"),), name="ld", allow_any_len=True),
    lambda gep, ld: ld.replace(dtype=ld.dtype.scalar().vec(gep.dtype.count), src=(gep.src[0],)+ld.src[1:]).gep(gep.arg)),
   # GEP on data of STORE
-  (UPat(Ops.STORE, src=(UPat(Ops.GEP, name="gep"), UPat.var("st")), name="sto"), gep_on_store),
+  (UPat(Ops.STORE, src=(UPat(Ops.GEP, name="gep"), UPat.var("st"))), gep_on_store),
   # put PTRCAT after LOAD
   (UPat(Ops.LOAD, src=(UPat(Ops.PTRCAT, name="cat"),), name="ld", allow_any_len=True),
    lambda cat,ld: UOp(Ops.VCAT, cat.dtype.base.vec(cat.dtype.vcount), tuple(ld.replace(dtype=x.dtype.base, src=(x,)+ld.src[1:]) for x in cat.src))),
   # put PTRCAT after STORE
-  (UPat(Ops.STORE, src=(UPat(Ops.PTRCAT, name="cat"), UPat(name="data")), name="sto"), cat_after_store),
+  (UPat(Ops.STORE, src=(UPat(Ops.PTRCAT, name="cat"), UPat(name="data"))), cat_after_store),
 ])
 
 # *** correct load/store ***
@@ -187,7 +188,7 @@ def split_load_store(ctx:Renderer|None, ls:UOp, idx:UOp):
       if global_offset+fold_length > sz: continue
       lidx = buf.index((offset + global_offset).valid(mask), ptr=True)
       if fold_length > 1: lidx = lidx.cast(buf.ptrdtype.base.vec(fold_length).ptr(size=buf.ptrdtype.size, addrspace=buf.ptrdtype.addrspace))
-      if ls.op is Ops.STORE: ret.append(ls.replace(src=(lidx,ls.src[1].gep(tuple(range(global_offset, global_offset+fold_length))))+ls.src[2:]))
+      if ls.op is Ops.STORE: ret.append(ls.replace(src=(lidx,ls.src[1].gep(tuple(range(global_offset, global_offset+fold_length))))))
       else: ret.append(ls.replace(src=(lidx,)+ls.src[1:], dtype=ls.dtype.scalar().vec(fold_length)))
       global_offset += fold_length
       break
@@ -197,8 +198,9 @@ def split_load_store(ctx:Renderer|None, ls:UOp, idx:UOp):
   return UOp(Ops.VCAT, ls.dtype, tuple(ret)) if ls.op is Ops.LOAD else UOp.group(*ret)
 
 def get_image_idx(idx:UOp, width:int):
-  oidx = UOp(Ops.STACK, dtypes.weakint.vec(2), (((x:=idx.src[1].get_idx()) // 4) % width, (x // (4*width))))
-  return idx.replace(src=(idx.src[0], oidx.valid(idx.src[1].get_valid())))
+  x, valid = idx.src[1].get_idx(), idx.src[1].get_valid()
+  idx_x, idx_y = (x // 4) % width, x // (4*width)
+  return idx.replace(src=(idx.src[0], UOp.vectorize(idx_x, idx_y).valid(valid)))
 
 def image_fixup(ls:UOp):
   # normal image load or store, with the CAST from expand_index
@@ -280,18 +282,13 @@ pm_render = PatternMatcher([
   (UPat(Ops.GEP, name='gep'), lambda gep: UOp(Ops.STACK, gep.dtype, tuple(gep.src[0].gep(x) for x in gep.arg)) if len(gep.arg) > 1 else None),
   (UPat(Ops.GEP, name='gep'), lambda gep: gep.src[0] if gep.src[0].dtype.vcount == 1 and gep.arg == (0,) else None),
   (UPat(Ops.STACK, src=(UPat(name='x'),)), lambda x: x),
-  # give any loads that are masked an alt value
-  (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat())).or_casted(),), allow_any_len=True, name="x"),
-    lambda x: x.replace(src=(x.src[0], x.const_like(0))+x.src[1:])
-      if len(x.src) == 1 or x.src[1].op in (Ops.CUSTOM, Ops.STORE, Ops.BARRIER) else None),
   # Where after gated load becomes alt value
   # NOTE: if a is CAST and a.src[0].dtype == l.dtype, use a.src[0] to avoid roundtrip cast (e.g. uint->float->uint)
-  (UPat.var("c").where(UPat(Ops.LOAD, src=(UPat().index(UPat(), UPat.var("c")).or_casted(),), allow_any_len=True, name="l").or_casted(),
-    UPat.var("a")), lambda c,l,a: l.replace(src=(l.src[0], a.src[0] if a.op is Ops.CAST and a.src[0].dtype == l.dtype else a.cast(l.dtype))+
-                                            l.src[2:]).cast(a.dtype)),
-  (UPat.var("c").where(UPat.var("a"), UPat(Ops.LOAD, src=(UPat().index(UPat(), UPat.var("c", dtype=dtypes.bool).logical_not()).or_casted(),),
-    allow_any_len=True, name="l").or_casted()), lambda c,l,a: l.replace(src=(l.src[0], a.src[0] if a.op is Ops.CAST and a.src[0].dtype == l.dtype
-                                                                             else a.cast(l.dtype))+l.src[2:]).cast(a.dtype)),
+  (UPat.var("gate").where(UPat(Ops.LOAD, src=(UPat(), UPat(), UPat.var("gate")), name="l").or_casted(), UPat.var("a")), lambda gate,l,a:
+   l.replace(src=(l.src[0], a.src[0] if a.op is Ops.CAST and a.src[0].dtype == l.dtype else a.cast(l.dtype), l.src[2])).cast(a.dtype)),
+  (UPat.var("gate").where(UPat.var("a"), UPat(Ops.LOAD,
+    src=(UPat(), UPat(), UPat.var("gate", dtype=dtypes.bool).logical_not()), name="l").or_casted()), lambda gate,l,a:
+   l.replace(src=(l.src[0], a.src[0] if a.op is Ops.CAST and a.src[0].dtype == l.dtype else a.cast(l.dtype), l.src[2])).cast(a.dtype)),
 ])
 
 # *** Ops.REDUCE -> Ops.DEFINE_ACC ***
@@ -367,7 +364,7 @@ pm_add_loads = PatternMatcher([
   # add loads to non ptr index
   (UPat(Ops.INDEX, name="idx"), add_load),
   # remove loads from stores
-  (UPat(Ops.STORE, src=(UPat(Ops.LOAD), UPat(name="val")), name="s"), lambda s,val: s.replace(src=(s.src[0].src[0], val))),
+  (UPat(Ops.STORE, src=(UPat(Ops.LOAD),), allow_any_len=True, name="s"), lambda s: s.replace(src=(s.src[0].src[0],)+s.src[1:])),
 ])
 
 # make images

tinygrad/codegen/late/expander.py

@@ -45,7 +45,7 @@ def do_expand(root:UOp):
     else:
       # non-UNROLL input
       if root.op in range_start and i >= range_start[root.op]:
-        # for any range args of STORE/REDUCE, pass them through
+        # for any range args of REDUCE/WMMA/END/etc., pass them through
         new_srcs.append(src)
       elif root.op is Ops.INDEX and i >= 1 and not isinstance(root.dtype, PtrDType):
         new_srcs.append(src)

tinygrad/codegen/late/gater.py

@@ -0,0 +1,13 @@
+# this transforms Invalid into gated load/stores
+
+from tinygrad.uop.ops import PatternMatcher, UPat
+from tinygrad.dtype import Invalid, dtypes
+
+pm_move_gates_from_index = PatternMatcher([
+  (UPat.var("buf").index(UPat.var("gate").where(UPat.var("idx"), UPat(arg=Invalid))).or_casted(name="cast").load(name="l"),
+   lambda buf,gate,idx,cast,l: buf.index(idx, ptr=True).cast(cast.dtype).load(l.const_like(0), gate, dtype=l.dtype)),
+  (UPat.var("buf").index(UPat.var("gate").where(UPat.var("idx"), UPat(arg=Invalid))).or_casted(name="cast").store(UPat.var("data")),
+   lambda buf,gate,idx,cast,data: buf.index(idx, ptr=True).cast(cast.dtype).store(data, gate)),
+  # remove hanging weakint casts
+  (UPat.var("buf").index(UPat.var("idx", dtypes.ints).cast()), lambda buf,idx: buf.index(idx, ptr=True)),
+])

tinygrad/codegen/opt/search.py

@@ -1,6 +1,7 @@
 import math, time, multiprocessing, traceback, signal, atexit
 from dataclasses import replace
-from tinygrad.uop.ops import sym_infer, AxisType, pyrender, UOp
+from tinygrad.uop.ops import sym_infer, AxisType, UOp
+from tinygrad.uop.render import pyrender
 from tinygrad.device import Device, Buffer
 from tinygrad.helpers import prod, flatten, DEBUG, CACHELEVEL, diskcache_get, diskcache_put, getenv, Context, colored, time_to_str, unwrap
 from tinygrad.helpers import IGNORE_BEAM_CACHE

tinygrad/codegen/simplify.py

@@ -14,10 +14,11 @@ def flatten_range(r:UOp) -> UOp|None:
 
 pm_flatten_range = PatternMatcher([
   # real ranges only
-  (UPat((Ops.REDUCE, Ops.STORE, Ops.END), name="r"), flatten_range),
+  (UPat((Ops.REDUCE, Ops.END), name="r"), flatten_range),
 ])
 
-def count_divmod(x:UOp) -> int: return sum(u.op in {Ops.IDIV, Ops.MOD} for u in x.backward_slice)
+# index/range arithmetic uses FLOORDIV/FLOORMOD prior to late rewrite
+def count_divmod(x:UOp) -> int: return sum(u.op in {Ops.FLOORDIV, Ops.FLOORMOD} for u in x.backward_slice)
 def simplify_merge_adjacent(u:UOp) -> UOp|None:
   reduce_ranges = [x.ranges for x in u.backward_slice_with_self if x.op is Ops.REDUCE]
   # on END we only want to merge adjacent ranges, on REDUCE we want to try all combinations

tinygrad/device.py

@@ -5,7 +5,7 @@ from typing import Any, Generic, TypeVar, Iterator, Generator, TYPE_CHECKING
 import importlib, inspect, functools, pathlib, os, platform, contextlib, sys, re, atexit, pickle, decimal
 from tinygrad.helpers import BENCHMARKS, CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored
 from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, suppress_finalizing
-from tinygrad.helpers import select_by_name, select_first_inited, DEV, EMULATED_DTYPES, IMAGE, FLOAT16, TracingKey, size_to_str, Target, VIZ
+from tinygrad.helpers import select_by_name, select_first_inited, DEV, EMULATED_DTYPES, IMAGE, FLOAT16, TracingKey, size_to_str, Target
 from tinygrad.helpers import pluralize
 from tinygrad.dtype import DType, PtrDType, dtypes, _to_np_dtype
 if TYPE_CHECKING: from tinygrad.renderer import Renderer
@@ -371,9 +371,8 @@ if PROFILE:
     with open(fn:=temp("profile.pkl", append_user=True), "wb") as f: pickle.dump(cpu_events+Compiled.profile_events+Buffer.profile_events, f)
 
     PROFILE.value = 0
-    if VIZ > 0:
-      from tinygrad.uop.ops import launch_viz
-      launch_viz("PROFILE", fn)
+    from tinygrad.uop.ops import launch_viz
+    launch_viz("PROFILE", fn)
 
 def enumerate_devices_str() -> Generator[str, None, None]:
   from tinygrad import Tensor, Device

tinygrad/engine/jit.py

@@ -25,7 +25,7 @@ def prune_linear(linear:UOp, needed:set[UOp]) -> tuple[UOp, UOp]:
 def create_graph_call(batch:list[UOp]) -> UOp:
   # all external inputs are PARAMs
   input_list = dedup(u for si in batch for b in si.src[1:] for u in b.toposort() if u.op is Ops.PARAM)
-  cf = UOp(Ops.CUSTOM_FUNCTION, dtypes.void, src=(UOp(Ops.LINEAR, src=tuple(batch)), *input_list), arg="graph")
+  cf = UOp(Ops.CUSTOM_FUNCTION, dtypes.void, src=(UOp(Ops.LINEAR, src=tuple(batch)),), arg="graph")
   return cf.call(*input_list, metadata=tuple(m for si in batch for m in si.arg.metadata))
 
 def graph_split_rewrite(linear:UOp, max_batch_size:int=0) -> UOp:
@@ -48,8 +48,8 @@ def graph_split_rewrite(linear:UOp, max_batch_size:int=0) -> UOp:
     devs = dedup([Device[x] for b in si.src[1:] if b.op is not Ops.BIND for x in (b.device if isinstance(b.device, tuple) else (b.device,))])
     graph_t = graph_class(devs[0]) if devs[0].graph is not None else None
 
-    can_graph = graph_t is not None and graph_t.supports_exec_item(devs, si)
-    can_extend = can_graph and graph_t is not None and (not current_batch_devs or graph_t.supports_exec_item(current_batch_devs, si)) \
+    can_graph = graph_t is not None and graph_t.supports_uop(devs, si)
+    can_extend = can_graph and graph_t is not None and (not current_batch_devs or graph_t.supports_uop(current_batch_devs, si)) \
       and (max_batch_size == 0 or len(current_batch) < max_batch_size)
     if not can_extend and current_batch: flush_batch()
 
@@ -166,13 +166,13 @@ class GraphRunner:
                  for x in (b.device if isinstance(b.device, tuple) else (b.device,))])
 
   @staticmethod
-  def supports_exec_item(batch_devs:list[Compiled], new_call:UOp) -> bool:
+  def supports_uop(batch_devs:list[Compiled], new_call:UOp) -> bool:
     return new_call.src[0].op is Ops.PROGRAM and len(GraphRunner._all_devs(batch_devs, new_call)) == 1
 
 # a marker for your graph supporting multiple devices of the same type
 class MultiGraphRunner(GraphRunner):
   @staticmethod
-  def supports_exec_item(batch_devs:list[Compiled], new_call:UOp) -> bool:
+  def supports_uop(batch_devs:list[Compiled], new_call:UOp) -> bool:
     # Devices must be the same type
     return new_call.src[0].op in (Ops.PROGRAM, Ops.COPY) and len(dedup([type(d) for d in GraphRunner._all_devs(batch_devs, new_call)])) == 1
 

tinygrad/gradient.py

@@ -75,8 +75,8 @@ pm_gradient = PatternMatcher([
   (UPat(Ops.COPY, name="ret"), lambda ctx, ret: (ctx.copy_to_device(ret.src[0].device), None)),
   (UPat(Ops.MULTI, name="ret"), lambda ctx, ret: ctx.shard(ret.device, ret.axis).src),
   (UPat(Ops.TUPLE), lambda ctx: ctx.src),
-  # NOTE: this is only correct when the KERNEL has a single output
-  (UPat(Ops.AFTER), lambda ctx: (ctx, ctx)),
+  (UPat(Ops.AFTER, src=(UPat.var("d"), UPat(Ops.CALL, name="k"))), lambda ctx, d, k:
+    (ctx, UOp.maketuple(*(ctx if i == k.src.index(d)-1 else UOp(Ops.NOOP) for i in range(len(k.src)-1))))),
   # there's no gradient for bitcast
   (UPat(Ops.BITCAST), lambda: (None,)),
 ])
@@ -113,7 +113,11 @@ def compute_gradient(root:UOp, root_grad:UOp, targets:set[UOp]) -> dict[UOp, UOp
     assert len(lgrads) == len(t0.src), f"got {len(lgrads)} gradient, expected {len(t0.src)}"
     for k,v in zip(t0.src, lgrads):
       if v is None: continue
-      if k in grads and grads[k].op is not Ops.NOOP: grads[k] = grads[k] + v
+      if k in grads and grads[k].op is not Ops.NOOP:
+        if v.op is Ops.TUPLE and grads[k].op is Ops.TUPLE:
+          grads[k] = UOp.maketuple(*(p + n if (p.op is not Ops.NOOP and n.op is not Ops.NOOP) else
+                                     n if p.op is Ops.NOOP else p for p, n in zip(grads[k].src, v.src)))
+        else: grads[k] = grads[k] + v
       else: grads[k] = v
       if len(forward_metadata:=all_metadata.get(t0, ())):
         backward_metadata = tuple(dataclasses.replace(x, backward=True) for x in forward_metadata)

tinygrad/helpers.py

@@ -64,6 +64,9 @@ def next_power2(x): return 1 if x == 0 else 1 << (x - 1).bit_length()
 # cstyle div and mod
 def cdiv(x:int, y:int) -> int: return abs(x)//abs(y)*(1,-1)[x*y<0] if y != 0 else 0
 def cmod(x:int, y:int) -> int: return x-cdiv(x,y)*y
+# python floor div and mod
+def floordiv(x:int, y:int) -> int: return x//y if y != 0 else 0
+def floormod(x:int, y:int) -> int: return x-floordiv(x,y)*y
 def lo32(x:Any) -> Any: return x & 0xFFFFFFFF # Any is sint
 def hi32(x:Any) -> Any: return x >> 32 # Any is sint
 def data64(data:Any) -> tuple[Any, Any]: return (data >> 32, data & 0xFFFFFFFF) # Any is sint
@@ -238,7 +241,7 @@ SPLIT_REDUCEOP, NO_MEMORY_PLANNER, LRU = ContextVar("SPLIT_REDUCEOP", 1), Contex
 RING, ALL2ALL, ALLREDUCE_CAST = ContextVar("RING", 1), ContextVar("ALL2ALL", 0), ContextVar("ALLREDUCE_CAST", 1)
 CACHELEVEL, IGNORE_BEAM_CACHE, DEVECTORIZE = ContextVar("CACHELEVEL", 2), ContextVar("IGNORE_BEAM_CACHE", 0), ContextVar("DEVECTORIZE", 1)
 VALIDATE_WITH_CPU, DISABLE_FAST_IDIV = ContextVar("VALIDATE_WITH_CPU", 0), ContextVar("DISABLE_FAST_IDIV", 0)
-CORRECT_DIVMOD_FOLDING, FUSE_OPTIM = ContextVar("CORRECT_DIVMOD_FOLDING", 0), ContextVar("FUSE_OPTIM", 0)
+FUSE_OPTIM = ContextVar("FUSE_OPTIM", 0)
 ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE = ContextVar("ALLOW_DEVICE_USAGE", 1), ContextVar("MAX_BUFFER_SIZE", 0)
 MAX_KERNEL_BUFFERS = ContextVar("MAX_KERNEL_BUFFERS", 0)
 EMULATED_DTYPES = ContextVar("EMULATED_DTYPES", "")
@@ -441,24 +444,28 @@ def _ensure_downloads_dir() -> pathlib.Path:
     return downloads_dir
   return pathlib.Path(cache_dir) / "downloads"
 
-def fetch(url:str, name:pathlib.Path|str|None=None, subdir:str|None=None, gunzip:bool=False,
-          allow_caching=not getenv("DISABLE_HTTP_CACHE"), headers:dict[str, str]={}) -> pathlib.Path:
+def fetch(url:str, name:pathlib.Path|str|None=None, subdir:str|None=None, gunzip:bool=False, allow_caching=not getenv("DISABLE_HTTP_CACHE"),
+          headers:dict[str, str]={}, sha256:str|None=None) -> pathlib.Path:
   import urllib.request
   if url.startswith(("/", ".")): return pathlib.Path(url)
   if name is not None and (isinstance(name, pathlib.Path) or '/' in name): fp = pathlib.Path(name)
   else:
     hh = "_"+hashlib.md5(("\n".join(f"{k.strip()}:{v.strip()}" for k,v in sorted(headers.items()))).encode("utf-8")).hexdigest() if headers else ""
     fp = _ensure_downloads_dir() / (subdir or "") / ((name or hashlib.md5(url.encode('utf-8')).hexdigest()) + hh + (".gunzip" if gunzip else ""))
-  if not fp.is_file() or not allow_caching:
+  if not fp.is_file() or not allow_caching or (sha256 and hashlib.sha256(fp.read_bytes()).hexdigest() != sha256):
     (_dir := fp.parent).mkdir(parents=True, exist_ok=True)
     with urllib.request.urlopen(urllib.request.Request(url, headers={"User-Agent": "tinygrad 0.12.0", **headers}), timeout=10) as r:
       assert r.status in {200, 206}, r.status
       length = int(r.headers.get('content-length', 0)) if not gunzip else None
       readfile = gzip.GzipFile(fileobj=r) if gunzip else r
       progress_bar:tqdm = tqdm(total=length, unit='B', unit_scale=True, desc=f"{url}", disable=CI)
+      h = hashlib.sha256() if sha256 else None
       with tempfile.NamedTemporaryFile(dir=_dir, delete=False) as f:
-        while chunk := readfile.read(16384): progress_bar.update(f.write(chunk))
+        while chunk := readfile.read(16384):
+          if h: h.update(chunk)
+          progress_bar.update(f.write(chunk))
         f.close()
+        if h and (actual_sha256:=h.hexdigest()) != sha256: raise RuntimeError(f"fetch sha mismatch, expected {sha256} but got {actual_sha256}")
         pathlib.Path(f.name).rename(fp)
       progress_bar.update(close=True)
       if length and (file_size:=os.stat(fp).st_size) < length: raise RuntimeError(f"fetch size incomplete, {file_size} < {length}")

tinygrad/mixin/elementwise.py

@@ -181,7 +181,7 @@ class ElementwiseMixin(DTypeMixin, CreationMixin):
     return self._binop(Ops.IDIV, x, reverse)
 
   def mod(self, x: Self | ConstType, reverse: bool = False) -> Self:
-    return self._binop(Ops.MOD, x, reverse)
+    return self._binop(Ops.FLOORMOD, x, reverse)
 
   def div(self, x: Self | ConstType, reverse: bool = False) -> Self:
     lhs, rhs = self._broadcasted(x, reverse)
@@ -206,7 +206,7 @@ class ElementwiseMixin(DTypeMixin, CreationMixin):
     return self.div(x)
 
   def __floordiv__(self, x: Self | ConstType) -> Self:
-    return self.idiv(x)  # TODO: idiv is trunc div, not floordiv
+    return self._binop(Ops.FLOORDIV, x, False)
 
   def __mod__(self, x: Self | ConstType) -> Self:
     return self.mod(x)
@@ -233,7 +233,7 @@ class ElementwiseMixin(DTypeMixin, CreationMixin):
     return self.div(x, True)
 
   def __rfloordiv__(self, x: Self | ConstType) -> Self:
-    return self.idiv(x, True)
+    return self._binop(Ops.FLOORDIV, x, True)
 
   def __rand__(self, x: Self | ConstType) -> Self:
     return self.bitwise_and(x, True)

tinygrad/nn/onnx.py

@@ -644,7 +644,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
   def BitwiseOr(x:Tensor,y:Tensor): return x | y
   def BitwiseXor(x:Tensor,y:Tensor): return x ^ y
   def BitwiseNot(x:Tensor): return ~x
-  def Mod(x:Tensor,y:Tensor,fmod=0): return x - x.div(y, rounding_mode="trunc") * y if fmod else x % y
+  def Mod(x:Tensor,y:Tensor,fmod=0): return x.fmod(y) if fmod else x % y
 
   # ***** Casting Ops *****
   # NOTE: saturate only applies to FP8 types

tinygrad/renderer/cstyle.py

@@ -44,12 +44,11 @@ base_rewrite = PatternMatcher([
   # default const render
   (UPat(Ops.CONST, name="x"), lambda ctx,x: str(x.arg)),
   # new load/store
-  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var('idx')), allow_any_len=True),
+  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var('idx'))),
    lambda ctx,buf,idx: f"({ctx[buf]}+{strip_parens(ctx[idx]) if idx.arg == Ops.ADD else ctx[idx]})"),
-  (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat.var("gate"))).or_casted("bidx"), UPat.var("var"))),
-   lambda ctx,bidx,var,gate: f"({ctx[gate]}?*{ctx[bidx]}:{ctx[var]})"),
   (UPat(Ops.LOAD, src=(UPat.var('bidx'),)), lambda ctx,bidx: f"(*{ctx[bidx]})"),
-  (UPat(Ops.STORE, src=(UPat.var('bidx'), UPat.var("var"))), lambda ctx,bidx,var: f"*{ctx[bidx]} = {ctx[var]};"),
+  (UPat(Ops.LOAD, src=(UPat.var("bidx"), UPat.var("var"), UPat.var("gate"))), lambda ctx,bidx,var,gate: f"({ctx[gate]}?*{ctx[bidx]}:{ctx[var]})"),
+  (UPat(Ops.STORE, src=(UPat.var('bidx'), UPat.var("var")), allow_any_len=True), lambda ctx,bidx,var: f"*{ctx[bidx]} = {ctx[var]};"),
   # alu/gep
   # TODO: look for left-associative
   (UPat(GroupOp.ALU, name="x"), lambda ctx,x: ctx.code_for_op[x.op](
@@ -302,11 +301,11 @@ class OpenCLRenderer(CStyleLanguage):
     (UPat(Ops.CONST, dtypes.bfloat16, name="x"),
       lambda ctx,x: f"{(struct.unpack('I', struct.pack('f', float_to_bf16(x.arg)))[0] >> 16)}u"),
     # load/store image (OpenCL)
-    (UPat(Ops.LOAD, dtype=dtypes.float.vec(4), src=(UPat.var('buf').index(UPat.var('idx', dtypes.int.vec(2)), UPat.var("gate")), UPat.var("var"))),
+    (UPat(Ops.LOAD, dtype=dtypes.float.vec(4), src=(UPat.var('buf').index(UPat.var('idx', dtypes.int.vec(2))), UPat.var("var"), UPat.var("gate"))),
       lambda ctx,buf,idx,var,gate: f"({ctx[gate]}?read_imagef({ctx[buf]}, smp, {ctx[idx]}):{ctx[var]})"),
     (UPat(Ops.LOAD, dtype=dtypes.float.vec(4), src=(UPat.var('buf').index(UPat.var('idx', dtypes.int.vec(2))),)),
       lambda ctx,buf,idx: f"read_imagef({ctx[buf]}, smp, {ctx[idx]})"),
-    (UPat(Ops.STORE, src=(UPat.var('buf').index(UPat.var('idx', dtypes.int.vec(2)), allow_any_len=True),
+    (UPat(Ops.STORE, src=(UPat.var('buf').index(UPat.var('idx', dtypes.int.vec(2))),
                           UPat.var("var", dtypes.float.vec(4))), allow_any_len=True),
       lambda ctx,buf,idx,var: f"write_imagef({ctx[buf]}, {ctx[idx]}, {ctx[var]});"),
   ]) + base_rewrite

tinygrad/renderer/llvmir.py

@@ -76,14 +76,14 @@ base_rewrite = PatternMatcher([
   # memory load/store
   (UPat(Ops.INDEX, name="x"), lambda ctx,x:
    f"  {ctx[x]} = getelementptr inbounds {ldt(x.dtype.base)}, {ldt(x.src[0].dtype)} {ctx[x.src[0]]}, {ldt(x.src[1].dtype)} {ctx[x.src[1]]}"),
-  (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat.var("mask"))).or_casted("idx"), UPat.var("alt")), allow_any_len=True, name="x"),
+  (UPat(Ops.LOAD, src=(UPat.var("idx"), UPat.var("alt"), UPat.var("mask")), name="x"),
    lambda ctx,x,idx,alt,mask:
    f"  br label {ctx[x]}_entry\n{ctx[x][1:]}_entry:\n"
    f"  br i1 {ctx[mask]}, label {ctx[x]}_load, label {ctx[x]}_exit\n{ctx[x][1:]}_load:\n"
    f"  {ctx[x]}_yes = load {ldt(x.dtype)}, {ldt(idx.dtype)} {ctx[idx]}\n"
    f"  br label {ctx[x]}_exit\n{ctx[x][1:]}_exit:\n"
    f"  {ctx[x]} = phi {ldt(x.dtype)} [{ctx[x]}_yes, {ctx[x]}_load], [{ctx[alt]}, {ctx[x]}_entry]"),
-  (UPat(Ops.LOAD, src=(UPat.var('idx'),), allow_any_len=True, name="x"),
+  (UPat(Ops.LOAD, src=(UPat.var('idx'),), name="x"),
    lambda ctx,x,idx: f"  {ctx[x]} = load {ldt(x.dtype)}, {ldt(idx.dtype)} {ctx[idx]}"),
   (UPat(Ops.STORE, name="x"), lambda ctx,x: f"  store {ldt(x.src[1].dtype)} {ctx[x.src[1]]}, {ldt(x.src[0].dtype)} {ctx[x.src[0]]}"),
 

tinygrad/renderer/nir.py

@@ -129,7 +129,7 @@ class NIRRenderer(Renderer):
     (UPat(Ops.LOAD, dtypes.bool, name="x"),
      lambda x: x.replace(dtype=dtypes.uint8, src=x.src[0:1]+((x.src[1].cast(dtypes.uint8),) if len(x.src)>=2 else ())+x.src[2:]).cast(dtypes.bool)),
     (UPat(Ops.STORE, src=(UPat(), UPat(dtype=dtypes.bool)), name="x", allow_any_len=True),
-     lambda x: x.replace(src=x.src[0:1] + (x.src[1].cast(dtypes.uint8),) + x.src[2:])),
+     lambda x: x.replace(src=(x.src[0], x.src[1].cast(dtypes.uint8))+x.src[2:])),
     # NIR requires shift amount to be 32 bit: https://docs.mesa3d.org/nir/alu.html#nir-alu-op-ishl
     (UPat((Ops.SHL, Ops.SHR), name="x"), lambda x: x.replace(src=(x.src[0], x.src[1].cast(dtypes.uint))) if x.src[1].dtype.bitsize != 32 else None),
     # OpConvertFToU is undefined if Result Type is not wide enough, cast through int32
@@ -146,12 +146,12 @@ class NIRRenderer(Renderer):
     (UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx.param(ctx.b, x, 8)),
     (UPat(Ops.DEFINE_VAR, name="x"), lambda ctx,x: ctx.param(ctx.b, x, 4)),
     (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: nchannel(ctx.b, {'g':ngid, 'l':nlid, 'i': nid}[x.arg[0]](ctx.b), int(x.arg[-1]))),
-    (UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat.var("buf"),UPat.var("off")), allow_any_len=True), UPat.var("val")), allow_any_len=True),
+    (UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat.var("buf"),UPat.var("off"))), UPat.var("val")), allow_any_len=True),
      lambda ctx,buf,off,val: nstore(ctx.b, buf.ptrdtype.addrspace, nidx(ctx.b, ctx.r[buf], ctx.r[off], buf.dtype), ctx.r[val], val.dtype)),
-    (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off"), UPat.var("gate"))), UPat.var("alt")), allow_any_len=True, name="x"),
+    (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off"))), UPat.var("alt"), UPat.var("gate")), name="x"),
      lambda ctx,x,buf,off,alt,gate: if_phi(ctx.b, ctx.r[gate],
       lambda: nload(ctx.b, buf.ptrdtype.addrspace, nidx(ctx.b, ctx.r[buf], ctx.r[off], buf.dtype, ctx.r[gate]), x.dtype), lambda: ctx.r[alt])),
-    (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off"))),), allow_any_len=True, name="x"),
+    (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off"))),), name="x"),
      lambda ctx,x,buf,off: nload(ctx.b, buf.ptrdtype.addrspace, nidx(ctx.b, ctx.r[buf], ctx.r[off], buf.dtype), x.dtype)),
     (UPat(Ops.STACK, name="x"), lambda ctx,x: nalu(ctx.b, f"vec{x.dtype.count}", *[ctx.r[src] for src in x.src])),
     (UPat(GroupOp.ALU, name="x"), lambda ctx,x: nalu(ctx.b, aop[x.src[0].dtype.scalar()][x.op], *[ctx.r[src] for src in x.src])),
@@ -268,9 +268,9 @@ class IR3Renderer(NIRRenderer, OpenCLRenderer):
     return _nload_img(ctx.b, ctx.r[img], ctx.r[coord], img.dtype)
 
   def_rewrite = PatternMatcher([
-    (UPat(Ops.STORE, src=(UPat.var('img').index(UPat.var('coord', dtypes.int.vec(2)), allow_any_len=True), UPat.var("val")),
-          allow_any_len=True), lambda ctx,img,coord,val: nstore_img(ctx.b, ctx.r[img], ctx.r[coord], ctx.r[val], val.dtype)),
-    (UPat(Ops.LOAD, src=(UPat.var('img').index(UPat.var('coord', dtypes.int.vec(2)), UPat.var("gate")), UPat.var("alt"))),
+    (UPat(Ops.STORE, src=(UPat.var('img').index(UPat.var('coord', dtypes.int.vec(2))), UPat.var("val")), allow_any_len=True),
+     lambda ctx,img,coord,val: nstore_img(ctx.b, ctx.r[img], ctx.r[coord], ctx.r[val], val.dtype)),
+    (UPat(Ops.LOAD, src=(UPat.var('img').index(UPat.var('coord', dtypes.int.vec(2))), UPat.var("alt"), UPat.var("gate"))),
      lambda ctx,img,coord,alt,gate: if_phi(ctx.b, ctx.r[gate], lambda: ctx.nload_img(img, coord), lambda: ctx.r[alt])),
     (UPat(Ops.LOAD, src=(UPat.var('img').index(UPat.var('coord', dtypes.int.vec(2))),)), nload_img),
   ]) + NIRRenderer.def_rewrite

tinygrad/renderer/ptx.py

@@ -49,9 +49,9 @@ ptx_matcher = PatternMatcher([
   (UPat(Ops.LOAD, dtypes.bool, src=(UPat(dtype=dtypes.int64),), name="x", allow_any_len=True),
    lambda x: UOp(x.op, dtypes.uint8, x.src[0:1] + ((x.src[1].cast(dtypes.uint8),) if len(x.src) >= 2 else ()) + x.src[2:]).cast(dtypes.bool)),
   (UPat(Ops.STORE, src=(UPat(dtype=dtypes.int64), UPat(dtype=dtypes.bool)), name="x", allow_any_len=True),
-   lambda x: UOp(x.op, dtypes.void, x.src[0:1] + (x.src[1].cast(dtypes.uint8),) + x.src[2:])),
+   lambda x: UOp(x.op, dtypes.void, (x.src[0], x.src[1].cast(dtypes.uint8))+x.src[2:])),
   # indexing on PTX is in uint64, we do the math while it's still in the graph
-  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx")), name="op", allow_any_len=True), lambda buf,idx,op:
+  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx")), name="op"), lambda buf,idx,op:
     UOp(Ops.INDEX, dtype=dtypes.int64, src=(buf, buf.cast(dtypes.int64)+idx.cast(dtypes.int64)*buf.dtype.itemsize)+op.src[2:]) \
       if op.dtype != dtypes.int64 and buf.dtype.addrspace != AddrSpace.REG else None),
   # load/store use pointer arithmetic, and the cast does nothing
@@ -102,11 +102,11 @@ string_rewrite = PatternMatcher([
   (UPat(Ops.CAST, name="x", src=(UPat.var("a"),)),
    lambda ctx, x, a: f"cvt{modifier(x.dtype, a.dtype)}.{ctx.cast_types[x.dtype]}.{ctx.cast_types[a.dtype]} {ctx.r[x]}, {ctx.r[a]};"),
   # store / gated load / load
-  (UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("loc")), allow_any_len=True), UPat.var("var"))),
+  (UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("loc"))), UPat.var("var")), allow_any_len=True),
    lambda ctx, loc, var, buf: f"st.{mem_type(buf)}" + \
     f"{f'.v{cnt}' if ((cnt:=var.dtype.count)>1) else ''}.{ctx.mem_types[var.dtype.scalar()]} " + \
     f"[{ctx.r[loc]}+0], {('{' + ', '.join(ctx.r[var]) + '}') if var.dtype.count > 1 else ctx.r[var]};"),
-  (UPat(Ops.LOAD, name="x", src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("loc"), UPat.var("gate"))), UPat.var("alt")), allow_any_len=True),
+  (UPat(Ops.LOAD, name="x", src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("loc"))), UPat.var("alt"), UPat.var("gate")), allow_any_len=True),
     lambda ctx, x, loc, alt, gate, buf: flatten([
     [f"mov.{ctx.mem_types[x.dtype.scalar()]} {v}, {render_val(0, x.dtype.scalar())};" for v in ctx.r[x]],
     [f"@{ctx.r[gate]} ld.{mem_type(buf)}.v{x.dtype.count}.{ctx.mem_types[x.dtype.scalar()]} {{{', '.join(ctx.r[x])}}}, [{ctx.r[loc]}+0];"]

tinygrad/renderer/wgsl.py

@@ -10,21 +10,20 @@ def sign_extend(val:UOp, sext_am:int):
         | val.bitcast(dtypes.uint32)).bitcast(dtypes.int)
 
 # store for char: buf[idx/4] <- (var << (idx%4)*8))
-def packed_store(bidx:UOp, var:UOp):
+def packed_store(bidx:UOp, var:UOp, gate:UOp|None=None):
   elems, mask = 4//var.dtype.itemsize, _mask(var.dtype)
   shift_am, div_idx = (bidx.src[1].cast(dtypes.uint32) % elems) * (8*var.dtype.itemsize), bidx.src[1] // elems
   new_v, wmask = (var & mask).cast(dtypes.uint32) << shift_am, ((mask << shift_am) ^ 0xFFFFFFFF).cast(dtypes.uint32)
-  # preserve valid condition (bidx.src[2]) if it exists for gated stores
-  idx_src = (bidx.src[0], div_idx) if len(bidx.src) == 2 else (bidx.src[0], div_idx, bidx.src[2])
-  buf = UOp.load(UOp(Ops.INDEX, bidx.dtype, idx_src), dtype=dtypes.uint32)
-  return UOp.store(UOp(Ops.INDEX, bidx.dtype, idx_src), (buf & wmask) | new_v)
+  idx = UOp(Ops.INDEX, bidx.dtype, (bidx.src[0], div_idx))
+  buf = UOp.load(idx, *((UOp.const(dtypes.uint32, 0), gate) if gate is not None else ()), dtype=dtypes.uint32)
+  return UOp.store(idx, (buf & wmask) | new_v, *((gate,) if gate is not None else ()))
 
 # load for char: sign_extend(buf[idx/4] >> ((idx%4)*8))
-def packed_load(root:UOp, bidx:UOp, dtype:DType, var:UOp|None=None):
+def packed_load(root:UOp, bidx:UOp, dtype:DType, var:UOp|None=None, gate:UOp|None=None):
   elems, mask = 4//dtype.itemsize, _mask(dtype)
   shift_am, div_idx = (bidx.src[1].cast(dtypes.uint32) % elems) * (8*dtype.itemsize), bidx.src[1] // elems
-  idx = UOp(Ops.INDEX, bidx.dtype, (bidx.src[0], div_idx, bidx.src[2]) if var is not None else (bidx.src[0], div_idx))
-  load = UOp.load(idx, *([var] if var is not None else root.src[1:]), dtype=dtypes.uint32, arg=root.arg)
+  idx = UOp(Ops.INDEX, bidx.dtype, (bidx.src[0], div_idx))
+  load = UOp.load(idx, *((var, gate) if var is not None and gate is not None else root.src[1:]), dtype=dtypes.uint32, arg=root.arg)
   val = (load.cast(dtypes.uint32) >> shift_am) & mask
   return sign_extend(val, 8*dtype.itemsize).cast(dtype) if dtype in [dtypes.char, dtypes.short] else val.cast(dtype)
 
@@ -41,10 +40,12 @@ wgsl_matcher = PatternMatcher([
   (UPat((Ops.CMPLT, Ops.XOR), src=(UPat(name="a", dtype=dtypes.bool), UPat.var("b")), name="c"),
    lambda a,b,c: a.cast(dtypes.int).alu(c.op, b.cast(dtypes.int)).cast(dtypes.bool)),
   # TODO: load alt value doesnt have to be a const
-  (UPat.load(UPat.var("b"), UPat.cvar("c"), allow_any_len=True, name="l"),
-   lambda l,b,c: packed_load(l,b,l.dtype,c.cast(dtypes.uint32)) if is_packed(l.dtype, b.dtype) else None),
-  (UPat.load(UPat.var("b"), name='l', allow_any_len=True), lambda l,b: packed_load(l, b, l.dtype) if is_packed(l.dtype, b.dtype) else None),
-  (UPat.store(UPat.var("bidx"), UPat.var("var"), allow_any_len=True),
+  (UPat.load(UPat.var("b"), UPat.cvar("c"), UPat.var("gate"), name="l"),
+   lambda l,b,c,gate: packed_load(l,b,l.dtype,c.cast(dtypes.uint32),gate) if is_packed(l.dtype, b.dtype) else None),
+  (UPat.load(UPat.var("b"), name='l'), lambda l,b: packed_load(l, b, l.dtype) if is_packed(l.dtype, b.dtype) else None),
+  (UPat.store(UPat.var("bidx"), UPat.var("var"), UPat.var("gate")),
+   lambda bidx,var,gate: packed_store(bidx,var,gate) if is_packed(var.dtype, bidx.dtype) else None),
+  (UPat.store(UPat.var("bidx"), UPat.var("var")),
    lambda bidx,var: packed_store(bidx,var) if is_packed(var.dtype, bidx.dtype) else None),
   (UPat.var("a") << UPat.var("b"),lambda a,b:(a.bitcast(dtypes.uint32)<<b.cast(dtypes.uint32)).bitcast(a.dtype) if b.dtype!=dtypes.uint32 else None),
   (UPat.var("x") >> UPat.var("y"), lambda x,y: UOp(Ops.SHR, x.dtype, (x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None),
@@ -82,14 +83,14 @@ class WGSLRenderer(CStyleLanguage):
      if x.src[0].dtype == dtypes.half else f"((i32({ctx[x.src[0]]}&0xFFFF)<<16)>>16)"),
     (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"bitcast<{ctx.type_map[x.dtype]}>({ctx[x.src[0]]})"),
     # TODO: load alt value doesnt have to be a const
-    (UPat.load(UPat.var("b"), UPat.cvar("v"), allow_any_len=True),
-      lambda ctx,b,v: f"select({ctx[v]}, {ctx.render_load(ctx[b],b.src[0].dtype)}, {ctx[b.src[2]]})"),
-    (UPat.load(UPat.var("b"), allow_any_len=True), lambda ctx, b: ctx.render_load(ctx[b], b.dtype)),
+    (UPat.load(UPat.var("b"), UPat.cvar("v"), UPat.var("gate")),
+      lambda ctx,b,v,gate: f"select({ctx[v]}, {ctx.render_load(ctx[b],b.src[0].dtype)}, {ctx[gate]})"),
+    (UPat.load(UPat.var("b")), lambda ctx, b: ctx.render_load(ctx[b], b.dtype)),
     (UPat.store(UPat.var("b"), UPat.var("v"), allow_any_len=True),lambda ctx,b,v:\
      # (load & mask) | var -> mask = v.src[0].src[1], var = v.src[1]
      f"atomicAnd(&{ctx[b]},{ctx[v.src[0].src[1]]});\n  atomicAdd(&{ctx[b]},{ctx[v.src[1]]});" if is_packed(b.src[0].dtype) \
       else f"{ctx[b]} = {ctx[v]};"),
-    (UPat(Ops.INDEX, src=(UPat.var("b"), UPat.var("idx")), allow_any_len=True),
+    (UPat(Ops.INDEX, src=(UPat.var("b"), UPat.var("idx"))),
      lambda ctx,b,idx: f"{ctx[b]}[{strip_parens(ctx[idx]) if idx.arg is Ops.ADD else ctx[idx]}]"),
   ]) + base_rewrite
 

tinygrad/runtime/autogen/__init__.py

@@ -10,6 +10,9 @@ linux_headers_deb = "https://snapshot.debian.org/archive/debian/20260207T145350Z
 linux_headers_kern_deb = "https://snapshot.debian.org/archive/debian/20260207T145350Z/pool/main/l/linux/linux-headers-6.18.9+deb14-common_6.18.9-1_all.deb"
 liburing_src = "https://raw.githubusercontent.com/axboe/liburing/refs/tags/liburing-2.14/src/include/liburing.h"
 ggml_common_src = "https://raw.githubusercontent.com/ggml-org/ggml/d4fcfe88a8bcf5c9840be14be6c2fbf1f5b3b2db/src/ggml-common.h"
+cudart_src = "https://developer.download.nvidia.com/compute/cuda/redist/cuda_cudart/linux-x86_64/cuda_cudart-linux-x86_64-12.0.146-archive.tar.xz"
+nvrtc_src = "https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvrtc/linux-x86_64/cuda_nvrtc-linux-x86_64-12.0.140-archive.tar.xz"
+opencl_src = "https://github.com/KhronosGroup/OpenCL-Headers/archive/2e30669d48718fd460f085b4b35b160dad51ce9d.tar.gz"
 macossdk = "/var/db/xcode_select_link/Platforms/MacOSX.platform/Developer/SDKs/MacOSX.sdk"
 
 llvm_lib = (
@@ -50,9 +53,9 @@ def __getattr__(nm):
       [i for i in system("dpkg -L libc6-dev").split() if 'sys/mman.h' in i or 'sys/syscall.h' in i] +
       ["/usr/include/string.h", "/usr/include/elf.h", "/usr/include/unistd.h", "/usr/include/asm-generic/mman-common.h"]), dll="'c'", errno=True)
     case "avcodec": return load("avcodec", ["{}/libavcodec/hevc/hevc.h", "{}/libavcodec/cbs_h265.h"], srcs=ffmpeg_src)
-    case "opencl": return load("opencl", ["/usr/include/CL/cl.h"], dll="'OpenCL'")
-    case "cuda": return load("cuda", ["/usr/include/cuda.h"], dll="'cuda'", args=["-D__CUDA_API_VERSION_INTERNAL"], macros=False)
-    case "nvrtc": return load("nvrtc", ["/usr/include/nvrtc.h"], dll="'nvrtc'", paths=nv_lib_path, prolog=["import sysconfig"])
+    case "opencl": return load("opencl", ["{}/CL/cl.h"], dll="'OpenCL'", args=["-I{}"], srcs=opencl_src)
+    case "cuda": return load("cuda", ["{}/include/cuda.h"], dll="'cuda'", args=["-D__CUDA_API_VERSION_INTERNAL"], srcs=cudart_src, macros=False)
+    case "nvrtc": return load("nvrtc", ["{}/include/nvrtc.h"], dll="'nvrtc'", paths=nv_lib_path, srcs=nvrtc_src, prolog=["import sysconfig"])
     case "nvjitlink": load("nvjitlink", [root/"extra/nvJitLink.h"], dll="'nvJitLink'", paths=nv_lib_path, prolog=["import sysconfig"])
     case "kfd": return load("kfd", [root/"extra/hip_gpu_driver/kfd_ioctl.h"])
     case "nv_570" | "nv_580":

tinygrad/runtime/autogen/am/__init__.py

@@ -1,11 +1,37 @@
-import pathlib, hashlib
+import pathlib, hashlib, re, itertools
 from tinygrad.runtime.autogen import load, root
 
+__all__ = ["am", "pm4_soc15", "pm4_nv", "sdma_4_0_0", "sdma_5_0_0", "sdma_6_0_0", "smu_13_0_0", "smu_13_0_6", "smu_13_0_12", "smu_14_0_2",
+           "fw", "navi_offsets", "vega_offsets", "regs", "soc_9", "soc_11", "soc_12"]
+
 am_src="https://github.com/ROCm/ROCK-Kernel-Driver/archive/33970e1351f5e511029602454979f3de7e22260f.tar.gz"
+rocm_src="https://github.com/ROCm/rocm-systems/archive/cccc350dc620e61ae2554978b62ab3532dc10bd9.tar.gz"
 AMD, AMDINC = "{}/drivers/gpu/drm/amd", "{}/drivers/gpu/drm/amd/include"
 inc, kern_rules = ["-include", "stdint.h"], [(r'le32_to_cpu', ''),]
 fw_src="https://gitlab.com/kernel-firmware/linux-firmware/-/archive/1e2c15348485939baf1b6d1f5a7a3b799d80703d/1e2c15348485939baf1b6d1f5a7a3b799d80703d.tar.gz"
 
+reg_files = {
+  "gc": [(9,4,3), (11,0,0), (11,0,3), (11,5,0), (12,0,0)],
+  "mmhub": [(1,8,0), (3,0,0), (3,0,1), (3,0,2), (3,3,0), (4,1,0)],
+  "nbio": [(4,3,0), (7,2,0), (7,7,0), (7,9,0), (7,11,0)], "nbif": [(6,3,1)],
+  "mp": [(11,0,0), (13,0,0), (14,0,2)], "hdp": [(4,4,2), (6,0,0), (7,0,0)],
+  "osssys": [(4,4,2), (6,0,0), (6,1,0), (7,0,0)], "sdma": [(4,4,2)]
+}
+
+reg_patterns = {
+  "gc": ["GCVM", "GCMC_VM", "CP_(HQD|MQD|MEC|ME_CNTL|PERFMON|RB_WPTR_POLL_CNTL|INT_CNTL|STAT|PFP_PRGRM|ME_PRGRM|COHER_START)", "COMPUTE_",
+         "(SQ|GL2C|TCC)_PERFCOUNTER", "SQ_THREAD_TRACE", "SPI_(CONFIG_CNTL|COMPUTE_QUEUE_RESET)", "GRBM", "SH_MEM", "RLC", "TCP", "GB_ADDR_CONFIG",
+         "SDMA[01]_(WATCHDOG_CNTL|UTCL1_(CNTL|PAGE)|MCU_CNTL|F32_CNTL|CNTL|QUEUE0_|RLC_CGCG_CTRL)", "SCRATCH_REG[67]"],
+  "mmhub": ["MMVM", "MMMC_VM", "MM_ATC_L2_MISC_CG"],
+  "nbio": (nbio:=["BIF_BX_PF[01]_GPU_HDP_FLUSH", "BIF_BX_PF0_RSMU", "BIF_BX0_(REMAP_HDP_MEM_FLUSH_CNTL|BIF_DOORBELL_INT_CNTL|PCIE_INDEX2|PCIE_DATA2)",
+                  "BIFC_(DOORBELL_ACCESS_EN_PF|GFX_INT_MONITOR_MASK)", "XCC_DOORBELL_FENCE", "DOORBELL0_CTRL_ENTRY", "GDC_S2A0_S2A_DOORBELL_ENTRY",
+                  "S2A_DOORBELL_ENTRY", "RCC_DEV0_EPF0_RCC_DOORBELL_APER_EN", "RCC_DEV0_EPF2_STRAP2"]),
+  "nbif": nbio,
+  "mp": ["MP([01]|ASP)_SMN_C2PMSG"], "hdp": ["HDP_MEM_POWER_CTRL"], "oss": ["IH_"], "sdma": ["SDMA_GFX", "SDMA_CNTL"]
+}
+
+soc_patterns = ["SQ_TT", "VGT_EVENT_TYPE", "CS", "MTYPE", "SH"]
+
 def __getattr__(nm):
   match nm:
     case "am": return load("am/am", [root/f"extra/amdpci/headers/{s}.h" for s in ["v11_structs", "v12_structs", "amdgpu_vm",
@@ -20,13 +46,13 @@ def __getattr__(nm):
                                    args=["-I/opt/rocm/include", "-x", "c++"], srcs=am_src)
     case "sdma_6_0_0": return load("am/sdma_6_0_0", [root/"extra/hip_gpu_driver/sdma_registers.h", f"{AMD}/amdgpu/sdma_v6_0_0_pkt_open.h"],
                                    args=["-I/opt/rocm/include", "-x", "c++"], srcs=am_src)
-    case "smu_v13_0_0": return load("am/smu_v13_0_0", [f"{AMD}/pm/swsmu/inc/pmfw_if/{s}.h" for s in ["smu_v13_0_0_ppsmc","smu13_driver_if_v13_0_0"]]
+    case "smu_13_0_0": return load("am/smu_13_0_0", [f"{AMD}/pm/swsmu/inc/pmfw_if/{s}.h" for s in ["smu_v13_0_0_ppsmc","smu13_driver_if_v13_0_0"]]
                                     +[root/"extra/amdpci/headers/amdgpu_smu.h"], args=inc, srcs=am_src)
-    case "smu_v13_0_6": return load("am/smu_v13_0_6", [f"{AMD}/pm/swsmu/inc/pmfw_if/{s}.h" for s in ["smu_v13_0_6_ppsmc","smu_v13_0_6_pmfw", \
+    case "smu_13_0_6": return load("am/smu_13_0_6", [f"{AMD}/pm/swsmu/inc/pmfw_if/{s}.h" for s in ["smu_v13_0_6_ppsmc","smu_v13_0_6_pmfw", \
       "smu13_driver_if_v13_0_6"]] +[root/"extra/amdpci/headers/amdgpu_smu.h"], args=inc, srcs=am_src)
-    case "smu_v13_0_12": return load("am/smu_v13_0_12", [f"{AMD}/pm/swsmu/inc/pmfw_if/{s}.h" for s in ["smu_v13_0_12_ppsmc","smu_v13_0_12_pmfw",
+    case "smu_13_0_12": return load("am/smu_13_0_12", [f"{AMD}/pm/swsmu/inc/pmfw_if/{s}.h" for s in ["smu_v13_0_12_ppsmc","smu_v13_0_12_pmfw",
       "smu13_driver_if_v13_0_6"]] +[root/"extra/amdpci/headers/amdgpu_smu.h"], args=inc, srcs=am_src)
-    case "smu_v14_0_2": return load("am/smu_v14_0_2", [f"{AMD}/pm/swsmu/inc/pmfw_if/{s}.h" for s in ["smu_v14_0_0_pmfw", "smu_v14_0_2_ppsmc",
+    case "smu_14_0_2": return load("am/smu_14_0_2", [f"{AMD}/pm/swsmu/inc/pmfw_if/{s}.h" for s in ["smu_v14_0_0_pmfw", "smu_v14_0_2_ppsmc",
                                     "smu14_driver_if_v14_0"]]+[root/"extra/amdpci/headers/amdgpu_smu.h"], args=inc, srcs=am_src)
     # firmware hashes
     case "fw":
@@ -34,4 +60,32 @@ def __getattr__(nm):
                                                                            for f in files if (p:=pathlib.Path(f)).is_file()] + ["}"])
       return load("am/fw", ["{}/amdgpu/psp_*_sos.bin", "{}/amdgpu/smu_*.bin", "{}/amdgpu/sdma_*.bin"] +
                            [f"{{}}/amdgpu/gc_*_{x}.bin" for x in ["pfp", "me", "mec", "imu", "rlc"]], srcs=fw_src, gen=genfw)
+    case "navi_offsets": return load("am/navi_offsets", [f"{AMD}/include/sienna_cichlid_ip_offset.h"], srcs=am_src)
+    case "vega_offsets": return load("am/vega_offsets", [f"{AMD}/include/vega20_ip_offset.h"], srcs=am_src)
+    case "regs":
+      def genreg(_, files, **kwargs):
+        out = ["__all__ = " + repr([file.split('/')[-1] for file in files])]
+        for file, nm in [(file.replace("mp_11_0_0", "mp_11_0"), file.split('/')[-1]) for file in files]:
+          pats = reg_patterns[prefix := {"osssys": "oss"}.get(x:=nm.split("_", 1)[0], x)]
+
+          def split_name(name): return name[:(pos:=next((i for i,c in enumerate(name) if c.isupper()), len(name)))], name[pos:]
+          # handle CDNA's different register names
+          def normalize(reg):
+            return s[0] + prefix.upper()[:2] + s[1] if prefix in ("gc", "mmhub") and (s:=split_name(reg))[1].startswith(("VM_", "MC_VM_")) else reg
+          def extract(lines, pat): return ((normalize(m.group(1)), int(m.group(2), 0)) for l in lines if (m:=re.match(pat, l)))
+
+          offset, sh_mask = pathlib.Path(f"{file}_offset.h").read_text().splitlines(), pathlib.Path(f"{file}_sh_mask.h").read_text().splitlines()
+          defs = {k:v for k,v in extract(offset, r'#define\s+((?:mm|reg)\S+)\s+(0x[\da-fA-F]+|\d+)') if any(re.match("(mm|reg)"+p, k) for p in pats)}
+          fields = {reg: {name.split('__')[1].lower(): ((mask & -mask).bit_length() - 1, mask.bit_length() - 1) for name, mask in fs}
+                    for reg, fs in itertools.groupby(extract(sh_mask, r'#define\s+(\S+)_MASK\s+(0x[\da-fA-F]+|\d+)'), lambda x: x[0].split('__')[0])}
+
+          regs = {reg: (off, defs[f"{reg}_BASE_IDX"], fields.get(split_name(reg)[1], {})) for reg,off in defs.items() if f"{reg}_BASE_IDX" in defs}
+          print(f"defined {len(regs)} registers for {nm}")
+          out.extend([f"{nm} = {{"] + [f"  {k!r}: {v!r}," for k,v in regs.items()] + ["}"])
+        return "\n".join(out)
+      return load("am/regs", [AMDINC + "/asic_reg/" + {"osssys":"oss"}.get(pre, pre) + f"/{pre}_{'_'.join(map(str, ver))}"
+                              for pre in reg_files for ver in sorted(reg_files[pre])], srcs=am_src, gen=genreg)
+    case "soc_9" | "soc_11" | "soc_12":
+      return load(f"am/{nm}", ["{}/projects/aqlprofile/linux/" + {9: "vega10", 11: "soc21", 12: "soc24"}[int(nm.split('_')[1])] + "_enum.h"],
+                  srcs=rocm_src, patterns=soc_patterns, macros=False)
     case _: raise AttributeError(f"no such autogen: {nm}")

tinygrad/runtime/autogen/am/navi_offsets.py

@@ -0,0 +1,823 @@
+# mypy: disable-error-code="empty-body"
+from __future__ import annotations
+import ctypes
+from typing import Literal, TypeAlias
+from tinygrad.runtime.support.c import _IO, _IOW, _IOR, _IOWR
+from tinygrad.runtime.support import c
+@c.record
+class struct_IP_BASE_INSTANCE(c.Struct):
+  SIZE = 20
+  segment: c.Array[ctypes.c_uint32, Literal[5]]
+struct_IP_BASE_INSTANCE.register_fields([('segment', c.Array[ctypes.c_uint32, Literal[5]], 0)])
+@c.record
+class struct_IP_BASE(c.Struct):
+  SIZE = 140
+  instance: c.Array[struct_IP_BASE_INSTANCE, Literal[7]]
+struct_IP_BASE.register_fields([('instance', c.Array[struct_IP_BASE_INSTANCE, Literal[7]], 0)])
+MAX_INSTANCE = 7
+MAX_SEGMENT = 5
+ATHUB_BASE__INST0_SEG0 = 0x00000C00
+ATHUB_BASE__INST0_SEG1 = 0x02408C00
+ATHUB_BASE__INST0_SEG2 = 0
+ATHUB_BASE__INST0_SEG3 = 0
+ATHUB_BASE__INST0_SEG4 = 0
+ATHUB_BASE__INST1_SEG0 = 0
+ATHUB_BASE__INST1_SEG1 = 0
+ATHUB_BASE__INST1_SEG2 = 0
+ATHUB_BASE__INST1_SEG3 = 0
+ATHUB_BASE__INST1_SEG4 = 0
+ATHUB_BASE__INST2_SEG0 = 0
+ATHUB_BASE__INST2_SEG1 = 0
+ATHUB_BASE__INST2_SEG2 = 0
+ATHUB_BASE__INST2_SEG3 = 0
+ATHUB_BASE__INST2_SEG4 = 0
+ATHUB_BASE__INST3_SEG0 = 0
+ATHUB_BASE__INST3_SEG1 = 0
+ATHUB_BASE__INST3_SEG2 = 0
+ATHUB_BASE__INST3_SEG3 = 0
+ATHUB_BASE__INST3_SEG4 = 0
+ATHUB_BASE__INST4_SEG0 = 0
+ATHUB_BASE__INST4_SEG1 = 0
+ATHUB_BASE__INST4_SEG2 = 0
+ATHUB_BASE__INST4_SEG3 = 0
+ATHUB_BASE__INST4_SEG4 = 0
+ATHUB_BASE__INST5_SEG0 = 0
+ATHUB_BASE__INST5_SEG1 = 0
+ATHUB_BASE__INST5_SEG2 = 0
+ATHUB_BASE__INST5_SEG3 = 0
+ATHUB_BASE__INST5_SEG4 = 0
+ATHUB_BASE__INST6_SEG0 = 0
+ATHUB_BASE__INST6_SEG1 = 0
+ATHUB_BASE__INST6_SEG2 = 0
+ATHUB_BASE__INST6_SEG3 = 0
+ATHUB_BASE__INST6_SEG4 = 0
+CLK_BASE__INST0_SEG0 = 0x00016C00
+CLK_BASE__INST0_SEG1 = 0x02401800
+CLK_BASE__INST0_SEG2 = 0
+CLK_BASE__INST0_SEG3 = 0
+CLK_BASE__INST0_SEG4 = 0
+CLK_BASE__INST1_SEG0 = 0x00016E00
+CLK_BASE__INST1_SEG1 = 0x02401C00
+CLK_BASE__INST1_SEG2 = 0
+CLK_BASE__INST1_SEG3 = 0
+CLK_BASE__INST1_SEG4 = 0
+CLK_BASE__INST2_SEG0 = 0x00017000
+CLK_BASE__INST2_SEG1 = 0x02402000
+CLK_BASE__INST2_SEG2 = 0
+CLK_BASE__INST2_SEG3 = 0
+CLK_BASE__INST2_SEG4 = 0
+CLK_BASE__INST3_SEG0 = 0x00017200
+CLK_BASE__INST3_SEG1 = 0x02402400
+CLK_BASE__INST3_SEG2 = 0
+CLK_BASE__INST3_SEG3 = 0
+CLK_BASE__INST3_SEG4 = 0
+CLK_BASE__INST4_SEG0 = 0x0001B000
+CLK_BASE__INST4_SEG1 = 0x0242D800
+CLK_BASE__INST4_SEG2 = 0
+CLK_BASE__INST4_SEG3 = 0
+CLK_BASE__INST4_SEG4 = 0
+CLK_BASE__INST5_SEG0 = 0x0001B200
+CLK_BASE__INST5_SEG1 = 0x0242DC00
+CLK_BASE__INST5_SEG2 = 0
+CLK_BASE__INST5_SEG3 = 0
+CLK_BASE__INST5_SEG4 = 0
+CLK_BASE__INST6_SEG0 = 0x0001B400
+CLK_BASE__INST6_SEG1 = 0x0242E000
+CLK_BASE__INST6_SEG2 = 0
+CLK_BASE__INST6_SEG3 = 0
+CLK_BASE__INST6_SEG4 = 0
+DF_BASE__INST0_SEG0 = 0x00007000
+DF_BASE__INST0_SEG1 = 0x0240B800
+DF_BASE__INST0_SEG2 = 0
+DF_BASE__INST0_SEG3 = 0
+DF_BASE__INST0_SEG4 = 0
+DF_BASE__INST1_SEG0 = 0
+DF_BASE__INST1_SEG1 = 0
+DF_BASE__INST1_SEG2 = 0
+DF_BASE__INST1_SEG3 = 0
+DF_BASE__INST1_SEG4 = 0
+DF_BASE__INST2_SEG0 = 0
+DF_BASE__INST2_SEG1 = 0
+DF_BASE__INST2_SEG2 = 0
+DF_BASE__INST2_SEG3 = 0
+DF_BASE__INST2_SEG4 = 0
+DF_BASE__INST3_SEG0 = 0
+DF_BASE__INST3_SEG1 = 0
+DF_BASE__INST3_SEG2 = 0
+DF_BASE__INST3_SEG3 = 0
+DF_BASE__INST3_SEG4 = 0
+DF_BASE__INST4_SEG0 = 0
+DF_BASE__INST4_SEG1 = 0
+DF_BASE__INST4_SEG2 = 0
+DF_BASE__INST4_SEG3 = 0
+DF_BASE__INST4_SEG4 = 0
+DF_BASE__INST5_SEG0 = 0
+DF_BASE__INST5_SEG1 = 0
+DF_BASE__INST5_SEG2 = 0
+DF_BASE__INST5_SEG3 = 0
+DF_BASE__INST5_SEG4 = 0
+DF_BASE__INST6_SEG0 = 0
+DF_BASE__INST6_SEG1 = 0
+DF_BASE__INST6_SEG2 = 0
+DF_BASE__INST6_SEG3 = 0
+DF_BASE__INST6_SEG4 = 0
+DIO_BASE__INST0_SEG0 = 0x02404000
+DIO_BASE__INST0_SEG1 = 0
+DIO_BASE__INST0_SEG2 = 0
+DIO_BASE__INST0_SEG3 = 0
+DIO_BASE__INST0_SEG4 = 0
+DIO_BASE__INST1_SEG0 = 0
+DIO_BASE__INST1_SEG1 = 0
+DIO_BASE__INST1_SEG2 = 0
+DIO_BASE__INST1_SEG3 = 0
+DIO_BASE__INST1_SEG4 = 0
+DIO_BASE__INST2_SEG0 = 0
+DIO_BASE__INST2_SEG1 = 0
+DIO_BASE__INST2_SEG2 = 0
+DIO_BASE__INST2_SEG3 = 0
+DIO_BASE__INST2_SEG4 = 0
+DIO_BASE__INST3_SEG0 = 0
+DIO_BASE__INST3_SEG1 = 0
+DIO_BASE__INST3_SEG2 = 0
+DIO_BASE__INST3_SEG3 = 0
+DIO_BASE__INST3_SEG4 = 0
+DIO_BASE__INST4_SEG0 = 0
+DIO_BASE__INST4_SEG1 = 0
+DIO_BASE__INST4_SEG2 = 0
+DIO_BASE__INST4_SEG3 = 0
+DIO_BASE__INST4_SEG4 = 0
+DIO_BASE__INST5_SEG0 = 0
+DIO_BASE__INST5_SEG1 = 0
+DIO_BASE__INST5_SEG2 = 0
+DIO_BASE__INST5_SEG3 = 0
+DIO_BASE__INST5_SEG4 = 0
+DIO_BASE__INST6_SEG0 = 0
+DIO_BASE__INST6_SEG1 = 0
+DIO_BASE__INST6_SEG2 = 0
+DIO_BASE__INST6_SEG3 = 0
+DIO_BASE__INST6_SEG4 = 0
+DCN_BASE__INST0_SEG0 = 0x00000012
+DCN_BASE__INST0_SEG1 = 0x000000C0
+DCN_BASE__INST0_SEG2 = 0x000034C0
+DCN_BASE__INST0_SEG3 = 0x00009000
+DCN_BASE__INST0_SEG4 = 0x02403C00
+DCN_BASE__INST1_SEG0 = 0
+DCN_BASE__INST1_SEG1 = 0
+DCN_BASE__INST1_SEG2 = 0
+DCN_BASE__INST1_SEG3 = 0
+DCN_BASE__INST1_SEG4 = 0
+DCN_BASE__INST2_SEG0 = 0
+DCN_BASE__INST2_SEG1 = 0
+DCN_BASE__INST2_SEG2 = 0
+DCN_BASE__INST2_SEG3 = 0
+DCN_BASE__INST2_SEG4 = 0
+DCN_BASE__INST3_SEG0 = 0
+DCN_BASE__INST3_SEG1 = 0
+DCN_BASE__INST3_SEG2 = 0
+DCN_BASE__INST3_SEG3 = 0
+DCN_BASE__INST3_SEG4 = 0
+DCN_BASE__INST4_SEG0 = 0
+DCN_BASE__INST4_SEG1 = 0
+DCN_BASE__INST4_SEG2 = 0
+DCN_BASE__INST4_SEG3 = 0
+DCN_BASE__INST4_SEG4 = 0
+DCN_BASE__INST5_SEG0 = 0
+DCN_BASE__INST5_SEG1 = 0
+DCN_BASE__INST5_SEG2 = 0
+DCN_BASE__INST5_SEG3 = 0
+DCN_BASE__INST5_SEG4 = 0
+DCN_BASE__INST6_SEG0 = 0
+DCN_BASE__INST6_SEG1 = 0
+DCN_BASE__INST6_SEG2 = 0
+DCN_BASE__INST6_SEG3 = 0
+DCN_BASE__INST6_SEG4 = 0
+DPCS_BASE__INST0_SEG0 = 0x00000012
+DPCS_BASE__INST0_SEG1 = 0x000000C0
+DPCS_BASE__INST0_SEG2 = 0x000034C0
+DPCS_BASE__INST0_SEG3 = 0x00009000
+DPCS_BASE__INST0_SEG4 = 0x02403C00
+DPCS_BASE__INST1_SEG0 = 0
+DPCS_BASE__INST1_SEG1 = 0
+DPCS_BASE__INST1_SEG2 = 0
+DPCS_BASE__INST1_SEG3 = 0
+DPCS_BASE__INST1_SEG4 = 0
+DPCS_BASE__INST2_SEG0 = 0
+DPCS_BASE__INST2_SEG1 = 0
+DPCS_BASE__INST2_SEG2 = 0
+DPCS_BASE__INST2_SEG3 = 0
+DPCS_BASE__INST2_SEG4 = 0
+DPCS_BASE__INST3_SEG0 = 0
+DPCS_BASE__INST3_SEG1 = 0
+DPCS_BASE__INST3_SEG2 = 0
+DPCS_BASE__INST3_SEG3 = 0
+DPCS_BASE__INST3_SEG4 = 0
+DPCS_BASE__INST4_SEG0 = 0
+DPCS_BASE__INST4_SEG1 = 0
+DPCS_BASE__INST4_SEG2 = 0
+DPCS_BASE__INST4_SEG3 = 0
+DPCS_BASE__INST4_SEG4 = 0
+DPCS_BASE__INST5_SEG0 = 0
+DPCS_BASE__INST5_SEG1 = 0
+DPCS_BASE__INST5_SEG2 = 0
+DPCS_BASE__INST5_SEG3 = 0
+DPCS_BASE__INST5_SEG4 = 0
+DPCS_BASE__INST6_SEG0 = 0
+DPCS_BASE__INST6_SEG1 = 0
+DPCS_BASE__INST6_SEG2 = 0
+DPCS_BASE__INST6_SEG3 = 0
+DPCS_BASE__INST6_SEG4 = 0
+FUSE_BASE__INST0_SEG0 = 0x00017400
+FUSE_BASE__INST0_SEG1 = 0x02401400
+FUSE_BASE__INST0_SEG2 = 0
+FUSE_BASE__INST0_SEG3 = 0
+FUSE_BASE__INST0_SEG4 = 0
+FUSE_BASE__INST1_SEG0 = 0
+FUSE_BASE__INST1_SEG1 = 0
+FUSE_BASE__INST1_SEG2 = 0
+FUSE_BASE__INST1_SEG3 = 0
+FUSE_BASE__INST1_SEG4 = 0
+FUSE_BASE__INST2_SEG0 = 0
+FUSE_BASE__INST2_SEG1 = 0
+FUSE_BASE__INST2_SEG2 = 0
+FUSE_BASE__INST2_SEG3 = 0
+FUSE_BASE__INST2_SEG4 = 0
+FUSE_BASE__INST3_SEG0 = 0
+FUSE_BASE__INST3_SEG1 = 0
+FUSE_BASE__INST3_SEG2 = 0
+FUSE_BASE__INST3_SEG3 = 0
+FUSE_BASE__INST3_SEG4 = 0
+FUSE_BASE__INST4_SEG0 = 0
+FUSE_BASE__INST4_SEG1 = 0
+FUSE_BASE__INST4_SEG2 = 0
+FUSE_BASE__INST4_SEG3 = 0
+FUSE_BASE__INST4_SEG4 = 0
+FUSE_BASE__INST5_SEG0 = 0
+FUSE_BASE__INST5_SEG1 = 0
+FUSE_BASE__INST5_SEG2 = 0
+FUSE_BASE__INST5_SEG3 = 0
+FUSE_BASE__INST5_SEG4 = 0
+FUSE_BASE__INST6_SEG0 = 0
+FUSE_BASE__INST6_SEG1 = 0
+FUSE_BASE__INST6_SEG2 = 0
+FUSE_BASE__INST6_SEG3 = 0
+FUSE_BASE__INST6_SEG4 = 0
+GC_BASE__INST0_SEG0 = 0x00001260
+GC_BASE__INST0_SEG1 = 0x0000A000
+GC_BASE__INST0_SEG2 = 0x0001C000
+GC_BASE__INST0_SEG3 = 0x02402C00
+GC_BASE__INST0_SEG4 = 0
+GC_BASE__INST1_SEG0 = 0
+GC_BASE__INST1_SEG1 = 0
+GC_BASE__INST1_SEG2 = 0
+GC_BASE__INST1_SEG3 = 0
+GC_BASE__INST1_SEG4 = 0
+GC_BASE__INST2_SEG0 = 0
+GC_BASE__INST2_SEG1 = 0
+GC_BASE__INST2_SEG2 = 0
+GC_BASE__INST2_SEG3 = 0
+GC_BASE__INST2_SEG4 = 0
+GC_BASE__INST3_SEG0 = 0
+GC_BASE__INST3_SEG1 = 0
+GC_BASE__INST3_SEG2 = 0
+GC_BASE__INST3_SEG3 = 0
+GC_BASE__INST3_SEG4 = 0
+GC_BASE__INST4_SEG0 = 0
+GC_BASE__INST4_SEG1 = 0
+GC_BASE__INST4_SEG2 = 0
+GC_BASE__INST4_SEG3 = 0
+GC_BASE__INST4_SEG4 = 0
+GC_BASE__INST5_SEG0 = 0
+GC_BASE__INST5_SEG1 = 0
+GC_BASE__INST5_SEG2 = 0
+GC_BASE__INST5_SEG3 = 0
+GC_BASE__INST5_SEG4 = 0
+GC_BASE__INST6_SEG0 = 0
+GC_BASE__INST6_SEG1 = 0
+GC_BASE__INST6_SEG2 = 0
+GC_BASE__INST6_SEG3 = 0
+GC_BASE__INST6_SEG4 = 0
+HDA_BASE__INST0_SEG0 = 0x004C0000
+HDA_BASE__INST0_SEG1 = 0x02404800
+HDA_BASE__INST0_SEG2 = 0
+HDA_BASE__INST0_SEG3 = 0
+HDA_BASE__INST0_SEG4 = 0
+HDA_BASE__INST1_SEG0 = 0
+HDA_BASE__INST1_SEG1 = 0
+HDA_BASE__INST1_SEG2 = 0
+HDA_BASE__INST1_SEG3 = 0
+HDA_BASE__INST1_SEG4 = 0
+HDA_BASE__INST2_SEG0 = 0
+HDA_BASE__INST2_SEG1 = 0
+HDA_BASE__INST2_SEG2 = 0
+HDA_BASE__INST2_SEG3 = 0
+HDA_BASE__INST2_SEG4 = 0
+HDA_BASE__INST3_SEG0 = 0
+HDA_BASE__INST3_SEG1 = 0
+HDA_BASE__INST3_SEG2 = 0
+HDA_BASE__INST3_SEG3 = 0
+HDA_BASE__INST3_SEG4 = 0
+HDA_BASE__INST4_SEG0 = 0
+HDA_BASE__INST4_SEG1 = 0
+HDA_BASE__INST4_SEG2 = 0
+HDA_BASE__INST4_SEG3 = 0
+HDA_BASE__INST4_SEG4 = 0
+HDA_BASE__INST5_SEG0 = 0
+HDA_BASE__INST5_SEG1 = 0
+HDA_BASE__INST5_SEG2 = 0
+HDA_BASE__INST5_SEG3 = 0
+HDA_BASE__INST5_SEG4 = 0
+HDA_BASE__INST6_SEG0 = 0
+HDA_BASE__INST6_SEG1 = 0
+HDA_BASE__INST6_SEG2 = 0
+HDA_BASE__INST6_SEG3 = 0
+HDA_BASE__INST6_SEG4 = 0
+HDP_BASE__INST0_SEG0 = 0x00000F20
+HDP_BASE__INST0_SEG1 = 0x0240A400
+HDP_BASE__INST0_SEG2 = 0
+HDP_BASE__INST0_SEG3 = 0
+HDP_BASE__INST0_SEG4 = 0
+HDP_BASE__INST1_SEG0 = 0
+HDP_BASE__INST1_SEG1 = 0
+HDP_BASE__INST1_SEG2 = 0
+HDP_BASE__INST1_SEG3 = 0
+HDP_BASE__INST1_SEG4 = 0
+HDP_BASE__INST2_SEG0 = 0
+HDP_BASE__INST2_SEG1 = 0
+HDP_BASE__INST2_SEG2 = 0
+HDP_BASE__INST2_SEG3 = 0
+HDP_BASE__INST2_SEG4 = 0
+HDP_BASE__INST3_SEG0 = 0
+HDP_BASE__INST3_SEG1 = 0
+HDP_BASE__INST3_SEG2 = 0
+HDP_BASE__INST3_SEG3 = 0
+HDP_BASE__INST3_SEG4 = 0
+HDP_BASE__INST4_SEG0 = 0
+HDP_BASE__INST4_SEG1 = 0
+HDP_BASE__INST4_SEG2 = 0
+HDP_BASE__INST4_SEG3 = 0
+HDP_BASE__INST4_SEG4 = 0
+HDP_BASE__INST5_SEG0 = 0
+HDP_BASE__INST5_SEG1 = 0
+HDP_BASE__INST5_SEG2 = 0
+HDP_BASE__INST5_SEG3 = 0
+HDP_BASE__INST5_SEG4 = 0
+HDP_BASE__INST6_SEG0 = 0
+HDP_BASE__INST6_SEG1 = 0
+HDP_BASE__INST6_SEG2 = 0
+HDP_BASE__INST6_SEG3 = 0
+HDP_BASE__INST6_SEG4 = 0
+MMHUB_BASE__INST0_SEG0 = 0x0001A000
+MMHUB_BASE__INST0_SEG1 = 0x02408800
+MMHUB_BASE__INST0_SEG2 = 0
+MMHUB_BASE__INST0_SEG3 = 0
+MMHUB_BASE__INST0_SEG4 = 0
+MMHUB_BASE__INST1_SEG0 = 0
+MMHUB_BASE__INST1_SEG1 = 0
+MMHUB_BASE__INST1_SEG2 = 0
+MMHUB_BASE__INST1_SEG3 = 0
+MMHUB_BASE__INST1_SEG4 = 0
+MMHUB_BASE__INST2_SEG0 = 0
+MMHUB_BASE__INST2_SEG1 = 0
+MMHUB_BASE__INST2_SEG2 = 0
+MMHUB_BASE__INST2_SEG3 = 0
+MMHUB_BASE__INST2_SEG4 = 0
+MMHUB_BASE__INST3_SEG0 = 0
+MMHUB_BASE__INST3_SEG1 = 0
+MMHUB_BASE__INST3_SEG2 = 0
+MMHUB_BASE__INST3_SEG3 = 0
+MMHUB_BASE__INST3_SEG4 = 0
+MMHUB_BASE__INST4_SEG0 = 0
+MMHUB_BASE__INST4_SEG1 = 0
+MMHUB_BASE__INST4_SEG2 = 0
+MMHUB_BASE__INST4_SEG3 = 0
+MMHUB_BASE__INST4_SEG4 = 0
+MMHUB_BASE__INST5_SEG0 = 0
+MMHUB_BASE__INST5_SEG1 = 0
+MMHUB_BASE__INST5_SEG2 = 0
+MMHUB_BASE__INST5_SEG3 = 0
+MMHUB_BASE__INST5_SEG4 = 0
+MMHUB_BASE__INST6_SEG0 = 0
+MMHUB_BASE__INST6_SEG1 = 0
+MMHUB_BASE__INST6_SEG2 = 0
+MMHUB_BASE__INST6_SEG3 = 0
+MMHUB_BASE__INST6_SEG4 = 0
+MP0_BASE__INST0_SEG0 = 0x00016000
+MP0_BASE__INST0_SEG1 = 0x00DC0000
+MP0_BASE__INST0_SEG2 = 0x00E00000
+MP0_BASE__INST0_SEG3 = 0x00E40000
+MP0_BASE__INST0_SEG4 = 0x0243FC00
+MP0_BASE__INST1_SEG0 = 0
+MP0_BASE__INST1_SEG1 = 0
+MP0_BASE__INST1_SEG2 = 0
+MP0_BASE__INST1_SEG3 = 0
+MP0_BASE__INST1_SEG4 = 0
+MP0_BASE__INST2_SEG0 = 0
+MP0_BASE__INST2_SEG1 = 0
+MP0_BASE__INST2_SEG2 = 0
+MP0_BASE__INST2_SEG3 = 0
+MP0_BASE__INST2_SEG4 = 0
+MP0_BASE__INST3_SEG0 = 0
+MP0_BASE__INST3_SEG1 = 0
+MP0_BASE__INST3_SEG2 = 0
+MP0_BASE__INST3_SEG3 = 0
+MP0_BASE__INST3_SEG4 = 0
+MP0_BASE__INST4_SEG0 = 0
+MP0_BASE__INST4_SEG1 = 0
+MP0_BASE__INST4_SEG2 = 0
+MP0_BASE__INST4_SEG3 = 0
+MP0_BASE__INST4_SEG4 = 0
+MP0_BASE__INST5_SEG0 = 0
+MP0_BASE__INST5_SEG1 = 0
+MP0_BASE__INST5_SEG2 = 0
+MP0_BASE__INST5_SEG3 = 0
+MP0_BASE__INST5_SEG4 = 0
+MP0_BASE__INST6_SEG0 = 0
+MP0_BASE__INST6_SEG1 = 0
+MP0_BASE__INST6_SEG2 = 0
+MP0_BASE__INST6_SEG3 = 0
+MP0_BASE__INST6_SEG4 = 0
+MP1_BASE__INST0_SEG0 = 0x00016000
+MP1_BASE__INST0_SEG1 = 0x00DC0000
+MP1_BASE__INST0_SEG2 = 0x00E00000
+MP1_BASE__INST0_SEG3 = 0x00E40000
+MP1_BASE__INST0_SEG4 = 0x0243FC00
+MP1_BASE__INST1_SEG0 = 0
+MP1_BASE__INST1_SEG1 = 0
+MP1_BASE__INST1_SEG2 = 0
+MP1_BASE__INST1_SEG3 = 0
+MP1_BASE__INST1_SEG4 = 0
+MP1_BASE__INST2_SEG0 = 0
+MP1_BASE__INST2_SEG1 = 0
+MP1_BASE__INST2_SEG2 = 0
+MP1_BASE__INST2_SEG3 = 0
+MP1_BASE__INST2_SEG4 = 0
+MP1_BASE__INST3_SEG0 = 0
+MP1_BASE__INST3_SEG1 = 0
+MP1_BASE__INST3_SEG2 = 0
+MP1_BASE__INST3_SEG3 = 0
+MP1_BASE__INST3_SEG4 = 0
+MP1_BASE__INST4_SEG0 = 0
+MP1_BASE__INST4_SEG1 = 0
+MP1_BASE__INST4_SEG2 = 0
+MP1_BASE__INST4_SEG3 = 0
+MP1_BASE__INST4_SEG4 = 0
+MP1_BASE__INST5_SEG0 = 0
+MP1_BASE__INST5_SEG1 = 0
+MP1_BASE__INST5_SEG2 = 0
+MP1_BASE__INST5_SEG3 = 0
+MP1_BASE__INST5_SEG4 = 0
+MP1_BASE__INST6_SEG0 = 0
+MP1_BASE__INST6_SEG1 = 0
+MP1_BASE__INST6_SEG2 = 0
+MP1_BASE__INST6_SEG3 = 0
+MP1_BASE__INST6_SEG4 = 0
+NBIO_BASE__INST0_SEG0 = 0x00000000
+NBIO_BASE__INST0_SEG1 = 0x00000014
+NBIO_BASE__INST0_SEG2 = 0x00000D20
+NBIO_BASE__INST0_SEG3 = 0x00010400
+NBIO_BASE__INST0_SEG4 = 0x0241B000
+NBIO_BASE__INST1_SEG0 = 0
+NBIO_BASE__INST1_SEG1 = 0
+NBIO_BASE__INST1_SEG2 = 0
+NBIO_BASE__INST1_SEG3 = 0
+NBIO_BASE__INST1_SEG4 = 0
+NBIO_BASE__INST2_SEG0 = 0
+NBIO_BASE__INST2_SEG1 = 0
+NBIO_BASE__INST2_SEG2 = 0
+NBIO_BASE__INST2_SEG3 = 0
+NBIO_BASE__INST2_SEG4 = 0
+NBIO_BASE__INST3_SEG0 = 0
+NBIO_BASE__INST3_SEG1 = 0
+NBIO_BASE__INST3_SEG2 = 0
+NBIO_BASE__INST3_SEG3 = 0
+NBIO_BASE__INST3_SEG4 = 0
+NBIO_BASE__INST4_SEG0 = 0
+NBIO_BASE__INST4_SEG1 = 0
+NBIO_BASE__INST4_SEG2 = 0
+NBIO_BASE__INST4_SEG3 = 0
+NBIO_BASE__INST4_SEG4 = 0
+NBIO_BASE__INST5_SEG0 = 0
+NBIO_BASE__INST5_SEG1 = 0
+NBIO_BASE__INST5_SEG2 = 0
+NBIO_BASE__INST5_SEG3 = 0
+NBIO_BASE__INST5_SEG4 = 0
+NBIO_BASE__INST6_SEG0 = 0
+NBIO_BASE__INST6_SEG1 = 0
+NBIO_BASE__INST6_SEG2 = 0
+NBIO_BASE__INST6_SEG3 = 0
+NBIO_BASE__INST6_SEG4 = 0
+OSSSYS_BASE__INST0_SEG0 = 0x000010A0
+OSSSYS_BASE__INST0_SEG1 = 0x0240A000
+OSSSYS_BASE__INST0_SEG2 = 0
+OSSSYS_BASE__INST0_SEG3 = 0
+OSSSYS_BASE__INST0_SEG4 = 0
+OSSSYS_BASE__INST1_SEG0 = 0
+OSSSYS_BASE__INST1_SEG1 = 0
+OSSSYS_BASE__INST1_SEG2 = 0
+OSSSYS_BASE__INST1_SEG3 = 0
+OSSSYS_BASE__INST1_SEG4 = 0
+OSSSYS_BASE__INST2_SEG0 = 0
+OSSSYS_BASE__INST2_SEG1 = 0
+OSSSYS_BASE__INST2_SEG2 = 0
+OSSSYS_BASE__INST2_SEG3 = 0
+OSSSYS_BASE__INST2_SEG4 = 0
+OSSSYS_BASE__INST3_SEG0 = 0
+OSSSYS_BASE__INST3_SEG1 = 0
+OSSSYS_BASE__INST3_SEG2 = 0
+OSSSYS_BASE__INST3_SEG3 = 0
+OSSSYS_BASE__INST3_SEG4 = 0
+OSSSYS_BASE__INST4_SEG0 = 0
+OSSSYS_BASE__INST4_SEG1 = 0
+OSSSYS_BASE__INST4_SEG2 = 0
+OSSSYS_BASE__INST4_SEG3 = 0
+OSSSYS_BASE__INST4_SEG4 = 0
+OSSSYS_BASE__INST5_SEG0 = 0
+OSSSYS_BASE__INST5_SEG1 = 0
+OSSSYS_BASE__INST5_SEG2 = 0
+OSSSYS_BASE__INST5_SEG3 = 0
+OSSSYS_BASE__INST5_SEG4 = 0
+OSSSYS_BASE__INST6_SEG0 = 0
+OSSSYS_BASE__INST6_SEG1 = 0
+OSSSYS_BASE__INST6_SEG2 = 0
+OSSSYS_BASE__INST6_SEG3 = 0
+OSSSYS_BASE__INST6_SEG4 = 0
+PCIE0_BASE__INST0_SEG0 = 0x00000000
+PCIE0_BASE__INST0_SEG1 = 0x00000014
+PCIE0_BASE__INST0_SEG2 = 0x00000D20
+PCIE0_BASE__INST0_SEG3 = 0x00010400
+PCIE0_BASE__INST0_SEG4 = 0x0241B000
+PCIE0_BASE__INST1_SEG0 = 0
+PCIE0_BASE__INST1_SEG1 = 0
+PCIE0_BASE__INST1_SEG2 = 0
+PCIE0_BASE__INST1_SEG3 = 0
+PCIE0_BASE__INST1_SEG4 = 0
+PCIE0_BASE__INST2_SEG0 = 0
+PCIE0_BASE__INST2_SEG1 = 0
+PCIE0_BASE__INST2_SEG2 = 0
+PCIE0_BASE__INST2_SEG3 = 0
+PCIE0_BASE__INST2_SEG4 = 0
+PCIE0_BASE__INST3_SEG0 = 0
+PCIE0_BASE__INST3_SEG1 = 0
+PCIE0_BASE__INST3_SEG2 = 0
+PCIE0_BASE__INST3_SEG3 = 0
+PCIE0_BASE__INST3_SEG4 = 0
+PCIE0_BASE__INST4_SEG0 = 0
+PCIE0_BASE__INST4_SEG1 = 0
+PCIE0_BASE__INST4_SEG2 = 0
+PCIE0_BASE__INST4_SEG3 = 0
+PCIE0_BASE__INST4_SEG4 = 0
+PCIE0_BASE__INST5_SEG0 = 0
+PCIE0_BASE__INST5_SEG1 = 0
+PCIE0_BASE__INST5_SEG2 = 0
+PCIE0_BASE__INST5_SEG3 = 0
+PCIE0_BASE__INST5_SEG4 = 0
+PCIE0_BASE__INST6_SEG0 = 0
+PCIE0_BASE__INST6_SEG1 = 0
+PCIE0_BASE__INST6_SEG2 = 0
+PCIE0_BASE__INST6_SEG3 = 0
+PCIE0_BASE__INST6_SEG4 = 0
+SDMA0_BASE__INST0_SEG0 = 0x00001260
+SDMA0_BASE__INST0_SEG1 = 0x0000A000
+SDMA0_BASE__INST0_SEG2 = 0x0001C000
+SDMA0_BASE__INST0_SEG3 = 0x02402C00
+SDMA0_BASE__INST0_SEG4 = 0
+SDMA0_BASE__INST1_SEG0 = 0
+SDMA0_BASE__INST1_SEG1 = 0
+SDMA0_BASE__INST1_SEG2 = 0
+SDMA0_BASE__INST1_SEG3 = 0
+SDMA0_BASE__INST1_SEG4 = 0
+SDMA0_BASE__INST2_SEG0 = 0
+SDMA0_BASE__INST2_SEG1 = 0
+SDMA0_BASE__INST2_SEG2 = 0
+SDMA0_BASE__INST2_SEG3 = 0
+SDMA0_BASE__INST2_SEG4 = 0
+SDMA0_BASE__INST3_SEG0 = 0
+SDMA0_BASE__INST3_SEG1 = 0
+SDMA0_BASE__INST3_SEG2 = 0
+SDMA0_BASE__INST3_SEG3 = 0
+SDMA0_BASE__INST3_SEG4 = 0
+SDMA0_BASE__INST4_SEG0 = 0
+SDMA0_BASE__INST4_SEG1 = 0
+SDMA0_BASE__INST4_SEG2 = 0
+SDMA0_BASE__INST4_SEG3 = 0
+SDMA0_BASE__INST4_SEG4 = 0
+SDMA0_BASE__INST5_SEG0 = 0
+SDMA0_BASE__INST5_SEG1 = 0
+SDMA0_BASE__INST5_SEG2 = 0
+SDMA0_BASE__INST5_SEG3 = 0
+SDMA0_BASE__INST5_SEG4 = 0
+SDMA0_BASE__INST6_SEG0 = 0
+SDMA0_BASE__INST6_SEG1 = 0
+SDMA0_BASE__INST6_SEG2 = 0
+SDMA0_BASE__INST6_SEG3 = 0
+SDMA0_BASE__INST6_SEG4 = 0
+SDMA1_BASE__INST0_SEG0 = 0x00001260
+SDMA1_BASE__INST0_SEG1 = 0x0000A000
+SDMA1_BASE__INST0_SEG2 = 0x0001C000
+SDMA1_BASE__INST0_SEG3 = 0x02402C00
+SDMA1_BASE__INST0_SEG4 = 0
+SDMA1_BASE__INST1_SEG0 = 0
+SDMA1_BASE__INST1_SEG1 = 0
+SDMA1_BASE__INST1_SEG2 = 0
+SDMA1_BASE__INST1_SEG3 = 0
+SDMA1_BASE__INST1_SEG4 = 0
+SDMA1_BASE__INST2_SEG0 = 0
+SDMA1_BASE__INST2_SEG1 = 0
+SDMA1_BASE__INST2_SEG2 = 0
+SDMA1_BASE__INST2_SEG3 = 0
+SDMA1_BASE__INST2_SEG4 = 0
+SDMA1_BASE__INST3_SEG0 = 0
+SDMA1_BASE__INST3_SEG1 = 0
+SDMA1_BASE__INST3_SEG2 = 0
+SDMA1_BASE__INST3_SEG3 = 0
+SDMA1_BASE__INST3_SEG4 = 0
+SDMA1_BASE__INST4_SEG0 = 0
+SDMA1_BASE__INST4_SEG1 = 0
+SDMA1_BASE__INST4_SEG2 = 0
+SDMA1_BASE__INST4_SEG3 = 0
+SDMA1_BASE__INST4_SEG4 = 0
+SDMA1_BASE__INST5_SEG0 = 0
+SDMA1_BASE__INST5_SEG1 = 0
+SDMA1_BASE__INST5_SEG2 = 0
+SDMA1_BASE__INST5_SEG3 = 0
+SDMA1_BASE__INST5_SEG4 = 0
+SDMA1_BASE__INST6_SEG0 = 0
+SDMA1_BASE__INST6_SEG1 = 0
+SDMA1_BASE__INST6_SEG2 = 0
+SDMA1_BASE__INST6_SEG3 = 0
+SDMA1_BASE__INST6_SEG4 = 0
+SMUIO_BASE__INST0_SEG0 = 0x00016800
+SMUIO_BASE__INST0_SEG1 = 0x00016A00
+SMUIO_BASE__INST0_SEG2 = 0x00440000
+SMUIO_BASE__INST0_SEG3 = 0x02401000
+SMUIO_BASE__INST0_SEG4 = 0
+SMUIO_BASE__INST1_SEG0 = 0
+SMUIO_BASE__INST1_SEG1 = 0
+SMUIO_BASE__INST1_SEG2 = 0
+SMUIO_BASE__INST1_SEG3 = 0
+SMUIO_BASE__INST1_SEG4 = 0
+SMUIO_BASE__INST2_SEG0 = 0
+SMUIO_BASE__INST2_SEG1 = 0
+SMUIO_BASE__INST2_SEG2 = 0
+SMUIO_BASE__INST2_SEG3 = 0
+SMUIO_BASE__INST2_SEG4 = 0
+SMUIO_BASE__INST3_SEG0 = 0
+SMUIO_BASE__INST3_SEG1 = 0
+SMUIO_BASE__INST3_SEG2 = 0
+SMUIO_BASE__INST3_SEG3 = 0
+SMUIO_BASE__INST3_SEG4 = 0
+SMUIO_BASE__INST4_SEG0 = 0
+SMUIO_BASE__INST4_SEG1 = 0
+SMUIO_BASE__INST4_SEG2 = 0
+SMUIO_BASE__INST4_SEG3 = 0
+SMUIO_BASE__INST4_SEG4 = 0
+SMUIO_BASE__INST5_SEG0 = 0
+SMUIO_BASE__INST5_SEG1 = 0
+SMUIO_BASE__INST5_SEG2 = 0
+SMUIO_BASE__INST5_SEG3 = 0
+SMUIO_BASE__INST5_SEG4 = 0
+SMUIO_BASE__INST6_SEG0 = 0
+SMUIO_BASE__INST6_SEG1 = 0
+SMUIO_BASE__INST6_SEG2 = 0
+SMUIO_BASE__INST6_SEG3 = 0
+SMUIO_BASE__INST6_SEG4 = 0
+THM_BASE__INST0_SEG0 = 0x00016600
+THM_BASE__INST0_SEG1 = 0x02400C00
+THM_BASE__INST0_SEG2 = 0
+THM_BASE__INST0_SEG3 = 0
+THM_BASE__INST0_SEG4 = 0
+THM_BASE__INST1_SEG0 = 0
+THM_BASE__INST1_SEG1 = 0
+THM_BASE__INST1_SEG2 = 0
+THM_BASE__INST1_SEG3 = 0
+THM_BASE__INST1_SEG4 = 0
+THM_BASE__INST2_SEG0 = 0
+THM_BASE__INST2_SEG1 = 0
+THM_BASE__INST2_SEG2 = 0
+THM_BASE__INST2_SEG3 = 0
+THM_BASE__INST2_SEG4 = 0
+THM_BASE__INST3_SEG0 = 0
+THM_BASE__INST3_SEG1 = 0
+THM_BASE__INST3_SEG2 = 0
+THM_BASE__INST3_SEG3 = 0
+THM_BASE__INST3_SEG4 = 0
+THM_BASE__INST4_SEG0 = 0
+THM_BASE__INST4_SEG1 = 0
+THM_BASE__INST4_SEG2 = 0
+THM_BASE__INST4_SEG3 = 0
+THM_BASE__INST4_SEG4 = 0
+THM_BASE__INST5_SEG0 = 0
+THM_BASE__INST5_SEG1 = 0
+THM_BASE__INST5_SEG2 = 0
+THM_BASE__INST5_SEG3 = 0
+THM_BASE__INST5_SEG4 = 0
+THM_BASE__INST6_SEG0 = 0
+THM_BASE__INST6_SEG1 = 0
+THM_BASE__INST6_SEG2 = 0
+THM_BASE__INST6_SEG3 = 0
+THM_BASE__INST6_SEG4 = 0
+UMC_BASE__INST0_SEG0 = 0x00014000
+UMC_BASE__INST0_SEG1 = 0x02425800
+UMC_BASE__INST0_SEG2 = 0
+UMC_BASE__INST0_SEG3 = 0
+UMC_BASE__INST0_SEG4 = 0
+UMC_BASE__INST1_SEG0 = 0x00054000
+UMC_BASE__INST1_SEG1 = 0x02425C00
+UMC_BASE__INST1_SEG2 = 0
+UMC_BASE__INST1_SEG3 = 0
+UMC_BASE__INST1_SEG4 = 0
+UMC_BASE__INST2_SEG0 = 0x00094000
+UMC_BASE__INST2_SEG1 = 0x02426000
+UMC_BASE__INST2_SEG2 = 0
+UMC_BASE__INST2_SEG3 = 0
+UMC_BASE__INST2_SEG4 = 0
+UMC_BASE__INST3_SEG0 = 0x000D4000
+UMC_BASE__INST3_SEG1 = 0x02426400
+UMC_BASE__INST3_SEG2 = 0
+UMC_BASE__INST3_SEG3 = 0
+UMC_BASE__INST3_SEG4 = 0
+UMC_BASE__INST4_SEG0 = 0x00114000
+UMC_BASE__INST4_SEG1 = 0x02426800
+UMC_BASE__INST4_SEG2 = 0
+UMC_BASE__INST4_SEG3 = 0
+UMC_BASE__INST4_SEG4 = 0
+UMC_BASE__INST5_SEG0 = 0x00154000
+UMC_BASE__INST5_SEG1 = 0x02426C00
+UMC_BASE__INST5_SEG2 = 0
+UMC_BASE__INST5_SEG3 = 0
+UMC_BASE__INST5_SEG4 = 0
+UMC_BASE__INST6_SEG0 = 0x00194000
+UMC_BASE__INST6_SEG1 = 0x02427000
+UMC_BASE__INST6_SEG2 = 0
+UMC_BASE__INST6_SEG3 = 0
+UMC_BASE__INST6_SEG4 = 0
+USB0_BASE__INST0_SEG0 = 0x0242A800
+USB0_BASE__INST0_SEG1 = 0x05B00000
+USB0_BASE__INST0_SEG2 = 0
+USB0_BASE__INST0_SEG3 = 0
+USB0_BASE__INST0_SEG4 = 0
+USB0_BASE__INST1_SEG0 = 0
+USB0_BASE__INST1_SEG1 = 0
+USB0_BASE__INST1_SEG2 = 0
+USB0_BASE__INST1_SEG3 = 0
+USB0_BASE__INST1_SEG4 = 0
+USB0_BASE__INST2_SEG0 = 0
+USB0_BASE__INST2_SEG1 = 0
+USB0_BASE__INST2_SEG2 = 0
+USB0_BASE__INST2_SEG3 = 0
+USB0_BASE__INST2_SEG4 = 0
+USB0_BASE__INST3_SEG0 = 0
+USB0_BASE__INST3_SEG1 = 0
+USB0_BASE__INST3_SEG2 = 0
+USB0_BASE__INST3_SEG3 = 0
+USB0_BASE__INST3_SEG4 = 0
+USB0_BASE__INST4_SEG0 = 0
+USB0_BASE__INST4_SEG1 = 0
+USB0_BASE__INST4_SEG2 = 0
+USB0_BASE__INST4_SEG3 = 0
+USB0_BASE__INST4_SEG4 = 0
+USB0_BASE__INST5_SEG0 = 0
+USB0_BASE__INST5_SEG1 = 0
+USB0_BASE__INST5_SEG2 = 0
+USB0_BASE__INST5_SEG3 = 0
+USB0_BASE__INST5_SEG4 = 0
+USB0_BASE__INST6_SEG0 = 0
+USB0_BASE__INST6_SEG1 = 0
+USB0_BASE__INST6_SEG2 = 0
+USB0_BASE__INST6_SEG3 = 0
+USB0_BASE__INST6_SEG4 = 0
+VCN_BASE__INST0_SEG0 = 0x00007800
+VCN_BASE__INST0_SEG1 = 0x00007E00
+VCN_BASE__INST0_SEG2 = 0x02403000
+VCN_BASE__INST0_SEG3 = 0
+VCN_BASE__INST0_SEG4 = 0
+VCN_BASE__INST1_SEG0 = 0x00007B00
+VCN_BASE__INST1_SEG1 = 0x00012000
+VCN_BASE__INST1_SEG2 = 0x02445000
+VCN_BASE__INST1_SEG3 = 0
+VCN_BASE__INST1_SEG4 = 0
+VCN_BASE__INST2_SEG0 = 0
+VCN_BASE__INST2_SEG1 = 0
+VCN_BASE__INST2_SEG2 = 0
+VCN_BASE__INST2_SEG3 = 0
+VCN_BASE__INST2_SEG4 = 0
+VCN_BASE__INST3_SEG0 = 0
+VCN_BASE__INST3_SEG1 = 0
+VCN_BASE__INST3_SEG2 = 0
+VCN_BASE__INST3_SEG3 = 0
+VCN_BASE__INST3_SEG4 = 0
+VCN_BASE__INST4_SEG0 = 0
+VCN_BASE__INST4_SEG1 = 0
+VCN_BASE__INST4_SEG2 = 0
+VCN_BASE__INST4_SEG3 = 0
+VCN_BASE__INST4_SEG4 = 0
+VCN_BASE__INST5_SEG0 = 0
+VCN_BASE__INST5_SEG1 = 0
+VCN_BASE__INST5_SEG2 = 0
+VCN_BASE__INST5_SEG3 = 0
+VCN_BASE__INST5_SEG4 = 0
+VCN_BASE__INST6_SEG0 = 0
+VCN_BASE__INST6_SEG1 = 0
+VCN_BASE__INST6_SEG2 = 0
+VCN_BASE__INST6_SEG3 = 0
+VCN_BASE__INST6_SEG4 = 0

tinygrad/runtime/autogen/am/soc_11.py

@@ -0,0 +1,44 @@
+# mypy: disable-error-code="empty-body"
+from __future__ import annotations
+import ctypes
+from typing import Literal, TypeAlias
+from tinygrad.runtime.support.c import _IO, _IOW, _IOR, _IOWR
+from tinygrad.runtime.support import c
+enum_MTYPE: dict[int, str] = {(MTYPE_C_RW_US:=0): 'MTYPE_C_RW_US', (MTYPE_RESERVED_1:=1): 'MTYPE_RESERVED_1', (MTYPE_C_RO_S:=2): 'MTYPE_C_RO_S', (MTYPE_UC:=3): 'MTYPE_UC', (MTYPE_C_RW_S:=4): 'MTYPE_C_RW_S', (MTYPE_RESERVED_5:=5): 'MTYPE_RESERVED_5', (MTYPE_C_RO_US:=6): 'MTYPE_C_RO_US', (MTYPE_RESERVED_7:=7): 'MTYPE_RESERVED_7'}
+MTYPE: TypeAlias = ctypes.c_uint32
+enum_SH_MEM_ADDRESS_MODE: dict[int, str] = {(SH_MEM_ADDRESS_MODE_64:=0): 'SH_MEM_ADDRESS_MODE_64', (SH_MEM_ADDRESS_MODE_32:=1): 'SH_MEM_ADDRESS_MODE_32'}
+SH_MEM_ADDRESS_MODE: TypeAlias = ctypes.c_uint32
+enum_SH_MEM_ALIGNMENT_MODE: dict[int, str] = {(SH_MEM_ALIGNMENT_MODE_DWORD:=0): 'SH_MEM_ALIGNMENT_MODE_DWORD', (SH_MEM_ALIGNMENT_MODE_DWORD_STRICT:=1): 'SH_MEM_ALIGNMENT_MODE_DWORD_STRICT', (SH_MEM_ALIGNMENT_MODE_STRICT:=2): 'SH_MEM_ALIGNMENT_MODE_STRICT', (SH_MEM_ALIGNMENT_MODE_UNALIGNED:=3): 'SH_MEM_ALIGNMENT_MODE_UNALIGNED'}
+SH_MEM_ALIGNMENT_MODE: TypeAlias = ctypes.c_uint32
+enum_SQ_TT_MODE: dict[int, str] = {(SQ_TT_MODE_OFF:=0): 'SQ_TT_MODE_OFF', (SQ_TT_MODE_ON:=1): 'SQ_TT_MODE_ON', (SQ_TT_MODE_GLOBAL:=2): 'SQ_TT_MODE_GLOBAL', (SQ_TT_MODE_DETAIL:=3): 'SQ_TT_MODE_DETAIL'}
+SQ_TT_MODE: TypeAlias = ctypes.c_uint32
+enum_SQ_TT_RT_FREQ: dict[int, str] = {(SQ_TT_RT_FREQ_NEVER:=0): 'SQ_TT_RT_FREQ_NEVER', (SQ_TT_RT_FREQ_1024_CLK:=1): 'SQ_TT_RT_FREQ_1024_CLK', (SQ_TT_RT_FREQ_4096_CLK:=2): 'SQ_TT_RT_FREQ_4096_CLK'}
+SQ_TT_RT_FREQ: TypeAlias = ctypes.c_uint32
+enum_SQ_TT_TOKEN_MASK_INST_EXCLUDE: dict[int, str] = {(SQ_TT_INST_EXCLUDE_VMEM_OTHER_SIMD_BIT:=1): 'SQ_TT_INST_EXCLUDE_VMEM_OTHER_SIMD_BIT', (SQ_TT_INST_EXCLUDE_EXPGNT234_BIT:=2): 'SQ_TT_INST_EXCLUDE_EXPGNT234_BIT'}
+SQ_TT_TOKEN_MASK_INST_EXCLUDE: TypeAlias = ctypes.c_uint32
+enum_SQ_TT_TOKEN_MASK_INST_EXCLUDE_SHIFT: dict[int, str] = {(SQ_TT_INST_EXCLUDE_VMEM_OTHER_SIMD_SHIFT:=0): 'SQ_TT_INST_EXCLUDE_VMEM_OTHER_SIMD_SHIFT', (SQ_TT_INST_EXCLUDE_EXPGNT234_SHIFT:=1): 'SQ_TT_INST_EXCLUDE_EXPGNT234_SHIFT'}
+SQ_TT_TOKEN_MASK_INST_EXCLUDE_SHIFT: TypeAlias = ctypes.c_uint32
+enum_SQ_TT_TOKEN_MASK_REG_EXCLUDE: dict[int, str] = {(SQ_TT_REG_EXCLUDE_USER_DATA_BIT:=1): 'SQ_TT_REG_EXCLUDE_USER_DATA_BIT', (SQ_TT_REG_EXCLUDE_CP_ME_MC_RADDR_BIT:=2): 'SQ_TT_REG_EXCLUDE_CP_ME_MC_RADDR_BIT', (SQ_TT_REG_EXCLUDE_GRBM_COMPUTE_EXCLUDE_BIT:=4): 'SQ_TT_REG_EXCLUDE_GRBM_COMPUTE_EXCLUDE_BIT'}
+SQ_TT_TOKEN_MASK_REG_EXCLUDE: TypeAlias = ctypes.c_uint32
+enum_SQ_TT_TOKEN_MASK_REG_EXCLUDE_SHIFT: dict[int, str] = {(SQ_TT_REG_EXCLUDE_USER_DATA_SHIFT:=0): 'SQ_TT_REG_EXCLUDE_USER_DATA_SHIFT', (SQ_TT_REG_EXCLUDE_CP_ME_MC_RADDR_SHIFT:=1): 'SQ_TT_REG_EXCLUDE_CP_ME_MC_RADDR_SHIFT', (SQ_TT_REG_EXCLUDE_GRBM_COMPUTE_EXCLUDE_SHIFT:=2): 'SQ_TT_REG_EXCLUDE_GRBM_COMPUTE_EXCLUDE_SHIFT'}
+SQ_TT_TOKEN_MASK_REG_EXCLUDE_SHIFT: TypeAlias = ctypes.c_uint32
+enum_SQ_TT_TOKEN_MASK_REG_INCLUDE: dict[int, str] = {(SQ_TT_TOKEN_MASK_SQDEC_BIT:=1): 'SQ_TT_TOKEN_MASK_SQDEC_BIT', (SQ_TT_TOKEN_MASK_SHDEC_BIT:=2): 'SQ_TT_TOKEN_MASK_SHDEC_BIT', (SQ_TT_TOKEN_MASK_GFXUDEC_BIT:=4): 'SQ_TT_TOKEN_MASK_GFXUDEC_BIT', (SQ_TT_TOKEN_MASK_COMP_BIT:=8): 'SQ_TT_TOKEN_MASK_COMP_BIT', (SQ_TT_TOKEN_MASK_CONTEXT_BIT:=16): 'SQ_TT_TOKEN_MASK_CONTEXT_BIT', (SQ_TT_TOKEN_MASK_CONFIG_BIT:=32): 'SQ_TT_TOKEN_MASK_CONFIG_BIT', (SQ_TT_TOKEN_MASK_ALL_BIT:=64): 'SQ_TT_TOKEN_MASK_ALL_BIT', (SQ_TT_TOKEN_MASK_RSVD_BIT:=128): 'SQ_TT_TOKEN_MASK_RSVD_BIT'}
+SQ_TT_TOKEN_MASK_REG_INCLUDE: TypeAlias = ctypes.c_uint32
+enum_SQ_TT_TOKEN_MASK_REG_INCLUDE_SHIFT: dict[int, str] = {(SQ_TT_TOKEN_MASK_SQDEC_SHIFT:=0): 'SQ_TT_TOKEN_MASK_SQDEC_SHIFT', (SQ_TT_TOKEN_MASK_SHDEC_SHIFT:=1): 'SQ_TT_TOKEN_MASK_SHDEC_SHIFT', (SQ_TT_TOKEN_MASK_GFXUDEC_SHIFT:=2): 'SQ_TT_TOKEN_MASK_GFXUDEC_SHIFT', (SQ_TT_TOKEN_MASK_COMP_SHIFT:=3): 'SQ_TT_TOKEN_MASK_COMP_SHIFT', (SQ_TT_TOKEN_MASK_CONTEXT_SHIFT:=4): 'SQ_TT_TOKEN_MASK_CONTEXT_SHIFT', (SQ_TT_TOKEN_MASK_CONFIG_SHIFT:=5): 'SQ_TT_TOKEN_MASK_CONFIG_SHIFT', (SQ_TT_TOKEN_MASK_ALL_SHIFT:=6): 'SQ_TT_TOKEN_MASK_ALL_SHIFT', (SQ_TT_TOKEN_MASK_RSVD_SHIFT:=7): 'SQ_TT_TOKEN_MASK_RSVD_SHIFT'}
+SQ_TT_TOKEN_MASK_REG_INCLUDE_SHIFT: TypeAlias = ctypes.c_uint32
+enum_SQ_TT_TOKEN_MASK_TOKEN_EXCLUDE_SHIFT: dict[int, str] = {(SQ_TT_TOKEN_EXCLUDE_VMEMEXEC_SHIFT:=0): 'SQ_TT_TOKEN_EXCLUDE_VMEMEXEC_SHIFT', (SQ_TT_TOKEN_EXCLUDE_ALUEXEC_SHIFT:=1): 'SQ_TT_TOKEN_EXCLUDE_ALUEXEC_SHIFT', (SQ_TT_TOKEN_EXCLUDE_VALUINST_SHIFT:=2): 'SQ_TT_TOKEN_EXCLUDE_VALUINST_SHIFT', (SQ_TT_TOKEN_EXCLUDE_WAVERDY_SHIFT:=3): 'SQ_TT_TOKEN_EXCLUDE_WAVERDY_SHIFT', (SQ_TT_TOKEN_EXCLUDE_WAVESTARTEND_SHIFT:=4): 'SQ_TT_TOKEN_EXCLUDE_WAVESTARTEND_SHIFT', (SQ_TT_TOKEN_EXCLUDE_IMMEDIATE_SHIFT:=5): 'SQ_TT_TOKEN_EXCLUDE_IMMEDIATE_SHIFT', (SQ_TT_TOKEN_EXCLUDE_REG_SHIFT:=6): 'SQ_TT_TOKEN_EXCLUDE_REG_SHIFT', (SQ_TT_TOKEN_EXCLUDE_EVENT_SHIFT:=7): 'SQ_TT_TOKEN_EXCLUDE_EVENT_SHIFT', (SQ_TT_TOKEN_EXCLUDE_INST_SHIFT:=8): 'SQ_TT_TOKEN_EXCLUDE_INST_SHIFT', (SQ_TT_TOKEN_EXCLUDE_UTILCTR_SHIFT:=9): 'SQ_TT_TOKEN_EXCLUDE_UTILCTR_SHIFT', (SQ_TT_TOKEN_EXCLUDE_WAVEALLOC_SHIFT:=10): 'SQ_TT_TOKEN_EXCLUDE_WAVEALLOC_SHIFT', (SQ_TT_TOKEN_EXCLUDE_PERF_SHIFT:=11): 'SQ_TT_TOKEN_EXCLUDE_PERF_SHIFT'}
+SQ_TT_TOKEN_MASK_TOKEN_EXCLUDE_SHIFT: TypeAlias = ctypes.c_uint32
+enum_SQ_TT_UTIL_TIMER: dict[int, str] = {(SQ_TT_UTIL_TIMER_100_CLK:=0): 'SQ_TT_UTIL_TIMER_100_CLK', (SQ_TT_UTIL_TIMER_250_CLK:=1): 'SQ_TT_UTIL_TIMER_250_CLK'}
+SQ_TT_UTIL_TIMER: TypeAlias = ctypes.c_uint32
+enum_SQ_TT_WAVESTART_MODE: dict[int, str] = {(SQ_TT_WAVESTART_MODE_SHORT:=0): 'SQ_TT_WAVESTART_MODE_SHORT', (SQ_TT_WAVESTART_MODE_ALLOC:=1): 'SQ_TT_WAVESTART_MODE_ALLOC', (SQ_TT_WAVESTART_MODE_PBB_ID:=2): 'SQ_TT_WAVESTART_MODE_PBB_ID'}
+SQ_TT_WAVESTART_MODE: TypeAlias = ctypes.c_uint32
+enum_SQ_TT_WTYPE_INCLUDE: dict[int, str] = {(SQ_TT_WTYPE_INCLUDE_PS_BIT:=1): 'SQ_TT_WTYPE_INCLUDE_PS_BIT', (SQ_TT_WTYPE_INCLUDE_RSVD0_BIT:=2): 'SQ_TT_WTYPE_INCLUDE_RSVD0_BIT', (SQ_TT_WTYPE_INCLUDE_GS_BIT:=4): 'SQ_TT_WTYPE_INCLUDE_GS_BIT', (SQ_TT_WTYPE_INCLUDE_RSVD1_BIT:=8): 'SQ_TT_WTYPE_INCLUDE_RSVD1_BIT', (SQ_TT_WTYPE_INCLUDE_HS_BIT:=16): 'SQ_TT_WTYPE_INCLUDE_HS_BIT', (SQ_TT_WTYPE_INCLUDE_RSVD2_BIT:=32): 'SQ_TT_WTYPE_INCLUDE_RSVD2_BIT', (SQ_TT_WTYPE_INCLUDE_CS_BIT:=64): 'SQ_TT_WTYPE_INCLUDE_CS_BIT'}
+SQ_TT_WTYPE_INCLUDE: TypeAlias = ctypes.c_uint32
+enum_SQ_TT_WTYPE_INCLUDE_SHIFT: dict[int, str] = {(SQ_TT_WTYPE_INCLUDE_PS_SHIFT:=0): 'SQ_TT_WTYPE_INCLUDE_PS_SHIFT', (SQ_TT_WTYPE_INCLUDE_RSVD0_SHIFT:=1): 'SQ_TT_WTYPE_INCLUDE_RSVD0_SHIFT', (SQ_TT_WTYPE_INCLUDE_GS_SHIFT:=2): 'SQ_TT_WTYPE_INCLUDE_GS_SHIFT', (SQ_TT_WTYPE_INCLUDE_RSVD1_SHIFT:=3): 'SQ_TT_WTYPE_INCLUDE_RSVD1_SHIFT', (SQ_TT_WTYPE_INCLUDE_HS_SHIFT:=4): 'SQ_TT_WTYPE_INCLUDE_HS_SHIFT', (SQ_TT_WTYPE_INCLUDE_RSVD2_SHIFT:=5): 'SQ_TT_WTYPE_INCLUDE_RSVD2_SHIFT', (SQ_TT_WTYPE_INCLUDE_CS_SHIFT:=6): 'SQ_TT_WTYPE_INCLUDE_CS_SHIFT'}
+SQ_TT_WTYPE_INCLUDE_SHIFT: TypeAlias = ctypes.c_uint32
+enum_CSCNTL_TYPE: dict[int, str] = {(CSCNTL_TYPE_TG:=0): 'CSCNTL_TYPE_TG', (CSCNTL_TYPE_STATE:=1): 'CSCNTL_TYPE_STATE', (CSCNTL_TYPE_EVENT:=2): 'CSCNTL_TYPE_EVENT', (CSCNTL_TYPE_PRIVATE:=3): 'CSCNTL_TYPE_PRIVATE'}
+CSCNTL_TYPE: TypeAlias = ctypes.c_uint32
+enum_CSDATA_TYPE: dict[int, str] = {(CSDATA_TYPE_TG:=0): 'CSDATA_TYPE_TG', (CSDATA_TYPE_STATE:=1): 'CSDATA_TYPE_STATE', (CSDATA_TYPE_EVENT:=2): 'CSDATA_TYPE_EVENT', (CSDATA_TYPE_PRIVATE:=3): 'CSDATA_TYPE_PRIVATE'}
+CSDATA_TYPE: TypeAlias = ctypes.c_uint32
+enum_VGT_EVENT_TYPE: dict[int, str] = {(Reserved_0x00:=0): 'Reserved_0x00', (SAMPLE_STREAMOUTSTATS1:=1): 'SAMPLE_STREAMOUTSTATS1', (SAMPLE_STREAMOUTSTATS2:=2): 'SAMPLE_STREAMOUTSTATS2', (SAMPLE_STREAMOUTSTATS3:=3): 'SAMPLE_STREAMOUTSTATS3', (CACHE_FLUSH_TS:=4): 'CACHE_FLUSH_TS', (CONTEXT_DONE:=5): 'CONTEXT_DONE', (CACHE_FLUSH:=6): 'CACHE_FLUSH', (CS_PARTIAL_FLUSH:=7): 'CS_PARTIAL_FLUSH', (VGT_STREAMOUT_SYNC:=8): 'VGT_STREAMOUT_SYNC', (Reserved_0x09:=9): 'Reserved_0x09', (VGT_STREAMOUT_RESET:=10): 'VGT_STREAMOUT_RESET', (END_OF_PIPE_INCR_DE:=11): 'END_OF_PIPE_INCR_DE', (END_OF_PIPE_IB_END:=12): 'END_OF_PIPE_IB_END', (RST_PIX_CNT:=13): 'RST_PIX_CNT', (BREAK_BATCH:=14): 'BREAK_BATCH', (VS_PARTIAL_FLUSH:=15): 'VS_PARTIAL_FLUSH', (PS_PARTIAL_FLUSH:=16): 'PS_PARTIAL_FLUSH', (FLUSH_HS_OUTPUT:=17): 'FLUSH_HS_OUTPUT', (FLUSH_DFSM:=18): 'FLUSH_DFSM', (RESET_TO_LOWEST_VGT:=19): 'RESET_TO_LOWEST_VGT', (CACHE_FLUSH_AND_INV_TS_EVENT:=20): 'CACHE_FLUSH_AND_INV_TS_EVENT', (WAIT_SYNC:=21): 'WAIT_SYNC', (CACHE_FLUSH_AND_INV_EVENT:=22): 'CACHE_FLUSH_AND_INV_EVENT', (PERFCOUNTER_START:=23): 'PERFCOUNTER_START', (PERFCOUNTER_STOP:=24): 'PERFCOUNTER_STOP', (PIPELINESTAT_START:=25): 'PIPELINESTAT_START', (PIPELINESTAT_STOP:=26): 'PIPELINESTAT_STOP', (PERFCOUNTER_SAMPLE:=27): 'PERFCOUNTER_SAMPLE', (FLUSH_ES_OUTPUT:=28): 'FLUSH_ES_OUTPUT', (BIN_CONF_OVERRIDE_CHECK:=29): 'BIN_CONF_OVERRIDE_CHECK', (SAMPLE_PIPELINESTAT:=30): 'SAMPLE_PIPELINESTAT', (SO_VGTSTREAMOUT_FLUSH:=31): 'SO_VGTSTREAMOUT_FLUSH', (SAMPLE_STREAMOUTSTATS:=32): 'SAMPLE_STREAMOUTSTATS', (RESET_VTX_CNT:=33): 'RESET_VTX_CNT', (BLOCK_CONTEXT_DONE:=34): 'BLOCK_CONTEXT_DONE', (CS_CONTEXT_DONE:=35): 'CS_CONTEXT_DONE', (VGT_FLUSH:=36): 'VGT_FLUSH', (TGID_ROLLOVER:=37): 'TGID_ROLLOVER', (SQ_NON_EVENT:=38): 'SQ_NON_EVENT', (SC_SEND_DB_VPZ:=39): 'SC_SEND_DB_VPZ', (BOTTOM_OF_PIPE_TS:=40): 'BOTTOM_OF_PIPE_TS', (FLUSH_SX_TS:=41): 'FLUSH_SX_TS', (DB_CACHE_FLUSH_AND_INV:=42): 'DB_CACHE_FLUSH_AND_INV', (FLUSH_AND_INV_DB_DATA_TS:=43): 'FLUSH_AND_INV_DB_DATA_TS', (FLUSH_AND_INV_DB_META:=44): 'FLUSH_AND_INV_DB_META', (FLUSH_AND_INV_CB_DATA_TS:=45): 'FLUSH_AND_INV_CB_DATA_TS', (FLUSH_AND_INV_CB_META:=46): 'FLUSH_AND_INV_CB_META', (CS_DONE:=47): 'CS_DONE', (PS_DONE:=48): 'PS_DONE', (FLUSH_AND_INV_CB_PIXEL_DATA:=49): 'FLUSH_AND_INV_CB_PIXEL_DATA', (SX_CB_RAT_ACK_REQUEST:=50): 'SX_CB_RAT_ACK_REQUEST', (THREAD_TRACE_START:=51): 'THREAD_TRACE_START', (THREAD_TRACE_STOP:=52): 'THREAD_TRACE_STOP', (THREAD_TRACE_MARKER:=53): 'THREAD_TRACE_MARKER', (THREAD_TRACE_DRAW:=54): 'THREAD_TRACE_DRAW', (THREAD_TRACE_FINISH:=55): 'THREAD_TRACE_FINISH', (PIXEL_PIPE_STAT_CONTROL:=56): 'PIXEL_PIPE_STAT_CONTROL', (PIXEL_PIPE_STAT_DUMP:=57): 'PIXEL_PIPE_STAT_DUMP', (PIXEL_PIPE_STAT_RESET:=58): 'PIXEL_PIPE_STAT_RESET', (CONTEXT_SUSPEND:=59): 'CONTEXT_SUSPEND', (OFFCHIP_HS_DEALLOC:=60): 'OFFCHIP_HS_DEALLOC', (ENABLE_NGG_PIPELINE:=61): 'ENABLE_NGG_PIPELINE', (ENABLE_LEGACY_PIPELINE:=62): 'ENABLE_LEGACY_PIPELINE', (DRAW_DONE:=63): 'DRAW_DONE'}
+VGT_EVENT_TYPE: TypeAlias = ctypes.c_uint32

tinygrad/runtime/autogen/am/soc_12.py

@@ -0,0 +1,24 @@
+# mypy: disable-error-code="empty-body"
+from __future__ import annotations
+import ctypes
+from typing import Literal, TypeAlias
+from tinygrad.runtime.support.c import _IO, _IOW, _IOR, _IOWR
+from tinygrad.runtime.support import c
+enum_MTYPE: dict[int, str] = {(MTYPE_C_RW_US:=0): 'MTYPE_C_RW_US', (MTYPE_RESERVED_1:=1): 'MTYPE_RESERVED_1', (MTYPE_C_RO_S:=2): 'MTYPE_C_RO_S', (MTYPE_UC:=3): 'MTYPE_UC', (MTYPE_C_RW_S:=4): 'MTYPE_C_RW_S', (MTYPE_RESERVED_5:=5): 'MTYPE_RESERVED_5', (MTYPE_C_RO_US:=6): 'MTYPE_C_RO_US', (MTYPE_RESERVED_7:=7): 'MTYPE_RESERVED_7'}
+MTYPE: TypeAlias = ctypes.c_uint32
+enum_CSCNTL_TYPE: dict[int, str] = {(CSCNTL_TYPE_TG:=0): 'CSCNTL_TYPE_TG', (CSCNTL_TYPE_STATE:=1): 'CSCNTL_TYPE_STATE', (CSCNTL_TYPE_EVENT:=2): 'CSCNTL_TYPE_EVENT', (CSCNTL_TYPE_PRIVATE:=3): 'CSCNTL_TYPE_PRIVATE'}
+CSCNTL_TYPE: TypeAlias = ctypes.c_uint32
+enum_CSDATA_TYPE: dict[int, str] = {(CSDATA_TYPE_TG:=0): 'CSDATA_TYPE_TG', (CSDATA_TYPE_STATE:=1): 'CSDATA_TYPE_STATE', (CSDATA_TYPE_EVENT:=2): 'CSDATA_TYPE_EVENT', (CSDATA_TYPE_PRIVATE:=3): 'CSDATA_TYPE_PRIVATE'}
+CSDATA_TYPE: TypeAlias = ctypes.c_uint32
+enum_VGT_EVENT_TYPE: dict[int, str] = {(Reserved_0x00:=0): 'Reserved_0x00', (SAMPLE_STREAMOUTSTATS1:=1): 'SAMPLE_STREAMOUTSTATS1', (SAMPLE_STREAMOUTSTATS2:=2): 'SAMPLE_STREAMOUTSTATS2', (SAMPLE_STREAMOUTSTATS3:=3): 'SAMPLE_STREAMOUTSTATS3', (CACHE_FLUSH_TS:=4): 'CACHE_FLUSH_TS', (CONTEXT_DONE:=5): 'CONTEXT_DONE', (CACHE_FLUSH:=6): 'CACHE_FLUSH', (CS_PARTIAL_FLUSH:=7): 'CS_PARTIAL_FLUSH', (VGT_STREAMOUT_SYNC:=8): 'VGT_STREAMOUT_SYNC', (EVENT_STATE_CHANGE:=9): 'EVENT_STATE_CHANGE', (VGT_STREAMOUT_RESET:=10): 'VGT_STREAMOUT_RESET', (END_OF_PIPE_INCR_DE:=11): 'END_OF_PIPE_INCR_DE', (END_OF_PIPE_IB_END:=12): 'END_OF_PIPE_IB_END', (RST_PIX_CNT:=13): 'RST_PIX_CNT', (BREAK_BATCH:=14): 'BREAK_BATCH', (VS_PARTIAL_FLUSH:=15): 'VS_PARTIAL_FLUSH', (PS_PARTIAL_FLUSH:=16): 'PS_PARTIAL_FLUSH', (FLUSH_HS_OUTPUT:=17): 'FLUSH_HS_OUTPUT', (FLUSH_DFSM:=18): 'FLUSH_DFSM', (RESET_TO_LOWEST_VGT:=19): 'RESET_TO_LOWEST_VGT', (CACHE_FLUSH_AND_INV_TS_EVENT:=20): 'CACHE_FLUSH_AND_INV_TS_EVENT', (WAIT_SYNC:=21): 'WAIT_SYNC', (CACHE_FLUSH_AND_INV_EVENT:=22): 'CACHE_FLUSH_AND_INV_EVENT', (PERFCOUNTER_START:=23): 'PERFCOUNTER_START', (PERFCOUNTER_STOP:=24): 'PERFCOUNTER_STOP', (PIPELINESTAT_START:=25): 'PIPELINESTAT_START', (PIPELINESTAT_STOP:=26): 'PIPELINESTAT_STOP', (PERFCOUNTER_SAMPLE:=27): 'PERFCOUNTER_SAMPLE', (FLUSH_ES_OUTPUT:=28): 'FLUSH_ES_OUTPUT', (BIN_CONF_OVERRIDE_CHECK:=29): 'BIN_CONF_OVERRIDE_CHECK', (SAMPLE_PIPELINESTAT:=30): 'SAMPLE_PIPELINESTAT', (SO_VGTSTREAMOUT_FLUSH:=31): 'SO_VGTSTREAMOUT_FLUSH', (SAMPLE_STREAMOUTSTATS:=32): 'SAMPLE_STREAMOUTSTATS', (RESET_VTX_CNT:=33): 'RESET_VTX_CNT', (BLOCK_CONTEXT_DONE:=34): 'BLOCK_CONTEXT_DONE', (CS_CONTEXT_DONE:=35): 'CS_CONTEXT_DONE', (VGT_FLUSH:=36): 'VGT_FLUSH', (TGID_ROLLOVER:=37): 'TGID_ROLLOVER', (SQ_NON_EVENT:=38): 'SQ_NON_EVENT', (SC_SEND_DB_VPZ:=39): 'SC_SEND_DB_VPZ', (BOTTOM_OF_PIPE_TS:=40): 'BOTTOM_OF_PIPE_TS', (FLUSH_SX_TS:=41): 'FLUSH_SX_TS', (DB_CACHE_FLUSH_AND_INV:=42): 'DB_CACHE_FLUSH_AND_INV', (FLUSH_AND_INV_DB_DATA_TS:=43): 'FLUSH_AND_INV_DB_DATA_TS', (FLUSH_AND_INV_DB_META:=44): 'FLUSH_AND_INV_DB_META', (FLUSH_AND_INV_CB_DATA_TS:=45): 'FLUSH_AND_INV_CB_DATA_TS', (FLUSH_AND_INV_CB_META:=46): 'FLUSH_AND_INV_CB_META', (CS_DONE:=47): 'CS_DONE', (PS_DONE:=48): 'PS_DONE', (FLUSH_AND_INV_CB_PIXEL_DATA:=49): 'FLUSH_AND_INV_CB_PIXEL_DATA', (SX_CB_RAT_ACK_REQUEST:=50): 'SX_CB_RAT_ACK_REQUEST', (THREAD_TRACE_START:=51): 'THREAD_TRACE_START', (THREAD_TRACE_STOP:=52): 'THREAD_TRACE_STOP', (THREAD_TRACE_MARKER:=53): 'THREAD_TRACE_MARKER', (THREAD_TRACE_DRAW:=54): 'THREAD_TRACE_DRAW', (THREAD_TRACE_FINISH:=55): 'THREAD_TRACE_FINISH', (PIXEL_PIPE_STAT_CONTROL:=56): 'PIXEL_PIPE_STAT_CONTROL', (PIXEL_PIPE_STAT_DUMP:=57): 'PIXEL_PIPE_STAT_DUMP', (PIXEL_PIPE_STAT_RESET:=58): 'PIXEL_PIPE_STAT_RESET', (CONTEXT_SUSPEND:=59): 'CONTEXT_SUSPEND', (OFFCHIP_HS_DEALLOC:=60): 'OFFCHIP_HS_DEALLOC', (ENABLE_NGG_PIPELINE:=61): 'ENABLE_NGG_PIPELINE', (ENABLE_PIPELINE_NOT_USED:=62): 'ENABLE_PIPELINE_NOT_USED', (DRAW_DONE:=63): 'DRAW_DONE'}
+VGT_EVENT_TYPE: TypeAlias = ctypes.c_uint32
+enum_SH_MEM_ADDRESS_MODE: dict[int, str] = {(SH_MEM_ADDRESS_MODE_64:=0): 'SH_MEM_ADDRESS_MODE_64', (SH_MEM_ADDRESS_MODE_32:=1): 'SH_MEM_ADDRESS_MODE_32'}
+SH_MEM_ADDRESS_MODE: TypeAlias = ctypes.c_uint32
+enum_SH_MEM_ALIGNMENT_MODE: dict[int, str] = {(SH_MEM_ALIGNMENT_MODE_DWORD:=0): 'SH_MEM_ALIGNMENT_MODE_DWORD', (SH_MEM_ALIGNMENT_MODE_DWORD_STRICT:=1): 'SH_MEM_ALIGNMENT_MODE_DWORD_STRICT', (SH_MEM_ALIGNMENT_MODE_STRICT:=2): 'SH_MEM_ALIGNMENT_MODE_STRICT', (SH_MEM_ALIGNMENT_MODE_UNALIGNED:=3): 'SH_MEM_ALIGNMENT_MODE_UNALIGNED'}
+SH_MEM_ALIGNMENT_MODE: TypeAlias = ctypes.c_uint32
+enum_SQ_TT_TOKEN_MASK_REG_INCLUDE: dict[int, str] = {(SQ_TT_TOKEN_MASK_SQDEC_BIT:=1): 'SQ_TT_TOKEN_MASK_SQDEC_BIT', (SQ_TT_TOKEN_MASK_SHDEC_BIT:=2): 'SQ_TT_TOKEN_MASK_SHDEC_BIT', (SQ_TT_TOKEN_MASK_GFXUDEC_BIT:=4): 'SQ_TT_TOKEN_MASK_GFXUDEC_BIT', (SQ_TT_TOKEN_MASK_COMP_BIT:=8): 'SQ_TT_TOKEN_MASK_COMP_BIT', (SQ_TT_TOKEN_MASK_CONTEXT_BIT:=16): 'SQ_TT_TOKEN_MASK_CONTEXT_BIT'}
+SQ_TT_TOKEN_MASK_REG_INCLUDE: TypeAlias = ctypes.c_uint32
+enum_SQ_TT_TOKEN_MASK_TOKEN_EXCLUDE_SHIFT: dict[int, str] = {(SQ_TT_TOKEN_EXCLUDE_VMEMEXEC_SHIFT:=0): 'SQ_TT_TOKEN_EXCLUDE_VMEMEXEC_SHIFT', (SQ_TT_TOKEN_EXCLUDE_ALUEXEC_SHIFT:=1): 'SQ_TT_TOKEN_EXCLUDE_ALUEXEC_SHIFT', (SQ_TT_TOKEN_EXCLUDE_WAVERDY_SHIFT:=3): 'SQ_TT_TOKEN_EXCLUDE_WAVERDY_SHIFT'}
+SQ_TT_TOKEN_MASK_TOKEN_EXCLUDE_SHIFT: TypeAlias = ctypes.c_uint32
+enum_SQ_TT_MODE: dict[int, str] = {(SQ_TT_MODE_OFF:=0): 'SQ_TT_MODE_OFF', (SQ_TT_MODE_ON:=1): 'SQ_TT_MODE_ON'}
+SQ_TT_MODE: TypeAlias = ctypes.c_uint32

tinygrad/runtime/autogen/am/soc_9.py

@@ -0,0 +1,16 @@
+# mypy: disable-error-code="empty-body"
+from __future__ import annotations
+import ctypes
+from typing import Literal, TypeAlias
+from tinygrad.runtime.support.c import _IO, _IOW, _IOR, _IOWR
+from tinygrad.runtime.support import c
+enum_MTYPE: dict[int, str] = {(MTYPE_NC:=0): 'MTYPE_NC', (MTYPE_WC:=1): 'MTYPE_WC', (MTYPE_RW:=1): 'MTYPE_RW', (MTYPE_CC:=2): 'MTYPE_CC', (MTYPE_UC:=3): 'MTYPE_UC'}
+MTYPE: TypeAlias = ctypes.c_uint32
+enum_SH_MEM_ADDRESS_MODE: dict[int, str] = {(SH_MEM_ADDRESS_MODE_64:=0): 'SH_MEM_ADDRESS_MODE_64', (SH_MEM_ADDRESS_MODE_32:=1): 'SH_MEM_ADDRESS_MODE_32'}
+SH_MEM_ADDRESS_MODE: TypeAlias = ctypes.c_uint32
+enum_SH_MEM_ALIGNMENT_MODE: dict[int, str] = {(SH_MEM_ALIGNMENT_MODE_DWORD:=0): 'SH_MEM_ALIGNMENT_MODE_DWORD', (SH_MEM_ALIGNMENT_MODE_DWORD_STRICT:=1): 'SH_MEM_ALIGNMENT_MODE_DWORD_STRICT', (SH_MEM_ALIGNMENT_MODE_STRICT:=2): 'SH_MEM_ALIGNMENT_MODE_STRICT', (SH_MEM_ALIGNMENT_MODE_UNALIGNED:=3): 'SH_MEM_ALIGNMENT_MODE_UNALIGNED'}
+SH_MEM_ALIGNMENT_MODE: TypeAlias = ctypes.c_uint32
+enum_CSDATA_TYPE: dict[int, str] = {(CSDATA_TYPE_TG:=0): 'CSDATA_TYPE_TG', (CSDATA_TYPE_STATE:=1): 'CSDATA_TYPE_STATE', (CSDATA_TYPE_EVENT:=2): 'CSDATA_TYPE_EVENT', (CSDATA_TYPE_PRIVATE:=3): 'CSDATA_TYPE_PRIVATE'}
+CSDATA_TYPE: TypeAlias = ctypes.c_uint32
+enum_VGT_EVENT_TYPE: dict[int, str] = {(Reserved_0x00:=0): 'Reserved_0x00', (SAMPLE_STREAMOUTSTATS1:=1): 'SAMPLE_STREAMOUTSTATS1', (SAMPLE_STREAMOUTSTATS2:=2): 'SAMPLE_STREAMOUTSTATS2', (SAMPLE_STREAMOUTSTATS3:=3): 'SAMPLE_STREAMOUTSTATS3', (CACHE_FLUSH_TS:=4): 'CACHE_FLUSH_TS', (CONTEXT_DONE:=5): 'CONTEXT_DONE', (CACHE_FLUSH:=6): 'CACHE_FLUSH', (CS_PARTIAL_FLUSH:=7): 'CS_PARTIAL_FLUSH', (VGT_STREAMOUT_SYNC:=8): 'VGT_STREAMOUT_SYNC', (Reserved_0x09:=9): 'Reserved_0x09', (VGT_STREAMOUT_RESET:=10): 'VGT_STREAMOUT_RESET', (END_OF_PIPE_INCR_DE:=11): 'END_OF_PIPE_INCR_DE', (END_OF_PIPE_IB_END:=12): 'END_OF_PIPE_IB_END', (RST_PIX_CNT:=13): 'RST_PIX_CNT', (BREAK_BATCH:=14): 'BREAK_BATCH', (VS_PARTIAL_FLUSH:=15): 'VS_PARTIAL_FLUSH', (PS_PARTIAL_FLUSH:=16): 'PS_PARTIAL_FLUSH', (FLUSH_HS_OUTPUT:=17): 'FLUSH_HS_OUTPUT', (FLUSH_DFSM:=18): 'FLUSH_DFSM', (RESET_TO_LOWEST_VGT:=19): 'RESET_TO_LOWEST_VGT', (CACHE_FLUSH_AND_INV_TS_EVENT:=20): 'CACHE_FLUSH_AND_INV_TS_EVENT', (ZPASS_DONE:=21): 'ZPASS_DONE', (CACHE_FLUSH_AND_INV_EVENT:=22): 'CACHE_FLUSH_AND_INV_EVENT', (PERFCOUNTER_START:=23): 'PERFCOUNTER_START', (PERFCOUNTER_STOP:=24): 'PERFCOUNTER_STOP', (PIPELINESTAT_START:=25): 'PIPELINESTAT_START', (PIPELINESTAT_STOP:=26): 'PIPELINESTAT_STOP', (PERFCOUNTER_SAMPLE:=27): 'PERFCOUNTER_SAMPLE', (Available_0x1c:=28): 'Available_0x1c', (Available_0x1d:=29): 'Available_0x1d', (SAMPLE_PIPELINESTAT:=30): 'SAMPLE_PIPELINESTAT', (SO_VGTSTREAMOUT_FLUSH:=31): 'SO_VGTSTREAMOUT_FLUSH', (SAMPLE_STREAMOUTSTATS:=32): 'SAMPLE_STREAMOUTSTATS', (RESET_VTX_CNT:=33): 'RESET_VTX_CNT', (BLOCK_CONTEXT_DONE:=34): 'BLOCK_CONTEXT_DONE', (CS_CONTEXT_DONE:=35): 'CS_CONTEXT_DONE', (VGT_FLUSH:=36): 'VGT_FLUSH', (TGID_ROLLOVER:=37): 'TGID_ROLLOVER', (SQ_NON_EVENT:=38): 'SQ_NON_EVENT', (SC_SEND_DB_VPZ:=39): 'SC_SEND_DB_VPZ', (BOTTOM_OF_PIPE_TS:=40): 'BOTTOM_OF_PIPE_TS', (FLUSH_SX_TS:=41): 'FLUSH_SX_TS', (DB_CACHE_FLUSH_AND_INV:=42): 'DB_CACHE_FLUSH_AND_INV', (FLUSH_AND_INV_DB_DATA_TS:=43): 'FLUSH_AND_INV_DB_DATA_TS', (FLUSH_AND_INV_DB_META:=44): 'FLUSH_AND_INV_DB_META', (FLUSH_AND_INV_CB_DATA_TS:=45): 'FLUSH_AND_INV_CB_DATA_TS', (FLUSH_AND_INV_CB_META:=46): 'FLUSH_AND_INV_CB_META', (CS_DONE:=47): 'CS_DONE', (PS_DONE:=48): 'PS_DONE', (FLUSH_AND_INV_CB_PIXEL_DATA:=49): 'FLUSH_AND_INV_CB_PIXEL_DATA', (SX_CB_RAT_ACK_REQUEST:=50): 'SX_CB_RAT_ACK_REQUEST', (THREAD_TRACE_START:=51): 'THREAD_TRACE_START', (THREAD_TRACE_STOP:=52): 'THREAD_TRACE_STOP', (THREAD_TRACE_MARKER:=53): 'THREAD_TRACE_MARKER', (THREAD_TRACE_FLUSH:=54): 'THREAD_TRACE_FLUSH', (THREAD_TRACE_FINISH:=55): 'THREAD_TRACE_FINISH', (PIXEL_PIPE_STAT_CONTROL:=56): 'PIXEL_PIPE_STAT_CONTROL', (PIXEL_PIPE_STAT_DUMP:=57): 'PIXEL_PIPE_STAT_DUMP', (PIXEL_PIPE_STAT_RESET:=58): 'PIXEL_PIPE_STAT_RESET', (CONTEXT_SUSPEND:=59): 'CONTEXT_SUSPEND', (OFFCHIP_HS_DEALLOC:=60): 'OFFCHIP_HS_DEALLOC', (ENABLE_NGG_PIPELINE:=61): 'ENABLE_NGG_PIPELINE', (ENABLE_LEGACY_PIPELINE:=62): 'ENABLE_LEGACY_PIPELINE', (Reserved_0x3f:=63): 'Reserved_0x3f'}
+VGT_EVENT_TYPE: TypeAlias = ctypes.c_uint32

tinygrad/runtime/autogen/am/vega_offsets.py

@@ -0,0 +1,774 @@
+# mypy: disable-error-code="empty-body"
+from __future__ import annotations
+import ctypes
+from typing import Literal, TypeAlias
+from tinygrad.runtime.support.c import _IO, _IOW, _IOR, _IOWR
+from tinygrad.runtime.support import c
+@c.record
+class struct_IP_BASE_INSTANCE(c.Struct):
+  SIZE = 24
+  segment: c.Array[ctypes.c_uint32, Literal[6]]
+struct_IP_BASE_INSTANCE.register_fields([('segment', c.Array[ctypes.c_uint32, Literal[6]], 0)])
+@c.record
+class struct_IP_BASE(c.Struct):
+  SIZE = 144
+  instance: c.Array[struct_IP_BASE_INSTANCE, Literal[6]]
+struct_IP_BASE.register_fields([('instance', c.Array[struct_IP_BASE_INSTANCE, Literal[6]], 0)])
+MAX_INSTANCE = 6
+MAX_SEGMENT = 6
+ATHUB_BASE__INST0_SEG0 = 0x00000C20
+ATHUB_BASE__INST0_SEG1 = 0
+ATHUB_BASE__INST0_SEG2 = 0
+ATHUB_BASE__INST0_SEG3 = 0
+ATHUB_BASE__INST0_SEG4 = 0
+ATHUB_BASE__INST0_SEG5 = 0
+ATHUB_BASE__INST1_SEG0 = 0
+ATHUB_BASE__INST1_SEG1 = 0
+ATHUB_BASE__INST1_SEG2 = 0
+ATHUB_BASE__INST1_SEG3 = 0
+ATHUB_BASE__INST1_SEG4 = 0
+ATHUB_BASE__INST1_SEG5 = 0
+ATHUB_BASE__INST2_SEG0 = 0
+ATHUB_BASE__INST2_SEG1 = 0
+ATHUB_BASE__INST2_SEG2 = 0
+ATHUB_BASE__INST2_SEG3 = 0
+ATHUB_BASE__INST2_SEG4 = 0
+ATHUB_BASE__INST2_SEG5 = 0
+ATHUB_BASE__INST3_SEG0 = 0
+ATHUB_BASE__INST3_SEG1 = 0
+ATHUB_BASE__INST3_SEG2 = 0
+ATHUB_BASE__INST3_SEG3 = 0
+ATHUB_BASE__INST3_SEG4 = 0
+ATHUB_BASE__INST3_SEG5 = 0
+ATHUB_BASE__INST4_SEG0 = 0
+ATHUB_BASE__INST4_SEG1 = 0
+ATHUB_BASE__INST4_SEG2 = 0
+ATHUB_BASE__INST4_SEG3 = 0
+ATHUB_BASE__INST4_SEG4 = 0
+ATHUB_BASE__INST4_SEG5 = 0
+ATHUB_BASE__INST5_SEG0 = 0
+ATHUB_BASE__INST5_SEG1 = 0
+ATHUB_BASE__INST5_SEG2 = 0
+ATHUB_BASE__INST5_SEG3 = 0
+ATHUB_BASE__INST5_SEG4 = 0
+ATHUB_BASE__INST5_SEG5 = 0
+CLK_BASE__INST0_SEG0 = 0x00016C00
+CLK_BASE__INST0_SEG1 = 0x00016E00
+CLK_BASE__INST0_SEG2 = 0x00017000
+CLK_BASE__INST0_SEG3 = 0x00017200
+CLK_BASE__INST0_SEG4 = 0x0001B000
+CLK_BASE__INST0_SEG5 = 0x0001B200
+CLK_BASE__INST1_SEG0 = 0
+CLK_BASE__INST1_SEG1 = 0
+CLK_BASE__INST1_SEG2 = 0
+CLK_BASE__INST1_SEG3 = 0
+CLK_BASE__INST1_SEG4 = 0
+CLK_BASE__INST1_SEG5 = 0
+CLK_BASE__INST2_SEG0 = 0
+CLK_BASE__INST2_SEG1 = 0
+CLK_BASE__INST2_SEG2 = 0
+CLK_BASE__INST2_SEG3 = 0
+CLK_BASE__INST2_SEG4 = 0
+CLK_BASE__INST2_SEG5 = 0
+CLK_BASE__INST3_SEG0 = 0
+CLK_BASE__INST3_SEG1 = 0
+CLK_BASE__INST3_SEG2 = 0
+CLK_BASE__INST3_SEG3 = 0
+CLK_BASE__INST3_SEG4 = 0
+CLK_BASE__INST3_SEG5 = 0
+CLK_BASE__INST4_SEG0 = 0
+CLK_BASE__INST4_SEG1 = 0
+CLK_BASE__INST4_SEG2 = 0
+CLK_BASE__INST4_SEG3 = 0
+CLK_BASE__INST4_SEG4 = 0
+CLK_BASE__INST4_SEG5 = 0
+CLK_BASE__INST5_SEG0 = 0
+CLK_BASE__INST5_SEG1 = 0
+CLK_BASE__INST5_SEG2 = 0
+CLK_BASE__INST5_SEG3 = 0
+CLK_BASE__INST5_SEG4 = 0
+CLK_BASE__INST5_SEG5 = 0
+DCE_BASE__INST0_SEG0 = 0x00000012
+DCE_BASE__INST0_SEG1 = 0x000000C0
+DCE_BASE__INST0_SEG2 = 0x000034C0
+DCE_BASE__INST0_SEG3 = 0
+DCE_BASE__INST0_SEG4 = 0
+DCE_BASE__INST0_SEG5 = 0
+DCE_BASE__INST1_SEG0 = 0
+DCE_BASE__INST1_SEG1 = 0
+DCE_BASE__INST1_SEG2 = 0
+DCE_BASE__INST1_SEG3 = 0
+DCE_BASE__INST1_SEG4 = 0
+DCE_BASE__INST1_SEG5 = 0
+DCE_BASE__INST2_SEG0 = 0
+DCE_BASE__INST2_SEG1 = 0
+DCE_BASE__INST2_SEG2 = 0
+DCE_BASE__INST2_SEG3 = 0
+DCE_BASE__INST2_SEG4 = 0
+DCE_BASE__INST2_SEG5 = 0
+DCE_BASE__INST3_SEG0 = 0
+DCE_BASE__INST3_SEG1 = 0
+DCE_BASE__INST3_SEG2 = 0
+DCE_BASE__INST3_SEG3 = 0
+DCE_BASE__INST3_SEG4 = 0
+DCE_BASE__INST3_SEG5 = 0
+DCE_BASE__INST4_SEG0 = 0
+DCE_BASE__INST4_SEG1 = 0
+DCE_BASE__INST4_SEG2 = 0
+DCE_BASE__INST4_SEG3 = 0
+DCE_BASE__INST4_SEG4 = 0
+DCE_BASE__INST4_SEG5 = 0
+DCE_BASE__INST5_SEG0 = 0
+DCE_BASE__INST5_SEG1 = 0
+DCE_BASE__INST5_SEG2 = 0
+DCE_BASE__INST5_SEG3 = 0
+DCE_BASE__INST5_SEG4 = 0
+DCE_BASE__INST5_SEG5 = 0
+DF_BASE__INST0_SEG0 = 0x00007000
+DF_BASE__INST0_SEG1 = 0
+DF_BASE__INST0_SEG2 = 0
+DF_BASE__INST0_SEG3 = 0
+DF_BASE__INST0_SEG4 = 0
+DF_BASE__INST0_SEG5 = 0
+DF_BASE__INST1_SEG0 = 0
+DF_BASE__INST1_SEG1 = 0
+DF_BASE__INST1_SEG2 = 0
+DF_BASE__INST1_SEG3 = 0
+DF_BASE__INST1_SEG4 = 0
+DF_BASE__INST1_SEG5 = 0
+DF_BASE__INST2_SEG0 = 0
+DF_BASE__INST2_SEG1 = 0
+DF_BASE__INST2_SEG2 = 0
+DF_BASE__INST2_SEG3 = 0
+DF_BASE__INST2_SEG4 = 0
+DF_BASE__INST2_SEG5 = 0
+DF_BASE__INST3_SEG0 = 0
+DF_BASE__INST3_SEG1 = 0
+DF_BASE__INST3_SEG2 = 0
+DF_BASE__INST3_SEG3 = 0
+DF_BASE__INST3_SEG4 = 0
+DF_BASE__INST3_SEG5 = 0
+DF_BASE__INST4_SEG0 = 0
+DF_BASE__INST4_SEG1 = 0
+DF_BASE__INST4_SEG2 = 0
+DF_BASE__INST4_SEG3 = 0
+DF_BASE__INST4_SEG4 = 0
+DF_BASE__INST4_SEG5 = 0
+DF_BASE__INST5_SEG0 = 0
+DF_BASE__INST5_SEG1 = 0
+DF_BASE__INST5_SEG2 = 0
+DF_BASE__INST5_SEG3 = 0
+DF_BASE__INST5_SEG4 = 0
+DF_BASE__INST5_SEG5 = 0
+FUSE_BASE__INST0_SEG0 = 0x00017400
+FUSE_BASE__INST0_SEG1 = 0
+FUSE_BASE__INST0_SEG2 = 0
+FUSE_BASE__INST0_SEG3 = 0
+FUSE_BASE__INST0_SEG4 = 0
+FUSE_BASE__INST0_SEG5 = 0
+FUSE_BASE__INST1_SEG0 = 0
+FUSE_BASE__INST1_SEG1 = 0
+FUSE_BASE__INST1_SEG2 = 0
+FUSE_BASE__INST1_SEG3 = 0
+FUSE_BASE__INST1_SEG4 = 0
+FUSE_BASE__INST1_SEG5 = 0
+FUSE_BASE__INST2_SEG0 = 0
+FUSE_BASE__INST2_SEG1 = 0
+FUSE_BASE__INST2_SEG2 = 0
+FUSE_BASE__INST2_SEG3 = 0
+FUSE_BASE__INST2_SEG4 = 0
+FUSE_BASE__INST2_SEG5 = 0
+FUSE_BASE__INST3_SEG0 = 0
+FUSE_BASE__INST3_SEG1 = 0
+FUSE_BASE__INST3_SEG2 = 0
+FUSE_BASE__INST3_SEG3 = 0
+FUSE_BASE__INST3_SEG4 = 0
+FUSE_BASE__INST3_SEG5 = 0
+FUSE_BASE__INST4_SEG0 = 0
+FUSE_BASE__INST4_SEG1 = 0
+FUSE_BASE__INST4_SEG2 = 0
+FUSE_BASE__INST4_SEG3 = 0
+FUSE_BASE__INST4_SEG4 = 0
+FUSE_BASE__INST4_SEG5 = 0
+FUSE_BASE__INST5_SEG0 = 0
+FUSE_BASE__INST5_SEG1 = 0
+FUSE_BASE__INST5_SEG2 = 0
+FUSE_BASE__INST5_SEG3 = 0
+FUSE_BASE__INST5_SEG4 = 0
+FUSE_BASE__INST5_SEG5 = 0
+GC_BASE__INST0_SEG0 = 0x00002000
+GC_BASE__INST0_SEG1 = 0x0000A000
+GC_BASE__INST0_SEG2 = 0
+GC_BASE__INST0_SEG3 = 0
+GC_BASE__INST0_SEG4 = 0
+GC_BASE__INST0_SEG5 = 0
+GC_BASE__INST1_SEG0 = 0
+GC_BASE__INST1_SEG1 = 0
+GC_BASE__INST1_SEG2 = 0
+GC_BASE__INST1_SEG3 = 0
+GC_BASE__INST1_SEG4 = 0
+GC_BASE__INST1_SEG5 = 0
+GC_BASE__INST2_SEG0 = 0
+GC_BASE__INST2_SEG1 = 0
+GC_BASE__INST2_SEG2 = 0
+GC_BASE__INST2_SEG3 = 0
+GC_BASE__INST2_SEG4 = 0
+GC_BASE__INST2_SEG5 = 0
+GC_BASE__INST3_SEG0 = 0
+GC_BASE__INST3_SEG1 = 0
+GC_BASE__INST3_SEG2 = 0
+GC_BASE__INST3_SEG3 = 0
+GC_BASE__INST3_SEG4 = 0
+GC_BASE__INST3_SEG5 = 0
+GC_BASE__INST4_SEG0 = 0
+GC_BASE__INST4_SEG1 = 0
+GC_BASE__INST4_SEG2 = 0
+GC_BASE__INST4_SEG3 = 0
+GC_BASE__INST4_SEG4 = 0
+GC_BASE__INST4_SEG5 = 0
+GC_BASE__INST5_SEG0 = 0
+GC_BASE__INST5_SEG1 = 0
+GC_BASE__INST5_SEG2 = 0
+GC_BASE__INST5_SEG3 = 0
+GC_BASE__INST5_SEG4 = 0
+GC_BASE__INST5_SEG5 = 0
+HDP_BASE__INST0_SEG0 = 0x00000F20
+HDP_BASE__INST0_SEG1 = 0
+HDP_BASE__INST0_SEG2 = 0
+HDP_BASE__INST0_SEG3 = 0
+HDP_BASE__INST0_SEG4 = 0
+HDP_BASE__INST0_SEG5 = 0
+HDP_BASE__INST1_SEG0 = 0
+HDP_BASE__INST1_SEG1 = 0
+HDP_BASE__INST1_SEG2 = 0
+HDP_BASE__INST1_SEG3 = 0
+HDP_BASE__INST1_SEG4 = 0
+HDP_BASE__INST1_SEG5 = 0
+HDP_BASE__INST2_SEG0 = 0
+HDP_BASE__INST2_SEG1 = 0
+HDP_BASE__INST2_SEG2 = 0
+HDP_BASE__INST2_SEG3 = 0
+HDP_BASE__INST2_SEG4 = 0
+HDP_BASE__INST2_SEG5 = 0
+HDP_BASE__INST3_SEG0 = 0
+HDP_BASE__INST3_SEG1 = 0
+HDP_BASE__INST3_SEG2 = 0
+HDP_BASE__INST3_SEG3 = 0
+HDP_BASE__INST3_SEG4 = 0
+HDP_BASE__INST3_SEG5 = 0
+HDP_BASE__INST4_SEG0 = 0
+HDP_BASE__INST4_SEG1 = 0
+HDP_BASE__INST4_SEG2 = 0
+HDP_BASE__INST4_SEG3 = 0
+HDP_BASE__INST4_SEG4 = 0
+HDP_BASE__INST4_SEG5 = 0
+HDP_BASE__INST5_SEG0 = 0
+HDP_BASE__INST5_SEG1 = 0
+HDP_BASE__INST5_SEG2 = 0
+HDP_BASE__INST5_SEG3 = 0
+HDP_BASE__INST5_SEG4 = 0
+HDP_BASE__INST5_SEG5 = 0
+MMHUB_BASE__INST0_SEG0 = 0x0001A000
+MMHUB_BASE__INST0_SEG1 = 0
+MMHUB_BASE__INST0_SEG2 = 0
+MMHUB_BASE__INST0_SEG3 = 0
+MMHUB_BASE__INST0_SEG4 = 0
+MMHUB_BASE__INST0_SEG5 = 0
+MMHUB_BASE__INST1_SEG0 = 0
+MMHUB_BASE__INST1_SEG1 = 0
+MMHUB_BASE__INST1_SEG2 = 0
+MMHUB_BASE__INST1_SEG3 = 0
+MMHUB_BASE__INST1_SEG4 = 0
+MMHUB_BASE__INST1_SEG5 = 0
+MMHUB_BASE__INST2_SEG0 = 0
+MMHUB_BASE__INST2_SEG1 = 0
+MMHUB_BASE__INST2_SEG2 = 0
+MMHUB_BASE__INST2_SEG3 = 0
+MMHUB_BASE__INST2_SEG4 = 0
+MMHUB_BASE__INST2_SEG5 = 0
+MMHUB_BASE__INST3_SEG0 = 0
+MMHUB_BASE__INST3_SEG1 = 0
+MMHUB_BASE__INST3_SEG2 = 0
+MMHUB_BASE__INST3_SEG3 = 0
+MMHUB_BASE__INST3_SEG4 = 0
+MMHUB_BASE__INST3_SEG5 = 0
+MMHUB_BASE__INST4_SEG0 = 0
+MMHUB_BASE__INST4_SEG1 = 0
+MMHUB_BASE__INST4_SEG2 = 0
+MMHUB_BASE__INST4_SEG3 = 0
+MMHUB_BASE__INST4_SEG4 = 0
+MMHUB_BASE__INST4_SEG5 = 0
+MMHUB_BASE__INST5_SEG0 = 0
+MMHUB_BASE__INST5_SEG1 = 0
+MMHUB_BASE__INST5_SEG2 = 0
+MMHUB_BASE__INST5_SEG3 = 0
+MMHUB_BASE__INST5_SEG4 = 0
+MMHUB_BASE__INST5_SEG5 = 0
+MP0_BASE__INST0_SEG0 = 0x00016000
+MP0_BASE__INST0_SEG1 = 0
+MP0_BASE__INST0_SEG2 = 0
+MP0_BASE__INST0_SEG3 = 0
+MP0_BASE__INST0_SEG4 = 0
+MP0_BASE__INST0_SEG5 = 0
+MP0_BASE__INST1_SEG0 = 0
+MP0_BASE__INST1_SEG1 = 0
+MP0_BASE__INST1_SEG2 = 0
+MP0_BASE__INST1_SEG3 = 0
+MP0_BASE__INST1_SEG4 = 0
+MP0_BASE__INST1_SEG5 = 0
+MP0_BASE__INST2_SEG0 = 0
+MP0_BASE__INST2_SEG1 = 0
+MP0_BASE__INST2_SEG2 = 0
+MP0_BASE__INST2_SEG3 = 0
+MP0_BASE__INST2_SEG4 = 0
+MP0_BASE__INST2_SEG5 = 0
+MP0_BASE__INST3_SEG0 = 0
+MP0_BASE__INST3_SEG1 = 0
+MP0_BASE__INST3_SEG2 = 0
+MP0_BASE__INST3_SEG3 = 0
+MP0_BASE__INST3_SEG4 = 0
+MP0_BASE__INST3_SEG5 = 0
+MP0_BASE__INST4_SEG0 = 0
+MP0_BASE__INST4_SEG1 = 0
+MP0_BASE__INST4_SEG2 = 0
+MP0_BASE__INST4_SEG3 = 0
+MP0_BASE__INST4_SEG4 = 0
+MP0_BASE__INST4_SEG5 = 0
+MP0_BASE__INST5_SEG0 = 0
+MP0_BASE__INST5_SEG1 = 0
+MP0_BASE__INST5_SEG2 = 0
+MP0_BASE__INST5_SEG3 = 0
+MP0_BASE__INST5_SEG4 = 0
+MP0_BASE__INST5_SEG5 = 0
+MP1_BASE__INST0_SEG0 = 0x00016000
+MP1_BASE__INST0_SEG1 = 0
+MP1_BASE__INST0_SEG2 = 0
+MP1_BASE__INST0_SEG3 = 0
+MP1_BASE__INST0_SEG4 = 0
+MP1_BASE__INST0_SEG5 = 0
+MP1_BASE__INST1_SEG0 = 0
+MP1_BASE__INST1_SEG1 = 0
+MP1_BASE__INST1_SEG2 = 0
+MP1_BASE__INST1_SEG3 = 0
+MP1_BASE__INST1_SEG4 = 0
+MP1_BASE__INST1_SEG5 = 0
+MP1_BASE__INST2_SEG0 = 0
+MP1_BASE__INST2_SEG1 = 0
+MP1_BASE__INST2_SEG2 = 0
+MP1_BASE__INST2_SEG3 = 0
+MP1_BASE__INST2_SEG4 = 0
+MP1_BASE__INST2_SEG5 = 0
+MP1_BASE__INST3_SEG0 = 0
+MP1_BASE__INST3_SEG1 = 0
+MP1_BASE__INST3_SEG2 = 0
+MP1_BASE__INST3_SEG3 = 0
+MP1_BASE__INST3_SEG4 = 0
+MP1_BASE__INST3_SEG5 = 0
+MP1_BASE__INST4_SEG0 = 0
+MP1_BASE__INST4_SEG1 = 0
+MP1_BASE__INST4_SEG2 = 0
+MP1_BASE__INST4_SEG3 = 0
+MP1_BASE__INST4_SEG4 = 0
+MP1_BASE__INST4_SEG5 = 0
+MP1_BASE__INST5_SEG0 = 0
+MP1_BASE__INST5_SEG1 = 0
+MP1_BASE__INST5_SEG2 = 0
+MP1_BASE__INST5_SEG3 = 0
+MP1_BASE__INST5_SEG4 = 0
+MP1_BASE__INST5_SEG5 = 0
+NBIO_BASE__INST0_SEG0 = 0x00000000
+NBIO_BASE__INST0_SEG1 = 0x00000014
+NBIO_BASE__INST0_SEG2 = 0x00000D20
+NBIO_BASE__INST0_SEG3 = 0x00010400
+NBIO_BASE__INST0_SEG4 = 0
+NBIO_BASE__INST0_SEG5 = 0
+NBIO_BASE__INST1_SEG0 = 0
+NBIO_BASE__INST1_SEG1 = 0
+NBIO_BASE__INST1_SEG2 = 0
+NBIO_BASE__INST1_SEG3 = 0
+NBIO_BASE__INST1_SEG4 = 0
+NBIO_BASE__INST1_SEG5 = 0
+NBIO_BASE__INST2_SEG0 = 0
+NBIO_BASE__INST2_SEG1 = 0
+NBIO_BASE__INST2_SEG2 = 0
+NBIO_BASE__INST2_SEG3 = 0
+NBIO_BASE__INST2_SEG4 = 0
+NBIO_BASE__INST2_SEG5 = 0
+NBIO_BASE__INST3_SEG0 = 0
+NBIO_BASE__INST3_SEG1 = 0
+NBIO_BASE__INST3_SEG2 = 0
+NBIO_BASE__INST3_SEG3 = 0
+NBIO_BASE__INST3_SEG4 = 0
+NBIO_BASE__INST3_SEG5 = 0
+NBIO_BASE__INST4_SEG0 = 0
+NBIO_BASE__INST4_SEG1 = 0
+NBIO_BASE__INST4_SEG2 = 0
+NBIO_BASE__INST4_SEG3 = 0
+NBIO_BASE__INST4_SEG4 = 0
+NBIO_BASE__INST4_SEG5 = 0
+NBIO_BASE__INST5_SEG0 = 0
+NBIO_BASE__INST5_SEG1 = 0
+NBIO_BASE__INST5_SEG2 = 0
+NBIO_BASE__INST5_SEG3 = 0
+NBIO_BASE__INST5_SEG4 = 0
+NBIO_BASE__INST5_SEG5 = 0
+OSSSYS_BASE__INST0_SEG0 = 0x000010A0
+OSSSYS_BASE__INST0_SEG1 = 0
+OSSSYS_BASE__INST0_SEG2 = 0
+OSSSYS_BASE__INST0_SEG3 = 0
+OSSSYS_BASE__INST0_SEG4 = 0
+OSSSYS_BASE__INST0_SEG5 = 0
+OSSSYS_BASE__INST1_SEG0 = 0
+OSSSYS_BASE__INST1_SEG1 = 0
+OSSSYS_BASE__INST1_SEG2 = 0
+OSSSYS_BASE__INST1_SEG3 = 0
+OSSSYS_BASE__INST1_SEG4 = 0
+OSSSYS_BASE__INST1_SEG5 = 0
+OSSSYS_BASE__INST2_SEG0 = 0
+OSSSYS_BASE__INST2_SEG1 = 0
+OSSSYS_BASE__INST2_SEG2 = 0
+OSSSYS_BASE__INST2_SEG3 = 0
+OSSSYS_BASE__INST2_SEG4 = 0
+OSSSYS_BASE__INST2_SEG5 = 0
+OSSSYS_BASE__INST3_SEG0 = 0
+OSSSYS_BASE__INST3_SEG1 = 0
+OSSSYS_BASE__INST3_SEG2 = 0
+OSSSYS_BASE__INST3_SEG3 = 0
+OSSSYS_BASE__INST3_SEG4 = 0
+OSSSYS_BASE__INST3_SEG5 = 0
+OSSSYS_BASE__INST4_SEG0 = 0
+OSSSYS_BASE__INST4_SEG1 = 0
+OSSSYS_BASE__INST4_SEG2 = 0
+OSSSYS_BASE__INST4_SEG3 = 0
+OSSSYS_BASE__INST4_SEG4 = 0
+OSSSYS_BASE__INST4_SEG5 = 0
+OSSSYS_BASE__INST5_SEG0 = 0
+OSSSYS_BASE__INST5_SEG1 = 0
+OSSSYS_BASE__INST5_SEG2 = 0
+OSSSYS_BASE__INST5_SEG3 = 0
+OSSSYS_BASE__INST5_SEG4 = 0
+OSSSYS_BASE__INST5_SEG5 = 0
+SDMA0_BASE__INST0_SEG0 = 0x00001260
+SDMA0_BASE__INST0_SEG1 = 0
+SDMA0_BASE__INST0_SEG2 = 0
+SDMA0_BASE__INST0_SEG3 = 0
+SDMA0_BASE__INST0_SEG4 = 0
+SDMA0_BASE__INST0_SEG5 = 0
+SDMA0_BASE__INST1_SEG0 = 0
+SDMA0_BASE__INST1_SEG1 = 0
+SDMA0_BASE__INST1_SEG2 = 0
+SDMA0_BASE__INST1_SEG3 = 0
+SDMA0_BASE__INST1_SEG4 = 0
+SDMA0_BASE__INST1_SEG5 = 0
+SDMA0_BASE__INST2_SEG0 = 0
+SDMA0_BASE__INST2_SEG1 = 0
+SDMA0_BASE__INST2_SEG2 = 0
+SDMA0_BASE__INST2_SEG3 = 0
+SDMA0_BASE__INST2_SEG4 = 0
+SDMA0_BASE__INST2_SEG5 = 0
+SDMA0_BASE__INST3_SEG0 = 0
+SDMA0_BASE__INST3_SEG1 = 0
+SDMA0_BASE__INST3_SEG2 = 0
+SDMA0_BASE__INST3_SEG3 = 0
+SDMA0_BASE__INST3_SEG4 = 0
+SDMA0_BASE__INST3_SEG5 = 0
+SDMA0_BASE__INST4_SEG0 = 0
+SDMA0_BASE__INST4_SEG1 = 0
+SDMA0_BASE__INST4_SEG2 = 0
+SDMA0_BASE__INST4_SEG3 = 0
+SDMA0_BASE__INST4_SEG4 = 0
+SDMA0_BASE__INST4_SEG5 = 0
+SDMA0_BASE__INST5_SEG0 = 0
+SDMA0_BASE__INST5_SEG1 = 0
+SDMA0_BASE__INST5_SEG2 = 0
+SDMA0_BASE__INST5_SEG3 = 0
+SDMA0_BASE__INST5_SEG4 = 0
+SDMA0_BASE__INST5_SEG5 = 0
+SDMA1_BASE__INST0_SEG0 = 0x00001860
+SDMA1_BASE__INST0_SEG1 = 0
+SDMA1_BASE__INST0_SEG2 = 0
+SDMA1_BASE__INST0_SEG3 = 0
+SDMA1_BASE__INST0_SEG4 = 0
+SDMA1_BASE__INST0_SEG5 = 0
+SDMA1_BASE__INST1_SEG0 = 0
+SDMA1_BASE__INST1_SEG1 = 0
+SDMA1_BASE__INST1_SEG2 = 0
+SDMA1_BASE__INST1_SEG3 = 0
+SDMA1_BASE__INST1_SEG4 = 0
+SDMA1_BASE__INST1_SEG5 = 0
+SDMA1_BASE__INST2_SEG0 = 0
+SDMA1_BASE__INST2_SEG1 = 0
+SDMA1_BASE__INST2_SEG2 = 0
+SDMA1_BASE__INST2_SEG3 = 0
+SDMA1_BASE__INST2_SEG4 = 0
+SDMA1_BASE__INST2_SEG5 = 0
+SDMA1_BASE__INST3_SEG0 = 0
+SDMA1_BASE__INST3_SEG1 = 0
+SDMA1_BASE__INST3_SEG2 = 0
+SDMA1_BASE__INST3_SEG3 = 0
+SDMA1_BASE__INST3_SEG4 = 0
+SDMA1_BASE__INST3_SEG5 = 0
+SDMA1_BASE__INST4_SEG0 = 0
+SDMA1_BASE__INST4_SEG1 = 0
+SDMA1_BASE__INST4_SEG2 = 0
+SDMA1_BASE__INST4_SEG3 = 0
+SDMA1_BASE__INST4_SEG4 = 0
+SDMA1_BASE__INST4_SEG5 = 0
+SDMA1_BASE__INST5_SEG0 = 0
+SDMA1_BASE__INST5_SEG1 = 0
+SDMA1_BASE__INST5_SEG2 = 0
+SDMA1_BASE__INST5_SEG3 = 0
+SDMA1_BASE__INST5_SEG4 = 0
+SDMA1_BASE__INST5_SEG5 = 0
+SMUIO_BASE__INST0_SEG0 = 0x00016800
+SMUIO_BASE__INST0_SEG1 = 0x00016A00
+SMUIO_BASE__INST0_SEG2 = 0
+SMUIO_BASE__INST0_SEG3 = 0
+SMUIO_BASE__INST0_SEG4 = 0
+SMUIO_BASE__INST0_SEG5 = 0
+SMUIO_BASE__INST1_SEG0 = 0
+SMUIO_BASE__INST1_SEG1 = 0
+SMUIO_BASE__INST1_SEG2 = 0
+SMUIO_BASE__INST1_SEG3 = 0
+SMUIO_BASE__INST1_SEG4 = 0
+SMUIO_BASE__INST1_SEG5 = 0
+SMUIO_BASE__INST2_SEG0 = 0
+SMUIO_BASE__INST2_SEG1 = 0
+SMUIO_BASE__INST2_SEG2 = 0
+SMUIO_BASE__INST2_SEG3 = 0
+SMUIO_BASE__INST2_SEG4 = 0
+SMUIO_BASE__INST2_SEG5 = 0
+SMUIO_BASE__INST3_SEG0 = 0
+SMUIO_BASE__INST3_SEG1 = 0
+SMUIO_BASE__INST3_SEG2 = 0
+SMUIO_BASE__INST3_SEG3 = 0
+SMUIO_BASE__INST3_SEG4 = 0
+SMUIO_BASE__INST3_SEG5 = 0
+SMUIO_BASE__INST4_SEG0 = 0
+SMUIO_BASE__INST4_SEG1 = 0
+SMUIO_BASE__INST4_SEG2 = 0
+SMUIO_BASE__INST4_SEG3 = 0
+SMUIO_BASE__INST4_SEG4 = 0
+SMUIO_BASE__INST4_SEG5 = 0
+SMUIO_BASE__INST5_SEG0 = 0
+SMUIO_BASE__INST5_SEG1 = 0
+SMUIO_BASE__INST5_SEG2 = 0
+SMUIO_BASE__INST5_SEG3 = 0
+SMUIO_BASE__INST5_SEG4 = 0
+SMUIO_BASE__INST5_SEG5 = 0
+THM_BASE__INST0_SEG0 = 0x00016600
+THM_BASE__INST0_SEG1 = 0
+THM_BASE__INST0_SEG2 = 0
+THM_BASE__INST0_SEG3 = 0
+THM_BASE__INST0_SEG4 = 0
+THM_BASE__INST0_SEG5 = 0
+THM_BASE__INST1_SEG0 = 0
+THM_BASE__INST1_SEG1 = 0
+THM_BASE__INST1_SEG2 = 0
+THM_BASE__INST1_SEG3 = 0
+THM_BASE__INST1_SEG4 = 0
+THM_BASE__INST1_SEG5 = 0
+THM_BASE__INST2_SEG0 = 0
+THM_BASE__INST2_SEG1 = 0
+THM_BASE__INST2_SEG2 = 0
+THM_BASE__INST2_SEG3 = 0
+THM_BASE__INST2_SEG4 = 0
+THM_BASE__INST2_SEG5 = 0
+THM_BASE__INST3_SEG0 = 0
+THM_BASE__INST3_SEG1 = 0
+THM_BASE__INST3_SEG2 = 0
+THM_BASE__INST3_SEG3 = 0
+THM_BASE__INST3_SEG4 = 0
+THM_BASE__INST3_SEG5 = 0
+THM_BASE__INST4_SEG0 = 0
+THM_BASE__INST4_SEG1 = 0
+THM_BASE__INST4_SEG2 = 0
+THM_BASE__INST4_SEG3 = 0
+THM_BASE__INST4_SEG4 = 0
+THM_BASE__INST4_SEG5 = 0
+THM_BASE__INST5_SEG0 = 0
+THM_BASE__INST5_SEG1 = 0
+THM_BASE__INST5_SEG2 = 0
+THM_BASE__INST5_SEG3 = 0
+THM_BASE__INST5_SEG4 = 0
+THM_BASE__INST5_SEG5 = 0
+UMC_BASE__INST0_SEG0 = 0x00014000
+UMC_BASE__INST0_SEG1 = 0
+UMC_BASE__INST0_SEG2 = 0
+UMC_BASE__INST0_SEG3 = 0
+UMC_BASE__INST0_SEG4 = 0
+UMC_BASE__INST0_SEG5 = 0
+UMC_BASE__INST1_SEG0 = 0
+UMC_BASE__INST1_SEG1 = 0
+UMC_BASE__INST1_SEG2 = 0
+UMC_BASE__INST1_SEG3 = 0
+UMC_BASE__INST1_SEG4 = 0
+UMC_BASE__INST1_SEG5 = 0
+UMC_BASE__INST2_SEG0 = 0
+UMC_BASE__INST2_SEG1 = 0
+UMC_BASE__INST2_SEG2 = 0
+UMC_BASE__INST2_SEG3 = 0
+UMC_BASE__INST2_SEG4 = 0
+UMC_BASE__INST2_SEG5 = 0
+UMC_BASE__INST3_SEG0 = 0
+UMC_BASE__INST3_SEG1 = 0
+UMC_BASE__INST3_SEG2 = 0
+UMC_BASE__INST3_SEG3 = 0
+UMC_BASE__INST3_SEG4 = 0
+UMC_BASE__INST3_SEG5 = 0
+UMC_BASE__INST4_SEG0 = 0
+UMC_BASE__INST4_SEG1 = 0
+UMC_BASE__INST4_SEG2 = 0
+UMC_BASE__INST4_SEG3 = 0
+UMC_BASE__INST4_SEG4 = 0
+UMC_BASE__INST4_SEG5 = 0
+UMC_BASE__INST5_SEG0 = 0
+UMC_BASE__INST5_SEG1 = 0
+UMC_BASE__INST5_SEG2 = 0
+UMC_BASE__INST5_SEG3 = 0
+UMC_BASE__INST5_SEG4 = 0
+UMC_BASE__INST5_SEG5 = 0
+UVD_BASE__INST0_SEG0 = 0x00007800
+UVD_BASE__INST0_SEG1 = 0x00007E00
+UVD_BASE__INST0_SEG2 = 0
+UVD_BASE__INST0_SEG3 = 0
+UVD_BASE__INST0_SEG4 = 0
+UVD_BASE__INST0_SEG5 = 0
+UVD_BASE__INST1_SEG0 = 0
+UVD_BASE__INST1_SEG1 = 0x00009000
+UVD_BASE__INST1_SEG2 = 0
+UVD_BASE__INST1_SEG3 = 0
+UVD_BASE__INST1_SEG4 = 0
+UVD_BASE__INST1_SEG5 = 0
+UVD_BASE__INST2_SEG0 = 0
+UVD_BASE__INST2_SEG1 = 0
+UVD_BASE__INST2_SEG2 = 0
+UVD_BASE__INST2_SEG3 = 0
+UVD_BASE__INST2_SEG4 = 0
+UVD_BASE__INST2_SEG5 = 0
+UVD_BASE__INST3_SEG0 = 0
+UVD_BASE__INST3_SEG1 = 0
+UVD_BASE__INST3_SEG2 = 0
+UVD_BASE__INST3_SEG3 = 0
+UVD_BASE__INST3_SEG4 = 0
+UVD_BASE__INST3_SEG5 = 0
+UVD_BASE__INST4_SEG0 = 0
+UVD_BASE__INST4_SEG1 = 0
+UVD_BASE__INST4_SEG2 = 0
+UVD_BASE__INST4_SEG3 = 0
+UVD_BASE__INST4_SEG4 = 0
+UVD_BASE__INST4_SEG5 = 0
+UVD_BASE__INST5_SEG0 = 0
+UVD_BASE__INST5_SEG1 = 0
+UVD_BASE__INST5_SEG2 = 0
+UVD_BASE__INST5_SEG3 = 0
+UVD_BASE__INST5_SEG4 = 0
+UVD_BASE__INST5_SEG5 = 0
+VCE_BASE__INST0_SEG0 = 0x00008800
+VCE_BASE__INST0_SEG1 = 0
+VCE_BASE__INST0_SEG2 = 0
+VCE_BASE__INST0_SEG3 = 0
+VCE_BASE__INST0_SEG4 = 0
+VCE_BASE__INST0_SEG5 = 0
+VCE_BASE__INST1_SEG0 = 0
+VCE_BASE__INST1_SEG1 = 0
+VCE_BASE__INST1_SEG2 = 0
+VCE_BASE__INST1_SEG3 = 0
+VCE_BASE__INST1_SEG4 = 0
+VCE_BASE__INST1_SEG5 = 0
+VCE_BASE__INST2_SEG0 = 0
+VCE_BASE__INST2_SEG1 = 0
+VCE_BASE__INST2_SEG2 = 0
+VCE_BASE__INST2_SEG3 = 0
+VCE_BASE__INST2_SEG4 = 0
+VCE_BASE__INST2_SEG5 = 0
+VCE_BASE__INST3_SEG0 = 0
+VCE_BASE__INST3_SEG1 = 0
+VCE_BASE__INST3_SEG2 = 0
+VCE_BASE__INST3_SEG3 = 0
+VCE_BASE__INST3_SEG4 = 0
+VCE_BASE__INST3_SEG5 = 0
+VCE_BASE__INST4_SEG0 = 0
+VCE_BASE__INST4_SEG1 = 0
+VCE_BASE__INST4_SEG2 = 0
+VCE_BASE__INST4_SEG3 = 0
+VCE_BASE__INST4_SEG4 = 0
+VCE_BASE__INST4_SEG5 = 0
+VCE_BASE__INST5_SEG0 = 0
+VCE_BASE__INST5_SEG1 = 0
+VCE_BASE__INST5_SEG2 = 0
+VCE_BASE__INST5_SEG3 = 0
+VCE_BASE__INST5_SEG4 = 0
+VCE_BASE__INST5_SEG5 = 0
+XDMA_BASE__INST0_SEG0 = 0x00003400
+XDMA_BASE__INST0_SEG1 = 0
+XDMA_BASE__INST0_SEG2 = 0
+XDMA_BASE__INST0_SEG3 = 0
+XDMA_BASE__INST0_SEG4 = 0
+XDMA_BASE__INST0_SEG5 = 0
+XDMA_BASE__INST1_SEG0 = 0
+XDMA_BASE__INST1_SEG1 = 0
+XDMA_BASE__INST1_SEG2 = 0
+XDMA_BASE__INST1_SEG3 = 0
+XDMA_BASE__INST1_SEG4 = 0
+XDMA_BASE__INST1_SEG5 = 0
+XDMA_BASE__INST2_SEG0 = 0
+XDMA_BASE__INST2_SEG1 = 0
+XDMA_BASE__INST2_SEG2 = 0
+XDMA_BASE__INST2_SEG3 = 0
+XDMA_BASE__INST2_SEG4 = 0
+XDMA_BASE__INST2_SEG5 = 0
+XDMA_BASE__INST3_SEG0 = 0
+XDMA_BASE__INST3_SEG1 = 0
+XDMA_BASE__INST3_SEG2 = 0
+XDMA_BASE__INST3_SEG3 = 0
+XDMA_BASE__INST3_SEG4 = 0
+XDMA_BASE__INST3_SEG5 = 0
+XDMA_BASE__INST4_SEG0 = 0
+XDMA_BASE__INST4_SEG1 = 0
+XDMA_BASE__INST4_SEG2 = 0
+XDMA_BASE__INST4_SEG3 = 0
+XDMA_BASE__INST4_SEG4 = 0
+XDMA_BASE__INST4_SEG5 = 0
+XDMA_BASE__INST5_SEG0 = 0
+XDMA_BASE__INST5_SEG1 = 0
+XDMA_BASE__INST5_SEG2 = 0
+XDMA_BASE__INST5_SEG3 = 0
+XDMA_BASE__INST5_SEG4 = 0
+XDMA_BASE__INST5_SEG5 = 0
+RSMU_BASE__INST0_SEG0 = 0x00012000
+RSMU_BASE__INST0_SEG1 = 0
+RSMU_BASE__INST0_SEG2 = 0
+RSMU_BASE__INST0_SEG3 = 0
+RSMU_BASE__INST0_SEG4 = 0
+RSMU_BASE__INST0_SEG5 = 0
+RSMU_BASE__INST1_SEG0 = 0
+RSMU_BASE__INST1_SEG1 = 0
+RSMU_BASE__INST1_SEG2 = 0
+RSMU_BASE__INST1_SEG3 = 0
+RSMU_BASE__INST1_SEG4 = 0
+RSMU_BASE__INST1_SEG5 = 0
+RSMU_BASE__INST2_SEG0 = 0
+RSMU_BASE__INST2_SEG1 = 0
+RSMU_BASE__INST2_SEG2 = 0
+RSMU_BASE__INST2_SEG3 = 0
+RSMU_BASE__INST2_SEG4 = 0
+RSMU_BASE__INST2_SEG5 = 0
+RSMU_BASE__INST3_SEG0 = 0
+RSMU_BASE__INST3_SEG1 = 0
+RSMU_BASE__INST3_SEG2 = 0
+RSMU_BASE__INST3_SEG3 = 0
+RSMU_BASE__INST3_SEG4 = 0
+RSMU_BASE__INST3_SEG5 = 0
+RSMU_BASE__INST4_SEG0 = 0
+RSMU_BASE__INST4_SEG1 = 0
+RSMU_BASE__INST4_SEG2 = 0
+RSMU_BASE__INST4_SEG3 = 0
+RSMU_BASE__INST4_SEG4 = 0
+RSMU_BASE__INST4_SEG5 = 0
+RSMU_BASE__INST5_SEG0 = 0
+RSMU_BASE__INST5_SEG1 = 0
+RSMU_BASE__INST5_SEG2 = 0
+RSMU_BASE__INST5_SEG3 = 0
+RSMU_BASE__INST5_SEG4 = 0
+RSMU_BASE__INST5_SEG5 = 0

tinygrad/runtime/graph/hcq.py

@@ -97,7 +97,7 @@ class HCQGraph(MultiGraphRunner):
 
       # set any fixedvars on the device
       self.device_vars[enqueue_dev] = merge_dicts([self.device_vars.get(enqueue_dev, {}), device_vars])
-      if runtime is not None: self.device_vars[enqueue_dev] = merge_dicts([self.device_vars[enqueue_dev], ast.arg.runtimevars])
+      if runtime is not None: self.device_vars[enqueue_dev] = merge_dicts([self.device_vars[enqueue_dev], {k: 0 for k in ast.arg.runtimevars}])
 
       if runtime is not None:
         enqueue_queue = self.comp_queues[enqueue_dev]
@@ -316,7 +316,7 @@ class HCQGraph(MultiGraphRunner):
     for fdev, buf in self.kernargs_bufs.items(): fdev.allocator._free(buf, BufferSpec(cpu_access=True))
 
   @staticmethod
-  def supports_exec_item(batch_devs:list[Compiled], new_call:UOp) -> bool:
+  def supports_uop(batch_devs:list[Compiled], new_call:UOp) -> bool:
     # Check if all devices are HCQ
     all_devs = cast(list[HCQCompiled], GraphRunner._all_devs(batch_devs, new_call))
     if not all(issubclass(type(d), HCQCompiled) for d in all_devs): return False

tinygrad/runtime/graph/metal.py

@@ -107,7 +107,7 @@ class MetalGraph(GraphRunner):
       self.collect_timestamps()
 
   @staticmethod
-  def supports_exec_item(batch_devs, new_call:UOp) -> bool:
+  def supports_uop(batch_devs, new_call:UOp) -> bool:
     # Metal ICB replay encodes offsets as uint32; reject if any Metal buffer offset exceeds 32-bit range.
     if any(b.op is Ops.BUFFER_VIEW and b.arg[1] * b.dtype.itemsize > 0xFFFFFFFF for b in new_call.src[1:]): return False
-    return GraphRunner.supports_exec_item(batch_devs, new_call)
+    return GraphRunner.supports_uop(batch_devs, new_call)

tinygrad/runtime/ops_amd.py

@@ -15,7 +15,7 @@ from tinygrad.runtime.autogen import kfd, hsa, sqtt, amdgpu_kd, amdgpu_drm
 from tinygrad.runtime.autogen.am import am
 from tinygrad.runtime.support.elf import elf_loader
 from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager
-from tinygrad.runtime.support.amd import AMDReg, AMDIP, import_module, import_soc, import_ip_offsets, import_pmc
+from tinygrad.runtime.support.amd import AMDReg, AMDIP, import_module, import_soc, import_pmc
 from tinygrad.runtime.support.system import System, PCIIfaceBase, PCIAllocationMeta, USBPCIDevice, MAP_FIXED, MAP_NORESERVE
 from tinygrad.runtime.support.usb import USB3
 from tinygrad.runtime.support.memory import AddrSpace
@@ -90,7 +90,7 @@ class AMDComputeQueue(HWQueue):
     return self
 
   def acquire_mem(self, addr=0x0, sz=(1 << 64)-1, gli=1, glm=1, glk=1, glv=1, gl1=1, gl2=1):
-    if self.dev.target >= (10,0,0):
+    if self.dev.target[0] != 9:
       cache_flags_dw = self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLI_INV(gli) \
                      | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLM_INV(glm) | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLM_WB(glm) \
                      | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLK_INV(glk) | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLK_WB(glk) \
@@ -108,7 +108,7 @@ class AMDComputeQueue(HWQueue):
     return self
 
   def release_mem(self, address=0x0, value=0, data_sel=0, int_sel=2, ctxid=0, cache_flush=False):
-    if self.dev.target >= (10,0,0):
+    if self.dev.target[0] != 9:
       cache_flags_dw = 0 if not cache_flush else (self.pm4.PACKET3_RELEASE_MEM_GCR_GLV_INV | self.pm4.PACKET3_RELEASE_MEM_GCR_GL1_INV \
                      | self.pm4.PACKET3_RELEASE_MEM_GCR_GL2_INV | self.pm4.PACKET3_RELEASE_MEM_GCR_GLM_WB \
                      | self.pm4.PACKET3_RELEASE_MEM_GCR_GLM_INV | self.pm4.PACKET3_RELEASE_MEM_GCR_GL2_WB | self.pm4.PACKET3_RELEASE_MEM_GCR_SEQ)
@@ -131,7 +131,7 @@ class AMDComputeQueue(HWQueue):
     return self
 
   def memory_barrier(self):
-    pf = '' if self.nbio.version[0] == 2 else '0' if self.nbio.version[:2] != (7, 11) else '1'
+    pf = '0' if self.nbio.version[:2] != (7, 11) else '1'
     self.wait_reg_mem(reg=getattr(self.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_REQ').addr[0],
                       reg_done=getattr(self.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_DONE').addr[0], value=0xffffffff)
     return self.acquire_mem()
@@ -151,7 +151,7 @@ class AMDComputeQueue(HWQueue):
   def pmc_start(self, counters):
     self.pmc_reset_counters(en=False)
     self.wreg(self.gc.regSQ_PERFCOUNTER_CTRL, cs_en=1, ps_en=1, gs_en=1, hs_en=1, **({'vmid_mask':0xffff} if (gfx9:=self.dev.target[0] == 9) else {}))
-    if self.dev.target[0] >= 11: self.wreg(self.gc.regSQ_PERFCOUNTER_CTRL2, force_en=1, vmid_en=0xffff)
+    if not gfx9: self.wreg(self.gc.regSQ_PERFCOUNTER_CTRL2, force_en=1, vmid_en=0xffff)
 
     end_off = 0
     block2pid:dict[str, itertools.count] = collections.defaultdict(lambda: itertools.count())
@@ -163,7 +163,7 @@ class AMDComputeQueue(HWQueue):
 
       # gfx11+ and later require even-numbered SQ *_SELECT registers
       regsample = f'reg{block}_PERFCOUNTER{(pcid:=next(block2pid[block]))}'
-      if (regsel:=getattr(self.gc, (f'reg{block}_PERFCOUNTER{(pcid*2) if self.dev.target[0]>=11 and block=="SQ" else pcid}_SELECT'), None)) is None:
+      if (regsel:=getattr(self.gc, (f'reg{block}_PERFCOUNTER{(pcid*2) if not gfx9 and block=="SQ" else pcid}_SELECT'), None)) is None:
         raise RuntimeError(f'{block} is out of perfcounter registers: ({regsample} is not found)')
 
       self.wreg(regsel, perf_sel=idx, **({'simd_mask':0xf, 'sqc_bank_mask':0xf, 'sqc_client_mask':0xf} if gfx9 and block == "SQ" else {}))
@@ -206,7 +206,7 @@ class AMDComputeQueue(HWQueue):
       cu_per_se = prod([x if isinstance(x, int) else 1 for x in global_size]) // ((self.dev.cu_cnt // self.dev.se_cnt) * 4)
       for xcc in range(self.dev.xccs):
         with self.pred_exec(xcc_mask=1 << xcc):
-          for i in range(8 if prg.dev.target >= (11,0,0) else 4):
+          for i in range(8 if prg.dev.target[0] != 9 else 4):
             if SQTT_LIMIT_SE > 1: mask = 1 if SQTT_ITRACE_SE_MASK.value & (1 << i) else 0 # only run unmasked shader engines
             else:
               sa_mask = (1 << (self.dev.iface.props['cu_per_simd_array'] // 2)) - 1
@@ -282,7 +282,7 @@ class AMDComputeQueue(HWQueue):
         self.sqtt_config(tracing=True)
 
     self.set_grbm()
-    if self.dev.target[0] > 9: self.wreg(self.gc.regCOMPUTE_THREAD_TRACE_ENABLE, 1)
+    if self.dev.target[0] != 9: self.wreg(self.gc.regCOMPUTE_THREAD_TRACE_ENABLE, 1)
     self.memory_barrier()
     return self
 
@@ -303,7 +303,7 @@ class AMDComputeQueue(HWQueue):
         self.set_grbm(se=se % self.dev.se_cnt, sh=0)
 
         regstatus = self.gc.regSQ_THREAD_TRACE_STATUS.addr[0] - (self.pm4.PACKET3_SET_UCONFIG_REG_START if self.dev.target[0] == 9 else 0)
-        if self.dev.target >= (10,0,0):
+        if self.dev.target[0] != 9:
           self.wait_reg_mem(reg=regstatus, mask=self.gc.regSQ_THREAD_TRACE_STATUS.fields_mask('finish_pending'), op=WAIT_REG_MEM_FUNCTION_EQ, value=0)
           self.sqtt_config(tracing=False)
         self.wait_reg_mem(reg=regstatus, mask=self.gc.regSQ_THREAD_TRACE_STATUS.fields_mask('busy'), op=WAIT_REG_MEM_FUNCTION_EQ, value=0)
@@ -313,7 +313,7 @@ class AMDComputeQueue(HWQueue):
         self.pkt3(self.pm4.PACKET3_COPY_DATA, 1 << 20 | 2 << 8 | 4, self.gc.regSQ_THREAD_TRACE_WPTR.addr[0], 0, *data64_le(wptrs.va_addr+(se*4)))
 
     self.set_grbm()
-    if self.dev.target[0] > 9: self.spi_config(tracing=False)
+    if self.dev.target[0] != 9: self.spi_config(tracing=False)
     self.memory_barrier()
     return self
 
@@ -348,22 +348,20 @@ class AMDComputeQueue(HWQueue):
     self.wreg(self.gc.regCOMPUTE_PGM_RSRC3, prg.rsrc3)
     self.wreg(self.gc.regCOMPUTE_TMPRING_SIZE, prg.dev.tmpring_size)
 
-    if prg.dev.has_scratch_base_registers:
-      for xcc_id in range(self.dev.xccs):
-        with self.pred_exec(xcc_mask=1<<xcc_id):
-          scratch_base = prg.dev.scratch.va_addr + (prg.dev.scratch.size // self.dev.xccs * xcc_id)
-          self.wreg(self.gc.regCOMPUTE_DISPATCH_SCRATCH_BASE_LO, *data64_le(scratch_base >> 8))
-
-    if (10,0,0) <= prg.dev.target < (11,0,0): self.wreg(self.gc.mmCP_COHER_START_DELAY, 0x20)
+    # this is what llvm refers to as "architected flat scratch"
+    for xcc_id in range(self.dev.xccs):
+      with self.pred_exec(xcc_mask=1<<xcc_id):
+        scratch_base = prg.dev.scratch.va_addr + (prg.dev.scratch.size // self.dev.xccs * xcc_id)
+        self.wreg(self.gc.regCOMPUTE_DISPATCH_SCRATCH_BASE_LO, *data64_le(scratch_base >> 8))
 
     self.wreg(self.gc.regCOMPUTE_RESTART_X, 0, 0, 0)
     self.wreg(self.gc.regCOMPUTE_USER_DATA_0, *user_regs)
     self.wreg(self.gc.regCOMPUTE_RESOURCE_LIMITS, waves_per_sh=getenv("WAVES_PER_SH"))
     self.wreg(self.gc.regCOMPUTE_START_X, 0, 0, 0, *local_size, 0, 0)
 
-    gfx10p = {'cs_w32_en': int(prg.wave32)} if prg.dev.target >= (10,0,0) else {}
     self.pkt3(self.pm4.PACKET3_DISPATCH_DIRECT, *global_size,
-              self.gc.regCOMPUTE_DISPATCH_INITIATOR.encode(**gfx10p, force_start_at_000=1, compute_shader_en=1))
+              self.gc.regCOMPUTE_DISPATCH_INITIATOR.encode(**({'cs_w32_en': int(prg.wave32)} if prg.dev.target[0] != 9 else {}),
+                                                           force_start_at_000=1, compute_shader_en=1))
 
     if prg.dev.sqtt_enabled: self.pkt3(self.pm4.PACKET3_EVENT_WRITE, self.pm4.EVENT_TYPE(self.soc.THREAD_TRACE_MARKER) | self.pm4.EVENT_INDEX(0))
     self.pkt3(self.pm4.PACKET3_EVENT_WRITE, self.pm4.EVENT_TYPE(self.soc.CS_PARTIAL_FLUSH) | self.pm4.EVENT_INDEX(EVENT_INDEX_PARTIAL_FLUSH))
@@ -486,7 +484,7 @@ class AMDCopyQueue(HWQueue):
     return self
 
   def signal(self, signal:AMDSignal, value:sint=0):
-    fence_flags = self.sdma.SDMA_PKT_FENCE_HEADER_MTYPE(3) if self.dev.target >= (10,0,0) else 0
+    fence_flags = self.sdma.SDMA_PKT_FENCE_HEADER_MTYPE(3) if self.dev.target[0] != 9 else 0
     self.q(self.sdma.SDMA_OP_FENCE | fence_flags, *data64_le(signal.value_addr), value)
 
     if (dev:=signal.owner) is not None and signal.is_timeline and not dev.is_am():
@@ -593,7 +591,7 @@ class AMDProgram(HCQProgram):
     self.wave32: bool = desc.kernel_code_properties & 0x400 == 0x400
 
     # Set rsrc1.priv=1 on gfx11 to workaround cwsr.
-    self.rsrc1: int = desc.compute_pgm_rsrc1 | ((1 << 20) if (11,0,0) <= self.dev.target < (12,0,0) else 0)
+    self.rsrc1: int = desc.compute_pgm_rsrc1 | ((1 << 20) if self.dev.target[0] == 11 else 0)
     self.rsrc2: int = desc.compute_pgm_rsrc2 | (lds_size << 15)
     self.rsrc3: int = desc.compute_pgm_rsrc3
     self.aql_prog_addr: int = self.lib_gpu.va_addr + rodata_entry
@@ -961,30 +959,26 @@ class AMDDevice(HCQCompiled):
 
     self.target:tuple[int, ...] = ((trgt:=self.iface.props['gfx_target_version']) // 10000, (trgt // 100) % 100, trgt % 100)
     self.arch = "gfx%d%x%x" % self.target
-    if self.target < (9,4,2) or self.target >= (13,0,0): raise RuntimeError(f"Unsupported arch: {self.arch}")
+    assert (self.target in ((9,4,2),(9,5,0))) or self.target[0] in (11, 12), f"Unsupported arch: {self.arch}"
     if DEBUG >= 1: print(f"AMDDevice: opening {self.device_id} with target {self.target} arch {self.arch}")
 
     self.xccs = self.iface.props.get('num_xcc', 1)
     self.se_cnt = self.iface.props['array_count'] // self.iface.props['simd_arrays_per_engine'] // self.xccs
     self.cu_cnt = self.iface.props['simd_count'] // self.iface.props['simd_per_cu'] // self.xccs
     self.waves_per_cu = self.iface.props['max_waves_per_simd'] * self.iface.props['simd_per_cu']
-    self.wave_cnt = (self.cu_cnt * self.waves_per_cu) if self.target >= (10,1,0) else min(self.cu_cnt * 40, self.se_cnt * self.xccs * 512)
-    # this is what llvm refers to as "architected flat scratch"
-    self.has_scratch_base_registers = self.target >= (11,0,0) or self.target in {(9,4,2), (9,5,0)}
+    self.wave_cnt = (self.cu_cnt * self.waves_per_cu) if self.target[0] != 9 else min(self.cu_cnt * 40, self.se_cnt * self.xccs * 512)
 
     # https://gitlab.freedesktop.org/agd5f/linux/-/blob/a1fc9f584c4aaf8bc1ebfa459fc57a3f26a290d8/drivers/gpu/drm/amd/amdkfd/kfd_queue.c#L391
-    sgrp_size_per_cu, lds_size_per_cu, hwreg_size_per_cu = 0x4000, 0x10000, 0x1000
-    if self.target[:2] == (9,5): lds_size_per_cu = self.iface.props["lds_size_in_kb"] << 10
-    vgpr_size_per_cu = 0x60000 if self.target in {(11,0,0), (11,0,1), (11,5,1), (12,0,0), (12,0,1)} else \
-                       0x80000 if (self.target[:2]) in {(9,4), (9,5)} or self.target in {(9,0,8), (9,0,10)} else 0x40000
+    sgrp_size_per_cu, hwreg_size_per_cu = 0x4000, 0x1000
+    lds_size_per_cu = self.iface.props["lds_size_in_kb"] << 10 if self.target[:2] == (9,5) else 0x10000
+    vgpr_size_per_cu = 0x60000 if self.target in {(11,0,0), (11,0,1), (11,5,1), (12,0,0), (12,0,1)} else 0x80000 if self.target[0] == 9 else 0x40000
     wg_data_size = round_up((vgpr_size_per_cu + sgrp_size_per_cu + lds_size_per_cu + hwreg_size_per_cu) * self.cu_cnt, mmap.PAGESIZE)
-    ctl_stack_size = round_up((12 if self.target >= (10,1,0) else 8) * self.wave_cnt + 8 + 40, mmap.PAGESIZE)
-    if self.target[0] == 10: ctl_stack_size = min(ctl_stack_size, 0x7000)
+    ctl_stack_size = round_up((12 if self.target[0] != 9 else 8) * self.wave_cnt + 8 + 40, mmap.PAGESIZE)
     debug_memory_size = round_up(self.wave_cnt * 32, 64)
 
-    self.ip_off = import_ip_offsets(self.target)
+    self.ip_off = importlib.import_module(f"tinygrad.runtime.autogen.am.{'vega' if self.target[0] == 9 else 'navi'}_offsets")
     self.soc = import_soc(self.target)
-    self.pm4 = importlib.import_module(f"tinygrad.runtime.autogen.am.pm4_{'nv' if self.target[0] >= 10 else 'soc15'}")
+    self.pm4 = importlib.import_module(f"tinygrad.runtime.autogen.am.pm4_{'soc15' if self.target[0] == 9 else 'nv'}")
     self.sdma = import_module('sdma', min(self.iface.ip_versions[am.SDMA0_HWIP], (6, 0, 0)))
     self.gc = AMDIP('gc', self.iface.ip_versions[am.GC_HWIP],
                     bases={i: tuple(getattr(self.ip_off, f'GC_BASE__INST{i}_SEG{s}', 0) for s in range(6)) for i in range(6)})
@@ -1017,7 +1011,6 @@ class AMDDevice(HCQCompiled):
 
     self.pmc_enabled:bool = PROFILE > 0 and PMC > 0
     if self.pmc_enabled:
-      if self.target[0] not in {9, 11, 12}: raise RuntimeError(f'PMC are not supported on gc:{self.target}')
       self.iface.require_profile_mode()
 
       self.pmc_sched:list[PMCSample] = []
@@ -1036,7 +1029,6 @@ class AMDDevice(HCQCompiled):
     # SQTT is disabled by default because of runtime overhead and big file sizes (~200mb to Tensor.full() two 4096x4096 tensors and matmul them)
     self.sqtt_enabled:bool = PROFILE > 0 and SQTT > 0
     if self.sqtt_enabled:
-      if self.target[0] not in {9, 11, 12}: raise RuntimeError(f'SQ Thread Tracing is not supported on gc:{self.target}')
       self.iface.require_profile_mode()
 
       SQTT_BUFFER_SIZE = getenv("SQTT_BUFFER_SIZE", 256) # in mb, per shader engine
@@ -1074,7 +1066,7 @@ class AMDDevice(HCQCompiled):
     if self.max_private_segment_size >= private_segment_size: return
 
     lanes_per_wave = 64 # wave64
-    mem_alignment_size = 256 if self.target >= (11,0,0) else 1024
+    mem_alignment_size = 256 if self.target[0] != 9 else 1024
     size_per_thread = round_up(private_segment_size, mem_alignment_size // lanes_per_wave)
     size_per_xcc = size_per_thread * lanes_per_wave * self.iface.props['max_slots_scratch_cu'] * self.cu_cnt
     self.scratch, ok = self._realloc(getattr(self, 'scratch', None), size_per_xcc * self.xccs)
@@ -1082,18 +1074,18 @@ class AMDDevice(HCQCompiled):
       # NOTE: xcc logic is correct only for GFX9.
       max_scratch_waves = self.cu_cnt * self.iface.props['max_slots_scratch_cu'] * self.xccs
       wave_scratch = ceildiv(lanes_per_wave * size_per_thread, mem_alignment_size)
-      num_waves = (size_per_xcc // (wave_scratch * mem_alignment_size)) // (self.se_cnt if self.target >= (11,0,0) else 1)
+      num_waves = (size_per_xcc // (wave_scratch * mem_alignment_size)) // (self.se_cnt if self.target[0] != 9 else 1)
 
-      tmpring_t = getattr(hsa, f'union_COMPUTE_TMPRING_SIZE{"_GFX"+str(self.target[0]) if self.target[0] >= 11 else ""}_bitfields')
+      tmpring_t = getattr(hsa, f'union_COMPUTE_TMPRING_SIZE{"_GFX"+str(self.target[0]) if self.target[0] != 9 else ""}_bitfields')
       self.tmpring_size = int.from_bytes(tmpring_t(WAVES=min(num_waves, max_scratch_waves), WAVESIZE=wave_scratch), 'little')
       self.max_private_segment_size = private_segment_size
 
       if hasattr(self, 'aql_desc'):
         gfx9_rsrc = {'NUM_FORMAT':hsa.BUF_NUM_FORMAT_UINT, 'DATA_FORMAT':hsa.BUF_DATA_FORMAT_32, 'ELEMENT_SIZE':1, 'INDEX_STRIDE':3}
         rsrc = {'DST_SEL_X':hsa.SQ_SEL_X, 'DST_SEL_Y':hsa.SQ_SEL_Y, 'DST_SEL_Z':hsa.SQ_SEL_Z, 'DST_SEL_W':hsa.SQ_SEL_W, 'ADD_TID_ENABLE':1,
-                'TYPE':hsa.SQ_RSRC_BUF, **(gfx9_rsrc if self.target[0] < 10 else {'FORMAT':hsa.BUF_FORMAT_32_UINT, 'OOB_SELECT':2})}
-        rsrc1_t = getattr(hsa, f'union_SQ_BUF_RSRC_WORD1{"_GFX11" if self.target[0] >= 11 else ""}_bitfields')
-        rsrc3_t = getattr(hsa, f'union_SQ_BUF_RSRC_WORD3{"_GFX"+str(self.target[0]) if self.target[0] >= 10 else ""}_bitfields')
+                'TYPE':hsa.SQ_RSRC_BUF, **(gfx9_rsrc if self.target[0] == 9 else {'FORMAT':hsa.BUF_FORMAT_32_UINT, 'OOB_SELECT':2})}
+        rsrc1_t = getattr(hsa, f'union_SQ_BUF_RSRC_WORD1{"_GFX11" if self.target[0] != 9 else ""}_bitfields')
+        rsrc3_t = getattr(hsa, f'union_SQ_BUF_RSRC_WORD3{"_GFX"+str(self.target[0]) if self.target[0] != 9 else ""}_bitfields')
 
         self.aql_desc.scratch_backing_memory_location = int(self.scratch.va_addr)
         self.aql_desc.scratch_wave64_lane_byte_size = self.max_private_segment_size * lanes_per_wave // 64

tinygrad/runtime/ops_metal.py

@@ -147,10 +147,10 @@ class MetalProgram:
     encoder.endEncoding()
     command_buffer.setLabel(to_ns_str(self.name)) # TODO: is this always needed?
     command_buffer.commit()
+    self.dev.mtl_buffers_in_flight.append(command_buffer)
     if wait:
       wait_check(command_buffer)
       return command_buffer.GPUEndTime() - command_buffer.GPUStartTime()
-    self.dev.mtl_buffers_in_flight.append(command_buffer)
 
 class MetalBuffer:
   def __init__(self, buf:metal.MTLBuffer, size:int, offset=0): self.buf, self.size, self.offset = buf, size, offset

tinygrad/runtime/ops_null.py

@@ -1,10 +1,10 @@
 import inspect, functools
-from tinygrad.device import Compiled, Allocator
+from tinygrad.device import Compiled, Allocator, ProfileGraphEntry, ProfileGraphEvent
 from tinygrad.engine.jit import MultiGraphRunner
 from tinygrad.renderer import Renderer, cstyle, nir, ptx, llvmir, wgsl
 from tinygrad.renderer.cstyle import CStyleLanguage
 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.helpers import cpu_profile, getenv, dedup, NULL_ALLOW_COPYOUT
+from tinygrad.helpers import cpu_profile, getenv, dedup, NULL_ALLOW_COPYOUT, PROFILE, cpu_events, perf_counter_us
 
 class NullRenderer(CStyleLanguage):
   has_local = False
@@ -27,11 +27,16 @@ class NullAllocator(Allocator['NullDevice']):
   def _copyout(self, dest:memoryview, src):
     if not NULL_ALLOW_COPYOUT: raise RuntimeError("no copyout on NULL")
   def _transfer(self, dest, src, sz:int, src_dev, dest_dev):
-    with cpu_profile(f"{src_dev.device} -> {dest_dev.device}", f"{self.dev.device}:COPY"): pass
+    with cpu_profile(f"{src_dev.device} -> {dest_dev.device}", f"{src_dev.device}:SDMA:0"): pass
   def _offset(self, buf, offset:int, size:int): pass
 
 class NullGraph(MultiGraphRunner):
-  def __call__(self, input_uops:tuple[UOp, ...], var_vals:dict[str, int], wait=False) -> float|None: return 1e-1
+  def __call__(self, input_uops:tuple[UOp, ...], var_vals:dict[str, int], wait=False) -> float|None:
+    # description based on command, copied from HCQ graph
+    if PROFILE: cpu_events.append(ProfileGraphEvent(ents:=[ProfileGraphEntry(runtime.device if runtime is not None else f"{bufs[1].device}:SDMA:0", \
+        runtime.name if runtime is not None else f"{bufs[1].device} -> {bufs[0].device}", i, i+1) \
+        for i,((_,_,bufs,_),runtime) in enumerate(zip(self.calls, self.runtimes))], [], [perf_counter_us() for _ in range(len(ents)+1)]))
+    return 1e-1
 
 class NullDevice(Compiled):
   def __init__(self, device:str):

tinygrad/runtime/ops_python.py

@@ -18,8 +18,8 @@ def _load(m, i, dtype: DType):
   return from_storage_scalar(m[i], dtype)
 
 def load(inp, j, dtype: DType):
-  if len(inp) == 2: return [_load(m, x+j if x is not None else None, dtype) if gate else default for (m,x,gate),default in zip(*inp)]
-  return [_load(m, x+j if x is not None else None, dtype) for m,x,_ in inp[0]]
+  if len(inp) >= 3: return [_load(m, x+j if x is not None else None, dtype) if gate else default for (m,x),default,gate in zip(*inp[:3])]
+  return [_load(m, x+j if x is not None else None, dtype) for m,x in inp[0]]
 
 def _store(m, i, v, dtype: DType):
   if i < 0 or i >= len(m): raise IndexError(f"store out of bounds, size is {len(m)}, access is {i}, value is {v}")
@@ -67,8 +67,9 @@ class PythonProgram:
           continue
         assert dtype is not None, f"{uop} is missing a dtype"
         if uop is Ops.STORE:
+          store_gate = src_values[2] if len(src_values) >= 3 else [True] * warp_size
           for j,val in enumerate(src_values[1] if src_dtypes[1].count > 1 else [src_values[1]]):
-            for (m,o,g),v in zip(src_values[0], val):
+            for (m,o),v,g in zip(src_values[0], val, store_gate):
               if g: _store(m, o+j, v, src_dtypes[1].scalar())
           i += 1
           continue
@@ -91,6 +92,7 @@ class PythonProgram:
           elif arg[0] == 'l': values[i] = [x[2-int(arg[-1])] for x in warp]
         elif uop is Ops.CONST: values[i] = [arg] * warp_size
         elif uop is Ops.INDEX:
+          if len(src_values) != 2: raise RuntimeError("gates must be on LOAD/STORE, not INDEX")
           ret:list = []
           if isinstance(src_dtypes[0], ImageDType):
             for m,ox,oy in zip(src_values[0], src_values[1][0], src_values[1][1]):
@@ -98,7 +100,7 @@ class PythonProgram:
               else: ret.append((m, ox*4 + oy*src_dtypes[0].shape[1]*4))
           else:
             for m,o in zip(src_values[0], src_values[1]): ret.append((m,o))
-          values[i] = [(m,o,g) for (m,o),g in zip(ret, src_values[2] if len(src_values) == 3 else [True]*len(ret))] # set the gate last
+          values[i] = ret
         elif uop is Ops.CAST and isinstance(dtype, PtrDType):
           values[i] = src_values[0]
         elif uop is Ops.RANGE:

tinygrad/runtime/support/am/amdev.py

@@ -112,7 +112,7 @@ class AMFirmware:
     if (sys.version_info >= (3,14) and (p:=pathlib.Path("/lib/firmware/amdgpu")/f"{fname}.zst").is_file() and
         hashlib.sha256(b:=zstd.decompress(p.read_bytes())).hexdigest() == fw.hashes[fname]): blob = memoryview(bytearray(b))
     else: blob = memoryview(bytearray(fetch(f"https://gitlab.com/kernel-firmware/linux-firmware/-/raw/1e2c15348485939baf1b6d1f5a7a3b799d80703d/amdgpu/{fname}",
-                                            subdir="fw").read_bytes()))
+                                            subdir="fw", sha256=fw.hashes[fname]).read_bytes()))
     if AM_DEBUG >= 1: print(f"am {self.adev.devfmt}: loading firmware {fname}: {hashlib.sha256(blob).hexdigest()}")
     if versioned_header:
       chdr = am.struct_common_firmware_header.from_address(mv_address(blob))
@@ -177,10 +177,12 @@ class AMDev:
     # Init hw for IP blocks where it is needed
     if not self.partial_boot:
       if self.psp.is_sos_alive() and self.smu.is_smu_alive():
+        self.pci_dev.write_config_flush(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) & ~pci.PCI_COMMAND_MASTER, 2)
         if self.is_hive():
           if reset_mode: return # in reset mode, do not raise
           raise RuntimeError("Malformed state. Use extra/amdpci/hive_reset.py to reset the hive")
         self.smu.mode1_reset()
+      self.pci_dev.write_config_flush(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)
       self.init_hw(self.soc, self.gmc, self.ih, self.psp, self.smu)
 
     # Booting done
@@ -188,7 +190,6 @@ class AMDev:
 
     # Re-initialize main blocks
     self.init_hw(self.gfx, self.sdma)
-    self.pci_dev.write_config(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)
 
     if (max_power:=getenv("AM_POWER_LIMIT", 0.0)) > 0:
       self.smu.set_power_limit(max_power)
@@ -325,7 +326,7 @@ class AMDev:
   @functools.cached_property
   def hwid_names(self) -> dict[int, str]: return {v:k.removesuffix('_HWID') for k,v in vars(am).items() if k.endswith('_HWID') and isinstance(v, int)}
 
-  def _ip_module(self, prefix:str, hwip, prever_prefix:str=""): return import_module(prefix, self.ip_ver[hwip], prever_prefix)
+  def _ip_module(self, prefix:str, hwip): return import_module(prefix, self.ip_ver[hwip])
 
   def _build_regs(self):
     mods = [("mp", am.MP0_HWIP), ("hdp", am.HDP_HWIP), ("gc", am.GC_HWIP), ("mmhub", am.MMHUB_HWIP), ("osssys", am.OSSSYS_HWIP),
@@ -334,4 +335,4 @@ class AMDev:
 
     for prefix, hwip in mods:
       self.__dict__.update(import_asic_regs(prefix, self.ip_ver[hwip], cls=functools.partial(AMRegister, adev=self, bases=self.regs_offset[hwip])))
-    self.__dict__.update(import_asic_regs('mp', (11, 0), cls=functools.partial(AMRegister, adev=self, bases=self.regs_offset[am.MP1_HWIP])))
+    self.__dict__.update(import_asic_regs('mp', (11, 0, 0), cls=functools.partial(AMRegister, adev=self, bases=self.regs_offset[am.MP1_HWIP])))

tinygrad/runtime/support/am/ip.py

@@ -173,7 +173,7 @@ class AM_GMC(AM_IP):
 
 class AM_SMU(AM_IP):
   def init_sw(self):
-    self.smu_mod = self.adev._ip_module("smu", am.MP1_HWIP, prever_prefix='v')
+    self.smu_mod = self.adev._ip_module("smu", am.MP1_HWIP)
     self.driver_table_paddr = self.adev.mm.palloc(0x4000, zero=False, boot=True)
 
   def init_hw(self):
@@ -304,9 +304,10 @@ class AM_GFX(AM_IP):
   def reset_mec(self):
     self._dequeue_hqds()
 
-    for xcc in range(self.xccs): self.adev.regGRBM_SOFT_RESET.write(soft_reset_cp=1, soft_reset_cpc=1, inst=xcc)
-    time.sleep(0.05)
-    for xcc in range(self.xccs): self.adev.regGRBM_SOFT_RESET.write(0x0, inst=xcc)
+    if self.adev.ip_ver[am.GC_HWIP] < (10,0,0): # gfx10+ uses mec_pipe0_reset
+      for xcc in range(self.xccs): self.adev.regGRBM_SOFT_RESET.write(soft_reset_cp=1, soft_reset_cpc=1, inst=xcc)
+      time.sleep(0.05)
+      for xcc in range(self.xccs): self.adev.regGRBM_SOFT_RESET.write(0x0, inst=xcc)
 
     self._config_mec()
     self._enable_mec()

tinygrad/runtime/support/amd.py

@@ -1,9 +1,7 @@
-import functools, re, urllib, tinygrad.runtime.autogen
-from collections import defaultdict
+import functools, re, tinygrad.runtime.autogen.am
 from dataclasses import dataclass
 from tinygrad.helpers import getbits, fetch
 
-AMDGPU_URL = "https://gitlab.com/linux-kernel/linux-next/-/raw/cf6d949a409e09539477d32dbe7c954e4852e744/drivers/gpu/drm/amd"
 ROCM_URL = "https://raw.githubusercontent.com/ROCm/rocm-systems/cccc350dc620e61ae2554978b62ab3532dc10bd9/projects"
 
 @dataclass
@@ -20,7 +18,6 @@ class AMDReg:
 @dataclass
 class AMDIP:
   name:str; version:tuple[int, ...]; bases:dict[int, tuple[int, ...]] # noqa: E702
-  def __post_init__(self): self.version = fixup_ip_version(self.name, self.version)[0]
 
   @functools.cached_property
   def regs(self): return import_asic_regs(self.name, self.version, cls=functools.partial(AMDReg, bases=self.bases))
@@ -30,39 +27,18 @@ class AMDIP:
     if (name10:=name.replace('reg', 'mm')) in self.regs: return self.regs[name10]
     raise AttributeError(f"{self.name.upper()} has no register {name}")
 
-def fixup_ip_version(ip:str, version:tuple[int, ...]) -> list[tuple[int, ...]]:
-  # override versions
-  def _apply_ovrd(ovrd:dict[tuple[int, ...], tuple[int, ...]]) -> tuple[int, ...]:
-    for ver, ovrd_ver in ovrd.items():
-      if version[:len(ver)] == ver: return ovrd_ver
-    return version
+# load the greatest module with matching major version that's less than or equal to the target version
+# this is not universally correct, see below for an example, but appears reliable for recent gpus
+# https://github.com/torvalds/linux/blob/9207d47f966be9f4d52e7e0119ac2b7a7e366f3e/drivers/gpu/drm/amd/amdgpu/amdgpu_discovery.c#L3163
+def import_module(name:str, target:tuple[int, ...], submod=""):
+  mod = getattr(tinygrad.runtime.autogen.am, submod) if submod else tinygrad.runtime.autogen.am
+  if (children:=[c for c in mod.__all__ if c.startswith(name) and (v:=tuple(map(int, c.split('_')[1:])))[0] == target[0] and v <= target]):
+    return getattr(mod, children[-1])
+  raise ImportError(f"Failed to import {submod+'.' if submod else ''}{name} {'.'.join(map(str, target))}")
 
-  if ip in ['nbio', 'nbif']: version = _apply_ovrd({(3,3): (2,3,0), (7,3): (7,2,0)})
-  elif ip in ['mp', 'smu']: version = _apply_ovrd({(14,0,3): (14,0,2)})
-  elif ip in ['gc']: version = _apply_ovrd({(9,5,0): (9,4,3)})
-  elif ip in ['sdma']: version = _apply_ovrd({(4,4,4): (4,4,2)})
+def header_download(file, url) -> str: return fetch(f"{url}/{file}", subdir="defines").read_text()
 
-  return [version, version[:2], version[:2]+(0,), version[:1]+(0, 0)]
-
-def header_download(file, name=None, subdir="defines", url=AMDGPU_URL) -> str: return fetch(f"{url}/{file}", name=name, subdir=subdir).read_text()
-
-def import_header(path:str, url=AMDGPU_URL):
-  t = re.sub(r'//.*|/\*.*?\*/','', header_download(path, subdir="defines", url=url), flags=re.S)
-  # TODO: refactor when clang2py is replaced
-  return {k:int(v,0) for k,v in re.findall(r'\b([A-Za-z_]\w*)\s*=\s*(0x[0-9A-Fa-f]+|\d+)', t) + \
-                                re.findall(r'^\s*#\s*define\s+([A-Za-z_0-9]\w*)\s+(0x[0-9A-Fa-f]+|\d+)', t, re.M)}
-
-def import_module(name:str, version:tuple[int, ...], version_prefix:str=""):
-  for ver in fixup_ip_version(name, version):
-    try: return getattr(tinygrad.runtime.autogen.am, f"{name}_{version_prefix}{'_'.join(map(str, ver))}")
-    except AttributeError: pass
-  raise ImportError(f"Failed to load autogen module for {name.upper()} {'.'.join(map(str, version))}")
-
-def import_soc(ip):
-  # rocm soc headers have more profiling enums than upstream linux
-  return type("SOC", (object,), import_header(f"aqlprofile/linux/{({9: 'vega10', 10: 'navi10', 11: 'soc21', 12: 'soc24'}[ip[0]])}_enum.h", ROCM_URL))
-
-def import_ip_offsets(ip): return type("IPOFF", (object,), import_header(f"include/{('sienna_cichlid' if ip[0] > 9 else 'vega20')}_ip_offset.h"))
+def import_soc(ip): return getattr(tinygrad.runtime.autogen.am, f"soc_{ip[0]}")
 
 def import_pmc(ip) -> dict[str, tuple[str, int]]:
   res:dict[str, tuple[str, int]] = {}
@@ -70,7 +46,7 @@ def import_pmc(ip) -> dict[str, tuple[str, int]]:
   # NOTE: precise arch for mi300+, generic for others, since rocm headers lack some archs
   arch = f"gfx{ip[0]}{ip[1]:x}{ip[2]:x}" if ip[0] == 9 else f"gfx{ip[0]}"
 
-  for sec in header_download("rocprofiler-compute/src/rocprof_compute_soc/profile_configs/counter_defs.yaml", url=ROCM_URL).split('- name: ')[1:]:
+  for sec in header_download("rocprofiler-compute/src/rocprof_compute_soc/profile_configs/counter_defs.yaml", ROCM_URL).split('- name: ')[1:]:
     for arch_spec in sec.split('- architectures:')[1:]:
       if arch in arch_spec and (block:=re.search(r'block:\s*([A-Za-z0-9_]+)', arch_spec)) and (ev:=re.search(r'event:\s*(\d+)', arch_spec)):
         res[sec.splitlines()[0].strip()] = (block.group(1), int(ev.group(1)))
@@ -78,35 +54,4 @@ def import_pmc(ip) -> dict[str, tuple[str, int]]:
   return res
 
 def import_asic_regs(prefix:str, version:tuple[int, ...], cls=AMDReg) -> dict[str, AMDReg]:
-  def _split_name(name): return name[:(pos:=next((i for i,c in enumerate(name) if c.isupper()), len(name)))], name[pos:]
-  def _extract_regs(txt):
-    x = {}
-    for k,v in {m.group(1): int(m.group(2), 0) for line in txt.splitlines() if (m:=re.match(r'#define\s+(\S+)\s+(0x[\da-fA-F]+|\d+)', line))}.items():
-      if k.startswith('VM_') or k.startswith('MC_'): x[prefix.upper()[:2]+k] = v
-      elif k.startswith('regVM_') or k.startswith('regMC_'): x["reg"+prefix.upper()[:2]+k[3:]] = v
-      else: x[k] = v
-    return x
-  def _download_file(ver, suff) -> str:
-    dir_prefix = {"osssys": "oss"}.get(prefix, prefix)
-    fetch_name = f"{prefix}_{'_'.join(map(str, ver))}_{suff}.h"
-    return header_download(f"include/asic_reg/{dir_prefix}/{fetch_name}", name=fetch_name, subdir="asic_regs")
-
-  for ver in fixup_ip_version(prefix, version):
-    try: offs, sh_masks = _extract_regs(_download_file(ver, "offset")), _extract_regs(_download_file(ver, "sh_mask"))
-    except urllib.error.HTTPError as e:
-      if e.code == 404: continue
-      raise
-
-    offsets = {k:v for k,v in offs.items() if _split_name(k)[0] in {'reg', 'mm'} and not k.endswith('_BASE_IDX')}
-    bases = {k[:-len('_BASE_IDX')]:v for k,v in offs.items() if _split_name(k)[0] in {'reg', 'mm'} and k.endswith('_BASE_IDX')}
-
-    fields: defaultdict[str, dict[str, tuple[int, int]]] = defaultdict(dict)
-    for field_name, field_mask in sh_masks.items():
-      if not ('__' in field_name and field_name.endswith('_MASK')): continue
-      reg_name, reg_field_name = field_name[:-len('_MASK')].split('__')
-      if reg_name.startswith('MC_') or reg_name.startswith('VM_'): reg_name = f"{prefix.upper()[:2]}{reg_name}"
-      fields[reg_name][reg_field_name.lower()] = ((field_mask & -field_mask).bit_length()-1, field_mask.bit_length()-1)
-
-    # NOTE: Some registers like regGFX_IMU_FUSESTRAP in gc_11_0_0 are missing base idx, just skip them
-    return {reg:cls(name=reg, offset=off, segment=bases[reg], fields=fields[_split_name(reg)[1]]) for reg,off in offsets.items() if reg in bases}
-  raise ImportError(f"Failed to load ASIC registers for {prefix.upper()} {'.'.join(map(str, version))}")
+  return {reg:cls(name=reg, offset=off, segment=seg, fields=fields) for reg,(off,seg,fields) in import_module(prefix, version, submod="regs").items()}

tinygrad/runtime/support/autogen.py

@@ -99,7 +99,8 @@ arc_families = ['alloc', 'copy', 'mutableCopy', 'new']
 
 def normalize(a): return ("_" + n if keyword.iskeyword(n:=nm(a)) else n)
 
-def gen(name, files, dll="", args=[], prolog=[], rules=[], epilog=[], recsym=False, errno=False, anon_names={}, types={}, macros=True, paths=[]):
+def gen(name, files, dll="", args=[], prolog=[], rules=[], epilog=[], recsym=False, errno=False, anon_names={}, types={}, macros=True, paths=[],
+        patterns=[]):
   extras, lines, anoncnt, types, objc, fns = [], [], itertools.count().__next__, {k:(v,True) for k,v in types.items()}, False, set()
 
   # ctypes automatically "unboxes" simple types
@@ -227,6 +228,7 @@ def gen(name, files, dll="", args=[], prolog=[], rules=[], epilog=[], recsym=Fal
     while q:
       c = q.pop()
       if loc_file(loc(c)) != str(f) and (not recsym or c.kind not in (clang.CXCursor_FunctionDecl,)): continue
+      if patterns and not any(re.match(p, nm(c)) for p in patterns): continue
       rollback = lines, types
       try:
         match c.kind:

tinygrad/runtime/support/c.py

@@ -94,7 +94,8 @@ class DLL(ctypes.CDLL):
     if nm == 'libc' and OSX: return '/usr/lib/libc.dylib'
     if pathlib.Path(path:=getenv(nm.replace('-', '_').upper()+"_PATH", '')).is_file(): return path
     for p in paths:
-      libpaths = {"posix": ["/usr/lib64", "/usr/lib", "/usr/local/lib"], "nt": os.environ['PATH'].split(os.pathsep),
+      libpaths = {"posix": [d for d in os.environ.get('LD_LIBRARY_PATH', '').split(os.pathsep) if d] + ["/usr/lib64", "/usr/lib", "/usr/local/lib"],
+                  "nt": os.environ['PATH'].split(os.pathsep),
                   "darwin": ["/opt/homebrew/lib", f"/System/Library/Frameworks/{p}.framework", f"/System/Library/PrivateFrameworks/{p}.framework"],
                   'linux': ['/lib', '/lib64', f"/lib/{sysconfig.get_config_var('MULTIARCH')}", "/usr/lib/wsl/lib/"]}
       if (pth:=pathlib.Path(p)).is_absolute():

tinygrad/runtime/support/nv/nvdev.py

@@ -73,7 +73,6 @@ class NVMemoryManager(MemoryManager):
 class NVDev:
   def __init__(self, pci_dev:PCIDevice):
     self.pci_dev, self.devfmt, self.mmio = pci_dev, pci_dev.pcibus, pci_dev.map_bar(0, fmt='I')
-    self.pci_dev.write_config(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)
 
     self.smi_dev, self.is_booting, self.is_err_state = False, True, False
     self._early_ip_init()
@@ -104,10 +103,12 @@ class NVDev:
     self.include("src/common/inc/swref/published/ampere/ga102/dev_gc6_island_addendum.h")
 
     if (needs_reset:=self.reg("NV_PFB_PRI_MMU_WPR2_ADDR_HI").read() != 0):
+      self.pci_dev.write_config_flush(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) & ~pci.PCI_COMMAND_MASTER, 2)
       if DEBUG >= 2: print(f"nv {self.devfmt}: WPR2 is up. Issuing a full reset.", flush=True)
       self.pci_dev.reset()
       time.sleep(0.1) # wait until device can respond again
 
+    self.pci_dev.write_config_flush(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)
     self.chip_id = self.reg("NV_PMC_BOOT_0").read()
     self.chip_details = self.reg("NV_PMC_BOOT_42").read_bitfields()
     self.chip_name = {0x17: "GA1", 0x19: "AD1", 0x1b: "GB2"}[self.chip_details['architecture']] + f"{self.chip_details['implementation']:02d}"

tinygrad/runtime/support/system.py

@@ -196,6 +196,9 @@ class PCIDevice:
   def reset(self): os.system(f"sudo sh -c 'echo 1 > /sys/bus/pci/devices/{self.pcibus}/reset'")
   def read_config(self, offset:int, size:int): return int.from_bytes(self.cfg_fd.read(size, binary=True, offset=offset), byteorder='little')
   def write_config(self, offset:int, value:int, size:int): self.cfg_fd.write(value.to_bytes(size, byteorder='little'), binary=True, offset=offset)
+  def write_config_flush(self, offset:int, value:int, size:int):
+    self.write_config(offset, value, size)
+    self.read_config(offset, size)
 
   @functools.cache
   def bar_fd(self, bar_idx:int) -> FileIOInterface:

tinygrad/schedule/indexing.py

@@ -5,7 +5,7 @@ from tinygrad.dtype import dtypes, AddrSpace
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType, profile_matches
 from tinygrad.uop.ops import consumer_map_from_toposort, gate_kernel_sink
 from tinygrad.uop.symbolic import symbolic, pm_simplify_valid, pm_drop_and_clauses
-from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored
+from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored, Context, SPEC
 
 ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.AFTER, Ops.COPY, Ops.BUFFER, Ops.BUFFER_VIEW,
                      Ops.CONST, Ops.BIND, Ops.DEVICE, Ops.MSELECT, Ops.MSTACK, Ops.PARAM,
@@ -265,7 +265,9 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
     # assign to the range map. rngs are the input ranges, out_rngs are the output ranges, from the x op.
     rctx.range_map[x] = (rngs, out_rngs)
 
-  tsink = graph_rewrite(tsink, pm_apply_rangeify, ctx=rctx, bottom_up=True, name="apply rangeify")
+  # NOTE: SPEC=3 is broken here with shape
+  with Context(SPEC=min(SPEC.value, 2)):
+    tsink = graph_rewrite(tsink, pm_apply_rangeify, ctx=rctx, bottom_up=True, name="apply rangeify")
   return tsink, rctx
 
 def render_ranges(*rngs_list, realized) -> str:

tinygrad/schedule/rangeify.py

@@ -442,7 +442,7 @@ pm_add_buffers = pm_mops+pm_flatten_bufferize+to_bufferview+PatternMatcher([
   (UPat(Ops.AFTER, src=(UPat.var("x"), UPat(Ops.AFTER, name="y"))), lambda x,y: x.after(*y.src[1:])),
 
   # remove invalid writes
-  (UPat(Ops.STORE, src=(UPat(), UPat(Ops.CONTIGUOUS, src=(UPat(Ops.CONST, arg=Invalid),))), allow_any_len=True), lambda: UOp(Ops.NOOP)),
+  (UPat(Ops.STORE, src=(UPat(), UPat(Ops.CONTIGUOUS, src=(UPat(Ops.CONST, arg=Invalid),)))), lambda: UOp(Ops.NOOP)),
   (UPat(Ops.AFTER, src=(UPat.var("x"), UPat(Ops.NOOP, src=()))), lambda x: x),
   (UPat(Ops.AFTER, src=(UPat.var("x"), UPat(Ops.END, src=(UPat(Ops.NOOP, src=()),), allow_any_len=True))), lambda x: x),
 ])

tinygrad/tensor.py

@@ -822,19 +822,27 @@ class Tensor(OpMixin):
     """
     Returns a tensor with `num_samples` indices sampled from a multinomial distribution weighted by `self`.
 
-    NOTE: `replacement=False` for `num_samples > 1` is not supported yet.
     ```python exec="true" source="above" session="tensor" result="python"
     Tensor.manual_seed(42)
     t = Tensor([1, 2, 3, 4])
     print(t.multinomial(20, replacement=True).numpy())
     ```
+    ```python exec="true" source="above" session="tensor" result="python"
+    Tensor.manual_seed(42)
+    t = Tensor([1, 2, 3, 4])
+    print(t.multinomial(3, replacement=False).numpy())
+    ```
     """
     assert 1 <= self.ndim <= 2 and num_samples > 0, f"{self.ndim=} must be 1 or 2 dim, {num_samples=} must be positive"
-    assert replacement or num_samples == 1, "no replacement only supports num_samples = 1"
     weight = self.unsqueeze(0) if self.ndim == 1 else self
-    cdf = (cw := weight.cumsum(1).float()) / cw[:, -1].unsqueeze(1)
-    unif_samples = Tensor.rand(num_samples, cdf.shape[0], 1).to(self.device)
-    indices = (unif_samples.expand((-1, -1, cdf.shape[1])) >= cdf).sum(2).permute((1, 0))
+    assert replacement or num_samples <= weight.shape[1], "no replacement samples must not exceed population size"
+    if replacement or num_samples == 1:
+      cdf = (cw := weight.cumsum(1).float()) / cw[:, -1].unsqueeze(1)
+      unif_samples = Tensor.rand(num_samples, cdf.shape[0], 1).to(self.device)
+      indices = (unif_samples.expand((-1, -1, cdf.shape[1])) >= cdf).sum(2).permute((1, 0))
+    else:
+      # Efraimidis–Spirakis
+      indices = (weight.rand_like(dtype=dtypes.float32).log2() / weight).topk(num_samples, dim=1)[1]
     return (indices.squeeze(0) if self.ndim == 1 else indices).cast(dtypes.int32)
 
   # ***** toposort and backward pass *****
@@ -1297,13 +1305,9 @@ class Tensor(OpMixin):
     """
     if rounding_mode is None: return super().div(x, reverse)  # type: ignore[arg-type]
     numerator, denominator = self._broadcasted(x, reverse)
-    if dtypes.is_int(dt:=least_upper_dtype(numerator.dtype, denominator.dtype)):
-      numerator, denominator = numerator.cast(dt), denominator.cast(dt)
+    if dtypes.is_int(numerator.dtype):
       if rounding_mode == "trunc": return numerator.idiv(denominator)
-      if rounding_mode == "floor":
-        truncate_div, truncate_mod = numerator.idiv(denominator), numerator._binop(Ops.MOD, denominator, False)
-        opposite_sign = ((numerator>0)&(denominator<0)) | ((numerator<0)&(denominator>0))
-        return (opposite_sign&(truncate_mod!=0)).where(truncate_div-1, truncate_div)
+      if rounding_mode == "floor": return numerator._binop(Ops.FLOORDIV, denominator, False)
     d = numerator.cast(least_upper_float(numerator.dtype)) * denominator.cast(least_upper_float(denominator.dtype)).reciprocal()
     output_dtype = numerator.dtype if dtypes.is_int(numerator.dtype) else d.dtype
     if rounding_mode == "trunc": return d.trunc().cast(output_dtype)
@@ -1321,8 +1325,21 @@ class Tensor(OpMixin):
     ```
     """
     a, b = self._broadcasted(x, reverse)
+    if dtypes.is_int(a.dtype): return a._binop(Ops.FLOORMOD, b, False)
     return a - a.div(b, rounding_mode="floor") * b
 
+  def fmod(self, x:Tensor|ConstType) -> Tensor:
+    """
+    C-style remainder of `self` divided by `x` (sign follows the dividend), using truncating division.
+    Differs from `mod`/`%`, which uses Python floor remainder.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([-4, 7, 5, 4, -7, 8]).fmod(Tensor([2, -3, 8, -2, 3, 5])).numpy())
+    ```
+    """
+    a, b = self._broadcasted(x)
+    return a - a.div(b, rounding_mode="trunc") * b
+
   def where(self:Tensor, x:Tensor|ConstType|sint, y:Tensor|ConstType|sint) -> Tensor:
     """
     Returns a tensor of elements selected from either `x` or `y`, depending on `self`.

tinygrad/uop/__init__.py

@@ -65,6 +65,7 @@ class Ops(FastEnum):
   CMPLT = auto(); CMPNE = auto(); CMPEQ = auto()
   XOR = auto(); OR = auto(); AND = auto()
   THREEFRY = auto(); SUB = auto(); FDIV = auto(); POW = auto()
+  FLOORDIV = auto(); FLOORMOD = auto()
 
   # TernaryOps
   WHERE = auto(); MULACC = auto()
@@ -110,7 +111,7 @@ class Ops(FastEnum):
 class GroupOp:
   Unary = {Ops.EXP2, Ops.LOG2, Ops.SIN, Ops.SQRT, Ops.RECIPROCAL, Ops.NEG, Ops.TRUNC}
   Binary = {Ops.ADD, Ops.MUL, Ops.IDIV, Ops.MAX, Ops.MOD, Ops.CMPLT, Ops.CMPNE, Ops.CMPEQ,
-            Ops.XOR, Ops.SHL, Ops.SHR, Ops.OR, Ops.AND, Ops.THREEFRY, Ops.SUB, Ops.FDIV, Ops.POW}
+            Ops.XOR, Ops.SHL, Ops.SHR, Ops.OR, Ops.AND, Ops.THREEFRY, Ops.SUB, Ops.FDIV, Ops.POW, Ops.FLOORDIV, Ops.FLOORMOD}
   Ternary = {Ops.WHERE, Ops.MULACC}
   ALU = set.union(Unary, Binary, Ternary)
 
@@ -137,6 +138,6 @@ class GroupOp:
   Comparison = {Ops.CMPLT, Ops.CMPNE, Ops.CMPEQ}
 
   # do not preserve f(0) = 0
-  UnsafePad = {Ops.RECIPROCAL, Ops.LOG2, Ops.EXP2, Ops.IDIV, Ops.POW}
+  UnsafePad = {Ops.RECIPROCAL, Ops.LOG2, Ops.EXP2, Ops.IDIV, Ops.POW, Ops.FLOORDIV}
 
   All = set(Ops)

tinygrad/uop/decompositions.py

@@ -290,8 +290,10 @@ def fast_idiv(target: Target, x: UOp, d: int, dont_cast=False) -> UOp|None:
   if m*vmin >= x.dtype.min and m*vmax <= x.dtype.max:
     return ((x*m) >> s) if is_unsigned else ((x*m) >> s) + (x<0).where(x.ufix(1), 0)
   # before we try casting to a larger dtype (slow), we see if there are powers of two in d we can shift to make x smaller
+  # use explicit Ops.IDIV (trunc) since the recursion assumes trunc semantics throughout
   if (largest_factor_of_two_in_d := (d & -d)) > 1:
-    if (ret:=fast_idiv(target, x//largest_factor_of_two_in_d, d//largest_factor_of_two_in_d, dont_cast=True)) is not None: return ret
+    if (ret:=fast_idiv(target, x.alu(Ops.IDIV, x.const_like(largest_factor_of_two_in_d)),
+                       d//largest_factor_of_two_in_d, dont_cast=True)) is not None: return ret
   if dont_cast: return None
   # promo_lattice needs to return an unsigned type if the type is unsigned
   if dtypes.is_int(next_dtype := promo_lattice[x.dtype.scalar()][-1]) and is_dtype_supported(next_dtype, target):
@@ -416,13 +418,15 @@ def f2f_clamp(val:UOp, dt:DType) -> UOp:
   # FIXME: CMPLT of nan is undefined
   return val.ne(val).where(val, (val < -mx).where(-sat, (mx < val).where(sat, val)))
 
-def f2f_load(x: UOp, fr:DType, to:DType) -> UOp:
-  if (n:=x.dtype.count) == 1: return f2f(x.replace(dtype=f2f_dt[fr]), fr, to)
-  return UOp.vectorize(*(f2f(x.replace(dtype=f2f_dt[fr], src=(reindex(x.src[0].src[0], i, 1),)), fr, to) for i in range(n)))
+def f2f_load(x:UOp, fr:DType, to:DType) -> UOp:
+  if (n:=x.dtype.count) == 1:
+    return f2f(x.replace(src=(x.src[0],)+((x.src[1].cast(f2f_dt[fr]), x.src[2]) if len(x.src) >= 3 else ()), dtype=f2f_dt[fr]), fr, to)
+  return UOp.vectorize(*(f2f(x.replace(dtype=f2f_dt[fr], src=(reindex(x.src[0].src[0], i, 1),) + \
+                                       ((x.src[1].gep(i).cast(f2f_dt[fr]), x.src[2]) if len(x.src) >= 3 else ())), fr, to) for i in range(n)))
 
-def f2f_store(st, idx, val, fr:DType, to:DType):
-  if (n:=val.dtype.count) == 1: return st.replace(src=(idx, f2f(val.bitcast(f2f_dt[to]), to, fr)))
-  return UOp.group(*(st.replace(src=(reindex(idx, i, 1), f2f(val.gep(i).bitcast(f2f_dt[to]), to, fr))) for i in range(n)))
+def f2f_store(st:UOp, idx, val, fr:DType, to:DType):
+  if (n:=val.dtype.count) == 1: return st.replace(src=(idx, f2f(val.bitcast(f2f_dt[to]), to, fr))+st.src[2:])
+  return UOp.group(*(st.replace(src=(reindex(idx, i, 1), f2f(val.gep(i).bitcast(f2f_dt[to]), to, fr))+st.src[2:]) for i in range(n)))
 
 # ***** decomposition patterns *****
 
@@ -438,28 +442,47 @@ def get_transcendental_patterns(ops:tuple[Ops, ...], force_transcendental:bool)
   if Ops.SQRT not in ops or force_transcendental: pat.append((UPat(Ops.SQRT, src=UPat.var("d")), lambda d: xpow(d, d.const_like(0.5))))
   return PatternMatcher(pat)
 
+def floordiv_to_idiv(a:UOp, b:UOp) -> UOp:
+  if (a.vmin >= 0 and b.vmin > 0) or (a.vmax <= 0 and b.vmax < 0): return a.alu(Ops.IDIV, b)
+  return a.alu(Ops.IDIV, b) - (a.alu(Ops.MOD, b).ne(0) & (a<0).ne(b<0)).cast(a.dtype)
+
+def floormod_to_mod(a:UOp, b:UOp) -> UOp:
+  if (a.vmin >= 0 and b.vmin > 0) or (a.vmax <= 0 and b.vmax < 0): return a.alu(Ops.MOD, b)
+  r = a.alu(Ops.MOD, b)
+  # use where instead of mul to avoid being fused into MULACC (which int64 long-decomp doesn't handle)
+  return r + (r.ne(0) & (a<0).ne(b<0)).where(b, b.const_like(0))
+
 powers_of_two: dict[int, int] = {2**i:i for i in range(64)}
 @functools.cache
 def get_late_rewrite_patterns(ops:tuple[Ops, ...], disable_fast_idiv:bool) -> PatternMatcher:
-  pat: list[tuple[UPat, Callable]] = []
+  pat: list[tuple[UPat, Callable]] = [(UPat.var("a")//UPat.var("b"), floordiv_to_idiv)]
+  # FLOORMOD by 2**y -> x & (2**y-1) (correct floor mod for any sign in two's complement); fires before floormod_to_mod
+  if Ops.AND in ops: pat.append((UPat.var("x", dtypes.ints)%UPat.cvar("c"), lambda x,c: x & (c.arg-1) if c.arg in powers_of_two else None))
+  pat.append((UPat.var("a")%UPat.var("b"), floormod_to_mod))
   # no real hardware supports THREEFRY, but NullRenderer does
   if Ops.THREEFRY not in ops: pat.append((UPat(Ops.THREEFRY, dtype=dtypes.uint64, src=(UPat.var("x"), UPat.var("key"))), threefry2x32))
   # MAX can be rewritten as CMPLT + WHERE (max function is annoying on many cstyle backends)
   if Ops.MAX not in ops and Ops.CMPLT in ops: pat.append((UPat(Ops.MAX, name="m"), lambda m: (m.src[0] < m.src[1]).where(m.src[1], m.src[0])))
-  # rewrite MOD to AND (which should always be supported, but not for generic in tests): x % (2**y) -> x & (2**y-1)
-  if Ops.AND in ops: pat += [(UPat.var("x", dtypes.ints)%UPat.cvar("c"), lambda x,c: x & (c.arg-1) if c.arg in powers_of_two else None)]
   if Ops.OR in ops: pat += [(UPat.var("x", dtypes.bool).logical_not()&UPat.var("y", dtypes.bool).logical_not(),
     lambda x,y: (x | y).logical_not())]
   # rewrite MUL/IDIV to SHL+SHR: x*(2**y) -> shl(x,y) and x//(2**y) -> shr(x,y)
   if Ops.SHL in ops: pat += [(UPat.var("x", dtypes.ints)*UPat.cvar("c"), lambda c,x: x << v if (v:=powers_of_two.get(c.arg, 0)) else None)]
   if Ops.SHR in ops:
-    # no reason to check x<0 for uints
-    pat += [(UPat.var("x", dtypes.uints)//UPat.cvar("c"), lambda x,c: x >> v if (v:=powers_of_two.get(c.arg, 0)) else None)]
-    pat += [(UPat.var("x", dtypes.ints)//UPat.cvar("c"), lambda x,c: (x+(l.const_like(l.vmin) if (l:=(x<0)).vmin==l.vmax else l).where(
-      c-1, 0)) >> v if (v:=powers_of_two.get(c.arg, 0)) else None)]  # (x+(x<0).where(c-1, 0)) >> v
+    # uint IDIV by 2**v -> x >> v (FLOORDIV is lowered to IDIV by the rule above before reaching here)
+    pat += [(UPat(Ops.IDIV, src=(UPat.var("x", dtypes.uints), UPat.cvar("c"))),
+      lambda x,c: x >> v if (v:=powers_of_two.get(c.arg, 0)) else None)]
+    # signed IDIV (trunc) by 2**v -> (x + (x<0 ? c-1 : 0)) >> v
+    pat += [(UPat(Ops.IDIV, src=(UPat.var("x", dtypes.ints), UPat.cvar("c"))),
+      lambda x,c: (x+(l.const_like(l.vmin) if (l:=(x<0)).vmin==l.vmax else l).where(c-1, 0)) >> v
+        if (v:=powers_of_two.get(c.arg, 0)) else None)]
     if not disable_fast_idiv:
-      pat += [(UPat.var("x", dtypes.ints)//UPat.cvar("d", vec=False), lambda ctx, x, d: fast_idiv(ctx, x, d.arg))]
-      pat += [(UPat.var("x", dtypes.ints)%UPat.var("d"), lambda x, d: x-d*(x//d))]
+      # fast_idiv handles non-pow2: only fire on non-negative inputs (signed magic-mul is unreliable for x<0)
+      pat += [(UPat(Ops.IDIV, src=(UPat.var("x", dtypes.ints), UPat.cvar("d", vec=False))),
+        lambda ctx, x, d: fast_idiv(ctx, x, d.arg) if x.vmin >= 0 or x.dtype in dtypes.uints else None)]
+      # rewrite raw MOD -> x - d*IDIV(x,d) so fast_idiv can pick up the IDIV. only on non-negative inputs;
+      # avoids disturbing floormod_to_mod's general-path output (which uses a trunc Ops.MOD as an implementation detail)
+      pat += [(UPat(Ops.MOD, src=(UPat.var("x", dtypes.ints), UPat.var("d"))),
+        lambda x, d: x - d * x.alu(Ops.IDIV, d) if x.vmin >= 0 or x.dtype in dtypes.uints else None)]
   if Ops.NEG in ops:
     pat += [(UPat.var('x')*-1, lambda ctx,x: x.alu(Ops.NEG))]
     if Ops.SUB in ops: pat += [(UPat.var('x')+UPat.var('y').alu(Ops.NEG), lambda ctx,x,y: x.alu(Ops.SUB, y))]
@@ -488,8 +511,8 @@ pm_long_decomp = PatternMatcher([
   (UPat((*GroupOp.Defines, Ops.INDEX), name="x"), lambda x:
    x.replace(dtype=l2i_dt[x.dtype.base].ptr(x.dtype.size * 2)) if hasattr(x.dtype, 'size') and x.dtype.base in l2i_dt else None),
   (UPat(Ops.INDEX, tuple(l2i_dt.keys()), name='x'), lambda x: reindex(x, x.tag).replace(dtype=l2i_dt[x.dtype])),
-  (UPat(Ops.STORE, src=(UPat.var('idx'), UPat.var('val', tuple(l2i_dt.keys()))), name='st'), lambda st,idx,val:
-   st.replace(src=(reindex(idx, 0), val.rtag(0))).group(st.replace(src=(reindex(idx, 1), val.rtag(1)))) if val.tag is None else None),
+  (UPat(Ops.STORE, src=(UPat.var('idx'), UPat.var('val', tuple(l2i_dt.keys()))), allow_any_len=True, name='st'), lambda st,idx,val:
+   st.replace(src=(reindex(idx, 0), val.rtag(0))+st.src[2:]).group(st.replace(src=(reindex(idx, 1), val.rtag(1)))) if val.tag is None else None),
   (UPat(GroupOp.Comparison, src=(UPat.var('a', tuple(l2i_dt.keys())), UPat.var('b', tuple(l2i_dt.keys()))), name="x"), lambda a,b,x:
    l2i(x.op, dt:=l2i_dt[a.dtype], a.rtag(0).cast(dt), a.rtag(1).cast(dt), b.rtag(0).cast(dt), b.rtag(1).cast(dt))),
   (UPat(Ops.CAST, tuple(l2i_dt.keys()), src=(UPat.var('a'),), name="x"), lambda a,x:
@@ -501,7 +524,10 @@ pm_long_decomp = PatternMatcher([
   (UPat((*(GroupOp.ALU - GroupOp.Comparison), Ops.BITCAST), tuple(l2i_dt.keys()), name="x"), lambda x:
    l2i(x.op, l2i_dt[x.dtype], *flatten((a.rtag(0).cast(dt:=l2i_dt[x.src[-1].dtype]), a.rtag(1).cast(dt))
                                        if a.dtype in l2i_dt else (a,) for a in x.src))[x.tag] if x.tag is not None else None),
-  (UPat(Ops.LOAD, tuple(l2i_dt.keys()), src=(UPat.var('idx'),), name='x'), lambda x,idx: x.replace(dtype=l2i_dt[x.dtype],src=(reindex(idx, x.tag),))),
+  (UPat(Ops.LOAD, tuple(l2i_dt.keys()), src=(UPat.var('idx'), UPat.var('alt'), UPat.var('gate')), name='x'), lambda x,idx,alt,gate:
+   x.replace(dtype=l2i_dt[x.dtype], src=(reindex(idx, x.tag), alt.cast(l2i_dt[x.dtype]), gate))),
+  (UPat(Ops.LOAD, tuple(l2i_dt.keys()), src=(UPat.var('idx'),), name='x'), lambda x,idx:
+   x.replace(dtype=l2i_dt[x.dtype], src=(reindex(idx, x.tag),))),
   (UPat(Ops.CONST, tuple(l2i_dt.keys()), name='x'), lambda x:
    UOp.const(dt:=l2i_dt[x.dtype], truncate[dt]((x.arg >> 32) if x.tag == 1 else (x.arg & 0xFFFFFFFF))))
 ])
@@ -525,9 +551,9 @@ pm_float_decomp = PatternMatcher([
   (UPat(GroupOp.All-{Ops.BITCAST}, dtypes.floats, name="x"), lambda ctx,x:
    x.replace(dtype=ctx[1].vec(x.dtype.count), src=tuple(s.cast(ctx[1]) if s.dtype == ctx[0] else s for s in x.src))
    if x.dtype.scalar() == ctx[0] else None),
-  (UPat(Ops.STORE, src=(UPat.var("idx"), UPat(Ops.BITCAST, dtypes.floats, name="val")), name='st'), lambda ctx,st,idx,val:
-   st.replace(src=(idx, val.replace(dtype=f2f_dt[ctx[0]]))) if val.dtype == ctx[0] and idx.tag == ctx[0] else None),
-  (UPat(Ops.STORE, src=(UPat.var("idx"), UPat.var("val", dtypes.floats)), name='st'), lambda ctx,st,idx,val:
+  (UPat(Ops.STORE, src=(UPat.var("idx"), UPat(Ops.BITCAST, dtypes.floats, name="val")), allow_any_len=True, name='st'), lambda ctx,st,idx,val:
+   st.replace(src=(idx, val.replace(dtype=f2f_dt[ctx[0]]))+st.src[2:]) if val.dtype == ctx[0] and idx.tag == ctx[0] else None),
+  (UPat(Ops.STORE, src=(UPat.var("idx"), UPat.var("val", dtypes.floats)), allow_any_len=True, name='st'), lambda ctx,st,idx,val:
    f2f_store(st, idx, val, *ctx) if val.dtype.scalar() == ctx[1] and (idx:=idx.src[0] if idx.op == Ops.CAST else idx).tag == ctx[0] else None),
 ])
 

tinygrad/uop/divandmod.py

@@ -1,19 +1,19 @@
 import functools, itertools, math
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp
 from tinygrad.dtype import dtypes
-from tinygrad.helpers import cdiv, cmod, CORRECT_DIVMOD_FOLDING, unwrap
+from tinygrad.helpers import floordiv, floormod, unwrap
 
 # NOTE: this cache is only on index UOps
 @functools.cache
-def fold_divmod_general(d: UOp, correct_divmod_folding: bool) -> UOp|None:
+def fold_divmod_general(d: UOp) -> UOp|None:
   x, y = d.src
 
   # cancel_divmod: simple cancel div/mod case when the range of the numerator lies within a single denominator interval
   x_min, x_max, y_min, y_max = x.vmin, x.vmax, y.vmin, y.vmax
   assert isinstance(x_min, int) and isinstance(x_max, int) and isinstance(y_min, int) and isinstance(y_max, int)
-  if y_min==y_max==0: raise ZeroDivisionError(f"{'Division' if d.op is Ops.IDIV else 'Mod'} by zero trying to rewrite {x.alu(d.op, y)}")
-  if y_min*y_max > 0 and (qv:=cdiv(x_min,y_min)) == cdiv(x_min,y_max) == cdiv(x_max,y_min) == cdiv(x_max,y_max):
-    return x - qv*y if d.op is Ops.MOD else d.const_like(qv)
+  if y_min==y_max==0: raise ZeroDivisionError(f"{'Division' if d.op is Ops.FLOORDIV else 'Mod'} by zero trying to rewrite {x.alu(d.op, y)}")
+  if y_min*y_max > 0 and (qv:=floordiv(x_min,y_min)) == floordiv(x_min,y_max) == floordiv(x_max,y_min) == floordiv(x_max,y_max):
+    return x - qv*y if d.op is Ops.FLOORMOD else d.const_like(qv)
 
   # split uops for the rest of the processing
   x_peeled, const = x.pop_const()
@@ -22,19 +22,20 @@ def fold_divmod_general(d: UOp, correct_divmod_folding: bool) -> UOp|None:
   # ** Constant Denominator Rules **
   # these rules strictly require y to be a scalar constant > 0
   if y.op is Ops.CONST and (c := y.arg) > 0:
-    # nested_div_mod: (x%(k*c))//c -> (x//c)%k, and (x%(k*c))%c -> x%c
-    if x.op is Ops.MOD and (k := x.src[1].divides(c)) is not None:
-      return x.src[0] // y % k if d.op is Ops.IDIV else x.src[0] % y
+    # nested_div_mod: (x%(k*c))//c -> (x//c)%k (requires k>0), and (x%(k*c))%c -> x%c
+    if x.op is Ops.FLOORMOD and (k := x.src[1].divides(c)) is not None:
+      if d.op is Ops.FLOORMOD: return x.src[0] % y
+      if k > 0: return x.src[0] // y % k
 
-    # remove_nested_mod in sum: (a%4 + b)%2 -> (a+b)%2, requires non-negative sums
-    if d.op is Ops.MOD and x.vmin >= 0:
+    # remove_nested_mod in sum: (a%4 + b)%2 -> (a+b)%2
+    if d.op is Ops.FLOORMOD:
       new_xs, changed = [], False
       for u in uops_no_const:
-        if u.op is Ops.MOD and u.src[1].divides(c) is not None:
+        if u.op is Ops.FLOORMOD and u.src[1].divides(c) is not None:
           u = u.src[0]
           changed = True
         new_xs.append(u)
-      if changed and (new_x:=(UOp.usum(*new_xs) + const)).vmin >= 0: return new_x % y
+      if changed: return (UOp.usum(*new_xs) + const) % y
 
     # Shared decomposition for folding rules
     decomp = [(u.divides(f:=u.const_factor()),f) for u in uops_no_const]
@@ -42,40 +43,39 @@ def fold_divmod_general(d: UOp, correct_divmod_folding: bool) -> UOp|None:
 
     # fold_binary_numerator: fold if expression has one non-constant term that takes on two values
     if len(terms)==1 and (v:=terms[0]).vmax-v.vmin == 1:
-      y1 = (cmod if d.op is Ops.MOD else cdiv)(factors[0]*v.vmin+const, c)
-      y2 = (cmod if d.op is Ops.MOD else cdiv)(factors[0]*v.vmax+const, c)
+      y1 = (floormod if d.op is Ops.FLOORMOD else floordiv)(factors[0]*v.vmin+const, c)
+      y2 = (floormod if d.op is Ops.FLOORMOD else floordiv)(factors[0]*v.vmax+const, c)
       return (y2-y1)*(v-v.vmin) + y1
 
     # fold_divmod_congruence: fold if a is congruent to an expression whose range is between 0 and c
-    if not (x.vmin<0 and correct_divmod_folding):
-      # when f%c == c//2, abs(r) == abs(r-c) is a tie, try both signs since either may fit in one period
-      rem_choices = [(r, r-c) if (r:=f%c)*2 == c else (min(r, r-c, key=abs),) for f in factors]
-      for rems in itertools.product(*rem_choices):
-        if (rem:=sum(r*v for r,v in zip(rems,terms))+const%c).vmin//c==rem.vmax//c:
-          if d.op is Ops.MOD: return rem - rem.vmin//c*c
-          return sum((f-r)//c * v for f,r,v in zip(factors,rems,terms)) + const//c + rem.vmin//c
+    # when f%c == c//2, abs(r) == abs(r-c) is a tie, try both signs since either may fit in one period
+    rem_choices = [(r, r-c) if (r:=f%c)*2 == c else (min(r, r-c, key=abs),) for f in factors]
+    for rems in itertools.product(*rem_choices):
+      if (rem:=sum(r*v for r,v in zip(rems,terms))+const%c).vmin//c==rem.vmax//c:
+        if d.op is Ops.FLOORMOD: return rem - rem.vmin//c*c
+        return sum((f-r)//c * v for f,r,v in zip(factors,rems,terms)) + const//c + rem.vmin//c
 
     # gcd_with_remainder: factor out common gcd from numerator
-    if x.vmin >= 0 and (g:=math.gcd(*factors, c)) > 1:
+    if (g:=math.gcd(*factors, c)) > 1:
       new_x = unwrap(x_peeled.divides(g)).simplify() + (const//g)%(c//g)
       if new_x.vmin >= 0:
-        if d.op is Ops.MOD: return new_x % (c//g) * g + const%g
+        if d.op is Ops.FLOORMOD: return new_x % (c//g) * g + const%g
         return new_x // (c//g) + const//c
 
     # nest_by_factor: x//c -> (x//f)//(c//f), x%c -> (x//f%(c//f))*f + b where b=x%f
-    if x.vmin >= 0:
-      results = []
-      for div in {abs(f) for u, f in zip(uops_no_const, factors) if u.op not in (Ops.CONST, Ops.VCONST) and 1 < abs(f) < c and (c%f)==0}:
-        if (newxs := fold_divmod_general(x//div, correct_divmod_folding)) is not None and newxs.vmin >= 0:
-          if d.op is Ops.IDIV:
-            results.append((len(newxs.backward_slice), newxs // (c // div)))
-          else:
-            b_parts = [f%div*t for f, t in zip(factors, terms) if f%div]
-            if const % div: b_parts.append(x.const_like(const % div))
-            b = UOp.usum(*b_parts) if b_parts else x.const_like(0)
-            if 0 <= b.vmin and b.vmax < div:
-              results.append((len((r:=(newxs % x.ufix(c//div))*div + b).backward_slice), r))
-      if results: return min(results, key=lambda r: r[0])[1]
+    # FLOORDIV identity holds for any sign of x; FLOORMOD reconstruction needs x.vmin>=0
+    results = []
+    for div in {abs(f) for u, f in zip(uops_no_const, factors) if u.op not in (Ops.CONST, Ops.VCONST) and 1 < abs(f) < c and (c%f)==0}:
+      if (newxs := fold_divmod_general(x//div)) is not None:
+        if d.op is Ops.FLOORDIV:
+          results.append((len(newxs.backward_slice), newxs // (c // div)))
+        elif x.vmin >= 0 and newxs.vmin >= 0:
+          b_parts = [f%div*t for f, t in zip(factors, terms) if f%div]
+          if const % div: b_parts.append(x.const_like(const % div))
+          b = UOp.usum(*b_parts) if b_parts else x.const_like(0)
+          if 0 <= b.vmin and b.vmax < div:
+            results.append((len((r:=(newxs % x.ufix(c//div))*div + b).backward_slice), r))
+    if results: return min(results, key=lambda r: r[0])[1]
 
   # ** Variable Denominator / Fallback Rules **
   # These rules apply to variables OR constants that failed the checks above.
@@ -86,7 +86,7 @@ def fold_divmod_general(d: UOp, correct_divmod_folding: bool) -> UOp|None:
   gcd = UOp.gcd(*all_uops, y).simplify()
   if not (gcd.op is Ops.CONST and gcd.arg==1):
     ret = unwrap(x.divide_exact(gcd)).alu(d.op, unwrap(y.divide_exact(gcd)))
-    return ret*gcd if d.op is Ops.MOD else ret
+    return ret*gcd if d.op is Ops.FLOORMOD else ret
 
   # factor_remainder: (d*x+y)//d -> x+y//d
   if y.vmin<0 or x.vmin<0: return None
@@ -95,29 +95,22 @@ def fold_divmod_general(d: UOp, correct_divmod_folding: bool) -> UOp|None:
     if (q:=u.divide_exact(y)) is not None: quo.append(q)
     elif y.op is Ops.CONST and (c:=u.const_factor())%y.arg!=c:
       rem.append(u.divides(c)*(c%y.arg))
-      quo.append(u.divides(c)*(c//y.arg) if d.op is Ops.IDIV else u.const_like(0))
+      quo.append(u.divides(c)*(c//y.arg) if d.op is Ops.FLOORDIV else u.const_like(0))
     else: rem.append(u)
 
   if not quo: return None
   new_x = sum(rem)+x.const_like(0)
   if new_x.vmin<0: return None
-  return new_x%y if d.op is Ops.MOD else new_x//y+sum(quo)
+  return new_x%y if d.op is Ops.FLOORMOD else new_x//y+sum(quo)
 
 div_and_mod_symbolic = PatternMatcher([
   # ** 1. Fast Inline Rules **
-  ((UPat.var("x")//UPat.cvar("c") + UPat.cvar("a"))//UPat.cvar("d"), lambda x,c,a,d: (x+a*c)//(c*d)
-    if c.vmin>0 and d.vmin>0 and x.vmin>=0 and a.vmin>=0 else None),  # (x//c+a)//d -> (x+a*c)//(c*d)
-  (UPat.var("x", dtypes.weakint) // UPat.var("d"), lambda x,d: -(x//(-d)) if d.vmax < 0 else None),
-  (UPat.var("x", dtypes.weakint) // UPat.var("d"), lambda x,d: -((-x)//d) if x.vmax <= 0 else None),
-  ((UPat.var("x", dtypes.weakint)+UPat.cvar("c", vec=False)).named("n")//UPat.cvar("d", vec=False),
-    lambda x,c,n,d: ((x+c.arg%d.arg)//d + c.arg//d.arg) if c.arg%d.arg!=c.arg and x.vmin>=0 and n.vmin>=0 and d.arg>0 else None),
-  ((UPat.var("x", dtypes.weakint)+UPat.cvar("c", vec=False)).named("n")//UPat.cvar("d", vec=False),
-    lambda x,c,n,d: (-(-(c.arg%d.arg + x - (d.arg-1))//d) + c.arg//d.arg) if x.vmax<=0 and n.vmin>=0 and d.arg>0 else None),
+  # (x//c+a)//d -> (x+a*c)//(c*d) for c>0, d>0
+  ((UPat.var("x")//UPat.cvar("c") + UPat.cvar("a"))//UPat.cvar("d"), lambda x,c,a,d: (x+a*c)//(c*d) if c.vmin>0 and d.vmin>0 else None),
+  # (x+c)//d -> (x+c%d)//d + c//d for d>0 (split out the multiple of d in the constant)
+  ((UPat.var("x", dtypes.weakint)+UPat.cvar("c", vec=False))//UPat.cvar("d", vec=False),
+    lambda x,c,d: (x+c.arg%d.arg)//d + c.arg//d.arg if c.arg%d.arg!=c.arg and d.arg>0 else None),
 
   # ** 2. Slow Rules **
-  (UPat((Ops.IDIV, Ops.MOD), dtypes.weakint, name="d"), lambda d: fold_divmod_general(d, bool(CORRECT_DIVMOD_FOLDING))),
-
-  # NOTE: these have to go at the bottom or TestSymbolicOps.test_var loops
-  (UPat.var("x", dtypes.weakint) % UPat.var("d"), lambda x,d: -((-x)%d) if x.vmax <= 0 else None),
-  (UPat.var("x", dtypes.weakint) % UPat.var("d"), lambda x,d: (x%(-d)) if d.vmax < 0 else None),
-])
+  (UPat((Ops.FLOORDIV, Ops.FLOORMOD), dtypes.weakint, name="d"), lambda d: fold_divmod_general(d)),
+])

tinygrad/uop/ops.py

@@ -8,8 +8,9 @@ from tinygrad.dtype import ConstType, ImageDType, dtypes, DType, DTypeLike, to_d
 from tinygrad.dtype import ConstFloat, PyConst, storage_fmt_for_dtype, to_storage_scalar, from_storage_scalar
 from tinygrad.device import Buffer, MultiBuffer, canonicalize_device
 from tinygrad.helpers import ContextVar, all_int, prod, getenv, all_same, Context, partition, temp, unwrap, T, argfix, Metadata, flatten, TRACEMETA
-from tinygrad.helpers import PROFILE, dedup, cdiv, cmod, diskcache_put, to_function_name, cpu_profile, TracingKey, VIZ, SPEC, CAPTURE_PROCESS_REPLAY
-from tinygrad.helpers import strip_parens, colored, ansilen, printable
+from tinygrad.helpers import PROFILE, dedup, cdiv, cmod, floordiv, floormod, diskcache_put, to_function_name, cpu_profile, TracingKey
+from tinygrad.helpers import VIZ, SPEC, CAPTURE_PROCESS_REPLAY
+from tinygrad.helpers import colored, ansilen, printable
 if TYPE_CHECKING:
   from tinygrad.renderer import Estimates
 
@@ -26,7 +27,7 @@ axis_colors = {AxisType.GLOBAL: "blue", AxisType.THREAD: "BLUE", AxisType.LOCAL:
 axis_to_pos = {AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2, AxisType.UPCAST: 3,
                AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5}
 
-range_start = {Ops.BUFFERIZE: 1, Ops.REDUCE: 1, Ops.STORE: 2, Ops.WMMA: 3, Ops.END: 1, Ops.CALL: 1, Ops.FUNCTION: 1,
+range_start = {Ops.BUFFERIZE: 1, Ops.REDUCE: 1, Ops.WMMA: 3, Ops.END: 1, Ops.CALL: 1, Ops.FUNCTION: 1,
                Ops.COPY: 2, Ops.BUFFER_VIEW: 1, Ops.LINEAR: 0}
 
 # https://en.wikipedia.org/wiki/Identity_element
@@ -77,16 +78,6 @@ def consumer_map_from_toposort(lst:Iterable[UOp]):
       if s in ret: ret[s][u] = None
   return ret
 
-def pretty_print(x:UOp, cache=None, d=0)->str:
-  def dfs(x:UOp, cache:dict):
-    for s in x.src:
-      cache.setdefault(s, [len(cache), 0, False])[1] += 1
-      if cache[s][1] == 1: dfs(s, cache)
-  if cache is None: dfs(x, cache:={})
-  if (cx:=cache.setdefault(x, [0,0,False]))[2]: return f"{' '*d}x{cx[0]}"
-  cx[2], srcs = True, (''.join(f'\n{pretty_print(s, cache, d+2)},' for s in x.src))
-  return f"{' '*d}{f'x{cx[0]}:=' * (cx[1]>1)}{type(x).__name__}({x.op}, {x.dtype}, arg={x.argstr()}{x.tagstr()}, src=({srcs}))"
-
 class UOpMetaClass(type):
   ucache:dict[tuple, weakref.ReferenceType[UOp]] = {}
   def __call__(cls, op:Ops, dtype:DType=dtypes.void, src:tuple[UOp,...]=tuple(), arg:Any=None, tag:Any=None,
@@ -100,13 +91,17 @@ class UOpMetaClass(type):
       buffers[created] = _buffer
     if SPEC > 1:
       from tinygrad.uop.spec import full_spec, test_pyrender
-      if SPEC > 2: test_pyrender(created)
+      if SPEC > 2:
+        # SPEC=3 checks the shape
+        _ = created._shape
+        if SPEC > 3:
+          test_pyrender(created)
       with Context(CHECK_OOB=0): fret = cast(bool|None, full_spec.rewrite(created))
       if fret is not True: raise RuntimeError(f"SPEC ISSUE {fret}: {created}")
     return created
 
 # some uops map to other stuff
-buffers:weakref.WeakKeyDictionary[UOp, Buffer|MultiBuffer] = weakref.WeakKeyDictionary() # this maps BUFFER uops to their device Buffers
+buffers:weakref.WeakKeyDictionary[UOp, Buffer|MultiBuffer] = weakref.WeakKeyDictionary() # this maps BUFFER/BUFFER_VIEW uops to their device Buffers
 all_metadata:weakref.WeakKeyDictionary[UOp, tuple[Metadata, ...]] = weakref.WeakKeyDictionary() # TODO: should this be here?
 
 # recursive_property replaces functools.cached_property in recursive UOp functions to prevent RecursionError
@@ -150,7 +145,9 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
   @recursive_property
   def key(self) -> bytes:
     return hashlib.sha256(str((self.op, self.dtype, self.arg)).encode() + b"".join([s.key for s in self.src])).digest()
-  def __repr__(self): return pretty_print(self)
+  def __repr__(self):
+    from tinygrad.uop.render import pretty_print
+    return pretty_print(self)
   def argstr(self):
     if self.op is Ops.REDUCE: return f'({", ".join(map(str, self.arg))})'
     return f"ConstFloat({float.__repr__(self.arg)})" if isinstance(self.arg, ConstFloat) else repr(self.arg)
@@ -212,7 +209,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     match self.op:
       # late ops don't have shape
       case Ops.UNIQUE | Ops.LUNIQUE | Ops.DEVICE | Ops.IF | Ops.BARRIER | Ops.CUSTOM | Ops.CUSTOMI | \
-           Ops.STACK | Ops.GEP | Ops.UNROLL | Ops.CONTRACT | Ops.SINK | Ops.END | Ops.REWRITE_ERROR | \
+           Ops.CONTRACT | Ops.SINK | Ops.END | Ops.REWRITE_ERROR | Ops.PTRCAT | Ops.ENDIF | \
            Ops.LINEAR | Ops.PROGRAM | Ops.SOURCE | Ops.BINARY | Ops.INS | Ops.TUPLE | Ops.CALL | Ops.FUNCTION:
         return None
 
@@ -228,24 +225,26 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
         return inner_shape
 
       case Ops.CAST:
-        # when PTX casts from ptr to non ptr, remove the shape
+        # when PTX casts from ptr to non ptr, remove the shape of the buffer
         if isinstance(self.src[0].dtype, PtrDType) and not isinstance(self.src[0].dtype, ImageDType) and not isinstance(self.dtype, PtrDType):
-          return None
+          return ()
 
       case Ops.INDEX:
-        # non pointer index doesn't have a shape
-        if not isinstance(self.dtype, PtrDType): return None
-        # fully indexed doesn't have a shape. TODO: remove this
-        if self.src[0]._shape is None or len(self.src[1:]) == len(self.src[0].shape): return None
-        # pointer index
-        return self.src[0].shape[len(self.src[1:]):]
+        shp:list[sint] = []
+        for s in self.src[1:]: shp.extend(list(s.shape))
+        return tuple(shp) + self.src[0].shape[len(self.src[1:]):]
+
+      # TODO: these should have the shape of the dtype.count
+      case Ops.CONST | Ops.DEFINE_VAR: return ()
+      case Ops.GEP | Ops.STACK | Ops.VCONST | Ops.VCAT: return ()
 
       # some ops init the shape
-      case Ops.CONST | Ops.DEFINE_VAR | Ops.BIND | Ops.RANGE | Ops.SPECIAL: return ()
-      # TODO: VCONST should have the shape of the arg
-      case Ops.VCONST: return ()
+      case Ops.BIND | Ops.RANGE | Ops.SPECIAL | Ops.UNROLL: return ()
       case Ops.BUFFER: return (self.arg,)
-      case Ops.BUFFER_VIEW: return (self.arg[0],)
+      case Ops.BUFFER_VIEW:
+        # HACK: BUFFER_VIEW is used inside kernels, so we set the shape to () if it's on an INDEX
+        if self.src[0].op is Ops.INDEX: return ()
+        return (self.arg[0],)
       case Ops.CUSTOM_FUNCTION: return None
       case Ops.BUFFERIZE: return tuple([int(r.vmax+1) for r in self.src[1:]])
       case Ops.DEFINE_LOCAL | Ops.DEFINE_REG: return (self.ptrdtype.size,)
@@ -280,7 +279,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     # NOTE: ssimplify is required because the shape needs to be canonical for broadcasting and same shape checking
     if self.op in GroupOp.Movement.union({Ops.MULTI, Ops.REDUCE}):
       ps = self.src[0]._shape
-      if ps is None: raise RuntimeError(f"movement op {self.op} requires shape")
+      if ps is None: raise RuntimeError(f"movement op {self.op} requires shape, {self.src[0].op} doesn't have one")
       match self.op:
         case Ops.RESHAPE:
           if not all(x >= 0 for x in self.marg): raise ValueError(f"shape can't contain negative numbers {self.marg}")
@@ -316,7 +315,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     if self.op in GroupOp.ALU.union({Ops.CAST, Ops.COPY, Ops.NOOP, Ops.GROUP, Ops.SINK, Ops.ALLREDUCE, Ops.STORE}):
       input_shapes = [x._shape for x in self.src if x._shape is not None]
       if len(input_shapes) == 0: return None
-      if not all_same(input_shapes): raise RuntimeError(f"shape mismatch at {self.op}: {input_shapes}")
+      if not all_same(input_shapes): raise RuntimeError(f"shape mismatch at {self.op}: {input_shapes} {[x.op for x in self.src]}")
       return input_shapes[0]
 
     # all Ops must be explicitly handled
@@ -419,8 +418,8 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
   def group(*srcs:UOp|None):  # pylint: disable=no-self-argument
     if len(srcs) == 1 and isinstance(srcs[0], UOp): return srcs[0]
     return UOp(Ops.GROUP, dtypes.void, tuple([x for x in srcs if x is not None]))
-  def vectorize(self, *srcs, **kwargs):
-    return UOp(Ops.STACK, self.dtype.vec(len(srcs)+1), (self,)+srcs, **kwargs)
+  def vectorize(self, *srcs):
+    return UOp(Ops.STACK, self.dtype.vec(len(srcs)+1), (self,)+srcs)
   def index(self, *srcs:UOp|None, ptr=False, **kwargs):
     return UOp(Ops.INDEX, kwargs.pop("dtype", self.dtype if ptr else self.dtype.base), (self,)+tuple([x for x in srcs if x is not None]), **kwargs)
   def __getitem__(self, idx):
@@ -464,8 +463,9 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
       i = (i,)
     return UOp(Ops.GEP, self.dtype.scalar().vec(len(i)) if len(i) > 1 else self.dtype.scalar(), (self,), i)
   def load(self, *src:UOp, **kwargs): return UOp(Ops.LOAD, dtype=kwargs.pop("dtype", self.dtype.base), src=(self,)+src, **kwargs)
-  def store(self, src:UOp|ConstType, **kwargs):
-    return UOp(Ops.STORE, dtypes.void, (self, self.const_like(src) if not isinstance(src, UOp) else src), **kwargs)
+  def store(self, src:UOp|ConstType, gate:UOp|None=None, **kwargs):
+    srcs = (self, self.const_like(src) if not isinstance(src, UOp) else src) + ((gate,) if gate is not None else ())
+    return UOp(Ops.STORE, dtypes.void, srcs, **kwargs)
   def end(self, *src:UOp): return UOp(Ops.END, src=(self,)+src) if len(src) else self
   def after(self, *src:UOp, **kwargs): return UOp(Ops.AFTER, self.dtype, (self,)+src, **kwargs) if len(src) else self
   def barrier(self, *src:UOp): return UOp(Ops.BARRIER, src=(self,)+src)
@@ -737,13 +737,16 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
       assert isinstance(buf, Buffer), "must be a Buffer for BITCAST"
       return buf.view(prod(self.max_shape), self.dtype, 0)
     if self.op is Ops.BUFFER_VIEW:
+      if (cret:=buffers.get(self)) is not None: return cret
       buf = self.src[0].buffer
       if isinstance(buf, MultiBuffer):
         mbuf = MultiBuffer.__new__(MultiBuffer)
         mbuf.bufs = [b.view(self.arg[0], self.dtype, self.arg[1] * self.dtype.itemsize) for b in buf.bufs]
+        buffers[self] = mbuf
         return mbuf
       assert isinstance(buf, Buffer), "must be a Buffer for BUFFER_VIEW"
-      return buf.view(self.arg[0], self.dtype, self.arg[1] * self.dtype.itemsize)
+      buffers[self] = bv = buf.view(self.arg[0], self.dtype, self.arg[1] * self.dtype.itemsize)
+      return bv
     if self.op is Ops.MSELECT:
       ret = self.src[0].buffer
       assert isinstance(ret, MultiBuffer)
@@ -804,7 +807,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     # is f a monotonically increasing function regards its input
     if self.op in GroupOp.Irreducible: return True
     if self.op is Ops.ADD: return self.src[0].is_increasing() and self.src[1].is_increasing()
-    if self.op in (Ops.MUL, Ops.IDIV) and self.src[1].op is Ops.CONST and self.src[1].arg >= 0: return self.src[0].is_increasing()
+    if self.op in (Ops.MUL, Ops.IDIV, Ops.FLOORDIV) and self.src[1].op is Ops.CONST and self.src[1].arg >= 0: return self.src[0].is_increasing()
     return False  # False if not sure
   def const_factor(self) -> int:
     """largest known int that divides self"""
@@ -865,6 +868,17 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
         assert isinstance(s0_vmin, int) and isinstance(s0_vmax, int) and isinstance(s1_vmin, int) and isinstance(s1_vmax, int)
         if s1_vmin*s1_vmax>0:
           return min(vals:=(cdiv(s0_vmin, s1_vmin), cdiv(s0_vmin, s1_vmax), cdiv(s0_vmax, s1_vmin), cdiv(s0_vmax, s1_vmax))), max(vals)
+      if self.op is Ops.FLOORDIV:
+        assert isinstance(s0_vmin, int) and isinstance(s0_vmax, int) and isinstance(s1_vmin, int) and isinstance(s1_vmax, int)
+        if s0_vmin > s0_vmax: return 0, 0  # numerator range is empty (e.g. RANGE with end=0)
+        if s1_vmin*s1_vmax>0: return min(vals:=(s0_vmin//s1_vmin, s0_vmin//s1_vmax, s0_vmax//s1_vmin, s0_vmax//s1_vmax)), max(vals)
+      if self.op is Ops.FLOORMOD:
+        assert isinstance(s0_vmin, int) and isinstance(s0_vmax, int) and isinstance(s1_vmin, int) and isinstance(s1_vmax, int)
+        if s0_vmin > s0_vmax: return 0, 0  # numerator range is empty (e.g. RANGE with end=0)
+        if (c:=s1_vmin) == s1_vmax > 0: return (s0_vmin%c, s0_vmax%c) if s0_vmin//c == s0_vmax//c else (0, c-1)
+        if (c:=s1_vmin) == s1_vmax < 0: return (s0_vmin%c, s0_vmax%c) if s0_vmin//c == s0_vmax//c else (c+1, 0)
+        if s1_vmin > 0: return (0, s1_vmax-1)
+        if s1_vmax < 0: return (s1_vmin+1, 0)
       if self.op is Ops.XOR and s1_vmin == s1_vmax == -1 and isinstance(s0_vmin, int) and isinstance(s0_vmax, int): return ~s0_vmax, ~s0_vmin
       if self.op is Ops.MAX: return max(s0_vmin, s1_vmin), max(s0_vmax, s1_vmax)
       if self.op is Ops.CMPLT: return (s0_vmax<s1_vmin, s0_vmin<s1_vmax)
@@ -889,12 +903,13 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
 
   @functools.cached_property
   def _sym_fxn(self):
+    from tinygrad.uop.render import _render_with_splits, renderer_infer
     sself = self.simplify()
     varnames = tuple(x.expr for x in sself.toposort() if x.op is Ops.DEFINE_VAR)
     # TODO: sanitize varnames, or don't use naked eval while staying fast
     ret = _render_with_splits(list(sself.toposort()), renderer_infer, {sself})
     lines = [f"  {k}={v}" for k,v in ret.items() if k != "ast"] + [f"  return {ret['ast']}"]
-    ns: dict[str, Any] = {"max": max, "cdiv": cdiv, "cmod": cmod, "bitcast": bitcast, "dtypes": dtypes}
+    ns: dict[str, Any] = {"max": max, "cdiv": cdiv, "cmod": cmod, "floordiv": floordiv, "floormod": floormod, "bitcast": bitcast, "dtypes": dtypes}
     exec(f"def _f({','.join(varnames)}):\n"+'\n'.join(lines), ns)  # pylint: disable=exec-used
     return ns["_f"], varnames
 
@@ -904,12 +919,15 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
 
   def render(self, simplify=True, pm:PatternMatcher|None=None) -> str:
     ctx: dict[UOp, str] = {}
+    from tinygrad.uop.render import renderer
     pm = renderer if pm is None else pm
     for u in (s:=self.simplify() if simplify else self).toposort():
       ctx[u] = cast(str, pm.rewrite(u, ctx=ctx))
     return ctx[s]
 
-  def pyrender(self): return pyrender(self)
+  def pyrender(self):
+    from tinygrad.uop.render import pyrender
+    return pyrender(self)
 
   # *** uop high level syntactic sugar ***
 
@@ -1047,7 +1065,8 @@ python_alu: dict[Ops, Callable]  = {
   Ops.SIN: lambda x: math.sin(x) if not math.isinf(x) else math.nan, Ops.POW: safe_pow, Ops.TRUNC: math.trunc,
   Ops.NEG: operator.neg, Ops.ADD: operator.add, Ops.SUB: operator.sub, Ops.MUL: operator.mul, Ops.CMPNE: operator.ne, Ops.CMPLT: operator.lt,
   Ops.XOR: operator.xor, Ops.OR: operator.or_, Ops.AND: operator.and_, Ops.SHR: operator.rshift, Ops.SHL: operator.lshift, Ops.MAX: max,
-  Ops.MOD: cmod, Ops.IDIV: cdiv, Ops.MULACC: lambda x,y,z: (x*y)+z, Ops.WHERE: lambda x,y,z: y if x else z, Ops.CMPEQ: operator.eq}
+  Ops.MOD: cmod, Ops.IDIV: cdiv, Ops.FLOORDIV: floordiv, Ops.FLOORMOD: floormod,
+  Ops.MULACC: lambda x,y,z: (x*y)+z, Ops.WHERE: lambda x,y,z: y if x else z, Ops.CMPEQ: operator.eq}
 
 def exec_alu(op:Ops, dtype:DType, operands, truncate_output=True):
   if dtype.count > 1:
@@ -1056,13 +1075,15 @@ def exec_alu(op:Ops, dtype:DType, operands, truncate_output=True):
   alu = python_alu[op](*operands)
   return truncate.get(dtype, lambda x: x)(alu) if truncate_output else alu
 
-# ***** uop helpers *****
-
-def print_uops(uops:list[UOp]):
-  uops_index = {u:i for i,u in enumerate(uops)}
-  for i,u in enumerate(uops):
-    formatted_srcs = [(uops_index[x] if x.op is not Ops.CONST else f"{x.arg}") if x in uops else "--" for x in u.src]
-    print(f"{i:4d} {str(u.op):20s}: {multirange_str(u.ranges, color=True, pad=10)} {str(u.dtype):40s} " f"{str(formatted_srcs):32s} {u.arg}")
+def bitcast(x, in_dtype:DType, out_dtype:DType):
+  assert in_dtype.itemsize == out_dtype.itemsize, "bitcast itemsize mismatch"
+  in_count, out_count = in_dtype.count, out_dtype.count
+  in_vals = (x,) if in_count == 1 else tuple(x)
+  assert len(in_vals) == in_count, f"bitcast expected {in_count} values, got {len(in_vals)}"
+  packed = struct.pack(f"{in_count}{storage_fmt_for_dtype(in_dtype.scalar())}", *[to_storage_scalar(v, in_dtype.scalar()) for v in in_vals])
+  out_vals = struct.unpack(f"{out_count}{storage_fmt_for_dtype(out_dtype.scalar())}", packed)
+  ret = tuple(from_storage_scalar(v, out_dtype.scalar()) for v in out_vals)
+  return ret[0] if out_count == 1 else ret
 
 # ***** pattern matcher *****
 
@@ -1080,7 +1101,7 @@ class UPat(OpMixin):
   def __init__(self, op:Ops|tuple[Ops, ...]|set[Ops]|None=None, dtype:DType|tuple[DType, ...]|set[DType]|None=None,
                src:tuple[UPat, ...]|list[UPat]|UPat|None=None, arg:Any=None,
                name:str|None=None, allow_any_len:bool=False, custom_early_reject:set[Ops]|None=None, location=None, is_any:bool=False):
-    assert op is None or isinstance(op, (Ops, tuple, set)), "op must be Ops or tuple of Ops"
+    assert op is None or isinstance(op, (Ops, tuple, set)), f"op must be Ops or tuple of Ops, not {op!r}"
     self.op: tuple[Ops, ...]|None = (op,) if isinstance(op, Ops) else (tuple(op) if isinstance(op, set) else op)
     self.match_dtype: tuple[DType, ...]|None = (dtype,) if isinstance(dtype, DType) else (tuple(dtype) if isinstance(dtype, set) else dtype)
     self.arg, self.name, self._in_src, self.custom_early_reject = arg, name, src, custom_early_reject
@@ -1136,8 +1157,8 @@ class UPat(OpMixin):
 
   # copied from UOp
   def sink(self, *srcs:UPat|None, **kwargs): return UPat(Ops.SINK, dtypes.void, (self,)+tuple([x for x in srcs if x is not None]), **kwargs)
-  def index(self, idx:UPat, valid:UPat|None=None, **kwargs):
-    return UPat(Ops.INDEX, self.match_dtype, (self,idx,valid) if valid is not None else (self,idx), **kwargs)
+  def index(self, *srcs:UPat|None, **kwargs):
+    return UPat(Ops.INDEX, self.match_dtype, (self,)+tuple(x for x in srcs if x is not None), **kwargs)
   def cast(self, dtype=None, **kwargs):
     if dtype is not None and self.match_dtype == (dtype,): return self
     return UPat(Ops.CAST, dtype, (self,), **kwargs)
@@ -1362,10 +1383,9 @@ if TRACK_MATCH_STATS or PROFILE:
       with open(fn:=temp("rewrites.pkl", append_user=True), "wb") as f:
         print(f"rewrote {len(tracked_ctxs)} graphs and matched {sum(len(r.matches) for x in tracked_ctxs for r in x)} times, saved to {fn}")
         pickle.dump(RewriteTrace(tracked_keys, tracked_ctxs, uop_fields), f)
-    if VIZ > 0:
-      TRACK_MATCH_STATS.value = 0
-      return launch_viz("REWRITE_DATA", temp("rewrites.pkl", append_user=True))
-    if getenv("PRINT_MATCH_STATS", TRACK_MATCH_STATS.value and VIZ.value>=0):
+    TRACK_MATCH_STATS.value = 0
+    launch_viz("REWRITE_DATA", temp("rewrites.pkl", append_user=True))
+    if getenv("PRINT_MATCH_STATS", TRACK_MATCH_STATS.value and not VIZ):
       ret = [0,0,0.0,0.0]
       for k,v in sorted(list(match_stats.items()), key=lambda x: x[1][2]+x[1][3]):
         loc_str = f"{k.location[0].split('/')[-1]}:{k.location[1]}"
@@ -1381,8 +1401,9 @@ if TRACK_MATCH_STATS or PROFILE:
       args = ['--rewrites-path', os.getenv("REWRITE_DATA", "")] if os.getenv("REWRITE_DATA", "") else []
       args += ['--profile-path', os.getenv("PROFILE_DATA", "")] if os.getenv("PROFILE_DATA", "") else []
       viz_path = pathlib.Path(__file__).resolve().parent.parent / "viz" / "serve.py"
-      if sys.stdout.isatty(): os.execv(sys.executable, [sys.executable, viz_path.as_posix()] + args)
-      else: print("Successfully saved VIZ files, view using: python -m tinygrad.viz.cli")
+      if VIZ > 0 and sys.stdout.isatty(): os.execv(sys.executable, [sys.executable, viz_path.as_posix()] + args)
+      if VIZ: print("saved viz files, view using: python -m tinygrad.viz.cli")
+      VIZ.value = 0
 
 # *** simple graph rewrite engine ***
 
@@ -1523,12 +1544,6 @@ pm_lower_index_dtype = PatternMatcher([
   (UPat(Ops.DEFINE_VAR, dtype=dtypes.weakint, name="u"), lambda u: u.replace(dtype=dtypes.int).cast(dtypes.weakint)),
   (UPat(Ops.BIND, src=(UPat.var("var").cast(dtypes.weakint), UPat.cvar("val").cast(dtypes.weakint))),
     lambda var,val: var.bind(val).cast(dtypes.weakint)),
-  # lower Invalid
-  (UPat.var("buf").index(UPat.var("cond").where(UPat.var("idx"), UPat(Ops.CONST, arg=Invalid))), lambda buf,idx,cond: buf.index(idx, cond, ptr=True)),
-  # remove hanging casts
-  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx", dtypes.ints).cast()),), lambda buf,idx: buf.index(idx, ptr=True)),
-  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx", dtypes.ints).cast(), UPat.var("valid"))),
-   lambda buf,idx,valid: buf.index(idx, valid, ptr=True)),
   (UPat((Ops.SINK, Ops.NOOP, Ops.END), name="n"),
    lambda n: n.replace(src=tuple(s.src[0] if s.op is Ops.CAST and s.dtype == dtypes.weakint else s for s in n.src))),
   # vectorized indexes (ie. images) must be int
@@ -1552,152 +1567,6 @@ def do_unbind(ctx:dict[Variable, int], x:UOp):
   return v
 pm_unbind = PatternMatcher([(UPat(Ops.BIND, name="x"), do_unbind)])
 
-# for debug
-syms = { Ops.ADD: "+", Ops.SUB: "-", Ops.IDIV: "//", Ops.MOD: "%", Ops.SHL: "<<", Ops.SHR: ">>",
-         Ops.MUL: "*", Ops.CMPLT: "<", Ops.CMPNE: "!=", Ops.AND: "&", Ops.OR: "|", Ops.XOR: "^"}
-# comparison operators are not in here because they are chained in python, not left-associative
-precedence = {Ops.MUL:1, Ops.IDIV:1, Ops.MOD:1, Ops.ADD:2, Ops.SUB:2, Ops.SHL:3, Ops.SHR:3, Ops.AND:4, Ops.XOR:5, Ops.OR:6}
-def strip_binary_parens(x:UOp, left:str, right:str, code_for_op) -> str:
-  if x.op not in precedence: return code_for_op(left, right)
-  return code_for_op(strip_parens(left) if precedence.get(x.src[0].op,99)<=precedence[x.op] else left, strip_parens(right) if
-    precedence.get(x.src[1].op,99)<precedence[x.op] else right)
-
-def bitcast(x, in_dtype:DType, out_dtype:DType):
-  assert in_dtype.itemsize == out_dtype.itemsize, "bitcast itemsize mismatch"
-  in_count, out_count = in_dtype.count, out_dtype.count
-  in_vals = (x,) if in_count == 1 else tuple(x)
-  assert len(in_vals) == in_count, f"bitcast expected {in_count} values, got {len(in_vals)}"
-  packed = struct.pack(f"{in_count}{storage_fmt_for_dtype(in_dtype.scalar())}", *[to_storage_scalar(v, in_dtype.scalar()) for v in in_vals])
-  out_vals = struct.unpack(f"{out_count}{storage_fmt_for_dtype(out_dtype.scalar())}", packed)
-  ret = tuple(from_storage_scalar(v, out_dtype.scalar()) for v in out_vals)
-  return ret[0] if out_count == 1 else ret
-
-renderer = PatternMatcher([
-  (UPat((Ops.DEFINE_VAR,), name="x"), lambda x: x.expr),
-  (UPat(Ops.PARAM, src=(UPat(), UPat(), UPat(), UPat(), UPat(Ops.NOOP, name="x"))), lambda x: x.arg),
-  (UPat((Ops.SPECIAL), name="x"), lambda x: x.arg),
-  (UPat(Ops.RANGE, name="x"), lambda x: f"r{range_str(x)}"),
-  (UPat((Ops.CONST, Ops.VCONST), name="x"), lambda x: str(x.arg)),
-  (UPat(Ops.UNROLL, name="x"), lambda ctx,x,u: f"UNROLL({ctx[x.src[0]]}, {u.arg})"),
-  (UPat(Ops.CAST, name="x"), lambda ctx,x: f"({str(x.dtype)[7:]})({ctx[x.src[0]]})"),
-  (UPat(Ops.BIND, name="x"), lambda ctx,x: ctx[x.src[0]]),
-  (UPat(Ops.NEG, name="x"), lambda ctx,x: f"(-{ctx[x.src[0]]})"),
-  (UPat(Ops.RECIPROCAL, name="x"), lambda ctx,x: f"(1/{ctx[x.src[0]]})"),
-  (UPat(Ops.MAX, name="x"), lambda ctx,x: f"max({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
-  (UPat(Ops.MULACC, name="x"), lambda ctx,x: f"({ctx[x.src[0]]}*{ctx[x.src[1]]}+{ctx[x.src[2]]})"),
-  (UPat(Ops.WHERE, name="x"), lambda ctx,x: f"({ctx[x.src[1]]} if {ctx[x.src[0]]} else {ctx[x.src[2]]})"),
-  (UPat(set(syms.keys()), name="x"), lambda ctx,x: strip_binary_parens(x, ctx[x.src[0]], ctx[x.src[1]], lambda a,b: f"({a}{syms[x.op]}{b})")),
-  (UPat((Ops.INDEX, Ops.BUFFERIZE), name="x"), lambda x, ctx: ''.join([f"[{strip_parens(ctx[y])}]" for y in x.src[1:]])),
-  (UPat(Ops.STACK, name="x"),
-   lambda ctx,x: f"{{{','.join([ctx[y] for y in x.src])}}}" if not x.src or not all_same(x.src) else f"{{{ctx[x.src[0]]}, ...}}"),
-  (UPat(GroupOp.All, name="x"), lambda x: str(x)),
-])
-
-renderer_infer = PatternMatcher([
-  (UPat(Ops.MOD, name="x"), lambda ctx,x: f"cmod({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
-  (UPat(Ops.IDIV, name="x"), lambda ctx,x: f"cdiv({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
-  (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"bitcast({ctx[x.src[0]]}, {x.src[0].dtype!r}, {x.dtype!r})"),
-]) + renderer
-
-# *** pyrender ***
-
-def srcs(ctx, src): return f"({ctx[src[0]]},)" if len(src) == 1 else f"({', '.join([ctx[x] for x in src])})"
-def render_marg(ctx,x:UOp):
-  if x.op is Ops.PERMUTE: return str(x.marg)
-  if x.op is Ops.FLIP: return str(tuple([i for i,x in enumerate(x.marg) if x]))
-  pieces = []
-  if x.op in {Ops.RESHAPE, Ops.EXPAND}:
-    pieces = [f"{ctx[a] if isinstance(a, UOp) else str(a)}" for a in x.marg]
-  if x.op in {Ops.PAD, Ops.SHRINK}:
-    pieces = [f"({ctx[a[0]] if isinstance(a[0], UOp) else str(a[0])}, {ctx[a[1]] if isinstance(a[1], UOp) else str(a[1])})" for a in x.marg]
-  return f"({','.join(pieces)})" if len(pieces) != 1 else f"({pieces[0]},)"
-
-sugar = {Ops.SINK, Ops.END, Ops.STORE, Ops.LOAD, Ops.UNIQUE, Ops.SQRT, Ops.INDEX, Ops.REDUCE, Ops.AFTER, Ops.THREEFRY,
-         Ops.WHERE, Ops.RECIPROCAL, Ops.EXP2, Ops.LOG2, Ops.SIN, Ops.CONTIGUOUS, Ops.BARRIER, Ops.DETACH}
-pm_pyrender_extra = PatternMatcher([
-  (UPat(Ops.CONST, src=(UPat(Ops.UNIQUE, name="u"), UPat(Ops.DEVICE, name="d")), name="x"),
-   lambda x,u,d: f"UOp.unique_const({x.arg}, dtype={x.dtype}, device={repr(d.arg)}, unique={u.arg})"),
-  (UPat(Ops.CONST, src=(UPat(Ops.DEVICE, name="d"),), name="x"), lambda x,d: f"UOp.const({x.dtype}, {x.arg}, device={repr(d.arg)})"),
-  (UPat(Ops.CONST, src=(), name="x"), lambda x: f"UOp.const({x.dtype}, {x.arg})"),
-  (UPat(Ops.DEFINE_VAR, src=(), name="x"), lambda x:
-    f"UOp.variable(\"{x.arg[0]}\", {x.arg[1]}, {x.arg[2]}{', dtype='+str(x.dtype) if x.dtype is not dtypes.weakint else ''})"),
-  (UPat((Ops.CAST, Ops.BITCAST), name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({x.dtype})"),
-  (UPat(Ops.SPECIAL, src=(UPat(Ops.CONST),), name="x"), lambda x: f"UOp.special({x.src[0].arg}, {repr(x.arg)}, dtype={x.dtype})"),
-  (UPat(Ops.BUFFER, src=(UPat(Ops.UNIQUE, name="u"), UPat(Ops.DEVICE, name="d")), name="x"), lambda x,u,d:
-    f"UOp.new_buffer({repr(d.arg)}, {x.arg}, {x.dtype}, {u.arg})"),
-  (UPat(Ops.COPY, src=(UPat(name="x"), UPat(Ops.DEVICE, name="d"))), lambda ctx,x,d: f"{ctx[x]}.copy_to_device({repr(d.arg)})"),
-  (UPat(Ops.CUSTOM_FUNCTION, name="x"), lambda ctx,x: f"UOp(Ops.CUSTOM_FUNCTION, {x.dtype}, src={srcs(ctx, x.src)}, arg={x.arg!r})"),
-  (UPat(Ops.REDUCE, name="r"), lambda ctx,r: f"{ctx[r.src[0]]}._rop({r.arg[0]}, {r.arg[1]})" if len(r.arg[1]) else None),
-  # NOTE: range has srcs sometimes after control flow
-  (UPat(Ops.RANGE, src=(UPat(Ops.CONST, name="c"),), allow_any_len=True, name="x"), lambda ctx,x,c:
-    "UOp.range("+', '.join([str(c.arg)] + [repr(y) for y in x.arg])+
-      (f', src={srcs(ctx, x.src[1:])}' if len(x.src) > 1 else '')+(', dtype='+str(x.dtype) if x.dtype is not dtypes.weakint else '')+")"),
-  # TODO: index shouldn't mismatch dtype
-  (UPat(Ops.INDEX, src=(UPat(), UPat()), allow_any_len=True, name="x"), lambda ctx,x:
-   f"{ctx[x.src[0]]}.index({ctx[x.src[1]]}, "+''.join([f"{ctx[xx]}, " for xx in x.src[2:]])+
-    (f"dtype={x.dtype})" if x.src[0].dtype != x.dtype else "ptr=True)") if x.src[0].dtype.base != x.dtype else None),
-  # TODO: movement ops simplify stuff, this can break SPEC=2
-  #(UPat(GroupOp.Movement, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({render_marg(ctx,x)})"),
-  # NOTE: CMPNE doesn't work cause there's no __rne__
-  # NOTE: only match CONSTs without UNIQUE (len(src)==1), unique_const needs explicit rendering
-  (UPat(set(syms.keys())-{Ops.SUB, Ops.CMPNE}, src=(UPat(Ops.CONST, src=(UPat(Ops.DEVICE),), name="y"), UPat(name="z")), name="x"),
-   lambda ctx,x,y,z: strip_binary_parens(x, str(y.arg), ctx[z], lambda a,b: f"({a}{syms[x.op]}{b})")),
-  # NOTE: sub doesn't work cause it's written as add/mul
-  (UPat(set(syms.keys())-{Ops.SUB}, src=(UPat(name="y"), UPat(Ops.CONST, src=(UPat(Ops.DEVICE),), name="z")), name="x"), lambda ctx,x,y,z:
-    strip_binary_parens(x, ctx[y], str(z.arg), lambda a,b: f"({a}{syms[x.op]}{b})")),
-  (UPat(set(syms.keys())-{Ops.SUB}, name="x"), lambda ctx,x:
-    strip_binary_parens(x, ctx[x.src[0]], ctx[x.src[1]], lambda a,b: f"({a}{syms[x.op]}{b})")),
-  (UPat(sugar, src=(), name="x"), lambda x: f"UOp.{x.op.name.lower()}("+', '.join(([f'arg={repr(x.arg)}'] if x.arg is not None else []))+")"),
-  (UPat(sugar, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}("+', '.join([ctx[y] for y in x.src[1:]] + \
-    ([f'arg={repr(x.arg)}'] if x.arg is not None else []))+")"),
-])
-
-# NOTE: you can remove pm_pyrender_extra and it'll still be correct
-pm_pyrender = pm_pyrender_extra+PatternMatcher([
-  (UPat(GroupOp.All, name="u"), lambda ctx,u: f"UOp({u.op}, {u.dtype}, {srcs(ctx,u.src)}"+(f", {repr(u.arg)})" if u.arg is not None else ")")),
-])
-
-def _render_with_splits(lst:list[UOp], pm:PatternMatcher, to_render:set[UOp], split_depth:int=100) -> dict[str, str]:
-  r: dict[UOp, str] = {}
-  ret: dict[str, str] = {}
-  depth: dict[UOp, int] = {}
-  for i,u in enumerate(lst):
-    # limit inline depth to avoid "too many nested parentheses" in Python parser
-    op_depth = 1 + max([depth.get(s, 0) for s in u.src], default=0)
-    if op_depth > split_depth: to_render.add(u)
-    depth[u] = 0 if u in to_render else op_depth
-    ren = cast(str, pm.rewrite(u, ctx=r))
-    assert isinstance(ren, str)
-    if u.tag is not None: ren += f".rtag({repr(u.tag)})"
-    if u not in to_render: r[u] = ren
-    else:
-      r[u] = f"c{i}" if u is not lst[-1] else "ast"
-      ret[r[u]] = ren
-  return ret
-
-def pyrender(ast:UOp) -> str:
-  lst = list(ast.toposort())
-
-  cmap = consumer_map_from_toposort(lst)
-  not_rendered = {Ops.CONST, Ops.VCONST, Ops.DEVICE}
-  always_rendered = {Ops.PARAM, Ops.LOAD, Ops.SPECIAL, Ops.RANGE, Ops.CONTIGUOUS, Ops.STACK,
-                     Ops.BUFFER, Ops.COPY, Ops.CALL, Ops.FUNCTION, Ops.WHERE, Ops.END}
-
-  to_render: set[UOp] = {ast}
-  for u in lst:
-    if u.op in {Ops.SINK}:
-      for s in u.src: to_render.add(s)
-    if u.op is Ops.STORE: to_render.add(u.src[1])
-    if u.op is Ops.REDUCE: to_render.add(u.src[0])
-    if u.op in {Ops.CALL, Ops.FUNCTION}: raise NotImplementedError("call can't be pyrendered")
-    if u.op in not_rendered: continue
-    # checking the consumers is not enough, you have to make sure it's not used twice by the one consumer
-    if len(cmap[u]) == 1 and len([x for x in list(cmap[u].keys())[0].src if x is u]) == 1 and u.op not in always_rendered: continue
-    to_render.add(u)
-
-  ret = _render_with_splits(lst, pm_pyrender, to_render)
-  return '\n'.join([f"{k} = {strip_parens(v)}" for k,v in ret.items()])
-
 # *** what was symbolic.py ***
 
 sint = int|UOp

tinygrad/uop/render.py

@@ -0,0 +1,166 @@
+from typing import cast
+from tinygrad.dtype import dtypes
+from tinygrad.uop import Ops, GroupOp
+from tinygrad.uop.ops import UOp, PatternMatcher, UPat, multirange_str, range_str, consumer_map_from_toposort
+from tinygrad.helpers import strip_parens, all_same
+
+def pretty_print(x:UOp, cache=None, d=0)->str:
+  def dfs(x:UOp, cache:dict):
+    for s in x.src:
+      cache.setdefault(s, [len(cache), 0, False])[1] += 1
+      if cache[s][1] == 1: dfs(s, cache)
+  if cache is None: dfs(x, cache:={})
+  if (cx:=cache.setdefault(x, [0,0,False]))[2]: return f"{' '*d}x{cx[0]}"
+  cx[2], srcs = True, (''.join(f'\n{pretty_print(s, cache, d+2)},' for s in x.src))
+  return f"{' '*d}{f'x{cx[0]}:=' * (cx[1]>1)}{type(x).__name__}({x.op}, {x.dtype}, arg={x.argstr()}{x.tagstr()}, src=({srcs}))"
+
+# ***** uop helpers *****
+
+def print_uops(uops:list[UOp]):
+  uops_index = {u:i for i,u in enumerate(uops)}
+  for i,u in enumerate(uops):
+    formatted_srcs = [(uops_index[x] if x.op is not Ops.CONST else f"{x.arg}") if x in uops else "--" for x in u.src]
+    print(f"{i:4d} {str(u.op):20s}: {multirange_str(u.ranges, color=True, pad=10)} {str(u.dtype):40s} " f"{str(formatted_srcs):32s} {u.arg}")
+
+# for debug
+syms = { Ops.ADD: "+", Ops.SUB: "-", Ops.FLOORDIV: "//", Ops.FLOORMOD: "%", Ops.SHL: "<<", Ops.SHR: ">>",
+         Ops.MUL: "*", Ops.CMPLT: "<", Ops.CMPNE: "!=", Ops.AND: "&", Ops.OR: "|", Ops.XOR: "^"}
+# comparison operators are not in here because they are chained in python, not left-associative
+precedence = {Ops.MUL:1, Ops.FLOORDIV:1, Ops.FLOORMOD:1, Ops.ADD:2, Ops.SUB:2, Ops.SHL:3, Ops.SHR:3, Ops.AND:4, Ops.XOR:5, Ops.OR:6}
+def strip_binary_parens(x:UOp, left:str, right:str, code_for_op) -> str:
+  if x.op not in precedence: return code_for_op(left, right)
+  return code_for_op(strip_parens(left) if precedence.get(x.src[0].op,99)<=precedence[x.op] else left, strip_parens(right) if
+    precedence.get(x.src[1].op,99)<precedence[x.op] else right)
+
+renderer = PatternMatcher([
+  (UPat((Ops.DEFINE_VAR,), name="x"), lambda x: x.expr),
+  (UPat(Ops.PARAM, src=(UPat(), UPat(), UPat(), UPat(), UPat(Ops.NOOP, name="x"))), lambda x: x.arg),
+  (UPat((Ops.SPECIAL), name="x"), lambda x: x.arg),
+  (UPat(Ops.RANGE, name="x"), lambda x: f"r{range_str(x)}"),
+  (UPat((Ops.CONST, Ops.VCONST), name="x"), lambda x: str(x.arg)),
+  (UPat(Ops.UNROLL, name="x"), lambda ctx,x,u: f"UNROLL({ctx[x.src[0]]}, {u.arg})"),
+  (UPat(Ops.CAST, name="x"), lambda ctx,x: f"({str(x.dtype)[7:]})({ctx[x.src[0]]})"),
+  (UPat(Ops.BIND, name="x"), lambda ctx,x: ctx[x.src[0]]),
+  (UPat(Ops.NEG, name="x"), lambda ctx,x: f"(-{ctx[x.src[0]]})"),
+  (UPat(Ops.RECIPROCAL, name="x"), lambda ctx,x: f"(1/{ctx[x.src[0]]})"),
+  (UPat(Ops.MAX, name="x"), lambda ctx,x: f"max({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
+  (UPat(Ops.MULACC, name="x"), lambda ctx,x: f"({ctx[x.src[0]]}*{ctx[x.src[1]]}+{ctx[x.src[2]]})"),
+  (UPat(Ops.WHERE, name="x"), lambda ctx,x: f"({ctx[x.src[1]]} if {ctx[x.src[0]]} else {ctx[x.src[2]]})"),
+  (UPat(Ops.IDIV, name="x"), lambda ctx,x: f"cdiv({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
+  (UPat(Ops.MOD, name="x"), lambda ctx,x: f"cmod({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
+  (UPat(set(syms.keys()), name="x"), lambda ctx,x: strip_binary_parens(x, ctx[x.src[0]], ctx[x.src[1]], lambda a,b: f"({a}{syms[x.op]}{b})")),
+  (UPat((Ops.INDEX, Ops.BUFFERIZE), name="x"), lambda x, ctx: ''.join([f"[{strip_parens(ctx[y])}]" for y in x.src[1:]])),
+  (UPat(Ops.STACK, name="x"),
+   lambda ctx,x: f"{{{','.join([ctx[y] for y in x.src])}}}" if not x.src or not all_same(x.src) else f"{{{ctx[x.src[0]]}, ...}}"),
+  (UPat(GroupOp.All, name="x"), lambda x: str(x)),
+])
+
+renderer_infer = PatternMatcher([
+  (UPat(Ops.MOD, name="x"), lambda ctx,x: f"cmod({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
+  (UPat(Ops.IDIV, name="x"), lambda ctx,x: f"cdiv({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
+  (UPat(Ops.FLOORMOD, name="x"), lambda ctx,x: f"floormod({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
+  (UPat(Ops.FLOORDIV, name="x"), lambda ctx,x: f"floordiv({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
+  (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"bitcast({ctx[x.src[0]]}, {x.src[0].dtype!r}, {x.dtype!r})"),
+]) + renderer
+
+# *** pyrender ***
+
+def srcs(ctx, src): return f"({ctx[src[0]]},)" if len(src) == 1 else f"({', '.join([ctx[x] for x in src])})"
+def render_marg(ctx,x:UOp):
+  if x.op is Ops.PERMUTE: return str(x.marg)
+  if x.op is Ops.FLIP: return str(tuple([i for i,x in enumerate(x.marg) if x]))
+  pieces = []
+  if x.op in {Ops.RESHAPE, Ops.EXPAND}:
+    pieces = [f"{ctx[a] if isinstance(a, UOp) else str(a)}" for a in x.marg]
+  if x.op in {Ops.PAD, Ops.SHRINK}:
+    pieces = [f"({ctx[a[0]] if isinstance(a[0], UOp) else str(a[0])}, {ctx[a[1]] if isinstance(a[1], UOp) else str(a[1])})" for a in x.marg]
+  return f"({','.join(pieces)})" if len(pieces) != 1 else f"({pieces[0]},)"
+
+sugar = {Ops.SINK, Ops.END, Ops.STORE, Ops.LOAD, Ops.UNIQUE, Ops.SQRT, Ops.INDEX, Ops.REDUCE, Ops.AFTER, Ops.THREEFRY,
+         Ops.WHERE, Ops.RECIPROCAL, Ops.EXP2, Ops.LOG2, Ops.SIN, Ops.CONTIGUOUS, Ops.BARRIER, Ops.DETACH}
+pm_pyrender_extra = PatternMatcher([
+  (UPat(Ops.CONST, src=(UPat(Ops.UNIQUE, name="u"), UPat(Ops.DEVICE, name="d")), name="x"),
+   lambda x,u,d: f"UOp.unique_const({x.arg}, dtype={x.dtype}, device={repr(d.arg)}, unique={u.arg})"),
+  (UPat(Ops.CONST, src=(UPat(Ops.DEVICE, name="d"),), name="x"), lambda x,d: f"UOp.const({x.dtype}, {x.arg}, device={repr(d.arg)})"),
+  (UPat(Ops.CONST, src=(), name="x"), lambda x: f"UOp.const({x.dtype}, {x.arg})"),
+  (UPat(Ops.DEFINE_VAR, src=(), name="x"), lambda x:
+    f"UOp.variable(\"{x.arg[0]}\", {x.arg[1]}, {x.arg[2]}{', dtype='+str(x.dtype) if x.dtype is not dtypes.weakint else ''})"),
+  (UPat((Ops.CAST, Ops.BITCAST), name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({x.dtype})"),
+  (UPat(Ops.SPECIAL, src=(UPat(Ops.CONST),), name="x"), lambda x: f"UOp.special({x.src[0].arg}, {repr(x.arg)}, dtype={x.dtype})"),
+  (UPat(Ops.BUFFER, src=(UPat(Ops.UNIQUE, name="u"), UPat(Ops.DEVICE, name="d")), name="x"), lambda x,u,d:
+    f"UOp.new_buffer({repr(d.arg)}, {x.arg}, {x.dtype}, {u.arg})"),
+  (UPat(Ops.COPY, src=(UPat(name="x"), UPat(Ops.DEVICE, name="d"))), lambda ctx,x,d: f"{ctx[x]}.copy_to_device({repr(d.arg)})"),
+  (UPat(Ops.CUSTOM_FUNCTION, name="x"), lambda ctx,x: f"UOp(Ops.CUSTOM_FUNCTION, {x.dtype}, src={srcs(ctx, x.src)}, arg={x.arg!r})"),
+  (UPat(Ops.REDUCE, name="r"), lambda ctx,r: f"{ctx[r.src[0]]}._rop({r.arg[0]}, {r.arg[1]})" if len(r.arg[1]) else None),
+  # NOTE: range has srcs sometimes after control flow
+  (UPat(Ops.RANGE, src=(UPat(Ops.CONST, name="c"),), allow_any_len=True, name="x"), lambda ctx,x,c:
+    "UOp.range("+', '.join([str(c.arg)] + [repr(y) for y in x.arg])+
+      (f', src={srcs(ctx, x.src[1:])}' if len(x.src) > 1 else '')+(', dtype='+str(x.dtype) if x.dtype is not dtypes.weakint else '')+")"),
+  # TODO: index shouldn't mismatch dtype
+  (UPat(Ops.INDEX, src=(UPat(), UPat()), allow_any_len=True, name="x"), lambda ctx,x:
+   f"{ctx[x.src[0]]}.index({ctx[x.src[1]]}, "+''.join([f"{ctx[xx]}, " for xx in x.src[2:]])+
+    (f"dtype={x.dtype})" if x.src[0].dtype != x.dtype else "ptr=True)") if x.src[0].dtype.base != x.dtype else None),
+  # TODO: movement ops simplify stuff, this can break SPEC=2
+  #(UPat(GroupOp.Movement, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({render_marg(ctx,x)})"),
+  # NOTE: CMPNE doesn't work cause there's no __rne__
+  # explicit trunc ops: `//` and `%` parse as FLOORDIV/FLOORMOD, so render IDIV/MOD via their named methods
+  (UPat(Ops.IDIV, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.idiv({ctx[x.src[1]]})"),
+  (UPat(Ops.MOD, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.alu(Ops.MOD, {ctx[x.src[1]]})"),
+  # NOTE: only match CONSTs without UNIQUE (len(src)==1), unique_const needs explicit rendering
+  (UPat(set(syms.keys())-{Ops.SUB, Ops.CMPNE, Ops.IDIV, Ops.MOD}, src=(UPat(Ops.CONST, src=(UPat(Ops.DEVICE),), name="y"), UPat(name="z")), name="x"),
+   lambda ctx,x,y,z: strip_binary_parens(x, str(y.arg), ctx[z], lambda a,b: f"({a}{syms[x.op]}{b})")),
+  # NOTE: sub doesn't work cause it's written as add/mul
+  (UPat(set(syms.keys())-{Ops.SUB, Ops.IDIV, Ops.MOD}, src=(UPat(name="y"), UPat(Ops.CONST, src=(UPat(Ops.DEVICE),), name="z")), name="x"),
+    lambda ctx,x,y,z: strip_binary_parens(x, ctx[y], str(z.arg), lambda a,b: f"({a}{syms[x.op]}{b})")),
+  (UPat(set(syms.keys())-{Ops.SUB, Ops.IDIV, Ops.MOD}, name="x"), lambda ctx,x:
+    strip_binary_parens(x, ctx[x.src[0]], ctx[x.src[1]], lambda a,b: f"({a}{syms[x.op]}{b})")),
+  (UPat(sugar, src=(), name="x"), lambda x: f"UOp.{x.op.name.lower()}("+', '.join(([f'arg={repr(x.arg)}'] if x.arg is not None else []))+")"),
+  (UPat(sugar, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}("+', '.join([ctx[y] for y in x.src[1:]] + \
+    ([f'arg={repr(x.arg)}'] if x.arg is not None else []))+")"),
+])
+
+# NOTE: you can remove pm_pyrender_extra and it'll still be correct
+pm_pyrender = pm_pyrender_extra+PatternMatcher([
+  (UPat(GroupOp.All, name="u"), lambda ctx,u: f"UOp({u.op}, {u.dtype}, {srcs(ctx,u.src)}"+(f", {repr(u.arg)})" if u.arg is not None else ")")),
+])
+
+def _render_with_splits(lst:list[UOp], pm:PatternMatcher, to_render:set[UOp], split_depth:int=100) -> dict[str, str]:
+  r: dict[UOp, str] = {}
+  ret: dict[str, str] = {}
+  depth: dict[UOp, int] = {}
+  for i,u in enumerate(lst):
+    # limit inline depth to avoid "too many nested parentheses" in Python parser
+    op_depth = 1 + max([depth.get(s, 0) for s in u.src], default=0)
+    if op_depth > split_depth: to_render.add(u)
+    depth[u] = 0 if u in to_render else op_depth
+    ren = cast(str, pm.rewrite(u, ctx=r))
+    assert isinstance(ren, str)
+    if u.tag is not None: ren += f".rtag({repr(u.tag)})"
+    if u not in to_render: r[u] = ren
+    else:
+      r[u] = f"c{i}" if u is not lst[-1] else "ast"
+      ret[r[u]] = ren
+  return ret
+
+def pyrender(ast:UOp) -> str:
+  lst = list(ast.toposort())
+
+  cmap = consumer_map_from_toposort(lst)
+  not_rendered = {Ops.CONST, Ops.VCONST, Ops.DEVICE}
+  always_rendered = {Ops.PARAM, Ops.LOAD, Ops.SPECIAL, Ops.RANGE, Ops.CONTIGUOUS, Ops.STACK,
+                     Ops.BUFFER, Ops.COPY, Ops.CALL, Ops.FUNCTION, Ops.WHERE, Ops.END}
+
+  to_render: set[UOp] = {ast}
+  for u in lst:
+    if u.op in {Ops.SINK}:
+      for s in u.src: to_render.add(s)
+    if u.op is Ops.STORE: to_render.add(u.src[1])
+    if u.op is Ops.REDUCE: to_render.add(u.src[0])
+    if u.op in {Ops.CALL, Ops.FUNCTION}: raise NotImplementedError("call can't be pyrendered")
+    if u.op in not_rendered: continue
+    # checking the consumers is not enough, you have to make sure it's not used twice by the one consumer
+    if len(cmap[u]) == 1 and len([x for x in list(cmap[u].keys())[0].src if x is u]) == 1 and u.op not in always_rendered: continue
+    to_render.add(u)
+
+  ret = _render_with_splits(lst, pm_pyrender, to_render)
+  return '\n'.join([f"{k} = {strip_parens(v)}" for k,v in ret.items()])

tinygrad/uop/spec.py

@@ -1,6 +1,7 @@
 import math
 from typing import cast, Any
-from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, print_uops, AxisType, KernelInfo, pyrender
+from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, AxisType, KernelInfo
+from tinygrad.uop.render import print_uops, pyrender
 from tinygrad.dtype import DType, ImageDType, dtypes, PtrDType, AddrSpace, Invalid, ConstFloat
 from tinygrad.helpers import DEBUG, Context, prod, SPEC, Metadata, panic, CHECK_OOB
 
@@ -44,7 +45,7 @@ shared_spec = PatternMatcher([
   (UPat((Ops.CMPLT, Ops.CMPNE, Ops.CMPEQ), dtype=dtypes.bool, src=(UPat.var("x"), UPat.var("y"))), lambda x,y: x.dtype.base == y.dtype.base),
   # and SHL/SHR, the shift distance can be an int
   (UPat((Ops.SHL, Ops.SHR), src=(UPat.var("x"), UPat.var("y")), name="a"), lambda a,x,y: a.dtype == x.dtype and y.dtype in (x.dtype, dtypes.uint)),
-  (UPat((Ops.IDIV, Ops.MOD), name="x"), lambda x: None if dtypes.is_int(x.dtype) else False),
+  (UPat((Ops.IDIV, Ops.MOD, Ops.FLOORDIV, Ops.FLOORMOD), name="x"), lambda x: None if dtypes.is_int(x.dtype) else False),
   (UPat(GroupOp.ALU, name="x"), lambda x: all(x.dtype.base == y.dtype.base for y in x.src)),
 
   # CAST
@@ -75,7 +76,7 @@ movement_ops = PatternMatcher([
 
   # inputs to movement ops
   (UPat((Ops.STACK, Ops.VCONST), dtype=dtypes.weakint), lambda: True),
-  (UPat({Ops.ADD, Ops.MUL, Ops.IDIV}, dtype=dtypes.weakint), lambda: True),
+  (UPat({Ops.ADD, Ops.MUL, Ops.IDIV, Ops.FLOORDIV}, dtype=dtypes.weakint), lambda: True),
 
   # AFTER on Movement Op, INDEX, BUFFER, COPY, or BITCAST
   (UPat(Ops.AFTER, src=(UPat(GroupOp.Movement.union({Ops.INDEX, Ops.MULTI, Ops.CONTIGUOUS, Ops.BUFFER, Ops.BITCAST, Ops.COPY})),),
@@ -173,10 +174,12 @@ shared_codegen_spec = PatternMatcher([
   (UPat(Ops.STACK, name="x"), lambda x: len(x.src)>1 and len(x.src) == x.dtype.vcount and all(x.dtype == y.dtype.vec(len(x.src)) for y in x.src)),
   (UPat(Ops.GEP, src=(UPat.var("src"),), name="gep"), lambda gep,src: gep.dtype == src.dtype.scalar()),
 
-  # LOAD(idx) / STORE(idx, val)
-  (UPat().index(UPat()).or_casted().load(), lambda: True),
-  (UPat().index(UPat(), UPat(dtype=dtypes.bool)).or_casted().load(), lambda: True),  # gated load (alt added in program_spec)
-  (UPat(Ops.INDEX).or_casted().store(UPat()), lambda: True),
+  # LOAD(idx) / STORE(idx, val) with gates on the LOAD/STORE
+  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx"))).or_casted().load(), validate_index),
+  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx"))).or_casted().load(UPat.var("alt"), UPat.var("gate", dtype=dtypes.bool), name="load"),
+   lambda buf,idx,gate,alt,load: validate_index(buf, idx, gate) if alt.dtype == load.dtype else False),
+  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx"))).or_casted().store(UPat()), validate_index),
+  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx"))).or_casted().store(UPat(), UPat.var("gate", dtype=dtypes.bool)), validate_index),
 
   # CUSTOM (inline and non inline)
   (UPat((Ops.CUSTOMI, Ops.CUSTOM)), lambda: True),
@@ -184,9 +187,8 @@ shared_codegen_spec = PatternMatcher([
   # assembly instruction
   (UPat(Ops.INS), lambda: True),
 
-  # INDEX (2-arg and 3-arg with bool gate)
-  (UPat(GroupOp.Defines|{Ops.AFTER}, name="buf").index(UPat.var("idx")), validate_index),
-  (UPat(Ops.INDEX, src=(UPat(GroupOp.Defines|{Ops.AFTER}, name="buf"), UPat.var("idx"), UPat.var("gate", dtype=dtypes.bool))), validate_index),
+  # INDEX is just address calculation. OOB validation is on LOAD/STORE where the gate is available.
+  (UPat(GroupOp.Defines|{Ops.AFTER}).index(UPat()), lambda: True),
 
   # SPECIAL
   (UPat(Ops.SPECIAL, src=(UPat.var("x", (dtypes.weakint, dtypes.int32)),), name="s"), lambda s,x: s.dtype == x.dtype and isinstance(s.arg, str)),
@@ -235,9 +237,6 @@ tensor_spec = PatternMatcher([
 # ***** UOp spec in linearized programs *****
 
 program_spec = PatternMatcher([
-  # LOAD (idx, alt_value), LOAD can have an alt value, but only if the index has a gate
-  (UPat().index(UPat(), UPat(dtype=dtypes.bool)).or_casted().load(UPat()), lambda: True),
-
   # END closes ranges
   (UPat(Ops.END, src=(UPat(), UPat(Ops.RANGE)), dtype=dtypes.void), lambda: True),
 
@@ -281,7 +280,7 @@ full_spec = PatternMatcher([
   # where on index in rhs position is fine
   (UPat(Ops.WHERE, dtype=dtypes.weakint, src=(UPat(dtype=dtypes.bool), UPat(), UPat(dtype=dtypes.weakint))), lambda: True),
   # allow index dtype on a restricted set of UOps
-  (UPat((Ops.ADD, Ops.MUL, Ops.MOD, Ops.IDIV, Ops.MAX,
+  (UPat((Ops.ADD, Ops.MUL, Ops.MOD, Ops.IDIV, Ops.FLOORDIV, Ops.FLOORMOD, Ops.MAX,
          Ops.SPECIAL, Ops.CAST, Ops.RANGE, Ops.VCONST, Ops.STACK), dtype=dtypes.weakint), lambda: True),
 
   # while BIND is being casted
@@ -292,7 +291,9 @@ full_spec = PatternMatcher([
 
   # temp VECTORIZE/INDEX during rewrite have the wrong dtype
   (UPat(Ops.STACK), lambda: True),
-  (UPat(Ops.INDEX), lambda: True),
+
+  # no more bool in index
+  (UPat(Ops.INDEX, name="idx"), lambda idx: not any([dtypes.is_bool(x.dtype) for x in idx.src[1:]])),
 
   # all loads/stores
   (UPat((Ops.LOAD, Ops.STORE)), lambda: True),