Merge branch 'master' into mac_pytest

after with empty src is self [pr] (#14496 )
assembly/amd: add RDNA4 support to emulator (#14341 )
2026-02-02 23:40:03 +08:00 · 2026-02-02 10:19:05 -05:00 · 2026-02-02 21:35:59 +08:00 · 2026-02-02 08:34:45 -05:00 · 2026-02-02 21:39:48 +09:00 · 2026-02-02 16:02:00 +08:00
119 changed files with 15233 additions and 3104 deletions
@@ -56,7 +56,15 @@ runs:

    # **** Caching packages ****

+    - name: Cache Python packages (PR)
+      if: github.event_name == 'pull_request'
+      id: restore-venv-pr
+      uses: actions/cache/restore@v4
+      with:
+        path: ${{ github.workspace }}/.venv
+        key: venv-${{ runner.os }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
    - name: Cache Python packages
+      if: github.event_name != 'pull_request'
      id: restore-venv
      uses: actions/cache@v4
      with:
@@ -65,23 +73,23 @@ runs:

    # **** Caching downloads ****

-    - name: Cache downloads (Linux)
-      if: inputs.key != '' && runner.os == 'Linux'
-      uses: actions/cache@v4
+    - name: Cache downloads (PR)
+      if: inputs.key != '' && github.event_name == 'pull_request'
+      uses: actions/cache/restore@v4
      with:
-        path: ~/.cache/tinygrad/downloads/
+        path: ${{ runner.os == 'Linux' && '~/.cache/tinygrad/downloads/' || '~/Library/Caches/tinygrad/downloads/' }}
        key: downloads-${{ github.job }}-${{ inputs.key }}-${{ env.CACHE_VERSION }}
-    - name: Cache downloads (macOS)
-      if: inputs.key != '' && runner.os == 'macOS'
+    - name: Cache downloads
+      if: inputs.key != '' && github.event_name != 'pull_request'
      uses: actions/cache@v4
      with:
-        path: ~/Library/Caches/tinygrad/downloads/
+        path: ${{ runner.os == 'Linux' && '~/.cache/tinygrad/downloads/' || '~/Library/Caches/tinygrad/downloads/' }}
        key: downloads-${{ github.job }}-${{ inputs.key }}-${{ env.CACHE_VERSION }}

    # **** Python deps ****

    - name: Install dependencies in venv (with extra)
-      if: inputs.deps != '' && steps.restore-venv.outputs.cache-hit != 'true'
+      if: inputs.deps != '' && steps.restore-venv-pr.outputs.cache-hit != 'true' && steps.restore-venv.outputs.cache-hit != 'true'
      shell: bash
      run: |
        python -m venv .venv
@@ -92,7 +100,7 @@ runs:
        fi
        python -m pip install -e ".[${{ inputs.deps }}]" ${{ inputs.pydeps }} --extra-index-url https://download.pytorch.org/whl/cpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/
    - name: Install dependencies in venv (without extra)
-      if: inputs.deps == '' && steps.restore-venv.outputs.cache-hit != 'true'
+      if: inputs.deps == '' && steps.restore-venv-pr.outputs.cache-hit != 'true' && steps.restore-venv.outputs.cache-hit != 'true'
      shell: bash
      run: |
        python -m venv .venv
@@ -182,8 +190,14 @@ runs:
        echo "pkgs=$pkgs" >> "$GITHUB_OUTPUT"
        echo "hash=$(echo -n "$pkgs" | sha256sum | cut -d' ' -f1)" >> "$GITHUB_OUTPUT"

+    - name: Cache apt (PR)
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true') && github.event_name == 'pull_request'
+      uses: actions/cache/restore@v4
+      with:
+        path: /var/cache/apt/archives/
+        key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
    - name: Cache apt
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true') && github.event_name != 'pull_request'
      uses: actions/cache@v4
      with:
        path: /var/cache/apt/archives/
@@ -239,8 +253,17 @@ runs:
        ln -s /opt/homebrew/opt/[email protected] /opt/homebrew/opt/boost || true
        ln -s /opt/homebrew/opt/boost/lib/libboost_atomic-mt.dylib /opt/homebrew/opt/boost/lib/libboost_atomic.dylib || true
        ln -s /opt/homebrew/opt/boost/lib/libboost_thread-mt.dylib /opt/homebrew/opt/boost/lib/libboost_thread.dylib || true
+    - name: Cache gpuocelot (PR)
+      if: inputs.ocelot == 'true' && github.event_name == 'pull_request'
+      id: cache-build-pr
+      uses: actions/cache/restore@v4
+      env:
+        cache-name: cache-gpuocelot-build-1
+      with:
+        path: ${{ github.workspace }}/gpuocelot/ocelot
+        key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.CACHE_VERSION }}
    - name: Cache gpuocelot
-      if: inputs.ocelot == 'true'
+      if: inputs.ocelot == 'true' && github.event_name != 'pull_request'
      id: cache-build
      uses: actions/cache@v4
      env:
@@ -249,7 +272,7 @@ runs:
        path: ${{ github.workspace }}/gpuocelot/ocelot
        key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.CACHE_VERSION }}
    - name: Clone/compile gpuocelot
-      if: inputs.ocelot == 'true' && steps.cache-build.outputs.cache-hit != 'true'
+      if: inputs.ocelot == 'true' && steps.cache-build-pr.outputs.cache-hit != 'true' && steps.cache-build.outputs.cache-hit != 'true'
      shell: bash
      run: |
        git clone --recurse-submodules https://github.com/gpuocelot/gpuocelot.git ${{ github.workspace }}/gpuocelot
@@ -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
+
  testmacbenchmark:
    name: Mac Benchmark
    env:
@@ -145,6 +175,10 @@ jobs:
      run: |
        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
+    - name: Kill stale pids
+      run: |
+        PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
+        PYTHONPATH=. ./extra/hcq/hcq_smi.py nv kill_pids
    - name: UsbGPU boot time
      run: sudo -E PYTHONPATH=. DEBUG=2 AM_RESET=1 AMD=1 AMD_IFACE=USB time python3.11 test/test_tiny.py TestTiny.test_plus
    - name: UsbGPU tiny tests
@@ -332,9 +366,9 @@ jobs:
    - name: Setcap to python
      run: ./extra/amdpci/setup_python_cap.sh
    - name: Remove amd modules
-      run: ./extra/hcq/hcq_smi.py amd rmmod
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
    - name: Kill stale pids
-      run: ./extra/hcq/hcq_smi.py amd kill_pids
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
    #- name: Insert amdgpu
    #  run: sudo modprobe amdgpu
    - name: Symlink models and datasets
@@ -444,9 +478,9 @@ jobs:
    - name: Setcap to python
      run: ./extra/amdpci/setup_python_cap.sh
    - name: Remove amd modules
-      run: ./extra/hcq/hcq_smi.py amd rmmod
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
    - name: Kill stale pids
-      run: ./extra/hcq/hcq_smi.py amd kill_pids
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
    - name: Symlink models and datasets
      run: |
        mkdir -p weights
@@ -496,9 +530,9 @@ jobs:
    - name: Setcap to python
      run: ./extra/amdpci/setup_python_cap.sh
    - name: Remove amd modules
-      run: ./extra/hcq/hcq_smi.py amd rmmod
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
    - name: Kill stale pids
-      run: ./extra/hcq/hcq_smi.py amd kill_pids
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
    - name: Symlink models and datasets
      run: |
        mkdir -p weights
@@ -587,9 +621,9 @@ jobs:
    - name: Setcap to python
      run: ./extra/amdpci/setup_python_cap.sh
    - name: Remove amd modules
-      run: ./extra/hcq/hcq_smi.py amd rmmod
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
    - name: Kill stale pids
-      run: ./extra/hcq/hcq_smi.py amd kill_pids
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
    - name: Symlink models and datasets
      run: |
        mkdir -p weights
@@ -651,9 +685,9 @@ jobs:
    - name: Setcap to python
      run: ./extra/amdpci/setup_python_cap.sh
    - name: Remove nv modules
-      run: ./extra/hcq/hcq_smi.py nv rmmod
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py nv rmmod
    - name: Kill stale pids
-      run: ./extra/hcq/hcq_smi.py nv kill_pids
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py nv kill_pids
    - name: Symlink models and datasets
      run: |
        mkdir -p weights
@@ -26,7 +26,7 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: llvm-speed
-        deps: testing_minimal
+        deps: testing_unit
        llvm: 'true'
    - name: Speed Test
      run: CPU=1 CPU_LLVM=1 python3 test/speed/external_test_speed_v_torch.py
@@ -98,7 +98,7 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: torch-backend-pillow-torchvision-et-pt
-        deps: testing_minimal
+        deps: testing_unit
        pydeps: "pillow torchvision expecttest"
        llvm: 'true'
    - name: Install ninja
@@ -134,7 +134,7 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: torch-backend-pillow-torchvision-et-pt
-        deps: testing_minimal
+        deps: testing_unit
        llvm: 'true'
    - name: Install ninja
      run: |
@@ -156,7 +156,7 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: be-minimal
-        deps: testing_minimal
+        deps: testing_unit
    - name: Test dtype with Python emulator
      run: DEBUG=1 PYTHON=1 python3 -m pytest -n=auto test/test_dtype.py test/test_dtype_alu.py
    - name: Test ops with Python emulator
@@ -239,6 +239,7 @@ jobs:
    - name: Run mypy with lineprecision report
      run: |
        python -m mypy --lineprecision-report .
+        grep -v autogen lineprecision.txt | awk 'NR>2 {lines+=$2; precise+=$3; imprecise+=$4; any+=$5; empty+=$6} END {t=lines-empty; printf "TOTAL: %d lines, %d precise (%.1f%%), %d imprecise (%.1f%%), %d any (%.1f%%)\n", t, precise, 100*precise/t, imprecise, 100*imprecise/t, any, 100*any/t}'
        cat lineprecision.txt
    - name: Run TYPED=1
      run: CHECK_OOB=0 DEV=CPU TYPED=1 python test/test_tiny.py
@@ -255,9 +256,10 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: unittest-13
-        pydeps: "pillow numpy ftfy regex pre-commit"
+        pydeps: "pillow ftfy regex pre-commit"
        deps: testing_unit
        llvm: 'true'
+        amd: 'true'
    - name: Run pre-commit test hooks
      run: SKIP=ruff,mypy pre-commit run --all-files
    - name: Check Device.DEFAULT
@@ -346,7 +348,7 @@ jobs:
        uses: ./.github/actions/setup-tinygrad
        with:
          key: gpu-image
-          deps: testing_minimal
+          deps: testing_unit
          opencl: 'true'
      - name: Test CL IMAGE=2 ops
        run: |
@@ -422,7 +424,7 @@ jobs:
        with:
          key: onnxoptc
          deps: testing
-          python-version: '3.11'
+          python-version: '3.12'
          llvm: 'true'
      - name: Test ONNX (CPU)
        run: CPU=1 CPU_LLVM=0 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
@@ -450,7 +452,7 @@ jobs:
          key: onnxoptl
          deps: testing
          pydeps: "tensorflow==2.19"
-          python-version: '3.11'
+          python-version: '3.12'
          opencl: 'true'
      - name: Test ONNX (CL)
        run: CL=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
@@ -524,7 +526,7 @@ jobs:
      with:
        key: metal
        deps: testing
-        python-version: '3.11'
+        python-version: '3.12'
    - name: Test models (Metal)
      run: METAL=1 python -m pytest -n=auto test/models --durations=20
    - name: Test LLaMA compile speed
@@ -543,7 +545,7 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: devectorize-minimal
-        deps: testing_minimal
+        deps: testing_unit
        pydeps: "pillow"
        llvm: "true"
    - name: Test LLVM=1 DEVECTORIZE=0
@@ -564,8 +566,8 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: dsp-minimal
-        deps: testing_minimal
-        pydeps: "onnx==1.18.0 onnxruntime pillow"
+        deps: testing_unit
+        pydeps: "onnx==1.18.0 onnxruntime"
        llvm: "true"
    - name: Set up Docker Buildx
      uses: docker/setup-buildx-action@v3
@@ -598,8 +600,8 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: webgpu-minimal
-        deps: testing_minimal
-        python-version: '3.11'
+        deps: testing_unit
+        python-version: '3.12'
        webgpu: 'true'
    - name: Check Device.DEFAULT (WEBGPU) and print some source
      run: |
@@ -632,7 +634,7 @@ jobs:
        uses: ./.github/actions/setup-tinygrad
        with:
          key: ${{ matrix.backend }}-minimal
-          deps: testing_minimal
+          deps: testing_unit
          amd: 'true'
          llvm: ${{ matrix.backend == 'amdllvm' && 'true' }}
      - name: Check Device.DEFAULT and print some source
@@ -674,9 +676,9 @@ jobs:
        uses: ./.github/actions/setup-tinygrad
        with:
          key: rdna3-emu
-          deps: testing_minimal
+          deps: testing_unit
          amd: 'true'
-          python-version: '3.13'
+          python-version: '3.14'
      - name: Verify AMD autogen is up to date
        run: |
          python -m extra.assembly.amd.generate
@@ -702,6 +704,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:
@@ -722,7 +726,7 @@ jobs:
        uses: ./.github/actions/setup-tinygrad
        with:
          key: ${{ matrix.backend }}-minimal
-          deps: testing_minimal
+          deps: testing_unit
          cuda: 'true'
          ocelot: 'true'
      - name: Set env
@@ -755,7 +759,7 @@ jobs:
        uses: ./.github/actions/setup-tinygrad
        with:
          key: ${{ matrix.backend }}-minimal
-          deps: testing_minimal
+          deps: testing_unit
          opencl: ${{ matrix.backend == 'opencl' && 'true' }}
          llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'lvp' }}
          mesa: ${{ matrix.backend == 'lvp' && 'true' }}
@@ -786,7 +790,7 @@ jobs:
      with:
        key: metal
        deps: testing
-        python-version: '3.11'
+        python-version: '3.12'
        amd: 'true'
        cuda: 'true'
        ocelot: 'true'
@@ -884,8 +888,7 @@ jobs:
        uses: ./.github/actions/setup-tinygrad
        with:
          key: macos-${{ matrix.backend }}-minimal
-          deps: testing_minimal
-          pydeps: "capstone"
+          deps: testing_unit
          llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'lvp' }}
          mesa: ${{ matrix.backend == 'lvp' && 'true' }}
      - name: Set env
@@ -952,7 +955,7 @@ jobs:
        uses: ./.github/actions/setup-tinygrad
        with:
          key: compile-${{ matrix.backend }}
-          deps: testing_minimal
+          deps: testing_unit
          mesa: ${{ (matrix.backend == 'ir3' || matrix.backend == 'nak') && 'true' }}
          python-version: '3.14'
      - name: Set env
@@ -3,7 +3,7 @@ from pathlib import Path
 import multiprocessing

 from tinygrad import Device, GlobalCounters, Tensor, TinyJit, dtypes
-from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, Profiling
+from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, Profiling, profile_marker
 from tinygrad.nn.state import get_parameters, get_state_dict, load_state_dict, safe_load, safe_save
 from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup, Adam, AdamW

@@ -1292,7 +1292,6 @@ def train_llama3():
  BASEDIR            = config["BASEDIR"]                = Path(getenv("BASEDIR", "/raid/datasets/c4/"))
  BS                 = config["BS"]                     = getenv("BS", 16)
  grad_acc           = config["GRADIENT_ACC_STEPS"]     = getenv("GRADIENT_ACC_STEPS", 1)
-  assert grad_acc == 1, f"{grad_acc=} is not supported"
  GBS                = config["GLOBAL_BATCH_SIZE"]      = BS * grad_acc
  SEED               = config["SEED"]                   = getenv("SEED", 5760)
  SEQLEN             = config["SEQLEN"]                 = getenv("SEQLEN", 8192)
@@ -1322,6 +1321,8 @@ def train_llama3():
  opt_base_learning_rate = LR
  opt_end_learning_rate = END_LR

+  Tensor.manual_seed(SEED)  # seed for weight initialization
+
  # ** init wandb **
  WANDB = getenv("WANDB")
  if WANDB:
@@ -1370,6 +1371,12 @@ def train_llama3():

  optim = AdamW(get_parameters(model), lr=0.0,
                b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay)
+
+  # init grads
+  for p in optim.params:
+    p.grad = p.zeros_like().contiguous().realize()
+  grads = [p.grad for p in optim.params]
+
  scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)

  if resume_ckpt := getenv("RESUME_CKPT"):
@@ -1382,9 +1389,7 @@ def train_llama3():
    load_state_dict(scheduler, safe_load(fn), realize=False)

  @TinyJit
-  @Tensor.train()
-  def train_step(model, tokens:Tensor):
-    optim.zero_grad()
+  def minibatch(tokens:Tensor):
    if (DP := getenv("DP", 1)) > 1:
      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
      tokens = tokens.shard(device, 0)
@@ -1394,27 +1399,40 @@ def train_llama3():
    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
    loss.backward()
+    assert all(p.grad is g for p,g in zip(optim.params, grads))
+    Tensor.realize(loss, *grads)
+    return loss
+
+  @TinyJit
+  def optim_step():
+    for p in optim.params:
+      p.grad.assign(p.grad / grad_acc)
+
    # L2 norm grad clip
    # https://github.com/NVIDIA/NeMo/blob/3368c3fc0b4a186ab33a1d68a504315100c0b2a6/nemo/collections/nlp/modules/common/megatron/clip_grads.py#L57
    # https://docs.pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_norm_.html
    if not getenv("DISABLE_GRAD_CLIP_NORM"):
      total_norm = Tensor(0.0, dtype=dtypes.float32, device=optim.params[0].device)
-      for p in optim.params:
-        total_norm += p.grad.float().square().sum()
-      total_norm = total_norm.sqrt().contiguous()
-      for p in optim.params:
-        p.grad = (p.grad * (opt_gradient_clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(p.grad.dtype)
+      for g in grads:
+        total_norm += g.float().square().sum()
+      total_norm = total_norm.sqrt().contiguous().realize()
+      for g in grads:
+        g.assign((g * (opt_gradient_clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(g.dtype)).realize()

    optim.step()
    scheduler.step()

+    for g in grads:
+      g.assign(g.zeros_like().contiguous()).realize()
+
    lr = optim.lr
-    loss.realize(lr)
-    return loss, lr
+    Tensor.realize(lr, *grads)
+
+    return lr

  @TinyJit
  @Tensor.train(False)
-  def eval_step(model, tokens:Tensor):
+  def eval_step(tokens:Tensor):
    if (DP := getenv("DP", 1)) > 1:
      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
      tokens = tokens.shard(device, 0)
@@ -1456,35 +1474,54 @@ def train_llama3():
  while i < MAX_STEPS:
    GlobalCounters.reset()
    if getenv("TRAIN", 1):
+      profile_marker(f"train @ {i}")
      st = time.perf_counter()
-      try: tokens = next(train_iter)
-      except StopIteration: break
-      dt = time.perf_counter()
-      loss, lr = train_step(model, tokens)
+
+      stopped = False
+      for _ in range(grad_acc):
+        ist = time.perf_counter()
+        try: tokens = next(train_iter)
+        except StopIteration:
+          stopped = True
+          break
+        dt = time.perf_counter()
+        loss = minibatch(tokens)
+      if stopped: break
+
+      gt = time.perf_counter()
+      lr = optim_step()
+      ot = time.perf_counter()
+
      loss = loss.float().item()
      lr = lr.item()

      et = time.perf_counter()
      step_time = et - st
-      dev_time = et - dt
-      data_time = dt - st
+      gbs_time = gt - st
+      optim_time = ot - gt
+      data_time = dt - ist
+      dev_time = step_time - data_time * grad_acc
      if BENCHMARK: step_times.append(step_time)

      i += 1
-      sequences_seen += tokens.shape[0]
+      sequences_seen += GBS

      mem_gb = GlobalCounters.mem_used / 1e9
      gflops = GlobalCounters.global_ops / 1e9 / dev_time
-      mfu = ((6 * num_params * SEQLEN * BS) / (dev_time * max(getenv("DP", 1), getenv("MP", 1)) * 2.3e15)) * 100
+      mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * max(getenv("DP", 1), getenv("MP", 1)) * 2.3e15)) * 100
      tqdm.write(
-          f"{i:5} {step_time:.3f} s run, {dev_time:.3f} s device, {data_time:.3f} s data, {loss:.4f} loss, {lr:.12f} LR, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")
+          f"{i:5} {step_time:.3f} s step, {gbs_time:.3f} s gbs, {optim_time:.3f} s optim, {data_time:.3f} s data, {loss:.4f} loss, " \
+          f"{lr:.12f} LR, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")

      if WANDB:
        wandb.log({
          "lr": lr, "train/loss": loss,
          "train/step_time": step_time,
+          "train/gbs_time": gbs_time,
+          "train/optim_time": optim_time,
          "train/dev_time": dev_time,
          "train/data_time": data_time,
+          "train/mem": mem_gb,
          "train/GFLOPS": gflops,
          "train/MFU": mfu,
          "train/sequences_seen": sequences_seen
@@ -1508,7 +1545,9 @@ def train_llama3():
              f"epoch global_mem: {GlobalCounters.global_mem:_}")

    if (sequences_seen % EVAL_FREQ == 0 and (i != 1 or EVAL_FREQ == 1)) or (BENCHMARK and i == BENCHMARK):
+      if EVAL_BS == 0: return
      tqdm.write(f"evaluating after {sequences_seen} sequences")
+      profile_marker(f"eval @ {i}")

      # run eval
      eval_losses = []
@@ -1516,8 +1555,8 @@ def train_llama3():
      tqdm.write(f"evaluating {5760//EVAL_BS} batches of {EVAL_BS} sequences")

      for j,tokens in tqdm(enumerate(eval_iter), total=EVAL_SAMPLES//EVAL_BS):
-        eval_losses += eval_step(model, tokens).tolist()
-        
+        eval_losses += eval_step(tokens).tolist()
+
        if BENCHMARK and (j+1) == min(BENCHMARK, EVAL_SAMPLES//EVAL_BS):
          return

@@ -1606,7 +1645,7 @@ def train_stable_diffusion():
    loss, out_lr = loss.detach().to("CPU"), optimizer.lr.to("CPU")
    Tensor.realize(loss, out_lr)
    return loss, out_lr
-    
+
  # checkpointing takes ~9 minutes without this, and ~1 minute with this
  @TinyJit
  def ckpt_to_cpu():
@@ -1645,7 +1684,7 @@ def train_stable_diffusion():
    if i == 3:
      for _ in range(3): ckpt_to_cpu() # do this at the beginning of run to prevent OOM surprises when checkpointing
      print("BEAM COMPLETE", flush=True) # allows wrapper script to detect BEAM search completion and retry if it failed
-      
+
    total_train_time = time.perf_counter() - train_start_time
    if WANDB:
      wandb.log({"train/loss": loss_item, "train/lr": lr_item, "train/loop_time_prev": loop_time, "train/dl_time": dl_time, "train/step": i,
@@ -10,7 +10,7 @@ export FLASH_ATTENTION=${FLASH_ATTENTION:-1}
 export ALL2ALL=${ALL2ALL:-1}

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

 export MODEL="llama3"
@@ -18,7 +18,7 @@ export BASEDIR="/raid/datasets/c4-8b/"
 export SMALL=1
 export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
 export EVAL_TARGET=3.3 EVAL_FREQ=12288
-export LR="4e-4" END_LR="4e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
+export LR="2.5e-4" END_LR="2.5e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
 export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
 export SAMPLES=$((MAX_STEPS * GBS))

@@ -2,15 +2,18 @@

 export PYTHONPATH="."
 export DEV=${DEV:-AMD}
+export EMULATE="AMD_CDNA4"
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000

 export DEBUG=${DEBUG:-0}
 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=1
+export DP=${DP:-8} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
 export GBS=$((BS * GRADIENT_ACC_STEPS))

 export MODEL="llama3"
@@ -18,13 +21,13 @@ export BASEDIR="/raid/datasets/c4-8b/"
 export SMALL=1
 export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
 export EVAL_TARGET=3.3 EVAL_FREQ=12288
-export LR="4e-4" END_LR="4e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
+export LR="2.5e-4" END_LR="2.5e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
 export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
 export SAMPLES=$((MAX_STEPS * GBS))

-export SEED=5760
+export SEED=${SEED:-5760}

-export JITBEAM=3
+export JITBEAM=${JITBEAM:-3}
 export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5

 python3 examples/mlperf/model_train.py
@@ -0,0 +1,9 @@
+#!/bin/bash
+export BENCHMARK=5
+export EVAL_BS=0
+export FAKEDATA=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
@@ -93,7 +93,7 @@ if __name__ == "__main__":
    forward: Any = None

  sub_steps = [
-    Step(name = "textModel", input = [Tensor.randn(1, 77)], forward = model.cond_stage_model.transformer.text_model),
+    Step(name = "textModel", input = [Tensor.randint(1, 77, low=0, high=49408, dtype=dtypes.int32)], forward = model.cond_stage_model.transformer.text_model),
    Step(name = "diffusor", input = [Tensor.randn(1, 77, 768), Tensor.randn(1, 77, 768), Tensor.randn(1,4,64,64), Tensor.rand(1), Tensor.randn(1), Tensor.randn(1), Tensor.randn(1)], forward = model),
    Step(name = "decoder", input = [Tensor.randn(1,4,64,64)], forward = model.decode),
    Step(name = "f16tof32", input = [Tensor.randn(2097120, dtype=dtypes.uint32)], forward = u32_to_f16)
@@ -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():
@@ -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)
@@ -402,17 +423,19 @@ class _Ctx:
    return UOp.sink(*self.scalar_stores(assigns, sdst_reg, sdst_size), *self.inc_pc())

  def compile_lane_pcode(self, op, inst) -> UOp:
-    """Compile READLANE/READFIRSTLANE/WRITELANE using pcode parser."""
+    """Compile cross-lane ops (READLANE/WRITELANE/PERMLANE) using pcode parser."""
    pcode = get_pcode(op)
    op_name = op.name if hasattr(op, 'name') else str(op)
    src0_off, vdst_off = self.inst_field(type(inst).src0), self.inst_field(type(inst).vdst)
    src0_reg = (src0_off >= _c(256)).where(src0_off - _c(256), _c(0))  # VGPR index or 0
    src1_off = self.inst_field(type(inst).src1) if hasattr(type(inst), 'src1') else None
+    src2_off = self.inst_field(type(inst).src2) if hasattr(type(inst), 'src2') else None
    exec_lo = self.rsgpr_dyn(_c(EXEC_LO.offset))
    srcs = {
      'SRC0': src0_reg, 'VDST': vdst_off, 'EXEC_LO': exec_lo, 'EXEC': exec_lo.cast(dtypes.uint64), '_vgpr': self.vgpr,
      'S0': self.rsrc_dyn(src0_off, _c(0, dtypes.int)) if 'WRITELANE' in op_name else src0_reg,
      'S1': self.rsrc_dyn(src1_off, _c(0, dtypes.int)) if src1_off is not None else _c(0),
+      'S2': self.rsrc_dyn(src2_off, _c(0, dtypes.int)) if src2_off is not None else _c(0),
    }
    _, assigns = parse_pcode(pcode, srcs)
    stores = []
@@ -427,7 +450,8 @@ class _Ctx:
    pcode = get_pcode(op)
    vcc_reg = sdst_reg if sdst_reg is not None else VCC_LO.offset
    if 'VCC' not in srcs: srcs['VCC'] = self.rsgpr_dyn(_c(vcc_reg))
-    srcs.update({'EXEC': exec_mask, 'SCC': self.rsgpr_dyn(_c(SCC.offset)), 'laneId': lane})
+    srcs.update({'EXEC': exec_mask, 'SCC': self.rsgpr_dyn(_c(SCC.offset)), 'laneId': lane,
+                 'ROUND_MODE': _c(0), 'ROUND_TOWARD_ZERO': _c(0)})  # rounding mode: 0=RNE, RTZ constant
    _, assigns = parse_pcode(pcode, srcs)

    raw_stores: list = []
@@ -479,13 +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 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))
@@ -497,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
@@ -549,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 == 'V_READFIRSTLANE_B32_E32': return ctx.compile_lane_pcode(inst.op, 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)
@@ -569,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)
@@ -583,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)
@@ -604,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

@@ -635,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)
@@ -644,18 +677,22 @@ def _compile_vop3(inst: VOP3, ctx: _Ctx) -> UOp:
  if op_name in ('V_READLANE_B32', 'V_READFIRSTLANE_B32', 'V_READFIRSTLANE_B32_E64', 'V_WRITELANE_B32'):
    return ctx.compile_lane_pcode(inst.op, inst)

+  # V_PERMLANE16_B32 / V_PERMLANEX16_B32: cross-lane swizzle via pcode
+  if 'PERMLANE16' in op_name or 'PERMLANEX16' in op_name:
+    return ctx.compile_lane_pcode(inst.op, inst)
+
  # VOP3 VOPC (v_cmp_*_e64) - delegate to unified VOPC handler
  if 'V_CMP' in op_name or 'V_CMPX' in op_name:
    return _compile_vopc(inst, ctx, opsel=opsel, abs_bits=getattr(inst, 'abs', 0) or 0, neg_bits=getattr(inst, 'neg', 0) or 0)

  # 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)
@@ -665,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
@@ -723,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']
@@ -756,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
@@ -787,7 +824,7 @@ def _compile_vop3p(inst: VOP3P, ctx: _Ctx) -> UOp:
    s0_mod = apply_neg_mix(apply_abs(src0, 1, 1, 1), 1, 1, 1)
    s1_mod = apply_neg_mix(apply_abs(src1, 2, 2, 2), 2, 2, 2)
    s2_mod = apply_neg_mix(apply_abs(src2, 4, 4, 4), 4, 4, 4)
-    srcs = {'S0': s0_mod, 'S1': s1_mod, 'S2': s2_mod,
+    srcs = {'S@0': s0_mod, 'S@1': s1_mod, 'S@2': s2_mod,
            'OPSEL_HI': UOp.const(dtypes.uint32, combined_opsel_hi), 'OPSEL': UOp.const(dtypes.uint32, opsel)}
  else:
    def get_half_bits(val: UOp, use_hi: bool, apply_neg: bool = False) -> UOp:
@@ -796,24 +833,28 @@ def _compile_vop3p(inst: VOP3P, ctx: _Ctx) -> UOp:
      return bits
    def build_remapped_src(src: UOp, opsel_lo_bit: int, opsel_hi_bit: int, neg_lo_bit: int, neg_hi_bit: int) -> UOp:
      return get_half_bits(src, bool(opsel_lo_bit), bool(neg_lo_bit)) | (get_half_bits(src, bool(opsel_hi_bit), bool(neg_hi_bit)) << UOp.const(dtypes.uint32, 16))
-    s0_new = build_remapped_src(src0, opsel & 1, opsel_hi & 1, neg & 1, neg_hi & 1)
-    s1_new = build_remapped_src(src1, opsel & 2, opsel_hi & 2, neg & 2, neg_hi & 2)
-    s2_new = build_remapped_src(src2, opsel & 4, 1 if opsel_hi2 else 0, neg & 4, neg_hi & 4)
-    srcs = {'S0': s0_new, 'S1': s1_new, 'S2': s2_new}
+    # DOT IU instructions use NEG bits for signed/unsigned selection, not fp16 negation
+    is_dot_iu = 'DOT' in op_name and 'IU' in op_name
+    n0, n1, n2, nh0, nh1, nh2 = (0, 0, 0, 0, 0, 0) if is_dot_iu else (neg & 1, neg & 2, neg & 4, neg_hi & 1, neg_hi & 2, neg_hi & 4)
+    srcs = {'S0': build_remapped_src(src0, opsel & 1, opsel_hi & 1, n0, nh0),
+            'S1': build_remapped_src(src1, opsel & 2, opsel_hi & 2, n1, nh1),
+            'S2': build_remapped_src(src2, opsel & 4, 1 if opsel_hi2 else 0, n2, nh2)}
+    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)
@@ -824,49 +865,64 @@ 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:
+    pcode = get_pcode(inst.op)
+    srcs = {'ADDR': addr_reg, 'DATA0': vdata_reg, 'VDST': vdst_reg, 'OFFSET': offset,
+            'EXEC': exec_mask.cast(dtypes.uint64), '_vgpr': ctx.vgpr}
+    _, assigns = parse_pcode(pcode, srcs)
+    stores = [ctx.vgpr.index(val[0].cast(dtypes.int)).store(val[1].cast(dtypes.uint32)) for dest, val in assigns if dest.startswith('VGPR[')]
+    return UOp.sink(*stores, *ctx.inc_pc())

  def make_addr(lane: UOp) -> UOp:
    if is_lds: return ctx.rvgpr_dyn(addr_reg, lane)
@@ -908,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

@@ -966,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,
 }

 # ═══════════════════════════════════════════════════════════════════════════════
@@ -979,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')

@@ -1004,21 +1077,21 @@ def _get_runner(inst_bytes: bytes):
  canonical_name = f"{_op_name(inst).lower()}_{base.to_bytes(size, 'little').hex()}"
  sink = sink.replace(arg=KernelInfo(name=canonical_name)).rtag(1)

-  with Context(NOOPT=1, CHECK_OOB=0, TUPLE_ORDER=0):
+  with Context(NOOPT=1, CHECK_OOB=0, TUPLE_ORDER=0, EMULATED_DTYPES=""):
    runner = get_runner('CPU', sink)
  _canonical_runner_cache.append((base, mask, size, runner))
  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}>"
@@ -1077,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
@@ -1107,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
@@ -1123,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")
@@ -94,13 +94,19 @@ def _trig_reduce(x, phase=0.0):
  return UOp(Ops.SIN, x.dtype, (x - n * _const(x.dtype, 6.283185307179586),))

 def _signext(val: UOp) -> UOp:
-  for bits, mask, ext in [(8, 0xFF, 0xFFFFFF00), (16, 0xFFFF, 0xFFFF0000)]:
+  for bits, mask, ext in [(4, 0xF, 0xFFFFFFF0), (8, 0xFF, 0xFFFFFF00), (16, 0xFFFF, 0xFFFF0000)]:
    if (val.op == Ops.AND and len(val.src) == 2 and val.src[1].op == Ops.CONST and val.src[1].arg == mask) or val.dtype.itemsize == bits // 8:
      v32 = val.cast(dtypes.uint32) if val.dtype != dtypes.uint32 else val
      sb = (v32 >> _u32(bits - 1)) & _u32(1)
      return sb.ne(_u32(0)).where(v32 | _u32(ext), v32).cast(dtypes.int)
  return val.cast(dtypes.int64) if val.dtype in (dtypes.int, dtypes.int32) else val

+def _signext_4bit(val: UOp) -> UOp:
+  """Sign extend a 4-bit value to 32-bit signed integer."""
+  v32 = val.cast(dtypes.uint32) if val.dtype != dtypes.uint32 else val
+  sb = (v32 >> _u32(3)) & _u32(1)  # sign bit at position 3
+  return sb.ne(_u32(0)).where(v32 | _u32(0xFFFFFFF0), v32).bitcast(dtypes.int)
+
 def _abs(val: UOp) -> UOp:
  if val.dtype not in (dtypes.float32, dtypes.float64, dtypes.half): return val
  _, _, _, _, shift = _float_info(val)
@@ -194,6 +200,17 @@ def _ff1(val: UOp, bits: int) -> UOp:
    result = cond.where(_const(dtypes.int, i), result)
  return result

+def _sad_u8(a: UOp, b: UOp, acc: UOp, masked: bool = False) -> UOp:
+  """Sum of absolute differences of 4 unsigned bytes + accumulator. If masked, skips bytes where a == 0."""
+  a, b, acc = a.cast(dtypes.uint32), b.cast(dtypes.uint32), acc.cast(dtypes.uint32)
+  result = acc
+  for i in range(4):
+    a_byte = (a >> _u32(i * 8)) & _u32(0xFF)
+    b_byte = (b >> _u32(i * 8)) & _u32(0xFF)
+    diff = (a_byte > b_byte).where(a_byte - b_byte, b_byte - a_byte)
+    result = result + (a_byte.ne(_u32(0)).where(diff, _u32(0)) if masked else diff)
+  return result
+
 _FUNCS: dict[str, Callable[..., UOp]] = {
  'sqrt': lambda a: UOp(Ops.SQRT, a.dtype, (a,)), 'trunc': lambda a: UOp(Ops.TRUNC, a.dtype, (a,)),
  'log2': lambda a: UOp(Ops.LOG2, a.dtype, (a,)), 'sin': lambda a: _trig_reduce(a),
@@ -227,11 +244,53 @@ _FUNCS: dict[str, Callable[..., UOp]] = {
  'signext_from_bit': _signext_from_bit, 'ldexp': _ldexp, 'frexp_mant': _frexp_mant, 'mantissa': _frexp_mant,
  'frexp_exp': _frexp_exp, 'trig_preop_result': _trig_preop,
  's_ff1_i32_b32': lambda a: _ff1(a, 32), 's_ff1_i32_b64': lambda a: _ff1(a, 64),
+  # Normalization conversions: map [-1,1] or [0,1] to integer range
+  # Use floor(x + 0.5) for round-to-nearest
+  # SNORM: round(value * 32767), range is [-32767, 32767] (hardware behavior)
+  'f16_to_snorm': lambda a: _floor(_f16_extract(a).cast(dtypes.float32) * _const(dtypes.float32, 32767) + _const(dtypes.float32, 0.5)).cast(dtypes.int).cast(dtypes.int16),
+  'f16_to_unorm': lambda a: _floor(_f16_extract(a).cast(dtypes.float32) * _const(dtypes.float32, 65535) + _const(dtypes.float32, 0.5)).cast(dtypes.uint16),
+  'f32_to_snorm': lambda a: _floor(a.bitcast(dtypes.float32) * _const(dtypes.float32, 32767) + _const(dtypes.float32, 0.5)).cast(dtypes.int).cast(dtypes.int16),
+  'f32_to_unorm': lambda a: _floor(a.bitcast(dtypes.float32) * _const(dtypes.float32, 65535) + _const(dtypes.float32, 0.5)).cast(dtypes.uint16),
+  'f32_to_u8': lambda a: _f_to_u(a.bitcast(dtypes.float32), dtypes.uint8),
+  # Integer truncation conversions
+  'i32_to_i16': lambda a: a.cast(dtypes.int).cast(dtypes.int16),
+  'u32_to_u16': lambda a: a.cast(dtypes.uint32).cast(dtypes.uint16),
+  'u16_to_u32': lambda a: (a.cast(dtypes.uint32) & _u32(0xFFFF)),
+  'u8_to_u32': lambda a: (a.cast(dtypes.uint32) & _u32(0xFF)),
+  'u4_to_u32': lambda a: (a.cast(dtypes.uint32) & _u32(0xF)),
+  # Signed extraction with sign extension for dot products
+  'i16_to_i32': lambda a: _signext(a.cast(dtypes.uint32) & _u32(0xFFFF)),
+  'i8_to_i32': lambda a: _signext(a.cast(dtypes.uint32) & _u32(0xFF)),
+  'i4_to_i32': lambda a: _signext_4bit(a.cast(dtypes.uint32) & _u32(0xF)),
+  # Float to int16 conversions
+  'v_cvt_i16_f32': lambda a: UOp(Ops.TRUNC, dtypes.float32, (a.bitcast(dtypes.float32),)).cast(dtypes.int16),
+  'v_cvt_u16_f32': lambda a: _f_to_u(a.bitcast(dtypes.float32), dtypes.uint16),
+  # SAD (Sum of Absolute Differences) - sum |a_i - b_i| for 4 bytes + accumulator
+  'v_sad_u8': lambda a, b, c: _sad_u8(a, b, c),
+  'v_msad_u8': lambda a, b, c: _sad_u8(a, b, c, masked=True),
+  # 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')]:
    _FUNCS[f'v_{name}_{sfx}'] = lambda *a, im=is_max, d=dt: _minmax_reduce(im, d, *a)
    _FUNCS[f'v_{name}3_{sfx}'] = lambda *a, im=is_max, d=dt: _minmax_reduce(im, d, *a)
+# f16 min/max/min3/max3/med3
+for is_max, name in [(False, 'min'), (True, 'max')]:
+  _FUNCS[f'v_{name}_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
+  _FUNCS[f'v_{name}3_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
+  _FUNCS[f'v_{name}_num_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
+  _FUNCS[f'v_{name}_num_f32'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.float32, *a)
+  _FUNCS[f'v_{name}3_num_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
+  _FUNCS[f'v_{name}3_num_f32'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.float32, *a)
+  _FUNCS[f'v_{name}imum_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
+  _FUNCS[f'v_{name}imum_f32'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.float32, *a)
+  _FUNCS[f'v_{name}imum3_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
+  _FUNCS[f'v_{name}imum3_f32'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.float32, *a)

 # ═══════════════════════════════════════════════════════════════════════════════
 # TOKENIZER/PARSER
@@ -239,7 +298,7 @@ for is_max, name in [(False, 'min'), (True, 'max')]:

 DTYPES = {'u32': dtypes.uint32, 'i32': dtypes.int, 'f32': dtypes.float32, 'b32': dtypes.uint32, 'u64': dtypes.uint64, 'i64': dtypes.int64,
          'f64': dtypes.float64, 'b64': dtypes.uint64, 'u16': dtypes.uint16, 'i16': dtypes.short, 'f16': dtypes.half, 'b16': dtypes.uint16,
-          'u8': dtypes.uint8, 'i8': dtypes.int8, 'b8': dtypes.uint8, 'u1': dtypes.uint32}
+          'u8': dtypes.uint8, 'i8': dtypes.int8, 'b8': dtypes.uint8, 'u4': dtypes.uint8, 'i4': dtypes.int8, 'u1': dtypes.uint32}
 _BITS_DT = {8: dtypes.uint8, 16: dtypes.uint16, 32: dtypes.uint32, 64: dtypes.uint64}
 _NUM_SUFFIXES = ('ULL', 'LL', 'UL', 'U', 'L', 'F', 'f')
 def _strip_suffix(num: str) -> tuple[str, str]:
@@ -396,7 +455,7 @@ class Parser:
        self.eat('DOT')
        dt_name = self.eat('IDENT').val
        return self._handle_mem_load(addr, DTYPES.get(dt_name, dtypes.uint32))
-      if name == 'VGPR':
+      if name == 'VGPR' and self.at('LBRACKET'):
        self.eat('LBRACKET')
        lane = self.parse()
        self.eat('RBRACKET')
@@ -423,7 +482,21 @@ class Parser:
      if self.try_eat('LBRACE'):
        idx = self.eat('NUM').val
        self.eat('RBRACE')
-        elem = self.vars.get(f'{name}{idx}', _u32(0))
+        # Handle VGPR{lane}[reg] - 2D array access after loop unrolling
+        if name == 'VGPR' and self.at('LBRACKET'):
+          self.eat('LBRACKET')
+          reg = self.parse()
+          self.eat('RBRACKET')
+          vgpr = self.vars.get('_vgpr')
+          if vgpr is None: return _u32(0)
+          return vgpr.index(_to_u32(reg) * _u32(32) + _u32(int(idx)), ptr=True).load()
+        elem = self.vars.get(f'{name}@{idx}', self.vars.get(f'{name}{idx}'))
+        if elem is None:
+          # Extract bit idx from base variable (like var[idx])
+          base = self.vars.get(name)
+          assert isinstance(base, UOp), f"unknown variable: {name}{idx}"
+          dt = dtypes.uint64 if base.dtype in (dtypes.uint64, dtypes.int64) else dtypes.uint32
+          elem = (base.cast(dt) >> _const(dt, int(idx))) & _const(dt, 1)
        if self.try_eat('DOT'):
          dt_name = self.eat('IDENT').val
          return _cast_to(elem, DTYPES.get(dt_name, dtypes.uint32))
@@ -432,27 +505,17 @@ class Parser:
        return elem
      if self.at('LBRACKET') and name not in self.vars:
        self.eat('LBRACKET')
-        if self.at('NUM'):
-          idx_num = int(self.peek().val)
-          if f'{name}{idx_num}' in self.vars:
-            self.eat('NUM')
-            self.eat('RBRACKET')
-            elem = self.vars[f'{name}{idx_num}']
-            if self.try_eat('DOT'): return _cast_to(elem, DTYPES.get(self.eat('IDENT').val, dtypes.uint32))
-            return elem
        first = self.parse()
        return self._handle_bracket_rest(first, _u32(0), name)
      if name in self.vars:
        v = self.vars[name]
-        return v if isinstance(v, UOp) else _u32(0) if isinstance(v, dict) else _u32(0)
+        assert isinstance(v, UOp), f"expected UOp for {name}, got {type(v)}"
+        return v
      raise RuntimeError(f"unknown variable: {name}")
    raise RuntimeError(f"unexpected token in primary: {self.peek()}")

  def _handle_dot(self, base, field: str) -> UOp:
-    if isinstance(base, str): return _u32(0)
-    if not isinstance(base, UOp):
-      if isinstance(base, dict): return base.get(field, _u32(0))
-      return _u32(0)
+    assert isinstance(base, UOp), f"expected UOp for dot access, got {type(base)}"
    if field == 'u64' and self.at('LBRACKET') and self.peek(1).type == 'IDENT' and self.peek(1).val == 'laneId':
      self.eat('LBRACKET')
      self.eat_val('laneId', 'IDENT')
@@ -467,6 +530,7 @@ class Parser:
    if dt == base.dtype: return base
    if dt.itemsize == 2 and base.dtype.itemsize == 4:
      return (base & _const(base.dtype, 0xFFFF)).cast(dtypes.uint16) if dt == dtypes.uint16 else (base & _const(base.dtype, 0xFFFF)).cast(dtypes.uint16).bitcast(dt)
+    if field == 'i4': return _signext_4bit(base)
    return _cast_to(base, dt)

  def _handle_bracket(self, base, var_name: str | None = None) -> UOp:
@@ -509,16 +573,18 @@ class Parser:
      var_name = self._find_var_name(base)
    if first.op == Ops.CONST:
      idx = int(first.arg)
-      if var_name and f'{var_name}{idx}' in self.vars:
-        v = self.vars[f'{var_name}{idx}']
+      # Check for array element (var@idx)
+      if var_name and f'{var_name}@{idx}' in self.vars:
+        v = self.vars[f'{var_name}@{idx}']
        return _cast_to(v, dt_suffix) if dt_suffix else v
+      # Bit extraction
      dt = dtypes.uint64 if base.dtype in (dtypes.uint64, dtypes.int64) else dtypes.uint32
      base_cast = base.cast(dt) if base.dtype != dt else base
      result = ((base_cast >> _const(dt, idx)) & _const(dt, 1))
      return _cast_to(result, dt_suffix) if dt_suffix else result
    if var_name:
      idx_u32 = _to_u32(first)
-      elems = [(i, self.vars[f'{var_name}{i}']) for i in range(256) if f'{var_name}{i}' in self.vars]
+      elems = [(i, self.vars[f'{var_name}@{i}']) for i in range(256) if f'{var_name}@{i}' in self.vars]
      if elems:
        result = elems[-1][1]
        for ei, ev in reversed(elems[:-1]):
@@ -537,7 +603,7 @@ class Parser:
    self.eat('RBRACE')
    var_name = self._find_var_name(base)
    if var_name:
-      elem = self.vars.get(f'{var_name}{idx}', _u32(0))
+      elem = self.vars.get(f'{var_name}@{idx}', _u32(0))  # use @ to avoid collision with temps like A4
      if self.try_eat('DOT'):
        dt_name = self.eat('IDENT').val
        return _cast_to(elem, DTYPES.get(dt_name, dtypes.uint32))
@@ -599,13 +665,14 @@ class Parser:
    raise RuntimeError(f"unexpected token after {bits}': {self.peek()}")

  def _parse_number(self, num: str) -> UOp:
-    if num.startswith('0x') or num.startswith('0X'): return _const(dtypes.uint64, int(num.rstrip('ULul'), 16))
-    suffix, num = _strip_suffix(num)
-    if '.' in num or suffix in ('F', 'f'):
-      return _const(dtypes.float32 if suffix in ('F', 'f') else dtypes.float64, float(num))
-    val = int(num)
-    if 'ULL' in suffix: return _const(dtypes.uint64, val)
-    if 'LL' in suffix or 'L' in suffix: return _const(dtypes.uint64, val)
+    if num.startswith('0x') or num.startswith('0X'):
+      is_u64 = num.upper().endswith('ULL') or num.upper().endswith('LL') or num.upper().endswith('UL')
+      return _const(dtypes.uint64 if is_u64 else dtypes.uint32, int(num.rstrip('ULul'), 16))
+    suffix, num_str = _strip_suffix(num)
+    if '.' in num_str or suffix in ('F', 'f'):
+      return _const(dtypes.float32 if suffix in ('F', 'f') else dtypes.float64, float(num_str))
+    val = int(num_str)
+    if 'ULL' in suffix or 'LL' in suffix or 'L' in suffix: return _const(dtypes.uint64, val)
    if 'U' in suffix: return _const(dtypes.uint32, val)
    return _const(dtypes.int if val < 0 else dtypes.uint32, val)

@@ -623,7 +690,8 @@ class Parser:
      if ';' in body or '\n' in body or 'return' in body.lower():
        lines = [l.strip() for l in body.replace(';', '\n').split('\n') if l.strip() and not l.strip().startswith('//')]
        _, _, result = parse_block(lines, 0, lv, self.funcs)
-        return result if result is not None else _u32(0)
+        assert result is not None, f"lambda {name} must return a value"
+        return result
      return parse_expr(body, lv, self.funcs)
    if name in self.funcs:
      return self.funcs[name](*args)
@@ -631,7 +699,7 @@ class Parser:

  def _handle_mem_load(self, addr: UOp, dt) -> UOp:
    mem = self.vars.get('_vmem') if '_vmem' in self.vars else self.vars.get('_lds')
-    if mem is None: return _const(dt, 0)
+    assert mem is not None, "memory load requires _vmem or _lds"
    adt = dtypes.uint64 if addr.dtype == dtypes.uint64 else dtypes.uint32
    active = self.vars.get('_active')
    gate = (active,) if active is not None else ()
@@ -693,29 +761,9 @@ def parse_tokens(toks: list[Token], vars: dict[str, VarVal], funcs: dict | None

 # Unified block parser for pcode
 def _subst_loop_var(line: str, loop_var: str, val: int) -> str:
-  """Substitute loop variable and evaluate bracket expressions.
-  Converts var[loop_var] to var{val} for array element access (like the old regex parser)."""
+  """Substitute loop variable with its value."""
  toks = tokenize(line)
-  # First pass: convert var[loop_var] to var{loop_var} to mark for array element assignment
-  result_toks: list[Token] = []
-  j = 0
-  while j < len(toks):
-    t = toks[j]
-    # Check for pattern: IDENT[loop_var] where it's not preceded by a dot (not .type[...])
-    if t.type == 'IDENT' and j+3 < len(toks) and toks[j+1].type == 'LBRACKET' and toks[j+2].type == 'IDENT' and toks[j+2].val == loop_var and toks[j+3].type == 'RBRACKET':
-      # Check that it's not .type[loop_var]
-      if not result_toks or result_toks[-1].type != 'DOT':
-        result_toks.append(t)
-        result_toks.append(Token('LBRACE', '{'))
-        result_toks.append(Token('NUM', str(val)))
-        result_toks.append(Token('RBRACE', '}'))
-        j += 4
-        continue
-    result_toks.append(t)
-    j += 1
-  # Second pass: substitute loop variable in remaining positions
-  subst_parts = [str(val) if t.type == 'IDENT' and t.val == loop_var else t.val for t in result_toks if t.type != 'EOF']
-  return ' '.join(subst_parts)
+  return ' '.join(str(val) if t.type == 'IDENT' and t.val == loop_var else t.val for t in toks if t.type != 'EOF')

 def _set_bits(old: UOp, val: UOp, width: int, offset: int) -> UOp:
  """Set bits [offset:offset+width) in old to val, masking and shifting appropriately."""
@@ -765,8 +813,9 @@ def parse_block(lines: list[str], start: int, vars: dict[str, VarVal], funcs: di
      def parse_bound():
        if p.at('NUM') and p.peek(1).type == 'QUOTE': p.eat('NUM'); p.eat('QUOTE')
        if p.at('NUM'): return int(p.eat('NUM').val.rstrip('UuLl'))
-        expr = p.parse()
-        return int(expr.arg) if expr.op == Ops.CONST else 0
+        expr = p.parse().simplify()
+        assert expr.op == Ops.CONST, f"loop bound must be constant, got {expr}"
+        return int(expr.arg)
      start_val = parse_bound()
      p.eat('COLON')
      end_val = parse_bound()
@@ -787,7 +836,7 @@ def parse_block(lines: list[str], start: int, vars: dict[str, VarVal], funcs: di
      if found_var: vars[found_var] = block_assigns[found_var] = _const(dtypes.bool, False)
      for loop_i in range(start_val, end_val + 1):
        subst_lines = [_subst_loop_var(bl, loop_var, loop_i) for bl in body_lines if not (has_break and bl.strip().lower() == 'break')]
-        _, iter_assigns, _ = parse_block(subst_lines, 0, vars, funcs, assigns)
+        _, iter_assigns, _ = parse_block(subst_lines, 0, {**vars, **block_assigns}, funcs, assigns)
        if has_break:
          assert found_var is not None
          found = block_assigns.get(found_var, vars.get(found_var))
@@ -812,7 +861,9 @@ def parse_block(lines: list[str], start: int, vars: dict[str, VarVal], funcs: di

    # declare
    if first == 'declare':
-      if '[' not in line and len(toks) >= 2 and toks[1].type == 'IDENT': vars[toks[1].val] = _u32(0)
+      # Initialize scalar declarations (skip arrays and vars already passed as srcs)
+      if '[' not in line and len(toks) >= 2 and toks[1].type == 'IDENT':
+        vars.setdefault(toks[1].val, _u32(0))
      i += 1; continue

    # lambda definition
@@ -870,6 +921,7 @@ def parse_block(lines: list[str], start: int, vars: dict[str, VarVal], funcs: di
      j, lane_toks = _match_bracket(toks, 1)
      if j < len(toks) and toks[j].type == 'LBRACKET':
        j, reg_toks = _match_bracket(toks, j)
+        if j < len(toks) and toks[j].type == 'DOT': j += 2  # skip .type suffix
        if j < len(toks) and toks[j].type == 'EQUALS': j += 1
        ln, rg, val = parse_tokens(lane_toks, vars, funcs), parse_tokens(reg_toks, vars, funcs), parse_tokens(toks[j:], vars, funcs)
        if assigns is not None: assigns.append((f'VGPR[{_tok_str(lane_toks)}][{_tok_str(reg_toks)}]', (_to_u32(rg) * _u32(32) + _to_u32(ln), val)))
@@ -933,19 +985,32 @@ def parse_block(lines: list[str], start: int, vars: dict[str, VarVal], funcs: di
            block_assigns[var] = vars[var] = _set_bit(existing, _to_u32(parse_tokens(bit_toks, vars, funcs)), parse_tokens(toks[j+1:], vars, funcs))
            i += 1; continue

-    # Array element: var{idx} = value
-    if len(toks) >= 5 and toks[0].type == 'IDENT' and toks[1].type == 'LBRACE' and toks[2].type == 'NUM':
-      var, idx = toks[0].val, int(toks[2].val)
-      j = 4
-      while j < len(toks) and toks[j].type != 'EQUALS': j += 1
-      if j < len(toks):
-        val = parse_tokens(toks[j+1:], vars, funcs)
-        existing = block_assigns.get(var, vars.get(var))
-        if existing is not None and isinstance(existing, UOp):
-          block_assigns[var] = vars[var] = _set_bit(existing, _u32(idx), val)
-        else:
-          block_assigns[f'{var}{idx}'] = vars[f'{var}{idx}'] = val
-        i += 1; continue
+    # Array element: var[idx] = value (static index) or var[expr] = value (dynamic)
+    if len(toks) >= 4 and toks[0].type == 'IDENT' and toks[1].type == 'LBRACKET':
+      var = toks[0].val
+      j, idx_toks = _match_bracket(toks, 1)
+      if j < len(toks) and toks[j].type == 'EQUALS':
+        # Static index: var[NUM] = value
+        if len(idx_toks) == 1 and idx_toks[0].type == 'NUM':
+          idx = int(idx_toks[0].val.rstrip('UuLl'))
+          val = parse_tokens(toks[j+1:], vars, funcs)
+          existing = block_assigns.get(var, vars.get(var))
+          if existing is not None and isinstance(existing, UOp):
+            block_assigns[var] = vars[var] = _set_bit(existing, _u32(idx), val)
+          else:
+            block_assigns[f'{var}@{idx}'] = vars[f'{var}@{idx}'] = val
+          i += 1; continue
+        # Dynamic index: var[expr] = value where var has @-elements
+        elems = [(k.split('@')[1], v) for k, v in {**vars, **block_assigns}.items() if k.startswith(f'{var}@') and isinstance(v, UOp)]
+        if elems:
+          idx_expr = parse_tokens(idx_toks, vars, funcs)
+          val = parse_tokens(toks[j+1:], vars, funcs)
+          for elem_idx_str, old_elem in elems:
+            elem_idx = int(elem_idx_str)
+            cond = _to_u32(idx_expr).eq(_u32(elem_idx))
+            new_val = cond.where(val.cast(old_elem.dtype) if val.dtype != old_elem.dtype else val, old_elem)
+            block_assigns[f'{var}@{elem_idx}'] = vars[f'{var}@{elem_idx}'] = new_val
+          i += 1; continue

    # Compound assignment: var += or var -=
    assign_op = next((j for j, t in enumerate(toks) if t.type == 'ASSIGN_OP'), None)
@@ -992,13 +1057,14 @@ def parse_block(lines: list[str], start: int, vars: dict[str, VarVal], funcs: di
      def parse_cond(s, kw):
        ll = s.lower()
        return _to_bool(parse_expr(s[ll.find(kw) + len(kw):ll.rfind('then')].strip(), vars, funcs))
-      def not_static_false(c): return c.op != Ops.CONST or c.arg is not False
+      def is_const(c, v): return c.op == Ops.CONST and c.arg is v
      cond = parse_cond(line, 'if')
-      conditions: list[tuple[UOp, UOp | dict[str, VarVal] | None]] = [(cond, None)] if not_static_false(cond) else []
+      conditions: list[tuple[UOp, UOp | dict[str, VarVal] | None]] = [(cond, None)] if not is_const(cond, False) else []
      else_branch: tuple[UOp | None, dict[str, VarVal]] = (None, {})
      vars_snap = dict(vars)
+      static_true = is_const(cond, True)  # track if any condition is statically true
      i += 1
-      i, branch, ret = parse_block(lines, i, vars, funcs, assigns)
+      i, branch, ret = parse_block(lines, i, vars, funcs, assigns if not is_const(cond, False) else None)
      if conditions: conditions[0] = (cond, ret if ret is not None else branch)
      vars.clear(); vars.update(vars_snap)
      while i < len(lines):
@@ -1007,12 +1073,16 @@ def parse_block(lines: list[str], start: int, vars: dict[str, VarVal], funcs: di
        lf = ltoks[0].val.lower()
        if lf == 'elsif':
          c = parse_cond(lines[i], 'elsif')
-          i += 1; i, branch, ret = parse_block(lines, i, vars, funcs, assigns)
-          if not_static_false(c): conditions.append((c, ret if ret is not None else branch))
+          take = not static_true and not is_const(c, False)
+          i += 1; i, branch, ret = parse_block(lines, i, vars, funcs, assigns if take else None)
+          if take:
+            conditions.append((c, ret if ret is not None else branch))
+            if is_const(c, True): static_true = True
          vars.clear(); vars.update(vars_snap)
        elif lf == 'else':
-          i += 1; i, branch, ret = parse_block(lines, i, vars, funcs, assigns)
-          else_branch = (ret, branch)
+          i += 1
+          i, branch, ret = parse_block(lines, i, vars, funcs, assigns if not static_true else None)
+          if not static_true: else_branch = (ret, branch)
          vars.clear(); vars.update(vars_snap)
        elif lf == 'endif': i += 1; break
        else: break
@@ -1024,17 +1094,21 @@ def parse_block(lines: list[str], start: int, vars: dict[str, VarVal], funcs: di
            if rv.dtype != result.dtype and rv.dtype.itemsize == result.dtype.itemsize: result = result.cast(rv.dtype)
            result = c.where(rv, result)
        return i, block_assigns, result
-      # Main style: merge variable assignments with WHERE
-      else_assigns = else_branch[1]
-      all_vars = set().union(*[ba.keys() for _, ba in conditions if isinstance(ba, dict)], else_assigns.keys())
-      for var in all_vars:
-        res: Any = else_assigns.get(var, block_assigns.get(var, vars.get(var, _u32(0))))
-        for cond, ba in reversed(conditions):
-          if isinstance(ba, dict) and var in ba:
-            tv = ba[var]
-            if isinstance(tv, UOp) and isinstance(res, UOp):
-              res = cond.where(tv, res.cast(tv.dtype) if tv.dtype != res.dtype and tv.dtype.itemsize == res.dtype.itemsize else res)
-        block_assigns[var] = vars[var] = res
+      # If statically true, use that branch directly; otherwise merge with WHERE
+      if static_true:
+        ba = next((b for c, b in conditions if is_const(c, True) and isinstance(b, dict)), {})
+        block_assigns.update(ba); vars.update(ba)
+      else:
+        else_assigns = else_branch[1]
+        all_vars = set().union(*[ba.keys() for _, ba in conditions if isinstance(ba, dict)], else_assigns.keys())
+        for var in all_vars:
+          res: Any = else_assigns.get(var, block_assigns.get(var, vars.get(var, _u32(0))))
+          for cond, ba in reversed(conditions):
+            if isinstance(ba, dict) and var in ba:
+              tv = ba[var]
+              if isinstance(tv, UOp) and isinstance(res, UOp):
+                res = cond.where(tv, res.cast(tv.dtype) if tv.dtype != res.dtype and tv.dtype.itemsize == res.dtype.itemsize else res)
+          block_assigns[var] = vars[var] = res
      continue

    # Regular assignment: var = value
@@ -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
@@ -13,7 +13,7 @@ def _i32(f: float) -> int: return struct.unpack('<I', struct.pack('<f', f))[0]
 def _f32(i: int) -> float: return struct.unpack('<f', struct.pack('<I', i & 0xFFFFFFFF))[0]

 # f16 conversion helpers
-def _f16(i: int) -> float: return struct.unpack('<e', struct.pack('<H', i & 0xFFFF))[0]
+def f16(i: int) -> float: return struct.unpack('<e', struct.pack('<H', i & 0xFFFF))[0]
 def f32_to_f16(f: float) -> int:
  f = float(f)
  if math.isnan(f): return 0x7e00
@@ -43,6 +43,23 @@ VCC = VCC_LO  # For VOP3SD sdst field (VCC_LO is exported from dsl)
 USE_HW = os.environ.get("USE_HW", "0") == "1"
 FLOAT_TOLERANCE = 1e-5

+def get_gpu_target() -> tuple[int, int, int]:
+  """Get the GPU target as (major, minor, stepping) tuple."""
+  if not USE_HW: return (0, 0, 0)
+  from tinygrad.device import Device
+  return Device["AMD"].target
+
+def skip_unless_gfx(min_major: int, min_minor: int = 0, reason: str = ""):
+  """Skip test if GPU target is below the minimum required version."""
+  import unittest
+  def decorator(test_func):
+    if not USE_HW: return test_func
+    target = get_gpu_target()
+    if target[0] < min_major or (target[0] == min_major and target[1] < min_minor):
+      return unittest.skip(reason or f"requires gfx{min_major}{min_minor}0+")(test_func)
+    return test_func
+  return decorator
+
 # Output buffer layout: vgpr[16][32], sgpr[16], vcc, scc, exec
 N_VGPRS, N_SGPRS, WAVE_SIZE = 16, 16, 32
 VGPR_BYTES = N_VGPRS * WAVE_SIZE * 4  # 16 regs * 32 lanes * 4 bytes = 2048
@@ -212,8 +229,12 @@ amdhsa.kernels:

  return parse_output(bytes(out_buf), n_lanes)

-def compare_wave_states(emu_st: WaveState, hw_st: WaveState, n_lanes: int, n_vgprs: int = N_VGPRS) -> list[str]:
-  """Compare two WaveStates and return list of differences."""
+def compare_wave_states(emu_st: WaveState, hw_st: WaveState, n_lanes: int, n_vgprs: int = N_VGPRS, ulp_tolerance: int = 0) -> list[str]:
+  """Compare two WaveStates and return list of differences.
+
+  Args:
+    ulp_tolerance: Allow up to this many ULPs difference for float comparisons (0 = exact match required)
+  """
  import math
  diffs = []
  for i in range(n_vgprs):
@@ -224,6 +245,11 @@ def compare_wave_states(emu_st: WaveState, hw_st: WaveState, n_lanes: int, n_vgp
        emu_f, hw_f = _f32(emu_val), _f32(hw_val)
        if math.isnan(emu_f) and math.isnan(hw_f):
          continue
+        # Check ULP difference for floats (only for same-sign values)
+        if ulp_tolerance > 0 and (emu_val < 0x80000000) == (hw_val < 0x80000000):
+          ulp_diff = abs(int(emu_val) - int(hw_val))
+          if ulp_diff <= ulp_tolerance:
+            continue
        diffs.append(f"v[{i}] lane {lane}: emu=0x{emu_val:08x} ({emu_f:.6g}) hw=0x{hw_val:08x} ({hw_f:.6g})")
  for i in range(N_SGPRS):
    emu_val = emu_st.sgpr[i]
@@ -236,16 +262,19 @@ def compare_wave_states(emu_st: WaveState, hw_st: WaveState, n_lanes: int, n_vgp
    diffs.append(f"scc: emu={emu_st.scc} hw={hw_st.scc}")
  return diffs

-def run_program(instructions: list, n_lanes: int = 1) -> WaveState:
+def run_program(instructions: list, n_lanes: int = 1, ulp_tolerance: int = 0) -> WaveState:
  """Run instructions and return WaveState.

  If USE_HW=1, runs on both emulator and hardware, compares results, and raises if they differ.
  Otherwise, runs only on emulator.
+
+  Args:
+    ulp_tolerance: Allow up to this many ULPs difference for float comparisons (0 = exact match required)
  """
  emu_st = run_program_emu(instructions, n_lanes)
  if USE_HW:
    hw_st = run_program_hw(instructions, n_lanes)
-    diffs = compare_wave_states(emu_st, hw_st, n_lanes)
+    diffs = compare_wave_states(emu_st, hw_st, n_lanes, ulp_tolerance=ulp_tolerance)
    if diffs:
      raise AssertionError(f"Emulator vs Hardware mismatch:\n" + "\n".join(diffs))
    return hw_st
@@ -719,5 +719,47 @@ class TestAtomicOrdering(unittest.TestCase):
    self.assertEqual(st.vgpr[0][4], 150, "Final value should be 150")


+class TestDsPermute(unittest.TestCase):
+  """Tests for DS_PERMUTE_B32 and DS_BPERMUTE_B32 instructions."""
+
+  def test_ds_permute_b32_identity(self):
+    """DS_PERMUTE_B32 with identity permutation (lane 0 sends to lane 0)."""
+    # For simplicity, test with single lane
+    instructions = [
+      v_mov_b32_e32(v[0], 0),  # addr = 0 (lane 0)
+      v_mov_b32_e32(v[1], 0xDEADBEEF),  # data
+      ds_permute_b32(v[2], v[0], v[1]),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    # Lane 0 sends to lane 0, so lane 0 gets 0xDEADBEEF
+    self.assertEqual(st.vgpr[0][2], 0xDEADBEEF)
+
+  def test_ds_bpermute_b32_identity(self):
+    """DS_BPERMUTE_B32 with identity permutation (each lane reads from itself)."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0),  # addr = 0 (read from lane 0)
+      v_mov_b32_e32(v[1], 0xCAFEBABE),  # data in lane 0
+      ds_bpermute_b32(v[2], v[0], v[1]),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    # Lane 0 reads from lane 0's v[1]
+    self.assertEqual(st.vgpr[0][2], 0xCAFEBABE)
+
+  def test_ds_permute_b32_broadcast(self):
+    """DS_PERMUTE_B32 broadcast - all lanes send to lane 0."""
+    # With 4 lanes, all sending to lane 0, highest lane wins
+    instructions = [
+      v_mov_b32_e32(v[0], 0),  # All lanes send to addr 0 (lane 0)
+      v_mov_b32_e32(v[1], 0x11111111),  # All lanes send same data
+      ds_permute_b32(v[2], v[0], v[1]),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=4)
+    # Lane 0 receives data (highest numbered active lane wins)
+    self.assertEqual(st.vgpr[0][2], 0x11111111)
+
+
 if __name__ == '__main__':
  unittest.main()
@@ -62,6 +62,7 @@ class TestBasicScalar(unittest.TestCase):
    st = run_program(instructions, n_lanes=1)
    self.assertEqual(st.sgpr[1], 0x80000000)

+  @skip_unless_gfx(11, 5, "SALU FP ops require gfx1150+")
  def test_s_fmamk_f32(self):
    """S_FMAMK_F32: D = S0 * literal + S1."""
    # 2.0 * 3.0 + 1.0 = 7.0
@@ -73,6 +74,7 @@ class TestBasicScalar(unittest.TestCase):
    st = run_program(instructions, n_lanes=1)
    self.assertEqual(st.sgpr[2], f2i(7.0))

+  @skip_unless_gfx(11, 5, "SALU FP ops require gfx1150+")
  def test_s_fmamk_f32_negative(self):
    """S_FMAMK_F32 with negative values."""
    # -2.0 * 4.0 + 10.0 = 2.0
@@ -255,7 +255,6 @@ class TestF16Conversions(unittest.TestCase):

  def test_v_cvt_f16_f32_small(self):
    """V_CVT_F16_F32 converts small f32 value."""
-    from extra.assembly.amd.test.hw.helpers import f32_to_f16
    instructions = [
      v_mov_b32_e32(v[0], 0.5),
      v_cvt_f16_f32_e32(v[1], v[0]),
@@ -293,7 +292,6 @@ class TestF16Conversions(unittest.TestCase):

  def test_v_cvt_f16_f32_reads_full_32bit_source(self):
    """V_CVT_F16_F32 must read full 32-bit f32 source."""
-    from extra.assembly.amd.test.hw.helpers import _f16
    instructions = [
      s_mov_b32(s[0], 0x3fc00000),  # f32 1.5
      v_mov_b32_e32(v[0], s[0]),
@@ -302,7 +300,7 @@ class TestF16Conversions(unittest.TestCase):
    st = run_program(instructions, n_lanes=1)
    result = st.vgpr[0][1]
    lo_bits = result & 0xffff
-    self.assertEqual(lo_bits, 0x3e00, f"Expected f16(1.5)=0x3e00, got 0x{lo_bits:04x} ({_f16(lo_bits)})")
+    self.assertEqual(lo_bits, 0x3e00, f"Expected f16(1.5)=0x3e00, got 0x{lo_bits:04x} ({f16(lo_bits)})")

  def test_v_cvt_i16_f16_zero(self):
    """V_CVT_I16_F16 converts f16 zero to i16 zero."""
@@ -696,7 +694,6 @@ class TestCvtF16Modifiers(unittest.TestCase):

  def test_v_cvt_f32_f16_abs_negative(self):
    """V_CVT_F32_F16 with |abs| on negative value."""
-    from extra.assembly.amd.test.hw.helpers import f32_to_f16
    f16_neg1 = f32_to_f16(-1.0)  # 0xbc00
    instructions = [
      s_mov_b32(s[0], f16_neg1),
@@ -709,7 +706,6 @@ class TestCvtF16Modifiers(unittest.TestCase):

  def test_v_cvt_f32_f16_abs_positive(self):
    """V_CVT_F32_F16 with |abs| on positive value (should stay positive)."""
-    from extra.assembly.amd.test.hw.helpers import f32_to_f16
    f16_2 = f32_to_f16(2.0)  # 0x4000
    instructions = [
      s_mov_b32(s[0], f16_2),
@@ -722,7 +718,6 @@ class TestCvtF16Modifiers(unittest.TestCase):

  def test_v_cvt_f32_f16_neg_positive(self):
    """V_CVT_F32_F16 with neg on positive value."""
-    from extra.assembly.amd.test.hw.helpers import f32_to_f16
    f16_2 = f32_to_f16(2.0)  # 0x4000
    instructions = [
      s_mov_b32(s[0], f16_2),
@@ -735,7 +730,6 @@ class TestCvtF16Modifiers(unittest.TestCase):

  def test_v_cvt_f32_f16_neg_negative(self):
    """V_CVT_F32_F16 with neg on negative value (double negative)."""
-    from extra.assembly.amd.test.hw.helpers import f32_to_f16
    f16_neg2 = f32_to_f16(-2.0)  # 0xc000
    instructions = [
      s_mov_b32(s[0], f16_neg2),
@@ -748,7 +742,6 @@ class TestCvtF16Modifiers(unittest.TestCase):

  def test_v_cvt_f16_f32_then_pack_for_wmma(self):
    """CVT F32->F16 followed by pack (common WMMA pattern)."""
-    from extra.assembly.amd.test.hw.helpers import _f16
    f32_val = 3.5
    instructions = [
      s_mov_b32(s[0], f2i(f32_val)),
@@ -757,8 +750,8 @@ class TestCvtF16Modifiers(unittest.TestCase):
      v_pack_b32_f16(v[2], v[1], v[1]),  # Pack same value
    ]
    st = run_program(instructions, n_lanes=1)
-    lo = _f16(st.vgpr[0][2] & 0xffff)
-    hi = _f16((st.vgpr[0][2] >> 16) & 0xffff)
+    lo = f16(st.vgpr[0][2] & 0xffff)
+    hi = f16((st.vgpr[0][2] >> 16) & 0xffff)
    self.assertAlmostEqual(lo, f32_val, places=1)
    self.assertAlmostEqual(hi, f32_val, places=1)

@@ -804,7 +797,6 @@ class TestConversionRounding(unittest.TestCase):

  def test_f16_to_f32_precision(self):
    """F16 to F32 conversion precision."""
-    from extra.assembly.amd.test.hw.helpers import f32_to_f16
    f16_val = f32_to_f16(1.5)
    instructions = [
      s_mov_b32(s[0], f16_val),
@@ -816,7 +808,6 @@ class TestConversionRounding(unittest.TestCase):

  def test_f16_denormal_to_f32(self):
    """F16 denormal converts to small positive f32."""
-    from extra.assembly.amd.test.hw.helpers import _f16
    f16_denorm = 0x0001  # Smallest positive f16 denormal
    instructions = [
      v_mov_b32_e32(v[0], f16_denorm),
@@ -1512,5 +1503,82 @@ class TestReciprocalF16(unittest.TestCase):
    self.assertAlmostEqual(result, 0.25, places=2, msg="1/4.0 should be 0.25")


+class TestCvtNormF16(unittest.TestCase):
+  """Tests for V_CVT_NORM_I16_F16 and V_CVT_NORM_U16_F16."""
+
+  def test_cvt_norm_i16_f16_positive(self):
+    """V_CVT_NORM_I16_F16: f16 1.0 -> i16 max (32767)."""
+    instructions = [
+      s_mov_b32(s[0], f32_to_f16(1.0)),
+      v_mov_b32_e32(v[0], s[0]),
+      v_cvt_norm_i16_f16_e32(v[1], v[0]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][1] & 0xffff
+    self.assertEqual(result, 32767)
+
+  def test_cvt_norm_i16_f16_negative(self):
+    """V_CVT_NORM_I16_F16: f16 -1.0 -> i16 -32767 (0x8001)."""
+    instructions = [
+      s_mov_b32(s[0], f32_to_f16(-1.0)),
+      v_mov_b32_e32(v[0], s[0]),
+      v_cvt_norm_i16_f16_e32(v[1], v[0]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][1] & 0xffff
+    self.assertEqual(result, 0x8001)  # -32767, hardware uses symmetric range
+
+  def test_cvt_norm_i16_f16_zero(self):
+    """V_CVT_NORM_I16_F16: f16 0.0 -> i16 0."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0),
+      v_cvt_norm_i16_f16_e32(v[1], v[0]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][1] & 0xffff
+    self.assertEqual(result, 0)
+
+  def test_cvt_norm_u16_f16_one(self):
+    """V_CVT_NORM_U16_F16: f16 1.0 -> u16 max (65535)."""
+    instructions = [
+      s_mov_b32(s[0], f32_to_f16(1.0)),
+      v_mov_b32_e32(v[0], s[0]),
+      v_cvt_norm_u16_f16_e32(v[1], v[0]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][1] & 0xffff
+    self.assertEqual(result, 65535)
+
+  def test_cvt_norm_u16_f16_half(self):
+    """V_CVT_NORM_U16_F16: f16 0.5 -> u16 ~32768."""
+    instructions = [
+      s_mov_b32(s[0], f32_to_f16(0.5)),
+      v_mov_b32_e32(v[0], s[0]),
+      v_cvt_norm_u16_f16_e32(v[1], v[0]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][1] & 0xffff
+    self.assertAlmostEqual(result, 32768, delta=1)
+
+
+class TestPermlane64(unittest.TestCase):
+  """Tests for V_PERMLANE64_B32 instruction (wave64 cross-half swap)."""
+
+  def test_v_permlane64_b32_is_nop_in_wave32(self):
+    """V_PERMLANE64_B32 is a NOP in wave32 mode.
+
+    Per AMD pcode: "if WAVE32 then s_nop(...) else ... endif"
+    The emulator runs in wave32 mode, so this instruction should not modify registers.
+    """
+    instructions = [
+      v_mov_b32_e32(v[0], 0xCAFEBABE),  # source
+      v_mov_b32_e32(v[1], 0x12345678),  # dest (should be preserved)
+      v_permlane64_b32_e32(v[1], v[0]),  # NOP in wave32
+    ]
+    st = run_program(instructions, n_lanes=1)
+    # Dest register should be unchanged (NOP behavior in wave32)
+    self.assertEqual(st.vgpr[0][1], 0x12345678)
+
+
 if __name__ == '__main__':
  unittest.main()
@@ -857,7 +857,6 @@ class TestF16Modifiers(unittest.TestCase):

  def test_v_fma_f16_inline_const_1_0(self):
    """V_FMA_F16: a*b + 1.0 should use f16 inline constant."""
-    from extra.assembly.amd.test.hw.helpers import f32_to_f16, _f16
    f16_a = f32_to_f16(0.325928)  # ~0x3537
    f16_b = f32_to_f16(-0.486572)  # ~0xb7c9
    instructions = [
@@ -868,13 +867,12 @@ class TestF16Modifiers(unittest.TestCase):
      v_fma_f16(v[4], v[4], v[6], 1.0),  # 1.0 is inline constant
    ]
    st = run_program(instructions, n_lanes=1)
-    result = _f16(st.vgpr[0][4] & 0xffff)
+    result = f16(st.vgpr[0][4] & 0xffff)
    expected = 0.325928 * (-0.486572) + 1.0
    self.assertAlmostEqual(result, expected, delta=0.01)

  def test_v_fma_f16_inline_const_0_5(self):
    """V_FMA_F16: a*b + 0.5 should use f16 inline constant."""
-    from extra.assembly.amd.test.hw.helpers import f32_to_f16, _f16
    f16_a = f32_to_f16(2.0)
    f16_b = f32_to_f16(3.0)
    instructions = [
@@ -885,13 +883,12 @@ class TestF16Modifiers(unittest.TestCase):
      v_fma_f16(v[2], v[0], v[1], 0.5),  # 0.5 is inline constant
    ]
    st = run_program(instructions, n_lanes=1)
-    result = _f16(st.vgpr[0][2] & 0xffff)
+    result = f16(st.vgpr[0][2] & 0xffff)
    expected = 2.0 * 3.0 + 0.5
    self.assertAlmostEqual(result, expected, delta=0.01)

  def test_v_fma_f16_inline_const_neg_1_0(self):
    """V_FMA_F16: a*b + (-1.0) should use f16 inline constant."""
-    from extra.assembly.amd.test.hw.helpers import f32_to_f16, _f16
    f16_a = f32_to_f16(2.0)
    f16_b = f32_to_f16(3.0)
    instructions = [
@@ -902,13 +899,12 @@ class TestF16Modifiers(unittest.TestCase):
      v_fma_f16(v[2], v[0], v[1], -1.0),  # -1.0 is inline constant
    ]
    st = run_program(instructions, n_lanes=1)
-    result = _f16(st.vgpr[0][2] & 0xffff)
+    result = f16(st.vgpr[0][2] & 0xffff)
    expected = 2.0 * 3.0 + (-1.0)
    self.assertAlmostEqual(result, expected, delta=0.01)

  def test_v_add_f16_abs_both(self):
    """V_ADD_F16 with abs on both operands."""
-    from extra.assembly.amd.test.hw.helpers import f32_to_f16, _f16
    f16_neg2 = f32_to_f16(-2.0)
    f16_neg3 = f32_to_f16(-3.0)
    instructions = [
@@ -919,12 +915,11 @@ class TestF16Modifiers(unittest.TestCase):
      v_add_f16_e64(v[2], abs(v[0]), abs(v[1])),  # |-2| + |-3| = 5
    ]
    st = run_program(instructions, n_lanes=1)
-    result = _f16(st.vgpr[0][2] & 0xffff)
+    result = f16(st.vgpr[0][2] & 0xffff)
    self.assertAlmostEqual(result, 5.0, delta=0.01)

  def test_v_mul_f16_neg_abs(self):
    """V_MUL_F16 with neg on one operand and abs on another."""
-    from extra.assembly.amd.test.hw.helpers import f32_to_f16, _f16
    f16_2 = f32_to_f16(2.0)
    f16_neg3 = f32_to_f16(-3.0)
    instructions = [
@@ -935,7 +930,7 @@ class TestF16Modifiers(unittest.TestCase):
      v_mul_f16_e64(v[2], -v[0], abs(v[1])),  # -(2) * |-3| = -6
    ]
    st = run_program(instructions, n_lanes=1)
-    result = _f16(st.vgpr[0][2] & 0xffff)
+    result = f16(st.vgpr[0][2] & 0xffff)
    self.assertAlmostEqual(result, -6.0, delta=0.01)

  def test_v_fmac_f16_hi_dest(self):
@@ -943,7 +938,6 @@ class TestF16Modifiers(unittest.TestCase):

    This tests the case from AMD_LLVM sin(0) where V_FMAC_F16 writes to v0.h.
    """
-    from extra.assembly.amd.test.hw.helpers import _f16
    instructions = [
      s_mov_b32(s[0], 0x38003c00),  # v0 = {hi=0.5, lo=1.0}
      v_mov_b32_e32(v[0], s[0]),
@@ -954,8 +948,8 @@ class TestF16Modifiers(unittest.TestCase):
    ]
    st = run_program(instructions, n_lanes=1)
    v0 = st.vgpr[0][0]
-    result_hi = _f16((v0 >> 16) & 0xffff)
-    result_lo = _f16(v0 & 0xffff)
+    result_hi = f16((v0 >> 16) & 0xffff)
+    result_lo = f16(v0 & 0xffff)
    self.assertAlmostEqual(result_hi, 0.5, delta=0.01, msg=f"Expected hi=0.5, got {result_hi}")
    self.assertAlmostEqual(result_lo, 1.0, delta=0.01, msg=f"Expected lo=1.0, got {result_lo}")

@@ -2955,5 +2949,394 @@ class TestVOP3Clamp(unittest.TestCase):
    self.assertAlmostEqual(i2f(st.vgpr[3][1]), 1.0, places=5, msg="lane 3: 2.5 should clamp to 1.0")


+class TestCvtPkF16(unittest.TestCase):
+  """Tests for V_CVT_PK_RTZ_F16_F32 - pack two f32 to f16 with round toward zero."""
+
+  def test_cvt_pk_rtz_f16_f32_basic(self):
+    """V_CVT_PK_RTZ_F16_F32: basic pack of two f32 values."""
+    instructions = [
+      v_mov_b32_e32(v[0], 1.0),
+      v_mov_b32_e32(v[1], 2.0),
+      v_cvt_pk_rtz_f16_f32_e64(v[2], v[0], v[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][2]
+    lo_f16 = f16(result & 0xffff)
+    hi_f16 = f16((result >> 16) & 0xffff)
+    self.assertAlmostEqual(lo_f16, 1.0, delta=0.01)
+    self.assertAlmostEqual(hi_f16, 2.0, delta=0.01)
+
+
+class TestCvtPkNorm(unittest.TestCase):
+  """Tests for V_CVT_PK_NORM_I16_F32 and V_CVT_PK_NORM_U16_F32."""
+
+  def test_cvt_pk_norm_i16_f32_basic(self):
+    """V_CVT_PK_NORM_I16_F32: pack two f32 to normalized i16."""
+    instructions = [
+      v_mov_b32_e32(v[0], 1.0),
+      v_mov_b32_e32(v[1], -1.0),
+      v_cvt_pk_norm_i16_f32(v[2], v[0], v[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][2]
+    lo = result & 0xffff
+    hi = (result >> 16) & 0xffff
+    self.assertEqual(lo, 32767)
+    self.assertEqual(hi, 0x8001)  # -32767, hardware uses symmetric range
+
+  def test_cvt_pk_norm_u16_f32_basic(self):
+    """V_CVT_PK_NORM_U16_F32: pack two f32 to normalized u16."""
+    instructions = [
+      v_mov_b32_e32(v[0], 1.0),
+      v_mov_b32_e32(v[1], 0.5),
+      v_cvt_pk_norm_u16_f32(v[2], v[0], v[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][2]
+    lo = result & 0xffff
+    hi = (result >> 16) & 0xffff
+    self.assertEqual(lo, 65535)
+    self.assertAlmostEqual(hi, 32768, delta=1)
+
+
+class TestCvtPkInt(unittest.TestCase):
+  """Tests for V_CVT_PK_I16_I32, V_CVT_PK_U16_U32, V_CVT_PK_I16_F32, V_CVT_PK_U16_F32."""
+
+  def test_cvt_pk_i16_i32_basic(self):
+    """V_CVT_PK_I16_I32: pack two i32 to i16."""
+    instructions = [
+      s_mov_b32(s[0], 100),
+      s_mov_b32(s[1], -100 & 0xffffffff),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_cvt_pk_i16_i32(v[2], v[0], v[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][2]
+    lo = result & 0xffff
+    hi = (result >> 16) & 0xffff
+    lo_signed = lo if lo < 32768 else lo - 65536
+    hi_signed = hi if hi < 32768 else hi - 65536
+    self.assertEqual(lo_signed, 100)
+    self.assertEqual(hi_signed, -100)
+
+  def test_cvt_pk_u16_u32_basic(self):
+    """V_CVT_PK_U16_U32: pack two u32 to u16."""
+    instructions = [
+      s_mov_b32(s[0], 1000),
+      s_mov_b32(s[1], 2000),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_cvt_pk_u16_u32(v[2], v[0], v[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][2]
+    lo = result & 0xffff
+    hi = (result >> 16) & 0xffff
+    self.assertEqual(lo, 1000)
+    self.assertEqual(hi, 2000)
+
+  def test_cvt_pk_i16_f32_basic(self):
+    """V_CVT_PK_I16_F32: convert two f32 to packed i16."""
+    instructions = [
+      v_mov_b32_e32(v[0], 100.5),
+      v_mov_b32_e32(v[1], -50.7),
+      v_cvt_pk_i16_f32(v[2], v[0], v[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][2]
+    lo = result & 0xffff
+    hi = (result >> 16) & 0xffff
+    lo_signed = lo if lo < 32768 else lo - 65536
+    hi_signed = hi if hi < 32768 else hi - 65536
+    self.assertEqual(lo_signed, 100)
+    self.assertEqual(hi_signed, -50)
+
+  def test_cvt_pk_u16_f32_basic(self):
+    """V_CVT_PK_U16_F32: convert two f32 to packed u16."""
+    instructions = [
+      v_mov_b32_e32(v[0], 100.9),
+      v_mov_b32_e32(v[1], 200.1),
+      v_cvt_pk_u16_f32(v[2], v[0], v[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][2]
+    lo = result & 0xffff
+    hi = (result >> 16) & 0xffff
+    self.assertEqual(lo, 100)
+    self.assertEqual(hi, 200)
+
+  def test_cvt_pk_u8_f32_basic(self):
+    """V_CVT_PK_U8_F32: convert f32 to u8 and pack at byte position."""
+    instructions = [
+      v_mov_b32_e32(v[0], 128.5),
+      v_mov_b32_e32(v[1], 0),
+      v_mov_b32_e32(v[2], 0),
+      v_cvt_pk_u8_f32(v[2], v[0], v[1], v[2]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][2]
+    byte0 = result & 0xff
+    self.assertEqual(byte0, 128)
+
+
+class TestDotProduct(unittest.TestCase):
+  """Tests for dot product instructions V_DOT4_U32_U8, V_DOT8_U32_U4."""
+
+  def test_v_dot4_u32_u8_basic(self):
+    """V_DOT4_U32_U8: 4-element dot product of u8 vectors."""
+    src0 = 0x04030201  # {4, 3, 2, 1}
+    src1 = 0x01010101  # {1, 1, 1, 1}
+    instructions = [
+      s_mov_b32(s[0], src0),
+      s_mov_b32(s[1], src1),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], 0),
+      v_dot4_u32_u8(v[2], v[0], v[1], v[2]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][2]
+    self.assertEqual(result, 10)
+
+  def test_v_dot4_u32_u8_with_accumulator(self):
+    """V_DOT4_U32_U8 with non-zero accumulator."""
+    src0 = 0x02020202  # {2, 2, 2, 2}
+    src1 = 0x03030303  # {3, 3, 3, 3}
+    instructions = [
+      s_mov_b32(s[0], src0),
+      s_mov_b32(s[1], src1),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], 100),
+      v_dot4_u32_u8(v[2], v[0], v[1], v[2]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][2]
+    self.assertEqual(result, 124)
+
+  def test_v_dot8_u32_u4_basic(self):
+    """V_DOT8_U32_U4: 8-element dot product of u4 vectors."""
+    # src0 = 8 nibbles: {1,2,3,4,5,6,7,8} packed as 0x87654321
+    # src1 = 8 nibbles: {1,1,1,1,1,1,1,1} packed as 0x11111111
+    # result = 1+2+3+4+5+6+7+8 = 36
+    src0 = 0x87654321
+    src1 = 0x11111111
+    instructions = [
+      s_mov_b32(s[0], src0),
+      s_mov_b32(s[1], src1),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], 0),
+      v_dot8_u32_u4(v[2], v[0], v[1], v[2]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][2]
+    self.assertEqual(result, 36)
+
+
+class TestMinMaxF16Vop3(unittest.TestCase):
+  """Tests for V_MIN3_F16, V_MAX3_F16, V_MED3_F16, V_MINMAX_F16, V_MAXMIN_F16."""
+
+  def test_v_min3_f16_basic(self):
+    """V_MIN3_F16: minimum of three f16 values."""
+    instructions = [
+      s_mov_b32(s[0], f32_to_f16(3.0)),
+      s_mov_b32(s[1], f32_to_f16(1.0)),
+      s_mov_b32(s[2], f32_to_f16(2.0)),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], s[2]),
+      v_min3_f16(v[3], v[0], v[1], v[2]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = f16(st.vgpr[0][3] & 0xffff)
+    self.assertAlmostEqual(result, 1.0, delta=0.01)
+
+  def test_v_max3_f16_basic(self):
+    """V_MAX3_F16: maximum of three f16 values."""
+    instructions = [
+      s_mov_b32(s[0], f32_to_f16(1.0)),
+      s_mov_b32(s[1], f32_to_f16(3.0)),
+      s_mov_b32(s[2], f32_to_f16(2.0)),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], s[2]),
+      v_max3_f16(v[3], v[0], v[1], v[2]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = f16(st.vgpr[0][3] & 0xffff)
+    self.assertAlmostEqual(result, 3.0, delta=0.01)
+
+  def test_v_med3_f16_basic(self):
+    """V_MED3_F16: median of three f16 values."""
+    instructions = [
+      s_mov_b32(s[0], f32_to_f16(3.0)),
+      s_mov_b32(s[1], f32_to_f16(1.0)),
+      s_mov_b32(s[2], f32_to_f16(2.0)),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], s[2]),
+      v_med3_f16(v[3], v[0], v[1], v[2]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = f16(st.vgpr[0][3] & 0xffff)
+    self.assertAlmostEqual(result, 2.0, delta=0.01)
+
+  def test_v_minmax_f16_basic(self):
+    """V_MINMAX_F16: clamp(src0, min=src1, max=src2)."""
+    instructions = [
+      s_mov_b32(s[0], f32_to_f16(2.5)),
+      s_mov_b32(s[1], f32_to_f16(1.0)),
+      s_mov_b32(s[2], f32_to_f16(2.0)),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], s[2]),
+      v_minmax_f16(v[3], v[0], v[1], v[2]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = f16(st.vgpr[0][3] & 0xffff)
+    self.assertAlmostEqual(result, 2.0, delta=0.01)
+
+  def test_v_maxmin_f16_basic(self):
+    """V_MAXMIN_F16: clamp(src0, min=src2, max=src1)."""
+    instructions = [
+      s_mov_b32(s[0], f32_to_f16(0.5)),
+      s_mov_b32(s[1], f32_to_f16(2.0)),
+      s_mov_b32(s[2], f32_to_f16(1.0)),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], s[2]),
+      v_maxmin_f16(v[3], v[0], v[1], v[2]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = f16(st.vgpr[0][3] & 0xffff)
+    self.assertAlmostEqual(result, 1.0, delta=0.01)
+
+  def test_v_min3_f16_with_neg(self):
+    """V_MIN3_F16 with neg modifier: min(-3, 1, 2) = -3."""
+    instructions = [
+      s_mov_b32(s[0], f32_to_f16(3.0)),
+      s_mov_b32(s[1], f32_to_f16(1.0)),
+      s_mov_b32(s[2], f32_to_f16(2.0)),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], s[2]),
+      v_min3_f16(v[3], -v[0], v[1], v[2]),  # neg on first operand
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = f16(st.vgpr[0][3] & 0xffff)
+    self.assertAlmostEqual(result, -3.0, delta=0.01)
+
+  def test_v_max3_f16_with_abs(self):
+    """V_MAX3_F16 with abs modifier: max(|-3|, 1, 2) = 3."""
+    instructions = [
+      s_mov_b32(s[0], f32_to_f16(-3.0)),
+      s_mov_b32(s[1], f32_to_f16(1.0)),
+      s_mov_b32(s[2], f32_to_f16(2.0)),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], s[2]),
+      v_max3_f16(v[3], abs(v[0]), v[1], v[2]),  # abs on first operand
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = f16(st.vgpr[0][3] & 0xffff)
+    self.assertAlmostEqual(result, 3.0, delta=0.01)
+
+  def test_v_med3_f16_opsel_hi(self):
+    """V_MED3_F16 with opsel reading from hi half."""
+    # Pack two f16 values: hi=5.0, lo=1.0
+    packed = (f32_to_f16(5.0) << 16) | f32_to_f16(1.0)
+    instructions = [
+      s_mov_b32(s[0], packed),
+      s_mov_b32(s[1], f32_to_f16(3.0)),
+      s_mov_b32(s[2], f32_to_f16(4.0)),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], s[2]),
+      # Read hi half of v[0] (5.0), med3(5, 3, 4) = 4
+      v_med3_f16(v[3], v[0].h, v[1], v[2]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = f16(st.vgpr[0][3] & 0xffff)
+    self.assertAlmostEqual(result, 4.0, delta=0.01)
+
+
+class TestSadHi(unittest.TestCase):
+  """Tests for V_SAD_HI_U8 instruction."""
+
+  def test_v_sad_hi_u8_basic(self):
+    """V_SAD_HI_U8: (sad << 16) + acc."""
+    # |1-5| + |2-6| + |3-7| + |4-8| = 16, << 16 = 0x100000, + 100 = 0x100064
+    instructions = [
+      v_mov_b32_e32(v[0], 0x04030201),
+      v_mov_b32_e32(v[1], 0x08070605),
+      v_mov_b32_e32(v[2], 100),
+      v_sad_hi_u8(v[3], v[0], v[1], v[2]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][3], (16 << 16) + 100)
+
+  def test_v_sad_hi_u8_zero_diff(self):
+    """V_SAD_HI_U8: identical inputs gives acc only."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0x12345678),
+      v_mov_b32_e32(v[2], 50),
+      v_sad_hi_u8(v[3], v[0], v[0], v[2]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][3], 50)
+
+
+class TestPermlane(unittest.TestCase):
+  """Tests for V_PERMLANE16_B32 and V_PERMLANEX16_B32 instructions."""
+
+  def test_v_permlane16_b32_identity(self):
+    """V_PERMLANE16_B32 with identity permutation (lane i reads from lane i within row)."""
+    # lanesel encodes 4 bits per position: position i gets lanesel[i*4+3:i*4]
+    # Identity: position 0->0, 1->1, ..., 15->15
+    # lanesel = 0xFEDCBA9876543210 (positions 15-0 in nibbles)
+    instructions = [
+      v_mov_b32_e32(v[0], 0xDEADBEEF),  # source data
+      s_mov_b32(s[0], 0x76543210),       # lanesel low (positions 0-7)
+      s_mov_b32(s[1], 0xFEDCBA98),       # lanesel high (positions 8-15)
+      v_permlane16_b32(v[1], v[0], s[0], s[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    # Lane 0 reads from lane 0 (position 0 -> lanesel[3:0] = 0)
+    self.assertEqual(st.vgpr[0][1], 0xDEADBEEF)
+
+  def test_v_permlane16_b32_broadcast(self):
+    """V_PERMLANE16_B32 broadcast lane 0 to all lanes in row."""
+    # lanesel = all zeros -> all positions read from lane 0 within row
+    instructions = [
+      v_mov_b32_e32(v[0], 0xCAFEBABE),  # source data
+      s_mov_b32(s[0], 0),                # lanesel low = 0 (all read lane 0)
+      s_mov_b32(s[1], 0),                # lanesel high = 0
+      v_permlane16_b32(v[1], v[0], s[0], s[1]),
+    ]
+    st = run_program(instructions, n_lanes=4)
+    # All lanes read from lane 0 of their row
+    for lane in range(4):
+      self.assertEqual(st.vgpr[lane][1], 0xCAFEBABE)
+
+  def test_v_permlanex16_b32_identity(self):
+    """V_PERMLANEX16_B32 cross-row read with identity selection."""
+    # In wave32: row 0 (lanes 0-15) reads from row 1 (lanes 16-31) and vice versa
+    # With single lane in row 0, it reads from lane 0 of row 1 (lane 16)
+    # But lane 16 doesn't exist in 1-lane test, so use 32 lanes
+    instructions = [
+      v_mov_b32_e32(v[0], 0x11111111),  # All lanes have this initially
+      s_mov_b32(s[0], 0x76543210),       # lanesel low
+      s_mov_b32(s[1], 0xFEDCBA98),       # lanesel high
+      v_permlanex16_b32(v[1], v[0], s[0], s[1]),
+    ]
+    st = run_program(instructions, n_lanes=32)
+    # Lane 0 in row 0 reads from lane 0 of row 1 (lane 16)
+    self.assertEqual(st.vgpr[0][1], 0x11111111)
+    # Lane 16 in row 1 reads from lane 0 of row 0 (lane 0)
+    self.assertEqual(st.vgpr[16][1], 0x11111111)
+
+
 if __name__ == '__main__':
  unittest.main()
@@ -149,7 +149,6 @@ class TestFmaMix(unittest.TestCase):

  def test_v_fma_mix_f32_src2_f16_lo(self):
    """V_FMA_MIX_F32 with src2 as f16 from lo bits."""
-    from extra.assembly.amd.test.hw.helpers import f32_to_f16
    f16_2 = f32_to_f16(2.0)
    instructions = [
      s_mov_b32(s[0], f2i(1.0)),
@@ -166,7 +165,6 @@ class TestFmaMix(unittest.TestCase):

  def test_v_fma_mix_f32_src2_f16_hi(self):
    """V_FMA_MIX_F32 with src2 as f16 from hi bits."""
-    from extra.assembly.amd.test.hw.helpers import f32_to_f16
    f16_2 = f32_to_f16(2.0)
    val = (f16_2 << 16) | 0
    instructions = [
@@ -199,7 +197,6 @@ class TestFmaMix(unittest.TestCase):

  def test_v_fma_mix_f32_with_abs_f16_src2_lo(self):
    """V_FMA_MIX_F32 with abs modifier on f16 src2 (lo half). Regression test for sin(1.0) bug."""
-    from extra.assembly.amd.test.hw.helpers import f32_to_f16
    f16_neg1 = f32_to_f16(-1.0)  # 0xbc00
    instructions = [
      s_mov_b32(s[0], f2i(0.0)),  # src0 = 0.0 (f32)
@@ -217,7 +214,6 @@ class TestFmaMix(unittest.TestCase):

  def test_v_fma_mix_f32_with_neg_f16_src2_lo(self):
    """V_FMA_MIX_F32 with neg modifier on f16 src2 (lo half)."""
-    from extra.assembly.amd.test.hw.helpers import f32_to_f16
    f16_1 = f32_to_f16(1.0)  # 0x3c00
    instructions = [
      s_mov_b32(s[0], f2i(0.0)),  # src0 = 0.0 (f32)
@@ -235,7 +231,6 @@ class TestFmaMix(unittest.TestCase):

  def test_v_fma_mix_f32_with_abs_f16_src2_hi(self):
    """V_FMA_MIX_F32 with abs modifier on f16 src2 (hi half)."""
-    from extra.assembly.amd.test.hw.helpers import f32_to_f16
    f16_neg1 = f32_to_f16(-1.0)  # 0xbc00
    val = (f16_neg1 << 16) | 0  # -1.0 in hi, 0 in lo
    instructions = [
@@ -254,7 +249,6 @@ class TestFmaMix(unittest.TestCase):

  def test_v_fma_mixlo_f16(self):
    """V_FMA_MIXLO_F16 writes to low 16 bits of destination."""
-    from extra.assembly.amd.test.hw.helpers import _f16
    instructions = [
      s_mov_b32(s[0], f2i(2.0)),
      v_mov_b32_e32(v[0], s[0]),
@@ -267,14 +261,13 @@ class TestFmaMix(unittest.TestCase):
      VOP3P(VOP3POp.V_FMA_MIXLO_F16, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=0, opsel_hi=0, opsel_hi2=0),
    ]
    st = run_program(instructions, n_lanes=1)
-    lo = _f16(st.vgpr[0][3] & 0xffff)
+    lo = f16(st.vgpr[0][3] & 0xffff)
    hi = (st.vgpr[0][3] >> 16) & 0xffff
    self.assertAlmostEqual(lo, 7.0, places=1)
    self.assertEqual(hi, 0xdead, f"hi should be preserved, got 0x{hi:04x}")

  def test_v_fma_mixlo_f16_all_f32_sources(self):
    """V_FMA_MIXLO_F16 with all f32 sources."""
-    from extra.assembly.amd.test.hw.helpers import _f16
    instructions = [
      s_mov_b32(s[0], f2i(1.0)),
      v_mov_b32_e32(v[0], s[0]),
@@ -286,13 +279,12 @@ class TestFmaMix(unittest.TestCase):
      VOP3P(VOP3POp.V_FMA_MIXLO_F16, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=0, opsel_hi=0, opsel_hi2=0),
    ]
    st = run_program(instructions, n_lanes=1)
-    lo = _f16(st.vgpr[0][3] & 0xffff)
+    lo = f16(st.vgpr[0][3] & 0xffff)
    # 1*2+3 = 5
    self.assertAlmostEqual(lo, 5.0, places=1)

  def test_v_fma_mixlo_f16_sin_case(self):
    """V_FMA_MIXLO_F16 case from sin kernel."""
-    from extra.assembly.amd.test.hw.helpers import _f16
    instructions = [
      s_mov_b32(s[0], 0x3f800000),  # f32 1.0
      v_mov_b32_e32(v[3], s[0]),
@@ -305,7 +297,7 @@ class TestFmaMix(unittest.TestCase):
      VOP3P(VOP3POp.V_FMA_MIXLO_F16, vdst=v[3], src0=v[3], src1=s[6], src2=v[5], opsel=0, opsel_hi=0, opsel_hi2=0),
    ]
    st = run_program(instructions, n_lanes=1)
-    lo = _f16(st.vgpr[0][3] & 0xffff)
+    lo = f16(st.vgpr[0][3] & 0xffff)
    self.assertAlmostEqual(lo, -3.14159, delta=0.01)


@@ -314,7 +306,6 @@ class TestVOP3P(unittest.TestCase):

  def test_v_pk_add_f16_basic(self):
    """V_PK_ADD_F16 adds two packed f16 values."""
-    from extra.assembly.amd.test.hw.helpers import _f16
    instructions = [
      s_mov_b32(s[0], 0x40003c00),  # hi=2.0, lo=1.0
      s_mov_b32(s[1], 0x44004200),  # hi=4.0, lo=3.0
@@ -324,14 +315,13 @@ class TestVOP3P(unittest.TestCase):
    ]
    st = run_program(instructions, n_lanes=1)
    result = st.vgpr[0][2]
-    lo = _f16(result & 0xffff)
-    hi = _f16((result >> 16) & 0xffff)
+    lo = f16(result & 0xffff)
+    hi = f16((result >> 16) & 0xffff)
    self.assertAlmostEqual(lo, 4.0, places=2)
    self.assertAlmostEqual(hi, 6.0, places=2)

  def test_v_pk_mul_f16_basic(self):
    """V_PK_MUL_F16 multiplies two packed f16 values."""
-    from extra.assembly.amd.test.hw.helpers import _f16
    instructions = [
      s_mov_b32(s[0], 0x42004000),  # hi=3.0, lo=2.0
      s_mov_b32(s[1], 0x45004400),  # hi=5.0, lo=4.0
@@ -341,14 +331,13 @@ class TestVOP3P(unittest.TestCase):
    ]
    st = run_program(instructions, n_lanes=1)
    result = st.vgpr[0][2]
-    lo = _f16(result & 0xffff)
-    hi = _f16((result >> 16) & 0xffff)
+    lo = f16(result & 0xffff)
+    hi = f16((result >> 16) & 0xffff)
    self.assertAlmostEqual(lo, 8.0, places=1)
    self.assertAlmostEqual(hi, 15.0, places=1)

  def test_v_pk_fma_f16_basic(self):
    """V_PK_FMA_F16: D = A * B + C for packed f16."""
-    from extra.assembly.amd.test.hw.helpers import _f16
    instructions = [
      s_mov_b32(s[0], 0x42004000),  # A: hi=3.0, lo=2.0
      s_mov_b32(s[1], 0x45004400),  # B: hi=5.0, lo=4.0
@@ -360,8 +349,8 @@ class TestVOP3P(unittest.TestCase):
    ]
    st = run_program(instructions, n_lanes=1)
    result = st.vgpr[0][3]
-    lo = _f16(result & 0xffff)
-    hi = _f16((result >> 16) & 0xffff)
+    lo = f16(result & 0xffff)
+    hi = f16((result >> 16) & 0xffff)
    self.assertAlmostEqual(lo, 9.0, places=1)   # 2*4+1
    self.assertAlmostEqual(hi, 16.0, places=0)  # 3*5+1

@@ -370,7 +359,6 @@ class TestVOP3P(unittest.TestCase):
    Inline constants for VOP3P are f16 values in the low 16 bits only.
    hi half of inline constant is 0, so hi result = v0.hi + 0 = 1.0.
    """
-    from extra.assembly.amd.test.hw.helpers import _f16
    instructions = [
      s_mov_b32(s[0], 0x3c003c00),  # packed f16: hi=1.0, lo=1.0
      v_mov_b32_e32(v[0], s[0]),
@@ -378,8 +366,8 @@ class TestVOP3P(unittest.TestCase):
    ]
    st = run_program(instructions, n_lanes=1)
    result = st.vgpr[0][1]
-    lo = _f16(result & 0xffff)
-    hi = _f16((result >> 16) & 0xffff)
+    lo = f16(result & 0xffff)
+    hi = f16((result >> 16) & 0xffff)
    # lo = 1.0 + 1.0 = 2.0, hi = 1.0 + 0.0 = 1.0 (inline const hi half is 0)
    self.assertAlmostEqual(lo, 2.0, places=2)
    self.assertAlmostEqual(hi, 1.0, places=2)
@@ -388,7 +376,6 @@ class TestVOP3P(unittest.TestCase):
    """V_PK_MUL_F16 with inline constant POS_TWO (2.0).
    Inline constant has value only in low 16 bits, hi is 0.
    """
-    from extra.assembly.amd.test.hw.helpers import _f16
    # v0 = packed (3.0, 4.0), multiply by POS_TWO
    # lo = 3.0 * 2.0 = 6.0, hi = 4.0 * 0.0 = 0.0 (inline const hi is 0)
    instructions = [
@@ -398,8 +385,8 @@ class TestVOP3P(unittest.TestCase):
    ]
    st = run_program(instructions, n_lanes=1)
    result = st.vgpr[0][1]
-    lo = _f16(result & 0xffff)
-    hi = _f16((result >> 16) & 0xffff)
+    lo = f16(result & 0xffff)
+    hi = f16((result >> 16) & 0xffff)
    self.assertAlmostEqual(lo, 6.0, places=1)
    self.assertAlmostEqual(hi, 0.0, places=1)

@@ -413,7 +400,6 @@ class TestWMMAF16(unittest.TestCase):

  def test_v_wmma_f16_16x16x16_f16_all_ones(self):
    """V_WMMA_F16_16X16X16_F16 with all ones produces 16.0 in f16."""
-    from extra.assembly.amd.test.hw.helpers import _f16
    instructions = []
    instructions.append(s_mov_b32(s[0], 0x3c003c00))  # packed f16 1.0
    # Initialize A matrix in v[16:23] (8 regs)
@@ -432,13 +418,12 @@ class TestWMMAF16(unittest.TestCase):
    for lane in range(32):
      for reg in range(8):
        result = st.vgpr[lane][reg]
-        lo = _f16(result & 0xffff)
+        lo = f16(result & 0xffff)
        self.assertAlmostEqual(lo, 16.0, places=1, msg=f"v[{reg}] lane {lane}: expected 16.0, got {lo}")
        self.assertEqual(result >> 16, 0, msg=f"v[{reg}] lane {lane}: hi bits should be 0")

  def test_v_wmma_f16_16x16x16_f16_with_accumulator(self):
    """V_WMMA_F16_16X16X16_F16 with non-zero accumulator."""
-    from extra.assembly.amd.test.hw.helpers import _f16
    instructions = []
    instructions.append(s_mov_b32(s[0], 0x3c003c00))  # packed f16 1.0
    instructions.append(s_mov_b32(s[1], 0x4500))  # f16 5.0 in lo bits only
@@ -458,7 +443,7 @@ class TestWMMAF16(unittest.TestCase):
    for lane in range(32):
      for reg in range(8):
        result = st.vgpr[lane][reg]
-        lo = _f16(result & 0xffff)
+        lo = f16(result & 0xffff)
        self.assertAlmostEqual(lo, 21.0, places=0, msg=f"v[{reg}] lane {lane}: expected 21.0, got {lo}")
        self.assertEqual(result >> 16, 0, msg=f"v[{reg}] lane {lane}: hi bits should be 0")

@@ -468,7 +453,6 @@ class TestWMMAF16(unittest.TestCase):
    Regression test: WMMA was using static register indices instead of dynamic.
    This test uses v[64:71] for A, v[80:87] for B, v[96:103] for C/D.
    """
-    from extra.assembly.amd.test.hw.helpers import _f16
    instructions = []
    instructions.append(s_mov_b32(s[0], 0x3c003c00))  # packed f16 1.0
    # Initialize A matrix in v[64:71] (8 regs)
@@ -490,7 +474,7 @@ class TestWMMAF16(unittest.TestCase):
    for lane in range(32):
      for reg in range(8):
        result = st.vgpr[lane][reg]
-        lo = _f16(result & 0xffff)
+        lo = f16(result & 0xffff)
        self.assertAlmostEqual(lo, 16.0, places=1, msg=f"v[{reg}] lane {lane}: expected 16.0, got {lo}")
        self.assertEqual(result >> 16, 0, msg=f"v[{reg}] lane {lane}: hi bits should be 0")

@@ -713,7 +697,6 @@ class TestPackedMixedSigns(unittest.TestCase):

  def test_pk_add_f16_mixed_signs(self):
    """V_PK_ADD_F16 with mixed positive/negative values."""
-    from extra.assembly.amd.test.hw.helpers import _f16
    instructions = [
      s_mov_b32(s[0], 0xc0003c00),  # packed: hi=-2.0, lo=1.0
      s_mov_b32(s[1], 0x3c003c00),  # packed: hi=1.0, lo=1.0
@@ -723,14 +706,13 @@ class TestPackedMixedSigns(unittest.TestCase):
    ]
    st = run_program(instructions, n_lanes=1)
    result = st.vgpr[0][2]
-    lo = _f16(result & 0xffff)
-    hi = _f16((result >> 16) & 0xffff)
+    lo = f16(result & 0xffff)
+    hi = f16((result >> 16) & 0xffff)
    self.assertAlmostEqual(lo, 2.0, places=2)   # 1.0 + 1.0
    self.assertAlmostEqual(hi, -1.0, places=2)  # -2.0 + 1.0

  def test_pk_mul_f16_zero(self):
    """V_PK_MUL_F16 with zero."""
-    from extra.assembly.amd.test.hw.helpers import _f16
    instructions = [
      s_mov_b32(s[0], 0x40004000),  # packed: 2.0, 2.0
      s_mov_b32(s[1], 0x00000000),  # packed: 0.0, 0.0
@@ -743,5 +725,277 @@ class TestPackedMixedSigns(unittest.TestCase):
    self.assertEqual(result, 0x00000000, "2.0 * 0.0 should be 0.0")


+class TestDot2F32F16(unittest.TestCase):
+  """Tests for V_DOT2_F32_F16 - dot product of f16 pairs producing f32."""
+
+  def test_v_dot2_f32_f16_basic(self):
+    """V_DOT2_F32_F16: dot product of two packed f16 pairs -> f32."""
+    # src0 = {hi=2.0, lo=1.0}, src1 = {hi=4.0, lo=3.0}
+    # result = 1.0*3.0 + 2.0*4.0 + 0 = 3 + 8 = 11.0
+    src0 = (f32_to_f16(2.0) << 16) | f32_to_f16(1.0)
+    src1 = (f32_to_f16(4.0) << 16) | f32_to_f16(3.0)
+    instructions = [
+      s_mov_b32(s[0], src0),
+      s_mov_b32(s[1], src1),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], 0),
+      v_dot2_f32_f16(v[3], v[0], v[1], v[2], opsel_hi=3, opsel_hi2=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = i2f(st.vgpr[0][3])
+    self.assertAlmostEqual(result, 11.0, places=2)
+
+  def test_v_dot2_f32_f16_with_accumulator(self):
+    """V_DOT2_F32_F16 with non-zero f32 accumulator."""
+    # src0 = {hi=1.0, lo=1.0}, src1 = {hi=1.0, lo=1.0}, acc = 5.0
+    # result = 1.0*1.0 + 1.0*1.0 + 5.0 = 7.0
+    src0 = (f32_to_f16(1.0) << 16) | f32_to_f16(1.0)
+    instructions = [
+      s_mov_b32(s[0], src0),
+      s_mov_b32(s[1], f2i(5.0)),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[0]),  # same as src0
+      v_mov_b32_e32(v[2], s[1]),
+      v_dot2_f32_f16(v[3], v[0], v[1], v[2], opsel_hi=3, opsel_hi2=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = i2f(st.vgpr[0][3])
+    self.assertAlmostEqual(result, 7.0, places=2)
+
+  def test_v_dot2_f32_f16_negative_values(self):
+    """V_DOT2_F32_F16 with negative f16 values."""
+    # src0 = {hi=-2.0, lo=3.0}, src1 = {hi=1.0, lo=2.0}
+    # result = 3.0*2.0 + (-2.0)*1.0 + 0 = 6 - 2 = 4.0
+    # NOTE: Hardware DOT2 may have up to 1 ULP difference due to internal implementation
+    src0 = (f32_to_f16(-2.0) << 16) | f32_to_f16(3.0)
+    src1 = (f32_to_f16(1.0) << 16) | f32_to_f16(2.0)
+    instructions = [
+      s_mov_b32(s[0], src0),
+      s_mov_b32(s[1], src1),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], 0),
+      v_dot2_f32_f16(v[3], v[0], v[1], v[2], opsel_hi=3, opsel_hi2=1),
+    ]
+    st = run_program(instructions, n_lanes=1, ulp_tolerance=1)
+    result = i2f(st.vgpr[0][3])
+    self.assertAlmostEqual(result, 4.0, places=2)
+
+
+class TestDot2F16F16(unittest.TestCase):
+  """Tests for V_DOT2_F16_F16 - dot product of f16 pairs producing f16."""
+
+  def test_v_dot2_f16_f16_basic(self):
+    """V_DOT2_F16_F16: dot product of two packed f16 pairs -> f16."""
+    # src0 = {hi=2.0, lo=1.0}, src1 = {hi=3.0, lo=2.0}
+    # result = 1.0*2.0 + 2.0*3.0 + 0 = 2 + 6 = 8.0 (f16)
+    src0 = (f32_to_f16(2.0) << 16) | f32_to_f16(1.0)
+    src1 = (f32_to_f16(3.0) << 16) | f32_to_f16(2.0)
+    instructions = [
+      s_mov_b32(s[0], src0),
+      s_mov_b32(s[1], src1),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], 0),
+      v_dot2_f16_f16(v[3], v[0], v[1], v[2]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = f16(st.vgpr[0][3] & 0xffff)
+    self.assertAlmostEqual(result, 8.0, places=1)
+
+  def test_v_dot2_f16_f16_with_accumulator(self):
+    """V_DOT2_F16_F16 with non-zero f16 accumulator."""
+    # src0 = {hi=1.0, lo=1.0}, src1 = {hi=1.0, lo=1.0}, acc = 3.0 (f16)
+    # result = 1.0*1.0 + 1.0*1.0 + 3.0 = 5.0 (f16)
+    src0 = (f32_to_f16(1.0) << 16) | f32_to_f16(1.0)
+    acc = f32_to_f16(3.0)
+    instructions = [
+      s_mov_b32(s[0], src0),
+      s_mov_b32(s[2], acc),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[0]),  # same as src0
+      v_mov_b32_e32(v[2], s[2]),
+      v_dot2_f16_f16(v[3], v[0], v[1], v[2]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = f16(st.vgpr[0][3] & 0xffff)
+    self.assertAlmostEqual(result, 5.0, places=1)
+
+
+class TestSignedDotProducts(unittest.TestCase):
+  """Tests for V_DOT4_I32_IU8 and V_DOT8_I32_IU4 with signed inputs."""
+
+  def test_v_dot4_i32_iu8_signed_both(self):
+    """V_DOT4_I32_IU8 with both inputs signed (neg=0b011)."""
+    # src0 = {-1, -2, 3, 4} as i8 = {0xff, 0xfe, 0x03, 0x04}
+    # src1 = {1, 1, 1, 1} as i8
+    # result = (-1)*1 + (-2)*1 + 3*1 + 4*1 = -1 - 2 + 3 + 4 = 4
+    src0 = (0xff << 24) | (0xfe << 16) | (0x03 << 8) | 0x04  # -1, -2, 3, 4
+    src1 = 0x01010101  # 1, 1, 1, 1
+    instructions = [
+      s_mov_b32(s[0], src0),
+      s_mov_b32(s[1], src1),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], 0),
+      v_dot4_i32_iu8(v[3], v[0], v[1], v[2], neg=0b011),  # both signed
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][3]
+    # Result is i32, interpret as signed
+    if result >= 0x80000000:
+      result = result - 0x100000000
+    self.assertEqual(result, 4)
+
+  def test_v_dot4_i32_iu8_src0_signed(self):
+    """V_DOT4_I32_IU8 with only src0 signed (neg=0b001)."""
+    # src0 = {-1, -1, -1, -1} as i8 = {0xff, 0xff, 0xff, 0xff}
+    # src1 = {2, 2, 2, 2} as u8
+    # result = (-1)*2 + (-1)*2 + (-1)*2 + (-1)*2 = -8
+    src0 = 0xffffffff  # -1, -1, -1, -1 (as i8)
+    src1 = 0x02020202  # 2, 2, 2, 2 (as u8)
+    instructions = [
+      s_mov_b32(s[0], src0),
+      s_mov_b32(s[1], src1),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], 0),
+      v_dot4_i32_iu8(v[3], v[0], v[1], v[2], neg=0b001),  # src0 signed
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][3]
+    if result >= 0x80000000:
+      result = result - 0x100000000
+    self.assertEqual(result, -8)
+
+  def test_v_dot4_i32_iu8_src1_signed(self):
+    """V_DOT4_I32_IU8 with only src1 signed (neg=0b010)."""
+    # src0 = {2, 2, 2, 2} as u8
+    # src1 = {-1, -1, -1, -1} as i8 = {0xff, 0xff, 0xff, 0xff}
+    # result = 2*(-1) + 2*(-1) + 2*(-1) + 2*(-1) = -8
+    src0 = 0x02020202  # 2, 2, 2, 2 (as u8)
+    src1 = 0xffffffff  # -1, -1, -1, -1 (as i8)
+    instructions = [
+      s_mov_b32(s[0], src0),
+      s_mov_b32(s[1], src1),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], 0),
+      v_dot4_i32_iu8(v[3], v[0], v[1], v[2], neg=0b010),  # src1 signed
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][3]
+    if result >= 0x80000000:
+      result = result - 0x100000000
+    self.assertEqual(result, -8)
+
+  def test_v_dot4_i32_iu8_unsigned_as_reference(self):
+    """V_DOT4_I32_IU8 with both unsigned (neg=0) - same as V_DOT4_U32_U8."""
+    # src0 = {0xff, 0xff, 0xff, 0xff} = 255 each as u8
+    # src1 = {1, 1, 1, 1}
+    # result = 255*1 + 255*1 + 255*1 + 255*1 = 1020
+    src0 = 0xffffffff
+    src1 = 0x01010101
+    instructions = [
+      s_mov_b32(s[0], src0),
+      s_mov_b32(s[1], src1),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], 0),
+      v_dot4_i32_iu8(v[3], v[0], v[1], v[2], neg=0),  # both unsigned
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][3], 1020)
+
+  def test_v_dot8_i32_iu4_signed_both(self):
+    """V_DOT8_I32_IU4 with both inputs signed (neg=0b011)."""
+    # src0 = 8 nibbles: {-1, -2, 3, 4, -1, -2, 3, 4} as i4
+    # i4 -1 = 0xf, -2 = 0xe, 3 = 0x3, 4 = 0x4
+    # src0 = 0xfe34fe34
+    # src1 = {1, 1, 1, 1, 1, 1, 1, 1} as i4 = 0x11111111
+    # result = 2 * ((-1)*1 + (-2)*1 + 3*1 + 4*1) = 2 * 4 = 8
+    src0 = 0xfe34fe34
+    src1 = 0x11111111
+    instructions = [
+      s_mov_b32(s[0], src0),
+      s_mov_b32(s[1], src1),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], 0),
+      v_dot8_i32_iu4(v[3], v[0], v[1], v[2], neg=0b011),  # both signed
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][3]
+    if result >= 0x80000000:
+      result = result - 0x100000000
+    self.assertEqual(result, 8)
+
+  def test_v_dot8_i32_iu4_all_negative(self):
+    """V_DOT8_I32_IU4 with all negative signed values."""
+    # src0 = 8 nibbles all -1 (0xf) = 0xffffffff
+    # src1 = 8 nibbles all 1 = 0x11111111
+    # result = 8 * ((-1)*1) = -8
+    src0 = 0xffffffff  # all -1 as i4
+    src1 = 0x11111111  # all 1
+    instructions = [
+      s_mov_b32(s[0], src0),
+      s_mov_b32(s[1], src1),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], 0),
+      v_dot8_i32_iu4(v[3], v[0], v[1], v[2], neg=0b011),  # both signed
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][3]
+    if result >= 0x80000000:
+      result = result - 0x100000000
+    self.assertEqual(result, -8)
+
+
+class TestPkMinMaxF16(unittest.TestCase):
+  """Tests for V_PK_MIN_F16 and V_PK_MAX_F16."""
+
+  def test_v_pk_min_f16_basic(self):
+    """V_PK_MIN_F16: packed min of two f16 pairs."""
+    # src0 = {hi=3.0, lo=1.0}, src1 = {hi=2.0, lo=4.0}
+    # result = {min(3,2)=2, min(1,4)=1}
+    src0 = (f32_to_f16(3.0) << 16) | f32_to_f16(1.0)
+    src1 = (f32_to_f16(2.0) << 16) | f32_to_f16(4.0)
+    instructions = [
+      s_mov_b32(s[0], src0),
+      s_mov_b32(s[1], src1),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_pk_min_f16(v[2], v[0], v[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][2]
+    lo = f16(result & 0xffff)
+    hi = f16((result >> 16) & 0xffff)
+    self.assertAlmostEqual(lo, 1.0, delta=0.01)
+    self.assertAlmostEqual(hi, 2.0, delta=0.01)
+
+  def test_v_pk_max_f16_basic(self):
+    """V_PK_MAX_F16: packed max of two f16 pairs."""
+    # src0 = {hi=3.0, lo=1.0}, src1 = {hi=2.0, lo=4.0}
+    # result = {max(3,2)=3, max(1,4)=4}
+    src0 = (f32_to_f16(3.0) << 16) | f32_to_f16(1.0)
+    src1 = (f32_to_f16(2.0) << 16) | f32_to_f16(4.0)
+    instructions = [
+      s_mov_b32(s[0], src0),
+      s_mov_b32(s[1], src1),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_pk_max_f16(v[2], v[0], v[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][2]
+    lo = f16(result & 0xffff)
+    hi = f16((result >> 16) & 0xffff)
+    self.assertAlmostEqual(lo, 4.0, delta=0.01)
+    self.assertAlmostEqual(hi, 3.0, delta=0.01)
+
+
 if __name__ == '__main__':
  unittest.main()
@@ -294,7 +294,7 @@ class TestAllPcode(unittest.TestCase):
            'ADDR': u32(), 'ADDR_BASE': u32(), 'TADDR': u32(), 'DATA': u32(), 'DATA0': u32(), 'DATA1': u32(), 'DATA2': u32(),
            'VDATA': u32(), 'VDATA0': u32(), 'VDATA1': u32(), 'VDATA2': u32(), 'VDATA3': u32(),
            'OPSEL': u32(), 'OPSEL_HI': u32(), 'NEG': u32(), 'NEG_HI': u32(), 'CLAMP': u32(),
-            'M0': u32(), 'PC': u64(), 'DENORM': u32(1), 'ROUND_MODE': u32(), 'WAVE_STATUS': u32(),
+            'M0': u32(), 'PC': u64(), 'DENORM': u32(1), 'ROUND_MODE': u32(), 'ROUND_TOWARD_ZERO': u32(), 'ROUND_NEAREST_EVEN': u32(), 'WAVE_STATUS': u32(),
            'MAX_FLOAT_F32': u32(0x7f7fffff), 'Unsigned': u32(1), 'clampedLOD': u32(),
            '_lds': lds, '_vmem': lds, '_active': UOp.const(dtypes.bool, True)}

@@ -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"
@@ -471,7 +471,7 @@ THREADS = 128

 def test_matmul():
  dev = Device[Device.DEFAULT]
-  print(f"Device arch: {dev.arch}")
+  print(f"Device arch: {dev.renderer.arch}")

  if getenv("STOCK", 0):
    # Load the stock kernel from amd_seb/kernel8_batched_gmem.s
@@ -479,7 +479,7 @@ def test_matmul():
    asm = stock_path.read_text()
    print(f"Loaded stock kernel from {stock_path}")
  else:
-    asm = build_kernel(dev.arch)
+    asm = build_kernel(dev.renderer.arch)

  binary = dev.compiler.compile(asm)
  print(f"Compiled! Binary size: {len(binary)} bytes")
@@ -0,0 +1,95 @@
+import atexit, functools
+from tinygrad.runtime.support.compiler_amd import HIPCompiler
+from tinygrad import Tensor, Device, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
+from tinygrad.renderer import Estimates
+from tinygrad.helpers import getenv, all_same, dedup
+from extra.gemm.asm.cdna.asm import build_kernel, GEMM_ARGS
+
+# ** CDNA4 assembly gemm
+
+WORKGROUP_SIZE = 256
+
+def custom_asm_gemm(C:UOp, A:UOp, B:UOp, dname:str, arch:str, wg:int) -> UOp:
+  batch, M, K = A.shape
+  K2, N = B.shape[(1 if B.ndim == 3 else 0):]
+  assert K == K2
+  lidx = UOp.special(WORKGROUP_SIZE, "lidx0")
+  gidx = UOp.special(wg, "gidx0")
+  k = build_kernel(batch, M, N, K, A.dtype.base)
+  sink = UOp.sink(C.base, A.base, B.base, lidx, gidx,
+                  arg=KernelInfo(name=k.name, estimates=Estimates(ops=2*batch*M*N*K, mem=(batch*M*K + K*N + batch*M*N)*2)))
+  binary = HIPCompiler(arch).compile(k.to_asm())
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+                               UOp(Ops.SOURCE, arg=k.to_text()), UOp(Ops.BINARY, arg=binary)))
+
+counters = {"used":0, "todos":[]}
+def todo(msg:str) -> bool: counters["todos"].append(msg); return False
+atexit.register(lambda: print(f'asm_gemm: {counters["used"]} used, {len(counters["todos"])} not used'))
+
+def can_use_asm_gemm(a:Tensor, b:Tensor) -> bool:
+  if a.dtype != b.dtype: return todo(f"dtypes must match {a.dtype} != {b.dtype}")
+  if a.dtype not in {dtypes.bfloat16, dtypes.float16}: return todo(f"only bfloat16/float16, got {a.dtype}")
+  # only sharding on the batch is tested, others might work too
+  if isinstance(a.device, tuple) and not (a.ndim == 3 and a.uop.axis == 0 and b.uop.axis is None):
+    return todo(f"sharding mismatch a.ndim={a.ndim} a.uop.axis={a.uop.axis} b.uop.axis={b.uop.axis}")
+  batch, M, K = (1, *a.shape) if a.ndim == 2 else a.shape
+  N = b.shape[1]
+  if isinstance(a.device, tuple): batch //= len(a.device)
+  if batch not in {1, 2}: return todo(f"GEMM batch size {batch}")
+  if (key:=(M, N, K)) not in GEMM_ARGS: return todo(f"GEMM shape not supported {key}")
+  return True
+
+# ** UOp gemm to test Tensor.custom_kernel multi and backward correctness on non cdna4
+# note: this can be removed after we have GEMM on mixins
+
+def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
+  M, K = A.shape[0]*A.shape[1], A.shape[2]
+  K2, N = B.shape[(1 if B.ndim == 3 else 0):]
+  assert K == K2
+  m = UOp.range(M, 1, AxisType.LOOP)
+  n = UOp.range(N, 2, AxisType.LOOP)
+  k = UOp.range(K, 0, AxisType.REDUCE)
+  mul = (A.index((m*UOp.const(dtypes.index, K)+k))*B.index((k*UOp.const(dtypes.index, N)+n))).cast(dtypes.float32)
+  red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype.base)
+  store = C.index((m*UOp.const(dtypes.index, N)+n), ptr=True).store(red).end(m, n)
+  return store.sink(arg=KernelInfo(name=f'uop_gemm_{M}_{N}_{K}'))
+
+# ** backward gemm, might use the asm gemm
+
+def custom_gemm_bw(gradient:UOp, kernel:UOp):
+  out, a, b = kernel.src
+  assert all_same([gradient.device, a.device, b.device, out.device])
+  a_t, b_t, g_t = Tensor(a, device=a.device), Tensor(b, device=a.device), Tensor(gradient, device=a.device)
+  grad_a = (g_t @ b_t.T).uop
+  a_T = a_t.transpose(-2, -1)
+  a_T = a_T.reshape(*a_T.shape[:-1], 1, a_T.shape[-1])
+  g_r = g_t.reshape(*g_t.shape[:-2], 1, *g_t.shape[-2:]).transpose(-1, -2)
+  grad_b = (a_T * g_r).sum((-1, 0)).uop
+  return (None, grad_a, grad_b)
+
+# ** main gemm function
+
+def asm_gemm(a:Tensor, b:Tensor) -> Tensor:
+  assert can_use_asm_gemm(a, b), f"{counters['todos'][-1]}"
+  counters["used"] += 1
+  squeeze = a.ndim == 2
+  if squeeze: a = a.unsqueeze(0)
+
+  batch, M, K = a.shape
+  N = b.shape[1]
+  is_multi = isinstance(a.device, tuple)
+
+  if is_multi:
+    out = Tensor(Tensor.empty(batch//len(a.device), M, N, dtype=a.dtype, device=a.device).uop.multi(0), device=a.device)
+  else:
+    out = Tensor.empty(batch, M, N, dtype=a.dtype, device=a.device)
+
+  dname = a.device[0] if is_multi else a.device
+  arch = getattr(Device[dname].renderer, "arch", None)
+  if arch.startswith("gfx950") and getenv("USE_ASM", 1):
+    numWG = GEMM_ARGS[(M, N, K)][0]
+    out = Tensor.custom_kernel(out, a, b, fxn=functools.partial(custom_asm_gemm, dname=dname, wg=numWG, arch=arch), grad_fxn=custom_gemm_bw)[0]
+  else:
+    out = Tensor.custom_kernel(out, a, b, fxn=custom_uop_gemm, grad_fxn=custom_gemm_bw)[0]
+  return out.squeeze(0) if squeeze else out
@@ -1,78 +0,0 @@
-.text
-.section	.text.
-.global	gemm
-.p2align	8
-.type	gemm,@function
-
-gemm:
-INSTRUCTIONS
-
-.section .rodata,"a",@progbits
-.p2align 6, 0x0
-.amdhsa_kernel gemm
-  # basic memory requirements
-  .amdhsa_group_segment_fixed_size 133120
-  .amdhsa_private_segment_fixed_size 0
-  .amdhsa_kernarg_size 28
-  # register usage (RSRC1)
-  .amdhsa_next_free_vgpr 504
-  .amdhsa_next_free_sgpr 96
-  # workgroup / workitem IDs (RSRC2)
-  .amdhsa_system_sgpr_workgroup_id_x 1
-  .amdhsa_system_sgpr_workgroup_id_y 1
-  .amdhsa_system_sgpr_workgroup_id_z 1
-  # user SGPRs, we only specify the kernel args ptr in s[0:1]
-  .amdhsa_user_sgpr_kernarg_segment_ptr 1
-  .amdhsa_user_sgpr_count 2
-  .amdhsa_user_sgpr_kernarg_preload_length 0
-  .amdhsa_user_sgpr_kernarg_preload_offset 0
-  # gfx90a / gfx940 specifics (RSRC3)
-  .amdhsa_accum_offset 248
-  .amdhsa_uses_dynamic_stack 0
-  .amdhsa_tg_split 0
-.end_amdhsa_kernel
-
-.amdgpu_metadata
---
-amdhsa.kernels:
-  - .name: gemm
-    .symbol: gemm.kd
-    .args:
-      - .name: C
-        .address_space: global
-        .offset: 0
-        .size: 8
-        .value_kind: global_buffer
-        .value_type: bf16
-      - .name: B
-        .address_space: global
-        .offset: 8
-        .size: 8
-        .value_kind: global_buffer
-        .value_type: bf16
-      - .name: A
-        .address_space: global
-        .offset: 16
-        .size: 8
-        .value_kind: global_buffer
-        .value_type: bf16
-      - .name: sz
-        .offset: 24
-        .size: 4
-        .value_kind: by_value
-        .value_type: u32
-    .group_segment_fixed_size: 133120
-    .private_segment_fixed_size: 0
-    .kernarg_segment_align: 8
-    .kernarg_segment_size: 28
-    .max_flat_workgroup_size: 256
-    .sgpr_count: 88
-    .sgpr_spill_count: 0
-    .vgpr_count: 248
-    .vgpr_spill_count: 0
-    .wavefront_size: 64
-amdhsa.version:
-  - 1
-  - 0
-...
-.end_amdgpu_metadata
@@ -1,73 +0,0 @@
-# Run assembly on the AMD runtime and check correctness
-# VIZ=2 to profile
-import pathlib
-from tinygrad import Tensor, Device, dtypes, Context
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
-from tinygrad.engine.realize import Estimates
-from tinygrad.helpers import getenv
-
-fp = pathlib.Path(__file__).parent/"gemm.s"
-
-N = getenv("N", 8192)
-THREADS_PER_WG = 256
-NUM_WG = N//THREADS_PER_WG * N//THREADS_PER_WG
-
-assert N % THREADS_PER_WG == 0, "N must be divisible by THREADS_PER_WG"
-
-# ** generate inputs on CPU
-
-scale = 10.0
-
-import torch
-torch.manual_seed(0)
-A = (torch.randn(N, N, dtype=torch.float32, device="cpu") / scale).to(torch.bfloat16).contiguous()
-B = (torch.randn(N, N, dtype=torch.float32, device="cpu") / scale).to(torch.bfloat16).contiguous()
-Bt = B.t().contiguous() # transpose B for the asm gemm
-C_torch = A@B
-
-# ** copy buffers to AMD
-
-# input creation and validation run on the copy engine for simpler tracing
-
-def from_torch(t:torch.Tensor) -> Tensor:
-  return Tensor.from_blob(t.data_ptr(), t.shape, dtype=dtypes.bfloat16, device="cpu").to(Device.DEFAULT).realize()
-
-C_tiny = from_torch(A) @ from_torch(B)
-C_asm = Tensor.empty_like(C_tiny)
-
-# ** assembly custom kernel
-
-def custom_asm_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
-  lidx = UOp.special(THREADS_PER_WG, "lidx0")
-  gidx = UOp.special(NUM_WG, "gidx0")
-
-  src = (pathlib.Path(__file__).parent/"template.s").read_text().replace("INSTRUCTIONS", fp.read_text())
-
-  sz = UOp.variable("SZ", 256, 8192)
-
-  sink = UOp.sink(C.base, A.base, B.base, sz, lidx, gidx, arg=KernelInfo(name="gemm", estimates=Estimates(ops=N*N*N*2, mem=N*N*4*3)))
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src)))
-
-C_asm = Tensor.custom_kernel(C_asm, from_torch(A), from_torch(Bt), fxn=custom_asm_gemm)[0]
-
-# ** run gemms
-
-sched = Tensor.schedule(C_tiny, C_asm)
-eis = [si.lower() for si in sched]
-
-with Context(DEBUG=2):
-  for ei in eis:
-    et = ei.run({"SZ":N}, wait=True)
-    print(f"{(N*N*N*2 / et)*1e-12:.2f} REAL TFLOPS")
-
-# ** correctness
-
-import ctypes
-
-def torch_bf16(t:Tensor) -> torch.tensor:
-  asm_out = t.to("cpu").realize().uop.buffer._buf
-  buf = (ctypes.c_uint16*C_asm.uop.size).from_address(asm_out.va_addr)
-  return torch.frombuffer(buf, dtype=torch.bfloat16, count=C_asm.uop.size).reshape(C_asm.shape)
-
-assert torch.allclose(torch_bf16(C_asm), C_torch, rtol=1e-2, atol=1e-3)
-assert torch.allclose(torch_bf16(C_tiny), C_torch, rtol=1e-2, atol=1e-3)
@@ -0,0 +1,46 @@
+import unittest
+from tinygrad import Tensor, Device, dtypes, Context
+from tinygrad.helpers import getenv
+from extra.gemm.asm.cdna.gemm import asm_gemm
+
+def verify_asm_gemm(batch:int, M:int, N:int, K:int, dtype=dtypes.bfloat16, multi=False) -> None:
+  Tensor.manual_seed(0)
+  a_rand = Tensor.randn((batch, M, K), dtype=dtypes.float).sub(0.5).cast(dtype)
+  b_rand = Tensor.randn((K, N), dtype=dtypes.float).sub(0.5).cast(dtype)
+  with Context(DEBUG=0):
+    Tensor.realize(a_rand, b_rand)
+
+  devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(8)) if multi else None
+
+  a, b = Tensor(a_rand.numpy(), requires_grad=True).cast(dtype), Tensor(b_rand.numpy(), requires_grad=True).cast(dtype)
+  if multi: a, b = a.shard(devs, axis=0), b.shard(devs, axis=None)
+  tst = asm_gemm(a, b)
+  tst.sum().backward()
+  Tensor.realize(tst, a.grad, b.grad)
+
+  a_ref, b_ref = Tensor(a_rand.numpy(), requires_grad=True).cast(dtype), Tensor(b_rand.numpy(), requires_grad=True).cast(dtype)
+  if multi: a_ref, b_ref = a_ref.shard(devs, axis=0), b_ref.shard(devs, axis=None)
+  with Context(ASM_GEMM=0): ref = a_ref @ b_ref
+  ref.sum().backward()
+  Tensor.realize(ref, a_ref.grad, b_ref.grad)
+
+  with Context(DEBUG=0):
+    assert (tst - ref).square().max().float().item() < 1e-6, "forward mismatch"
+    assert (a.grad - a_ref.grad).square().max().float().item() < 1e-3, "grad_a mismatch"
+    assert (b.grad - b_ref.grad).square().max().float().item() < 1e-3, "grad_b mismatch"
+
+class TestGemm(unittest.TestCase):
+  def test_simple(self): verify_asm_gemm(1, N:=getenv("N", 4096), N, N, dtype=dtypes.half)
+
+  def test_gemm1(self): verify_asm_gemm(8, 8192, 4096, 14336, multi=True)
+  def test_gemm2(self): verify_asm_gemm(8, 8192, 128256, 4096, multi=True)
+  def test_gemm3(self): verify_asm_gemm(8, 8192, 14336, 4096, multi=True)
+  def test_gemm4(self): verify_asm_gemm(8, 4096, 14336, 4096, multi=True)
+  def test_gemm5(self): verify_asm_gemm(8, 4096, 4096, 14336, multi=True)
+  def test_gemm6(self): verify_asm_gemm(16, 4096, 4096, 14336, multi=True)
+  def test_gemm_unsupported(self):
+    with self.assertRaisesRegex(AssertionError, "shape not supported"):
+      verify_asm_gemm(8, 8192, 1024, 4096, multi=True)
+
+if __name__ == "__main__":
+  unittest.main()
@@ -1,14 +1,15 @@
 #!/usr/bin/env python3

 import argparse, glob, os, time, subprocess, sys
+from tinygrad.helpers import temp

 def scan_devs_based_on_lock(prefix:str, args) -> list[str]:
  target_dev = args.pci_bus if 'pci_bus' in args.__dir__() else ""

  devs = []
-  for dev in glob.glob(f'/tmp/{prefix}_*.lock'):
-    dev_id = dev[8:-5]
-    if os.path.exists(f"/sys/bus/pci/devices/{dev_id}") and dev_id.startswith(target_dev): devs.append(dev_id)
+  for dev in glob.glob(temp(f'{prefix}_*.lock')):
+    dev_id = dev.split('/')[-1][len(prefix)+1:-5]
+    if dev_id.startswith(target_dev): devs.append(dev_id)
  return devs

 def _do_reset_device(pci_bus): os.system(f"sudo sh -c 'echo 1 > /sys/bus/pci/devices/{pci_bus}/reset'")
@@ -53,16 +54,7 @@ def cmd_show_pids(args):

  for dev in devs:
    try:
-      pid = subprocess.check_output(['sudo', 'lsof', f'/tmp/{prefix}_{dev}.lock']).decode('utf-8').strip().split('\n')[1].split()[1]
-      print(f"{dev}: {pid}")
-    except subprocess.CalledProcessError: print(f"{dev}: No processes found using this device")
-
-def cmd_kill_pids(args):
-  devs = scan_devs_based_on_lock(prefix:={"amd":"am", "nv":"nv"}[args.backend], args)
-
-  for dev in devs:
-    try:
-      pid = subprocess.check_output(['sudo', 'lsof', f'/tmp/{prefix}_{dev}.lock']).decode('utf-8').strip().split('\n')[1].split()[1]
+      pid = subprocess.check_output(['sudo', 'lsof', temp(f'{prefix}_{dev}.lock')]).decode('utf-8').strip().split('\n')[1].split()[1]
      print(f"{dev}: {pid}")
    except subprocess.CalledProcessError: print(f"{dev}: No processes found using this device")

@@ -74,7 +66,7 @@ def cmd_kill_pids(args):
      if i > 0: time.sleep(0.2)

      try:
-        try: pid = subprocess.check_output(['sudo', 'lsof', f'/tmp/{prefix}_{dev}.lock']).decode('utf-8').strip().split('\n')[1].split()[1]
+        try: pid = subprocess.check_output(['sudo', 'lsof', temp(f'{prefix}_{dev}.lock')]).decode('utf-8').strip().split('\n')[1].split()[1]
        except subprocess.CalledProcessError: break

        print(f"Killing process {pid} (which uses {dev})")
@@ -202,7 +202,7 @@ def ioctl(fd, request, argp):
        if s.hClass == nv_gpu.NV1_MEMORY_SYSTEM: dump_struct(get_struct(s.pAllocParms, nv_gpu.NV_MEMORY_ALLOCATION_PARAMS))
        if s.hClass == nv_gpu.GT200_DEBUGGER: dump_struct(get_struct(s.pAllocParms, nv_gpu.NV83DE_ALLOC_PARAMETERS))
        if s.hClass == nv_gpu.MAXWELL_PROFILER_DEVICE: dump_struct(get_struct(s.pAllocParms, nv_gpu.NVB2CC_ALLOC_PARAMETERS))
-        if s.hClass == nv_gpu.AMPERE_CHANNEL_GPFIFO_A:
+        if s.hClass in {nv_gpu.AMPERE_CHANNEL_GPFIFO_A, nv_gpu.BLACKWELL_CHANNEL_GPFIFO_A}:
          sx = get_struct(s.pAllocParms, nv_gpu.NV_CHANNELGPFIFO_ALLOCATION_PARAMETERS)
          dump_struct(sx)
          gpus_fifo.append((sx.gpFifoOffset, sx.gpFifoEntries))
@@ -0,0 +1,193 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+import enum, collections
+from typing import Iterator
+from tinygrad.helpers import colored
+from extra.assembly.amd.sqtt import PacketType, bits
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# STALL REASONS
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class StallReason(enum.IntEnum):
+  # Based on CUpti_ActivityPCSamplingStallReason
+  INVALID = 0
+  NONE = 1              # selected, selected_not_issued
+  INST_FETCH = 2        # branch_resolving, no_instructions
+  EXEC_DEPENDENCY = 3   # short_scoreboard, wait
+  MEMORY_DEPENDENCY = 4 # long_scoreboard
+  TEXTURE = 5           # tex_throttle
+  SYNC = 6              # barrier, membar
+  CONSTANT_MEMORY = 7   # imc_miss
+  PIPE_BUSY = 8         # mio_throttle, math_pipe_throttle
+  MEMORY_THROTTLE = 9   # drain, lg_throttle
+  NOT_SELECTED = 10     # not_selected
+  OTHER = 11            # misc, dispatch_stall
+  SLEEPING = 12         # sleeping
+
+STALL_KEY_MAP_AMPERE: dict[int, StallReason] = {
+  1: StallReason.MEMORY_THROTTLE, 15: StallReason.MEMORY_THROTTLE,
+  2: StallReason.CONSTANT_MEMORY,
+  3: StallReason.SYNC,
+  6: StallReason.INST_FETCH, 11: StallReason.INST_FETCH,
+  7: StallReason.EXEC_DEPENDENCY, 10: StallReason.EXEC_DEPENDENCY,
+  9: StallReason.MEMORY_DEPENDENCY,
+  12: StallReason.PIPE_BUSY,
+  17: StallReason.OTHER, 20: StallReason.OTHER,
+  18: StallReason.NONE,
+}
+
+STALL_KEY_MAP_BLACKWELL: dict[int, StallReason] = {
+  0x01: StallReason.MEMORY_THROTTLE, 0x0e: StallReason.MEMORY_THROTTLE,
+  0x02: StallReason.SYNC,
+  0x05: StallReason.INST_FETCH, 0x0a: StallReason.INST_FETCH,
+  0x06: StallReason.EXEC_DEPENDENCY, 0x09: StallReason.EXEC_DEPENDENCY,
+  0x08: StallReason.MEMORY_DEPENDENCY,
+  0x0b: StallReason.PIPE_BUSY, 0x0f: StallReason.PIPE_BUSY,
+  0x10: StallReason.OTHER, 0x13: StallReason.OTHER,
+  0x11: StallReason.NONE,
+}
+
+# Lookup table for extracting sample bytes from 32-byte packet (bytes 0-3, 8-31, skipping header at 4-7)
+LOOKUP_28B = [0, 1, 2, 3, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31]
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# PACKET HEADER
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class PMAHeader(PacketType):
+  num_bytes  = bits[4:0]    # number of sample bytes in this packet
+  tpc_id_lo  = bits[15:8]   # TPC identifier low 8 bits
+  tpc_id_hi  = bits[27:25]  # TPC identifier high 3 bits
+  dropped    = bits[28:28]  # dropped flag (resets byte accumulator)
+  @property
+  def tpc_id(self) -> int: return self.tpc_id_lo | (self.tpc_id_hi << 8)
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# 8-BYTE SAMPLE FORMAT (Ampere/Ada/Hopper)
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class PMASampleAmpere8B(PacketType):
+  pc_raw     = bits[44:0]   # raw PC value (pc_offset = pc_raw << 4)
+  stall_key  = bits[49:45]  # stall reason key
+  wave_id    = bits[55:50]  # warp/wave identifier
+  active     = bits[62:62]  # 1 if warp was executing, 0 if scheduled but not issued
+  @property
+  def pc_offset(self) -> int: return self.pc_raw << 4
+  @property
+  def stall_reason(self) -> StallReason: return STALL_KEY_MAP_AMPERE.get(self.stall_key, StallReason.OTHER)
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# 9-BYTE SAMPLE FORMAT (Blackwell+)
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class PMASampleBlackwell9B(PacketType):
+  stall_key  = bits[5:0]    # stall reason key
+  pc_raw     = bits[60:8]   # raw PC value (pc_offset = pc_raw << 4)
+  wave_hi    = bits[7:6]    # wave_id high 2 bits
+  wave_lo    = bits[71:68]  # wave_id low 4 bits
+  active     = bits[67:67]  # 1 if warp was executing, 0 if scheduled but not issued
+  @property
+  def pc_offset(self) -> int: return self.pc_raw << 4
+  @property
+  def stall_reason(self) -> StallReason: return STALL_KEY_MAP_BLACKWELL.get(self.stall_key, StallReason.OTHER)
+  @property
+  def wave_id(self) -> int: return (self.wave_hi << 4) | self.wave_lo
+
+PMASample = PMASampleAmpere8B|PMASampleBlackwell9B
+
+def decode(data: bytes, sm_version: int = 0x800) -> Iterator[tuple[PMASample, int]]:
+  use_9byte = sm_version >= 0xa04
+  record_size = 9 if use_9byte else 8
+  sample_cls = PMASampleBlackwell9B if use_9byte else PMASampleAmpere8B
+
+  tpc_state: dict[int, list[int]] = collections.defaultdict(list)
+  for pkt_idx in range(len(data) // 32):
+    pkt = data[pkt_idx * 32:(pkt_idx + 1) * 32]
+    hdr = PMAHeader.from_raw(int.from_bytes(pkt[4:8], 'little'))
+
+    if hdr.dropped: tpc_state[hdr.tpc_id].clear()
+
+    for i in range(hdr.num_bytes):
+      tpc_state[hdr.tpc_id].append(pkt[LOOKUP_28B[i]])
+
+    while len(tpc_state[hdr.tpc_id]) >= record_size:
+      yield sample_cls.from_raw(int.from_bytes(bytes(tpc_state[hdr.tpc_id][:record_size]), 'little')), hdr.tpc_id
+      del tpc_state[hdr.tpc_id][:record_size]
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# CLI
+# ═══════════════════════════════════════════════════════════════════════════════
+
+STALL_COLORS = {
+  StallReason.NONE: "green", StallReason.INST_FETCH: "yellow", StallReason.EXEC_DEPENDENCY: "cyan",
+  StallReason.MEMORY_DEPENDENCY: "red", StallReason.SYNC: "magenta", StallReason.CONSTANT_MEMORY: "blue",
+  StallReason.PIPE_BUSY: "yellow", StallReason.MEMORY_THROTTLE: "RED", StallReason.OTHER: "white",
+}
+
+def decode_tpc_id(tpc_id:int) -> tuple[int, int, int]:
+  # NOTE: valid only for ops_nv, cuda encoding is different
+  return (tpc_id >> 5, (tpc_id >> 1) & 0xf, tpc_id & 1)
+
+def print_samples(samples:list[tuple[PMASample, int]]) -> None:
+  if not samples: return
+  base_pc = min(s.pc_offset for s, _ in samples)
+  for s, tpc_id in samples:
+    gpc, tpc, sm = decode_tpc_id(tpc_id)
+    stall_str = colored(f"{s.stall_reason.name:17}", STALL_COLORS.get(s.stall_reason, "white"))
+    print(f"pc=0x{s.pc_offset - base_pc:06x} {stall_str} ev={s.stall_key:2d} active={s.active} wave={s.wave_id:2d} gpc={gpc} tpc={tpc} sm={sm}")
+
+def print_packets(data:bytes, sm_version:int=0x800) -> None:
+  record_size = 9 if sm_version >= 0x890 else 8
+  tpc_state: dict[int, list[int]] = collections.defaultdict(list)
+  for i in range(len(data) // 32):
+    pkt = data[i * 32:(i + 1) * 32]
+    hdr = PMAHeader.from_raw(int.from_bytes(pkt[4:8], 'little'))
+    if hdr.dropped: tpc_state[hdr.tpc_id].clear()
+    for j in range(hdr.num_bytes): tpc_state[hdr.tpc_id].append(pkt[LOOKUP_28B[j]])
+    # Show complete records extracted from this packet
+    records = []
+    while len(tpc_state[hdr.tpc_id]) >= record_size:
+      records.append(bytes(tpc_state[hdr.tpc_id][:record_size]).hex())
+      del tpc_state[hdr.tpc_id][:record_size]
+    leftover = len(tpc_state[hdr.tpc_id])
+    print(f"Pkt {i:3d}: tpc={hdr.tpc_id:4d} n={hdr.num_bytes:2d} drop={hdr.dropped} left={leftover} | {' '.join(records)}")
+
+def print_aggregated(samples:list[tuple[PMASample, int]]) -> None:
+  if not samples: return
+  base_pc = min(s.pc_offset for s, _ in samples)
+  counter: collections.Counter[tuple[int, StallReason]] = collections.Counter((s.pc_offset, s.stall_reason) for s, _ in samples)
+  print(f"\nAggregated samples (base_pc=0x{base_pc:x}):")
+  for (pc, reason), cnt in sorted(counter.items()):
+    stall_str = colored(f"{reason.name:17}", STALL_COLORS.get(reason, "white"))
+    print(f"  pc=0x{pc - base_pc:06x} {stall_str} samples={cnt:4d}")
+
+if __name__ == "__main__":
+  import sys, pickle
+
+  if len(sys.argv) < 2:
+    print("Usage: python decode.py <pkl_file> [--raw] [--sm=0xNNN]")
+    sys.exit(1)
+
+  with open(sys.argv[1], "rb") as f:
+    data = pickle.load(f)
+
+  if isinstance(data, dict):
+    sm_version = 0x800  # default to Ampere
+    for arg in sys.argv:
+      if arg.startswith("--sm="): sm_version = int(arg[5:], 0)
+    dumps = [(i, x, sm_version) for i, x in enumerate(data["pma_raw_dumps"])]
+  else:
+    devs = {e.device: e for e in data if type(e).__name__ == "ProfileDeviceEvent"}
+    dumps = []
+    for i, e in enumerate(e for e in data if type(e).__name__ == "ProfilePMAEvent"):
+      dumps.append((i, e.blob, devs[e.device].props.get('sm_version', 0x800)))
+
+  for dump_idx, raw, sm_ver in dumps:
+    print(f"\n{'='*60}\nDump {dump_idx} ({len(raw)} bytes, {len(raw)//32} packets)\n{'='*60}")
+    if "--raw" in sys.argv: print_packets(raw, sm_ver)
+    else:
+      samples = list(decode(raw, sm_ver))
+      print(f"\nDecoded {len(samples)} samples:")
+      print_samples(samples)
+      print_aggregated(samples)
@@ -0,0 +1,76 @@
+import pickle, unittest
+from collections import Counter
+from pathlib import Path
+
+from extra.nv_pma.decode import decode
+from tinygrad.helpers import DEBUG
+
+EXAMPLES_DIR = Path(__file__).parent.parent / "examples"
+EXAMPLES_5090_DIR = Path(__file__).parent.parent / "examples_5090"
+
+def decode_and_aggregate(raw_dumps: list[bytes], sm_version: int = 0x800) -> Counter[tuple[int, int]]:
+  """Decode all PMA buffers and aggregate by (relative_pc, stall_reason). Each dump is normalized separately."""
+  result: Counter[tuple[int, int]] = Counter()
+  for raw in raw_dumps:
+    samples = [s for s, _ in decode(raw, sm_version)]
+    if not samples: continue
+    base_pc = min(s.pc_offset for s in samples)
+    result += Counter((s.pc_offset - base_pc, int(s.stall_reason)) for s in samples)
+  return result
+
+def cupti_to_counter(cupti_records: list[dict]) -> Counter[tuple[int, int]]:
+  """Convert CUPTI records to Counter[(pcOffset, stallReason)]."""
+  counter: Counter[tuple[int, int]] = Counter()
+  for r in cupti_records:
+    counter[(r['pcOffset'], r['stallReason'])] += r['samples']
+  return counter
+
+class TestNVProf(unittest.TestCase):
+  def _test_example(self, name: str, sm_version: int = 0x800, examples_dir: Path = EXAMPLES_DIR):
+    pkl_file = examples_dir / f"{name}.pkl"
+    if not pkl_file.exists():
+      self.skipTest(f"Example data not found: {pkl_file}. Run collect.py first.")
+
+    with open(pkl_file, "rb") as f:
+      data = pickle.load(f)
+
+    self.assertEqual(data["test_name"], name)
+    pma_agg = decode_and_aggregate(data["pma_raw_dumps"], sm_version)
+    cupti_agg = cupti_to_counter(data["cupti_pc_samples"])
+
+    if DEBUG >= 2:
+      total = sum(cupti_agg.values())
+      mismatched = sum(abs(pma_agg.get(k, 0) - v) for k, v in cupti_agg.items())
+      mismatched += sum(v for k, v in pma_agg.items() if k not in cupti_agg)
+      mismatched //= 2
+
+      print(f"\n=== Test: {name} ===")
+      print(f"Total samples: {total}, Mismatched: {mismatched} ({mismatched/total*100 if total else 0:.1f}%)")
+
+    self.assertEqual(pma_agg, cupti_agg, f"PMA: {dict(pma_agg)}\nCUPTI: {dict(cupti_agg)}")
+
+  # Ampere tests (8-byte format)
+  def test_decode_test_plus(self): self._test_example("test_plus")
+  def test_decode_test_reduce_sum(self): self._test_example("test_reduce_sum")
+  def test_decode_test_broadcast(self): self._test_example("test_broadcast")
+  def test_decode_test_matmul(self): self._test_example("test_matmul")
+  def test_decode_test_plus_big(self): self._test_example("test_plus_big")
+  def test_decode_test_elementwise_chain(self): self._test_example("test_elementwise_chain")
+  def test_decode_test_conv2d(self): self._test_example("test_conv2d")
+  def test_decode_test_large_matmul(self): self._test_example("test_large_matmul")
+
+  # Blackwell/5090 tests (9-byte format)
+  def test_5090_test_plus(self): self._test_example("test_plus", 0xa04, EXAMPLES_5090_DIR)
+  def test_5090_test_plus_big(self): self._test_example("test_plus_big", 0xa04, EXAMPLES_5090_DIR)
+  def test_5090_test_broadcast(self): self._test_example("test_broadcast", 0xa04, EXAMPLES_5090_DIR)
+  def test_5090_test_matmul(self): self._test_example("test_matmul", 0xa04, EXAMPLES_5090_DIR)
+  def test_5090_test_large_matmul(self): self._test_example("test_large_matmul", 0xa04, EXAMPLES_5090_DIR)
+  def test_5090_test_reduce_sum(self): self._test_example("test_reduce_sum", 0xa04, EXAMPLES_5090_DIR)
+  def test_5090_test_reduce_max(self): self._test_example("test_reduce_max", 0xa04, EXAMPLES_5090_DIR)
+  def test_5090_test_elementwise_chain(self): self._test_example("test_elementwise_chain", 0xa04, EXAMPLES_5090_DIR)
+  def test_5090_test_conv2d(self): self._test_example("test_conv2d", 0xa04, EXAMPLES_5090_DIR)
+  def test_5090_test_exp(self): self._test_example("test_exp", 0xa04, EXAMPLES_5090_DIR)
+  def test_5090_test_softmax(self): self._test_example("test_softmax", 0xa04, EXAMPLES_5090_DIR)
+
+if __name__ == "__main__":
+  unittest.main()
@@ -9,8 +9,8 @@ from extra.thunder.tiny.tk.kernel import Kernel
 from extra.thunder.tiny.tk.tiles import GL, TileLayout

 NUM_WORKERS = 1
-Q_BLOCK_SIZE = 16
-KV_BLOCK_SIZE = 16
+Q_BLOCK_SIZE = 32
+KV_BLOCK_SIZE = 32

 def _sharded_empty(shape:Tensor, ref:Tensor, axis:int|None) -> Tensor:
  if not isinstance(ref.device, tuple): return Tensor.empty(*shape, dtype=ref.dtype, device=ref.device)
@@ -70,10 +70,10 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
      mask_reg = ker.rt((Q_BLOCK_SIZE, KV_BLOCK_SIZE), dtypes.float32)
      mask_reg_transposed = ker.rt((KV_BLOCK_SIZE, Q_BLOCK_SIZE), dtypes.float32, TileLayout.COL)

-      max_vec_last = ker.rv(KV_BLOCK_SIZE, dtypes.float32)
-      max_vec = ker.rv(KV_BLOCK_SIZE, dtypes.float32)
-      norm_vec = ker.rv(KV_BLOCK_SIZE, dtypes.float32)
-      scale_vec = ker.rv(KV_BLOCK_SIZE, dtypes.float32)
+      max_vec_last = ker.rv(Q_BLOCK_SIZE, dtypes.float32)
+      max_vec = ker.rv(Q_BLOCK_SIZE, dtypes.float32)
+      norm_vec = ker.rv(Q_BLOCK_SIZE, dtypes.float32)
+      scale_vec = ker.rv(Q_BLOCK_SIZE, dtypes.float32)

      max_vec = warp.neg_inf(max_vec)
      norm_vec = warp.zero(norm_vec)
@@ -105,7 +105,7 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False

        # softmax
        max_vec_last = warp.copy(max_vec_last.after(kv_idx), max_vec)
-        max_vec = warp.row_reduce(max_vec.after(max_vec_last), att_block, lambda a, b: a.maximum(b), init_value=-math.inf)
+        max_vec = warp.col_reduce(max_vec.after(max_vec_last), att_block, lambda a, b: a.maximum(b), init_value=-math.inf)

        scale_vec = warp.map(scale_vec.after(max_vec_last, max_vec), lambda _, idx: max_vec_last[*idx] - max_vec[*idx])
        scale_vec = scale_vec.exp2()
@@ -116,7 +116,7 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
        att_block -= max_vec
        att_block = att_block.exp2()

-        norm_vec = warp.row_reduce(norm_vec.after(scale_vec), att_block, lambda a, b: a + b)
+        norm_vec = warp.col_reduce(norm_vec.after(scale_vec), att_block, lambda a, b: a + b)

        # mma av
        att_block_mma = warp.copy(att_block_mma.after(kv_idx, norm_vec), att_block)
@@ -313,7 +313,7 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
          att_block_transposed = warp.transpose(att_block_transposed, att_block_mma)
          att_smem = warp.store(att_smem, att_block_transposed)
          att_block_row = warp.load(att_block_row, att_smem)
-          dv_reg_ = warp.mma_AB(dv_reg, att_block_row, do_reg_col)
+          dv_reg_ = warp.mma_AtB(dv_reg, att_block_row, do_reg_col)

          dp_block = warp.zero(dp_block.after(g, q_idx, dv_reg_))
          dp_block = warp.mma_ABt(dp_block, v_reg, do_reg)
@@ -325,7 +325,7 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
          att_block_transposed = warp.transpose(att_block_transposed, att_block_mma)
          att_smem = warp.store(att_smem, att_block_transposed)
          att_block_row = warp.load(att_block_row, att_smem)
-          dk_reg = warp.mma_AB(dk_reg, att_block_row, q_reg_col)
+          dk_reg = warp.mma_AtB(dk_reg, att_block_row, q_reg_col)
      dk_reg = ker.endrange(2)
      dv_reg = dv_reg.after(dk_reg)

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
-import sys, os, zlib, struct, hashlib
-from tinygrad.helpers import DEBUG, getenv, fetch
+import os, zlib, struct, hashlib
+from tinygrad.helpers import getenv
 from tinygrad.runtime.support.usb import USB3

 SUPPORTED_CONTROLLERS = [
@@ -50,7 +50,7 @@ patched_fw = patch(file_path, file_hash, patches)
 dev = None
 for vendor, device in SUPPORTED_CONTROLLERS:
  try:
-    dev = USB3(vendor, device, 0x81, 0x83, 0x02, 0x04)
+    dev = USB3(vendor, device, 0x81, 0x83, 0x02, 0x04, use_bot=True)
    break
  except RuntimeError: pass
 if dev is None:
@@ -10,7 +10,11 @@ def optional_eq(val:dict, arg:str|None) -> bool: return arg is None or ansistrip
 def print_data(data:dict) -> None:
  if isinstance(data.get("value"), Iterator):
    for m in data["value"]:
+      if m.get("uop"):
+        print("Input UOp:")
+        print(m["uop"])
      if not m["diff"]: continue
+      print("Rewrites:")
      fp = pathlib.Path(m["upat"][0][0])
      print(f"{fp.parent.name}/{fp.name}:{m['upat'][0][1]}")
      print(m["upat"][1])
@@ -0,0 +1,137 @@
+# ruff: noqa: F405
+import unittest, subprocess, os
+from extra.assembly.amd.autogen.rdna3.ins import *  # noqa: F403
+from extra.assembly.amd.dsl import s, v, Inst, NULL
+
+def assemble_kernel(insts:list[Inst], 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}
+  disasm = "\n".join(inst.disasm() for inst in insts)
+  hsasrc = f".text\n.globl {name}\n.p2align 8\n.type {name},@function\n{name}:\n{disasm}\n"
+  return hsasrc + f".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"
+
+def _run(code:str, timeout:float=15.0) -> subprocess.CompletedProcess:
+  # TODO: AM_RESET is required for now, so subprocesses
+  return subprocess.run(["python", "-c", code], env={**os.environ, "AMD": "1"}, capture_output=True, text=True, timeout=timeout)
+
+def _run_asm(asm_src:str) -> subprocess.CompletedProcess:
+  return _run('from tinygrad.device import Device; from tinygrad.runtime.ops_amd import AMDProgram; '
+              'from tinygrad.runtime.support.compiler_amd import HIPCompiler; dev = Device["AMD"]; '
+              f'AMDProgram(dev, "test", HIPCompiler(dev.arch).compile("""{asm_src}"""))('
+              'dev.allocator.alloc(64), global_size=(1,1,1), local_size=(1,1,1), wait=True)')
+
+def _verify_recovery() -> subprocess.CompletedProcess:
+  return _run('from tinygrad import Tensor; t = Tensor([1.0, 2.0], device="AMD").realize(); assert (t + 1).numpy().tolist() == [2.0, 3.0]')
+
+_ILLEGAL_INST_ASM = ".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\n.byte 0xff,0xff,0xff,0xff\ns_endpgm\n" \
+  ".rodata\n.p2align 6\n.amdhsa_kernel test\n.amdhsa_next_free_vgpr 8\n.amdhsa_next_free_sgpr 8\n" \
+  ".amdhsa_wavefront_size32 1\n.amdhsa_user_sgpr_kernarg_segment_ptr 1\n.amdhsa_kernarg_size 8\n.end_amdhsa_kernel"
+
+@unittest.skipIf(os.environ.get("AMD") != "1" or os.environ.get("MOCKGPU") == "1", "AMD with AM driver required")
+class TestAMFaultRecovery(unittest.TestCase):
+  def _run_kernel(self, insts: list[Inst]) -> subprocess.CompletedProcess: return _run_asm(assemble_kernel(insts))
+
+  def _assert_fault_and_recovery(self, result:subprocess.CompletedProcess):
+    if result.stdout.strip(): print(f"\nstdout: {result.stdout.strip()}")
+    if result.stderr.strip(): print(f"\nstderr: {result.stderr.strip()}")
+    self.assertNotEqual(result.returncode, 0, f"Expected fault but succeeded: {result.stdout}")
+    self.assertEqual(_verify_recovery().returncode, 0)
+
+
+class TestGlobalMemoryFaults(TestAMFaultRecovery):
+  def test_global_load_unmapped(self):
+    insts = [v_mov_b32_e32(v[0], 0xBEEF0000), v_mov_b32_e32(v[1], 0xDEAD),
+             global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(vmcnt=0), s_endpgm()]
+    self._assert_fault_and_recovery(self._run_kernel(insts))
+
+  def test_global_store_unmapped(self):
+    insts = [v_mov_b32_e32(v[0], 0xBEEF0000), 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(vmcnt=0), s_endpgm()]
+    self._assert_fault_and_recovery(self._run_kernel(insts))
+
+  def test_global_null_ptr(self):
+    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(vmcnt=0), s_endpgm()]
+    self._assert_fault_and_recovery(self._run_kernel(insts))
+
+  def test_global_misaligned_b64(self):
+    insts = [v_mov_b32_e32(v[0], 0xBEEF0001), v_mov_b32_e32(v[1], 0xDEAD),
+             global_load_b64(v[2:3], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(vmcnt=0), s_endpgm()]
+    self._assert_fault_and_recovery(self._run_kernel(insts))
+
+  def test_global_misaligned_b128(self):
+    insts = [v_mov_b32_e32(v[0], 0xBEEF0004), v_mov_b32_e32(v[1], 0xDEAD),
+             global_load_b128(v[2:5], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(vmcnt=0), s_endpgm()]
+    self._assert_fault_and_recovery(self._run_kernel(insts))
+
+
+class TestSMEMFaults(TestAMFaultRecovery):
+  def test_smem_null_base(self):
+    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(lgkmcnt=0), s_endpgm()]
+    self._assert_fault_and_recovery(self._run_kernel(insts))
+
+  def test_smem_unmapped_address(self):
+    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(lgkmcnt=0), s_endpgm()]
+    self._assert_fault_and_recovery(self._run_kernel(insts))
+
+  def test_smem_misaligned_b64(self):
+    insts = [s_mov_b32(s[2], 0xBEEF0004), s_mov_b32(s[3], 0xDEAD),
+             s_load_b64(s[4:5], s[2:3], 0, soffset=NULL), s_waitcnt(lgkmcnt=0), s_endpgm()]
+    self._assert_fault_and_recovery(self._run_kernel(insts))
+
+  def test_smem_misaligned_b128(self):
+    insts = [s_mov_b32(s[2], 0xBEEF0004), s_mov_b32(s[3], 0xDEAD),
+             s_load_b128(s[4:7], s[2:3], 0, soffset=NULL), s_waitcnt(lgkmcnt=0), s_endpgm()]
+    self._assert_fault_and_recovery(self._run_kernel(insts))
+
+
+class TestIllegalInstruction(TestAMFaultRecovery):
+  def test_malformed_encoding(self):
+    self._assert_fault_and_recovery(_run_asm(_ILLEGAL_INST_ASM))
+
+
+class TestFlatFaults(TestAMFaultRecovery):
+  def test_flat_load_unmapped(self):
+    insts = [v_mov_b32_e32(v[0], 0xBEEF0000), v_mov_b32_e32(v[1], 0xDEAD),
+             flat_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(vmcnt=0, lgkmcnt=0), s_endpgm()]
+    self._assert_fault_and_recovery(self._run_kernel(insts))
+
+  def test_flat_store_unmapped(self):
+    insts = [v_mov_b32_e32(v[0], 0xBEEF0000), v_mov_b32_e32(v[1], 0xDEAD), v_mov_b32_e32(v[2], 0x12345678),
+             flat_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(vmcnt=0, lgkmcnt=0), s_endpgm()]
+    self._assert_fault_and_recovery(self._run_kernel(insts))
+
+
+class TestAtomicFaults(TestAMFaultRecovery):
+  def test_global_atomic_unmapped(self):
+    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(vmcnt=0), s_endpgm()]
+    self._assert_fault_and_recovery(self._run_kernel(insts))
+
+  def test_flat_atomic_unmapped(self):
+    insts = [v_mov_b32_e32(v[0], 0xBEEF0000), v_mov_b32_e32(v[1], 0xDEAD), v_mov_b32_e32(v[2], 1),
+             flat_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(vmcnt=0, lgkmcnt=0), s_endpgm()]
+    self._assert_fault_and_recovery(self._run_kernel(insts))
+
+
+class TestRecovery(TestAMFaultRecovery):
+  def test_recovery_after_memviol(self):
+    insts = [v_mov_b32_e32(v[0], 0xBEEF0000), v_mov_b32_e32(v[1], 0xDEAD),
+             global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(vmcnt=0), s_endpgm()]
+    self.assertNotEqual(self._run_kernel(insts).returncode, 0)
+    self.assertEqual(_verify_recovery().returncode, 0)
+
+  def test_recovery_after_illegal_inst(self):
+    self.assertNotEqual(_run_asm(_ILLEGAL_INST_ASM).returncode, 0)
+    self.assertEqual(_verify_recovery().returncode, 0)
+
+  def test_multiple_faults_recovery(self):
+    insts = [v_mov_b32_e32(v[0], 0xBEEF0000), v_mov_b32_e32(v[1], 0xDEAD),
+             global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(vmcnt=0), s_endpgm()]
+    for _ in range(3):
+      self.assertNotEqual(self._run_kernel(insts).returncode, 0)
+      self.assertEqual(_verify_recovery().returncode, 0)
+
+if __name__ == "__main__":
+  unittest.main()
@@ -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
@@ -42,7 +42,9 @@ def _memoryview(cls, mem):
      for st,en,rcb,wcb in d.tracked_addresses:
        if st <= addr <= en: return TrackedMemoryView(mem, rcb, wcb)
  return original_memoryview(mem)
-builtins.memoryview = type("memoryview", (), {'__new__': _memoryview}) # type: ignore
+class _MockMemoryviewMeta(type):
+  def __instancecheck__(cls, instance): return isinstance(instance, (original_memoryview, TrackedMemoryView))
+builtins.memoryview = _MockMemoryviewMeta("memoryview", (), {'__new__': _memoryview}) # type: ignore

 def _open(path, flags):
  for d in drivers:
@@ -163,5 +163,30 @@ class TestIndexing(unittest.TestCase):
  # at least the arange is being fused
  def test_llama_embedding_opt(self): self.test_llama_embedding(0, 1_736_704_000)

+  # NOTE: call doesn't work with SPEC=2
+  @unittest.skipIf(Device.DEFAULT not in ("CPU", "AMD"), "atomics only on AMD/CPU")
+  @Context(USE_ATOMICS=1, SPEC=1)
+  def test_llama_8b_embedding_backward(self):
+    from tinygrad.renderer.cstyle import CStyleLanguage
+    if Device.DEFAULT == "CPU" and not isinstance(Device["CPU"].renderer, CStyleLanguage): self.skipTest("CPU needs Clang renderer")
+    vocab_size, embed_size = 1000, 128
+    bs, seqlen = 4, 256
+    idx = Tensor.randint(bs, seqlen, high=vocab_size)
+    emb = nn.Embedding(vocab_size, embed_size)
+    emb.weight = Tensor.ones(vocab_size, embed_size, requires_grad=True)
+    gt = Tensor.zeros(bs, seqlen, embed_size)
+    Tensor.realize(idx, emb.weight, gt)
+    GlobalCounters.reset()
+    loss = (emb(idx)-gt).square().sum()
+    loss.backward()
+    emb.weight.grad.realize()
+    bwd_ops = GlobalCounters.global_ops
+    print(f"embedding bwd: {GlobalCounters.kernel_count} kernels, {bwd_ops:,} ops")
+    self.assertLess(bwd_ops, bs*seqlen*embed_size*20, f"backward ops {bwd_ops:,} should be less than 20 per with atomic scatter-add")
+    # correctness check
+    expected_grad = np.zeros((vocab_size, embed_size), dtype=np.float32)
+    for i in idx.flatten().numpy(): expected_grad[i] += 2
+    np.testing.assert_allclose(emb.weight.grad.numpy(), expected_grad, rtol=1e-5, atol=1e-5)
+
 if __name__ == "__main__":
  unittest.main()
@@ -247,5 +247,39 @@ class TestCustomKernel(unittest.TestCase):
    err = (O_custom - O_ref).square().max()
    self.assertLess(err.item(), 1e-6)

+  def test_multi_after_schedule_order(self):
+    """Test correct scheduling order when custom_kernel has multiple outputs.
+
+    custom_kernel with 4 arguments creates 4 AFTERs from the same kernel.
+    The custom_kernel depends on both A2 and B2, so it must be scheduled after both.
+    E only depends on A2, so E can run before custom_kernel finishes waiting for B2.
+
+    Expected schedule order: [A2, B2, E, custom_addmul, final_sum]
+    The custom_addmul kernel should be at index 3.
+    """
+    from tinygrad.engine.schedule import create_schedule
+    from tinygrad.schedule.rangeify import get_rangeify_map
+
+    A, B = Tensor.empty(4, 4), Tensor.empty(4, 4)
+    A2 = (A + 1).contiguous()                      # kernel 0: depends on A
+    B2 = (B * 2).contiguous()                      # kernel 1: depends on B
+    C, D = Tensor.empty(4, 4), Tensor.empty(4, 4)
+    C, D, _, _ = Tensor.custom_kernel(C, D, A2, B2, fxn=custom_elementwise_addmul_kernel)  # depends on A2 AND B2
+    E = (A2 * 3).contiguous()                      # kernel 2: depends only on A2
+    result = (C + D + E).sum()                     # kernel 3: custom_addmul, then kernel 4: sum
+
+    big_sink = result.uop.sink()
+    tensor_map = get_rangeify_map(big_sink)
+    sched_sink = big_sink.substitute(tensor_map)
+    schedule, _ = create_schedule(sched_sink)
+
+    # Find the custom_addmul kernel position
+    custom_idx = next((i for i, item in enumerate(schedule)
+                       if hasattr(item.ast, "arg") and hasattr(item.ast.arg, "name")
+                       and "custom_addmul" in item.ast.arg.name), None)
+
+    self.assertIsNotNone(custom_idx, "custom_addmul kernel not found in schedule")
+    self.assertEqual(custom_idx, 3, f"custom_addmul should be at index 3, got {custom_idx}")
+
 if __name__ == '__main__':
  unittest.main()
@@ -1,13 +1,14 @@
-import unittest, math
+import contextlib, unittest, math
 import numpy as np
 import torch
 from typing import Any, List
 from tinygrad.device import is_dtype_supported
-from tinygrad.helpers import getenv, DEBUG, CI
+from tinygrad.helpers import getenv, DEBUG, CI, EMULATED_DTYPES
 from tinygrad.dtype import DType, DTYPES_DICT, least_upper_dtype, fp8_to_float, float_to_fp8, _to_np_dtype, _to_torch_dtype, truncate
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.nir import NIRRenderer
-from tinygrad import Device, Tensor, dtypes
+from tinygrad import Context, Device, Tensor, dtypes
+from tinygrad.uop import Ops
 from hypothesis import given, settings, strategies as strat
 from test.helpers import rand_for_dtype
 from test.unit.test_dtype_spec import _assert_eq, core_dtypes, dtype_ints, dtype_floats, FP8E4M3_MAX, FP8E5M2_MAX
@@ -18,9 +19,12 @@ settings.register_profile("my_profile", max_examples=200, deadline=None, derando
 settings.load_profile("my_profile")

 def get_available_cast_dtypes(dtype: DType) -> List[DType]:
-  if not is_dtype_supported(dtype): return []
  # dont cast internal dtypes
-  return [v for k, v in DTYPES_DICT.items() if v != dtype and is_dtype_supported(v) and not k.startswith("_")]
+  dts = [v for k, v in DTYPES_DICT.items() if v != dtype and is_dtype_supported(v) and not k.startswith("_")]
+  if not is_dtype_supported(dtype) or dtypes.long in EMULATED_DTYPES.tolist(dtypes):
+    if dtype in (dtypes.long, dtypes.ulong): return [dt for dt in dts if dt != dtypes.double] # can't bitcast with no 64-bit support
+    else: return []
+  return dts

 def _to_torch_storage_type(dtype:DType):
  if dtype == dtypes.bfloat16: return torch.float32
@@ -334,11 +338,36 @@ class TestInt32DType(TestDType): DTYPE = dtypes.int32
 class TestUint32DType(TestDType): DTYPE = dtypes.uint32

 class TestInt64DType(TestDType): DTYPE = dtypes.int64
+
+@unittest.skipUnless(Ops.SHL in Device[Device.DEFAULT].renderer.code_for_op, "long decomp requires bitshift")
+@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
+class TestEmulatedInt64DType(TestInt64DType):
+  @classmethod
+  def setUpClass(cls):
+    cls.stack = contextlib.ExitStack()
+    cls.stack.enter_context(Context(EMULATED_DTYPES="long"))
+    cls.DATA = rand_for_dtype(cls.DTYPE, 10)
+
+  @classmethod
+  def tearDownClass(cls): cls.stack.close()
+
 class TestUint64DType(TestDType):
  DTYPE = dtypes.uint64
  def test_uint64_load(self):
    assert Tensor(2**64 - 1, dtype=dtypes.uint64).numpy() == 2**64 - 1

+@unittest.skipUnless(Ops.SHL in Device[Device.DEFAULT].renderer.code_for_op, "long decomp requires bitshift")
+@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
+class TestEmulatedUInt64DType(TestUint64DType):
+  @classmethod
+  def setUpClass(cls):
+    cls.stack = contextlib.ExitStack()
+    cls.stack.enter_context(Context(EMULATED_DTYPES="long"))
+    cls.DATA = rand_for_dtype(cls.DTYPE, 10)
+
+  @classmethod
+  def tearDownClass(cls): cls.stack.close()
+
 class TestBoolDType(TestDType): DTYPE = dtypes.bool

 class TestBFloat16Type(TestDType): DTYPE = dtypes.bfloat16
@@ -1,5 +1,5 @@
 import unittest, operator, math
-from tinygrad import Tensor, dtypes, Device
+from tinygrad import Context, Tensor, dtypes, Device
 from tinygrad.dtype import DType, truncate
 from tinygrad.helpers import CI, getenv
 from tinygrad.tensor import _to_np_dtype
@@ -7,6 +7,7 @@ from tinygrad.device import is_dtype_supported
 from tinygrad.runtime.ops_python import from_storage_scalar
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.nir import NIRRenderer
+from tinygrad.uop import Ops
 import numpy as np
 import pytest
 from hypothesis import assume, given, strategies as strat, settings
@@ -169,6 +170,12 @@ class TestDTypeALU(unittest.TestCase):
  @given(ht.uint64, ht.uint64, strat.sampled_from(integer_binary_operations))
  def test_uint64(self, a, b, op): universal_test(a, b, dtypes.uint64, op)

+  @unittest.skipUnless(Ops.SHL in Device[Device.DEFAULT].renderer.code_for_op, "long decomp requires bitshift")
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
+  @given(ht.uint64, ht.uint64, strat.sampled_from(integer_binary_operations))
+  @Context(EMULATED_DTYPES="long")
+  def test_emulated_uint64(self, a, b, op): universal_test(a, b, dtypes.uint64, op)
+
  @given(ht.int8, ht.int8, strat.sampled_from(integer_binary_operations))
  def test_int8(self, a, b, op): universal_test(a, b, dtypes.int8, op)

@@ -182,6 +189,12 @@ class TestDTypeALU(unittest.TestCase):
  @given(ht.int64, ht.int64, strat.sampled_from(integer_binary_operations))
  def test_int64(self, a, b, op): universal_test(a, b, dtypes.int64, op)

+  @unittest.skipUnless(Ops.SHL in Device[Device.DEFAULT].renderer.code_for_op, "long decomp requires bitshift")
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
+  @given(ht.int64, ht.int64, strat.sampled_from(integer_binary_operations))
+  @Context(EMULATED_DTYPES="long")
+  def test_emulated_int64(self, a, b, op): universal_test(a, b, dtypes.int64, op)
+
  @given(ht.uint8, strat.sampled_from(integer_unary_operations))
  def test_uint8_unary(self, a, op): universal_test_unary(a, dtypes.uint8, op)

@@ -197,6 +210,12 @@ class TestDTypeALU(unittest.TestCase):
  @given(ht.uint64, strat.sampled_from(integer_unary_operations))
  def test_uint64_unary(self, a, op): universal_test_unary(a, dtypes.uint64, op)

+  @unittest.skipUnless(Ops.SHL in Device[Device.DEFAULT].renderer.code_for_op, "long decomp requires bitshift")
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
+  @given(ht.uint64, strat.sampled_from(integer_unary_operations))
+  @Context(EMULATED_DTYPES="long")
+  def test_emulated_uint64_unary(self, a, op): universal_test_unary(a, dtypes.uint64, op)
+
  @given(ht.int8, strat.sampled_from(integer_unary_operations))
  def test_int8_unary(self, a, op): universal_test_unary(a, dtypes.int8, op)

@@ -210,6 +229,12 @@ class TestDTypeALU(unittest.TestCase):
  @given(ht.int64, strat.sampled_from(integer_unary_operations))
  def test_int64_unary(self, a, op): universal_test_unary(a, dtypes.int64, op)

+  @unittest.skipUnless(Ops.SHL in Device[Device.DEFAULT].renderer.code_for_op, "long decomp requires bitshift")
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
+  @given(ht.int64, strat.sampled_from(integer_unary_operations))
+  @Context(EMULATED_DTYPES="long")
+  def test_emulated_int64_unary(self, a, op): universal_test_unary(a, dtypes.int64, op)
+
  @given(ht.bool, ht.bool, strat.sampled_from(((operator.add, operator.add), (operator.mul, operator.mul))))
  def test_bool(self, a, b, op): universal_test(a, b, dtypes.bool, op)

@@ -26,7 +26,7 @@ import unittest
 import numpy as np
 import torch
 from tinygrad import Tensor, dtypes, nn
-from tinygrad.device import Device, is_dtype_supported
+from tinygrad.device import Device
 from tinygrad.helpers import getenv
 from tinygrad.renderer.nir import NIRRenderer

@@ -195,8 +195,10 @@ class TestAssignIssues(unittest.TestCase):
    t.shrink(((1, 3), (1, 3))).assign(Tensor.ones(2, 2))
    np.testing.assert_allclose(t.numpy(), torch_tensor.numpy())

+  @unittest.expectedFailure
  def test_assign_broadcast(self):
    # broadcasting during assign should behave like PyTorch
+    # NOTE: we don't want implicit dtype casting (int64 -> float32 loses precision), so this fails
    torch_tensor = torch.zeros(3, 5)
    torch_tensor[:] = torch.arange(5)
    t = Tensor.zeros(3, 5)
@@ -207,8 +209,7 @@ class TestUOpValidationIssue(unittest.TestCase):
  # these fail with UOp verification error.
  # we want more of these with diverse errors!

-  @unittest.skipIf((not is_dtype_supported(dtypes.long)) or MOCKGPU or isinstance(Device[Device.DEFAULT].renderer, NIRRenderer),
-                   "hangs gpuocelot, NIR cannot render")
+  @unittest.skipIf(MOCKGPU or isinstance(Device[Device.DEFAULT].renderer, NIRRenderer), "hangs gpuocelot, NIR cannot render")
  def test_tensor_index_overflow(self):
    val = Tensor([1])
    big = val.expand(2**31 + 3)
@@ -256,6 +256,18 @@ class TestMultiTensor(unittest.TestCase):
      a,b = _test_allreduce(Tensor.rand(256, 256))
      np.testing.assert_almost_equal(a.numpy(), b.numpy(), decimal=5)

+  def test_multiple_to_single_device_naive(self):
+    with Context(RING=0):
+      t = Tensor.arange(32).shard(devices_4, 0).to(Device.DEFAULT).realize()
+    self.assertEqual(t.device, Device.DEFAULT)
+    np.testing.assert_equal(t.numpy(), np.arange(32))
+
+  def test_multiple_to_single_device_ring(self):
+    with Context(RING=2):
+      t = Tensor.arange(32).shard(devices_4, 0).to(Device.DEFAULT).realize()
+    self.assertEqual(t.device, Device.DEFAULT)
+    np.testing.assert_equal(t.numpy(), np.arange(32))
+
  def test_allreduce_all2all(self):
    with Context(ALL2ALL=2):
      a,b = _test_allreduce(Tensor.rand(256, 256))
@@ -409,6 +421,28 @@ class TestMultiTensor(unittest.TestCase):

    np.testing.assert_allclose(z.numpy(), z_shard.numpy(), atol=1e-6, rtol=1e-6)

+  def test_embedding_backward(self, shard_weight_axis=None):
+    B, T, embed_size, vocab_size = 4, 10, 20, 28
+
+    layer = nn.Embedding(vocab_size, embed_size)
+    layer.weight.requires_grad = True
+    x = Tensor(np.random.randint(0, vocab_size, (B, T), dtype=np.int32))
+    z = layer(x)
+    z.sum().backward()
+    grad = layer.weight.grad.numpy()
+
+    layer_sharded = nn.Embedding(vocab_size, embed_size)
+    layer_sharded.weight.replace(layer.weight.shard(devices_2, axis=shard_weight_axis)).realize()
+    layer_sharded.weight.requires_grad = True
+    x_sharded = x.shard(devices_2, axis=None)
+    z_shard = layer_sharded(x_sharded)
+    z_shard.sum().backward()
+    grad_shard = layer_sharded.weight.grad.numpy()
+
+    np.testing.assert_allclose(grad, grad_shard, atol=1e-6, rtol=1e-6)
+
+  def test_embedding_backward_shard_weight(self): self.test_embedding_backward(shard_weight_axis=1)
+
  def test_rmsnorm(self):
    B, T, embed_size = 4, 10, 20

@@ -1251,6 +1285,21 @@ class TestMultiRamUsage(unittest.TestCase):
    _ = Tensor.zeros(self.N, self.N).contiguous().shard(devices_2, axis=0).contiguous().realize()
    self.assertUsed(self.N*self.N*4) # sharding should not increase total ram usage

+  def _test_matmul_half(self, dev_count:int):
+    N = 32
+    total_mem = {}
+    devs = tuple(f"NULL:{i}" for i in range(dev_count))
+    for dtype in {dtypes.float, dtypes.half}:
+      GlobalCounters.reset()
+      a = Tensor.empty((N, N), dtype=dtype, device=devs[0]).shard(devs, axis=0)
+      b = Tensor.empty((N, N), dtype=dtype, device=devs[0]).shard(devs, axis=None)
+      (a @ b).realize()
+      total_mem[dtype] = GlobalCounters.global_mem
+    self.assertEqual(total_mem[dtypes.half], total_mem[dtypes.float] // 2)
+
+  def test_matmul_half(self): self._test_matmul_half(dev_count=2)
+  def test_matmul_half_alt(self): self._test_matmul_half(dev_count=4)
+
@unittest.skipIf(not_support_multi_device(), "need multi")
 class TestMultiFromUnrenderable(unittest.TestCase):
  @needs_second_gpu
@@ -698,6 +698,7 @@ class TestOps(unittest.TestCase):
        tiny_out = get_tiny_gradient(x, c)
        torch_out = get_torch_gradient(x, c)
        if math.isnan(tiny_out):
+          if Device.DEFAULT == "WEBGPU": continue # TODO: WEBGPU issue with nan
          assert math.isnan(torch_out)
        else:
          self.assertAlmostEqual(tiny_out, torch_out, msg=f"{x}, {c}")
@@ -949,6 +950,8 @@ class TestOps(unittest.TestCase):
    helper_test_op([(45,65), (45,1)], torch.copysign, Tensor.copysign)
    helper_test_op([(45,1), (1,65)], torch.copysign, Tensor.copysign)
    helper_test_op([(), ()], torch.copysign, Tensor.copysign)
+
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "fails locally")
  def test_copysign_exact(self):
    # NOTE: -nan (negative nan) is not tested because we can't detect its sign bit without bitcast
    v = [-1., -0., 0., 1., math.inf, -math.inf, math.nan]
@@ -2184,6 +2184,14 @@ class TestBufferUOp(unittest.TestCase):
    run_schedule(check_schedule(a, 0))
    self.assertIsNone(a.uop.base.realized)

+  def test_unused_var_not_in_var_vals(self):
+    # unused variable should not appear in var_vals even when there's other work
+    a = Tensor(UOp.variable("unused", 0, 10).bind(1))
+    b = Tensor.empty(3) + 1
+    _, var_vals = Tensor.schedule_with_vars(a, b)
+    self.assertEqual(var_vals, {})
+    self.assertIsNone(a.uop.base.realized)
+
  def test_view_does_not_realize(self):
    a = Tensor.randn(1, 4).expand(4, 4)
    a.realize()
@@ -169,7 +169,6 @@ class TestSymbolicOps(unittest.TestCase):
      vi = Variable("i", 1, 10).bind(i)
      a = Tensor.rand(7, 11)
      symbolic = a[3:5, vi:vi+2]
-      print(symbolic.shape)
      symbolic = symbolic.numpy()
      expected = a[3:5, i:i+2].numpy()
      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
@@ -2,8 +2,7 @@ import numpy as np
 import torch
 import unittest, copy, mmap, random, math, array
 from tinygrad import Tensor, Device, dtypes, nn
-from tinygrad.tensor import _METADATA
-from tinygrad.helpers import Context, getenv, temp, mv_address
+from tinygrad.helpers import getenv, temp, mv_address
 from extra.gradcheck import numerical_jacobian, jacobian, gradcheck
 from hypothesis import given, settings, strategies as strat
 from tinygrad.device import is_dtype_supported
@@ -796,92 +795,6 @@ class TestInferenceMode(unittest.TestCase):
      assert W.grad is None
    f(x, m, W)

-class TestTensorMetadata(unittest.TestCase):
-  def setUp(self) -> None: _METADATA.set(None)
-
-  # NOOPs are not included in kernel metadata
-  @unittest.skip("why would this be true?")
-  def test_exclude_noop_metadata(self):
-    a = Tensor.rand(4, 4)*1
-    self.assertEqual(a.uop.metadata[0].name, "__mul__")
-    k = a.schedule()[-1]
-    self.assertEqual([m.name for m in k.metadata], ["rand"])
-
-  # we exclude const from kernel metadata because tensor methods can share the same CONST UOp
-  @unittest.skip("TODO: flaky")
-  def test_exclude_const_metadata(self):
-    a = Tensor.arange(4)
-    b = Tensor.full((4,), -1, dtype=dtypes.int).contiguous()
-    sched = Tensor.schedule(a, b)
-    self.assertEqual([m.name for m in sched[0].metadata], ["arange"])
-    self.assertEqual([m.name for m in sched[1].metadata], ["contiguous"])
-
-  def test_matmul(self):
-    x = Tensor.rand(3, requires_grad=True)
-    W = Tensor.rand(3, 3, requires_grad=True)
-    out = x.matmul(W)
-    self.assertEqual(out.uop.metadata[0].name, "matmul")
-    si = out.schedule()[-1]
-    self.assertEqual(len(si.metadata), 1)
-    self.assertEqual(si.metadata[0].name, "matmul")
-
-  def test_relu(self):
-    x = Tensor.rand(3, requires_grad=True)
-    out = x.relu()
-    self.assertEqual(out.uop.metadata[0].name, "relu")
-    si = out.schedule()[-1]
-    self.assertEqual(len(si.metadata), 1)
-    self.assertEqual(si.metadata[0].name, "relu")
-
-  @unittest.skip("this no longer works")
-  def test_assign(self):
-    x = Tensor.empty(10, 10).realize()
-    x.assign(Tensor.ones(10, 10).contiguous())
-    si = x.schedule()[-1]
-    self.assertEqual(len(si.metadata), 1)
-    self.assertEqual(si.metadata[0].name, "assign")
-
-  def test_complex(self):
-    x = Tensor.rand(3, requires_grad=True)
-    y = Tensor.rand(3, requires_grad=True)
-    out = x.relu() * y.sigmoid()
-    self.assertEqual(out.uop.metadata[0].name, "__mul__")
-    self.assertEqual(out.uop.src[0].metadata[0].name, "relu")
-    self.assertEqual(out.uop.src[1].metadata[0].name, "sigmoid")
-    si = out.schedule()[-1]
-    self.assertEqual(len(si.metadata), 3)
-    self.assertEqual(set(m.name for m in si.metadata), {"relu", "sigmoid", "__mul__"})
-
-  @unittest.skip("metadata is no longer promised to be exact with schedulecache")
-  def test_complex_backward(self):
-    x = Tensor.rand(3, requires_grad=True).realize()
-    y = Tensor.rand(3, requires_grad=True).realize()
-    out = (x.relu() * y.sigmoid()).sum()
-    self.assertEqual(out.uop.metadata[0].name, "sum")
-    out.backward()
-    self.assertEqual(x.grad.uop.metadata[0].name, "relu")
-    self.assertTrue(x.grad.uop.metadata[0].backward)
-    self.assertEqual(y.grad.uop.metadata[0].name, "sigmoid")
-    self.assertTrue(y.grad.uop.metadata[0].backward)
-    si = Tensor.schedule(out, x.grad, y.grad)[-1]
-    #self.assertEqual(len(si.metadata), 3, f"failed with {si.metadata}")
-    # skip numpy, this is schedule cache
-    self.assertSetEqual(set(m.name for m in si.metadata if m.name != "numpy"), {"sigmoid", "relu"})
-    #bw = [m for m in si.metadata if m.backward]
-    #self.assertEqual(len(bw), 1)
-    #self.assertEqual(bw[0].name, "sigmoid")
-
-  @unittest.skip("metadata is no longer promised to be exact with schedulecache")
-  def test_tracemeta_0(self):
-    with Context(TRACEMETA=0):
-      x = Tensor.rand(3, requires_grad=True)
-      y = Tensor.rand(3, requires_grad=True)
-      out = (x.relu() * y.sigmoid()).sum()
-      self.assertIsNone(out.uop.metadata)
-      self.assertIsNone(out.uop.src[0].metadata)
-      si = out.schedule()[-1]
-      self.assertEqual(si.metadata, ())
-
 class TestIdxUpcast(unittest.TestCase):
  def _find_op(self, ast: UOp, op: Ops):
    if ast.op is op: return ast
@@ -2,82 +2,40 @@
 # allow define from star imports

 import unittest
-import textwrap, functools
+import functools

 from tinygrad import Device, Tensor
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
-from tinygrad.helpers import getenv
 from tinygrad.device import Compiler
+from tinygrad.runtime.support.compiler_amd import HIPCompiler
 from tinygrad.viz.serve import amdgpu_cfg

 from extra.assembly.amd.autogen.rdna3.ins import *
-from extra.assembly.amd.dsl import Inst
+from extra.assembly.amd.dsl import s

-template = """.text
-.globl fn_name
-.p2align 8
-.type fn_name,@function
-fn_name:
-  INSTRUCTION
-
-.rodata
-.p2align 6
-.amdhsa_kernel fn_name
-  .amdhsa_kernarg_size 8
-  .amdhsa_user_sgpr_kernarg_segment_ptr 1
-  .amdhsa_next_free_vgpr .amdgcn.next_free_vgpr
-  .amdhsa_next_free_sgpr .amdgcn.next_free_sgpr
-  .amdhsa_wavefront_size32 1
-.end_amdhsa_kernel
-
-.amdgpu_metadata
---
-amdhsa.version:
-  - 1
-  - 0
-amdhsa.kernels:
-  - .name: fn_name
-    .symbol: fn_name.kd
-    .group_segment_fixed_size: 0
-    .private_segment_fixed_size: 0
-    .wavefront_size: 32
-    .sgpr_count: 8
-    .vgpr_count: 8
-    .max_flat_workgroup_size: 1024
-    .kernarg_segment_align: 8
-    .kernarg_segment_size: 8
-    .args:
-      - .address_space:  global
-        .name:           a
-        .offset:         0
-        .size:           8
-        .type_name:      'float*'
-        .value_kind:     global_buffer
-...
-.end_amdgpu_metadata
-"""
+# TODO: this belongs to the dsl infrastructure
+from extra.gemm.amd_asm_matmul import Kernel

 # TODO: shouldn't need compiler once we can output ELF
 # outputs a text disassembly for humans and a machine readable binary
-def assemble(name:str, insts:list[str|Inst], compiler:Compiler) -> tuple[str, bytes]:
-  asm = "\n".join([inst if isinstance(inst, str) else inst.disasm() for inst in insts])
-  src = template.replace("fn_name", name).replace("INSTRUCTION", textwrap.dedent(asm))
+def assemble(name:str, k:Kernel, compiler:Compiler) -> tuple[str, bytes]:
+  src = k.to_asm()
  return (src, compiler.compile(src))

-def asm_kernel(out:UOp, insts:list[str|Inst], name:str, device:str, compiler:Compiler, n_threads:int=1, n_workgroups:int=1) -> UOp:
+def asm_kernel(out:UOp, k:Kernel, name:str, device:str, compiler:Compiler, n_threads:int=1, n_workgroups:int=1) -> UOp:
  lidx = UOp.special(n_threads, "lidx0")
  gidx = UOp.special(n_workgroups, "gidx0")
  sink = UOp.sink(out, lidx, gidx, arg=KernelInfo(name=name))
-  src, lib = assemble(name, insts, compiler)
+  src, lib = assemble(name, k, compiler)
  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=device), UOp(Ops.LINEAR, src=(*sink.src, sink)),
                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))

-def run_asm(name:str, insts:list) -> None:
-  fxn = functools.partial(asm_kernel, insts=insts, name=name, device=Device.DEFAULT, compiler=Device[Device.DEFAULT].compiler)
+def run_asm(name:str, k:Kernel) -> None:
+  fxn = functools.partial(asm_kernel, k=k, name=name, device=Device.DEFAULT, compiler=HIPCompiler(Device[Device.DEFAULT].renderer.arch))
  out = Tensor.custom_kernel(Tensor.empty(1), fxn=fxn)[0]
  out.realize()

-@unittest.skipUnless(Device.DEFAULT == "AMD" and not getenv("AMD_LLVM"), "only on AMD with comgr")
+@unittest.skipUnless(Device.DEFAULT == "AMD", "only on AMD")
 class TestCfg(unittest.TestCase):
  def setUp(self):
    arch = Device["AMD"].arch
@@ -85,32 +43,32 @@ class TestCfg(unittest.TestCase):
      self.skipTest(f"tests written for RDNA, got arch {arch}")

  def test_simple(self):
-    run_asm("simple", [
-      "entry:",
-        "s_branch bb1",
-      "bb1:",
-        s_endpgm(),
-        s_code_end(),
-    ])
+    k = Kernel(arch=Device["AMD"].arch)
+    k.label("entry")
+    k.emit(s_branch(), target="bb1")
+    k.label("bb1")
+    k.emit(s_endpgm())
+    k.emit(s_code_end())
+    run_asm("simple", k)

  def test_diamond(self):
-    run_asm("diamond", insts:=[
-      "entry:",
-        s_mov_b32(s[0], 0),
-        s_mov_b32(s[1], 0),
-        s_cmp_eq_u64(s[0:1], 0),
-        "s_cbranch_scc1 if",
-        "s_branch else",
-      "if:",
-        s_nop(1),
-        "s_branch end",
-      "else:",
-        s_nop(0),
-      "end:",
-        s_endpgm(),
-        s_code_end(),
-    ])
-    _, lib = assemble("diamond", insts, Device[Device.DEFAULT].compiler)
+    k = Kernel(arch=Device["AMD"].arch)
+    k.label("entry")
+    k.emit(s_mov_b32(s[0], 0))
+    k.emit(s_mov_b32(s[1], 0))
+    k.emit(s_cmp_eq_u64(s[0:1], 0))
+    k.emit(s_cbranch_scc1(), target="if")
+    k.emit(s_branch(), target="else")
+    k.label("if")
+    k.emit(s_nop(1))
+    k.emit(s_branch(), target="end")
+    k.label("else")
+    k.emit(s_nop(0))
+    k.label("end")
+    k.emit(s_endpgm())
+    k.emit(s_code_end())
+    run_asm("diamond", k)
+    _, lib = assemble("diamond", k, HIPCompiler(Device[Device.DEFAULT].arch))
    cfg = amdgpu_cfg(lib, Device[Device.DEFAULT].device_props()["gfx_target_version"])["data"]
    self.assertEqual(len(cfg["blocks"]), 5)
    edge_count = sum(len(v) for v in cfg["paths"].values())
@@ -124,119 +82,138 @@ class TestCfg(unittest.TestCase):
    self.assertEqual(insts, ['s_mov_b32', 's_cmp_eq_u64'])

  def test_loop(self):
-    run_asm("simple_loop", [
-      "entry:",
-        s_mov_b32(s[1], 4),
-      "loop:",
-        s_add_u32(s[1], s[1], -1),
-        s_cmp_eq_i32(s[1], 0),
-        "s_cbranch_scc0 loop",
-        s_endpgm(),
-        s_code_end(),
-    ])
+    k = Kernel(arch=Device["AMD"].arch)
+    k.label("entry")
+    k.emit(s_mov_b32(s[1], 4))
+    k.label("loop")
+    k.emit(s_add_u32(s[1], s[1], -1))
+    k.emit(s_cmp_eq_i32(s[1], 0))
+    k.emit(s_cbranch_scc0(), target="loop")
+    k.emit(s_endpgm())
+    k.emit(s_code_end())
+    run_asm("simple_loop", k)

  def test_loop_branch(self):
-    run_asm("loop_if", [
-      "entry:",
-        s_mov_b32(s[1], 4),
-      "loop:",
-        s_add_u32(s[1], s[1], -1),
-        s_cmp_eq_i32(s[1], 2),
-        "s_cbranch_scc1 cond",
-        "s_branch cont",
-      "cond:",
-        s_add_u32(s[1], s[1], -2),
-      "cont:",
-        s_cmp_eq_i32(s[1], 0),
-        "s_cbranch_scc0 loop",
-        s_endpgm(),
-        s_code_end(),
-    ])
+    k = Kernel(arch=Device["AMD"].arch)
+    k.label("entry")
+    k.emit(s_mov_b32(s[1], 4))
+    k.label("loop")
+    k.emit(s_add_u32(s[1], s[1], -1))
+    k.emit(s_cmp_eq_i32(s[1], 2))
+    k.emit(s_cbranch_scc1(), target="cond")
+    k.emit(s_branch(), target="cont")
+    k.label("cond")
+    k.emit(s_add_u32(s[1], s[1], -2))
+    k.label("cont")
+    k.emit(s_cmp_eq_i32(s[1], 0))
+    k.emit(s_cbranch_scc0(), target="loop")
+    k.emit(s_endpgm())
+    k.emit(s_code_end())
+    run_asm("loop_if", k)

  def test_loop_break(self):
-    run_asm("loop_break", [
-      "entry:",
-        s_mov_b32(s[1], 8),
-      "loop:",
-        s_add_u32(s[1], s[1], -1),
-        s_cmp_eq_i32(s[1], 5),
-        "s_cbranch_scc1 break",
-        s_cmp_eq_i32(s[1], 0),
-        "s_cbranch_scc0 loop",
-      "break:",
-        s_endpgm(),
-        s_code_end(),
-    ])
+    k = Kernel(arch=Device["AMD"].arch)
+    k.label("entry")
+    k.emit(s_mov_b32(s[1], 8))
+    k.label("loop")
+    k.emit(s_add_u32(s[1], s[1], -1))
+    k.emit(s_cmp_eq_i32(s[1], 5))
+    k.emit(s_cbranch_scc1(), target="break")
+    k.emit(s_cmp_eq_i32(s[1], 0))
+    k.emit(s_cbranch_scc0(), target="loop")
+    k.label("break")
+    k.emit(s_endpgm())
+    k.emit(s_code_end())
+    run_asm("loop_break", k)

  def test_switch(self):
-    run_asm("switch_case", [
-      "entry:",
-        s_cmp_eq_i32(s[0], 0),
-        "s_cbranch_scc1 case0",
-        s_cmp_eq_i32(s[0], 1),
-        "s_cbranch_scc1 case1",
-        "s_branch case2",
-      "case0:",
-        s_nop(0),
-        "s_branch join",
-      "case1:",
-        s_nop(1),
-        "s_branch join",
-      "case2:",
-        s_nop(2),
-        "s_branch join",
-      "join:",
-        s_endpgm(),
-        s_code_end(),
-    ])
+    k = Kernel(arch=Device["AMD"].arch)
+    k.label("entry")
+    k.emit(s_cmp_eq_i32(s[0], 0))
+    k.emit(s_cbranch_scc1(), target="case0")
+    k.emit(s_cmp_eq_i32(s[0], 1))
+    k.emit(s_cbranch_scc1(), target="case1")
+    k.emit(s_branch(), target="case2")
+    k.label("case0")
+    k.emit(s_nop(0))
+    k.emit(s_branch(), target="join")
+    k.label("case1")
+    k.emit(s_nop(1))
+    k.emit(s_branch(), target="join")
+    k.label("case2")
+    k.emit(s_nop(2))
+    k.emit(s_branch(), target="join")
+    k.label("join")
+    k.emit(s_endpgm())
+    k.emit(s_code_end())
+    run_asm("switch_case", k)

  def test_ping_pong(self):
-    run_asm("ping_pong", [
-      "entry:",
-        s_cmp_eq_i32(s[0], 0),
-        "s_cbranch_scc1 ping",
-        "s_branch pong",
-      "ping:",
-        s_cmp_eq_i32(s[1], 0),
-        "s_cbranch_scc1 pong",
-        "s_branch end",
-      "pong:",
-        s_cmp_eq_i32(s[2], 0),
-        "s_cbranch_scc1 ping",
-      "end:",
-        s_endpgm(),
-        s_code_end(),
-    ])
+    k = Kernel(arch=Device["AMD"].arch)
+    k.label("entry")
+    k.emit(s_cmp_eq_i32(s[0], 0))
+    k.emit(s_cbranch_scc1(), target="ping")
+    k.emit(s_branch(), target="pong")
+    k.label("ping")
+    k.emit(s_cmp_eq_i32(s[1], 0))
+    k.emit(s_cbranch_scc1(), target="pong")
+    k.emit(s_branch(), target="end")
+    k.label("pong")
+    k.emit(s_cmp_eq_i32(s[2], 0))
+    k.emit(s_cbranch_scc1(), target="ping")
+    k.label("end")
+    k.emit(s_endpgm())
+    k.emit(s_code_end())
+    run_asm("ping_pong", k)

  def test_colored_blocks(self):
    N = 10
-    asm = ["entry:", "s_branch init0"]
+    k = Kernel(arch=Device["AMD"].arch)
+    k.label("entry")
+    k.emit(s_branch(), target="init0")
    for i in range(N):
-      asm += [f"init{i}:", s_mov_b32(s[1], i + 1), f"s_branch {(loop:=f'loop{i}')}"]
-      asm += [
-        f"{loop}:",
-          s_nop(i & 7),
-          s_add_u32(s[1], s[1], -1),
-          s_cmp_eq_i32(s[1], 0),
-          f"s_cbranch_scc0 {loop}",
-          f"s_branch {'init' + str(i+1) if i + 1 < N else 'end'}",
-      ]
-    asm += ["end:", s_endpgm(), s_code_end()]
-    run_asm("test_colored_blocks", asm)
+      loop = f"loop{i}"
+      k.label(f"init{i}")
+      k.emit(s_mov_b32(s[1], i + 1))
+      k.emit(s_branch(), target=loop)
+      k.label(loop)
+      k.emit(s_nop(i & 7))
+      k.emit(s_add_u32(s[1], s[1], -1))
+      k.emit(s_cmp_eq_i32(s[1], 0))
+      k.emit(s_cbranch_scc0(), target=loop)
+      k.emit(s_branch(), target=f"init{i+1}" if i + 1 < N else "end")
+    k.label("end")
+    k.emit(s_endpgm())
+    k.emit(s_code_end())
+    run_asm("test_colored_blocks", k)

  def test_jump_back_to_end(self):
-    run_asm("jump_back_to_end", [
-      "entry:",
-        s_mov_b32(s[1], 2),
-        "s_cbranch_execz loop",
-      "end:",
-        s_endpgm(),
-      "loop:",
-        s_add_u32(s[1], s[1], -1),
-        s_cmp_eq_i32(s[1], 0),
-        "s_branch end",
-        s_code_end(),
-    ])
+    k = Kernel(arch=Device["AMD"].arch)
+    k.label("entry")
+    k.emit(s_mov_b32(s[1], 2))
+    k.emit(s_cbranch_execz(), target="loop")
+    k.label("end")
+    k.emit(s_endpgm())
+    k.label("loop")
+    k.emit(s_add_u32(s[1], s[1], -1))
+    k.emit(s_cmp_eq_i32(s[1], 0))
+    k.emit(s_branch(), target="end")
+    k.emit(s_code_end())
+    run_asm("jump_back_to_end", k)
+
+  def test_hit_count(self):
+    k = Kernel(arch=Device["AMD"].arch)
+    k.label("entry")
+    k.emit(s_mov_b32(s[1], 1))
+    k.emit(s_branch(), target="alt")
+    k.label("continue")
+    k.emit(s_mov_b32(s[2], 2))
+    k.emit(s_add_u32(s[1], s[1], s[2]))
+    k.label("alt")
+    k.emit(s_add_u32(s[1], s[1], -1))
+    k.emit(s_endpgm())
+    k.emit(s_code_end())
+    run_asm("test_hit_count", k)

 if __name__ == "__main__":
  unittest.main()
@@ -456,6 +456,32 @@ class TestAssign(unittest.TestCase):
    assign.realize()
    np.testing.assert_allclose(a.numpy(), [2., 2., 2., 2., 1., 1., 1., 1.])

+  def test_setitem_list(self):
+    a = Tensor.zeros(8).contiguous().realize()
+    a[2:5] = [1, 2, 3]
+    np.testing.assert_allclose(a.numpy(), [0., 0., 1., 2., 3., 0., 0., 0.])
+
+  def test_assign_bitcast(self):
+    # assign to a bitcast view should modify the underlying buffer
+    a = Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32).realize()
+    # IEEE 754: 1.0f = 0x3f800000, 2.0f = 0x40000000, 3.0f = 0x40400000, 4.0f = 0x40800000
+    a.bitcast(dtypes.uint32).assign(Tensor([0x40800000, 0x40400000, 0x40000000, 0x3f800000], dtype=dtypes.uint32)).realize()
+    np.testing.assert_allclose(a.numpy(), [4.0, 3.0, 2.0, 1.0])
+    # double bitcast
+    b = Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32).realize()
+    b.bitcast(dtypes.uint32).bitcast(dtypes.int32).assign(Tensor([0x40800000, 0x40400000, 0x40000000, 0x3f800000], dtype=dtypes.int32)).realize()
+    np.testing.assert_allclose(b.numpy(), [4.0, 3.0, 2.0, 1.0])
+    # shrink then bitcast
+    c = Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32).realize()
+    c[0:2].bitcast(dtypes.uint32).assign(Tensor([0x40800000, 0x40400000], dtype=dtypes.uint32)).realize()
+    np.testing.assert_allclose(c.numpy(), [4.0, 3.0, 3.0, 4.0])
+
+  def test_assign_bitcast_different_size(self):
+    # different-size bitcast creates a new tensor, not a view, so assign doesn't modify the original
+    a = Tensor([0]*8, dtype=dtypes.uint8).realize()
+    a.bitcast(dtypes.int64).assign(Tensor([12345], dtype=dtypes.int64)).realize()
+    np.testing.assert_equal(a.numpy(), [0]*8)
+
  @unittest.skip("don't use output buffer, and mismatch dtype no longer supported")
  def test_cast_assignment(self):
    a = Tensor(np.arange(N*N, dtype=np.float32)).reshape(N,N)
@@ -467,6 +493,38 @@ class TestAssign(unittest.TestCase):
    assert oba1 is None and oba2 is None
    np.testing.assert_allclose(a.numpy(), np.arange(N*N,dtype=np.int32).reshape((N,N)))

+  def test_assign_dtype_mismatch(self):
+    # assign should not implicitly cast dtypes - this can lose precision
+    a = Tensor.zeros(4, dtype=dtypes.float32).contiguous().realize()
+    b = Tensor([1, 2, 3, 4], dtype=dtypes.int32)
+    with self.assertRaisesRegex(RuntimeError, "assign dtype mismatch"):
+      a.assign(b)
+
+  def test_assign_dtype_mismatch_int64_to_float32(self):
+    # int64 -> float32 loses precision for large values, should not be implicit
+    a = Tensor.zeros(1, dtype=dtypes.float32).contiguous().realize()
+    b = Tensor([16777217], dtype=dtypes.int64)  # 2^24 + 1, not exactly representable in float32
+    with self.assertRaisesRegex(RuntimeError, "assign dtype mismatch"):
+      a.assign(b)
+
+  def test_assign_shape_broadcast(self):
+    # shape broadcasting should work when dtypes match
+    a = Tensor.zeros(3, 5, dtype=dtypes.float32).contiguous().realize()
+    b = Tensor([1., 2., 3., 4., 5.], dtype=dtypes.float32)
+    a.assign(b)
+    a.realize()
+    expected = np.array([[1., 2., 3., 4., 5.]] * 3)
+    np.testing.assert_allclose(a.numpy(), expected)
+
+  def test_assign_shape_broadcast_2d(self):
+    # broadcast (1, 5) to (3, 5)
+    a = Tensor.zeros(3, 5, dtype=dtypes.float32).contiguous().realize()
+    b = Tensor([[1., 2., 3., 4., 5.]], dtype=dtypes.float32)
+    a.assign(b)
+    a.realize()
+    expected = np.array([[1., 2., 3., 4., 5.]] * 3)
+    np.testing.assert_allclose(a.numpy(), expected)
+
  def test_disk_assignment(self):
    a = Tensor.empty(5, device=f"disk:{temp('disk_assignment')}").assign(Tensor.ones(5)).numpy()
    np.testing.assert_equal(a, np.ones(5))
@@ -561,12 +619,12 @@ class TestAssignOrdering(unittest.TestCase):
  def test_slice_write_then_full_read(self):
    """Write to slice, then read full buffer."""
    # without .realize(): orphan slice assign not triggered by .numpy()
-    buf = Tensor.zeros(4).contiguous().realize()
+    buf = Tensor.zeros(4, dtype=dtypes.int32).contiguous().realize()
    buf[1:3].assign(Tensor([5, 6]))
    np.testing.assert_equal(buf.numpy(), [0, 0, 0, 0])  # TODO: wrong! should be [0, 5, 6, 0]

    # with .realize(): assign executes
-    buf = Tensor.zeros(4).contiguous().realize()
+    buf = Tensor.zeros(4, dtype=dtypes.int32).contiguous().realize()
    buf[1:3].assign(Tensor([5, 6])).realize()
    np.testing.assert_equal(buf.numpy(), [0, 5, 6, 0])

@@ -648,7 +706,7 @@ class TestAssignOrdering(unittest.TestCase):

  def test_three_buffer_chain(self):
    """Chain: A depends on B, B depends on C - ordering matters."""
-    a = Tensor.zeros(4).contiguous().realize()
+    a = Tensor.zeros(4, dtype=dtypes.int32).contiguous().realize()
    b = Tensor([1, 2, 3, 4]).contiguous().realize()
    c = Tensor([10, 10, 10, 10]).contiguous().realize()
    # b reads from c, a reads from b
@@ -660,8 +718,8 @@ class TestAssignOrdering(unittest.TestCase):

  def test_interleaved_assign_read_patterns(self):
    """Complex interleaved pattern: write A, read A into B, write B, read B."""
-    a = Tensor.zeros(4).contiguous().realize()
-    b = Tensor.zeros(4).contiguous().realize()
+    a = Tensor.zeros(4, dtype=dtypes.int32).contiguous().realize()
+    b = Tensor.zeros(4, dtype=dtypes.int32).contiguous().realize()

    a.assign(Tensor([1, 2, 3, 4]))
    b.assign(a.contiguous())       # b should get [1,2,3,4]
@@ -5,24 +5,13 @@ from tinygrad.runtime.support.c import DLL, record, init_records
 from tinygrad.runtime.support import c
 from tinygrad.runtime.support.autogen import gen

-class TestAutogen(unittest.TestCase):
+@unittest.skipIf(WIN, "doesn't compile on windows")
+class TestC(unittest.TestCase):
  def compile(self, src):
    with tempfile.NamedTemporaryFile(suffix=".so") as f:
      subprocess.check_output(('clang', '-x', 'c', '-fPIC', '-shared', '-', '-o', f.name), input=src.encode())
      return DLL("test", f.name)

-  def run_gen(self, contents):
-    with tempfile.NamedTemporaryFile(mode='w', suffix='.h') as f:
-      f.write(contents)
-      f.flush()
-
-      generated_code = gen(name="test_header", dll=None, files=[f.name])
-
-      namespace = {}
-      exec(generated_code, namespace)
-      return namespace
-
-  @unittest.skipIf(WIN, "doesn't compile on windows")
  def test_packed_struct(self):
    @record
    class Baz:
@@ -45,7 +34,6 @@ class TestAutogen(unittest.TestCase):
    assert b.c == 1
    assert b.d == 0

-  @unittest.skipIf(WIN, "doesn't compile on windows")
  def test_packed_struct_interop(self):
    @record
    class Baz:
@@ -75,7 +63,6 @@ class TestAutogen(unittest.TestCase):
    self.assertEqual(test(b), b.a + b.b + b.c + b.d)

  # https://github.com/python/cpython/issues/90914
-  @unittest.skipIf(WIN, "doesn't compile on windows")
  def test_bitfield_interop(self):
    @record
    class Baz:
@@ -103,7 +90,6 @@ class TestAutogen(unittest.TestCase):
    def test(x:Baz) -> ctypes.c_int: ...
    for i in range(8): self.assertEqual(test(Baz(*(j==i for j in range(8)))), i==2)

-  @unittest.skipIf(WIN, "doesn't compile on windows")
  def test_struct_interop(self):
    @record
    class Baz:
@@ -131,7 +117,6 @@ class TestAutogen(unittest.TestCase):
    def test(x:Baz) -> Baz: ...
    self.assertEqual(bytes(test(Baz(*range(8)))), struct.pack("8i", *range(7, -1, -1)))

-  @unittest.skipIf(WIN, "doesn't compile on windows")
  def test_aos_interop(self):
    @record
    class Item:
@@ -151,7 +136,6 @@ class TestAutogen(unittest.TestCase):
    def test(arr:(Item * 3)) -> ctypes.c_int: ...
    self.assertEqual(test((Item * 3)(Item(10), Item(20), Item(30))), 60)

-  @unittest.skipIf(WIN, "doesn't compile on windows")
  def test_soa_interop(self):
    @record
    class Row:
@@ -173,7 +157,6 @@ class TestAutogen(unittest.TestCase):
    self.assertEqual(r.data[1], 20)
    self.assertEqual(r.data[2], 10)

-  @unittest.skipIf(WIN, "doesn't compile on windows")
  def test_soa_ptr_interop(self):
    @record
    class Row:
@@ -191,7 +174,6 @@ class TestAutogen(unittest.TestCase):
    def test(x:Row) -> ctypes.c_int: ...
    assert test(Row((ctypes.c_int * 3)(10, 20, 30))) == 60

-  @unittest.skipIf(WIN, "doesn't compile on windows")
  def test_nested_struct_interop(self):
    @record
    class Inner:
@@ -217,7 +199,6 @@ class TestAutogen(unittest.TestCase):
    self.assertEqual(o.inner.a, 20)
    self.assertEqual(o.b, 10)

-  @unittest.skipIf(WIN, "doesn't compile on windows")
  def test_struct_pointer_interop(self):
    @record
    class Foo:
@@ -242,7 +223,88 @@ class TestAutogen(unittest.TestCase):
    self.assertEqual(out.contents.a, 20)
    self.assertEqual(out.contents.b, 10)

-  @unittest.skipIf(WIN, "doesn't compile on windows")
+  def test_pointer_field_roundtrip(self):
+    # This tests storing a pointer in a record struct field and passing it to C
+    # Mimics how mesa.struct_lp_build_tgsi_params.mask is used
+    from tinygrad.runtime.support.c import POINTER
+    @record
+    class Inner:
+      SIZE = 8
+      value: Annotated[ctypes.c_int, 0]
+      flag: Annotated[ctypes.c_int, 4]
+    @record
+    class Outer:
+      SIZE = 16
+      x: Annotated[ctypes.c_int, 0]
+      inner_ptr: Annotated[POINTER[Inner], 8]
+    init_records()
+
+    src = """
+      struct inner { int value; int flag; };
+      struct outer { int x; struct inner *inner_ptr; };
+      int test(struct inner *p) {
+        return p->value + p->flag;
+      }
+    """
+    dll = self.compile(src)
+    @dll.bind
+    def test(p:POINTER[Inner]) -> ctypes.c_int: ...
+
+    inner = Inner(value=42, flag=10)
+    outer = Outer(x=1, inner_ptr=ctypes.pointer(inner))
+    # Retrieve pointer from struct field and pass to C
+    self.assertEqual(test(outer.inner_ptr), 52)
+
+  def test_pointer_field_loses_reference(self):
+    # BUG: When a pointer is stored in a record struct field, only the address bytes are saved.
+    # The pointer's _objects dict (which prevents GC of the pointed-to object) is lost.
+    # This causes the pointed-to object to be garbage collected, leading to use-after-free.
+    from tinygrad.runtime.support.c import POINTER
+    @record
+    class MaskContext:
+      SIZE = 16
+      value: Annotated[ctypes.c_int, 0]
+      initialized: Annotated[ctypes.c_int, 4]
+      ptr: Annotated[ctypes.c_void_p, 8]
+    @record
+    class Params:
+      SIZE = 16
+      x: Annotated[ctypes.c_int, 0]
+      mask: Annotated[POINTER[MaskContext], 8]
+    init_records()
+
+    src = """
+      struct mask_ctx { int value; int initialized; void *ptr; };
+      void mask_begin(struct mask_ctx *m, int val) { m->value = val; m->initialized = 1; }
+      int mask_end(struct mask_ctx *m) { return m->value + m->initialized; }
+    """
+    dll = self.compile(src)
+    @dll.bind
+    def mask_begin(m:POINTER[MaskContext], val:ctypes.c_int) -> None: ...
+    @dll.bind
+    def mask_end(m:POINTER[MaskContext]) -> ctypes.c_int: ...
+
+    # When MaskContext() is created inline, it gets garbage collected after the pointer
+    # is stored because only the address bytes are saved, not the _objects reference.
+    params = Params(x=1, mask=ctypes.pointer(MaskContext()))
+    mask_begin(params.mask, 42)
+    result = mask_end(params.mask)
+    self.assertEqual(result, 43)  # 42 + 1
+
+@unittest.skipIf(OSX and ('MTLCompiler' in DLL._loaded_ or 'llvm' in DLL._loaded_), "libclang can't be loaded after MTLCompiler or llvm on OSX")
+@unittest.skipIf(WIN, "doesn't compile on windows")
+class TestAutogen(unittest.TestCase):
+  def run_gen(self, contents):
+    with tempfile.NamedTemporaryFile(mode='w', suffix='.h') as f:
+      f.write(contents)
+      f.flush()
+
+      generated_code = gen(name="test_header", dll=None, files=[f.name])
+
+      namespace = {}
+      exec(generated_code, namespace)
+      return namespace
+
  def test_packed_structs(self):
    ns = self.run_gen("""
 typedef unsigned NvU32;
@@ -292,47 +354,6 @@ typedef struct
    assert frts_cmd.readVbiosDesc.__class__ is FWSECLIC_READ_VBIOS_DESC
    assert frts_cmd.frtsRegionDesc.__class__ is FWSECLIC_FRTS_REGION_DESC

-  @unittest.skipIf(WIN, "doesn't compile on windows")
-  @unittest.skipIf(OSX, "can't find stdint?")
-  def test_packed_fields(self):
-    ns = self.run_gen("""#include <stdint.h>
-typedef struct die_info
- {
-	 uint16_t die_id;
-	 uint16_t die_offset; /* Points to the corresponding die_header structure */
- } die_info;
-
-typedef struct ip_discovery_header
- {
-	 uint32_t signature;    /* Table Signature */
-	 uint16_t version;      /* Table Version */
-	 uint16_t size;         /* Table Size */
-	 uint32_t id;           /* Table ID */
-	 uint16_t num_dies;     /* Number of Dies */
-	 die_info die_info[16]; /* list die information for up to 16 dies */
-	 union {
-		 uint16_t padding[1];	/* version <= 3 */
-		 struct {		/* version == 4 */
-			 uint8_t base_addr_64_bit : 1; /* ip structures are using 64 bit base address */
-			 uint8_t reserved : 7;
-			 uint8_t reserved2;
-		 };
-	 };
- } ip_discovery_header;
-""")
-
-    ip_discovery_header = ns['ip_discovery_header']
-
-    hdr = b'IPDS\x04\x00|\x1d\x80\x1a\xffd\x01\x00\x00\x00\x8c\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00' # noqa: E501
-    ihdr = ip_discovery_header.from_buffer_copy(hdr)
-
-    assert ctypes.sizeof(ihdr) == 80
-    assert ihdr.signature == 0x53445049
-    assert ihdr.version == 0x0004
-    assert ihdr.num_dies == 1
-    assert ihdr.base_addr_64_bit == 1
-
-  @unittest.skipIf(WIN, "doesn't compile on windows")
  def test_gen_from_header(self):
    namespace = self.run_gen("""
    typedef struct {
@@ -378,7 +399,6 @@ typedef struct ip_discovery_header
    self.assertTrue(hasattr(rect, 'height'))
    self.assertTrue(hasattr(rect, 'color'))

-  @unittest.skipIf(WIN, "doesn't compile on windows")
  def test_struct_ordering(self):
    namespace = self.run_gen("""
    struct A;
@@ -408,77 +428,6 @@ typedef struct ip_discovery_header
    self.assertTrue(hasattr(b, 'c_ptr'))
    self.assertTrue(hasattr(c, 'a_ptr'))

-  @unittest.skipIf(WIN, "doesn't compile on windows")
-  def test_pointer_field_roundtrip(self):
-    # This tests storing a pointer in a record struct field and passing it to C
-    # Mimics how mesa.struct_lp_build_tgsi_params.mask is used
-    from tinygrad.runtime.support.c import POINTER
-    @record
-    class Inner:
-      SIZE = 8
-      value: Annotated[ctypes.c_int, 0]
-      flag: Annotated[ctypes.c_int, 4]
-    @record
-    class Outer:
-      SIZE = 16
-      x: Annotated[ctypes.c_int, 0]
-      inner_ptr: Annotated[POINTER[Inner], 8]
-    init_records()
-
-    src = """
-      struct inner { int value; int flag; };
-      struct outer { int x; struct inner *inner_ptr; };
-      int test(struct inner *p) {
-        return p->value + p->flag;
-      }
-    """
-    dll = self.compile(src)
-    @dll.bind
-    def test(p:POINTER[Inner]) -> ctypes.c_int: ...
-
-    inner = Inner(value=42, flag=10)
-    outer = Outer(x=1, inner_ptr=ctypes.pointer(inner))
-    # Retrieve pointer from struct field and pass to C
-    self.assertEqual(test(outer.inner_ptr), 52)
-
-  @unittest.skipIf(WIN, "doesn't compile on windows")
-  def test_pointer_field_loses_reference(self):
-    # BUG: When a pointer is stored in a record struct field, only the address bytes are saved.
-    # The pointer's _objects dict (which prevents GC of the pointed-to object) is lost.
-    # This causes the pointed-to object to be garbage collected, leading to use-after-free.
-    from tinygrad.runtime.support.c import POINTER
-    @record
-    class MaskContext:
-      SIZE = 16
-      value: Annotated[ctypes.c_int, 0]
-      initialized: Annotated[ctypes.c_int, 4]
-      ptr: Annotated[ctypes.c_void_p, 8]
-    @record
-    class Params:
-      SIZE = 16
-      x: Annotated[ctypes.c_int, 0]
-      mask: Annotated[POINTER[MaskContext], 8]
-    init_records()
-
-    src = """
-      struct mask_ctx { int value; int initialized; void *ptr; };
-      void mask_begin(struct mask_ctx *m, int val) { m->value = val; m->initialized = 1; }
-      int mask_end(struct mask_ctx *m) { return m->value + m->initialized; }
-    """
-    dll = self.compile(src)
-    @dll.bind
-    def mask_begin(m:POINTER[MaskContext], val:ctypes.c_int) -> None: ...
-    @dll.bind
-    def mask_end(m:POINTER[MaskContext]) -> ctypes.c_int: ...
-
-    # When MaskContext() is created inline, it gets garbage collected after the pointer
-    # is stored because only the address bytes are saved, not the _objects reference.
-    params = Params(x=1, mask=ctypes.pointer(MaskContext()))
-    mask_begin(params.mask, 42)
-    result = mask_end(params.mask)
-    self.assertEqual(result, 43)  # 42 + 1
-
-  @unittest.skipIf(WIN, "doesn't compile on windows")
  def test_anonymous_children(self):
    namespace = self.run_gen("""
      struct foo {
@@ -491,7 +440,6 @@ typedef struct ip_discovery_header
    self.assertIn('struct_foo', namespace)
    self.assertIn('struct_foo_bar', namespace)

-  @unittest.skipIf(WIN, "doesn't compile on windows")
  def test_enums(self):
    namespace = self.run_gen("""
      enum Foo { A, B, C };
@@ -511,4 +459,43 @@ typedef struct ip_discovery_header
    assert namespace["enum_Bar"].get(1) == "Y"
    assert namespace["enum_Bar"].get(2) == "Z"

+  @unittest.skipIf(OSX, "can't find stdint?")
+  def test_packed_fields(self):
+    ns = self.run_gen("""#include <stdint.h>
+typedef struct die_info
+ {
+	 uint16_t die_id;
+	 uint16_t die_offset; /* Points to the corresponding die_header structure */
+ } die_info;
+
+typedef struct ip_discovery_header
+ {
+	 uint32_t signature;    /* Table Signature */
+	 uint16_t version;      /* Table Version */
+	 uint16_t size;         /* Table Size */
+	 uint32_t id;           /* Table ID */
+	 uint16_t num_dies;     /* Number of Dies */
+	 die_info die_info[16]; /* list die information for up to 16 dies */
+	 union {
+		 uint16_t padding[1];	/* version <= 3 */
+		 struct {		/* version == 4 */
+			 uint8_t base_addr_64_bit : 1; /* ip structures are using 64 bit base address */
+			 uint8_t reserved : 7;
+			 uint8_t reserved2;
+		 };
+	 };
+ } ip_discovery_header;
+""")
+
+    ip_discovery_header = ns['ip_discovery_header']
+
+    hdr = b'IPDS\x04\x00|\x1d\x80\x1a\xffd\x01\x00\x00\x00\x8c\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00' # noqa: E501
+    ihdr = ip_discovery_header.from_buffer_copy(hdr)
+
+    assert ctypes.sizeof(ihdr) == 80
+    assert ihdr.signature == 0x53445049
+    assert ihdr.version == 0x0004
+    assert ihdr.num_dies == 1
+    assert ihdr.base_addr_64_bit == 1
+
 if __name__ == "__main__": unittest.main()
@@ -0,0 +1,62 @@
+import unittest
+import numpy as np
+from tinygrad import Tensor
+from tinygrad.dtype import dtypes
+from tinygrad.uop.ops import UOp
+
+class TestCall(unittest.TestCase):
+  def test_call_plus(self):
+    a = Tensor.randn(10, 10)
+    b = Tensor.randn(10, 10)
+    Tensor.realize(a,b)
+
+    # we define a plus function
+    plus_fxn = UOp.param(0, dtypes.float, (10,10)) + UOp.param(1, dtypes.float, (10,10))
+
+    c = Tensor.call(a, b, fxn=plus_fxn)
+    np.testing.assert_equal(c.numpy(), (a+b).numpy())
+
+  def test_call_plus_backward(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
+
+    # this is the gradient for +
+    def grad_fxn(grad:UOp, call:UOp): return (grad, grad)
+
+    # we define a plus function
+    plus_fxn = UOp.param(0, dtypes.float, (10,10)) + UOp.param(1, dtypes.float, (10,10))
+    c = Tensor.call(a, b, fxn=plus_fxn, grad_fxn=grad_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)
+    b = Tensor.randn(K, N)
+    Tensor.realize(a, b)
+    c = Tensor.call(a, b, fxn=a.as_param(0) @ b.as_param(1))
+    np.testing.assert_allclose(c.numpy(), a.numpy() @ b.numpy(), rtol=1e-5, atol=1e-6)
+
+  @unittest.skip("needs GEMM on mixins")
+  def test_call_gemm_uop(self):
+    M, K, N = 4, 8, 4
+    a = Tensor.randn(M, K)
+    b = Tensor.randn(K, N)
+    Tensor.realize(a, b)
+
+    # we define a gemm function
+    x = UOp.param(0, dtypes.float, shape=(M, K))
+    y = UOp.param(1, dtypes.float, shape=(K, N))
+    c = Tensor.call(a, b, fxn=x@y)
+
+    np.testing.assert_allclose(c.numpy(), a.numpy() @ b.numpy(), rtol=1e-5, atol=1e-6)
+
+if __name__ == '__main__':
+  unittest.main()
@@ -59,13 +59,13 @@ class TestRawDiskBuffer(unittest.TestCase):
    _test_bitcasted(t, dtypes.float32, 0.0)
    _test_bitcasted(t, dtypes.uint32, 0)
    # pi in float16 stored via int16
-    t.bitcast(dtypes.uint16).assign(Tensor.full((128, 64), 0x4248, dtype=dtypes.uint16)).realize()
+    t.assign(Tensor.full((128, 64), 0x4248, dtype=dtypes.uint16).bitcast(dtypes.uint8)).realize()
    _test_bitcasted(t, dtypes.float16, 3.140625)
    _test_bitcasted(t, dtypes.float32, 50.064727)
    _test_bitcasted(t, dtypes.uint16, 0x4248)
    _test_bitcasted(t, dtypes.uint32, 0x42484248)
    # pi in float32 stored via float32
-    t.bitcast(dtypes.float32).assign(Tensor.full((128, 32), 3.1415927, dtype=dtypes.float32)).realize()
+    t.assign(Tensor.full((128, 32), 3.1415927, dtype=dtypes.float32).bitcast(dtypes.uint8)).realize()
    _test_bitcasted(t, dtypes.float32, 3.1415927)
    _test_bitcasted(t, dtypes.uint32, 0x40490FDB)
    # doesn't suport normal cast
@@ -250,6 +250,24 @@ class TestDiskTensor(unittest.TestCase):
    tout = [(x//256, x%256) for x in out]
    assert tout == list([(x+1,x) for x in range(32,64,2)])

+  def test_strided_read(self):
+    # test non-contiguous (strided) read - should read elements at indices 0, 2, 4
+    pathlib.Path(temp(fn:="dt_strided_read")).unlink(missing_ok=True)
+    dt = Tensor([0, 1, 2, 3, 4, 5]).to(f"disk:{temp(fn)}")
+    result = dt[::2].tolist()
+    # TODO: dt[::2] selects indices 0, 2, 4, so result should be [0, 2, 4]
+    # self.assertEqual(result, [0, 2, 4])
+    self.assertEqual(result, [0, 1, 2])  # wrong!
+
+  def test_permuted_read(self):
+    # test non-contiguous (permuted) read - should read transposed
+    pathlib.Path(temp(fn:="dt_permuted_read")).unlink(missing_ok=True)
+    dt = Tensor([[0, 1, 2], [3, 4, 5]]).to(f"disk:{temp(fn)}")
+    result = dt.T.tolist()
+    # TODO: transpose should give [[0, 3], [1, 4], [2, 5]]
+    # self.assertEqual(result, [[0, 3], [1, 4], [2, 5]])
+    self.assertEqual(result, [[0, 1], [2, 3], [4, 5]])  # wrong!
+
  def test_write_ones(self):
    pathlib.Path(temp("dt_write_ones")).unlink(missing_ok=True)

@@ -276,6 +294,15 @@ class TestDiskTensor(unittest.TestCase):
    dt[1] = [3]
    self.assertEqual(dt.tolist(), [[1], [3]])

+  def test_strided_setitem(self):
+    # test non-contiguous (strided) setitem - should set elements at indices 0, 2, 4
+    pathlib.Path(temp(fn:="dt_strided_setitem")).unlink(missing_ok=True)
+    dt = Tensor([1, 2, 3, 4, 5, 6]).to(f"disk:{temp(fn)}")
+    dt[::2] = Tensor([10, 20, 30])
+    # TODO: dt[::2] selects indices 0, 2, 4, so result should be [10, 2, 20, 4, 30, 6]
+    # self.assertEqual(dt.tolist(), [10, 2, 20, 4, 30, 6])
+    self.assertEqual(dt.tolist(), [10, 20, 30, 4, 5, 6])  # wrong!
+
  def test_assign_const_to_disk(self):
    # assign from CONST (Tensor.full) to disk - source has no buffer, needs contiguous first
    pathlib.Path(temp(fn:="dt_assign_const")).unlink(missing_ok=True)
@@ -321,13 +348,25 @@ class TestDiskTensor(unittest.TestCase):

  def test_assign_with_bitcast(self):
    # bitcast assign is used in safe_save for writing header length
-    # this tests the synchronous disk assign hack handles bitcast correctly
+    # bitcast on source side works, bitcast on target side raises
    pathlib.Path(temp(fn:="dt_assign_bitcast")).unlink(missing_ok=True)
    t = Tensor.empty(16, device=f"disk:{temp(fn)}", dtype=dtypes.uint8)
-    t[0:8].bitcast(dtypes.int64).assign([12345])
-    # verify the data was written correctly
+    # correct way: bitcast the source to match target dtype
+    t[0:8].assign(Tensor([12345], dtype=dtypes.int64, device="CPU").bitcast(dtypes.uint8))
    val = int.from_bytes(t[0:8].data(), 'little')
    self.assertEqual(val, 12345)
+    # bitcast on target with non-broadcastable dtype raises
+    with self.assertRaises(RuntimeError):
+      t[0:4].bitcast(dtypes.int32).assign(Tensor([12345], dtype=dtypes.int64))
+
+  def test_assign_to_bitcast_view(self):
+    # assign float values to a float32 view of a uint8 disk buffer (used by safe_save)
+    pathlib.Path(temp(fn:="dt_bitcast_view_assign")).unlink(missing_ok=True)
+    t = Tensor.empty(32, device=f"disk:{temp(fn)}", dtype=dtypes.uint8)
+    # create float32 view of bytes 8-24 (4 floats)
+    float_view = t[8:24].bitcast(dtypes.float32)
+    float_view.assign(Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32, device="CPU"))
+    np.testing.assert_array_equal(float_view.numpy(), [1.0, 2.0, 3.0, 4.0])

  def test_assign_cross_device(self):
    # disk assign allows cross-device (source on GPU/CPU, target on disk)
@@ -339,7 +339,7 @@ class TestIndexing(unittest.TestCase):
      numpy_testing_assert_equal_helper(output, input_list)
  '''

-  @unittest.skipUnless(is_dtype_supported(dtypes.long), f"long dtype not supported on {Device.DEFAULT}")
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU doesn't support long indexing: #13624")
  def test_index_ind_dtype(self):
    x = Tensor.randn(4, 4)
    # ind_long = torch.randint(4, (4,), dtype=torch.long)
@@ -1,5 +1,5 @@
 import unittest
-from tinygrad import dtypes, Device
+from tinygrad import dtypes
 from tinygrad.device import Buffer
 from tinygrad.engine.memory import _internal_memory_planner

@@ -7,7 +7,7 @@ global_map = {}
 def b(i, base=None, offset=0, pin=False, size=16):
  global global_map
  if i in global_map: return global_map[i]
-  global_map[i] = Buffer(Device.DEFAULT, size, dtypes.int8, base=global_map[base] if base is not None else None, offset=offset)
+  global_map[i] = Buffer("NULL", size, dtypes.int8, base=global_map[base] if base is not None else None, offset=offset)
  if pin: global_map[i].ref(1)
  return global_map[i]

@@ -35,7 +35,6 @@ class TestScheduleCache(unittest.TestCase):
    _, var_vals = t.schedule_with_vars()
    self.assertEqual(var_vals, {'pos': 42})

-  @Context(SPEC=0)
  def test_custom_kernel(self):
    for i in range(4):
      a = Tensor.empty(1)
@@ -43,7 +42,6 @@ class TestScheduleCache(unittest.TestCase):
      a.realize()
      self.assertEqual(a.item(), i)

-  @Context(SPEC=0)
  def test_same_custom_function_reuses_cache(self):
    schedule_cache.clear()
    fxn = functools.partial(custom_set0_kernel, num=10)
@@ -470,5 +470,19 @@ class TestUnfoldableImageChannelSelection(unittest.TestCase):
    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)))
    self.assertEqual(self._count_nans(load), 1)

+class TestDropTrueGate(unittest.TestCase):
+  def test_drop_true_gate_on_index(self):
+    # test that INDEX with a constant True gate gets simplified to drop the gate
+    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)
+    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))
+    # apply the optimization
+    result = graph_rewrite(index_with_gate, load_store_indexing)
+    # the True gate should be dropped (INDEX should only have 2 sources)
+    self.assertEqual(len(result.src), 2, "True gate should be dropped from INDEX")
+
 if __name__ == '__main__':
  unittest.main()
@@ -0,0 +1,94 @@
+import unittest
+from tinygrad import Tensor, dtypes
+from tinygrad.tensor import _METADATA
+from tinygrad.helpers import Context
+
+class TestTensorMetadata(unittest.TestCase):
+  def setUp(self) -> None:
+    _METADATA.set(None)
+    self._ctx = Context(SCACHE=0)
+    self._ctx.__enter__()
+  def tearDown(self) -> None:
+    self._ctx.__exit__(None, None, None)
+
+  @unittest.skip("why would this be true?")
+  def test_exclude_noop_metadata(self):
+    a = Tensor.rand(4, 4)*1
+    self.assertEqual(a.uop.metadata[0].name, "__mul__")
+    k = a.schedule()[-1]
+    self.assertEqual([m.name for m in k.metadata], ["rand"])
+
+  @unittest.skip("metadata not reaching kernel schedule")
+  def test_exclude_const_metadata(self):
+    a = Tensor.arange(4)
+    b = Tensor.full((4,), -1, dtype=dtypes.int).contiguous()
+    sched = Tensor.schedule(a, b)
+    self.assertEqual([m.name for m in sched[0].metadata], ["arange"])
+    self.assertEqual([m.name for m in sched[1].metadata], ["contiguous"])
+
+  def test_matmul(self):
+    x = Tensor.rand(3, requires_grad=True)
+    W = Tensor.rand(3, 3, requires_grad=True)
+    out = x.matmul(W)
+    self.assertEqual(out.uop.metadata[0].name, "matmul")
+    si = out.schedule()[-1]
+    self.assertEqual(len(si.metadata), 1)
+    self.assertEqual(si.metadata[0].name, "matmul")
+
+  def test_relu(self):
+    x = Tensor.rand(3, requires_grad=True)
+    out = x.relu()
+    self.assertEqual(out.uop.metadata[0].name, "relu")
+    si = out.schedule()[-1]
+    self.assertEqual(len(si.metadata), 1)
+    self.assertEqual(si.metadata[0].name, "relu")
+
+  @unittest.skip("assign metadata no longer captured")
+  def test_assign(self):
+    x = Tensor.empty(10, 10).realize()
+    x.assign(Tensor.ones(10, 10).contiguous())
+    si = x.schedule()[-1]
+    self.assertEqual(len(si.metadata), 1)
+    self.assertEqual(si.metadata[0].name, "assign")
+
+  def test_complex(self):
+    x = Tensor.rand(3, requires_grad=True)
+    y = Tensor.rand(3, requires_grad=True)
+    out = x.relu() * y.sigmoid()
+    self.assertEqual(out.uop.metadata[0].name, "__mul__")
+    self.assertEqual(out.uop.src[0].metadata[0].name, "relu")
+    self.assertEqual(out.uop.src[1].metadata[0].name, "sigmoid")
+    si = out.schedule()[-1]
+    self.assertEqual(len(si.metadata), 3)
+    self.assertEqual(set(m.name for m in si.metadata), {"relu", "sigmoid", "__mul__"})
+
+  def test_complex_backward(self):
+    x = Tensor.rand(3, requires_grad=True).realize()
+    y = Tensor.rand(3, requires_grad=True).realize()
+    out = (x.relu() * y.sigmoid()).sum()
+    self.assertEqual(out.uop.metadata[0].name, "sum")
+    out.backward()
+    self.assertEqual(x.grad.uop.metadata[0].name, "relu")
+    #self.assertTrue(x.grad.uop.metadata[0].backward)  # TODO: backward flag is False
+    self.assertEqual(y.grad.uop.metadata[0].name, "sigmoid")
+    #self.assertTrue(y.grad.uop.metadata[0].backward)  # TODO: backward flag is False
+    si = Tensor.schedule(out, x.grad, y.grad)[-1]
+    #self.assertEqual(len(si.metadata), 3, f"failed with {si.metadata}")
+    # skip numpy, this is schedule cache
+    self.assertSetEqual(set(m.name for m in si.metadata if m.name != "numpy"), {"sigmoid", "relu"})
+    #bw = [m for m in si.metadata if m.backward]
+    #self.assertEqual(len(bw), 1)
+    #self.assertEqual(bw[0].name, "sigmoid")
+
+  def test_tracemeta_0(self):
+    with Context(TRACEMETA=0):
+      x = Tensor.rand(3, requires_grad=True)
+      y = Tensor.rand(3, requires_grad=True)
+      out = (x.relu() * y.sigmoid()).sum()
+      self.assertIsNone(out.uop.metadata)
+      self.assertIsNone(out.uop.src[0].metadata)
+      si = out.schedule()[-1]
+      self.assertEqual(si.metadata, ())
+
+if __name__ == '__main__':
+  unittest.main()
@@ -479,6 +479,26 @@ class TestUOpGraph(unittest.TestCase):
    for u in uops:
      self.assertNotEqual(u.dtype, dtypes.long)

+  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=())
+    c1 = UOp.range(UOp.const(dtypes.index, 512), 1, AxisType.LOOP)
+    c2 = UOp.range(UOp.const(dtypes.index, 250), 2, AxisType.LOOP)
+    c3 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(512), arg=1, src=())
+    c4 = c3.index(c1)  # 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=())
+    c8 = c7.index(c6)
+    # (loaded + range) < const pattern - loaded value shouldn't be promoted to long
+    loaded_idx = c4.cast(dtypes.index)
+    comparison = (loaded_idx + c5) < UOp.const(dtypes.index, 60000)
+    c9 = comparison.where(c8.cast(dtypes.uint).cast(dtypes.uchar), 0).reduce(c5, arg=Ops.ADD)
+    c10 = c0.index(((c1*UOp.const(dtypes.index, 250))+c2)).store(c9).end(c1, c2)
+    uops = to_uops_list([c10])
+    for u in uops:
+      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)
@@ -686,6 +706,97 @@ class TestExpander(unittest.TestCase):
    sink = expander_rewrite(sink)
    print(sink)

+class TestReduceCollapse(unittest.TestCase):
+  def test_multi_range_reduce_add(self):
+    """Test that (x + y).reduce(r1, r2) distributes over multiple ranges"""
+    from tinygrad.codegen.simplify import pm_reduce_collapse
+    # Create two ranges
+    r1 = UOp.range(3, 0)
+    r2 = UOp.range(4, 1)
+    # Create x + y where x and y depend on different ranges
+    x = r1.cast(dtypes.float)
+    y = r2.cast(dtypes.float)
+    # (x + y).reduce(r1, r2) should be rewritten
+    red = (x + y).reduce(r1, r2, arg=Ops.ADD)
+    self.assertEqual(len(red.src), 3)  # value + 2 ranges
+    result = graph_rewrite(red, pm_reduce_collapse, name='test')
+    # Should become add of two separate reduces
+    self.assertEqual(result.op, Ops.ADD)
+
+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)
+    idx = UOp.const(dtypes.int, 0)
+    gate = UOp.const(dtypes.bool, True)
+    gated_index = buf.index(idx, gate)
+    gep = gated_index.gep(0)
+    alt = UOp.const(dtypes.float, 42.0)
+    gated_load = gep.load(alt)
+    self.assertEqual(len(gated_load.src), 2)  # GEP + alt
+    result = graph_rewrite(gated_load, load_store_folding, name='test')
+    # After rewrite, should still have alt value preserved
+    self.assertEqual(result.op, Ops.GEP)
+    inner_load = result.src[0]
+    self.assertEqual(inner_load.op, Ops.LOAD)
+    self.assertEqual(len(inner_load.src), 2)  # INDEX + alt
+
+  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)
+    idx = UOp.const(dtypes.int, 0)
+    idx1 = buf1.index(idx)
+    idx2 = buf2.index(idx)
+    ptrcat = UOp(Ops.PTRCAT, dtypes.float.ptr().vec(2), (idx1, idx2))
+    alt = UOp.const(dtypes.float.vec(2), 42.0)
+    gated_load = ptrcat.load(alt)
+    self.assertEqual(len(gated_load.src), 2)  # PTRCAT + alt
+    result = graph_rewrite(gated_load, load_store_folding, name='test')
+    # After rewrite, should be CAT of LOADs, each preserving alt
+    self.assertEqual(result.op, Ops.CAT)
+    for inner_load in result.src:
+      self.assertEqual(inner_load.op, Ops.LOAD)
+      self.assertEqual(len(inner_load.src), 2)  # INDEX + alt
+    self.assertEqual(inner_load.src[1].arg, 42.0)  # alt value preserved
+
+class TestConstBufferize(unittest.TestCase):
+  def test_const_bufferize_with_ranges(self):
+    """Test that CONST.BUFFERIZE with ranges is folded correctly.
+
+    BUFFERIZE can have ranges as additional sources beyond the value.
+    The pattern at rangeify.py uses allow_any_len=True because
+    CONST doesn't depend on ranges (constant is same value everywhere).
+    """
+    from tinygrad.schedule.rangeify import pm_const_buffer_folding, BufferizeOpts
+    c = UOp.const(dtypes.float, 42.0)
+    r1 = UOp.range(3, 0)
+    bufferize_with_range = UOp(Ops.BUFFERIZE, dtypes.float, (c, r1), arg=BufferizeOpts(device="CPU"))
+    self.assertEqual(len(bufferize_with_range.src), 2)  # const + 1 range
+
+    result = graph_rewrite(bufferize_with_range, pm_const_buffer_folding, name='test')
+    # BUFFERIZE should be removed, result is const broadcast to shape
+    self.assertNotEqual(result.op, Ops.BUFFERIZE)
+    const_vals = [u.arg for u in result.toposort() if u.op is Ops.CONST and u.dtype == dtypes.float]
+    self.assertIn(42.0, const_vals)
+
+  def test_const_bufferize_with_multiple_ranges(self):
+    """Test CONST.BUFFERIZE with multiple ranges is also folded."""
+    from tinygrad.schedule.rangeify import pm_const_buffer_folding, BufferizeOpts
+    c = UOp.const(dtypes.float, 3.14)
+    r1 = UOp.range(3, 0)
+    r2 = UOp.range(4, 1)
+    bufferize_with_ranges = UOp(Ops.BUFFERIZE, dtypes.float, (c, r1, r2), arg=BufferizeOpts(device="CPU"))
+    self.assertEqual(len(bufferize_with_ranges.src), 3)  # const + 2 ranges
+
+    result = graph_rewrite(bufferize_with_ranges, pm_const_buffer_folding, name='test')
+    # BUFFERIZE should be removed
+    self.assertNotEqual(result.op, Ops.BUFFERIZE)
+    const_vals = [u.arg for u in result.toposort() if u.op is Ops.CONST and u.dtype == dtypes.float]
+    self.assertIn(3.14, const_vals)
+
 class TestUOpTags(unittest.TestCase):
  def test_inc_by_one(self):
    g = UOp.const(dtypes.int, 1) + UOp.const(dtypes.int, 1)
@@ -1074,6 +1074,24 @@ class TestGatedUopGivenValid(unittest.TestCase):
    expected_vec = UOp(Ops.VECTORIZE, dtypes.index.vec(2), (uconst(0), r0))
    self.assertEqual(idx, (r0 < 3).where(expected_vec, UOp.invalid()))

+class TestRangeSplitting(unittest.TestCase):
+  def test_range_split_on_mod(self):
+    # test that mark_range_mod splits RANGE(8) into RANGE(4)*2 + RANGE(2) when used with %2
+    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)
+    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)),))
+    # count RANGEs before
+    ranges_before = len([u for u in sink.toposort() if u.op is Ops.RANGE])
+    # apply the range splitting optimization
+    sink_after = graph_rewrite(sink, pm_split_ranges+pm_flatten_range, ctx={}, name="test split ranges")
+    # count RANGEs after - should have more due to splitting
+    ranges_after = len([u for u in sink_after.toposort() if u.op is Ops.RANGE])
+    self.assertGreater(ranges_after, ranges_before, "RANGE should be split when used with mod of divisible constant")
+
 class TestBounds(unittest.TestCase):
  def test_unrolled_arange(self):
    # #include <metal_stdlib>
@@ -139,7 +139,7 @@ class TransformerBlock:
    v = self.cache_kv[1, :, :, 0:start_pos+T, :]

    # NOTE: this mask is causal_lower_right, not the causal_upper_left generated by is_casual = True
-    mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, device=x.device).triu(start_pos+1) if T > 1 else None
+    mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, device=x.device).triu(int(start_pos)+1) if T > 1 else None
    attn = q.scaled_dot_product_attention(k, v, attn_mask=mask, enable_gqa=True)     # (B,H,T,Hd)
    attn = attn.transpose(1, 2).reshape(B, T, -1)                                    # back to (B,T,D)
    attn = self.attn_output(attn)
@@ -1,6 +1,6 @@
 from typing import cast
 import itertools
-from tinygrad.helpers import DISABLE_FAST_IDIV, DEVECTORIZE, TRANSCENDENTAL, SPEC, DEBUG, getenv, TracingKey, Context
+from tinygrad.helpers import DISABLE_FAST_IDIV, EMULATED_DTYPES, DEVECTORIZE, TRANSCENDENTAL, SPEC, DEBUG, getenv, TracingKey, Context
 from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, pyrender
 from tinygrad.uop.spec import type_verify, program_spec, kernel_spec
 from tinygrad.renderer import Renderer, ProgramSpec
@@ -95,7 +95,8 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -

  # decompositions
  supported_ops = tuple(ren.code_for_op.keys())
-  pm_decomp = symbolic_simple+get_late_rewrite_patterns(supported_ops, TRANSCENDENTAL>=2, bool(DISABLE_FAST_IDIV))
+  pm_decomp = symbolic_simple+get_late_rewrite_patterns(supported_ops, ren.device, TRANSCENDENTAL>=2, bool(DISABLE_FAST_IDIV),
+                                                        tuple(EMULATED_DTYPES.tolist(dtypes)))
  sink = graph_rewrite(sink, pm_decomp, ctx=ren.device, name="decompositions")

  # final rules for the renderer (without sym)
@@ -114,8 +115,8 @@ pm_linearize_cleanups = PatternMatcher([
  # if statements are not allowed in the graph
  (UPat((Ops.IF, Ops.ENDIF)), lambda: panic(RuntimeError("if not allowed in graph"))),
  # gated INDEX becomes IF-STORE-ENDIF. this is the only use of IF-ENDIF
-  (UPat(Ops.STORE, name="u", src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat(name="gate", dtype=dtypes.bool))).or_casted(), UPat()),
-        allow_any_len=True), lambda u, gate: (u, [mif:=UOp(Ops.IF, src=(gate, u.src[0])), u, UOp(Ops.ENDIF, src=(mif,))]))
+  (UPat(Ops.STORE, name="u", src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat(name="gate", dtype=dtypes.bool))).or_casted(), UPat())),
+   lambda u, gate: (u, [mif:=UOp(Ops.IF, src=(gate, u.src[0])), u, UOp(Ops.ENDIF, src=(mif,))]))
 ])

 # requires lst be toposorted. like graph rewrite, but for lines
@@ -123,12 +123,12 @@ load_store_folding = PatternMatcher([
  (UPat(Ops.LOAD, src=(UPat(Ops.GEP, name="gep"),), name="ld", allow_any_len=True),
   lambda gep, ld: ld.replace(dtype=ld.dtype.scalar().vec(gep.dtype.count), src=(gep.src[0],)+ld.src[1:]).gep(gep.arg)),
  # GEP on data of STORE
-  (UPat(Ops.STORE, src=(UPat(Ops.GEP, name="gep"), UPat.var("st")), allow_any_len=True, name="sto"), gep_on_store),
+  (UPat(Ops.STORE, src=(UPat(Ops.GEP, name="gep"), UPat.var("st")), name="sto"), gep_on_store),
  # put PTRCAT after LOAD
  (UPat(Ops.LOAD, src=(UPat(Ops.PTRCAT, name="cat"),), name="ld", allow_any_len=True),
   lambda cat,ld: UOp(Ops.CAT, cat.dtype.base.vec(cat.dtype.vcount), tuple(ld.replace(dtype=x.dtype.base, src=(x,)+ld.src[1:]) for x in cat.src))),
  # put PTRCAT after STORE
-  (UPat(Ops.STORE, src=(UPat(Ops.PTRCAT, name="cat"), UPat(name="data")), allow_any_len=True, name="sto"), cat_after_store),
+  (UPat(Ops.STORE, src=(UPat(Ops.PTRCAT, name="cat"), UPat(name="data")), name="sto"), cat_after_store),
 ])

 # *** correct load/store ***
@@ -201,7 +201,7 @@ def image_fixup(ls:UOp):
    x_mod_4 = x % 4
    def sel(ret, i): return x_mod_4.ne(i).where(ret, vec_load.gep(i))
    # if x is non-negative, x % 4 is in [0, 3] and we can skip NAN fallback
-    if x_mod_4.vmin >= 0: return functools.reduce(sel, range(x_mod_4.vmin+1, x_mod_4.vmax+1), vec_load.gep(x_mod_4.vmin))
+    if x_mod_4.vmin >= 0: return functools.reduce(sel, range(int(x_mod_4.vmin)+1, int(x_mod_4.vmax)+1), vec_load.gep(int(x_mod_4.vmin)))
    return functools.reduce(sel, range(4), ls.const_like(float('nan')))

  return None
@@ -319,8 +319,7 @@ 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)
@@ -342,6 +341,6 @@ pm_add_loads = PatternMatcher([
  (UPat(Ops.INDEX, name="idx"), lambda idx: None if isinstance(idx.dtype, (PtrDType, ImageDType)) else
    idx.replace(dtype=idx.src[0].dtype).load(dtype=idx.dtype.base)),
  # remove loads from stores
-  (UPat(Ops.STORE, src=(UPat(Ops.LOAD),), allow_any_len=True, name="s"), lambda s: s.replace(src=(s.src[0].src[0],)+s.src[1:])),
+  (UPat(Ops.STORE, src=(UPat(Ops.LOAD), UPat(name="val")), name="s"), lambda s,val: s.replace(src=(s.src[0].src[0], val))),
 ])

@@ -1,3 +1,4 @@
+from __future__ import annotations
 import math, itertools
 from collections import defaultdict
 from typing import cast, Final
@@ -25,11 +26,11 @@ class Scheduler:
    # always in order by axistype
    return sorted([u for u in self.ast.backward_slice if u.op is Ops.RANGE and u.vmax > 0], key=lambda x: (axis_to_pos[x.arg[-1]],) + x.arg[0:-1])
  @property
-  def shape_len(self): return len(self.rngs)
+  def shape_len(self) -> int: return len(self.rngs)
  @property
  def full_shape(self): return [ssimplify(x.src[0]) for x in self.rngs]
  @property
-  def axis_types(self): return [x.arg[-1] for x in self.rngs]
+  def axis_types(self) -> list[AxisType]: return [x.arg[-1] for x in self.rngs]

  # strings like ['g0', 'g1', 'l0', 'l1', 'l2', 'l3', 'l4', 'l5', 'R0', 'r0', 'r1', 'r2', 'u0', 'u1', 'u2']
  def shape_str(self) -> list[str]:
@@ -41,7 +42,7 @@ class Scheduler:
    return ret
  def shape_str_to_axis(self, nms:list[str]) -> tuple[int, ...]: return tuple([self.shape_str().index(x) for x in nms])

-  def copy(self):
+  def copy(self) -> Scheduler:
    ret = Scheduler(self.ast, self.ren)
    ret.dont_use_locals = self.dont_use_locals
    ret.applied_opts = self.applied_opts[:]
@@ -49,7 +50,7 @@ class Scheduler:
    return ret

  kernel_cnt: Final[defaultdict[str, int]] = defaultdict(int)
-  def get_optimized_ast(self, name_override:str|None=None):
+  def get_optimized_ast(self, name_override:str|None=None) -> UOp:
    if name_override is not None: name = name_override
    else:
      k_type = "r" if self.reduceop is not None else "E"
@@ -72,8 +73,8 @@ class Scheduler:
        ret = [r for r in ret if r in x.ranges]
    return ret

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

    globalizible_rngs = self._globalizable_rngs()
    rng = [x.replace(arg=x.arg[0:-1]+(AxisType.GLOBAL,)) if x in globalizible_rngs else x for x in self.rngs]
@@ -92,7 +93,7 @@ class Scheduler:
    return ret
  def colored_shape(self) -> str: return ' '.join([colored(f'{x.src[0].render():>4s}', color) for x,color in zip(self.rngs, self.colors())])

-  def shift_to(self, rng:UOp, amount:int, new_type:AxisType, top:bool=False, input_new_rng=None):
+  def shift_to(self, rng:UOp, amount:int, new_type:AxisType, top:bool=False, input_new_rng:UOp|None=None):
    if (old_sz:=rng.src[0].divides(amount)) is None:
      raise KernelOptError(f"{amount} can't divide {rng.src[0]} in {self.colored_shape()}")
    new_rng = UOp.range(amount, next(self.opt_range), new_type) if input_new_rng is None else input_new_rng
@@ -104,7 +105,7 @@ class Scheduler:
  def ranges_of(self, *axis_type:AxisType) -> list[UOp]: return [r for r in self.rngs if r.arg[-1] in axis_type]
  def axes_of(self, *axis_type:AxisType) -> list[int]: return [i for i,t in enumerate(self.axis_types) if t in axis_type]

-  def upcast_size(self) -> int: return prod(self.full_shape[a] for a in self.axes_of(AxisType.UPCAST, AxisType.UNROLL))
+  def upcast_size(self): return prod(self.full_shape[a] for a in self.axes_of(AxisType.UPCAST, AxisType.UNROLL))

  # copied from kernel.py
  @property
@@ -114,7 +115,7 @@ class Scheduler:
  def unrollable_dims(self) -> list[int]: return [i for i in self.axes_of(AxisType.GROUP_REDUCE, AxisType.REDUCE) \
                                                  if isinstance(s:=self.full_shape[i], int) and s > 1]

-  def real_axis(self, op:OptOps, axis:int|None):
+  def real_axis(self, op:OptOps, axis:int|None) -> int:
    try:
      if axis is None or op is OptOps.TC: return -1
      if op is OptOps.UNROLL: return self.unrollable_dims[axis]
@@ -4,7 +4,7 @@ from tinygrad.uop.symbolic import symbolic
 from tinygrad.helpers import partition, dedup
 from tinygrad.dtype import dtypes, ImageDType

-def flatten_range(r:UOp):
+def flatten_range(r:UOp) -> UOp|None:
  off = range_start[r.op]
  rngs = r.src[off:]
  if not len(rngs): return None
@@ -16,7 +16,7 @@ pm_flatten_range = PatternMatcher([
  (UPat((Ops.REDUCE, Ops.STORE, Ops.END), name="r"), flatten_range),
 ])

-def count_divmod(x:UOp): return len([u for u in x.toposort() if u.op in {Ops.IDIV, Ops.MOD}])
+def count_divmod(x:UOp) -> int: return len([u for u in x.toposort() if u.op in {Ops.IDIV, Ops.MOD}])
 def simplify_merge_adjacent(u:UOp) -> UOp|None:
  reduce_ranges = [x.ranges for x in u.backward_slice_with_self if x.op is Ops.REDUCE]
  # on END we only want to merge adjacent ranges, on REDUCE we want to try all combinations
@@ -40,10 +40,10 @@ pm_simplify_ranges = PatternMatcher([
  (UPat((Ops.END, Ops.REDUCE), name="u"), simplify_merge_adjacent),
 ])

-def mark_range_mod(ctx, r:UOp, c:UOp):
+def mark_range_mod(ctx:dict[UOp, UOp|None], r:UOp, c:UOp) -> None:
  if r not in ctx and r.src[0].op is Ops.CONST and r.src[0].divides(c.arg) is not None: ctx[r] = c

-def do_substitute(ctx, x: UOp):
+def do_substitute(ctx:dict[UOp, UOp|None], x: UOp) -> UOp|None:
  subs = {}
  for k,v in ctx.items():
    if v is not None:
@@ -53,7 +53,7 @@ def do_substitute(ctx, x: UOp):
  ctx.clear()
  return ret

-def dont_sub_ranges_for_image(ctx, x:UOp):
+def dont_sub_ranges_for_image(ctx:dict[UOp, UOp|None], x:UOp) -> None:
  if isinstance(x.src[0].src[0].dtype, ImageDType):
    for s in x.src[0].ranges: ctx[s] = None

@@ -67,7 +67,7 @@ pm_split_ranges = PatternMatcher([

 def no_range(u:UOp) -> bool: return not any(x.op is Ops.RANGE for x in u.backward_slice_with_self)

-def reduce_unparented(red:UOp):
+def reduce_unparented(red:UOp) -> UOp|None:
  if red.arg not in {Ops.ADD, Ops.MAX, Ops.MUL}: return None
  assert all(x.op is Ops.RANGE for x in red.src[1:]), "some reduce srcs aren't ranges"
  reduce_parented, reduce_unparented = partition(red.src[1:], lambda x: x in red.src[0].ranges)
@@ -119,7 +119,7 @@ pm_reduce_load_collapse = pm_reduce_collapse + PatternMatcher([
   lambda r,idx,expr: (v:=(idx.cast(r.dtype) >= 0) & (idx.cast(r.dtype) < r.src[0])).where(expr.substitute({r:idx.cast(r.dtype).valid(v)}),0)),
 ])

-def reduce_collapse(red:UOp, u:UOp, pm=pm_reduce_collapse):
+def reduce_collapse(red:UOp, u:UOp, pm:PatternMatcher=pm_reduce_collapse) -> UOp|None:
  for r in red.src[1:]:
    included = u.toposort(gate=lambda x: r in x.ranges)
    if any(x.op in {Ops.STORE, Ops.REDUCE} for x in included): return None
@@ -134,7 +134,7 @@ def reduce_collapse(red:UOp, u:UOp, pm=pm_reduce_collapse):
    u = sink.substitute({v:k for k,v in replaces.items()})
  return u

-def reduce_load_collapse(red:UOp, u:UOp): return reduce_collapse(red, u, pm=pm_reduce_load_collapse)
+def reduce_load_collapse(red:UOp, u:UOp) -> UOp|None: return reduce_collapse(red, u, pm=pm_reduce_load_collapse)

 # remove REDUCE without loads (generic arange opt / indexing).
 pm_reduce_simplify = pm_reduce_unparented + PatternMatcher([
@@ -148,7 +148,7 @@ pm_load_collapse = PatternMatcher([
  ((UPat.var("x", dtypes.index)+UPat.var("y"))<UPat.var("c"), lambda x,y,c: x < c-y if no_load(y) and no_load(c) and not no_load(x) else None),
 ])

-def cut_store_range(ctx, store:UOp, r:UOp):
+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
@@ -6,6 +6,7 @@ import importlib, inspect, functools, pathlib, os, platform, contextlib, sys, re
 from tinygrad.helpers import CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored
 from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, dedup, ContextVar
 from tinygrad.helpers import unwrap_class_type, suppress_finalizing, select_first_inited, VIZ, CPU_LLVM, CPU_LVP, NV_PTX, CUDA_PTX, NV_NAK
+from tinygrad.helpers import EMULATED_DTYPES
 from tinygrad.dtype import DType, ImageDType, PtrDType, dtypes, _to_np_dtype
 from tinygrad.renderer import Renderer

@@ -146,8 +147,9 @@ class Buffer:
  def deallocate(self):
    assert hasattr(self, '_buf'), "buffer must be allocated to deallocate"
    if DEBUG is not None and DEBUG >= 7: print(f"buffer: deallocate {self.nbytes} bytes on {self.device}")
-    if self._base is None and (self.options is None or self.options.external_ptr is None):
-      if GlobalCounters is not None and not self.device.startswith("DISK"): GlobalCounters.mem_used -= self.nbytes
+    if self._base is None:
+      if GlobalCounters is not None and not self.device.startswith("DISK") and (self.options is None or self.options.external_ptr is None):
+        GlobalCounters.mem_used -= self.nbytes
      if PROFILE: Buffer.profile_events.append(ProfilePointEvent(self.device, "free", self.trace_num))
      self.allocator.free(self._buf, self.nbytes, self.options)
    elif self._base is not None: self._base.allocated_views -= 1
@@ -262,7 +264,7 @@ class LRUAllocator(Allocator, Generic[DeviceType]):
      for opaque in opaques: super().free(opaque, sz, options)
      opaques.clear()
  def free(self, opaque:Any, size:int, options:BufferSpec|None=None):
-    if LRU and (options is None or not options.nolru): self.cache[(size, options)].append(opaque)
+    if LRU and (options is None or (not options.nolru and options.external_ptr is None)): self.cache[(size, options)].append(opaque)
    else: super().free(opaque, size, options)

 # **************** for Compiled Devices ****************
@@ -367,13 +369,15 @@ def is_dtype_supported(dtype:DType, device:str|None=None) -> bool:
  # for CI LLVM, it segfaults because it can't link to the casting function
  # CI CUDA architecture is sm_35 but we need at least sm_70 to run fp16 ALUs
  # PYTHON supports half memoryview in 3.12+ https://github.com/python/cpython/issues/90751
+  # double can't be bitcast to anything without long support
  if dtype == dtypes.half:
    if device == "CL": return not CI and not OSX
    if device == "QCOM": return False # QCOM compiler is flaky with half
    if device in ["CUDA", "NV"]: return not CI
    if device == "CPU" and CPU_LLVM: return OSX
    if device == "PYTHON": return sys.version_info >= (3, 12)
-  if dtype == dtypes.float64: return device not in {"METAL", "QCOM"} and not (OSX and device == "CL") and not getenv("NULL_IR3")
+  if dtype == dtypes.float64: return (device not in {"METAL", "QCOM"} and not (OSX and device == "CL") and not getenv("NULL_IR3")
+                                      and dtypes.long not in EMULATED_DTYPES.tolist(dtypes))
  return True

 if PROFILE:
@@ -3,6 +3,7 @@ import time, pprint, random, itertools, math
 from dataclasses import dataclass, replace, field
 from tinygrad.helpers import all_same, colored, DEBUG, GlobalCounters, ansilen, BEAM, NOOPT, all_int, CAPTURING, Metadata, TRACEMETA, TracingKey
 from tinygrad.helpers import DEVECTORIZE, time_to_str, VALIDATE_WITH_CPU, cpu_profile, PROFILE, ProfilePointEvent, cpu_events, prod, Context, unwrap
+from tinygrad.helpers import EMULATED_DTYPES
 from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, sym_infer
 from tinygrad.device import Device, Buffer
 from tinygrad.renderer import ProgramSpec, Estimates
@@ -74,7 +75,7 @@ class BufferCopy(Runner):
      getattr(src.allocator.dev, 'fd', None) is not None and dest.allocator.supports_copy_from_disk
    if disk_supports_fast_copyout and hasattr(dest.allocator, 'copy_from_disk') and src.nbytes >= 4096:
      dest.allocator.copy_from_disk(dest._buf, src._buf, src.nbytes)
-    elif (src.device.startswith("DISK") or src.device.startswith("TINYFS")) and hasattr(dest.allocator, '_as_buffer'):
+    elif isinstance(src.device, str) and src.device.startswith(("DISK", "TINYFS")) and hasattr(dest.allocator, '_as_buffer'):
      # fast(ish) path, uses readinto in diskbuffers
      src.allocator._copyout(dest.allocator._as_buffer(dest._buf), src._buf)
    else:
@@ -106,10 +107,10 @@ class EncDec(Runner):

 # **************** method cache ****************

-method_cache: dict[tuple[str, type, bytes, tuple[int, ...], bool], CompiledRunner] = {}
+method_cache: dict[tuple[str, type, bytes, tuple, bool], CompiledRunner] = {}
 def get_runner(device:str, ast:UOp) -> CompiledRunner:
  # TODO: this should be all context relevant to rendering
-  context = (BEAM.value, NOOPT.value, DEVECTORIZE.value)
+  context = (BEAM.value, NOOPT.value, DEVECTORIZE.value, EMULATED_DTYPES.value)
  ckey = (device, type(Device[device].compiler), ast.key, context, False)
  if cret:=method_cache.get(ckey): return cret
  bkey = (device.split(":")[0], type(Device[device].compiler), ast.key, context, True)
@@ -158,7 +159,7 @@ class ExecItem:
    var_vals = self.fixedvars if _var_vals is None else (_var_vals|self.fixedvars)
    # reorder bufs to match program globals if needed
    _bufs = [self.bufs[i] for i in self.prg.p.globals] if isinstance(self.prg, CompiledRunner) else self.bufs
-    bufs = cast(list[Buffer], [unwrap(x) for x in _bufs] if jit else [unwrap(x).ensure_allocated() for x in _bufs])
+    bufs = [unwrap(x) for x in _bufs] if jit else [unwrap(x).ensure_allocated() for x in _bufs]
    if PROFILE:
      payload = {"metadata":self.metadata, "var_vals":var_vals, "bufs":[b.trace_num for b in bufs], "name":self.prg.display_name}
      payload["outputs"], payload["inputs"] = (self.prg.p.outs, self.prg.p.ins) if isinstance(self.prg, CompiledRunner) else ([0], [1])
@@ -1,21 +1,25 @@
 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
+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.device import Buffer, MultiBuffer
-from tinygrad.helpers import DEBUG, cpu_profile, TracingKey, SPEC, flatten, pluralize, SCACHE
+from tinygrad.helpers import DEBUG, cpu_profile, TracingKey, SPEC, flatten, pluralize, SCACHE, Metadata
 from tinygrad.engine.realize import ExecItem

 # **** schedule linearizer

+# ScheduleItem = tuple[AST, buffer UOps, metadata, fixedvars, bound_ranges]
+ScheduleItem = tuple[UOp, tuple[UOp, ...], tuple[Metadata, ...], dict[str, int], tuple[UOp, ...]]
+
+# unwrap VIEW/CAST/etc to find the actual data source (kernel output, buffer, or multi-device op)
 def _unwrap_src(s: UOp) -> UOp:
  while len(s.src) and s.op not in {Ops.AFTER, Ops.BUFFER, Ops.MSELECT, Ops.MSTACK, Ops.BIND}: s = s.src[0]
  return s

 def create_schedule(sched_sink:UOp) -> tuple[list[ExecItem], UOp]:
  with cpu_profile(TracingKey("toposort sched_sink")):
-    # construct the KERNEL children graph based on assigns
+    # build kernel dependency graph: edges from producer kernel to consumer kernels
    children: dict[UOp, list[UOp]] = {}
    in_degree: dict[UOp, int] = {}
    for u in sched_sink.toposort():
@@ -26,81 +30,77 @@ def create_schedule(sched_sink:UOp) -> tuple[list[ExecItem], UOp]:
      k = u.src[1]
      in_degree.setdefault(k, 0)
      for s in k.src[0].src if k.op is Ops.END else k.src:
-        s = _unwrap_src(s)
-        if s.op is Ops.AFTER:
-          children.setdefault(s.src[1], []).append(k)
-          in_degree[k] += 1
-        elif s.op in {Ops.MSELECT, Ops.MSTACK}:
-          for ss in s.src:
-            if ss.op is Ops.MSELECT: ss = ss.src[0]
-            if ss.op is not Ops.BUFFER:
-              assert ss.op is Ops.AFTER, f"ss.op is not AFTER, it's {ss.op}"
-              children.setdefault(ss.src[1], []).append(k)
-              in_degree[k] += 1
-        elif s.op in {Ops.BUFFER, Ops.BIND}:
-          pass  # a BUFFER is already realized, BINDs are handled in complete_create_schedule_with_vars
-        else:
-          raise RuntimeError(f"input to kernel must be AFTER or BUFFER, not {s.op}")
+        match (s := _unwrap_src(s)).op:
+          case Ops.AFTER:
+            children.setdefault(s.src[1], []).append(k)
+            in_degree[k] += 1
+          case Ops.MSELECT | Ops.MSTACK:
+            for ss in s.src:
+              if ss.op is Ops.MSELECT: ss = ss.src[0]
+              if ss.op is not Ops.BUFFER:
+                assert ss.op is Ops.AFTER, f"ss.op is not AFTER, it's {ss.op}"
+                children.setdefault(ss.src[1], []).append(k)
+                in_degree[k] += 1
+          case Ops.BUFFER | Ops.BIND:
+            pass  # BUFFER is already realized, BIND is outer range (handled via bound_ranges below)
+          case _:
+            raise RuntimeError(f"input to kernel must be AFTER, BUFFER, MSELECT, MSTACK, or BIND, not {s.op}")

  with cpu_profile(TracingKey("linearize schedule")):
    queue: deque[UOp] = deque()
    for k,v in in_degree.items():
      if v == 0: queue.append(k)

-    schedule: list[tuple|UOp] = []
+    schedule: list[ScheduleItem|UOp] = []  # ScheduleItem for kernels, UOp for RANGE/END
    while len(queue):
      k = rk = queue.popleft()
      if k.op is Ops.END: k = k.src[0]
+      assert k.op in {Ops.RANGE, Ops.KERNEL}, f"unexpected op in queue: {k.op}"
      if k.op is Ops.RANGE: schedule.append(k)
      elif k.op is Ops.KERNEL:
-        ast = k.arg.ast
+        ast = (kernel:=cast(Kernel, k.arg)).ast
        buf_uops = tuple(_unwrap_src(s).buf_uop for s in k.src if s.op is not Ops.BIND)
        bound_ranges = tuple(s for s in k.src if s.op is Ops.BIND and len(s.src) > 1 and s.src[1].op is Ops.RANGE)
-        schedule.append((ast, buf_uops, k.arg.metadata, {}, bound_ranges))
+        schedule.append((ast, buf_uops, kernel.metadata, {}, bound_ranges))
        if rk.op is Ops.END: schedule.append(rk)
-      else:
-        raise RuntimeError(f"can't schedule {k.op}")
      for x in children.get(rk, []):
        in_degree[x] -= 1
        if in_degree[x] == 0: queue.append(x)

-  with cpu_profile(TracingKey("expand ranges")):
-    pre_schedule: list[ExecItem] = []
-    buf_uops_list: list[UOp] = []
-    sched_ptr = 0
-    in_ranges: dict[UOp, int] = {}
-    range_ptrs: dict[UOp, int] = {}
-    while sched_ptr < len(schedule):
-      si = schedule[sched_ptr]
-      if isinstance(si, UOp):
-        if si.op is Ops.RANGE:
-          in_ranges[si] = 0
-          range_ptrs[si] = sched_ptr + 1
-        elif si.op is Ops.END:
-          if in_ranges[si.src[1]] < si.src[1].vmax:
-            in_ranges[si.src[1]] += 1
-            sched_ptr = range_ptrs[si.src[1]]
-            continue
-      else:
-        ast, buf_uops, metadata, fixedvars, bound_ranges = si
-        fixedvars = fixedvars | {s.src[0].arg[0]:in_ranges[s.src[1]] for s in bound_ranges}
-        pre_schedule.append(ExecItem(ast, [], metadata, fixedvars))
-        buf_uops_list.append(UOp.sink(*buf_uops))
-      sched_ptr += 1
+  with cpu_profile(TracingKey("unroll outer ranges")):
+    pre_schedule, buf_uops_list = unroll_outer_ranges(schedule)
  return pre_schedule, UOp.sink(*buf_uops_list)

+def unroll_outer_ranges(schedule:list[ScheduleItem|UOp]) -> tuple[list[ExecItem], list[UOp]]:
+  pre_schedule: list[ExecItem] = []
+  buf_uops_list: list[UOp] = []
+  sched_ptr, in_ranges, range_ptrs = 0, dict[UOp, int](), dict[UOp, int]()
+  while sched_ptr < len(schedule):
+    if isinstance(si := schedule[sched_ptr], UOp):
+      if si.op is Ops.RANGE:
+        in_ranges[si] = 0
+        range_ptrs[si] = sched_ptr + 1
+      elif si.op is Ops.END:
+        if in_ranges[si.src[1]] < si.src[1].vmax:
+          in_ranges[si.src[1]] += 1
+          sched_ptr = range_ptrs[si.src[1]]
+          continue
+    else:
+      ast, buf_uops, metadata, _, bound_ranges = si
+      fixedvars = {s.src[0].arg[0]:in_ranges[s.src[1]] for s in bound_ranges}
+      pre_schedule.append(ExecItem(ast, [], metadata, fixedvars))
+      buf_uops_list.append(UOp.sink(*buf_uops))
+    sched_ptr += 1
+  return pre_schedule, buf_uops_list
+
 from tinygrad.engine.memory import memory_planner
 from tinygrad.schedule.rangeify import get_rangeify_map
 from tinygrad.schedule.multi import get_multi_map

 def replace_input_buffer(ctx:tuple[dict[UOp, UOp], dict[str, int]], b:UOp):
  if (ret:=ctx[0].get(b, None)) is None:
-    if b.op is Ops.BUFFER:
-      ctx[0][b] = ret = b.replace(src=(UOp(Ops.LUNIQUE, arg=len(ctx[0])), b.src[1]))
-    else:
-      # TODO: flip args in CONST
-      assert b.op is Ops.CONST
-      ctx[0][b] = ret = b.replace(src=(b.src[0], UOp(Ops.LUNIQUE, arg=len(ctx[0]))))
+    # both BUFFER and CONST have src=(UNIQUE, DEVICE), replace UNIQUE with LUNIQUE
+    ctx[0][b] = ret = b.replace(src=(UOp(Ops.LUNIQUE, arg=len(ctx[0])), b.src[1]))
  return ret

 def strip_bind(ctx:tuple[dict[UOp, UOp], dict[str, int]], b:UOp):
@@ -110,10 +110,8 @@ def strip_bind(ctx:tuple[dict[UOp, UOp], dict[str, int]], b:UOp):
  return ctx[0].setdefault(b, b.replace(src=(b.src[0],)))

 pm_pre_sched_cache = PatternMatcher([
-  # replace input buffers
-  (UPat(Ops.BUFFER, src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE)), name="b"), replace_input_buffer),
-  # remove unique consts
-  (UPat(Ops.CONST, src=(UPat(Ops.DEVICE), UPat(Ops.UNIQUE)), name="b"), replace_input_buffer),
+  # replace UNIQUE with LUNIQUE for cache key normalization
+  (UPat((Ops.BUFFER, Ops.CONST), src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE)), name="b"), replace_input_buffer),
  # strip value from BIND for cache key normalization, so different values hit same cache
  (UPat(Ops.BIND, src=(UPat(Ops.DEFINE_VAR), UPat(Ops.CONST)), name="b"), strip_bind),
 ])
@@ -126,8 +124,8 @@ def replace_input_buffer_back(ctx:dict[UOp, UOp], b:UOp):
  return ret

 pm_post_sched_cache = PatternMatcher([
-  (UPat(Ops.BUFFER, src=(UPat(Ops.LUNIQUE), UPat(Ops.DEVICE)), name="b"), replace_input_buffer_back),
-  (UPat(Ops.CONST, src=(UPat(Ops.DEVICE), UPat(Ops.LUNIQUE)), name="b"), replace_input_buffer_back),
+  # restore LUNIQUE back to UNIQUE
+  (UPat((Ops.BUFFER, Ops.CONST), src=(UPat(Ops.LUNIQUE), UPat(Ops.DEVICE)), name="b"), replace_input_buffer_back),
  # restore BIND value stripped in pm_pre_sched_cache
  (UPat(Ops.BIND, src=(UPat(Ops.DEFINE_VAR),), name="b"), lambda ctx,b: ctx.get(b)),
 ])
@@ -175,8 +173,8 @@ def complete_create_schedule_with_vars(big_sink:UOp) -> tuple[dict[UOp, UOp], li
    pre_schedule, combined_sink = sc_ret

  # replace all the LUNIQUEs with UNIQUEs (single graph_rewrite for everything)
-  input_buffers_reverse = {v:k for k,v in input_buffers.items()}
-  combined = graph_rewrite(combined_sink, pm_post_sched_cache, ctx=input_buffers_reverse, name="unrewrite combined")
+  input_buffers_inverse = {v:k for k,v in input_buffers.items()}
+  combined = graph_rewrite(combined_sink, pm_post_sched_cache, ctx=input_buffers_inverse, name="unrewrite combined")
  tensor_map_sink, buf_uops_sink = combined.src
  tm_src = tensor_map_sink.src
  tensor_map = {tm_src[i]:tm_src[i+1] for i in range(0, len(tm_src), 2)}
@@ -205,4 +203,6 @@ def complete_create_schedule_with_vars(big_sink:UOp) -> tuple[dict[UOp, UOp], li
    print(f"scheduled {len(schedule):4d} kernels in {(time.perf_counter()-st)*1000:8.2f} ms"+\
          f" | {' cache hit' if SCACHE and sc_ret is not None else 'CACHE MISS'} {sched_cache_key.hex()[:8]}"+\
          f" | {len(UOpMetaClass.ucache)} uops in cache")
-  return tensor_map, schedule, var_vals if schedule else {}
+
+  used_vars = set().union(*[{v.arg[0] for v in si.ast.variables()} for si in schedule])
+  return tensor_map, schedule, {k:v for k,v in var_vals.items() if k in used_vars}
@@ -44,6 +44,8 @@ pm_gradient = PatternMatcher([
  # NOTE: this is only correct when the KERNEL has a single output
  (UPat(Ops.AFTER), lambda ctx: (ctx, ctx)),
  (UPat(Ops.CUSTOM_KERNEL, name="k"), lambda ctx, k: k.arg.grad_fxn(ctx, k)),
+  # gradient on CALL is a custom function
+  (UPat(Ops.CALL, name="k"), lambda ctx, k: (None,)+k.arg(ctx, k)),
  # there's no gradient for bitcast
  (UPat(Ops.BITCAST), lambda: (None,)),
 ])
@@ -165,6 +165,9 @@ class ContextVar(Generic[T]):
  def __ge__(self, x): return self.value >= x
  def __gt__(self, x): return self.value > x
  def __lt__(self, x): return self.value < x
+  def tolist(self, obj=None):
+    assert isinstance(self.value, str)
+    return [getattr(obj, x) if obj else x for x in self.value.split(',') if x]

 DEBUG, IMAGE, BEAM, NOOPT = ContextVar("DEBUG", 0), ContextVar("IMAGE", 0), ContextVar("BEAM", 0), ContextVar("NOOPT", 0)
 JIT, JIT_BATCH_SIZE = ContextVar("JIT", 2 if OSX and ARCH_X86 else 1), ContextVar("JIT_BATCH_SIZE", 32)
@@ -177,7 +180,7 @@ CACHELEVEL, IGNORE_BEAM_CACHE, DEVECTORIZE = ContextVar("CACHELEVEL", 2), Contex
 VALIDATE_WITH_CPU, DISABLE_FAST_IDIV = ContextVar("VALIDATE_WITH_CPU", 0), ContextVar("DISABLE_FAST_IDIV", 0)
 CORRECT_DIVMOD_FOLDING, FUSE_OPTIM = ContextVar("CORRECT_DIVMOD_FOLDING", 0), ContextVar("FUSE_OPTIM", 0)
 ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE = ContextVar("ALLOW_DEVICE_USAGE", 1), ContextVar("MAX_BUFFER_SIZE", 0)
-EMULATE = ContextVar("EMULATE", "")
+EMULATE, EMULATED_DTYPES = ContextVar("EMULATE", ""), ContextVar("EMULATED_DTYPES", "")
 CPU_COUNT = ContextVar("CPU_COUNT", max(1, len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1)))
 # Compilers
 CPU_LLVM, CPU_LVP, AMD_LLVM = ContextVar("CPU_LLVM", 0), ContextVar("CPU_LVP", 0), ContextVar("AMD_LLVM", 0)
@@ -201,6 +204,10 @@ CCACHE = ContextVar("CCACHE", 1)
 ALLOW_TF32 = ContextVar("ALLOW_TF32", 0)
 # set to 0 to disable the scheduler cache
 SCACHE = ContextVar("SCACHE", 1)
+# allow use of atomics for embedding backward
+USE_ATOMICS = ContextVar("USE_ATOMICS", 0)
+# allow use of assembly for gemm
+ASM_GEMM = ContextVar("ASM_GEMM", 0)

@dataclass(frozen=True)
 class Metadata:
@@ -19,22 +19,22 @@ class MathMixin:
  def _binop(self, op: Ops, x: Self | ConstType, reverse: bool) -> Self:
    return self.ufix(x).alu(op, self) if reverse else self.alu(op, self.ufix(x))

-  def logical_not(self):
+  def logical_not(self) -> Self:
    return self.ne(True)

-  def neg(self):
+  def neg(self) -> Self:
    if (dtype := getattr(self, "dtype")) is None:
      raise TypeError(f"MathTraits __neg__ requires a dtype, {self=}")
    return self.logical_not() if dtype.scalar() == dtypes.bool else self * (-1)

-  def _check_dtype(self):
+  def _check_dtype(self) -> None:
    if (dtype := getattr(self, "dtype")) is not None:
      if isinstance(dtype, tuple):
        dtype = dtype[0]
      if not (dtypes.is_bool(dtype) or dtypes.is_int(dtype)):
        raise RuntimeError(f"{dtype} is not supported")

-  def add(self, x: Self | ConstType, reverse: bool = False):
+  def add(self, x: Self | ConstType, reverse: bool = False) -> Self:
    """
    Adds `self` and `x`.
    Equivalent to `self + x`.
@@ -53,7 +53,7 @@ class MathMixin:
    """
    return self._binop(Ops.ADD, x, reverse)

-  def mul(self, x: Self | ConstType, reverse: bool = False):
+  def mul(self, x: Self | ConstType, reverse: bool = False) -> Self:
    """
    Multiplies `self` and `x`.
    Equivalent to `self * x`.
@@ -73,7 +73,7 @@ class MathMixin:
    """
    return self._binop(Ops.MUL, x, reverse)

-  def bitwise_and(self, x: Self | ConstType, reverse: bool = False):
+  def bitwise_and(self, x: Self | ConstType, reverse: bool = False) -> Self:
    """
    Computes the bitwise AND of `self` and `x`.
    Equivalent to `self & x`.
@@ -88,7 +88,7 @@ class MathMixin:
    self._check_dtype()
    return self._binop(Ops.AND, x, reverse)

-  def bitwise_or(self, x: Self | ConstType, reverse: bool = False):
+  def bitwise_or(self, x: Self | ConstType, reverse: bool = False) -> Self:
    """
    Computes the bitwise OR of `self` and `x`.
    Equivalent to `self | x`.
@@ -103,7 +103,7 @@ class MathMixin:
    self._check_dtype()
    return self._binop(Ops.OR, x, reverse)

-  def bitwise_xor(self, x: Self | ConstType, reverse: bool = False):
+  def bitwise_xor(self, x: Self | ConstType, reverse: bool = False) -> Self:
    """
    Computes bitwise xor of `self` and `x`.
    Equivalent to `self ^ x`.
@@ -119,7 +119,7 @@ class MathMixin:
    self._check_dtype()
    return self._binop(Ops.XOR, x, reverse)

-  def idiv(self, x: Self | ConstType, reverse: bool = False):
+  def idiv(self, x: Self | ConstType, reverse: bool = False) -> Self:
    """
    Divides `self` by `x`.
    Equivalent to `self // x`.
@@ -132,133 +132,133 @@ class MathMixin:
    """
    return self._binop(Ops.IDIV, x, reverse)

-  def mod(self, x: Self | ConstType, reverse: bool = False):
+  def mod(self, x: Self | ConstType, reverse: bool = False) -> Self:
    return self._binop(Ops.MOD, x, reverse)

-  def sub(self, x: Self | ConstType, reverse: bool = False):
+  def sub(self, x: Self | ConstType, reverse: bool = False) -> Self:
    return self.ufix(x).alu(Ops.ADD, -self) if reverse else self.alu(Ops.ADD, -self.ufix(x))

-  def div(self, x: Self | ConstType, reverse: bool = False):
+  def div(self, x: Self | ConstType, reverse: bool = False) -> Self:
    return (self.ufix(x) * self.alu(Ops.RECIPROCAL)) if reverse else (self * self.ufix(x).alu(Ops.RECIPROCAL))

-  def __neg__(self):
+  def __neg__(self) -> Self:
    return self.neg()

-  def __add__(self, x: Self | ConstType):
+  def __add__(self, x: Self | ConstType) -> Self:
    return self.add(x)

-  def __sub__(self, x: Self | ConstType):
+  def __sub__(self, x: Self | ConstType) -> Self:
    return self.sub(x)

-  def __mul__(self, x: Self | ConstType):
+  def __mul__(self, x: Self | ConstType) -> Self:
    return self.mul(x)

-  def __truediv__(self, x: Self | ConstType):
+  def __truediv__(self, x: Self | ConstType) -> Self:
    return self.div(x)

-  def __floordiv__(self, x: Self | ConstType):
+  def __floordiv__(self, x: Self | ConstType) -> Self:
    return self.idiv(x)  # TODO: idiv is trunc div, not floordiv

-  def __mod__(self, x: Self | ConstType):
+  def __mod__(self, x: Self | ConstType) -> Self:
    return self.mod(x)

-  def __and__(self, x: Self | ConstType):
+  def __and__(self, x: Self | ConstType) -> Self:
    return self.bitwise_and(x)

-  def __or__(self, x: Self | ConstType):
+  def __or__(self, x: Self | ConstType) -> Self:
    return self.bitwise_or(x)

-  def __xor__(self, x: Self | ConstType):
+  def __xor__(self, x: Self | ConstType) -> Self:
    return self.bitwise_xor(x)

-  def __radd__(self, x: Self | ConstType):
+  def __radd__(self, x: Self | ConstType) -> Self:
    return self.add(x, True)

-  def __rsub__(self, x: Self | ConstType):
+  def __rsub__(self, x: Self | ConstType) -> Self:
    return self.sub(x, True)

-  def __rmul__(self, x: Self | ConstType):
+  def __rmul__(self, x: Self | ConstType) -> Self:
    return self.mul(x, True)

-  def __rtruediv__(self, x: Self | ConstType):
+  def __rtruediv__(self, x: Self | ConstType) -> Self:
    return self.div(x, True)

-  def __rfloordiv__(self, x: Self | ConstType):
+  def __rfloordiv__(self, x: Self | ConstType) -> Self:
    return self.idiv(x, True)

-  def __rand__(self, x: Self | ConstType):
+  def __rand__(self, x: Self | ConstType) -> Self:
    return self.bitwise_and(x, True)

-  def __ror__(self, x: Self | ConstType):
+  def __ror__(self, x: Self | ConstType) -> Self:
    return self.bitwise_or(x, True)

-  def __rxor__(self, x: Self | ConstType):
+  def __rxor__(self, x: Self | ConstType) -> Self:
    return self.bitwise_xor(x, True)

-  def __rmod__(self, x: Self | ConstType):
+  def __rmod__(self, x: Self | ConstType) -> Self:
    return self.mod(x, True)

-  def __lt__(self, x: Self | ConstType):
+  def __lt__(self, x: Self | ConstType) -> Self:
    return self.alu(Ops.CMPLT, self.ufix(x))

-  def __gt__(self, x: Self | ConstType):
+  def __gt__(self, x: Self | ConstType) -> Self:
    return self.ufix(x).alu(Ops.CMPLT, self)

-  def __ge__(self, x: Self | ConstType):
+  def __ge__(self, x: Self | ConstType) -> Self:
    return (self < x).logical_not()

-  def __le__(self, x: Self | ConstType):
+  def __le__(self, x: Self | ConstType) -> Self:
    return (self > x).logical_not()

-  def ne(self, x: Self | ConstType):
+  def ne(self, x: Self | ConstType) -> Self:
    return self.alu(Ops.CMPNE, self.ufix(x))

-  def eq(self, x: Self | ConstType):
+  def eq(self, x: Self | ConstType) -> Self:
    return self.ne(x).logical_not()

-  def __ne__(self, x: Self | ConstType):  # type: ignore[override]
+  def __ne__(self, x: Self | ConstType) -> Self:  # type: ignore[override]
    return self.ne(x)

  # NOTE: __eq__ isn't overridden, and means the same thing as is by default

-  def lshift(self, x: Self | int, reverse: bool = False):
+  def lshift(self, x: Self | int, reverse: bool = False) -> Self:
    return self._binop(Ops.SHL, x, reverse)

-  def rshift(self, x: Self | int, reverse: bool = False):
+  def rshift(self, x: Self | int, reverse: bool = False) -> Self:
    return self._binop(Ops.SHR, x, reverse)

-  def __lshift__(self, x: Self | int):
+  def __lshift__(self, x: Self | int) -> Self:
    return self.lshift(x)

-  def __rshift__(self, x: Self | int):
+  def __rshift__(self, x: Self | int) -> Self:
    return self.rshift(x)

-  def __rlshift__(self, x: Self | int):
+  def __rlshift__(self, x: Self | int) -> Self:
    return self.lshift(x, True)

-  def __rrshift__(self, x: Self | int):
+  def __rrshift__(self, x: Self | int) -> Self:
    return self.rshift(x, True)

-  def maximum(self, x: Self | ConstType):
+  def maximum(self, x: Self | ConstType) -> Self:
    return self.alu(Ops.MAX, self.ufix(x))

-  def minimum(self, x: Self | ConstType):
+  def minimum(self, x: Self | ConstType) -> Self:
    return -(-self).maximum(-self.ufix(x))

-  def where(self, x: Self | ConstType, y: Self | ConstType):
+  def where(self, x: Self | ConstType, y: Self | ConstType) -> Self:
    if isinstance(x, type(self)):
      return self.alu(Ops.WHERE, x, x.ufix(y))
    if isinstance(y, type(self)):
      return self.alu(Ops.WHERE, y.ufix(x), y)
    raise RuntimeError("where needs at least one UOp arg")

-  def threefry(self, seed: Self):
+  def threefry(self, seed: Self) -> Self:
    return self.alu(Ops.THREEFRY, seed)

-  def reciprocal(self):
+  def reciprocal(self) -> Self:
    return self.alu(Ops.RECIPROCAL)

-  def trunc(self):
+  def trunc(self) -> Self:
    """
    Truncates the tensor element-wise.

@@ -268,25 +268,25 @@ class MathMixin:
    """
    return self.alu(Ops.TRUNC)

-  def sqrt(self):
+  def sqrt(self) -> Self:
    return self.alu(Ops.SQRT)

-  def sin(self):
+  def sin(self) -> Self:
    return self.alu(Ops.SIN)

-  def log2(self):
+  def log2(self) -> Self:
    return self.alu(Ops.LOG2)

-  def exp2(self):
+  def exp2(self) -> Self:
    return self.alu(Ops.EXP2)

-  def pow(self, x: Self | ConstType):
+  def pow(self, x: Self | ConstType) -> Self:
    return self.alu(Ops.POW, self.ufix(x))

-  def __pow__(self, x: Self | ConstType):
+  def __pow__(self, x: Self | ConstType) -> Self:
    return self.pow(x)

-  def square(self):
+  def square(self) -> Self:
    """
    Squares the tensor element-wise.
    Equivalent to `self*self`.
@@ -297,7 +297,7 @@ class MathMixin:
    """
    return self * self

-  def clamp(self, min_=None, max_=None):
+  def clamp(self, min_=None, max_=None) -> Self:
    """
    Clips (clamps) the values in the tensor between `min_` and `max_` element-wise.
    If `min_` is `None`, there is no lower bound. If `max_` is None, there is no upper bound.
@@ -310,11 +310,11 @@ class MathMixin:
    ret = (self < min_).where(min_, self) if min_ is not None else self
    return (ret > max_).where(max_, ret) if max_ is not None else ret

-  def clip(self, min_=None, max_=None):
+  def clip(self, min_=None, max_=None) -> Self:
    """Alias for `Tensor.clamp`."""
    return self.clamp(min_, max_)

-  def isnan(self):
+  def isnan(self) -> Self:
    """
    Checks the tensor element-wise to return True where the element is NaN, otherwise returns False

@@ -324,7 +324,7 @@ class MathMixin:
    """
    return self != self

-  def isinf(self, detect_positive: bool = True, detect_negative: bool = True):
+  def isinf(self, detect_positive: bool = True, detect_negative: bool = True) -> Self:
    """
    Checks the tensor element-wise to return True where the element is infinity, otherwise returns False

@@ -334,7 +334,7 @@ class MathMixin:
    """
    return self.eq(float("inf")) * detect_positive + self.eq(float("-inf")) * detect_negative

-  def isfinite(self):
+  def isfinite(self) -> Self:
    """
    Checks the tensor element-wise to return True where the element is finite, otherwise returns False

@@ -344,7 +344,7 @@ class MathMixin:
    """
    return (self.isinf() | self.isnan()).logical_not()

-  def ceil(self):
+  def ceil(self) -> Self:
    """
    Rounds the tensor element-wise towards positive infinity.

@@ -354,7 +354,7 @@ class MathMixin:
    """
    return (self > (b := self.trunc())).where(b+1, b)

-  def floor(self):
+  def floor(self) -> Self:
    """
    Rounds the tensor element-wise towards negative infinity.

@@ -364,7 +364,7 @@ class MathMixin:
    """
    return (self < (b := self.trunc())).where(b-1, b)

-  def relu(self):
+  def relu(self) -> Self:
    """
    Applies the Rectified Linear Unit (ReLU) function element-wise.

@@ -375,7 +375,7 @@ class MathMixin:
    # NOTE: if you write this as self.maximum(0) the gradient is wrong, passing through half when self is 0
    return (self > 0).where(self, 0)

-  def sigmoid(self):
+  def sigmoid(self) -> Self:
    """
    Applies the Sigmoid function element-wise.

@@ -387,7 +387,7 @@ class MathMixin:
    """
    return (1 + (self * (-1/math.log(2))).exp2()).reciprocal()

-  def relu6(self):
+  def relu6(self) -> Self:
    """
    Applies the ReLU6 function element-wise.

@@ -399,7 +399,7 @@ class MathMixin:
    """
    return self.relu() - (self-6).relu()

-  def hardswish(self):
+  def hardswish(self) -> Self:
    """
    Applies the Hardswish function element-wise.

@@ -411,7 +411,7 @@ class MathMixin:
    """
    return self * (self+3).relu6() * (1/6)

-  def hardsigmoid(self, alpha: float = 1/6, beta: float = 0.5):
+  def hardsigmoid(self, alpha: float = 1/6, beta: float = 0.5) -> Self:
    """
    Applies the Hardsigmoid function element-wise.
    NOTE: default `alpha` and `beta` values are taken from torch
@@ -424,7 +424,7 @@ class MathMixin:
    """
    return (alpha * self + beta).relu() - (alpha * self + beta - 1).relu()

-  def hardtanh(self, min_val=-1, max_val=1):
+  def hardtanh(self, min_val=-1, max_val=1) -> Self:
    """
    Applies the Hardtanh function element-wise.

@@ -434,7 +434,7 @@ class MathMixin:
    """
    return self.clip(min_val, max_val)

-  def leaky_relu(self, neg_slope=0.01):
+  def leaky_relu(self, neg_slope=0.01) -> Self:
    """
    Applies the Leaky ReLU function element-wise.

@@ -447,7 +447,7 @@ class MathMixin:
    """
    return (self < 0).where(neg_slope*self, self)

-  def tanh(self):
+  def tanh(self) -> Self:
    """
    Applies the Hyperbolic Tangent (tanh) function element-wise.

@@ -459,7 +459,7 @@ class MathMixin:
    """
    return 2.0 * ((2.0 * self).sigmoid()) - 1.0

-  def quick_gelu(self):
+  def quick_gelu(self) -> Self:
    """
    Applies the Sigmoid GELU approximation element-wise.

@@ -469,7 +469,7 @@ class MathMixin:
    """
    return self * (self * 1.702).sigmoid()

-  def gelu(self):
+  def gelu(self) -> Self:
    """
    Applies the Gaussian Error Linear Unit (GELU) function element-wise.

@@ -481,7 +481,7 @@ class MathMixin:
    """
    return 0.5 * self * (1 + (math.sqrt(2 / math.pi) * (self + 0.044715 * self ** 3)).tanh())

-  def swish(self):
+  def swish(self) -> Self:
    """
    See `.silu()`

@@ -493,7 +493,7 @@ class MathMixin:
    """
    return self * self.sigmoid()

-  def silu(self):
+  def silu(self) -> Self:
    """
    Applies the Sigmoid Linear Unit (SiLU) function element-wise.

@@ -505,7 +505,7 @@ class MathMixin:
    """
    return self.swish()  # The SiLU function is also known as the swish function.

-  def rsqrt(self):
+  def rsqrt(self) -> Self:
    """
    Computes the reciprocal of the square root of the tensor element-wise.

@@ -1,12 +1,12 @@
 # mixins add syntactic sugar to Tensor and UOp
-from typing import TypeAlias, TYPE_CHECKING, Self
+from __future__ import annotations
+from typing import TYPE_CHECKING, Self
 from tinygrad.uop import Ops
 from tinygrad.helpers import prod, argfix, flatten, dedup, make_tuple, ceildiv
 from tinygrad.uop.ops import resolve, smax

 if TYPE_CHECKING:
-  from tinygrad.uop.ops import UOp
-sint: TypeAlias = "UOp | int"
+  from tinygrad.uop.ops import sint


 def _align_left(*shapes: tuple[sint, ...]) -> tuple[tuple[sint, ...], ...]:
@@ -3,7 +3,7 @@ import math
 from tinygrad.tensor import Tensor
 from tinygrad.dtype import dtypes
 from tinygrad.device import is_dtype_supported
-from tinygrad.helpers import prod, make_tuple, flatten
+from tinygrad.helpers import prod, make_tuple, flatten, USE_ATOMICS
 from tinygrad.nn import optim, state, datasets  # noqa: F401

 class BatchNorm:
@@ -304,6 +304,46 @@ class RMSNorm:
    x = self._norm(x.float()).cast(x.dtype)
    return x if self.weight is None else x * self.weight

+from tinygrad.uop.ops import UOp, KernelInfo, Ops
+def _embedding_bwd(grad_emb:UOp, call:UOp) -> tuple:
+  weight, idx = call.src[1:]
+  # for multi-device: unshard inputs to one device
+  if isinstance(weight.device, tuple):
+    assert weight.axis is None, "sharded weights on Embedding not supported with USE_ATOMICS"
+    grad_emb = grad_emb.copy_to_device(weight.device)
+    idx = idx.copy_to_device(weight.device)
+  # weight is replicated, grad_weight should match
+  grad_weight_uop = Tensor.empty(weight.shape, dtype=dtypes.float, device=weight.device).uop
+
+  # TODO: how do we remove this dumb kernel and use Tensor.zeros?
+  def _zero_kernel(out:UOp) -> UOp:
+    i = UOp.range(out.size, 0)
+    return out.flatten()[i].store(0).end(i).sink(arg=KernelInfo(name="zero"))
+  grad_weight_uop = grad_weight_uop.custom_kernel(fxn=_zero_kernel)[0]
+
+  # TODO: do we have a universal helper for this?
+  device = call.device.split(":")[0] if not isinstance(call.device, tuple) else call.device[0].split(":")[0]
+
+  # this is the real atomic kernel
+  def _embedding_bwd_kernel(grad_weight:UOp, grad_emb:UOp, idx:UOp) -> UOp:
+    idx_flat, grad_emb_flat = idx.flatten(), grad_emb.reshape((idx.size, grad_weight.shape[-1]))
+    i = UOp.range(grad_emb_flat.shape[0], 0)  # batch_size * sequence_length
+    j = UOp.range(grad_emb_flat.shape[1], 1)  # embed_size
+    token_id = idx_flat[i].clip(0, grad_weight.shape[0]-1).cast(dtypes.index)
+    # atomic scatter-add: grad_weight[token_id, j] += grad_emb_flat[i, j]
+    if device in ("CPU", "NULL"): atomic_arg = "__atomic_fetch_add({0}, {1}, __ATOMIC_RELAXED);"
+    elif device == "AMD": atomic_arg = "__hip_atomic_fetch_add({0}, {1}, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);"
+    else: raise NotImplementedError(f"no atomics for device {device}")
+    atomic = UOp(Ops.CUSTOM, dtypes.void, (grad_weight.index(token_id, j, ptr=True), grad_emb_flat[i, j].cast(dtypes.float)), arg = atomic_arg)
+    return atomic.end(i, j).sink(arg=KernelInfo(name="embedding_bwd", opts_to_apply=()))
+  grad_weight_uop = grad_weight_uop.custom_kernel(grad_emb, idx, fxn=_embedding_bwd_kernel)[0]
+
+  return (grad_weight_uop.cast(weight.dtype), None)
+
+def _embedding_fwd(weight:Tensor, idx:Tensor) -> Tensor:
+  arange = Tensor.arange(weight.shape[0], requires_grad=False, device=weight.device)
+  return (arange == idx.unsqueeze(-1)).unsqueeze(-1).where(weight, 0).sum(-2, dtype=weight.dtype)
+
 class Embedding:
  """
  A simple lookup table that stores embeddings of a fixed dictionary and size.
@@ -316,14 +356,12 @@ class Embedding:
  ```
  """
  def __init__(self, vocab_size:int, embed_size:int):
-    self.vocab_sz, self.embed_sz, self.weight = vocab_size, embed_size, Tensor.glorot_uniform(vocab_size, embed_size)
+    self.weight = Tensor.glorot_uniform(vocab_size, embed_size)

  def __call__(self, idx:Tensor) -> Tensor:
-    if not hasattr(self, 'arange'): self.arange = Tensor.arange(self.vocab_sz, requires_grad=False, device=self.weight.device).unsqueeze(-1)
    if not dtypes.is_int(idx.dtype): raise TypeError(f"Expected integer dtype for index in embedding, got {idx.dtype}")
-    big_shp = idx.shape+(self.vocab_sz, self.embed_sz)
-    arange, idx, vals = self.arange.expand(big_shp), idx.reshape(idx.shape+(1, 1)).expand(big_shp), self.weight.expand(big_shp)
-    return (arange == idx).where(vals, 0).sum(-2, dtype=vals.dtype)
+    if USE_ATOMICS: return Tensor.call(self.weight, idx, fxn=_embedding_fwd(self.weight.as_param(0), idx.as_param(1)), grad_fxn=_embedding_bwd)
+    return _embedding_fwd(self.weight, idx)

 class LSTMCell:
  """
@@ -1,12 +1,12 @@
 # pylint: disable=possibly-unused-variable
 from typing import Any, Sequence, cast, Literal, NamedTuple, Generator
 import dataclasses, functools, io, math, types, warnings, pathlib, sys, os, struct, enum
-from io import BufferedReader
 from tinygrad.nn.state import TensorIO
 from tinygrad.tensor import Tensor, _broadcast_shape, ReductionStr
 from tinygrad.helpers import getenv, all_same, prod, flatten, make_tuple, argsort, is_numpy_ndarray, get_single_element, polyN
 from tinygrad.dtype import DType, ConstType, dtypes, _from_np_dtype, truncate, least_upper_dtype, DTYPES_DICT
 from tinygrad.device import is_dtype_supported, Device
+from tinygrad.uop.ops import sint

 # ***** protobuf definitions ******
 class WireType(enum.IntEnum):
@@ -75,7 +75,7 @@ class OnnxNode:
  opts: dict[str, Any]

 # ***** protobuf parsing ******
-class PBBufferedReader(BufferedReader):
+class PBBufferedReader(io.BufferedReader):
  def __init__(self, tensor: Tensor):
    assert tensor.dtype == dtypes.uint8, tensor
    super().__init__(TensorIO(tensor))
@@ -677,7 +677,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
  def ReduceLogSumExp(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
    return ReduceSum(data.exp(), axes, keepdims, noop_with_empty_axes).log()
  def ArgMax(x:Tensor, axis:int=0, keepdims:int=1, select_last_index:int=0):
-    if select_last_index: return ((x.shape[axis]-1) - x.flip(axis).argmax(axis, keepdim=keepdims)).cast(dtypes.int64)
+    if select_last_index: return ((int(x.shape[axis])-1) - x.flip(axis).argmax(axis, keepdim=keepdims)).cast(dtypes.int64)
    return x.argmax(axis, keepdim=keepdims).cast(dtypes.int64)
  def ArgMin(x, axis:int=0, keepdims:int=1, select_last_index:int=0):
    return ArgMax(-x, axis=axis, keepdims=keepdims, select_last_index=select_last_index)
@@ -704,7 +704,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
    return data[tuple(slices)]

  def Split(data:Tensor, split:list[int]|None=None, num_outputs:int=0, axis:int=0):
-    sz = data.shape[axis]
+    sz = int(data.shape[axis])
    if split is None: split = [sz // num_outputs + (1 if i < sz % num_outputs else 0) for i in range(num_outputs)]
    return data.split(split, axis)

@@ -717,8 +717,8 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
    return x.pad(padding=_onnx_pads_to_tiny_pads(real_pads), mode={"edge":"replicate", "wrap":"circular"}.get(mode, mode), value=value)

  def CenterCropPad(t:Tensor, shape:list[int], axes:list[int]|None=None):
-    shrink_arg:list[None|tuple[int,int]] = [None] * t.ndim
-    pad_arg:list[None|tuple[int,int]] = [None] * t.ndim
+    shrink_arg:list[None|tuple[sint,sint]] = [None] * t.ndim
+    pad_arg:list[None|tuple[sint,sint]] = [None] * t.ndim
    for s, x in zip(shape, axes or range(t.ndim)):
      tx = t.shape[x]
      if s < tx: shrink_arg[x] = (tx//2 - (s+1)//2, tx//2 + s//2)
@@ -752,8 +752,8 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
      pads = _auto_pad([s_*(i-1) + op_ + ((k_-1)*d_+1) - os for s_,i,op_,k_,d_,os in
                        zip(strides_, input_shape_, output_padding_, kernel_shape_, dilations_, output_shape)], auto_pad)
    if pads is None: # we generate pads
-      output_shape = output_shape or [X.shape[i+2] * strides_[i] for i in range(len(strides_))]
-      pads = [strides_[i]*(input_shape_[i]-1)+output_padding_[i]+((kernel_shape_[i]-1)*dilations_[i]+1)-output_shape[i]
+      output_shape = output_shape or [int(X.shape[i+2]) * strides_[i] for i in range(len(strides_))]
+      pads = [int(strides_[i]*(input_shape_[i]-1)+output_padding_[i]+((kernel_shape_[i]-1)*dilations_[i]+1)-output_shape[i])
              for i in range(len(input_shape_))]
      pads = _auto_pad(pads, auto_pad) if auto_pad != "NOTSET" else [0] * len(input_shape_) * 2
    pads = _onnx_pads_to_tiny_pads(pads)
@@ -1015,7 +1015,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
                        num_heads:int|None=None, past_present_share_buffer:int|None=None, qkv_hidden_sizes:list[int]|None=None,
                        rotary_embedding_dim:int|None=None, scale:float|None=None, unidirectional:int=0):
    assert not do_rotary and not attention_bias, "TODO"
-    if qkv_hidden_sizes is None: qkv_hidden_sizes = [weights.shape[1] // 3] * 3
+    if qkv_hidden_sizes is None: qkv_hidden_sizes = [int(weights.shape[1] // 3)] * 3
    qkv = x.linear(weights, bias)
    q, k, v = qkv.split(qkv_hidden_sizes, dim=2)

@@ -1113,7 +1113,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
    if X.ndim == 4: X = X.permute(0, 2, 1, 3)
    elif X.ndim == 3:
      assert num_heads is not None, "num_heads must be provided for 3D input"
-      X = X.unflatten(-1, (num_heads, X.shape[-1] // num_heads))
+      X = X.unflatten(-1, (num_heads, int(X.shape[-1]) // num_heads))

    head_size = cast(int, X.shape[-1])
    rot_dim = rotary_embedding_dim or head_size
@@ -1183,7 +1183,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT

  def TensorScatter(data: Tensor, updates: Tensor, indices: Tensor, mode: str = 'default'):
    # scatter updates along axis -2 at positions given by indices, for each batch
-    B, U, D = indices.shape[0], updates.shape[-2], data.shape[-2]
+    B, U, D = indices.shape[0], updates.shape[-2], int(data.shape[-2])
    orig_shape, data_flat, updates_flat = data.shape, data.reshape(-1, D, data.shape[-1]), updates.reshape(-1, U, updates.shape[-1])
    B_total = data_flat.shape[0]
    batch_idx = Tensor.arange(B_total, device=data.device).reshape(B_total, 1).expand(B_total, U)
@@ -78,7 +78,7 @@ def safe_save(tensors:dict[str, Tensor], fn:str, metadata:dict[str, Any]|None=No
  j += "\x20"*(round_up(len(j),8)-len(j))
  pathlib.Path(fn).unlink(missing_ok=True)
  t = Tensor.empty(8+len(j)+offset, dtype=dtypes.uint8, device=f"disk:{fn}")
-  t[0:8].bitcast(dtypes.int64).assign([len(j)])
+  t[0:8].assign(Tensor([len(j)], dtype=dtypes.int64, device="CPU").bitcast(dtypes.uint8))
  t[8:8+len(j)].assign(list(j.encode('utf-8')))
  for k,v in safe_load(t).items(): v.assign(tensors[k])

@@ -46,10 +46,10 @@ base_rewrite = PatternMatcher([
  # new load/store
  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var('idx')), allow_any_len=True),
   lambda ctx,buf,idx: f"({ctx[buf]}+{strip_parens(ctx[idx]) if idx.arg == Ops.ADD else ctx[idx]})"),
-  (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat.var("gate"))).or_casted("bidx"), UPat.var("var")), allow_any_len=True),
+  (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat.var("gate"))).or_casted("bidx"), UPat.var("var"))),
   lambda ctx,bidx,var,gate: f"({ctx[gate]}?*{ctx[bidx]}:{ctx[var]})"),
-  (UPat(Ops.LOAD, src=(UPat.var('bidx'),), allow_any_len=True), lambda ctx,bidx: f"(*{ctx[bidx]})"),
-  (UPat(Ops.STORE, src=(UPat.var('bidx'), UPat.var("var")), allow_any_len=True), lambda ctx,bidx,var: f"*{ctx[bidx]} = {ctx[var]};"),
+  (UPat(Ops.LOAD, src=(UPat.var('bidx'),)), lambda ctx,bidx: f"(*{ctx[bidx]})"),
+  (UPat(Ops.STORE, src=(UPat.var('bidx'), UPat.var("var"))), lambda ctx,bidx,var: f"*{ctx[bidx]} = {ctx[var]};"),
  # alu/gep
  # TODO: look for left-associative
  (UPat(GroupOp.ALU, name="x"), lambda ctx,x: ctx.code_for_op[x.op](
@@ -87,7 +87,8 @@ def create_non_native_float_pats(dts:tuple[DType, ...], casting:bool=True):
 def cast_float_to_bf16(x: UOp) -> UOp:
  assert x.dtype == dtypes.float, "cast float -> bf16 must start with float"
  x = x.bitcast(dtypes.uint)
-  x = ((-x & 0x7f800000) != 0).where(x + ((x >> 16) & 1) + 0x7fff, ((x & 0xffff) != 0).where((x | 0x10000), x))
+  # NOTE: != returns UOp, not bool, issue with mypy
+  x = ((-x & 0x7f800000) != 0).where(x + ((x >> 16) & 1) + 0x7fff, ((x & 0xffff) != 0).where((x | 0x10000), x))  # type: ignore[comparison-overlap]
  return (x >> 16).cast(dtypes.ushort).bitcast(dtypes.bfloat16)

 # manual bfloat16 casting patterns (shared between LLVM, Clang, and AMD renderers to avoid compiler intrinsics)
@@ -528,7 +528,7 @@ class AMDCopyQueue(HWQueue):
    # USB devices run in single-step mode, so they can't overrun the queue.
    total_bytes = (tail_blit_dword * 4 if rem_packet_cnt == 0 else -sdma_queue.put_value % sdma_queue.ring.nbytes) + rem_packet_cnt * 4
    assert total_bytes < sdma_queue.ring.nbytes, "SDMA queue overrun"
-    while not dev.is_usb() and sdma_queue.put_value + total_bytes - sdma_queue.read_ptr > sdma_queue.ring.nbytes: pass
+    while not dev.is_usb() and sdma_queue.put_value + total_bytes - sdma_queue.read_ptr[0] > sdma_queue.ring.nbytes: pass

    start_idx = (sdma_queue.put_value % sdma_queue.ring.nbytes) // 4
    sdma_queue.ring[start_idx : start_idx + tail_blit_dword] = array.array('I', cmds[:tail_blit_dword])
@@ -640,24 +640,21 @@ class AMDAllocator(HCQAllocator['AMDDevice']):
@dataclass
 class AMDQueueDesc:
  ring: MMIOInterface
-  read_ptrs: list[MMIOInterface]
-  write_ptrs: list[MMIOInterface]
-  doorbells: list[MMIOInterface]
+  read_ptr: MMIOInterface
+  write_ptr: MMIOInterface
+  doorbell: MMIOInterface
  put_value: int = 0

-  @property
-  def read_ptr(self): return min(p[0] for p in self.read_ptrs)
-
  def signal_doorbell(self, dev, doorbell_value:int|None=None):
    try:
-      for write_ptr in self.write_ptrs: write_ptr[0] = self.put_value
+      self.write_ptr[0] = self.put_value

      # Ensure all prior writes are visible to the GPU.
      System.memory_barrier()

      # Flush hdp if queue is in dev mem.
      if dev.is_am() and not dev.is_usb(): dev.iface.dev_impl.gmc.flush_hdp()
-      for doorbell in self.doorbells: doorbell[0] = self.put_value if doorbell_value is None else doorbell_value
+      self.doorbell[0] = self.put_value if doorbell_value is None else doorbell_value
    except Exception as e:
      dev.error_state = e
      raise
@@ -753,7 +750,7 @@ class KFDIface:

  def as_dmaref(self, mem:HCQBuffer) -> DMAFdRef:
    base = mem._base if mem._base is not None else mem
-    dmaref = DMAFdRef(kfd.AMDKFD_IOC_EXPORT_DMABUF(KFDIface.kfd, handle=base.meta.handle, flags=0).dmabuf_fd, mem.va_addr-base.va_addr, mem.size)
+    dmaref = DMAFdRef(kfd.AMDKFD_IOC_EXPORT_DMABUF(KFDIface.kfd, handle=base.meta.handle, flags=0).dmabuf_fd, int(mem.va_addr-base.va_addr), mem.size)
    weakref.finalize(dmaref, os.close, dmaref.fd)
    return dmaref

@@ -776,9 +773,9 @@ class KFDIface:
      self.doorbells_base = queue.doorbell_offset & (~0x1fff) # doorbell is two pages
      self.doorbells = cast(FileIOInterface, KFDIface.kfd).mmap(0, 0x2000, mmap.PROT_READ|mmap.PROT_WRITE, mmap.MAP_SHARED, self.doorbells_base)

-    return AMDQueueDesc(ring=MMIOInterface(ring.va_addr, ring.size, fmt='I'), read_ptrs=[MMIOInterface(queue.read_pointer_address, 8, fmt='Q')],
-                        write_ptrs=[MMIOInterface(queue.write_pointer_address, 8, fmt='Q')],
-                        doorbells=[MMIOInterface(self.doorbells + queue.doorbell_offset - self.doorbells_base, 8, fmt='Q')])
+    return AMDQueueDesc(ring=MMIOInterface(ring.va_addr, ring.size, fmt='I'), read_ptr=MMIOInterface(queue.read_pointer_address, 8, fmt='Q'),
+                        write_ptr=MMIOInterface(queue.write_pointer_address, 8, fmt='Q'),
+                        doorbell=MMIOInterface(self.doorbells + queue.doorbell_offset - self.doorbells_base, 8, fmt='Q'))

  def sleep(self, tm:int) -> bool:
    kfd.AMDKFD_IOC_WAIT_EVENTS(KFDIface.kfd, events_ptr=self.queue_event_arr_ptr, num_events=1, wait_for_all=1, timeout=tm)
@@ -857,8 +854,8 @@ class PCIIface(PCIIfaceBase):
        wptr_addr=gart.va_addr+wptr, eop_addr=eop_buffer.va_addr, eop_size=eop_buffer.size,
        idx=int(is_aql:=(queue_type==kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL)), aql=is_aql)

-    return AMDQueueDesc(ring=ring.cpu_view().view(fmt='I'), doorbells=[self.dev_impl.doorbell64.view(doorbell_index * 8, 8, fmt='Q')],
-      read_ptrs=[gart.cpu_view().view(offset=rptr, size=8, fmt='Q')], write_ptrs=[gart.cpu_view().view(offset=wptr, size=8, fmt='Q')], put_value=pv)
+    return AMDQueueDesc(ring=ring.cpu_view().view(fmt='I'), doorbell=self.dev_impl.doorbell64.view(doorbell_index * 8, 8, fmt='Q'),
+      read_ptr=gart.cpu_view().view(offset=rptr, size=8, fmt='Q'), write_ptr=gart.cpu_view().view(offset=wptr, size=8, fmt='Q'), put_value=pv)

  def sleep(self, timeout) -> bool:
    if hasattr(self.pci_dev, 'irq_poller') and self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))):
@@ -53,7 +53,7 @@ class CLProgram:
    try: check(cl.clReleaseProgram(self.program))
    except (TypeError, AttributeError): pass

-  def __call__(self, *bufs:tuple[ctypes._CData, BufferSpec], global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]|None=None,
+  def __call__(self, *bufs:tuple[cl.cl_mem, BufferSpec], global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]|None=None,
               vals:tuple[int, ...]=(), wait=False) -> float|None:
    for i,(b,_) in enumerate(bufs):
      if isinstance(dt:=self.buf_dtypes[i], ImageDType):
@@ -75,13 +75,13 @@ class CLProgram:
    return None

 class CLAllocator(LRUAllocator['CLDevice']):
-  def _alloc(self, size:int, options:BufferSpec) -> tuple[ctypes._CData, BufferSpec]:
+  def _alloc(self, size:int, options:BufferSpec) -> tuple[cl.cl_mem, BufferSpec]:
    # Recalculate real size for texture
    if options.image is not None: size = options.image.pitch * options.image.shape[0]
    return (checked(cl.clCreateBuffer(self.dev.context, cl.CL_MEM_READ_WRITE, size, None, status := ctypes.c_int32()), status), options)
  @suppress_finalizing
-  def _free(self, opaque:tuple[ctypes._CData, BufferSpec], options:BufferSpec): check(cl.clReleaseMemObject(opaque[0]))
-  def _copyin(self, dest:tuple[ctypes._CData, BufferSpec], src:memoryview):
+  def _free(self, opaque:tuple[cl.cl_mem, BufferSpec], options:BufferSpec): check(cl.clReleaseMemObject(opaque[0]))
+  def _copyin(self, dest:tuple[cl.cl_mem, BufferSpec], src:memoryview):
    if mv_address(src) % 16: src = memoryview(bytearray(src))
    if (img:=dest[1].image):
      stride = img.shape[1]*img.itemsize*4
@@ -89,7 +89,7 @@ class CLAllocator(LRUAllocator['CLDevice']):
        check(cl.clEnqueueWriteBuffer(self.dev.queue, dest[0], False, i*img.pitch, stride, mv_address(src)+(i*stride), 0, None, None))
    else: check(cl.clEnqueueWriteBuffer(self.dev.queue, dest[0], False, 0, len(src)*src.itemsize, from_mv(src), 0, None, None))
    self.dev.pending_copyin.append(src)    # NOTE: these can't be freed until the GPU actually executes this command
-  def _copyout(self, dest:memoryview, src:tuple[ctypes._CData, BufferSpec]):
+  def _copyout(self, dest:memoryview, src:tuple[cl.cl_mem, BufferSpec]):
    if (img:=src[1].image):
      stride = img.shape[1]*img.itemsize*4
      for i in range(img.shape[0]):
@@ -69,11 +69,11 @@ class CUDAAllocator(LRUAllocator['CUDADevice']):
    if options.external_ptr: return cuda.CUdeviceptr_v2(options.external_ptr)
    if options.host: return init_c_var(ctypes.c_void_p, lambda x: check(cuda.cuMemHostAlloc(ctypes.byref(x), size, 0x01)))
    return init_c_var(cuda.CUdeviceptr, lambda x: check(cuda.cuMemAlloc_v2(ctypes.byref(x), size)))
+  @suppress_finalizing
  def _free(self, opaque, options:BufferSpec):
-    try:
-      if options.host: check(cuda.cuMemFreeHost(opaque))
-      else: check(cuda.cuMemFree_v2(opaque))
-    except (TypeError, AttributeError): pass
+    if options.external_ptr: return
+    if options.host: check(cuda.cuMemFreeHost(opaque))
+    else: check(cuda.cuMemFree_v2(opaque))
  def _copyin(self, dest, src:memoryview):
    check(cuda.cuCtxSetCurrent(self.dev.context))
    host_mem = self.alloc(len(src), BufferSpec(host=True))
@@ -1,15 +1,16 @@
-import subprocess, pathlib, struct, ctypes, tempfile, functools, contextlib, decimal, platform, sys
-from tinygrad.helpers import prod, to_mv, getenv, round_up, cache_dir, PROFILE, ProfileRangeEvent, cpu_profile, unwrap
+import subprocess, pathlib, struct, ctypes, tempfile, functools, contextlib, decimal, platform
+from tinygrad.helpers import prod, to_mv, getenv, round_up, cache_dir, PROFILE, ProfileRangeEvent, cpu_profile, unwrap, suppress_finalizing
 import tinygrad.runtime.support.objc as objc
 from tinygrad.device import Compiled, Compiler, CompileError, LRUAllocator, ProfileDeviceEvent, CompilerSet, CompilerPair
 from tinygrad.renderer.cstyle import MetalRenderer
 from tinygrad.runtime.autogen import metal
+from tinygrad.runtime.support.c import DLL

 # 13 is requestType that metal uses to compile source code into MTLB, there aren't any docs or symbols.
 REQUEST_TYPE_COMPILE = 13

 # Must be loaded for default Metal Device: https://developer.apple.com/documentation/metal/1433401-mtlcreatesystemdefaultdevice?language=objc
-ctypes.CDLL("/System/Library/Frameworks/CoreGraphics.framework/CoreGraphics")
+DLL("CoreGraphics", "CoreGraphics")

 # FIXME: these need autogen to support objc categories
 # https://developer.apple.com/library/archive/documentation/Cocoa/Conceptual/ObjectiveC/Chapters/ocCategories.html
@@ -67,7 +68,7 @@ class MetalCompiler(Compiler):
  # doesn't seem to be anything we can do.
  with contextlib.suppress(FileNotFoundError, ModuleNotFoundError):
    import tinygrad.runtime.autogen.llvm # noqa: F401
-  support = ctypes.CDLL("/System/Library/PrivateFrameworks/MTLCompiler.framework/MTLCompiler")
+  support = DLL("MTLCompiler", "MTLCompiler")
  support.MTLCodeGenServiceCreate.restype = ctypes.c_void_p

  def __init__(self):
@@ -167,8 +168,9 @@ class MetalAllocator(LRUAllocator[MetalDevice]):
    ret.retain = False
    if ret.value is None: raise MemoryError(f"Metal OOM while allocating {size=}")
    return MetalBuffer(ret, size)
+  @suppress_finalizing
  def _free(self, opaque:MetalBuffer, options):
-    if not sys.is_finalizing(): opaque.buf.release
+    if not options.external_ptr: opaque.buf.release
  def _transfer(self, dest:MetalBuffer, src:MetalBuffer, sz:int, src_dev:MetalDevice, dest_dev:MetalDevice):
    dest_dev.synchronize()
    src_command_buffer = src_dev.mtl_queue.commandBuffer().retained()
@@ -1,11 +1,10 @@
 import functools
 from tinygrad.device import Compiled, Compiler, Allocator, CompilerSet, CompilerPair
 from tinygrad.engine.jit import MultiGraphRunner
-from tinygrad.renderer.cstyle import Renderer, CStyleLanguage
+from tinygrad.renderer.cstyle import Renderer, CStyleLanguage, AMDHIPRenderer
 from tinygrad.uop.ops import Ops
 from tinygrad.helpers import cpu_profile, EMULATE, NULL_IR3, NULL_NAK
 from tinygrad.renderer.nir import IR3Renderer, NAKRenderer
-from tinygrad.renderer.llvmir import AMDLLVMRenderer

 class NullRenderer(CStyleLanguage):
  device = "NULL"
@@ -34,9 +33,9 @@ class NullDevice(Compiled):
  def __init__(self, device:str):
    renderer:functools.partial|type[Renderer]
    match str(EMULATE.value):
-      case "AMD": renderer = functools.partial(AMDLLVMRenderer, "gfx1100")
-      case "AMD_RDNA4": renderer = functools.partial(AMDLLVMRenderer, "gfx1201")
-      case "AMD_CDNA4": renderer = functools.partial(AMDLLVMRenderer, "gfx950")
+      case "AMD": renderer = functools.partial(AMDHIPRenderer, "gfx1100")
+      case "AMD_RDNA4": renderer = functools.partial(AMDHIPRenderer, "gfx1201")
+      case "AMD_CDNA4": renderer = functools.partial(AMDHIPRenderer, "gfx950")
      case "": renderer = NullRenderer
      case _: raise RuntimeError(f"can't EMULATE device: {EMULATE.value}")
    compilers = CompilerSet([CompilerPair(renderer, Compiler), CompilerPair(functools.partial(IR3Renderer, 0x6030001), None, NULL_IR3), # adreno 630
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
George HotzandGitHub	1019a3d8f8	Merge branch 'master' into mac_pytest	2026-02-02 23:40:03 +08: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
George HotzandGitHub	965149a46d	assembly/amd: add ds perm instructions (#14486 ) * assembly/amd: add ds perm instructions * NO SKIP * fix preexisting RDNA3 issues * pcode * assert * asserts * unify * simp * good fix	2026-02-02 16:02:00 +08:00
qazalandGitHub	1746d1f997	remove SPEC=0 context in custom_kernel tests, pyrender always skips it (#14489 )	2026-02-02 16:32:01 +09:00
George HotzandGitHub	d4007f36e0	remove DEFINE_GLOBAL (it is PARAM now) (#14488 )	2026-02-02 14:56:37 +08:00
qazalandGitHub	6c487656f9	viz: kernel metadata from rodata entry (#14487 )	2026-02-02 15:41:42 +09:00
Robbe DerksandGitHub	d75a1b0d5a	usbgpu: use BOT interface for `patch.py` (#13644 ) * BOT usage * cleanup * fix lint * fix ruff * fix -7?	2026-02-02 11:54:46 +08:00
sirhcmandGitHub	2931b52875	skip autogen if MTLCompiler is loaded (#14466 )	2026-02-01 22:12:27 -05:00
George HotzandGitHub	9a32d6e090	add depth limit for SPEC=2 (#14485 ) * make SPEC=2 work for everything * that's a horrible fix * add depth limit	2026-02-02 10:43:28 +08:00
George HotzandGitHub	368a692e1a	make SPEC=2 work for everything (#14476 ) * make SPEC=2 work for everything * that's a horrible fix	2026-02-02 10:30:56 +08:00
chenyuandGitHub	ea1f1d2b9d	test_assign_to_bitcast_view (#14483 ) currently disk allows assign same size dtype into a bitcasted view	2026-02-01 16:46:04 -05:00
chenyuandGitHub	6deeccc192	fix RING with single dest (#14482 )	2026-02-01 12:14:46 -05:00
chenyuandGitHub	3ff390159b	don't implicitly change dtype in assign (#14481 ) broadcast shape is fine, but implicitly cast dtype is hard to find	2026-02-01 11:48:54 -05:00
imaolo GitHub chenyu	2111762a48	failed test case for RING output device (#14191 ) * Add enable/disable scheduler cache ContextVar * add allreduce ring and naive to() tests * clearer test comparing native vs ring allreduce * split tests, add helper * removing trailing whitespace --------- Co-authored-by: chenyu <[email protected]>	2026-02-01 11:48:43 -05:00
chenyuandGitHub	02afae04f4	atol in test_call_gemm (#14480 ) flaky	2026-02-01 11:24:58 -05:00
chenyuandGitHub	5705398a1f	assign cleanup [pr] (#14479 ) share more code path between disk and non-disk. also raise RuntimeError instead of Assert for mismatches	2026-02-01 09:10:22 -05:00
chenyuandGitHub	da500dbe06	simplify late_buffer_view [pr] (#14478 ) check the only allowed Ops in the chain, and offset cannot be negative	2026-01-31 22:38:40 -05:00
chenyuandGitHub	b4f96301e0	remove unused rules [pr] (#14477 )	2026-01-31 21:29:30 -05:00
qazalandGitHub	54e78dbec8	viz: remove hardcoded strings in cfg tests (#14462 )	2026-02-01 09:30:43 +09:00
chenyuandGitHub	5d38db9da6	generic bitcast assign (#14474 ) a.bitcast(X).assign(src) -> a.assign(src.bitcast(a.dtype))	2026-01-31 17:29:20 -05:00
chenyuandGitHub	b38fc43b07	assert assign dtype mismatch for disk [pr] (#14473 ) the disk hack is generally wrong, now force bitcast on the source before assign	2026-01-31 17:08:54 -05:00
chenyuandGitHub	ced886f26c	failed test case for assign into bitcast (#14469 ) * failed test case for assign into bitcast DISK assign has custom hack for this. need to fix before we can unify assign * test_assign_bitcast_different_size	2026-01-31 14:26:47 -05:00
chenyuandGitHub	81eee5b30a	unused spec [pr] (#14468 ) no BUFFER_VIEW in tensor, and no CONTIGUOUS in KERNEL	2026-01-31 13:53:16 -05:00
nimlgenandGitHub	f873c7b6c5	amd: fetch_name is file_name (#14465 )	2026-01-31 20:11:07 +03:00
chenyuandGitHub	c765641215	remove unused allow_any_len [pr] (#14464 ) STORE has 2 src, RESHAPE has 2 src, BUFFER has 2 src added some tests for the untested allow_any_len	2026-01-31 11:05:42 -05:00
chenyuandGitHub	b4f5a51ebb	move tests to unit (#14463 ) test_uop_graph does not need device, test_memory_planner can use NULL	2026-01-31 10:49:31 -05:00
qazalandGitHub	616e9c1483	CDNA assembly gemm in tensor.py with flag (#14310 ) * work * work * the assembly * remove the old one * remove ws bufs, assert splitk * notes cleanup * work * gemm args * gemm in mixins would be nice * add gemm gradient * print counters * the realize is for DEBUG=2 aesthetics * dedup * rewrite to python dsl, no list copies * leave that * add B, M, N, K to gemm name * it's M0 not NULL * fp16 support * test cleanup + more gemms * work from viz * more work * gemm batch_size * xccg path work * tiny comments on the label naming * s_waitcnt	2026-01-31 22:34:14 +09:00
chenyuandGitHub	55f806b713	tighter late_buffer_view match [pr] (#14456 ) src must be len 2 at that point	2026-01-31 07:28:26 -05:00
qazalandGitHub	d69bc5aa1a	make DEV=NULL EMULATE=AMD amd_asm_matmul run (#14460 )	2026-01-31 20:45:24 +09:00
geohot	4d7b16f330	in the pyenv	2026-01-31 13:37:31 +08:00
geohot	814a1d59e2	fresh db	2026-01-31 13:22:03 +08:00
geohot	de52fa6116	comment	2026-01-31 13:19:15 +08:00
geohot	067747560b	setup env	2026-01-31 13:18:21 +08:00
qazalandGitHub	4976544bf9	multi ram usage tests on the NULL device (#14457 )	2026-01-31 14:14:53 +09:00
geohot	d3c808d90b	3 min	2026-01-31 13:09:03 +08:00
geohot	dde297fd47	3 minute timeout	2026-01-31 13:08:09 +08:00
geohot	84b3d8117f	add pytest -nauto to benchmark	2026-01-31 13:06:48 +08:00
chenyuandGitHub	99b44121bc	failed test case for non-consecutive disk read (#14455 ) silently fail now	2026-01-30 23:44:04 -05:00
George HotzandGitHub	b705c9143c	assembly/amd: test more instructions (#14365 ) * assembly/amd: test more instructions * more * passing * revert * no const fold * remove junk * cleaner	2026-01-31 12:40:22 +08:00
George HotzandGitHub	c9a3ddb341	benchmark llama walltime script (#14454 ) * benchmark llama walltime script * adj layers	2026-01-31 10:21:54 +08:00
George HotzandGitHub	f5346d6a1a	fix USE_ATOMICS for non float dtypes and make it the default (#14444 ) * embedded multistep test * complex test * with jit * fix dtypes and reenable USE_ATOMICS * that test didn't catch anything	2026-01-31 09:44:16 +08:00
sirhcmandGitHub	e575dd8275	prevent UB in long decomp and more emulated tests (#14447 )	2026-01-30 19:38:41 -05:00
chenyuandGitHub	3204f94454	correct var_vals schedule filter (#14451 ) complete_create_schedule_with_vars returns var_vals that's used in schedule	2026-01-30 17:10:07 -05:00
chenyuandGitHub	cfcd1debb5	test schedule with multiple AFTER (#14449 )	2026-01-30 15:59:00 -05:00
nimlgenandGitHub	486d53d646	device: call free for external_ptr (#14448 ) * device: call free for external_ptr * lin	2026-01-30 23:53:17 +03:00
nimlgenandGitHub	e0978498dc	amd: read_ptr/write_ptr/doorbells are not lists (#14445 )	2026-01-30 23:11:57 +03:00
sirhcmandGitHub	1803ee939d	EMULATED_DTYPES=long works with CPU_LLVM (#14446 )	2026-01-30 13:54:43 -05:00
chenyuandGitHub	03613e83ad	update TestTensorMetadata (#14443 ) run with SCACHE=0 some more TODOs	2026-01-30 12:39:01 -05:00
geohot	cbb1eed57b	hotfix: partial revert of `9eb449f88`, caused llama NaN	2026-01-30 17:19:27 +00:00
chenyuandGitHub	26f5c00265	move TestTensorMetadata to unit (#14442 )	2026-01-30 12:14:21 -05:00
chenyuandGitHub	c05a0b85ae	flip unique const src order [pr] (#14441 ) * flip unique const src order [pr] matches buffer, simplifies replace_input_buffer * combine rules	2026-01-30 11:44:18 -05:00
geohot	ee2c78709d	mlperf/llama: disable USE_ATOMICS for now	2026-01-31 00:42:08 +08:00
chenyuandGitHub	beecac4d85	expand ranges -> unroll outer ranges [pr] (#14440 )	2026-01-30 11:26:05 -05:00
chenyuandGitHub	9eb449f882	clean up toposort sched_sink [pr] (#14439 )	2026-01-30 10:18:28 -05:00
George HotzandGitHub	838cd078bc	use atomics for embedding backward (#14400 ) * embedding is slow * failing * float is fine * null * it fails * simplify embedding with broadcasting * ATOMIC_ADD incoming * min change * simpler test * better test * fix test * real test * simpler * cleanups * types and names * _zero_kernel * grad multi * hack * none * multi unshard * more for call * don't tag in call * good * call_multi * call_multi wow claude is useless * embedding backward mutli test * test passes * fix as_param * shape_to_shape_arg * add clip * before cast * fix spec=2, use atomics	2026-01-30 18:10:59 +08:00
nimlgenandGitHub	1998e0bb28	nv: add prof props to dev (#14437 )	2026-01-30 12:51:43 +03:00
George HotzandGitHub	7a9dee4e50	add call/param UOps (#14433 ) * add call/param UOps * resolve call * skip that for now * grad on call * fix tests	2026-01-30 14:51:45 +08:00
qazalandGitHub	66d6a68016	viz: sqtt work from cdna gemm (#14434 ) * it's the tag * initialize rows based on the disasm * test_cfg with Ops.BINARY * pyremu wants s_code_end? * test_diamond * diff cleanup	2026-01-30 14:00:56 +09:00
sirhcmandGitHub	88caf57ef4	ci: unify python versions (#14430 )	2026-01-29 21:42:03 -05:00
chenyuandGitHub	86a204d22a	allow Tensor setitem input to be list/tuple (#14432 ) matches assign, and generally matches numpy	2026-01-29 21:26:58 -05:00
chenyuandGitHub	4a80319093	clean up split_store final logic [pr] (#14429 ) explicitly check the structure	2026-01-29 18:40:07 -05:00
sirhcmandGitHub	e47f12f671	ci: replace testing_minimal with testing_unit (#14427 )	2026-01-29 18:02:43 -05:00
wozeparrotandGitHub	c2fb8b208f	fa: 32 block size (#14416 )	2026-01-29 13:59:13 -08:00
chenyuandGitHub	a979fafae5	cleanup around disk buffer [pr] (#14428 ) style change, prep for refactor	2026-01-29 16:18:44 -05:00
nimlgenandGitHub	dc977a03b0	nv_pma: bw decoder (#14424 ) * nv_pma: bw decoder * decoder fix * better	2026-01-30 00:12:39 +03:00
chenyuandGitHub	ddc041854b	failed test case for disk setitem (#14426 ) strided setitem is wrong	2026-01-29 14:54:19 -05:00
chenyuandGitHub	31706bf6bc	add few more types [pr] (#14425 )	2026-01-29 14:04:09 -05:00
nimlgenandGitHub	2d5c24879f	nv: pma for 5090 (#14420 ) * nv: pma for 5090 * hm * 4090	2026-01-29 20:06:01 +03:00
nimlgenandGitHub	c8dc6332d2	memory: read_fields is not universal (#14348 )	2026-01-29 20:00:00 +03:00
chenyuandGitHub	dbe8f034a7	pass z3.Context in validate ctx [pr] (#14423 ) does not need to pass the whole solver	2026-01-29 11:11:47 -05:00
chenyuandGitHub	033ce1b885	types for validate.py (#14422 )	2026-01-29 10:56:50 -05:00
nimlgenandGitHub	230d08ec70	test for am recovery and faults handling (#14421 ) * test for am recovery and faults handling * linter	2026-01-29 17:11:24 +03:00
George HotzandGitHub	793afbd473	simplify nn.Embedding, support AFTER in CUSTOM_KERNEL (#14419 )	2026-01-29 17:22:13 +08:00
sirhcmandGitHub	0c855d6149	ci: remove unused pydeps (#14418 )	2026-01-29 01:51:26 -05:00
wozeparrotandGitHub	4845e42135	llama3 gradacc fixes (#14414 )	2026-01-28 19:12:39 -08:00
chenyuandGitHub	37cde4a01a	add one line mypy report (#14415 )	2026-01-28 20:39:32 -05:00
chenyuandGitHub	15aed51544	return types for all math.py function (#14413 ) calling int() on sint -> int, i think it's better support since some UOp can be safely cast to int	2026-01-28 20:10:11 -05:00
nimlgenandGitHub	aec1ae0de1	llama: set manual_seed (#14409 )	2026-01-28 14:40:00 -08:00
chenyuandGitHub	0870ed28b1	add Self type to MathMixin (#14411 ) these don't cause error	2026-01-28 16:59:38 -05:00
chenyuandGitHub	079f33c208	fix type in Tensor.mean and Tensor.var (#14410 ) use Tensor.from_uop to wrap UOp from symbolic shape, kernels are the same	2026-01-28 15:24:02 -05:00
chenyuandGitHub	2b5e99ccc1	minor type cleanups [pr] (#14408 ) mypy --warn-redundant-casts has false negative	2026-01-28 14:11:50 -05:00
chenyuandGitHub	726415dbc8	import sint directly in movement.py TYPE_CHECKING (#14406 ) avoid creating string TypeAlias, fixed warning in `TYPED=1 python test/test_tiny.py`	2026-01-28 12:47:26 -05:00
nimlgenandGitHub	acb2fc36ba	nv_pma: add decoder (#14404 ) * nv_pma: add decoder * cl	2026-01-28 20:44:02 +03:00
chenyuandGitHub	7b9bc1d8cf	_MockMemoryviewMeta for mockgpu (#14405 ) fixed `PYTHONPATH=. TYPED=1 DEV=AMD MOCKGPU=1 python test/test_tiny.py`. basically make `isinstance(TrackedMemoryView_instance, memoryview)` true	2026-01-28 11:59:00 -05:00
chenyuandGitHub	93793a645b	use cl.cl_mem instead of internal ctypes._CData (#14403 ) fixed `CHECK_OOB=0 DEV=CL TYPED=1 python test/test_tiny.py`	2026-01-28 10:56:41 -05:00
chenyuandGitHub	a9b44070a8	fix webgpu runtime types (#14402 ) `CHECK_OOB=0 DEV=WEBGPU TYPED=1 python test/test_tiny.py` passed, also skip tests that failed locally	2026-01-28 10:37:25 -05:00
George HotzandGitHub	0c6b3f50aa	add marker to llama training (#14401 )	2026-01-28 22:44:28 +08:00
Jakob SachsandGitHub	2b7c00d3d2	fix sd-example dtype for CLIP embeddings (#14397 )	2026-01-28 09:07:19 -05:00
qazalandGitHub	a5a9ce3fdf	viz: disasm cleanups from null emulate (#14399 ) * it's AMDHIPRenderer * don't need that indent * less assignment stuff * that arg order did not make sense * pmc	2026-01-28 22:03:30 +09:00
nimlgenandGitHub	544928766d	hcq_smi: kill mac pids (#14398 )	2026-01-28 15:00:28 +03:00
George HotzandGitHub	202b74b369	assembly/amd: continue refactors (#14386 ) * simpler * merge * flat * no ctx * use the correct apis * dup code * write clean code * remove bad helpers * bits junk remove * junk remove * smem test * fix tests * correct fix + tests * Fmt matters it seems * wmma refactor * a lil more * kimi cleanups * line	2026-01-28 17:33:03 +08:00
qazalandGitHub	5bffa17f82	llama train: better NULL=1 EMULATE=AMD_CDNA4 dev experience (#14395 ) * beam opens devices * switch to hip renderer * amd: true? * llvm true is for test_autogen	2026-01-28 17:31:22 +09:00
qazalandGitHub	0294014108	fix bufferize cost function for multi, improve VIZ=-1 cli (#14394 ) * improve cli * remove_bufferize change	2026-01-28 15:53:18 +09:00
qazalandGitHub	c158acea29	failing multi ram usage test from llama gemm (#14392 )	2026-01-28 14:32:32 +09:00
sirhcmandGitHub	067e27857e	nested composite actions don't work (#14393 )	2026-01-28 00:13:30 -05:00
sirhcmandGitHub	9dddf3d478	don't save caches for PRs, try 2 (#14391 )	2026-01-27 23:30:17 -05:00
sirhcmandGitHub	68fe5d8b36	Revert "don't save caches for PRs (#14389 )" (#14390 )	2026-01-27 23:22:26 -05:00
sirhcmandGitHub	4ab228b498	don't save caches for PRs (#14389 )	2026-01-27 23:21:31 -05:00
sirhcmandGitHub	5e36482314	decompose long to ints where unsupported, try 2 (#14383 )	2026-01-27 23:20:43 -05:00
wozeparrotandGitHub	e496547720	llama3 gradacc (#14291 )	2026-01-27 19:48:10 -08:00