CALL with return value is FUNCTION

2026-08-14 07:18:27 +00:00 · 2026-04-16 12:36:14 +08:00
800 changed files with 78652 additions and 164147 deletions
@@ -5,7 +5,6 @@ runs:
  steps:
    - name: Run process replay tests
      shell: bash
-      if: env.CAPTURE_PROCESS_REPLAY == '1'
      run: |
        export PR_TITLE=$(jq -r .pull_request.title "$GITHUB_EVENT_PATH")
        export CURRENT_SHA=${{ github.event.pull_request && github.event.pull_request.head.sha || github.sha }}
@@ -4,13 +4,13 @@ inputs:
  python-version:
    description: 'Python version to use'
    required: false
-    default: '' # if you don't set a version, the native python version will be used
+    default: '3.12'
  key:
    description: 'Key for the python cache'
    required: false
    default: ''  # if you don't set a key, it doesn't cache
  deps:
-    description: 'Extra dependency groups (space separated)'
+    description: 'Extra dependency groups (comma separated)'
    required: false
    default: ''
  pydeps:
@@ -41,33 +41,20 @@ inputs:
    description: "Install LLVM?"
    required: false
    default: 'false'
-  tinydreno:
-    description: "Install tinydreno"
+  mesa:
+    description: "Install mesa"
    required: false
    default: 'false'
-  qemu:
-    description: "Install qemu"
+  tinydreno:
+    description: "Install tinydreno"
    required: false
    default: 'false'
 runs:
  using: "composite"
  steps:
-    - name: Setup environment
-      shell: bash
-      run: |
-        echo "UV_CACHE_DIR=/tmp/.uv-cache" >> "$GITHUB_ENV"
-        echo "OMP_NUM_THREADS=1" >> "$GITHUB_ENV"
-        # no buffers should be over 300MB in CI
-        echo "MAX_BUFFER_SIZE=300000000" >> "$GITHUB_ENV"
-
-    - name: Set up uv
-      uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b
-      with:
-        enable-cache: 'false' # see below for manual caching
-
    - name: Set up Python ${{ inputs.python-version }}
+      id: setup-python
      uses: actions/setup-python@v6
-      if: inputs.python-version != ''
      with:
        python-version: ${{ inputs.python-version }}

@@ -76,23 +63,23 @@ runs:
    - name: Cache Python packages (PR)
      if: github.event_name == 'pull_request'
      id: restore-venv-pr
-      uses: actions/cache/restore@v5
+      uses: actions/cache/restore@v4
      with:
-        path: /tmp/.uv-cache
-        key: uv-${{ runner.os }}-${{ runner.arch }}-python-${{ inputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
+        path: ${{ github.workspace }}/.venv
+        key: venv-${{ runner.os }}-${{ runner.arch }}-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@v5
      with:
-        path: /tmp/.uv-cache
-        key: uv-${{ runner.os }}-${{ runner.arch }}-python-${{ inputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
+        path: ${{ github.workspace }}/.venv
+        key: venv-${{ runner.os }}-${{ runner.arch }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}

    # **** Caching downloads ****

    - name: Cache downloads (PR)
      if: inputs.key != '' && github.event_name == 'pull_request'
-      uses: actions/cache/restore@v5
+      uses: actions/cache/restore@v4
      with:
        path: ${{ runner.os == 'Linux' && '~/.cache/tinygrad/downloads/' || '~/Library/Caches/tinygrad/downloads/' }}
        key: downloads-${{ github.job }}-${{ inputs.key }}-${{ env.CACHE_VERSION }}
@@ -106,26 +93,34 @@ runs:
    # **** Python deps ****

    - name: Install dependencies in venv (with extra)
-      if: inputs.deps != ''
+      if: inputs.deps != '' && steps.restore-venv-pr.outputs.cache-hit != 'true' && steps.restore-venv.outputs.cache-hit != 'true'
      shell: bash
      run: |
-        uv venv .venv
-        DEPS="${{ inputs.deps }}"
-        uv pip install --python .venv -e ".[${DEPS// /,}]" ${{ inputs.pydeps }} --torch-backend cpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/
+        python -m venv .venv
+        if [[ "$RUNNER_OS" == "Windows" ]]; then
+          source .venv/Scripts/activate
+        else
+          . .venv/bin/activate
+        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 == ''
+      if: inputs.deps == '' && steps.restore-venv-pr.outputs.cache-hit != 'true' && steps.restore-venv.outputs.cache-hit != 'true'
      shell: bash
      run: |
-        uv venv .venv
-        uv pip install --python .venv -e . ${{ inputs.pydeps }}
-    - name: Prune uv cache
-      if: github.event_name != 'pull_request'
-      shell: bash
-      run: uv cache prune --ci
-    - name: Configure venv
+        python -m venv .venv
+        if [[ "$RUNNER_OS" == "Windows" ]]; then
+          source .venv/Scripts/activate
+        else
+          . .venv/bin/activate
+        fi
+        python -m pip install -e . ${{ inputs.pydeps }}
+    - name: Set up venv environment
      shell: bash
      run: |
        echo "VIRTUAL_ENV=${{ github.workspace }}/.venv" >> "$GITHUB_ENV"
+        echo "OMP_NUM_THREADS=1" >> "$GITHUB_ENV"
+        # no buffers should be over 300MB in CI
+        echo "MAX_BUFFER_SIZE=300000000" >> "$GITHUB_ENV"
        if [[ "$RUNNER_OS" == "Windows" ]]; then
          echo "${{ github.workspace }}/.venv/Scripts" >> "$GITHUB_PATH"
        else
@@ -134,7 +129,7 @@ runs:

    # ******************* apt *******************
    - name: Setup apt
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true')
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
      shell: bash
      run: |
        sudo chown -R $USER:$USER /var/cache/apt/archives
@@ -143,6 +138,11 @@ runs:
        echo 'Acquire::http::Pipeline-Depth "5";' | sudo tee -a /etc/apt/apt.conf.d/99parallel
        echo 'Binary::apt::APT::Keep-Downloaded-Packages "true";' | sudo tee -a /etc/apt/apt.conf.d/99keep-debs

+    - name: Add OpenCL Repo
+      if: inputs.opencl == 'true' && runner.os == 'Linux'
+      shell: bash
+      run: echo "deb [ allow-insecure=yes ] https://apt.repos.intel.com/oneapi all main" | sudo tee /etc/apt/sources.list.d/oneAPI.list
+
    - name: Add AMD Repo (Linux)
      if: inputs.amd == 'true' && runner.os == 'Linux'
      shell: bash
@@ -161,50 +161,54 @@ runs:
        echo "deb http://apt.llvm.org/$(lsb_release -cs)/ llvm-toolchain-$(lsb_release -cs)-20 main" | sudo tee /etc/apt/sources.list.d/llvm.list

    - name: Compute Package List + Hash
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true')
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
      id: apt-pkgs
      shell: bash
      run: |
        pkgs=""
        # **** OpenCL ****
        if [[ "${{ inputs.opencl }}" == "true" ]]; then
-          pkgs+=" ocl-icd-opencl-dev"
+          pkgs+=" opencl-headers \
+            intel-oneapi-runtime-openmp=2023.2.1-16 intel-oneapi-runtime-compilers-common=2023.2.1-16 intel-oneapi-runtime-compilers=2023.2.1-16 \
+            intel-oneapi-runtime-dpcpp-sycl-opencl-cpu=2023.2.1-16 intel-oneapi-runtime-tbb-common=2021.10.0-49541 \
+            intel-oneapi-runtime-tbb=2021.10.0-49541 intel-oneapi-runtime-opencl=2023.2.1-16"
        fi
        # **** AMD ****
        if [[ "${{ inputs.amd }}" == "true" ]]; then
-          pkgs+=" comgr"
+          pkgs+=" hsa-rocr comgr hsa-rocr-dev liburing-dev libibverbs-dev libc6-dev"
+        fi
+        # **** CUDA ****
+        if [[ "${{ inputs.cuda }}" == "true" ]]; then
+          pkgs+=" git g++ cmake ninja-build llvm-15-dev zlib1g-dev libglew-dev \
+            flex bison libfl-dev libboost-thread-dev libboost-filesystem-dev nvidia-cuda-toolkit-gcc libzstd-dev"
        fi
        # **** WebGPU (dependencies for software-based vulkan) ****
        if [[ "${{ inputs.webgpu }}" == "true" ]]; then
-          pkgs+=" mesa-vulkan-drivers"
+          pkgs+=" libgl1 libglx-mesa0 libgl1-mesa-dri libxcb-xfixes0-dev mesa-vulkan-drivers"
        fi
        # **** LLVM ****
        if [[ "${{ inputs.llvm }}" == "true" ]]; then
          pkgs+=" libllvm20 clang-20 lld-20"
        fi
-        # **** QEMU ****
-        if [[ "${{ inputs.qemu }}" == "true" ]]; then
-          pkgs+=" qemu-user-static"
-        fi

        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.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true') && github.event_name == 'pull_request'
-      uses: actions/cache/restore@v5
+      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 }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
    - name: Cache apt
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true') && github.event_name != 'pull_request'
+      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@v5
      with:
        path: /var/cache/apt/archives/
        key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}

    - name: Run apt Update + Install
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true')
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
      shell: bash
      run: |
        sudo apt -qq update || true
@@ -216,11 +220,6 @@ runs:

        sudo chown -R $USER:$USER /var/cache/apt/archives/

-    - name: Add clang to PATH (Linux)
-      if: inputs.llvm == 'true' && runner.os == 'Linux'
-      shell: bash
-      run: echo "/usr/lib/llvm-20/bin" >> "$GITHUB_PATH"
-
    # **** AMD ****
    - name: Setup AMD (Linux)
      if: inputs.amd == 'true' && runner.os == 'Linux'
@@ -240,33 +239,78 @@ runs:
          jq -r '.assets[] | select(.name == "libamd_comgr.dylib").browser_download_url' | \
          sudo xargs curl -fL -o /usr/local/lib/libamd_comgr.dylib

-    # **** CUDA ****
-    - name: Install CUDA
-      if: inputs.cuda == 'true'
+    # **** gpuocelot ****
+
+    - name: Install gpuocelot dependencies (MacOS)
+      if: inputs.ocelot == 'true' && runner.os == 'macOS'
      shell: bash
      run: |
-        sudo mkdir -p /usr/local/cuda/targets/x86_64-linux
-        curl -fL https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvrtc/linux-x86_64/cuda_nvrtc-linux-x86_64-11.5.119-archive.tar.xz \
-          | sudo tar -xJ -C /usr/local/cuda/targets/x86_64-linux --strip-components=1
-        echo /usr/local/cuda/targets/x86_64-linux/lib | sudo tee /etc/ld.so.conf.d/cuda-nvrtc.conf
-        sudo ldconfig
+        pkgs=(cmake ninja llvm@15 zlib glew flex bison [email protected] zstd ncurses)
+        for f in "${pkgs[@]}"; do
+          brew ls --versions "$f" >/dev/null 2>&1 || brew install --quiet "$f"
+        done

-    # **** gpuocelot ****
+        # Fix boost 1.85 for gpuocelot
+        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' && github.event_name != 'pull_request'
+      id: cache-build
+      uses: actions/cache@v5
+      env:
+        cache-name: cache-gpuocelot-build-1
+      with:
+        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-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
+        cd ${{ github.workspace }}/gpuocelot/ocelot
+        git checkout b16039dc940dc6bc4ea0a98380495769ff35ed99
+        mkdir build
+        cd build
+
+        CMAKE_ARGS="-Wno-dev -G Ninja -DOCELOT_BUILD_TOOLS=OFF -DCMAKE_BUILD_ALWAYS=0 -DBUILD_TESTS_CUDA=OFF -DCMAKE_POLICY_VERSION_MINIMUM=3.5"
+        if [[ "${{ runner.os }}" == "macOS" ]]; then
+          sudo xcode-select -s /Applications/Xcode_16.2.app/Contents/Developer
+          CMAKE_ARGS="$CMAKE_ARGS -DBoost_INCLUDE_DIR=$(brew --prefix boost)/include -DBoost_LIBRARY_DIR=$(brew --prefix boost)/lib"
+        fi
+
+        cmake .. $CMAKE_ARGS
+        ninja
    - name: Install gpuocelot
      if: inputs.ocelot == 'true'
      shell: bash
      run: |
-        sudo mkdir -p /usr/local/lib
-        sudo curl --output-dir /usr/local/lib -fLO https://github.com/tinygrad/gpuocelot/releases/download/v0.1.0/libgpuocelot.${{ runner.os == 'Linux' && 'so' || 'dylib' }}
+        cd ${{ github.workspace }}/gpuocelot/ocelot/build
+        sudo cp libgpuocelot.${{ runner.os == 'macOS' && 'dylib' || 'so' }} /usr/${{ runner.os == 'macOS' && 'local/' || '' }}lib/

    # **** WebGPU ****

-    - name: Install WebGPU dawn
-      if: inputs.webgpu == 'true'
+    - name: Install WebGPU dawn (Linux)
+      if: inputs.webgpu == 'true' && runner.os == 'Linux'
      shell: bash
      run: |
-        sudo mkdir -p /usr/local/lib
-        sudo curl --output-dir /usr/local/lib -fLO https://github.com/wpmed92/pydawn/releases/download/v0.1.6/libwebgpu_dawn.${{ runner.os == 'Linux' && 'so' || 'dylib' }}
+        sudo curl -fL https://github.com/wpmed92/pydawn/releases/download/v0.1.6/libwebgpu_dawn.so -o /usr/local/lib/libwebgpu_dawn.so
+        sudo ldconfig
+    - name: Install WebGPU dawn (macOS)
+      if: inputs.webgpu == 'true' && runner.os == 'macOS'
+      shell: bash
+      run: |
+        brew tap wpmed92/dawn
+        brew install dawn

    # **** LLVM ****

@@ -275,18 +319,18 @@ runs:
      shell: bash
      run: brew install llvm@20

+    # **** mesa ****
+    - name: Install mesa (linux)
+      if: inputs.mesa == 'true' && runner.os == 'Linux'
+      shell: bash
+      run: sudo curl -fL https://github.com/sirhcm/tinymesa/releases/download/v1/libtinymesa_cpu-mesa-25.2.7-linux-amd64.so -o /usr/lib/libtinymesa_cpu.so
+    - name: Install mesa (macOS)
+      if: inputs.mesa == 'true' && runner.os == 'macOS'
+      shell: bash
+      run: brew install sirhcm/tinymesa/tinymesa_cpu
+
    # *** tinydreno ***
    - name: Install tinydreno (linux)
      if: inputs.tinydreno == 'true' && runner.os == 'Linux'
      shell: bash
      run: sudo curl -fL https://github.com/sirhcm/tinydreno/raw/refs/heads/master/libllvm-qcom.so -o /usr/lib/libllvm-qcom.so
-
-    # *** OpenCL ***
-    - name: Install rusticl
-      if: inputs.opencl == 'true'
-      shell: bash
-      run: |
-        sudo curl -fL https://github.com/sirhcm/tinymesa/releases/download/rusticl-v1/libRusticlOpenCL.so.1.0.0 -o /usr/lib/libRusticlOpenCL.so
-        sudo mkdir -p /etc/OpenCL/vendors
-        echo "/usr/lib/libRusticlOpenCL.so" | sudo tee /etc/OpenCL/vendors/rusticl.icd
-        echo "RUSTICL_ENABLE=llvmpipe" >> "$GITHUB_ENV"
@@ -33,20 +33,23 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: 'autogen'
+        opencl: 'true'
        amd: 'true'
+        cuda: 'true'
        llvm: 'true'
+        webgpu: 'true'
+        mesa: 'true'
        pydeps: 'pyyaml mako'
    - name: Install autogen support packages
-      run: sudo apt-get install -y --no-install-recommends libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev libdrm-dev liburing-dev
+      run: sudo apt-get install -y --no-install-recommends libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev libdrm-dev
    - name: Regenerate autogen files
      run: |
        find tinygrad/runtime/autogen -type f -name "*.py" -not -path "*/amd/*" -not -name "__init__.py" -not -name "comgr.py" -not -name "metal.py" -not -name "iokit.py" -not -name "corefoundation.py" -not -name "libclang.py" -delete
        python3 -c "from tinygrad.runtime.autogen import opencl"
-        python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv_580, nv_610, nv"
+        python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv_580, nv"
        python3 -c "from tinygrad.runtime.autogen import comgr_3, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd, amdgpu_drm"
-        python3 -c "from tinygrad.runtime.autogen.am import *"
-        python3 -c "from tinygrad.runtime.autogen.nv_regs import *"
-        python3 -c "from tinygrad.runtime.autogen import libc, kfd, io_uring, pci, vfio"
+        python3 -c "from tinygrad.runtime.autogen.am import am, pm4_soc15, pm4_nv, sdma_4_0_0, sdma_5_0_0, sdma_6_0_0, smu_v13_0_0, smu_v13_0_6, smu_v13_0_12, smu_v14_0_2"
+        python3 -c "from tinygrad.runtime.autogen import libc, kfd, io_uring, ib, pci, vfio"
        python3 -c "from tinygrad.runtime.autogen import llvm"
        python3 -c "from tinygrad.runtime.autogen import webgpu"
        python3 -c "from tinygrad.runtime.autogen import kgsl, qcom_dsp"
@@ -55,7 +58,6 @@ jobs:
        python3 -c "from tinygrad.runtime.autogen import avcodec"
        python3 -c "from tinygrad.runtime.autogen import llvm_qcom"
        python3 -c "from tinygrad.runtime.autogen import mlx5"
-        python3 -c "from tinygrad.runtime.autogen import ggml_common"
        REGEN=1 python3 -c "from tinygrad.runtime.autogen import libclang"
    - name: Check for differences
      run: |
@@ -1,8 +1,8 @@
 name: Run MLPerf Training

 on:
-  #schedule:
-  #  - cron: '5 8 * * *'  # Runs at 08:05 UTC (12:05 AM Pacific Time)
+  schedule:
+    - cron: '5 8 * * *'  # Runs at 08:05 UTC (12:05 AM Pacific Time)
  push:
    branches:
      - update_mlperf
@@ -68,4 +68,3 @@ mutants
 .mutmut-cache
 dagre/
 graphlib/
-uv.lock
@@ -1,5 +0,0 @@
-# Notes
-
- Run tests with `-n12` for speed (e.g. `python -m pytest test/null/test_dtype.py -x -q -n12`)
- Run `python -m mypy tinygrad/` to typecheck
- Run `python -m ruff check .` to lint
@@ -72,7 +72,7 @@ As it turns out, 90% of what you need for neural networks are a decent autograd/
 Throw in an optimizer, a data loader, and some compute, and you have all you need.

 ```python
-from tinygrad import Tensor, nn, Context
+from tinygrad import Tensor, nn

 class LinearNet:
  def __init__(self):
@@ -86,7 +86,7 @@ optim = nn.optim.Adam([model.l1, model.l2], lr=0.001)

 x, y = Tensor.rand(4, 1, 28, 28), Tensor([2,4,3,7])  # replace with real mnist dataloader

-with Context(TRAINING=1):
+with Tensor.train():
  for i in range(10):
    optim.zero_grad()
    loss = model(x).sparse_categorical_crossentropy(y).backward()
@@ -140,8 +140,8 @@ Documentation along with a quick start guide can be found on the [docs website](
 ```python
 from tinygrad import Tensor

-x = Tensor.eye(3)
-y = Tensor([[2.0,0,-2.0]])
+x = Tensor.eye(3, requires_grad=True)
+y = Tensor([[2.0,0,-2.0]], requires_grad=True)
 z = y.matmul(x).sum()
 z.backward()

@@ -164,9 +164,7 @@ print(y.grad.tolist())  # dz/dy

 ## Contributing

-There has been a lot of interest in tinygrad lately. Following these guidelines will help your PR get accepted. If you do submit a PR, please include a sentence or two about why you want this merged and why you think it will improve the project.
-
-If you are a new contributor with something that looks even close to AI written, it will be closed without feedback and you may be banned from our GitHub. No human should waste time reading AI slop. And for everyone, if you used AI, disclose what you used it for.
+There has been a lot of interest in tinygrad lately. Following these guidelines will help your PR get accepted.

 We'll start with what will get your PR closed with a pointer to this section:

@@ -198,8 +196,6 @@ python3 test/backend/test_ops.py        # just the ops tests
 python3 -m pytest test/                 # whole test suite
 ```

-For agents, always run tests with `-n12` for speed.
-
 #### Process replay tests

 [Process replay](https://github.com/tinygrad/tinygrad/blob/master/test/external/process_replay/README.md) compares your PR's generated kernels against master. If your PR is a refactor or speedup without any expected behavior change, It should include [pr] in the pull request title.
@@ -1,4 +1,6 @@
 # abstractions2 goes from back to front, here we will go from front to back
+from typing import List
+from tinygrad.helpers import tqdm

 # *****
 # 0. Load mnist on the device
@@ -11,7 +13,7 @@ X_train -= X_train.mean()
 # *****
 # 1. Define an MNIST model.

-from tinygrad import Tensor, Context
+from tinygrad import Tensor

 l1 = Tensor.kaiming_uniform(128, 784)
 l2 = Tensor.kaiming_uniform(10, 128)
@@ -24,28 +26,28 @@ l1n, l2n = l1.numpy(), l2.numpy()
 from tinygrad.nn.optim import SGD
 optim = SGD([l1, l2])

-with Context(TRAINING=1):
-  X, Y = X_train[(samples:=Tensor.randint(128, high=X_train.shape[0]))], Y_train[samples]
-  optim.zero_grad()
-  model(X).sparse_categorical_crossentropy(Y).backward()
-  optim.schedule_step()   # this will step the optimizer without running realize
+Tensor.training = True
+X, Y = X_train[(samples:=Tensor.randint(128, high=X_train.shape[0]))], Y_train[samples]
+optim.zero_grad()
+model(X).sparse_categorical_crossentropy(Y).backward()
+optim.schedule_step()   # this will step the optimizer without running realize

 # *****
-# 3. Create a schedule (linear uop).
+# 3. Create a schedule.

 # The weight Tensors have been assigned to, but not yet realized. Everything is still lazy at this point
 # l1.uop and l2.uop define a computation graph

-from tinygrad.engine.realize import run_linear
-linear = Tensor.schedule_linear(l1, l2)
+from tinygrad.schedule import ExecItem
+schedule: List[ExecItem] = Tensor.schedule(l1, l2)

-print(f"The schedule contains {len(linear.src)} items.")
-for call in linear.src: print(str(call)[:80])
+print(f"The schedule contains {len(schedule)} items.")
+for si in schedule: print(str(si)[:80])

 # *****
-# 4. Lower and run the schedule (linear uop).
+# 4. Lower and run the schedule.

-run_linear(linear)
+for si in tqdm(schedule): si.run()

 # *****
 # 5. Print the weight change
@@ -1,9 +1,9 @@
 # tinygrad allows you to write kernels at many different abstractions levels.
 # This is for RDNA3, but if you don't have one you can run with the emulator
-# PYTHONPATH="." DEV=MOCKPCI+AMD
+# PYTHONPATH="." MOCKGPU=1 DEV=AMD

 from tinygrad import Tensor, Context, GlobalCounters, UOp, Device
-from tinygrad.helpers import DEV, DEBUG, getenv
+from tinygrad.helpers import DEBUG, getenv
 from tinygrad.uop.ops import AxisType, KernelInfo, Ops
 from tinygrad.dtype import AddrSpace, dtypes
 from tinygrad.runtime.autogen.amd.rdna3.ins import *
@@ -16,7 +16,7 @@ def eval_harness(name, tensor, fxn, check=None):
  print(f"computed in {GlobalCounters.time_sum_s*1000:.2f} ms, {(a.nbytes()/1e9)/GlobalCounters.time_sum_s:.2f} GB/s")
  return out

-SZ = 256*1024 if DEV.interface.startswith("MOCK") else 1024*1024*1024
+SZ = 256*1024 if getenv("MOCKGPU") else 1024*1024*1024

 def example_2_hip(a:Tensor, correct):
  GLOBALS = 1024
@@ -67,7 +67,8 @@ def example_2_hip(a:Tensor, correct):
    # the sink specifies the GLOBAL and LOCAL sizes, along with the input buffers and name
    sink = UOp.sink(UOp.special(GLOBALS, 'gidx0'), UOp.special(THREADS, 'lidx0'), out, buf,
                    arg=KernelInfo(name="hip_reduce_sum_kernel"))
-    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT),
+                UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
  eval_harness("HIP kernel", a, lambda x: Tensor.empty(GLOBALS).custom_kernel(x, fxn=hip_reduce_sum)[0].sum(), check=correct)

 def example_3_custom_uop(a:Tensor, correct):
@@ -104,7 +105,7 @@ def example_3_custom_uop(a:Tensor, correct):
 def example_5_custom_assembly(a:Tensor, correct):
  # Kernel class copied from amd_asm_matmul
  class Kernel:
-    def __init__(self): self.instructions, self.labels, self.pos = [], {}, 0
+    def __init__(self, arch='gfx1100'): self.instructions, self.labels, self.pos, self.arch = [], {}, 0, arch
    def label(self, name): self.labels[name] = self.pos
    def emit(self, inst, target=None):
      self.instructions.append(inst)
@@ -122,7 +123,8 @@ def example_5_custom_assembly(a:Tensor, correct):
        offset_dwords = (self.labels[inst._target] - inst._pos - inst.size()) // 4
        if not -32768 <= offset_dwords <= 32767: raise ValueError(f"branch to '{inst._target}' offset {offset_dwords} exceeds simm16 range")
        inst.simm16 = offset_dwords
-      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in self.instructions]))))
+      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT),
+                                   UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in self.instructions]))))

  CU_COUNT = 32
  LANES = 64
@@ -17,13 +17,15 @@ The `UOp` graph specifies the compute in terms of low level tinygrad ops. Not al

 ## Scheduling

-The [scheduler](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/schedule/__init__.py) converts the graph of UOps into a `LINEAR` UOp whose `src` is a list of `CALL` UOps. One `CALL` is one kernel on the GPU, and the scheduler is responsible for breaking the large compute graph into subgraphs that can fit in a kernel. The `CALL`'s `src[0]` (a `SINK` ast) specifies what compute to run, and the remaining `src` are the buffers to run it on.
+The [scheduler](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/schedule/__init__.py) converts the graph of UOps into a list of `ExecItem`. One `ExecItem` is one kernel on the GPU, and the scheduler is responsible for breaking the large compute graph into subgraphs that can fit in a kernel. `ast` specifies what compute to run, and `bufs` specifies what buffers to run it on.
+
+::: tinygrad.schedule.ExecItem

 ## Lowering

-The code in [realize](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/engine/realize.py) lowers each `CALL` by compiling its ast into a `PROGRAM` and running it.
+The code in [realize](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/engine/realize.py) lowers `ExecItem` by populating its `prg` field with

-::: tinygrad.engine.realize.run_linear
+::: tinygrad.engine.realize.run_schedule

 There's a ton of complexity hidden behind this, see the `codegen/` directory.

@@ -33,7 +35,13 @@ Then we render the UOps into code with a `Renderer`, then we compile the code to

 ## Execution

-`run_linear` walks the `LINEAR` UOp, dispatching each `CALL` to a runner (kernel, copy, view, encdec, or graph).
+Creating `ExecItem`, which has a run method
+
+::: tinygrad.engine.realize.ExecItem
+    options:
+        members: true
+
+Lists of `ExecItem` can be condensed into a single ExecItem with the Graph API (rename to Queue?)

 ## Runtime

@@ -28,7 +28,7 @@ Transforms the ast into an optimized ast. This is where BEAM search and heuristi

 Transform the optimized ast into a linearized and rendered program.

-::: tinygrad.codegen.to_program
+::: tinygrad.codegen.get_program
    options:
        members: false
        show_labels: false
@@ -53,7 +53,7 @@ Transform the linearized list of UOps into a program, represented as a string.

 Abstracted high level interface to the runtimes.

-::: tinygrad.engine.realize.to_program
+::: tinygrad.engine.realize.get_program
    options:
        members: false
        show_labels: false
@@ -62,7 +62,7 @@ A lot of work can still be done here. For example, we never copy the inputs to o

 Many accelerators have Tensor Cores / MAC arrays / systolic arrays. The main value of these is that, since they are 2-D, they create an n^2 ratio between the compute and the input data.

-GPUs use Tensor Cores instead of MAC arrays to fit better in the GPU warp paradigm. This is because the output of Tensor Cores is O(n) wrt the input, while the output of MAC arrays is O(n^2)
+GPUs use Tensor Cores instead of MAC arrays to fit better in the GPU warp paradigm. This is because the output of Tensor Cores is O(n) wrt the input, while the output of MAC arrays like the AMX is O(n^2)

 We have a simple framework in tinygrad for adding these ALU blocks and achieving good performance from them.

@@ -57,8 +57,6 @@ AMD:LLVM      | use the AMD device with the LLVM renderer
 NV:CUDA:sm_70 | use the NV device with the CUDA renderer targetting sm_70
 AMD::gfx950   | use the AMD device targetting gfx950
 USB+AMD       | use the AMD device over the USB interface
-CPU:LLVM      | use the CPU device with the LLVM renderer
-CPU:LLVM:x86_64,znver2,avx2,-avx512f | use the CPU device with the LLVM renderer, with [additional arch flags](runtime.md#cpu-arch)

 ### Debug breakdown

@@ -24,7 +24,7 @@ You will see `CUDA` here on a GPU instance, or `CPU` here on a CPU instance.
 We'll use the model from [the Keras tutorial](https://keras.io/examples/vision/mnist_convnet/).

 ```python
-from tinygrad import Tensor, nn, Context
+from tinygrad import Tensor, nn

 class Model:
  def __init__(self):
@@ -74,8 +74,8 @@ We'll use the Adam optimizer. The `nn.state.get_parameters` will walk the model
 ```python
 optim = nn.optim.Adam(nn.state.get_parameters(model))
 batch_size = 128
-@Context(TRAINING=1)
 def step():
+  Tensor.training = True  # makes dropout work
  samples = Tensor.randint(batch_size, high=X_train.shape[0])
  X, Y = X_train[samples], Y_train[samples]
  optim.zero_grad()
@@ -143,6 +143,7 @@ Since we are just randomly sampling from the dataset, there's no real concept of
 for step in range(7000):
  loss = jit_step()
  if step%100 == 0:
+    Tensor.training = False
    acc = (model(X_test).argmax(axis=1) == Y_test).mean().item()
    print(f"step {step:4d}, loss {loss.item():.2f}, acc {acc*100.:.2f}%")
 ```
@@ -37,4 +37,4 @@
    options:
        show_signature: false
        separate_signature: false
-::: tinygrad.llm.gguf.gguf_load
+::: tinygrad.nn.state.gguf_load
@@ -133,7 +133,7 @@ For our loss function we will be using sparse categorical cross entropy loss. Th
 ```python
 def sparse_categorical_crossentropy(self, Y, ignore_index=-1) -> Tensor:
    loss_mask = Y != ignore_index
-    y_counter = Tensor.arange(self.shape[-1], dtype=dtypes.int32).unsqueeze(0).expand(Y.numel(), self.shape[-1])
+    y_counter = Tensor.arange(self.shape[-1], dtype=dtypes.int32, requires_grad=False, device=self.device).unsqueeze(0).expand(Y.numel(), self.shape[-1])
    y = ((y_counter == Y.flatten().reshape(-1, 1)).where(-1.0, 0) * loss_mask.reshape(-1, 1)).reshape(*Y.shape, self.shape[-1])
    return self.log_softmax().mul(y).sum() / loss_mask.sum()
 ```
@@ -165,18 +165,17 @@ from extra.datasets import fetch_mnist
 Now we have everything we need to start training our neural network.
 We will be training for 1000 steps with a batch size of 64.

-We use `with Context(TRAINING=1)` to enable training mode.
+We use `with Tensor.train()` to set the internal flag `Tensor.training` to `True` during training.
 Upon exit, the flag is restored to its previous value by the context manager.

 ```python
-from tinygrad import Context
 X_train, Y_train, X_test, Y_test = fetch_mnist()

-with Context(TRAINING=1):
+with Tensor.train():
  for step in range(1000):
    # random sample a batch
    samp = np.random.randint(0, X_train.shape[0], size=(64))
-    batch = Tensor(X_train[samp])
+    batch = Tensor(X_train[samp], requires_grad=False)
    # get the corresponding labels
    labels = Tensor(Y_train[samp])

@@ -214,7 +213,7 @@ with Timing("Time: "):
  for step in range(1000):
    # random sample a batch
    samp = np.random.randint(0, X_test.shape[0], size=(64))
-    batch = Tensor(X_test[samp])
+    batch = Tensor(X_test[samp], requires_grad=False)
    # get the corresponding labels
    labels = Y_test[samp]

@@ -258,7 +257,7 @@ with Timing("Time: "):
  for step in range(1000):
    # random sample a batch
    samp = np.random.randint(0, X_test.shape[0], size=(64))
-    batch = Tensor(X_test[samp])
+    batch = Tensor(X_test[samp], requires_grad=False)
    # get the corresponding labels
    labels = Y_test[samp]

@@ -5,12 +5,12 @@ tinygrad supports various runtimes, enabling your code to scale across a wide ra
 | Runtime | Description | Compiler Options | Requirements |
 |---------|-------------|------------------|--------------|
 | [NV](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_nv.py) | Provides acceleration for NVIDIA GPUs | nvrtc (default)<br>PTX (`DEV=NV:PTX`) | Ampere/Ada/Blackwell series GPUs.<br>You can select an interface via [the `DEV` variable](env_vars.md#dev-variable). See [NV interfaces](#nv-interfaces) for details. |
-| [AMD](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_amd.py) | Provides acceleration for AMD GPUs | LLVM (`DEV=AMD:LLVM`)<br>HIP/COMGR (`DEV=AMD:HIP`) | CDNA3, CDNA4, RDNA3 or RDNA4 GPUs.<br>You can select an interface via [the `DEV` variable](env_vars.md#dev-variable). See [AMD interfaces](#amd-interfaces) for details. |
+| [AMD](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_amd.py) | Provides acceleration for AMD GPUs | LLVM (`DEV=AMD:LLVM`)<br>HIP/COMGR (`DEV=AMD:HIP`) | RDNA2 or newer GPUs.<br>You can select an interface via [the `DEV` variable](env_vars.md#dev-variable). See [AMD interfaces](#amd-interfaces) for details. |
 | [QCOM](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_qcom.py) | Provides acceleration for QCOM GPUs | - | 6xx series GPUs |
 | [METAL](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_metal.py) | Utilizes Metal for acceleration on Apple devices | - | M1+ Macs; Metal 3.0+ for `bfloat` support |
 | [CUDA](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cuda.py) | Utilizes CUDA for acceleration on NVIDIA GPUs | nvrtc (default)<br> PTX (`DEV=CUDA:PTX`) | NVIDIA GPU with CUDA support |
 | [CL](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cl.py) | Accelerates computations using OpenCL on GPUs | - | OpenCL 2.0 compatible device |
-| [CPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cpu.py) | Runs on CPU using the clang or llvm compiler | Clang JIT (default)<br>LLVM IR (`DEV=CPU:LLVM`) | `clang` compiler in system `PATH`<br>You can specify additional arch parameters via [the `DEV` variable](env_vars.md#dev-variable). See [CPU arch](#cpu-arch) for details. |
+| [CPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cpu.py) | Runs on CPU using the clang or llvm compiler | Clang JIT (default)<br>LLVM IR (`DEV=CPU:LLVM`) | `clang` compiler in system `PATH` |
 | [WEBGPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_webgpu.py) | Runs on GPU using the Dawn WebGPU engine (used in Google Chrome) | - | Dawn library installed and discoverable. Binaries: [pydawn v0.3.0](https://github.com/wpmed92/pydawn/releases/tag/v0.3.0) |


@@ -79,9 +79,3 @@ NV backend supports several interfaces for communicating with devices:

 * `NVK`: uses the nvidia driver
 * `PCI`: uses the [NV driver](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/support/nv/nvdev.py)
-
-## CPU Arch
-The CPU renderers may be additionally configured using the arch component of [the `DEV` environment variable](env_vars.md#dev-variable).
-CPU arch should be specified as a comma-separated list of parameters, and must contain at least two values: the architecture family (ie. x86_64, arm64, or riscv64) and the cpu type (as accepted by `clang`'s `-march`).
-If native is specified as the cpu type, tinygrad (or delegate compiler) will query the host cpu type. Additional comma-separated values are interpreted as cpu feature flags. When a value is preceded by a `-` character, the corresponding feature flag will be disabled, otherwise the flag will be enabled.
-Note that enabled feature flags should not be preceded by a `+`.
@@ -66,8 +66,8 @@ Elementwise ops operate on a per element basis. They don't change the shape of t
 ::: tinygrad.Tensor.sub
 ::: tinygrad.Tensor.mul
 ::: tinygrad.Tensor.div
+::: tinygrad.Tensor.idiv
 ::: tinygrad.Tensor.mod
-::: tinygrad.Tensor.fmod
 ::: tinygrad.Tensor.bitwise_xor
 ::: tinygrad.Tensor.bitwise_and
 ::: tinygrad.Tensor.bitwise_or
@@ -19,8 +19,8 @@

 ## tinygrad ops

-::: tinygrad.Tensor.linear_with_vars
-::: tinygrad.Tensor.schedule_linear
+::: tinygrad.Tensor.schedule_with_vars
+::: tinygrad.Tensor.schedule
 ::: tinygrad.Tensor.realize
 ::: tinygrad.Tensor.replace
 ::: tinygrad.Tensor.assign
@@ -4,7 +4,7 @@ TinyGPU app lets you use AMD and NVIDIA GPUs on macOS over USB4/Thunderbolt with

 ## Requirements

- macOS (13.0+)
+- macOS (12.1+)
 - USB4/Thunderbolt port
 - A supported GPU (AMD RDNA3+ or NVIDIA Ampere+)

@@ -55,7 +55,7 @@ export PATH="$HOME/.local/bin:$PATH"
 ### 5. Use it!

 ```bash
-DEV={AMD|NV} python3 -m tinygrad.llm
+DEV={AMD|NV} python3 tinygrad/apps/llm.py
 ```

 **Note:** Use `JITBEAM=2` to search for faster kernels (one-time search cost, results cached).
@@ -0,0 +1,196 @@
+from tinygrad import Tensor, dtypes, Context, getenv, UOp, fetch
+from tinygrad.uop.ops import Ops, PatternMatcher, UPat
+from tinygrad.uop.symbolic import symbolic
+from tinygrad.codegen import Renderer
+from tinygrad.codegen.opt import Opt, OptOps
+
+# ************************* implementation of the problem ************************
+
+def myhash(a: Tensor) -> Tensor:
+  a = (a + 0x7ED55D16) + (a << 12)
+  a = (a ^ 0xC761C23C) ^ (a >> 19)
+  a = (a + 0x165667B1) + (a << 5)
+  a = (a + 0xD3A2646C) ^ (a << 9)
+  a = (a + 0xFD7046C5) + (a << 3)
+  a = (a ^ 0xB55A4F09) ^ (a >> 16)
+  return a
+
+def select_with_where_tree(values: Tensor, relative_idx: Tensor) -> Tensor:
+  n = values.shape[0]
+  if n == 1: return values[0].expand(relative_idx.shape)
+
+  mid = n // 2
+  left = select_with_where_tree(values[:mid], relative_idx)
+  right = select_with_where_tree(values[mid:], relative_idx - mid)
+
+  go_left = relative_idx < mid
+  return go_left.where(left, right)
+
+def tree_traversal(forest: Tensor, val: Tensor, height: int, rounds: int, where_tree_threshold=3) -> Tensor:
+  # All walkers start at idx=0
+  idx = Tensor.zeros(val.shape, device=val.device, dtype=dtypes.uint32)
+
+  for r in range(rounds):
+    level = r % (height + 1)
+    level_start = (1 << level) - 1
+    level_size = 1 << level
+
+    if level == 0:
+      # At root (level 0), all walkers are at idx=0
+      # No gather needed, just broadcast the root value
+      node_val = forest[0].expand(val.shape)
+      idx = idx * 0  # Reset to 0
+    elif level <= where_tree_threshold:
+      # Small level: use where-tree
+      level_values = forest[level_start : level_start + level_size]
+      relative_idx = (idx - level_start)
+      node_val = select_with_where_tree(level_values, relative_idx)
+    else:
+      # Large level: use gather
+      node_val = forest.gather(0, idx)
+
+    val = myhash(val ^ node_val)
+    idx = (idx << 1) + (1 + (val & 1))
+
+    # No wrap check needed! At round 10 (level becomes 0), we reset idx above.
+
+  return val.contiguous(arg=(Opt(OptOps.UPCAST, 0, 8),))
+
+# ************************* renderer for VLIW machine *************************
+
+def loop_unrolling(sink:UOp):
+  rng = [x for x in sink.toposort() if x.op is Ops.RANGE]
+  if len(rng) == 0: return None
+  print(f"unrolling loop with size {rng[0].vmax+1}")
+  unrolled_sinks = [sink.substitute({rng[0]:rng[0].const_like(i)}).src[0] for i in range(rng[0].vmax+1)]
+  return UOp.sink(*unrolled_sinks, arg=sink.arg)
+
+global_addrs = []
+vliw_prepare = PatternMatcher([
+  # loop unrolling (should be a part of tinygrad)
+  (UPat(Ops.SINK, name="sink"), loop_unrolling),
+  # cast is fake
+  (UPat(Ops.CAST, name="c"), lambda c: c.src[0]),
+  # rewrites to hardcode the addresses in memory
+  (UPat(Ops.PARAM, name="dg"), lambda dg: UOp.const(dtypes.uint, global_addrs[dg.arg])),
+  # INDEX is just plus
+  (UPat(Ops.INDEX, name="i"), lambda i: i.src[0]+i.src[1]),
+])+symbolic
+
+class VLIWRenderer(Renderer):
+  has_local = False  # TODO: this should be the default / cleaned up
+  # this says this backend supports MULACC + more. decompositions uses this
+  code_for_op: dict = {Ops.MULACC: None, Ops.ADD: "+", Ops.MUL: "*",
+                       Ops.XOR: "^", Ops.AND: "&", Ops.OR: "|",
+                       Ops.SHL: "<<", Ops.SHR: ">>", Ops.CMPLT: "<"}
+  # this matcher runs while still in graph form
+  pre_matcher = vliw_prepare
+
+  def render(self, uops:list[UOp]):
+
+    # TODO: this is a minimal renderer. for low cycle count, make it good
+    # to get speed, you need to add VLIW packing
+    # to get under 1536 regs, you need to add a register allocator
+    # we left the fun parts to you
+
+    print(f"rendering with {len(uops)} uops")
+    reg, inst = 0, []
+    r: dict[UOp, int] = {}
+    for u in uops:
+      assert u.dtype.count in (1,8), "dtype count must be 1 or 8"
+
+      # dumb register allocator
+      if u.op not in {Ops.STORE, Ops.SINK, Ops.GEP}:
+        r[u] = reg
+        reg += u.dtype.count
+
+      # render UOps to instructions
+      match u.op:
+        case Ops.SINK:
+          inst.append({"flow": [("halt",)]})
+        case Ops.CONST:
+          inst.append({"load": [("const", r[u], u.arg)]})
+        case Ops.GEP:
+          # a GEP is just an alias to a special register in the vector
+          r[u] = r[u.src[0]] + u.arg[0]
+        case Ops.VECTORIZE:
+          if all(s == u.src[0] for s in u.src):
+            # if all sources are the same, we can broadcast
+            inst.append({"valu": [("vbroadcast", r[u], r[u.src[0]])]})
+          else:
+            # this is a copy into a contiguous chunk of registers
+            inst.extend({"flow": [("add_imm", r[u]+i, r[s], 0)]} for i,s in enumerate(u.src) if r[s] != r[u]+i)
+        case Ops.LOAD:
+          op = "vload" if u.dtype.count > 1 else "load"
+          inst.append({"load": [(op, r[u], r[u.src[0]])]})
+        case Ops.STORE:
+          op = "vstore" if u.src[1].dtype.count > 1 else "store"
+          inst.append({"store": [(op, r[u.src[0]], r[u.src[1]])]})
+        case Ops.MULACC:
+          assert u.dtype.count == 8
+          inst.append({"valu": [("multiply_add", r[u], r[u.src[0]], r[u.src[1]], r[u.src[2]])]})
+        case Ops.WHERE:
+          assert u.dtype.count == 8
+          inst.append({"flow": [("vselect", r[u], r[u.src[0]], r[u.src[1]], r[u.src[2]])]})
+        case _ if u.op in self.code_for_op:
+          cat = "valu" if u.dtype.count > 1 else "alu"
+          inst.append({cat: [(self.code_for_op[u.op], r[u], r[u.src[0]], r[u.src[1]])]})
+        case _:
+          raise NotImplementedError(f"unhandled op {u.op}")
+    return repr(inst)
+
+# ************************* test and render *************************
+
+import sys, types
+PROBLEM_URL = "https://raw.githubusercontent.com/anthropics/original_performance_takehome/refs/heads/main/tests/frozen_problem.py"
+sys.modules["problem"] = problem = types.ModuleType("problem")
+exec(fetch(PROBLEM_URL).read_text(), problem.__dict__)
+
+if __name__ == "__main__":
+  batch_size = getenv("BS", 256)
+  height = 10
+  rounds = getenv("ROUNDS", 16)
+
+  # build problem
+  tree = problem.Tree.generate(height)
+  inp = problem.Input.generate(tree, batch_size, rounds)
+  mem = problem.build_mem_image(tree, inp)
+  global_addrs.extend([mem[6], mem[6], mem[4]])  # output, input, forest
+
+  # *** verify the kernel in tinygrad compared to reference ***
+
+  forest_t = Tensor(tree.values, dtype=dtypes.uint32)
+  val_t = Tensor(inp.values, dtype=dtypes.uint32)
+
+  if getenv("VERIFY", 1):
+    # verify on normal tinygrad device
+    with Context(PCONTIG=2):
+      out = tree_traversal(forest_t, val_t, height, rounds)
+      val_out = out.tolist()
+    problem.reference_kernel(tree, inp)
+    assert val_out == inp.values
+    print("verification passed")
+
+  # *** render to device ***
+
+  from tinygrad.codegen import get_program
+  with Context(PCONTIG=2, DEVECTORIZE=2, SPEC=0):
+    out = tree_traversal(forest_t, val_t, height, rounds)
+    sink = out.schedule()[-1].ast
+    prg = get_program(sink, VLIWRenderer())
+
+  # *** run on Machine and compare ***
+
+  # NOTE: the scratch size needs to be reduced to 1536 when you have a register allocator
+  src = eval(prg.src)
+  max_regs = max(t[1] for instr in src for v in instr.values() for t in v if len(t) > 1) + 8
+  print(f"{max_regs:5d} regs used" + ("" if max_regs <= 1536 else "       <-- WARNING: TOO MANY REGISTERS, MUST BE <= 1536"))
+  machine = problem.Machine(mem, src, problem.DebugInfo(scratch_map={}), n_cores=1, trace=False, scratch_size=max_regs)
+  machine.run()
+  print(f"ran for {machine.cycle:5d} cycles" + ("" if machine.cycle <= 1363 else "  <-- EVEN CLAUDE GOT 1363"))
+
+  # compare to reference
+  ref_mem = mem.copy()
+  for _ in problem.reference_kernel2(ref_mem, {}): pass
+  assert machine.mem[mem[6]:mem[6]+mem[2]] == ref_mem[mem[6]:mem[6]+mem[2]]
+  print("compare passed!")
@@ -4,10 +4,10 @@ from tinygrad.dtype import DTypeLike, dtypes
 import math

 # rewritten from numpy
-def rfftfreq(n: int, d: float = 1.0) -> Tensor:
+def rfftfreq(n: int, d: float = 1.0, device=None) -> Tensor:
  val = 1.0 / (n * d)
  N = n // 2 + 1
-  results = Tensor.arange(N)
+  results = Tensor.arange(N, device=device)
  return results * val

 # just like in librosa
@@ -1,6 +1,6 @@
 from typing import Tuple
 import time
-from tinygrad import Tensor, TinyJit, nn, Context
+from tinygrad import Tensor, TinyJit, nn
 import gymnasium as gym
 from tinygrad.helpers import trange
 import numpy as np  # TODO: remove numpy import
@@ -55,7 +55,7 @@ if __name__ == "__main__":

  @TinyJit
  def train_step(x:Tensor, selected_action:Tensor, reward:Tensor, old_log_dist:Tensor) -> Tuple[Tensor, Tensor, Tensor]:
-    with Context(TRAINING=1):
+    with Tensor.train():
      log_dist, value = model(x)
      action_mask = (selected_action.reshape(-1, 1) == Tensor.arange(log_dist.shape[1]).reshape(1, -1).expand(selected_action.shape[0], -1)).float()

@@ -67,8 +67,8 @@ class ConvGroup:
    self.conv2 = nn.Conv2d(channels_out, channels_out, kernel_size=3, padding=1, bias=False)
    self.norm1 = nn.BatchNorm(channels_out, track_running_stats=False, eps=1e-12, momentum=hyp['net']['batch_norm_momentum'])
    self.norm2 = nn.BatchNorm(channels_out, track_running_stats=False, eps=1e-12, momentum=hyp['net']['batch_norm_momentum'])
-    cast(Tensor, self.norm1.weight).is_param_(False)
-    cast(Tensor, self.norm2.weight).is_param_(False)
+    cast(Tensor, self.norm1.weight).requires_grad = False
+    cast(Tensor, self.norm2.weight).requires_grad = False
  def __call__(self, x:Tensor) -> Tensor:
    x =    self.norm1(self.conv1(x).max_pool2d().float()).cast(dtypes.default_float).quick_gelu()
    return self.norm2(self.conv2(x).float()).cast(dtypes.default_float).quick_gelu() + x
@@ -122,7 +122,7 @@ if __name__ == "__main__":
    return ret.mul(hyp['opt']['loss_scale_scaler']*loss_batchsize_scaler).sum().div(hyp['opt']['loss_scale_scaler'])

  @TinyJit
-  @Context(TRAINING=1)
+  @Tensor.train()
  def train_step(idxs:Tensor) -> Tensor:
    X, Y = X_train[idxs], Y_train[idxs]
    if len(GPUS) > 1:
@@ -1,6 +1,6 @@
 # model based off https://medium.com/data-science/going-beyond-99-mnist-handwritten-digits-recognition-cfff96337392
 from typing import Callable
-from tinygrad import Tensor, TinyJit, nn, GlobalCounters, function, Context
+from tinygrad import Tensor, TinyJit, nn, GlobalCounters, function
 from tinygrad.helpers import getenv, colored, trange
 from tinygrad.nn.datasets import mnist

@@ -19,7 +19,7 @@ class Model:
  def __call__(self, x:Tensor) -> Tensor: return x.sequential(self.layers)

  @TinyJit
-  @Context(TRAINING=1)
+  @Tensor.train()
  def train_step(self, X_train:Tensor, Y_train:Tensor) -> Tensor:
    opt.zero_grad()
    samples = Tensor.randint(getenv("BS", 512), high=X_train.shape[0])
@@ -1,6 +1,6 @@
 # model based off https://towardsdatascience.com/going-beyond-99-mnist-handwritten-digits-recognition-cfff96337392
 from typing import List, Callable
-from tinygrad import Tensor, TinyJit, nn, GlobalCounters, Device, Context
+from tinygrad import Tensor, TinyJit, nn, GlobalCounters, Device
 from tinygrad.helpers import getenv, colored, trange
 from tinygrad.nn.datasets import mnist

@@ -31,7 +31,7 @@ if __name__ == "__main__":

  @TinyJit
  def train_step() -> Tensor:
-    with Context(TRAINING=1):
+    with Tensor.train():
      opt.zero_grad()
      samples = Tensor.randint(getenv("BS", 512), high=X_train.shape[0])
      Xt, Yt = X_train[samples].shard_(GPUS, axis=0), Y_train[samples].shard_(GPUS, axis=0)  # we shard the data on axis 0
@@ -35,11 +35,12 @@ def compile_onnx_model(onnx_model):
  tinyonnx = TinyOnnx(onnx_model)
  the_input = Tensor.randn(1,32)

-  linear, output_bufs = jit_model(tinyonnx, the_input)
-  the_output = [tinyonnx.forward(the_input)]
+  run, special_names = jit_model(tinyonnx, the_input)

-  functions, statements, bufs, bufs_to_save = compile_net(linear, output_bufs)
+  functions, statements, bufs, bufs_to_save = compile_net(run, special_names)
  prg = export_model_clang(functions, statements, bufs, {}, ["input0"], ["output0"])
+
+  the_output = run(the_input)
  cprog = ["#include <string.h>", "#include <stdio.h>", "#include <stdlib.h>"]
  cprog.append(prg)

@@ -5,9 +5,8 @@ with contextlib.suppress(ImportError): import tiktoken
 from tinygrad import Tensor, TinyJit, Device, GlobalCounters, Variable, dtypes
 from tinygrad.uop.ops import UOp
 from tinygrad.helpers import Timing, DEBUG, JIT, getenv, fetch, colored, trange
-from tinygrad.llm.gguf import gguf_load
 from tinygrad.nn import Embedding, Linear, LayerNorm
-from tinygrad.nn.state import torch_load, load_state_dict, get_state_dict
+from tinygrad.nn.state import gguf_load, torch_load, load_state_dict, get_state_dict
 from extra.bench_log import BenchEvent, WallTimeEvent

 MAX_CONTEXT = getenv("MAX_CONTEXT", 128)
@@ -1,6 +1,6 @@
 import itertools
 from typing import Callable
-from tinygrad import nn, Tensor, dtypes, Device, TinyJit, Context
+from tinygrad import nn, Tensor, dtypes, Device, TinyJit
 from tinygrad.helpers import getenv, trange, partition

 class Model:
@@ -35,21 +35,22 @@ if __name__ == "__main__":

  params = nn.state.get_parameters(model)

-  # init params
+  # init params, set requires grad on the ones we need gradients of
  for x in params:
+    if x.requires_grad is None: x.requires_grad_()
    x.replace(x.contiguous())
  Tensor.realize(*params)

  # split params (with grads) and buffers (without)
-  params, buffers = partition(params, lambda x: x.is_param)
+  params, buffers = partition(params, lambda x: x.requires_grad)
  print(f"params: {len(params)} buffers: {len(buffers)}")

  # optim params
  pos_params = list(itertools.accumulate(params, lambda x,y: x+y.numel(), initial=0))
  adam_m = Tensor.zeros(pos_params[-1], device="CPU").contiguous()
  adam_v = Tensor.zeros(pos_params[-1], device="CPU").contiguous()
-  adam_b1_t = Tensor.ones((1,), dtype=dtypes.float32, device="CPU").contiguous()
-  adam_b2_t = Tensor.ones((1,), dtype=dtypes.float32, device="CPU").contiguous()
+  adam_b1_t = Tensor.ones((1,), dtype=dtypes.float32, device="CPU", requires_grad=False).contiguous()
+  adam_b2_t = Tensor.ones((1,), dtype=dtypes.float32, device="CPU", requires_grad=False).contiguous()
  adam_params = [adam_m, adam_v, adam_b1_t, adam_b2_t]

  # create loss and grads. init all state so the JIT works on microbatch
@@ -59,7 +60,7 @@ if __name__ == "__main__":
  Tensor.realize(*params, *buffers, *adam_params, loss, grads)

  @TinyJit
-  @Context(TRAINING=1)
+  @Tensor.train()
  def microbatch():
    samples = Tensor.randint(BS // ACC_STEPS, high=X_train.shape[0])
    for t in params: t.grad = None
@@ -10,7 +10,7 @@ from extra.lr_scheduler import OneCycleLR
 from tinygrad import nn, dtypes, Tensor, Device, GlobalCounters, TinyJit, Variable
 from tinygrad.nn.state import get_state_dict
 from tinygrad.nn import optim
-from tinygrad.helpers import Context, BEAM, WINO, getenv, colored, prod, TRAINING
+from tinygrad.helpers import Context, BEAM, WINO, getenv, colored, prod
 from extra.bench_log import BenchEvent, WallTimeEvent

 cifar_mean = [0.4913997551666284, 0.48215855929893703, 0.4465309133731618]
@@ -30,9 +30,9 @@ class UnsyncedBatchNorm:
    if affine: self.weight, self.bias = Tensor.ones(sz, dtype=dtypes.float32), Tensor.zeros(sz, dtype=dtypes.float32)
    else: self.weight, self.bias = None, None

-    self.running_mean = Tensor.zeros(num_devices, sz, dtype=dtypes.float32).is_param_(False)
-    self.running_var = Tensor.ones(num_devices, sz, dtype=dtypes.float32).is_param_(False)
-    self.num_batches_tracked = Tensor.zeros(1, dtype=dtypes.int).is_param_(False)
+    self.running_mean = Tensor.zeros(num_devices, sz, dtype=dtypes.float32, requires_grad=False)
+    self.running_var = Tensor.ones(num_devices, sz, dtype=dtypes.float32, requires_grad=False)
+    self.num_batches_tracked = Tensor.zeros(1, dtype=dtypes.int, requires_grad=False)

  def __call__(self, x:Tensor):
    xr = x.reshape(self.num_devices, -1, *x.shape[1:]).cast(dtypes.float32)
@@ -44,7 +44,7 @@ class UnsyncedBatchNorm:
    return ret.reshape(x.shape).cast(x.dtype)

  def calc_stats(self, x:Tensor):
-    if TRAINING:
+    if Tensor.training:
      # This requires two full memory accesses to x
      # https://github.com/pytorch/pytorch/blob/c618dc13d2aa23625cb0d7ada694137532a4fa33/aten/src/ATen/native/cuda/Normalization.cuh
      # There's "online" algorithms that fix this, like https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_Online_algorithm
@@ -68,7 +68,8 @@ class UnsyncedBatchNorm:
 class BatchNorm(nn.BatchNorm2d if getenv("SYNCBN") else UnsyncedBatchNorm):
  def __init__(self, num_features):
    super().__init__(num_features, track_running_stats=False, eps=1e-12, momentum=0.85, affine=True)
-    self.weight.is_param_(False)
+    self.weight.requires_grad = False
+    self.bias.requires_grad = True

 class ConvGroup:
  def __init__(self, channels_in, channels_out):
@@ -152,21 +153,26 @@ def train_cifar():

  # ========== Model ==========
  def whitening(X, kernel_size=hyp['net']['kernel_size']):
-    def _patches(data:Tensor, patch_size=(kernel_size,kernel_size)):
+    def _cov(X):
+      return (X.T @ X) / (X.shape[0] - 1)
+
+    def _patches(data, patch_size=(kernel_size,kernel_size)):
      h, w = patch_size
-      _, c, _, _ = data.shape
-      return data._pool((h, w)).permute(1, 4, 5, 0, 3, 2).reshape(c*h*w, -1)
+      c = data.shape[1]
+      axis = (2, 3)
+      return np.lib.stride_tricks.sliding_window_view(data, window_shape=(h,w), axis=axis).transpose((0,3,2,1,4,5)).reshape((-1,c,h,w))

    def _eigens(patches):
-      cov = ((patches @ patches.T) / (patches.shape[1] - 1)).numpy()
-      eigvals, eigvecs = np.linalg.eigh(cov, UPLO='U')
-      return np.flip(eigvals, 0), np.flip(eigvecs.T.reshape(patches.shape[0], X.shape[1], kernel_size, kernel_size), 0)
+      n,c,h,w = patches.shape
+      Σ = _cov(patches.reshape(n, c*h*w))
+      Λ, V = np.linalg.eigh(Σ, UPLO='U')
+      return np.flip(Λ, 0), np.flip(V.T.reshape(c*h*w, c, h, w), 0)

    # NOTE: np.linalg.eigh only supports float32 so the whitening layer weights need to be converted to float16 manually
-    eigvals, eigvecs = _eigens(_patches(X.float()))
-    W = eigvecs/np.sqrt(eigvals+1e-2)[:,None,None,None]
+    Λ, V = _eigens(_patches(X.float().numpy()))
+    W = V/np.sqrt(Λ+1e-2)[:,None,None,None]

-    return Tensor(W.astype(np.float32)).cast(dtypes.default_float).is_param_(False)
+    return Tensor(W.astype(np.float32), requires_grad=False).cast(dtypes.default_float)

  # ========== Loss ==========
  def cross_entropy(x:Tensor, y:Tensor, reduction:str='mean', label_smoothing:float=0.0) -> Tensor:
@@ -218,7 +224,7 @@ def train_cifar():

  @TinyJit
  def augmentations(X:Tensor, Y:Tensor):
-    perms = Tensor.randperm(X.shape[0], device=X.device) # We reuse perms for cutmix, because they are expensive to generate
+    perms = Tensor.randperm(X.shape[0], device=X.device) # We reuse perms for cutmix, because they are expensivne to generate
    if getenv("RANDOM_CROP", 1):
      X = random_crop(X, crop_size=32)
    if getenv("RANDOM_FLIP", 1):
@@ -258,6 +264,7 @@ def train_cifar():
      # self.model_ema = copy.deepcopy(net) # won't work for opencl due to unpickeable pyopencl._cl.Buffer
      self.net_ema = SpeedyResNet(w)
      for net_ema_param, net_param in zip(get_state_dict(self.net_ema).values(), get_state_dict(net).values()):
+        net_ema_param.requires_grad = False
        net_ema_param.assign(net_param.numpy())

    @TinyJit
@@ -300,7 +307,7 @@ def train_cifar():
  params_bias = []
  params_non_bias = []
  for params in params_dict:
-    if params_dict[params].is_param:
+    if params_dict[params].requires_grad is not False:
      if 'bias' in params:
        params_bias.append(params_dict[params])
      else:
@@ -309,9 +316,6 @@ def train_cifar():
  opt_bias     = optim.SGD(params_bias,     lr=0.01, momentum=hyp['opt']['momentum'], nesterov=True, weight_decay=hyp['opt']['bias_decay'])
  opt_non_bias = optim.SGD(params_non_bias, lr=0.01, momentum=hyp['opt']['momentum'], nesterov=True, weight_decay=hyp['opt']['non_bias_decay'])

-  # realize model params and optimizer state before JIT to avoid cache misses
-  Tensor.realize(*params_dict.values(), *opt_bias.b, *opt_non_bias.b)
-
  # NOTE taken from the hlb_CIFAR repository, might need to be tuned
  initial_div_factor = hyp['opt']['initial_div_factor']
  final_lr_ratio = hyp['opt']['final_lr_ratio']
@@ -328,7 +332,9 @@ def train_cifar():
      # index 0 for bias and 1 for non-bias
      optimizer.zero_grad()
      loss.backward()
-      return loss.realize(*optimizer.schedule_step(), *lr_scheduler[0].schedule_step(), *lr_scheduler[1].schedule_step())
+      optimizer.step()
+      lr_scheduler[0].step()
+      lr_scheduler[1].step()
    return loss.realize()

  train_step_jitted = TinyJit(train_step)
@@ -355,11 +361,11 @@ def train_cifar():
  i = 0
  eval_acc_pct = 0.0
  batcher = fetch_batches(X_train, Y_train, BS=BS, is_train=True)
-  with Context(TRAINING=1):
+  with Tensor.train():
    st = time.monotonic()
    while i <= STEPS:
      if i % getenv("EVAL_STEPS", STEPS) == 0 and i > 1 and not getenv("DISABLE_BACKWARD"):
-        # Using Context(TRAINING=0) here actually bricks batchnorm, even with track_running_stats=True
+        # Use Tensor.training = False here actually bricks batchnorm, even with track_running_stats=True
        corrects = []
        corrects_ema = []
        losses = []
@@ -2,8 +2,7 @@ from pathlib import Path
 from typing import List
 import json, argparse, random, time, os
 from extra.models.llama import Transformer, convert_from_huggingface, convert_from_gguf, fix_bf16
-from tinygrad.llm.gguf import gguf_load
-from tinygrad.nn.state import safe_load, torch_load, load_state_dict, get_parameters
+from tinygrad.nn.state import safe_load, torch_load, load_state_dict, get_parameters, gguf_load
 from tinygrad import Tensor, dtypes, nn, Context, Device, GlobalCounters
 from tinygrad.helpers import Profiling, Timing, DEBUG, colored, fetch, tqdm
 from extra.bench_log import BenchEvent, WallTimeEvent
@@ -102,7 +101,7 @@ class Int8Embedding:
    self.weight, self.scale = Tensor.ones(vocab_size, embed_size, dtype=dtypes.int8), Tensor.ones(vocab_size, dtype=dtypes.half)

  def __call__(self, idx:Tensor) -> Tensor:
-    if not hasattr(self, 'arange'): self.arange = Tensor.arange(self.vocab_sz).unsqueeze(-1)
+    if not hasattr(self, 'arange'): self.arange = Tensor.arange(self.vocab_sz, requires_grad=False, device=self.weight.device).unsqueeze(-1)
    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.cast(self.scale.dtype).T*self.scale).T
    return (arange == idx).mul(vals).sum(-2, dtype=vals.dtype)
@@ -123,7 +122,7 @@ def NF4Linear(block_size):
    def __call__(self, x: Tensor) -> Tensor:
      high_bits = self.weight
      low_bits = (self.weight * 2 ** 4).contiguous()
-      unpacked = Tensor.stack(high_bits, low_bits, dim=-1).div(2 ** 4, rounding_mode="trunc")
+      unpacked = Tensor.stack(high_bits, low_bits, dim=-1).idiv(2 ** 4)
      unscaled = CODE[unpacked].to(x.device).reshape(-1, block_size) * self.scale
      return x.linear(unscaled.reshape(self.out_features, self.in_features).T)

@@ -3,7 +3,7 @@ import os
 if "NOOPT" not in os.environ: os.environ["NOOPT"] = "1"
 from tinygrad import Device, nn, Tensor, dtypes
 from train_gpt2 import GPT, GPTConfig
-from tinygrad.helpers import DEV, dedup, flatten, getenv, GlobalCounters, to_function_name, Context
+from tinygrad.helpers import DEV, dedup, flatten, getenv, GlobalCounters, to_function_name
 from tinygrad.engine.realize import get_kernel
 from tinygrad.schedule.memory import memory_planner
 from tinygrad.uop.ops import Ops
@@ -23,23 +23,23 @@ if __name__ == "__main__":
  #B, T = Variable("B", 1, 128).bind(4), 64 #Variable("T", 1, 1024).bind(64)
  B, T = 4, 64

+  Tensor.training = True
  optimizer = nn.optim.Adam(nn.state.get_parameters(model), lr=1e-4)
  warmup_count = getenv("WARMUP", 3)
-  with Context(TRAINING=1):
-    for i in range(warmup_count):  # TODO: why does it take three and not two to stabilize
-      GlobalCounters.reset()
-      X = Tensor.empty(4, 64, dtype=dtypes.int).reshape(B, T)
-      Y = Tensor.empty(4, 64, dtype=dtypes.int).reshape(B, T)
-      _, loss = model(X, Y)
-      optimizer.zero_grad()
-      if getenv("BACKWARD", 1):
-        loss.backward()
-        tensors = optimizer.schedule_step()
-      else:
-        tensors = []
-      sched = loss.schedule(*tensors)
-      print(f"calls {i}:", len(sched))
-      #run_schedule(sched[:])
+  for i in range(warmup_count):  # TODO: why does it take three and not two to stabilize
+    GlobalCounters.reset()
+    X = Tensor.empty(4, 64, dtype=dtypes.int).reshape(B, T)
+    Y = Tensor.empty(4, 64, dtype=dtypes.int).reshape(B, T)
+    _, loss = model(X, Y)
+    optimizer.zero_grad()
+    if getenv("BACKWARD", 1):
+      loss.backward()
+      tensors = optimizer.schedule_step()
+    else:
+      tensors = []
+    sched = loss.schedule(*tensors)
+    print(f"calls {i}:", len(sched))
+    #run_schedule(sched[:])
  sched = memory_planner(sched)
  ast_dedup = dedup([si.ast for si in sched if si.ast.op is Ops.SINK])
  srcs = {}
@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 import os, math, time
 import numpy as np
-from tinygrad import Tensor, nn, fetch, Device, TinyJit, GlobalCounters, Context
+from tinygrad import Tensor, nn, fetch, Device, TinyJit, GlobalCounters
 from dataclasses import dataclass

@dataclass
@@ -25,7 +25,7 @@ class CausalSelfAttention:
    self.n_embd = config.n_embd
    # not really a 'bias', more of a mask, but following the OpenAI/HF naming though
    self.bias = Tensor.ones(1, 1, config.block_size, config.block_size).tril()
-    self.bias.is_param_(False)
+    self.bias.requires_grad = False

  def __call__(self, x:Tensor):
    B, T, C = x.shape
@@ -99,7 +99,7 @@ class GPT:

  def __call__(self, idx:Tensor, targets=None):
    b, t = idx.shape
-    pos = Tensor.arange(0, t)
+    pos = Tensor.arange(0, t, device=idx.device)

    tok_emb = self.wte(idx) # token embeddings of shape (b, t, n_embd)
    pos_emb = self.wpe(pos) # position embeddings of shape (t, n_embd)
@@ -177,7 +177,7 @@ if __name__ == "__main__":
  if args.gpus > 1: x, y = x.shard(GPUS, axis=0), y.shard(GPUS, axis=0)

  @TinyJit
-  @Context(TRAINING=1)
+  @Tensor.train()
  def step(x:Tensor, y:Tensor) -> Tensor:
    _, loss = model(x, y)
    optimizer.zero_grad()
@@ -204,3 +204,4 @@ if __name__ == "__main__":
    top_k = 40
    y = model.generate(x, max_new_tokens, temperature=temperature, top_k=top_k)
    print(decode(y[0].tolist()))
+
@@ -1,5 +1,5 @@
 # much taken from https://github.com/cloneofsimo/minRF
-from tinygrad import Tensor, nn, GlobalCounters, TinyJit, Context
+from tinygrad import Tensor, nn, GlobalCounters, TinyJit
 from tinygrad.helpers import getenv, trange
 from extra.models.llama import Attention, FeedForward, precompute_freqs_cis

@@ -135,7 +135,7 @@ if __name__ == "__main__":
  optimizer = nn.optim.Adam(nn.state.get_parameters(model), lr=5e-4)

  @TinyJit
-  @Context(TRAINING=1)
+  @Tensor.train()
  def train_step():
    if getenv("OVERFIT"): samples = Tensor.zeros(getenv("BS", 256), dtype='int')
    else: samples = Tensor.randint(getenv("BS", 256), high=X_train.shape[0])
@@ -1,6 +1,6 @@
 import functools, argparse, pathlib
 from tinygrad import Tensor, nn, Device, GlobalCounters, Variable
-from tinygrad.helpers import Timing, Profiling, tqdm
+from tinygrad.helpers import Timing, Profiling, CI, tqdm
 from tinygrad.nn.state import torch_load, get_state_dict
 from extra.models.llama import FeedForward, Transformer
 from extra.bench_log import BenchEvent, WallTimeEvent
@@ -36,7 +36,7 @@ if __name__ == "__main__":
    model = Transformer(n_layers=32, dim=4096, hidden_dim=14336, n_heads=32, n_kv_heads=8, norm_eps=1e-5, vocab_size=32000, feed_forward=functools.partial(MixtureFeedForward, 8), jit=False)
    model_state_dict = get_state_dict(model)

-    for k in (t := tqdm(state, disable=None)):
+    for k in (t := tqdm(state, disable=CI)):
      if 'feed_forward.experts.' in k:
        expert_no = int(k.split('feed_forward.experts.')[1].split('.')[0])
        device = Device.DEFAULT + ":" + str((expert_no//2)+1)
@@ -44,7 +44,7 @@ if __name__ == "__main__":
        device = Device.DEFAULT
      t.set_description(f"ram used: {GlobalCounters.mem_used/1e9:5.2f} GB, loading {k} to {device}")
      model_state_dict[k].replace(state[k].to(device).half()).realize()
-  if t.disable: print(f"ram used: {GlobalCounters.mem_used/1e9:5.2f} GB")
+  if CI: print(f"ram used: {GlobalCounters.mem_used/1e9:5.2f} GB")

  from sentencepiece import SentencePieceProcessor
  spp = SentencePieceProcessor(model_file=args.weights + "/tokenizer.model")
@@ -2,7 +2,7 @@ import math
 from typing import Union

 from tinygrad import Tensor, nn, dtypes
-from tinygrad.helpers import prod, argfix, Context, TRAINING
+from tinygrad.helpers import prod, argfix, Context
 from tinygrad.nn.state import get_parameters
 from extra.models.unet import UNetModel

@@ -57,7 +57,7 @@ class EmbeddingBert(nn.Embedding):
  def __call__(self, idx:Tensor) -> Tensor:
    if idx.numel() == 0: return Tensor.empty(idx.shape+(self.embed_sz,), dtype=self.weight.dtype, device=self.weight.device)
    arange_shp, weight_shp, big_shp = (1, 1, self.vocab_sz, 1), (1, 1, self.vocab_sz, self.embed_sz), idx.shape+(self.vocab_sz, self.embed_sz,)
-    if not hasattr(self, 'arange'): self.arange = Tensor.arange(self.vocab_sz).reshape(arange_shp)
+    if not hasattr(self, 'arange'): self.arange = Tensor.arange(self.vocab_sz, requires_grad=False, device=self.weight.device).reshape(arange_shp)
    arange, idx, vals = self.arange.expand(big_shp), idx.reshape(idx.shape+(1, 1,)).expand(big_shp), self.weight.cast(dtypes.default_float).reshape(weight_shp).expand(big_shp)
    return (arange == idx).where(vals, 0).sum(2, dtype=vals.dtype)

@@ -77,15 +77,15 @@ class FrozenBatchNorm2dRetinaNet(nn.BatchNorm2d):
  def __init__(self, sz:int, eps=1e-5, affine=True, track_running_stats=True, momentum=0.1):
    self.eps, self.track_running_stats, self.momentum = eps, track_running_stats, momentum

-    self.weight = Tensor.ones(sz, dtype=dtypes.float32).is_param_(False) if affine else None
-    self.bias = Tensor.zeros(sz, dtype=dtypes.float32).is_param_(False) if affine else None
+    self.weight = Tensor.ones(sz, dtype=dtypes.float32, requires_grad=False) if affine else None
+    self.bias = Tensor.zeros(sz, dtype=dtypes.float32, requires_grad=False) if affine else None

-    if track_running_stats: self.running_mean, self.running_var = Tensor.zeros(sz, dtype=dtypes.float32).is_param_(False), Tensor.ones(sz, dtype=dtypes.float32).is_param_(False)
-    self.num_batches_tracked = Tensor.zeros(1, dtype=dtypes.long).is_param_(False)
+    if track_running_stats: self.running_mean, self.running_var = Tensor.zeros(sz, dtype=dtypes.float32, requires_grad=False), Tensor.ones(sz, dtype=dtypes.float32, requires_grad=False)
+    self.num_batches_tracked = Tensor.zeros(1, dtype=dtypes.long, requires_grad=False)

  def __call__(self, x:Tensor) -> Tensor:
    batch_mean, batch_var = super().calc_stats(x.cast(dtypes.float32))
-    if self.track_running_stats and TRAINING:
+    if self.track_running_stats and Tensor.training:
      self.running_mean.assign((1-self.momentum) * self.running_mean + self.momentum * batch_mean.detach().cast(self.running_mean.dtype))
      self.running_var.assign((1-self.momentum) * self.running_var + self.momentum * x.numel()/(x.numel()-x.shape[1]) * batch_var.detach().cast(self.running_var.dtype))
      self.num_batches_tracked += 1
@@ -358,7 +358,7 @@ def eval_stable_diffusion():
      batch = batch.cat(batch[-1:].expand(bs - unpadded_bs, *batch[-1].shape))
    return batch, unpadded_bs 

-  @Context(TRAINING=0)
+  @Tensor.train(mode=False)
  def eval_unet(eval_inputs:list[dict], unet:UNetModel, cond_stage:FrozenOpenClipEmbedder, first_stage:AutoencoderKL,
                inception:FidInceptionV3, clip:OpenClipEncoder) -> tuple[float, float]:
    # Eval is divided into 5 jits, one per model
@@ -498,10 +498,11 @@ def eval_stable_diffusion():

 if __name__ == "__main__":
  # inference only
+  Tensor.training = False
+
  models = getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert").split(",")
-  with Context(TRAINING=0):
-    for m in models:
-      nm = f"eval_{m}"
-      if nm in globals():
-        print(f"eval {m}")
-        globals()[nm]()
+  for m in models:
+    nm = f"eval_{m}"
+    if nm in globals():
+      print(f"eval {m}")
+      globals()[nm]()
@@ -1,6 +1,6 @@
 # load each model here, quick benchmark
 from tinygrad import Tensor, GlobalCounters
-from tinygrad.helpers import getenv, Context
+from tinygrad.helpers import getenv
 import numpy as np

 def test_model(model, *inputs):
@@ -59,10 +59,11 @@ def spec_mrcnn():

 if __name__ == "__main__":
  # inference only for now
-  with Context(TRAINING=0):
-    for m in getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert,mrcnn").split(","):
-      nm = f"spec_{m}"
-      if nm in globals():
-        print(f"testing {m}")
-        globals()[nm]()
+  Tensor.training = False
+
+  for m in getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert,mrcnn").split(","):
+    nm = f"spec_{m}"
+    if nm in globals():
+      print(f"testing {m}")
+      globals()[nm]()

@@ -2,7 +2,7 @@ import os, time, math, functools, random, contextlib
 from pathlib import Path
 import multiprocessing

-from tinygrad import Device, GlobalCounters, Tensor, TinyJit, dtypes, Context
+from tinygrad import Device, GlobalCounters, Tensor, TinyJit, dtypes
 from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, Profiling, profile_marker, DEBUG
 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
@@ -157,7 +157,6 @@ def train_resnet():
  # input_std = Tensor([0.229, 0.224, 0.225], device=GPUS, dtype=dtypes.float32).reshape(1, -1, 1, 1)
  def normalize(x): return (x.permute([0, 3, 1, 2]) - input_mean).cast(dtypes.default_float)
  @TinyJit
-  @Context(TRAINING=1)
  def train_step(X, Y):
    optimizer_group.zero_grad()
    X = normalize(X)
@@ -171,7 +170,6 @@ def train_resnet():
    return loss.realize(), top_1.realize()

  @TinyJit
-  @Context(TRAINING=0)
  def eval_step(X, Y):
    X = normalize(X)
    out = model.forward(X)
@@ -182,11 +180,11 @@ def train_resnet():
  def fake_data_get(batch_size):
    x = Tensor.zeros(batch_size, 224, 224, 3, dtype=dtypes.uchar).contiguous()
    y = [0] * batch_size
-    return x.shard(GPUS, axis=0).realize(), Tensor(y).shard(GPUS, axis=0), y, None
+    return x.shard(GPUS, axis=0).realize(), Tensor(y, requires_grad=False).shard(GPUS, axis=0), y, None

  def data_get(it):
    x, y, cookie = next(it)
-    return x.shard(GPUS, axis=0).realize(), Tensor(y).shard(GPUS, axis=0), y, cookie
+    return x.shard(GPUS, axis=0).realize(), Tensor(y, requires_grad=False).shard(GPUS, axis=0), y, cookie

  # ** epoch loop **
  step_times = []
@@ -194,6 +192,7 @@ def train_resnet():
    # ** train loop **
    if MLLOGGER and RUNMLPERF:
      MLLOGGER.start(key=mllog_constants.EPOCH_START, value=e+1, metadata=dict(epoch_num=e+1))
+    Tensor.training = True
    BEAM.value = TRAIN_BEAM

    if INITMLPERF:
@@ -272,6 +271,7 @@ def train_resnet():
      eval_loss = 0.0
      eval_top_1 = 0
      eval_num_samples = 0
+      Tensor.training = False
      BEAM.value = EVAL_BEAM

      if INITMLPERF:
@@ -413,7 +413,7 @@ def train_retinanet():
    layers_to_train = ["layer4", "layer3", "layer2", "layer1", "conv1"][:trainable_layers]
    for k, v in get_state_dict(backbone).items():
      if all([not k.startswith(layer) for layer in layers_to_train]):
-        v.is_param_(False)
+        v.requires_grad = False

  def _data_get(it:Iterator[tuple[Tensor, ...]], val:bool=False):
    if val:
@@ -614,7 +614,7 @@ def train_retinanet():

      if getenv("RESET_STEP", 1): _train_step.reset()

-      with Context(TRAINING=0):
+      with Tensor.train(mode=False):
        if not RUNMLPERF:
          i, proc = 0, _fake_data_get(EVAL_BS, val=(val:=True))
        else:
@@ -784,7 +784,7 @@ def train_unet3d():
    return x.shard(GPUS, axis=0).realize(), y.shard(GPUS, axis=0), cookie

  @TinyJit
-  @Context(TRAINING=1)
+  @Tensor.train()
  def train_step(model, x, y):
    optim.zero_grad()

@@ -795,10 +795,10 @@ def train_unet3d():
    optim.step()
    return loss.realize()

-  @Context(TRAINING=0)
+  @Tensor.train(mode=False)
  def eval_step(model, x, y):
    y_hat, y = sliding_window_inference(model, x, y, gpus=GPUS)
-    y_hat, y = Tensor(y_hat), Tensor(y)
+    y_hat, y = Tensor(y_hat), Tensor(y, requires_grad=False)
    loss = dice_ce_loss(y_hat, y)
    score = dice_score(y_hat, y)
    return loss.realize(), score.realize()
@@ -919,7 +919,6 @@ def train_rnnt():
  pass

@TinyJit
-@Context(TRAINING=0)
 def eval_step_bert(model, input_ids:Tensor, segment_ids:Tensor, attention_mask:Tensor, masked_positions:Tensor, masked_lm_ids:Tensor,
                   masked_lm_weights:Tensor, next_sentence_labels:Tensor, GPUS):
  for t in [input_ids, segment_ids, attention_mask, masked_positions, masked_lm_ids, masked_lm_weights, next_sentence_labels]:
@@ -1107,7 +1106,6 @@ def train_bert():
      MLLOGGER.start(key=mllog_constants.EPOCH_START, value=i*GBS, metadata={"epoch_num": i*GBS})

  @TinyJit
-  @Context(TRAINING=1)
  def train_step_bert(input_ids:Tensor, segment_ids:Tensor, attention_mask:Tensor,
                      masked_positions:Tensor, masked_lm_ids:Tensor, masked_lm_weights:Tensor, next_sentence_labels:Tensor):
    for t in [input_ids, segment_ids, attention_mask, masked_positions, masked_lm_ids, masked_lm_weights, next_sentence_labels]:
@@ -1135,6 +1133,7 @@ def train_bert():

  while train_data is not None and i < train_steps and not achieved:
    if getenv("TRAIN", 1):
+      Tensor.training = True
      BEAM.value = TRAIN_BEAM
      st = time.perf_counter()
      GlobalCounters.reset()
@@ -1187,6 +1186,7 @@ def train_bert():
      eval_lm_accs = []
      eval_clsf_accs = []
      eval_times = []
+      Tensor.training = False
      BEAM.value = EVAL_BEAM

      for j in tqdm(range(max_eval_steps), desc="Evaluating", total=max_eval_steps, disable=BENCHMARK):
@@ -1282,7 +1282,7 @@ def train_bert():
        previous_step = i

 def train_llama3():
-  from examples.mlperf.models.flat_llama import FlatTransformer, apply_grad, FP8_DTYPE, MXFP8
+  from examples.mlperf.models.flat_llama import FlatTransformer, apply_grad, FP8
  from examples.llama3 import MODEL_PARAMS
  from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
  from examples.mlperf.optim import GradAccClipAdamW
@@ -1357,7 +1357,6 @@ def train_llama3():
      MLLOGGER.event(key=mllog_constants.OPT_LR_WARMUP_STEPS, value=WARMUP_STEPS)
      MLLOGGER.event(key=mllog_constants.NUM_WARMUP_STEPS, value=WARMUP_STEPS)
      MLLOGGER.event(key=mllog_constants.OPT_LR_DECAY_STEPS, value=MAX_STEPS - WARMUP_STEPS)
-      MLLOGGER.event(key=mllog_constants.OPT_LR_DECAY_SCHEDULE, value="cosine with linear warmup")
      MLLOGGER.event(key=mllog_constants.OPT_GRADIENT_CLIP_NORM, value=1.0)
  else:
    MLLOGGER = None
@@ -1396,7 +1395,7 @@ def train_llama3():

  params = get_parameters(model)

-  if getenv("EMPTYWEIGHT"):
+  if getenv("FAKEDATA"):
    for v in get_parameters(model):
      v = v.assign(Tensor.empty(v.shape, dtype=v.dtype))

@@ -1417,9 +1416,9 @@ def train_llama3():
  optim = GradAccClipAdamW(params, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2,
                           eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=grad_acc, device=optim_device)

+  # init grads
  for p in optim.params:
-    grad_dtype = dtypes.bfloat16 if p.dtype == FP8_DTYPE else p.dtype
-    p.grad = p.zeros_like(dtype=grad_dtype).contiguous()
+    p.grad = Tensor.zeros(p.shape, dtype=p.dtype, device=p.device).contiguous()
  grads = [p.grad for p in optim.params]

  scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)
@@ -1433,68 +1432,37 @@ def train_llama3():
    print(f"loading optim checkpoint from {fn}")
    load_state_dict(scheduler, safe_load(fn), realize=False)

-  fp8_amax = [t for ts in model._fp8_amax.values() for t in ts]
-  fp8_next_amax = [t for ts in model._fp8_next_amax.values() for t in ts] if hasattr(model, "_fp8_next_amax") else []
-  fp8_grad_amax = [t for ts in model._fp8_grad_amax.values() for t in ts] if hasattr(model, "_fp8_grad_amax") else []
-  fp8_next_grad_amax = [t for ts in model._fp8_next_grad_amax.values() for t in ts] if hasattr(model, "_fp8_next_grad_amax") else []
-  fp8_inv_scales = list(model._fp8_inv_scale.values()) + list(model._fp8_next_inv_scale.values())
-
-  from tinygrad.nn.state import get_state_dict
-  model_state = get_state_dict(model)
-  for wname in model._fp8_inv_scale:
-    w = model_state[wname]
-    w._inv_scale = model._fp8_inv_scale[wname]
-    w._next_inv_scale = model._fp8_next_inv_scale[wname]
-    if optim.master_params:
-      idx = next(j for j, p in enumerate(optim.params) if p is w)
-      master = optim.master_params[idx]
-      inv = w._inv_scale if w._inv_scale.device == master.device else w._inv_scale.to(master.device)
-      if MXFP8:
-        from extra.gemm.cdna_asm_gemm import _mx_block_scale
-        bs = _mx_block_scale(inv.reshape(-1, inv.shape[-1])).reshape(w.shape)
-        master.assign((master * bs).contiguous())
-      else:
-        master.assign((master * inv.reshape(*inv.shape, *([1]*(w.ndim-inv.ndim)))).contiguous())
-
-  # realize everything here
-  if optim.master_params: Tensor.realize(*optim.master_params)
-  Tensor.realize(*optim.params, *fp8_inv_scales, *fp8_amax, *fp8_next_amax, *fp8_grad_amax, *fp8_next_grad_amax)
+  fp8_amax = [t for ts in model._fp8_amax.values() for t in ts] if FP8 else []

  @TinyJit
  def minibatch(tokens:Tensor):
    if is_dp: tokens = tokens.to(None).shard(device, 0)
    if is_mp: tokens = tokens.shard(device)
    if not is_sharding: tokens = tokens.to(None)
-    logits:Tensor = model(tokens[:, :-1], save=bool(SMALL))
-    if getenv("FAST_CE", 0):
-      from extra.llama_kernels.fused_ce import fused_ce_loss
-      loss = fused_ce_loss(logits.cast(dtypes.bfloat16), tokens[:, 1:], label_smoothing=0.0)
-    else:
-      loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
+    logits:Tensor = model(tokens[:, :-1])
+    loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])

    for g, new_g in zip(grads, loss.gradient(*optim.params)):
      apply_grad(g, new_g.uop)

    loss_cpu = loss.flatten().float().to("CPU")
-    return loss_cpu.realize(*grads, *fp8_amax, *fp8_next_amax, *fp8_grad_amax, *fp8_next_grad_amax)
+    return loss_cpu.realize(*grads, *fp8_amax)

  @TinyJit
  def optim_step():
    grad_norm = optim.fstep(grads)
    scheduler.step()

-    for g in grads: g.assign(0)
-    for cur, nxt in zip(fp8_amax, fp8_next_amax): cur.assign(nxt)
-    for cur, nxt in zip(fp8_grad_amax, fp8_next_grad_amax): cur.assign(nxt)
+    for g in grads: g.assign(g.zeros_like())

    lr_cpu = optim.lr.float().to("CPU")
    grad_norm_cpu = grad_norm.float().to("CPU")
-    Tensor.realize(lr_cpu, grad_norm_cpu, *grads, *fp8_inv_scales, *fp8_amax, *fp8_grad_amax)
+    Tensor.realize(lr_cpu, grad_norm_cpu, *grads)

    return lr_cpu, grad_norm_cpu

  @TinyJit
-  @Context(TRAINING=0)
+  @Tensor.train(False)
  def eval_step(tokens:Tensor):
    if is_dp: tokens = tokens.to(None).shard(device, 0)
    if is_mp: tokens = tokens.shard(device)
@@ -1507,7 +1475,7 @@ def train_llama3():
  def fake_data(bs, samples):
    import numpy as np
    for _ in range(samples // bs):
-      fake_data_np = np.random.randint(0, real_vocab_size, size=(bs, SEQLEN + 1), dtype=np.int32)
+      fake_data_np = np.random.randint(0, model_params["vocab_size"], size=(bs, SEQLEN + 1), dtype=np.int32)
      yield Tensor(fake_data_np, device="NPY")

  def get_train_iter():
@@ -1576,7 +1544,7 @@ def train_llama3():

      mem_gb = GlobalCounters.mem_used / 1e9
      gflops = GlobalCounters.global_ops / 1e9 / dev_time
-      mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * device_count * 4.6e15)) * 100
+      mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * device_count * (4.6e15 if FP8 else 2.3e15))) * 100
      tqdm.write(
          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, {grad_norm:.6f} grad_norm, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")
@@ -1653,6 +1621,7 @@ def train_llama3():
        tqdm.write(f"target achieved after {sequences_seen} sequences")
        if MLLOGGER and RUNMLPERF:
          MLLOGGER.end(key=mllog_constants.EPOCH_STOP, metadata={mllog_constants.SAMPLES_COUNT: sequences_seen})
+          MLLOGGER.event(key=mllog_constants.TRAIN_SAMPLES, value=sequences_seen)
          MLLOGGER.end(key=mllog_constants.RUN_STOP, metadata={mllog_constants.STATUS: mllog_constants.SUCCESS})
        if getenv("CKPT"):
          if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)
@@ -1662,276 +1631,6 @@ def train_llama3():
      if MLLOGGER and RUNMLPERF:
        MLLOGGER.start(key=mllog_constants.BLOCK_START, metadata={mllog_constants.SAMPLES_COUNT: sequences_seen})

-def train_gptoss():
-  from examples.mlperf.models.gpt_oss import GPTOSS, GPT_OSS_20B, apply_grad, FP8_DTYPE
-  from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
-  from examples.mlperf.optim import GradAccClipAdamW
-
-  BENCHMARK = getenv("BENCHMARK")
-
-  config = {}
-  BASEDIR            = config["BASEDIR"]                = Path(getenv("BASEDIR", "/raid/datasets/c4-8b/"))
-  BS                 = config["BS"]                     = getenv("BS", 16)
-  grad_acc           = config["GRADIENT_ACC_STEPS"]     = getenv("GRADIENT_ACC_STEPS", 1)
-  GBS                = config["GLOBAL_BATCH_SIZE"]      = BS * grad_acc
-  SEED               = config["SEED"]                   = getenv("SEED", 5760)
-  DATA_SEED          = config["DATA_SEED"]              = getenv("DATA_SEED", SEED)
-  SEQLEN             = config["SEQLEN"]                 = getenv("SEQLEN", 8192)
-  TRAIN_ON_VAL       = config["TRAIN_ON_VAL"]           = getenv("TRAIN_ON_VAL", 0)
-  MAX_STEPS          = config["MAX_STEPS"]              = getenv("MAX_STEPS", 1_200_000)
-  SAMPLES            = config["SAMPLES"]                = getenv("SAMPLES", 5_760 if TRAIN_ON_VAL else MAX_STEPS * GBS)
-  EVAL_SAMPLES       = config["EVAL_SAMPLES"]           = getenv("EVAL_SAMPLES", 1024)
-  WARMUP_STEPS       = config["WARMUP_STEPS"]           = getenv("WARMUP_STEPS", 128)
-  LR                 = config["LR"]                     = getenv("LR", 4e-4 * GBS / 16)
-  END_LR             = config["END_LR"]                 = getenv("END_LR", 4e-5)
-  EVAL_FREQ          = config["EVAL_FREQ"]              = getenv("EVAL_FREQ", 12288)
-  EVAL_BS            = config["EVAL_BS"]                = getenv("EVAL_BS", 16)
-  EVAL_TARGET        = config["EVAL_TARGET"]            = getenv("EVAL_TARGET", 3.34)
-
-  opt_adamw_beta_1 = 0.9
-  opt_adamw_beta_2 = 0.95
-  opt_adamw_epsilon = 1e-5
-  opt_adamw_weight_decay = 0.1
-
-  opt_learning_rate_warmup_steps = WARMUP_STEPS
-  opt_learning_rate_decay_steps = MAX_STEPS - opt_learning_rate_warmup_steps
-  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:
-    import wandb
-    wandb_args = {"id": wandb_id, "resume": "must"} if (wandb_id := getenv("WANDB_RESUME", "")) else {}
-    wandb.init(config=config, **wandb_args, project="MLPerf-gpt-oss")
-
-  model_params = GPT_OSS_20B
-  model_params['vocab_size'] = 128256
-  real_vocab_size = model_params['vocab_size']
-  if (layers:=getenv("LAYERS")) != 0: model_params['n_layers'] = layers
-  print(f"model parameters: {model_params}")
-
-  model = GPTOSS(**model_params, max_context=SEQLEN)
-
-  params = get_parameters(model)
-
-  if getenv("EMPTYWEIGHT"):
-    for v in get_parameters(model):
-      v = v.assign(Tensor.empty(v.shape, dtype=v.dtype))
-
-  is_dp = (DP := getenv("DP", 1)) > 1
-  is_sharding = is_dp
-  device_count = DP
-  device = tuple(f"{Device.DEFAULT}:{i}" for i in range(device_count))
-
-  model.shard(device, False)
-
-  is_offload_optim = bool(getenv("OFFLOAD_OPTIM"))
-  is_fake_offload = Device.DEFAULT == "NULL"
-  optim_device = ("CPU" if not is_fake_offload else "NULL:99") if is_offload_optim else None
-  optim = GradAccClipAdamW(params, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=grad_acc, device=optim_device)
-
-  for p in optim.params:
-    grad_dtype = dtypes.bfloat16 if p.dtype == FP8_DTYPE else p.dtype
-    p.grad = p.zeros_like(dtype=grad_dtype).contiguous()
-  grads = [p.grad for p in optim.params]
-
-  from extra.gemm.cdna_asm_gemm import _mx_block_scale
-  model_state = get_state_dict(model)
-  fp8_scale_names = {n: f"{n}_scale" for n, t in model_state.items() if t.dtype == FP8_DTYPE}
-  fp8_inv_scales = [model_state[sname] for sname in fp8_scale_names.values()]
-  for wname, sname in fp8_scale_names.items():
-    w, scale = model_state[wname], model_state[sname]
-    w._inv_scale = scale
-    if optim.master_params:
-      master = optim.master_params[next(j for j, p in enumerate(optim.params) if p is w)]
-      inv = scale if scale.device == master.device else scale.to(master.device)
-      bs = _mx_block_scale(inv.reshape(-1, inv.shape[-1])).reshape(w.shape)
-      master.assign((master * bs).contiguous())
-
-  scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)
-
-  # realize everything here
-  if optim.master_params: Tensor.realize(*optim.master_params)
-  Tensor.realize(*optim.params, *fp8_inv_scales)
-
-  @TinyJit
-  def minibatch(tokens:Tensor):
-    if is_dp: tokens = tokens.to(None).shard(device, 0)
-    if not is_sharding: tokens = tokens.to(None)
-    logits:Tensor = model(tokens[:, :-1], save=True)
-    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
-
-    for g, new_g in zip(grads, loss.gradient(*optim.params)):
-      apply_grad(g, new_g.uop)
-
-    loss_cpu = loss.flatten().float().to("CPU")
-    return loss_cpu.realize(*grads)
-
-  @TinyJit
-  def optim_step():
-    grad_norm = optim.fstep(grads)
-    scheduler.step()
-
-    for g in grads: g.assign(0)
-
-    lr_cpu = optim.lr.float().to("CPU")
-    grad_norm_cpu = grad_norm.float().to("CPU")
-    Tensor.realize(lr_cpu, grad_norm_cpu, *grads, *fp8_inv_scales)
-
-    return lr_cpu, grad_norm_cpu
-
-  @TinyJit
-  @Context(TRAINING=0)
-  def eval_step(tokens:Tensor):
-    if is_dp: tokens = tokens.to(None).shard(device, 0)
-    if not is_sharding: tokens = tokens.to(None)
-    logits:Tensor = model(tokens[:, :-1])
-    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
-    return loss.flatten().float().to("CPU")
-
-  # ** data iters **
-  def fake_data(bs, samples):
-    import numpy as np
-    for _ in range(samples // bs):
-      fake_data_np = np.random.randint(0, real_vocab_size, size=(bs, SEQLEN + 1), dtype=np.int32)
-      yield Tensor(fake_data_np, device="NPY")
-
-  def get_train_iter():
-    if getenv("FAKEDATA", 0):
-      return fake_data(BS, SAMPLES)
-    else:
-      from examples.mlperf.dataloader import batch_load_llama3
-      return batch_load_llama3(BS, SAMPLES, SEQLEN, BASEDIR, seed=DATA_SEED, val=bool(TRAIN_ON_VAL), small=True)
-
-  if getenv("FAKEDATA", 0):
-    eval_dataset = None
-  else:
-    from examples.mlperf.dataloader import get_llama3_dataset
-    eval_dataset = get_llama3_dataset(EVAL_SAMPLES, SEQLEN, BASEDIR, val=True, small=True)
-
-  def get_eval_iter():
-    if eval_dataset is None:
-      return fake_data(EVAL_BS, EVAL_SAMPLES)
-    from examples.mlperf.dataloader import iterate_llama3_dataset
-    return iterate_llama3_dataset(eval_dataset, EVAL_BS)
-
-  num_params = sum(p.numel() for p in params) - model_params["vocab_size"]*model_params["dim"]
-  train_iter = get_train_iter()
-  i, sequences_seen = 0, 0
-  step_times = []
-
-  while i < MAX_STEPS:
-    GlobalCounters.reset()
-    actual_gbs = GBS if i >= 2 else BS
-    if getenv("TRAIN", 1):
-      profile_marker(f"train @ {i}")
-      st = time.perf_counter()
-
-      stopped = False
-      losses, data_time, dev_time = [], 0, 0
-      for _ in range(grad_acc if i >= 2 else 1):
-        ist = time.perf_counter()
-        try: tokens = next(train_iter)
-        except StopIteration:
-          stopped = True
-          break
-        mst = time.perf_counter()
-        data_time += mst - ist
-        losses.append(minibatch(tokens).item())
-        dev_time += time.perf_counter() - mst
-      if stopped: break
-
-      gt = time.perf_counter()
-      ret = optim_step()
-      lr, grad_norm = ret[0].item(), ret[1].item()
-      et = time.perf_counter()
-
-      loss = sum(losses) / len(losses)
-      optim_time = et - gt
-      dev_time += optim_time
-      step_time = et - st
-      gbs_time = gt - st
-      if BENCHMARK: step_times.append(step_time)
-
-      i += 1
-      sequences_seen += actual_gbs
-
-      mem_gb = GlobalCounters.mem_used / 1e9
-      gflops = GlobalCounters.global_ops / 1e9 / dev_time
-      mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * device_count * 4.6e15)) * 100
-      tqdm.write(
-          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, {grad_norm:.6f} grad_norm, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")
-      if DEBUG >= 1: tqdm.write("  mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
-
-      if WANDB:
-        wandb.log({
-          "train/loss": loss,
-          "train/lr": lr,
-          "train/grad_norm": grad_norm,
-          "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
-        })
-
-      if (ckpt_freq := getenv("CKPT")) and (i % ckpt_freq == 0 and (i != 1 or ckpt_freq == 1)):
-        tqdm.write("saving checkpoint")
-        if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)
-        fn = f"{ckpt_dir}/gptoss_{i}.safe"
-        safe_save(get_state_dict(model), fn)
-
-        tqdm.write("saving optim checkpoint")
-        fn = f"{ckpt_dir}/gptoss_{i}_optim.safe"
-        safe_save(get_state_dict(scheduler), fn)
-
-      if i == BENCHMARK:
-        median_step_time = sorted(step_times)[BENCHMARK // 2]
-        estimated_steps = MAX_STEPS
-        estimated_total_minutes = int(median_step_time * estimated_steps / 60)
-        print(f"Estimated training time: {estimated_total_minutes // 60}h{estimated_total_minutes % 60}m")
-        print(f"epoch global_ops: {GlobalCounters.global_ops:_}, "
-              f"epoch global_mem: {GlobalCounters.global_mem:_}")
-
-    if (sequences_seen // EVAL_FREQ != (sequences_seen - actual_gbs) // EVAL_FREQ 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 = []
-      eval_iter = get_eval_iter()
-      tqdm.write(f"evaluating {EVAL_SAMPLES//EVAL_BS} batches of {EVAL_BS} sequences")
-
-      for j,tokens in tqdm(enumerate(eval_iter), total=EVAL_SAMPLES//EVAL_BS):
-        eval_losses += eval_step(tokens).tolist()
-
-        if BENCHMARK and (j+1) == min(BENCHMARK, EVAL_SAMPLES//EVAL_BS):
-          return
-
-      log_perplexity = sum(eval_losses) / len(eval_losses)
-
-      tqdm.write(f"eval log perplexity: {log_perplexity:.4f}")
-
-      if WANDB:
-        wandb.log({"eval/log_perplexity": log_perplexity, "eval/sequences_seen": sequences_seen})
-
-      if log_perplexity < EVAL_TARGET:
-        tqdm.write(f"target achieved after {sequences_seen} sequences")
-        if getenv("CKPT"):
-          if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)
-          fn = f"{ckpt_dir}/gptoss.safe"
-          safe_save(get_state_dict(model), fn)
-        break
-
 def train_stable_diffusion():
  from extra.models.unet import UNetModel
  from examples.mlperf.dataloader import batch_load_train_stable_diffusion
@@ -2010,7 +1709,7 @@ def train_stable_diffusion():
    # move to CPU first so more GPU bufs aren't created (can trigger OOM)
    for k,v in ckpt.items(): ckpt[k] = v.detach().to("CPU")
    Tensor.realize(*[v for v in ckpt.values()])
-    for k,v in ckpt.items(): ckpt[k] = v.cast(v.dtype).contiguous()
+    for k,v in ckpt.items(): ckpt[k] = v.cast(v.dtype.base).contiguous()
    Tensor.realize(*[v for v in ckpt.values()])
    return ckpt

@@ -2077,7 +1776,7 @@ if __name__ == "__main__":
  elif getenv("RUNMLPERF"): bench_log_manager = WallTimeEvent(BenchEvent.MLPERF_RUN)
  else: bench_log_manager = contextlib.nullcontext()

-  with Context(TRAINING=1):
+  with Tensor.train():
    for m in getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert,maskrcnn,stable_diffusion").split(","):
      nm = f"train_{m}"
      if nm in globals():
@@ -1,9 +1,10 @@
-import math, os
+import math, os, functools
 if __name__ == "__main__":
  os.environ["DEFAULT_FLOAT"] = "bfloat16"
  os.environ["OPTIM_DTYPE"] = "bfloat16"
-  if "DEV" not in os.environ: os.environ["DEV"] = "NULL::gfx950"
+  if "DEV" not in os.environ: os.environ["DEV"] = "NULL"
  # CDNA
+  os.environ["EMULATE"] = "AMD_CDNA4"
  os.environ["DEVICE_IN_FUNCTION_BUG"] = "1"
  os.environ["ALL2ALL"] = "1"
  os.environ["USE_ATOMICS"] = "1"
@@ -12,113 +13,68 @@ if __name__ == "__main__":
    if "ASM_GEMM" not in os.environ:
      os.environ["ASM_GEMM"] = "1"
 from tinygrad import Tensor, nn, function, getenv, dtypes, TinyJit
-from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker, round_up
+from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
 from tinygrad.uop.ops import Ops, UOp
 from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
-from extra.llama_kernels.rmsnorm import rmsnorm
-from extra.llama_kernels import FP8_MAX, local_abs_max

-ASM_GEMM = getenv("ASM_GEMM", 0)
-FUSED_INPUT_QUANTIZE = getenv("FUSED_INPUT_QUANTIZE", 0)
-FUSED_ADD_NORM_MUL_QUANTIZE = getenv("FUSED_ADD_NORM_MUL_QUANTIZE", 0)
-FUSED_SILU_W13 = getenv("FUSED_SILU_W13", 0)
-SPLIT_W13 = getenv("SPLIT_W13", 0)
-COLUMNWISE_WEIGHT_SCALE = getenv("COLUMNWISE_WEIGHT_SCALE", 0)
-MXFP8 = getenv("MXFP8", 0)
+FP8 = getenv("FP8", 0)

 FP8_DTYPE = dtypes.fp8e4m3
 FP8_GRAD_DTYPE = dtypes.fp8e5m2
+FP8_MAX = 448.0
+
+# per-device abs max without allreduce (matches TE delayed scaling behavior)
+@functools.cache
+def _local_abs_max_fxn(x_p, device):
+  x = Tensor(x_p, device=device)
+  inner = Tensor(x.uop.src[0]) if x.uop.op is Ops.MULTI else x
+  return (inner.abs().max(),)
+
+def _local_abs_max(x:Tensor) -> Tensor:
+  param = x.as_param(0)
+  fxn = _local_abs_max_fxn(param.uop, x.device)
+  return Tensor(fxn[0].uop.call(x.uop).gettuple(0))

 def quantize_fp8(x:Tensor, amax_state:Tensor|None=None):
-  new_amax = (local_abs_max(x) if isinstance(x.device, tuple) else x.abs().max()).detach().cast(dtypes.float32)
+  new_amax = (_local_abs_max(x) if isinstance(x.device, tuple) else x.abs().max()).detach()
  scale = FP8_MAX / ((amax_state if amax_state is not None else new_amax) + 1e-8)
  x_scaled = x * scale
  x_clamped = x_scaled + (x_scaled.detach().clamp(-FP8_MAX, FP8_MAX) - x_scaled.detach())  # STE
  return x_clamped.cast(FP8_DTYPE), scale.float().reciprocal(), new_amax

-def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_scale:Tensor|None=None,
-           x_fp8:Tensor|None=None, x_new_amax:Tensor|None=None,
-           grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None, x_prequant_mx:tuple|None=None) -> tuple[Tensor,...]:
+def matmul(x:Tensor, w:Tensor, fp8=FP8, amax_x:Tensor|None=None, amax_w:Tensor|None=None) -> tuple[Tensor,...]:
  if not fp8:
-    if ASM_GEMM:
+    if getenv("ASM_GEMM"):
      from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
      if can_use_asm_gemm(x, w.T): return (asm_gemm(x, w.T),)
    return (x @ w.T,)
-  assert w_inv_scale is not None, "fp8 matmul requires w_inv_scale (weights must be stored in fp8 with per-tensor scale)"
-  if MXFP8:
-    from extra.gemm.cdna_asm_gemm import asm_gemm, quantize_mxfp8, mx_pack, can_use_asm_gemm, _mx_block_scale
-    if x_prequant_mx is not None: x_q, x_e8, x_si = x_prequant_mx       # fused producer already quantized (2d)
-    else: x_q, x_e8, x_si = quantize_mxfp8(x.reshape(-1, x.shape[-1]))
-    l_shape = x.shape[:-1] if x is not None else x_q.shape[:-1]
-    if can_use_asm_gemm(x_q, w.T):
-      out = asm_gemm(x_q, w.T, mx=True, mx_scales=(x_si, x_e8, mx_pack(w_inv_scale), w_inv_scale),
-                     mx_w_stored=True).reshape(*l_shape, w.shape[0])
-    else:
-      x_phys = (x_q.cast(dtypes.bfloat16) * _mx_block_scale(x_e8)).reshape(*l_shape, x_q.shape[-1])
-      out = x_phys @ (w.cast(dtypes.bfloat16) * _mx_block_scale(w_inv_scale)).T
-    return out, (amax_x.detach() if amax_x is not None else None), x_q
-  if x_fp8 is None:
-    if FUSED_INPUT_QUANTIZE and amax_x is not None:
-      from extra.llama_kernels.quantize_fp8_delayed import quantize_fp8_delayed
-      x_fp8, _, x_new_amax, _ = quantize_fp8_delayed(x, amax_x, FP8_DTYPE)
-    else:
-      x_fp8, _, x_new_amax = quantize_fp8(x, amax_state=amax_x)
-  if ASM_GEMM:
+  x_fp8, x_scale, x_new_amax = quantize_fp8(x, amax_state=amax_x)
+  w_fp8, w_scale, w_new_amax = quantize_fp8(w, amax_state=amax_w)
+  combined_scale = x_scale * w_scale
+  if getenv("ASM_GEMM"):
    from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
-    if can_use_asm_gemm(x_fp8, w.T):
-      assert amax_x is not None
-      if COLUMNWISE_WEIGHT_SCALE:
-        out = asm_gemm(x_fp8, w.T, x_scale=amax_x, grad_amax_state=grad_amax_state,
-                       next_grad_amax_state=next_grad_amax_state, w_post_scale=w_inv_scale)
-      else:
-        out = asm_gemm(x_fp8, w.T, x_scale=amax_x, w_scale=w_inv_scale, grad_amax_state=grad_amax_state,
-                       next_grad_amax_state=next_grad_amax_state)
-      return out, x_new_amax, x_fp8
-  return (x_fp8.dot(w.T, dtype=dtypes.float) * ((amax_x.float() + 1e-8) / FP8_MAX) * w_inv_scale).cast(dtypes.bfloat16), x_new_amax, x_fp8
+    if can_use_asm_gemm(x_fp8, w_fp8.T): return asm_gemm(x_fp8, w_fp8.T, combined_scale=combined_scale), x_new_amax, w_new_amax, x_fp8, w_fp8
+  return x_fp8.dot(w_fp8.T, dtype=dtypes.float) * combined_scale, x_new_amax, w_new_amax, x_fp8, w_fp8

-def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor,
-                         grad_amax_state:Tensor, next_grad_amax_state:Tensor):
-  if FUSED_ADD_NORM_MUL_QUANTIZE:
-    from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_rmsnorm_mul_quantize_fp8
-    x_fp8, new_amax, x_normed, rrms = fused_rmsnorm_mul_quantize_fp8(x, norm, amax_x, eps, FP8_DTYPE)
-    out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x, x_new_amax=new_amax,
-                       grad_amax_state=grad_amax_state, next_grad_amax_state=next_grad_amax_state)
-    return out, x_normed, rrms, ret
-  x_normed, rrms = rmsnorm(x, eps)
-  out, *ret = matmul(x_normed * norm, w, amax_x=amax_x, w_inv_scale=w_inv_scale, grad_amax_state=grad_amax_state,
-                     next_grad_amax_state=next_grad_amax_state)
-  return out, x_normed, rrms, ret
+def _rmsnorm_fwd(x_in:Tensor, eps:float) -> tuple[Tensor, Tensor]:
+  x = x_in.float()
+  rrms = (x.square().mean(-1, keepdim=True) + eps).rsqrt()
+  return (x * rrms).cast(x_in.dtype), rrms

-def add_norm_quantize_matmul(x:Tensor, residual:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor,
-                             grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None):
-  if FUSED_ADD_NORM_MUL_QUANTIZE:
-    from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_add_rmsnorm_mul_quantize_fp8
-    x_fp8, new_amax, h, x_normed, rrms = fused_add_rmsnorm_mul_quantize_fp8(x, residual, norm, amax_x, eps, FP8_DTYPE)
-    out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x, x_new_amax=new_amax,
-                       grad_amax_state=grad_amax_state, next_grad_amax_state=next_grad_amax_state)
-    return out, h, x_normed, rrms, ret
-  h = x + residual
-  x_normed, rrms = rmsnorm(h, eps)
-  out, *ret = matmul(x_normed * norm, w, amax_x=amax_x, w_inv_scale=w_inv_scale, grad_amax_state=grad_amax_state,
-                     next_grad_amax_state=next_grad_amax_state)
-  return out, h, x_normed, rrms, ret
+@functools.cache
+def _rmsnorm_fwd_fxn(x_in_p, eps, device):
+  return _rmsnorm_fwd(Tensor(x_in_p, device=device), eps)

-def silu_w13_quantize_matmul(x_w13:Tensor, w2:Tensor, s_2:Tensor,
-                             amax_x2:Tensor,
-                             grad_amax_xw13:Tensor, next_grad_amax_xw13:Tensor,
-                             grad_amax_xout:Tensor, next_grad_amax_xout:Tensor):
-  if FUSED_SILU_W13:
-    from extra.llama_kernels.cast_amax import fused_quantize_fp8_w13
-    x2_fp8, new_amax_x2 = fused_quantize_fp8_w13(x_w13, amax_x2, FP8_DTYPE, grad_amax_state=grad_amax_xw13,
-                                                 next_grad_amax_state=next_grad_amax_xw13)
-    out, *ret = matmul(None, w2, w_inv_scale=s_2, x_fp8=x2_fp8, amax_x=amax_x2, x_new_amax=new_amax_x2,
-                       grad_amax_state=grad_amax_xout, next_grad_amax_state=next_grad_amax_xout)
-    return out, ret
-  hidden = x_w13.shape[-1] // 2
-  x_w1, x_w3 = x_w13[..., :hidden], x_w13[..., hidden:]
-  out, *ret = matmul(x_w1.silu() * x_w3, w2, amax_x=amax_x2, w_inv_scale=s_2, grad_amax_state=grad_amax_xout,
-                     next_grad_amax_state=next_grad_amax_xout)
-  return out, ret
+def _rmsnorm_bwd(grad:UOp, call:UOp) -> tuple:
+  x_normed = Tensor(call.gettuple(0)).float()
+  do_float = Tensor(grad).float()
+  d_x = Tensor(call.gettuple(1)) * (do_float - x_normed * (do_float * x_normed).mean(-1, keepdim=True))
+  return (d_x.cast(call.src[1].dtype).uop,)
+
+def rmsnorm(x_in:Tensor, eps:float) -> tuple[Tensor, Tensor]:
+  fxn = _rmsnorm_fwd_fxn(x_in.as_param(0).uop, eps, x_in.device)
+  call = UOp.maketuple(fxn[0].uop, fxn[1].uop).call(x_in.uop, grad_fxn=_rmsnorm_bwd)
+  return Tensor(call.gettuple(0)), Tensor(call.gettuple(1))

 class FlatTransformer:
  def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_layers:int, norm_eps:float, vocab_size:int, n_kv_heads:int|None=None,
@@ -129,21 +85,17 @@ class FlatTransformer:
    self.n_kv_heads = n_kv_heads if n_kv_heads is not None else n_heads # n_kv_heads != n_heads implies MQA [arxiv/2307.09288, A.2.1]
    self.head_dim = dim // n_heads
    self.n_rep = self.n_heads // self.n_kv_heads
-    self.hidden_dim = hidden_dim

    scaled_std = 0.02 / math.sqrt(2 * n_layers)

    # Attention
-    self.wqkv, s_qkv = self.lin_per_layer(dim, self.n_heads * self.head_dim + self.n_kv_heads * self.head_dim * 2)
-    self.wo, s_o = self.lin_per_layer(self.n_heads * self.head_dim, dim, std=scaled_std)
+    self.wqkv = self.lin_per_layer(dim, self.n_heads * self.head_dim + self.n_kv_heads * self.head_dim * 2)
+    self.wo = self.lin_per_layer(self.n_heads * self.head_dim, dim, std=scaled_std)

    # FeedForward
-    if SPLIT_W13:
-      self.w1, s_1 = self.lin_per_layer(dim, hidden_dim)
-      self.w3, s_3 = self.lin_per_layer(dim, hidden_dim)
-    else:
-      self.w13, s_13 = self.lin_per_layer(dim, hidden_dim * 2)
-    self.w2, s_2 = self.lin_per_layer(hidden_dim, dim, std=scaled_std)
+    self.w1 = self.lin_per_layer(dim, hidden_dim)
+    self.w2 = self.lin_per_layer(hidden_dim, dim, std=scaled_std)
+    self.w3 = self.lin_per_layer(dim, hidden_dim)

    self.norm_eps = norm_eps
    self.attention_norm = Tensor.ones(n_layers, dim).contiguous()
@@ -154,121 +106,87 @@ class FlatTransformer:
    self.tok_embeddings = nn.Embedding(vocab_size, dim)
    self.tok_embeddings.weight = Tensor.normal(vocab_size, dim, mean=0.0, std=0.02, dtype=dtypes.bfloat16)
    self.output = Tensor.normal(1, vocab_size, dim, mean=0.0, std=0.02, dtype=dtypes.bfloat16)
-    self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).contiguous().is_param_(False)
+    self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).contiguous().requires_grad_(False)

-    def _amax(): return Tensor.full((), FP8_MAX, dtype=dtypes.float32).contiguous().is_param_(False)
-    names = ["xqkv", "xo", "x2"]
-    names += ["x1", "x3"] if SPLIT_W13 else ["x13"]
-    self._fp8_amax = {name: [_amax() for _ in range(n_layers)] for name in names}
-    self._fp8_next_amax = {name: [_amax() for _ in range(n_layers)] for name in names}
-    grad_names = ["xqkv", "xo", "xout"]
-    grad_names += ["xw1", "xw3"] if SPLIT_W13 else ["xw13"]
-    self._fp8_grad_amax = {name: [_amax() for _ in range(n_layers)] for name in grad_names}
-    self._fp8_next_grad_amax = {name: [_amax() for _ in range(n_layers)] for name in grad_names}
-    w_scales = [("wqkv", s_qkv), ("wo", s_o), ("w2", s_2)]
-    w_scales += [("w1", s_1), ("w3", s_3)] if SPLIT_W13 else [("w13", s_13)]
-    self._fp8_inv_scale = {name: (s if MXFP8 else s.float()).contiguous().is_param_(False) for name, s in w_scales}
-    self._fp8_next_inv_scale = {name: (s if MXFP8 else s.float()).contiguous().is_param_(False) for name, s in w_scales}
+    if FP8:
+      def _amax(): return Tensor.full((), FP8_MAX).contiguous().requires_grad_(False)
+      names = ["xqkv", "wqkv", "xo", "wo", "x1", "w1", "x2", "w2", "x3", "w3"]
+      # _fp8_amax[name][layer_idx] = scalar amax tensor
+      self._fp8_amax = {name: [_amax() for _ in range(n_layers)] for name in names}
+      self._fp8_amax["xout"] = [_amax()]
+      self._fp8_amax["wout"] = [_amax()]

-  def lin_per_layer(self, in_features:int, out_features:int, std:float=0.02, w:Tensor|None=None):
-    if w is None:
-      if getenv("ZEROS"): w = Tensor.zeros(self.n_layers, out_features, in_features)
-      else: w = Tensor.normal(self.n_layers, out_features, in_features, mean=0.0, std=std)
-    if MXFP8:
-      from extra.gemm.cdna_asm_gemm import quantize_mxfp8
-      w_q, w_e8, _ = quantize_mxfp8(w.reshape(self.n_layers * out_features, in_features))
-      return w_q.reshape(self.n_layers, out_features, in_features), w_e8.reshape(self.n_layers, out_features, in_features // 32)
-    amax = (w.abs().max(axis=2) if COLUMNWISE_WEIGHT_SCALE else w.abs().flatten(1).max(1)).detach()
-    scale = FP8_MAX / (amax + 1e-8)
-    inv_scale = (amax + 1e-8) / FP8_MAX
-    scale_b = scale.reshape(self.n_layers, out_features, 1) if COLUMNWISE_WEIGHT_SCALE else scale.reshape(-1, 1, 1)
-    return (w * scale_b).clamp(-FP8_MAX, FP8_MAX).cast(FP8_DTYPE), inv_scale
+  def lin_per_layer(self, in_features:int, out_features:int, std:float=0.02):
+    if getenv("ZEROS"): return Tensor.zeros(self.n_layers, out_features, in_features)
+    return Tensor.normal(self.n_layers, out_features, in_features, mean=0.0, std=std)

-  def attention(self, x:Tensor, freqs_cis:Tensor, *, attention_norm:Tensor, wqkv:Tensor, wo:Tensor,
-                amax_xqkv:Tensor, amax_xo:Tensor, s_qkv:Tensor, s_o:Tensor,
-                grad_amax_xqkv:Tensor, grad_amax_xo:Tensor, next_grad_amax_xqkv:Tensor, next_grad_amax_xo:Tensor):
+  def attention(self, x:Tensor, freqs_cis:Tensor, attention_norm:Tensor, wqkv:Tensor, wo:Tensor,
+                amax_xqkv=None, amax_wqkv=None, amax_xo=None, amax_wo=None):
    bsz, seqlen, _ = x.shape
-    amaxs, saves = [], []
+    new_amaxs, saves = [], []

-    xqkv, x_normed, rrms, (new_amax, *s) = norm_quantize_matmul(x, attention_norm, wqkv, s_qkv, self.norm_eps,
-                                                                  amax_x=amax_xqkv, grad_amax_state=grad_amax_xqkv,
-                                                                  next_grad_amax_state=next_grad_amax_xqkv)
-    amaxs.append(new_amax)
-    saves.extend([x_normed, rrms, *s, xqkv])
+    x, rrms = rmsnorm(x, self.norm_eps)
+    saves.extend([x, rrms])
+    x = x * attention_norm
+
+    xqkv, *ret = matmul(x, wqkv, amax_x=amax_xqkv, amax_w=amax_wqkv)
+    new_amaxs.extend(ret[:2])
+    saves.extend(ret[2:] + [xqkv])
    xqkv = xqkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
    xq = xqkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
    xk = xqkv[:, :, :, self.n_rep].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
    xv = xqkv[:, :, :, self.n_rep+1].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)

    xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
-    xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16)
+    if FP8: xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16)
+    xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
    if getenv("HK_FLASH_ATTENTION"):
      from extra.thunder.amd.fa import flash_attention
-      attn, *save = flash_attention(xq, xk, xv, is_causal=True, write_flat=True)
+      attn, *save = flash_attention(xq, xk, xv, is_causal=True)
      saves.extend(save)
    else:
-      xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
-      attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True).transpose(1, 2)
-    attn = attn.reshape(bsz, seqlen, -1)
+      attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True)
+    attn = attn.transpose(1, 2).reshape(bsz, seqlen, -1)

-    out, new_amax, *s = matmul(attn, wo, amax_x=amax_xo, w_inv_scale=s_o, grad_amax_state=grad_amax_xo,
-                               next_grad_amax_state=next_grad_amax_xo)
-    amaxs.append(new_amax)
-    saves.extend([*s, out])
-    return out, amaxs, saves
+    out, *ret = matmul(attn, wo, amax_x=amax_xo, amax_w=amax_wo)
+    new_amaxs.extend(ret[:2])
+    saves.extend(ret[2:] + [out])
+    return (out, *new_amaxs, *saves)

-  def feed_forward(self, x:Tensor, residual:Tensor, **kwargs):
-    amaxs, saves = [], []
+  def feed_forward(self, x:Tensor, ffn_norm:Tensor, w1:Tensor, w2:Tensor, w3:Tensor,
+                   amax_x1=None, amax_w1=None, amax_x2=None, amax_w2=None, amax_x3=None, amax_w3=None):
+    new_amaxs, saves = [], []

-    if SPLIT_W13:
-      h = x + residual
-      x_normed, rrms = rmsnorm(h, self.norm_eps)
-      saves.extend([x_normed, rrms])
-      inp = x_normed * kwargs["ffn_norm"]
-      x_w1, new_amax, *s = matmul(inp, kwargs["w1"], amax_x=kwargs["amax_x1"], w_inv_scale=kwargs["s_1"],
-                                  grad_amax_state=kwargs["grad_amax_xw1"], next_grad_amax_state=kwargs["next_grad_amax_xw1"])
-      amaxs.append(new_amax)
-      saves.extend([*s, x_w1])
-      x_w3, new_amax, *s = matmul(inp, kwargs["w3"], amax_x=kwargs["amax_x3"], w_inv_scale=kwargs["s_3"],
-                                  grad_amax_state=kwargs["grad_amax_xw3"], next_grad_amax_state=kwargs["next_grad_amax_xw3"])
-      amaxs.append(new_amax)
-      saves.extend([*s, x_w3])
-      if FUSED_SILU_W13 and MXFP8:
-        from extra.llama_kernels.fused_silu_mul_quantize_mxfp8 import fused_silu_mul_quantize_mxfp8
-        aq, ae8, asi = fused_silu_mul_quantize_mxfp8(x_w1.reshape(-1, x_w1.shape[-1]), x_w3.reshape(-1, x_w3.shape[-1]))
-        out, new_amax, *s = matmul(None, kwargs["w2"], x_prequant_mx=(aq, ae8, asi), amax_x=kwargs["amax_x2"],
-                                   w_inv_scale=kwargs["s_2"], grad_amax_state=kwargs["grad_amax_xout"],
-                                   next_grad_amax_state=kwargs["next_grad_amax_xout"])
-        out = out.reshape(*x_w1.shape[:-1], kwargs["w2"].shape[0])
-      else:
-        out, new_amax, *s = matmul(x_w1.silu() * x_w3, kwargs["w2"], amax_x=kwargs["amax_x2"], w_inv_scale=kwargs["s_2"],
-                                   grad_amax_state=kwargs["grad_amax_xout"], next_grad_amax_state=kwargs["next_grad_amax_xout"])
-      amaxs.append(new_amax)
-      saves.extend([*s, out])
-    else:
-      x_w13, h, x_normed, rrms, (new_amax, *s) = add_norm_quantize_matmul(x, residual, kwargs["ffn_norm"], kwargs["w13"], kwargs["s_13"],
-                                                                          self.norm_eps, amax_x=kwargs["amax_x13"],
-                                                                          grad_amax_state=kwargs["grad_amax_xw13"],
-                                                                          next_grad_amax_state=kwargs["next_grad_amax_xw13"])
-      amaxs.append(new_amax)
-      saves.extend([x_normed, rrms, *s, x_w13])
-      out, (new_amax, *s) = silu_w13_quantize_matmul(x_w13, kwargs["w2"], kwargs["s_2"], amax_x2=kwargs["amax_x2"],
-                                                     grad_amax_xw13=kwargs["grad_amax_xw13"],
-                                                     next_grad_amax_xw13=kwargs["next_grad_amax_xw13"],
-                                                     grad_amax_xout=kwargs["grad_amax_xout"],
-                                                     next_grad_amax_xout=kwargs["next_grad_amax_xout"])
-      amaxs.append(new_amax)
-      saves.extend([*s, out])
-    return out, h, amaxs, saves
+    x, rrms = rmsnorm(x, self.norm_eps)
+    saves.extend([x, rrms])
+    x = x * ffn_norm
+
+    x_w1, *ret = matmul(x, w1, amax_x=amax_x1, amax_w=amax_w1)
+    new_amaxs.extend(ret[:2])
+    saves.extend(ret[2:] + [x_w1])
+    x_w3, *ret = matmul(x.contiguous_backward(), w3, amax_x=amax_x3, amax_w=amax_w3)
+    new_amaxs.extend(ret[:2])
+    saves.extend(ret[2:] + [x_w3])
+    out, *ret = matmul(x_w1.silu() * x_w3, w2, amax_x=amax_x2, amax_w=amax_w2)
+    new_amaxs.extend(ret[:2])
+    saves.extend(ret[2:] + [out])
+    return (out, *new_amaxs, *saves)

  @function(precompile=True, precompile_backward=True)
-  def run_layer(self, x:Tensor, freqs_cis:Tensor, attn_kwargs:dict, ffn_kwargs:dict, save:bool=True):
-    attn, attn_amaxs, attn_saves = self.attention(x, freqs_cis, **attn_kwargs)
-    ffn, h, ffn_amaxs, ffn_saves = self.feed_forward(x, attn, **ffn_kwargs)
+  def run_layer(self, x:Tensor, freqs_cis:Tensor,
+                attention_norm:Tensor, wqkv:Tensor, wo:Tensor,
+                ffn_norm:Tensor, w1:Tensor, w2:Tensor, w3:Tensor,
+                amax_xqkv=None, amax_wqkv=None, amax_xo=None, amax_wo=None,
+                amax_x1=None, amax_w1=None, amax_x2=None, amax_w2=None, amax_x3=None, amax_w3=None):
+    attn, *attn_ret = self.attention(x, freqs_cis, attention_norm, wqkv, wo,
+                                     amax_xqkv=amax_xqkv, amax_wqkv=amax_wqkv, amax_xo=amax_xo, amax_wo=amax_wo)
+    attn_amaxs, attn_saves = attn_ret[:4], attn_ret[4:]
+    h = x + attn
+    ffn, *ffn_ret = self.feed_forward(h, ffn_norm, w1, w2, w3,
+                                      amax_x1=amax_x1, amax_w1=amax_w1, amax_x2=amax_x2, amax_w2=amax_w2, amax_x3=amax_x3, amax_w3=amax_w3)
+    ffn_amaxs, ffn_saves = ffn_ret[:6], ffn_ret[6:]
    h = h + ffn
-    amaxs = tuple(a.detach() for a in (*attn_amaxs, *ffn_amaxs))
-    if save: return (h, *amaxs, *attn_saves, *ffn_saves)
-    else: return (h, *amaxs)
+    return (h, *attn_amaxs, *ffn_amaxs, *attn_saves, *ffn_saves)

  def shard(self, device:tuple[str, ...], mp:bool=False):
    from tinygrad.nn.state import get_parameters
@@ -276,65 +194,43 @@ class FlatTransformer:
      for v in get_parameters(self): v.shard_(device, axis=None)
    else:
      # flat per-layer weights: axis 0 is n_layers, so shard axes are +1 vs per-layer Transformer
-      def _shard_fp8(name:str, axis:int, std:float=0.02):
-        w = getattr(self, name)
-        if MXFP8:
-          from extra.gemm.cdna_asm_gemm import quantize_mxfp8
-          w_bf16 = Tensor.empty(self.n_layers, w.shape[1], w.shape[2], dtype=dtypes.bfloat16).shard(device, axis=axis).randn_like() * std
-          w_q, w_e8, _ = quantize_mxfp8(w_bf16)
-          w.replace(w_q)
-          self._fp8_inv_scale[name].replace(w_e8.contiguous()).is_param_(False)
-          self._fp8_next_inv_scale[name].replace(w_e8.contiguous()).is_param_(False)
-        else:
-          w.shard_(device, axis=axis)
-          scale_axis = (1 if axis == 1 else None) if COLUMNWISE_WEIGHT_SCALE else None
-          self._fp8_inv_scale[name] = self._fp8_inv_scale[name].shard(device, axis=scale_axis).contiguous().is_param_(False)
-          self._fp8_next_inv_scale[name] = self._fp8_next_inv_scale[name].shard(device, axis=scale_axis).contiguous().is_param_(False)
-        Tensor.realize(w, self._fp8_inv_scale[name], self._fp8_next_inv_scale[name])
-      sstd = 0.02 / math.sqrt(2 * self.n_layers)
-      _shard_fp8("wqkv", 1)          # (n_layers, out, dim) shard out
-      _shard_fp8("wo", 2, sstd)      # (n_layers, dim, in) shard in
-      if SPLIT_W13:
-        _shard_fp8("w1", 1)
-        _shard_fp8("w3", 1)
-      else:
-        _shard_fp8("w13", 1)         # (n_layers, hidden*2, dim) shard out
-      _shard_fp8("w2", 2, sstd)      # (n_layers, dim, hidden) shard in
+      self.wqkv.shard_(device, axis=1).realize()          # (n_layers, out, dim) shard out
+      self.wo.shard_(device, axis=2).realize()             # (n_layers, dim, in) shard in
+      self.w1.shard_(device, axis=1).realize()             # (n_layers, hidden, dim) shard out
+      self.w2.shard_(device, axis=2).realize()             # (n_layers, dim, hidden) shard in
+      self.w3.shard_(device, axis=1).realize()             # (n_layers, hidden, dim) shard out
      self.attention_norm.shard_(device, axis=None).realize()
      self.ffn_norm.shard_(device, axis=None).realize()
      self.norm.weight.shard_(device, axis=None).realize()
      self.tok_embeddings.weight.shard_(device, axis=0).realize()
      self.output.shard_(device, axis=1).realize()
      self.freqs_cis.shard_(device, axis=None).realize()
-      for amax_dict in (self._fp8_amax, self._fp8_next_amax, self._fp8_grad_amax, self._fp8_next_grad_amax):
-        for name in amax_dict:
-          for i in range(len(amax_dict[name])):
-            amax_dict[name][i] = amax_dict[name][i].to(device).contiguous().is_param_(False)
+      if FP8:
+        for name in self._fp8_amax:
+          for i in range(len(self._fp8_amax[name])):
+            self._fp8_amax[name][i] = self._fp8_amax[name][i].to(device).contiguous().requires_grad_(False)

-  def __call__(self, tokens:Tensor, save:bool=True):
+  def __call__(self, tokens:Tensor):
    h = self.tok_embeddings(tokens)
    freqs_cis = self.freqs_cis.cast(h.dtype)[:, :tokens.shape[1], :, :, :]
-    a, na, ga, nga, s = self._fp8_amax, self._fp8_next_amax, self._fp8_grad_amax, self._fp8_next_grad_amax, self._fp8_inv_scale
+    a = self._fp8_amax if FP8 else None
    for i in range(self.n_layers):
-      attn_kwargs = dict(attention_norm=self.attention_norm[i], wqkv=self.wqkv[i], wo=self.wo[i],
-                         amax_xqkv=a["xqkv"][i], amax_xo=a["xo"][i], s_qkv=s["wqkv"][i], s_o=s["wo"][i],
-                         grad_amax_xqkv=ga["xqkv"][i], grad_amax_xo=ga["xo"][i],
-                         next_grad_amax_xqkv=nga["xqkv"][i], next_grad_amax_xo=nga["xo"][i])
-      ffn_kwargs = dict(ffn_norm=self.ffn_norm[i], w2=self.w2[i],
-                        amax_x2=a["x2"][i], s_2=s["w2"][i], grad_amax_xout=ga["xout"][i], next_grad_amax_xout=nga["xout"][i])
-      if SPLIT_W13:
-        ffn_kwargs.update(w1=self.w1[i], w3=self.w3[i], amax_x1=a["x1"][i], amax_x3=a["x3"][i],
-                          s_1=s["w1"][i], s_3=s["w3"][i], grad_amax_xw1=ga["xw1"][i], grad_amax_xw3=ga["xw3"][i],
-                          next_grad_amax_xw1=nga["xw1"][i], next_grad_amax_xw3=nga["xw3"][i])
-      else:
-        ffn_kwargs.update(w13=self.w13[i], amax_x13=a["x13"][i], s_13=s["w13"][i], grad_amax_xw13=ga["xw13"][i],
-                          next_grad_amax_xw13=nga["xw13"][i])
-      h, *ret = self.run_layer(h, freqs_cis, attn_kwargs, ffn_kwargs, save=save)
-      amax_names = ["xqkv", "xo"] + (["x1", "x3"] if SPLIT_W13 else ["x13"]) + ["x2"]
-      for name, new_val in zip(amax_names, ret[:len(amax_names)]):
-        na[name][i].assign(new_val)
+      amax_layer = {"amax_xqkv": a["xqkv"][i], "amax_wqkv": a["wqkv"][i],
+                    "amax_xo": a["xo"][i], "amax_wo": a["wo"][i],
+                    "amax_x1": a["x1"][i], "amax_w1": a["w1"][i],
+                    "amax_x2": a["x2"][i], "amax_w2": a["w2"][i],
+                    "amax_x3": a["x3"][i], "amax_w3": a["w3"][i]} if a else {}
+      h, *ret = self.run_layer(h, freqs_cis,
+                               self.attention_norm[i], self.wqkv[i], self.wo[i],
+                               self.ffn_norm[i], self.w1[i], self.w2[i], self.w3[i],
+                               **amax_layer)
+      if a:
+        amaxs = ret[:10]
+        amax_names = ["xqkv", "wqkv", "xo", "wo", "x1", "w1", "x3", "w3", "x2", "w2"]
+        for name, new_val in zip(amax_names, amaxs):
+          a[name][i].assign(new_val)

-    logits = matmul(self.norm(h), self.output[0], fp8=False)[0]
+    logits = matmul(self.norm(h).contiguous().contiguous_backward(), self.output[0], fp8=False)[0].contiguous_backward()
    return logits

 def _get_pads(uop:UOp) -> list[UOp]:
@@ -343,61 +239,37 @@ def _get_pads(uop:UOp) -> list[UOp]:

 def apply_grad(grad_buf:Tensor, new_grad:UOp):
  pads = _get_pads(new_grad)
+  new_grad = new_grad.cast(grad_buf.dtype)
  if len(pads) <= 1:
-    new_grad = new_grad.cast(grad_buf.dtype)
-    grad_buf.uop = grad_buf.uop.after(grad_buf.uop.store(grad_buf.uop + new_grad))
+    store = grad_buf.uop.store(grad_buf.uop + new_grad)
+    grad_buf.uop = grad_buf.uop.after(store)
    return
-  cur = grad_buf.uop
-  for pad in sorted(pads, key=lambda p: p.marg[0][0] if p.op == Ops.PAD else 0, reverse=True):
-    if pad.op == Ops.PAD:
-      grad_shrink = tuple([(p[0], s+p[0]) for s,p in zip(pad.src[0].shape, pad.marg)])
-      buf_slice = cur.shrink(grad_shrink)
-      cur = cur.after(buf_slice.store(buf_slice + pad.src[0].cast(cur.dtype)))
-    else:
-      cur = cur.after(cur.store(cur + pad.cast(cur.dtype)))
-  grad_buf.uop = cur
+  sorted_pads = sorted(pads, key=lambda p: p.marg[0][0] if p.op == Ops.PAD else 0)
+  inners = [Tensor(p.src[0] if p.op == Ops.PAD else p, device=grad_buf.device).cast(grad_buf.dtype) for p in sorted_pads]
+  grad_buf.assign(grad_buf + inners[0].cat(*inners[1:], dim=0))

 if __name__ == "__main__":
  config = {}
  BS                 = config["BS"]                     = getenv("BS", 16)
  SEQLEN             = config["SEQLEN"]                 = getenv("SEQLEN", 8192)
-  SMALL              = config["SMALL"]                  = getenv("SMALL", 0)

  from examples.llama3 import MODEL_PARAMS
-  model_params = MODEL_PARAMS[llama_size:=getenv("LLAMA3_SIZE", "8B")]["args"]
-  # vocab_size from mixtral tokenizer
-  if not SMALL: model_params |= {"vocab_size": 32000}
-  real_vocab_size = model_params['vocab_size']
-  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: model_params["n_layers"] = llama_layers
-
-  # pad vocab
-  if (MP := getenv("MP", 1)) > 1: model_params["vocab_size"] = round_up(model_params["vocab_size"], 256 * MP)
-  vocab_mask:Tensor = Tensor.arange(model_params["vocab_size"]).reshape(1, 1, -1) >= real_vocab_size
-
+  model_params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
+  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: model_params['n_layers'] = llama_layers
  model = FlatTransformer(**model_params, max_context=SEQLEN)
-
  state = nn.state.get_state_dict(model)
  print("tensor count:", len(state))

  # shard the model
  from tinygrad import Device
-  is_dp = (DP := getenv("DP", 1)) > 1
-  is_mp = (MP := getenv("MP", 1)) > 1
-  is_sharding = is_dp or is_mp
-  device_count = max(DP, MP)
-  device = tuple(f"{Device.DEFAULT}:{i}" for i in range(device_count))
-
-  model.shard(device, is_mp)
-
-  if is_dp: vocab_mask.shard_(device, axis=None).realize()
-  if is_mp: vocab_mask.shard_(device, axis=2).realize()
+  if (DP := getenv("DP", 1)) > 1:
+    model.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(DP)))
+  if (MP := getenv("MP", 1)) > 1:
+    model.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(MP)), mp=True)

  # preallocate all the grad buffers and zero them out
-  grad_dtype = lambda x: dtypes.bfloat16 if x.dtype in dtypes.fp8s else x.dtype
-  grads = {x:x.zeros_like(dtype=grad_dtype(x)).contiguous() for x in state.values() if x.is_param}
-
-  fp8_amax = [t for ts in model._fp8_amax.values() for t in ts]
-  fp8_grad_amax = [t for ts in model._fp8_grad_amax.values() for t in ts]
+  grads = {x:Tensor.zeros(x.shape, dtype=x.dtype, device=x.device).contiguous()
+           for x in state.values() if x.requires_grad is None}

  # print model size
  sz = 0
@@ -406,31 +278,23 @@ if __name__ == "__main__":
    sz += v.nbytes()
  print(f"total sz: {sz/1e9:.2f} GB")

-  with Timing("fake data: "): tokens = Tensor.randint(BS, SEQLEN+1, low=0, high=real_vocab_size, dtype=dtypes.int)
+  with Timing("fake data: "): tokens = Tensor.randint(BS, SEQLEN+1, low=0, high=model.vocab_size, dtype=dtypes.int)
  with Timing("realize weights/grads/data: "): Tensor.realize(*state.values(), *grads.values(), tokens)
  print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
  if DP > 1: tokens = tokens.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(DP)), axis=0)
  if MP > 1: tokens = tokens.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(MP)))

  @TinyJit
-  def fwd_bwd(tokens:Tensor):
-    with Timing("python forward: "):
-      logits = model(tokens[:, :-1], save=llama_size=="8B")
-      loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
+  def jit_step(tokens:Tensor):
+    with Timing("python forward: "): loss = model(tokens[:, :-1]).sparse_categorical_crossentropy(tokens[:, 1:])
    with Timing("python backward: "):
      for t,g in zip(grads, loss.gradient(*grads)):
        apply_grad(grads[t], g.uop)
-    with Timing("run fwd_bwd: "): loss.realize(*grads.values(), *fp8_amax, *fp8_grad_amax)
-
-  @TinyJit
-  def optim_step():
-    for g in grads.values(): g.assign(g.zeros_like())
-    Tensor.realize(*grads.values())
+    with Timing("run step: "): loss.realize(*grads.values())

  for i in range(6):
    GlobalCounters.reset()
    profile_marker(f"step {i}")
    with Timing(colored(f"*** step {i}: ", "red")):
-      fwd_bwd(tokens)
-      optim_step()
+      jit_step(tokens)
  print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
@@ -1,275 +0,0 @@
-import math, os, functools
-if __name__ == "__main__":
-  os.environ["DEFAULT_FLOAT"] = "bfloat16"
-  os.environ["OPTIM_DTYPE"] = "bfloat16"
-  if "DEV" not in os.environ: os.environ["DEV"] = "NULL::gfx950"
-  # CDNA
-  os.environ["DEVICE_IN_FUNCTION_BUG"] = "1"
-  os.environ["ALL2ALL"] = "1"
-  os.environ["USE_ATOMICS"] = "1"
-from tinygrad import Tensor, nn, function, getenv, dtypes, TinyJit
-from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
-from tinygrad.uop.ops import Ops, UOp
-from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
-from extra.llama_kernels.rmsnorm import rmsnorm
-from extra.gemm.cdna_asm_gemm import _mx_block_scale, quantize_mxfp8
-
-FP8_DTYPE = dtypes.fp8e4m3
-FP8_MAX = 448.0
-INIT_STD = 0.008
-
-def _quant_dequant_fwd(x:Tensor) -> Tensor:
-  # x (2d bf16) -> bf16 value after an mxfp8 round-trip (1x32 block scaling on the last axis)
-  M, K = x.shape
-  scale_K = K // 32
-  amax = x.float().reshape(M, scale_K, 32).abs().max(axis=-1)
-  e8 = (amax.maximum(1e-38).log2().floor() + 127).clamp(0, 254).cast(dtypes.uint8)
-  qscale = (127.0 - e8.cast(dtypes.float32)).exp2().reshape(M, scale_K, 1).expand(M, scale_K, 32).reshape(M, K)
-  x_fp8 = (x.float() * qscale).clamp(-FP8_MAX, FP8_MAX).cast(FP8_DTYPE).cast(dtypes.float32)
-  return (x_fp8 * _mx_block_scale(e8)).cast(dtypes.bfloat16)
-
-@functools.cache
-def _quant_dequant_fwd_fxn(x_p, device):
-  return _quant_dequant_fwd(Tensor(x_p, device=device))
-
-def _quant_dequant_bwd(grad:UOp, call:UOp) -> tuple:
-  return (Tensor(grad).cast(dtypes.bfloat16).uop,)
-
-def quant_dequant_mx(x:Tensor) -> Tensor:
-  fxn = _quant_dequant_fwd_fxn(x.as_param(0).uop, x.device)
-  return Tensor(UOp.maketuple(fxn.uop).call(x.uop, grad_fxn=_quant_dequant_bwd).gettuple(0))
-
-def _dequant_fwd(w_q:Tensor, w_scale:Tensor) -> Tensor:
-  return w_q.cast(dtypes.bfloat16) * _mx_block_scale(w_scale)
-
-@functools.cache
-def _dequant_fwd_fxn(wq_p, ws_p, device):
-  return _dequant_fwd(Tensor(wq_p, device=device), Tensor(ws_p, device=device))
-
-def _dequant_bwd(grad:UOp, call:UOp) -> tuple:
-  w_scale = Tensor(call.src[2])
-  return ((Tensor(grad).cast(dtypes.bfloat16) * _mx_block_scale(w_scale).cast(dtypes.bfloat16)).uop, None)
-
-def dequant_weight(w_q:Tensor, w_scale:Tensor) -> Tensor:
-  fxn = _dequant_fwd_fxn(w_q.as_param(0).uop, w_scale.as_param(1).uop, w_q.device)
-  call = UOp.maketuple(fxn.uop).call(w_q.uop, w_scale.uop, grad_fxn=_dequant_bwd)
-  return Tensor(call.gettuple(0))
-
-def matmul_mx(x:Tensor, w_q:Tensor, w_scale:Tensor) -> Tensor:
-  l_shape = x.shape[:-1]
-  x_phys = quant_dequant_mx(x.reshape(-1, x.shape[-1])).reshape(*l_shape, x.shape[-1])
-  w_phys = dequant_weight(w_q, w_scale)
-  return (x_phys @ w_phys.T).cast(dtypes.bfloat16)
-
-def swiglu(x:Tensor, limit:float=7.0, alpha:float=1.702) -> Tensor:
-  x_glu, x_linear = x[..., ::2], x[..., 1::2]
-  x_glu = x_glu.clamp(max_=limit)
-  x_linear = x_linear.clamp(-limit, limit)
-  return (x_glu * (alpha * x_glu).sigmoid()) * (x_linear + 1)
-
-class GPTOSS:
-  def __init__(self, dim:int, n_layers:int, n_heads:int, n_kv_heads:int, head_dim:int, n_experts:int, experts_per_tok:int,
-               intermediate_size:int, vocab_size:int, norm_eps:float=1e-5, rope_theta:int=150000, sliding_window:int=128,
-               swiglu_limit:float=7.0, max_context:int=8192):
-    self.dim, self.n_layers, self.n_heads, self.n_kv_heads, self.head_dim = dim, n_layers, n_heads, n_kv_heads, head_dim
-    self.n_rep = n_heads // n_kv_heads
-    self.n_experts, self.experts_per_tok, self.intermediate_size = n_experts, experts_per_tok, intermediate_size
-    self.vocab_size, self.norm_eps, self.sliding_window, self.swiglu_limit = vocab_size, norm_eps, sliding_window, swiglu_limit
-    self.sm_scale = 1.0 / math.sqrt(head_dim)
-
-    scaled_std = INIT_STD / math.sqrt(2 * n_layers)
-    q_dim, qkv_dim = n_heads * head_dim, head_dim * (n_heads + 2 * n_kv_heads)
-
-    # attn
-    self.wqkv, self.wqkv_scale = self._quant_weight(n_layers, qkv_dim, dim)
-    self.wqkv_bias = Tensor.zeros(n_layers, qkv_dim, dtype=dtypes.bfloat16).contiguous()
-    self.wo, self.wo_scale = self._quant_weight(n_layers, dim, q_dim, std=scaled_std)
-    self.wo_bias = Tensor.zeros(n_layers, dim, dtype=dtypes.bfloat16).contiguous()
-    self.sinks = Tensor.zeros(n_layers, n_heads, dtype=dtypes.bfloat16).contiguous()
-    self.attention_norm = Tensor.ones(n_layers, dim).contiguous()
-
-    # moe ffn
-    self.ffn_norm = Tensor.ones(n_layers, dim).contiguous()
-    self.gate = Tensor.normal(n_layers, n_experts, dim, mean=0.0, std=INIT_STD, dtype=dtypes.bfloat16)
-    self.gate_bias = Tensor.zeros(n_layers, n_experts, dtype=dtypes.bfloat16).contiguous()
-    self.w_gate_up, self.w_gate_up_scale = self._quant_weight(n_layers, n_experts, intermediate_size * 2, dim)
-    self.w_gate_up_bias = Tensor.zeros(n_layers, n_experts, intermediate_size * 2, dtype=dtypes.bfloat16).contiguous()
-    self.w_down, self.w_down_scale = self._quant_weight(n_layers, n_experts, dim, intermediate_size, std=scaled_std)
-    self.w_down_bias = Tensor.zeros(n_layers, n_experts, dim, dtype=dtypes.bfloat16).contiguous()
-
-    # output
-    self.norm = nn.RMSNorm(dim, norm_eps)
-    self.tok_embeddings = nn.Embedding(vocab_size, dim)
-    self.tok_embeddings.weight = Tensor.normal(vocab_size, dim, mean=0.0, std=INIT_STD, dtype=dtypes.bfloat16)
-    self.output = Tensor.normal(vocab_size, dim, mean=0.0, std=INIT_STD, dtype=dtypes.bfloat16)
-    self.freqs_cis = precompute_freqs_cis(head_dim, max_context * 2, rope_theta).contiguous().is_param_(False)
-
-  def _quant_weight(self, *shape:int, std:float=INIT_STD):
-    w = Tensor.zeros(*shape) if getenv("ZEROS") else Tensor.normal(*shape, mean=0.0, std=std)
-    w_q, w_e8, _ = quantize_mxfp8(w)
-    return w_q, w_e8.is_param_(False)
-
-  def _attn_mask(self, seqlen:int, sliding:bool, dtype) -> Tensor:
-    i, j = Tensor.arange(seqlen).reshape(seqlen, 1), Tensor.arange(seqlen).reshape(1, seqlen)
-    allowed = j <= i
-    if sliding: allowed = allowed & (i - j < self.sliding_window)
-    return allowed.where(0.0, -1e30).cast(dtype).contiguous()
-
-  def attention(self, x:Tensor, freqs_cis:Tensor, mask:Tensor, *, attention_norm:Tensor, wqkv:Tensor, wqkv_scale:Tensor,
-                wqkv_bias:Tensor, wo:Tensor, wo_scale:Tensor, wo_bias:Tensor, sinks:Tensor):
-    bsz, seqlen, _ = x.shape
-    x_normed, rrms = rmsnorm(x, self.norm_eps)
-    qkv = matmul_mx(x_normed * attention_norm, wqkv, wqkv_scale) + wqkv_bias
-    qkv = qkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
-    xq = qkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
-    xk, xv = qkv[:, :, :, self.n_rep], qkv[:, :, :, self.n_rep + 1]
-    xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
-
-    xq = xq.cast(dtypes.bfloat16).reshape(bsz, seqlen, self.n_kv_heads, self.n_rep, self.head_dim).permute(0, 2, 3, 1, 4)
-    xk = xk.cast(dtypes.bfloat16).permute(0, 2, 1, 3).unsqueeze(2)
-    xv = xv.cast(dtypes.bfloat16).permute(0, 2, 1, 3).unsqueeze(2)
-    scores = (xq @ xk.transpose(-2, -1)).float() * self.sm_scale + mask
-    sink = sinks.reshape(1, self.n_kv_heads, self.n_rep, 1, 1).float()
-    m = scores.max(-1, keepdim=True).maximum(sink)
-    e = (scores - m).exp()
-    w = (e / (e.sum(-1, keepdim=True) + (sink - m).exp())).cast(dtypes.bfloat16)
-    attn = (w @ xv).permute(0, 3, 1, 2, 4).reshape(bsz, seqlen, self.n_heads * self.head_dim)
-
-    out = matmul_mx(attn, wo, wo_scale) + wo_bias
-    return out, [x_normed, rrms, attn]
-
-  def feed_forward(self, x:Tensor, *, ffn_norm:Tensor, gate:Tensor, gate_bias:Tensor,
-                   w_gate_up:Tensor, w_gate_up_scale:Tensor, w_gate_up_bias:Tensor,
-                   w_down:Tensor, w_down_scale:Tensor, w_down_bias:Tensor):
-    x_normed, rrms = rmsnorm(x, self.norm_eps)
-    inp = x_normed * ffn_norm
-
-    logits = inp.float() @ gate.float().T + gate_bias.float()
-    thresh = logits.topk(self.experts_per_tok)[0][..., -1:]
-    weights = (logits >= thresh).where(logits, -float("inf")).softmax(-1)
-
-    out = None
-    for e in range(self.n_experts):
-      gate_up = matmul_mx(inp, w_gate_up[e], w_gate_up_scale[e]) + w_gate_up_bias[e]
-      y = (matmul_mx(swiglu(gate_up, self.swiglu_limit), w_down[e], w_down_scale[e]) + w_down_bias[e]).contiguous()
-      contrib = weights[..., e:e+1].cast(y.dtype) * y
-      out = contrib if out is None else out + contrib
-    return out, [x_normed, rrms]
-
-  @function(precompile=True, precompile_backward=True)
-  def run_layer(self, x:Tensor, freqs_cis:Tensor, mask:Tensor, attn_kwargs:dict, ffn_kwargs:dict, save:bool=True):
-    attn, attn_saves = self.attention(x, freqs_cis, mask, **attn_kwargs)
-    h = x + attn
-    ffn, ffn_saves = self.feed_forward(h, **ffn_kwargs)
-    h = h + ffn
-    if save: return (h, *attn_saves, *ffn_saves)
-    return (h,)
-
-  def shard(self, device:tuple[str, ...], mp:bool=False):
-    assert not mp, "MP not supported"
-    from tinygrad.nn.state import get_parameters
-    for v in get_parameters(self): v.shard_(device, axis=None)
-    Tensor.realize(*get_parameters(self))
-
-  def __call__(self, tokens:Tensor, save:bool=True):
-    h = self.tok_embeddings(tokens)
-    bsz, seqlen = tokens.shape
-    freqs_cis = self.freqs_cis.cast(h.dtype)[:, :seqlen, :, :, :]
-    mask_full = self._attn_mask(seqlen, False, dtypes.float32)
-    mask_sliding = self._attn_mask(seqlen, True, dtypes.float32)
-    for i in range(self.n_layers):
-      attn_kwargs = dict(attention_norm=self.attention_norm[i], wqkv=self.wqkv[i], wqkv_scale=self.wqkv_scale[i],
-                         wqkv_bias=self.wqkv_bias[i], wo=self.wo[i], wo_scale=self.wo_scale[i], wo_bias=self.wo_bias[i],
-                         sinks=self.sinks[i])
-      ffn_kwargs = dict(ffn_norm=self.ffn_norm[i], gate=self.gate[i], gate_bias=self.gate_bias[i],
-                        w_gate_up=self.w_gate_up[i], w_gate_up_scale=self.w_gate_up_scale[i], w_gate_up_bias=self.w_gate_up_bias[i],
-                        w_down=self.w_down[i], w_down_scale=self.w_down_scale[i], w_down_bias=self.w_down_bias[i])
-      mask = mask_sliding if i % 2 == 0 else mask_full
-      h, *_ = self.run_layer(h, freqs_cis, mask, attn_kwargs, ffn_kwargs, save=save)
-
-    logits = self.norm(h) @ self.output.T
-    return logits
-
-def _get_pads(uop:UOp) -> list[UOp]:
-  if uop.op == Ops.ADD: return _get_pads(uop.src[0]) + _get_pads(uop.src[1])
-  return [uop]
-
-def apply_grad(grad_buf:Tensor, new_grad:UOp):
-  pads = _get_pads(new_grad)
-  if len(pads) <= 1:
-    new_grad = new_grad.cast(grad_buf.dtype)
-    grad_buf.uop = grad_buf.uop.after(grad_buf.uop.store(grad_buf.uop + new_grad))
-    return
-  cur = grad_buf.uop
-  for pad in sorted(pads, key=lambda p: p.marg[0][0] if p.op == Ops.PAD else 0, reverse=True):
-    if pad.op == Ops.PAD:
-      grad_shrink = tuple([(p[0], s+p[0]) for s,p in zip(pad.src[0].shape, pad.marg)])
-      buf_slice = cur.shrink(grad_shrink)
-      cur = cur.after(buf_slice.store(buf_slice + pad.src[0].cast(cur.dtype)))
-    else:
-      cur = cur.after(cur.store(cur + pad.cast(cur.dtype)))
-  grad_buf.uop = cur
-
-GPT_OSS_20B = dict(dim=2880, n_layers=24, n_heads=64, n_kv_heads=8, head_dim=64, n_experts=32, experts_per_tok=4,
-                   intermediate_size=2880, vocab_size=128256, norm_eps=1e-5, rope_theta=150000, sliding_window=128,
-                   swiglu_limit=7.0)
-
-if __name__ == "__main__":
-  config = {}
-  BS      = config["BS"]      = getenv("BS", 16)
-  SEQLEN  = config["SEQLEN"]  = getenv("SEQLEN", 8192)
-
-  model_params = GPT_OSS_20B
-  real_vocab_size = model_params["vocab_size"]
-  if (layers := getenv("LAYERS")) != 0: model_params["n_layers"] = layers
-
-  model = GPTOSS(**model_params, max_context=SEQLEN)
-
-  state = nn.state.get_state_dict(model)
-  print("tensor count:", len(state))
-
-  from tinygrad import Device
-  is_dp = (DP := getenv("DP", 1)) > 1
-  device_count = DP
-  device = tuple(f"{Device.DEFAULT}:{i}" for i in range(device_count))
-
-  if is_dp: model.shard(device)
-
-  # preallocate all the grad buffers and zero them out
-  grad_dtype = lambda x: dtypes.bfloat16 if x.dtype in dtypes.fp8s else x.dtype
-  grads = {x:x.zeros_like(dtype=grad_dtype(x)).contiguous() for x in state.values() if x.is_param}
-
-  # print model size
-  sz = 0
-  for k,v in state.items():
-    print(f"{colored(k, 'green' if v in grads else 'white'):30s} {str(v.shape):30s} {str(v.dtype):20s} {v.device}  {v.nbytes()/1e9:.2f} GB")
-    sz += v.nbytes()
-  print(f"total sz: {sz/1e9:.2f} GB")
-
-  with Timing("fake data: "): tokens = Tensor.randint(BS, SEQLEN+1, low=0, high=real_vocab_size, dtype=dtypes.int)
-  with Timing("realize weights/grads/data: "): Tensor.realize(*state.values(), *grads.values(), tokens)
-  print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
-  if is_dp: tokens = tokens.shard(device, axis=0)
-
-  @TinyJit
-  def fwd_bwd(tokens:Tensor):
-    with Timing("python forward: "):
-      logits = model(tokens[:, :-1], save=True)
-      loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
-    with Timing("python backward: "):
-      for t,g in zip(grads, loss.gradient(*grads)):
-        apply_grad(grads[t], g.uop)
-    with Timing("run fwd_bwd: "): loss.realize(*grads.values())
-
-  @TinyJit
-  def optim_step():
-    for g in grads.values(): g.assign(g.zeros_like())
-    Tensor.realize(*grads.values())
-
-  for i in range(6):
-    GlobalCounters.reset()
-    profile_marker(f"step {i}")
-    with Timing(colored(f"*** step {i}: ", "red")):
-      fwd_bwd(tokens)
-      optim_step()
-  print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
@@ -1,68 +0,0 @@
-import unittest
-from tinygrad import Tensor, TinyJit
-from tinygrad.nn.state import get_parameters
-from examples.mlperf.models.flat_llama import apply_grad
-
-class FlatModel:
-  def __init__(self, n_layers:int, dim:int, hidden:int):
-    self.n_layers = n_layers
-    self.w1 = Tensor.uniform(n_layers, dim, hidden, low=-0.1, high=0.1)
-    self.w2 = Tensor.uniform(n_layers, hidden, dim, low=-0.1, high=0.1)
-    self.scale = Tensor.uniform(dim, low=0.9, high=1.1)
-    self.bias = Tensor.zeros(dim).contiguous()
-
-  def __call__(self, x:Tensor) -> Tensor:
-    h = x
-    for i in range(self.n_layers):
-      h = (h @ self.w1[i]).relu() @ self.w2[i] + h
-    return (h * self.scale + self.bias).sum()
-
-class TestApplyGradE2E(unittest.TestCase):
-  def _run_with_apply_grad(self, model, xs):
-    grads = {p: Tensor.zeros(p.shape, dtype=p.dtype).contiguous().realize() for p in get_parameters(model)}
-    for x in xs:
-      loss = model(x)
-      for p, g in zip(grads, loss.gradient(*grads)):
-        apply_grad(grads[p], g.uop)
-      Tensor.realize(loss, *grads.values())
-    return [grads[p] for p in get_parameters(model)]
-
-  def _run_reference(self, model, xs):
-    for x in xs: model(x).backward()
-    return [p.grad for p in get_parameters(model)]
-
-  def _assert_close(self, got, expected, atol, rtol):
-    for g, e in zip(got, expected):
-      self.assertTrue(g.allclose(e, atol=atol, rtol=rtol).item(), f"grad mismatch (max abs diff {(g - e).abs().max().item()})")
-
-  def _assert_match(self, model, xs, atol, rtol):
-    self._assert_close(self._run_with_apply_grad(model, xs), self._run_reference(model, xs), atol, rtol)
-
-  def test_e2e_single_step(self):
-    model = FlatModel(n_layers=3, dim=8, hidden=16)
-    Tensor.realize(*get_parameters(model))
-    self._assert_match(model, [Tensor.randn(2, 8).realize()], atol=1e-4, rtol=1e-4)
-
-  def test_e2e_multi_step_accumulation(self):
-    model = FlatModel(n_layers=4, dim=8, hidden=16)
-    Tensor.realize(*get_parameters(model))
-    self._assert_match(model, [Tensor.randn(2, 8).realize() for _ in range(3)], atol=1e-4, rtol=1e-4)
-
-  def test_e2e_jit(self):
-    model = FlatModel(n_layers=3, dim=8, hidden=16)
-    Tensor.realize(*get_parameters(model))
-    grads = {p: Tensor.zeros(p.shape, dtype=p.dtype).contiguous().realize() for p in get_parameters(model)}
-
-    @TinyJit
-    def fwd_bwd(x:Tensor):
-      loss = model(x)
-      for p, g in zip(grads, loss.gradient(*grads)): apply_grad(grads[p], g.uop)
-      Tensor.realize(loss, *grads.values())
-
-    xs = [Tensor.randn(2, 8).realize() for _ in range(3)]
-    for x in xs: fwd_bwd(x)
-    self._assert_close([grads[p] for p in get_parameters(model)], self._run_reference(model, xs), atol=1e-3, rtol=1e-3)
-
-
-if __name__ == "__main__":
-  unittest.main()
@@ -3,7 +3,8 @@ os.environ["WQKV"] = "1"
 import unittest
 import numpy as np
 from tinygrad import Tensor, nn, dtypes
-from tinygrad.device import Device
+from tinygrad.nn.state import get_parameters
+from tinygrad.device import is_dtype_supported, Device
 from examples.mlperf.models.llama import Transformer
 from examples.mlperf.models.flat_llama import FlatTransformer

@@ -44,6 +45,8 @@ class TestFlatLlama(unittest.TestCase):
    flat = FlatTransformer(**params)
    copy_weights(flat, ref)

+    for p in get_parameters(ref): p.requires_grad_(True)
+    for p in get_parameters(flat): p.requires_grad_(True)
    Tensor.realize(*nn.state.get_state_dict(flat).values())

    tokens = Tensor([[1, 50, 100, 999, 2, 10]])
@@ -111,7 +114,7 @@ class TestFlatLlama(unittest.TestCase):
    self.assertEqual(ref_logits.shape, flat_logits.shape)
    np.testing.assert_allclose(flat_logits, ref_logits, atol=1e-4, rtol=1e-4)

-  @unittest.skipUnless(dtypes.fp8e4m3 in Device[Device.DEFAULT].renderer.supported_dtypes(), "fp8 not supported on this device")
+  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3), "fp8 not supported on this device")
  def test_forward_fp8(self):
    import examples.mlperf.models.flat_llama as flat_llama_mod
    old_fp8 = flat_llama_mod.FP8
@@ -6,10 +6,6 @@ from tinygrad.uop.ops import UOp, Ops

 STOCHASTIC_ROUND = getenv("STOCHASTIC_ROUND", 0)
 MASTER_WEIGHTS = getenv("MASTER_WEIGHTS", 0)
-ZERO_OPTIM = getenv("ZERO_OPTIM", 0)
-FP8_AMAX_MARGIN = getenv("FP8_AMAX_MARGIN", 1.1)
-IMMEDIATE_SCALE = getenv("IMMEDIATE_SCALE", 0)
-MXFP8 = getenv("MXFP8", 0)

 def stochastic_round_bf16(x:Tensor) -> Tensor:
  bits = x.bitcast(dtypes.uint32)
@@ -25,24 +21,11 @@ class GradAccClipAdamW(Optimizer):
  def __init__(self, params:list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-6, weight_decay=0.0, grad_acc=1, clip_norm=1.0, device=None, fused=FUSE_OPTIM):
    super().__init__(params, lr, device, fused)
    self.b1, self.b2, self.eps, self.wd = b1, b2, eps, weight_decay
-    self.b1_t, self.b2_t = (Tensor.ones((1,), dtype=dtypes.float32, device=self.device) for _ in [b1, b2])
-    self.zero = bool(ZERO_OPTIM) and isinstance(self.device, tuple) and not self.fused
-    self.m = [self._zero_shard(x) for x in self._new_optim_param()]
-    self.v = [self._zero_shard(x) for x in self._new_optim_param()]
+    self.b1_t, self.b2_t = (Tensor.ones((1,), dtype=dtypes.float32, device=self.device, requires_grad=False) for _ in [b1, b2])
+    self.m = self._new_optim_param()
+    self.v = self._new_optim_param()
    self.grad_acc, self.clip_norm = grad_acc, clip_norm
-    if MASTER_WEIGHTS and self.params[0].dtype != dtypes.float32:
-      self.master_params:list[Tensor]|None = [self._zero_shard(p.to(self.device).float().contiguous()) for p in self.params]
-    else:
-      self.master_params = None
-
-  def _zero_shard(self, t:Tensor) -> Tensor:
-    if not self.zero or (t.shape[0] % len(self.device)) != 0: return t
-    return Tensor(t.uop._shard(0, len(self.device)).multi(0)).clone()
-
-  def _zero_gather(self, t:Tensor) -> Tensor:
-    if not isinstance(t.device, tuple) or t.uop.axis != 0: return t
-    n, sz = len(t.device), t.shape[0] // len(t.device)
-    return Tensor.cat(*[t[p*sz:(p+1)*sz] for p in range(n)], dim=0)
+    self.master_params:list[Tensor]|None = [p.float().contiguous() for p in self.params] if MASTER_WEIGHTS and self.params[0].dtype != dtypes.float32 else None

  def fstep(self, grads:list[Tensor]):
    if self.fused:
@@ -51,10 +34,7 @@ class GradAccClipAdamW(Optimizer):
    else:
      updates, extra = self._step([], grads)
    for i, tt in enumerate(self.params): tt.assign(self._apply_update(tt, updates[i], self.master_params[i] if self.master_params else None))
-    # collect inv_scale tensors attached to fp8 params (set by _apply_update)
-    fp8_inv_scales = [tt._inv_scale for tt in self.params if hasattr(tt, '_inv_scale')]
-    fp8_next_inv_scales = [tt._next_inv_scale for tt in self.params if hasattr(tt, '_next_inv_scale')]
-    to_realize = extra+self.params+self.buffers+(self.master_params or [])+fp8_inv_scales+fp8_next_inv_scales
+    to_realize = extra+self.params+self.buffers+(self.master_params or [])

    Tensor.realize(*to_realize)
    return extra[-1]
@@ -96,38 +76,5 @@ class GradAccClipAdamW(Optimizer):
    up = up.float().shard_like(w) + self.lr.to(w.device) * wd * w.detach()
    new_w = w.detach() - up
    if master is not None: master.assign(new_w)
-    if self.zero: new_w = self._zero_gather(new_w)
-    # when master is offloaded to a different device than the param, results are resharded back onto the param's (sharded) device
-    offloaded = master is not None and master.device != t.device
-    if STOCHASTIC_ROUND and t.dtype == dtypes.bfloat16:
-      out = stochastic_round_bf16(new_w)
-      return out.shard_like(t) if offloaded else out
-    if t.dtype in dtypes.fp8s:
-      if MXFP8:
-        from extra.gemm.cdna_asm_gemm import quantize_mxfp8
-        w_q, w_e8, _ = quantize_mxfp8(new_w.reshape(-1, new_w.shape[-1]))
-        new_e8 = w_e8.reshape(t._inv_scale.shape)
-        t._inv_scale.assign(new_e8.shard_like(t._inv_scale) if offloaded else new_e8)
-        ret = w_q.reshape(new_w.shape)
-        return ret.shard_like(t) if offloaded else ret
-      from examples.mlperf.models.flat_llama import FP8_MAX
-      if IMMEDIATE_SCALE:
-        amax_axis = tuple(range(t._inv_scale.ndim, new_w.ndim))
-        new_inv = ((new_w.float().abs().max(axis=amax_axis).detach() + 1e-8) / FP8_MAX).cast(t._inv_scale.dtype)
-        t._inv_scale.assign(new_inv.shard_like(t._inv_scale) if offloaded else new_inv)
-        scale = new_inv.reciprocal().reshape(*new_inv.shape, *([1]*(new_w.ndim-new_inv.ndim)))
-        ret = (new_w * scale).clamp(-FP8_MAX, FP8_MAX).cast(t.dtype)
-        return ret.shard_like(t) if offloaded else ret
-      # delayed scaling: reuse previous step's inv_scale
-      t._inv_scale.assign(t._next_inv_scale)
-      inv_scale = t._inv_scale.to(new_w.device) if offloaded else t._inv_scale
-      scale = inv_scale.reciprocal().reshape(*inv_scale.shape, *([1]*(new_w.ndim-inv_scale.ndim)))
-      scaled = (new_w * scale).clamp(-FP8_MAX, FP8_MAX)
-      ret = scaled.cast(t.dtype)
-      # update inv_scale for next step from quantized result
-      new_amax = (ret.float().abs().max(axis=tuple(range(inv_scale.ndim, ret.ndim))) * inv_scale * FP8_AMAX_MARGIN).detach()
-      new_inv = ((new_amax + 1e-8) / FP8_MAX).cast(t._inv_scale.dtype)
-      t._next_inv_scale.assign(new_inv.shard_like(t._next_inv_scale) if offloaded else new_inv)
-      return ret.shard_like(t) if offloaded else ret
-    out = new_w.cast(t.dtype)
-    return out.shard_like(t) if offloaded else out
+    if STOCHASTIC_ROUND and t.dtype == dtypes.bfloat16: return stochastic_round_bf16(new_w)
+    return new_w.cast(t.dtype)
@@ -1,44 +0,0 @@
-#!/usr/bin/env bash
-
-export PYTHONPATH="."
-export PATH="/opt/rocm-7.1.1/bin:$PATH"
-export ROCM_PATH="/opt/rocm-7.1.1"
-export DEV=${DEV:-AMD}
-export CHECK_OOB=0
-export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
-export DEVICE_IN_FUNCTION_BUG=1
-
-export DEBUG=${DEBUG:-2}
-export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
-export ALL2ALL=${ALL2ALL:-1}
-export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
-export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
-export USE_ATOMICS=${USE_ATOMICS:-1}
-export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
-export MXFP8=${MXFP8:-1}
-export ZERO_OPTIM=${ZERO_OPTIM:-1}
-export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}
-
-export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=${DP:-8} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
-export GBS=$((BS * GRADIENT_ACC_STEPS))
-
-export MODEL="gptoss"
-export BASEDIR="/raid/datasets/c4-8b/"
-export EVAL_TARGET=3.34 EVAL_FREQ=12288
-export END_LR="4e-5" WARMUP_STEPS=128 MAX_STEPS=1200000
-export SAMPLES=$((MAX_STEPS * GBS))
-export SEQLEN=${SEQLEN:-8192}
-
-export SEED=${SEED:-5760}
-export DATA_SEED=${DATA_SEED:-5760}
-
-export JITBEAM=${JITBEAM:-3}
-export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
-
-export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
-if [ -z "$FULL_LAYERS" ]; then
-  export LAYERS=${LAYERS:-2}
-fi
-
-python3 examples/mlperf/model_train.py
@@ -1,39 +0,0 @@
-#!/usr/bin/env bash
-
-export PYTHONPATH="."
-export PATH="/opt/rocm-7.1.1/bin:$PATH"
-export ROCM_PATH="/opt/rocm-7.1.1"
-export DEV=${DEV:-AMD}
-export CHECK_OOB=0
-export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
-export DEVICE_IN_FUNCTION_BUG=1
-
-export DEBUG=${DEBUG:-0}
-export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
-export ALL2ALL=${ALL2ALL:-1}
-export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
-export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
-export USE_ATOMICS=${USE_ATOMICS:-1}
-export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
-export MXFP8=${MXFP8:-1}
-export ZERO_OPTIM=${ZERO_OPTIM:-1}
-export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}
-
-export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=${DP:-8} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
-export GBS=$((BS * GRADIENT_ACC_STEPS))
-
-export MODEL="gptoss"
-export BASEDIR="/raid/datasets/c4-8b/"
-export EVAL_TARGET=3.34 EVAL_FREQ=12288
-export END_LR="4e-5" WARMUP_STEPS=128 MAX_STEPS=1200000
-export SAMPLES=$((MAX_STEPS * GBS))
-export SEQLEN=${SEQLEN:-8192}
-
-export SEED=${SEED:-$RANDOM}
-export DATA_SEED=${DATA_SEED:-5760}
-
-export JITBEAM=${JITBEAM:-3}
-export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
-
-python3 examples/mlperf/model_train.py
@@ -1,6 +0,0 @@
-#!/bin/bash
-export BENCHMARK=${BENCHMARK:-5}
-export EVAL_BS=0
-VIZ=${VIZ:--1} FULL_LAYERS=1 DEBUG=${DEBUG:--0} examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh
-SRC="AMD"; [[ $DEV == NULL* ]] && SRC="NULL"
-[ "$BENCHMARK" -le 3 ] || python -m tinygrad.viz.cli -s "$SRC" -t --interval "train @ 2" "train @ 3"
@@ -1,9 +1,8 @@
 #!/usr/bin/env bash

 export PYTHONPATH="."
-export PATH="/opt/rocm-7.1.1/bin:$PATH"
-export ROCM_PATH="/opt/rocm-7.1.1"
 export DEV=${DEV:-AMD}
+export EMULATE="AMD_CDNA4"
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
 export DEVICE_IN_FUNCTION_BUG=1
@@ -11,24 +10,14 @@ export DEVICE_IN_FUNCTION_BUG=1
 export DEBUG=${DEBUG:-2}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
 export ALL2ALL=${ALL2ALL:-1}
-export LATE_ALLREDUCE=${LATE_ALLREDUCE:-1}
-export USE_ATOMICS=${USE_ATOMICS:-1}
+export USE_ATOMICS=${USE_ATOMICS:-0}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
-export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
-export FP8=${FP8:-1}
-export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
-export FAST_CE=${FAST_CE:-0}
-export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0}
-export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0}
-export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0}
-export FUSED_SILU_W13=${FUSED_SILU_W13:-0}
-export SPLIT_W13=${SPLIT_W13:-1}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}

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

 export MODEL="llama3"
 export BASEDIR="/raid/datasets/c4/"
@@ -41,9 +30,9 @@ export DATA_SEED=${DATA_SEED:-5760}
 export JITBEAM=${JITBEAM:-3}
 export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1

-export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
+export FAKEDATA=1 BENCHMARK=10
 if [ -z "$FULL_LAYERS" ]; then
-  export LLAMA_LAYERS=${LLAMA_LAYERS:-2}
+  export LLAMA_LAYERS=2
 fi

 python3 examples/mlperf/model_train.py
@@ -1,34 +1,22 @@
 #!/usr/bin/env bash

 export PYTHONPATH="."
-export PATH="/opt/rocm-7.1.1/bin:$PATH"
-export ROCM_PATH="/opt/rocm-7.1.1"
 export DEV=${DEV:-AMD}
+export EMULATE="AMD_CDNA4"
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
-export DEVICE_IN_FUNCTION_BUG=1

 export DEBUG=${DEBUG:-0}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
 export ALL2ALL=${ALL2ALL:-1}
-export LATE_ALLREDUCE=${LATE_ALLREDUCE:-1}
-export USE_ATOMICS=${USE_ATOMICS:-1}
+export USE_ATOMICS=${USE_ATOMICS:-0}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
-export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
-export FP8=${FP8:-1}
-export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
-export FAST_CE=${FAST_CE:-0}
-export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0}
-export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0}
-export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0}
-export FUSED_SILU_W13=${FUSED_SILU_W13:-0}
-export SPLIT_W13=${SPLIT_W13:-1}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}

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

 export MODEL="llama3"
 export BASEDIR="/raid/datasets/c4/"
@@ -1,8 +1,6 @@
 #!/usr/bin/env bash

 export PYTHONPATH="."
-export PATH="/opt/rocm-7.1.1/bin:$PATH"
-export ROCM_PATH="/opt/rocm-7.1.1"
 export DEV=${DEV:-AMD}
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
@@ -11,23 +9,15 @@ export DEVICE_IN_FUNCTION_BUG=1
 export DEBUG=${DEBUG:-2}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
 export ALL2ALL=${ALL2ALL:-1}
-export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
 export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
 export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
 export FP8=${FP8:-1}
 export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
-export FAST_CE=${FAST_CE:-1}
-export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
-export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-1}
-export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
-export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
-export SPLIT_W13=${SPLIT_W13:-0}
-export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}

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

 export MODEL="llama3"
@@ -46,9 +36,9 @@ export DATA_SEED=${DATA_SEED:-5760}
 export JITBEAM=${JITBEAM:-3}
 export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1

-export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
+export FAKEDATA=1 BENCHMARK=${BENCHMARK:-10}
 if [ -z "$FULL_LAYERS" ]; then
-  export LLAMA_LAYERS=${LLAMA_LAYERS:-2}
+  export LLAMA_LAYERS=2
 fi

 python3 examples/mlperf/model_train.py
@@ -1,9 +1,8 @@
 #!/usr/bin/env bash

 export PYTHONPATH="."
-export PATH="/opt/rocm-7.1.1/bin:$PATH"
-export ROCM_PATH="/opt/rocm-7.1.1"
 export DEV=${DEV:-AMD}
+export EMULATE="AMD_CDNA4"
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
 export DEVICE_IN_FUNCTION_BUG=1
@@ -11,19 +10,9 @@ export DEVICE_IN_FUNCTION_BUG=1
 export DEBUG=${DEBUG:-2}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
 export ALL2ALL=${ALL2ALL:-1}
-export LATE_ALLREDUCE=${LATE_ALLREDUCE:-1}
-export USE_ATOMICS=${USE_ATOMICS:-1}
+export USE_ATOMICS=${USE_ATOMICS:-0}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
-export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
-export FP8=${FP8:-1}
-export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
-export FAST_CE=${FAST_CE:-0}
-export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0}
-export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0}
-export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0}
-export FUSED_SILU_W13=${FUSED_SILU_W13:-0}
-export SPLIT_W13=${SPLIT_W13:-1}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}

 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
@@ -46,9 +35,9 @@ export DATA_SEED=${DATA_SEED:-5760}
 export JITBEAM=${JITBEAM:-3}
 export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1

-export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
+export FAKEDATA=1 BENCHMARK=10
 if [ -z "$FULL_LAYERS" ]; then
-  export LLAMA_LAYERS=${LLAMA_LAYERS:-2}
+  export LLAMA_LAYERS=2
 fi

 python3 examples/mlperf/model_train.py
@@ -1,8 +1,6 @@
 #!/usr/bin/env bash

 export PYTHONPATH="."
-export PATH="/opt/rocm-7.1.1/bin:$PATH"
-export ROCM_PATH="/opt/rocm-7.1.1"
 export DEV=${DEV:-AMD}
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
@@ -11,23 +9,15 @@ export DEVICE_IN_FUNCTION_BUG=1
 export DEBUG=${DEBUG:-0}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
 export ALL2ALL=${ALL2ALL:-1}
-export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
 export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
 export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
 export FP8=${FP8:-1}
 export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
-export FAST_CE=${FAST_CE:-1}
-export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
-export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-1}
-export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
-export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
-export SPLIT_W13=${SPLIT_W13:-0}
-export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}

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

 export MODEL="llama3"
@@ -1,9 +1,8 @@
 #!/usr/bin/env bash

 export PYTHONPATH="."
-export PATH="/opt/rocm-7.1.1/bin:$PATH"
-export ROCM_PATH="/opt/rocm-7.1.1"
 export DEV=${DEV:-AMD}
+export EMULATE="AMD_CDNA4"
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
 export DEVICE_IN_FUNCTION_BUG=1
@@ -11,19 +10,9 @@ export DEVICE_IN_FUNCTION_BUG=1
 export DEBUG=${DEBUG:-0}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
 export ALL2ALL=${ALL2ALL:-1}
-export LATE_ALLREDUCE=${LATE_ALLREDUCE:-1}
-export USE_ATOMICS=${USE_ATOMICS:-1}
+export USE_ATOMICS=${USE_ATOMICS:-0}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
-export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
-export FP8=${FP8:-1}
-export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
-export FAST_CE=${FAST_CE:-0}
-export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0}
-export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0}
-export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0}
-export FUSED_SILU_W13=${FUSED_SILU_W13:-0}
-export SPLIT_W13=${SPLIT_W13:-1}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}

 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
@@ -0,0 +1,6 @@
+#!/bin/bash
+export BENCHMARK=5
+export EVAL_BS=0
+VIZ=${VIZ:--1} FULL_LAYERS=1 DEBUG=0 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_beam.sh
+SRC="AMD"; [[ $DEV == NULL* ]] && SRC="NULL"
+extra/viz/cli.py --profile -s "$SRC"
@@ -3,8 +3,6 @@ set -e  # Exit on any error
 set -o pipefail  # Make pipeline fail if any command fails

 export PYTHONPATH="."
-export PATH="/opt/rocm-7.1.1/bin:$PATH"
-export ROCM_PATH="/opt/rocm-7.1.1"
 export DEV=AMD
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
@@ -12,22 +10,15 @@ export DEVICE_IN_FUNCTION_BUG=1

 export HK_FLASH_ATTENTION=1
 export ALL2ALL=1
-export LATE_ALLREDUCE=0
 export USE_ATOMICS=1
 export ASM_GEMM=1
 export WQKV=1
 export MASTER_WEIGHTS=1
 export FP8=1
 export ALLREDUCE_CAST=1
-export FAST_CE=1
-export FUSED_INPUT_QUANTIZE=1
-export FUSED_GRAD_QUANTIZE=1
-export FUSED_ADD_NORM_MUL_QUANTIZE=1
-export FUSED_SILU_W13=1
-export SPLIT_W13=0

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

 export MODEL="llama3"
@@ -4,7 +4,7 @@ export EVAL_BS=0
 export FAKEDATA=1
 export NULL_ALLOW_COPYOUT=1
 export HIP_VISIBLE_DEVICES=""
-export DEV=NULL:HIP:gfx950
+export DEV=NULL
 export JITBEAM=0
 export LLAMA_LAYERS=${LLAMA_LAYERS:-"2"}
-time examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh
+time examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh
@@ -3,7 +3,7 @@ import torch
 from torchvision.utils import make_grid, save_image
 from tinygrad.nn.state import get_parameters
 from tinygrad.tensor import Tensor
-from tinygrad.helpers import trange, Context
+from tinygrad.helpers import trange
 from tinygrad.nn import optim
 from tinygrad.nn.datasets import mnist

@@ -71,7 +71,7 @@ def train_generator(optimizer, data_fake):
 if __name__ == "__main__":
  # data for training and validation
  X_train, _, _, _ = mnist()
-  ds_noise = Tensor.randn(64, 128)
+  ds_noise = Tensor.randn(64, 128, requires_grad=False)
  # parameters
  epochs, batch_size, k = 300, 512, 1
  sample_interval = epochs // 10
@@ -86,7 +86,7 @@ if __name__ == "__main__":
  optim_g = optim.Adam(get_parameters(generator), lr=0.0002, b1=0.5)  # 0.0002 for equilibrium!
  optim_d = optim.Adam(get_parameters(discriminator), lr=0.0002, b1=0.5)
  # training loop
-  with Context(TRAINING=1):
+  with Tensor.train():
    for epoch in (t := trange(epochs)):
      loss_g, loss_d = 0.0, 0.0
      for _ in range(n_steps):
@@ -4,7 +4,7 @@ if "JIT_BATCH_SIZE" not in os.environ: os.environ["JIT_BATCH_SIZE"] = "0"

 from tinygrad import fetch, Tensor, TinyJit, Context, GlobalCounters, Device, dtypes
 from tinygrad.helpers import DEBUG, getenv
-from tinygrad.uop.ops import Ops
+from tinygrad.engine.realize import CompiledRunner
 from tinygrad.nn.onnx import OnnxRunner

 OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx"
@@ -21,8 +21,6 @@ def compile(onnx_file):
  # TODO this seems dumb
  input_types = {k:(dtypes.float32 if v is dtypes.float16 else v) for k,v in input_types.items()}
  Tensor.manual_seed(100)
-  # replace symbolic dimensions (e.g. 'b' for dynamic batch) with 1
-  input_shapes = {k:tuple(s if isinstance(s, int) else 1 for s in shp) for k,shp in input_shapes.items()}
  inputs = {k:Tensor(Tensor.randn(*shp, dtype=input_types[k]).mul(8).realize().numpy(), device='NPY') for k,shp in sorted(input_shapes.items())}
  if not getenv("NPY_IMG"):
    inputs = {k:Tensor(v.numpy(), device=Device.DEFAULT).realize() if 'img' in k else v for k,v in inputs.items()}
@@ -37,26 +35,21 @@ def compile(onnx_file):
      ret = run_onnx_jit(**inputs).numpy()
    # copy i == 1 so use of JITBEAM is okay
    if i == 1: test_val = np.copy(ret)
-  # iterate kernel CALLs in the captured LINEAR UOp; toposort descends into batched graph CUSTOM_FUNCTIONs
-  kernel_asts = {Ops.PROGRAM}
-  kernel_calls = [u for u in run_onnx_jit.captured.linear.toposort(gate=lambda x: x.op not in kernel_asts)
-                  if u.op is Ops.CALL and u.src[0].op in kernel_asts]
-  print(f"captured {len(kernel_calls)} kernels")
-  if getenv("TEST", 1): np.testing.assert_equal(test_val, ret, "JIT run failed")
+  print(f"captured {len(run_onnx_jit.captured.jit_cache)} kernels")
+  np.testing.assert_equal(test_val, ret, "JIT run failed")
  print("jit run validated")

  # check gated read_image usage
  kernel_count = 0
  read_image_count = 0
  gated_read_image_count = 0
-  for call in kernel_calls:
-    _, _, source, _ = call.src[0].src
-    src = source.arg
-    kernel_count += 1
-    read_image_count += src.count("read_image")
-    gated_read_image_count += src.count("?read_image")
-    for v in [m.group(1) for m in re.finditer(r'(val\d+)\s*=\s*read_imagef\(', src)]:
-      if len(re.findall(fr'[\?\:]{v}\.[xyzw]', src)) > 0: gated_read_image_count += 1
+  for ei in run_onnx_jit.captured.jit_cache:
+    if isinstance(ei.prg, CompiledRunner):
+      kernel_count += 1
+      read_image_count += ei.prg.p.src.count("read_image")
+      gated_read_image_count += ei.prg.p.src.count("?read_image")
+      for v in [m.group(1) for m in re.finditer(r'(val\d+)\s*=\s*read_imagef\(', ei.prg.p.src)]:
+        if len(re.findall(fr'[\?\:]{v}\.[xyzw]', ei.prg.p.src)) > 0: gated_read_image_count += 1
  print(f"{kernel_count=},  {read_image_count=}, {gated_read_image_count=}")
  if (allowed_kernel_count:=getenv("ALLOWED_KERNEL_COUNT", -1)) != -1:
    assert kernel_count == allowed_kernel_count, f"different kernels! {kernel_count=}, {allowed_kernel_count=}"
@@ -65,30 +58,6 @@ def compile(onnx_file):
  if (allowed_gated_read_image:=getenv("ALLOWED_GATED_READ_IMAGE", -1)) != -1:
    assert gated_read_image_count == allowed_gated_read_image, f"different gated read_image! {gated_read_image_count=}, {allowed_gated_read_image=}"

-  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_RPT", 1):
-    from extra.gemm.qcom_openpilot_vision_fp16 import patch_fp32_rpt
-    if (patched:=patch_fp32_rpt(run_onnx_jit)): print(f"repeat-packed {patched} QCOM vision kernels")
-  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_SCHEDULE", 1):
-    from extra.gemm.qcom_openpilot_schedule_projection import patch_projection
-    if (patched:=patch_projection(run_onnx_jit)): print(f"rescheduled {patched} QCOM vision kernels")
-  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_FULL_RPT", 1):
-    from extra.gemm.qcom_openpilot_inverse_full_rpt import patch_model as patch_full_rpt
-    if (patched:=patch_full_rpt(run_onnx_jit)): print(f"fully repeat-packed {patched} QCOM vision kernels")
-  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_DEDUPE", 1):
-    from extra.gemm.qcom_openpilot_dedupe_head import dedupe_identical_calls
-    if (removed:=dedupe_identical_calls(run_onnx_jit)): print(f"deduplicated {len(removed)} QCOM kernels")
-  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_PACK_CONV", 1):
-    from extra.gemm.qcom_openpilot_pack_conv_weights import patch_conv
-    if (patched:=patch_conv(run_onnx_jit)): print(f"packed weights for {patched} QCOM convolution kernels")
-  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_LEVEL_SCHEDULE", 1):
-    from extra.gemm.qcom_openpilot_level_schedule import schedule_levels
-    if (moved:=schedule_levels(run_onnx_jit)): print(f"rescheduled {moved} QCOM kernels by dependency level")
-  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_BATCH_HEAD", 1):
-    from extra.gemm.qcom_openpilot_batch_head import batch_head
-    if (combined:=batch_head(run_onnx_jit)): print(f"batched {combined} groups of QCOM head kernels")
-  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_INPUT_PACK", 1):
-    from extra.gemm.qcom_openpilot_input_pack import patch_input_pack
-    if (patched:=patch_input_pack(run_onnx_jit)): print(f"vectorized {patched} QCOM input kernel")
  with open(OUTPUT, "wb") as f:
    pickle.dump(run_onnx_jit, f)
  mdl_sz = os.path.getsize(onnx_file)
@@ -96,7 +65,7 @@ def compile(onnx_file):
  print(f"mdl size is {mdl_sz/1e6:.2f}M")
  print(f"pkl size is {pkl_sz/1e6:.2f}M")
  print("**** compile done ****")
-  return run_onnx_jit, inputs, test_val
+  return inputs, test_val

 def test_vs_compile(run, inputs, test_val=None):

@@ -111,7 +80,7 @@ def test_vs_compile(run, inputs, test_val=None):
    step_times.append((et-st)*1e3)
    print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {step_times[-1]:6.2f} ms")

-  if (assert_time:=getenv("ASSERT_MIN_STEP_TIME", 0.0)):
+  if (assert_time:=getenv("ASSERT_MIN_STEP_TIME")):
    min_time = min(step_times)
    assert min_time < assert_time, f"Speed regression, expected min step time of < {assert_time} ms but took: {min_time} ms"

@@ -128,7 +97,7 @@ def test_vs_compile(run, inputs, test_val=None):
 def test_vs_onnx(new_inputs, test_val, onnx_file, tol):
  import onnx
  import onnxruntime as ort
-
+  
  onnx_inputs = {k:v.numpy() for k,v in new_inputs.items()}
  onnx_model = onnx.load(onnx_file)

@@ -159,21 +128,14 @@ def bench(run, inputs):
      run(**inputs).numpy()

 if __name__ == "__main__":
-  if getenv("RUN_PICKLE"):
-    with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f)
-    inputs = {name: Tensor(Tensor.randn(*view.shape, dtype=dtype).numpy(), device=device)
-              for name, (view, _vars, dtype, device) in zip(pickle_loaded.captured.expected_names, pickle_loaded.captured.expected_input_info)}
-    test_vs_compile(pickle_loaded, inputs)
-  else:
-    onnx_file = fetch(OPENPILOT_MODEL)
-    pickle_loaded, inputs, outputs = compile(onnx_file)
+  onnx_file = fetch(OPENPILOT_MODEL)
+  inputs, outputs = compile(onnx_file)

-    if OUTPUT != os.devnull:
-      with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f)
+  with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f)

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

  if getenv("BENCHMARK_LOG", ""):
    bench(pickle_loaded, inputs)
@@ -5,7 +5,7 @@
 #  - symbolic removal

 from examples.beautiful_mnist import Model
-from tinygrad import Tensor, nn, getenv, GlobalCounters, Variable, Context
+from tinygrad import Tensor, nn, getenv, GlobalCounters, Variable
 from tinygrad.nn.datasets import mnist
 from tinygrad.helpers import trange

@@ -26,7 +26,7 @@ if __name__ == "__main__":
  X_samp, Y_samp = X_train[samples], Y_train[samples]
  print("*** got samples")

-  with Context(TRAINING=1):
+  with Tensor.train():
    """
    i = UOp.range(samples.shape[0])  # TODO: fix range function on UOp
    losses = model(X_samp[i]).sparse_categorical_crossentropy(Y_samp[i]).backward().contract(i)
@@ -1,6 +1,7 @@
 from tinygrad import Tensor, Device, TinyJit, dtypes
+from tinygrad.helpers import getenv

-GPUS = Device[Device.DEFAULT].count()
+GPUS = getenv("GPUS", 4) # TODO: expose a way in tinygrad to access this
 N = 6144

@TinyJit
@@ -164,8 +164,8 @@ elif cmd == "train":
    x_img = image_load(samples_base + "/" + str(sample_idx) + "a.png")
    y_img = image_load(samples_base + "/" + str(sample_idx) + "b.png")

-    sample_x = Tensor(x_img)
-    sample_y = Tensor(y_img)
+    sample_x = Tensor(x_img, requires_grad = False)
+    sample_y = Tensor(y_img, requires_grad = False)

    # magic code roughly from readme example
    # An explanation, in case anyone else has to go down this path:
@@ -111,19 +111,19 @@ if __name__ == "__main__":
    return code

  def compile_step(model, step: Step):
-    linear, output_bufs = jit_model(step, *step.input)
-    functions, statements, bufs, _ = compile_net(linear, output_bufs)
+    run, special_names = jit_model(step, *step.input)
+    functions, statements, bufs, _ = compile_net(run, special_names)
    state = get_state_dict(model)
-    weights = {(id(b), b.offset, b.size, b.dtype): name for name, x in state.items() if (b:=x.uop.base.realized) is not None}
+    weights = {id(x.uop.base.realized): name for name, x in state.items()}
    kernel_code = '\n\n'.join([f"const {key} = `{fixup_code(code, key)}`;" for key, code in functions.items()])
    kernel_names = ', '.join([name for (name, _, _, _) in statements])
-    input_names = [f"input{i}" for i in range(len(step.input))]
-    output_names = [f"output{i}" for i in range(len(output_bufs))]
+    input_names = [name for _,name in special_names.items() if "input" in name]
+    output_names = [name for _,name in special_names.items() if "output" in name]
    input_buf_types = [dtype_to_js_type(bufs[inp_name][1]) for inp_name in input_names]
    output_buf_types = [dtype_to_js_type(bufs[out_name][1]) for out_name in output_names]
    kernel_calls = '\n        '.join([f"addComputePass(device, commandEncoder, piplines[{i}], [{', '.join(args)}], {global_size});" for i, (_name, args, global_size, _local_size) in enumerate(statements) ])
    exported_bufs =  '\n    '.join([f"const {name} = " + (f"createEmptyBuf(device, {size});" if _key not in weights else f"createWeightBuf(device, {size}, getTensorBuffer(safetensor, metadata['{weights[_key]}'], '{weights[_key]}'))") + ";"  for name,(size,dtype,_key) in bufs.items()])
-    gpu_write_bufs =  '\n    '.join([f"const gpuWriteBuffer{i} = device.createBuffer({{size:input{i}.size, usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.MAP_WRITE }});" for i in range(len(input_names))])
+    gpu_write_bufs =  '\n    '.join([f"const gpuWriteBuffer{i} = device.createBuffer({{size:input{i}.size, usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.MAP_WRITE }});" for i,(_,value) in enumerate(special_names.items()) if "output" not in value])
    input_writer = '\n    '.join([f"await gpuWriteBuffer{i}.mapAsync(GPUMapMode.WRITE);\n    new {input_buf_types[i]}(gpuWriteBuffer{i}.getMappedRange()).set(" + f'data{i});' + f"\n    gpuWriteBuffer{i}.unmap();\ncommandEncoder.copyBufferToBuffer(gpuWriteBuffer{i}, 0, input{i}, 0, gpuWriteBuffer{i}.size);"  for i,_ in enumerate(input_names)])
    return f"""\n    var {step.name} = function() {{

@@ -141,7 +141,7 @@ if __name__ == "__main__":
        const kernels = [{kernel_names}];
        const piplines = await Promise.all(kernels.map(name => device.createComputePipelineAsync({{layout: "auto", compute: {{ module: device.createShaderModule({{ code: name }}), entryPoint: "main" }}}})));

-        return async ({",".join([f'data{i}' for i in range(len(input_names))])}) => {{
+        return async ({",".join([f'data{i}' for i,(k,v) in enumerate(special_names.items()) if v != "output0"])}) => {{
            const commandEncoder = device.createCommandEncoder();

            {input_writer}
@@ -193,8 +193,8 @@ class SPPF:
    self.cv1 = Conv_Block(c1, c_, 1, 1, padding=None)
    self.cv2 = Conv_Block(c_ * 4, c2, 1, 1, padding=None)

-    # Pad with -inf to match PyTorch's MaxPool2d behavior.
-    self.maxpool = lambda x : x.pad((k // 2, k // 2, k // 2, k // 2), value=float('-inf')).max_pool2d(kernel_size=k, stride=1)
+    # TODO: this pads with 0s, whereas torch function pads with -infinity. This results in a < 2% difference in prediction which does not make a difference visually.
+    self.maxpool = lambda x : x.pad((k // 2, k // 2, k // 2, k // 2)).max_pool2d(kernel_size=k, stride=1)

  def __call__(self, x):
    x = self.cv1(x)
@@ -1,14 +1,14 @@
 #!/usr/bin/env python3

 import time, mmap, sys, shutil, os, glob, subprocess, argparse, collections
-from tinygrad.helpers import DEBUG, NO_COLOR, colored, ansilen
+from tinygrad.helpers import DEBUG, colored, ansilen
 from tinygrad.runtime.autogen import libc
 from tinygrad.runtime.autogen.am import am
 from tinygrad.runtime.support.hcq import MMIOInterface
 from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager, AMPageTableEntry
 from tinygrad.runtime.support.am.ip import AM_SOC, AM_GMC, AM_IH, AM_PSP, AM_SMU, AM_GFX, AM_SDMA

-def bold(s): return s if NO_COLOR else f"\033[1m{s}\033[0m"
+def bold(s): return f"\033[1m{s}\033[0m"

 def trim(s:str, length:int) -> str:
  if len(s) > length: return s[:length-3] + "..."
@@ -64,7 +64,7 @@ def get_bar0_size(pcibus):

 class AMSMI(AMDev):
  def __init__(self, pcibus, vram_bar:MMIOInterface, doorbell_bar:MMIOInterface, mmio_bar:MMIOInterface):
-    self.pcibus, self.devfmt = pcibus, pcibus
+    self.pcibus = pcibus
    self.vram, self.doorbell64, self.mmio = vram_bar, doorbell_bar, mmio_bar
    self.pci_state = self.read_pci_state()
    if self.pci_state == "D0": self._init_from_d0()
@@ -91,7 +91,6 @@ class SMICtx:
    self.prev_lines_cnt = 0
    self.prev_terminal_width = 0
    self.prev_terminal_height = 0
-    self.prev_metrics = {}

    remove_parts = ["Advanced Micro Devices, Inc. [AMD/ATI]", "VGA compatible controller:", "Processing accelerators:"]
    lspci = subprocess.check_output(["lspci"]).decode("utf-8").splitlines()
@@ -236,29 +235,6 @@ class SMICtx:
      case (13,0,12): return self._smuq10_round(metrics.SocketPower), self._smuq10_round(metrics.SocketPowerLimit)
      case _: return metrics.SmuMetrics.AverageSocketPower, metrics.SmuMetrics.dGPU_W_MAX

-  def get_throttle_info(self, dev, metrics):
-    match dev.ip_ver[am.MP1_HWIP]:
-      case (13,0,6)|(13,0,12):
-        throttle_fields = [('ProchotResidencyAcc', 'Prochot'), ('PptResidencyAcc', 'PPT'),
-                           ('SocketThmResidencyAcc', 'Socket Thm'), ('VrThmResidencyAcc', 'VR Thm'), ('HbmThmResidencyAcc', 'HBM Thm')]
-        prev = self.prev_metrics.get(dev.pcibus)
-        active = []
-        if prev is not None:
-          acc_delta = metrics.AccumulationCounter - prev.AccumulationCounter
-          if acc_delta > 0:
-            for field, name in throttle_fields:
-              delta = getattr(metrics, field) - getattr(prev, field)
-              if delta > 0 and (pct := min(100, (delta * 100 + acc_delta // 2) // acc_delta)) > 0: active.append((name, pct))
-        return active
-      case _:
-        smu_mod = dev.smu.smu_mod
-        throttler_names = {getattr(smu_mod, a): a[len('THROTTLER_'):-len('_BIT')]
-                           for a in dir(smu_mod) if a.startswith('THROTTLER_') and a.endswith('_BIT')}
-        active = []
-        for i, pct in enumerate(metrics.SmuMetrics.ThrottlingPercentage):
-          if pct > 0: active.append((throttler_names.get(i, f"UNK_{i}"), int(pct)))
-        return active
-
  def get_mem_usage(self, dev):
    usage = 0
    pt_stack = [dev.mm.root_page_table]
@@ -276,7 +252,7 @@ class SMICtx:
    return usage

  def draw(self, once):
-    terminal_width, terminal_height = shutil.get_terminal_size(fallback=(231, 24))
+    terminal_width, terminal_height = shutil.get_terminal_size()
    if not once and (self.prev_terminal_width != terminal_width or self.prev_terminal_height != terminal_height):
      os.system('clear')
    self.prev_terminal_width, self.prev_terminal_height = terminal_width, terminal_height
@@ -305,13 +281,6 @@ class SMICtx:
                    + [f"MEM Activity {draw_bar(self.get_mem_activity(dev, metrics) / 100, activity_line_width)}"] \
                    + [f"MEM Usage    {draw_bar(mem_used / mem_total, activity_line_width, opt_text=mem_fmt)}"] \

-      throttle_info = self.get_throttle_info(dev, metrics)
-      if throttle_info:
-        throttle_text = colored(', '.join(f"{name} {pct}%" for name, pct in throttle_info), "red")
-      else:
-        throttle_text = colored("None", "green")
-      activity_line += [f"Throttle     {throttle_text}" + " " * (activity_line_width + 2)]
-
      temps_data, temps_data_compact = self.get_temps(dev, metrics), self.get_temps(dev, metrics, compact=True)
      temps_table = ["=== Temps (°C) ==="] + [f"{name:<16}: {color_temp(val)}" for name, val in temps_data.items()]
      temps_table_compact = ["Temps (°C):" + '/'.join([f"{color_temp(val)} {name}" for name, val in temps_data_compact.items()])]
@@ -355,8 +324,6 @@ class SMICtx:

      dev_content.append(device_line + activity_line + same_line([temps_table, power_table, frequency_table]))

-    self.prev_metrics = {dev.pcibus: m for dev, m in dev_metrics.items() if m is not None}
-
    raw_text = 'AM Monitor'.center(terminal_width) + "\n" + "=" * terminal_width + "\n\n"
    for i in range(0, len(dev_content), 2):
      if i + 1 < len(dev_content): raw_text += '\n'.join(same_line([dev_content[i], dev_content[i+1]], split=padding))
@@ -28,7 +28,15 @@
 // #include "soc15_ih_clientid.h"
 // #include "amdgpu_ih.h"

+#define int32_t int
+#define uint32_t unsigned int
+#define int8_t signed char
+#define uint8_t unsigned char
+#define uint16_t unsigned short
+#define int16_t short
+#define uint64_t unsigned long long
 #define bool _Bool
+#define u32 unsigned int

 #define AMDGPU_MAX_IRQ_SRC_ID		0x100
 #define AMDGPU_MAX_IRQ_CLIENT_ID	0x100
@@ -22,7 +22,15 @@
 #ifndef __AMDGPU_SMU_H__
 #define __AMDGPU_SMU_H__

+#define int32_t int
+#define uint32_t unsigned int
+#define int8_t signed char
+#define uint8_t unsigned char
+#define uint16_t unsigned short
+#define int16_t short
+#define uint64_t unsigned long long
 #define bool _Bool
+#define u32 unsigned int

 #define SMU_THERMAL_MINIMUM_ALERT_TEMP		0
 #define SMU_THERMAL_MAXIMUM_ALERT_TEMP		255
@@ -24,7 +24,15 @@
 #define __AMDGPU_UCODE_H__

 // #include "amdgpu_socbb.h"
+#define int32_t int
+#define uint32_t unsigned int
+#define int8_t signed char
+#define uint8_t unsigned char
+#define uint16_t unsigned short
+#define int16_t short
+#define uint64_t unsigned long long
 #define bool _Bool
+#define u32 unsigned int

 struct common_firmware_header {
 	uint32_t size_bytes; /* size of the entire header+image(s) in bytes */
@@ -1,50 +1,47 @@
 from typing import Tuple, Dict, List, Optional
-from tinygrad.dtype import DType, dtypes, AddrSpace
+from tinygrad.dtype import DType, dtypes
+from tinygrad.renderer import ProgramSpec
 from tinygrad.tensor import Tensor
-from tinygrad.device import Device, Buffer
+from tinygrad.device import Device
 from tinygrad.engine.jit import TinyJit
 from tinygrad.nn.state import get_state_dict
-from tinygrad.helpers import Context, to_mv, prod
-from tinygrad.uop.ops import Ops, UOp
-from tinygrad.codegen import to_program
+from tinygrad.helpers import Context, to_mv
+from tinygrad.uop.ops import Ops
 import json
 from collections import OrderedDict

 EXPORT_SUPPORTED_DEVICE = ["WEBGPU", "CPU", "CUDA", "CL"]

-_KERNEL_ASTS = {Ops.SINK, Ops.PROGRAM}
-def iter_kernel_calls(linear:UOp):
-  """Yield kernel CALLs from a LINEAR UOp. Toposort descends naturally into CUSTOM_FUNCTION graph batches; gate stops at kernel ASTs."""
-  return (u for u in linear.toposort(gate=lambda x: x.op not in _KERNEL_ASTS) if u.op is Ops.CALL and u.src[0].op in _KERNEL_ASTS)
+def compile_net(run:TinyJit, special_names:Dict[int,str]) -> Tuple[Dict[str,str],List[Tuple[str,List[str],List[int]]],Dict[str,Tuple[int,DType,int]],Dict[str,Tensor]]:
+  # memory-planned subbuffers can have multiple Buffer objects for the same memory region
+  canon, _seen = {}, {}
+  for ji in run.jit_cache:
+    for b in ji.bufs:
+      if b is not None: canon[id(b)] = _seen.setdefault((id(b.base._buf), b.offset, b.size, b.dtype), b)
+  special_names = {id(canon[k]): v for k, v in special_names.items() if k in canon}

-def compile_net(linear:UOp, output_bufs:List[Buffer]) -> Tuple[Dict[str,str], List, Dict[str,Tuple[int,DType,int]], Dict[str,Buffer]]:
-  output_name = {id(b): f"output{i}" for i, b in enumerate(output_bufs)}
-  functions, bufs, bufs_to_save, statements, n = {}, {}, {}, [], 0
+  functions, bufs, bufs_to_save, statements, bufnum = {}, {}, {}, [], 0
+  for ji in run.jit_cache:
+    fxn: ProgramSpec = ji.prg.p
+    functions[fxn.function_name] = fxn.src   # NOTE: this assumes all with the same name are the same
+    cargs = []
+    for i,arg in enumerate(ji.bufs):
+      arg = canon[id(arg)]
+      key = id(arg)
+      if key not in bufs:
+        if key in special_names:
+          bufs[key] = (special_names[key], arg.size*arg.dtype.itemsize, arg.dtype, key)
+        else:
+          bufs[key] = (f"buf_{bufnum}", arg.size*arg.dtype.itemsize, arg.dtype, key)
+          bufnum += 1
+          if i > 0: bufs_to_save[bufs[key][0]] = arg   # if first usage of a buffer is not an output, and it's not a special name
+      cargs.append(bufs[key][0])
+    cargs += [var for var in fxn.vars if getattr(var, "op", None) is Ops.DEFINE_VAR] # symbolic vars; is it necessary or sufficient to check for DEFINE_VAR?
+    statements.append((fxn.function_name, cargs, fxn.global_size, fxn.local_size))

-  def name_of(bu:UOp, is_out:bool) -> str:
-    nonlocal n
-    if bu.op is Ops.PARAM: key, name, size = ("in", bu.arg.slot), f"input{bu.arg.slot}", prod(bu.shape)*bu.dtype.itemsize
-    else:
-      b = bu.buffer
-      key, size = (id(b.base), b.offset, b.size, b.dtype), b.size*b.dtype.itemsize
-      if key in bufs: return bufs[key][0]
-      if (name:=output_name.get(id(b))) is None:
-        name, n = f"buf_{n}", n+1
-        if not is_out: bufs_to_save[name] = b
-    bufs[key] = (name, size, bu.dtype, key)
-    return name
+  return functions, statements, {name:(size, dtype, key) for (name,size,dtype,key) in bufs.values()}, bufs_to_save

-  for call in iter_kernel_calls(linear):
-    arg_uops = [b for b in call.src[1:] if b.op is not Ops.BIND]
-    prg = to_program(call.src[0], Device[arg_uops[0].device].renderer)
-    info = prg.arg
-    functions[info.function_name] = prg.src[2].arg
-    cargs = [name_of(bu, i == 0) for i, bu in enumerate(arg_uops)] + list(info.vars)
-    statements.append((info.function_name, cargs, info.global_size, info.local_size))
-
-  return functions, statements, {name:(size, dtype, key) for name, size, dtype, key in bufs.values()}, bufs_to_save
-
-def jit_model(model, *args) -> Tuple[UOp, List[Buffer]]:
+def jit_model(model, *args) -> Tuple[TinyJit,Dict[int,str]]:
  assert hasattr(model, "forward") or callable(model), "model needs a forward function"
  @TinyJit
  def run(*x):
@@ -53,10 +50,20 @@ def jit_model(model, *args) -> Tuple[UOp, List[Buffer]]:
    out = [out] if isinstance(out, Tensor) else out
    return [o.realize() for o in out]

-  # run twice to trigger JIT capture
+  # twice to run the JIT
  for _ in range(2): the_output = run(*args)
-  assert run.captured is not None
-  return run.captured.linear, [o.uop.base.realized for o in the_output]
+  special_names = {}
+
+  # hack to put the inputs back
+  for (j,i),idx in run.input_replace.items():
+    realized_input = args[idx].uop.base.realized
+    run.jit_cache[j].bufs[i] = realized_input
+    special_names[id(realized_input)] = f'input{idx}'
+
+  # TODO: fetch this from the jit in self.input_replace and self.ret (hint: use get_parameters on self.ret)
+  for i, output in enumerate(the_output):
+    special_names[id(output.uop.base.realized)] = f'output{i}'
+  return run, special_names

 def export_model_clang(functions:Dict[str,str], statements:Dict[str,Tuple[str,int,int]], bufs:Dict[str,Tuple[str,int,int]],
  bufs_to_save:Dict[str,Tensor], input_names:List[str], output_names:List[str], weight_names={}, model_name="model", symbolic_vars={}, wasm=False) -> str:
@@ -242,29 +249,28 @@ def export_model(model, target:str, *inputs, model_name: Optional[str] = "model"
  assert Device.DEFAULT in EXPORT_SUPPORTED_DEVICE, f"only {', '.join(EXPORT_SUPPORTED_DEVICE)} are supported"

  # NOTE: CPU_COUNT=1, since export does not support threading
-  with Context(JIT=2, CPU_COUNT=1): linear, output_bufs = jit_model(model, *inputs)
-  functions, statements, bufs, bufs_to_save = compile_net(linear, output_bufs)
+  with Context(JIT=2, CPU_COUNT=1): run,special_names = jit_model(model, *inputs)
+  functions, statements, bufs, bufs_to_save = compile_net(run, special_names)
  state = get_state_dict(model)
-  weight_names = {(id(b), b.offset, b.size, b.dtype): name for name, x in state.items() if (b:=x.uop.base.realized) is not None}
-  input_names = [f"input{i}" for i in range(len(inputs))]
-  output_names = [f"output{i}" for i in range(len(output_bufs))]
+  weight_names = {id(x.uop.base.realized): name for name, x in state.items()}
+  input_names = [name for _,name in special_names.items() if "input" in name]
+  output_names = [name for _,name in special_names.items() if "output" in name]

  # handle symbolic variables; TODO: refactor to fix some of this stuff upstream in tinygrad
  symbolic_vars = OrderedDict()
  for i, (_, args, global_size, _) in enumerate(statements):
    for j, var in enumerate(args):
-      if getattr(var, "op", None) is Ops.PARAM and var.addrspace is AddrSpace.ALU and var.arg.name is not None:
+      if getattr(var, "op", None) is Ops.DEFINE_VAR and isinstance(getattr(var, "arg", None), tuple) and isinstance(var.arg[0], str):
        if var not in symbolic_vars:
-          symbolic_vars[var] = var.expr
+          symbolic_vars[var] = var.arg[0]
          bufs[symbolic_vars[var]] = (var.dtype.itemsize, var.dtype, symbolic_vars[var])
        statements[i][1][j] = symbolic_vars[var]

    if global_size:
      for j, dim in enumerate(global_size):
-        if getattr(dim, "op", None) is Ops.ADD and len(dim.src) == 2 and \
-           any(s.op is Ops.PARAM and s.addrspace is AddrSpace.ALU for s in dim.src) and any(s.op is Ops.CONST for s in dim.src):
+        if getattr(dim, "op", None) is Ops.ADD and len(dim.src) == 2 and {dim.src[0].op, dim.src[1].op} == {Ops.DEFINE_VAR, Ops.CONST}:
          name, val = dim.src if dim.src[1].op is Ops.CONST else reversed(dim.src)
-          global_size[j] = f"_{name.expr}[0] + {val.arg}"
+          global_size[j] = f"_{name.arg[0]}[0] + {val.arg}"

  prg = ""
  if target == "clang":
@@ -24,7 +24,7 @@ def custom_matmul(output: UOp, inp: UOp, weight: UOp) -> UOp:
  reduce_idx = UOp.range(IN, 0, AxisType.REDUCE)
  product = (inp.index((seq_idx*IN+reduce_idx+batch_idx*IN*SEQ)) * weight.index((out_idx*IN+reduce_idx))).cast(dtypes.float)
  reduced = product.reduce(reduce_idx, arg=Ops.ADD)
-  store_op = output.index((seq_idx*OUT+out_idx+batch_idx*OUT*SEQ)).store(reduced).end(batch_idx, seq_idx, out_idx)
+  store_op = output.index((seq_idx*OUT+out_idx+batch_idx*OUT*SEQ), ptr=True).store(reduced).end(batch_idx, seq_idx, out_idx)
  return store_op.sink(arg=KernelInfo(name=f"fp8_matmul_{inp.shape}x{weight.shape}"))

 def custom_matmul_backward(gradient: UOp, kernel: UOp) -> tuple[UOp, UOp]:
@@ -1,973 +0,0 @@
-# Adreno 630 (Snapdragon 845) FP16 GEMM Optimization
-
-## Device Access
-
-```bash
-ssh tc3
-cd /data/openpilot/tinygrad_repo
-pkill -9 python3   # recover from GPU hangs (no reboot needed)
-```
-
-## Running the benchmarks
-
-```bash
-# Patched compiled kernel (~190 GFLOPS)
-PYTHONPATH=. DEV=QCOM python3 extra/gemm/qcom_gemm.py
-PYTHONPATH=. DEV=QCOM python3 extra/gemm/qcom_gemm.py --m 512 --n 512 --k 512
-
-# Hand-assembled kernel tests (pure ALU, pure load, patched GEMM)
-PYTHONPATH=. DEV=QCOM python3 extra/gemm/qcom_asm_gemm.py
-
-# Subgroup/quad broadcast probes
-PYTHONPATH=. DEV=QCOM python3 extra/gemm/qcom_shfl_probe.py
-PYTHONPATH=. DEV=QCOM python3 extra/gemm/qcom_shfl_probe.py --bench throughput --ops-per-iter 16
-PYTHONPATH=. DEV=QCOM python3 extra/gemm/qcom_shfl_probe.py --op quad --bench throughput --ops-per-iter 16
-
-# Direct texture/isam bandwidth sweep
-PYTHONPATH=. DEV=QCOM python3 extra/gemm/qcom_texture_bw.py --threads 128 --loads 32
-```
-
-## Current Findings: THREAD128 Runtime
-
-The QCOM runtime used to hardcode `mesa.THREAD64` in compute dispatch state. Adding
-`THREAD128=1` to `tinygrad/runtime/ops_qcom.py` selects `mesa.THREAD128` for:
-
- `A6XX_SP_CS_WGE_CNTL`
- `A6XX_SP_CS_CNTL_0`
- the NIR `A6XX_SP_CS_WGE_CNTL` path
-
-This matches OpenCL's FP16 MAD peak on A630:
-
-| Command | Result |
-|---------|--------|
-| `PYTHONPATH=. DEV=QCOM python3 extra/mmapeak/qcom_fp16_mad_peak.py` | `345.64 GFLOPS` |
-| `PYTHONPATH=. DEV=QCOM THREAD128=1 python3 extra/mmapeak/qcom_fp16_mad_peak.py` | `690.35 GFLOPS` |
-| `PYTHONPATH=. DEV=CL python3 extra/mmapeak/qcom_fp16_mad_peak.py` | `690.76 GFLOPS` |
-
-For hand GEMM kernels, use `THREAD128=1` for all new measurements.
-
-### ALU-Only GEMM-Shape Measurements
-
-Measured on `tc3` with `THREAD128=1`, scalar `8x8` GEMM shape:
-
-| Kernel/profile | Registers | Result | Notes |
-|----------------|-----------|--------|-------|
-| Compiler vector16 `mmapeak` | compiler | `~690 GFLOPS` | Not GEMM-shaped; vector-vector MAD stream |
-| Hand compiler-pattern ALU stream | `f9 h8` | `714-718 GFLOPS` | Mirrors OpenCL vec16 lowering; `x=mad(x,y,y)`, `y=mad(x,y,x)` |
-| True GEMM ALU body, `4x12`, distinct B, `row_col_kk` | `f8 h28` | `676.1 GFLOPS` | `acc=A_scalar*B_half4+acc`, one-shot unrolled body |
-| True GEMM ALU body, `4x8`, distinct B, `row_col_kk` | `f8 h24` | `662.3 GFLOPS` | `acc=A_scalar*B_half4+acc`, four wave-pairs |
-| True GEMM ALU body, `4x16`, reused B, `row_col_kk` | `f8 h32` | `679.3 GFLOPS` | Valid FMA form, but B columns are reused for ALU stress |
-| Generic hand ALU stream, bad source pattern | `f8 h48` | `~357 GFLOPS` | Repeatedly reads same `hr0.x/hr4.x` |
-| Generic hand ALU stream with source1 relative `(r)` | `f8 h32` | `~519 GFLOPS` | Best at 3-4 wave-pair occupancy |
-| Correct high-reg `8x8 --profile alu` | `f28 h32` | `454.6 GFLOPS` | GEMM scalar-broadcast schedule |
-| Low-reg `8x8 --experimental-twopass --profile alu` | `f15 h32` | `452.4 GFLOPS` | Donor/two-pass profile remains occupancy-limited |
-| Low-reg `8x8 serial --profile alu` | `f8 h32` | `681.7 GFLOPS` | Four wave-pair ALU profile; not a correct full GEMM path yet |
-| Serial `8x16 --profile alu` | `f8 h48` | `467.5 GFLOPS` | More accumulators, but lower occupancy |
-
-Takeaways:
-
- Raw hand ALU can exceed `600 GFLOPS` when it uses the compiler vec16 source pattern and a low register footprint: `qcom_alu_peak.py --compiler-pattern --pairs 8 --loops 64` measured `714.0 GFLOPS`.
- The >600 pattern is not the GEMM accumulation form. It writes `dst=src1` and uses the other vector as addend, while GEMM needs `dst += A*B` (`dst=src3`).
- A true scalar-broadcast GEMM FMA body can also exceed `600 GFLOPS` if scheduled as `row_col_kk` and measured as a one-shot unrolled body: `qcom_alu_peak.py --gemm-pattern --rows 4 --ncols 3 --bmode percol --order row_col_kk --unroll 16 --loops 1` measured `676.1 GFLOPS`.
- The `row_col_kk` ordering is the key ALU finding: consume all four K components for one output vector accumulator before moving to the next accumulator.
- Repeating the synthetic GEMM ALU body in a loop is not a valid source-preserving benchmark unless A/B sources are reloaded or loop-control registers are kept out of their half-register aliases; use `--loops 1` for `--gemm-pattern`.
- Arithmetic intensity is not the current ALU issue limit.
- The old high-reg/donor-style `8x8` GEMM ALU profiles are capped around `452-455 GFLOPS`, but the low-freg serial profile reaches `681.7 GFLOPS`; the `8x8` ALU body is not inherently capped.
- MAD instruction order and source1-relative encoding did not materially improve the donor-style `8x8` profiles.
- Occupancy/register footprint, texture-sync placement, and a correct low-reg store path matter more than the specific legal MAD order.
-
-### Current Correct GEMM Results
-
-All entries below are full-output all-ones checked unless noted otherwise.
-
-| Kernel | THREAD128 | Result | Notes |
-|--------|-----------|--------|-------|
-| Correct scalar `8x4` donor-store | yes | `255.9 GFLOPS` | `f12 h24`, texture-roof limited by AI 2.67 |
-| Correct high-reg scalar `8x8` donor-store | yes | `196.8 GFLOPS` | `f28 h32`, store/loop not improved by THREAD128 |
-| Low-reg scalar `8x8` two-pass store | yes | `188.8 GFLOPS` | Now correctness-stable under THREAD128 but slower |
-| Low-reg scalar `8x8` serial + donor8 store | yes | `189.1-191.1 GFLOPS` | Correct; `f12 h32`, proves low-reg serial compute is valid when store is fixed |
-| Low-reg scalar `8x8` split-A + add256 donor store | yes | `360.2-378.6 GFLOPS` | Correct; `f10 h28`, four wave-pairs, pre-unroll baseline |
-| Low-reg scalar `8x8` split-A + K-unroll 4 + add256 donor store | yes | `425.8-436.0 GFLOPS` | Correct; `f10 h28`, four wave-pairs, previous best 8x8 path |
-| Low-reg scalar `8x8` split-A + K-unroll 8 + next-B prefetch + tight add256 store | yes | `467.9-468.8 GFLOPS` | Correct; `f8/f9 h28`, four wave-pairs, first verified >460 path |
-| Low-reg scalar `8x8` pipelined A/B | yes | `287.2 GFLOPS` | Correct; double-buffered inputs, `f15 h48` |
-| Low-reg scalar `8x8` pipelined A/B, no next-buffer sync | yes | `288.4 GFLOPS` | Correct; `--b-coord-delay -1 --no-next-sy`, current-buffer sync still required |
-| Low-reg scalar `8x8` pipeline4 | yes | `287.9 GFLOPS` | Correct; 4x K4 unroll needs larger donor envelope, does not improve throughput |
-| Low-reg scalar `8x8` batch2 | yes | `222.0 GFLOPS` | Correct but slower; loading two K steps then computing loses overlap |
-| Pipelined scalar `8x4` | yes | `200.7 GFLOPS` | Correct with `--a-coord-delay 0`; lower AI plus extra buffering is slower than baseline `8x4` |
-| Direct `4x8` low-reg donor-store | yes | `271.5 GFLOPS` | Correct; repeated `4x4` compiler donor store, `--coord-delay 0` or `-1` |
-| Direct `4x16` native-store | yes | `184.2 GFLOPS` | Correct; native `4x16` compiler store fixes coverage but needs high full-register footprint |
-| Direct `4x16` low-reg donor-store | yes | `331.4 GFLOPS` | Correct; stride dependency waits fixed full coverage, `f8 h32`, `--coord-delay 4` |
-| Direct `4x16` compact-acc hand ASM store | yes | `~334-336 GFLOPS` | Correct; accumulators start at `hr12`, `f8 h28`, `--k-unroll 4`, no runtime donor-store slicing |
-| Direct `4x16` compact-acc hand ASM store, reduced K-sync | yes | `~382-388 GFLOPS` | Correct; `--stable-bx --k-unroll 4 --first-sync-only`, `f8 h28`, `sy=2` |
-| Direct `4x16` compact-acc hand ASM store, persistent coords | yes | `400.5-402.1 GFLOPS` | Correct; `--stable-bx --stable-ay --inc-coords --persistent-coords --first-sync-only`, `f10 h28`, loop `421 -> 417` |
-| Direct `4x16` persistent coords, B-first schedule | yes | `421.2-434.8 GFLOPS` | Correct; same `f10 h28` and loop size, but loads first B pair before A to hide B texture latency |
-| Direct `4x16` B-first with low A coords | yes | `424.4-429.6 GFLOPS` | Correct; lowers metadata to `f8 h28`, but speed is flat vs `f10 h28` |
-
-The split-A `8x8` K-unroll-8 path with next-B prefetch and tight add256 stores is the fastest correct hand path so far and is the first verified path above 460 GFLOPS. The compact-acc direct low-register `4x16` kernel with reduced per-unroll sync, persistent coordinates, and B-first scheduling remains the fastest correct 4x16 hand path.
-
-#### FP32 Accumulate From FP16 Images
-
-The standalone hand FP32 path in `qcom_8x4_gemm.py` is correctness-stable but not competitive with the compiler-shaped assembly patch. The original scalar `8x4` route reads FP16 images with `isam.f16`, converts with `cov.f16f32`, accumulates with `(rpt3)mad.f32`, and writes a float C buffer with `stg.f32`.
-
-```bash
-PYTHONPATH=. DEV=QCOM IMAGE=1 FLOAT16=1 THREAD128=1 \
-  python3 extra/gemm/qcom_8x4_gemm.py --fp32-accum --variant serial \
-  --ncols 1 --threads 128 --b-coord-delay 5 --check
-```
-
-Current checked result:
-
-```text
-serial:fp32 ncols=1 scalar_tile=8x4 threads=128 fregs=28 hregs=1 reg_count=29 wave_pairs=3 intensity=2.67 flop/B mad_density=1.03 shader_instrs=273 loop_instrs=124 bytes=2184 envelope_bytes=2832
-mad.f16=0 mad.f32=32 rpt3=32 isam=12 sy=14 serial_syncs=all
-CHECK PASS all 1048576 float outputs are 1024.0
-```
-
-Latest direct-load probes added `emit_isam_f32_vec`, `--direct-f32-loads`, `--sampler-per-texture`, and `--fp32-accum --ncols 2`. Correct checked timings were still low: ncols1 conversion path `43.7 GFLOPS`, ncols1 direct `110.5 GFLOPS`, ncols2 direct `146.6 GFLOPS`, and ncols2 direct no-store `145.1 GFLOPS`. Direct `isam.f32` from `imageh` is therefore valid with the sampler-per-texture path, but this full hand-assembled route is too slow for the 250 GFLOPS target.
-
-The lower-register `4x4` FP32 prototype in `qcom_intensity_gemm.py` is now verified with full-output float checks. It must use the direct FP32 donor prologue; the older half donor prologue made B loads miss 1-2 K contributions in row/column-dependent regions even though post-constant stores passed.
-
-```bash
-PYTHONPATH=. DEV=QCOM IMAGE=1 FLOAT16=1 THREAD128=1 \
-  python3 extra/gemm/qcom_intensity_gemm.py --fp32-accum --ncols 1 \
-  --threads 128 --coord-delay 4 --direct-f32-loads \
-  --sampler-per-texture --check
-```
-
-Current checked/timed result:
-
-```text
-ncols=1 covered_N=1024 fregs=20 hregs=1 waves=96 intensity=2.00 flop/B mad_density=1.36 shader_instrs=161 loop_instrs=47 bytes=1288 envelope_bytes=2792
-mad.f16=0 mad.f32=16 rpt3=16 isam=8 qbc=0 sy=2
-CHECK PASS all 1048576 float outputs are 1024.0
-best observed timing: 154.8 GFLOPS (13.872 ms)
-```
-
-Direct `isam.f32` from `imageh` is correct in this direct-prologue `4x4` path. With sampler 0 for both textures it reached `137.7 GFLOPS`; using sampler index equal to texture index reached `151.6-154.8 GFLOPS`. Probe timings for the faster direct-load shape: no-store `150.8 GFLOPS`, skip A loads `171.4 GFLOPS`, skip B loads `233.9 GFLOPS`, skip A+B loads `292.8 GFLOPS`. The scalar-MAD variant (`64` scalar `mad.f32`, no `rpt3`) is correct but slower at `104.6 GFLOPS`.
-
-THREAD128 compact-register `4x4` FP32 probes in `qcom_intensity_gemm.py` are correct but not a 300 route. `--compact-fp32` streams one A vector at a time and lowers metadata to `f12`; it passes full-output float checks with the donor float-store epilogue but only measured `120.5 GFLOPS` full and `114.0 GFLOPS` no-store. `--compact-fp32-preload` preloads A/B into `r0-r7` and keeps state in `r12`; a short wait is required before the donor store when copying state back to `r7`, and the checked full kernel measured `217.2 GFLOPS` at `f13`. `--compact-fp32-hybrid` keeps row/col/K state in `r7`, places A3 in `r12`, and is the cleanest low-register variant: full-output checks pass, `--coord-delay 3` is valid and measured `209.9 GFLOPS`, while delays `1` and `2` are invalid (`1020.0` outputs). At `--coord-delay 4`, the hybrid path measured `205.6 GFLOPS` full, `205.4 GFLOPS` no-store, `233.8 GFLOPS` no-store skip-A, `277.3 GFLOPS` no-store skip-B, and `312.9 GFLOPS` no-store skip-A+B. The generic hand `STG_F32` store path produced mostly zero output; the compiler-donor float epilogue is still required for reliable stores. Lowering the full-register footprint alone is therefore insufficient: real A/B texture scheduling remains the limiter.
-
-Low-register `4x8` FP32 A-reuse now works correctly in `qcom_intensity_gemm.py`, and the fastest checked version uses default dispatch rather than `THREAD128=1`. The useful version is `--low-4x8-fp32 --preload-b`, which keeps both B column blocks live, uses `r12-r19` for accumulators, and keeps state in `r20` (`f21`). Reusing the 4-row donor float epilogue twice was invalid because the donor slice carried an `end` and needed store-spacing; the working full-store path uses the compiler's `ncols=2` float donor epilogue with a low-copy repack through dead input registers, avoiding the old `r24-r31` temp copy and preserving `f21`. Best checked command so far:
-
-```bash
-PYTHONUNBUFFERED=1 PYTHONPATH=. DEV=QCOM IMAGE=1 FLOAT16=1 HCQ2=1 \
-  python3 extra/gemm/qcom_intensity_gemm.py --fp32-accum --low-4x8-fp32 \
-    --preload-b --batch-coords --ncols 2 --threads 128 --sampler-per-texture \
-    --coord-delay -1 --alu-order kk_col_row --check
-```
-
-It passes all `1048576` float outputs. A 120-iteration full benchmark measured `232.3 GFLOPS` (`f21`, loop `73`, `12` direct `isam.f32`, `32` `(rpt3)mad.f32`, `sy=2`). The same shape without `HCQ2=1` measured `230.4 GFLOPS`; with `THREAD128=1 HCQ2=1` it only measured `199.5 GFLOPS`, so this hand FP32 path should currently use default dispatch. Correct no-store with the best order is `230.3 GFLOPS`, skip-A is `247.5 GFLOPS`, skip-B is `286.9 GFLOPS`, and skip-A+B is `309.2 GFLOPS`, showing B texture latency is still the primary limiter and the FP32 MAD body itself is only slightly above 300 in this schedule.
-
-Negative `4x8` FP32 follow-ups: the original `f17` non-preload path is correct only as a diagnostic and remains slow (`~120 GFLOPS` no-store under default dispatch, `117.1` under `THREAD128=1`). Half-image `isam.f16` plus explicit `cov.f16f32` collapses to `56.5 GFLOPS` no-store, so direct `isam.f32` is still the right input path. `--stream-b` without a sync reaches `241.0 GFLOPS` no-store under default dispatch but fails full checks; adding the required sync makes it correct but only `202.2 GFLOPS`. Fixed NOP waits before consuming streamed B1 do not fix correctness. Double-buffered software pipeline variants are slower (`f30` B-only pipeline `~151 GFLOPS`, `f34` A+B pipeline `~162 GFLOPS` no-store), so the extra live registers cost more than the overlap buys. Underdeclaring the working `f21` kernel as `f20` hangs, so the metadata cannot be lowered. The explicit hand `STG_F32` path remains mostly zero/sparse output. The remaining limiter is real B texture scheduling, not store correctness.
-
-Wider hand FP32 attempts in `qcom_intensity_gemm.py` are still not promising. The `--fp32-accum --ncols 2` path now has a correct compiler-donor `ncols=2` float-store epilogue and passes full-output checks, but the real conversion-load path is only `30.3 GFLOPS` under `THREAD128=1` (`fregs=32`, `loop_instrs=124`). No-store is still only `28-29 GFLOPS`; skip-A, skip-B, and skip-A+B no-store probes measured `37.8`, `112.0`, and `235.5 GFLOPS`, respectively. Direct `isam.f32` loads raise the ncols2 no-store probe to about `111 GFLOPS`, but full-output checks remain unstable/incorrect for ncols2, so those timings are diagnostics only. The full hand 4x4 direct path did pass a coordinate-delay sweep, with the best historical run around `153.1 GFLOPS` at `--coord-delay 1`, but that remains far below the compiler-shaped assembly patch.
-
-The compiler-generated `simple_matmul.py` path with `DEV=QCOM:IR3 DEBUG=2 IMAGE=1 FLOAT16=1 N=1024 HALF=1` reaches about `196-199 GFLOPS` in the main `r_32_16_8_16_4_4_256_4` kernel. Disassembly shows a `4x4` FP32 accumulator tile with `max_reg=12`, `64` scalar `mad.f32`, `8` direct `isam.f32`, `1` `(sy)`, and typed image-float stores. That is the current practical compiler baseline for FP32 accumulate from FP16 images.
-
-The best verified compiler-side FP32 patch is now `qcom_ir3_matmul_patch.py --n 704 --patch rpt3_l25_postinc_unroll22`. It keeps tinygrad's normal packed image layout, rewrites the compiler's `l25` loop into `(rpt3)mad.f32` accumulator groups, increments the K loop counter after the texture loads, and compares only once per 22-way unrolled group. The first l25 rewrite missed the original `end` instruction and hung; the fixed epilogue includes `instrs[119:134]`.
-
-```bash
-PYTHONUNBUFFERED=1 PYTHONPATH=. DEV=QCOM:IR3 IMAGE=1 FLOAT16=1 HCQ2=1 \
-  python3 extra/gemm/qcom_ir3_matmul_patch.py --n 704 --dtype half \
-  --acc-dtype none --patch rpt3_l25_postinc_unroll22 --check --bench --iters 40
-```
-
-Verified result on `tc3`, `HCQ2=1` with default THREAD64 dispatch:
-
-```text
-main=r_22_11_8_16_4_4_176_4 image_bytes=7208 instrs=901 fregs=16 hregs=0
-mad.f32=352 rpt_mad=352 isam=176 stores=4
-CHECK PASS all 495616 outputs are 704.0
-BENCH main 269.9 GFLOPS (2.585 ms)
-```
-
-The previous long-run l25 best was `rpt3_l25_unroll16_nosnop` at `258.8 GFLOPS`; `rpt3_l25_unroll16_nosnop_lastcmp0` reached `262.7 GFLOPS` by comparing only in the last unrolled body. The post-increment rewrite removes the explicit `mov r2.y, r10.x` loop-counter copy, drops obsolete loop nops, and moves the increment under the MAD body. Long checked results for the post-increment form: default dispatch `269.4 GFLOPS`, `HCQ2=1` `269.9 GFLOPS`, and `THREAD128=1` `256.2 GFLOPS`.
-
-For the same post-increment l25 shape, `THREAD128=1` is still required for a possible 300+ path even though the full kernel is currently slower. Under `THREAD128=1`, no-store is only `253.8 GFLOPS`, but no-store with skipped A loads reaches `294.4 GFLOPS`, skipped B loads reaches `266.2 GFLOPS`, and skipped A+B loads reaches `316.7 GFLOPS`. This shows the THREAD128 control/ALU ceiling can cross 300, but the current A/B texture schedule cannot. A is the larger limiter on this shape.
-
-Nearby checked post-increment probes did not beat N=704: N=736/unroll23 reached `262.8 GFLOPS`, N=800/unroll25 reached `263.9 GFLOPS` on a long run, N=608/unroll19 reached `266.1 GFLOPS` on a short run, and N=832/unroll26 fell to `203.0 GFLOPS`. For N=704, unroll22 is best so far; unroll16 was `268.9 GFLOPS`, unroll11 was `268.7 GFLOPS`, and unroll44 fell to `206.4 GFLOPS`. MAD accumulator reorderings were flat (`acc3210` long `269.6 GFLOPS` with `HCQ2=1`), and reverse `k3210` remained slower (`258.6 GFLOPS`).
-
-THREAD128-specific l25 probes were negative: unroll4/8/11/16/22 measured about `247.7/250.9/253.2/250.1/249.8 GFLOPS`, while unroll44 fell to `181.3 GFLOPS`; `THREAD128=1 HCQ2=1` was also flat at `253.6 GFLOPS`. Correct load-order variants (`a0early`, `bfirst`) remained around `250-252 GFLOPS`, single-coordinate hoisting was either slower or invalid, and an A `isam.f16` plus `cov.f16f32` path was correct but collapsed to `94.2 GFLOPS`. A0 prefetch into `r6.w` after the current A0 MADs was invalid even with waits, so source-overwrite hazards are stricter than the logical liveness suggests. Follow-up prefetch diagnostics confirmed the constraint: moving A0 to `r6.x` corrupts accumulator registers, moving it to `r15.w..r16.z` is correct but drops to `182.4 GFLOPS` from `fregs=17`, fregs16 coordinate-pair rewrites for A0 still fail checks, and A2/A3 prefetch fail even when delayed until after all current MADs. B0-low remaps are not THREAD128-safe: waits around the late B0 reload and an extra `(sy)` after it still fail checks; the symmetric B0-first low-register schedule also fails.
-
-Additional 300 push checks: `QCOM_PRIORITY=15` did not improve the current best (`THREAD128=1` remained `250.9-253.5 GFLOPS`, `HCQ2=1` default stayed `269.9 GFLOPS`). A short THREAD128 shape sweep around N704 left N704 as the only useful l25 candidate: N608/unroll19 passed but was only `203.3 GFLOPS`, N736/unroll23 passed but was `196.4 GFLOPS`, and N576/N640/N672/N768 did not match the l25 patch shape. Hand FP32 8x8 remains structurally register-heavy (`fregs` in the high 30s for ncols=2), so it is not a near-term 300 route without a major register-layout rewrite.
-
-More N704 THREAD128/300-route probes were also negative. Reversing local-axis priority produced `r_11_22_16_8_4_4_176_4`, but noop was only `193.1 GFLOPS` and the l25 postinc patch was `250.7 GFLOPS`; skip-A/skip-B/skip-A+B no-store ceilings were `286.7`, `268.2`, and `316.7 GFLOPS`, so the load balance did not improve. Applying locals unsorted changed the prologue but kept A driven by `r48.x`, and noop fell to `183.1 GFLOPS`. Image upcast 8 collapsed to `79.0 GFLOPS`, image upcast 2 collapsed to `11.9 GFLOPS`, and nearby N640/N896 l23 patches stayed around `219-222 GFLOPS`. Corrected quad-A with `r48.x&3`, quad-A with an explicit texture wait, and quad-B one-load-per-quad with an explicit post-broadcast wait all failed checks with zero output. Low-register B0 remaps into `r1.y`, `r0.z`, and aligned `r1.x` failed (`352`, `4`, and `352` at idx0), so the low coordinate registers are not a usable f15 escape hatch for this l25 schedule. A bounded `BEAM=2` run again hit `OSError: [Errno 35] Resource deadlock avoided`; avoid longer BEAM on this device for this route.
-
-Follow-up THREAD128 l25 scheduling checks also did not find a 300 route. Splitting the texture wait by delaying A0/A1 loads until after the first A2 MAD was only correct if the `r2.y` loop-counter increment stayed after the delayed A loads; the corrected variants passed but dropped to `166.3 GFLOPS` and `173.8 GFLOPS`, while moving the increment immediately after B0 failed (`idx=16 got=700.0`). Runtime local-size overrides were invalid for this compiled shape: `16,8,1` does not divide the total launch, while `4,32,1` and `8,8,1` failed checks with zero-output regions. Additional checked K/accumulator orders were flat or slower under THREAD128: `k2301` `234.4`, `k2310` `213.4`, `k1023` `242.7`, `k0132` `250.7`, `k0213` `251.2` on a longer run, `k3210` `209.5`, `acc3210` `253.2`, `acc1230` `253.3`, and accumulator-major `239.6 GFLOPS`; `a1mid` load order failed (`idx=32 got=700.0`).
-
-THREAD128 runtime-state probes were also negative and the env hooks were removed. Mesa-like `QCOM_TSIZE=2` was flat on a sequential long run (`251.3 GFLOPS`), `QCOM_TSIZE=1` failed with zero output, `QCOM_TSIZE=4` and `QCOM_USIZE=1` were flat, `QCOM_WGE_SCALAR=1` was flat, `QCOM_SINGLE_SP=1` dropped to `130.2 GFLOPS`, `QCOM_CONSTLEN=128/192` only produced short-run noise and long `CONSTLEN=128` was `252.0 GFLOPS`, `QCOM_THREADMODE=1` dropped to `50.5 GFLOPS`, `QCOM_MERGEDREGS=1` failed with zero output, `QCOM_ISAMMODE_CL=1` was flat, and TPL1 destination datatype override dropped to `240.5 GFLOPS`. Underdeclaring the normal l25 kernel as `f15` failed at idx0, so THREAD128 needs a real lower-register schedule rather than metadata-only occupancy tricks. New f15 B0-streaming attempts into old B1/B2 slots failed checks (`700.0` outputs), and the old `b0low` f15 schedule still fails THREAD128 even at shorter unrolls and stronger waits.
-
-Additional THREAD128-focused follow-up remained negative. Rebaselining current code gave `rpt3_l25_postinc_unroll22` at `253.7 GFLOPS` on a short checked run and `253.7 GFLOPS` on a 30-iter run, while default dispatch stayed around `269.8 GFLOPS`. Setting `SP_PS_WAVE_CNTL.THREADSIZE` through a temporary `QCOM_PS_WAVE_THREADSIZE=1` runtime hook was flat/slower (`251.3 GFLOPS`), so the hook was removed. Moving A1 earlier is not safe: `a1copyearly`, `a1copyearly_wait`, and the f17 coordinate-copy version all failed full-output checks at `idx=32 got=700.0`, even when B2/B3 coordinates were copied away from `r8.*`. Combining `a0early` with K orders where late A1 is consumed last was correctness-safe but not a stable speedup: best short run was `a0early_k0231` at `254.9 GFLOPS`, but a 30-iter comparison fell to `252.5 GFLOPS`. A full 24-permutation `a0early_k####` sweep did not produce a clear winner. An in-unroll coordinate-increment rewrite, intended to avoid recomputing A/B coordinates after the first body of `unroll22`, failed checks (`idx=0 got=440.0`, then `606.0/611.0` after safer recomputation attempts) because A0/A1 texture destinations clobber the apparent persistent coordinate registers. Current conclusion is unchanged: THREAD128 is blocked by texture scheduling/register-liveness constraints in this l25 shape, not by stores or dispatch bits.
-
-The lower-register `b0low` post-increment schedule removed all `r15.*` B-vector use and passed at `fregs=15` under default dispatch, but it was slower (`~255.6 GFLOPS`) and failed correctness under `THREAD128=1`. Lower metadata alone is therefore not enough; the MAD/load order must also be THREAD128-safe.
-
-Important diagnostics: for the earlier `rpt3_l25_unroll16_nosnop` loop, no-store measured only about `256.8 GFLOPS`; no-store with skipped A loads reached `277.1 GFLOPS`, skipped B loads `264.1 GFLOPS`, and skipped A+B loads `284.4 GFLOPS`. That ceiling is still below 300, so load/store deletion alone is not enough; the remaining FP32 gap is dominated by full-register pressure/control scheduling rather than the typed image stores.
-
-The best verified N=1024 compiler-side patch remains `rpt3_accum_f32_unroll8`. It keeps tinygrad's normal packed image layout and rewrites only the default main IR3 kernel. The compact `rpt3_accum_f32_default` patch moves the A0 vector to `r13`, raises the declared full-register footprint to `f14`, replaces the compiler's `64` scalar `mad.f32` ops with `16` `(rpt3)mad.f32` groups, and removes the now-dead `r5.z/r5.w` saves from the loop prefix. Full-output all-ones checks pass.
-
-Same-session patch-harness comparison on `tc3` with `THREAD128=1`, `IMAGE=1`, `FLOAT16=1`, `dtype=half`, and `acc_dtype=float`:
-
-| Patch | Main GFLOPS | Notes |
-|-------|-------------|-------|
-| `noop` | `159.7` | Compiler default: `f13`, `64` scalar `mad.f32` |
-| `reorder_rpt_f32_compact` | `196.3` | `f13`, `36` scalar `mad.f32`, `16` rpt groups |
-| `rpt3_accum_f32_default` | `204.5` | `f14`, `16` `(rpt3)mad.f32`, dead saves removed |
-| `rpt3_accum_f32_unroll8` | `208.0` | `f14`, unrolled checked best for N=1024 |
-
-The same tightened `rpt3_accum_f32_default` patch measured `155.1 GFLOPS` in the harness full-flow timer and `198 GFLOPS` for the main kernel in a single `DEBUG=2 --stats-run` run where noop measured `155 GFLOPS`; the device was in a throttled/low-clock state for that comparison. Negative but correct probes: `rpt3_accum_f32_accmajor` (`199.2 GFLOPS`) and `rpt3_accum_f32_nosnop` (`199.5 GFLOPS`) were slower than the default ordering with the compiler `(ss)nop` retained.
-
-For N=512, the best checked path so far is `rpt3_n512_b0low_k3210_unroll16_nosnop`, which moves the B0 texture vector below `r15`, declares `f15` instead of `f16`, uses `rpt3` accumulator groups, unrolls the K loop by 16, and drops the `(ss)nop`. Fresh checked runs after killing stale remote Python measured `234.7-234.8 GFLOPS` on default THREAD64. The same patch with `THREAD128=1` measured about `231.0 GFLOPS`; `HCQ2=1` measured `234.6 GFLOPS`. Earlier `rpt3_n512_b0low_k3210_unroll16` measured `229.4-230.4 GFLOPS`, and `rpt3_n512_unroll16` measured `225.5 GFLOPS`.
-
-Important negative probes: deleting or NOPing the apparent N=512 dead coordinate copies corrupts output, so those packed sampler-coordinate writes are semantically required. `rpt3_n512_b0low_unroll32` is not reliable (`idx=33216 got=508.0`), `rpt3_n512_b0low_k3210_unroll32_nosnop` is also wrong (`idx=65664 got=508.0`), and f14 N=512 repacks fail even when declared as f15, so the shifted-load schedule is wrong rather than merely underdeclared. The N=1024 `f13pack` attempt also fails all-ones checks (`got=1020.0`), so the current N=1024 verified ceiling remains around `208 GFLOPS`. `BEAM=2/4` hit QCOM deadlocks during beam-search timing and should be avoided for this route.
-
-`BEAM=1` found an alternate compiler schedule (`r_2_32_16_4_4_4_2_4_2_256_4`), but it is not a valid improvement candidate: with real filled inputs it measured only `~92-94 GFLOPS` despite passing all-ones correctness. Earlier higher BEAM timings came from an uninitialized/zero-like input state and should not be counted.
-
-Follow-up performance probes showed the current `8x8` FP32 shape is not the route to 400 GFLOPS:
-
-| Probe | Result | Finding |
-|-------|--------|---------|
-| Compiler imageh input, FP32 output, `ncols=2` | `82.9 GFLOPS` | Correct but far below target |
-| Hand FP32 `ncols=2`, donor `ncols=2` store, post-constant | correct | Reusing the compiler 16-vector `stg.f32` epilogue can cover the full output |
-| Hand FP32 `ncols=2`, real B loads | invalid | B texture path produces sparse/row-group-dependent output; not countable |
-| Hand FP32 `ncols=2`, skip A/B loads, no store | `199.0 GFLOPS` | Upper bound for this 8x8 register footprint/schedule is about stock FP32 speed |
-| Raw hand FP32 MAD microbench | `~355 GFLOPS` | Device can issue more FP32 ALU than the GEMM-shaped loop, but still below the nominal 468 note here |
-
-Implication: pushing FP32 GEMM above 400 needs a different tile/schedule, not incremental fixes to this `8x8` path. The likely next candidate is a lower-register `4x16` FP32-accumulate shape that keeps more wave-pairs resident while amortizing B loads; the current `8x8` FP32 footprint (`f37`) is boxed in around 200 even before real texture loads.
-
-#### Latest 420+ GFLOPS Run
-
-Measured on `tc3` with `THREAD128=1`, `IMAGE=1`, `FLOAT16=1`:
-
-```bash
-PYTHONPATH=. DEV=QCOM IMAGE=1 FLOAT16=1 THREAD128=1 \
-  python3 extra/gemm/qcom_intensity_gemm.py --threads 128 --ncols 4 \
-  --direct --compact-acc --stable-bx --stable-ay --inc-coords \
-  --persistent-coords --alu-order row_col_kk --coord-delay -1 \
-  --k-unroll 4 --first-sync-only --b-first --check
-```
-
-Final check:
-
-```text
-ncols=4 covered_N=1024 fregs=10 hregs=28 waves=3 intensity=3.20 flop/B mad_density=2.46 shader_instrs=677 loop_instrs=417 bytes=5416 envelope_bytes=15744
-mad.f16=256 rpt3=256 isam=80 qbc=0 sy=2
-CHECK PASS all 1048576 outputs are 1024.0
-```
-
-Benchmark command:
-
-```bash
-PYTHONPATH=. DEV=QCOM IMAGE=1 FLOAT16=1 THREAD128=1 \
-  python3 extra/gemm/qcom_intensity_gemm.py --threads 128 --ncols 4 \
-  --direct --compact-acc --stable-bx --stable-ay --inc-coords \
-  --persistent-coords --alu-order row_col_kk --coord-delay -1 \
-  --k-unroll 4 --first-sync-only --b-first --iters 220
-```
-
-Final checked benchmark runs:
-
-| Run | GFLOPS | Time |
-|-----|--------|------|
-| 1 | `430.8` | `4.984 ms` |
-| 2 | `430.0` | `4.994 ms` |
-| 3 | `434.8` | `4.939 ms` |
-| 4 | `425.6` | `5.046 ms` |
-| 5 | `421.2` | `5.099 ms` |
-
-#### How 420 Was Reached
-
-Starting point was the previous fastest verified kernel:
-
-```bash
-PYTHONPATH=. DEV=QCOM IMAGE=1 FLOAT16=1 THREAD128=1 \
-  python3 extra/gemm/qcom_intensity_gemm.py --threads 128 --ncols 4 \
-  --direct --compact-acc --stable-bx --stable-ay --inc-coords \
-  --persistent-coords --alu-order row_col_kk --coord-delay -1 \
-  --k-unroll 4 --first-sync-only
-```
-
-Rebaseline before tuning was noisy but centered around 390-401 GFLOPS:
-
-| Run | GFLOPS | Time |
-|-----|--------|------|
-| 1 | `399.4` | `5.377 ms` |
-| 2 | `387.3` | `5.544 ms` |
-| 3 | `385.7` | `5.567 ms` |
-| 4 | `401.4` | `5.349 ms` |
-| 5 | `394.1` | `5.449 ms` |
-
-The bottleneck probes showed stores were not limiting:
-
-| Probe | Result | Finding |
-|-------|--------|---------|
-| Same kernel, `--no-store` | `382.9-403.9 GFLOPS` | Removing stores did not materially improve throughput |
-| Same kernel, `--post-constant` | `390.7-394.4 GFLOPS` | Store path plus loop remained in the same range |
-| Same kernel, `--store-constant` | `~0.089 ms` | Store-only epilogue is tiny vs `~5.0 ms` full GEMM |
-
-The ALU/load probes showed the checked kernel was not ALU-issue limited:
-
-| Probe | Result | Finding |
-|-------|--------|---------|
-| Same kernel, `--no-store --alu-reps 2` | `536.4-545.2 GFLOPS` | More ALU per same loads immediately beats 420 |
-| Same kernel, `--no-store --alu-reps 3` | `594.1-595.3 GFLOPS` | Load/setup overhead is being amortized |
-| Same kernel, `--no-store --alu-reps 4` | `538.2-549.1 GFLOPS` | Too much body/envelope pressure; not useful as a real path |
-| Same kernel, `--no-store --skip-a-loads` | `448.5-452.7 GFLOPS` | A loads have cost but are not dominant |
-| Same kernel, `--no-store --skip-b-loads` | `587.6-595.7 GFLOPS` | B texture loads/setup dominate the gap |
-| Same kernel, `--no-store --skip-a-loads --skip-b-loads` | `673.3-673.4 GFLOPS` | ALU/control ceiling for this loop shape |
-
-The successful change was `--b-first`: load the first B pair before issuing A loads. This keeps the same `f10 h28`, same `loop_instrs=417`, same `mad.f16=256`, same `isam=80`, and same `sy=2`, but lets the A texture loads hide part of first-pair B texture latency. That moved the full-output checked kernel from `~400 GFLOPS` to `421.2-434.8 GFLOPS`.
-
-Robustness checks around `--b-first`:
-
-| Variant | Result | Finding |
-|---------|--------|---------|
-| `--coord-delay -1` | correct, `421.2-434.8 GFLOPS` | Best path |
-| `--coord-delay 0/1/2/4` | correct, slower | Extra NOPs reduce MAD density from `2.46` to `2.06` |
-| `--store-shlg-offsets` | correct, `426.5-428.7 GFLOPS` | Store variant is flat; default hand store is fine |
-| `--store-scalar-offsets` | correct, no speedup | Store math is not bottleneck |
-| `--donor-store` | correct, no speedup | Donor-store slicing is not needed |
-| `--threads 128` | correct, best | Best balance for this schedule |
-| `--threads 256` | correct, `421.1-423.0 GFLOPS` | Works but slightly slower |
-| `--threads 64` | invalid | Sparse wrong outputs; do not use with `--b-first` |
-| `--low-a-coords` | correct, `424.4-429.6 GFLOPS` | Reduces metadata to `f8 h28`; not faster, so full-register metadata is not limiting |
-| `--low-a-coords --threads 64` | correct, `257.9-261.6 GFLOPS` | Lower fregs fixes 64-thread correctness but remains slow |
-| `--low-a-coords --threads 256` | correct, `426.3-428.1 GFLOPS` | Flat vs 128-thread path |
-| `--k-unroll 2 --first-sync-only` | correct, `369.6-375.3 GFLOPS` | Too little latency hiding |
-| `--k-unroll 4` without `--first-sync-only` | correct, `326.5-341.1 GFLOPS` | Extra MAD syncs dominate |
-| `--k-unroll 8 --b-first --first-sync-only` | correct, `411.1-413.1 GFLOPS` | B-first fixes the old sparse-output failure but the larger body is slower |
-| `--stream-b --stream-b-no-sync` variants | correct, `349-375 GFLOPS` | Hides some latency but adds too many instructions |
-
-#### 460 GFLOPS Attempt
-
-The current 4x16 tile appears boxed in below 460 GFLOPS without reducing B ingress or changing tile shape.
-
-Hard upper-bound probes on the current B-first path:
-
-| Probe | Result | Finding |
-|-------|--------|---------|
-| `--b-first --no-store --skip-a-loads` | `453.3-453.5 GFLOPS` | Even deleting all A loads stays below 460 |
-| `--b-first --low-a-coords --no-store --skip-a-loads` | `458.0-459.4 GFLOPS` | Best A-free upper bound; still below target |
-| `--b-first --no-store --skip-b-loads` | `590.5-590.8 GFLOPS` | B ingress remains the dominant limiter |
-| `--b-first --no-store --skip-a-loads --skip-b-loads` | `673.4 GFLOPS` | ALU/control body has enough headroom |
-
-Additional 460-path probes:
-
-| Probe | Result | Finding |
-|-------|--------|---------|
-| Raise KGSL `devfreq/min_freq` to `710000000` | permission denied | Cannot lock max clock from this user |
-| `--b-first` MAD order sweep | `row_col_kk` still best | Other legal orders were `~385-397 GFLOPS`; `kk_col_row` was invalid |
-| Col2 prefetch into `hr28..hr31` | correct only with targeted waits, `~240 GFLOPS` | Extra high half regs / waits destroy throughput; probe removed from script |
-| Tail column split schedule | correct, `426.0-427.8 GFLOPS` | Same loop size, no improvement; probe removed from script |
-| Partial `ncols=5` B-first probe | `~250 GFLOPS` with `f10 h32`, `~388-393 GFLOPS` with `f8 h32` | Wider 4-row tile is not promising; probe removed from parser |
-| Low-freg `8x8 serial --profile alu` | `681.7 GFLOPS` | Strong ALU headroom, but full serial path still lacks a correct low-reg store/prologue combination |
-| Correct `8x8 --experimental-twopass` with `--fregs-override 8` | hung | High full-register use cannot be hidden by lowering metadata |
-| `8x8 serial --donor8-store` | correct, `189.1-191.1 GFLOPS` | Known-good 8-row donor store fixes correctness at `f12 h32`, but remains slow |
-| `8x8 --split-a --donor8-add256-store --no-next-sy` | correct, `360.2-378.6 GFLOPS` | Pre-unroll split-A baseline; `f10 h28`, four wave-pairs |
-| `8x8 --split-a --split-k-unroll 2 --donor8-add256-store` | correct, `404.2 GFLOPS` | K-unroll starts to hide texture/setup cost |
-| `8x8 --split-a --split-k-unroll 4 --b-coord-delay 3 --donor8-add256-store` | correct, `425.8-436.0 GFLOPS` | Previous best 8x8 path; `f10 h28`, `loop_instrs=110`, `isam=64`, `sy=2` |
-| `8x8 --split-a --split-k-unroll 8 --b-coord-delay 3 --donor8-add256-store` | correct, `415.8 GFLOPS` | Same register footprint but larger shader; instruction-cache/body size likely hurts |
-| `8x8 split-A K-unroll-4 --split-prefetch-next-b --split-fast-coords --fregs-override 8` | correct, `446.2 GFLOPS` | Refills dead B registers for next K step; first real improvement after K-unroll-4 |
-| `8x8 split-A K-unroll-8 --split-prefetch-next-b --split-fast-coords --fregs-override 8` | correct, `449.5-453.3 GFLOPS` | Next-B prefetch makes unroll-8 viable; best before store tightening |
-| Same K-unroll-8 prefetch path, `--no-store` | `466.7 GFLOPS` | Shows store epilogue became the final blocker for 460 |
-| Same K-unroll-8 prefetch path, `--add256-store-mode pairs` | correct, `451.9 GFLOPS` | Generated store slice with fewer nops; correct but not enough |
-| Same K-unroll-8 prefetch path, `--add256-store-mode tight` | correct, `467.9-468.8 GFLOPS` | First verified >460 path; generated SAD + back-to-back stores |
-| Same tight path, `--b-coord-delay 0` | correct, `468.5 GFLOPS` | Flat vs delay 1; delay `-1` is still invalid |
-| Same tight path, `--split-hoist-b0-coord --fregs-override 9` | correct, `468.8 GFLOPS` long run, `469.3 GFLOPS` short run | Hoisting first next-B0 coord into `r8.x/r8.y` is correct but essentially flat |
-| Same tight path, no-store/skip probes | `466.7 / 529.1 / 535.5 / 562.4 GFLOPS` | no-store / skip-A / skip-B / skip-both; remaining 500 gap is A+B texture ingress, not ALU |
-| Same tight path, `--threads 64` | correct, `277.0 GFLOPS` | Lower thread count is much slower |
-| Same tight path, `--threads 256` | correct, `464.7-468.6 GFLOPS` | Fixed 8-row prologue row-log for 256 threads; no speedup vs 128 |
-| Same tight path, `--fregs-override 7` | invalid | Full-register metadata below 8 corrupts output |
-| Same tight path, `--fregs-override 6` | hung | Recover with `pkill -9 python3`; do not use |
-| Same tight path, `--add256-gap <16` | invalid | Tight store still needs the old inter-column gap |
-| Same tight path, `--add256-direct-sources` | invalid | Direct stores from accumulator hregs still violate the low-reg store-source convention |
-| Same tight path, `--split-buffer-a` | invalid | Both `hr28..hr31` A buffering and low `hr12..hr15` A buffering with accumulators at `hr16` corrupt output |
-| Same tight path, `--split-prefetch-next-a` | correct, `465.2 GFLOPS`; swapped before B1 `445.8 GFLOPS` | Moving A0-next earlier hurts texture issue balance |
-| Same tight path, `--split-interleave-next-b` | correct, `460.4 GFLOPS` | Splitting B0-next refill around col1 MADs is slower |
-| Same tight path, `--split-hoist-b0-coord` with `fregs=8` | invalid | Hoisted coord in `r4.y/r4.z` is clobbered before ISAM |
-| Same tight path, `--split-inline-b-wait --split-inline-b-nop 1..7` | invalid | Inline `add.s(nop)` cannot replace the explicit coordinate wait NOP |
-| Same tight path, `--split-add-a-rows` | invalid | A row coordinate formation must stay `or.b` for this schedule |
-| Same tight path, `--split-prefetch-loop-b` | correct, `442.6 GFLOPS` | Predicate-skipped final prefetch fixes correctness, but loop-boundary B prefetch is much slower |
-| Same tight path, `--split-quad-a` | hung/invalid | Row-per-quad A sharing with full-register quad broadcasts is not a valid path yet; early high/default layouts hung |
-| Same tight path, `--split-high-a` | correct, `468.6 GFLOPS` | Moves A to `hr24..hr27` and accumulators to `hr8..hr23`; register layout is flat |
-| Same tight path, `--split-high-a --split-hoist-b0-coord --fregs-override 9` | correct, `469.0 GFLOPS` | Flat vs non-high-A B0 hoist |
-| Same tight path, `--split-low-a` | correct, `459.2-461.9 GFLOPS` | Moves A to `hr0..hr3` and B to `hr4..hr11`; needs declared `fregs=10`, while `fregs=8` corrupts output |
-| Same tight path, `--split-low-a --split-quad-a` | invalid | Single-component quad broadcasts avoid the earlier hang but rows sourced through qbc are mixed/NaN; `shader_instrs=1127`, `loop_instrs=122`, `isam=80`, `sy=18` |
-| Same tight path, `--split-high-a --split-quad-a --fregs-override 14` | invalid | Same row pattern as low-A qbc: directly loaded rows are ok, broadcast-derived rows are mixed; register placement/freg declaration is not the fix |
-| Same tight path, branch-gated low-A quad load | invalid, not kept | Lane-0-only A loading plus qbc kept `isam=128`, grew to `shader_instrs=1231`, and corrupted every row; divergent hand branch form is not usable here |
-| Same tight path, `--split-pair-b-coords` | initially correct but slower, then invalid when tightened | Pairing two B coordinates per wait did not improve texture issue; tightened `1006`-instruction version corrupts output |
-| Same tight path, `--split-base-b-y` | invalid | Keeping B y as a base multiple of 4 and forming kk offsets with `or.b` corrupts first outputs |
-| Same tight path, `--split-stream-next-b0` | correct, `458.2 GFLOPS` | Per-K component streaming of next B0 frees texture issue earlier but hurts MAD order enough to lose speed |
-| Same tight path, `--split-stream-next-b1` | correct, `456.1 GFLOPS` | Same result for B1 streaming; earlier B issue does not offset disrupted row/col/kk order |
-| Same tight path, `--swap-grid` | invalid at `--b-coord-delay 0`, correct but `441.1 GFLOPS` at delay 3 | Swapping row/column group IDs can make the store map correct, but fast B delay loses contributions and safe delay is slower |
-| Same tight path, `--hregs-override 27/24/22/20/18` | hung before check | Underdeclaring half-register metadata is unsafe; recover with `pkill -9 python3` |
-| Same tight path after FP16 peak warmup | `455.4 GFLOPS` | Governor/preheat did not help; long warmup can be slower |
-| `8x16 split-A K-unroll-4` | correct, `307.3 GFLOPS` | Higher arithmetic intensity is overwhelmed by `hregs=48` / 3-wave occupancy; threads 64 is slower and threads 256 is invalid |
-| `8x8 split-A K-unroll-16` with next-B prefetch | correct, `391.3 GFLOPS` | Fits larger envelope but instruction-cache/body size dominates |
-| `8x8 split-A --no-store` | `380.0 GFLOPS` | Store overhead is modest; same `f10 h28` metadata |
-| `8x8 split-A --no-store --skip-a-loads` | `425.0 GFLOPS` | A texture path costs about 45 GFLOPS from no-store baseline |
-| `8x8 split-A --no-store --skip-b-loads` | `480.8 GFLOPS` | B texture path is the main limiter and has enough headroom for 4x16 parity if reduced |
-| `8x8 split-A --no-store --skip-a-loads --skip-b-loads` | `583.2 GFLOPS` | ALU/control ceiling for this split-A loop; not ALU-limited |
-| `8x8 split-A K-unroll-4 --strip-mad-sy` | invalid | First outputs become `-inf` / `NaN`; keep the current two MAD syncs |
-| `8x8 split-A K-unroll-4 --grouped-b --b-coord-delay -1` | correct, `405.5 GFLOPS` | Fewer B coord waits but worse texture issue pattern |
-| `8x8 split-A K-unroll-4 --grouped-b-cols --b-coord-delay -1` | correct, `409.9 GFLOPS` | Also slower than the scalar B setup path |
-| `8x8 split-A K-unroll-8 --grouped-b --b-coord-delay -1` | correct, `379.9 GFLOPS` | Larger body plus grouped B is a dead end |
-| `8x8 split-A K-unroll-4 --no-store` | `429.8 GFLOPS` | Store is not the main remaining limiter in the unrolled path |
-| `8x8 split-A K-unroll-4 --no-store --skip-a-loads` | `508.2 GFLOPS` | A texture path still costs significant throughput |
-| `8x8 split-A K-unroll-4 --no-store --skip-b-loads` | `524.6 GFLOPS` | B texture path is still the larger limiter |
-| `8x8 split-A K-unroll-4 --no-store --skip-a-loads --skip-b-loads` | `567.0 GFLOPS` | ALU/control ceiling for the unrolled split-A shape |
-| `8x8 split-A K-unroll-4 --b-coord-delay 2/1/0/-1` | invalid | `--b-coord-delay 3` is still required |
-| `8x8 split-A K-unroll-4 --threads 64` | correct, `259.8 GFLOPS` | Lower occupancy/parallelism is much slower |
-| `8x8 split-A K-unroll-4 --threads 256` | correct, `429.1 GFLOPS` | Fixed by 8-row prologue row-log update; still flat vs 128 |
-| `8x8 split-A --add256-gap <16` | invalid | Gap 16 is required; smaller gaps corrupt row 7 / first column |
-| `8x8 split-A --stream-b1` | invalid | Tried to load second B group during first-column MADs; still misses contributions even with waits and syncs; parser flag removed |
-
-Conclusion: 460 was reached by combining real B-ingress overlap with an epilogue reduction. The next-B prefetch schedule moves B loads for the next unrolled K step into dead B registers after current group-4 col0/col1 use, and tight generated add256 stores remove the donor store nops that became visible once the loop reached the mid-450s. The first 500 push did not find a valid faster schedule; the best verified long run remains `468.8 GFLOPS`, with skip-A and skip-B probes showing that another real A/B texture-ingress reduction is needed.
-
-The key fix was adding dependency waits while widening the donor prologue's
-column base from `gid.x*32+tid` to `gid.x*128+tid`; without waits, the repeated
-adds did not chain and the kernel overlapped columns instead of covering the
-tail.
-The later compact-acc improvement moves the accumulator base from `hr16` to
-`hr12`, reducing metadata from `hregs=32` to `hregs=28`. This is store-safe
-because the 4-row donor pack only overwrites `hr12` after `row0,col0` has already
-been copied into the store scratch registers.
-The current default direct store is now explicit hand ASM rather than a runtime
-slice from a donor binary. It hard-codes the compiler-style four-row address
-schedule and `stg.f16` sequence, then packs output rows into `hr0..hr3` before
-stores. A naive single-address `stg` path was invalid because it used dependent
-scalar address math too aggressively and did not follow the compiler's low-register
-store-source convention.
-
-Useful commands:
-
-```bash
-# FP16 MAD peak parity with OpenCL
-PYTHONPATH=. DEV=QCOM THREAD128=1 python3 extra/mmapeak/qcom_fp16_mad_peak.py
-
-# GEMM-shaped ALU-only profile
-PYTHONPATH=. DEV=QCOM IMAGE=1 FLOAT16=1 THREAD128=1 \
-  python3 extra/gemm/qcom_8x4_gemm.py --ncols 2 --threads 128 --profile alu
-
-# Fastest correct 4x16 full GEMM path
-PYTHONPATH=. DEV=QCOM IMAGE=1 FLOAT16=1 THREAD128=1 \
-  python3 extra/gemm/qcom_intensity_gemm.py --threads 128 --ncols 4 \
-  --direct --compact-acc --stable-bx --stable-ay --inc-coords \
-  --persistent-coords --alu-order row_col_kk --coord-delay -1 \
-  --k-unroll 4 --first-sync-only --b-first --check
-PYTHONPATH=. DEV=QCOM IMAGE=1 FLOAT16=1 THREAD128=1 \
-  python3 extra/gemm/qcom_intensity_gemm.py --threads 128 --ncols 4 \
-  --direct --compact-acc --stable-bx --stable-ay --inc-coords \
-  --persistent-coords --alu-order row_col_kk --coord-delay -1 \
-  --k-unroll 4 --first-sync-only --b-first --iters 220
-
-# Fastest correct 8x8 path so far
-PYTHONPATH=. DEV=QCOM IMAGE=1 FLOAT16=1 THREAD128=1 \
-  python3 extra/gemm/qcom_8x4_gemm.py --variant serial --ncols 2 \
-  --threads 128 --split-a --split-k-unroll 8 --b-coord-delay 0 \
-  --donor8-add256-store --split-prefetch-next-b --split-fast-coords \
-  --fregs-override 8 --add256-store-mode tight --check
-PYTHONPATH=. DEV=QCOM IMAGE=1 FLOAT16=1 THREAD128=1 \
-  python3 extra/gemm/qcom_8x4_gemm.py --variant serial --ncols 2 \
-  --threads 128 --split-a --split-k-unroll 8 --b-coord-delay 0 \
-  --donor8-add256-store --split-prefetch-next-b --split-fast-coords \
-  --fregs-override 8 --add256-store-mode tight --warmup 10 --iters 500
-
-# Previous fastest pipelined 8x8 path
-PYTHONPATH=. DEV=QCOM IMAGE=1 FLOAT16=1 THREAD128=1 \
-  python3 extra/gemm/qcom_8x4_gemm.py --ncols 2 --threads 128 --pipeline --a-coord-delay 4 --b-coord-delay -1 --no-next-sy --check
-PYTHONPATH=. DEV=QCOM IMAGE=1 FLOAT16=1 THREAD128=1 \
-  python3 extra/gemm/qcom_8x4_gemm.py --ncols 2 --threads 128 --pipeline --a-coord-delay 4 --b-coord-delay -1 --no-next-sy --warmup 5 --iters 30
-```
-
-Recent results and negative checks:
-
- Grouped A/B coordinate scheduling can pass some all-ones runs but is flaky under full scan/check; do not count its timings.
- Removing the current-buffer pipeline `(sy)` is incorrect; removing only the next-buffer `(sy)` is correct and gives a small speedup.
- `4x16` direct constant-store diagnostics still fail with the hand store path, proving that path is store/address incorrect before GEMM math is considered.
- Reusing a sliced donor `4x4` store epilogue is correct for direct `4x8` and direct `4x16` only after adding waits between every dependent widened-column stride add, including the final wait before B-coordinate setup.
- The earlier `4x16` tail-zero pattern was not primarily a store-epilogue limit: the donor prologue stride adds were reading the old `r7.y`, effectively using a `4x8` column stride and overlapping workgroups.
- The native direct `4x16` compiler store epilogue fixes full-output coverage, but it uses high full registers and drops the verified full kernel to about `184 GFLOPS`.
- Hybrid hand-tail stores and scalar/shlg-offset donor-store diagnostics did not beat the fixed low-reg donor-store path.
- A pipelined direct `4x16` no-store experiment was slower (`~220 GFLOPS`) because the extra double-buffer registers reduced occupancy (`hregs=40`).
- `ncols=3`/`4x12` probes now report `covered_N=768`; after correcting for partial coverage the donor-store path is only `295.0 GFLOPS`, so a split `4x12 + tail` plan is not promising.
- `threads=256` is correctness-clean for direct `4x8`, but slower (`~252.6 GFLOPS`) than `threads=128`.
- The experimental `8x16` donor-store path in `qcom_8x4_gemm.py` also needed stride-add waits; this fixes tail coverage but it still has sparse row failures and remains invalid.
- Semantic K-unroll for direct `4x16` is correct for `--k-unroll 2` and `4`, but mostly flat (`~330-332 GFLOPS` without compact accumulators). `--k-unroll 8` produced sparse zero output chunks and is invalid.
- Direct `4x16 --preload-b` is full-output correct but slow (`176.6 GFLOPS`) because `hregs=36` drops occupancy.
- Direct `4x16 --stream-b` and `--stream-b --stream-b-no-sync` are full-output correct, but did not beat the baseline (`~310 GFLOPS` with sync, `~329 GFLOPS` without the extra pair sync).
- Direct `4x16 --compact-acc` is correct and is the best small improvement so far (`~334-336 GFLOPS` with hand ASM stores and `--k-unroll 4`).
- Direct `4x16 --compact-acc --first-sync-only` is the main current improvement. Full-output checks pass with only the first MAD sync in each unrolled K loop (`sy=2` total), and `--stable-bx --k-unroll 4` measures `~382-388 GFLOPS`.
- Direct `4x16 --compact-acc --stable-bx --stable-ay --inc-coords --persistent-coords --first-sync-only` is the previous best verified 4x16 path. It keeps A row coords in `r8/r9`, increments A/B coords across unrolled K steps, preserves them across loop iterations, and has measured `400.5-402.1 GFLOPS` with full-output checks.
- Direct `4x16 --compact-acc --stable-bx --stable-ay --inc-coords --persistent-coords --first-sync-only --b-first` is the current best verified 4x16 path. It preserves the same loop instruction count/register footprint as the persistent-coordinate path but loads the first B pair before A, hiding part of the B texture latency under the A loads. Final checked runs measured `421.2-434.8 GFLOPS`.
- Direct `4x16 --b-first --low-a-coords` is correct and reduces metadata to `f8 h28`, but it remains flat at `424.4-429.6 GFLOPS`; metadata pressure is not the current limiter.
- Current 4x16 upper-bound probes put the practical scheduling ceiling below 460: `--b-first --no-store --skip-a-loads` is only `453.3-453.5 GFLOPS`, and `--b-first --low-a-coords --no-store --skip-a-loads` is only `458.0-459.4 GFLOPS`.
- The 460-specific schedule probes were negative: col2 B prefetch was either invalid or `~240 GFLOPS`, tail column split was flat at `426.0-427.8 GFLOPS`, and temporary partial `ncols=5` was slow and not full-output coverage.
- Low-freg `8x8 serial --profile alu` reaches `681.7 GFLOPS`, so `8x8` has ALU headroom if it stays at `f8 h32`; the full serial path still fails full-output checks because the naive scalar store is unsafe, the 4x16 hand epilogue mismatches the 8-row prologue, and the dynamic 4-row store remains invalid.
- `8x8 --experimental-twopass --fregs-override 8` hung on `tc3`; recover with `pkill -9 python3`. Keep the correct two-pass path at its declared `f15 h32` metadata.
- `8x8 serial --donor8-store` proves the low-reg serial compute loop is correct once stores are fixed, but only reaches `189.1-191.1 GFLOPS` at `f12 h32`.
- `8x8 --split-a --donor8-add256-store --no-next-sy` is the correct pre-unroll split-A baseline. It preloads both B groups, computes two 4-row A groups, and uses a low-freg donor store that forms the second column by adding `+256` bytes to the first column's row addresses. Full-output checks pass at `f10 h28`; benchmark range is `360.2-378.6 GFLOPS`.
- `8x8 --split-a --split-k-unroll 4 --b-coord-delay 3 --donor8-add256-store` was the previous best verified 8x8 path. Full-output checks pass with `f10 h28`, `reg_count=24`, `shader_instrs=602`, `loop_instrs=110`, `isam=64`, `sy=2`; benchmark range is `425.8-436.0 GFLOPS`.
- `8x8 --split-a --split-k-unroll 8 --b-coord-delay 0 --split-prefetch-next-b --split-fast-coords --fregs-override 8 --add256-store-mode tight` is the current best practical path. Full-output checks pass with `f8 h28`, `reg_count=22`, `shader_instrs=1022`, `loop_instrs=109`, `isam=128`, `sy=2`; benchmark range is `467.9-468.6 GFLOPS` on long runs, with prior short runs at `468.1-468.6 GFLOPS`.
- The best 500-push variant, `--b-coord-delay 0 --split-hoist-b0-coord --fregs-override 9`, also full-output checks and measured `468.8 GFLOPS` on a long run (`469.3 GFLOPS` short run). It needs `f9` for `r8.x/r8.y` hoisted B0 coords and is effectively tied with the `f8` path.
- The winning path depends on both parts. K-unroll-8 plus next-B prefetch but donor store mode topped out at `449.5-453.3 GFLOPS`; `--no-store` reached `466.7 GFLOPS`, exposing the epilogue as the last blocker. `--add256-store-mode tight` replaces the donor store slice with generated SAD plus back-to-back stores and raises the verified full kernel above 460.
- The post-460 bottleneck is A+B texture ingress. On the tight K-unroll-8 path, no-store is `466.7 GFLOPS`, skip-A is `529.1 GFLOPS`, skip-B is `535.5 GFLOPS`, and skip-both is `562.4 GFLOPS`.
- The 500-specific low-register schedule probes were negative: direct accumulator store sources are invalid, smaller add256 gaps are invalid, next-A prefetch is correct but slower, buffered-A variants are invalid, high-A and low-A layouts are flat/slower, paired/base B-coordinate forms are invalid or slower, interleaved next-B refill is slower, per-component B0/B1 streaming is slower, swapped-grid B-cache reuse is invalid or slow, predicated loop-boundary B prefetch is correct but slower, inline B wait encoding is invalid, row-per-quad A sharing is invalid/hung, and fregs/hregs below the known-safe footprint corrupt or hang.
- K-unroll-16 with the same next-B prefetch is correct but slow (`391.3 GFLOPS`) despite fitting a larger envelope; do not continue in that direction unless instruction-cache behavior changes.
- `8x8 --split-a --split-k-unroll 8 --b-coord-delay 3 --donor8-add256-store` is correct and still fits the enlarged donor envelope (`8336 / 13064` bytes), but it is slower at `415.8 GFLOPS`; doubling the body does not pay for reduced loop control.
- Split-A K-unroll robustness is narrow. The original K-unroll-4 path requires `--b-coord-delay 3`; lower B coordinate delays corrupt output. `--threads 64` is correct but slow at `259.8 GFLOPS`; `--threads 256` is now correct after the row-log fix but flat (`429.1 GFLOPS`). The add256 donor store still needs `--add256-gap 16`; smaller gaps corrupt row 7 / first column.
- Split-A K-unroll grouped-B modes are correct with `--b-coord-delay -1`, but slower: K-unroll-4 `--grouped-b` is `405.5 GFLOPS`, K-unroll-4 `--grouped-b-cols` is `409.9 GFLOPS`, and K-unroll-8 `--grouped-b` is `379.9 GFLOPS`. The scalar B setup with explicit delay remains best.
- Removing MAD syncs with `--strip-mad-sy` is invalid on K-unroll-4; first outputs become `-inf` / `NaN`.
- Split-A K-unroll-4 bottleneck probes show the path is still ISAM/texture limited, not ALU limited: no-store is `429.8 GFLOPS`, skipping A loads is `508.2 GFLOPS`, skipping B loads is `524.6 GFLOPS`, and skipping both reaches `567.0 GFLOPS`.
- Experimental split-A `--stream-b1` did not become correct. It still misses one contribution in the streamed column even after adding B-coordinate waits, col1 syncs, and hard gaps; the parser flag was removed.
- Direct `4x16 --b-kk-pipeline` is invalid: it repeatedly missed 1-2 FP16 contributions even with strong MAD sync diagnostics.
- Direct `4x16 --compact-acc --stable-bx --first-sync-only --k-unroll 8` is full-output correct but slower (`~355-358 GFLOPS`); non-stable `k-unroll 8` still has sparse zero chunks and is invalid.
- Experimental `8x16` is full-output correct at `--threads 64` and `--threads 256`, but slow (`144.7` and `~260 GFLOPS` respectively); `--threads 128` still has sparse failures.
- Experimental low-freg `8x8 --donor4-store` is invalid. The two 4-row donor chunks do not match the 8-row prologue/store convention; observed failures include `1020.0` outputs and zero rows.
-
-### Combined ISAM + Real GEMM ALU Probes
-
-These are throughput probes, not valid GEMM results when `--no-store` or
-`--alu-reps > 1` is used. They combine real texture `isam` loads with the legal
-GEMM FMA form `acc = A_scalar * B_half4 + acc`.
-
-| Probe | Result | Notes |
-|-------|--------|-------|
-| `4x16 --direct --no-store --row-col-kk --alu-reps 1` | `328.4 GFLOPS` | One real ALU body per loaded A/B tile; no stores |
-| `4x16 --direct --no-store --row-col-kk --alu-reps 2` | `471.2 GFLOPS` | First combined ISAM+real-GEMM-ALU probe over 400 |
-| `4x16 --direct --no-store --row-col-kk --alu-reps 4` | `577.1 GFLOPS` | Load overhead amortized further |
-| `4x16 --direct --no-store --row-col-kk --alu-reps 8` | `639.8 GFLOPS` | Approaches true GEMM ALU body ceiling |
-| `4x16 --direct --no-store --row-col-kk --quad-a --alu-reps 2` | `444.0 GFLOPS` | Quad-A path is slower than normal A loads here |
-| `4x16 --direct --donor-store --row-col-kk --coord-delay 4` | `331.4 GFLOPS` | Correct full-output GEMM after stride-dependency fix |
-| `4x16 --direct --compact-acc --stable-bx --first-sync-only --k-unroll 4` | `~382-388 GFLOPS` | Correct full-output GEMM; previous reduced-sync path |
-| `4x16 --direct --compact-acc --stable-bx --stable-ay --inc-coords --persistent-coords --first-sync-only --k-unroll 4` | `400.5-402.1 GFLOPS` | Correct full-output GEMM; previous persistent-coordinate path |
-| Same persistent-coordinate path with `--b-first` | `421.2-434.8 GFLOPS` | Correct full-output GEMM; current best verified path |
-| Same B-first path with `--low-a-coords` | `424.4-429.6 GFLOPS` | Correct full-output GEMM; fregs drops to 8 but speed is flat |
-| Same persistent-coordinate path, `--no-store` | `382.9-403.9 GFLOPS` | Store path is not the bottleneck |
-| Same persistent-coordinate path, `--store-constant` | `~0.089 ms` | Store-only lower bound; epilogue is negligible vs `~5.0 ms` GEMM |
-| Same persistent-coordinate path, `--no-store --alu-reps 2` | `536.4-545.2 GFLOPS` | Load/setup amortization probe |
-| Same persistent-coordinate path, `--no-store --alu-reps 3` | `594.1-595.3 GFLOPS` | Confirms the checked kernel is not ALU-issue limited |
-| Same persistent-coordinate path, `--no-store --skip-a-loads` | `448.5-452.7 GFLOPS` | A loads cost measurable time but are not dominant |
-| Same persistent-coordinate path, `--no-store --skip-b-loads` | `587.6-595.7 GFLOPS` | B texture loads/setup are the dominant bottleneck |
-| Same persistent-coordinate path, `--no-store --skip-a-loads --skip-b-loads` | `673.3-673.4 GFLOPS` | ALU/control ceiling for this loop shape |
-| Same B-first path, `--no-store --skip-a-loads` | `453.3-453.5 GFLOPS` | A-free upper bound for current B ingress is still below 460 |
-| Same B-first low-A path, `--no-store --skip-a-loads` | `458.0-459.4 GFLOPS` | Best current 4x16 upper-bound probe, still below 460 |
-| `4x16 --direct --native-store --row-col-kk` | `184.2 GFLOPS` | Correct full coverage, high full-register store epilogue |
-
-Useful command for the first over-400 combined probe:
-
-```bash
-PYTHONPATH=. DEV=QCOM IMAGE=1 FLOAT16=1 THREAD128=1 \
-  python3 extra/gemm/qcom_intensity_gemm.py --threads 128 --ncols 4 \
-  --direct --no-store --row-col-kk --alu-reps 2 --iters 40
-```
-
-## Hardware: Adreno 630
-
- **SP**: Shader Processor, Qualcomm's shader core/cluster; roughly analogous to an NVIDIA SM or AMD CU
- **2 SPs**, each with 64 ALUs, 128 total
- **Clock**: ~400 MHz (thermal-dependent)
- **FP16 MAD peak**: 690 GFLOPS (measured via mmapeak with same-register repeated MAD)
- **FP16 MAD sustained**: 590 GFLOPS (realistic with `(rpt3)mad.f16`, 16 groups in a tight loop)
- **FP32 MAD peak**: 468 GFLOPS
- **Texture bandwidth**: 168 GB/s (measured, isam throughput)
- **Register file**: 192 KiB per SP on A630 (`reg_size_vec4=96`, `threadsize_base=64`, `wave_granularity=2`)
- **Wave sizes**: THREAD128 (128 fibers/wave) or THREAD64 (64 fibers/wave)
-
-### Register File Constraints
-
-The `fregs` and `hregs` fields in the shader binary are **vec4 footprints**, not scalar component counts.
-
- `r0` is one full vec4: `r0.x/r0.y/r0.z/r0.w`, four 32-bit components.
- `hr0` is one half vec4: `hr0.x/hr0.y/hr0.z/hr0.w`, four 16-bit components.
- In split OpenCL mode, one full vec4 costs the same storage as two half vec4s.
- The useful GPR namespace is `r0..r47` for full regs and `hr0..hr47` for half regs. `hr48+` reaches special/non-GPR names and is not usable for hand accumulators on A630.
-
-For split full/half allocation, the full-equivalent per-fiber footprint is:
-
-`reg_count = fregs + ceil(hregs / 2)`
-
-Mesa reports A630 as `reg_size_vec4=96`, so a single 128-fiber wave-pair can hold up to 96 full-equivalent vec4 registers per fiber. The physical storage per SP is:
-
-`96 vec4/fiber * 64 fibers * 2 wave-granularity * 16 bytes/vec4 = 196608 bytes = 192 KiB`
-
-Older notes used `floor(12288 / (hregs * threads))`, which is only a rough half-only shortcut and is wrong once full regs or split full/half accounting matters.
-
-### 8x4 Half Register Budget
-
-`hr0..hr47` is 48 half4 registers = 192 FP16 scalar values per fiber. That is enough only for the serial-B 8x4 schedule:
-
-| Live data | half4 regs | FP16 values |
-|-----------|------------|-------------|
-| 8x4 accumulators | 32 | 128 |
-| 8 A texels | 8 | 32 |
-| 4 B texels, one column group | 4 | 16 |
-| **Serial-B subtotal** | **44** | **176** |
-| Spare before scratch/alias pressure | **4** | **16** |
-
-The preload-B schedule does not fit:
-
-| Live data | half4 regs | FP16 values |
-|-----------|------------|-------------|
-| 8x4 accumulators | 32 | 128 |
-| 8 A texels | 8 | 32 |
-| 16 B texels, all column groups | 16 | 64 |
-| **Preload-B subtotal** | **56** | **224** |
-
-So 8x4 is not register-impossible, but only the serial-B form fits the addressable half register file. Preloading all B values needs at least 56 half4 registers before any scratch or store epilogue, beyond the usable `hr0..hr47` range.
-
-| hregs | Max waves | Total fibers |
-|-------|-----------|-------------|
-| 24    | 4         | 512         |
-| 31    | 3         | 384         |
-| 48    | 2         | 256         |
-
-Full registers and half registers **share the same physical storage**:
-`r0.x` = `{hr0.x, hr0.y}`, `r0.y` = `{hr0.z, hr0.w}`, etc.
-Writing a full register clobbers the aliased half registers and vice versa.
-
-## Architecture of the GEMM Kernel
-
-### Tiling
-
- **128 threads/workgroup** = 4 subgroups of 32 threads
- Each thread computes **4 rows x 1 col4** (4 output half4 vectors)
- Grid: `(N/128, M/16, 1)` — 16 rows per WG (4 subgroups x 4 rows)
- A is stored as `image2d_t` shape `(M, K/4)`, each pixel = half4
- B is stored as `image2d_t` shape `(K, N/4)`, each pixel = half4
- Per K iteration: 4 A loads + 4 B loads = 8 `isam.1d` texture fetches
-
-### Loop Body (compiled, before patching)
-
-```
-mov r2.y, r6.z      ;; k4 -> A coord x (row0)
-(rpt5)nop            ;; wait for mov
-isam hr3.x, r2.y, t#0   ;; A[k4, row0] -> hr3
-mov r2.w, r6.z
-(rpt5)nop
-isam hr2.x, r2.w, t#0   ;; A[k4, row1] -> hr2
-mov r3.y, r6.z
-(rpt5)nop
-isam hr1.x, r3.y, t#0   ;; A[k4, row2] -> hr1
-mov r3.w, r6.z
-(rpt5)nop
-isam hr0.x, r3.w, t#0   ;; A[k4, row3] -> hr0
-
-add.s r4.z, r6.y, -3
-(rpt5)nop
-isam hr4.x, r4.y, t#1   ;; B[col4, k4*4+0] -> hr4
-(sy)mad.f16 ...          ;; 16 scalar MADs for B[0] x 4 rows
-;; ... repeat for B[1], B[2], B[3] with more isam + (sy) + MADs
-```
-
-**Problems**: 5 `(sy)` syncs per iteration (~100 cycles each), 4 `(rpt5)nop` waits
-(6 wasted cycles each), scalar MADs instead of packed `(rpt3)`.
-
-### Binary Patching (`patch_kernel` in `qcom_gemm.py`)
-
-1. **Strip redundant `(sy)`**: Keep only the first `(sy)` on a MAD instruction per loop
-   iteration. The QCOM compiler inserts `(sy)` before every MAD that follows an isam,
-   but only one sync is needed to wait for all pending texture results.
-
-2. **Convert scalar MADs to `(rpt3)mad.f16`**: When 4 consecutive MAD instructions have
-   the same `src1`, sequential `dst/src2/src3`, the pattern matches `(rpt3)` repeat
-   encoding. Each `(rpt3)` packs 4 MADs into 1 instruction slot.
-
-3. **Merge `(rpt1)+(rpt1)` into `(rpt3)`**: Two adjacent `(rpt1)mad.f16` with compatible
-   register sequences combine into a single `(rpt3)`.
-
-Result: **5 `(sy)` → 2**, **41 scalar MADs → 15 `(rpt3)` + 2 `(rpt1)`**.
-Speedup: **78 → 190 GFLOPS** (2.4x).
-
-### Hand-Assembled Optimized Loop
-
-Best verified kernel places B texels into 4 separate registers (hr4-hr7 instead of
-all-hr4), enabling all 8 isam to be issued back-to-back with a single `(sy)`:
-
-```
-;; Coord setup (8 instructions)
-mov r2.y, r6.z       ;; A coords
-mov r2.w, r6.z
-mov r3.y, r6.z
-mov r3.w, r6.z
-add.s r4.z, r6.y, -3 ;; B coords
-add.s r5.x, r6.y, -2
-add.s r5.z, r6.y, -1
-mov r6.x, r6.y
-
-;; 8 isam back-to-back (no nops between)
-isam hr3.x, r2.y, t#0   ;; A row0
-isam hr2.x, r2.w, t#0   ;; A row1
-isam hr1.x, r3.y, t#0   ;; A row2
-isam hr0.x, r3.w, t#0   ;; A row3
-isam hr4.x, r4.y, t#1   ;; B k0
-isam hr5.x, r4.w, t#1   ;; B k1
-isam hr6.x, r5.y, t#1   ;; B k2
-isam hr7.x, r5.w, t#1   ;; B k3
-
-;; Single (sy) + 15 (rpt3)mad.f16 + 2 (rpt1)mad.f16 = 64 MADs
-(sy)(rpt3)mad.f16 hr20.z, hr3.x, (r)hr4.x, (r)hr20.z  ;; row0 x B0
-(rpt3)mad.f16 hr24.z, hr2.x, (r)hr4.x, (r)hr24.z       ;; row1 x B0
-... ;; 13 more (rpt3) groups
-(rpt1)mad.f16 hr13.z, hr0.w, (r)hr7.x, (r)hr13.z       ;; row3 x B3 (noncontiguous)
-(rpt1)mad.f16 hr15.x, hr0.w, (r)hr7.z, (r)hr15.x
-
-;; Loop control
-cmps.s.eq p0.x, r6.z, 255
-add.s r6.z, r6.z, 1
-add.s r6.y, r6.y, 4
-(rpt3)nop
-br !p0.x, #loop_top
-```
-
-Result: **200 GFLOPS** (verified correct), limited by 3-wave occupancy (`hregs=31`).
-
-## ir3 Assembler (`ir3asm.py`)
-
-Hand-assembles Adreno a6xx (ir3 ISA) instructions. Uses a compiled OpenCL kernel as
-a "donor" for the binary envelope (headers, buffer descriptors, sampler info, constant
-tables) and replaces the shader instructions and register counts.
-
-### Key functions
-
-| Function | Description |
-|----------|-------------|
-| `get_envelope(dev, src)` | Compile OpenCL, return `(lib, img_off, img_sz, reg_off)` |
-| `inject(lib, ..., shader, fregs, hregs)` | Replace shader + reg counts in binary |
-| `assemble(instr_list)` | Concatenate instruction bytes |
-| `disasm(shader_bytes)` | Disassemble via Mesa `ir3_isa_disasm` |
-| `MAD_F16(dst, src1, src2, src3, rpt, sy, r)` | Encode `(sy?)(rptN?)mad.f16` |
-| `ISAM_F16(dst, coord, tex)` | Encode `isam.1d (f16)(xyzw)` |
-| `STG_F16(addr, data_hreg)` | Encode `stg.f16 g[rADDR], hrDATA, 4` |
-
-### Instruction encoding (64-bit, little-endian)
-
-Each instruction is 8 bytes stored as two 32-bit words `[lo, hi]`:
-
- **hi[31:24]**: Opcode category (0x00=nop/br, 0x20=mov, 0x42=add.s, 0x40=add.f,
-  0x63/0x73=mad.f16, 0xa0=isam, 0xc0=stg)
- **hi[23:16]**: Sub-opcode and flags (e.g., `(sy)` sets bit 28 → 0x73 vs 0x63)
- **hi[15:8]**: Repeat count and register flags (`rpt` in bits [6:0], `r` flag in bit 7)
- **hi[7:0]**: Destination register index
- **lo**: Source registers and immediates (layout varies by category)
-
-## Measured Performance
-
-| Configuration | GFLOPS | Notes |
-|---------------|--------|-------|
-| Pure ALU ceiling (16 rpt3, T128) | 590 | No texture, just MADs |
-| Pure texture ceiling (8 isam/iter) | 168 GB/s ≈ 335 GFLOPS equiv | No MADs |
-| Compiled 4-row GEMM (unpatched) | 78 | 5 (sy), scalar MADs |
-| Patched 4-row GEMM (sy-strip + rpt3) | 190 | 2 (sy), 15 rpt3 |
-| Hand-assembled (separate B, hregs=31) | 200 | 1 (sy), 16 rpt3, 3 waves |
-| Hand-assembled (hregs=24, WRONG output) | 240 | Register aliasing, 4 waves |
-| Direct 4x16 compact persistent coords | 400.5-402.1 | Correct full-output GEMM, `f10 h28`, `loop_instrs=417` |
-| Direct 4x16 compact persistent coords + B-first | 421.2-434.8 | Current fastest checked full GEMM |
-
-## Legacy Bottleneck Analysis (200 GFLOPS Kernel)
-
-This older analysis explains the first hand-assembled 200 GFLOPS kernel. The current 420+ kernel bottleneck analysis is in the `How 420 Was Reached` section above.
-
-At 200 GFLOPS with 3 waves and `hregs=31`:
-
- **Loop body**: 8 coord setup + 8 isam + 17 MAD instrs + 5 loop ctrl = **38 instructions**
- **Effective**: 64 MADs / 38 total = 1.68 MADs/instruction
- **Texture-limited peak**: 168 GB/s / (64 bytes/iter) × 128 FLOPS/iter = **336 GFLOPS**
- **Achieved/peak**: 200/336 = **60%** — the gap is `(sy)` stall time not hidden by 3 waves
-
-### Why 300+ GFLOPS requires 4 waves
-
-With 4 waves, the GPU can switch to another wave during the `(sy)` stall, keeping ALUs
-busy. But 4 waves requires `hregs ≤ 24` (24 × 128 × 4 = 12288 = register file size).
-
-The compiled kernel uses `hregs=31` because its accumulator layout spans hreg indices
-54-121 (max index 121, requiring ≥31 vec4 slots). A clean layout using indices 32-95
-(max 95, requiring 24 slots) fits in 4 waves but needs a **custom store epilogue**.
-
-The store epilogue is difficult because:
-1. Full registers (r0-r3) alias half registers (hr0-hr7) in the same physical file
-2. The QCOM runtime uses 64-bit buffer addresses requiring `cmps.u.lt` + `sad.s32`
-   for carry propagation, which references constant registers `c20.x/c20.y`
-3. The address computation and accumulator reduction must be sequenced to avoid
-   clobbering results through register aliasing
-
-## Approaches Tried
-
-| Approach | Result | Why |
-|----------|--------|-----|
-| Strip `(sy)` + rpt3 patching | 190 GFLOPS | Baseline, 2.4x over compiled |
-| Separate B texture registers | 200 GFLOPS | Single `(sy)`, 3 waves |
-| Remove coord nops | +5 GFLOPS | Nops not needed between mov and isam |
-| Fast B coords (increment vs recompute) | Same | Saves instructions but not cycles |
-| 8-row kernel (2 waves) | 53 GFLOPS | Too few waves, 4 `(sy)` after patching |
-| Software pipelining (double buffer) | N/A | Requires hregs>31 for double A+B, ≤2 waves |
-| Interleaved B (4x sy) | 84 GFLOPS | 4 `(sy)` stalls kill throughput |
-| 2x K-unroll | GPU hang | Immediate overflow (256 > 8-bit) in CMPS |
-| Clean acc layout + custom epilogue | Close | Full/half reg aliasing in epilogue |
-| hregs=24 with compiled epilogue | 240 GFLOPS wrong | Acc indices > 95 alias across fibers |
-| Local-memory staging | 99 GFLOPS | Barriers/local-memory path are slower than direct texture fetch here |
-| Buffer/global loads | 87 GFLOPS | `ldg.f16` path measured far below texture throughput |
-| Compiler 4x2 col tile | 47 GFLOPS | Higher arithmetic intensity, but register allocation destroys `(rpt3)` MAD packing |
-| Hand 4x2 col tile, 4 partial accs | 204 GFLOPS wrong | Intended 12 isam + 32 rpt3 loop, custom epilogue still writes partial output |
-| Hand 4x2 direct acc, hregs=24 | 247 GFLOPS wrong | Faster occupancy, but repeated accumulator dependencies produce NaNs/infs |
-| `shfl.rdown.u32` A broadcast probe | 9.0 G lane-shuffles/s | Too slow to replace texture ingress |
-| `quad_shuffle.brcst.u32` probe | 22.7 G lane-broadcasts/s | Fast enough for quad-level A sharing on paper |
-| 4x2 direct baseline, T128 | 249.6 GFLOPS wrong | 61-instruction loop, 32 `(rpt3)` MADs |
-| 4x2 quad-A, 8 scalar qbc | 210.9 GFLOPS wrong | Branch + 8 broadcasts cost more than saved A ingress |
-| 4x2 quad-A, 4 `(xy)` qbc | 225.4 GFLOPS wrong | Wrmask cuts qbc count but still below baseline |
-| 4x2 quad-A, 2 `(xyzw)` qbc | 232.9 GFLOPS wrong | Best quad-A result so far, still slower than baseline |
-
-### Direct Texture Bandwidth Sweep
-
-`qcom_texture_bw.py` measures logical half4 `isam.1d` bytes issued by a hand shader.
-Each load is 8 bytes. The best stable point measured on tc3 is ~148 GB/s.
-
-| Threads | Loads/K step | hregs | waves | GB/s | Notes |
-|---------|--------------|-------|-------|------|-------|
-| 128 | 4 | 20 | 4 | 96.9 | Too few independent loads per sync |
-| 128 | 8 | 24 | 4 | 127.4 | 4-wave 4x1-like load count |
-| 128 | 12 | 28 | 3 | 143.5 | Good balance |
-| 128 | 16 | 32 | 3 | 75.1 | Stable slow point; not enough load depth after occupancy drop |
-| 128 | 20 | 36 | 2 | 72.1 | Stable slow point |
-| 128 | 24 | 40 | 2 | 144.5 | Recovers with deeper load stream |
-| 128 | 28 | 44 | 2 | 146.4 | Near roof |
-| 128 | 32 | 48 | 2 | 147.8 | Best measured |
-
-If the ALU target is 717 GFLOPS, the texture path requires arithmetic intensity
-`717 / 147.8 = 4.85 FLOP/byte`. With the 590 GFLOPS sustained ALU number, the
-requirement is `590 / 147.8 = 3.99 FLOP/byte`.
-
-For an `R x C` per-thread tile, where `C` is the number of col4 output vectors:
-
-`AI = 32*R*C / (8*R + 32*C) = 4*R*C / (R + 4*C)`.
-
-This explains why widening only columns helps slowly:
-
-| Tile | AI |
-|------|----|
-| 4x2 | 2.67 |
-| 4x8 | 3.56 |
-| 8x2 | 4.00 |
-| 8x4 | 5.33 |
-| 16x2 | 5.33 |
-
-So 4x8 cannot feed a 717 GFLOPS target from the measured texture path. 8x4 or
-16x2 is the first class of tiles with enough texture arithmetic intensity.
-
-## Current 4x2 Intensity Experiment
-
-`qcom_intensity_gemm.py` is an experimental hand-assembled 4-row x 2-col4 tile:
-
- Per K iteration: 4 A `isam` + 8 B `isam` = 96 bytes/thread
- Work per K iteration: 8 output half4 vectors x 4 K lanes = 128 MADs = 256 FLOPs/thread
- Texture roof: `168 GB/s / 96 bytes * 256 FLOPs` = **448 GFLOPS**
- The loop assembles as 12 `isam`, 32 `(rpt3)mad.f16`, one `(sy)`-bearing MAD, plus loop/control overhead.
-
-Important pitfalls found while building this:
-
-1. `BR(offset)` is relative to the branch instruction, not the next instruction.
-   The old `loop_start - loop_end - 1` form jumps back one instruction too far.
-2. `(rpt3)mov.f16f16 hrX.x, hrX.x` does **not** broadcast an immediate to `xyzw`.
-   Use `mov imm hrX.x` then `mov hrX.y, hrX.x (rpt2)`, or copy from a known scalar into a different destination base.
-3. The `SAD_S32` encoding only matched the observed odd component forms initially.
-   Using `r6.x` decoded as `(neg)r6.y`; use/check disassembly for every new source register.
-4. Patching `shlg` from immediate 5 to 6 is not a safe way to compute `gid.x*64 + lane`.
-   Use the raw group id (`r51.w`) and integer adds, then refresh duplicated B coordinate registers.
-5. Direct accumulation into one output vector is too dependent: updating the same accumulator four times inside one loop iteration produced NaNs/infs even though it lowers `hregs` to 24.
-6. The custom store epilogue is still not correct. With all-one inputs, row 0 starts correctly but most output locations remain zero, so the 4x2 GFLOPS numbers are throughput probes only.
-
-## Subgroup / Quad Broadcast Findings
-
-`extra/gemm/qcom_shfl_probe.py` tests register-to-register data movement across
-fibers using the hand assembler.
-
-Measured on tc3:
-
-| Operation | Result | Notes |
-|-----------|--------|-------|
-| `shfl.rdown.u32` immediate 1 | Works, ~9.0 G lane-ops/s | Other tested immediates/register xor read back as zero in the current probe |
-| `quad_shuffle.brcst.u32` | Works, ~22.7 G lane-ops/s | Requires cat5 FULL bit set for u32 sources; supports wrmask `(xy)`/`(xyzw)` |
-| `getfiberid.u32` | Hangs in injected envelope | Do not use in GEMM kernels until the required envelope/control setup is understood |
-| Simple hand divergent `br` | Not reliable | Uniform loop branch encoding is not enough for divergent control flow |
-| Compiler image branch | Emits `br !p0.x` around `isam` plus `(ss)(jp)` join target | Use this pattern before hand-assembling conditional A loads |
-
-Implication: full-subgroup `shfl` is not the right ingress path. Quad broadcast is
-fast enough as an instruction by itself, but the first 4x2 GEMM integration is
-slower than the direct 4x2 baseline because the branch/join and broadcast
-instructions reduce MAD issue density.
-
-Arithmetic intensity if quad-level A sharing works:
-
-| Tile | Texture bytes/thread/K | FLOPs/thread/K | Intensity | Texture roof @168 GB/s |
-|------|------------------------|----------------|-----------|------------------------|
-| Current 4x1 | 64 | 128 | 2.00 FLOP/B | 336 GFLOPS |
-| 4x1 + quad A sharing | 40 | 128 | 3.20 FLOP/B | 538 GFLOPS |
-| Current 4x2 | 96 | 256 | 2.67 FLOP/B | 448 GFLOPS |
-| 4x2 + quad A sharing | 72 | 256 | 3.56 FLOP/B | 597 GFLOPS |
-
-Quad broadcast itself is not the limiting roof for 4x1: 2 `(xyzw)` quad broadcasts
-per K iteration gives roughly `22.7 / 2 * 128 = 1453 GFLOPS` of broadcast capacity.
-The limiting issue is the extra loop instructions. In 4x2 direct mode, the loop
-grew from 61 to 70 instructions while keeping the same 32 `(rpt3)` MADs, so static
-MAD density fell from `128/61 = 2.10` to `128/70 = 1.83` MADs/instruction. Even if
-the divergent branch suppresses 3/4 of A texture lanes, this does not compensate
-at the 4x2 tile size.
-
-Next implication: do not use quad-A sharing for 4x2. If this path is tried again,
-it needs a wider in-register tile where the 2 qbc + branch/join overhead is
-amortized across more B columns/MADs, or a way to suppress A loads without a
-divergent branch sequence.
-
-## Key ISA Details
-
-### `(sy)` — Texture Sync
-
-Stalls until all pending texture results have arrived. Costs ~80-100 cycles per
-occurrence. With 4 waves, other waves execute during the stall. With 3 waves,
-the stall is only partially hidden.
-
-### `(rpt3)mad.f16` — Packed 4x MAD
-
-Executes 4 MAD operations in a single instruction slot. Requires consecutive
-`dst`, `src2`, `src3` registers. The `(r)` flag enables auto-increment on
-`src2` and `src3`. Throughput: 1 `(rpt3)` per cycle → 4 MADs/cycle/ALU.
-
-### `isam.1d (f16)(xyzw)` — Integer-Sampled Texture Fetch
-
-Reads a half4 from an image using integer coordinates packed in a full register pair.
-Latency ~100 cycles. Multiple isam can be pipelined (issued back-to-back); `(sy)`
-waits for all of them.
-
-### `shlg` / `shrm` — Shift with Merge
-
-Used for packing workgroup/thread IDs into coordinate registers.
-`shlg(imm, src1, src2)` ≈ `(src1 << imm) | (src2 & ((1<<imm)-1))`.
-
-### `stg.f16` — Global Store (FP16)
-
-`stg.f16 g[rADDR], hrDATA, 4` stores 4 consecutive half-registers (8 bytes) to the
-address in a full register pair. The data hreg index in the encoding is `hreg * 2`
-(byte offset within the register file).
-
-### `quad_shuffle.brcst` — Quad Register Broadcast
-
-`quad_shuffle.brcst (u32)(x)rD, rS, rI` broadcasts one source lane inside a 4-lane
-quad. For full-width types the cat5 FULL bit must be set; otherwise Mesa disassembles
-the sources as half registers. The cat5 wrmask works: `(xy)` and `(xyzw)` forms
-disassemble and run, allowing two A half4 rows to be broadcast with one u32 `(xyzw)`
-instruction. Measured throughput is ~22.7 G lane-broadcasts/s.
-
-### `shfl` — Subgroup Shuffle
-
-`shfl.rdown.u32` encodes and executes, but measured throughput is only ~9.0 G
-lane-shuffles/s on this device. That is below the texture-ingress rate it would need
-to replace, so it is not the preferred A broadcast primitive.
-
-## Files
-
-| File | Description |
-|------|-------------|
-| `extra/gemm/ir3asm.py` | ir3 instruction assembler + binary envelope injection |
-| `extra/gemm/qcom_gemm.py` | Compiled GEMM + binary patching benchmark |
-| `extra/gemm/qcom_asm_gemm.py` | Hand-assembled GEMM test suite (ALU, load, full) |
-| `extra/gemm/qcom_shfl_probe.py` | `shfl`, `quad_shuffle.brcst`, and branch/join probes |
-| `extra/gemm/qcom_texture_bw.py` | Direct hand-assembled `isam.1d` texture GB/s benchmark |
@@ -13,7 +13,7 @@ from tinygrad import Tensor, Device, Context, GlobalCounters
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.helpers import getenv, colored
 from tinygrad.dtype import dtypes, AddrSpace
-from tinygrad.engine.realize import Estimates, run_linear
+from tinygrad.engine.realize import Estimates
 from tinygrad.renderer.amd.dsl import s, v, VCC_LO, NULL
 from tinygrad.runtime.autogen.amd.rdna3.ins import *

@@ -167,7 +167,7 @@ PREFETCH_LOADS = [(V_LDS_A_DATA[4+2*i], V_LDS_A_DATA[4+2*i+1], V_GLOBAL_B_ADDR,
 # =============================================================================

 class Kernel:
-  def __init__(self): self.instructions, self.labels, self.pos = [], {}, 0
+  def __init__(self, arch='gfx1100'): self.instructions, self.labels, self.pos, self.arch = [], {}, 0, arch
  def label(self, name): self.labels[name] = self.pos

  def emit(self, inst, target=None):
@@ -196,10 +196,10 @@ class Kernel:
 # Kernel builder
 # =============================================================================

-def build_kernel(N):
+def build_kernel(N, arch='gfx1100'):
  assert N % 128 == 0, f"N must be a multiple of 128 (tile size), got {N}"
  assert N >= 256, f"N must be >= 256 (prefetch pipeline requires at least 2 K-blocks), got {N}"
-  k = Kernel()
+  k = Kernel(arch)

  # ===========================================================================
  # PROLOGUE: Load kernel arguments, compute tile coordinates and addresses
@@ -443,7 +443,7 @@ def test_matmul():
  dev = Device[Device.DEFAULT]
  print(f"Device arch: {dev.renderer.target.arch}")

-  insts = build_kernel(N)
+  insts = build_kernel(N, dev.renderer.target.arch)

  rng = np.random.default_rng(42)
  a = Tensor(rng.random((N, N), dtype=np.float32) - 0.5)
@@ -458,20 +458,16 @@ def test_matmul():
  def asm_kernel(A:UOp, B:UOp, C:UOp) -> UOp:
    gidxs = [UOp.special(n, f"gidx{i}") for i,n in enumerate(grid)]
    lidxs = [UOp.special(n, f"lidx{i}") for i,n in enumerate(local)]
-    lds_size = max(LDS_SIZE, 65536//getenv("LIMIT_OCC", 65536))
-    lds = UOp.placeholder((lds_size,), dtypes.uint8, 0, AddrSpace.LOCAL)
+    lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=max(LDS_SIZE, 65536//getenv("LIMIT_OCC", 65536)), addrspace=AddrSpace.LOCAL), (), 'lds')
    sink = UOp.sink(A.base, B.base, C.base, lds, *gidxs, *lidxs, arg=KernelInfo(name=colored("kernel", "cyan"),
                                                                                  estimates=Estimates(ops=N*N*N*2, mem=N*N*4*3)))
-    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
  c = Tensor.custom_kernel(a, b, c, fxn=asm_kernel)[2]
-  linear = c.schedule_linear()
+  ei = c.schedule()[0].lower()

  ets = []
  with Context(DEBUG=2):
-    for _ in range(getenv("CNT", 5)):
-      start = GlobalCounters.time_sum_s
-      run_linear(linear)
-      ets.append(GlobalCounters.time_sum_s - start)
+    for _ in range(getenv("CNT", 5)): ets.append(ei.run(wait=True))
  print(f"REAL TFLOPS {N * N * N * 2 / min(ets) * 1e-12:.2f}")

  if getenv("VERIFY", 1):
@@ -1,5 +1,5 @@
 from tinygrad import Device, UOp, getenv
-from tinygrad.uop.ops import AxisType, KernelInfo
+from tinygrad.uop.ops import AxisType, KernelInfo, Ops
 from tinygrad.dtype import AddrSpace, dtypes

 N = getenv("N", 4096)
@@ -46,8 +46,8 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
  # -- GLOBAL -> LOCAL --
  # wmma: spatial outer, k inner (k contiguous for vectorized WMMA tile loads)
  # gemm: k outer, spatial inner
-  A_local = UOp.placeholder((BLOCK_M, BLOCK_K) if use_wmma else (BLOCK_K, BLOCK_M), a.dtype, slot=0, addrspace=AddrSpace.LOCAL)
-  B_local = UOp.placeholder((BLOCK_N, BLOCK_K) if use_wmma else (BLOCK_K, BLOCK_N), b.dtype, slot=1, addrspace=AddrSpace.LOCAL)
+  A_local = UOp.placeholder((BLOCK_M, BLOCK_K) if use_wmma else (BLOCK_K, BLOCK_M), a.dtype.base, slot=0, addrspace=AddrSpace.LOCAL)
+  B_local = UOp.placeholder((BLOCK_N, BLOCK_K) if use_wmma else (BLOCK_K, BLOCK_N), b.dtype.base, slot=1, addrspace=AddrSpace.LOCAL)

  a = a.reshape(K // BLOCK_K, BLOCK_K, BLOCK_M)
  b = b.reshape(K // BLOCK_K, BLOCK_K, BLOCK_N)
@@ -66,7 +66,7 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:

  # accumulator (unified: both paths use (TM, TN) with scalar dtypes.float)
  acc = UOp.placeholder((TM, TN), dtypes.float, slot=2, addrspace=AddrSpace.REG)
-  acc = acc.after(acc.store(acc.zeros_like(buffer=False)))
+  acc = acc.after(acc.store(acc.zeros_like()))

  if use_wmma:
    k = UOp.range(BLOCK_K // WMMA_K, 101, AxisType.REDUCE)
@@ -80,7 +80,7 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
      # NOTE: since this is part of K, these 2 can be anywhere in the frags and long as a and b match
      a_frag = a_frag.reshape(2, 8)[lane_m, :]
      b_frag = b_frag.reshape(2, 8)[lane_m, :]
-    wmma = UOp.wmma(a_frag, b_frag, acc_frag.after(k), ((16, 16, 16), 'AMD', 32))
+    wmma = UOp(Ops.SHAPED_WMMA, dtypes.float, (a_frag, b_frag, acc_frag.after(k)), arg=((16, 16, 16), 'AMD', 32))
    acc_store = acc_frag.store(wmma).end(tile_m, tile_n)
  else:
    # registers for LOCAL -> REG
@@ -19,7 +19,6 @@ LOG2E = math.log2(math.e)
 def warp_shfl_xor(val, offset, lane):
  """Read val from lane ^ offset using ds_bpermute."""
  idx = ((lane ^ offset) * 4).cast(dtypes.int)
-  if val.op is Ops.INDEX and val.addrspace == AddrSpace.REG: val = val.load()
  return UOp(Ops.CUSTOM, dtypes.float, (idx, val),
             arg="__builtin_bit_cast(float, __builtin_amdgcn_ds_bpermute({0}, __builtin_bit_cast(int, {1})))")

@@ -97,7 +96,7 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
  S_frag = S_reg.reshape(TM // WMMA_ACC, WMMA_ACC, TN).permute(0, 2, 1)[tm1, tn1]
  q_frag = Q_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, D // WMMA_K, WMMA_K)[wave_m, tm1, lane_n, k_qk]
  k_frag = KV_lds_k.reshape(WAVES_N, TN, WMMA_N, D // WMMA_K, WMMA_K)[wave_n, tn1, lane_n, k_qk]
-  qk = UOp.wmma(q_frag, k_frag, S_frag.after(k_qk), WMMA_ARG)
+  qk = UOp(Ops.SHAPED_WMMA, dtypes.float, (q_frag, k_frag, S_frag.after(k_qk)), arg=WMMA_ARG)
  qk_done = S_frag.store(qk).end(tm1, tn1).end(k_qk)
  S_reg = S_reg.after(qk_done)

@@ -127,7 +126,10 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
  P_lds = QP_lds[:, :BLOCK_N]
  P_write = P_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_ACC, LANES_PER_WAVE_M, WAVES_N, TN, LANES_PER_WAVE_N)
  P_write = P_write.permute((0, 4, 3, 6, 1, 2, 5)).reshape(THREADS_PER_BLOCK, TM, TN)
-  P_store = P_write[tid].store(S_reg.cast(dtypes.half))
+  # TODO: P_write[tid].store(S_reg.cast(dtypes.half)) — shaped store fails due to RESHAPE(DEFINE_LOCAL) surviving linearization
+  rw1 = UOp.range(TM, 296, AxisType.LOOP)
+  rw2 = UOp.range(TN, 297, AxisType.LOOP)
+  P_store = P_write[tid, rw1, rw2].store(S_reg[rw1, rw2].cast(dtypes.half)).end(rw1, rw2)

  # -- online softmax correction --
  ri4 = UOp.range(TM, 330, AxisType.LOOP)
@@ -158,7 +160,7 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
  acc_frag = acc.reshape(TM // WMMA_ACC, WMMA_ACC, TD).permute(0, 2, 1)[tm2, tn2]
  p_frag = P_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, BLOCK_N // WMMA_K, WMMA_K)[wave_m, tm2, lane_n, k_pv]
  v_frag = KV_lds_v.reshape(WAVES_N, TD, WMMA_N, BLOCK_N // WMMA_K, WMMA_K)[wave_n, tn2, lane_n, k_pv]
-  pv = UOp.wmma(p_frag, v_frag, acc_frag.after(k_pv), WMMA_ARG)
+  pv = UOp(Ops.SHAPED_WMMA, dtypes.float, (p_frag, v_frag, acc_frag.after(k_pv)), arg=WMMA_ARG)

  # end KV tile loop
  n_tile_end = acc_frag.store(pv).end(tm2, tn2).end(k_pv).barrier().end(n_tile)
@@ -1,39 +1,31 @@
 # kernel8_batched_gmem.s from https://seb-v.github.io/optimization/update/2025/01/20/Fast-GPU-Matrix-multiplication.html
 # sudo PATH=/opt/homebrew/Cellar/llvm/20.1.6/bin:$PATH AMD_LLVM=0 AMD=1 DEBUG=2 python3 extra/gemm/amd_matmul.py
 import pathlib
+from dataclasses import replace
 from tinygrad import Tensor, Device, Context, GlobalCounters
 from tinygrad.helpers import getenv
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
-from tinygrad.renderer import Estimates
-from tinygrad.engine.realize import run_linear
+from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program

 N = 4096
 run_count = 5

-def make_matmul_kernel(name:str, src:str, local_size:int):
-  def fxn(a:UOp, b:UOp, c:UOp) -> UOp:
-    threads = UOp.special(local_size, "lidx0")
-    wg_x = UOp.special(N//128, "gidx0")
-    wg_y = UOp.special(N//128, "gidx1")
-    sink = UOp.sink(a.base, b.base, c.base, threads, wg_x, wg_y, arg=KernelInfo(name, estimates=Estimates(ops=2*N**3, mem=3*N*N*4)))
-    lib = Device[Device.DEFAULT].compiler.compile_cached(src)
-    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)),
-                                 UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
-  return fxn
-
 if __name__ == "__main__":
+  ast = (Tensor.empty(N, N)@Tensor.empty(N, N)).schedule()[-1].ast
+  prg = get_program(ast, Device.default.renderer)
+
  if getenv("ASM") == 1:
    src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel8_batched_gmem.s").read_text()
-    name, local_size = "kernel", 128
+    prgfast = replace(prg, name="kernel", src=src, global_size=[N//128, N//128, 1], local_size=[128, 1, 1])
  elif getenv("ASM") == -1:
    src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel3_registers.cpp").read_text()
-    name, local_size = "kernel3_registers", 256
+    prgfast = replace(prg, name="kernel3_registers", src=src, global_size=[N//128, N//128, 1], local_size=[256, 1, 1])
  elif getenv("ASM") == -2:
    src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel4_gmem_df.cpp").read_text()
-    name, local_size = "kernel4_gmem_db", 256
+    prgfast = replace(prg, name="kernel4_gmem_db", src=src, global_size=[N//128, N//128, 1], local_size=[256, 1, 1])
  else:
    src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel5_lds_optim.cpp").read_text()
-    name, local_size = "kernel5_lds_optim", 128
+    prgfast = replace(prg, name="kernel5_lds_optim", src=src, global_size=[N//128, N//128, 1], local_size=[128, 1, 1])
+  runner = CompiledRunner(prgfast)

  a = Tensor.randn(N, N).realize()
  b = Tensor.randn(N, N).realize()
@@ -43,8 +35,8 @@ if __name__ == "__main__":
  with Context(DEBUG=2):
    for _ in range(run_count): tc = (a@b).realize()

-  linear = Tensor.custom_kernel(a, b, c, fxn=make_matmul_kernel(name, src, local_size))[2].schedule_linear()
  GlobalCounters.reset()
+  ei = ExecItem(ast, [a.uop.buffer, b.uop.buffer, c.uop.buffer], prg=runner)
  with Context(DEBUG=2):
-    for _ in range(run_count): run_linear(linear)
+    for _ in range(run_count): ei.run(wait=True)
  print(f"custom  {(c-tc).square().mean().item()}")
@@ -122,7 +122,7 @@ def eval_custom_matmul(fxn, dt=dtypes.float):
  with Context(DEBUG=0): Tensor.realize(a, b)

  ets = []
-  with Context(DEBUG=max(2, DEBUG.value)):
+  with Context(DEBUG=max(2, DEBUG.value), DEVECTORIZE=2 if dt == dtypes.half else 0):
    for _ in range(NUM_RUNS):
      GlobalCounters.reset()
      tst = Tensor.custom_kernel(c, a, b, fxn=fxn)[0].realize()
@@ -0,0 +1,180 @@
+#!/usr/bin/env python3
+import numpy as np
+import time
+import sys
+np.set_printoptions(linewidth=160)
+np.set_printoptions(linewidth=1000, threshold=10000000000, suppress=False)
+from tinygrad.runtime.ops_llvm import LLVMDevice, LLVMProgram, LLVMCompiler
+from llvmlite import ir  # type: ignore
+from tinygrad.helpers import flat_mv
+from tinygrad.device import MallocAllocator
+
+# https://github.com/corsix/amx/blob/main/Instructions.md
+# 12 lines for AMX support
+from functools import partialmethod
+class AMX:
+  @staticmethod
+  def nop_op_imm5(op, imm5, builder): builder.asm(ir.FunctionType(ir.VoidType(), []), f".word (0x201000 + ({op} << 5) + {imm5}); amx op {op} imm {imm5}", "", tuple(), True)
+  @staticmethod
+  def op_gpr(op, builder, gpr): builder.asm(ir.FunctionType(ir.VoidType(), [ir.IntType(64)]), f".word (0x201000 + ({op} << 5) + 0$0 - ((0$0 >> 4) * 6)); amx op {op} reg $0", "r", (gpr,), True)
+  set, clr = partialmethod(nop_op_imm5, 17, 0), partialmethod(nop_op_imm5, 17, 1)
+  ldx, ldy, stx, sty = partialmethod(op_gpr, 0), partialmethod(op_gpr, 1), partialmethod(op_gpr, 2), partialmethod(op_gpr, 3)
+  ldz, stz, ldzi, stzi = partialmethod(op_gpr, 4), partialmethod(op_gpr, 5), partialmethod(op_gpr, 6), partialmethod(op_gpr, 7)
+  extrx, extry = partialmethod(op_gpr, 8), partialmethod(op_gpr, 9)
+  fma64, fms64, fma32, fms32 = partialmethod(op_gpr, 10), partialmethod(op_gpr, 11), partialmethod(op_gpr, 12), partialmethod(op_gpr, 13)
+  mac16, fma16, fms16 = partialmethod(op_gpr, 14), partialmethod(op_gpr, 15), partialmethod(op_gpr, 16)
+  vecint, vecfp, matint, matfp, genlut = partialmethod(op_gpr, 18), partialmethod(op_gpr, 19), partialmethod(op_gpr, 20), partialmethod(op_gpr, 21), partialmethod(op_gpr, 22)
+
+def int_const(x): return ir.Constant(ir.IntType(64), x)
+
+
+N = 4096
+# N = 1024
+# N = 64
+
+BW = N*N*4
+
+# matrix is 64M, max load bandwidth is 57 GB/s
+# cache line looks like 256 bytes (64 floats)
+
+na = np.zeros((256), dtype=np.float32)
+# na = np.zeros((N, N), dtype=np.float32)
+nb = np.random.randn(N, N).astype(np.float32)
+nc = np.random.randn(N, N).astype(np.float32)
+
+ns = nb.reshape(-1, 32).sum(axis=0)
+
+a = MallocAllocator.alloc(na.nbytes)
+b = MallocAllocator.alloc(nb.nbytes)
+c = MallocAllocator.alloc(nc.nbytes)
+
+MallocAllocator._copyin(b, flat_mv(nb.data))
+MallocAllocator._copyin(c, flat_mv(nc.data))
+
+module = ir.Module(name=__file__)
+func = ir.Function(module, ir.FunctionType(ir.IntType(64), [ir.FloatType().as_pointer()]*3), name='exec')
+
+# load all
+entry = ir.IRBuilder(func.append_basic_block(name="entry"))
+zm, xm, ym = [entry.ptrtoint(func.args[i], ir.IntType(64)) for i in range(3)]
+
+loop_1 = ir.IRBuilder(func.append_basic_block(name="loop_y"))
+loop_1_exit = ir.IRBuilder(func.append_basic_block(name="loop_y_exit"))
+exit = ir.IRBuilder(func.append_basic_block(name="exit"))
+
+y = loop_1.phi(ir.IntType(64), name="y")
+y.add_incoming(int_const(0), entry._block)
+yp = loop_1_exit.add(y, int_const(32*2))
+y.add_incoming(yp, loop_1_exit._block)
+
+prefetch_function = ir.Function(module, ir.FunctionType(ir.VoidType(), [ir.PointerType(ir.FloatType()), ir.IntType(32), ir.IntType(32), ir.IntType(32)]), name="llvm.prefetch")
+
+xptr = y
+addr = loop_1_exit.add(xm, loop_1_exit.mul(int_const(4), xptr))
+
+#prefetch_ptr = loop_1_exit.inttoptr(loop_1_exit.add(addr, int_const(128)), ir.PointerType(ir.FloatType()))
+#loop_1_exit.call(prefetch_function, [prefetch_ptr, ir.IntType(32)(0), ir.IntType(32)(2), ir.IntType(32)(1)])
+
+AMX.ldx(loop_1_exit, loop_1_exit.add(int_const(1<<62), addr))
+xptr = loop_1_exit.add(xptr, int_const(32))
+AMX.ldy(loop_1_exit, loop_1_exit.add(int_const(1<<62), loop_1_exit.add(xm, loop_1_exit.mul(int_const(4), xptr))))
+
+AMX.fma32(loop_1_exit, int_const(1 << 63 | 1 << 28))
+AMX.fma32(loop_1_exit, int_const(1 << 63 | 1 << 28 | 1 << 20 | (16*4)<<10))
+AMX.fma32(loop_1_exit, int_const(1 << 63 | 1 << 29))
+AMX.fma32(loop_1_exit, int_const(1 << 63 | 1 << 29 | 1 << 20 | (16*4)))
+
+AMX.set(entry)
+
+AMX.stz(exit, exit.add(zm, int_const(1 << 62 | (0 << 56) | 0)))
+AMX.clr(exit)
+
+entry.branch(loop_1._block)
+loop_1.branch(loop_1_exit._block)
+loop_1_exit.cbranch(loop_1_exit.icmp_unsigned("==", yp, int_const(N*N)), exit._block, loop_1._block)
+exit.ret(int_const(0))
+
+device = LLVMDevice("llvm")
+prog = LLVMProgram(device, "exec", LLVMCompiler(device).compile(str(module)))
+
+"""
+loop_1 = ir.IRBuilder(func.append_basic_block(name="loop_y"))
+loop_2 = ir.IRBuilder(func.append_basic_block(name="loop_x"))
+loop_3 = ir.IRBuilder(func.append_basic_block(name="loop_k"))
+loop_3_exit = ir.IRBuilder(func.append_basic_block(name="loop_k_exit"))
+loop_2_exit = ir.IRBuilder(func.append_basic_block(name="loop_x_exit"))
+loop_1_exit = ir.IRBuilder(func.append_basic_block(name="loop_y_exit"))
+
+y = loop_1.phi(ir.IntType(64), name="y")
+x = loop_2.phi(ir.IntType(64), name="x")
+k = loop_3.phi(ir.IntType(64), name="k")
+
+exit = ir.IRBuilder(func.append_basic_block(name="exit"))
+
+AMX.set(loop_2)
+
+# stride
+xptr = loop_3_exit.add(x, loop_3_exit.mul(k, int_const(N)))
+yptr = loop_3_exit.add(y, loop_3_exit.mul(k, int_const(N)))
+
+# if you are okay with the wrong answer, this is faster
+#xptr = loop_3_exit.add(x, loop_3_exit.mul(k, int_const(32)))
+#yptr = loop_3_exit.add(y, loop_3_exit.mul(k, int_const(32)))
+
+# double loads load 32 floats
+AMX.ldx(loop_3_exit, loop_3_exit.add(int_const(1<<62), loop_3_exit.add(xm, loop_3_exit.mul(int_const(4), xptr))))
+AMX.ldy(loop_3_exit, loop_3_exit.add(int_const(1<<62), loop_3_exit.add(ym, loop_3_exit.mul(int_const(4), yptr))))
+
+# <Z row> <X offset> <Y offset>
+AMX.fma32(loop_3_exit, int_const(0<<20 | (0*16*4)<<10 | (0*16*4)))
+AMX.fma32(loop_3_exit, int_const(1<<20 | (1*16*4)<<10 | (0*16*4)))
+AMX.fma32(loop_3_exit, int_const(2<<20 | (0*16*4)<<10 | (1*16*4)))
+AMX.fma32(loop_3_exit, int_const(3<<20 | (1*16*4)<<10 | (1*16*4)))
+
+# store
+gptr = loop_2_exit.mul(loop_2_exit.add(loop_2.mul(y, int_const(N)), x), int_const(4))
+zmp = loop_2_exit.add(zm, gptr)
+for j in range(2):
+  for r in range(16):
+    z_row = j*2
+    ptr = ((j*16)+r)*N
+    AMX.stz(loop_2_exit, loop_2_exit.add(zmp, int_const(1 << 62 | ((r*4+z_row) << 56) | ptr*4)))
+AMX.clr(loop_2_exit)
+
+yp = loop_1_exit.add(y, int_const(32))
+xp = loop_2_exit.add(x, int_const(32))
+kp = loop_3_exit.add(k, int_const(1))
+
+y.add_incoming(int_const(0), entry._block)
+x.add_incoming(int_const(0), loop_1._block)
+k.add_incoming(int_const(0), loop_2._block)
+y.add_incoming(yp, loop_1_exit._block)
+x.add_incoming(xp, loop_2_exit._block)
+k.add_incoming(kp, loop_3_exit._block)
+
+entry.branch(loop_1._block)
+loop_1.branch(loop_2._block)
+loop_2.branch(loop_3._block)
+loop_3.branch(loop_3_exit._block)
+loop_3_exit.cbranch(loop_3_exit.icmp_unsigned("==", kp, int_const(N)), loop_2_exit._block, loop_3._block)
+loop_2_exit.cbranch(loop_2_exit.icmp_unsigned("==", xp, int_const(N)), loop_1_exit._block, loop_2._block)
+loop_1_exit.cbranch(loop_1_exit.icmp_unsigned("==", yp, int_const(N)), exit._block, loop_1._block)
+exit.ret(int_const(0))
+
+device = LLVMDevice("llvm")
+prog = LLVMProgram(device, "exec", LLVMCompiler(device).compile(str(module)))
+"""
+
+def timeit(fxn):
+  st = time.perf_counter()
+  et = fxn()
+  return time.perf_counter() - st
+
+tm = min([timeit(lambda: prog(a, b, c, N**2)) for _ in range(20)])
+MallocAllocator._copyout(flat_mv(na.data), a)
+print(f"{N*N:10d} {tm*1e6:9.2f} us, {BW*1e-9/tm:.2f} GB/s")
+
+np.testing.assert_allclose(na[:ns.shape[0]], ns, atol=1e-4, rtol=1e-4)
+
+# comp = (nb.T @ nc).T
+# np.testing.assert_allclose(na, comp, atol=1e-4, rtol=1e-5)
@@ -1,141 +0,0 @@
-#!/usr/bin/env python3
-"""Standalone correctness/throughput harness for the Hexagon HVX int8 GEMM."""
-import argparse
-
-import numpy as np
-
-from tinygrad import Device, dtypes
-from tinygrad.device import Buffer
-
-
-KERNEL = r"""
-typedef int int32x32 __attribute__((aligned(128),vector_size(128)));
-typedef unsigned char uchar4 __attribute__((aligned(4),vector_size(4)));
-typedef signed char char128 __attribute__((aligned(128),vector_size(128)));
-typedef unsigned char uchar128 __attribute__((aligned(128),vector_size(128)));
-typedef unsigned char uchar256 __attribute__((aligned(256),vector_size(256)));
-union V256 { uchar256 vec256; struct { uchar128 lo128, hi128; }; };
-
-__attribute__((noinline)) void gemm(unsigned char * restrict __attribute__((align_value(128))) out,
-                                    unsigned char * restrict __attribute__((align_value(128))) weight,
-                                    signed char * restrict __attribute__((align_value(128))) activation) {
-  for (int n = 0; n < 512; n++) {
-    int noff = n << 9;
-    for (int mb = 0; mb < 4; mb++) {
-      int moff = mb << 7;
-      int32x32 acc0 = __builtin_HEXAGON_V6_vd0_128B();
-      int32x32 acc1 = __builtin_HEXAGON_V6_vd0_128B();
-      int32x32 acc2 = __builtin_HEXAGON_V6_vd0_128B();
-      int32x32 acc3 = __builtin_HEXAGON_V6_vd0_128B();
-      for (int k4 = 0; k4 < 128; k4++) {
-        uchar4 w4 = *((uchar4 *)(weight + noff + (k4 << 2)));
-        int aoff = moff + (k4 << 11);
-        char128 x0 = *((char128 *)(activation + aoff));
-        char128 x1 = *((char128 *)(activation + aoff + 512));
-        char128 x2 = *((char128 *)(activation + aoff + 1024));
-        char128 x3 = *((char128 *)(activation + aoff + 1536));
-        union V256 s01, s23, slo, shi;
-        s01.vec256 = __builtin_HEXAGON_V6_vshufoeb_128B(x1, x0);
-        s23.vec256 = __builtin_HEXAGON_V6_vshufoeb_128B(x3, x2);
-        slo.vec256 = __builtin_HEXAGON_V6_vdealvdd_128B(s23.lo128, s01.lo128, 2);
-        shi.vec256 = __builtin_HEXAGON_V6_vdealvdd_128B(s23.hi128, s01.hi128, 2);
-        uchar128 w = __builtin_HEXAGON_V6_lvsplatw_128B(*((unsigned int *)&w4));
-        acc0 = __builtin_HEXAGON_V6_vrmpybusv_acc_128B(acc0, w, slo.lo128);
-        acc1 = __builtin_HEXAGON_V6_vrmpybusv_acc_128B(acc1, w, shi.lo128);
-        acc2 = __builtin_HEXAGON_V6_vrmpybusv_acc_128B(acc2, w, slo.hi128);
-        acc3 = __builtin_HEXAGON_V6_vrmpybusv_acc_128B(acc3, w, shi.hi128);
-      }
-      acc0 /= 1000; acc1 /= 1000; acc2 /= 1000; acc3 /= 1000;
-      uchar128 packed = __builtin_HEXAGON_V6_vpackhub_sat_128B(
-        __builtin_HEXAGON_V6_vpackwh_sat_128B(acc3, acc2),
-        __builtin_HEXAGON_V6_vpackwh_sat_128B(acc1, acc0));
-      packed = __builtin_HEXAGON_V6_vshuffb_128B(packed);
-      packed = __builtin_HEXAGON_V6_vshuffb_128B(packed);
-      *((uchar128 *)(out + noff + moff)) = packed;
-    }
-  }
-}
-
-struct dcvs_v2_req { int type; int pad; _Bool dcvs_enable; char dcvs_option; _Bool set_latency; int latency;
-  _Bool set_dcvs_params; short pad2; char target_corner; char min_corner; char max_corner; int pad3[3]; };
-typedef union { struct { void *pv; unsigned int len; } buf; struct { int fd; unsigned int offset; } dma; } remote_arg;
-int HAP_power_set(void *, void *);
-void *HAP_mmap(void *, int, int, int, int, long);
-int HAP_munmap(void *, int);
-unsigned long long HAP_perf_get_time_us(void);
-
-int entry(unsigned long long handle, unsigned int sc, remote_arg *pra) {
-  struct dcvs_v2_req req = {.type=7, .dcvs_enable=0, .set_latency=1, .latency=100,
-    .set_dcvs_params=1, .target_corner=6};
-  HAP_power_set((void *)handle, (void *)&req);
-  if ((sc >> 24) != 2) return 0;
-  int *sizes = (int *)pra[0].buf.pv, *offs = (int *)pra[1].buf.pv;
-  void *out = HAP_mmap(0, sizes[0], 3, 0, pra[3].dma.fd, 0) + offs[0];
-  void *weight = HAP_mmap(0, sizes[1], 3, 0, pra[4].dma.fd, 0) + offs[1];
-  void *activation = HAP_mmap(0, sizes[2], 3, 0, pra[5].dma.fd, 0) + offs[2];
-  unsigned long long start = HAP_perf_get_time_us();
-  gemm(out, weight, activation);
-  *(unsigned long long *)pra[2].buf.pv = HAP_perf_get_time_us() - start;
-  HAP_munmap(out-offs[0], sizes[0]); HAP_munmap(weight-offs[1], sizes[1]); HAP_munmap(activation-offs[2], sizes[2]);
-  return 0;
-}
-"""
-
-
-def main() -> None:
-  parser = argparse.ArgumentParser()
-  parser.add_argument("--iters", type=int, default=5)
-  parser.add_argument("--check", action="store_true")
-  parser.add_argument("--raw", action="store_true", help="store raw int32 accumulators without requantization")
-  args = parser.parse_args()
-  dev = Device["DSP"]
-  source = KERNEL
-  if args.raw:
-    source = source.replace(
-      "unsigned char * restrict __attribute__((align_value(128))) out,\n                                    unsigned char * restrict",
-      "int * restrict __attribute__((align_value(128))) out,\n                                    unsigned char * restrict", 1)
-    old = """      acc0 /= 1000; acc1 /= 1000; acc2 /= 1000; acc3 /= 1000;
-      uchar128 packed = __builtin_HEXAGON_V6_vpackhub_sat_128B(
-        __builtin_HEXAGON_V6_vpackwh_sat_128B(acc3, acc2),
-        __builtin_HEXAGON_V6_vpackwh_sat_128B(acc1, acc0));
-      packed = __builtin_HEXAGON_V6_vshuffb_128B(packed);
-      packed = __builtin_HEXAGON_V6_vshuffb_128B(packed);
-      *((uchar128 *)(out + noff + moff)) = packed;"""
-    new = """      int base = noff + moff;
-      *((int32x32 *)(out + base + 0)) = acc0;
-      *((int32x32 *)(out + base + 32)) = acc1;
-      *((int32x32 *)(out + base + 64)) = acc2;
-      *((int32x32 *)(out + base + 96)) = acc3;"""
-    if old not in source: raise RuntimeError("raw-kernel source pattern not found")
-    source = source.replace(old, new)
-  lib = dev.compiler.compile(source)
-  prg = dev.runtime("entry", lib)
-  rng = np.random.default_rng(0)
-  # Kernel contract is weight[N,K] and activation[K,M], both contiguous.
-  weight_np = rng.integers(0, 16, (512, 512), dtype=np.uint8)
-  activation_np = rng.integers(-8, 8, (512, 512), dtype=np.int8)
-  out_dtype = dtypes.int if args.raw else dtypes.uint8
-  bufs = [Buffer("DSP", 512*512, dt, preallocate=True) for dt in (out_dtype, dtypes.uint8, dtypes.int8)]
-  bufs[1].copyin(memoryview(weight_np).cast("B"))
-  bufs[2].copyin(memoryview(activation_np).cast("B"))
-  for _ in range(2): prg(*(x._buf for x in bufs), wait=True)
-  times = [prg(*(x._buf for x in bufs), wait=True) for _ in range(args.iters)]
-  best = min(times)
-  print(f"{2*512**3/best/1e9:.1f} GOPS ({best*1e3:.3f} ms)")
-  if args.check:
-    raw = bytearray(bufs[0].nbytes)
-    bufs[0].copyout(memoryview(raw))
-    expected_dot = weight_np.astype(np.int32) @ activation_np.astype(np.int32)
-    if args.raw:
-      got = np.frombuffer(raw, dtype=np.int32).reshape(512, 4, 4, 32).transpose(0, 1, 3, 2).reshape(512, 512)
-      expected = expected_dot
-      delta = np.abs(got.astype(np.int64)-expected.astype(np.int64))
-    else:
-      got = np.frombuffer(raw, dtype=np.uint8).reshape(512, 512)
-      expected = (expected_dot // 1000).clip(0, 255).astype(np.uint8)
-      delta = np.abs(got.astype(np.int16)-expected.astype(np.int16))
-    print(f"check={np.array_equal(got, expected)} max_abs={delta.max()} mismatches={np.count_nonzero(delta)}")
-    if not np.array_equal(got, expected): raise SystemExit(1)
-
-
-if __name__ == "__main__": main()
@@ -0,0 +1,43 @@
+#!/usr/bin/env python3
+import numpy as np
+from tinygrad.runtime.ops_cl import CLProgram, CLCompiler
+from tinygrad import Device, dtypes
+from tinygrad.device import Buffer
+from hexdump import hexdump
+
+# https://github.com/intel/intel-graphics-compiler/blob/master/documentation/visa/instructions/DPAS.md
+# https://registry.khronos.org/OpenCL/extensions/intel/cl_intel_subgroups.html
+# https://registry.khronos.org/OpenCL/extensions/intel/cl_intel_subgroup_matrix_multiply_accumulate.html
+# https://registry.khronos.org/OpenCL/extensions/intel/cl_intel_subgroup_split_matrix_multiply_accumulate.html
+# https://hc34.hotchips.org/assets/program/conference/day1/GPU%20HPC/Intel_s%20Ponte%20Vecchio%20GPU%20-%20Architecture%20Systems%20and%20Software%20FINAL.pdf
+
+device = Device["CL"]
+
+# NOTE: only the subgroup type 8 ones work
+prog = CLProgram(device, "test", CLCompiler(device, "test").compile(f"""
+__attribute__((intel_reqd_sub_group_size(8)))
+__kernel void test(__global float* data0, const __global int* data1, const __global int8* data2) {{
+  int lidx0 = get_local_id(0);
+  int a = data1[lidx0];
+  int8 b = data2[lidx0];
+  float out = intel_sub_group_f16_f16_matrix_mad_k16(a, b, 0.0f);
+  data0[lidx0] = out;
+}}
+"""))
+#with open("/tmp/test.elf", "wb") as f: f.write(prog.lib)
+
+a = Buffer("CL", 8, dtypes.float32).allocate()
+b = Buffer("CL", 0x10, dtypes.float16).allocate()
+c = Buffer("CL", 8*0x10, dtypes.float16).allocate()
+
+row = np.array([1,2,3,4,5,6,7,8,1,2,3,4,5,6,7,8], np.float16)
+mat = np.random.random((8, 0x10)).astype(np.float16)
+
+b.copyin(row.data)
+c.copyin(mat.data)
+ret = prog(a._buf, b._buf, c._buf, global_size=[1,1,1], local_size=[8,1,1], wait=True)
+print(ret)
+out = np.frombuffer(a.as_memoryview(), np.float32)
+real = row.astype(np.float32)@mat.T.astype(np.float32)
+print("out:", out)
+print("real", real)
@@ -1,440 +0,0 @@
-"""ir3 assembler for Adreno a6xx (A630).
-
-Constructs complete QCOM shader binaries from instruction listings.
-Uses a compiled "donor" kernel for the binary envelope (header, metadata,
-buffer descriptors, sampler info) and replaces the shader instructions
-and register counts.
-
-Encoding reference: derived from Mesa ir3 disassembly of known-good shaders.
-Instruction format: 64 bits (8 bytes) stored as two little-endian 32-bit words.
-"""
-import struct
-
-# ============================================================
-# HELPERS
-# ============================================================
-
-def _hreg(name):
-    """Parse 'hr3.z' -> half-register number 14."""
-    if isinstance(name, int): return name
-    r, c = name.replace('hr','').replace('r','').split('.')
-    return int(r) * 4 + 'xyzw'.index(c)
-
-def _freg(name):
-    """Parse 'r3.z' -> full-register number 14."""
-    if isinstance(name, int): return name
-    r, c = name.replace('r','').split('.')
-    return int(r) * 4 + 'xyzw'.index(c)
-
-def _pack(lo, hi):
-    return struct.pack('<II', lo & 0xFFFFFFFF, hi & 0xFFFFFFFF)
-
-# ============================================================
-# CAT0: FLOW CONTROL
-# ============================================================
-
-def NOP(rpt=0):
-    """(rptN)nop"""
-    return _pack(0, (rpt & 0x7F) << 8)
-
-def NOP_SS(rpt=0):
-    """(ss)(rptN)nop -- wait until prior instructions have consumed their sources."""
-    return _pack(0, 0x1000 | ((rpt & 0x7F) << 8))
-
-def END():
-    """end"""
-    return _pack(0, 0x03000000)
-
-def BR(offset, inv=True):
-    """br !p0.x, #offset  (inv=True means branch when predicate is FALSE)
-    offset is signed, relative to the branch instruction."""
-    return struct.pack('<iI', offset, 0x00900000 if inv else 0x00800000)
-
-def JUMP(offset):
-    """jump #offset. Offset is signed, relative to the jump instruction."""
-    return struct.pack('<iI', offset, 0x01000000)
-
-# ============================================================
-# CAT1: MOVE / CONVERT
-# ============================================================
-
-def MOV_S32(dst, imm, sy=False):
-    """(sy?)mov.s32s32 rDST, #imm"""
-    return _pack(imm, ((0x30 if sy else 0x20) << 24) | (0x55 << 16) | (0x40 << 8) | (_freg(dst) & 0xFF))
-
-def MOV_F32(dst, src, rpt=0, sy=False, ss=False, r=False):
-    """(sy?)(ss?)(rptN?)mov.f32f32 rDST, (r?)rSRC"""
-    return _pack(_freg(src), (0x30044000 if sy else 0x20044000) | (0x1000 if ss else 0) |
-                 (0x800 if r else 0) | ((rpt & 0x7F) << 8) | (_freg(dst) & 0xFF))
-
-def MOV_H(dst, src, rpt=0, r=False):
-    """(rptN?)mov.f16f16 hrDST, (r?)hrSRC."""
-    return _pack(_hreg(src), 0x20000000 | (0x800 if r else 0) | ((rpt & 0x7F) << 8) | (_hreg(dst) & 0xFF))
-
-def MOV_H_IMM(dst, imm_u16=0, rpt=0):
-    """(rptN?)mov.f16f16 hrDST, h(imm) -- imm is raw fp16 bits (0=zero, 0x3c00=1.0)."""
-    return _pack(imm_u16, 0x20400000 | ((rpt & 0x7F) << 8) | (_hreg(dst) & 0xFF))
-
-def COV_F16F32(dst, src, sy=False, rpt=0, r=False):
-    """(sy?)(rptN?)cov.f16f32 rDST, (r?)hrSRC"""
-    return _pack(_hreg(src), ((0x30 if sy else 0x20) << 24) | 0x004000 | (0x800 if r else 0) |
-                 ((rpt & 0x7f) << 8) | (_freg(dst) & 0xFF))
-
-# ============================================================
-# CAT2: INTEGER / FLOAT ALU (2 operands)
-# ============================================================
-
-def ADD_S(dst, src1, imm, nop=0, ss=False):
-    """(ss?)(nopN?)add.s rDST, rSRC1, #imm  (signed immediate add)"""
-    d, s = _freg(dst), _freg(src1)
-    hi_base = 0x42300000 | (d & 0xFF)
-    if nop > 0:
-        hi_base = (hi_base & 0xFF00FFFF) | (0x38 << 16) | ((nop & 0x7) << 11)
-    if ss: hi_base |= 0x1000
-    lo = ((0x27 if imm < 0 else 0x20) << 24) | ((imm & 0xFF) << 16) | (s & 0xFF)
-    return _pack(lo, hi_base)
-
-def ADD_S_REG(dst, src1, src2, nop=0):
-    """(nopN?)add.s rDST, rSRC1, rSRC2"""
-    d, s1, s2 = _freg(dst), _freg(src1), _freg(src2)
-    hi_base = 0x42300000 | (d & 0xFF)
-    if nop > 0:
-        hi_base = (hi_base & 0xFF00FFFF) | (0x38 << 16) | ((nop & 0x7) << 11)
-    return _pack(((s2 & 0xFF) << 16) | (s1 & 0xFF), hi_base)
-
-def ADD_S_CONST_REG(dst, const_src, src2, nop=0):
-    """(nopN?)add.s rDST, cSRC1, rSRC2"""
-    d, c1, s2 = _freg(dst), _freg(const_src.replace('c', 'r', 1)), _freg(src2)
-    hi_base = 0x42300000 | (d & 0xFF)
-    if nop > 0:
-        hi_base = (hi_base & 0xFF00FFFF) | (0x38 << 16) | ((nop & 0x7) << 11)
-    return _pack(((s2 & 0xFF) << 16) | 0x1000 | (c1 & 0xFF), hi_base)
-
-def ADD_F(dst, src1, src2, rpt=0, r1=False, r2=False, sy=False):
-    """Vector-capable add.f; full registers use the same scalar indices."""
-    hi = (0x50100000 if sy else 0x40100000) | (0x800 if r1 else 0) | (0x80000 if r2 else 0)
-    return _pack(((_hreg(src2) & 0xFF) << 16) | (_hreg(src1) & 0xFF),
-                 hi | ((rpt & 0x7f) << 8) | (_hreg(dst) & 0xFF))
-
-def SUB_F(dst, src1, src2, rpt=0, r1=False, r2=False, sy=False):
-    """Vector-capable add.f with a negated second source."""
-    hi = (0x50100000 if sy else 0x40100000) | (0x800 if r1 else 0) | (0x80000 if r2 else 0)
-    return _pack(0x40000000 | ((_hreg(src2) & 0xFF) << 16) | (_hreg(src1) & 0xFF),
-                 hi | ((rpt & 0x7f) << 8) | (_hreg(dst) & 0xFF))
-
-def ADD_U(dst, src1_const, src2):
-    """add.u rDST, cSRC1, rSRC2  -- src1 is constant register"""
-    # From: 42100008_00031050 = add.u r2.x, c20.x, r0.w
-    return _pack((_freg(src2) << 16) | 0x1050, 0x42100000 | (_freg(dst) & 0xFF))
-
-def CMPS_S_EQ(src1, imm, nop=0):
-    """(nopN?)cmps.s.eq p0.x, rSRC1, #imm"""
-    hi = 0x42b400f8
-    if nop > 0:
-        hi = (hi & 0xFF00FFFF) | (0xb4 << 16) | ((nop & 0x7) << 11)
-    # Integer immediates use the low bits of the source descriptor for bits 8+.
-    # Keeping this fixed at 0x20 silently truncated loop bounds above 255.
-    lo = ((0x20 | (imm >> 8)) << 24) | ((imm & 0xFF) << 16) | (_freg(src1) & 0xFF)
-    return _pack(lo, hi)
-
-def CMPS_S_LT_REG(src1, src2, nop=0):
-    """(nopN?)cmps.s.lt p0.x, rSRC1, rSRC2"""
-    hi = 0x42b000f8
-    if nop > 0: hi = (hi & 0xFF00FFFF) | (0xb0 << 16) | ((nop & 0x7) << 11)
-    return _pack(((_freg(src2) & 0xff) << 16) | (_freg(src1) & 0xff), hi)
-
-def SHL_B(dst, src, imm, jp=False, ss=False, nop=0):
-    """(ss?)(jp?)(nopN?)shl.b rDST, rSRC, #imm"""
-    hi = (0x4ed00000 if jp else 0x46d00000) | (_freg(dst) & 0xFF)
-    if ss: hi |= 1 << 12
-    if nop & 1: hi |= 1 << 11
-    if nop & 2: hi |= 1 << 19
-    return _pack((0x20 << 24) | ((imm & 0xFF) << 16) | (_freg(src) & 0xFF), hi)
-
-def SHR_B(dst, src, imm):
-    """shr.b rDST, rSRC, #imm"""
-    return _pack((0x20 << 24) | ((imm & 0xFF) << 16) | (_freg(src) & 0xFF), 0x46f00000 | (_freg(dst) & 0xFF))
-
-def AND_B(dst, src, imm, nop=0):
-    """(nopN?)and.b rDST, rSRC, #imm"""
-    hi = 0x43900000 | (_freg(dst) & 0xFF)
-    if nop & 1: hi |= 1 << 11
-    if nop & 2: hi |= 1 << 19
-    return _pack((0x20 << 24) | ((imm & 0xFF) << 16) | (_freg(src) & 0xFF), hi)
-
-def AND_B_CONST(dst, src, const_src, nop=0):
-    """(nopN?)and.b rDST, rSRC, cSRC2"""
-    d, s, c = _freg(dst), _freg(const_src.replace('c', 'r', 1)), _freg(src)
-    hi = 0x43900000 | (d & 0xFF)
-    if nop & 1: hi |= 1 << 11
-    if nop & 2: hi |= 1 << 19
-    return _pack((0x10 << 24) | ((c & 0xFF) << 16) | (s & 0xFF), hi)
-
-def OR_B(dst, src, imm, ss=False):
-    """(ss?)or.b rDST, rSRC, #imm"""
-    return _pack((0x20 << 24) | ((imm & 0xFF) << 16) | (_freg(src) & 0xFF),
-                 0x43b00000 | (0x1000 if ss else 0) | (_freg(dst) & 0xFF))
-
-def CMPS_U_LT(dst, src1, src2_const):
-    """cmps.u.lt rDST, rSRC1, cSRC2"""
-    # From: 42900010_10500008 = cmps.u.lt r4.x, r2.x, c20.x
-    return _pack(0x10500000 | (_freg(src1) & 0xFF), 0x42900000 | (_freg(dst) & 0xFF))
-
-def CMPS_U_LT_REG(dst, src1, src2, sy=False):
-    """(sy?)cmps.u.lt rDST, rSRC1, rSRC2"""
-    hi = (0x52900000 if sy else 0x42900000) | (_freg(dst) & 0xff)
-    return _pack(((_freg(src2) & 0xff) << 16) | (_freg(src1) & 0xff), hi)
-
-# ============================================================
-# CAT3: MAD (3 operands)
-# ============================================================
-
-def MAD_F16(dst, src1, src2, src3, rpt=0, sy=False, r=False, r1=False, r3=False):
-    """(sy?)(rptN?)mad.f16 hrDST, (r1?)hrSRC1, (r?)hrSRC2, (r?)hrSRC3
-    When rpt>0, r1 auto-increments src1 and r auto-increments src2/src3/dst."""
-    d, s1, s2, s3 = _hreg(dst), _hreg(src1), _hreg(src2), _hreg(src3)
-    hi = ((0x73 if sy else 0x63) << 24) | ((s2 >> 1) << 16) | ((((s2 & 1) << 7) | (0x08 if r1 else 0) | (rpt & 0x7F)) << 8) | (d & 0xFF)
-    lo = (0x20000000 if (r or r3) else 0) | ((s3 & 0xFF) << 16) | (0x8000 if r else 0) | (s1 & 0xFF)
-    return _pack(lo, hi)
-
-def MAD_F32(dst, src1, src2, src3, rpt=0, sy=False, r=False, r1=False):
-    """(sy?)(rptN?)mad.f32 rDST, rSRC1, (r?)rSRC2, (r?)rSRC3"""
-    d, s1, s2, s3 = _freg(dst), _freg(src1), _freg(src2), _freg(src3)
-    hi = ((0x73 if sy else 0x63) << 24) | (0x80 << 16) | ((s2 >> 1) << 16) | \
-         ((((s2 & 1) << 7) | (0x08 if r1 else 0) | (rpt & 0x7F)) << 8) | (d & 0xFF)
-    lo = (0x20000000 if r else 0) | ((s3 & 0xFF) << 16) | (0x8000 if r else 0) | (s1 & 0xFF)
-    return _pack(lo, hi)
-
-def DP4ACC(dst, src1, src2, src3, sy=False, mixed=False, signed=None):
-    """A6xx packed 4x int8 dot product accumulated into a full int32 register.
-
-    ``mixed=False`` selects unsigned*unsigned. ``mixed=True`` selects the
-    pre-A7xx mixed signedness mode used by A630 (signed lhs, unsigned rhs).
-    The instruction has no repeat form on this generation.
-    """
-    if signed is not None: mixed = signed
-    d, s1, s2, s3 = _freg(dst), _freg(src1), _freg(src2), _freg(src3)
-    hi = ((0x76 if sy else 0x66) << 24) | (0x80 << 16) | ((s2 >> 1) << 16)
-    hi |= (((s2 & 1) << 7) | 0x40) << 8
-    # AL-OP is bit 13 and the pre-A7 signed/unsigned selector is bit 14.
-    lo = ((s3 & 0xff) << 16) | 0x2000 | (0x4000 if mixed else 0) | (s1 & 0xff)
-    return _pack(lo, hi | (d & 0xff))
-
-# ============================================================
-# CAT3: SHLG / SHRM (shift with merge)
-# ============================================================
-
-def SHLG(dst, imm, src1, src2, nop=0):
-    """(nopN?)shlg rDST, #imm, rSRC1, rSRC2.
-
-    This covers the packed image-coordinate forms emitted by the a6xx compiler
-    for GEMM kernels. The low byte encodes the shift immediate and bits 23:16
-    encode src2; the remaining source mode bits are pattern-specific.
-    """
-    d, s1, s2 = _freg(dst), _freg(src1), _freg(src2)
-    if (s1, s2) in ((_freg('r0.y'), _freg('r0.z')), (_freg('r0.z'), _freg('r0.x'))):
-        hi_mid, lo_mid = 0x80, 0xb0
-        if (s1, s2) == (_freg('r0.z'), _freg('r0.x')): hi_mid = 0x81
-    elif (s1, s2) in ((_freg('r0.w'), _freg('r0.x')), (_freg('r0.w'), _freg('r0.y'))):
-        hi_mid, lo_mid = 0x81, 0x30
-    else:
-        raise ValueError('unsupported SHLG source pattern %s, %s' % (src1, src2))
-    hi = (0x65 << 24) | (hi_mid << 16) | (0x84 << 8) | (d & 0xFF)
-    lo = ((s2 & 0xFF) << 16) | (lo_mid << 8) | (imm & 0xFF)
-    return _pack(lo, hi)
-
-def SHLG_IMM(dst, imm, src, merge):
-    """shlg rDST, #imm, rSRC, #merge.
-
-    Observed in compiler address generation for widened column stores, e.g.
-    65b08402_10803002 = shlg r0.z, 2, r24.y, 128.
-    """
-    d, s = _freg(dst), _freg(src)
-    hi = (0x65 << 24) | ((0x80 | ((s >> 1) & 0x7f)) << 16) | (0x84 << 8) | (d & 0xff)
-    lo = (0x10 << 24) | ((merge & 0xffff) << 16) | 0x3000 | (imm & 0xff)
-    return _pack(lo, hi)
-
-def SHRM(dst, shift, src1, merge):
-    """shrm rDST, #shift, rSRC1, #merge.
-
-    Observed compiler form for subgroup row offsets, e.g.
-    64000402_100c3003 = shrm r0.z, 3, r0.x, 12.
-    """
-    d, s1 = _freg(dst), _freg(src1)
-    if s1 != _freg('r0.x'):
-        raise ValueError('unsupported SHRM source %s' % src1)
-    hi = 0x64000400 | (d & 0xFF)
-    lo = (0x10 << 24) | ((merge & 0xFF) << 16) | 0x3000 | (shift & 0xFF)
-    return _pack(lo, hi)
-
-# ============================================================
-# CAT5: TEXTURE (ISAM)
-# ============================================================
-
-def ISAM_F16(dst, coord, tex=0, samp=0, sy=False, wrmask=0xf):
-    """isam.1d (f16)(xyzw) hrDST, rCOORD, s#SAMP, t#TEX
-    dst: first half-register of the xyzw quad
-    coord: full-register containing the (int2) coordinate pair"""
-    return _pack((tex * 2) << 24 | ((samp & 0x7) << 21) | (_freg(coord) * 2 + 1),
-                 (0xb0000000 if sy else 0xa0000000) | ((wrmask & 0xf) << 8) | (_hreg(dst) & 0xFF))
-
-def ISAM_F32(dst, coord, tex=0, samp=0):
-    """isam.1d (f32)(xyzw) rDST, rCOORD, s#SAMP, t#TEX"""
-    return _pack((tex * 2) << 24 | ((samp & 0x7) << 21) | (_freg(coord) * 2 + 1), 0xa0001f00 | (_freg(dst) & 0xFF))
-
-def ISAM_U32(dst, coord, tex=0, samp=0):
-    """isam.1d (u32)(xyzw) rDST, rCOORD, s#SAMP, t#TEX"""
-    return _pack((tex * 2) << 24 | ((samp & 0x7) << 21) | (_freg(coord) * 2 + 1), 0xa0003f00 | (_freg(dst) & 0xFF))
-
-def COV_S32S16(dst, src, rpt=0, r=False, sy=False):
-    """cov.s32s16 hDST, rSRC, optionally repeating over four packed lanes."""
-    hi = (0x30150000 if sy else 0x20150000) | ((rpt & 0x7) << 8) | (0x800 if r else 0) | (_hreg(dst) & 0xff)
-    return _pack(_freg(src) & 0xff, hi)
-
-def SHRG_H(dst, src, shift=16, rpt=0, r=False):
-    """shrg hDST, #shift, rSRC, #0 for extracting packed high half lanes."""
-    s = _freg(src)
-    hi = 0x65004400 | (((s >> 1) & 0x7f) << 16) | ((rpt & 0x7) << 8) | (_hreg(dst) & 0xff)
-    lo = 0x10003000 | (0x8000 if r else 0) | (shift & 0xff)
-    return _pack(lo, hi)
-
-def QUAD_BRCST(dst, src, idx, typ=3, wrmask=1, sy=False, jp=False):
-    """quad_shuffle.brcst.{typ} DST, SRC, IDX"""
-    half = typ in (0, 2, 4, 6)
-    d = _hreg(dst) if half else _freg(dst)
-    s = _hreg(src) if half else _freg(src)
-    i = _hreg(idx) if half else _freg(idx)
-    lo = (0 if half else 1) | ((s & 0xff) << 1) | ((i & 0xff) << 9)
-    hi = 0xa7e00000 | ((typ & 7) << 12) | ((wrmask & 0xf) << 8) | (d & 0xff)
-    if jp: hi |= 1 << 27
-    if sy: hi |= 1 << 28
-    return _pack(lo, hi)
-
-# ============================================================
-# CAT6: LOAD / STORE
-# ============================================================
-
-def STG_F16(addr, data_hreg, count=4, sy=False):
-    """(sy?)stg.f16 g[rADDR], hrDATA, count"""
-    # Encoding from compiled kernels:
-    # c0c01100_04800000 = stg.f16 g[r2.x], hr0.x, 4
-    # c0c01500_04800008 = stg.f16 g[r2.z], hr1.x, 4
-    # c0c01900_04800010 = stg.f16 g[r3.x], hr2.x, 4
-    # c0c01d00_04800018 = stg.f16 g[r3.z], hr3.x, 4
-    # hi pattern: c0c0XX00 where XX encodes the address register
-    # lo pattern: 048000YY where YY encodes the data register
-    a, d = _freg(addr), _hreg(data_hreg)
-    # addr encoding: r2.x=8 -> 0x11, r2.z=10 -> 0x15, r3.x=12 -> 0x19, r3.z=14 -> 0x1d
-    # Pattern: (addr * 2 + 1) = 17,21,25,29 = 0x11,0x15,0x19,0x1d
-    addr_enc = a * 2 + 1
-    hi = (0xd0c00000 if sy else 0xc0c00000) | (addr_enc << 8)
-    lo = 0x04800000 | ((d * 2) & 0xFF)
-    return _pack(lo, hi)
-
-def STG_U32(addr, data_reg, count=1, sy=False):
-    """(sy?)stg.u32 g[rADDR], rDATA, count"""
-    a, d = _freg(addr), _freg(data_reg)
-    hi = (0xd0c00000 if sy else 0xc0c00000) | (3 << 17) | ((a * 2 + 1) << 8)
-    lo = ((count & 0x7) << 24) | 0x00800000 | ((d << 1) & 0x1FE)
-    return _pack(lo, hi)
-
-def STG_F32(addr, data_reg, count=4, sy=False):
-    """(sy?)stg.f32 g[rADDR], rDATA, count"""
-    a, d = _freg(addr), _freg(data_reg)
-    hi = (0xd0c00000 if sy else 0xc0c00000) | (1 << 17) | ((a * 2 + 1) << 8)
-    lo = ((count & 0x7) << 24) | 0x00800000 | ((d << 1) & 0x1FE)
-    return _pack(lo, hi)
-
-def STIB_F32(data_reg, coord_reg, sy=False):
-    """Typed 2D image store of float4 data to integer (x,y) coordinates."""
-    hi = (0xd0220000 if sy else 0xc0220000) | (_freg(data_reg) & 0xff)
-    lo = ((_freg(coord_reg) & 0xff) << 24) | 0x00677a00
-    return _pack(lo, hi)
-
-def GETFIBERID(dst):
-    """getfiberid.u32 rDST"""
-    return _pack(0x00c98000, 0xc0260000 | (_freg(dst) & 0xff))
-
-def SHFL(dst, src, idx, mode=7, typ=2, sy=False, jp=False):
-    """shfl.{mode}.{typ} DST, SRC, IDX
-
-    mode: xor=1, up=2, down=3, rup=6, rdown=7.
-    typ: f16=0, f32=1, u16=2, u32=3, s16=4, s32=5.
-    idx can be an immediate int or a full register. For half types, dst/src are
-    half-register indices; SRC2 is always a full register/immediate per Mesa.
-    """
-    d = _hreg(dst) if typ in (0, 2, 4, 6) else _freg(dst)
-    s = _hreg(src) if typ in (0, 2, 4, 6) else _freg(src)
-    if isinstance(idx, int):
-        idx_im, idx_bits = 1, idx & 0xff
-    else:
-        idx_im, idx_bits = 0, _freg(idx) & 0xff
-    lo = ((s & 0xff) << 1) | (idx_im << 23) | (idx_bits << 24)
-    hi = (0xc0000000 | (0x1b << 22) | (2 << 20) | ((typ & 7) << 17) |
-          ((mode & 7) << 13) | (d & 0xff))
-    if jp: hi |= 1 << 27
-    if sy: hi |= 1 << 28
-    return _pack(lo, hi)
-
-# ============================================================
-# CAT3 SPECIAL: SAD.S32
-# ============================================================
-
-def SAD_S32(dst, src1_const, src2, src3, nop=0):
-    """(nopN?)sad.s32 rDST, cSRC1, (neg)rSRC2, rSRC3"""
-    # From: 67888009_40101051 = sad.s32 r2.y, c20.y, (neg)r4.y, r4.x
-    d, s2, s3 = _freg(dst), _freg(src2), _freg(src3)
-    hi_src2 = 0x80 | ((s2 >> 1) & 0xF)
-    # Observed nop3 form uses 0x88 in the third byte; plain sad.s32 uses 0x80.
-    hi_nop = 0x88 if nop > 0 else 0x80
-    hi = (0x67 << 24) | (hi_src2 << 16) | (hi_nop << 8) | (d & 0xFF)
-    lo = 0x40000000 | (s3 << 16) | 0x1051
-    return _pack(lo, hi)
-
-# ============================================================
-# BINARY ENVELOPE
-# ============================================================
-
-def get_envelope(dev, src):
-    """Compile an OpenCL kernel and return the binary as a mutable envelope."""
-    lib = bytearray(dev.compiler.compile_cached(src))
-    img_off = struct.unpack_from('<I', lib, 0xc0)[0]
-    img_sz = struct.unpack_from('<I', lib, 0x100)[0]
-    reg_off = struct.unpack_from('<I', lib, 0x34)[0]
-    return lib, img_off, img_sz, reg_off
-
-def inject(lib, img_off, img_sz, reg_off, shader_bytes, fregs, hregs, mergedregs=None):
-    """Replace shader binary and register counts in the envelope."""
-    lib = bytearray(lib)
-    shader = bytearray(shader_bytes)
-    if len(shader) > img_sz:
-        raise ValueError(f"shader is {len(shader)} bytes but donor image is only {img_sz} bytes")
-    # Pad to original size
-    while len(shader) < img_sz:
-        shader += NOP()
-    lib[img_off:img_off+img_sz] = shader[:img_sz]
-    if mergedregs is True: fregs |= 1 << 31
-    if mergedregs is False: hregs |= 1 << 31
-    struct.pack_into('<I', lib, reg_off + 0x14, fregs)
-    struct.pack_into('<I', lib, reg_off + 0x18, hregs)
-    return bytes(lib)
-
-def disasm(shader_bytes, gpu_id=630):
-    """Disassemble shader binary using Mesa's ir3_isa_disasm."""
-    import ctypes, tempfile
-    from tinygrad.runtime.autogen import mesa
-    from tinygrad.helpers import data64
-    with tempfile.TemporaryFile('w+', buffering=1) as tf:
-        @ctypes.CFUNCTYPE(None, ctypes.c_void_p, ctypes.c_uint32, ctypes.c_void_p)
-        def hd(data, n, instr):
-            fst, snd = data64(ctypes.cast(instr, ctypes.POINTER(ctypes.c_uint64)).contents.value)
-            print(f"{n:04} [{fst:08x}_{snd:08x}] ", end="", flush=True, file=tf)
-        libc = ctypes.CDLL(None)
-        libc.setlinebuf(fp:=ctypes.cast(libc.fdopen(tf.fileno(), b"w"), ctypes.POINTER(mesa.struct__IO_FILE)))
-        mesa.ir3_isa_disasm(bytes(shader_bytes), len(shader_bytes), fp, mesa.struct_isa_decode_options(gpu_id, True, 0, True, pre_instr_cb=hd))
-        tf.seek(0)
-        return tf.read()
-
-def assemble(instr_list):
-    """Assemble a list of instruction bytes into a shader binary."""
-    return b''.join(instr_list)
@@ -1,6 +1,7 @@
 import numpy as np, os
 from tinygrad.helpers import getenv, flat_mv
 from tinygrad import dtypes
+from tinygrad.engine.realize import get_program

 # for copied uops
 from tinygrad import dtypes
@@ -20,8 +20,8 @@ def hand_spec_tc_cores():

  gk = UOp.range(N // 8, 0, AxisType.REDUCE)

-  a_tc = UOp.stack(*[mat_idx(a, gx, gk, warp, i) for i in range(2)])
-  b_tc = UOp.stack(*[mat_idx(b, gk, gy, warp, i) for i in range(2)])
+  a_tc = UOp.vectorize(*[mat_idx(a, gx, gk, warp, i) for i in range(2)])
+  b_tc = UOp.vectorize(*[mat_idx(b, gk, gy, warp, i) for i in range(2)])

  acc = UOp.placeholder((2,), dtypes.float, slot=0, addrspace=AddrSpace.REG)
  acc = acc[0].set(0.0)
@@ -30,10 +30,10 @@ def hand_spec_tc_cores():
  # TODO: make this simple
  wmma_arg = ('WMMA_8_8_8_float_float', (8, 8, 8), dtypes.float, dtypes.float, 'METAL', 32, (((3, 2),), ((3, 2),), ((3, 2),)), ())

-  acc_load = UOp.stack(acc.after(gk)[0], acc.after(gk)[1])
-  out = UOp(Ops.WMMA, dtypes.float, (a_tc, b_tc, acc_load), arg=wmma_arg)
+  acc_load = UOp.vectorize(acc.after(gk)[0], acc.after(gk)[1])
+  out = UOp(Ops.WMMA, dtypes.float.vec(2), (a_tc, b_tc, acc_load), arg=wmma_arg)

-  end_loop = UOp.group(*[acc[i].store(out.index(i)) for i in range(2)]).end(gk)
+  end_loop = UOp.group(*[acc[i].store(out.gep(i)) for i in range(2)]).end(gk)

  sink = UOp.group(*[mat_idx(c.after(end_loop), gx, gy, warp, i).store(acc[i]) for i in range(2)])
  return sink.sink(arg=KernelInfo(name="custom_metal_matmul", opts_to_apply=())).simplify()
@@ -77,9 +77,9 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
  B = B.reshape((K//BLOCK_K, BLOCK_K, N//BLOCK_N, BLOCK_N))

  # this is the big accumulator
-  acc = UOp.placeholder((BLOCK_N//TC_N, BLOCK_M//TC_M//WARPGROUP_SIZE), dtypes.float, 0, AddrSpace.REG)
+  acc = UOp.placeholder((BLOCK_N//TC_N, BLOCK_M//TC_M//WARPGROUP_SIZE), dtypes.float.vec(4), 0, AddrSpace.REG)
  assert acc.size*WARP_SIZE*WARPGROUP_SIZE*4 == BLOCK_M*BLOCK_N
-  acc = acc[init_l:=UOp.range(acc.size, 500)].set(UOp.const(dtypes.float, (0.0,)*4), end=init_l)
+  acc = acc[init_l:=UOp.range(acc.size, 500)].set(UOp.const(dtypes.float.vec(4), 0.0), end=init_l)

  # create locals (note A is permuted, and the stride is changed to avoid bank conflicts)
  def make_locals(slot) -> tuple[UOp, UOp]:
@@ -114,8 +114,8 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
    K_inner_loop = UOp.range(BLOCK_K//TC_K, rng, AxisType.REDUCE)

    # load from locals into registers
-    Ar = UOp.placeholder((BLOCK_M//TC_M//WARPGROUP_SIZE,), dtypes.half, slot=1, addrspace=AddrSpace.REG)
-    Br = UOp.placeholder((BLOCK_N//TC_N,), dtypes.half, slot=2, addrspace=AddrSpace.REG)
+    Ar = UOp.placeholder((BLOCK_M//TC_M//WARPGROUP_SIZE,), dtypes.half.vec(8), slot=1, addrspace=AddrSpace.REG)
+    Br = UOp.placeholder((BLOCK_N//TC_N,), dtypes.half.vec(8), slot=2, addrspace=AddrSpace.REG)

    M_load_loop = UOp.range(BLOCK_M//TC_M//WARPGROUP_SIZE, rng+10)
    Asl = Asl.reshape((BLOCK_K//TC_K, TC_K, BLOCK_M//TC_M//WARPGROUP_SIZE, WARPGROUP_SIZE, TC_M))
@@ -138,7 +138,7 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:

    # do WMMA
    wmma_arg = ('WMMA_16_16_32_half_float', (16, 16, 32), dtypes.half, dtypes.float, 'AMD', 64, ((), (), ((3, 2), (2, 2))), ())
-    out = UOp(Ops.WMMA, dtypes.float, (Ar[M_inner_loop], Br[N_inner_loop], acc_load), arg=wmma_arg)
+    out = UOp(Ops.WMMA, dtypes.float.vec(4), (Ar[M_inner_loop], Br[N_inner_loop], acc_load), arg=wmma_arg)

    # store back the acc
    acc_store = acc[N_inner_loop, M_inner_loop].store(out)
@@ -180,7 +180,7 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
  # store the acc into gmem
  cp_i, cp_j = UOp.range(BLOCK_M//TC_M//WARPGROUP_SIZE, 10004), UOp.range(BLOCK_N//TC_N, 10005)
  c_load = lambda i: C[gx, cp_i*TC_M*WARPGROUP_SIZE + warpgroup*TC_M + (warp//16)*4+i, gy, cp_j*TC_N + warp%16]
-  store = UOp.group(*[c_load(i).store(acc[cp_j, cp_i].index(i)) for i in range(4)])
+  store = UOp.group(*[c_load(i).store(acc[cp_j, cp_i].gep(i)) for i in range(4)])
  store = store.end(cp_i, cp_j)

  return store.sink(arg=KernelInfo(name="custom_gemm", opts_to_apply=())).simplify()
@@ -192,12 +192,12 @@ acc = UOp.placeholder((4,), dtypes.float, 0, AddrSpace.REG)
 acc = acc[init_l:=UOp.range(4, 1)].set(0.0, end=init_l)

 # do the wmma
-acc_load = UOp.stack(*[acc.after(K_loop)[i] for i in range(4)])
+acc_load = UOp.vectorize(*[acc.after(K_loop)[i] for i in range(4)])
 wmma_arg = ('WMMA_16_16_32_half_float', (16, 16, 32), dtypes.half, dtypes.float, 'AMD', 64, ((), (), ((3, 2), (2, 2))), ())
-out = UOp(Ops.WMMA, dtypes.float, (A_in, B_in, acc_load), arg=wmma_arg)
+out = UOp(Ops.WMMA, dtypes.float.vec(4), (A_in, B_in, acc_load), arg=wmma_arg)

 # store back the acc
-acc = acc.after(UOp.group(*[acc[i].store(out.index(i)) for i in range(4)]).end(K_loop))
+acc = acc.after(UOp.group(*[acc[i].store(out.gep(i)) for i in range(4)]).end(K_loop))

 # store the acc into gmem
 store = UOp.group(*[C[gx, (warp//16)*4+i, gy, warp%16].store(acc[i]) for i in range(4)])
@@ -218,7 +218,7 @@ if __name__ == "__main__":
  ref.realize()

  GlobalCounters.reset()
-  with Context(DEBUG=max(2, DEBUG.value)):
+  with Context(DEBUG=max(2, DEBUG.value), DEVECTORIZE=2):
    tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0]
    tst.realize()
  print(f"{(N*M*K*2 / GlobalCounters.time_sum_s)*1e-12:.2f} REAL TFLOPS")
@@ -37,8 +37,8 @@ def compute_on_locals(acc:UOp, Asl:UOp, Bsl:UOp, rng:int, afters:tuple[UOp, ...]
  K_inner_loop = UOp.range(BLOCK_K//TC_K, rng, AxisType.REDUCE)

  # load from locals into registers
-  Ar = UOp.placeholder((BLOCK_M//TC_M//WARPGROUP_SIZE,), dtypes.half, slot=1, addrspace=AddrSpace.REG)
-  Br = UOp.placeholder((BLOCK_N//TC_N,), dtypes.half, slot=2, addrspace=AddrSpace.REG)
+  Ar = UOp.placeholder((BLOCK_M//TC_M//WARPGROUP_SIZE,), dtypes.half.vec(8), slot=1, addrspace=AddrSpace.REG)
+  Br = UOp.placeholder((BLOCK_N//TC_N,), dtypes.half.vec(8), slot=2, addrspace=AddrSpace.REG)

  M_load_loop = UOp.range(BLOCK_M//TC_M//WARPGROUP_SIZE, rng+10)
  Asl = Asl.reshape(BLOCK_K//TC_K, TC_K, BLOCK_M//TC_M//WARPGROUP_SIZE, WARPGROUP_SIZE, TC_M)
@@ -61,7 +61,7 @@ def compute_on_locals(acc:UOp, Asl:UOp, Bsl:UOp, rng:int, afters:tuple[UOp, ...]

  # do WMMA
  wmma_arg = ('WMMA_16_16_32_half_float', (16, 16, 32), dtypes.half, dtypes.float, 'AMD', 64, ((), (), ((3, 2), (2, 2))), ())
-  out = UOp(Ops.WMMA, dtypes.float, (Ar[M_inner_loop], Br[N_inner_loop], acc_load), arg=wmma_arg)
+  out = UOp(Ops.WMMA, dtypes.float.vec(4), (Ar[M_inner_loop], Br[N_inner_loop], acc_load), arg=wmma_arg)

  # store back the acc
  acc_store = acc[N_inner_loop, M_inner_loop].store(out)
@@ -72,7 +72,7 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
  K_outer_loop = UOp.range(K//BLOCK_K, 0, AxisType.REDUCE)

  # split out the globals into blocks
-  C = C.src[0].cast(dtypes.float).reshape((M//BLOCK_M, BLOCK_M, N//BLOCK_N, BLOCK_N))
+  C = C.src[0].cast(dtypes.float.vec(4).ptr(C.ptrdtype.size)).reshape((M//BLOCK_M, BLOCK_M, N//BLOCK_N, BLOCK_N))
  A = A.reshape((M//BLOCK_M, BLOCK_M, K//BLOCK_K, BLOCK_K))[gx, :, K_outer_loop, :]
  B = B.reshape((K//BLOCK_K, BLOCK_K, N//BLOCK_N, BLOCK_N))[K_outer_loop, :, gy, :]

@@ -107,7 +107,7 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
  if getenv("COMPUTE"):
    As, Bs = As.after(barrier), Bs.after(barrier)

-    acc = UOp.placeholder((BLOCK_N//TC_N, BLOCK_M//TC_M//WARPGROUP_SIZE), dtypes.float, 0, AddrSpace.REG)
+    acc = UOp.placeholder((BLOCK_N//TC_N, BLOCK_M//TC_M//WARPGROUP_SIZE), dtypes.float.vec(4), 0, AddrSpace.REG)

    sink = compute_on_locals(acc, As, Bs, 200, afters=(barrier,), warpgroup=warpgroup, warp=warp)
    sink = sink.end(K_outer_loop)
@@ -127,7 +127,7 @@ if __name__ == "__main__":


  GlobalCounters.reset()
-  with Context(DEBUG=max(2, DEBUG.value)):
+  with Context(DEBUG=max(2, DEBUG.value), DEVECTORIZE=2):
    tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0]
    tst.realize()
  print(f"{(N*M*K*2 / GlobalCounters.time_sum_s)*1e-12:.2f} REAL TFLOPS")
@@ -1,58 +0,0 @@
-#!/usr/bin/env python3
-"""Rapidly measure ONNX output sensitivity to FP16-rounded initializers."""
-import argparse, copy
-
-import numpy as np
-import onnx
-import onnxruntime as ort
-from onnx import numpy_helper
-
-
-def main() -> None:
-  ap = argparse.ArgumentParser()
-  ap.add_argument("model")
-  ap.add_argument("corpus")
-  ap.add_argument("--case", type=int, default=9)
-  ap.add_argument("--chunks", type=int, default=8)
-  ap.add_argument("--start", type=int, default=0)
-  ap.add_argument("--stop", type=int)
-  ap.add_argument("--list", action="store_true")
-  args = ap.parse_args()
-
-  model = onnx.load(args.model)
-  consumers: dict[str, list[str]] = {}
-  for node in model.graph.node:
-    for name in node.input: consumers.setdefault(name, []).append(node.op_type)
-  initializers = [(init, numpy_helper.to_array(init)) for init in model.graph.initializer]
-  selected = [(init, arr) for init, arr in initializers if arr.dtype == np.float32 and arr.ndim >= 2 and
-              any(op in {"Conv", "Gemm", "MatMul"} for op in consumers.get(init.name, []))]
-  if args.list:
-    for i, (init, arr) in enumerate(selected): print(i, init.name, arr.shape, consumers.get(init.name))
-
-  # Listing an initializer as a graph input lets one ORT session override it at run time.
-  known_inputs = {x.name for x in model.graph.input}
-  for init, _ in selected:
-    if init.name not in known_inputs:
-      model.graph.input.append(copy.deepcopy(onnx.helper.make_tensor_value_info(init.name, init.data_type, init.dims)))
-  session_options = ort.SessionOptions()
-  session_options.log_severity_level = 3
-  session = ort.InferenceSession(model.SerializeToString(), session_options, providers=["CPUExecutionProvider"])
-  corpus = np.load(args.corpus)
-  feeds = {spec.name: corpus[f"case{args.case}:input:{spec.name}"] for spec in session.get_inputs()
-           if f"case{args.case}:input:{spec.name}" in corpus}
-  expected = corpus[f"case{args.case}:output"].astype(np.float32)
-
-  def check(indices: list[int]) -> tuple[float, float]:
-    overrides = {selected[i][0].name: selected[i][1].astype(np.float16).astype(np.float32) for i in indices}
-    got = session.run(None, feeds | overrides)[0].astype(np.float32)
-    delta = np.abs(expected.reshape(got.shape)-got)
-    return float(delta.max()), float(delta.mean())
-
-  scan = list(range(args.start, len(selected) if args.stop is None else args.stop))
-  print(f"selected={len(selected)} scan={scan[0]}..{scan[-1]} baseline={check([])} scan_error={check(scan)}")
-  for chunk in np.array_split(np.asarray(scan), args.chunks):
-    ids = [int(x) for x in chunk]
-    print(f"range={ids[0]}..{ids[-1]} count={len(ids)} error={check(ids)}")
-
-
-if __name__ == "__main__": main()
@@ -1,266 +0,0 @@
-#!/usr/bin/env python3
-"""Random-matrix oracle for the hand IR3 4x16 FP16 GEMM."""
-import os, struct
-
-import numpy as np
-
-from tinygrad import Device, dtypes
-from tinygrad.device import Buffer
-from extra.gemm import qcom_intensity_gemm as q
-from extra.gemm.ir3asm import disasm, get_envelope, inject
-
-
-def upload(x: np.ndarray) -> Buffer:
-  ret = Buffer("QCOM", x.size, dtypes.half).allocate()
-  ret.copyin(memoryview(np.ascontiguousarray(x)).cast("B"))
-  return ret
-
-
-def strip_redundant_mad_sy(lib: bytes) -> bytes:
-  ret = bytearray(lib)
-  io, sz = struct.unpack_from('<I', ret, 0xc0)[0], struct.unpack_from('<I', ret, 0x100)[0]
-  seen = False
-  for off in range(io, io+sz, 8):
-    hi = struct.unpack_from('<I', ret, off+4)[0]
-    if (hi >> 24) == 0x73:
-      if seen: struct.pack_into('<I', ret, off+4, (hi & 0x0fffffff) | 0x60000000)
-      else: seen = True
-  return bytes(ret)
-
-
-def restore_all_mad_sy(lib: bytes) -> bytes:
-  ret = bytearray(lib)
-  io, sz = struct.unpack_from('<I', ret, 0xc0)[0], struct.unpack_from('<I', ret, 0x100)[0]
-  for off in range(io, io+sz, 8):
-    hi = struct.unpack_from('<I', ret, off+4)[0]
-    if (hi >> 24) == 0x63: struct.pack_into('<I', ret, off+4, (hi & 0x0fffffff) | 0x70000000)
-  return bytes(ret)
-
-
-def restore_original_mad_sy(lib: bytes, original: bytes) -> bytes:
-  ret = bytearray(lib)
-  io, sz = struct.unpack_from('<I', ret, 0xc0)[0], struct.unpack_from('<I', ret, 0x100)[0]
-  for off in range(io, io+sz, 8):
-    hi = struct.unpack_from('<I', ret, off+4)[0]
-    old_hi = struct.unpack_from('<I', original, off+4)[0]
-    if (hi >> 24) == 0x63 and (old_hi >> 24) == 0x73:
-      struct.pack_into('<I', ret, off+4, (hi & 0x0fffffff) | 0x70000000)
-  return bytes(ret)
-
-
-def main() -> None:
-  m, n, k = int(os.getenv("M", "128")), int(os.getenv("N", "1024")), int(os.getenv("K", "384"))
-  stride = int(os.getenv("STRIDE", str(n)))
-  ncols = int(os.getenv("NCOLS", "4"))
-  threads = int(os.getenv("THREADS", "128"))
-  rng = np.random.default_rng(int(os.getenv("SEED", "4")))
-  a = (rng.standard_normal((m, k))*0.05).astype(np.float16)
-  b = (rng.standard_normal((k, n))*0.05).astype(np.float16)
-  if pattern := os.getenv("PATTERN", ""):
-    a.fill(0)
-    b.fill(0)
-    if pattern == "row":
-      a[:, 0] = np.arange(1, m+1)
-      b[0, :] = 1
-    elif pattern == "col":
-      a[:, 0] = 1
-      b[0, :] = (np.arange(n) % 251) + 1
-    elif pattern.startswith("k"):
-      kk = int(pattern[1:])
-      a[:, kk] = np.arange(1, m+1)
-      b[kk, :] = 1
-    else: raise ValueError(f"unknown PATTERN={pattern!r}")
-  q.M, q.N, q.K, q.K4 = m, stride, k, k//4
-  dev = Device["QCOM"]
-  compiler = bool(int(os.getenv("COMPILER", "0")))
-  env_ncols = int(os.getenv("ENV_NCOLS", str(ncols)))
-  direct_env = compiler or bool(int(os.getenv("ENV_DIRECT", "0")))
-  image_store = bool(int(os.getenv("IMAGE_STORE", "0")))
-  output_float = bool(int(os.getenv("OUTPUT_FLOAT", "0")))
-  dynamic_splits = int(os.getenv("DYNAMIC_SPLIT", "0"))
-  env_src = q.make_direct_image_donor_src(env_ncols, threads) if image_store else \
-            q.make_direct_donor_src(env_ncols if direct_env else ncols, threads) if direct_env else q.make_donor_src(env_ncols, threads)
-  env, io, sz, ro = get_envelope(dev, env_src)
-  fast = bool(int(os.getenv("FAST", "1")))
-  preserve_coords = bool(int(os.getenv("PRESERVE_COORDS", "0")))
-  high_inputs = bool(int(os.getenv("HIGH_INPUTS", "0")))
-  high_store = bool(int(os.getenv("HIGH_STORE", "0")))
-  low_a = bool(int(os.getenv("LOW_A", "0")))
-  inc = bool(int(os.getenv("INC", str(int(fast and not preserve_coords)))))
-  persistent = bool(int(os.getenv("PERSISTENT", str(int(fast and inc and not preserve_coords)))))
-  unroll = int(os.getenv("K_UNROLL", str(4 if k % 16 == 0 else 1)))
-  k_count = int(os.getenv("K_COUNT", str(k//4)))
-  if compiler:
-    patch_mode = os.getenv("PATCH_COMPILER", "0")
-    if patch_mode == "sync": lib = strip_redundant_mad_sy(env)
-    elif patch_mode != "0":
-      from extra.gemm.qcom_gemm import patch_kernel
-      lib = patch_kernel(env)
-      if patch_mode == "rpt": lib = restore_all_mad_sy(lib)
-      elif patch_mode == "original": lib = restore_original_mad_sy(lib, env)
-    else: lib = env
-    shader = bytes(env[io:io+sz])
-  else:
-    safe_store = bool(int(os.getenv("SAFE_STORE", "0")))
-    compact = bool(int(os.getenv("COMPACT", str(int(not safe_store)))))
-    isolated = bool(int(os.getenv("ISOLATED", "0")))
-    shader, _ = q.build_4x16_isolated_shader(dev, threads, k_unroll=unroll) if isolated else q.build_4xn_shader(
-      dev, threads, ncols=ncols, direct=True, compact_acc=compact,
-      store_constant=bool(int(os.getenv("STORE_CONSTANT", "0"))),
-      donor_store=bool(int(os.getenv("DONOR_STORE", "0"))),
-      native_store=bool(int(os.getenv("NATIVE_STORE", "0"))),
-      safe_store=safe_store,
-      linear_store=bool(int(os.getenv("LINEAR_STORE", "0"))),
-      image_store=image_store,
-      preserve_coords=preserve_coords,
-      preload_b=bool(int(os.getenv("PRELOAD_B", "0"))),
-      preload_b_safe_coords=bool(int(os.getenv("PRELOAD_B_SAFE_COORDS", "0"))),
-      high_inputs=high_inputs,
-      high_store=high_store,
-      copy_b_probe=bool(int(os.getenv("COPY_B_PROBE", "0"))),
-      thread_store=bool(int(os.getenv("THREAD_STORE", "0"))),
-      repeat_first_store=bool(int(os.getenv("REPEAT_FIRST_STORE", "0"))),
-      repair_row1_store=bool(int(os.getenv("REPAIR_ROW1_STORE", "0"))),
-      repeat_each_store=bool(int(os.getenv("REPEAT_EACH_STORE", "0"))),
-      post_constant=bool(int(os.getenv("POST", "0"))),
-      stable_bx=fast and not preserve_coords, stable_ay=fast, low_a_coords=low_a,
-      inc_coords=inc, persistent_coords=persistent,
-      serial_b_cols=bool(int(os.getenv("SERIAL", "0"))),
-      single_cols_all=bool(int(os.getenv("SINGLE_COLS_ALL", "0"))),
-      first_sync_only=bool(int(os.getenv("FIRST_SYNC_ONLY", str(int(fast))))),
-      no_store=bool(int(os.getenv("NO_STORE", "0"))),
-      skip_a_loads=bool(int(os.getenv("SKIP_A_LOADS", "0"))),
-      skip_b_loads=bool(int(os.getenv("SKIP_B_LOADS", "0"))),
-      k_unroll=unroll, b_first=fast and ncols == 4 and not preserve_coords,
-      k_count=None if dynamic_splits else k_count,
-      coord_delay=int(os.getenv("COORD_DELAY", "-1" if fast else "4")),
-      stable_settle_delay=int(os.getenv("STABLE_SETTLE_DELAY", "5")),
-      row_sync=bool(int(os.getenv("ROW_SYNC", "0"))), store_row_shift=int(os.getenv("STORE_ROW_SHIFT", "10")),
-      store_gap=int(os.getenv("STORE_GAP", "-1")),
-      safe_b_y=bool(int(os.getenv("SAFE_B_Y", "0"))), sync_b_y=bool(int(os.getenv("SYNC_B_Y", "0"))),
-      separate_b_coords=bool(int(os.getenv("SEPARATE_B_COORDS", "0"))),
-      high_b_coords=bool(int(os.getenv("HIGH_B_COORDS", "0"))),
-      reuse_separate_b_y=bool(int(os.getenv("REUSE_SEPARATE_B_Y", "0"))),
-      persistent_b_coords=bool(int(os.getenv("PERSISTENT_B_COORDS", "0"))),
-      interleave_second_pair=bool(int(os.getenv("INTERLEAVE_SECOND_PAIR", "0"))),
-      pipeline=bool(int(os.getenv("PIPELINE", "0"))),
-      acc_hr=int(os.getenv("ACC_HR")) if os.getenv("ACC_HR") else None,
-      high_a_only=bool(int(os.getenv("HIGH_A_ONLY", "0"))),
-      save_output_coords=bool(int(os.getenv("SAVE_OUTPUT_COORDS", "0"))),
-      vector_init=bool(int(os.getenv("VECTOR_INIT", "0"))),
-      dynamic_split_k=dynamic_splits,
-      alu_order=os.getenv("ALU_ORDER", "auto"),
-      first_cols_only=bool(int(os.getenv("FIRST_COLS_ONLY", "0"))), first_cols_offset=int(os.getenv("FIRST_COLS_OFFSET", "0")))
-    merged_opt = os.getenv("MERGEDREGS")
-    mergedregs = None if merged_opt is None else bool(int(merged_opt))
-    native = bool(int(os.getenv("NATIVE_STORE", "0")))
-    save_output = bool(int(os.getenv("SAVE_OUTPUT_COORDS", "0")))
-    persistent_b = bool(int(os.getenv("PERSISTENT_B_COORDS", "0")))
-    default_fregs = (23 if isolated else 30 if high_store else 28 if native else 19 if save_output and persistent_b else
-                     18 if bool(int(os.getenv("HIGH_B_COORDS", "0"))) else
-                     16 if bool(int(os.getenv("SAFE_B_Y", "0"))) else 11 if save_output or bool(int(os.getenv("THREAD_STORE", "0"))) else
-                     8 if high_inputs and low_a else 10)
-    acc_hr = int(os.getenv("ACC_HR", "0"))
-    default_hregs = (28 if isolated else max(acc_hr + 4*ncols, 36 if bool(int(os.getenv("HIGH_A_ONLY", "0"))) else
-                     44 if high_inputs and low_a else 48 if high_inputs else 32 if not compact else 12 + 4*ncols))
-    lib = inject(env, io, sz, ro, shader, fregs=int(os.getenv("FREGS", str(default_fregs))),
-                 hregs=int(os.getenv("HREGS", str(default_hregs))), mergedregs=mergedregs)
-  if int(os.getenv("PRINT_META", "0")):
-    asm = disasm(shader)
-    print(f"shader_instrs={len(shader)//8} mad_f16={asm.count('mad.f16')} isam={asm.count('isam')} sy={asm.count('(sy)')}")
-  if int(os.getenv("DUMP", "0")):
-    print(disasm(shader))
-    return
-  ab, bb = upload(a), upload(b.reshape(k, n//4, 4))
-  cb = Buffer("QCOM", max(1, dynamic_splits)*m*stride, dtypes.float if output_float else dtypes.half).allocate()
-  cb.copyin(memoryview(np.zeros((max(1, dynamic_splits)*m, stride), np.float32 if output_float else np.float16)).cast("B"))
-  specs = ([((0, dtypes.half, (m, stride//4, 4)),), ((0, dtypes.half, (m, k//4, 4)),),
-            ((1, dtypes.half, (k, n//4, 4)),)] if image_store and not output_float else
-           [((0, dtypes.float, (m, stride//4, 4)),), ((0, dtypes.half, (m, k//4, 4)),),
-            ((1, dtypes.half, (k, n//4, 4)),)] if image_store else
-           [((0, dtypes.half, (m, k//4, 4)),), ((1, dtypes.half, (k, n//4, 4)),),
-            ((2, dtypes.half, (max(1, dynamic_splits)*m*stride,)),)])
-  prg = dev.runtime("gemm_h", lib, buf_dtypes=specs)
-  call_bufs = (cb._buf, ab._buf, bb._buf) if image_store else (ab._buf, bb._buf, cb._buf)
-  tile_m = (threads//32)*4
-  # Each workgroup covers 32 lanes * ncols half4 vectors = 128*ncols scalar columns;
-  # thread count changes only the number of 4-row subtiles in Y.
-  times = [prg(*call_bufs, global_size=(n//(128*ncols), (m//tile_m)*max(1, dynamic_splits), 1), local_size=(threads, 1, 1), wait=True)*1e3
-           for _ in range(10)]
-  if int(os.getenv("NO_STORE", "0")):
-    print(f"K={k} ncols={ncols} compute_only_ms={min(times):.4f}")
-    return
-  got = np.empty((max(1, dynamic_splits)*m, stride), np.float32 if output_float else np.float16)
-  cb.copyout(memoryview(got).cast("B"))
-  if dynamic_splits:
-    split_got = got.reshape(dynamic_splits, m, stride).astype(np.float32)
-    if int(os.getenv("SPLIT_STATS", "0")):
-      chunk = k//dynamic_splits
-      split_expected = np.stack([a[:, s*chunk:(s+1)*chunk].astype(np.float32) @
-                                 b[s*chunk:(s+1)*chunk].astype(np.float32) for s in range(dynamic_splits)])
-      print("split_err", [[float(np.abs(split_got[x, :, :n]-split_expected[y]).mean())
-                           for y in range(dynamic_splits)] for x in range(dynamic_splits)])
-      print("split_norm", [float(np.abs(split_got[x, :, :n]).mean()) for x in range(dynamic_splits)])
-    got = split_got.sum(axis=0)
-  if int(os.getenv("RAW_STATS", "0")):
-    nz = np.flatnonzero(got.reshape(-1))
-    print("c_va", hex(cb._buf.va_addr), "raw_nonzero", len(nz), "head", nz[:128].tolist(), "tail", nz[-32:].tolist())
-  if int(os.getenv("THREAD_STORE", "0")):
-    raw, got = got.reshape(-1, 4, ncols, 4), np.empty_like(got)
-    nz = np.flatnonzero(raw.reshape(-1))
-    print("thread_nonzero_head", nz[:64].tolist(), "threads", np.unique(nz//(16*ncols))[:64].tolist(),
-          "thread_count", len(np.unique(nz//(16*ncols))))
-    # The thread-major kernel reserves tile slots using the physical output
-    # stride, even when only a logical prefix of columns is launched.
-    storage_gx_count = stride//(128*ncols)
-    launched_gx_count = n//(128*ncols)
-    for gy in range(m//16):
-      for gx in range(launched_gx_count):
-        for lid in range(128):
-          tm, tid = lid//32, lid%32
-          thread = (gy*storage_gx_count+gx)*128+lid
-          col_base = gx*32*ncols+tid
-          for row in range(4):
-            for col in range(ncols): got[gy*16+tm*4+row, (col_base+col*32)*4:(col_base+col*32+1)*4] = raw[thread, row, col]
-  got = got[:, :n]
-  expected = (np.full((m, n), 1024.0, np.float32) if int(os.getenv("POST", "0")) else
-              np.broadcast_to(b[0].astype(np.float32), (m, n)) if int(os.getenv("COPY_B_PROBE", "0")) else
-              np.full((m, n), float(k), np.float32) if int(os.getenv("SKIP_A_LOADS", "0")) and int(os.getenv("SKIP_B_LOADS", "0")) else
-              a[:, :k_count*4].astype(np.float32) @ b[:k_count*4].astype(np.float32))
-  delta = np.abs(got.astype(np.float32)-expected)
-  checked = np.ones((m, n), dtype=bool)
-  if int(os.getenv("FIRST_COLS_ONLY", "0")):
-    selected_parity = int(os.getenv("FIRST_COLS_OFFSET", "0")) & 1
-    for block in range(n//128):
-      if (block % ncols) % 2 != selected_parity:
-        delta[:, block*128:(block+1)*128] = 0
-        checked[:, block*128:(block+1)*128] = False
-  correct = np.allclose(got[checked], expected[checked], rtol=2e-2, atol=2e-2)
-  print(f"K={k} fast={fast} ms={min(times):.4f} max={delta.max():.8g} mean={delta.mean():.8g} "
-        f"finite={np.isfinite(got[checked]).all()} allclose={correct}")
-  print("samples", got[0, :16].tolist(), expected[0, :16].tolist())
-  if pattern:
-    print("pattern_blocks", [(x, got[0, x:x+8].tolist()) for x in range(0, n, 32)])
-  print("worst", np.unravel_index(int(np.nanargmax(delta)), delta.shape),
-        "col_means", [float(delta[:, x:x+128].mean()) for x in range(0, n, 128)],
-        "row_means", [float(delta[x:x+16].mean()) for x in range(0, m, 16)])
-  for out_block in (1, 3):
-    x = out_block * 128
-    print("block_match", out_block,
-          [float(np.abs(got[:, x:x+128].astype(np.float32)-expected[:, y:y+128]).mean())
-           for y in range(0, n, 128)])
-  bad = np.argwhere(delta > 0.02)
-  print("bad_count", len(bad), "bad_head", bad[:32].tolist())
-  print("bad_rows", [(int(r), int((bad[:, 0] == r).sum())) for r in np.unique(bad[:, 0])],
-        "bad_col_range", (int(bad[:, 1].min()), int(bad[:, 1].max())) if len(bad) else None)
-  print("row1_match", [float(np.abs(got[1].astype(np.float32)-expected[r]).mean()) for r in range(16)])
-  for probe_row in (127, 128, m-1):
-    if probe_row >= m: continue
-    probe_cols = checked[probe_row]
-    row_delta = np.abs(expected[:, probe_cols] - got[probe_row, probe_cols].astype(np.float32)).mean(axis=1)
-    nearest = np.argsort(row_delta)[:4]
-    print("row_match", probe_row, [(int(r), float(row_delta[r])) for r in nearest])
-  if not correct: raise SystemExit(1)
-
-
-if __name__ == "__main__": main()
@@ -1,277 +0,0 @@
-#!/usr/bin/env python3
-"""Random-matrix oracle for the high-throughput 8x8 IR3 GEMM."""
-import os, hashlib
-import numpy as np
-
-from tinygrad import Device, dtypes
-from tinygrad.device import Buffer
-from extra.gemm import qcom_8x4_gemm as q8
-from extra.gemm import qcom_intensity_gemm as q4
-from extra.gemm.ir3asm import disasm, get_envelope, inject
-
-
-def main():
-  m, n, k = int(os.getenv("M", "128")), int(os.getenv("N", "512")), int(os.getenv("K", "192"))
-  batch = int(os.getenv("BATCH", "1"))
-  threads = int(os.getenv("THREADS", "128"))
-  stride = int(os.getenv("STRIDE", str(max(1024, n))))
-  k_start, k_count = int(os.getenv("K_START", "0")), int(os.getenv("K_COUNT", str(k//4)))
-  rng = np.random.default_rng(int(os.getenv("SEED", "0")))
-  a_np = (rng.standard_normal((batch*m, k))*0.05).astype(np.float16)
-  b_np = (rng.standard_normal((batch*k, n))*0.05).astype(np.float16)
-  batch_horizontal = batch > 1 and bool(int(os.getenv("BATCH_HORIZONTAL", "1")))
-  batch_repeat_b = batch > 1 and not batch_horizontal and bool(int(os.getenv("BATCH_REPEAT_B", "0")))
-  batch_repeat_b_x = batch > 1 and not batch_horizontal and bool(int(os.getenv("BATCH_REPEAT_B_X", "0")))
-  b_storage = (np.concatenate([b_np[x*k:(x+1)*k] for x in range(batch) for _ in range(m//8)], axis=1)
-               if batch_repeat_b_x else
-               np.concatenate([b_np[x*k:(x+1)*k] for x in range(batch)], axis=1) if batch_horizontal else
-               np.concatenate([b_np[x*k:(x+1)*k] for x in range(batch) for _ in range(m//8)])
-               if batch_repeat_b else b_np)
-  pattern = os.getenv("PATTERN", "")
-  if pattern:
-    a_np.fill(0)
-    b_np.fill(0)
-    if pattern == "row":
-      a_np[:, 0] = np.arange(1, m+1, dtype=np.float16)
-      b_np[0, :] = 1
-    elif pattern == "col":
-      a_np[:, 0] = 1
-      b_np[0, :] = (np.arange(n, dtype=np.float16) % 251) + 1
-    elif pattern.startswith("k"):
-      kk = int(pattern[1:])
-      a_np[:, kk] = np.arange(1, m+1, dtype=np.float16)
-      b_np[kk, :] = 1
-    elif pattern.startswith("cross"):
-      ak, bk = map(int, pattern[5:].split("_"))
-      a_np[:, ak] = np.arange(1, m+1, dtype=np.float16)
-      b_np[bk, :] = 1
-    elif pattern == "ones":
-      a_np.fill(1)
-      b_np.fill(1)
-    else: raise ValueError(f"unknown PATTERN={pattern!r}")
-  q8.M, q8.N, q8.K, q8.K4 = batch*m, stride, k, k//4
-  dev = Device["QCOM"]
-  image_store = bool(int(os.getenv("IMAGE_STORE", "0")))
-  batch_const_mask = batch > 1 and bool(int(os.getenv("BATCH_CONST_MASK", "0")))
-  batch_z = batch > 1 and bool(int(os.getenv("BATCH_Z", "0")))
-  loop_instrs = -1
-  if bool(int(os.getenv("COMPILER", "0"))):
-    lib, _, _, _ = get_envelope(dev, q8.make_donor_src8(2, 128))
-  else:
-    wide = bool(int(os.getenv("WIDE", "0")))
-    tri = bool(int(os.getenv("TRI", "0")))
-    env_src = q4.make_direct_image_donor_src(4, threads) if image_store else q8.make_donor_src8(4, threads)
-    if batch_const_mask:
-      if not image_store: raise ValueError("BATCH_CONST_MASK currently requires IMAGE_STORE=1")
-      groups_per_batch = m // ((threads//32)*8)
-      env_src = env_src.replace("for(int k4=0", f"int batch=get_group_id(1)/{groups_per_batch};for(int k4=0")
-      env_src = env_src.replace("k4*4+", f"batch*{k}+k4*4+")
-    env, io, sz, ro = get_envelope(dev, env_src)
-    if int(os.getenv("PERSISTENT8", "0")):
-      shader, hregs, fregs, loop_instrs = q8.build_8x8_persistent_shader(dev, threads, k_count=k_count,
-        store_row_shift=int(os.getenv("STORE_ROW_SHIFT", "10")), pipeline_b=bool(int(os.getenv("PERSISTENT_PIPELINE", "0"))),
-        b_reuse_gap=int(os.getenv("B_REUSE_GAP", "0")), double_b=bool(int(os.getenv("PERSISTENT_DOUBLE_B", "0"))),
-        rotate_b=bool(int(os.getenv("PERSISTENT_ROTATE_B", "0"))), pipeline_a=bool(int(os.getenv("PERSISTENT_PIPELINE_A", "0"))),
-        one_sync=bool(int(os.getenv("PERSISTENT_ONE_SYNC", "0"))), one_sync_wait=int(os.getenv("PERSISTENT_ONE_SYNC_WAIT", "0")),
-        stagger_b=bool(int(os.getenv("PERSISTENT_STAGGER_B", "0"))),
-        stagger_rows=int(os.getenv("STAGGER_ROWS", "2")), masked_prefetch_a4=bool(int(os.getenv("MASKED_PREFETCH_A4", "0"))),
-        lagged_a4=bool(int(os.getenv("LAGGED_A4", "0"))), dual_a_tile=bool(int(os.getenv("DUAL_A_TILE", "0"))),
-        stream_a4_gap=int(os.getenv("STREAM_A4_GAP", "-1")),
-        dynamic_a4_dual=bool(int(os.getenv("DYNAMIC_A4_DUAL", "0"))),
-        dynamic_a4_wait=int(os.getenv("DYNAMIC_A4_WAIT", "0")),
-        dynamic_b_prefetch=bool(int(os.getenv("DYNAMIC_B_PREFETCH", "0"))),
-        dynamic_b_rows=int(os.getenv("DYNAMIC_B_ROWS", "1")),
-        dynamic_b_gap=int(os.getenv("DYNAMIC_B_GAP", "0")),
-        rotate_low_banks=bool(int(os.getenv("ROTATE_LOW_BANKS", "0"))),
-        rotate_no_prefetch=bool(int(os.getenv("ROTATE_NO_PREFETCH", "0"))),
-        batch_m=m if batch > 1 and bool(int(os.getenv("BATCH_SHADER", "1"))) else 0,
-        batch_n=n if batch > 1 and bool(int(os.getenv("BATCH_SHADER", "1"))) else 0,
-        batch_k=k if batch > 1 and bool(int(os.getenv("BATCH_SHADER", "1"))) else 0,
-        batch_b_offset=bool(int(os.getenv("BATCH_B_OFFSET", "1"))),
-        batch_row_offset=bool(int(os.getenv("BATCH_ROW_OFFSET", "1"))),
-        batch_horizontal=batch_horizontal,
-        batch_repeat_b=batch_repeat_b,
-        batch_repeat_b_x=batch_repeat_b_x,
-        batch_fixed_b=-2 if batch_z else int(os.getenv("BATCH_FIXED_B", "-1")),
-        batch_const_mask=batch_const_mask,
-        image_store_gap=int(os.getenv("IMAGE_STORE_GAP", "16")),
-        image_store=image_store)
-    elif int(os.getenv("PACKED_B8", "0")):
-      shader, hregs, fregs, loop_instrs = q8.build_8x8_bpacked_shader(dev, threads, coord_delay=int(os.getenv("ADELAY", "5")),
-        merged_alias=bool(int(os.getenv("PACKED_B_ALIAS", "0"))))
-    elif int(os.getenv("PACKED8", "0")):
-      shader, hregs, fregs, loop_instrs = q8.build_8x8_packed8_shader(dev, threads, coord_delay=int(os.getenv("ADELAY", "2")))
-    elif tri:
-      shader, hregs, fregs, loop_instrs = q8.build_8x4_shader(dev, 128, os.getenv("TRI_VARIANT", "serial"), 3,
-        a_coord_delay=int(os.getenv("ADELAY", "-1")), b_coord_delay=int(os.getenv("BDELAY", "-1")),
-        post_constant=bool(int(os.getenv("POST_CONSTANT", "0"))), image_store=image_store)
-    elif wide:
-      if image_store: raise ValueError("WIDE image store is not implemented")
-      shader, hregs, fregs, loop_instrs = q8.build_8x16_split_a_unroll_shader(dev, 128, k_unroll=int(os.getenv("KUNROLL", "4")),
-        b_coord_delay=int(os.getenv("BDELAY", "0")), fast_coords=True, safe_coords=bool(int(os.getenv("SAFE_COORDS", "1"))),
-        add256_store_mode=os.getenv("STORE_MODE", "tight"), alu_order=os.getenv("ALU_ORDER", "row_col_kk"),
-        post_constant=bool(int(os.getenv("POST_CONSTANT", "0"))),
-        skip_a_loads=bool(int(os.getenv("SKIP_A_LOADS", "0"))), skip_b_loads=bool(int(os.getenv("SKIP_B_LOADS", "0"))),
-        store_row_shift=int(os.getenv("STORE_ROW_SHIFT", "10")))
-    elif int(os.getenv("LIFETIME", "0")):
-      shader, hregs, fregs, loop_instrs = q8.build_8x8_lifetime_shader(dev, 128, k_unroll=int(os.getenv("KUNROLL", "4")),
-        b_coord_delay=int(os.getenv("BDELAY", "0")), a_coord_delay=int(os.getenv("ADELAY", "0")),
-        k_start=k_start, k_count=k_count, post_sequence=bool(int(os.getenv("POST_SEQUENCE", "0"))))
-    elif int(os.getenv("SELF_COORDS", "0")):
-      shader, hregs, fregs, loop_instrs = q8.build_8x8_selfcoord_shader(
-        dev, 128, coord_delay=int(os.getenv("ADELAY", "0")), post_sequence=bool(int(os.getenv("POST_SEQUENCE", "0"))))
-    elif int(os.getenv("BASE", "0")):
-      shader, hregs, fregs, loop_instrs = q8.build_8x8_split_a_shader(dev, 128,
-        a_coord_delay=int(os.getenv("ADELAY", "4")), b_coord_delay=int(os.getenv("BDELAY", "4")),
-        post_constant=bool(int(os.getenv("POST_CONSTANT", "0"))),
-        thread_store_gx=n//256 if int(os.getenv("THREAD_STORE", "0")) else 0,
-        thread_store_lid_reg=None if os.getenv("SAVE_REG", "r28.x") == "none" else os.getenv("SAVE_REG", "r28.x"),
-        thread_store_group_regs=("r36.y", "r36.z") if int(os.getenv("SAVE_GROUPS", "0")) else None,
-        row_sync=bool(int(os.getenv("ROW_SYNC", "0"))), reserved_out=int(os.getenv("RESERVED_OUT", "-1")))
-    else:
-      hist = bool(int(os.getenv("HIST", "0")))
-      common = dict(k_unroll=int(os.getenv("KUNROLL", "8")), b_coord_delay=int(os.getenv("BDELAY", "0")),
-        fast_coords=True, prefetch_next_b=bool(int(os.getenv("PREFETCH", "0"))), add256_store_mode=os.getenv("STORE_MODE", "tight"),
-        prefetch_next_a=bool(int(os.getenv("PREFETCH_A", "0"))),
-        grouped_b=bool(int(os.getenv("GROUPED_B", "0"))), grouped_b_cols=bool(int(os.getenv("GROUPED_B_COLS", "0"))),
-        stream_col1=bool(int(os.getenv("STREAM_COL1", "0"))), stream_col1_sync=bool(int(os.getenv("STREAM_COL1_SYNC", "0"))),
-        add256_gap=int(os.getenv("ADD256_GAP", "16")),
-        add256_offset_before_gap=bool(int(os.getenv("ADD256_OFFSET_BEFORE_GAP", "0"))),
-        alu_order=os.getenv("ALU_ORDER", "row_col_kk"),
-        post_constant=bool(int(os.getenv("POST_CONSTANT", "0"))))
-      if hist:
-        shader, hregs, fregs, loop_instrs = q8.build_8x8_split_a_unroll_shader(dev, 128, **common)
-      else:
-        shader, hregs, fregs, loop_instrs = q8.build_8x8_split_a_unroll_shader(dev, threads, **common, k_start=k_start, k_count=k_count,
-        thread_store_gx=n//256 if int(os.getenv("THREAD_STORE", "0")) else 0,
-        post_sequence=bool(int(os.getenv("POST_SEQUENCE", "0"))),
-        a_coord_delay=int(os.getenv("ADELAY", "4")), unroll_gap=int(os.getenv("GAP", "0")),
-        relaxed_sync=bool(int(os.getenv("RELAXED_SYNC", "0"))), sync_mask=int(os.getenv("SYNC_MASK", "7"), 0),
-        sync_wait=int(os.getenv("SYNC_WAIT", "0")), high_inputs=bool(int(os.getenv("HIGH_INPUTS", "0"))), image_store=image_store,
-        mid_acc=bool(int(os.getenv("MID_ACC", "0"))),
-        safe_coords=bool(int(os.getenv("SAFE_COORDS", "0"))), low_stable_coords=bool(int(os.getenv("LOW_STABLE_COORDS", "0"))),
-        triple_coords=bool(int(os.getenv("TRIPLE_COORDS", "0"))),
-        dual_a_coords=bool(int(os.getenv("DUAL_A_COORDS", "0"))),
-        high_pair_coords=bool(int(os.getenv("HIGH_PAIR_COORDS", "0"))),
-        high_a=bool(int(os.getenv("HIGH_A", "0"))),
-        low_a=bool(int(os.getenv("LOW_A", "0"))),
-        high_pair_b=bool(int(os.getenv("HIGH_PAIR_B", "0"))), high_pair_a=bool(int(os.getenv("HIGH_PAIR_A", "0"))),
-        serial_safe_coords=bool(int(os.getenv("SERIAL_SAFE_COORDS", "0"))),
-        separate_coords=bool(int(os.getenv("SEPARATE_COORDS", "0"))), buffer_a=bool(int(os.getenv("BUFFER_A", "0"))),
-        prefetch_loop_b=bool(int(os.getenv("PREFETCH_LOOP_B", "0"))), preload_a8=bool(int(os.getenv("PRELOAD_A8", "0"))),
-        reuse_b=bool(int(os.getenv("REUSE_B", "0"))), row_stream=bool(int(os.getenv("ROW_STREAM", "0"))),
-        phase_stream=bool(int(os.getenv("PHASE_STREAM", "0"))), split_low_pairs=bool(int(os.getenv("SPLIT_LOW_PAIRS", "0"))),
-        quad_a=bool(int(os.getenv("QUAD_A", "0"))), quad_map=os.getenv("QUAD_MAP", "0123"),
-        sampler_source_sync=bool(int(os.getenv("SOURCE_SYNC", "0"))),
-        stream_b_a8=bool(int(os.getenv("STREAM_B_A8", "0"))), store_row_shift=int(os.getenv("STORE_ROW_SHIFT", "10")),
-        source_hold_delay=int(os.getenv("SOURCE_HOLD_DELAY", "-1")), one_sync_tile=bool(int(os.getenv("ONE_SYNC_TILE", "0"))),
-        interleave_a4=bool(int(os.getenv("INTERLEAVE_A4", "0"))), interleave_a_reuse_gap=int(os.getenv("A_REUSE_GAP", "0")),
-        single_high_coord=bool(int(os.getenv("SINGLE_HIGH_COORD", "0"))))
-    assert len(shader) <= sz
-    if int(os.getenv("DISASM", "0")): print(disasm(shader))
-    lib = inject(env, io, sz, ro, shader, fregs=int(os.getenv("FREGS", str(fregs))), hregs=int(os.getenv("HREGS", str(hregs))),
-                 mergedregs=False if bool(int(os.getenv("SEPARATE_REGS", "0"))) else None)
-    if int(os.getenv("PRINT_META", "0")): print("shader_meta", fregs, hregs, len(shader), loop_instrs, hashlib.sha1(lib).hexdigest()[:8])
-  a, b = Buffer("QCOM", a_np.size, dtypes.half).allocate(), Buffer("QCOM", b_storage.size, dtypes.half).allocate()
-  c = Buffer("QCOM", batch*m*stride, dtypes.half).allocate()
-  q8.buf_copyin(a, memoryview(a_np).cast("B"))
-  q8.buf_copyin(b, memoryview(b_storage).cast("B"))
-  if not int(os.getenv("NO_INIT", "0")):
-    q8.buf_copyin(c, memoryview(np.zeros(batch*m*stride, dtype=np.float16)).cast("B"))
-  packed8 = bool(int(os.getenv("PACKED8", "0")))
-  packed_b8 = bool(int(os.getenv("PACKED_B8", "0")))
-  specs = ([((0, dtypes.half, (batch*m, stride//4, 4)),), ((0, dtypes.half, (batch*m, k//4, 4)),),
-            ((1, dtypes.half, (k, batch*(m//8)*n//4, 4)),) if batch_repeat_b_x else
-            ((1, dtypes.half, (k, batch*n//4, 4)),) if batch_horizontal else
-            ((1, dtypes.half, ((batch*k*(m//8) if batch_repeat_b else batch*k), n//4, 4)),)] if image_store else
-           [((0, dtypes.uint32, (m, k//8, 4)),), ((0, dtypes.uint32, (k, n//8, 4)),), ((0, dtypes.half, None),)] if packed8 else
-           [((0, dtypes.half, (m, k//4, 4)),), ((0, dtypes.uint32, (k, n//8, 4)),), ((0, dtypes.half, None),)] if packed_b8 else
-           [((0, dtypes.half, (batch*m, k//4, 4)),),
-            ((0, dtypes.half, (k, batch*n//4, 4)),) if batch_horizontal else ((0, dtypes.half, (batch*k, n//4, 4)),),
-            ((0, dtypes.half, None),)])
-  prg = dev.runtime("gemm_h", lib, buf_dtypes=specs)
-  if int(os.getenv("PRINT_META", "0")):
-    print("buffer_specs", specs)
-    print("runtime_meta", {k:v for k,v in vars(prg).items() if k in ("wgid", "lid", "max_threads", "prg")})
-  call_bufs = (c._buf, a._buf, b._buf) if image_store else (a._buf, b._buf, c._buf)
-  tile_n = 384 if bool(int(os.getenv("TRI", "0"))) else 512 if bool(int(os.getenv("WIDE", "0"))) else 256
-  tile_m = (threads//32)*8
-  global_size = ((n//tile_n, m//tile_m, batch) if batch_z else
-                 (batch*n//tile_n, m//tile_m, 1) if batch_horizontal else
-                 (n//tile_n, batch*m//tile_m, 1))
-  times = [prg(*call_bufs, global_size=global_size,
-               local_size=(threads, 1, 1), wait=True) for _ in range(int(os.getenv("BENCH_RUNS", "10")))]
-  elapsed = min(times)
-  got = np.empty(batch*m*stride, dtype=np.float16)
-  q8.buf_copyout(c, memoryview(got).cast("B"))
-  if int(os.getenv("RAW_STATS", "0")):
-    nz = np.flatnonzero(got)
-    print("raw_nonzero", nz.size, "first", nz[:64].tolist(), "last", nz[-64:].tolist(),
-          "values", np.unique(got[nz])[:16].tolist())
-  if int(os.getenv("THREAD_STORE", "0")) or int(os.getenv("DECODE_THREAD", "0")):
-    raw, matrix = got[:m*n].reshape(-1, 8, 2, 4), np.empty((m, n), np.float16)
-    if pattern:
-      print("raw_lids=", [[float(raw[lid, row, 0, 0]) for row in range(8)] for lid in range(0, 128, 8)])
-    gx_count = n//256
-    for gy in range(m//tile_m):
-      for gx in range(gx_count):
-        for lid in range(threads):
-          tm, tid = lid//32, lid%32
-          thread = (gy*gx_count+gx)*threads+lid
-          for row in range(8):
-            for col in range(2):
-              x = (gx*64+tid+col*32)*4
-              matrix[gy*tile_m+tm*8+row, x:x+4] = raw[thread, row, col]
-    got = matrix.astype(np.float32)
-  else: got = got.reshape(batch*m, stride)[:, :n].astype(np.float32)
-  if int(os.getenv("POST_SEQUENCE", "0")):
-    tile = np.empty((8, 256), np.float32)
-    for row in range(8):
-      for col in range(2): tile[row, col*128:(col+1)*128] = row*2+col+1
-    expected = np.tile(tile, (m//8, n//256))
-  else: expected = (np.full((batch*m, n), 1024, np.float32) if int(os.getenv("POST_CONSTANT", "0")) else
-              np.concatenate([a_np[x*m:(x+1)*m, k_start*4:(k_start+k_count)*4].astype(np.float32) @
-                              b_np[x*k+k_start*4:x*k+(k_start+k_count)*4].astype(np.float32)
-                              for x in range(batch)]))
-  delta = np.abs(expected-got)
-  if (reserved_out := int(os.getenv("RESERVED_OUT", "-1"))) >= 0:
-    row, col = divmod(reserved_out, 2)
-    delta[row::8, col*128:(col+1)*128] = 0
-  correct = np.allclose(expected, got, rtol=2e-2, atol=2e-2)
-  gflops = batch*2*m*n*(k_count*4)/elapsed/1e9
-  print(f"shape={batch}x{m}x{n}x{k_count*4} accumulate=fp16 elapsed_ms={elapsed*1e3:.3f} gflops={gflops:.1f} "
-        f"max_abs={delta.max():.9g} mean_abs={delta.mean():.9g} allclose={correct}")
-  bad = ~np.isfinite(got) | (delta > .02)
-  bad_idx = np.argwhere(bad)
-  print(f"bad_count={bad_idx.shape[0]}")
-  if int(os.getenv("VERBOSE", "0")):
-    for r in range(8): print(f"row{r} expected={expected[r,:8].tolist()} got={got[r,:8].tolist()}")
-    print("block_max=", [[float(delta[r:r+8, c:c+128].max()) for c in range(0, n, 128)] for r in range(0, m, 8)])
-    print("local_rows=", [(lr, float(delta[lr::8].max()), float(delta[lr::8].mean())) for lr in range(8)])
-    print("bad_by_row=", [(int(r), int(bad[r].sum())) for r in np.flatnonzero(bad.any(axis=1))])
-    print("bad_first=", [(int(r), int(c), float(expected[r, c]), float(got[r, c])) for r, c in bad_idx[:64]])
-  if int(os.getenv("POST_SEQUENCE", "0")):
-    print("sequence_blocks=", [[np.unique(got[row, col:col+128], return_counts=True) for col in range(0, n, 128)] for row in range(8)])
-  if int(os.getenv("VERBOSE", "0")):
-    row0_matches = np.abs(expected-got[0]).mean(axis=1)
-    print("row0_matches=", [(int(i), float(row0_matches[i])) for i in np.argsort(row0_matches)[:8]])
-    if batch > 1:
-      for row in range(0, batch*m, (threads//32)*8):
-        candidates = [a_np[row].astype(np.float32) @ b_np[x*k:(x+1)*k].astype(np.float32) for x in range(batch)]
-        print("batch_map=", row, [(x, float(np.abs(c-got[row]).mean())) for x, c in enumerate(candidates)])
-  if int(os.getenv("VERBOSE", "0")) and not int(os.getenv("POST_SEQUENCE", "0")) and not int(os.getenv("POST_CONSTANT", "0")):
-    contrib = np.stack([a_np[0, kk*4:kk*4+4].astype(np.float32) @
-                        b_np[kk*4:kk*4+4].astype(np.float32) for kk in range(k_start, k_start+k_count)])
-    excluded = np.abs((expected[0][None, :]-contrib)-got[0]).mean(axis=1)
-    prefixes = np.abs(np.cumsum(contrib, axis=0)-got[0]).mean(axis=1)
-    print("row0_k=", "exclude", [(k_start+int(i), float(excluded[i])) for i in np.argsort(excluded)[:4]],
-          "prefix", [(k_start+int(i)+1, float(prefixes[i])) for i in np.argsort(prefixes)[:4]])
-    if k_count == 1:
-      cs = np.stack([a_np[:, k_start*4+j:k_start*4+j+1].astype(np.float32) @
-                     b_np[k_start*4+j:k_start*4+j+1].astype(np.float32) for j in range(4)])
-      subset = [(mask, float(np.abs(sum((cs[j] for j in range(4) if mask & (1<<j)), np.zeros_like(got))-got).mean())) for mask in range(16)]
-      print("component_subsets=", sorted(subset, key=lambda x:x[1])[:8])
-  if not correct: raise SystemExit(1)
-
-
-if __name__ == "__main__": main()
@@ -1,111 +0,0 @@
-#!/usr/bin/env python3
-"""Dependency-free lane-mapping probe for the thread-major 8x8 shader."""
-import ctypes, os, random, struct
-
-from tinygrad import Device, dtypes
-from tinygrad.device import Buffer
-from extra.gemm import qcom_8x4_gemm as q8
-from extra.gemm.ir3asm import get_envelope, inject
-
-
-def half_bytes(values): return bytearray(struct.pack(f"<{len(values)}e", *values))
-
-
-def main():
-  m, n, k = int(os.getenv("M", "32")), int(os.getenv("N", "256")), int(os.getenv("K", "192"))
-  pattern = os.getenv("PATTERN", "row")
-  int8_b = bool(int(os.getenv("INT8_B", "0")))
-  a = [0.0] * (m*k)
-  b = [0.0] * (k*n)
-  if pattern == "row":
-    for row in range(m): a[row*k] = row+1
-    for col in range(n): b[col] = 1
-  elif pattern == "col":
-    for row in range(m): a[row*k] = 1
-    for col in range(n): b[col] = col % 251 + 1
-  elif pattern == "random":
-    rng = random.Random(int(os.getenv("SEED", "0")))
-    a = [rng.uniform(-0.05, 0.05) for _ in a]
-    b = [rng.uniform(-0.05, 0.05) for _ in b]
-  else: raise ValueError(pattern)
-  # The oracle must use the exact FP16 values consumed by the images.
-  a = list(struct.unpack(f"<{len(a)}e", half_bytes(a)))
-  if int8_b:
-    bq = [max(-127, min(127, round(x*127))) for x in b]
-    b = [x/127.0 for x in bq]
-    b_bytes = bytearray((x & 0xff) for x in bq)
-  else:
-    b = list(struct.unpack(f"<{len(b)}e", half_bytes(b)))
-    b_bytes = half_bytes(b)
-
-  q8.M, q8.N, q8.K, q8.K4 = m, n, k, k//4
-  dev = Device["QCOM"]
-  compiler = bool(int(os.getenv("COMPILER", "0")))
-  tight_store = bool(int(os.getenv("TIGHT_STORE", "0")))
-  mode = os.getenv("MODE", "")
-  if compiler:
-    lib, _, _, _ = get_envelope(dev, q8.make_donor_src8(2, 128))
-  elif mode:
-    env, io, sz, ro = get_envelope(dev, q8.make_donor_src8(4, 128))
-    if mode == "pipeline": shader, hregs, fregs, _, _ = q8.build_8x8_pipelined_shader(dev, 128, 4, 4, thread_store_gx=n//256)
-    elif mode == "pipeline4": shader, hregs, fregs, _, _ = q8.build_8x8_pipeline4_shader(dev, 128, 4, 4)
-    elif mode == "batch2": shader, hregs, fregs, _, _ = q8.build_8x8_batch2_shader(dev, 128, 4, 4)
-    else: raise ValueError(mode)
-    assert len(shader) <= sz
-    lib = inject(env, io, sz, ro, shader, fregs=fregs, hregs=hregs)
-  elif int(os.getenv("BASE", "0")):
-    env, io, sz, ro = get_envelope(dev, q8.make_donor_src8(4, 128))
-    shader, hregs, fregs, _ = q8.build_8x8_split_a_shader(dev, 128,
-      a_coord_delay=int(os.getenv("ADELAY", "3")), b_coord_delay=int(os.getenv("BDELAY", "3")),
-      pre_mad_nops=int(os.getenv("PMAD", "-1")), grouped_b=bool(int(os.getenv("GROUPED_B", "0"))),
-      grouped_b_cols=bool(int(os.getenv("GROUPED_COLS", "0"))), thread_store_gx=0 if tight_store else 1,
-      add256_store_mode="tight" if tight_store else "donor")
-    lib = inject(env, io, sz, ro, shader, fregs=fregs, hregs=hregs)
-  else:
-    env, io, sz, ro = get_envelope(dev, q8.make_donor_src8(4, 128))
-    shader, hregs, fregs, _ = q8.build_8x8_split_a_unroll_shader(dev, 128,
-      k_unroll=int(os.getenv("KUNROLL", "8")), b_coord_delay=int(os.getenv("BDELAY", "0")),
-      fast_coords=bool(int(os.getenv("FAST", "1"))), prefetch_next_b=bool(int(os.getenv("PREFETCH", "0"))),
-      thread_store_gx=0 if tight_store else 1, add256_store_mode="tight" if tight_store else "donor",
-      post_sequence=bool(int(os.getenv("POST_SEQUENCE", "0"))), a_coord_delay=int(os.getenv("ADELAY", "4")),
-      unroll_gap=int(os.getenv("GAP", "0")), relaxed_sync=bool(int(os.getenv("RELAXED_SYNC", "0"))),
-      sync_mask=int(os.getenv("SYNC_MASK", "7"), 0), sync_wait=int(os.getenv("SYNC_WAIT", "0")))
-    lib = inject(env, io, sz, ro, shader, fregs=fregs, hregs=hregs)
-  ab = Buffer("QCOM", len(a), dtypes.half).allocate()
-  bb = Buffer("QCOM", len(b), dtypes.int8 if int8_b else dtypes.half).allocate()
-  cb = Buffer("QCOM", m*n, dtypes.half).allocate()
-  for buf, raw in ((ab, half_bytes(a)), (bb, b_bytes), (cb, bytearray(m*n*2))):
-    src = (ctypes.c_ubyte * len(raw)).from_buffer(raw)
-    ctypes.memmove(int(buf._buf.va_addr), ctypes.addressof(src), len(raw))
-  specs = [((0, dtypes.half, (m, k//4, 4)),), ((0, dtypes.int8 if int8_b else dtypes.half, (k, n//4, 4)),),
-           ((0, dtypes.half, None),)]
-  prg = dev.runtime("gemm_h", lib, buf_dtypes=specs)
-  times = [prg(ab._buf, bb._buf, cb._buf, global_size=(n//256, m//32, 1), local_size=(128, 1, 1), wait=True)*1e3 for _ in range(5)]
-  print("elapsed_ms=", min(times))
-  out = bytearray(m*n*2)
-  ctypes.memmove(ctypes.addressof((ctypes.c_ubyte * len(out)).from_buffer(out)), int(cb._buf.va_addr), len(out))
-  raw = struct.unpack(f"<{m*n}e", out)
-  if int(os.getenv("DUMP_RAW", "0")):
-    for row in range(min(m, 16)): print("raw", row, list(raw[row*n:row*n+min(n, 64)]))
-    return
-  if pattern == "random":
-    worst = total = 0.0
-    worst_at = None
-    for row in range(m):
-      tm, rr = row//8, row%8
-      for col in range(n):
-        tid, cc, lane = (col//4)%32, col//128, col%4
-        got = raw[row*n+col] if compiler or tight_store or mode in ("pipeline4", "batch2") else raw[(tm*32+tid)*64 + rr*8 + cc*4 + lane]
-        expected = sum(a[row*k+kk] * b[kk*n+col] for kk in range(k))
-        delta = abs(got-expected)
-        if delta > worst: worst, worst_at = delta, (row, col, got, expected)
-        total += delta
-    print("max_abs=", worst, "mean_abs=", total/(m*n), "worst_at=", worst_at)
-    if worst > 0.02: raise SystemExit(1)
-    return
-  for lid in range(min(128, m*n//64)):
-    vals = [raw[lid*64+row*8] for row in range(8)]
-    print(lid, vals)
-
-
-if __name__ == "__main__": main()
@@ -1,277 +0,0 @@
-#!/usr/bin/env python3
-"""Adreno 630 FP16 MAD throughput benchmark."""
-import argparse, ctypes, struct
-
-from tinygrad import Device, dtypes
-from tinygrad.device import Buffer
-from extra.gemm.ir3asm import *
-from extra.gemm.ir3asm import _hreg
-from extra.gemm.qcom_intensity_gemm import M, N, K4, make_donor_src, prologue_4x2, store_output
-
-
-def make_bufs(dev):
-    a = Buffer(dev.device, (K4)*M*4, dtypes.half, preallocate=True)
-    b = Buffer(dev.device, (N//4)*(K4*4)*4, dtypes.half, preallocate=True)
-    c = Buffer(dev.device, M*N, dtypes.half, preallocate=True)
-    ctypes.memset(int(a._buf.va_addr), 0, a.nbytes)
-    ctypes.memset(int(b._buf.va_addr), 0, b.nbytes)
-    ctypes.memset(int(c._buf.va_addr), 0, c.nbytes)
-    return a, b, c
-
-
-def emit_mov_h_block(instrs, start, end, src):
-    pos = start
-    while pos < end:
-        rpt = min(3, end - pos - 1)
-        instrs.append(MOV_H(pos, src, rpt=rpt))
-        pos += rpt + 1
-
-
-def build_compiler_pattern_shader(dev, threads, loops, pairs, store):
-    if pairs < 2: raise ValueError('compiler-pattern needs at least two x/y MAD pairs; smaller shaders have caused QCOM hangs')
-    instrs = prologue_4x2(dev, threads)
-    instrs += [MOV_S32('r8.x', 0, sy=True), MOV_H_IMM('hr0.x', 0x3c00)]
-    emit_mov_h_block(instrs, 1, _hreg('hr8.x'), 0)
-
-    loop_start = len(instrs)
-    for _ in range(pairs):
-        # This mirrors the vec16 OpenCL MAD peak lowering: one vector MAD into x,
-        # then one vector MAD into y. The split scalar lane avoids clobbering hr0.y.
-        instrs += [
-            MAD_F16('hr0.z', 'hr0.z', 'hr4.y', 'hr4.y', rpt=3, r=True, r1=True),
-            MAD_F16('hr1.z', 'hr1.z', 'hr5.y', 'hr5.y', rpt=3, r=True, r1=True),
-            MAD_F16('hr2.z', 'hr2.z', 'hr6.y', 'hr6.y', rpt=3, r=True, r1=True),
-            MAD_F16('hr3.z', 'hr3.z', 'hr7.y', 'hr7.y', rpt=2, r=True, r1=True),
-            MAD_F16('hr0.x', 'hr0.x', 'hr0.y', 'hr0.y'),
-            MAD_F16('hr4.y', 'hr0.z', 'hr4.y', 'hr0.z', rpt=3, r=True, r1=True),
-            MAD_F16('hr5.y', 'hr1.z', 'hr5.y', 'hr1.z', rpt=3, r=True, r1=True),
-            MAD_F16('hr6.y', 'hr2.z', 'hr6.y', 'hr2.z', rpt=3, r=True, r1=True),
-            MAD_F16('hr7.y', 'hr3.z', 'hr7.y', 'hr3.z', rpt=2, r=True, r1=True),
-            MAD_F16('hr0.y', 'hr0.x', 'hr0.y', 'hr0.x'),
-        ]
-    instrs += [
-        ADD_S('r8.y', 'r8.x', 1),
-        CMPS_S_EQ('r8.x', loops - 1, nop=1),
-        MOV_F32('r8.x', 'r8.y'),
-        NOP(rpt=3),
-    ]
-    loop_end = len(instrs)
-    instrs.append(BR(loop_start - loop_end))
-    if store: store_output(instrs, 'r7.x', 'r7.y', 0)
-    instrs.append(END())
-    return assemble(instrs), loop_end - loop_start, 8, 9, pairs * 32 * 2
-
-
-def build_alu_shader(dev, threads, groups, rpt, loops, unroll, independent, r1):
-    if rpt > 3: raise ValueError('mad.f16 repeat counts above rpt3 encode other flags on A630, not more FP16 lanes')
-    if not (1 <= loops <= 256): raise ValueError('loops must be in 1..256; current immediate compare encodes only 8 bits')
-    width = rpt + 1
-    instrs = prologue_4x2(dev, threads)
-    instrs += [
-        MOV_S32('r6.z', 0, sy=True),
-        MOV_H_IMM('hr0.x', 0x3c00),
-        MOV_H_IMM('hr16.x', 0), MOV_H('hr16.y', 'hr16.x', rpt=2),
-    ]
-    emit_mov_h_block(instrs, 1, max(width, 4), 0)
-    emit_mov_h_block(instrs, _hreg('hr4.x'), _hreg('hr4.x') + max(width, 4), 0)
-    acc0 = _hreg('hr16.x')
-    hregs = (acc0 + groups * width + 3) // 4
-    emit_mov_h_block(instrs, acc0 + 4, acc0 + groups * width, acc0)
-
-    loop_start = len(instrs)
-    for _ in range(unroll):
-        for g in range(groups):
-            src1 = (g * width) % max(width, 4)
-            src2 = _hreg('hr4.x') + ((g * width) % max(width, 4))
-            src3 = src1 if independent else acc0 + g * width
-            instrs.append(MAD_F16(acc0 + g * width, src1, src2, src3, rpt=rpt, r=True, r1=r1))
-    instrs += [
-        ADD_S('r0.x', 'r6.z', 1),
-        CMPS_S_EQ('r6.z', loops - 1, nop=1),
-        MOV_F32('r6.z', 'r0.x'),
-        NOP(rpt=3),
-    ]
-    loop_end = len(instrs)
-    instrs.append(BR(loop_start - loop_end))
-    store_output(instrs, 'r7.x', 'r7.y', acc0)
-    instrs.append(END())
-    return assemble(instrs), loop_end - loop_start, hregs, None, unroll * groups * width * 2
-
-
-def gemm_check_inputs(rows, ncols):
-    a = [[((row * 3 + kk) % 4 + 1) / 8 for kk in range(4)] for row in range(rows)]
-    b = [[[[((col * 7 + kk * 3 + lane) % 4 + 1) / 8 for lane in range(4)] for kk in range(4)] for col in range(ncols)]][0]
-    return a, b
-
-
-def half_raw(value):
-    return struct.unpack('<H', struct.pack('<e', value))[0]
-
-
-def build_gemm_pattern_shader(dev, threads, loops, rows, ncols, unroll, order, bmode, r1, check_pattern=False, store_group=0):
-    if loops != 1: raise ValueError('gemm-pattern is a one-shot ALU body benchmark; use --loops 1 so loop-control regs do not clobber A/B sources')
-    if rows not in (4, 8): raise ValueError('rows must be 4 or 8')
-    if ncols < 1: raise ValueError('ncols must be positive')
-    instrs = prologue_4x2(dev, threads)
-    instrs += [MOV_S32('r6.z', 0, sy=True)]
-
-    # A lives in hr0..hr(rows-1). B either reuses one 4-texel column group or
-    # allocates one 4-texel group per output col4. Accumulators start at hr16 to
-    # match the working GEMM kernels and avoid low full-register aliases.
-    a_base = 0
-    b_base = rows * 4
-    b_groups = ncols if bmode == 'percol' else 1
-    b_end = b_base + b_groups * 16
-    acc0 = max(_hreg('hr16.x'), ((b_end + 3) // 4) * 4)
-    if check_pattern:
-        check_a, check_b = gemm_check_inputs(rows, ncols)
-        for row in range(rows):
-            for kk in range(4): instrs.append(MOV_H_IMM(a_base + row * 4 + kk, half_raw(check_a[row][kk])))
-        for col in range(b_groups):
-            for kk in range(4):
-                for lane in range(4): instrs.append(MOV_H_IMM(b_base + col * 16 + kk * 4 + lane, half_raw(check_b[col][kk][lane])))
-    else:
-        instrs.append(MOV_H_IMM('hr0.x', 0x3c00))
-        emit_mov_h_block(instrs, 1, rows * 4, 0)
-        emit_mov_h_block(instrs, b_base, b_end, 0)
-    for lane in range(acc0, acc0 + rows * ncols * 4): instrs.append(MOV_H_IMM(lane, 0))
-    hregs = (max(b_end, acc0 + rows * ncols * 4) + 3) // 4
-
-    loop_start = len(instrs)
-    def emit(row, kk, col):
-        b_col = col if bmode == 'percol' else 0
-        instrs.append(MAD_F16(acc0 + (row * ncols + col) * 4, a_base + row * 4 + kk, b_base + b_col * 16 + kk * 4,
-                              acc0 + (row * ncols + col) * 4, rpt=3, r=True, r1=r1))
-    for _ in range(unroll):
-        if order == 'kk_row_col':
-            for kk in range(4):
-                for row in range(rows):
-                    for col in range(ncols): emit(row, kk, col)
-        elif order == 'kk_col_row':
-            for kk in range(4):
-                for col in range(ncols):
-                    for row in range(rows): emit(row, kk, col)
-        elif order == 'col_kk_row':
-            for col in range(ncols):
-                for kk in range(4):
-                    for row in range(rows): emit(row, kk, col)
-        elif order == 'row_kk_col':
-            for row in range(rows):
-                for kk in range(4):
-                    for col in range(ncols): emit(row, kk, col)
-        elif order == 'row_col_kk':
-            for row in range(rows):
-                for col in range(ncols):
-                    for kk in range(4): emit(row, kk, col)
-        else: raise ValueError('unknown order %s' % order)
-    instrs += [
-        ADD_S('r0.x', 'r6.z', 1),
-        CMPS_S_EQ('r6.z', loops - 1, nop=1),
-        MOV_F32('r6.z', 'r0.x'),
-        NOP(rpt=3),
-    ]
-    loop_end = len(instrs)
-    instrs.append(BR(loop_start - loop_end))
-    if check_pattern:
-        # The per-column B register bank aliases the donor prologue's r7 output
-        # coordinates. Every lane computes the same diagnostic tile, so use one
-        # common output address and bit-check the selected accumulator vector.
-        instrs += [MOV_S32('r7.x', 0), MOV_S32('r7.y', 0), NOP(rpt=2)]
-    store_output(instrs, 'r7.x', 'r7.y', acc0 + store_group * 4)
-    instrs.append(END())
-    return assemble(instrs), loop_end - loop_start, hregs, None, unroll * rows * ncols * 4 * 4 * 2
-
-
-def run(args):
-    dev = Device['QCOM']
-    env_ncols = max(4 if args.gemm_pattern else 2, args.ncols if args.gemm_pattern else 2)
-    envelope, img_off, img_sz, reg_off = get_envelope(dev, make_donor_src(env_ncols, args.threads))
-    if args.compiler_pattern:
-        shader, loop_instrs, hregs, fregs, flops_per_thread_loop = build_compiler_pattern_shader(dev, args.threads, args.loops, args.pairs, args.store)
-    elif args.gemm_pattern:
-        shader, loop_instrs, hregs, fregs, flops_per_thread_loop = build_gemm_pattern_shader(
-            dev, args.threads, args.loops, args.rows, args.ncols, args.unroll, args.order, args.bmode, args.r1, args.check_gemm)
-    else:
-        shader, loop_instrs, hregs, fregs, flops_per_thread_loop = build_alu_shader(dev, args.threads, args.groups, args.rpt, args.loops, args.unroll, args.independent, args.r1)
-        width = args.rpt + 1
-    if fregs is None: fregs = args.fregs
-    if hregs > 48 and not args.allow_invalid_regs:
-        print('skipped: groups=%d needs hregs=%d, but A630 addressable GPR half registers stop at hr47 (hregs=48).' % (args.groups, hregs))
-        return
-    if len(shader) > img_sz:
-        print('skipped: shader is %d bytes but envelope has only %d bytes.' % (len(shader), img_sz))
-        return
-    lib = inject(envelope, img_off, img_sz, reg_off, shader, fregs=fregs, hregs=hregs)
-    asm = disasm(shader)
-    reg_count = fregs + (hregs + 1) // 2
-    wave_pairs = 96 // reg_count
-    mode = 'compiler-pattern' if args.compiler_pattern else ('gemm-pattern' if args.gemm_pattern else ('independent' if args.independent else 'accumulate'))
-    print('mode=%s r1=%d rows=%d ncols=%d bmode=%s order=%s groups=%d rpt=%d width=%d unroll=%d pairs=%d fregs=%d hregs=%d reg_count=%d wave_pairs=%d loop_instrs=%d shader_instrs=%d mad=%d rpt3=%d' % (
-        mode, args.r1, args.rows, args.ncols, args.bmode, args.order, args.groups, args.rpt, args.rpt + 1, args.unroll, args.pairs, fregs, hregs, reg_count, wave_pairs, loop_instrs, len(shader)//8, asm.count('mad.f16'), asm.count('(rpt3)mad.f16')))
-    if args.disasm: print(asm)
-
-    a, b, c = make_bufs(dev)
-    # Runtime buffer metadata now carries image shape separately from the scalar dtype.
-    buf_dtypes = [((0, dtypes.half, (M, K4, 4)),), ((0, dtypes.half, (K4*4, N//4, 4)),), ((0, dtypes.half, None),)]
-    prg = dev.runtime('gemm_h', lib, buf_dtypes=buf_dtypes)
-    tile_m = (args.threads // 32) * 4
-    gs, ls = (8, M // tile_m, 1), (args.threads, 1, 1)
-    for _ in range(5): prg(a._buf, b._buf, c._buf, global_size=gs, local_size=ls, wait=True)
-    times = []
-    for _ in range(args.iters):
-        t = prg(a._buf, b._buf, c._buf, global_size=gs, local_size=ls, wait=True)
-        if t: times.append(t)
-    best = min(times)
-    median = sorted(times)[len(times) // 2]
-    total_threads = gs[0] * gs[1] * args.threads
-    flops = total_threads * args.loops * flops_per_thread_loop
-    print('%.1f GFLOPS best (%.3f ms), %.1f GFLOPS median (%.3f ms), flops=%d runs=%d' %
-          (flops / best / 1e9, best * 1e3, flops / median / 1e9, median * 1e3, flops, len(times)))
-    if args.check_gemm:
-        if not args.gemm_pattern or args.bmode != 'percol' or args.loops != 1:
-            raise ValueError('--check-gemm requires --gemm-pattern --bmode percol --loops 1')
-        check_a, check_b = gemm_check_inputs(args.rows, args.ncols)
-        checked = 0
-        for group in range(args.rows * args.ncols):
-            check_shader, _, check_hregs, _, _ = build_gemm_pattern_shader(
-                dev, args.threads, args.loops, args.rows, args.ncols, args.unroll, args.order, args.bmode, args.r1,
-                check_pattern=True, store_group=group)
-            check_lib = inject(envelope, img_off, img_sz, reg_off, check_shader, fregs=fregs, hregs=check_hregs)
-            check_prg = dev.runtime('gemm_h', check_lib, buf_dtypes=buf_dtypes)
-            check_prg(a._buf, b._buf, c._buf, global_size=gs, local_size=ls, wait=True)
-            raw = c.copyout(memoryview(bytearray(c.nbytes))).cast('H')
-            row, col = divmod(group, args.ncols)
-            expected = [half_raw(args.unroll * sum(check_a[row][kk] * check_b[col][kk][lane] for kk in range(4))) for lane in range(4)]
-            bad = next((i for i, value in enumerate(raw[:4]) if value != expected[i]), None)
-            if bad is not None:
-                got = struct.unpack('<e', struct.pack('<H', raw[bad]))[0]
-                want = struct.unpack('<e', struct.pack('<H', expected[bad]))[0]
-                raise RuntimeError('GEMM CHECK FAIL group=%d index=%d got=%r expected=%r' % (group, bad, got, want))
-            checked += 4
-        print('GEMM CHECK PASS groups=%d scalar_outputs=%d bit_exact=true shape_per_thread=%dx%dx%d' %
-              (args.rows * args.ncols, checked, args.rows, args.ncols * 4, args.unroll * 4))
-
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--groups', type=int, default=16)
-    parser.add_argument('--rpt', type=int, choices=(0, 1, 3), default=3)
-    parser.add_argument('--loops', type=int, default=K4)
-    parser.add_argument('--unroll', type=int, default=1)
-    parser.add_argument('--pairs', type=int, default=8, help='compiler-pattern vector MAD pairs per loop')
-    parser.add_argument('--independent', action='store_true', help='remove loop-carried accumulator dependency for raw FMA issue peak')
-    parser.add_argument('--r1', action='store_true', help='auto-increment mad.f16 source1 across repeat lanes')
-    parser.add_argument('--compiler-pattern', action='store_true', help='use the vec16 OpenCL peak MAD source/destination pattern')
-    parser.add_argument('--gemm-pattern', action='store_true', help='use true GEMM-style acc=A_scalar*B_half4+acc MADs')
-    parser.add_argument('--check-gemm', action='store_true', help='use nonuniform exact inputs and bit-check every GEMM accumulator')
-    parser.add_argument('--rows', type=int, choices=(4, 8), default=4)
-    parser.add_argument('--ncols', type=int, default=4)
-    parser.add_argument('--bmode', choices=('reuse', 'percol'), default='reuse')
-    parser.add_argument('--order', choices=('kk_row_col', 'kk_col_row', 'col_kk_row', 'row_kk_col', 'row_col_kk'), default='kk_row_col')
-    parser.add_argument('--store', action='store_true', help='store one result after the ALU loop')
-    parser.add_argument('--threads', type=int, choices=(64, 128, 256), default=128)
-    parser.add_argument('--fregs', type=int, default=8)
-    parser.add_argument('--iters', type=int, default=20)
-    parser.add_argument('--allow-invalid-regs', action='store_true')
-    parser.add_argument('--disasm', action='store_true')
-    run(parser.parse_args())
@@ -1,440 +0,0 @@
-#!/usr/bin/env python3
-"""Hand-assembled GEMM kernels for Adreno 630.
-
-Tests:
-1. Pure ALU kernel (MAD throughput ceiling)
-2. Pure LOAD kernel (texture throughput ceiling)
-3. Full GEMM with optimal isam/mad interleaving
-"""
-import struct, ctypes, math
-from tinygrad import Device, dtypes
-from tinygrad.device import Buffer
-from extra.gemm.ir3asm import *
-
-dev = Device['QCOM']
-
-# ============================================================
-# DONOR KERNEL: compile the 4-row GEMM for the binary envelope
-# ============================================================
-DONOR_SRC = (
-    '#pragma OPENCL EXTENSION cl_khr_fp16 : enable\n'
-    'const sampler_t smp=CLK_NORMALIZED_COORDS_FALSE|CLK_ADDRESS_CLAMP|CLK_FILTER_NEAREST;\n'
-    '__attribute__((reqd_work_group_size(128,1,1)))\n'
-    '__kernel void gemm_h(read_only image2d_t A, read_only image2d_t B, __global half *C) {\n'
-    '  int lid=get_local_id(0); int tm=lid>>5; int tn=lid&31;\n'
-    '  int row=get_group_id(1)*16+tm*4; int col4=get_group_id(0)*32+tn;\n'
-    '  half4 r0c0=(half4)(0); for(int k4=0;k4<256;k4++){\n'
-    '    half4 a=read_imageh(A,smp,(int2)(k4,row));\n'
-    '    half4 b0=read_imageh(B,smp,(int2)(col4,k4*4));\n'
-    '    r0c0+=a.xxxx*b0;\n'
-    '  }\n'
-    '  vstore4(r0c0, 0, C+row*1024+col4*4);\n'
-    '}\n'
-)
-
-# Use the 4-row GEMM as donor since it has the right metadata for image textures
-_DONOR4 = (
-    '#pragma OPENCL EXTENSION cl_khr_fp16 : enable\n'
-    'const sampler_t smp=CLK_NORMALIZED_COORDS_FALSE|CLK_ADDRESS_CLAMP|CLK_FILTER_NEAREST;\n'
-    '__attribute__((reqd_work_group_size(128,1,1)))\n'
-    '__kernel void gemm_h(read_only image2d_t A, read_only image2d_t B, __global half *C) {\n'
-    '  int lid=get_local_id(0); int tm=lid>>5; int tn=lid&31;\n'
-    '  int row=get_group_id(1)*16+tm*4; int col4=get_group_id(0)*32+tn;\n'
-    '  half4 r0c0=(half4)(0),r0c1=(half4)(0),r0c2=(half4)(0),r0c3=(half4)(0);\n'
-    '  half4 r1c0=(half4)(0),r1c1=(half4)(0),r1c2=(half4)(0),r1c3=(half4)(0);\n'
-    '  half4 r2c0=(half4)(0),r2c1=(half4)(0),r2c2=(half4)(0),r2c3=(half4)(0);\n'
-    '  half4 r3c0=(half4)(0),r3c1=(half4)(0),r3c2=(half4)(0),r3c3=(half4)(0);\n'
-    '  for (int k4=0;k4<256;k4++) {\n'
-    '    half4 ar0=read_imageh(A,smp,(int2)(k4,row));\n'
-    '    half4 ar1=read_imageh(A,smp,(int2)(k4,row+1));\n'
-    '    half4 ar2=read_imageh(A,smp,(int2)(k4,row+2));\n'
-    '    half4 ar3=read_imageh(A,smp,(int2)(k4,row+3));\n'
-    '    half4 b0=read_imageh(B,smp,(int2)(col4,k4*4));\n'
-    '    half4 b1=read_imageh(B,smp,(int2)(col4,k4*4+1));\n'
-    '    half4 b2=read_imageh(B,smp,(int2)(col4,k4*4+2));\n'
-    '    half4 b3=read_imageh(B,smp,(int2)(col4,k4*4+3));\n'
-    '    r0c0+=ar0.xxxx*b0; r0c1+=ar0.yyyy*b1; r0c2+=ar0.zzzz*b2; r0c3+=ar0.wwww*b3;\n'
-    '    r1c0+=ar1.xxxx*b0; r1c1+=ar1.yyyy*b1; r1c2+=ar1.zzzz*b2; r1c3+=ar1.wwww*b3;\n'
-    '    r2c0+=ar2.xxxx*b0; r2c1+=ar2.yyyy*b1; r2c2+=ar2.zzzz*b2; r2c3+=ar2.wwww*b3;\n'
-    '    r3c0+=ar3.xxxx*b0; r3c1+=ar3.yyyy*b1; r3c2+=ar3.zzzz*b2; r3c3+=ar3.wwww*b3;\n'
-    '  }\n'
-    '  vstore4(r0c0+r0c1+r0c2+r0c3, 0, C+row*1024+col4*4);\n'
-    '  vstore4(r1c0+r1c1+r1c2+r1c3, 0, C+(row+1)*1024+col4*4);\n'
-    '  vstore4(r2c0+r2c1+r2c2+r2c3, 0, C+(row+2)*1024+col4*4);\n'
-    '  vstore4(r3c0+r3c1+r3c2+r3c3, 0, C+(row+3)*1024+col4*4);\n'
-    '}\n'
-)
-envelope, img_off, img_sz, reg_off = get_envelope(dev, _DONOR4)
-
-M, N, K = 1024, 1024, 1024
-K4 = K // 4  # 256
-
-def make_bufs():
-    a = Buffer(dev.device, (K//4)*M*4, dtypes.half, preallocate=True)
-    b = Buffer(dev.device, (N//4)*K*4, dtypes.half, preallocate=True)
-    c = Buffer(dev.device, M*N, dtypes.half, preallocate=True)
-    ctypes.memset(int(a._buf.va_addr), 0, a.nbytes)
-    ctypes.memset(int(b._buf.va_addr), 0, b.nbytes)
-    return a, b, c
-
-def bench(lib, gs, ls, label, flops=2*1024*1024*1024, iters=20):
-    a, b, c = make_bufs()
-    try:
-        prg = dev.runtime('gemm_h', lib, buf_dtypes=[((0, dtypes.half, (M, K//4, 4)),),
-                                                     ((1, dtypes.half, (K, N//4, 4)),),
-                                                     ((2, dtypes.half, None),)])
-        for _ in range(5):
-            prg(a._buf, b._buf, c._buf, global_size=gs, local_size=ls, wait=True)
-        times = []
-        for _ in range(iters):
-            t = prg(a._buf, b._buf, c._buf, global_size=gs, local_size=ls, wait=True)
-            if t: times.append(t)
-        if times:
-            best = min(times)
-            gflops = flops / best / 1e9
-            print("  %s: %.1f GFLOPS (%.0fus)" % (label, gflops, best*1e6))
-            return gflops
-    except Exception as e:
-        print("  %s: ERROR %s" % (label, str(e)[:80]))
-    return 0
-
-# ============================================================
-# Register plan for 4-row GEMM (matching the compiled kernel):
-#
-# Address/coordinate registers (full):
-#   r0.x(0)   = lid (hardware input)
-#   r0.y(1)   = group_id(1) + lid_row_offset
-#   r0.z(2)   = tm = lid >> 5
-#   r0.w(3)   = group_id(0) + lid_col_offset
-#   r2.y(9)   = A coord (k4 value for isam)
-#   r2.z(10)  = A row0 coord
-#   r2.w(11)  = A coord duplicate
-#   r3.x(12)  = A row0+1 coord
-#   r3.y(13)  = A coord dup
-#   r3.z(14)  = A row0+2 coord
-#   r3.w(15)  = A coord dup
-#   r4.x(16)  = A row0+3 coord
-#   r4.y(17)  = B col coord
-#   r4.z(18)  = B K offset
-#   r4.w(19)  = B K offset
-#   r5.y(21)  = B col coord dup
-#   r5.w(23)  = B col coord dup
-#   r6.x(24)  = temp
-#   r6.y(25)  = k4*4 base
-#   r6.z(26)  = k4 counter
-#   r7.x(28)  = row base addr
-#   r7.y(29)  = col4 base addr
-#
-# Texture result registers (half):
-#   hr0(0-3)   = A row3 texel (or temp)
-#   hr1(4-7)   = A row2 texel
-#   hr2(8-11)  = A row1 texel
-#   hr3(12-15) = A row0 texel
-#   hr4(16-19) = B texel (shared across all rows)
-#
-# Accumulator registers (half): 64 values = 16 groups of 4
-#   Row0: hr13.z(54)-hr16.y(65) = 4 groups: K0-K3
-#   Row1: hr17.z(70)-hr20.y(81) = 4 groups  [WRONG, let me read the actual mapping]
-#
-# Actually, from the disasm the accumulator mapping is:
-#   Row0 K0: hr20.z(82),hr20.w(83),hr21.x(84),hr21.y(85)
-#   Row0 K1: hr21.z(86),hr21.w(87),hr22.x(88),hr22.y(89) 
-#   Row0 K2: hr22.z(90),hr22.w(91),hr23.x(92),hr23.y(93)
-#   Row0 K3: hr23.z(94),hr23.w(95),hr24.x(96),hr24.y(97)
-#   Row1 K0: hr24.z(98),hr24.w(99),hr25.x(100),hr25.y(101)
-#   Row1 K1: hr25.z(102),hr25.w(103),hr26.x(104),hr26.y(105)
-#   Row1 K2: hr26.z(106),hr26.w(107),hr27.x(108),hr27.y(109)
-#   Row1 K3: hr27.z(110),hr27.w(111),hr28.x(112),hr28.y(113)
-#   Row2 K0: hr28.z(114),hr28.w(115),hr29.x(116),hr29.y(117)
-#   Row2 K1: hr29.z(118),hr29.w(119),hr30.x(120),hr30.y(121)
-#   Row2 K2: (from rpt1+rpt1, noncontiguous)
-#   Row2 K3: (from rpt3)
-#   Row3 K0: hr17.z(70),hr17.w(71),hr18.x(72),hr18.y(73)
-#   ... etc
-# This is messy. Let me use a CLEAN register plan instead.
-# ============================================================
-
-# ============================================================
-# TEST 1: PURE ALU - 16 (rpt3)mad.f16 in a loop, no texture loads
-# ============================================================
-
-print("=== TEST 1: Pure ALU (MAD throughput ceiling) ===")
-
-# Accumulator regs: hr20.x(80) through hr35.w(143) = 64 half-regs = 16 groups of 4
-# Source A: hr0.x(0) - hr0.w(3)
-# Source B: hr4.x(16) - hr7.w(31) (unused, just for mad operands)
-
-alu_instrs = [
-    MOV_S32('r6.z', 0, sy=True),     # counter = 0
-    MOV_H_IMM('hr0.x', 0x3c00),      # hr0.x = 1.0 (fp16)
-    MOV_H('hr0.y', 'hr0.x', rpt=2),  # hr0.y,z,w = 1.0
-    MOV_H_IMM('hr20.x', 0),           # zero first acc
-]
-# Zero all 64 accumulator regs (hr20.x=80 through hr35.w=143)
-for base in range(84, 144, 4):
-    alu_instrs.append(MOV_H(base, 80, rpt=3))
-# Set source B regs to 1.0
-for base in range(16, 32, 4):
-    alu_instrs.append(MOV_H(base, 0, rpt=3))
-
-# Loop label will be here
-loop_start = len(alu_instrs)
-
-# 16x (rpt3)mad.f16 = 64 MADs per iteration
-for g in range(16):
-    acc = 80 + g * 4           # accumulator base: hr20.x + g*4
-    src1 = g % 4               # hr0.x, hr0.y, hr0.z, hr0.w (cycling)
-    src2 = 16 + (g % 4) * 4   # hr4.x, hr5.x, hr6.x, hr7.x
-    alu_instrs.append(MAD_F16(acc, src1, src2, acc, rpt=3, r=True))
-
-# Loop control
-alu_instrs.append(ADD_S('r6.z', 'r6.z', 1))
-alu_instrs.append(CMPS_S_EQ('r6.z', K4 - 1))
-
-loop_end = len(alu_instrs)
-alu_instrs.append(BR(loop_start - loop_end))
-
-# Epilogue: sum and store (minimal - just write something)
-alu_instrs.append(ADD_F('hr0.x', 80, 84))
-alu_instrs.append(ADD_F('hr0.y', 88, 92))
-alu_instrs.append(ADD_F('hr0.z', 96, 100))
-alu_instrs.append(ADD_F('hr0.w', 104, 108))
-alu_instrs.append(NOP(rpt=5))
-alu_instrs.append(STG_F16('r0.z', 'hr0.x'))
-alu_instrs.append(END())
-
-shader_alu = assemble(alu_instrs)
-lib_alu = inject(envelope, img_off, img_sz, reg_off, shader_alu, fregs=8, hregs=64)
-
-print("  Shader: %d instrs (loop body: %d)" % (len(alu_instrs), loop_end - loop_start))
-print("  Disasm loop body:")
-asm = disasm(shader_alu)
-lines = asm.strip().split('\n')
-for line in lines[loop_start:loop_end+2]:
-    print("    " + line[:120])
-
-total_mads = 64 * K4  # 64 MADs per iter * 256 iters
-total_threads = 128 * (M // 128) * (M // 16)  # same grid as GEMM
-total_flops = total_mads * 2 * total_threads
-bench(lib_alu, (M//128, M//16, 1), (128, 1, 1), "PURE ALU", flops=total_flops)
-
-# ============================================================
-# TEST 2: PURE LOAD - 8 isam per iteration, accumulate results
-# ============================================================
-
-print("\n=== TEST 2: Pure LOAD (texture throughput ceiling) ===")
-
-# Same coordinate setup as the real GEMM but no MAD - just isam + add
-# We reuse the donor kernel's prologue for coordinate setup.
-# Actually let's just build it from scratch with minimal coord math.
-
-load_instrs = [
-    MOV_S32('r6.y', 3, sy=True),     # k4*4 base = 3 (initial)
-    MOV_S32('r6.z', 0),               # k4 counter = 0
-    MOV_H_IMM('hr20.x', 0),           # zero accumulator
-    MOV_H('hr20.y', 'hr20.x', rpt=2), # hr20.y,z,w = 0
-    # Compute row and col4 from lid
-    MOV_F32('r0.y', 'r52.x'),         # gid1 (from hardware constant)
-    NOP(rpt=2),
-    ADD_S('r0.y', 'r0.y', 0),         # r0.y = gid1 (simplified; real kernel adds c7.y)
-]
-# Copy the coordinate setup from the compiled kernel (lines 0-20)
-# Actually this is getting complex. Let me just build a simple version:
-# Use the compiled kernel verbatim but NOP out all the MADs.
-
-# Load the full 4-row donor kernel
-donor4 = (
-    '#pragma OPENCL EXTENSION cl_khr_fp16 : enable\n'
-    'const sampler_t smp=CLK_NORMALIZED_COORDS_FALSE|CLK_ADDRESS_CLAMP|CLK_FILTER_NEAREST;\n'
-    '__attribute__((reqd_work_group_size(128,1,1)))\n'
-    '__kernel void gemm_h(read_only image2d_t A, read_only image2d_t B, __global half *C) {\n'
-    '  int lid=get_local_id(0); int tm=lid>>5; int tn=lid&31;\n'
-    '  int row=get_group_id(1)*16+tm*4; int col4=get_group_id(0)*32+tn;\n'
-    '  half4 r0c0=(half4)(0),r0c1=(half4)(0),r0c2=(half4)(0),r0c3=(half4)(0);\n'
-    '  half4 r1c0=(half4)(0),r1c1=(half4)(0),r1c2=(half4)(0),r1c3=(half4)(0);\n'
-    '  half4 r2c0=(half4)(0),r2c1=(half4)(0),r2c2=(half4)(0),r2c3=(half4)(0);\n'
-    '  half4 r3c0=(half4)(0),r3c1=(half4)(0),r3c2=(half4)(0),r3c3=(half4)(0);\n'
-    '  for (int k4=0;k4<256;k4++) {\n'
-    '    half4 ar0=read_imageh(A,smp,(int2)(k4,row));\n'
-    '    half4 ar1=read_imageh(A,smp,(int2)(k4,row+1));\n'
-    '    half4 ar2=read_imageh(A,smp,(int2)(k4,row+2));\n'
-    '    half4 ar3=read_imageh(A,smp,(int2)(k4,row+3));\n'
-    '    half4 b0=read_imageh(B,smp,(int2)(col4,k4*4));\n'
-    '    half4 b1=read_imageh(B,smp,(int2)(col4,k4*4+1));\n'
-    '    half4 b2=read_imageh(B,smp,(int2)(col4,k4*4+2));\n'
-    '    half4 b3=read_imageh(B,smp,(int2)(col4,k4*4+3));\n'
-    '    r0c0+=ar0.xxxx*b0; r0c1+=ar0.yyyy*b1; r0c2+=ar0.zzzz*b2; r0c3+=ar0.wwww*b3;\n'
-    '    r1c0+=ar1.xxxx*b0; r1c1+=ar1.yyyy*b1; r1c2+=ar1.zzzz*b2; r1c3+=ar1.wwww*b3;\n'
-    '    r2c0+=ar2.xxxx*b0; r2c1+=ar2.yyyy*b1; r2c2+=ar2.zzzz*b2; r2c3+=ar2.wwww*b3;\n'
-    '    r3c0+=ar3.xxxx*b0; r3c1+=ar3.yyyy*b1; r3c2+=ar3.zzzz*b2; r3c3+=ar3.wwww*b3;\n'
-    '  }\n'
-    '  vstore4(r0c0+r0c1+r0c2+r0c3, 0, C+row*1024+col4*4);\n'
-    '  vstore4(r1c0+r1c1+r1c2+r1c3, 0, C+(row+1)*1024+col4*4);\n'
-    '  vstore4(r2c0+r2c1+r2c2+r2c3, 0, C+(row+2)*1024+col4*4);\n'
-    '  vstore4(r3c0+r3c1+r3c2+r3c3, 0, C+(row+3)*1024+col4*4);\n'
-    '}\n'
-)
-
-lib4, io4, isz4, ro4 = get_envelope(dev, donor4)
-shader4 = bytearray(lib4[io4:io4+isz4])
-total4 = isz4 // 8
-
-# NOP out all MAD instructions
-for i in range(total4):
-    lo, hi = struct.unpack_from('<II', shader4, i*8)
-    if (hi >> 24) in (0x63, 0x73) and ((hi >> 24) & 0xF) == 3:
-        struct.pack_into('<Q', shader4, i*8, 0)
-
-lib_load = inject(lib4, io4, isz4, ro4, shader4, fregs=8, hregs=31)
-bench(lib_load, (M//128, M//16, 1), (128, 1, 1), "PURE LOAD")
-
-# ============================================================
-# TEST 3: FULL GEMM - patched 4-row kernel (sy-stripped + rpt3)
-# ============================================================
-
-print("\n=== TEST 3: Patched GEMM (sy-stripped + rpt3) ===")
-
-# Take the compiled 4-row kernel, strip extra (sy), convert to rpt3
-shader_gemm = bytearray(lib4[io4:io4+isz4])
-
-# Strip extra (sy) flags - keep only the first one
-first_sy = False
-for i in range(total4):
-    lo, hi = struct.unpack_from('<II', shader_gemm, i*8)
-    if (hi >> 24) in (0x63, 0x73) and ((hi >> 24) & 0xF) == 3 and (hi >> 28) == 7:
-        if first_sy:
-            struct.pack_into('<I', shader_gemm, i*8+4, (hi & 0x0FFFFFFF) | 0x60000000)
-        else:
-            first_sy = True
-
-# Convert eligible 4-scalar MAD groups to (rpt3)
-i = 0
-while i < total4 - 3:
-    lo0, hi0 = struct.unpack_from('<II', shader_gemm, i*8)
-    if not ((hi0 >> 24) in (0x63, 0x73) and ((hi0 >> 24) & 0xF) == 3) or (hi0 >> 8) & 0x7F > 0 or (hi0 & 0xFF) != ((lo0 >> 16) & 0xFF):
-        i += 1; continue
-    d0, s1_0 = hi0 & 0xFF, lo0 & 0xFF
-    s2_0 = ((hi0 >> 16) & 0xFF) * 2 + (((hi0 >> 8) & 0xFF) >> 7)
-    ok = True
-    for j in range(1, 4):
-        lj, hj = struct.unpack_from('<II', shader_gemm, (i+j)*8)
-        if not ((hj >> 24) in (0x63, 0x73) and ((hj >> 24) & 0xF) == 3): ok = False; break
-        dj, rpj = hj & 0xFF, (hj >> 8) & 0x7F
-        s1j, s3j = lj & 0xFF, (lj >> 16) & 0xFF
-        s2j = ((hj >> 16) & 0xFF) * 2 + (((hj >> 8) & 0xFF) >> 7)
-        if rpj != 0 or s1j != s1_0 or dj != d0+j or s2j != s2_0+j or s3j != d0+j: ok = False; break
-    if ok:
-        rb = ((hi0 >> 8) & 0x80) | 3
-        struct.pack_into('<I', shader_gemm, i*8+4, (hi0 & 0xFFFF00FF) | (rb << 8))
-        struct.pack_into('<I', shader_gemm, i*8, lo0 | 0x20000000)
-        for j in range(1, 4): struct.pack_into('<Q', shader_gemm, (i+j)*8, 0)
-        i += 4
-    else:
-        i += 1
-
-# Merge (rpt1)+(rpt1) -> (rpt3)
-for i in range(total4 - 1):
-    lo0, hi0 = struct.unpack_from('<II', shader_gemm, i*8)
-    lo1, hi1 = struct.unpack_from('<II', shader_gemm, (i+1)*8)
-    if hi0 == 0 or hi1 == 0: continue
-    if not ((hi0 >> 24) in (0x63, 0x73) and ((hi0 >> 24) & 0xF) == 3): continue
-    if not ((hi1 >> 24) in (0x63, 0x73) and ((hi1 >> 24) & 0xF) == 3): continue
-    if (hi0 >> 8) & 0x7F != 1 or (hi1 >> 8) & 0x7F != 1: continue
-    d0, d1 = hi0 & 0xFF, hi1 & 0xFF
-    s10, s11 = lo0 & 0xFF, lo1 & 0xFF
-    s20 = ((hi0 >> 16) & 0xFF) * 2 + (((hi0 >> 8) & 0xFF) >> 7)
-    s21 = ((hi1 >> 16) & 0xFF) * 2 + (((hi1 >> 8) & 0xFF) >> 7)
-    if s10 != s11 or d1 != d0 + 2 or s21 != s20 + 2: continue
-    rb = ((hi0 >> 8) & 0x80) | 3
-    struct.pack_into('<I', shader_gemm, i*8+4, (hi0 & 0xFFFF00FF) | (rb << 8))
-    struct.pack_into('<Q', shader_gemm, (i+1)*8, 0)
-
-lib_gemm = inject(lib4, io4, isz4, ro4, shader_gemm, fregs=8, hregs=31)
-
-# Count stats
-asm_gemm = disasm(shader_gemm)
-print("  mad.f16: %d, (rpt3): %d, isam: %d, (sy): %d" % (
-    asm_gemm.count('mad.f16'), asm_gemm.count('(rpt3)mad.f16'),
-    asm_gemm.count('isam'), asm_gemm.count('(sy)')))
-
-bench(lib_gemm, (M//128, M//16, 1), (128, 1, 1), "PATCHED GEMM")
-
-# ============================================================
-# TEST 4: FULL GEMM at different sizes
-# ============================================================
-
-print("\n=== TEST 4: Patched GEMM at various sizes ===")
-for dim in [512, 768, 1024, 2048]:
-    if dim % 128 != 0 or dim % 16 != 0: continue
-    K4d = dim // 4
-    src_d = donor4.replace('k4<256', 'k4<%d' % K4d)
-    for s in ['row*1024', '(row+1)*1024', '(row+2)*1024', '(row+3)*1024']:
-        src_d = src_d.replace(s, s.replace('1024', str(dim)))
-    lib_d, io_d, isz_d, ro_d = get_envelope(dev, src_d)
-    s_d = bytearray(lib_d[io_d:io_d+isz_d])
-    t_d = isz_d // 8
-    # Apply same patches
-    fsy = False
-    for i in range(t_d):
-        lo, hi = struct.unpack_from('<II', s_d, i*8)
-        if (hi >> 24) in (0x63, 0x73) and ((hi >> 24) & 0xF) == 3 and (hi >> 28) == 7:
-            if fsy: struct.pack_into('<I', s_d, i*8+4, (hi & 0x0FFFFFFF) | 0x60000000)
-            else: fsy = True
-    i = 0
-    while i < t_d - 3:
-        lo0, hi0 = struct.unpack_from('<II', s_d, i*8)
-        if not ((hi0 >> 24) in (0x63, 0x73) and ((hi0 >> 24) & 0xF) == 3) or (hi0 >> 8) & 0x7F > 0 or (hi0 & 0xFF) != ((lo0 >> 16) & 0xFF):
-            i += 1; continue
-        d0, s1_0 = hi0 & 0xFF, lo0 & 0xFF
-        s2_0 = ((hi0 >> 16) & 0xFF) * 2 + (((hi0 >> 8) & 0xFF) >> 7)
-        ok = True
-        for j in range(1, 4):
-            lj, hj = struct.unpack_from('<II', s_d, (i+j)*8)
-            if not ((hj >> 24) in (0x63, 0x73) and ((hj >> 24) & 0xF) == 3): ok = False; break
-            dj, rpj = hj & 0xFF, (hj >> 8) & 0x7F
-            s1j, s3j = lj & 0xFF, (lj >> 16) & 0xFF
-            s2j = ((hj >> 16) & 0xFF) * 2 + (((hj >> 8) & 0xFF) >> 7)
-            if rpj != 0 or s1j != s1_0 or dj != d0+j or s2j != s2_0+j or s3j != d0+j: ok = False; break
-        if ok:
-            rb = ((hi0 >> 8) & 0x80) | 3
-            struct.pack_into('<I', s_d, i*8+4, (hi0 & 0xFFFF00FF) | (rb << 8))
-            struct.pack_into('<I', s_d, i*8, lo0 | 0x20000000)
-            for j in range(1, 4): struct.pack_into('<Q', s_d, (i+j)*8, 0)
-            i += 4
-        else: i += 1
-    for i in range(t_d - 1):
-        lo0, hi0 = struct.unpack_from('<II', s_d, i*8)
-        lo1, hi1 = struct.unpack_from('<II', s_d, (i+1)*8)
-        if hi0 == 0 or hi1 == 0: continue
-        if not ((hi0 >> 24) in (0x63, 0x73) and ((hi0 >> 24) & 0xF) == 3): continue
-        if not ((hi1 >> 24) in (0x63, 0x73) and ((hi1 >> 24) & 0xF) == 3): continue
-        if (hi0 >> 8) & 0x7F != 1 or (hi1 >> 8) & 0x7F != 1: continue
-        d0v, d1v = hi0 & 0xFF, hi1 & 0xFF
-        s10, s11 = lo0 & 0xFF, lo1 & 0xFF
-        s20 = ((hi0 >> 16) & 0xFF) * 2 + (((hi0 >> 8) & 0xFF) >> 7)
-        s21 = ((hi1 >> 16) & 0xFF) * 2 + (((hi1 >> 8) & 0xFF) >> 7)
-        if s10 != s11 or d1v != d0v + 2 or s21 != s20 + 2: continue
-        rb = ((hi0 >> 8) & 0x80) | 3
-        struct.pack_into('<I', s_d, i*8+4, (hi0 & 0xFFFF00FF) | (rb << 8))
-        struct.pack_into('<Q', s_d, (i+1)*8, 0)
-    ld = inject(lib_d, io_d, isz_d, ro_d, s_d, fregs=8, hregs=31)
-    M2 = N2 = K2 = dim
-    a2 = Buffer(dev.device, (K2//4)*M2*4, dtypes.half, preallocate=True)
-    b2 = Buffer(dev.device, (N2//4)*K2*4, dtypes.half, preallocate=True)
-    c2 = Buffer(dev.device, M2*N2, dtypes.half, preallocate=True)
-    ctypes.memset(int(a2._buf.va_addr), 0, a2.nbytes)
-    ctypes.memset(int(b2._buf.va_addr), 0, b2.nbytes)
-    try:
-        prg_d = dev.runtime('gemm_h', ld, [[(0, dtypes.imageh((M2, K2//4)))], [(1, dtypes.imageh((K2, N2//4)))], [(2, dtypes.half.ptr())]])
-        gs_d = (dim//128, dim//16, 1)
-        for _ in range(5): prg_d(a2._buf, b2._buf, c2._buf, global_size=gs_d, local_size=(128,1,1), wait=True)
-        ts = []
-        for _ in range(20):
-            t = prg_d(a2._buf, b2._buf, c2._buf, global_size=gs_d, local_size=(128,1,1), wait=True)
-            if t: ts.append(t)
-        if ts:
-            best = min(ts)
-            gf = 2*dim*dim*dim / best / 1e9
-            print("  %dx%d: %.1f GFLOPS (%.1fms)" % (dim, dim, gf, best*1e3))
-    except Exception as e:
-        print("  %d: ERROR %s" % (dim, str(e)[:60]))
@@ -1,106 +0,0 @@
-#!/usr/bin/env python3
-"""Patch openpilot's 4x16 FP32 GEMM with FP16 K4 partials and FP32 totals."""
-import argparse, itertools, pickle, struct
-
-from tinygrad.engine.jit import create_graph_call
-from tinygrad.uop.ops import Ops, UOp
-from extra.gemm.ir3asm import BR, CMPS_S_EQ, COV_F16F32, ISAM_F16, JUMP, MAD_F16, MAD_F32, MOV_F32, MOV_H_IMM, MOV_S32, NOP, inject
-from extra.gemm.qcom_ir3_matmul_patch import plain_name
-
-
-def blocked_image(image:bytes, block:int=1, direct_branch:bool=False, outer_iters:int|None=None, no_back_edge:bool=False) -> bytes:
-  instrs = [image[i:i+8] for i in range(0, len(image), 8)]
-  if len(instrs) != 349: raise ValueError(f"expected 349 instructions, got {len(instrs)}")
-  if block not in (1, 2, 4): raise ValueError(f"block must be 1, 2, or 4, got {block}")
-  # The compiler's loop is 47..100. Preserve its coordinate arithmetic and
-  # loop control, but sample native half vectors into a disjoint register bank.
-  # Each partial vector contains four output columns. Accumulate four scalar K
-  # terms per substep. Several substeps can share one partial before promotion.
-  body = [MOV_H_IMM(f"hr{34+row}.x", 0, rpt=3) for row in range(4)]
-  for substep in range(block):
-    # Drain the preceding half MADs before reusing their texture-source
-    # registers. Reissuing ISAM into a still-live half register can deadlock.
-    if substep: body += [MOV_F32("r0.x", "r0.x", sy=True), NOP(rpt=2)]
-    body += instrs[47:55]
-    for dst, coord in zip(("hr26.x", "hr27.x", "hr28.x", "hr29.x"), ("r0.x", "r1.x", "r2.x", "r3.x")):
-      body.append(ISAM_F16(dst, coord, 1, 1))
-    body += instrs[63:71]
-    for dst, coord in zip(("hr30.x", "hr31.x", "hr32.x", "hr33.x"), ("r4.x", "r5.x", "r6.x", "r7.x")):
-      body.append(ISAM_F16(dst, coord, 0, 0))
-    first = True
-    for kk in range(4):
-      for row in range(4):
-        body.append(MAD_F16(f"hr{34+row}.x", 4*(30+row)+kk, f"hr{26+kk}.x", f"hr{34+row}.x",
-                              rpt=3, sy=first, r=True))
-        first = False
-    # Keep the compare even between substeps: besides setting p0 it provides
-    # the latency slot needed by add r0.x -> mov r12.w. The final compare below
-    # overwrites p0 before loop control.
-    if substep != block-1: body += instrs[95:100]
-  # r4 is dead after all texture operations and supplies scalar 1.0 to vector
-  # MADs, giving FP32 total += promoted_partial without a separate add opcode.
-  body.append(MOV_S32("r4.x", 0x3f800000))
-  for row in range(4): body.append(COV_F16F32(f"r{row}.x", f"hr{34+row}.x", sy=(row == 0), rpt=3, r=True))
-  for row in range(4): body.append(MAD_F32(f"r{8+row}.x", "r4.x", f"r{row}.x", f"r{8+row}.x", rpt=3, r=True))
-  loop_limit = 95 if outer_iters is None else outer_iters*block-1
-  body += instrs[95:97] + [CMPS_S_EQ("r12.w", loop_limit, nop=1)] + instrs[98:100]
-
-  out = instrs[:47] + body
-  if no_back_edge:
-    pass
-  elif block == 1 or direct_branch:
-    out.append(BR(47-len(out), inv=True))
-  else:
-    # A6xx conditional branches have a much shorter reliable backward range
-    # than unconditional jumps. Branch past a long-range jump when complete.
-    branch_index = len(out)
-    out += [BR(2, inv=False), JUMP(47-(branch_index+1))]
-  out += instrs[101:]
-  while len(out) > len(instrs) and out[-1] == NOP(): out.pop()
-  if len(out) > len(instrs): raise ValueError(f"patched shader grew beyond envelope: {len(out)} > {len(instrs)}")
-  out += [NOP()] * (len(instrs)-len(out))
-  return b"".join(out)
-
-
-def patch_lib(lib:bytes, block:int, direct_branch:bool=False, outer_iters:int|None=None, no_back_edge:bool=False) -> bytes:
-  image_off = struct.unpack_from("<I", lib, 0xc0)[0]
-  image_size = struct.unpack_from("<I", lib, 0x100)[0]
-  reg_off = struct.unpack_from("<I", lib, 0x34)[0]
-  image = blocked_image(lib[image_off:image_off+image_size], block, direct_branch, outer_iters, no_back_edge)
-  return inject(lib, image_off, image_size, reg_off, image, fregs=13, hregs=38)
-
-
-def main() -> None:
-  parser = argparse.ArgumentParser()
-  parser.add_argument("input")
-  parser.add_argument("output")
-  parser.add_argument("--global-size", default="12,8,1")
-  parser.add_argument("--block", type=int, default=1)
-  parser.add_argument("--direct-branch", action="store_true")
-  parser.add_argument("--outer-iters", type=int, help="diagnostic loop limit; normal model execution requires 96/block iterations")
-  parser.add_argument("--no-back-edge", action="store_true", help="diagnostic: execute one outer body with no loop branch")
-  args = parser.parse_args()
-  target_global = tuple(int(x) for x in args.global_size.split(","))
-  with open(args.input, "rb") as f: jit = pickle.load(f)
-  slots = [x.arg.slot for x in jit.captured.linear.toposort()
-           if x.op is Ops.BUFFER and hasattr(x.arg, "slot") and x.arg.slot >= 0]
-  UOp.unique_num = itertools.count(max(slots, default=-1)+1)
-  outer = jit.captured.linear.src[0]
-  batch = outer.src[0].src[0].src
-  cache, replacements = {}, {}
-  for call in batch:
-    if call.op is not Ops.CALL or call.src[0].op is not Ops.PROGRAM: continue
-    program = call.src[0]
-    if plain_name(program.arg.name) != "gemm_h" or tuple(program.arg.global_size) != target_global: continue
-    old_lib = program.src[3].arg
-    new_lib = cache.setdefault(old_lib, patch_lib(old_lib, args.block, args.direct_branch, args.outer_iters, args.no_back_edge))
-    replacements[call] = call.replace(src=(program.replace(src=program.src[:3]+(program.src[3].replace(arg=new_lib),)), *call.src[1:]))
-  if not replacements: raise ValueError(f"no gemm_h calls with global size {target_global}")
-  new_outer = create_graph_call([replacements.get(call, call) for call in batch])
-  jit.captured._linear = jit.captured.linear.substitute({outer:new_outer}, walk=True)
-  jit.captured.__dict__.pop("linear", None)
-  with open(args.output, "wb") as f: pickle.dump(jit, f)
-  print(f"patched {len(replacements)} calls across {len(cache)} binaries with block={args.block} direct_branch={args.direct_branch}")
-
-
-if __name__ == "__main__": main()
@@ -1,90 +0,0 @@
-#!/usr/bin/env python3
-"""Validate and time an OpenCL blocked-half/FP32 GEMM on QCOM."""
-import argparse
-
-import numpy as np
-
-from tinygrad import Device, dtypes
-from tinygrad.device import Buffer
-
-
-def upload(values:np.ndarray, dtype) -> Buffer:
-  return Buffer("QCOM", values.size, dtype, initial_value=np.ascontiguousarray(values).tobytes())
-
-
-def source(m:int, n:int, k:int, stride:int, block4:int, linear:bool=False, ldib:bool=False) -> str:
-  assert m % 16 == 0 and n % 128 == 0 and k % (block4*4) == 0
-  image_type = "read_write image2d_t" if ldib else "read_only image1d_buffer_t" if linear else "read_only image2d_t"
-  def coord(index:str) -> str: return f"(int2)(({index})&16383,({index})>>14)"
-  def a_load(row:str) -> str: return coord(f"({row})*{k//4}+k4") if ldib else f"{row}*{k//4}+k4" if linear else f"(int2)(k4,{row})"
-  def b_load(krow:str) -> str: return coord(f"({krow})*{n//4}+col4") if ldib else f"{krow}*{n//4}+col4" if linear else f"(int2)(col4,{krow})"
-  image_args = "," if (linear or ldib) else ",smp,"
-  return f"""#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-const sampler_t smp=CLK_NORMALIZED_COORDS_FALSE|CLK_ADDRESS_CLAMP|CLK_FILTER_NEAREST;
-__attribute__((reqd_work_group_size(128,1,1)))
-__kernel void gemm_blocked({image_type} A,{image_type} B,__global float *C) {{
-  int lid=get_local_id(0), row=get_group_id(1)*16+(lid>>5)*4;
-  int col4=get_group_id(0)*32+(lid&31);
-  float4 t0=(float4)(0),t1=(float4)(0),t2=(float4)(0),t3=(float4)(0);
-  for(int kb=0;kb<{k//4};kb+={block4}) {{
-    half4 h0=(half4)(0),h1=(half4)(0),h2=(half4)(0),h3=(half4)(0);
-    #pragma unroll
-    for(int q=0;q<{block4};q++) {{
-      int k4=kb+q;
-      half4 a0=read_imageh(A{image_args}{a_load('row+0')});
-      half4 a1=read_imageh(A{image_args}{a_load('row+1')});
-      half4 a2=read_imageh(A{image_args}{a_load('row+2')});
-      half4 a3=read_imageh(A{image_args}{a_load('row+3')});
-      half4 b0=read_imageh(B{image_args}{b_load('k4*4+0')});
-      half4 b1=read_imageh(B{image_args}{b_load('k4*4+1')});
-      half4 b2=read_imageh(B{image_args}{b_load('k4*4+2')});
-      half4 b3=read_imageh(B{image_args}{b_load('k4*4+3')});
-      h0+=a0.xxxx*b0+a0.yyyy*b1+a0.zzzz*b2+a0.wwww*b3;
-      h1+=a1.xxxx*b0+a1.yyyy*b1+a1.zzzz*b2+a1.wwww*b3;
-      h2+=a2.xxxx*b0+a2.yyyy*b1+a2.zzzz*b2+a2.wwww*b3;
-      h3+=a3.xxxx*b0+a3.yyyy*b1+a3.zzzz*b2+a3.wwww*b3;
-    }}
-    t0+=convert_float4(h0);t1+=convert_float4(h1);t2+=convert_float4(h2);t3+=convert_float4(h3);
-  }}
-  vstore4(t0,0,C+(row+0)*{stride}+col4*4);vstore4(t1,0,C+(row+1)*{stride}+col4*4);
-  vstore4(t2,0,C+(row+2)*{stride}+col4*4);vstore4(t3,0,C+(row+3)*{stride}+col4*4);
-}}"""
-
-
-def main() -> None:
-  ap = argparse.ArgumentParser()
-  ap.add_argument("--m", type=int, default=128)
-  ap.add_argument("--n", type=int, default=1536)
-  ap.add_argument("--k", type=int, default=384)
-  ap.add_argument("--stride", type=int, default=2048)
-  ap.add_argument("--block4", type=int, default=4)
-  ap.add_argument("--seed", type=int, default=0)
-  ap.add_argument("--float-a", action="store_true", help="sample an FP32 activation image with read_imageh")
-  ap.add_argument("--linear", action="store_true", help="use image1d_buffer_t with explicit flattened indices")
-  ap.add_argument("--ldib", action="store_true", help="use read-write image2d_t and LDIB with flattened 2D indices")
-  args = ap.parse_args()
-  rng = np.random.default_rng(args.seed)
-  av = (rng.standard_normal((args.m, args.k))*0.05).astype(np.float32 if args.float_a else np.float16)
-  bv = (rng.standard_normal((args.k, args.n))*0.05).astype(np.float16)
-  a, b = upload(av, dtypes.float if args.float_a else dtypes.half), upload(bv, dtypes.half)
-  c = upload(np.zeros(args.m*args.stride, np.float32), dtypes.float)
-  src = source(args.m, args.n, args.k, args.stride, args.block4, args.linear, args.ldib)
-  if args.ldib:
-    ashape = ((args.m*(args.k//4)+16383)//16384, 16384, 4)
-    bshape = ((args.k*(args.n//4)+16383)//16384, 16384, 4)
-  else:
-    ashape = (1, args.m*(args.k//4), 4) if args.linear else (args.m, args.k//4, 4)
-    bshape = (1, args.k*(args.n//4), 4) if args.linear else (args.k, args.n//4, 4)
-  specs = [((0, dtypes.float if args.float_a else dtypes.half, ashape),),
-           ((1, dtypes.half, bshape),), ((2, dtypes.float, (args.m*args.stride,)),)]
-  program = Device["QCOM"].runtime("gemm_blocked", Device["QCOM"].compiler.compile(src), buf_dtypes=specs)
-  times = [program(a._buf, b._buf, c._buf, global_size=(args.n//128, args.m//16, 1),
-                   local_size=(128, 1, 1), wait=True)*1e3 for _ in range(8)]
-  storage = c.numpy().reshape(args.m, args.stride)
-  got, expected = storage[:, :args.n], av.astype(np.float32) @ bv.astype(np.float32)
-  delta = np.abs(got-expected)
-  print(f"block4={args.block4} ms={min(times):.4f} max_abs={float(delta.max()):.9g} "
-        f"mean_abs={float(delta.mean()):.9g} allclose={np.allclose(got, expected, rtol=1e-2, atol=1e-2)}")
-
-
-if __name__ == "__main__": main()
@@ -1,39 +0,0 @@
-#!/usr/bin/env python3
-"""Sweep QCOM compute texture/UAV partition registers on one captured model."""
-import argparse, os, pickle, time
-
-import numpy as np
-
-from tinygrad import Tensor
-from tinygrad.engine.realize import graph_cache
-
-
-def main() -> None:
-  parser = argparse.ArgumentParser()
-  parser.add_argument("model")
-  parser.add_argument("corpus")
-  parser.add_argument("--case", type=int, default=9)
-  parser.add_argument("--pairs", default="128:64,1:1,1:64,64:1,32:32,64:32,128:32,64:64")
-  parser.add_argument("--runs", type=int, default=5)
-  args = parser.parse_args()
-  with open(args.model, "rb") as f: model = pickle.load(f)
-  corpus = np.load(args.corpus)
-  inputs = {}
-  for name, (view, _vars, dtype, device) in zip(model.captured.expected_names, model.captured.expected_input_info):
-    key=f"case{args.case}:input:{name}"
-    inputs[name] = Tensor(corpus[key if key in corpus else name].astype(np.dtype(dtype.fmt), copy=False), device=device).realize()
-  output_key=f"case{args.case}:output"
-  expected = corpus[output_key if output_key in corpus else "out"]
-  for pair in args.pairs.split(","):
-    tsize, usize = pair.split(":")
-    os.environ["QCOM_TSIZE"], os.environ["QCOM_USIZE"] = tsize, usize
-    graph_cache.clear()
-    for _ in range(2): got = model(**inputs).numpy()
-    start = time.perf_counter()
-    for _ in range(args.runs): got = model(**inputs).numpy()
-    elapsed = (time.perf_counter()-start)*1000/args.runs
-    delta = np.abs(got.astype(np.float32)-expected.reshape(got.shape).astype(np.float32))
-    print(f"tsize={tsize} usize={usize} ms={elapsed:.3f} max_abs={float(delta.max()):.9g}")
-
-
-if __name__ == "__main__": main()
@@ -1,64 +0,0 @@
-#!/usr/bin/env python3
-"""Compare the exact cached target-3 GEMM with its graph replacement."""
-import argparse, pickle
-
-import numpy as np
-
-from tinygrad import Device, dtypes
-from tinygrad.device import Buffer
-from tinygrad.uop.ops import Ops
-from extra.gemm.qcom_ir3_matmul_patch import plain_name
-
-
-def upload(x:np.ndarray, dtype) -> Buffer:
-  ret = Buffer("QCOM", x.size, dtype).allocate()
-  ret.copyin(memoryview(np.ascontiguousarray(x)).cast("B"))
-  return ret
-
-
-def read(buf:Buffer, count:int, dtype) -> np.ndarray:
-  ret = np.empty(count, dtype=dtype)
-  buf.copyout(memoryview(ret).cast("B"))
-  return ret
-
-
-def batch(model): return model.captured.linear.src[0].src[0].src[0].src
-
-
-def main() -> None:
-  parser = argparse.ArgumentParser()
-  parser.add_argument("reference")
-  parser.add_argument("candidate")
-  args = parser.parse_args()
-  with open(args.reference, "rb") as f: reference = pickle.load(f)
-  with open(args.candidate, "rb") as f: candidate = pickle.load(f)
-  ref_call = next(c for c in batch(reference) if c.op is Ops.CALL and c.src[0].op is Ops.PROGRAM and
-                  plain_name(c.src[0].arg.name) == "gemm_h" and tuple(c.src[0].arg.global_size) == (12, 8, 1))
-  cbatch = batch(candidate)
-  cand_call = next(cbatch[i] for i in range(len(cbatch)-1) if cbatch[i].op is Ops.CALL and
-                   cbatch[i].src[0].op is Ops.PROGRAM and plain_name(cbatch[i+1].src[0].arg.name) == "cached_epi3")
-  rng = np.random.default_rng(7)
-  a_np = (rng.standard_normal((128, 384))*0.05).astype(np.float16)
-  a = upload(a_np, dtypes.half)
-  ref_out = upload(np.zeros(128*2048, np.float32), dtypes.float)
-  cand_out = upload(np.zeros(128*2048, np.float16), dtypes.half)
-  dev = Device["QCOM"]
-  ref_runtime = dev.runtime("ref", ref_call.src[0].src[3].arg, buf_dtypes=ref_call.src[0].arg.aux[0])
-  cand_runtime = dev.runtime("cand", cand_call.src[0].src[3].arg, buf_dtypes=cand_call.src[0].arg.aux[0])
-  ref_runtime(a._buf, ref_call.src[2].buffer._buf, ref_out._buf,
-              global_size=ref_call.src[0].arg.global_size, local_size=ref_call.src[0].arg.local_size, wait=True)
-  cand_runtime(a._buf, cand_call.src[2].buffer._buf, cand_out._buf,
-               global_size=cand_call.src[0].arg.global_size, local_size=cand_call.src[0].arg.local_size, wait=True)
-  ref = read(ref_out, 128*2048, np.float32).reshape(128, 2048)[:, :1536]
-  got = read(cand_out, 128*2048, np.float16).reshape(128, 2048)[:, :1536].astype(np.float32)
-  weight = np.asarray(ref_call.src[2].buffer.numpy()).reshape(384, 1536)
-  cpu = a_np.astype(np.float32) @ weight.astype(np.float32)
-  delta = np.abs(got-ref)
-  at = np.unravel_index(int(delta.argmax()), delta.shape)
-  print(f"max_abs={float(delta[at]):.9g} mean_abs={float(delta.mean()):.9g} at={at} "
-        f"got={float(got[at]):.9g} reference={float(ref[at]):.9g}")
-  print(f"weight_max={float(np.max(np.abs(weight))):.9g} cpu_max={float(np.max(np.abs(cpu))):.9g} "
-        f"reference_max={float(np.max(np.abs(ref))):.9g} candidate_max={float(np.max(np.abs(got))):.9g}")
-
-
-if __name__ == "__main__": main()
@@ -1,67 +0,0 @@
-#!/usr/bin/env python3
-"""Compare the cached exact target-3 GEMM against the THREAD128 FP16 hand kernel."""
-import pickle
-import numpy as np
-
-from tinygrad import Device, dtypes
-from tinygrad.device import Buffer
-from tinygrad.uop.ops import Ops
-from extra.gemm import qcom_intensity_gemm as q
-from extra.gemm.ir3asm import get_envelope, inject
-from extra.gemm.qcom_ir3_matmul_patch import plain_name
-
-
-def upload(x:np.ndarray, dtype) -> Buffer:
-  ret = Buffer("QCOM", x.size, dtype).allocate()
-  ret.copyin(memoryview(np.ascontiguousarray(x)).cast("B"))
-  return ret
-
-
-def read(buf:Buffer, count:int, dtype) -> np.ndarray:
-  ret = np.empty(count, dtype=dtype)
-  buf.copyout(memoryview(ret).cast("B"))
-  return ret
-
-
-with open("/data/openpilot_p3_rpt245679.pkl", "rb") as f: model = pickle.load(f)
-batch = model.captured.linear.src[0].src[0].src[0].src
-call = next(x for x in batch if x.op is Ops.CALL and x.src[0].op is Ops.PROGRAM and
-            plain_name(x.src[0].arg.name) == "gemm_h" and tuple(x.src[0].arg.global_size) == (12, 8, 1))
-dev, rng = Device["QCOM"], np.random.default_rng(7)
-a = upload((rng.standard_normal(128*384)*0.05).astype(np.float16), dtypes.half)
-a_np = read(a, 128*384, np.float16).reshape(128, 384)
-w_np = np.array(call.src[2].buffer.numpy(), copy=True).reshape(384, 384, 4).reshape(384, 1536)
-exact_out = upload(np.zeros(128*2048, np.float32), dtypes.float)
-hand_out = upload(np.zeros(128*2048, np.float16), dtypes.half)
-
-exact = dev.runtime("gemm_h", call.src[0].src[3].arg, buf_dtypes=call.src[0].arg.aux[0])
-exact(a._buf, call.src[2].buffer._buf, exact_out._buf,
-      global_size=call.src[0].arg.global_size, local_size=call.src[0].arg.local_size, wait=True)
-
-q.M, q.N, q.K, q.K4 = 128, 1536, 384, 96
-env, io, sz, ro = get_envelope(dev, q.make_direct_image_donor_src(4, 128))
-shader, _ = q.build_4xn_shader(dev, 128, ncols=4, direct=True, compact_acc=True,
-  stable_bx=True, stable_ay=True, inc_coords=True, persistent_coords=True,
-  first_sync_only=True, k_unroll=4, b_first=True, coord_delay=-1, stable_settle_delay=0,
-  store_row_shift=11, image_store=True, high_inputs=True)
-lib = inject(env, io, sz, ro, shader, fregs=10, hregs=48)
-hand = dev.runtime("gemm_h", lib, buf_dtypes=[((0, dtypes.half, (128, 512, 4)),),
-  ((0, dtypes.half, (128, 96, 4)),), ((1, dtypes.half, (384, 384, 4)),)])
-hand(hand_out._buf, a._buf, call.src[2].buffer._buf,
-     global_size=(3, 8, 1), local_size=(128, 1, 1), wait=True)
-
-expected = read(exact_out, 128*2048, np.float32).reshape(128, 2048)[:, :1536]
-got = read(hand_out, 128*2048, np.float16).reshape(128, 2048)[:, :1536].astype(np.float32)
-delta = np.abs(got-expected)
-cpu0 = a_np[0].astype(np.float32) @ w_np.astype(np.float32)
-for name, value in (("exact", expected[0]), ("hand", got[0])):
-  d_cpu = np.abs(value-cpu0)
-  print(name+"_cpu0", "max_abs", float(d_cpu.max()), "mean_abs", float(d_cpu.mean()))
-at = np.unravel_index(int(np.argmax(delta)), delta.shape)
-print("exact_hand", "max_abs", float(delta[at]), "mean_abs", float(delta.mean()), "at", at,
-      "got", float(got[at]), "expected", float(expected[at]))
-for tile in range(3):
-  d = np.abs(got[:, tile*512:(tile+1)*512]-expected[:, tile*512:(tile+1)*512])
-  print("tile", tile, "max_abs", float(d.max()), "mean_abs", float(d.mean()))
-print("timing_ms", min(hand(hand_out._buf, a._buf, call.src[2].buffer._buf,
-  global_size=(3,8,1), local_size=(128,1,1), wait=True) for _ in range(20))*1e3)
@@ -1,234 +0,0 @@
-#!/usr/bin/env python3
-"""Compact the GEMM loop by removing NOP instructions and adjusting branch offsets."""
-import struct, ctypes, tempfile
-from tinygrad import Device, dtypes
-from tinygrad.device import Buffer
-from tinygrad.runtime.autogen import mesa
-from tinygrad.helpers import data64
-
-dev = Device['QCOM']
-
-src = (
-    '#pragma OPENCL EXTENSION cl_khr_fp16 : enable\n'
-    'const sampler_t smp = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_CLAMP | CLK_FILTER_NEAREST;\n'
-    '__attribute__((reqd_work_group_size(128, 1, 1)))\n'
-    '__kernel void gemm_h(read_only image2d_t A, read_only image2d_t B, __global half *C) {\n'
-    '  int lid = get_local_id(0);\n'
-    '  int row = get_group_id(1) * 4 + (lid >> 5);\n'
-    '  int col4 = get_group_id(0) * 32 + (lid & 31);\n'
-    '  half4 acc0=(half4)(0), acc1=(half4)(0), acc2=(half4)(0), acc3=(half4)(0);\n'
-    '  for (int k4 = 0; k4 < 256; k4++) {\n'
-    '    half4 a = read_imageh(A, smp, (int2)(k4, row));\n'
-    '    half4 b0 = read_imageh(B, smp, (int2)(col4, k4*4));\n'
-    '    half4 b1 = read_imageh(B, smp, (int2)(col4, k4*4+1));\n'
-    '    half4 b2 = read_imageh(B, smp, (int2)(col4, k4*4+2));\n'
-    '    half4 b3 = read_imageh(B, smp, (int2)(col4, k4*4+3));\n'
-    '    acc0 += a.xxxx * b0;\n'
-    '    acc1 += a.yyyy * b1;\n'
-    '    acc2 += a.zzzz * b2;\n'
-    '    acc3 += a.wwww * b3;\n'
-    '  }\n'
-    '  half4 r = acc0 + acc1 + acc2 + acc3;\n'
-    '  vstore4(r, 0, C + row*1024 + col4*4);\n'
-    '}\n'
-)
-
-lib = bytearray(dev.compiler.compile_cached(src))
-image_offset = struct.unpack_from('<I', lib, 0xc0)[0]
-image_size_orig = struct.unpack_from('<I', lib, 0x100)[0]
-shader = bytearray(lib[image_offset:image_offset+image_size_orig])
-total = image_size_orig // 8
-
-def ri(buf, line):
-    off = line * 8
-    return struct.unpack_from('<I', buf, off+4)[0], struct.unpack_from('<I', buf, off)[0]
-
-def wi(buf, line, hi, lo):
-    off = line * 8
-    struct.pack_into('<I', buf, off, lo)
-    struct.pack_into('<I', buf, off+4, hi)
-
-def rn(r):
-    return "hr%d.%s" % (r // 4, "xyzw"[r % 4])
-
-def get_disasm(binary):
-    with tempfile.TemporaryFile('w+', buffering=1) as tf:
-        @ctypes.CFUNCTYPE(None, ctypes.c_void_p, ctypes.c_uint32, ctypes.c_void_p)
-        def hd(data, n, instr):
-            fst, snd = data64(ctypes.cast(instr, ctypes.POINTER(ctypes.c_uint64)).contents.value)
-            print(f"{n:04} [{fst:08x}_{snd:08x}] ", end="", flush=True, file=tf)
-        libc = ctypes.CDLL(None)
-        libc.setlinebuf(fp:=ctypes.cast(libc.fdopen(tf.fileno(), b"w"), ctypes.POINTER(mesa.struct__IO_FILE)))
-        mesa.ir3_isa_disasm(bytes(binary), len(binary), fp, mesa.struct_isa_decode_options(630, True, 0, True, pre_instr_cb=hd))
-        tf.seek(0)
-        return tf.read()
-
-# Step 1: Apply register remap (48->44, 49->45)
-for old_r, new_r in [(48, 44), (49, 45)]:
-    for i in range(total):
-        hi, lo = ri(shader, i)
-        if hi == 0 and lo == 0: continue
-        changed = False
-        if (hi & 0xFF) == old_r: hi = (hi & 0xFFFFFF00) | new_r; changed = True
-        if (lo & 0xFF) == old_r: lo = (lo & 0xFFFFFF00) | new_r; changed = True
-        if ((lo >> 16) & 0xFF) == old_r: lo = (lo & 0xFF00FFFF) | (new_r << 16); changed = True
-        if changed: wi(shader, i, hi, lo)
-
-# Step 2: Convert all eligible MAD groups to (rpt3)
-# First convert 4x scalar -> rpt3
-i = 0
-while i < total - 3:
-    hi0, lo0 = ri(shader, i)
-    if not ((hi0 >> 24) in (0x63, 0x73) and ((hi0 >> 24) & 0x0F) == 0x3): i += 1; continue
-    dst0, rpt0 = hi0 & 0xFF, (hi0 >> 8) & 0x7F
-    src1_0, src3_0 = lo0 & 0xFF, (lo0 >> 16) & 0xFF
-    src2_0 = ((hi0 >> 16) & 0xFF) * 2 + (((hi0 >> 8) & 0xFF) >> 7)
-    if rpt0 > 0 or dst0 != src3_0: i += 1; continue
-    ok = True
-    for j in range(1, 4):
-        hj, lj = ri(shader, i+j)
-        if not ((hj >> 24) in (0x63, 0x73) and ((hj >> 24) & 0x0F) == 0x3): ok = False; break
-        dj, rpj = hj & 0xFF, (hj >> 8) & 0x7F
-        s1j, s3j = lj & 0xFF, (lj >> 16) & 0xFF
-        s2j = ((hj >> 16) & 0xFF) * 2 + (((hj >> 8) & 0xFF) >> 7)
-        if rpj != 0 or s1j != src1_0 or dj != dst0+j or s2j != src2_0+j or s3j != dst0+j: ok = False; break
-    if ok:
-        rpt_byte_new = ((hi0 >> 8) & 0x80) | 3
-        hi_new = (hi0 & 0xFFFF00FF) | (rpt_byte_new << 8)
-        wi(shader, i, hi_new, lo0 | 0x20000000)
-        for j in range(1, 4): wi(shader, i+j, 0, 0)
-        i += 4
-    else:
-        i += 1
-
-# Merge (rpt1)+(rpt1) -> (rpt3)
-for i in range(total - 1):
-    hi0, lo0 = ri(shader, i)
-    hi1, lo1 = ri(shader, i+1)
-    if hi0 == 0 or hi1 == 0: continue
-    if not ((hi0 >> 24) in (0x63, 0x73) and ((hi0 >> 24) & 0x0F) == 0x3): continue
-    if not ((hi1 >> 24) in (0x63, 0x73) and ((hi1 >> 24) & 0x0F) == 0x3): continue
-    rpt0 = (hi0 >> 8) & 0x7F
-    rpt1v = (hi1 >> 8) & 0x7F
-    if rpt0 != 1 or rpt1v != 1: continue
-    dst0, dst1 = hi0 & 0xFF, hi1 & 0xFF
-    src1_0, src1_1 = lo0 & 0xFF, lo1 & 0xFF
-    src2_0 = ((hi0 >> 16) & 0xFF) * 2 + (((hi0 >> 8) & 0xFF) >> 7)
-    src2_1 = ((hi1 >> 16) & 0xFF) * 2 + (((hi1 >> 8) & 0xFF) >> 7)
-    if src1_0 != src1_1 or dst1 != dst0 + 2 or src2_1 != src2_0 + 2: continue
-    rpt_byte_new = ((hi0 >> 8) & 0x80) | 3
-    wi(shader, i, (hi0 & 0xFFFF00FF) | (rpt_byte_new << 8), lo0)
-    wi(shader, i+1, 0, 0)
-
-# Step 3: COMPACT - remove NOP instructions from the loop body
-# Find the branch and loop target
-branch_line = None
-for i in range(total):
-    hi, lo = ri(shader, i)
-    if (hi >> 20) == 0x009:
-        branch_line = i
-        br_offset_raw = lo
-        br_offset = struct.unpack('<i', struct.pack('<I', lo))[0]
-        target_line = i + 1 + br_offset
-
-if branch_line is None:
-    print("ERROR: no branch found")
-    exit(1)
-
-print("Branch at line %d, target line %d (offset %d)" % (branch_line, target_line, br_offset))
-
-# Count NOPs in the LOOP (between target_line and branch_line inclusive)
-loop_nops = []
-for i in range(target_line, branch_line + 1):
-    hi, lo = ri(shader, i)
-    if hi == 0 and lo == 0:
-        loop_nops.append(i)
-
-print("Loop body: lines %d-%d (%d instrs), %d NOPs to remove" % (
-    target_line, branch_line, branch_line - target_line + 1, len(loop_nops)))
-
-# Build new instruction stream: remove NOPs from the loop body
-# Also need to handle: some "NOPs" are actually (nop2), (nop3) etc which are
-# instruction modifiers, not standalone NOPs. Only remove pure 00000000_00000000 NOPs.
-new_instrs = []
-old_to_new = {}  # map old line numbers to new line numbers
-
-for i in range(total):
-    hi, lo = ri(shader, i)
-    # Remove pure NOPs that are inside the loop
-    if hi == 0 and lo == 0 and target_line <= i <= branch_line:
-        continue  # skip this NOP
-    old_to_new[i] = len(new_instrs)
-    new_instrs.append((hi, lo))
-
-new_total = len(new_instrs)
-print("Compacted: %d -> %d instructions (removed %d)" % (total, new_total, total - new_total))
-
-# Fix the branch offset
-if branch_line in old_to_new and target_line in old_to_new:
-    new_branch = old_to_new[branch_line]
-    new_target = old_to_new[target_line]
-    new_br_offset = new_target - new_branch - 1
-    # Update the branch instruction
-    br_hi, br_lo = new_instrs[new_branch]
-    new_instrs[new_branch] = (br_hi, struct.unpack('<I', struct.pack('<i', new_br_offset))[0])
-    print("Branch: old offset %d -> new offset %d" % (br_offset, new_br_offset))
-
-# Build new shader binary - KEEP SAME SIZE by padding with NOPs at the end
-new_shader = bytearray()
-for hi, lo in new_instrs:
-    new_shader += struct.pack('<II', lo, hi)
-
-# Pad to original size with end + nop instructions
-while len(new_shader) < image_size_orig:
-    new_shader += struct.pack('<II', 0x00000000, 0x00000000)  # nop padding
-
-new_image_size = image_size_orig  # keep same size!
-print("New shader: %d bytes = %d real instrs + %d padding" % (new_image_size, new_total, (image_size_orig - new_total*8)//8))
-
-# Don't resize - just replace shader in-place
-lib_new = bytearray(lib)
-lib_new[image_offset:image_offset+image_size_orig] = new_shader
-# image_size stays the same - no need to update
-
-# Verify disassembly
-print("\n=== COMPACTED KERNEL ===")
-asm = get_disasm(bytes(new_shader))
-mad_count = asm.count('mad.f16')
-rpt3_count = asm.count('(rpt3)mad.f16')
-isam_count = asm.count('isam')
-nop_count = asm.count('nop')
-print("instrs=%d mad=%d rpt3=%d isam=%d nop=%d" % (new_total, mad_count, rpt3_count, isam_count, nop_count))
-
-for line in asm.strip().split('\n'):
-    if line.strip():
-        print(line[:120])
-
-# Benchmark
-a_imgdt = dtypes.imageh((1024, 256))
-b_imgdt = dtypes.imageh((1024, 256))
-a_buf = Buffer(dev.device, 256*1024*4, dtypes.half, preallocate=True)
-b_buf = Buffer(dev.device, 256*1024*4, dtypes.half, preallocate=True)
-c_buf = Buffer(dev.device, 1024*1024, dtypes.half, preallocate=True)
-ctypes.memset(int(a_buf._buf.va_addr), 0, a_buf.nbytes)
-ctypes.memset(int(b_buf._buf.va_addr), 0, b_buf.nbytes)
-
-try:
-    prg = dev.runtime('gemm_h', bytes(lib_new), [[(0, a_imgdt)], [(1, b_imgdt)], [(2, dtypes.half.ptr())]])
-    gs = (1024 // 128, 1024 // 4, 1)
-    ls = (128, 1, 1)
-
-    for _ in range(5):
-        prg(a_buf._buf, b_buf._buf, c_buf._buf, global_size=gs, local_size=ls, wait=True)
-
-    times = []
-    for _ in range(30):
-        t = prg(a_buf._buf, b_buf._buf, c_buf._buf, global_size=gs, local_size=ls, wait=True)
-        if t: times.append(t)
-
-    if times:
-        best = min(times)
-        gflops = 2 * 1024 * 1024 * 1024 / best / 1e9
-        print("\n*** COMPACTED: %.1f GFLOPS (%.0fus) ***" % (gflops, best * 1e6))
-except Exception as e:
-    print("ERROR: %s" % str(e)[:200])
@@ -1,59 +0,0 @@
-#!/usr/bin/env python3
-"""Compare captured buffers after selected calls in two OpenPilot pickles."""
-import argparse
-import pickle
-
-import numpy as np
-
-from tinygrad import Tensor
-from tinygrad.engine.jit import _prepare_jit_inputs, create_graph_call
-from tinygrad.engine.realize import resolve_params, run_linear
-from tinygrad.uop.ops import Ops, UOp
-from extra.gemm.qcom_ir3_matmul_patch import plain_name
-
-
-def capture(model, corpus, name: str, arg_index: int) -> np.ndarray:
-  inputs = {key: Tensor(corpus[key], device=device).realize()
-            for key, (_view, _vars, _dtype, device) in zip(model.captured.expected_names, model.captured.expected_input_info)}
-  input_uops, var_vals, _names, _info = _prepare_jit_inputs((), inputs)
-  batch = model.captured.linear.src[0].src[0].src[0].src
-  index, call = next((i, call) for i, call in enumerate(batch) if call.op is Ops.CALL and
-                     call.src[0].op is Ops.PROGRAM and plain_name(call.src[0].arg.name) == name)
-  run_linear(UOp(Ops.LINEAR, src=(create_graph_call(list(batch[:index+1])),)), var_vals,
-             input_uops=input_uops, jit=True, wait=True)
-  resolved = resolve_params(call, tuple(input_uops))
-  output = resolved[call.src[0].arg.outs[0] if arg_index < 0 else arg_index]
-  return output.buffer.numpy().copy()
-
-
-def main() -> None:
-  parser = argparse.ArgumentParser()
-  parser.add_argument("left")
-  parser.add_argument("left_name")
-  parser.add_argument("right")
-  parser.add_argument("right_name")
-  parser.add_argument("corpus")
-  parser.add_argument("--left-arg", type=int, default=-1)
-  parser.add_argument("--right-arg", type=int, default=-1)
-  args = parser.parse_args()
-  with open(args.left, "rb") as f:
-    left = pickle.load(f)
-  with open(args.right, "rb") as f:
-    right = pickle.load(f)
-  corpus = np.load(args.corpus)
-  a = capture(left, corpus, args.left_name, args.left_arg)
-  b = capture(right, corpus, args.right_name, args.right_arg)
-  if a.size == 32*1088*4 and b.size == 2048*16*4:
-    image, expected = a.reshape(32, 1088, 4), np.empty((2048, 16, 4), dtype=a.dtype)
-    for row in range(2048):
-      idx1, block = row >> 2, row & 3
-      expected[row] = image[idx1 >> 4, (idx1 & 15)*68+block*17:(idx1 & 15)*68+block*17+16]
-    a = expected.reshape(-1)
-  delta = np.abs(a.astype(np.float32)-b.astype(np.float32))
-  print("shape", a.shape, b.shape, "max", float(delta.max()), "mean", float(delta.mean()))
-  print("left", a[:32])
-  print("right", b[:32])
-
-
-if __name__ == "__main__":
-  main()
--- a/Show More
+++ b/Show More