codegen

allow kernel_spec in tensor for custom kernels
allow after on param
2026-08-19 07:58:27 +00:00 · 2026-02-03 18:09:15 +08:00 · 2026-02-03 18:01:07 +08:00 · 2026-02-03 17:41:54 +08:00 · 2026-02-03 16:28:53 +08:00 · 2026-02-03 15:09:38 +08:00
121 changed files with 942 additions and 1489 deletions
@@ -16,6 +16,36 @@ on:
  workflow_dispatch:

 jobs:
+  # the goal of this test is to replicate a normal person on a laptop running the test
+  # no process replay, no benchmarks, no CI, just a normal laptop person
+  # the 3 minute timeout should not be raised
+  testmacpytest:
+    name: Mac pytest
+    runs-on: [self-hosted, macOS]
+    timeout-minutes: 3
+    defaults:
+      run:
+        shell: bash -e -o pipefail {0}
+    if: github.repository_owner == 'tinygrad'
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    # brew install uv
+    - name: setup python environment
+      run: |
+        rm -rf /tmp/tinygrad_pytest_ci
+        uv venv /tmp/tinygrad_pytest_ci
+        source /tmp/tinygrad_pytest_ci/bin/activate
+        uv pip install .[testing]
+    - name: setup staging db
+      run: |
+        echo "CACHEDB=/tmp/pytest-db-ci.db" >> $GITHUB_ENV
+        rm -f /tmp/pytest-db-ci*
+    - name: Run pytest -nauto
+      run: |
+        source /tmp/tinygrad_pytest_ci/bin/activate
+        pytest -nauto --durations=20
+
  testmacbenchmark:
    name: Mac Benchmark
    env:
@@ -266,17 +266,19 @@ jobs:
      run: python -c "from tinygrad import Device; assert Device.DEFAULT == 'CPU', Device.DEFAULT"
    - name: Run unit tests
      run: |
-        CPU=1 python test/unit/test_device.py TestRunAsModule.test_module_runs
-        CPU=1 python -m pytest -n=auto test/unit/ --durations=20 --deselect=test/unit/test_device.py::TestRunAsModule::test_module_runs
+        CPU=1 python test/null/test_device.py TestRunAsModule.test_module_runs
+        CPU=1 python -m pytest -n=auto test/unit/ --durations=20
+    - name: Run NULL backend tests
+      run: NULL=1 python -m pytest -n=auto test/null/ --durations=20
    - name: Run targetted tests on NULL backend
-      run: NULL=1 python3 -m unittest test.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step test/device/test_null.py
+      run: NULL=1 python3 -m unittest test.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step
    # TODO: too slow
    # - name: Run SDXL on NULL backend
    #   run: NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
    - name: Run Clip tests for SD MLPerf on NULL backend
      run: NULL=1 python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
    - name: Run AMD emulated BERT training on NULL backend
-      run: EMULATE=AMD_RDNA4 NULL=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+      run: EMULATE=AMD_RDNA4 NULL=1 NULL_ALLOW_COPYOUT=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
    # TODO: support fake weights
    #- name: Run LLaMA 7B on 4 fake devices
    #  run: NULL=1 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing
@@ -314,7 +316,7 @@ jobs:
        deps: testing_unit
        python-version: '3.14'
    - name: Test SPEC=2
-      run: SPEC=2 pytest --maxfail=10 -n auto --durations=30 --ignore=test/models --ignore test/test_custom_kernel.py --ignore test/unit/test_hashing.py --ignore test/unit/test_autogen.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}
+      run: SPEC=2 pytest --maxfail=10 -n auto --durations=30 --ignore=test/models --ignore=test/null --ignore test/test_custom_kernel.py --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}

  fuzzing:
    name: Fuzzing
@@ -465,11 +467,11 @@ jobs:
      - name: Test MLPerf stuff
        run: CL=1 python -m pytest -n=auto test/external/external_test_optim.py test/external/external_test_losses.py test/external/external_test_metrics.py test/external/external_test_datasets.py --durations=20
      - name: NULL=1 beautiful_mnist_multigpu
-        run: NULL=1 python examples/beautiful_mnist_multigpu.py
+        run: NULL=1 NULL_ALLOW_COPYOUT=1 python examples/beautiful_mnist_multigpu.py
      - name: Test Bert training
-        run: NULL=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+        run: NULL=1 NULL_ALLOW_COPYOUT=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
      - name: Test llama 3 training
-        run: NULL=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=1 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
+        run: NULL=1 NULL_ALLOW_COPYOUT=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=1 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -609,7 +611,7 @@ jobs:
        WEBGPU=1 DEBUG=4 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
    - name: Run selected webgpu tests
      run: |
-          WEBGPU=1 WEBGPU_BACKEND="WGPUBackendType_Vulkan" python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
+          WEBGPU=1 WEBGPU_BACKEND="WGPUBackendType_Vulkan" python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore=test/null --durations=20
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

@@ -704,6 +706,8 @@ jobs:
      # TODO: run all once emulator is faster
      - name: Run RDNA3 ops tests
        run: SKIP_SLOW_TEST=1 AMD_LLVM=0 pytest -n=auto test/test_ops.py -k "test_sparse_categorical_crossentropy or test_tril or test_nonzero or test_softmax_argmax" --durations 20
+      - name: Run RDNA4 emulator tests
+        run: MOCKGPU_ARCH=rdna4 python -m pytest test/test_tiny.py -v --durations 20

  testnvidia:
    strategy:
@@ -735,7 +739,7 @@ jobs:
          DEBUG=5 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
      - name: Run pytest (cuda)
        # skip multitensor because it's slow
-        run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore test/test_gc.py --ignore test/test_multitensor.py --durations=20
+        run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore=test/null --ignore test/test_gc.py --ignore test/test_multitensor.py --durations=20
      - name: Run TestOps.test_add with PMA
        run: VIZ=-1 PMA=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
      - name: Run process replay tests
@@ -768,7 +772,7 @@ jobs:
          python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CPU','CL'], Device.DEFAULT"
          DEBUG=5 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
      - name: Run pytest (${{ matrix.backend }})
-        run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
+        run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore=test/null --durations=20
      - name: Run TRANSCENDENTAL math
        run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
      - name: Run process replay tests
@@ -795,6 +799,8 @@ jobs:
        llvm: 'true'
    - name: Run unit tests
      run: METAL=1 python -m pytest -n=auto test/unit/ --durations=20
+    - name: Run NULL backend tests
+      run: NULL=1 python -m pytest -n=auto test/null/ --durations=20
    - name: Run ONNX
      run: METAL=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
    - name: Test tensor core ops (fake)
@@ -896,7 +902,7 @@ jobs:
          python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU','LVP':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
          DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
      - name: Run pytest (${{ matrix.backend }})
-        run: python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
+        run: python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore=test/null --durations=20
      - name: Run process replay tests
        uses: ./.github/actions/process-replay
      - name: Run macOS-specific unit test
@@ -929,7 +935,11 @@ jobs:
      - name: Run unit tests
        if: matrix.backend=='llvm'
        # test_newton_schulz hits RecursionError
-        run: python -m pytest -n=auto test/unit/ --ignore=test/unit/test_disk_tensor.py --ignore=test/unit/test_elf.py --ignore=test/unit/test_tar.py --ignore=test/unit/test_linalg.py --durations=20
+        run: python -m pytest -n=auto test/unit/ --ignore=test/unit/test_disk_tensor.py --ignore=test/unit/test_tar.py --ignore=test/unit/test_linalg.py --durations=20
+      - name: Run NULL backend tests
+        if: matrix.backend=='llvm'
+        shell: bash
+        run: CPU=0 CPU_LLVM=0 NULL=1 python -m pytest -n=auto test/null/ --ignore=test/null/test_elf.py --durations=20
      - name: Run pytest (${{ matrix.backend }})
        shell: bash
        run: |
@@ -955,10 +965,10 @@ jobs:
          key: compile-${{ matrix.backend }}
          deps: testing_unit
          mesa: ${{ (matrix.backend == 'ir3' || matrix.backend == 'nak') && 'true' }}
-          python-version: '3.14'
+          python-version: '3.12'
      - name: Set env
        shell: bash
-        run: printf "NULL=1\n${{ matrix.backend == 'ir3' && 'NULL_IR3=1' || matrix.backend == 'nak' && 'NULL_NAK=1' }}" >> $GITHUB_ENV
+        run: printf "NULL=1\nNULL_ALLOW_COPYOUT=1\n${{ matrix.backend == 'ir3' && 'NULL_IR3=1' || matrix.backend == 'nak' && 'NULL_NAK=1' }}" >> $GITHUB_ENV
      - name: Run test_ops
        shell: bash
        run: |
@@ -28,7 +28,7 @@ repos:
        pass_filenames: false
      - id: tests
        name: comprehensive test suite
-        entry: env OMP_NUM_THREADS=1 SKIP_SLOW_TEST=1 PYTHONPATH="." python3 -m pytest -n=6 test/test_ops.py test/test_schedule.py test/unit/test_assign.py test/test_tensor.py test/test_jit.py test/unit/test_schedule_cache.py test/unit/test_pattern_matcher.py test/unit/test_uop_symbolic.py test/unit/test_helpers.py
+        entry: env OMP_NUM_THREADS=1 SKIP_SLOW_TEST=1 PYTHONPATH="." python3 -m pytest -n=6 test/test_ops.py test/test_schedule.py test/unit/test_assign.py test/test_tensor.py test/test_jit.py test/unit/test_schedule_cache.py test/null/test_pattern_matcher.py test/null/test_uop_symbolic.py test/unit/test_helpers.py
        language: system
        always_run: true
        pass_filenames: false
@@ -72,7 +72,7 @@ vliw_prepare = PatternMatcher([
  # cast is fake
  (UPat(Ops.CAST, name="c"), lambda c: c.src[0]),
  # rewrites to hardcode the addresses in memory
-  (UPat(Ops.DEFINE_GLOBAL, name="dg"), lambda dg: UOp.const(dtypes.uint, global_addrs[dg.arg])),
+  (UPat(Ops.PARAM, name="dg"), lambda dg: UOp.const(dtypes.uint, global_addrs[dg.arg])),
  # INDEX is just plus
  (UPat(Ops.INDEX, name="i"), lambda i: i.src[0]+i.src[1]),
 ])+symbolic
@@ -8,9 +8,11 @@ export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
 export DEBUG=${DEBUG:-2}
 export FLASH_ATTENTION=${FLASH_ATTENTION:-1}
 export ALL2ALL=${ALL2ALL:-1}
+export USE_ATOMICS=${USE_ATOMICS:-1}
+export ASM_GEMM=${ASM_GEMM:-1}

 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=8 BS=8 EVAL_BS=8 GRADIENT_ACC_STEPS=2
+export DP=8 BS=16 EVAL_BS=8 GRADIENT_ACC_STEPS=1
 export GBS=$((BS * GRADIENT_ACC_STEPS))

 export MODEL="llama3"
@@ -13,7 +13,7 @@ export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}

 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=${DP:-8} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
+export DP=${DP:-8} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
 export GBS=$((BS * GRADIENT_ACC_STEPS))

 export MODEL="llama3"
@@ -2,8 +2,9 @@
 export BENCHMARK=5
 export EVAL_BS=0
 export FAKEDATA=1
+export NULL_ALLOW_COPYOUT=1
 export HIP_VISIBLE_DEVICES=""
 export DEV=NULL
 export JITBEAM=0
 export LLAMA_LAYERS=${LLAMA_LAYERS:-"2"}
-time examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh
+time examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh
@@ -49,10 +49,11 @@ from tinygrad.helpers import Context, DEBUG, colored
 from tinygrad.engine.realize import get_runner

 from extra.assembly.amd import decode_inst
-from extra.assembly.amd.autogen.rdna3.str_pcode import PCODE
-from extra.assembly.amd.autogen.rdna3.ins import (SOP1, SOP2, SOPC, SOPK, SOPP, SMEM, VOP1, VOP1_SDST, VOP2, VOP3, VOP3_SDST, VOP3SD, VOP3P, VOPC,
-  DS, FLAT, GLOBAL, SCRATCH, VOPD, SOPPOp, SMEMOp, VOP1Op, VOP2Op, VOP3Op, VOPDOp)
-from extra.assembly.amd.dsl import VCC_LO, EXEC_LO, SCC
+from extra.assembly.amd.autogen.rdna3.str_pcode import PCODE as PCODE_RDNA3
+from extra.assembly.amd.autogen.rdna4.str_pcode import PCODE as PCODE_RDNA4
+from extra.assembly.amd.autogen.rdna3 import ins as ir3
+from extra.assembly.amd.autogen.rdna4 import ins as ir4
+from extra.assembly.amd.dsl import VCC_LO, EXEC_LO, SCC, ttmp
 from extra.assembly.amd.autogen.common import Fmt, OpType
 from extra.assembly.amd.pcode import parse_block, _FUNCS

@@ -79,15 +80,23 @@ def _apply_src_mods(val: UOp, mod_bit: int, abs_bits: int, neg_bits: int, bits:
  if neg_bits & (1 << mod_bit): fv = fv.neg()
  return fv.bitcast(ut).cast(dtypes.uint32) if bits == 16 else fv.bitcast(ut)

-# Map VOPD ops to VOP2 ops for pcode lookup
+# Map VOPD ops to VOP2 ops for pcode lookup (both RDNA3 and RDNA4)
 VOPD_TO_VOP2 = {
-  VOPDOp.V_DUAL_FMAC_F32: VOP2Op.V_FMAC_F32_E32, VOPDOp.V_DUAL_MUL_F32: VOP2Op.V_MUL_F32_E32,
-  VOPDOp.V_DUAL_ADD_F32: VOP2Op.V_ADD_F32_E32, VOPDOp.V_DUAL_SUB_F32: VOP2Op.V_SUB_F32_E32,
-  VOPDOp.V_DUAL_SUBREV_F32: VOP2Op.V_SUBREV_F32_E32, VOPDOp.V_DUAL_MAX_F32: VOP2Op.V_MAX_F32_E32,
-  VOPDOp.V_DUAL_MIN_F32: VOP2Op.V_MIN_F32_E32, VOPDOp.V_DUAL_ADD_NC_U32: VOP2Op.V_ADD_NC_U32_E32,
-  VOPDOp.V_DUAL_LSHLREV_B32: VOP2Op.V_LSHLREV_B32_E32, VOPDOp.V_DUAL_AND_B32: VOP2Op.V_AND_B32_E32,
-  VOPDOp.V_DUAL_MOV_B32: VOP1Op.V_MOV_B32_E32, VOPDOp.V_DUAL_CNDMASK_B32: VOP2Op.V_CNDMASK_B32_E32,
-  VOPDOp.V_DUAL_FMAAK_F32: VOP2Op.V_FMAAK_F32_E32, VOPDOp.V_DUAL_FMAMK_F32: VOP2Op.V_FMAMK_F32_E32,
+  ir3.VOPDOp.V_DUAL_FMAC_F32: ir3.VOP2Op.V_FMAC_F32_E32, ir3.VOPDOp.V_DUAL_MUL_F32: ir3.VOP2Op.V_MUL_F32_E32,
+  ir3.VOPDOp.V_DUAL_ADD_F32: ir3.VOP2Op.V_ADD_F32_E32, ir3.VOPDOp.V_DUAL_SUB_F32: ir3.VOP2Op.V_SUB_F32_E32,
+  ir3.VOPDOp.V_DUAL_SUBREV_F32: ir3.VOP2Op.V_SUBREV_F32_E32, ir3.VOPDOp.V_DUAL_MAX_F32: ir3.VOP2Op.V_MAX_F32_E32,
+  ir3.VOPDOp.V_DUAL_MIN_F32: ir3.VOP2Op.V_MIN_F32_E32, ir3.VOPDOp.V_DUAL_ADD_NC_U32: ir3.VOP2Op.V_ADD_NC_U32_E32,
+  ir3.VOPDOp.V_DUAL_LSHLREV_B32: ir3.VOP2Op.V_LSHLREV_B32_E32, ir3.VOPDOp.V_DUAL_AND_B32: ir3.VOP2Op.V_AND_B32_E32,
+  ir3.VOPDOp.V_DUAL_MOV_B32: ir3.VOP1Op.V_MOV_B32_E32, ir3.VOPDOp.V_DUAL_CNDMASK_B32: ir3.VOP2Op.V_CNDMASK_B32_E32,
+  ir3.VOPDOp.V_DUAL_FMAAK_F32: ir3.VOP2Op.V_FMAAK_F32_E32, ir3.VOPDOp.V_DUAL_FMAMK_F32: ir3.VOP2Op.V_FMAMK_F32_E32,
+  # RDNA4 mappings (same VOP1/VOP2 targets, RDNA4 uses _NUM_ suffix for min/max)
+  ir4.VOPDOp.V_DUAL_FMAC_F32: ir3.VOP2Op.V_FMAC_F32_E32, ir4.VOPDOp.V_DUAL_MUL_F32: ir3.VOP2Op.V_MUL_F32_E32,
+  ir4.VOPDOp.V_DUAL_ADD_F32: ir3.VOP2Op.V_ADD_F32_E32, ir4.VOPDOp.V_DUAL_SUB_F32: ir3.VOP2Op.V_SUB_F32_E32,
+  ir4.VOPDOp.V_DUAL_SUBREV_F32: ir3.VOP2Op.V_SUBREV_F32_E32, ir4.VOPDOp.V_DUAL_MAX_NUM_F32: ir3.VOP2Op.V_MAX_F32_E32,
+  ir4.VOPDOp.V_DUAL_MIN_NUM_F32: ir3.VOP2Op.V_MIN_F32_E32, ir4.VOPDOp.V_DUAL_ADD_NC_U32: ir3.VOP2Op.V_ADD_NC_U32_E32,
+  ir4.VOPDOp.V_DUAL_LSHLREV_B32: ir3.VOP2Op.V_LSHLREV_B32_E32, ir4.VOPDOp.V_DUAL_AND_B32: ir3.VOP2Op.V_AND_B32_E32,
+  ir4.VOPDOp.V_DUAL_MOV_B32: ir3.VOP1Op.V_MOV_B32_E32, ir4.VOPDOp.V_DUAL_CNDMASK_B32: ir3.VOP2Op.V_CNDMASK_B32_E32,
+  ir4.VOPDOp.V_DUAL_FMAAK_F32: ir3.VOP2Op.V_FMAAK_F32_E32, ir4.VOPDOp.V_DUAL_FMAMK_F32: ir3.VOP2Op.V_FMAMK_F32_E32,
 }
 WAVE_SIZE = 32
 # Special registers stored after inline constants (256-259)
@@ -146,11 +155,15 @@ _pcode_fixes = {
  'V_TRIG_PREOP_F64': ("result = 64'F((1201'B(2.0 / PI)[1200 : 0] << shift.u32) & 1201'0x1fffffffffffff)", "result = trig_preop_result(shift)"),
 }

+def _get_pcode_dict(op) -> dict:
+  """Return the PCODE dictionary for the given opcode based on its architecture."""
+  return PCODE_RDNA4 if 'rdna4' in type(op).__module__ else PCODE_RDNA3
+
 # Pcode parser
@functools.cache
 def get_pcode(op) -> str:
  op_name = op.name
-  pcode = PCODE[op]
+  pcode = _get_pcode_dict(op)[op]
  if op_name in _pcode_fixes: pcode = pcode.replace(*_pcode_fixes[op_name])
  if 'V_DIV_SCALE' in op_name:
    dt, exp_lim, ldexp_val = ('f32', '23', '64') if 'F32' in op_name else ('f64', '52', '128')
@@ -174,7 +187,12 @@ def get_pcode(op) -> str:
 def parse_pcode(pcode: str, srcs: dict[str, UOp] | None = None) -> tuple[dict, list[tuple[str, UOp]]]:
  vars: dict = srcs.copy() if srcs else {}
  assigns: list[tuple[str, UOp]] = []
-  lines = [l.strip().rstrip(';') for l in pcode.split('\n') if l.strip() and not l.strip().startswith('//')]
+  raw_lines = [l.strip().rstrip(';') for l in pcode.split('\n') if l.strip() and not l.strip().startswith('//')]
+  # TODO: pcode.py should tokenize full pcode string instead of line-by-line, then this hack can be removed
+  lines: list[str] = []
+  for l in raw_lines:
+    if lines and lines[-1].endswith('&&'): lines[-1] = lines[-1] + ' ' + l
+    else: lines.append(l)
  _, final, _ = parse_block(lines, 0, vars, assigns=assigns)
  sliced = set(d.split('[')[0] for d, _ in assigns if '[' in d)
  for var, val in final.items():
@@ -247,11 +265,11 @@ def _collect_data_slices(assigns: list[tuple[str, UOp]], data_prefix: str, pcode
 class _Ctx:
  """Context for instruction compilation - holds buffers and helpers."""
  __slots__ = ('inst_size', 'dyn_fields', '_axis_id')
-  sgpr = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(SGPR_COUNT), arg=0)
-  vgpr = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(VGPR_SIZE), arg=1)
-  vmem = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(1 << 46), arg=2)
-  lds = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(16384), arg=3)
-  scratch = UOp(Ops.DEFINE_GLOBAL, dtypes.uint8.ptr(1 << 30), arg=4)
+  sgpr = UOp(Ops.PARAM, dtypes.uint32.ptr(SGPR_COUNT), arg=0)
+  vgpr = UOp(Ops.PARAM, dtypes.uint32.ptr(VGPR_SIZE), arg=1)
+  vmem = UOp(Ops.PARAM, dtypes.uint32.ptr(1 << 46), arg=2)
+  lds = UOp(Ops.PARAM, dtypes.uint32.ptr(16384), arg=3)
+  scratch = UOp(Ops.PARAM, dtypes.uint8.ptr(1 << 30), arg=4)

  def __init__(self, inst_size: int):
    self.inst_size, self._axis_id = inst_size, 0
@@ -317,9 +335,9 @@ class _Ctx:
    return base, mask, size

  # Dynamic register access (takes UOp index instead of int)
-  def rsgpr_dyn(self, reg: UOp) -> UOp:
+  def rsgpr_dyn(self, reg: UOp, valid: UOp | None = None) -> UOp:
    """Read SGPR with dynamic register index."""
-    return self.sgpr.index(reg.cast(dtypes.int), ptr=True).load()
+    return self.sgpr.index(reg.cast(dtypes.int), valid, ptr=True).load() if valid is not None else 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. Writes to NULL (124) are discarded."""
@@ -341,15 +359,18 @@ class _Ctx:
    If lane is None, only scalar access is supported (off must be < 256).
    is_f64: True for F64 operations where 64-bit literals go in high 32 bits."""
    is_float_const = (off >= _c(240)) & (off <= _c(248))
-    sgpr_lo = self.rsgpr_dyn(off)
+    is_vgpr = off >= _c(256)
+    is_sgpr = is_vgpr.ne(True)
+    sgpr_lo = self.rsgpr_dyn(off, is_sgpr)

    if lane is not None:
-      is_vgpr, vgpr_reg = off >= _c(256), off - _c(256)
+      vgpr_reg = off - _c(256)
      vgpr_lo = self.rvgpr_dyn(vgpr_reg, lane, is_vgpr)
      vgpr_val = _u64(vgpr_lo, self.rvgpr_dyn(vgpr_reg + _c(1), lane, is_vgpr)) if bits == 64 else vgpr_lo

    if bits == 64:
-      sgpr_val = _u64(sgpr_lo, self.rsgpr_dyn(off + _c(1)))
+      sgpr_hi = self.rsgpr_dyn(off + _c(1), is_sgpr)
+      sgpr_val = _u64(sgpr_lo, sgpr_hi)
      # Integer inline constants: sign-extend 32-bit value from buffer to 64-bit
      # Float constants: cast F32 to F64
      int_inline = sgpr_lo.cast(dtypes.int32).cast(dtypes.int64)
@@ -482,14 +503,14 @@ class _Ctx:
 # INSTRUCTION HANDLERS
 # ═══════════════════════════════════════════════════════════════════════════════

-def _compile_sopp(inst: SOPP, ctx: _Ctx) -> UOp:
-  simm16 = ctx.inst_field_signed(SOPP.simm16).cast(dtypes.int16)
-  if inst.op == SOPPOp.S_ENDPGM:
+def _compile_sopp(inst: ir3.SOPP | ir4.SOPP, ctx: _Ctx) -> UOp:
+  simm16 = ctx.inst_field_signed(type(inst).simm16).cast(dtypes.int16)
+  if inst.op in (ir3.SOPPOp.S_ENDPGM, ir4.SOPPOp.S_ENDPGM):
    return UOp.sink(ctx.wsgpr_dyn(_c(PC_LO_IDX), UOp.const(dtypes.uint32, 0xFFFFFFFF)),
                          ctx.wsgpr_dyn(_c(PC_HI_IDX), UOp.const(dtypes.uint32, 0xFFFFFFFF)))
-  if inst.op == SOPPOp.S_NOP: return UOp.sink(*ctx.inc_pc())  # S_NOP is a no-op
+  if inst.op in (ir3.SOPPOp.S_NOP, ir4.SOPPOp.S_NOP): return UOp.sink(*ctx.inc_pc())  # S_NOP is a no-op
  # NOTE: we ignore SOPPs without PCODE
-  if inst.op in PCODE:
+  if inst.op in _get_pcode_dict(inst.op):
    pcode = get_pcode(inst.op)
    pc_bytes = ctx.rpc()  # PC is already 64-bit byte address
    vcc, exec_lo = ctx.rsgpr_dyn(_c(VCC_LO.offset)), ctx.rsgpr_dyn(_c(EXEC_LO.offset))
@@ -501,50 +522,57 @@ def _compile_sopp(inst: SOPP, ctx: _Ctx) -> UOp:
        return UOp.sink(ctx.wsgpr_dyn(_c(PC_LO_IDX), lo), ctx.wsgpr_dyn(_c(PC_HI_IDX), hi))
  return UOp.sink(*ctx.inc_pc())

-def _compile_smem(inst: SMEM, ctx: _Ctx) -> UOp:
+def _compile_smem(inst: ir3.SMEM | ir4.SMEM, ctx: _Ctx) -> UOp:
  # Cache invalidation instructions are no-ops in the emulator (we don't model caches)
-  if inst.op in (SMEMOp.S_GL1_INV, SMEMOp.S_DCACHE_INV):
+  cache_inv_ops = [ir3.SMEMOp.S_GL1_INV, ir3.SMEMOp.S_DCACHE_INV, ir4.SMEMOp.S_DCACHE_INV]
+  if hasattr(ir4.SMEMOp, 'S_GL1_INV'): cache_inv_ops.append(ir4.SMEMOp.S_GL1_INV)
+  if inst.op in cache_inv_ops:
    return UOp.sink(*ctx.inc_pc())
  # Dynamic sbase field (bits 5:0) - SGPR pair, field value * 2 = register offset
-  sbase = ctx.inst_field(SMEM.sbase) * _c(2)
+  sbase = ctx.inst_field(type(inst).sbase) * _c(2)
  # Dynamic sdata field (bits 12:6) - destination SGPR
-  sdata_reg = ctx.inst_field(SMEM.sdata)
-  offset = ctx.inst_field_signed(SMEM.offset)  # 21-bit signed immediate
-  # Dynamic soffset field (bits 63:57) - SGPR for additional offset (NULL=124 reads as 0)
-  soffset = ctx.inst_field(SMEM.soffset)
+  sdata_reg = ctx.inst_field(type(inst).sdata)
+  # RDNA4 uses 'ioffset', RDNA3 uses 'offset' - use type(inst) to get correct field
+  offset_field = type(inst).ioffset if hasattr(type(inst), 'ioffset') else type(inst).offset
+  offset = ctx.inst_field_signed(offset_field)  # signed immediate
+  # Dynamic soffset field - SGPR for additional offset (NULL=124 reads as 0)
+  soffset = ctx.inst_field(type(inst).soffset)
  addr = _u64(ctx.rsgpr_dyn(sbase), ctx.rsgpr_dyn(sbase + _c(1))) + offset.cast(dtypes.uint64) + ctx.rsgpr_dyn(soffset).cast(dtypes.uint64)
-  ndwords = {SMEMOp.S_LOAD_B32: 1, SMEMOp.S_LOAD_B64: 2, SMEMOp.S_LOAD_B128: 4, SMEMOp.S_LOAD_B256: 8, SMEMOp.S_LOAD_B512: 16}.get(inst.op, 1)
+  _SMEM_NDWORDS = {ir3.SMEMOp.S_LOAD_B32: 1, ir3.SMEMOp.S_LOAD_B64: 2, ir3.SMEMOp.S_LOAD_B128: 4,
+    ir3.SMEMOp.S_LOAD_B256: 8, ir3.SMEMOp.S_LOAD_B512: 16, ir4.SMEMOp.S_LOAD_B32: 1, ir4.SMEMOp.S_LOAD_B64: 2,
+    ir4.SMEMOp.S_LOAD_B96: 3, ir4.SMEMOp.S_LOAD_B128: 4, ir4.SMEMOp.S_LOAD_B256: 8, ir4.SMEMOp.S_LOAD_B512: 16}
+  ndwords = _SMEM_NDWORDS[inst.op]
  stores = [ctx.wsgpr_dyn(sdata_reg + _c(i), ctx.vmem.index((addr + UOp.const(dtypes.uint64, i * 4) >> UOp.const(dtypes.uint64, 2)).cast(dtypes.int)))
            for i in range(ndwords)]
  return UOp.sink(*stores, *ctx.inc_pc())

-def _compile_sop(inst: SOP1 | SOP2 | SOPC | SOPK, ctx: _Ctx) -> UOp:
+def _compile_sop(inst: ir3.SOP1 | ir3.SOP2 | ir3.SOPC | ir3.SOPK | ir4.SOP1 | ir4.SOP2 | ir4.SOPC | ir4.SOPK, ctx: _Ctx) -> UOp:
  bits = inst.canonical_op_bits
  literal = ctx.inst_field(type(inst).literal) if hasattr(type(inst), 'literal') else None

-  if isinstance(inst, SOPK):
-    sdst_off = ctx.inst_field(SOPK.sdst)
-    simm16 = ctx.inst_field(SOPK.simm16)
+  if isinstance(inst, (ir3.SOPK, ir4.SOPK)):
+    sdst_off = ctx.inst_field(type(inst).sdst)
+    simm16 = ctx.inst_field(type(inst).simm16)
    # Sign-extend simm16
    simm16_sext = simm16.cast(dtypes.int16).cast(dtypes.int32)
    srcs = {'S0': ctx.rsgpr_dyn(sdst_off), 'SIMM16': simm16_sext, 'D0': ctx.rsgpr_dyn(sdst_off)}
    dst_off, dst_size = sdst_off, 1
-  elif isinstance(inst, SOP1):
-    sdst_off = ctx.inst_field(SOP1.sdst)
-    ssrc0_off = ctx.inst_field(SOP1.ssrc0)
+  elif isinstance(inst, (ir3.SOP1, ir4.SOP1)):
+    sdst_off = ctx.inst_field(type(inst).sdst)
+    ssrc0_off = ctx.inst_field(type(inst).ssrc0)
    srcs = {'S0': ctx.rsrc_dyn(ssrc0_off, None, bits['s0'], literal)}
    dst_off, dst_size = sdst_off, bits['d'] // 32
-  elif isinstance(inst, SOP2):
-    sdst_off = ctx.inst_field(SOP2.sdst)
-    ssrc0_off = ctx.inst_field(SOP2.ssrc0)
-    ssrc1_off = ctx.inst_field(SOP2.ssrc1)
+  elif isinstance(inst, (ir3.SOP2, ir4.SOP2)):
+    sdst_off = ctx.inst_field(type(inst).sdst)
+    ssrc0_off = ctx.inst_field(type(inst).ssrc0)
+    ssrc1_off = ctx.inst_field(type(inst).ssrc1)
    srcs = {'S0': ctx.rsrc_dyn(ssrc0_off, None, bits['s0'], literal),
            'S1': ctx.rsrc_dyn(ssrc1_off, None, bits['s1'], literal)}
    if literal is not None: srcs['SIMM32'] = literal
    dst_off, dst_size = sdst_off, bits['d'] // 32
-  elif isinstance(inst, SOPC):
-    ssrc0_off = ctx.inst_field(SOPC.ssrc0)
-    ssrc1_off = ctx.inst_field(SOPC.ssrc1)
+  elif isinstance(inst, (ir3.SOPC, ir4.SOPC)):
+    ssrc0_off = ctx.inst_field(type(inst).ssrc0)
+    ssrc1_off = ctx.inst_field(type(inst).ssrc1)
    srcs = {'S0': ctx.rsrc_dyn(ssrc0_off, None, bits['s0'], literal),
            'S1': ctx.rsrc_dyn(ssrc1_off, None, bits['s1'], literal)}
    dst_off, dst_size = _c(0), 0  # SOPC writes to SCC, not sdst
@@ -553,18 +581,18 @@ def _compile_sop(inst: SOP1 | SOP2 | SOPC | SOPK, ctx: _Ctx) -> UOp:

  return ctx.compile_sop_pcode(inst.op, srcs, dst_off, dst_size)

-def _compile_vop12(inst: VOP1 | VOP1_SDST | VOP2, ctx: _Ctx) -> UOp:
+def _compile_vop12(inst: ir3.VOP1 | ir3.VOP1_SDST | ir3.VOP2 | ir4.VOP1 | ir4.VOP1_SDST | ir4.VOP2, ctx: _Ctx) -> UOp:
  op_name = _op_name(inst)
  if op_name in ('V_READFIRSTLANE_B32_E32', 'V_PERMLANE64_B32_E32'): return ctx.compile_lane_pcode(inst.op, inst)
  lane, exec_mask, bits = ctx.range(), ctx.rsgpr_dyn(_c(EXEC_LO.offset)), inst.canonical_op_bits
  literal = ctx.inst_field(type(inst).literal) if hasattr(type(inst), 'literal') else None
-  vdst_reg = ctx.inst_field(VOP1.vdst)
+  vdst_reg = ctx.inst_field(type(inst).vdst)
  write_hi_half = bits['d'] == 16 and (vdst_reg >= _c(128))
  if isinstance(write_hi_half, UOp): vdst_reg = write_hi_half.where(vdst_reg - _c(128), vdst_reg)
  elif write_hi_half: vdst_reg -= 128
-  if isinstance(inst, VOP1):
+  if isinstance(inst, (ir3.VOP1, ir4.VOP1)):
    # Handle VOP1 hi-half source operand (src0 >= v[128] for 16-bit ops)
-    src0_off = ctx.inst_field(VOP1.src0)
+    src0_off = ctx.inst_field(type(inst).src0)
    s0 = ctx.rsrc_dyn(src0_off, lane, bits['s0'], literal)
    if bits['s0'] == 16:
      src0_hi = src0_off >= _c(384)
@@ -573,13 +601,13 @@ def _compile_vop12(inst: VOP1 | VOP1_SDST | VOP2, ctx: _Ctx) -> UOp:
      s0 = src0_hi.where(_hi16(ctx.rvgpr_dyn(src0_reg, lane)), s0)
    srcs = {'S0': s0}
  else:
-    vsrc1_reg = ctx.inst_field(VOP2.vsrc1)
+    vsrc1_reg = ctx.inst_field(type(inst).vsrc1)
    vsrc1_hi = bits['s0'] == 16 and (vsrc1_reg >= _c(128))
    vsrc1_actual = _cond(vsrc1_hi, vsrc1_reg - _c(128), vsrc1_reg)
    s1 = _cond_hi16(vsrc1_hi, ctx.rvgpr_dyn(vsrc1_actual, lane))
    d0 = _cond_hi16(write_hi_half, ctx.rvgpr_dyn(vdst_reg, lane))  # FMAC/FMAMK hi-half dest needs hi-half accumulator
    # Handle VOP2 hi-half src0 operand (src0 >= v[128] for 16-bit ops)
-    src0_off = ctx.inst_field(VOP2.src0)
+    src0_off = ctx.inst_field(type(inst).src0)
    s0 = ctx.rsrc_dyn(src0_off, lane, bits['s0'], literal)
    if bits['s0'] == 16:
      src0_hi = src0_off >= _c(384)
@@ -587,19 +615,20 @@ def _compile_vop12(inst: VOP1 | VOP1_SDST | VOP2, ctx: _Ctx) -> UOp:
      src0_reg = src0_hi.where(src0_off - _c(384), _c(0))
      s0 = src0_hi.where(_hi16(ctx.rvgpr_dyn(src0_reg, lane)), s0)
    srcs = {'S0': s0, 'S1': s1, 'D0': d0}
-    if inst.op in (VOP2Op.V_FMAAK_F32_E32, VOP2Op.V_FMAMK_F32_E32, VOP2Op.V_FMAAK_F16_E32, VOP2Op.V_FMAMK_F16_E32):
+    if inst.op in (ir3.VOP2Op.V_FMAAK_F32_E32, ir3.VOP2Op.V_FMAMK_F32_E32, ir3.VOP2Op.V_FMAAK_F16_E32,
+                   ir3.VOP2Op.V_FMAMK_F16_E32):
      assert literal is not None
      srcs['SIMM32'] = literal
  return ctx.compile_vop_pcode(inst.op, srcs, lane, vdst_reg, exec_mask, opsel_dst_hi=write_hi_half)

-def _compile_vopc(inst: VOPC | VOP3, ctx: _Ctx, opsel: int = 0, abs_bits: int = 0, neg_bits: int = 0) -> UOp:
+def _compile_vopc(inst: ir3.VOPC | ir3.VOP3 | ir4.VOPC | ir4.VOP3, ctx: _Ctx, opsel: int = 0, abs_bits: int = 0, neg_bits: int = 0) -> UOp:
  exec_mask, op_name, bits = ctx.rsgpr_dyn(_c(EXEC_LO.offset)), _op_name(inst), inst.canonical_op_bits
  is_cmpx, is_vopc = 'CMPX' in op_name, hasattr(inst, 'vsrc1')  # is_vopc: e32 vs e64

  # Handle both VOPC (vsrc1) and VOP3 (src1) instruction formats - read operands dynamically
  if is_vopc:
-    src0_off = ctx.inst_field(VOPC.src0)
-    vsrc1_off = ctx.inst_field(VOPC.vsrc1)
+    src0_off = ctx.inst_field(type(inst).src0)
+    vsrc1_off = ctx.inst_field(type(inst).vsrc1)
    # For 16-bit ops, vsrc1 >= 128 means hi-half of v[vsrc1-128]
    if bits['s0'] == 16:
      vsrc1_hi = vsrc1_off >= _c(128)
@@ -608,9 +637,9 @@ def _compile_vopc(inst: VOPC | VOP3, ctx: _Ctx, opsel: int = 0, abs_bits: int =
      vsrc1_hi = False
      src1_off = _c(256) + vsrc1_off
  else:
-    src0_off = ctx.inst_field(VOP3.src0)
-    src1_off = ctx.inst_field(VOP3.src1)
-    dst_off = ctx.inst_field(VOP3.vdst)
+    src0_off = ctx.inst_field(type(inst).src0)
+    src1_off = ctx.inst_field(type(inst).src1)
+    dst_off = ctx.inst_field(type(inst).vdst)
    vsrc1_hi = False
  literal = ctx.inst_field(type(inst).literal) if hasattr(type(inst), 'literal') else None

@@ -639,7 +668,7 @@ def _compile_vopc(inst: VOPC | VOP3, ctx: _Ctx, opsel: int = 0, abs_bits: int =
    stores = [ctx.wsgpr_dyn(dst_off, new_result)] if not is_vopc else [ctx.wsgpr_dyn(_c(VCC_LO.offset), new_result)]
  return UOp.sink(*stores, *ctx.inc_pc())

-def _compile_vop3(inst: VOP3, ctx: _Ctx) -> UOp:
+def _compile_vop3(inst: ir3.VOP3 | ir4.VOP3, ctx: _Ctx) -> UOp:
  exec_mask = ctx.rsgpr_dyn(_c(EXEC_LO.offset))
  bits = inst.canonical_op_bits
  opsel, op_name = getattr(inst, 'opsel', 0) or 0, _op_name(inst)
@@ -658,12 +687,12 @@ def _compile_vop3(inst: VOP3, ctx: _Ctx) -> UOp:

  # Regular VOP3 - read operands dynamically
  lane = ctx.range()
-  vdst_reg = ctx.inst_field(VOP3.vdst)
+  vdst_reg = ctx.inst_field(type(inst).vdst)
  literal = ctx.inst_field(type(inst).literal) if hasattr(type(inst), 'literal') else None
  ops = inst.canonical_operands
-  src0 = ctx.rsrc_dyn(ctx.inst_field(VOP3.src0), lane, bits['s0'], literal, 's0' in ops and ops['s0'][0] == Fmt.FMT_NUM_F64)
-  src1 = ctx.rsrc_dyn(ctx.inst_field(VOP3.src1), lane, bits['s1'], literal, 's1' in ops and ops['s1'][0] == Fmt.FMT_NUM_F64)
-  src2 = ctx.rsrc_dyn(ctx.inst_field(VOP3.src2), lane, bits['s2'], literal, 's2' in ops and ops['s2'][0] == Fmt.FMT_NUM_F64)
+  src0 = ctx.rsrc_dyn(ctx.inst_field(type(inst).src0), lane, bits['s0'], literal, 's0' in ops and ops['s0'][0] == Fmt.FMT_NUM_F64)
+  src1 = ctx.rsrc_dyn(ctx.inst_field(type(inst).src1), lane, bits['s1'], literal, 's1' in ops and ops['s1'][0] == Fmt.FMT_NUM_F64)
+  src2 = ctx.rsrc_dyn(ctx.inst_field(type(inst).src2), lane, bits['s2'], literal, 's2' in ops and ops['s2'][0] == Fmt.FMT_NUM_F64)
  if bits['s0'] == 16:
    src0 = _apply_opsel(src0, 0, opsel)
    src1 = _apply_opsel(src1, 1, opsel)
@@ -673,19 +702,19 @@ def _compile_vop3(inst: VOP3, ctx: _Ctx) -> UOp:
  src1 = _apply_src_mods(src1, 1, abs_bits, neg_bits, bits['s1'])
  src2 = _apply_src_mods(src2, 2, abs_bits, neg_bits, bits['s2'])
  srcs = {'S0': src0, 'S1': src1, 'S2': src2}
-  if inst.op in (VOP3Op.V_CNDMASK_B32_E64, VOP3Op.V_CNDMASK_B16) and src2 is not None: srcs['VCC'] = src2
+  if inst.op in (ir3.VOP3Op.V_CNDMASK_B32_E64, ir3.VOP3Op.V_CNDMASK_B16) and src2 is not None: srcs['VCC'] = src2
  # FMAC instructions need D0 (accumulator) from destination register
  if 'FMAC' in op_name: srcs['D0'] = ctx.rvgpr_dyn(vdst_reg, lane)
  opsel_dst_hi = bool(opsel & 0b1000) and bits['d'] == 16
  return ctx.compile_vop_pcode(inst.op, srcs, lane, vdst_reg, exec_mask, opsel_dst_hi=opsel_dst_hi, clmp=getattr(inst, 'clmp', 0))

-def _compile_vop3sd(inst: VOP3SD, ctx: _Ctx) -> UOp:
+def _compile_vop3sd(inst: ir3.VOP3SD | ir4.VOP3SD, ctx: _Ctx) -> UOp:
  exec_mask = ctx.rsgpr_dyn(_c(EXEC_LO.offset))
  bits, pcode, ops = inst.canonical_op_bits, get_pcode(inst.op), inst.canonical_operands

  # Read operands dynamically from instruction encoding
-  vdst_reg, sdst_off = ctx.inst_field(VOP3SD.vdst), ctx.inst_field(VOP3SD.sdst)
-  src0_off, src1_off, src2_off = ctx.inst_field(VOP3SD.src0), ctx.inst_field(VOP3SD.src1), ctx.inst_field(VOP3SD.src2)
+  vdst_reg, sdst_off = ctx.inst_field(type(inst).vdst), ctx.inst_field(type(inst).sdst)
+  src0_off, src1_off, src2_off = ctx.inst_field(type(inst).src0), ctx.inst_field(type(inst).src1), ctx.inst_field(type(inst).src2)
  literal = ctx.inst_field(type(inst).literal) if hasattr(type(inst), 'literal') else None

  has_carry_in = 's2' in ops and ops['s2'][2] == OpType.OPR_SREG
@@ -731,13 +760,13 @@ def _compile_vop3sd(inst: VOP3SD, ctx: _Ctx) -> UOp:
  else:
    return ctx.compile_vop_pcode(inst.op, srcs, lane, vdst_reg, exec_mask, sdst_reg=inst.sdst.offset)

-def _compile_wmma(inst: VOP3P, ctx: _Ctx) -> UOp:
+def _compile_wmma(inst: ir3.VOP3P | ir4.VOP3P, ctx: _Ctx) -> UOp:
  op_name = _op_name(inst)
  exec_mask = ctx.rsgpr_dyn(_c(EXEC_LO.offset))
-  vdst_reg = ctx.inst_field(VOP3P.vdst)
-  src0_r = ctx.inst_field(VOP3P.src0) - _c(256)
-  src1_r = ctx.inst_field(VOP3P.src1) - _c(256)
-  src2_r = ctx.inst_field(VOP3P.src2) - _c(256)
+  vdst_reg = ctx.inst_field(type(inst).vdst)
+  src0_r = ctx.inst_field(type(inst).src0) - _c(256)
+  src1_r = ctx.inst_field(type(inst).src1) - _c(256)
+  src2_r = ctx.inst_field(type(inst).src2) - _c(256)
  is_f16_output = 'F16_16X16X16_F16' in op_name or 'BF16_16X16X16_BF16' in op_name  # F16/BF16 output vs F32 output
  is_bf16 = 'BF16' in op_name
  cvt = _FUNCS['bf16_to_f32'] if is_bf16 else _FUNCS['f16_to_f32']
@@ -764,16 +793,16 @@ def _compile_wmma(inst: VOP3P, ctx: _Ctx) -> UOp:
    stores = [ctx.wvgpr_dyn(vdst_reg + _c(i // 32), UOp.const(dtypes.int, i % 32), mat_d[i].bitcast(dtypes.uint32), exec_mask) for i in range(256)]
  return UOp.sink(*stores, *ctx.inc_pc())

-def _compile_vop3p(inst: VOP3P, ctx: _Ctx) -> UOp:
+def _compile_vop3p(inst: ir3.VOP3P | ir4.VOP3P, ctx: _Ctx) -> UOp:
  op_name = _op_name(inst)
  if 'WMMA' in op_name and ('16X16X16_F16' in op_name or '16X16X16_BF16' in op_name): return _compile_wmma(inst, ctx)

  lane = ctx.range()
  exec_mask = ctx.rsgpr_dyn(_c(EXEC_LO.offset))
-  vdst_reg = ctx.inst_field(VOP3P.vdst)
-  src0 = ctx.rsrc_dyn(ctx.inst_field(VOP3P.src0), lane, 16)
-  src1 = ctx.rsrc_dyn(ctx.inst_field(VOP3P.src1), lane, 16)
-  src2 = ctx.rsrc_dyn(ctx.inst_field(VOP3P.src2), lane, 16)
+  vdst_reg = ctx.inst_field(type(inst).vdst)
+  src0 = ctx.rsrc_dyn(ctx.inst_field(type(inst).src0), lane, 16)
+  src1 = ctx.rsrc_dyn(ctx.inst_field(type(inst).src1), lane, 16)
+  src2 = ctx.rsrc_dyn(ctx.inst_field(type(inst).src2), lane, 16)
  opsel, opsel_hi = getattr(inst, 'opsel', 0) or 0, getattr(inst, 'opsel_hi', 3) if getattr(inst, 'opsel_hi', 3) is not None else 3
  opsel_hi2 = getattr(inst, 'opsel_hi2', 1) if getattr(inst, 'opsel_hi2', 1) is not None else 1
  neg, neg_hi = getattr(inst, 'neg', 0) or 0, getattr(inst, 'neg_hi', 0) or 0
@@ -813,18 +842,19 @@ def _compile_vop3p(inst: VOP3P, ctx: _Ctx) -> UOp:
    if is_dot_iu: srcs['NEG'] = UOp.const(dtypes.uint32, neg)
  return ctx.compile_vop_pcode(inst.op, srcs, lane, vdst_reg, exec_mask)

-def _compile_vopd(inst: VOPD, ctx: _Ctx) -> UOp:
+def _compile_vopd(inst: ir3.VOPD | ir4.VOPD, ctx: _Ctx) -> UOp:
  exec_mask = ctx.rsgpr_dyn(_c(EXEC_LO.offset))
-  # Read operands dynamically
-  vdstx_reg = ctx.inst_field(VOPD.vdstx)
+  # Read operands dynamically - use type(inst) to get correct field descriptors
+  inst_type = type(inst)
+  vdstx_reg = ctx.inst_field(inst_type.vdstx)
  # vdsty has complex encoding: actual = (raw << 1) | ((vdstx & 1) ^ 1)
-  vdsty_raw = ctx.inst_field(VOPD.vdsty)
+  vdsty_raw = ctx.inst_field(inst_type.vdsty)
  vdsty_reg = (vdsty_raw << _c(1)) | ((vdstx_reg & _c(1)) ^ _c(1))
-  srcx0_off = ctx.inst_field(VOPD.srcx0)
-  srcy0_off = ctx.inst_field(VOPD.srcy0)
-  vsrcx1_reg = ctx.inst_field(VOPD.vsrcx1)
-  vsrcy1_reg = ctx.inst_field(VOPD.vsrcy1)
-  literal = ctx.inst_field(type(inst).literal) if hasattr(type(inst), 'literal') else None
+  srcx0_off = ctx.inst_field(inst_type.srcx0)
+  srcy0_off = ctx.inst_field(inst_type.srcy0)
+  vsrcx1_reg = ctx.inst_field(inst_type.vsrcx1)
+  vsrcy1_reg = ctx.inst_field(inst_type.vsrcy1)
+  literal = ctx.inst_field(inst_type.literal) if hasattr(inst_type, 'literal') else None

  lane = ctx.range()
  srcy0, srcy1 = ctx.rsrc_dyn(srcy0_off, lane, literal=literal), ctx.rvgpr_dyn(vsrcy1_reg, lane)
@@ -835,49 +865,55 @@ def _compile_vopd(inst: VOPD, ctx: _Ctx) -> UOp:
    assert vop is not None, f"no VOP mapping for VOPD {label}: {op}"
    if label == 'Y': srcs = {'S0': srcy0, 'S1': srcy1, 'D0': ctx.rvgpr_dyn(vdst_reg, lane)}
    else: srcs = {'S0': ctx.rsrc_dyn(src0_off, lane, literal=literal), 'S1': ctx.rvgpr_dyn(vsrc1_reg, lane), 'D0': ctx.rvgpr_dyn(vdst_reg, lane)}
-    if op in (VOPDOp.V_DUAL_FMAAK_F32, VOPDOp.V_DUAL_FMAMK_F32):
+    if op in (ir3.VOPDOp.V_DUAL_FMAAK_F32, ir3.VOPDOp.V_DUAL_FMAMK_F32, ir4.VOPDOp.V_DUAL_FMAAK_F32, ir4.VOPDOp.V_DUAL_FMAMK_F32):
      assert literal is not None
      srcs['SIMM32'] = literal
-    if op == VOPDOp.V_DUAL_CNDMASK_B32: srcs['VCC'] = ctx.rsgpr_dyn(_c(VCC_LO.offset))
+    if op in (ir3.VOPDOp.V_DUAL_CNDMASK_B32, ir4.VOPDOp.V_DUAL_CNDMASK_B32): srcs['VCC'] = ctx.rsgpr_dyn(_c(VCC_LO.offset))
    pcode = get_pcode(vop)
    srcs.update({'VCC': ctx.rsgpr_dyn(_c(VCC_LO.offset)), 'EXEC': exec_mask, 'SCC': ctx.rsgpr_dyn(_c(SCC.offset)), 'laneId': lane})
    for dest, val in parse_pcode(pcode, srcs)[1]:
      if dest.startswith('D0'): all_stores.append(ctx.wvgpr_dyn(vdst_reg, lane, _val_to_u32(val), exec_mask, after=srcy1))
  return UOp.sink(UOp.group(*all_stores).end(lane), *ctx.inc_pc())

-def _compile_mem_op(inst: DS | FLAT | GLOBAL | SCRATCH, ctx: _Ctx) -> UOp:
+def _compile_mem_op(inst: ir3.DS | ir3.FLAT | ir3.GLOBAL | ir3.SCRATCH | ir4.DS | ir4.VFLAT | ir4.VGLOBAL | ir4.VSCRATCH, ctx: _Ctx) -> UOp:
  """Unified memory operation compiler for DS, FLAT, GLOBAL, SCRATCH."""
  exec_mask, op_name = ctx.rsgpr_dyn(_c(EXEC_LO.offset)), _op_name(inst)
  pcode = get_pcode(inst.op)

-  is_lds = isinstance(inst, DS)
-  is_scratch = isinstance(inst, SCRATCH)
+  is_lds = isinstance(inst, (ir3.DS, ir4.DS))
+  is_scratch = isinstance(inst, (ir3.SCRATCH, ir4.VSCRATCH))
  mem = ctx.lds if is_lds else ctx.scratch if is_scratch else ctx.vmem
  addr_shift = UOp.const(dtypes.uint32 if is_lds else dtypes.uint64, 2)

  # Extract register info - all dynamic for deduplication
  if is_lds:
-    addr_reg = ctx.inst_field(DS.addr)
-    vdata_reg = ctx.inst_field(DS.data0)
-    vdst_reg = ctx.inst_field(DS.vdst)
-    offset0 = ctx.inst_field(DS.offset0)
-    offset1 = ctx.inst_field(DS.offset1)
+    addr_reg = ctx.inst_field(type(inst).addr)
+    vdata_reg = ctx.inst_field(type(inst).data0)
+    vdst_reg = ctx.inst_field(type(inst).vdst)
+    offset0 = ctx.inst_field(type(inst).offset0)
+    offset1 = ctx.inst_field(type(inst).offset1)
    offset = offset0  # DS uses offset0 as primary offset
    saddr_reg = None
-  else:
+  elif isinstance(inst, (ir4.VGLOBAL, ir4.VSCRATCH, ir4.VFLAT)):  # RDNA4: vaddr, vsrc, ioffset
+    addr_reg = ctx.inst_field(type(inst).vaddr)
+    vdata_reg = ctx.inst_field(type(inst).vsrc)
+    vdst_reg = ctx.inst_field(type(inst).vdst)
+    offset = ctx.inst_field_signed(type(inst).ioffset)
+    offset0, offset1 = _c(0), _c(0)
+    saddr_reg = ctx.inst_field(type(inst).saddr) if hasattr(type(inst), 'saddr') else None
+  else:  # RDNA3: addr, data, offset
    addr_reg = ctx.inst_field(type(inst).addr)
    vdata_reg = ctx.inst_field(type(inst).data)
    vdst_reg = ctx.inst_field(type(inst).vdst)
    offset = ctx.inst_field_signed(type(inst).offset)
    offset0, offset1 = _c(0), _c(0)
-    # Dynamic saddr - read field, NULL (124) or >= 128 means no saddr
-    saddr_reg = ctx.inst_field(type(inst).saddr) if hasattr(inst, 'saddr') else None
+    saddr_reg = ctx.inst_field(type(inst).saddr) if hasattr(type(inst), 'saddr') else None

  # Data width from canonical_op_bits (32/64/96/128), default to 32 for untyped ops
  data_bits_mem = inst.canonical_op_bits.get('data', 32)
  is_atomic, glc = 'ATOMIC' in op_name, getattr(inst, 'glc', 0)
  has_data1 = is_lds and hasattr(inst, 'data1') and inst.data1 is not None
-  data1_reg = ctx.inst_field(DS.data1) if is_lds else _c(0)
+  data1_reg = ctx.inst_field(type(inst).data1) if is_lds else _c(0)

  # DS_PERMUTE/DS_BPERMUTE: cross-lane VGPR access via pcode
  if is_lds and 'PERMUTE' in op_name:
@@ -928,14 +964,26 @@ def _compile_mem_op(inst: DS | FLAT | GLOBAL | SCRATCH, ctx: _Ctx) -> UOp:
      else:
        data = {'DATA': _u64(ctx.rvgpr_dyn(vdata_reg, lane), ctx.rvgpr_dyn(vdata_reg + _c(1), lane)),
                'DATA2': _u64(ctx.rvgpr_dyn(data1_reg, lane), ctx.rvgpr_dyn(data1_reg + _c(1), lane)) if has_data1 else UOp.const(dtypes.uint64, 0)}
-      return {'ADDR': addr, 'ADDR_BASE': addr, 'OFFSET': offset, 'OFFSET0': offset0, 'OFFSET1': offset1, '_lds': mem, 'laneId': lane, **data}
+      # RDNA3 uses ADDR/OFFSET, RDNA4 uses vgpr_a/offset (lowercase) + CalcDsAddr function
+      return {'ADDR': addr, 'ADDR_BASE': addr, 'OFFSET': offset, 'OFFSET0': offset0, 'OFFSET1': offset1, '_lds': mem, 'laneId': lane,
+              'vgpr_a': ctx.rvgpr_dyn(addr_reg, lane), 'offset': offset, **data}
    active = _lane_active(exec_mask, lane)
+    # saddr < 124 means valid SGPR pair, otherwise use 0 (NULL means no saddr contribution)
+    use_saddr = (saddr_reg < _c(124)) if saddr_reg is not None else UOp.const(dtypes.bool, False)
+    saddr_raw = _u64(ctx.rsgpr_dyn(saddr_reg), ctx.rsgpr_dyn(saddr_reg + _c(1))) if saddr_reg is not None else UOp.const(dtypes.uint64, 0)
+    saddr_base = use_saddr.where(saddr_raw, UOp.const(dtypes.uint64, 0))
+    # Sign-extend offset to 64-bit for the final address calculation
+    ioffset64 = offset.cast(dtypes.int64).cast(dtypes.uint64)
+    # v_addr for CalcGlobalAddr: when saddr valid, use low 32 bits as offset; otherwise full 64-bit address. Include ioffset.
+    vaddr_full = _u64(ctx.rvgpr_dyn(addr_reg, lane), ctx.rvgpr_dyn(addr_reg + _c(1), lane))
+    vaddr_lo = ctx.rvgpr_dyn(addr_reg, lane).cast(dtypes.uint64)
+    vaddr_base = use_saddr.where(vaddr_lo + ioffset64, vaddr_full + ioffset64)
    if is_atomic:
      return {'ADDR': addr, 'DATA': _u64(ctx.rvgpr_dyn(vdata_reg, lane), ctx.rvgpr_dyn(vdata_reg + _c(1), lane)) if data_bits_mem == 64 else ctx.rvgpr_dyn(vdata_reg, lane),
-              '_vmem': mem, '_active': active, 'laneId': lane}
+              '_vmem': mem, '_active': active, 'laneId': lane, 'v_addr': vaddr_base, 's_saddr': saddr_base}
    vdata = ctx.rvgpr_dyn(vdata_reg, lane).cast(dtypes.uint64) if 'STORE' in op_name else ctx.rvgpr_dyn(vdst_reg, lane) if 'D16' in op_name else UOp.const(dtypes.uint32, 0)
    if 'STORE' in op_name and data_bits_mem >= 64: vdata = vdata | (ctx.rvgpr_dyn(vdata_reg + _c(1), lane).cast(dtypes.uint64) << UOp.const(dtypes.uint64, 32))
-    srcs = {'ADDR': addr, 'VDATA': vdata, '_vmem': mem, '_active': active, 'laneId': lane}
+    srcs = {'ADDR': addr, 'VDATA': vdata, '_vmem': mem, '_active': active, 'laneId': lane, 'v_addr': vaddr_base, 's_saddr': saddr_base}
    for i in range(data_bits_mem // 32): srcs[f'VDATA{i}'] = ctx.rvgpr_dyn(vdata_reg + _c(i), lane) if 'STORE' in op_name else UOp.const(dtypes.uint32, 0)
    return srcs

@@ -986,10 +1034,15 @@ def _compile_mem_op(inst: DS | FLAT | GLOBAL | SCRATCH, ctx: _Ctx) -> UOp:

 # Dispatch table: instruction type -> handler function
 _INST_HANDLERS: dict[type, Callable[..., UOp]] = {
-  SOPP: _compile_sopp, SMEM: _compile_smem, SOP1: _compile_sop, SOP2: _compile_sop, SOPC: _compile_sop, SOPK: _compile_sop,
-  VOP1: _compile_vop12, VOP1_SDST: _compile_vop12, VOP2: _compile_vop12, VOPC: _compile_vopc, VOP3: _compile_vop3, VOP3_SDST: _compile_vop3,
-  VOP3SD: _compile_vop3sd, VOP3P: _compile_vop3p, VOPD: _compile_vopd,
-  DS: _compile_mem_op, FLAT: _compile_mem_op, GLOBAL: _compile_mem_op, SCRATCH: _compile_mem_op,
+  ir3.SOPP: _compile_sopp, ir3.SMEM: _compile_smem, ir3.SOP1: _compile_sop, ir3.SOP2: _compile_sop, ir3.SOPC: _compile_sop, ir3.SOPK: _compile_sop,
+  ir3.VOP1: _compile_vop12, ir3.VOP1_SDST: _compile_vop12, ir3.VOP2: _compile_vop12, ir3.VOPC: _compile_vopc, ir3.VOP3: _compile_vop3,
+  ir3.VOP3_SDST: _compile_vop3, ir3.VOP3SD: _compile_vop3sd, ir3.VOP3P: _compile_vop3p, ir3.VOPD: _compile_vopd,
+  ir3.DS: _compile_mem_op, ir3.FLAT: _compile_mem_op, ir3.GLOBAL: _compile_mem_op, ir3.SCRATCH: _compile_mem_op,
+  # RDNA4 instruction classes
+  ir4.SOPP: _compile_sopp, ir4.SMEM: _compile_smem, ir4.SOP1: _compile_sop, ir4.SOP2: _compile_sop, ir4.SOPC: _compile_sop, ir4.SOPK: _compile_sop,
+  ir4.VOP1: _compile_vop12, ir4.VOP1_SDST: _compile_vop12, ir4.VOP2: _compile_vop12, ir4.VOPC: _compile_vopc, ir4.VOP3: _compile_vop3,
+  ir4.VOP3_SDST: _compile_vop3, ir4.VOP3SD: _compile_vop3sd, ir4.VOP3P: _compile_vop3p, ir4.VOPD: _compile_vopd,
+  ir4.DS: _compile_mem_op, ir4.VFLAT: _compile_mem_op, ir4.VGLOBAL: _compile_mem_op, ir4.VSCRATCH: _compile_mem_op,
 }

 # ═══════════════════════════════════════════════════════════════════════════════
@@ -999,9 +1052,9 @@ _INST_HANDLERS: dict[type, Callable[..., UOp]] = {
 _canonical_runner_cache: list[tuple[int, int, int, object]] = []  # [(base, mask, size, runner), ...]

@functools.cache
-def _get_runner(inst_bytes: bytes):
+def _get_runner(inst_bytes: bytes, arch: str = "rdna3"):
  """Build and compile instruction to CompiledRunner. Cached by instruction bytes, with canonical dedup."""
-  inst = decode_inst(inst_bytes)
+  inst = decode_inst(inst_bytes, arch)
  inst_size = inst.size()
  inst_int = int.from_bytes(inst_bytes[:inst_size], 'little')

@@ -1030,15 +1083,15 @@ def _get_runner(inst_bytes: bytes):
  return runner, True

@functools.cache
-def decode_program(data: bytes) -> dict[int, tuple[str, Callable, list[int], Any]]:
+def decode_program(data: bytes, arch: str = "rdna3") -> dict[int, tuple[str, Callable, list[int], Any]]:
  """Decode program to {pc: (name, fxn, globals, runner)}."""
  result: dict[int, tuple[str, Callable, list[int], Any]] = {}
  i = 0
  while i < len(data):
-    inst = decode_inst(data[i:])
-    if isinstance(inst, SOPP) and inst.op == SOPPOp.S_CODE_END: break
+    inst = decode_inst(data[i:], arch)
+    if hasattr(inst, 'op') and inst.op in (ir3.SOPPOp.S_CODE_END, ir4.SOPPOp.S_CODE_END): break
    try:
-      runner, is_new = _get_runner(bytes(data[i:i + inst.size() + 4]))
+      runner, is_new = _get_runner(bytes(data[i:i + inst.size() + 4]), arch)
      if DEBUG >= 3:
        try: inst_str = repr(inst)
        except Exception: inst_str = f"<{type(inst).__name__} at PC={i}>"
@@ -1097,9 +1150,9 @@ class WaveState:
 # ═══════════════════════════════════════════════════════════════════════════════

 def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int, lz: int, args_ptr: int, rsrc2: int = 0x19c,
-            scratch_size: int = 0) -> int:
+            scratch_size: int = 0, arch: str = "rdna3") -> int:
  """Execute AMD assembly program. scratch_size is private_segment_fixed_size from kernel descriptor (per-lane)."""
-  program_raw = decode_program(bytes((ctypes.c_char * lib_sz).from_address(lib).raw))
+  program_raw = decode_program(bytes((ctypes.c_char * lib_sz).from_address(lib).raw), arch)
  program = {lib + offset: val for offset, val in program_raw.items()}  # Remap to actual addresses
  lds_size = ((rsrc2 & hsa.AMD_COMPUTE_PGM_RSRC_TWO_GRANULATED_LDS_SIZE) >> hsa.AMD_COMPUTE_PGM_RSRC_TWO_GRANULATED_LDS_SIZE_SHIFT) * 512
  total_threads = lx * ly * lz
@@ -1127,6 +1180,12 @@ def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int,
                                 (hsa.AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_Z, gidz)]:
              if rsrc2 & enabled: st._write_sgpr(sgpr_idx, gid); sgpr_idx += 1

+            # RDNA4 uses TTMP registers for workgroup IDs: ttmp[9]=gidx, ttmp[10]=gidy, ttmp[11]=gidz
+            if arch == "rdna4":
+              st._write_sgpr(ttmp[9].offset, gidx)
+              st._write_sgpr(ttmp[10].offset, gidy)
+              st._write_sgpr(ttmp[11].offset, gidz)
+
            # v0 = packed workitem IDs, scratch stride in secret SGPR
            for lane in range(n_lanes):
              tid = wave_start + lane
@@ -1143,7 +1202,7 @@ def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int,
              assert fxn is not None, f"[emu] No fxn for {name} at PC={pc}"
              assert 4 not in globals_list or scratch_buf, f"SCRATCH instruction {name} but scratch_size=0"
              if DEBUG >= 6:
-                inst = decode_inst(bytes((ctypes.c_char * 12).from_address(pc).raw))
+                inst = decode_inst(bytes((ctypes.c_char * 12).from_address(pc).raw), arch)
                print(f"[emu] exec PC={pc:X}: {inst!r}")
              fxn(*[c_bufs[g] for g in globals_list])
            else: raise RuntimeError("exceeded 1M instructions, likely infinite loop")
@@ -271,6 +271,9 @@ _FUNCS: dict[str, Callable[..., UOp]] = {
  # System NOPs - these are scheduling hints, no effect on emulation
  'MIN': lambda a, b: (a < b).where(a, b),
  's_nop': lambda a: _u32(0),
+  # Address calculation for memory operations
+  'CalcDsAddr': lambda a, o, *r: a.cast(dtypes.uint32) + o.cast(dtypes.uint32),
+  'CalcGlobalAddr': lambda v, s, *r: v.cast(dtypes.uint64) + s.cast(dtypes.uint64),
 }
 for is_max, name in [(False, 'min'), (True, 'max')]:
  for dt, sfx in [(dtypes.float32, 'f32'), (dtypes.int, 'i32'), (dtypes.uint32, 'u32'), (dtypes.int16, 'i16'), (dtypes.uint16, 'u16')]:
@@ -2,11 +2,8 @@
 from dataclasses import dataclass
 from typing import Iterator

-from tinygrad.runtime.support.elf import elf_loader
-
 from extra.assembly.amd.sqtt import decode, print_packets, INST, VALUINST, IMMEDIATE, WAVESTART, WAVEEND, InstOp, PacketType, IMMEDIATE_MASK
 from extra.assembly.amd.dsl import Inst
-from extra.assembly.amd import decode_inst
 from extra.assembly.amd.autogen.rdna3.ins import SOPP, s_endpgm
 from extra.assembly.amd.autogen.rdna3.enum import SOPPOp

@@ -16,19 +13,11 @@ class InstructionInfo:
  wave: int
  inst: Inst

-def map_insts(data:bytes, lib:bytes) -> Iterator[tuple[PacketType, InstructionInfo|None]]:
+def map_insts(data:bytes, lib:bytes, target:int) -> Iterator[tuple[PacketType, InstructionInfo|None]]:
  """maps SQTT packets to instructions, yields (packet, instruction_info or None)"""
  # map pcs to insts
-  pc_map:dict[int, Inst] = {}
-  image, sections, _ = elf_loader(lib)
-  text = next((sh for sh in sections if sh.name == ".text"), None)
-  assert text is not None, "no .text section found"
-  text_off, text_size = text.header.sh_addr, text.header.sh_size
-  offset = text_off
-  while offset < text_off + text_size:
-    inst = decode_inst(image[offset:])
-    pc_map[offset-text_off] = inst
-    offset += inst.size()
+  from tinygrad.viz.serve import amd_decode
+  pc_map = amd_decode(lib, target)

  wave_pc:dict[int, int] = {}
  # only processing packets on one [CU, SIMD] unit
@@ -37,7 +26,7 @@ def map_insts(data:bytes, lib:bytes) -> Iterator[tuple[PacketType, InstructionIn
    if not simd_select(p): continue
    if isinstance(p, WAVESTART):
      assert p.wave not in wave_pc, "only one inflight wave per unit"
-      wave_pc[p.wave] = 0
+      wave_pc[p.wave] = next(iter(pc_map))
      continue
    if isinstance(p, WAVEEND):
      pc = wave_pc.pop(p.wave)
@@ -80,22 +69,22 @@ def map_insts(data:bytes, lib:bytes) -> Iterator[tuple[PacketType, InstructionIn
 # test to compare every packet with the rocprof decoder

 def test_rocprof_inst_traces_match(sqtt, prg, target):
-  from tinygrad.viz.serve import llvm_disasm
+  from tinygrad.viz.serve import amd_decode
  from extra.sqtt.roc import decode as roc_decode, InstExec
-  disasm = {addr+prg.base:inst_disasm for addr, inst_disasm in llvm_disasm(target, prg.lib).items()}
-  rctx = roc_decode([sqtt], {prg.name:disasm})
-  rwaves = rctx.inst_execs[(sqtt.kern, sqtt.exec_tag)]
+  addr_table = amd_decode(prg.lib, target)
+  disasm = {addr+prg.base:(inst.disasm(), inst.size()) for addr,inst in addr_table.items()}
+  rctx = roc_decode([sqtt], {prg.tag:disasm})
+  rwaves = rctx.inst_execs.get((sqtt.kern, sqtt.exec_tag), [])
  rwaves_iter:dict[int, list[Iterator[InstExec]]] = {} # wave unit (0-15) -> list of inst trace iterators for all executions on that unit
  for w in rwaves: rwaves_iter.setdefault(w.wave_id, []).append(w.unpack_insts())
-  rwaves_base = next(iter(disasm)) # base program counter

  passed_insts = 0
-  for pkt, info in map_insts(sqtt.blob, prg.lib):
+  for pkt, info in map_insts(sqtt.blob, prg.lib, target):
    if DEBUG >= 2: print_packets([pkt])
    if info is None: continue
    if DEBUG >= 2: print(f"{' '*29}{info.inst.disasm()}")
    rocprof_inst = next(rwaves_iter[info.wave][0])
-    ref_pc = rocprof_inst.pc-rwaves_base
+    ref_pc = rocprof_inst.pc-prg.base
    # always check pc matches
    assert ref_pc == info.pc, f"pc mismatch {ref_pc}:{disasm[rocprof_inst.pc][0]} != {info.pc}:{info.inst.disasm()}"
    # special handling for s_endpgm, it marks the wave completion.
@@ -110,7 +99,8 @@ def test_rocprof_inst_traces_match(sqtt, prg, target):
  for k,v in rwaves_iter.items():
    assert len(v) == 0, f"incomplete wave {k}"

-  print(f"passed for {passed_insts} instructions across {len(rwaves)} waves scheduled on {len(rwaves_iter)} wave units")
+  if len(rwaves):
+    print(f"passed for {passed_insts} instructions across {len(rwaves)} waves scheduled on {len(rwaves_iter)} wave units")

 if __name__ == "__main__":
  import argparse, pickle, pathlib
@@ -123,7 +113,7 @@ if __name__ == "__main__":
  with open(args.profile, "rb") as f:
    data = pickle.load(f)
  sqtt_events = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"]
-  kern_events = {e.name:e for e in data if type(e).__name__ == "ProfileProgramEvent"}
+  kern_events = {e.tag:e for e in data if type(e).__name__ == "ProfileProgramEvent"}
  target = next((e for e in data if type(e).__name__ == "ProfileDeviceEvent" and e.device.startswith("AMD"))).props["gfx_target_version"]
  for e in sqtt_events:
    if args.kernel is not None and args.kernel != e.kern: continue
@@ -180,7 +180,7 @@ class TestDSPcodePatterns(unittest.TestCase):
  def test_mem_read_parsing(self):
    """Test MEM[addr].type read expression parsing."""
    # Create a mock LDS buffer
-    lds = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(16384), arg=3)
+    lds = UOp(Ops.PARAM, dtypes.uint32.ptr(16384), arg=3)
    addr = UOp.const(dtypes.uint32, 0)
    vars = {'_lds': lds, 'ADDR': addr, 'OFFSET': UOp.const(dtypes.uint32, 0)}

@@ -213,7 +213,7 @@ class TestDSPcodePatterns(unittest.TestCase):
    """Test DS_LOAD_2ADDR_B32 pcode parsing produces RETURN_DATA assignments."""
    pcode = PCODE.get(DSOp.DS_LOAD_2ADDR_B32)
    self.assertIsNotNone(pcode)
-    lds = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(16384), arg=3)
+    lds = UOp(Ops.PARAM, dtypes.uint32.ptr(16384), arg=3)
    srcs = {
      'ADDR': UOp.const(dtypes.uint32, 0),
      'OFFSET0': UOp.const(dtypes.uint32, 0),
@@ -286,7 +286,7 @@ class TestAllPcode(unittest.TestCase):
  def _make_srcs(self):
    """Create dummy source variables for pcode parsing."""
    u32, u64 = lambda v=0: UOp.const(dtypes.uint32, v), lambda v=0: UOp.const(dtypes.uint64, v)
-    lds = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(16384), arg=3)
+    lds = UOp(Ops.PARAM, dtypes.uint32.ptr(16384), arg=3)
    return {'laneId': u32(), 'laneID': u32(), 'S0': u32(), 'S1': u32(), 'S2': u32(), 'S3': u32(), 'SRC0': u32(),
            'D0': u32(), 'D1': u32(), 'DST': u32(), 'VDST': u32(), 'SDST': u32(),
            'VCC': u64(), 'VCCZ': u32(), 'EXEC': u64(), 'EXEC_LO': u32(), 'EXECZ': u32(), 'SCC': u32(),
@@ -0,0 +1,98 @@
+import unittest, ctypes
+from extra.assembly.amd.autogen.rdna4 import ins as ir4
+from extra.assembly.amd.dsl import v, s
+from extra.assembly.amd.emu import WaveState, decode_program
+from tinygrad.device import Buffer, BufferSpec
+from tinygrad.dtype import dtypes
+
+class TestRDNA4Emu(unittest.TestCase):
+  def _run(self, insts: list, sgprs: dict[int, int] = None, vgprs: dict[tuple[int, int], int] = None) -> WaveState:
+    """Run instructions and return final WaveState."""
+    # Add S_ENDPGM if not present
+    if not any(isinstance(i, ir4.SOPP) and i.op == ir4.SOPPOp.S_ENDPGM for i in insts):
+      insts = list(insts) + [ir4.SOPP(ir4.SOPPOp.S_ENDPGM, simm=0)]
+
+    # Assemble and decode
+    code = b''.join(i.to_bytes() for i in insts)
+    code_buf = (ctypes.c_uint8 * len(code)).from_buffer_copy(code)
+    code_addr = ctypes.addressof(code_buf)
+    program_raw = decode_program(code, "rdna4")
+    program = {code_addr + offset: val for offset, val in program_raw.items()}
+
+    # Setup wave state
+    st = WaveState(n_lanes=1)
+    st.pc = code_addr
+    if sgprs:
+      for idx, val in sgprs.items(): st._write_sgpr(idx, val)
+    if vgprs:
+      for (reg, lane), val in vgprs.items(): st._write_vgpr(reg, lane, val)
+
+    # Setup vmem buffer with external_ptr=0 (maps to address 0, allows any pointer access)
+    vmem_buf = Buffer('CPU', 1 << 40, dtypes.uint32, options=BufferSpec(external_ptr=0)).ensure_allocated()
+
+    # Execute
+    c_bufs = [ctypes.c_uint64(st.sgpr_buf._buf.va_addr), ctypes.c_uint64(st.vgpr_buf._buf.va_addr),
+              ctypes.c_uint64(vmem_buf._buf.va_addr), ctypes.c_uint64(0), ctypes.c_uint64(0)]
+    for _ in range(100):
+      if (pc := st.pc) == 0xFFFFFFFFFFFFFFFF or pc not in program: break
+      _, fxn, globals_list, _ = program[pc]
+      fxn(*[c_bufs[g] for g in globals_list])
+    return st
+
+  def test_vopd_dual_mov(self):
+    """Test VOPD with two V_DUAL_MOV_B32 operations: v[1]=s[1], v[2]=s[2]."""
+    insts = [ir4.VOPD(ir4.VOPDOp.V_DUAL_MOV_B32, ir4.VOPDOp.V_DUAL_MOV_B32,
+                      vdstx=v[1], vdsty=v[2], srcx0=s[1], srcy0=s[2], vsrcx1=v[0], vsrcy1=v[0])]
+    st = self._run(insts, sgprs={1: 0x40e00000, 2: 0x41100000})  # 7.0f, 9.0f
+    self.assertEqual(st._read_vgpr(1, 0), 0x40e00000)  # v[1] = 7.0
+    self.assertEqual(st._read_vgpr(2, 0), 0x41100000)  # v[2] = 9.0
+
+  def test_vopd_dual_mov_after_other_vopd(self):
+    """Test VOPD reuse: first VOPD(v[3]=0, v[0]=?), then VOPD(v[1]=s[1], v[2]=s[2])."""
+    # This matches the BEAM kernel sequence that fails
+    insts = [
+      ir4.VOPD(ir4.VOPDOp.V_DUAL_MOV_B32, ir4.VOPDOp.V_DUAL_MOV_B32,
+               vdstx=v[3], vdsty=v[0], srcx0=0, srcy0=s[0], vsrcx1=v[0], vsrcy1=v[0]),  # v[3]=0, v[0]=s[0]
+      ir4.VOPD(ir4.VOPDOp.V_DUAL_MOV_B32, ir4.VOPDOp.V_DUAL_MOV_B32,
+               vdstx=v[1], vdsty=v[2], srcx0=s[1], srcy0=s[2], vsrcx1=v[0], vsrcy1=v[0]),  # v[1]=s[1], v[2]=s[2]
+    ]
+    st = self._run(insts, sgprs={0: 0x40a00000, 1: 0x40e00000, 2: 0x41100000})  # 5.0f, 7.0f, 9.0f
+    self.assertEqual(st._read_vgpr(1, 0), 0x40e00000)  # v[1] = 7.0
+    self.assertEqual(st._read_vgpr(2, 0), 0x41100000)  # v[2] = 9.0
+
+  def test_vopd_with_s_add_f32_sequence(self):
+    """Test full BEAM kernel sequence: s_add_f32 then VOPD."""
+    # This is the exact sequence from the failing BEAM kernel
+    insts = [
+      ir4.SOP2(ir4.SOP2Op.S_ADD_F32, sdst=s[0], ssrc0=s[0], ssrc1=s[8]),   # s[0] = s[0] + s[8]
+      ir4.SOP2(ir4.SOP2Op.S_ADD_F32, sdst=s[1], ssrc0=s[1], ssrc1=s[9]),   # s[1] = s[1] + s[9]
+      ir4.SOP2(ir4.SOP2Op.S_ADD_F32, sdst=s[2], ssrc0=s[2], ssrc1=s[10]),  # s[2] = s[2] + s[10]
+      ir4.VOPD(ir4.VOPDOp.V_DUAL_MOV_B32, ir4.VOPDOp.V_DUAL_MOV_B32,
+               vdstx=v[3], vdsty=v[0], srcx0=0, srcy0=s[0], vsrcx1=v[0], vsrcy1=v[0]),
+      ir4.VOPD(ir4.VOPDOp.V_DUAL_MOV_B32, ir4.VOPDOp.V_DUAL_MOV_B32,
+               vdstx=v[1], vdsty=v[2], srcx0=s[1], srcy0=s[2], vsrcx1=v[0], vsrcy1=v[0]),
+    ]
+    # Input: s[0:2] = [1,2,3], s[8:10] = [4,5,6]
+    # After s_add_f32: s[0:2] = [5,7,9]
+    st = self._run(insts, sgprs={0: 0x3f800000, 1: 0x40000000, 2: 0x40400000,  # 1.0, 2.0, 3.0
+                                  8: 0x40800000, 9: 0x40a00000, 10: 0x40c00000})  # 4.0, 5.0, 6.0
+    self.assertEqual(st._read_vgpr(1, 0), 0x40e00000)  # v[1] = 7.0
+    self.assertEqual(st._read_vgpr(2, 0), 0x41100000)  # v[2] = 9.0
+
+  def test_s_mov_b32_then_vopd(self):
+    """Test s_mov_b32 followed by VOPD - simulates BEAM kernel sequence."""
+    # Use s_mov_b32 with SGPR source (copy from pre-initialized SGPRs)
+    # s[10:12] will have values set by test harness, copy to s[0:2], then VOPD to VGPRs
+    insts = [
+      ir4.SOP1(ir4.SOP1Op.S_MOV_B32, sdst=s[0], ssrc0=s[10]),  # s[0] = s[10]
+      ir4.SOP1(ir4.SOP1Op.S_MOV_B32, sdst=s[1], ssrc0=s[11]),  # s[1] = s[11]
+      ir4.SOP1(ir4.SOP1Op.S_MOV_B32, sdst=s[2], ssrc0=s[12]),  # s[2] = s[12]
+      ir4.VOPD(ir4.VOPDOp.V_DUAL_MOV_B32, ir4.VOPDOp.V_DUAL_MOV_B32,
+               vdstx=v[1], vdsty=v[2], srcx0=s[1], srcy0=s[2], vsrcx1=v[0], vsrcy1=v[0]),
+    ]
+    st = self._run(insts, sgprs={10: 0x40a00000, 11: 0x40e00000, 12: 0x41100000})  # 5.0, 7.0, 9.0
+    self.assertEqual(st._read_vgpr(1, 0), 0x40e00000)  # v[1] = 7.0
+    self.assertEqual(st._read_vgpr(2, 0), 0x41100000)  # v[2] = 9.0
+
+if __name__ == '__main__':
+  unittest.main()
@@ -203,12 +203,12 @@ class SQTTExamplesTestBase(unittest.TestCase):
 class TestSQTTExamplesRDNA3(SQTTExamplesTestBase):
  target = "gfx1100"
  expected = {
-    "profile_empty_run_0": [1803, 1908, 1928, 1979, 2006, 1912],
-    "profile_empty_run_1": [1803, 1908, 1928, 1979, 2006, 1912],
-    "profile_gemm_run_0": [2531, 1844, 1864, 1915, 1942, 1848, 3074, 1919, 1939, 1990, 2017, 1923, 19026, 1919, 1939, 1990, 2017, 1929],
-    "profile_gemm_run_1": [2554, 1844, 1864, 1915, 1942, 1848, 3084, 1919, 1939, 1990, 2017, 1923, 19010, 1919, 1939, 1990, 2017, 1923],
-    "profile_plus_run_0": [1900, 1908, 1928, 1979, 2006, 1912],
-    "profile_plus_run_1": [1856, 1908, 1928, 1979, 2006, 1912],
+    "profile_empty_run_0": [1844, 1885, 1905, 1956, 1983, 1889],
+    "profile_empty_run_1": [1780, 1885, 1905, 1956, 1983, 1889],
+    "profile_gemm_run_0": [2656, 2025, 2045, 2096, 2123, 2029, 3183, 2019, 2039, 2090, 2117, 2023, 19119, 2013, 2033, 2084, 2111, 2017],
+    "profile_gemm_run_1": [2662, 2025, 2045, 2096, 2123, 2029, 3179, 2019, 2039, 2090, 2117, 2023, 19113, 2071, 2091, 2142, 2169, 2075],
+    "profile_plus_run_0": [1886, 2013, 2033, 2084, 2111, 2017],
+    "profile_plus_run_1": [1988, 2071, 2091, 2142, 2169, 2075],
  }

 class TestSQTTExamplesRDNA4(SQTTExamplesTestBase): target = "gfx1200"
@@ -1,31 +0,0 @@
-import random
-from extra.optimization.helpers import load_worlds, ast_str_to_lin
-from tinygrad.codegen.opt.search import actions
-from tinygrad.codegen.opt.kernel import Kernel
-from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
-from tinygrad.helpers import tqdm
-
-tactions = set()
-def test_rebuild(lin):
-  linr = Kernel(lin.ast)
-  for o in lin.applied_opts:
-    assert o in actions, f"{o} is not in actions"
-    tactions.add(o)
-    linr.apply_opt(o)
-
-  assert len(lin.sts) == len(linr.sts)
-  for st1,st2 in zip(lin.sts, linr.sts):
-    assert st1 == st2, f"{st1} != {st2}"
-
-if __name__ == "__main__":
-  ast_strs = load_worlds(False, False, False)
-  random.shuffle(ast_strs)
-  ast_strs = ast_strs[:2000]
-  for ast_str in tqdm(ast_strs):
-    lin = ast_str_to_lin(ast_str)
-    #if not lin.apply_tensor_cores():
-    lin.apply_opts(hand_coded_optimizations(lin))
-    test_rebuild(lin)
-
-  print(len(tactions), len(actions))
-  print(sorted(list(tactions)))
@@ -1,76 +0,0 @@
-import os
-import numpy as np
-import math, random
-from tinygrad.tensor import Tensor
-from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
-from tinygrad.codegen.opt.search import actions, bufs_from_lin, get_kernel_actions
-from tinygrad.nn.optim import Adam
-from extra.optimization.extract_policynet import PolicyNet
-from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats, time_linearizer
-
-if __name__ == "__main__":
-  net = PolicyNet()
-  if os.path.isfile("/tmp/policynet.safetensors"): load_state_dict(net, safe_load("/tmp/policynet.safetensors"))
-  optim = Adam(get_parameters(net))
-
-  ast_strs = load_worlds()
-
-  # select a world
-  all_feats, all_acts, all_rews = [], [], []
-  while 1:
-    Tensor.training = False
-    lin = ast_str_to_lin(random.choice(ast_strs))
-    rawbufs = bufs_from_lin(lin)
-    tm = last_tm = base_tm = time_linearizer(lin, rawbufs)
-
-    # take actions
-    feats, acts, rews = [], [], []
-    while 1:
-      feat = lin_to_feats(lin)
-      feats.append(feat)
-      probs = net(Tensor([feat])).exp()[0].numpy()
-
-      # mask valid actions
-      valid_action_mask = np.zeros((len(actions)+1), dtype=np.float32)
-      for x in get_kernel_actions(lin): valid_action_mask[x] = 1
-      probs *= valid_action_mask
-      probs /= sum(probs)
-
-      act = np.random.choice(len(probs), p=probs)
-      acts.append(act)
-      if act == 0:
-        rews.append(0)
-        break
-      try:
-        lin.apply_opt(actions[act-1])
-        tm = time_linearizer(lin, rawbufs)
-        if math.isinf(tm): raise Exception("failed")
-        rews.append(((last_tm-tm)/base_tm))
-        last_tm = tm
-      except Exception:
-        rews.append(-0.5)
-        break
-      #print(f"{tm*1e6:10.2f}", lin.colored_shape())
-
-    assert len(feats) == len(acts) and len(acts) == len(rews)
-    #print(rews)
-    print(f"***** EPISODE {len(rews)} steps, {sum(rews):5.2f} reward, {base_tm*1e6:12.2f} -> {tm*1e6:12.2f} : {lin.colored_shape()}")
-    all_feats += feats
-    all_acts += acts
-    # rewards to go
-    for i in range(len(rews)-2, -1, -1): rews[i] += rews[i+1]
-    all_rews += rews
-
-    BS = 32
-    if len(all_feats) >= BS:
-      Tensor.training = True
-      x = Tensor(all_feats[:BS])
-      mask = np.zeros((BS, len(actions)+1), dtype=np.float32)
-      mask[range(BS), all_acts[:BS]] = all_rews[:BS]
-      loss = -(net(x) * Tensor(mask)).mean()
-      optim.zero_grad()
-      loss.backward()
-      optim.step()
-      all_feats = all_feats[BS:]
-      all_acts = all_acts[BS:]
-      all_rews = all_rews[BS:]
@@ -1,32 +0,0 @@
-from typing import List, Tuple
-from tinygrad.codegen.opt.kernel import Kernel
-from tinygrad.codegen.opt.search import get_kernel_actions, actions
-
-_net = None
-def beam_q_estimate(beam:List[Tuple[Kernel, float]]) -> List[Tuple[Kernel, float]]:
-  global _net
-  if _net is None:
-    from tinygrad.nn.state import load_state_dict, safe_load
-    from extra.optimization.pretrain_valuenet import ValueNet
-    _net = ValueNet(1021+len(actions), 2)
-    load_state_dict(_net, safe_load("/tmp/qnet.safetensors"), verbose=False)
-  from tinygrad.tensor import Tensor
-  from tinygrad.helpers import Context
-  from extra.optimization.helpers import lin_to_feats
-  import numpy as np
-  feats = []
-  lins = []
-  base_tms = []
-  for lin,tm in beam:
-    lin_feats = lin_to_feats(lin)
-    for a,v in get_kernel_actions(lin, include_0=False).items():
-      acts = np.zeros(len(actions))
-      acts[a-1] = 1.0
-      feats.append(np.concatenate([lin_feats, acts]))
-      lins.append(v)
-      base_tms.append(tm)
-  with Context(BEAM=0):
-    with Tensor.train(False):
-      preds = _net(Tensor(feats)).numpy()
-  pred_time = np.array(base_tms) / np.exp(preds[:, 0])
-  return sorted(zip(lins, pred_time), key=lambda x: x[1])
@@ -1,34 +0,0 @@
-import argparse
-from extra.optimization.helpers import ast_str_to_lin, time_linearizer
-
-from tinygrad import dtypes
-from tinygrad.helpers import BEAM, getenv
-from tinygrad.device import Device, Compiled
-from tinygrad.codegen.opt.kernel import Kernel
-from tinygrad.codegen.opt.search import beam_search, bufs_from_lin
-
-
-if __name__ == '__main__':
-  parser = argparse.ArgumentParser(description="Run a search for the optimal opts for a kernel", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
-  parser.add_argument("--ast", type=str, default=None, help="the ast for the kernel to be optimized")
-  parser.add_argument("--file", type=str, default=None, help="a file containing asts to be optimized, one per line")
-  args = parser.parse_args()
-
-  device: Compiled = Device[Device.DEFAULT]
-  print(f"optimizing for {Device.DEFAULT}")
-
-  if args.ast is not None:
-    ast_strs = [args.ast]
-  elif args.file is not None:
-    with open(args.file, 'r') as file:
-      ast_strs = file.readlines()
-
-  for i, ast_str in enumerate(ast_strs):
-    print(f"optimizing {i}/{len(ast_strs)}\nast={ast_str}")
-    lin = ast_str_to_lin(ast_str, opts=device.renderer)
-    rawbufs = bufs_from_lin(lin)
-    lin = beam_search(lin, rawbufs, getenv("BEAM", 8), bool(getenv("BEAM_ESTIMATE", 1)))
-
-    tm = time_linearizer(lin, rawbufs, allow_test_size=False, cnt=10)
-    print(f"final time {tm*1e6:9.0f} us: {lin.colored_shape()}")
-    print(lin.applied_opts)
@@ -1,19 +0,0 @@
-import unittest
-
-from extra.optimization.helpers import load_worlds
-
-class TestKernelDataset(unittest.TestCase):
-  def test_load_worlds_filters(self):
-    all_kernels = load_worlds(filter_reduce=False, filter_noimage=False, filter_novariable=False)
-
-    reduce_kernels = load_worlds(filter_reduce=True, filter_noimage=False, filter_novariable=False)
-    self.assertGreater(len(all_kernels), len(reduce_kernels))
-
-    image_kernels = load_worlds(filter_reduce=False, filter_noimage=True, filter_novariable=False)
-    self.assertGreater(len(all_kernels), len(image_kernels))
-
-    variable_kernels = load_worlds(filter_reduce=False, filter_noimage=False, filter_novariable=True)
-    self.assertGreater(len(all_kernels), len(variable_kernels))
-
-if __name__ == '__main__':
-  unittest.main()
@@ -1,67 +0,0 @@
-import numpy as np
-import math
-import random
-np.set_printoptions(suppress=True)
-from copy import deepcopy
-from tinygrad.helpers import getenv, colored
-from tinygrad.tensor import Tensor
-from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
-from tinygrad.codegen.opt.search import bufs_from_lin, actions, get_kernel_actions
-from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
-from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats, time_linearizer
-from extra.optimization.extract_policynet import PolicyNet
-from extra.optimization.pretrain_valuenet import ValueNet
-
-VALUE = getenv("VALUE")
-
-if __name__ == "__main__":
-  if VALUE:
-    net = ValueNet()
-    load_state_dict(net, safe_load("/tmp/valuenet.safetensors"))
-  else:
-    net = PolicyNet()
-    load_state_dict(net, safe_load("/tmp/policynet.safetensors"))
-
-  ast_strs = load_worlds()
-
-  # real randomness
-  random.seed()
-  random.shuffle(ast_strs)
-
-  wins = 0
-  for ep_num,ast_str in enumerate(ast_strs):
-    print("\nEPISODE", ep_num, f"win {wins*100/max(1,ep_num):.2f}%")
-    lin = ast_str_to_lin(ast_str)
-    rawbufs = bufs_from_lin(lin)
-
-    linhc = deepcopy(lin)
-    linhc.applied_opts(hand_coded_optimizations(linhc))
-    tmhc = time_linearizer(linhc, rawbufs)
-    print(f"{tmhc*1e6:10.2f}     HC    ", linhc.colored_shape())
-
-    pred_time = float('nan')
-    tm = float('inf')
-    while 1:
-      if VALUE:
-        acts,feats = [], []
-        for k,v in get_kernel_actions(lin).items():
-          acts.append(k)
-          feats.append(lin_to_feats(v))
-        preds = net(Tensor(feats))
-        pred_time = math.exp(preds.numpy().min())
-        act = acts[preds.numpy().argmin()]
-      else:
-        probs = net(Tensor([lin_to_feats(lin)]))
-        dist = probs.exp().numpy()
-        act = dist.argmax()
-      if act == 0: break
-      try:
-        lin.apply_opt(actions[act-1])
-      except Exception:
-        print("FAILED")
-        break
-      tm = time_linearizer(lin, rawbufs)
-      print(f"{tm*1e6:10.2f} {pred_time*1e6:10.2f}", lin.colored_shape())
-
-    print(f"{colored('BEAT', 'green') if tm < tmhc else colored('lost', 'red')} hand coded {tmhc/tm:5.2f}x")
-    wins += int(tm < tmhc)
@@ -1,21 +0,0 @@
-from extra.optimization.helpers import load_worlds, ast_str_to_lin, time_linearizer
-from tinygrad.codegen.opt.search import bufs_from_lin, get_kernel_actions
-
-if __name__ == "__main__":
-  ast_strs = load_worlds()
-  for i, ast_str in enumerate(ast_strs):
-    lin = ast_str_to_lin(ast_str)
-    rawbufs = bufs_from_lin(lin)
-    test_tm = time_linearizer(lin, rawbufs)
-    if test_tm < 1e-2: continue
-    print(f"EXAMPLE {i}")
-    acted_lins = get_kernel_actions(lin)
-    ok_avg, short_avg = 0, 0
-    for k,v in acted_lins.items():
-      tm1 = time_linearizer(v, rawbufs)
-      tm2 = time_linearizer(v, rawbufs)
-      tm3 = time_linearizer(v, rawbufs, False)
-      print(v.colored_shape(50), f"{tm1*1e3:10.2f} {tm2*1e3:10.2f} {tm3*1e3:10.2f}  :  {((tm1-tm2)/tm1)*100:5.2f}%  vs  {((tm1-tm3)/tm1)*100:5.2f}%")
-      ok_avg += (tm1-tm2)/tm1
-      short_avg += (tm1-tm3)/tm1
-    print(f"{ok_avg/len(acted_lins)*100:5.2f}% vs {short_avg/len(acted_lins)*100:5.2f}%")
@@ -2,7 +2,7 @@ import math

 from tinygrad import Tensor, dtypes
 from tinygrad.helpers import DEBUG
-from tinygrad.uop.ops import UOp
+from tinygrad.uop.ops import UOp, Ops

 from extra.thunder.tiny.tk import WARP_THREADS
 from extra.thunder.tiny.tk.kernel import Kernel
@@ -43,11 +43,12 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
  B_local = B // num_devices
  if DEBUG >= 2: print(f"Flash Attention {B=} {B_local=} {N=} {H=} {D=} {H_KV=} {GROUP_SIZE=}")

-  def custom_forward(ou:UOp, l_vecu:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp) -> UOp:
+  def _custom_forward_impl(ou:UOp, l_vecu:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp|None) -> UOp:
    with Kernel("fa_custom_forward", (H, N // (Q_BLOCK_SIZE*NUM_WORKERS), B_local), NUM_WORKERS * WARP_THREADS) as ker:
      warp = ker.warp

-      o, q, k, v, mask, l_vec = GL(ou, ker), GL(qu, ker), GL(ku, ker), GL(vu, ker), GL(masku, ker), GL(l_vecu, ker)
+      o, q, k, v, l_vec = GL(ou, ker), GL(qu, ker), GL(ku, ker), GL(vu, ker), GL(l_vecu, ker)
+      mask = GL(masku, ker) if masku is not None else None

      head = ker.blockIdx_x
      head_kv = head // GROUP_SIZE
@@ -86,7 +87,8 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
      q_reg = warp.copy(q_reg, q_reg_fl)
      q_reg_transposed = warp.transpose(q_reg_transposed, q_reg)

-      for kv_idx in ker.range(N // KV_BLOCK_SIZE):
+      num_kv_blocks = (q_seq + 1) if is_causal else (N // KV_BLOCK_SIZE)
+      for kv_idx in ker.range(num_kv_blocks):
        k_smem = warp.load(k_smem, k, (), (batch, kv_idx, head_kv, 0), axis=1)
        v_smem = warp.load(v_smem, v, (), (batch, kv_idx, head_kv, 0), axis=1)

@@ -99,9 +101,16 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
        att_block = warp.mma_AtB(att_block, k_reg_transposed, q_reg_transposed)

        # apply attention mask
-        mask_reg = warp.load(mask_reg, mask, (), (batch, 0, q_seq, kv_idx), axis=2)
-        mask_reg_transposed = warp.transpose(mask_reg_transposed, mask_reg)
-        att_block += mask_reg_transposed
+        if is_causal:
+          bs_rows, bs_cols, bs_stride = att_block.base_shape.rows, att_block.base_shape.cols, att_block.base_shape.stride
+          q_base = q_seq * Q_BLOCK_SIZE + (warp.laneid % bs_cols)
+          kv_base = kv_idx * KV_BLOCK_SIZE + (warp.laneid // bs_cols) * bs_stride
+          att_block = warp.map(att_block,
+            lambda x, idx: ((kv_base + idx[0]*bs_rows + idx[2]) > (q_base + idx[1]*bs_cols)).alu(Ops.WHERE, UOp.ufix(x._uop, -math.inf), x))
+        elif mask is not None:
+          mask_reg = warp.load(mask_reg, mask, (), (batch, 0, q_seq, kv_idx), axis=2)
+          mask_reg_transposed = warp.transpose(mask_reg_transposed, mask_reg)
+          att_block += mask_reg_transposed

        # softmax
        max_vec_last = warp.copy(max_vec_last.after(kv_idx), max_vec)
@@ -141,11 +150,18 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False

      return ker.finish()

-  def custom_backward_q(dqu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp, l_vecu:UOp, delta_vecu:UOp) -> UOp:
+  def custom_forward_causal(ou:UOp, l_vecu:UOp, qu:UOp, ku:UOp, vu:UOp) -> UOp:
+    return _custom_forward_impl(ou, l_vecu, qu, ku, vu, None)
+
+  def custom_forward_masked(ou:UOp, l_vecu:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp) -> UOp:
+    return _custom_forward_impl(ou, l_vecu, qu, ku, vu, masku)
+
+  def _custom_backward_q_impl(dqu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp|None, l_vecu:UOp, delta_vecu:UOp) -> UOp:
    with Kernel("fa_custom_backward_q", (H, N // (Q_BLOCK_SIZE*NUM_WORKERS), B_local), NUM_WORKERS * WARP_THREADS) as ker:
      warp = ker.warp

-      dq, do, q, k, v, mask = GL(dqu, ker), GL(dou, ker), GL(qu, ker), GL(ku, ker), GL(vu, ker), GL(masku, ker)
+      dq, do, q, k, v = GL(dqu, ker), GL(dou, ker), GL(qu, ker), GL(ku, ker), GL(vu, ker)
+      mask = GL(masku, ker) if masku is not None else None
      l_vec, delta_vec = GL(l_vecu, ker), GL(delta_vecu, ker)

      head = ker.blockIdx_x
@@ -194,7 +210,8 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
      l_vec_reg *= 1.0 / math.log(2)
      delta_vec_reg = warp.load(delta_vec_reg, delta_vec, (), (batch, head, 0, q_seq), axis=2)

-      for kv_idx in ker.range(N // KV_BLOCK_SIZE):
+      num_kv_blocks = (q_seq + 1) if is_causal else (N // KV_BLOCK_SIZE)
+      for kv_idx in ker.range(num_kv_blocks):
        k_smem = warp.load(k_smem, k, (), (batch, kv_idx, head_kv, 0), axis=1)
        v_smem = warp.load(v_smem, v, (), (batch, kv_idx, head_kv, 0), axis=1)

@@ -209,9 +226,16 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
        att_block = warp.mma_AtB(att_block, k_reg_t, q_reg_t)

        # apply attention mask
-        mask_reg = warp.load(mask_reg, mask, (), (batch, 0, q_seq, kv_idx), axis=2)
-        mask_reg_transposed = warp.transpose(mask_reg_transposed, mask_reg)
-        att_block += mask_reg_transposed
+        if is_causal:
+          bs_rows, bs_cols, bs_stride = att_block.base_shape.rows, att_block.base_shape.cols, att_block.base_shape.stride
+          q_base = q_seq * Q_BLOCK_SIZE + (warp.laneid % bs_cols)
+          kv_base = kv_idx * KV_BLOCK_SIZE + (warp.laneid // bs_cols) * bs_stride
+          att_block = warp.map(att_block,
+            lambda x, idx: ((kv_base + idx[0]*bs_rows + idx[2]) > (q_base + idx[1]*bs_cols)).alu(Ops.WHERE, UOp.ufix(x._uop, -math.inf), x))
+        elif mask is not None:
+          mask_reg = warp.load(mask_reg, mask, (), (batch, 0, q_seq, kv_idx), axis=2)
+          mask_reg_transposed = warp.transpose(mask_reg_transposed, mask_reg)
+          att_block += mask_reg_transposed

        att_block -= l_vec_reg
        att_block = att_block.exp2()
@@ -231,11 +255,18 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False

      return ker.finish()

-  def custom_backward_kv(dku:UOp, dvu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp, l_vecu:UOp, delta_vecu:UOp):
+  def custom_backward_q_causal(dqu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, l_vecu:UOp, delta_vecu:UOp) -> UOp:
+    return _custom_backward_q_impl(dqu, dou, qu, ku, vu, None, l_vecu, delta_vecu)
+
+  def custom_backward_q_masked(dqu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp, l_vecu:UOp, delta_vecu:UOp) -> UOp:
+    return _custom_backward_q_impl(dqu, dou, qu, ku, vu, masku, l_vecu, delta_vecu)
+
+  def _custom_backward_kv_impl(dku:UOp, dvu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp|None, l_vecu:UOp, delta_vecu:UOp):
    with Kernel("fa_custom_backward_kv", (H_KV, N // (KV_BLOCK_SIZE*NUM_WORKERS), B_local), NUM_WORKERS * WARP_THREADS) as ker:
      warp = ker.warp

-      dk, dv, do, q, k, v, mask = GL(dku, ker), GL(dvu, ker), GL(dou, ker), GL(qu, ker), GL(ku, ker), GL(vu, ker), GL(masku, ker)
+      dk, dv, do, q, k, v = GL(dku, ker), GL(dvu, ker), GL(dou, ker), GL(qu, ker), GL(ku, ker), GL(vu, ker)
+      mask = GL(masku, ker) if masku is not None else None
      l_vec, delta_vec = GL(l_vecu, ker), GL(delta_vecu, ker)

      head_kv = ker.blockIdx_x
@@ -302,9 +333,16 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
          att_block *= (1.0 / math.sqrt(D)) * (1.0 / math.log(2))

          # apply attention mask
-          mask_reg = warp.load(mask_reg, mask, (), (batch, 0, q_idx, kv_seq), axis=2)
-          mask_reg_transposed = warp.transpose(mask_reg_transposed, mask_reg)
-          att_block += mask_reg_transposed
+          if is_causal:
+            bs_rows, bs_cols, bs_stride = att_block.base_shape.rows, att_block.base_shape.cols, att_block.base_shape.stride
+            q_base = q_idx * Q_BLOCK_SIZE + (warp.laneid % bs_cols)
+            kv_base = kv_seq * KV_BLOCK_SIZE + (warp.laneid // bs_cols) * bs_stride
+            att_block = warp.map(att_block,
+              lambda x, idx: ((kv_base + idx[0]*bs_rows + idx[2]) > (q_base + idx[1]*bs_cols)).alu(Ops.WHERE, UOp.ufix(x._uop, -math.inf), x))
+          elif mask is not None:
+            mask_reg = warp.load(mask_reg, mask, (), (batch, 0, q_idx, kv_seq), axis=2)
+            mask_reg_transposed = warp.transpose(mask_reg_transposed, mask_reg)
+            att_block += mask_reg_transposed

          att_block -= l_vec_reg
          att_block = att_block.exp2()
@@ -336,24 +374,31 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False

      return ker.finish(2)

+  def custom_backward_kv_causal(dku:UOp, dvu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, l_vecu:UOp, delta_vecu:UOp):
+    return _custom_backward_kv_impl(dku, dvu, dou, qu, ku, vu, None, l_vecu, delta_vecu)
+
+  def custom_backward_kv_masked(dku:UOp, dvu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp, l_vecu:UOp, delta_vecu:UOp):
+    return _custom_backward_kv_impl(dku, dvu, dou, qu, ku, vu, masku, l_vecu, delta_vecu)
+
  single_device = xq.device[0] if isinstance(xq.device, tuple) else xq.device

  if is_causal:
    if attn_mask is not None: raise RuntimeError("cannot set attn_mask when is_causal=True")
-    attn_mask = Tensor.ones((B, 1, N, N), requires_grad=False, device=single_device, dtype=dtypes.bool).tril()
-  if attn_mask is not None:
+  elif attn_mask is not None:
    if attn_mask.dtype == dtypes.bool: attn_mask = attn_mask.where(0, -float("inf"))
+    if attn_mask.shape != (B, 1, N, N):
+      attn_mask = attn_mask.expand(B, 1, N, N)
+    if isinstance(xq.device, tuple) and not isinstance(attn_mask.device, tuple):
+      attn_mask = attn_mask.shard(xq.device, axis=0)
  else:
    attn_mask = Tensor.zeros((B, 1, N, N), requires_grad=False, device=single_device, dtype=dtypes.float32)
-  if attn_mask.shape != (B, 1, N, N):
-    attn_mask = attn_mask.expand(B, 1, N, N)
-  if isinstance(xq.device, tuple) and not isinstance(attn_mask.device, tuple):
-    attn_mask = attn_mask.shard(xq.device, axis=0)
+    if isinstance(xq.device, tuple):
+      attn_mask = attn_mask.shard(xq.device, axis=0)

  attn = _sharded_empty_like(xq, axis=0)
  l_vec = _sharded_empty((B, H, 1, N), xq, axis=0)

-  def grad(gradu:UOp, _) -> tuple[None, None, UOp, UOp, UOp, None]:
+  def grad_causal(gradu:UOp, _) -> tuple[None, None, UOp, UOp, UOp]:
    grad = Tensor(gradu, device=gradu.device)
    grad_q = _sharded_empty_like(xq, axis=0)
    grad_k = _sharded_empty_like(xk, axis=0)
@@ -361,11 +406,26 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False

    delta_vec = (grad * attn).sum(-1, dtype=dtypes.float32).transpose(1, 2).unsqueeze(-2).detach()

-    grad_q = Tensor.custom_kernel(grad_q, grad, xq, xk, xv, attn_mask, l_vec, delta_vec, fxn=custom_backward_q)[0]
-    grad_k, grad_v = Tensor.custom_kernel(grad_k, grad_v, grad, xq, xk, xv, attn_mask, l_vec, delta_vec, fxn=custom_backward_kv)[:2]
+    grad_q = Tensor.custom_kernel(grad_q, grad, xq, xk, xv, l_vec, delta_vec, fxn=custom_backward_q_causal)[0]
+    grad_k, grad_v = Tensor.custom_kernel(grad_k, grad_v, grad, xq, xk, xv, l_vec, delta_vec, fxn=custom_backward_kv_causal)[:2]
+    return (None, None, grad_q.uop, grad_k.uop, grad_v.uop)
+
+  def grad_masked(gradu:UOp, _) -> tuple[None, None, UOp, UOp, UOp, None]:
+    grad = Tensor(gradu, device=gradu.device)
+    grad_q = _sharded_empty_like(xq, axis=0)
+    grad_k = _sharded_empty_like(xk, axis=0)
+    grad_v = _sharded_empty_like(xv, axis=0)
+
+    delta_vec = (grad * attn).sum(-1, dtype=dtypes.float32).transpose(1, 2).unsqueeze(-2).detach()
+
+    grad_q = Tensor.custom_kernel(grad_q, grad, xq, xk, xv, attn_mask, l_vec, delta_vec, fxn=custom_backward_q_masked)[0]
+    grad_k, grad_v = Tensor.custom_kernel(grad_k, grad_v, grad, xq, xk, xv, attn_mask, l_vec, delta_vec, fxn=custom_backward_kv_masked)[:2]
    return (None, None, grad_q.uop, grad_k.uop, grad_v.uop, None)

-  attn, l_vec = Tensor.custom_kernel(attn, l_vec, xq, xk, xv, attn_mask, fxn=custom_forward, grad_fxn=grad)[:2]
+  if is_causal:
+    attn, l_vec = Tensor.custom_kernel(attn, l_vec, xq, xk, xv, fxn=custom_forward_causal, grad_fxn=grad_causal)[:2]
+  else:
+    attn, l_vec = Tensor.custom_kernel(attn, l_vec, xq, xk, xv, attn_mask, fxn=custom_forward_masked, grad_fxn=grad_masked)[:2]
  attn_ = attn[:, :N_, :, :D_]

  return attn_.transpose(1, 2).cast(odtype)
@@ -1,43 +0,0 @@
-import random
-from tinygrad.helpers import getenv
-from tinygrad.codegen.opt.search import beam_search, bufs_from_lin
-from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
-from extra.optimization.helpers import load_worlds, ast_str_to_lin, time_linearizer
-
-def optimize_kernel(k):
-  # TODO: update this
-  return hand_coded_optimizations(k)
-
-if __name__ == '__main__':
-  hcopt_wins = beam_wins = tie = 0
-  hcopt_total = beam_total = 0.0
-
-  worlds = load_worlds(filter_reduce=False, filter_noimage=True, filter_novariable=False)
-  random.seed(0)
-  random.shuffle(worlds)
-
-  for world in worlds[:500]:
-    k = ast_str_to_lin(world)
-    rawbufs = bufs_from_lin(k)
-
-    k_hcopt = k.copy()
-    k_hcopt.apply_opts(optimize_kernel(k_hcopt))
-    k_beam = beam_search(k.copy(), rawbufs, getenv("BEAM", 2))
-
-    disable_cache = bool(getenv("NOCACHE", 0))
-    t_hcopt = time_linearizer(k_hcopt, rawbufs, allow_test_size=False, cnt=10, disable_cache=disable_cache, clear_l2=True) * 1e6
-    t_beam = time_linearizer(k_beam, rawbufs, allow_test_size=False, cnt=10, disable_cache=disable_cache, clear_l2=True) * 1e6
-
-    if t_hcopt == t_beam: tie += 1
-    elif t_hcopt < t_beam: hcopt_wins += 1
-    else: beam_wins += 1
-    hcopt_total += t_hcopt
-    beam_total += t_beam
-
-    print(f"{t_hcopt=:5.2f} {k_hcopt.applied_opts=}")
-    print("")
-    print(f"{t_beam=:5.2f} {k_beam.applied_opts=}")
-    print("*"*20)
-
-  print(f"{hcopt_wins=}, {beam_wins=}, {tie=}")
-  print(f"{hcopt_total=:.2f}, {beam_total=:.2f}")
@@ -11,7 +11,7 @@ from tinygrad.dtype import ImageDType, Invalid
 # PYTHONPATH="." DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx

 def vision_conv_143():
-  c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((16, 1024, 4)), (), 0)
+  c0 = UOp(Ops.PARAM, dtypes.imageh((16, 1024, 4)), (), 0)
  c2 = UOp.range(32, 3, AxisType.LOOP)
  c5 = UOp.range(128, 4, AxisType.LOOP)
  c8 = UOp.range(16, 2, AxisType.LOOP)
@@ -21,13 +21,13 @@ def vision_conv_143():
  c26 = UOp.range(7, 1, AxisType.REDUCE)
  c27 = c2*2+c26
  c32 = ((c27<3)!=True)&(c27<67)
-  c34 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 1024, 4)), (), 1)
+  c34 = UOp(Ops.PARAM, dtypes.imageh((32, 1024, 4)), (), 1)
  c38 = c5//2
  c45 = (c32&c24).where((c27*64+c38+c17*4096+-12480), UOp.const(dtypes.index, Invalid))
  c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
-  c49 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((64, 49, 4)), (), 2)
+  c49 = UOp(Ops.PARAM, dtypes.imageh((64, 49, 4)), (), 2)
  c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
-  c63 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(128), (), 3)
+  c63 = UOp(Ops.PARAM, dtypes.float.ptr(128), (), 3)
  c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5)
  c67 = c0.index((c2*128+c5+c8*4096), ptr=True).store(c65).end(c8, c2, c5)

@@ -37,7 +37,7 @@ def vision_conv_143():
  return c67.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))

 def vision_conv_153():
-  c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((8, 1024, 4)), (), 0)
+  c0 = UOp(Ops.PARAM, dtypes.imageh((8, 1024, 4)), (), 0)
  c2 = UOp.range(16, 3, AxisType.LOOP)
  c5 = UOp.range(256, 4, AxisType.LOOP)
  c8 = UOp.range(8, 2, AxisType.LOOP)
@@ -47,13 +47,13 @@ def vision_conv_153():
  c26 = UOp.range(7, 1, AxisType.REDUCE)
  c27 = c2*2+c26
  c32 = ((c27<3)!=True)&(c27<35)
-  c34 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((16, 1024, 4)), (), 1)
+  c34 = UOp(Ops.PARAM, dtypes.imageh((16, 1024, 4)), (), 1)
  c38 = c5//2
  c45 = (c32&c24).where((c27*128+c38+c17*4096+-12672), UOp.const(dtypes.index, Invalid))
  c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
-  c49 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((128, 49, 4)), (), 2)
+  c49 = UOp(Ops.PARAM, dtypes.imageh((128, 49, 4)), (), 2)
  c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
-  c63 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(256), (), 3)
+  c63 = UOp(Ops.PARAM, dtypes.float.ptr(256), (), 3)
  c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5)
  c67 = c0.index((c2*256+c5+c8*4096), ptr=True).store(c65).end(c8, c2, c5)

@@ -63,16 +63,16 @@ def vision_conv_153():
  return c67.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))

 def dm_conv_172():
-  c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 240, 4)), (), 0)
+  c0 = UOp(Ops.PARAM, dtypes.imageh((1, 240, 4)), (), 0)
  c2 = UOp.range(960, 4, AxisType.LOOP)
-  c5 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((8, 384, 4)), (), 1)
+  c5 = UOp(Ops.PARAM, dtypes.imageh((8, 384, 4)), (), 1)
  c7 = UOp.range(32, 0, AxisType.REDUCE)
  c10 = UOp.range(4, 1, AxisType.REDUCE)
  c13 = UOp.range(12, 3, AxisType.REDUCE)
  c18 = UOp.range(8, 2, AxisType.REDUCE)
-  c23 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((240, 128, 4)), (), 2)
+  c23 = UOp(Ops.PARAM, dtypes.imageh((240, 128, 4)), (), 2)
  c35 = c5.index((c7*4+c10+c13*128+c18*1536))*c23.index((c10*4+c2%4+c7*16+c2//4*512))
-  c37 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(960), (), 3)
+  c37 = UOp(Ops.PARAM, dtypes.float.ptr(960), (), 3)
  c39 = c35.reduce(c7, c10, arg=Ops.ADD)+c37.index(c2)
  c50 = (1.0+((c39+0.044708251953125*(c39*(c39*c39)))*-2.3021129851685216).exp2()).reciprocal()*c39
  c53 = c50.reduce(c18, c13, arg=Ops.ADD)*0.010416666666666666
@@ -91,7 +91,7 @@ allocator = Device.default.allocator
 ps = get_program(ast, renderer)
 cr = CompiledRunner(replace(ps, device=Device.DEFAULT))

-gs = sorted(dedup([u for u in ast.toposort() if u.op is Ops.DEFINE_GLOBAL]), key=lambda u: u.arg)
+gs = sorted(dedup([u for u in ast.toposort() if u.op is Ops.PARAM]), key=lambda u: u.arg)
 # print(len(gs))
 # print([g.dtype for g in gs])
 bufs = [Buffer(ps.device, g.size, g.dtype if isinstance(g.dtype, ImageDType) else g.dtype._base).ensure_allocated() for g in gs]
@@ -0,0 +1,126 @@
+# ruff: noqa: F405
+"""Tests for GPU crash scenarios using AMD assembly to trigger invalid operations.
+
+These tests intentionally cause GPU faults to verify error handling.
+Run with: AMD=1 python -m pytest test/external/external_test_gpu_crash.py -v
+"""
+import unittest, re
+from tinygrad.device import Device
+from extra.assembly.amd.autogen.rdna3.ins import *  # noqa: F403
+from extra.assembly.amd.dsl import s, v, Inst, NULL
+
+def assemble(code:str, name:str="test") -> str:
+  kd = {"next_free_vgpr": 8, "next_free_sgpr": 8, "wavefront_size32": 1, "user_sgpr_kernarg_segment_ptr": 1, "kernarg_size": 8}
+  return f".text\n.globl {name}\n.p2align 8\n.type {name},@function\n{name}:\n{code}\n.rodata\n.p2align 6\n.amdhsa_kernel {name}\n" + \
+         "\n".join(f".amdhsa_{k} {v}" for k,v in kd.items()) + "\n.end_amdhsa_kernel"
+
+@unittest.skipIf(Device.DEFAULT != "AMD", "AMD required")
+class TestGPUCrash(unittest.TestCase):
+  @classmethod
+  def setUpClass(cls):
+    from tinygrad.runtime.support.compiler_amd import HIPCompiler
+    cls.dev = Device["AMD"]
+    cls.compiler = HIPCompiler(cls.dev.arch)
+
+  def setUp(self):
+    # Verify device works before each test
+    from tinygrad import Tensor
+    try:
+      t = Tensor([1.0, 2.0], device="AMD").realize()
+      assert (t + 1).numpy().tolist() == [2.0, 3.0]
+    except Exception:
+      self.fail("Device not working before test")
+
+  def _run(self, code: str):
+    from tinygrad.runtime.ops_amd import AMDProgram
+    prg = AMDProgram(self.dev, "test", self.compiler.compile(assemble(code)))
+    prg(self.dev.allocator.alloc(64), global_size=(1,1,1), local_size=(1,1,1), wait=True)
+
+  def _run_insts(self, insts: list[Inst]): self._run("\n".join(i.disasm() for i in insts))
+
+  def _assert_gpu_fault(self, func):
+    """Assert that func raises a RuntimeError indicating a GPU fault (not a setup error)."""
+    with self.assertRaises(RuntimeError) as cm:
+      func()
+    err_msg = str(cm.exception).lower()
+    # Verify it's a GPU fault, not a setup/device initialization error
+    self.assertTrue(
+      re.search(r'fault|hang|timeout|illegal|memviol', err_msg),
+      f"Expected GPU fault error, got: {cm.exception}"
+    )
+
+
+class TestOutOfBoundsMemoryAccess(TestGPUCrash):
+  """Tests for out-of-bounds memory accesses."""
+
+  def test_global_load_null_ptr(self):
+    """Global load from NULL pointer."""
+    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0),
+             global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
+    self._assert_gpu_fault(lambda: self._run_insts(insts))
+
+  def test_global_store_null_ptr(self):
+    """Global store to NULL pointer."""
+    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0), v_mov_b32_e32(v[2], 0xDEADBEEF),
+             global_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
+    self._assert_gpu_fault(lambda: self._run_insts(insts))
+
+  def test_global_load_unmapped_high_address(self):
+    """Global load from high unmapped address (0xDEAD00000000)."""
+    insts = [v_mov_b32_e32(v[0], 0x00000000), v_mov_b32_e32(v[1], 0xDEAD),
+             global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
+    self._assert_gpu_fault(lambda: self._run_insts(insts))
+
+  def test_global_store_unmapped_high_address(self):
+    """Global store to high unmapped address."""
+    insts = [v_mov_b32_e32(v[0], 0x00000000), v_mov_b32_e32(v[1], 0xDEAD), v_mov_b32_e32(v[2], 0x12345678),
+             global_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
+    self._assert_gpu_fault(lambda: self._run_insts(insts))
+
+  def test_global_atomic_unmapped(self):
+    """Atomic operation on unmapped memory."""
+    insts = [v_mov_b32_e32(v[0], 0xBEEF0000), v_mov_b32_e32(v[1], 0xDEAD), v_mov_b32_e32(v[2], 1),
+             global_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
+    self._assert_gpu_fault(lambda: self._run_insts(insts))
+
+
+class TestSMEMFaults(TestGPUCrash):
+  """Tests for scalar memory (SMEM) faults."""
+
+  def test_smem_load_null(self):
+    """SMEM load from NULL base."""
+    insts = [s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
+             s_load_b32(s[4], s[2:3], 0, soffset=NULL), s_waitcnt(0), s_endpgm()]
+    self._assert_gpu_fault(lambda: self._run_insts(insts))
+
+  def test_smem_load_unmapped(self):
+    """SMEM load from unmapped address."""
+    insts = [s_mov_b32(s[2], 0xBEEF0000), s_mov_b32(s[3], 0xDEAD),
+             s_load_b32(s[4], s[2:3], 0, soffset=NULL), s_waitcnt(0), s_endpgm()]
+    self._assert_gpu_fault(lambda: self._run_insts(insts))
+
+
+class TestFlatMemoryFaults(TestGPUCrash):
+  """Tests for FLAT memory instruction faults."""
+
+  def test_flat_load_null(self):
+    """FLAT load from NULL address."""
+    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0),
+             flat_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
+    self._assert_gpu_fault(lambda: self._run_insts(insts))
+
+  def test_flat_store_null(self):
+    """FLAT store to NULL address."""
+    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0), v_mov_b32_e32(v[2], 0xDEADBEEF),
+             flat_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
+    self._assert_gpu_fault(lambda: self._run_insts(insts))
+
+  def test_flat_atomic_null(self):
+    """FLAT atomic on NULL address."""
+    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0), v_mov_b32_e32(v[2], 1),
+             flat_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
+    self._assert_gpu_fault(lambda: self._run_insts(insts))
+
+
+if __name__ == "__main__":
+  unittest.main()
@@ -1,339 +0,0 @@
-import random, traceback, ctypes, argparse, os
-from typing import Any
-import numpy as np
-from collections import defaultdict
-from extra.optimization.helpers import load_worlds, ast_str_to_lin, kern_str_to_lin
-from tinygrad.engine.realize import get_program
-
-# We need to insert ioctl before opening devices.
-if os.getenv("VALIDATE_HCQ", 0) != 0:
-  try:
-    import extra.nv_gpu_driver.nv_ioctl
-    from tinygrad import Device
-    _, _ = Device["NV"], Device["CUDA"]
-  except Exception: pass
-
-  try:
-    import extra.qcom_gpu_driver.opencl_ioctl
-    from tinygrad import Device
-    _, _ = Device["QCOM"], Device["CL"]
-  except Exception: pass
-
-from tinygrad import Tensor, Device, dtypes
-from tinygrad.tensor import _to_np_dtype
-from tinygrad.codegen.opt.kernel import Kernel
-from tinygrad.codegen.opt import Opt, OptOps
-from tinygrad.codegen.opt.search import get_kernel_actions, bufs_from_lin
-from tinygrad.engine.realize import CompiledRunner
-from tinygrad.helpers import getenv, from_mv, prod, colored, Context, DEBUG, Timing
-from tinygrad.uop.ops import UOp, Ops
-from tinygrad.device import is_dtype_supported
-
-def on_linearizer_will_run(): pass
-def on_linearizer_did_run(): pass
-def compare_states(x, y): return (True, "")
-
-if getenv("VALIDATE_HCQ"):
-  if Device.DEFAULT == "NV":
-    print("VALIDATE_HCQ: Comparing NV to CUDA")
-    import extra.nv_gpu_driver.nv_ioctl
-    validate_device = Device["CUDA"]
-    on_linearizer_will_run = extra.nv_gpu_driver.nv_ioctl.before_launch
-    on_linearizer_did_run = extra.nv_gpu_driver.nv_ioctl.collect_last_launch_state
-    compare_states = extra.nv_gpu_driver.nv_ioctl.compare_launch_state
-  elif Device.DEFAULT == "QCOM":
-    print("VALIDATE_HCQ: Comparing QCOM to CL")
-    import extra.qcom_gpu_driver.opencl_ioctl
-    validate_device = Device["CL"]
-    on_linearizer_will_run = extra.qcom_gpu_driver.opencl_ioctl.before_launch
-    on_linearizer_did_run = extra.qcom_gpu_driver.opencl_ioctl.collect_last_launch_state
-    compare_states = extra.qcom_gpu_driver.opencl_ioctl.compare_launch_state
-  else:
-    print(colored("VALIDATE_HCQ options is ignored", 'red'))
-
-def tuplize_uops(uops:list[UOp]) -> tuple:
-  return tuple([(x.op, x.dtype, tuple(uops.index(x) for x in x.src), x.arg) for x in uops])
-
-def get_fuzz_rawbufs(lin):
-  rawbufs = bufs_from_lin(lin)
-
-  # Reallocate output buffer with additional area to detect out-of-bounds writes.
-  RED_AREA_SIZE = 1024
-  # setting output  # TODO: multi-output kernel
-  rawbufs[0] = get_fuzz_rawbuf_like(rawbufs[0], zero=True, size=rawbufs[0].size+RED_AREA_SIZE)
-  # setting inputs
-  with Context(DEBUG=0):
-    for rawbuf in rawbufs[1:]:
-      if dtypes.is_unsigned(rawbuf.dtype):
-        data = np.random.randint(0, 100, size=rawbuf.size, dtype=_to_np_dtype(rawbuf.dtype))
-      elif dtypes.is_int(rawbuf.dtype):
-        data = np.random.randint(-100, 100, size=rawbuf.size, dtype=_to_np_dtype(rawbuf.dtype))
-      elif rawbuf.dtype == dtypes.bool:
-        data = np.random.choice([True, False], size=rawbuf.size)
-      elif rawbuf.dtype == dtypes.half:
-        data = np.random.uniform(-1, 1, size=rawbuf.size).astype(dtype=_to_np_dtype(rawbuf.dtype))
-      else:
-        data = np.random.uniform(-10, 10, size=rawbuf.size).astype(dtype=_to_np_dtype(rawbuf.dtype))
-      rawbuf.copyin(Tensor(data, device=lin.opts.device).realize().uop.base.realized.as_buffer())
-  return rawbufs
-
-def get_fuzz_rawbuf_like(old_rawbuf, zero=False, copy=False, size=None, force_device=None):
-  rawbuf = type(old_rawbuf)(force_device or old_rawbuf.device, old_rawbuf.size if size is None else size, old_rawbuf.dtype).allocate()
-  if copy:
-    with Context(DEBUG=0): rawbuf.copyin(old_rawbuf.as_buffer())
-  elif zero:
-    with Context(DEBUG=0):
-      mv = memoryview(bytearray(rawbuf.size * rawbuf.dtype.itemsize))
-      ctypes.memset(from_mv(mv), 0, len(mv))
-      rawbuf.copyin(mv)
-  return rawbuf
-
-def run_linearizer(lin: Kernel, rawbufs=None, var_vals=None) -> tuple[str, Any]: # (error msg, run state)
-  if rawbufs is None: rawbufs = bufs_from_lin(lin)
-  if var_vals is None: var_vals = {v.expr: v.min for v in lin.vars}
-
-  # TODO: images needs required_optimization
-  try:
-    prg = CompiledRunner(get_program(lin.get_optimized_ast(), lin.opts))
-  except KeyboardInterrupt: raise
-  except Exception:
-    traceback.print_exc()
-    return "COMPILE_ERROR", None
-
-  if getenv("VALIDATE_HCQ"): on_linearizer_will_run()
-  try:
-    prg(rawbufs, var_vals, wait=True)
-  except KeyboardInterrupt: raise
-  except Exception:
-    traceback.print_exc()
-    return "EXEC_ERROR", None
-
-  if getenv("VALIDATE_HCQ"): run_state = on_linearizer_did_run()
-  else: run_state = None
-
-  return "PASS", run_state
-
-def compare_linearizer(lin: Kernel, rawbufs=None, var_vals=None, ground_truth=None, rtol=1e-2, atol=1e-2):
-  # TODO: for bfloat16 it compiles linearizer, but it does not run because numpy cannot generate bf16 buffer.
-  has_bf16 = any(b.dtype.base == dtypes.bfloat16 for b in lin.bufs)
-
-  # TODO: raise specific fuzzing errors instead of str, and propagate the error message
-  try:
-    if rawbufs is None:
-      rawbufs = get_fuzz_rawbufs(lin)
-    else:
-      rawbufs[0] = get_fuzz_rawbuf_like(rawbufs[0], zero=True) # get a new output buffer
-  except KeyboardInterrupt: raise
-  except BaseException:
-    return ("RAWBUFS_ERROR", rawbufs, var_vals, ground_truth, None)
-
-  if var_vals is None:
-    # TODO: handle symbolic max case
-    var_vals = {v.expr: random.randint(v.vmin, v.vmax) for v in lin.ast.variables()}
-
-  if ground_truth is None and not has_bf16:
-    unoptimized = Kernel(lin.ast)
-    if run_linearizer(unoptimized, rawbufs, var_vals)[0] != "PASS":
-      return ("BASELINE_ERROR", rawbufs, var_vals, ground_truth, None)
-    ground_truth = np.frombuffer(rawbufs[0].as_buffer(), _to_np_dtype(rawbufs[0].dtype)).copy()
-
-  rawbufs[0] = get_fuzz_rawbuf_like(rawbufs[0], zero=True) # get a new output buffer
-  run_msg, run_state = run_linearizer(lin, rawbufs, var_vals)
-  if run_msg != "PASS": return (run_msg, rawbufs, var_vals, ground_truth, run_state)
-
-  try:
-    if not has_bf16:
-      result = np.frombuffer(rawbufs[0].as_buffer(), _to_np_dtype(rawbufs[0].dtype))
-      np.testing.assert_allclose(result, ground_truth, rtol=rtol, atol=atol)
-  except KeyboardInterrupt: raise
-  except AssertionError as e:
-    if DEBUG >= 2:
-      print(f"COMPARE_ERROR details: {e}")
-      if getenv("DEBUG_VALUES") > 0:
-        mismatch_indices = np.where(~np.isclose(result, ground_truth, rtol=rtol, atol=atol))
-        mismatched_result = result[mismatch_indices]
-        mismatched_ground_truth = ground_truth[mismatch_indices]
-        for i, idx in enumerate(mismatch_indices[0]):
-          print(f"mismatch at {idx=}: result={mismatched_result[i]} <> ground_truth={mismatched_ground_truth[i]}")
-    return ("COMPARE_ERROR", rawbufs, var_vals, ground_truth, run_state)
-
-  return ("PASS", rawbufs, var_vals, ground_truth, run_state)
-
-def fuzz_linearizer(lin: Kernel, rtol=1e-2, atol=1e-2, opts_list=None):
-  SEED = getenv("SEED", 42)
-  random.seed(SEED)
-  np.random.seed(SEED)
-  print(lin.ast)
-  print(lin.colored_shape())
-  seen_uops = {}
-  last_lins = [lin]
-  failures:defaultdict[str, list[tuple[tuple[UOp, ...], list[Opt]]]] = defaultdict(list)
-  rawbufs, var_vals, ground_truth, validate_rawbufs = None, None, None, None
-
-  FUZZ_ALL_ACTIONS = getenv("FUZZ_ALL_ACTIONS", 0)
-  FUZZ_MAX_SIZE = getenv("FUZZ_MAX_SIZE", 0)
-  FUZZ_IGNORE_SIMPLE_OPS = getenv("FUZZ_IGNORE_SIMPLE_OPS", 1)
-
-  if FUZZ_MAX_SIZE > 0 and prod(lin.full_shape) > FUZZ_MAX_SIZE:
-    print("skipping large kernel")
-    return failures
-  if FUZZ_IGNORE_SIMPLE_OPS and _is_simple(lin):
-    print("skipping simple kernel")
-    return failures
-
-  test_depth = 1 if opts_list is not None else getenv("DEPTH", 1 if FUZZ_ALL_ACTIONS else 10)
-  for depth in range(test_depth):
-    next_lins = []
-    for lin in last_lins:
-      if opts_list is None: actions = get_kernel_actions(lin, include_0=False)
-      else:
-        actions = {}
-        for oi,opts in enumerate(opts_list):
-          lin2 = lin.copy()
-          for o in opts: lin2.apply_opt(o)
-          actions[oi] = lin2
-
-      if not actions: continue
-      if depth == 0 and getenv("FUZZ_REQUIRE_TC", 0):
-        tc_acts = {i: k for k in actions.values() if k.applied_opts[0].op == OptOps.TC}
-        if len(tc_acts) == 0: return failures
-        else: actions = tc_acts
-
-      test_lins = list(actions.values())
-      if FUZZ_ALL_ACTIONS: print(f"testing {lin.applied_opts=} with {len(actions)} actions")
-      elif opts_list is None: test_lins = [random.choice(test_lins)]
-
-      for test_lin in test_lins:
-        if not FUZZ_ALL_ACTIONS and test_lin.applied_opts: print(f"applied opts: {test_lin.applied_opts}")
-
-        # stop if kernel uops repeat
-        try: tuops = tuplize_uops(get_program(test_lin.get_optimized_ast(), test_lin.ren).uops)
-        except KeyboardInterrupt: raise
-        except BaseException as e:
-          print(test_lin.ast)
-          print(test_lin.applied_opts)
-          print(e)
-          failures["LINEARIZE_ERROR"].append((test_lin.ast, test_lin.applied_opts))
-          continue
-
-        if tuops in seen_uops: continue
-        seen_uops[tuops] = tuple(test_lin.applied_opts)
-
-        if not FUZZ_ALL_ACTIONS: print(test_lin.colored_shape())
-
-        (msg, rawbufs, var_vals, ground_truth, state1) = compare_linearizer(test_lin, rawbufs, var_vals, ground_truth, rtol=rtol, atol=atol)
-        if state1 is not None and validate_device is not None:
-          validate_lin = test_lin.copy()
-          validate_lin.ren = validate_device.renderer
-          if validate_rawbufs is None:
-            validate_rawbufs = [get_fuzz_rawbuf_like(x, copy=True, force_device=validate_device.device) for x in rawbufs]
-          (_msg, _, _, _, state2) = compare_linearizer(validate_lin, validate_rawbufs, var_vals, ground_truth, rtol=rtol, atol=atol)
-
-          if _msg != "PASS": failures[f"VALIDATE_DEV_{_msg}"].append((validate_lin.ast, validate_lin.applied_opts))
-
-          ok, err_msg = compare_states(state1, state2)
-          if not ok: failures["HCQ_COMPARE_FAILURE"].append((err_msg, test_lin.ast, test_lin.applied_opts, state1, state2))
-
-        if msg != "PASS":
-          print(test_lin.ast)
-          print(test_lin.applied_opts)
-          print(msg)
-          failures[msg].append((test_lin.ast, test_lin.applied_opts))
-          continue
-
-        next_lins.append(test_lin)
-
-    last_lins = next_lins
-    if FUZZ_ALL_ACTIONS: print(f"depth={depth} total_lins={len(last_lins)} {failures=}")
-  return failures
-
-def _is_simple(lin: Kernel) -> bool:
-  if len(lin.ast.src) > 1: return False
-  ast:UOp = lin.ast.src[0]
-  if ast.src[0].op is Ops.CAST and ast.src[0].src[0].op is Ops.LOAD: return True
-  return False
-
-if __name__ == "__main__":
-  parser = argparse.ArgumentParser(description="Run a fuzz testing on one or more kernels", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
-  parser.add_argument("--ast", type=str, default=None, help="the ast for the kernel to be optimized")
-  parser.add_argument("--file", type=str, default=None, help="a file containing asts to be optimized, one per line")
-  parser.add_argument("--beamreplay", type=str, default=None, help="replay asts and opts got from beam with CAPTURE_BEAM")
-  parser.add_argument("--logfile", type=str, default=None, help="a file containing a tuple of ast and applied_opts, one per line")
-  parser.add_argument("--expected-failures", type=int, default=0, help="the number of expected failed kernels")
-  parser.add_argument("--rtol", type=float, default=1e-2, help="relative tolerance for numerical comparison")
-  parser.add_argument("--atol", type=float, default=1e-2, help="absolute tolerance for numerical comparison")
-  args = parser.parse_args()
-
-  opts_list = None
-  if args.ast is not None:
-    print("loaded AST from CLI")
-    ast_strs = [args.ast]
-  elif args.file is not None:
-    print(f"loading ASTs from file '{args.file}'")
-    with open(args.file, 'r') as file:
-      ast_strs = file.readlines()
-  elif args.beamreplay is not None:
-    print(f"loading BEAM replay from file '{args.beamreplay}'")
-    with open(args.beamreplay, 'r') as file: fdata = file.readlines()
-    ast_strs, opts_list = [x.split(' :: ')[0] for x in fdata if not x.startswith("#")], [x.split(' :: ')[1] for x in fdata if not x.startswith("#")]
-
-    # dedup ast_strs and opts_list
-    dct = defaultdict(list)
-    for i in range(len(ast_strs)): dct[ast_strs[i]].append(eval(opts_list[i]))
-    ast_strs_items = list(dct.keys())
-    opts_list = [dct[c] for c in ast_strs_items]
-  elif args.logfile is not None:
-    print(f"loading ASTs from LOGKERNS file '{args.file}'")
-    with open(args.logfile, 'r') as file:
-      kern_strs = file.readlines()
-      test_lins = [kern_str_to_lin(kern_str) for kern_str in kern_strs]
-      ast_strs = [f"{lin.ast}" for lin in test_lins]
-  else:
-    print("loading ASTs from world")
-    ast_strs = load_worlds(filter_reduce=False, filter_novariable=False)
-
-  print(f"{len(ast_strs)=}")
-  tested = 0
-  failed_ids = []
-  failures = defaultdict(list)
-  seen_ast_strs = set()
-
-  try:
-    for i, ast in enumerate(ast_strs[:getenv("FUZZ_N", len(ast_strs))]):
-      if (nth := getenv("FUZZ_NTH", -1)) != -1 and i != nth: continue
-      if getenv("FUZZ_IMAGEONLY") and "dtypes.image" not in ast: continue
-      if "dtypes.image" in ast and Device.DEFAULT not in {"CL", "QCOM"}: continue  # IMAGE is only for CL
-      if ast in seen_ast_strs: continue
-      seen_ast_strs.add(ast)
-
-      lin = ast_str_to_lin(ast)
-      if not all(is_dtype_supported(buf.dtype) for buf in lin.bufs):
-        print("skipping kernel due to not supported dtype")
-        continue
-
-      with Timing(f"tested ast {i}: "):
-        tested += 1
-        fuzz_failures = fuzz_linearizer(lin, rtol=args.rtol, atol=args.atol, opts_list=(opts_list[i] if opts_list else None))
-        if fuzz_failures: failed_ids.append(i)
-        for k, v in fuzz_failures.items():
-          for f in v:
-            failures[k].append(f)
-  except KeyboardInterrupt: print(colored("STOPPING...", 'red'))
-
-  for msg, errors in failures.items():
-    for i, payload in enumerate(errors):
-      print(f"{msg} {i} kernel: {payload}") # easier to use with output with verify_kernel.py
-
-  print(f"{tested=}")
-  if failures:
-    print(f"{failed_ids=}")
-    for msg, errors in failures.items():
-      print(f"{msg}: {len(errors)}")
-    if len(failed_ids) == args.expected_failures:
-      print(colored(f"{len(failed_ids)} failed as expected", "yellow"))
-  if len(failed_ids) != args.expected_failures:
-    print(colored(f"failed on {len(failed_ids)} kernels, expected {args.expected_failures}", "red"))
-    # TODO: fix this
-    # raise RuntimeError(f"failed on {len(failed_ids)} kernels, expected {args.expected_failures}")
-  else:
-    print(colored("all passed", "green"))
@@ -1,157 +0,0 @@
-from tinygrad import Device, dtypes
-from tinygrad.helpers import getenv, colorize_float, DEBUG
-from extra.optimization.helpers import load_worlds, ast_str_to_lin
-from test.external.fuzz_linearizer import get_fuzz_rawbufs
-from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
-from tinygrad.codegen.opt.search import bufs_from_lin
-from tinygrad.engine.realize import CompiledRunner, get_program
-from tinygrad.tensor import _to_np_dtype
-from tinygrad.runtime.ops_amd import AMDDevice
-from contextlib import contextmanager
-import numpy as np
-import os, random, statistics
-
-am_signal_pages, am_signal_pool, am_devices = [], [], []
-amd_signal_pages, amd_signal_pool, amd_devices = [], [], []
-
-def rebind_vfio(pcibus="0000:44:00.0"):
-  print("rebind ", pcibus)
-  os.system("sudo rmmod amdgpu")
-  os.system("sudo modprobe vfio-pci")
-
-  base = f"/sys/bus/pci/devices/{pcibus}"
-  if os.path.exists(f"{base}/driver"):
-    with open(f"{base}/driver/unbind", "w") as f: f.write(pcibus)
-  with open(f"{base}/driver_override", "w") as f: f.write("vfio-pci")
-  with open("/sys/bus/pci/drivers_probe", "w") as f: f.write(pcibus)
-
-  os.system("sudo modprobe amdgpu")
-  os.system("rocm-smi --setprofile compute")
-  os.system("rocm-smi --setperflevel high")
-
-@contextmanager
-def run_amd():
-  global amd_signal_pages, amd_signal_pool, amd_devices
-  AMDDevice.driverless = False
-  AMDDevice.signal_pages, AMDDevice.signal_pool, AMDDevice.devices = amd_signal_pages, amd_signal_pool, amd_devices
-  yield
-  amd_signal_pages, amd_signal_pool, amd_devices = AMDDevice.signal_pages, AMDDevice.signal_pool, AMDDevice.devices
-  AMDDevice.signal_pages, AMDDevice.signal_pool, AMDDevice.devices = [], [], []
-
-@contextmanager
-def run_am():
-  global am_signal_pages, am_signal_pool, am_devices
-  AMDDevice.driverless = True
-  AMDDevice.signal_pages, AMDDevice.signal_pool, AMDDevice.devices = am_signal_pages, am_signal_pool, am_devices
-  yield
-  am_signal_pages, am_signal_pool, am_devices = AMDDevice.signal_pages, AMDDevice.signal_pool, AMDDevice.devices
-  AMDDevice.signal_pages, AMDDevice.signal_pool, AMDDevice.devices = [], [], []
-
-if __name__ == "__main__":
-  CHECK_CPU = getenv("CHECK_CPU", 0)
-  SEED = getenv("SEED", 42)
-  CNT = getenv("CNT", 7)
-  random.seed(SEED)
-  np.random.seed(SEED)
-
-  # TODO: NUM=780 is super slow
-  # NUM=1907 is broken on AMD and AM have some mismatches (0 vs 1)
-  # kfd feels so bad when taking gpu out while it's running... Need hacks to rebind it before running.
-  rebind_vfio(pcibus="0000:44:00.0")
-
-  ast_strs = load_worlds(filter_reduce=False, filter_novariable=True)
-
-  with run_am():
-    amdev = Device["AMD:1"]
-
-  with run_amd():
-    amddev = Device["AMD"]
-
-  if CHECK_CPU: cpudev = Device["CPU"]
-
-  single = getenv("NUM", -1)
-  if single != -1: ast_strs = ast_strs[single:single+1]
-
-  average_tm_amd, average_tm_am = 0, 0
-  for num,ast in enumerate(ast_strs):
-    with run_amd():
-      amdlin = ast_str_to_lin(ast, opts=amddev.renderer)
-      amdlin.apply_opts(hand_coded_optimizations(amdlin))
-      has_bf16 = any(b.dtype == dtypes.bfloat16 for b in amdlin.bufs)
-
-      amd_prg = CompiledRunner(get_program(amdlin.get_optimized_ast(), amdlin.opts))
-      amdbufs = bufs_from_lin(amdlin)
-      test_amdbufs = get_fuzz_rawbufs(amdlin) if not has_bf16 else amdbufs
-      if not has_bf16: contents = [buf.as_buffer() for buf in test_amdbufs]
-
-    with run_am():
-      rdr = amdev.renderer
-      rdr.device = "AMD:1"
-      amlin = ast_str_to_lin(ast, opts=amdev.renderer)
-      amlin.apply_opts(hand_coded_optimizations(amlin))
-      am_prg = CompiledRunner(get_program(amlin.get_optimized_ast(), amlin.opts))
-      ambufs = bufs_from_lin(amlin)
-      test_ambufs = get_fuzz_rawbufs(amlin) if not has_bf16 else ambufs
-      if not has_bf16:
-        for i,rawbuf in enumerate(test_ambufs): rawbuf.copyin(contents[i])
-
-    if CHECK_CPU:
-      cpu_rdr = cpudev.renderer
-      cpu_rdr.device = "CPU"
-      cpulin = ast_str_to_lin(ast, opts=cpu_rdr)
-      cpulin.apply_opts(hand_coded_optimizations(cpulin))
-      cpu_prg = CompiledRunner(get_program(cpulin.get_optimized_ast(), cpulin.opts))
-      cpubufs = bufs_from_lin(cpulin)
-      test_cpubufs = get_fuzz_rawbufs(cpulin) if not has_bf16 else ambufs
-      if not has_bf16:
-        for i,rawbuf in enumerate(test_cpubufs): rawbuf.copyin(contents[i])
-
-    # warmup
-    tm_amd, tm_am, failed = [], [], False
-    with run_amd():
-      try:
-        amd_prg(test_amdbufs, {}, wait=True)
-        for i in range(CNT): tm_amd.append(amd_prg(amdbufs, {}, wait=True))
-      except RuntimeError:
-        print("AMD FAILED")
-        tm_amd = [1e9]
-        failed = True
-
-    with run_am():
-      try:
-        am_prg(test_ambufs, {}, wait=True)
-        for i in range(CNT): tm_am.append(am_prg(ambufs, {}, wait=True))
-      except RuntimeError:
-        print("AM FAILED")
-        tm_am = [1e9]
-        failed = True
-
-    if CHECK_CPU:
-      cpu_prg(test_cpubufs, {}, wait=True)
-      for i in range(1): cpu_prg(cpubufs, {}, wait=True)
-
-    if not failed and not has_bf16:
-      with run_amd():
-        curesult = np.frombuffer(test_amdbufs[0].as_buffer(), _to_np_dtype(test_amdbufs[0].dtype))
-
-      with run_am():
-        amresult = np.frombuffer(test_ambufs[0].as_buffer(), _to_np_dtype(test_ambufs[0].dtype))
-
-      if CHECK_CPU:
-        cpuresult = np.frombuffer(test_cpubufs[0].as_buffer(), _to_np_dtype(test_cpubufs[0].dtype))
-        np.testing.assert_allclose(amresult, cpuresult, rtol=1e-2, atol=1e-2)
-        np.testing.assert_allclose(curesult, cpuresult, rtol=1e-2, atol=1e-2)
-
-      try:
-        np.testing.assert_allclose(curesult, amresult, rtol=1e-2, atol=1e-2)
-      except AssertionError as e:
-        print("AM and AMD results do not match")
-        print(e)
-
-    bam = statistics.median(tm_am)
-    bamd = statistics.median(tm_amd)
-    average_tm_amd += bamd
-    average_tm_am += bam
-    ratio = bam/bamd
-    print(f"{average_tm_am/average_tm_amd:5.2f}x -- {num:4d} {colorize_float(ratio)} {bam*1e6:7.2f} vs {bamd*1e6:7.2f} us", amlin.name)
-    if DEBUG > 3 and ratio > 1.04: print(f"AM slower {ratio}", amlin.ast, amlin.applied_opts)
@@ -1,69 +0,0 @@
-from tinygrad import Device, dtypes
-from tinygrad.helpers import getenv, colorize_float
-from extra.optimization.helpers import load_worlds, ast_str_to_lin
-from test.external.fuzz_linearizer import get_fuzz_rawbufs
-from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
-from tinygrad.codegen.opt.search import bufs_from_lin
-from tinygrad.engine.realize import CompiledRunner, get_program
-from tinygrad.tensor import _to_np_dtype
-import numpy as np
-
-if __name__ == "__main__":
-  ast_strs = load_worlds(filter_reduce=False, filter_novariable=True)
-  cudev = Device["CUDA"]
-  nvdev = Device["NV"]
-
-  # NUM=112 python3 test/external/speed_compare_cuda_nv.py
-
-  single = getenv("NUM", -1)
-  if single != -1: ast_strs = ast_strs[single:single+1]
-
-  average_tm_cuda, average_tm_nv = 0, 0
-  for num,ast in enumerate(ast_strs):
-    # cuda compile
-    culin = ast_str_to_lin(ast, opts=cudev.renderer)
-    culin.apply_opts(hand_coded_optimizations(culin))
-    has_bf16 = any(b.dtype == dtypes.bfloat16 for b in culin.bufs)
-
-    cuda_prg = CompiledRunner(get_program(culin.get_optimized_ast(), culin.opts))
-    cubufs = bufs_from_lin(culin)
-    test_cubufs = get_fuzz_rawbufs(culin) if not has_bf16 else cubufs
-
-    rdr = nvdev.renderer
-    rdr.device = "NV"
-    nvlin = ast_str_to_lin(ast, opts=rdr)
-    nvlin.apply_opts(hand_coded_optimizations(nvlin))
-    nv_prg = CompiledRunner(get_program(nvlin.get_optimized_ast(), nvlin.opts))
-    nvbufs = bufs_from_lin(nvlin)
-    test_nvbufs = get_fuzz_rawbufs(nvlin) if not has_bf16 else nvbufs
-    if not has_bf16:
-      for i,rawbuf in enumerate(test_nvbufs): rawbuf.copyin(test_cubufs[i].as_buffer())
-
-    # warmup
-    tm_cuda, tm_nv, failed = [], [], False
-    try:
-      cuda_prg(test_cubufs, {}, wait=True)
-      for i in range(5): tm_cuda.append(cuda_prg(cubufs, {}, wait=True))
-    except RuntimeError:
-      print("CUDA FAILED")
-      tm_cuda = [1e9]
-      failed = True
-
-    try:
-      nv_prg(test_nvbufs, {}, wait=True)
-      for i in range(5): tm_nv.append(nv_prg(nvbufs, {}, wait=True))
-    except RuntimeError:
-      print("NV FAILED")
-      tm_nv = [1e9]
-      failed = True
-
-    if not failed and not has_bf16:
-      curesult = np.frombuffer(test_cubufs[0].as_buffer(), _to_np_dtype(test_cubufs[0].dtype))
-      nvresult = np.frombuffer(test_nvbufs[0].as_buffer(), _to_np_dtype(test_nvbufs[0].dtype))
-      np.testing.assert_allclose(curesult, nvresult, rtol=1e-2, atol=1e-2)
-
-    average_tm_cuda += min(tm_cuda)
-    average_tm_nv += min(tm_nv)
-    ratio = min(tm_nv)/min(tm_cuda)
-    print(f"{average_tm_nv/average_tm_cuda:5.2f}x -- {num:4d} {colorize_float(ratio)}  {min(tm_nv)*1e6:7.2f} us", nvlin.name)
-    if ratio > 1.04: print(f"NV slower {ratio}", nvlin.ast, nvlin.applied_opts)
@@ -1,60 +0,0 @@
-import itertools
-from tinygrad import Device
-from tinygrad.engine.realize import CompiledRunner, get_program
-from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
-from tinygrad.helpers import getenv, colorize_float
-from extra.optimization.helpers import load_worlds, ast_str_to_lin
-from tinygrad.codegen.opt.search import bufs_from_lin
-from tinygrad.runtime.ops_cuda import PTXCompiler, PTXRenderer, CUDACompiler
-
-if __name__ == "__main__":
-  ast_strs = load_worlds(filter_reduce=False, filter_novariable=True)
-  # no bfloat16 for ptx at the moment
-  ast_strs = [x for x in ast_strs if "dtypes.bfloat16" not in x]
-  dev = Device["CUDA"]
-  ptx = PTXRenderer(dev.arch)
-
-  # NUM=112 python3 test/external/speed_compare_cuda_ptx.py
-
-  single = getenv("NUM", -1)
-  if single != -1: ast_strs = ast_strs[single:single+1]
-
-  average_tm_cuda, average_tm_ptx = 0, 0
-  for num,ast in enumerate(ast_strs):
-    # cuda compile
-    dev.compiler = CUDACompiler(dev.arch)
-    lin = ast_str_to_lin(ast, opts=dev.renderer)
-    lin.apply_opts(hand_coded_optimizations(lin))
-    cuda_prg = CompiledRunner(get_program(lin.get_optimized_ast(), lin.opts))
-
-    bufs = bufs_from_lin(lin)
-
-    # ptx compile
-    dev.compiler = PTXCompiler(dev.arch)
-    lin = ast_str_to_lin(ast, opts=ptx)
-    lin.apply_opts(hand_coded_optimizations(lin))
-    ptx_prg = CompiledRunner(get_program(lin.get_optimized_ast(), lin.opts))
-
-    # warmup
-    try:
-      cuda_prg(bufs, {}, wait=True)
-    except RuntimeError:
-      print("cuda failed ast:", num)
-      continue
-    ptx_prg(bufs, {}, wait=True)
-
-    tm_cuda, tm_ptx = [], []
-    for i in range(5):
-      tm_cuda.append(cuda_prg(bufs, {}, wait=True))
-      tm_ptx.append(ptx_prg(bufs, {}, wait=True))
-    average_tm_cuda += min(tm_cuda)
-    average_tm_ptx += min(tm_ptx)
-    ratio = min(tm_ptx)/min(tm_cuda)
-    print(f"{average_tm_ptx/average_tm_cuda:5.2f}x -- {num:4d} {colorize_float(ratio)}  {min(tm_ptx)*1e6:7.2f} us", lin.name)
-    if ratio > 1.5:
-      def fix(x): return x.replace('\t', ' ').strip()
-      ll1, ll2 = cuda_prg.lib.decode().split('\n'), ptx_prg.lib.decode().split('\n')
-      if single != -1:
-        for ln, (l1, l2) in enumerate(itertools.zip_longest(ll1, ll2, fillvalue='')):
-          print(f"{ln:5d} | {fix(l1):80s} | {fix(l2):80s}")
-      print(len(ll1), len(ll2), "RATIO", ratio, "us", min(tm_ptx)*1e6)
@@ -1,78 +0,0 @@
-import argparse
-from collections import defaultdict
-from extra.optimization.helpers import kern_str_to_lin, time_linearizer
-from test.external.fuzz_linearizer import compare_linearizer
-from tinygrad.helpers import colored
-from tinygrad.codegen.opt.kernel import Kernel
-
-# Use this with the LOGKERNS options to verify that all executed kernels are valid and evaluate to the same ground truth results
-
-# Example for GPT2:
-# 1) Run the model to log all kernels: `PYTHONPATH=. LOGKERNS=/tmp/gpt2_kerns.txt JIT=1 HALF=1 BEAM=2 CACHELEVEL=0 python3 examples/gpt2.py --count 10 --temperature 0 --timing`   # noqa: E501
-# 2) Validate the kernel correctness: `PYTHONPATH=. python3 ./test/external/verify_kernel.py --file /tmp/gpt2_kerns.txt`
-
-if __name__ == "__main__":
-  parser = argparse.ArgumentParser(description="Verify the correctness of one or more kernel", formatter_class=argparse.ArgumentDefaultsHelpFormatter)    # noqa: E501
-  parser.add_argument("--kernel", type=str, default=None, help="a string of a tuple of (ast, applied_opts,)")
-  parser.add_argument("--file", type=str, default=None, help="a file containing a tuple of ast and applied_opts, one per line")
-  parser.add_argument("--pkl", type=str, default=None, help="a pickle file containing a single tuple of ast and applied_opts")
-  parser.add_argument("--rtol", type=float, default=1e-2, help="relative tolerance for numerical comparison")
-  parser.add_argument("--atol", type=float, default=1e-2, help="absolute tolerance for numerical comparison")
-  parser.add_argument("--timing", action='store_true', help="show final timing for the kernel")
-  parser.add_argument("--expected-failures", type=int, default=0, help="the number of expected failed kernels")
-  args = parser.parse_args()
-
-  if args.kernel is not None:
-    print("loading kernel from args")
-    test_lins = [kern_str_to_lin(args.kernel)]
-  elif args.file is not None:
-    print(f"loading kernel from file '{args.file}'")
-    with open(args.file, 'r') as file:
-      kern_strs = file.readlines()
-      test_lins = [kern_str_to_lin(kern_str) for kern_str in kern_strs]
-  elif args.pkl is not None:
-    print(f"loading kernel from pickle file '{args.file}'")
-    import pickle
-    with open(args.pkl, 'rb') as file:
-      (ast, applied_opts,) = pickle.load(file)
-      lin = Kernel(ast)
-      lin.apply_opts(applied_opts)
-      test_lins = [lin]
-
-  else:
-    raise RuntimeError("no kernel specified; use --kernel, --file, or --pkl options")
-
-  print(f"verifying {len(test_lins)} kernels")
-
-  failed_ids = []
-  failures = defaultdict(list)
-  for i, test_lin in enumerate(test_lins):
-    print(f"testing kernel {i}")
-    print(test_lin.ast)
-    print(test_lin.applied_opts)
-    unoptimized_lin = Kernel(test_lin.ast)
-    print(f"{unoptimized_lin.colored_shape()} -> {test_lin.colored_shape()}")
-    (msg,rb,vv,gt) = compare_linearizer(test_lin, None, None, None, rtol=args.rtol, atol=args.atol)
-    if msg != "PASS":
-      failed_ids.append(i)
-      failures[msg].append((test_lin.ast, test_lin.applied_opts))
-    if args.timing:
-      tm = time_linearizer(test_lin, rb, allow_test_size=False, cnt=10)
-      print(f"final time {tm*1e6:9.0f} us")
-
-  for msg, errors in failures.items():
-    for i, (ast, opts) in enumerate(errors):
-      print(f"{msg} {i} AST: {ast}")
-      print(f"{msg} {i} OPTS: {opts}\n")
-
-  print(f"tested {len(test_lins)} kernels")
-  if failures:
-    print(f"{failed_ids=}")
-    for msg, errors in failures.items():
-      print(f"{msg}: {len(errors)}")
-    if len(failed_ids) == args.expected_failures:
-      print(colored(f"{len(failed_ids)} failed as expected", "yellow"))
-  if len(failed_ids) != args.expected_failures:
-    raise RuntimeError(f"failed on {len(failed_ids)} kernels, expected {args.expected_failures}")
-  else:
-    print(colored("all passed", "green"))
@@ -61,7 +61,7 @@ def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None):
  for buf_dt, data in inputs or []:
    bufs.append(buf:=allocator.alloc(len(data) * buf_dt.itemsize))
    allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + (buf_dt.fmt or ""), *data)))
-  g = UOp(Ops.DEFINE_GLOBAL, uop.dtype.ptr(), arg=0, src=())
+  g = UOp(Ops.PARAM, uop.dtype.ptr(), arg=0, src=())
  prg = get_program(UOp.store(g.index(UOp.const(dtypes.int, 0)), uop).sink(), PythonRenderer())
  prog = PythonProgram("run", PythonCompiler().compile(prg.src))
  prog(out_buf:=allocator.alloc(uop.dtype.itemsize), *bufs)
@@ -4,7 +4,7 @@ import tinygrad.runtime.autogen.am.am as am
 import tinygrad.runtime.autogen.amdgpu_drm as amdgpu_drm
 from tinygrad.helpers import from_mv
 from test.mockgpu.driver import VirtDriver, VirtFileDesc, TextFileDesc, DirFileDesc, VirtFile
-from test.mockgpu.amd.amdgpu import AMDGPU, gpu_props
+from test.mockgpu.amd.amdgpu import AMDGPU, gpu_props, GFX_TARGET_VERSION, MOCKGPU_ARCH

 libc = ctypes.CDLL(ctypes.util.find_library("c"))
 libc.mmap.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_long]
@@ -90,35 +90,30 @@ class AMDDriver(VirtDriver):
  def _prepare_gpu(self, gpu_id):
    self.doorbells[gpu_id] = memoryview(bytearray(0x2000))
    self.gpus[gpu_id] = AMDGPU(gpu_id)
+    # IP versions: rdna3 = GC 11.0.0, NBIF 4.3.0; rdna4 = GC 12.0.0, NBIF 6.3.1
+    ip_versions = {"rdna3": {"gc": (11, 0, 0), "sdma": (6, 0, 0), "nbif": (4, 3, 0)},
+                   "rdna4": {"gc": (12, 0, 0), "sdma": (6, 0, 0), "nbif": (6, 3, 1)}}[MOCKGPU_ARCH]
+    def ip_discovery_files(hwid, ver, base_addr):
+      p = f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{hwid}/0'
+      return [VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{hwid}', functools.partial(DirFileDesc, child_names=['0'])),
+              VirtFile(f'{p}/major', functools.partial(TextFileDesc, text=str(ver[0]))),
+              VirtFile(f'{p}/minor', functools.partial(TextFileDesc, text=str(ver[1]))),
+              VirtFile(f'{p}/revision', functools.partial(TextFileDesc, text=str(ver[2]))),
+              VirtFile(f'{p}/base_addr', functools.partial(TextFileDesc, text=base_addr))]
    self.tracked_files += [
      VirtFile('/sys/module/amdgpu', functools.partial(TextFileDesc, text="1")),
      VirtFile('/sys/module/amdgpu/parameters/ppfeaturemask', functools.partial(TextFileDesc, text="0xffff3fff")),
      VirtFile(f'/sys/devices/virtual/kfd/kfd/topology/nodes/{gpu_id}', functools.partial(DirFileDesc, child_names=['gpu_id', 'properties'])),
      VirtFile(f'/sys/devices/virtual/kfd/kfd/topology/nodes/{gpu_id}/gpu_id', functools.partial(TextFileDesc, text=f"{gpu_id}")),
      VirtFile(f'/sys/devices/virtual/kfd/kfd/topology/nodes/{gpu_id}/properties',
-        functools.partial(TextFileDesc, text=gpu_props.format(drm_render_minor=gpu_id))),
+        functools.partial(TextFileDesc, text=gpu_props.format(drm_render_minor=gpu_id, gfx_target_version=GFX_TARGET_VERSION))),
      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/power_dpm_force_performance_level',
               functools.partial(TextFileDesc, text='profile_standard\n')),
      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0',
               functools.partial(DirFileDesc, child_names=[str(am.GC_HWID), str(am.SDMA0_HWID), str(am.NBIF_HWID)])),
-      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.GC_HWID}', functools.partial(DirFileDesc, child_names=['0'])),
-      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.GC_HWID}/0/major', functools.partial(TextFileDesc, text='11')),
-      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.GC_HWID}/0/minor', functools.partial(TextFileDesc, text='0')),
-      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.GC_HWID}/0/revision', functools.partial(TextFileDesc, text='0')),
-      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.GC_HWID}/0/base_addr',
-               functools.partial(TextFileDesc, text='0x00001260\n0x0000A000\n0x0001C000\n0x02402C00')),
-      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.SDMA0_HWID}', functools.partial(DirFileDesc, child_names=['0'])),
-      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.SDMA0_HWID}/0/major', functools.partial(TextFileDesc, text='6')),
-      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.SDMA0_HWID}/0/minor', functools.partial(TextFileDesc, text='0')),
-      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.SDMA0_HWID}/0/revision', functools.partial(TextFileDesc, text='0')),
-      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.SDMA0_HWID}/0/base_addr',
-               functools.partial(TextFileDesc, text='0x00001260\n0x0000A000\n0x0001C000\n0x02402C00')),
-      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.NBIF_HWID}', functools.partial(DirFileDesc, child_names=['0'])),
-      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.NBIF_HWID}/0/major', functools.partial(TextFileDesc, text='4')),
-      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.NBIF_HWID}/0/minor', functools.partial(TextFileDesc, text='3')),
-      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.NBIF_HWID}/0/revision', functools.partial(TextFileDesc, text='0')),
-      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.NBIF_HWID}/0/base_addr',
-               functools.partial(TextFileDesc, text='0x00000000\n0x00000014\n0x00000D20\n0x00010400\n0x0241B000\n0x04040000')),
+      *ip_discovery_files(am.GC_HWID, ip_versions["gc"], '0x00001260\n0x0000A000\n0x0001C000\n0x02402C00'),
+      *ip_discovery_files(am.SDMA0_HWID, ip_versions["sdma"], '0x00001260\n0x0000A000\n0x0001C000\n0x02402C00'),
+      *ip_discovery_files(am.NBIF_HWID, ip_versions["nbif"], '0x00000000\n0x00000014\n0x00000D20\n0x00010400\n0x0241B000\n0x04040000'),
      VirtFile(f'/dev/dri/renderD{gpu_id}', functools.partial(DRMFileDesc, driver=self, gpu=f"{self.gpus[gpu_id]}")),
    ]

@@ -1,8 +1,11 @@
 import ctypes, time
 from test.mockgpu.gpu import VirtGPU
 from test.mockgpu.helpers import _try_dlopen_remu
-from tinygrad.helpers import getbits, to_mv
+from tinygrad.helpers import getbits, to_mv, getenv
 from tinygrad.runtime.support import c
+
+MOCKGPU_ARCH = getenv("MOCKGPU_ARCH", "rdna3")
+GFX_TARGET_VERSION = {"rdna3": 110000, "rdna4": 120000}[MOCKGPU_ARCH]
 import tinygrad.runtime.autogen.amd_gpu as amd_gpu, tinygrad.runtime.autogen.am.pm4_nv as pm4

 SDMA_MAX_COPY_SIZE = 0x400000
@@ -194,10 +197,11 @@ class PM4Executor(AMDQueue):
    scratch_size = wavesize * 4  # This gives the scratch size per thread (lane)

    assert prg_sz > 0, "Invalid prg ptr (not found in mapped ranges)"
-    # Pass valid memory ranges, rsrc2, and scratch_size to Python emulator
+    # Pass valid memory ranges, rsrc2, scratch_size and arch to Python emulator
    if hasattr(remu, 'valid_mem_ranges'): remu.valid_mem_ranges = self.gpu.mapped_ranges
    if hasattr(remu, 'rsrc2'): remu.rsrc2 = rsrc2
    if hasattr(remu, 'scratch_size'): remu.scratch_size = scratch_size
+    if hasattr(remu, 'arch'): remu.arch = self.gpu.arch
    err = remu.run_asm(prg_addr, prg_sz, *gl, *lc, args_addr)
    if err != 0: raise RuntimeError("remu does not support the new instruction introduced in this kernel")

@@ -314,6 +318,7 @@ class AMDGPU(VirtGPU):
    self.regs = AMDGPURegisters()
    self.mapped_ranges = set()
    self.queues = []
+    self.arch = MOCKGPU_ARCH

  def map_range(self, vaddr, size): self.mapped_ranges.add((vaddr, size))
  def unmap_range(self, vaddr, size): self.mapped_ranges.remove((vaddr, size))
@@ -342,7 +347,7 @@ simd_arrays_per_engine 2
 cu_per_simd_array 8
 simd_per_cu 2
 max_slots_scratch_cu 32
-gfx_target_version 110000
+gfx_target_version {gfx_target_version}
 vendor_id 4098
 device_id 29772
 location_id 34304
@@ -16,14 +16,15 @@ def _try_dlopen_gpuocelot():
  return None

 class PythonRemu:
-  """Python RDNA3 emulator wrapper that matches the libremu.so interface."""
+  """Python RDNA3/RDNA4 emulator wrapper that matches the libremu.so interface."""
  valid_mem_ranges: set[tuple[int, int]] = set()
  rsrc2: int = 0x19c  # Default: USER_SGPR_COUNT=14, enable X and Y workgroup IDs
  scratch_size: int = 0  # private_segment_fixed_size from kernel descriptor
+  arch: str = "rdna3"  # Architecture: rdna3 or rdna4

  def run_asm(self, lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int, lz: int, args_ptr: int) -> int:
    from extra.assembly.amd.emu import run_asm
-    return run_asm(lib, lib_sz, gx, gy, gz, lx, ly, lz, args_ptr, self.rsrc2, self.scratch_size)
+    return run_asm(lib, lib_sz, gx, gy, gz, lx, ly, lz, args_ptr, self.rsrc2, self.scratch_size, self.arch)

 def _try_dlopen_remu():
  # Use Python emulator only if PYTHON_REMU=1
@@ -1,30 +1,20 @@
 # ruff: noqa: E501
 import unittest
-from tinygrad import Tensor
 from tinygrad.uop.ops import UOp, Ops, AxisType
 from tinygrad.dtype import dtypes
 from tinygrad.engine.realize import get_program
 from tinygrad.device import Device

 class TestLinearizerFailures(unittest.TestCase):
-  def test_cumsum_repeat_reshape_multiply(self):
-    # cumsum + repeat + reshape + multiply fails when step > 512
-    step, num_steps = 513, 10
-    t = Tensor.arange(step).float()
-    phase = t.cumsum()
-    tiled = phase.repeat((num_steps,)).reshape(num_steps, step)
-    pattern = Tensor([1,0,0,1,0,0,0,0,1,0]).reshape(num_steps, 1)
-    result = (tiled * pattern).flatten()
-    result.numpy()  # should not raise AssertionError in CFGContext
  def test_fail_1(self):
-    c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(64), arg=0, src=())
+    c0 = UOp(Ops.PARAM, dtypes.float.ptr(64), arg=0, src=())
    c1 = UOp.range(UOp.const(dtypes.index, 2), 1, AxisType.LOOP)
    c2 = UOp.range(UOp.const(dtypes.index, 32), 2, AxisType.LOOP)
    c3 = ((c1*UOp.const(dtypes.index, 32))+c2)
-    c4 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(163840), arg=1, src=())
+    c4 = UOp(Ops.PARAM, dtypes.float.ptr(163840), arg=1, src=())
    c5 = UOp.range(UOp.const(dtypes.index, 2560), 0, AxisType.REDUCE)
    c6 = c4.index(((((((c5//UOp.const(dtypes.index, 8))%UOp.const(dtypes.index, 8))*UOp.const(dtypes.index, 8))+(c5%UOp.const(dtypes.index, 8)))+(((c2*UOp.const(dtypes.index, 40))+(c5//UOp.const(dtypes.index, 64)))*UOp.const(dtypes.index, 64)))+(c1*UOp.const(dtypes.index, 81920))))
-    c7 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(64), arg=2, src=())
+    c7 = UOp(Ops.PARAM, dtypes.float.ptr(64), arg=2, src=())
    c8 = c7.index(c3)
    c9 = ((((c6+(c8*UOp.const(dtypes.float, -1.0)))*(c6+(c8*UOp.const(dtypes.float, -1.0)))).reduce(c5, arg=Ops.ADD)*UOp.const(dtypes.float, 0.000390625))+UOp.const(dtypes.float, 1e-05)).sqrt().reciprocal()
    c10 = c0.index(c3).store(c9).end(c1, c2)
@@ -5,17 +5,17 @@ from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.uop.symbolic import simplify_valid
 from tinygrad.helpers import Context
-from test.unit.test_uop_symbolic import check_uop_against_string
+from test.null.test_uop_symbolic import check_uop_against_string

 def get_gated_load_uop(valid:UOp, idx:UOp):
  return UOp(Ops.LOAD, dtypes.float, (
-    UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=0).index(idx.valid(valid), ptr=True),
+    UOp(Ops.PARAM, dtypes.float.ptr(), arg=0).index(idx.valid(valid), ptr=True),
    UOp.const(dtypes.float, 0.0)
  ))

 def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UOp]):
  return UOp(Ops.LOAD, dtypes.float.vec(4), (
-    UOp(Ops.DEFINE_GLOBAL, dtypes.imagef(image_shape), arg=0).index(UOp(Ops.VECTORIZE, dtypes.index.vec(2), idx).valid(valid), ptr=True),
+    UOp(Ops.PARAM, dtypes.imagef(image_shape), arg=0).index(UOp(Ops.VECTORIZE, dtypes.index.vec(2), idx).valid(valid), ptr=True),
    UOp(Ops.VECTORIZE, dtypes.float.vec(4), src=(UOp.const(dtypes.float, 0.0),) * 4)
  ))

@@ -461,13 +461,13 @@ class TestUnfoldableImageChannelSelection(unittest.TestCase):
  def test_bounded_channel_no_nan(self):
    # unfoldable image load with bounded idx % 4 range [0,1] -> no NAN fallback needed
    lidx = Special("lidx", 2)
-    load = UOp(Ops.LOAD, dtypes.float, (UOp(Ops.DEFINE_GLOBAL, dtypes.imagef((10, 10, 4)), arg=0).index(lidx, ptr=True), UOp.const(dtypes.float, 0)))
+    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)))
    self.assertEqual(self._count_nans(load), 0)

  def test_unbounded_channel_has_nan(self):
    # variable with negative range -> x % 4 can be negative -> needs NAN fallback
    x = Variable("x", -10, 10)
-    load = UOp(Ops.LOAD, dtypes.float, (UOp(Ops.DEFINE_GLOBAL, dtypes.imagef((10, 10, 4)), arg=0).index(x, ptr=True), UOp.const(dtypes.float, 0)))
+    load = UOp(Ops.LOAD, dtypes.float, (UOp(Ops.PARAM, dtypes.imagef((10, 10, 4)), arg=0).index(x, ptr=True), UOp.const(dtypes.float, 0)))
    self.assertEqual(self._count_nans(load), 1)

 class TestDropTrueGate(unittest.TestCase):
@@ -475,7 +475,7 @@ class TestDropTrueGate(unittest.TestCase):
    # test that INDEX with a constant True gate gets simplified to drop the gate
    from tinygrad.codegen.late.devectorizer import load_store_indexing
    from tinygrad.uop.ops import graph_rewrite
-    buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), arg=0)
+    buf = UOp(Ops.PARAM, dtypes.int.ptr(), arg=0)
    idx = UOp.const(dtypes.index, 0)
    true_gate = UOp.const(dtypes.bool, True)
    index_with_gate = UOp(Ops.INDEX, dtypes.int.ptr(), (buf, idx, true_gate))
@@ -10,7 +10,7 @@ class TestTranscendentalFunctions(unittest.TestCase):
  def test_payne_hanek_reduction(self):
    # TODO: Test constant input when constant folding is fixed (or maybe test both variants)
    # Load input value from a buffer to prevent constant folding
-    input_buf = UOp(Ops.DEFINE_GLOBAL, dtypes.double.ptr(), arg=1, src=())
+    input_buf = UOp(Ops.PARAM, dtypes.double.ptr(), arg=1, src=())
    loaded_value = input_buf.index(UOp.const(dtypes.int, 0))
    def eval_payne_hanek_reduction(v:float) -> tuple[float, int]:
      return tuple(eval_uop(u, [(dtypes.float64, [v])]) for u in payne_hanek_reduction(loaded_value))
@@ -253,7 +253,7 @@ class TestUOpGraph(unittest.TestCase):

  @unittest.skip("this test isn't valid uops")
  def test_noop_vectorize_fold(self):
-    d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=0)
+    d0 = UOp(Ops.PARAM, dtypes.float.ptr(), arg=0)
    idx = UOp.const(dtypes.int, 0)
    ld = UOp(Ops.LOAD, dtypes.float.vec(2), (d0, idx))
    vec = UOp(Ops.VECTORIZE, dtypes.float.vec(2), (ld,))
@@ -265,9 +265,9 @@ class TestUOpGraph(unittest.TestCase):

  @unittest.skip("this test isn't valid uops")
  def test_gep_vec_fold(self):
-    d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
-    d1 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 1)
-    d2 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 2)
+    d0 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 0)
+    d1 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 1)
+    d2 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 2)
    idx = UOp.const(dtypes.int, 0)
    def _test_vec(geps, count=4):
      vec = UOp(Ops.VECTORIZE, dtypes.float.vec(count), geps)
@@ -373,8 +373,8 @@ class TestUOpGraph(unittest.TestCase):
      self.assertEqual(uops[-2], wmma)  # -2 to skip SINK

  def test_cast_alu_fold(self):
-    d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.bool.ptr(), arg=0)
-    d1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), arg=1)
+    d0 = UOp(Ops.PARAM, dtypes.bool.ptr(), arg=0)
+    d1 = UOp(Ops.PARAM, dtypes.int.ptr(), arg=1)
    idx = UOp.const(dtypes.int, 0)
    ld = d1.index(idx)
    alu = (ld<1).cast(dtypes.bool)
@@ -383,8 +383,8 @@ class TestUOpGraph(unittest.TestCase):
    self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 0)

  def test_double_cast_fold(self):
-    d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=0)
-    d1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), arg=1)
+    d0 = UOp(Ops.PARAM, dtypes.float.ptr(), arg=0)
+    d1 = UOp(Ops.PARAM, dtypes.int.ptr(), arg=1)
    idx = UOp.const(dtypes.int, 0)
    ld = d1.index(idx)
    alu = ld.cast(dtypes.float).cast(dtypes.float)
@@ -407,7 +407,7 @@ class TestUOpGraph(unittest.TestCase):

  def test_bitcast_to_same_dtype_fold(self):
    for dt in dtypes.ints + dtypes.floats + (dtypes.bool,):
-      d0 = UOp(Ops.DEFINE_GLOBAL, dt.ptr(), arg=0)
+      d0 = UOp(Ops.PARAM, dt.ptr(), arg=0)
      v = d0.index(UOp.const(dtypes.int, 0))
      uops = to_uops_list([v.bitcast(dt)])
      self.assertEqual(len([x for x in uops if x.op is Ops.BITCAST]), 0, f"dtype = {dt}")
@@ -420,7 +420,7 @@ class TestUOpGraph(unittest.TestCase):

  def test_where_on_gated_load_fold(self):
    ridx0 = UOp.range(100, 0)
-    d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 0)
+    d0 = UOp(Ops.PARAM, dtypes.long.ptr(), (), 0)
    ld = d0.index(ridx0.valid(ridx0<50))
    w = (ridx0<50).where(ld, 5)
    uops = to_uops_list([w])
@@ -430,7 +430,7 @@ class TestUOpGraph(unittest.TestCase):

  def test_where_on_gated_load_folds_swapped_branches(self):
    ridx0 = UOp.range(100, 0)
-    d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 0)
+    d0 = UOp(Ops.PARAM, dtypes.long.ptr(), (), 0)
    ld = d0.index(ridx0.valid((ridx0<50).logical_not()))
    w = (ridx0<50).where(5, ld)
    uops = to_uops_list([w])
@@ -440,7 +440,7 @@ class TestUOpGraph(unittest.TestCase):

  def test_where_on_gated_load_with_cast(self):
    ridx0 = UOp.range(100, 0)
-    d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
+    d0 = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0)
    gate_idx = ridx0.valid((ridx0<50))
    ld = d0.index(gate_idx).cast(dtypes.float)
    w = (ridx0<50).where(ld, 5.0)
@@ -451,7 +451,7 @@ class TestUOpGraph(unittest.TestCase):

  def test_where_in_store_becomes_gate(self):
    ridx0 = UOp.range(100, 0)
-    d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 0)
+    d0 = UOp(Ops.PARAM, dtypes.long.ptr(), (), 0)
    idx = d0.index(ridx0)
    ld = idx.load()
    val = (ridx0<50).where(5, ld)
@@ -464,14 +464,14 @@ class TestUOpGraph(unittest.TestCase):
  def test_load_idx_becomes_int(self):
    # mnist indexing with split reduceop
    # Make sure we are not doign math on the loaded index, which would promote it to long
-    c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.uchar.ptr(128000), arg=0, src=())
+    c0 = UOp(Ops.PARAM, dtypes.uchar.ptr(128000), arg=0, src=())
    c1 = UOp.range(UOp.const(dtypes.index, 512), 1, AxisType.LOOP)
    c2 = UOp.range(UOp.const(dtypes.index, 250), 2, AxisType.LOOP)
-    c3 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(512), arg=1, src=())
+    c3 = UOp(Ops.PARAM, dtypes.int.ptr(512), arg=1, src=())
    c4 = c3.index(c1)
    c5 = UOp.range(UOp.const(dtypes.index, 240), 0, AxisType.REDUCE)
    c6 = ((c2*UOp.const(dtypes.index, 240))+c5)
-    c7 = UOp(Ops.DEFINE_GLOBAL, dtypes.uchar.ptr(60000), arg=2, src=())
+    c7 = UOp(Ops.PARAM, dtypes.uchar.ptr(60000), arg=2, src=())
    c8 = c7.index(c6)
    c9 = ((c4<0).where((c4+60000), c4)!=c6.cast(dtypes.int)).where(0, c8.cast(dtypes.uint).cast(dtypes.uchar)).reduce(c5, arg=Ops.ADD)
    c10 = c0.index(((c1*UOp.const(dtypes.index, 250))+c2)).store(c9).end(c1, c2)
@@ -481,14 +481,14 @@ class TestUOpGraph(unittest.TestCase):

  def test_load_idx_no_math_on_loaded(self):
    # test the (x+y)<c pattern where x has loads - we shouldn't do math on loaded indices
-    c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.uchar.ptr(128000), arg=0, src=())
+    c0 = UOp(Ops.PARAM, dtypes.uchar.ptr(128000), arg=0, src=())
    c1 = UOp.range(UOp.const(dtypes.index, 512), 1, AxisType.LOOP)
    c2 = UOp.range(UOp.const(dtypes.index, 250), 2, AxisType.LOOP)
-    c3 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(512), arg=1, src=())
+    c3 = UOp(Ops.PARAM, dtypes.int.ptr(512), arg=1, src=())
    c4 = c3.index(c1)  # c4 is a load
    c5 = UOp.range(UOp.const(dtypes.index, 240), 0, AxisType.REDUCE)
    c6 = ((c2*UOp.const(dtypes.index, 240))+c5)
-    c7 = UOp(Ops.DEFINE_GLOBAL, dtypes.uchar.ptr(60000), arg=2, src=())
+    c7 = UOp(Ops.PARAM, dtypes.uchar.ptr(60000), arg=2, src=())
    c8 = c7.index(c6)
    # (loaded + range) < const pattern - loaded value shouldn't be promoted to long
    loaded_idx = c4.cast(dtypes.index)
@@ -500,9 +500,9 @@ class TestUOpGraph(unittest.TestCase):
      self.assertNotEqual(u.dtype, dtypes.long)

  def test_fold_gated_load(self):
-    glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
-    glbl1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 1)
-    glbl2 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 2)
+    glbl0 = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0)
+    glbl1 = UOp(Ops.PARAM, dtypes.int.ptr(), (), 1)
+    glbl2 = UOp(Ops.PARAM, dtypes.int.ptr(), (), 2)
    idx = UOp.const(dtypes.int, 0)
    ld0 = glbl1.index(UOp.invalid())
    ld1 = glbl2.index(idx.valid(UOp.const(dtypes.bool, True)))
@@ -512,7 +512,7 @@ class TestUOpGraph(unittest.TestCase):
    self.assertEqual(ld0, UOp.load(glbl2.index(idx, ptr=True), dtype=dtypes.int))

  def test_fold_gated_load_local(self):
-    glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
+    glbl0 = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0)
    smem = UOp(Ops.DEFINE_LOCAL, dtypes.int.ptr(size=18, addrspace=AddrSpace.LOCAL), (), "temp")
    lidx = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 16),), "lidx0")
    st = UOp(Ops.STORE, dtypes.void, (smem.index(lidx, ptr=True), glbl0.index(lidx, ptr=True).load()))
@@ -526,7 +526,7 @@ class TestUOpGraph(unittest.TestCase):
    self.assertEqual(ld0.src[0], smem.after(barrier).index(lidx+2, ptr=True))

  def test_fold_gated_store(self):
-    glbl = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
+    glbl = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0)
    idx0 = UOp.const(dtypes.int, 0)
    idx1 = UOp.const(dtypes.int, 0)
    val = UOp.const(dtypes.int, 42)
@@ -539,7 +539,7 @@ class TestUOpGraph(unittest.TestCase):

  @unittest.skip("this is a uop type error")
  def test_asserts_bad_gate(self):
-    glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
+    glbl0 = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0)
    idx = UOp.const(dtypes.int, 0)
    bad_gate = UOp.const(dtypes.int, 1)
    with self.assertRaises(AssertionError): to_uops_list([UOp(Ops.STORE, dtypes.void, (glbl0, idx, UOp.const(dtypes.int, 42), bad_gate))])
@@ -727,7 +727,7 @@ class TestLoadStoreFolding(unittest.TestCase):
  def test_gated_load_gep_preserves_alt(self):
    """Test that LOAD(GEP, alt) preserves alt value after rewrite"""
    from tinygrad.codegen.late.devectorizer import load_store_folding
-    buf = UOp(Ops.DEFINE_GLOBAL, dtypes.float.vec(4).ptr(), (), 0)
+    buf = UOp(Ops.PARAM, dtypes.float.vec(4).ptr(), (), 0)
    idx = UOp.const(dtypes.int, 0)
    gate = UOp.const(dtypes.bool, True)
    gated_index = buf.index(idx, gate)
@@ -745,8 +745,8 @@ class TestLoadStoreFolding(unittest.TestCase):
  def test_gated_load_ptrcat_preserves_alt(self):
    """Test that LOAD(PTRCAT, alt) preserves alt value after rewrite"""
    from tinygrad.codegen.late.devectorizer import load_store_folding
-    buf1 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
-    buf2 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 1)
+    buf1 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 0)
+    buf2 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 1)
    idx = UOp.const(dtypes.int, 0)
    idx1 = buf1.index(idx)
    idx2 = buf2.index(idx)
@@ -746,7 +746,7 @@ class TestSymbolic(unittest.TestCase):
    expr = cond.where(a, b).cast(dtypes.half)

    # TODO: copied from render, render does not support cast
-    glbl = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), arg=0)
+    glbl = UOp(Ops.PARAM, dtypes.int.ptr(), arg=0)
    uops = get_uops(UOp(Ops.STORE, dtypes.void, (glbl.index(UOp.const(dtypes.int, 0)), expr)).sink())
    rewritten_uop = [uop for uop in uops if uop.op is Ops.STORE][0].src[1]

@@ -1028,7 +1028,7 @@ class TestStoreLoadFolding(unittest.TestCase):
  """Tests for store(index, load(index)) -> NOOP rule. This rule matches patterns that EMERGE during simplification."""
  def test_store_load_folding(self):
    # store(idx, load(idx)) -> NOOP, including emergent patterns like store(idx, load(idx) + 0)
-    buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), arg=0)
+    buf = UOp(Ops.PARAM, dtypes.int.ptr(), arg=0)
    index = buf.index(UOp.const(dtypes.index, 0))
    # Direct: store(idx, load(idx)) -> NOOP
    self.assertEqual(graph_rewrite(index.store(index.load()), sym).op, Ops.NOOP)
@@ -1080,7 +1080,7 @@ class TestRangeSplitting(unittest.TestCase):
    from tinygrad.codegen.simplify import pm_split_ranges, pm_flatten_range
    r0 = UOp.range(uconst(8), 0)
    # create a simple expression using the range with mod: store range%2 to a buffer
-    buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), arg=0)
+    buf = UOp(Ops.PARAM, dtypes.int.ptr(), arg=0)
    val = (r0 % uconst(2)).cast(dtypes.int)
    store = UOp(Ops.STORE, dtypes.void, (buf.index(uconst(0)), val))
    sink = UOp(Ops.SINK, dtypes.void, (UOp(Ops.END, dtypes.void, (store, r0)),))
@@ -82,7 +82,7 @@ class TestVminVmaxProperties(unittest.TestCase):
  def test_vmin_vmax_multiplication_0_inf(self):
    # vmin and vmax for multiplication with a variable
    x = UOp.const(dtypes.float, 0.0)
-    y = UOp.load(UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0), UOp.const(dtypes.int, 0), dtype=dtypes.float)
+    y = UOp.load(UOp(Ops.PARAM, dtypes.float.ptr(), (), 0), UOp.const(dtypes.int, 0), dtype=dtypes.float)
    uop = x * y
    # TODO: these should be 0, but definitely should not be nan
    self.assertEqual(uop.vmin, -math.inf)
@@ -279,7 +279,7 @@ class TestVminVmaxVConst(unittest.TestCase):

  def test_vmin_vmax_vector_with_gep(self):
    # vmin and vmax for a vector constant of bool values
-    d1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 1)
+    d1 = UOp(Ops.PARAM, dtypes.int.ptr(), (), 1)
    idx = UOp.const(dtypes.int, 0)
    val = UOp(Ops.LOAD, dtypes.int.vec(2), (d1.index(idx),))
    uop = (val // 32).gep(0)
@@ -11,7 +11,7 @@ class TestValidateOOB(unittest.TestCase):
  # basic index patterns
  def test_const_index(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
+      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
      to_uops_list([buf.index(UOp.const(dtypes.int, 0), ptr=True).load(dtype=dtypes.int)])  # valid
      to_uops_list([buf.index(UOp.const(dtypes.int, 15), ptr=True).load(dtype=dtypes.int)])  # valid (last element)
      with self.assertRaises(RuntimeError):
@@ -21,7 +21,7 @@ class TestValidateOOB(unittest.TestCase):

  def test_variable_index(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
+      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
      to_uops_list([buf.index(Variable("i", 0, 15), ptr=True).load(dtype=dtypes.int)])  # valid
      with self.assertRaises(RuntimeError):
        to_uops_list([buf.index(Variable("i", 0, 20), ptr=True).load(dtype=dtypes.int)])  # oob
@@ -30,7 +30,7 @@ class TestValidateOOB(unittest.TestCase):

  def test_range_with_mask(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
+      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
      r = UOp.range(42, 0, AxisType.GLOBAL)
      to_uops_list([buf.index(r.valid(r < 16), ptr=True).load(dtype=dtypes.int)])  # valid
      with self.assertRaises(RuntimeError):
@@ -38,7 +38,7 @@ class TestValidateOOB(unittest.TestCase):

  def test_variable_with_mask(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
+      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
      v = Variable("v", -5, 80)
      to_uops_list([buf.index(v.valid((v >= 0) & (v < 16)), ptr=True).load(dtype=dtypes.int)])  # valid
      with self.assertRaises(RuntimeError):
@@ -46,7 +46,7 @@ class TestValidateOOB(unittest.TestCase):

  def test_gated_store(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
+      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
      v = Variable("v", 0, 20)
      to_uops_list([buf.index(v.valid(v < 16)).store(0)])  # valid
      with self.assertRaises(RuntimeError):
@@ -55,14 +55,14 @@ class TestValidateOOB(unittest.TestCase):
  # ALU ops in index
  def test_idiv(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
+      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
      to_uops_list([buf.index(UOp.range(32, 0, AxisType.GLOBAL) // 2, ptr=True).load(dtype=dtypes.int)])  # 0..15 valid
      with self.assertRaises(RuntimeError):
        to_uops_list([buf.index(UOp.range(34, 0, AxisType.GLOBAL) // 2, ptr=True).load(dtype=dtypes.int)])  # 0..16 oob

  def test_mod(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
+      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
      r = UOp.range(100, 0, AxisType.GLOBAL)
      to_uops_list([buf.index(r % 16, ptr=True).load(dtype=dtypes.int)])  # 0..15 valid
      with self.assertRaises(RuntimeError):
@@ -70,14 +70,14 @@ class TestValidateOOB(unittest.TestCase):

  def test_shr(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
+      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
      to_uops_list([buf.index(UOp.range(64, 0, AxisType.GLOBAL) >> 2, ptr=True).load(dtype=dtypes.int)])  # 0..15 valid
      with self.assertRaises(RuntimeError):
        to_uops_list([buf.index(UOp.range(128, 0, AxisType.GLOBAL) >> 2, ptr=True).load(dtype=dtypes.int)])  # 0..31 oob

  def test_shl(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(64), (), 0)
+      buf = UOp(Ops.PARAM, dtypes.int.ptr(64), (), 0)
      r = UOp.range(8, 0, AxisType.GLOBAL)
      to_uops_list([buf.index(r << 2, ptr=True).load(dtype=dtypes.int)])  # 0..28 valid
      with self.assertRaises(RuntimeError):
@@ -85,7 +85,7 @@ class TestValidateOOB(unittest.TestCase):

  def test_and(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
+      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
      r = UOp.range(100, 0, AxisType.GLOBAL)
      to_uops_list([buf.index(r & 15, ptr=True).load(dtype=dtypes.int)])  # 0..15 valid
      with self.assertRaises(RuntimeError):
@@ -93,14 +93,14 @@ class TestValidateOOB(unittest.TestCase):

  def test_max(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
+      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
      to_uops_list([buf.index(Variable("v", -10, 15).maximum(0), ptr=True).load(dtype=dtypes.int)])  # 0..15 valid
      with self.assertRaises(RuntimeError):
        to_uops_list([buf.index(Variable("v2", -10, 20).maximum(0), ptr=True).load(dtype=dtypes.int)])  # 0..20 oob

  def test_xor_in_mask(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
+      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
      r = UOp.range(32, 0, AxisType.GLOBAL)
      to_uops_list([buf.index(r.valid((r < 8) ^ ((r >= 8) & (r < 16))), ptr=True).load(dtype=dtypes.int)])  # 0..15 valid
      with self.assertRaises(RuntimeError):
@@ -109,22 +109,22 @@ class TestValidateOOB(unittest.TestCase):
  # cast patterns
  def test_float_cast_in_index(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
+      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
      r = UOp.range(20, 0)
      i = (r.cast(dtypes.float) * 0.68).trunc().cast(dtypes.int)
      to_uops_list([buf.index(i.valid((i >= 0) & (i < 16)), ptr=True).load(dtype=dtypes.int)])

  def test_bool_cast_in_mask(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1), (), 0)
+      buf = UOp(Ops.PARAM, dtypes.int.ptr(1), (), 0)
      r = UOp.range(20, 0)
      to_uops_list([buf.index(r.valid(r.cast(dtypes.bool).logical_not()), ptr=True).load(dtype=dtypes.int)])  # only r=0 valid

  # load result as index/mask
  def test_load_as_index(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
-      buf1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(64), (), 1)
+      buf0 = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
+      buf1 = UOp(Ops.PARAM, dtypes.int.ptr(64), (), 1)
      r = UOp.range(42, 0, AxisType.GLOBAL)
      ld0 = buf0.index(r.valid(r < 8), ptr=True).load(dtype=dtypes.int).cast(dtypes.index)
      to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 32)), ptr=True).load(dtype=dtypes.int)])  # valid
@@ -133,8 +133,8 @@ class TestValidateOOB(unittest.TestCase):

  def test_load_bool_as_mask(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf_bool = UOp(Ops.DEFINE_GLOBAL, dtypes.bool.ptr(16), (), 0)
-      buf_int = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(8), (), 1)
+      buf_bool = UOp(Ops.PARAM, dtypes.bool.ptr(16), (), 0)
+      buf_int = UOp(Ops.PARAM, dtypes.int.ptr(8), (), 1)
      gidx = UOp(Ops.SPECIAL, dtypes.index, (UOp.const(dtypes.index, 16),), "gidx0")
      ld_bool = buf_bool.index(gidx, ptr=True).load()
      with self.assertRaises(RuntimeError):
@@ -145,7 +145,7 @@ class TestValidateOOB(unittest.TestCase):
  def test_in_bounds_access_gated_local(self):
    with Context(CHECK_OOB=1):
      # Define buffers
-      gbuf = UOp(Ops.DEFINE_GLOBAL, dtypes.uint.ptr(400), (), 0)
+      gbuf = UOp(Ops.PARAM, dtypes.uint.ptr(400), (), 0)
      sbuf = UOp(Ops.DEFINE_LOCAL, dtypes.uint.ptr(8, addrspace=AddrSpace.LOCAL), (), "temp0")

      # Define indices, valids and barrier
@@ -169,8 +169,8 @@ class TestValidateOOB(unittest.TestCase):
  @unittest.skip("Bool load is not supported yet")
  def test_load_mask(self):
    with Context(CHECK_OOB=1):
-      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
-      mask = UOp(Ops.DEFINE_GLOBAL, dtypes.bool.ptr(16), (), 0)
+      glbl0 = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
+      mask = UOp(Ops.PARAM, dtypes.bool.ptr(16), (), 0)
      ridx = UOp.range(20, 0)
      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(UOp.const(ridx, ridx<16&mask), ptr=True)))
      to_uops_list([ld0])
@@ -188,9 +188,8 @@ class TestLinearizer(unittest.TestCase):
    assert any(x.op is Ops.DEFINE_LOCAL for x in stores[0].toposort())
    # the second store is to gds with no upcasts
    assert stores[1].src[1].dtype == dtypes.float
-    assert any(x.op is Ops.DEFINE_GLOBAL for x in stores[1].toposort())
+    assert any(x.op is Ops.PARAM for x in stores[1].toposort())

-  @unittest.skipIf(Device.DEFAULT=="CPU", "CPU splits the cat so cant upcast")
  def test_zero_fold(self):
    a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize()
    r = Tensor.stack(a, b)
@@ -450,7 +449,7 @@ class TestLinearizer(unittest.TestCase):
    def get_recursive(uop): return set.union(set(uop.src), [uop], *[get_recursive(v) for v in uop.src])
    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=opt).uops
    local_stores = [u for u in uops if u.op is Ops.STORE and any(x.op is Ops.DEFINE_LOCAL for x in get_recursive(u.src[0]))]
-    global_stores = [u for u in uops if u.op is Ops.STORE and any(x.op is Ops.DEFINE_GLOBAL for x in get_recursive(u.src[0]))]
+    global_stores = [u for u in uops if u.op is Ops.STORE and any(x.op is Ops.PARAM for x in get_recursive(u.src[0]))]
    barrier = [u for u in uops if u.op is Ops.BARRIER]
    assert len(barrier) == 1
    # check that the float4 cast collapses for all stores
@@ -11,16 +11,16 @@ from tinygrad.engine.realize import get_program
 class TestLinearizerFailure(unittest.TestCase):
  @unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL")
  def test_failure_beam_mnist(self):
-    c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.uchar.ptr(4014080), arg=0, src=())
+    c0 = UOp(Ops.PARAM, dtypes.uchar.ptr(4014080), arg=0, src=())
    c1 = UOp.range(UOp.const(dtypes.index, 512), 0, AxisType.GLOBAL)
    c2 = UOp.range(UOp.const(dtypes.index, 784), 1, AxisType.GLOBAL)
    c3 = UOp.range(UOp.const(dtypes.index, 10), 3, AxisType.GLOBAL)
-    c4 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(512), arg=1, src=())
+    c4 = UOp(Ops.PARAM, dtypes.int.ptr(512), arg=1, src=())
    c5 = c4.index(c1.valid(UOp.const(dtypes.bool, True)))
    c6 = UOp.range(UOp.const(dtypes.index, 6000), 1004, AxisType.REDUCE)
    c7 = UOp.range(UOp.const(dtypes.index, 3750), 2006, AxisType.REDUCE)
    c8 = UOp.range(UOp.const(dtypes.index, 16), 2007, AxisType.GROUP_REDUCE)
-    c9 = UOp(Ops.DEFINE_GLOBAL, dtypes.uchar.ptr(47040000), arg=2, src=())
+    c9 = UOp(Ops.PARAM, dtypes.uchar.ptr(47040000), arg=2, src=())
    c10 = c9.index((((c3*UOp.const(dtypes.index, 4704000))+c2)+(c6*UOp.const(dtypes.index, 784))).valid(UOp.const(dtypes.bool, True)))
    c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(dtypes.index, 6000))+c6)+((c7*UOp.const(dtypes.index, 16))+c8)).alu(Ops.CMPLT, UOp.const(dtypes.index, 59999)).where(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(dtypes.int, -1))).where(UOp.const(dtypes.uchar, 0), c10).reduce(c6, arg=Ops.ADD)
    c12 = c0.index((((c1*UOp.const(dtypes.index, 7840))+(c2*UOp.const(dtypes.index, 10)))+c3).valid(UOp.const(dtypes.bool, True))).store(c11).end(c1, c2, c3)
@@ -26,8 +26,8 @@ def _test_uop_result(inputs:list[Tensor], prg, local_size=None):

 def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp):
  dtype = alu_src_uops[0].dtype
-  a = UOp(Ops.DEFINE_GLOBAL, dtype.ptr(), (), 0)
-  b = UOp(Ops.DEFINE_GLOBAL, dtype.ptr(), (), 1)
+  a = UOp(Ops.PARAM, dtype.ptr(), (), 0)
+  b = UOp(Ops.PARAM, dtype.ptr(), (), 1)
  idx = UOp.const(dtypes.int, 0)
  ld = b.index(idx)
  alu = ld.alu(alu_op, *alu_src_uops)
@@ -39,7 +39,7 @@ def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp):
 class TestRendererFailures(unittest.TestCase):
  @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
  def test_gated_store_with_alu(self):
-    a = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
+    a = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0)
    gate_alu = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'lidx0')).ne(0)
    gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index(lidx0.valid(gate_alu)), UOp.const(dtypes.int, 1)))
    sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,))
@@ -49,7 +49,7 @@ class TestRendererFailures(unittest.TestCase):

  @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
  def test_gated_store_with_alu_2d(self):
-    a = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
+    a = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0)
    gate_alu_0 = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'lidx0')).ne(0)
    gate_alu_1 = (lidx1:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 2),), 'lidx1')).ne(0)
    gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index((lidx0+lidx1*4).valid(gate_alu_0&gate_alu_1)), UOp.const(dtypes.int, 1)))
@@ -94,7 +94,7 @@ class TestWGSLFailures(unittest.TestCase):
 class TestPTXFailures(unittest.TestCase):
  @unittest.skip("INDEX can only have a gate ALU parent, not an IF")
  def test_gated_store_with_if(self):
-    a = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
+    a = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0)
    gate_alu = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'lidx0')).ne(0)
    val = UOp.const(dtypes.int, 1)
    if_uop = UOp(Ops.IF, dtypes.void, (gate_alu,))
@@ -1092,6 +1092,18 @@ class TestSchedule(unittest.TestCase):
    np.testing.assert_allclose(out[0].numpy(), np.sqrt(np.square(x.numpy() - np_mu).sum(-1)/x.shape[-1]), atol=1e-4, rtol=1e-4)
    np.testing.assert_allclose(out[1].numpy(), np.sqrt(np.square(y.numpy() - np_mu).sum(-1)/y.shape[-1]), atol=1e-4, rtol=1e-4)

+  def test_cumsum_parallel_reduce_fused(self):
+    # two-stage cumsum + ops triggers parallel REDUCEs in one kernel that must share an END
+    step, num_steps = 513, 10
+    t = Tensor.arange(step).float().realize()
+    phase = t.cumsum()
+    tiled = phase.repeat((num_steps,)).reshape(num_steps, step)
+    pattern = Tensor([1,0,0,1,0,0,0,0,1,0]).reshape(num_steps, 1)
+    out = (tiled * pattern).flatten()
+    expected = np.tile(np.arange(step).astype(np.float32).cumsum(), num_steps).reshape(num_steps, step)
+    expected = (expected * np.array([1,0,0,1,0,0,0,0,1,0]).reshape(num_steps, 1)).flatten()
+    np.testing.assert_allclose(out.numpy(), expected, atol=1e-4, rtol=1e-4)
+
  def test_multimatmul_fusion(self):
    Tensor.manual_seed(0)
    a,b = Tensor.randn(4, 64).realize(), Tensor.rand(64,8).realize()
@@ -36,8 +36,8 @@ def uop(uops:list[UOp], op:Ops, dtype:Optional[DType], src:tuple[UOp, ...], arg:
 def _test_single_value(vals, op, dts):
  uops = []
  output_dtype = dtypes.bool if op in (Ops.CMPLT, Ops.CMPNE) else dts[-1]
-  buf_store = uop(uops, Ops.DEFINE_GLOBAL, output_dtype.ptr(), (), 0)
-  buf_loads = [uop(uops, Ops.DEFINE_GLOBAL, dtype.ptr(), (), i+1) for i,dtype in enumerate(dts)]
+  buf_store = uop(uops, Ops.PARAM, output_dtype.ptr(), (), 0)
+  buf_loads = [uop(uops, Ops.PARAM, dtype.ptr(), (), i+1) for i,dtype in enumerate(dts)]
  loads = (buf_loads[i].index(uop(uops, Ops.CONST, dtypes.int32, (), 0)) for i, dtype in enumerate(dts))
  alu = uop(uops, op, output_dtype, loads)
  out = uop(uops, Ops.STORE, dtypes.void, (buf_store.index(uop(uops, Ops.CONST, dtypes.int32, (), 0), ptr=True), alu))
@@ -52,7 +52,7 @@ def _test_single_value(vals, op, dts):
 def _test_single_value_const(vals, op, dts):
  uops = []
  output_dtype = dtypes.bool if op in (Ops.CMPLT, Ops.CMPNE) else dts[-1]
-  buf_store = uop(uops, Ops.DEFINE_GLOBAL, output_dtype.ptr(), (), 0)
+  buf_store = uop(uops, Ops.PARAM, output_dtype.ptr(), (), 0)
  loads = (uop(uops, Ops.CONST, dtype, [], a) for a,dtype in zip(vals, dts))
  alu = uop(uops, op, output_dtype, loads)
  out = buf_store[UOp.const(dtypes.int32, 0)].store(alu)
@@ -65,7 +65,7 @@ def _test_single_value_const(vals, op, dts):

 def _test_uops_result(output_dtype, uops, res):
  # uops = []
-  buf_store = uop(uops, Ops.DEFINE_GLOBAL, output_dtype.ptr(), (), 0)
+  buf_store = uop(uops, Ops.PARAM, output_dtype.ptr(), (), 0)
  # res = output_fn(uops)
  out = uop(uops, Ops.STORE, dtypes.void, (buf_store.index(uop(uops, Ops.CONST, dtypes.int32, (), 0)), res))
  buf = Buffer(Device.DEFAULT, 1, output_dtype).allocate()
@@ -273,7 +273,7 @@ class TestConstantFolding(unittest.TestCase):

 class TestGatedStoreRewrite(unittest.TestCase):
  def test_tiny_gate_store(self):
-    gmem = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
+    gmem = UOp(Ops.PARAM, dtypes.float.ptr(), (), 0)
    gidx0 = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'gidx0')
    gate = gidx0<UOp.const(dtypes.int, 1)
    idx = UOp(Ops.INDEX, dtypes.float.ptr(), (gmem, (gidx0 * UOp.const(dtypes.int, 2)).valid(gate)))
@@ -289,8 +289,8 @@ class TestGatedStoreRewrite(unittest.TestCase):
    self.assertIs(gated_uops[-1].op, Ops.STORE)

  def test_gate_some_stores(self):
-    gmem0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
-    gmem1 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 1)
+    gmem0 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 0)
+    gmem1 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 1)
    gidx0 = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'gidx0')
    idx = gidx0 * UOp.const(dtypes.int, 2)
    idx0 = UOp(Ops.INDEX, dtypes.float.ptr(), (gmem0, idx.valid(gidx0<UOp.const(dtypes.int, 1))))
@@ -309,8 +309,8 @@ class TestGatedStoreRewrite(unittest.TestCase):
  # scaled down version of TestLinearizerDumb.test_unmerged_ifs
  @unittest.skip("we don't merge ifs anymore")
  def test_merge_ifs_alt(self):
-    gmem0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
-    gmem1 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 1)
+    gmem0 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 0)
+    gmem1 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 1)
    gidx0 = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'gidx0')
    idx = gidx0*UOp.const(dtypes.int, 2)
    gate = gidx0<UOp.const(dtypes.int, 1)
@@ -380,7 +380,7 @@ class TestLocalAccess(unittest.TestCase):
 class TestFastIdiv(unittest.TestCase):
  def test_division_power_of_two(self):
    for dt in (dtypes.int32, dtypes.uint32):
-      g = UOp(Ops.DEFINE_GLOBAL, dt.ptr(), (), 0)
+      g = UOp(Ops.PARAM, dt.ptr(), (), 0)
      c = UOp.const(dt, 2)
      l = g.index(c)
      a = UOp(Ops.IDIV, dt, (l, c))
@@ -392,7 +392,7 @@ class TestFastIdiv(unittest.TestCase):

  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU doesn't support long")
  def test_fast_idiv_and_mod(self):
-    g = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(), (), 0)
+    g = UOp(Ops.PARAM, dtypes.uint32.ptr(), (), 0)
    c = UOp.const(dtypes.uint, 3)
    l = g.index(c)
    a = UOp(Ops.IDIV, dtypes.uint, (l, c))
@@ -420,7 +420,7 @@ class TestFastIdiv(unittest.TestCase):
  @unittest.expectedFailure
  def test_fast_idiv_overflow(self):
    # This will be possible with a slightly different method for fast_idiv
-    g = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(), (), 0)
+    g = UOp(Ops.PARAM, dtypes.uint32.ptr(), (), 0)
    c = UOp.const(dtypes.uint, 7)
    l = UOp(Ops.LOAD, dtypes.uint, (g.index(c),))
    a = UOp(Ops.IDIV, dtypes.uint, (l, c))
@@ -431,7 +431,7 @@ class TestFastIdiv(unittest.TestCase):
    self.assertNotIn(Ops.IDIV, ops)

  def test_disable_fast_idiv(self):
-    g = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(), (), 0)
+    g = UOp(Ops.PARAM, dtypes.uint32.ptr(), (), 0)
    c = UOp.const(dtypes.uint, 3)
    l = g.index(c)
    a = UOp(Ops.IDIV, dtypes.uint, (l, c))
@@ -445,7 +445,7 @@ class TestFastIdiv(unittest.TestCase):
@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "This only tests assembly backends")
 class TestAssembly(unittest.TestCase):
  def test_bitshift_left(self):
-    g1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 0)
+    g1 = UOp(Ops.PARAM, dtypes.int32.ptr(), (), 0)
    c1 = UOp.const(dtypes.int, 2)
    c2 = UOp.const(dtypes.int, 3)
    l1 = g1.index(c1)
@@ -471,7 +471,7 @@ class TestAssembly(unittest.TestCase):
    self.assertEqual(len([x.op for x in uops if x.op is Ops.MULACC]), 4)

  def test_use_cmpeq(self):
-    g = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(), (), 0)
+    g = UOp(Ops.PARAM, dtypes.uint32.ptr(), (), 0)
    c = UOp.const(dtypes.uint, 7)
    comp = g.index(c).ne(c).ne(True)
    uops = to_uops_list([comp], ren=Device[Device.DEFAULT].renderer)
@@ -507,7 +507,7 @@ class TestUOpMethod(unittest.TestCase):
    self.assertEqual((gidx0*3+1).const_factor(), 1)

  def test_replace(self):
-    x = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
+    x = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0)
    self.assertIs(x.replace(arg=None).arg, None)
    with self.assertRaises(AssertionError): x.replace(field="a")

@@ -137,7 +137,7 @@ class TestUOpsStats(unittest.TestCase):

  #MULACC should have the same stats as MUL + ADD
  def test_mulacc(self):
-    globl = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), tuple())
+    globl = UOp(Ops.PARAM, dtypes.int.ptr(), tuple())
    o1 = UOp(Ops.CONST, dtypes.int, tuple(), 1)
    o2 = UOp(Ops.CONST, dtypes.int, tuple(), 2)
    u1 = globl.index(o1)
@@ -147,7 +147,7 @@ class TestUOpsStats(unittest.TestCase):
    u5 = UOp(Ops.ADD, dtypes.int, (u4,u3))
    uops = list(u5.toposort())

-    globl = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), tuple())
+    globl = UOp(Ops.PARAM, dtypes.int.ptr(), tuple())
    o1 = UOp(Ops.CONST, dtypes.int, tuple(), 1)
    o2 = UOp(Ops.CONST, dtypes.int, tuple(), 2)
    u1 = globl.index(o1)
@@ -750,6 +750,35 @@ class TestTK(unittest.TestCase):

    fa_jitted = TinyJit(flash_attention)

+    for _ in range(10):
+      st = time.perf_counter()
+      out = fa_jitted(q, k, v, is_causal=False)
+      et = time.perf_counter() - st
+      attn_flops = 2 * B * H * N * N * D + \
+                   4 * B * H * N * N + \
+                   2 * B * H * N * N * D
+      print(f"{attn_flops/(et*1e9):2f} GFLOPS")
+    out = out.float().transpose(1, 2)
+
+    ref = q.scaled_dot_product_attention(k, v, is_causal=False, enable_gqa=True).float().transpose(1, 2)
+
+    np.testing.assert_allclose(out.numpy(), ref.numpy(), atol=2e-2, rtol=2e-2)
+
+  def test_fast_fa_causal(self):
+    from extra.thunder.tiny.fa import flash_attention
+
+    B, N, H, H_KV, D = 2, 8192, 32, 8, 128
+
+    with Context(DEBUG=0):
+      q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16).contiguous()
+      k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
+      v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
+      Tensor.realize(q, k, v)
+
+    q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
+
+    fa_jitted = TinyJit(flash_attention)
+
    for _ in range(10):
      st = time.perf_counter()
      out = fa_jitted(q, k, v, is_causal=True)
@@ -838,7 +867,7 @@ class TestTK(unittest.TestCase):

    np.testing.assert_allclose(q.grad.numpy(), q_ref.grad.numpy(), atol=2e-2, rtol=2e-2)
    np.testing.assert_allclose(v.grad.numpy(), v_ref.grad.numpy(), atol=2e-2, rtol=2e-2)
-    np.testing.assert_allclose(k.grad.numpy(), k_ref.grad.numpy(), atol=5e-2, rtol=2e-2)
+    np.testing.assert_allclose(k.grad.numpy(), k_ref.grad.numpy(), atol=6e-2, rtol=2e-2)

  def test_fast_fa_bwd_causal_jitted(self):
    from extra.thunder.tiny.fa import flash_attention
@@ -36,6 +36,22 @@ class TestCall(unittest.TestCase):
    np.testing.assert_allclose(a.grad.numpy(), gt_a_grad, rtol=1e-5)
    np.testing.assert_allclose(b.grad.numpy(), gt_b_grad, rtol=1e-5)

+  def test_call_plus_backward_auto(self):
+    a = Tensor.ones(10, 10, requires_grad=True)
+    b = Tensor.ones(10, 10, requires_grad=True)
+
+    (a+b).mean().backward()
+    gt_a_grad = a.grad.numpy()
+    gt_b_grad = b.grad.numpy()
+    a.grad, b.grad = None, None
+
+    plus_fxn = UOp.param(0, dtypes.float, (10,10)) + UOp.param(1, dtypes.float, (10,10))
+    c = Tensor.call(a, b, fxn=plus_fxn)
+    c.mean().backward()
+
+    np.testing.assert_allclose(a.grad.numpy(), gt_a_grad, rtol=1e-5)
+    np.testing.assert_allclose(b.grad.numpy(), gt_b_grad, rtol=1e-5)
+
  def test_call_gemm(self):
    M, K, N = 4, 8, 4
    a = Tensor.randn(M, K)
@@ -58,5 +74,23 @@ class TestCall(unittest.TestCase):

    np.testing.assert_allclose(c.numpy(), a.numpy() @ b.numpy(), rtol=1e-5, atol=1e-6)

+  def test_call_complex_backward_auto(self):
+    # complex chain: (a*b + a).exp2() * b.reciprocal() - tests mul, add, exp2, reciprocal, param reuse
+    a = Tensor.randn(10, 10, requires_grad=True)
+    b = Tensor.randn(10, 10, requires_grad=True) + 2  # avoid div by zero
+    Tensor.realize(a, b)
+
+    ((a*b + a).exp2() * b.reciprocal()).mean().backward()
+    gt_a_grad, gt_b_grad = a.grad.numpy(), b.grad.numpy()
+    a.grad, b.grad = None, None
+
+    p0, p1 = UOp.param(0, dtypes.float, (10,10)), UOp.param(1, dtypes.float, (10,10))
+    complex_fxn = (p0*p1 + p0).exp2() * p1.reciprocal()
+    c = Tensor.call(a, b, fxn=complex_fxn)
+    c.mean().backward()
+
+    np.testing.assert_allclose(a.grad.numpy(), gt_a_grad, rtol=1e-5)
+    np.testing.assert_allclose(b.grad.numpy(), gt_b_grad, rtol=1e-5)
+
 if __name__ == '__main__':
  unittest.main()
@@ -16,7 +16,7 @@ from tinygrad.codegen.late.expander import expander, pm_pre_expander, pm_group_f
 from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_indexing, devectorize, pm_reduce, \
  ReduceContext, correct_load_store, pm_render, pm_add_loads
 from tinygrad.codegen.opt.postrange import apply_opts, make_images
-from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_split_ranges, pm_load_collapse, pm_split_store
+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
 from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow, linearize

@@ -50,9 +50,6 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -
    # optimize (schedule) the AST
    sink = graph_rewrite(sink, pm_simplify_ranges, name="simplify ranges")

-    # split store range (only on CPU for now)
-    sink = graph_rewrite(sink, pm_split_store, ctx=ren.device, name="cut store ranges")
-
    # create image buffers
    sink = make_images(sink, ren)

@@ -1,7 +1,7 @@
 from typing import Any, cast
 import functools, operator, itertools
 from collections import defaultdict
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from tinygrad.dtype import dtypes, ImageDType, DType, AddrSpace, Invalid, PtrDType
 from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, GroupOp, identity_element
 from tinygrad.uop.symbolic import uop_given_valid, parse_valid, invalid_gate
@@ -299,6 +299,8 @@ pm_render = PatternMatcher([
@dataclass
 class ReduceContext:
  acc_num: int = 0
+  # track ENDs by range for merging parallel reduces
+  range_to_ends: dict[tuple[UOp, ...], list[UOp]] = field(default_factory=dict)

 def horizontal_reduce(inp:UOp, out_dtype:DType) -> list[UOp]:
  # if this has a horizontal reduction component, do that first
@@ -319,17 +321,24 @@ def reduce_to_acc(ctx:ReduceContext, red:UOp):
    input_ranges = tuple([x for x in topo if x.op is Ops.RANGE and x not in reduce_range and x not in ended_ranges])
    identity = red.const(red.dtype, identity_element(red.arg, red.dtype.scalar()))
    acc = UOp(Ops.DEFINE_REG, red.dtype.ptr(size=1, addrspace=AddrSpace.REG), arg=ctx.acc_num)
-    acc_init = acc.after(*input_ranges).index(UOp.const(dtypes.int, 0)).store(identity) if len(input_ranges) else \
-               acc.index(UOp.const(dtypes.int, 0)).store(identity)
+    acc_init = acc.after(*input_ranges).index(UOp.const(dtypes.int, 0)).store(identity)
    lst = [acc.after(acc_init, *reduce_range).index(UOp.const(dtypes.int, 0))] + lst  # put acc as the first element
    ctx.acc_num += 1
  ret = functools.reduce(lambda x,y: x.alu(red.arg, y), lst)
  if len(reduce_range) == 0: return ret
-  return acc.after(acc.index(UOp.const(dtypes.int, 0)).store(ret).end(*reduce_range)).index(UOp.const(dtypes.int, 0))
+  end = acc.index(UOp.const(dtypes.int, 0)).store(ret).end(*reduce_range)
+  ctx.range_to_ends.setdefault(reduce_range, []).append(end)
+  return acc.after(end).index(UOp.const(dtypes.int, 0))
+
+def merge_reduce_ends(ctx:ReduceContext, sink:UOp):
+  # merge ENDs that share the same range
+  subs = {e: UOp.group(*(e.src[0] for e in ends)).end(*r) for r, ends in ctx.range_to_ends.items() if len(ends) > 1 for e in ends}
+  return sink.substitute(subs) if subs else None

 pm_reduce = PatternMatcher([
-  # REDUCE -> DEFINE_ACC+ASSIGN
+  # REDUCE -> DEFINE_ACC+ASSIGN, then merge ENDs with same range
  (UPat(Ops.REDUCE, name="red"), reduce_to_acc),
+  (UPat(Ops.SINK, name="sink"), merge_reduce_ends),
  # tensor core built in accumulate
  (UPat(Ops.WMMA, name="wmma") + UPat.var("add"),
    lambda add, wmma: UOp(wmma.op, wmma.dtype, (wmma.src[0], wmma.src[1], wmma.src[2]+add), wmma.arg)),
@@ -26,7 +26,7 @@ def linearize(sink:UOp) -> list[UOp]:
    extra = None
    match u.op:
      # the order and placement of these defines is important
-      case Ops.DEFINE_GLOBAL: priority, extra = -20, u.arg
+      case Ops.PARAM: priority, extra = -20, u.arg
      case Ops.DEFINE_VAR: priority, extra = -19, u.arg
      case Ops.DEFINE_LOCAL: priority = -18
      case Ops.DEFINE_REG: priority = -17
@@ -331,7 +331,7 @@ class Scheduler:
  def group_for_reduces(self) -> int: return len(self.axes_of(AxisType.GROUP_REDUCE))

 def bufs_from_ast(ast:UOp, dname:str) -> list[Buffer]:
-  glbls = sorted([x for x in ast.backward_slice if x.op is Ops.DEFINE_GLOBAL], key=lambda x: x.arg)
+  glbls = sorted([x for x in ast.backward_slice if x.op is Ops.PARAM], key=lambda x: x.arg)
  return [Buffer(dname, x.ptrdtype.size, x.dtype.base if not isinstance(x.dtype, ImageDType) else x.dtype) for x in glbls]

 def apply_opts(ast:UOp, ren:Renderer) -> UOp:
@@ -362,7 +362,7 @@ def make_images(ast:UOp, ren:Renderer) -> UOp:
        ctx[dg.arg] = dt
        return dg.replace(dtype=dtypes.imagef((1, dt.size // 4, 4), dt.nbytes()))

-    ast = graph_rewrite(ast, PatternMatcher([(UPat(Ops.DEFINE_GLOBAL, name="dg"), make_image)]), ctx=dg_types, name="create image buffers")
+    ast = graph_rewrite(ast, PatternMatcher([(UPat(Ops.PARAM, name="dg"), make_image)]), ctx=dg_types, name="create image buffers")

    # undo unfoldable stores
    def undo_image_store(ctx, st, idx, dg):
@@ -370,6 +370,6 @@ def make_images(ast:UOp, ren:Renderer) -> UOp:
        return st.replace(src=(idx.replace(src=(dg.replace(dtype=ctx[dg.arg]),)+idx.src[1:]),)+st.src[1:])

    ast = graph_rewrite(ast, PatternMatcher([
-      (UPat(Ops.DEFINE_GLOBAL, name="dg").index(UPat(), name="idx").store(UPat(), name="st"), undo_image_store)
+      (UPat(Ops.PARAM, name="dg").index(UPat(), name="idx").store(UPat(), name="st"), undo_image_store)
    ]), ctx=dg_types, name="remove unfoldable image stores")
  return ast
@@ -1,7 +1,7 @@
 import itertools
 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, graph_rewrite, _substitute, range_start
 from tinygrad.uop.symbolic import symbolic
-from tinygrad.helpers import partition, dedup
+from tinygrad.helpers import partition
 from tinygrad.dtype import dtypes, ImageDType

 def flatten_range(r:UOp) -> UOp|None:
@@ -126,7 +126,7 @@ def reduce_collapse(red:UOp, u:UOp, pm:PatternMatcher=pm_reduce_collapse) -> UOp
    replaces: dict[UOp, UOp] = {}
    for u in included:
      for s in u.src:
-        if s in included or s in replaces or s.op in {Ops.CONST, Ops.VCONST, Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_VAR}: continue
+        if s in included or s in replaces or s.op in {Ops.CONST, Ops.VCONST, Ops.PARAM, Ops.DEFINE_LOCAL, Ops.DEFINE_VAR}: continue
        replaces[s] = UOp(Ops.DEFINE_VAR, dtype=s.dtype, arg=(f'in{len(replaces)}', s.vmin, s.vmax))
    collapse_fxn = u.substitute(replaces).reduce(r, arg=Ops.ADD)
    sink = graph_rewrite(collapse_fxn, pm, name="reduce_collapse")
@@ -147,16 +147,3 @@ pm_load_collapse = PatternMatcher([
  # we want to make sure we dont do math on a loaded index since that can cause overflow, this undoes the rule in pm_reduce_load_collapse
  ((UPat.var("x", dtypes.index)+UPat.var("y"))<UPat.var("c"), lambda x,y,c: x < c-y if no_load(y) and no_load(c) and not no_load(x) else None),
 ])
-
-def cut_store_range(ctx:str, store:UOp, r:UOp) -> UOp|None:
-  # only cut ranges on CPU for now
-  if r.src[0].op is not Ops.CONST or ctx!="CPU": return None
-  if not (cuts:=[c.src[1].arg for c in store.get_consumer_map()[r] if c.op is Ops.CMPLT and r is c.src[0] and c.src[1].op is Ops.CONST]): return None
-  cuts = sorted(dedup([0] + cuts + [r.src[0].arg]))
-  ranges = [UOp.range((end-start), *(r.arg[0:-1]+(i,r.arg[-1]))) for i,(start,end) in enumerate(zip(cuts[:-1], cuts[1:]))]
-
-  return UOp.group(*[store.substitute({r: new_r+start}).end(new_r) for new_r, start in zip(ranges, cuts[:-1])])
-
-pm_split_store = pm_flatten_range+PatternMatcher([
-  (UPat(Ops.END, src=(UPat(Ops.STORE, name="store"), UPat.var("r"))), cut_store_range),
-])
@@ -2,7 +2,7 @@ import time
 from typing import cast
 from collections import deque
 from tinygrad.uop.ops import UOp, Ops, buffers, UOpMetaClass, track_rewrites, PatternMatcher, UPat, graph_rewrite, graph_rewrite_map, Kernel
-from tinygrad.uop.spec import type_verify, tensor_spec
+from tinygrad.uop.spec import type_verify, tensor_spec, kernel_spec
 from tinygrad.device import Buffer, MultiBuffer
 from tinygrad.helpers import DEBUG, cpu_profile, TracingKey, SPEC, flatten, pluralize, SCACHE, Metadata
 from tinygrad.engine.realize import ExecItem
@@ -144,7 +144,7 @@ def complete_create_schedule_with_vars(big_sink:UOp) -> tuple[dict[UOp, UOp], li

  if not SCACHE or (sc_ret:=schedule_cache.get(sched_cache_key, None)) is None:
    # verify Tensors match the spec (on big_sink, we only need to do this if cache misses)
-    if SPEC: type_verify(big_sink, tensor_spec)
+    if SPEC: type_verify(big_sink, tensor_spec+kernel_spec)

    # hack to preserve metadata
    graph_rewrite_map(big_sink, pm_pre_sched_cache, ctx=({}, {}), name="preserve metadata")
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
geohot	308eb13eae	codegen	2026-02-03 18:09:15 +08:00
geohot	1b1d81e3d5	allow kernel_spec in tensor for custom kernels	2026-02-03 18:01:07 +08:00
geohot	50b7b283dc	allow after on param	2026-02-03 17:41:54 +08:00
geohot	39da624581	new custom_kernel function in tensor	2026-02-03 16:28:53 +08:00
George HotzandGitHub	dd2de4f838	rename all DEFINE_GLOBAL to PARAM (#14511 )	2026-02-03 15:09:38 +08:00
George HotzandGitHub	dc77b3318b	move files that pass with NULL=1 to test/null (#14508 ) * move files that pass with NULL=1 to test/null * fix windows * cpu 0 * bugfix + durations	2026-02-03 13:52:36 +08:00
George HotzandGitHub	888819ee09	call autodiff gradient (#14510 )	2026-02-03 13:51:02 +08:00
wozeparrotandGitHub	bbcd3d67a3	fa: faster (#14453 )	2026-02-02 21:34:17 -08:00
sirhcmandGitHub	e579613b90	IR3 has aux (#14509 )	2026-02-02 23:46:41 -05:00
George HotzandGitHub	85c7b23160	add pytest -nauto to benchmark for mac (#14458 ) * add pytest -nauto to benchmark * 3 minute timeout * 3 min * setup env * comment * fresh db * in the pyenv	2026-02-03 12:26:09 +08:00
sirhcmandGitHub	a5d7eb37db	IR3 works on versions earlier than 3.14 (#14507 )	2026-02-02 23:10:19 -05:00
George HotzandGitHub	33c886cafa	disable copyout on NULL backend by default (#14506 ) * disable copyout on NULL backend * gate it * allow copyout on some tests	2026-02-03 11:57:47 +08:00
chenyuandGitHub	3c5845e8a5	remove cut_store_range (#14505 ) special scheduling for CPU	2026-02-02 21:58:36 -05:00
chenyuandGitHub	4f2e7aed24	fix multiple REDUCE on same RANGE (#14504 ) each RANGE maps to one END, but reduce_to_acc is local and would not know this	2026-02-02 20:42:09 -05:00
chenyuandGitHub	93c41a78fa	clean up NOOP [pr] (#14503 ) should not be used as a COPY, started with removing from ALWAYS_RUN_OPS	2026-02-02 19:46:45 -05:00
chenyuandGitHub	66d2b02f11	delete files that depends on extra.optimization.helpers (#14499 )	2026-02-02 13:33:33 -05:00
George HotzandGitHub	ec0398fceb	test amd gpu crashes (#14459 ) * test amd gpu crashes * cleanup * less sketch tests	2026-02-02 18:57:47 +03:00
nimlgenandGitHub	6e4238c016	amd: recovery (#14461 ) * rec * ? * rv * cleaner * post merge * not used * um * clnr * x * x * d * move	2026-02-02 18:57:35 +03:00
chenyuandGitHub	61ca19ff24	after with empty src is self [pr] (#14496 )	2026-02-02 10:19:05 -05:00
George HotzandGitHub	6e958dbfd4	assembly/amd: add RDNA4 support to emulator (#14341 ) * start new rdna4 * work * plus works * more pass * rdna4 * assembly/amd: fix RDNA4 emulator for float16 and VOP3 clamp * stale * rev * rr * rdna4 emu tests * cleanup * cleanup * simp * works * better factorizaion * hacks * fix mockgpu * guard both * cleaner * gate * bug fix and a few tests * all test_tiny	2026-02-02 21:35:59 +08:00
chenyuandGitHub	a908f447d5	remove disk special case in mstack_early_shrink [pr] (#14494 )	2026-02-02 08:34:45 -05:00
qazalandGitHub	965940dd00	sqtt: update examples after event field change (#14493 ) * regen sqtt examples * cdna * rdna4 * packet counts for rdna3 * sqttmap work	2026-02-02 21:39:48 +09:00