Merge branch 'master' into migrate_calL_expand

* hcq2: retain linked buffers with after

* hcq2: trim ref attachment comment

* x

.github/actions/setup-tinygrad/action.yml

@@ -61,6 +61,7 @@ runs:
         echo "MAX_BUFFER_SIZE=300000000" >> "$GITHUB_ENV"
         if [[ "$RUNNER_OS" == "Linux" ]]; then
           echo "VIRTUAL_ENV=/opt/venv/${{ inputs.python-version }}" >> "$GITHUB_ENV"
+          echo "UV_PYTHON_INSTALL_DIR=/opt/python" >> "$GITHUB_ENV"
         else
           echo "VIRTUAL_ENV=${{ github.workspace }}/.venv" >> "$GITHUB_ENV"
         fi
@@ -70,11 +71,6 @@ runs:
       with:
         enable-cache: 'false' # see below for manual caching
 
-    - name: Set up Python ${{ inputs.python-version }}
-      uses: actions/setup-python@v6
-      with:
-        python-version: ${{ inputs.python-version }}
-
     # **** Caching packages ****
 
     - name: Cache Python packages (PR)
@@ -198,21 +194,29 @@ runs:
         echo "pkgs=$pkgs" >> "$GITHUB_OUTPUT"
         echo "hash=$(echo -n "$pkgs" | sha256sum | cut -d' ' -f1)" >> "$GITHUB_OUTPUT"
 
+        installed=true
+        for pkg in $pkgs; do
+          info=$(dpkg-query -W -f='${db:Status-Abbrev} ${Version}' "$pkg" 2> /dev/null || true)
+          echo "${pkg}: ${info:-not in dpkg database}"
+          [[ "$info" == ii* ]] || installed=false
+        done
+        echo "installed=$installed" >> "$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' || inputs.ninja == 'true') && github.event_name == 'pull_request'
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && github.event_name == 'pull_request' && steps.apt-pkgs.outputs.installed == 'false'
       uses: actions/cache/restore@v5
       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' || inputs.ninja == 'true') && github.event_name != 'pull_request'
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && github.event_name != 'pull_request' && steps.apt-pkgs.outputs.installed == 'false'
       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' || inputs.ninja == 'true')
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && steps.apt-pkgs.outputs.installed == 'false'
       shell: bash
       run: |
         sudo apt -qq update || true

.github/workflows/autogen.yml

@@ -40,10 +40,10 @@ jobs:
       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
     - 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
+        find tinygrad/runtime/autogen -type f -name "*.py" -not -path "*/amd/*" -not -name "__init__.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 comgr_3, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd, amdgpu_drm"
+        python3 -c "from tinygrad.runtime.autogen import comgr, 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"
@@ -54,7 +54,7 @@ jobs:
         python3 -c "from tinygrad.runtime.autogen import mesa"
         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 mlx5, bnxt"
         python3 -c "from tinygrad.runtime.autogen import ggml_common"
         REGEN=1 python3 -c "from tinygrad.runtime.autogen import libclang"
     - name: Check for differences
@@ -102,42 +102,3 @@ jobs:
       with:
         name: autogen-macos-patch
         path: autogen-macos.patch
-
-  autogen-comgr-2:
-    name: In-tree Autogen (comgr 2)
-    runs-on: ubuntu-24.04
-    timeout-minutes: 15
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: 'autogen-comgr'
-    - name: Install autogen support packages
-      run: |
-        wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null
-        sudo tee /etc/apt/sources.list.d/rocm.list <<EOF
-        deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/6.2 $(lsb_release -cs) main
-        EOF
-        echo -e 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600' | sudo tee /etc/apt/preferences.d/rocm-pin-600
-        sudo apt -qq update || true
-        sudo apt-get install -y --no-install-recommends libclang-20-dev comgr
-    - name: Regenerate autogen files
-      run: |
-        rm tinygrad/runtime/autogen/comgr.py
-        python3 -c "from tinygrad.runtime.autogen import comgr"
-    - name: Check for differences
-      run: |
-        if ! git diff --quiet; then
-          git diff
-          git diff > autogen-comgr2.patch
-          echo "Autogen mismatch detected. Patch available at: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}#artifacts"
-          exit 1
-        fi
-    - name: Upload patch artifact
-      if: failure()
-      uses: actions/upload-artifact@v7
-      with:
-        name: autogen-comgr2-patch
-        path: autogen-comgr2.patch

.github/workflows/benchmark.yml

@@ -82,19 +82,22 @@ jobs:
   #      pytest -nauto --durations=20
 
   llmbenchmark:
-    name: LLM (DEV=${{ matrix.dev }})
+    name: Benchmark ${{ matrix.model }} (DEV=${{ matrix.dev }})
     runs-on: [self-hosted, "${{ matrix.dev == 'METAL' && 'macOS' || matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
     strategy:
       fail-fast: false
       matrix:
         dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 60
+        model: ['llama3.2:3b-f16', 'qwen3.8:27b', 'olmoe']
+        # qwen3.8:27b doesn't fit on mac
+        exclude: [{ dev: 'METAL', model: 'qwen3.8:27b' }, { dev: 'AMD', model: 'olmoe' }, { dev: 'NV', model: 'olmoe' }]
+    timeout-minutes: 15
     defaults:
       run:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: "0"
+      HCQ2: '0'
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -102,12 +105,11 @@ jobs:
     - name: Setup (AMD)
       if: ${{ matrix.dev == 'AMD' }}
       run: |
-        ./extra/amdpci/setup_python_cap.sh
-        ./extra/hcq/hcq_smi.py amd rmmod
-        ./extra/hcq/hcq_smi.py amd kill_pids
+        ./extra/hcq/hcq_smi.py amd rmmod --expect
+        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
     - name: Setup (NV)
       if: ${{ matrix.dev == 'NV' }}
-      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
+      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
     - name: setup staging db
       if: github.ref == 'refs/heads/update_benchmark_staging'
       run: |
@@ -115,16 +117,10 @@ jobs:
         rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
     - name: reset process replay
       run: python3 test/external/process_replay/reset.py
-    - name: Run llama3.2
-      run: BENCHMARK_LOG=llama32_3b-f16 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m llama3.2:3b-f16 --benchmark --warmup
-    - name: Run qwen3.6
-      # qwen3.6:35b-a3b doesn't fit on mac
-      if: ${{ matrix.dev != 'METAL' }}
-      run: BENCHMARK_LOG=qwen36_35b-a3b JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m qwen3.6:35b-a3b --benchmark --warmup
-    - name: Run olmoe
-      # just metal for now
-      if: ${{ matrix.dev == 'METAL' }}
-      run: BENCHMARK_LOG=olmoe JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m olmoe --benchmark --warmup
+    - name: Run ${{ matrix.model }}
+      run: |
+        MODEL=${{ matrix.model }}
+        BENCHMARK_LOG=${MODEL//./} JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m $MODEL --benchmark --warmup
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
 
@@ -135,13 +131,13 @@ jobs:
       fail-fast: false
       matrix:
         dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 60
+    timeout-minutes: 10
     defaults:
       run:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: "0"
+      HCQ2: '0'
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -149,12 +145,11 @@ jobs:
     - name: Setup (AMD)
       if: ${{ matrix.dev == 'AMD' }}
       run: |
-        ./extra/amdpci/setup_python_cap.sh
-        ./extra/hcq/hcq_smi.py amd rmmod
-        ./extra/hcq/hcq_smi.py amd kill_pids
+        ./extra/hcq/hcq_smi.py amd rmmod --expect
+        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
     - name: Setup (NV)
       if: ${{ matrix.dev == 'NV' }}
-      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
+      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
     - name: setup staging db
       if: github.ref == 'refs/heads/update_benchmark_staging'
       run: |
@@ -184,13 +179,13 @@ jobs:
       fail-fast: false
       matrix:
         dev: ['AMD', 'NV']
-    timeout-minutes: 60
+    timeout-minutes: 5
     defaults:
       run:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: "0"
+      HCQ2: '0'
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -198,12 +193,11 @@ jobs:
     - name: Setup (AMD)
       if: ${{ matrix.dev == 'AMD' }}
       run: |
-        ./extra/amdpci/setup_python_cap.sh
-        ./extra/hcq/hcq_smi.py amd rmmod
-        ./extra/hcq/hcq_smi.py amd kill_pids
+        ./extra/hcq/hcq_smi.py amd rmmod --expect
+        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
     - name: Setup (NV)
       if: ${{ matrix.dev == 'NV' }}
-      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
+      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
     - name: Symlink models and datasets
       run: |
         mkdir -p extra/datasets
@@ -227,13 +221,13 @@ jobs:
       fail-fast: false
       matrix:
         dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 60
+    timeout-minutes: 15
     defaults:
       run:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: "0"
+      HCQ2: '0'
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -241,12 +235,11 @@ jobs:
     - name: Setup (AMD)
       if: ${{ matrix.dev == 'AMD' }}
       run: |
-        ./extra/amdpci/setup_python_cap.sh
-        ./extra/hcq/hcq_smi.py amd rmmod
-        ./extra/hcq/hcq_smi.py amd kill_pids
+        ./extra/hcq/hcq_smi.py amd rmmod --expect
+        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
     - name: Setup (NV)
       if: ${{ matrix.dev == 'NV' }}
-      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
+      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
     - name: setup staging db
       if: github.ref == 'refs/heads/update_benchmark_staging'
       run: |
@@ -273,13 +266,13 @@ jobs:
       fail-fast: false
       matrix:
         dev: ['AMD', 'NV']
-    timeout-minutes: 60
+    timeout-minutes: 20
     defaults:
       run:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: "0"
+      HCQ2: '0'
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -287,12 +280,11 @@ jobs:
     - name: Setup (AMD)
       if: ${{ matrix.dev == 'AMD' }}
       run: |
-        ./extra/amdpci/setup_python_cap.sh
-        ./extra/hcq/hcq_smi.py amd rmmod
-        ./extra/hcq/hcq_smi.py amd kill_pids
+        ./extra/hcq/hcq_smi.py amd rmmod --expect
+        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
     - name: Setup (NV)
       if: ${{ matrix.dev == 'NV' }}
-      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
+      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
     - name: Symlink models and datasets
       run: |
         mkdir -p weights
@@ -327,7 +319,7 @@ jobs:
       fail-fast: false
       matrix:
         dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 60
+    timeout-minutes: 10
     defaults:
       run:
         shell: bash -e -o pipefail {0}
@@ -340,12 +332,11 @@ jobs:
     - name: Setup (AMD)
       if: ${{ matrix.dev == 'AMD' }}
       run: |
-        ./extra/amdpci/setup_python_cap.sh
-        ./extra/hcq/hcq_smi.py amd rmmod
-        ./extra/hcq/hcq_smi.py amd kill_pids
+        ./extra/hcq/hcq_smi.py amd rmmod --expect
+        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
     - name: Setup (NV)
       if: ${{ matrix.dev == 'NV' }}
-      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
+      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
     - name: setup staging db
       if: github.ref == 'refs/heads/update_benchmark_staging'
       run: |
@@ -407,9 +398,10 @@ jobs:
     - name: Test benchmark allreduce
       if: ${{ matrix.dev == 'NV' }}
       run: python test/external/external_benchmark_multitensor_allreduce.py
-    - name: HEVC Decode Benchmark
-      if: ${{ matrix.dev == 'NV' }}
-      run: VALIDATE=1 MAX_FRAMES=100 ASSERT_FPS=1400 JITBEAM=1 PYTHONPATH=. python3 extra/hevc/decode.py
+    # TODO: HEVC decode timing test
+    # - name: HEVC Decode Benchmark
+    #   if: ${{ matrix.dev == 'NV' }}
+    #   run: IGNORE_BEAM_CACHE=1 VALIDATE=1 MAX_FRAMES=100 ASSERT_FPS=1400 JITBEAM=1 PYTHONPATH=. python3 extra/hevc/decode.py
     - uses: actions/upload-artifact@v7
       if: ${{ matrix.dev != 'AMD' }}
       with:
@@ -422,7 +414,7 @@ jobs:
   testusbgpu:
     name: UsbGPU Benchmark
     runs-on: [self-hosted, macOS]
-    timeout-minutes: 10
+    timeout-minutes: 3
     defaults:
       run:
         shell: bash -e -o pipefail {0}
@@ -437,32 +429,70 @@ jobs:
         rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
     - name: Kill stale pids
       run: |
-        PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
-        PYTHONPATH=. ./extra/hcq/hcq_smi.py nv kill_pids
-    # since sudo is required for usbgpu on macos, do not write bytecode, as some of the files are owned by root
+        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
+        ./extra/hcq/hcq_smi.py nv kill_pids --sudoless
     - name: UsbGPU boot time
-      run: sudo -E PYTHONDONTWRITEBYTECODE=1 PYTHONPATH=. GMMU=0 DEBUG=2 AM_RESET=1 DEV=USB+AMD time python3.11 test/test_tiny.py TestTiny.test_plus
+      run: GMMU=0 DEBUG=2 AM_RESET=1 DEV=USB+AMD time python3.11 test/test_tiny.py TestTiny.test_plus
     - name: UsbGPU tiny tests
-      run: sudo -E PYTHONDONTWRITEBYTECODE=1 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/test_tiny.py
+      run: GMMU=0 DEV=USB+AMD python3.11 test/test_tiny.py
     - name: UsbGPU copy speeds
-      run: sudo -E PYTHONDONTWRITEBYTECODE=1 SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
-    #- name: UsbGPU openpilot test
-    #  run: sudo -E PYTHONPATH=. GMMU=0 DEV=USB+AMD GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+      run: SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
     - name: UsbGPU (USB4/TB) install script
-      run: PYTHONPATH=. sh extra/setup_tinygpu_osx.sh
+      run: sh extra/setup_tinygpu_osx.sh
     - name: UsbGPU (USB4/TB) boot time
-      run: PYTHONPATH=. DEBUG=3 DEV=PCI+NV:NAK time python3.11 test/test_tiny.py TestTiny.test_plus
+      run: DEBUG=3 DEV=PCI+NV:NAK time python3.11 test/test_tiny.py TestTiny.test_plus
     - name: UsbGPU (USB4/TB) tiny tests
-      run: PYTHONPATH=. DEV=PCI+NV:NAK python3.11 test/test_tiny.py
+      run: DEV=PCI+NV:NAK python3.11 test/test_tiny.py
 
-  testcommalatest:
-    name: comma Benchmark (0.11.2)
-    runs-on: [self-hosted, Linux, comma]
-    timeout-minutes: 12
+  testcomma:
+    strategy:
+      matrix:
+        dev: ['QCOM', 'QCOM:IR3']
+        version: ['0.11.0', '0.11.2']
+        model: ['vision', 'policy', 'supercombo', 'dmonitoring']
+        # exclude non-existent models
+        exclude: [{ version: '0.11.0', model: supercombo }, { version: '0.11.2', model: vision }, { version: '0.11.2', model: policy }]
+        include:
+          - version: '0.11.0'
+            model: vision
+            url: https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
+            timing: 18
+          - version: '0.11.0'
+            model: policy
+            url: https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_policy.onnx
+            timing: 3.4
+          - version: '0.11.0'
+            model: dmonitoring
+            url: https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/dmonitoring_model.onnx
+            timing: 13
+          - version: '0.11.2'
+            model: supercombo
+            url: https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/433f85f956837606ad1f1cbee4aa7e2158ad23c768dea914b20436c97232741b
+            timing: 28
+          - dev: QCOM:IR3
+            version: '0.11.2'
+            model: supercombo
+            timing: 29
+          - version: '0.11.2'
+            model: dmonitoring
+            url: https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/3e7b31dfbc0a5234f1baf196513b77fc6af12204b8a8ffe8ee0417e48352f316
+            timing: 12.5
+          # IR3 dmonitoring is slightly slower
+          - dev: QCOM:IR3
+            model: dmonitoring
+            timing: 13.0
+      fail-fast: false
+    name: openpilot ${{ matrix.version }} compile3 ${{ matrix.model }} (DEV=${{ matrix.dev }})
+    runs-on: [self-hosted, Linux, comma4]
+    timeout-minutes: 5
     defaults:
       run:
         shell: bash -e -o pipefail {0}
     if: github.repository_owner == 'tinygrad'
+    env:
+      DEV: ${{ matrix.dev }}
+      ASSERT_MIN_STEP_TIME: ${{ matrix.timing }}
+      BENCHMARK_LOG: ${{ matrix.dev == 'QCOM:IR3' && 'ir3_' || '' }}openpilot_${{ matrix.version }}_${{ matrix.model }}
     steps:
     - name: Checkout Code
       uses: actions/checkout@v6
@@ -473,45 +503,10 @@ jobs:
         rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
     - name: reset process replay
       run: test/external/process_replay/reset.py
-    - name: openpilot compile3 0.11.2 supercombo
-      run: BENCHMARK_LOG=openpilot_0_11_2_supercombo PYTHONPATH="." ASSERT_MIN_STEP_TIME=26 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/433f85f956837606ad1f1cbee4aa7e2158ad23c768dea914b20436c97232741b
-    - name: openpilot compile3 0.11.2 supercombo (from pickle)
-      run: BENCHMARK_LOG=openpilot_0_11_2_supercombo_run_pickle RUN_PICKLE=1 PYTHONPATH="." ASSERT_MIN_STEP_TIME=26 DEV=QCOM taskset -c 4-7 python3 examples/openpilot/compile3.py
-    - name: IR3 openpilot compile3 0.11.2 supercombo
-      run: BENCHMARK_LOG=ir3_openpilot_0_11_2_supercombo PYTHONPATH="." ASSERT_MIN_STEP_TIME=41 DEV=QCOM:IR3 FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/433f85f956837606ad1f1cbee4aa7e2158ad23c768dea914b20436c97232741b
-    - name: openpilot compile3 0.11.2 dmonitoring
-      run: BENCHMARK_LOG=openpilot_0_11_2_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/3e7b31dfbc0a5234f1baf196513b77fc6af12204b8a8ffe8ee0417e48352f316
-    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
-
-  testcommaold:
-    name: comma Benchmark (0.11.0)
-    runs-on: [self-hosted, Linux, comma]
-    timeout-minutes: 12
-    defaults:
-      run:
-        shell: bash -e -o pipefail {0}
-    if: github.repository_owner == 'tinygrad'
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: setup staging db
-      if: github.ref == 'refs/heads/update_benchmark_staging'
-      run: |
-        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
-        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
-    - name: reset process replay
-      run: test/external/process_replay/reset.py
-    - name: openpilot compile3 0.11.0 driving_vision
-      run: BENCHMARK_LOG=openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
-    - name: openpilot compile3 0.11.0 driving_vision (from pickle)
-      run: BENCHMARK_LOG=openpilot_0_11_0_vision_run_pickle RUN_PICKLE=1 PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM taskset -c 4-7 python3 examples/openpilot/compile3.py
-    - name: IR3 openpilot compile3 0.11.0 driving_vision
-      run: BENCHMARK_LOG=ir3_openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=18 DEV=QCOM:IR3 FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
-    - name: openpilot compile3 0.11.0 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_11_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3.2 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_policy.onnx
-    - name: openpilot compile3 0.11.0 dmonitoring
-      run: BENCHMARK_LOG=openpilot_0_11_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/dmonitoring_model.onnx
+    - name: compile
+      run: FLOAT16=1 IMAGE=1 taskset -c 4-7 python3 examples/openpilot/compile3.py ${{ matrix.url }} openpilot.pkl
+    - name: run pickle
+      run: BENCHMARK_LOG="${BENCHMARK_LOG}_run_pickle" RUN_PICKLE=1 taskset -c 4-7 python3 examples/openpilot/compile3.py - openpilot.pkl
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
 
@@ -524,15 +519,6 @@ jobs:
         shell: bash -e -o pipefail {0}
     if: github.repository_owner == 'tinygrad'
     steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: setup staging db
-      if: github.ref == 'refs/heads/update_benchmark_staging'
-      run: |
-        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
-        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
-    - name: reset process replay
-      run: test/external/process_replay/reset.py
     - name: Checkout Code
       uses: actions/checkout@v6
     - name: setup staging db
@@ -545,8 +531,8 @@ jobs:
     - name: benchmark MobileNetV2 on DSP
       run: |
         # generate quantized weights
-        ln -s /data/home/tiny/tinygrad/extra/datasets/imagenet extra/datasets/imagenet
-        ln -s /data/home/tiny/tinygrad/testsig-*.so .
+        ln -s ~/tinygrad/extra/datasets/imagenet extra/datasets/imagenet
+        ln -s ~/tinygrad/testsig-*.so .
         PYTHONPATH=. DEV=CPU QUANT=1 CNT=0 python3 examples/test_onnx_imagenet.py https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx /tmp/model.quant.onnx
         # benchmark on DSP with NOOPT=1, the devectorizer has issues
         PYTHONPATH=. DEV=DSP NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
@@ -556,7 +542,7 @@ jobs:
   testcommausbgpubenchmark:
     name: UsbGPU Benchmark (comma)
     runs-on: [self-hosted, Linux, comma4]
-    timeout-minutes: 20
+    timeout-minutes: 10
     defaults:
       run:
         shell: bash -e -o pipefail {0}
@@ -585,7 +571,7 @@ jobs:
       fail-fast: false
       matrix:
         dev: ['AMD', 'NV']
-    timeout-minutes: 20
+    timeout-minutes: 5
     defaults:
       run:
         shell: bash -e -o pipefail {0}
@@ -597,9 +583,8 @@ jobs:
       uses: actions/checkout@v6
     - name: Setup
       run: |
-        ./extra/amdpci/setup_python_cap.sh
-        ./extra/hcq/hcq_smi.py ${{ matrix.dev }} rmmod
-        ./extra/hcq/hcq_smi.py ${{ matrix.dev }} kill_pids
+        ./extra/hcq/hcq_smi.py ${{ matrix.dev }} rmmod --expect
+        ./extra/hcq/hcq_smi.py ${{ matrix.dev }} kill_pids --sudoless
         mkdir -p extra/datasets
         ln -s /raid/datasets/imagenet extra/datasets/imagenet
     - name: setup staging db
@@ -634,6 +619,10 @@ jobs:
       run: |
         GRAPH_ONE_KERNEL=1 NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
         GRAPH_ONE_KERNEL=1 NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
+    # TODO: HEVC decode timing test
+    # - name: HEVC Decode Benchmark
+    #   if: ${{ matrix.dev == 'NV' }}
+    #   run: IGNORE_BEAM_CACHE=1 VALIDATE=1 MAX_FRAMES=100 ASSERT_FPS=1400 JITBEAM=1 PYTHONPATH=. python3 extra/hevc/decode.py
     - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
       if: ${{ matrix.dev == 'NV' }}
       run: BENCHMARK_LOG=resnet_10steps MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
@@ -656,12 +645,12 @@ jobs:
   llvmspeed:
     name: LLVM Speed
     runs-on: [self-hosted, Linux, tinyboxrandom]
-    timeout-minutes: 20
+    timeout-minutes: 10
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
       uses: actions/checkout@v6
     - name: Speed Test
-      run: DEV=CPU:LLVM THREADS=0 python3 test/speed/external_test_speed_v_torch.py
+      run: DEV=CPU:LLVM python3 test/speed/external_test_speed_v_torch.py
     - name: Speed Test (BEAM=2)
-      run: BEAM=2 DEV=CPU:LLVM THREADS=0 python3 test/speed/external_test_speed_v_torch.py
+      run: IGNORE_BEAM_CACHE=1 BEAM=2 DEV=CPU:LLVM python3 test/speed/external_test_speed_v_torch.py

.github/workflows/platform.yml

@@ -36,6 +36,8 @@ jobs:
         deps: testing_unit
     - name: Run unit tests
       run: DEV=METAL python -m pytest -n=auto test/unit/ --durations=20
+    - name: Run opt tests
+      run: DEV=METAL python -m pytest -n=auto test/opt --durations=20
     - name: Test tensor core ops (fake)
       run: DEV=METAL DEBUG=3 TC=2 python test/backend/test_ops.py TestOps.test_gemm
     - name: Test tensor core ops (real)

.github/workflows/test.yml

@@ -233,7 +233,7 @@ jobs:
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
     - name: Repo line count <= 26000 lines
-      run: MAX_LINE_COUNT=26000 python sz.py
+      run: MAX_LINE_COUNT=26500 python sz.py
 
   spec:
     strategy:
@@ -390,10 +390,10 @@ jobs:
         run: |
           parallel --link --tagstring '[{1}]' '{2}' \
             ::: llama 'llama q4' qwen3.5 qwen \
-            ::: $'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --model llama3.2:1b | tee /dev/stderr | grep -i rooster' \
-                $'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --model llama3.2:1b-q4 | tee /dev/stderr | grep -i rooster' \
-                $'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --model qwen3.5:0.8b | tee /dev/stderr | grep -i rooster' \
-                $'echo "What\'s a female chicken called? Answer with only one word." | python3 -m tinygrad.llm --model qwen3:0.6b | tee /dev/stderr | grep -i hen'
+            ::: $'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --no_chat_template --model llama3.2:1b | tee /dev/stderr | grep -i rooster' \
+                $'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --no_chat_template --model llama3.2:1b-q4 | tee /dev/stderr | grep -i rooster' \
+                $'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --no_chat_template --model qwen3.5:0.8b | tee /dev/stderr | grep -i rooster' \
+                $'echo "What\'s a female chicken called? Answer with only one word." | python3 -m tinygrad.llm --no_chat_template --model qwen3:0.6b | tee /dev/stderr | grep -i hen'
                 # NOTE: qwen is dumb and only knows about female chickens
 
 # ****** Models Tests ******
@@ -504,7 +504,7 @@ jobs:
       - name: Run AMD renderer tests (AMD:LLVM)
         run: DEV=MOCKKFD+AMD:LLVM python -m pytest -n=auto test/amd/ --durations 20
       - name: Run SQTT profiling tests
-        run: PROFILE=1 SQTT=1 python3 -m pytest -n=auto test/amd/test_sqtt_profiler.py
+        run: VIZ=-2 python3 -m pytest -n=auto test/amd/test_sqtt_profiler.py
       - name: Run AMD emulated tests on NULL backend
         env:
           AMD: 0
@@ -518,34 +518,15 @@ jobs:
       - name: Run LLVM test
         run: DEV=MOCKKFD+AMD:LLVM python test/device/test_amd_llvm.py
 
-  hcq2:
-    name: hcq2
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
-    timeout-minutes: 5
-    steps:
-      - name: Checkout Code
-        uses: actions/checkout@v6
-      - name: Setup Environment
-        uses: ./.github/actions/setup-tinygrad
-        with:
-          key: hcq2
-          deps: testing_unit
-          amd: 'true'
-      - name: Run HCQ2 tests
-        run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/test_tiny.py
-      - name: Run HCQ2 multi-device tests
-        run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python -m pytest -n=auto test/backend/test_multitensor.py
-      - name: Run HCQ2 JIT tests
-        run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/unit/test_jit.py
-      - name: Run HCQ2 unit tests
-        run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python -m pytest test/device/test_hcq2.py
-
   testmockam:
     name: Linux (am)
     runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 15
     env:
       DEV: MOCKPCI+AMD
+      HCQ2: 1
+      HCQ_RUNTIME_DEV: PYTHON
+      PYTHONPATH: .
     steps:
       - name: Checkout Code
         uses: actions/checkout@v6
@@ -555,15 +536,17 @@ jobs:
           key: mockam
           deps: testing_unit
           amd: 'true'
-      - name: Run test_tiny on MOCKAM
-        run: python test/test_tiny.py
+      - name: Run tests on MOCKAM
+        run: python -m pytest test/test_tiny.py test/unit/test_jit.py
       - name: Run test_tiny on MOCKUSB
-        run: GMMU=0 DEV=MOCKUSB+AMD python test/test_tiny.py
-      - name: Run test_hcq on MOCKPCI
-        run: python -m pytest test/device/test_hcq.py
+        run: HCQ2=0 GMMU=0 DEV=MOCKUSB+AMD python test/test_tiny.py
+      - name: Run test_hcq2 on MOCKPCI
+        run: python -m pytest test/device/test_hcq2.py
       - name: Run disk copy tests on MOCKPCI
         run: python -m pytest test/unit/test_disk_tensor.py -k test_copy_from_disk
       - name: Run test_tiny on MOCKPCI Remote
+        env:
+          HCQ2: 0
         run: |
           python extra/remote/serve.py 6667 &
           sleep 2
@@ -583,6 +566,9 @@ jobs:
     env:
       DEV: MOCKKFD+AMD:${{ matrix.backend == 'amdllvm' && 'LLVM' || '' }}:${{ matrix.arch }}
       SKIP_SLOW_TEST: 1
+      HCQ2: 1
+      HCQ_RUNTIME_DEV: PYTHON
+      PYTHONPATH: .
     steps:
       - name: Checkout Code
         uses: actions/checkout@v6
@@ -599,9 +585,11 @@ jobs:
           DEBUG=5 FORWARD_ONLY=1 python3 test/test_tiny.py TestTiny.test_plus
       - name: Run MXFP4 Llama training on NULL backend
         if: ${{ matrix.backend == 'amd' && matrix.arch == 'gfx950' }}
-        run: PYTHONPATH=. DEV=NULL:HIP:gfx950 MXFP4=1 LLAMA_LAYERS=2 BENCHMARK=3 NULL_ALLOW_COPYOUT=1 NO_HIPCC=1 ROCM_PATH=/opt/rocm JITBEAM=0 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/profile.sh
+        run: HCQ2=0 PYTHONPATH=. DEV=NULL:HIP:gfx950 MXFP4=1 LLAMA_LAYERS=2 BENCHMARK=3 NULL_ALLOW_COPYOUT=1 NO_HIPCC=1 ROCM_PATH=/opt/rocm JITBEAM=0 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/profile.sh
       - name: Run pytest (amd)
-        run: python -m pytest -n=auto test/backend/test_ops.py test/backend/test_dtype.py test/backend/test_dtype_alu.py test/backend/test_linearizer.py test/backend/test_randomness.py test/backend/test_jit.py test/backend/test_graph.py test/backend/test_multitensor.py test/device/test_hcq.py test/external/external_test_am.py test/backend/test_asm_gemm.py::TestAsmGEMM --durations=20
+        run: python -m pytest -n=auto test/backend/test_ops.py test/backend/test_dtype.py test/backend/test_dtype_alu.py test/backend/test_linearizer.py test/backend/test_randomness.py test/backend/test_jit.py test/backend/test_graph.py test/backend/test_multitensor.py test/device/test_hcq2.py test/external/external_test_am.py test/backend/test_asm_gemm.py::TestAsmGEMM --durations=20
+      - name: Run opt tests
+        run: python -m pytest -n=auto test/opt --durations=20
       - name: Run disk copy tests
         run: python -m pytest test/unit/test_disk_tensor.py -k test_copy_from_disk
       - name: Run TRANSCENDENTAL math
@@ -679,4 +667,5 @@ jobs:
         run: |
           python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
           DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
-          python -m pytest -n=auto test/backend/test_ops.py --durations=20
+          # QCOMCL compiles in qemu, too slow for parallel workers
+          ${{ contains(matrix.dev, 'QCOMCL') && 'PARALLEL=0' || '' }} python -m pytest -n=auto test/backend/test_ops.py --durations=20

.gitignore

@@ -69,3 +69,4 @@ mutants
 dagre/
 graphlib/
 uv.lock
+pi_session_window0.jsonl

.pre-commit-config.yaml

@@ -28,7 +28,7 @@ repos:
         pass_filenames: false
       - id: tests
         name: comprehensive test suite
-        entry: env OMP_NUM_THREADS=1 SKIP_SLOW_TEST=1 PYTHONPATH="." python3 -m pytest -n=6 test/backend/test_ops.py test/backend/test_schedule.py test/unit/test_assign.py test/backend/test_tensor.py test/backend/test_jit.py test/unit/test_schedule_cache.py test/null/test_pattern_matcher.py test/null/test_uop_symbolic.py test/unit/test_helpers.py
+        entry: env OMP_NUM_THREADS=1 SKIP_SLOW_TEST=1 PYTHONPATH="." python3 -m pytest -n=6 test/backend/test_ops.py test/backend/test_schedule.py test/backend/test_assign.py test/backend/test_tensor.py test/backend/test_jit.py test/unit/test_schedule_cache.py test/null/test_pattern_matcher.py test/null/test_uop_symbolic.py test/unit/test_helpers.py
         language: system
         always_run: true
         pass_filenames: false

AGENTS.md

@@ -4,3 +4,4 @@
 - Run `python -m mypy tinygrad/` to typecheck
 - Run `python -m ruff check .` to lint
 - Read `./tinygrad/viz/README.md` for profiling and debugging rewrite rules
+- Do not do amend commits. Always do a new commit if a force push to origin would be required.

conftest.py

@@ -2,7 +2,7 @@ import os, pytest, signal, threading
 
 @pytest.hookimpl(wrapper=True)
 def pytest_runtest_call(item):
-  t = threading.Timer(int(os.getenv("TEST_TIMEOUT", 300)), os.kill, args=(os.getpid(), signal.SIGABRT))
+  t = threading.Timer(int(os.getenv("TEST_TIMEOUT", 90)), os.kill, args=(os.getpid(), signal.SIGABRT))
   t.start()
   try: yield
   finally:

docs/abstractions4.py

@@ -122,7 +122,7 @@ 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.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in self.instructions]))))
 
   CU_COUNT = 32
   LANES = 64

docs/developer/speed.md

@@ -52,7 +52,7 @@ In `kernel.py` we have a set of `OptOps`, these control the parameters of the sp
 
 The main bottleneck in most kernels is accessing memory. In a freshman algorithms class, you'll learn about cache aware matrix multiplication, and this is all forms of that. While the same math is run, the order in which you run it can have large impacts on the speed depending on if the data you are loading. OptOps will change this order.
 
-Memory, even cache, is often much slower than accessing the register file. The amount of times data is used in math is called the "arithmetic intensity". For operations like BS=1 GEMV, the arithmetic intensity is 1, but for GEMMs and convs it can be much higher. OptOps like UPCAST and UNROLL can increase this, but be careful of making them too large, as if there's too much register pressure on the GPU the warp scheduler may not be able to fit many warps, or even worse, it could be spilling to local memory.
+Memory, even cache, is often much slower than accessing the register file. The amount of times data is used in math is called the "arithmetic intensity". For operations like BS=1 GEMV, the arithmetic intensity is 1, but for GEMMs and convs it can be much higher. Splitting an axis into UPCAST can increase this, but be careful of making them too large, as if there's too much register pressure on the GPU the warp scheduler may not be able to fit many warps, or even worse, it could be spilling to local memory.
 
 4090s have 1 TB/s of ram bandwidth and ~160 TFLOPS of compute, so you need to use each loaded value ~100 times. The L1 cache has around 40 TB/s of bandwidth, so in order to get full compute utilization you need to use each value ~4 times.
 

examples/mlperf/dataloader.py

@@ -5,7 +5,7 @@ from multiprocessing import Queue, Process, shared_memory, connection, Lock
 
 import numpy as np
 from tinygrad import dtypes, Tensor
-from tinygrad.helpers import getenv, prod, Context, round_up, tqdm, OSX, NUM_CPU_THREADS
+from tinygrad.helpers import getenv, prod, Context, round_up, tqdm, OSX, CPU_COUNT
 from tinygrad.nn.state import TensorIO
 
 ### ResNet
@@ -131,7 +131,7 @@ def batch_load_resnet(batch_size=64, val=False, shuffle=True, seed=None, pad_fir
     else: X = Tensor.empty(*sz, dtype=dtypes.uint8, device=f"disk:/dev/shm/{shm_name}")
     Y = [None] * (batch_size*BATCH_COUNT)
 
-    for _ in range(NUM_CPU_THREADS.value):
+    for _ in range(CPU_COUNT):
       p = Process(target=loader_process, args=(q_in, q_out, X, seed))
       p.daemon = True
       p.start()
@@ -212,7 +212,7 @@ def batch_load_train_bert(BS:int, seed:int|None=None):
     rng.shuffle(fs)
     train_files.append(fs.pop(0))
 
-  cycle_length = min(NUM_CPU_THREADS.value, len(train_files))
+  cycle_length = min(CPU_COUNT, len(train_files))
   assert cycle_length > 0, "cycle_length must be greater than 0"
 
   dataset = InterleavedDataset(train_files, cycle_length)
@@ -301,7 +301,7 @@ def batch_load_unet3d(preprocessed_dataset_dir:Path, batch_size:int=6, val:bool=
     X = Tensor.empty(*sz, dtype=dtypes.float32, device=f"disk:/dev/shm/{shm_name_x}")
     Y = Tensor.empty(*sz, dtype=dtypes.uint8, device=f"disk:/dev/shm/{shm_name_y}")
 
-    for _ in range(NUM_CPU_THREADS.value):
+    for _ in range(CPU_COUNT):
       proc = Process(target=load_unet3d_data, args=(preprocessed_dataset_dir, seed, queue_in, queue_out, X, Y))
       proc.daemon = True
       proc.start()
@@ -437,7 +437,7 @@ def batch_load_retinanet(dataset, val:bool, base_dir:Path, batch_size:int=32, sh
   dataset_iter = iter(image_ids)
 
   try:
-    for _ in range(NUM_CPU_THREADS.value):
+    for _ in range(CPU_COUNT):
       proc = Process(
         target=load_retinanet_data,
         args=(base_dir, val, queue_in, queue_out, imgs, boxes, labels),

examples/mlperf/model_train.py

@@ -1667,15 +1667,14 @@ def train_llama3():
 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, GradAccClipAdamWGroup, clip_grads
+  from examples.mlperf.optim import GradAccClipAdamW, GradAccClipAdamWGroup, fclip_grads
 
   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
+  GBS                = config["GLOBAL_BATCH_SIZE"]      = BS
   SEED               = config["SEED"]                   = getenv("SEED", 5760)
   DATA_SEED          = config["DATA_SEED"]              = getenv("DATA_SEED", SEED)
   SEQLEN             = config["SEQLEN"]                 = getenv("SEQLEN", 8192)
@@ -1737,8 +1736,8 @@ def train_gptoss():
   params_wd = [p for p in params if p.ndim >= 3]
   params_no_wd = [p for p in params if p.ndim < 3]
   optim = GradAccClipAdamWGroup(
-    GradAccClipAdamW(params_wd, 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),
-    GradAccClipAdamW(params_no_wd, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=0.0, grad_acc=grad_acc, device=optim_device),
+    GradAccClipAdamW(params_wd, 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=1, device=optim_device),
+    GradAccClipAdamW(params_no_wd, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=0.0, grad_acc=1, device=optim_device),
   )
 
   for p in optim.params:
@@ -1770,31 +1769,32 @@ def train_gptoss():
 
   @TinyJit
   @Context(TRAINING=1)
-  def minibatch(tokens:Tensor):
+  def 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], save=True)
-    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
+    if getenv("FUSED_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 = 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)
+    Tensor.realize(loss, *grads)
 
-  @TinyJit
-  def optim_step():
-    grad_norm = clip_grads(grads, grad_acc, 1.0)
-    optim.fstep(grads, grad_norm)
+    clipped_grads, grad_norm = fclip_grads(grads, 1.0)
+    optim.fstep(clipped_grads, grad_norm)
     scheduler.step()
 
-    for g in grads: g.assign(0)
-
+    loss_cpu = loss.flatten().float().to("CPU")
     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)
+    Tensor.realize(loss_cpu, lr_cpu, grad_norm_cpu, *grads, *fp8_inv_scales)
 
-    return lr_cpu, grad_norm_cpu
+    return loss_cpu, lr_cpu, grad_norm_cpu
 
   @TinyJit
   @Context(TRAINING=0)
@@ -1843,30 +1843,20 @@ def train_gptoss():
       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
+      ist = time.perf_counter()
 
-      gt = time.perf_counter()
-      ret = optim_step()
-      lr, grad_norm = ret[0].item(), ret[1].item()
+      try: tokens = next(train_iter)
+      except StopIteration: break
+      mst = time.perf_counter()
+      data_time = mst - ist
+
+      ret = step(tokens)
+      dev_time = time.perf_counter() - mst
+
+      loss, lr, grad_norm = ret[0].item(), ret[1].item(), ret[2].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
@@ -1876,7 +1866,7 @@ def train_gptoss():
       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"{i:5} {step_time:.3f} s step, {dev_time:.3f} s dev, {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())))
 
@@ -1886,8 +1876,6 @@ def train_gptoss():
           "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,

examples/mlperf/models/gpt_oss.py

@@ -14,7 +14,7 @@ from extra.models.llama import apply_rotary_emb
 from extra.llama_kernels.rmsnorm import rmsnorm
 from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8, asm_gemm, can_use_asm_gemm
 from extra.gemm.moe_gemm import grouped_mx_gemm
-from extra.gemm.moe_routing import route, dispatch, combine
+from extra.gemm.moe_routing import route, dispatch, combine, router_mfma
 
 FP8_DTYPE = dtypes.fp8e4m3
 FP8_MAX = 448.0
@@ -41,7 +41,7 @@ def _quant_dequant_bwd(grad:UOp, call:UOp) -> tuple:
 
 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))
+  return Tensor(fxn.uop.call_with_output(x.uop, grad_fxn=_quant_dequant_bwd))
 
 def _mx_scale(e8:Tensor) -> Tensor:
   return _mx_block_scale(e8) if e8.ndim == 2 else _mx_block_scale_3d(e8)
@@ -58,10 +58,27 @@ def _dequant_bwd(grad:UOp, call:UOp) -> tuple:
 
 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))
+  return Tensor(fxn.uop.call_with_output(w_q.uop, w_scale.uop, grad_fxn=_dequant_bwd))
 
-def matmul_mx(x:Tensor, w_q:Tensor, w_scale:Tensor) -> Tensor:
+def matmul_mx(x:Tensor|tuple[Tensor, Tensor], w_q:Tensor, w_scale:Tensor) -> Tensor:
+  if isinstance(x, tuple):
+    assert ASM_GEMM, "pre-quantized MXFP8 input requires ASM_GEMM"
+    from extra.gemm.cdna_asm_gemm import asm_gemm, can_use_asm_gemm, mx_pack
+    x_q, x_e8 = x
+    l_shape, padded = x_q.shape[:-1], x_q.shape[-1]
+    x_q, x_e8 = x_q.reshape(-1, padded), x_e8.reshape(-1, padded // 32)
+    K, N = w_q.shape[1], w_q.shape[0]
+    assert padded >= K and (padded - K) % 32 == 0 and x_e8.shape[-1] == padded // 32
+    wq, ws = w_q, w_scale
+    if (pad := padded - K):
+      wq = wq.pad(((0, 0), (0, pad)))
+      ws = ws.pad(((0, 0), (0, pad // 32)), value=127).cast(dtypes.uint8)
+    if (npad := (-N) % 256):
+      wq = wq.pad(((0, npad), (0, 0)))
+      ws = ws.pad(((0, npad), (0, 0)), value=127).cast(dtypes.uint8)
+    assert can_use_asm_gemm(x_q, wq.T)
+    out = asm_gemm(x_q, wq.T, mx=True, mx_scales=(mx_pack(x_e8), x_e8, mx_pack(ws), ws), mx_w_stored=True)
+    return (out[:, :N] if npad else out).reshape(*l_shape, N).cast(dtypes.bfloat16)
   l_shape = x.shape[:-1]
   if ASM_GEMM:
     from extra.gemm.cdna_asm_gemm import asm_gemm, can_use_asm_gemm, mx_pack
@@ -175,8 +192,22 @@ class GPTOSS:
   def attention(self, x:Tensor, freqs_cis:Tensor, mask:Tensor, sliding:bool, *, 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
+
+    if getenv("FUSED_RMSNORM_MX", 0):
+      from extra.gptoss_kernels.rmsnorm import rmsnorm_mul_quantize_mxfp8
+      x_q, x_e8, rrms = rmsnorm_mul_quantize_mxfp8(x, attention_norm, self.norm_eps)
+      qkv = matmul_mx((x_q, x_e8), wqkv, wqkv_scale) + wqkv_bias
+      norm_saves = [x_q, x_e8, rrms]
+    if getenv("FUSED_RMSNORM_MUL", 0):
+      from extra.gptoss_kernels.rmsnorm import rmsnorm_mul
+      x_normed, rrms = rmsnorm_mul(x, attention_norm, self.norm_eps)
+      qkv = matmul_mx(x_normed, wqkv, wqkv_scale) + wqkv_bias
+      norm_saves = [x_normed, rrms]
+    else:
+      x_normed, rrms = rmsnorm(x, self.norm_eps)
+      qkv = matmul_mx(x_normed * attention_norm, wqkv, wqkv_scale) + wqkv_bias
+      norm_saves = [x_normed, rrms]
+
     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]
@@ -202,14 +233,20 @@ class GPTOSS:
       attn = (w @ xvm).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] + fa_saves
+    return out, [attn] + norm_saves + fa_saves
 
   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()
+    if getenv("FUSED_RMSNORM_MUL", 0):
+      from extra.gptoss_kernels.rmsnorm import rmsnorm_mul
+      x_normed, rrms = rmsnorm_mul(x, ffn_norm, self.norm_eps)
+      inp = x_normed
+    else:
+      x_normed, rrms = rmsnorm(x, self.norm_eps)
+      inp = x_normed * ffn_norm
+
+    logits = router_mfma(inp, gate, gate_bias) if getenv("ROUTER_MFMA", 0) else inp.float() @ gate.float().T + gate_bias.float()
     dim, inter = self.dim, self.intermediate_size
 
     if getenv("GROUPED_MOE", 0):
@@ -282,14 +319,14 @@ 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))
+    grad_buf.uop = grad_buf.uop.after(grad_buf.uop.store(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)])
+      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)))
+      cur = cur.after(buf_slice.store(pad.src[0].cast(cur.dtype)))
     else:
       cur = cur.after(cur.store(cur + pad.cast(cur.dtype)))
   grad_buf.uop = cur

examples/mlperf/optim.py

@@ -15,7 +15,7 @@ def stochastic_round_bf16(x:Tensor) -> Tensor:
   bits = x.bitcast(dtypes.uint32)
   if isinstance(x.device, tuple):
     shape = x.uop.shard_shape if x.uop.axis is not None else x.shape
-    noise = Tensor(UOp(Ops.MSTACK, dtypes.default_float, tuple(Tensor.rand(*shape, device=d).uop for d in x.device)))
+    noise = Tensor(UOp(Ops.MSTACK, src=tuple(Tensor.rand(*shape, device=d).uop for d in x.device)))
   else:
     noise = x.rand_like()
   noise = (noise * 0xFFFF).cast(dtypes.uint32)
@@ -27,6 +27,11 @@ def clip_grads(grads:list[Tensor], grad_acc, clip_norm) -> Tensor:
   for g in grads: g.assign((g * (clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(g.dtype))
   return total_norm
 
+def fclip_grads(grads:list[Tensor], clip_norm) -> Tensor:
+  total_norm = Tensor.stack(*[g.float().square().sum() for g in grads]).sum().sqrt().contiguous()
+  scale = (clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)
+  return [(g * scale).cast(g.dtype) for g in grads], total_norm
+
 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)

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/buid_dataset_cache.py

@@ -0,0 +1,13 @@
+from pathlib import Path
+
+from examples.mlperf.dataloader import get_llama3_dataset
+from tinygrad.helpers import getenv
+
+BASEDIR = Path(getenv("BASEDIR", "/raid/datasets/c4-8b/"))
+SAMPLES = getenv("SAMPLES", 1_200_000 * 32)
+EVAL_SAMPLES = getenv("EVAL_SAMPLES", 1024)
+SEQLEN = getenv("SEQLEN", 8192)
+DATA_SEED = getenv("DATA_SEED", 5760)
+
+get_llama3_dataset(SAMPLES, SEQLEN, BASEDIR, seed=DATA_SEED, val=False, small=True)
+get_llama3_dataset(EVAL_SAMPLES, SEQLEN, BASEDIR, seed=0, val=True, small=True)

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

@@ -35,7 +35,7 @@ export BASEDIR="/raid/datasets/c4-8b/"
 export SMALL=1
 export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
 export EVAL_TARGET=3.3 EVAL_FREQ=12288
-export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
+export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=2048 MAX_STEPS=1200000
 export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
 export SAMPLES=$((MAX_STEPS * GBS))
 export SEQLEN=${SEQLEN:-8192}
@@ -44,9 +44,9 @@ 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 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
 
-export FAKEDATA=${FAKEDATA:-$([[ "$DEV" == NULL:* ]] && echo 1 || echo 0)} BENCHMARK=${BENCHMARK:-10}
+export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
 if [ -z "$FULL_LAYERS" ]; then
   export LLAMA_LAYERS=${LLAMA_LAYERS:-2}
 fi

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

@@ -35,7 +35,7 @@ export BASEDIR="/raid/datasets/c4-8b/"
 export SMALL=1
 export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
 export EVAL_TARGET=3.3 EVAL_FREQ=12288
-export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
+export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=2048 MAX_STEPS=1200000
 export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
 export SAMPLES=$((MAX_STEPS * GBS))
 export SEQLEN=${SEQLEN:-8192}
@@ -44,6 +44,6 @@ 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
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
 
 python3 examples/mlperf/model_train.py

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

@@ -17,7 +17,7 @@ export USE_ATOMICS=1
 export ASM_GEMM=1
 export WQKV=1
 export MASTER_WEIGHTS=1
-export FP8=1
+export MXFP4=1
 export ALLREDUCE_CAST=1
 export FAST_CE=1
 export FUSED_INPUT_QUANTIZE=1
@@ -26,7 +26,7 @@ export FUSED_ADD_NORM_MUL_QUANTIZE=1
 export FUSED_SILU_W13=1
 export SPLIT_W13=0
 
-export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="float32"
 export DP=8 MP=1 BS=16 EVAL_BS=8 GRADIENT_ACC_STEPS=2
 export GBS=$((BS * GRADIENT_ACC_STEPS))
 
@@ -44,7 +44,7 @@ export SEED=$RANDOM
 export DATA_SEED=$SEED
 
 export JITBEAM=3
-export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
 
 export LOGMLPERF=1
 

examples/openpilot/compile3.py

@@ -107,14 +107,21 @@ def compile(onnx_file):
   return inputs, test_val
 
 def test_vs_compile(run, inputs, test_val=None):
+  if (log:=bool(getenv("BENCHMARK_LOG", ""))): from extra.bench_log import WallTimeEvent, BenchEvent
 
   # run 20 times
   step_times = []
   for _ in range(20):
     st = time.perf_counter()
-    out = run(**inputs)
-    mt = time.perf_counter()
-    val = out.numpy()
+    if log:
+      with WallTimeEvent(BenchEvent.STEP):
+        out = run(**inputs)
+        mt = time.perf_counter()
+        val = out.numpy()
+    else:
+      out = run(**inputs)
+      mt = time.perf_counter()
+      val = out.numpy()
     et = time.perf_counter()
     step_times.append((et-st)*1e3)
     print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {step_times[-1]:6.2f} ms")
@@ -160,12 +167,6 @@ def test_vs_onnx(new_inputs, test_val, onnx_file, tol):
   print("test vs onnx passed")
   return timings
 
-def bench(run, inputs):
-  from extra.bench_log import WallTimeEvent, BenchEvent
-  for _ in range(10):
-    with WallTimeEvent(BenchEvent.STEP):
-      run(**inputs).numpy()
-
 if __name__ == "__main__":
   if getenv("RUN_PICKLE"):
     with open(OUTPUT, "rb") as f: pickle_loaded = load_pickle(f)
@@ -181,6 +182,3 @@ if __name__ == "__main__":
     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)

extra/amdpci/am_smi.py

@@ -84,7 +84,8 @@ class AMSMI(AMDev):
     with open(f"/sys/bus/pci/devices/{self.pcibus}/power_state", "r") as f: return f.read().strip().rstrip()
 
 class SMICtx:
-  def __init__(self):
+  def __init__(self, dev_filter=None):
+    self.dev_filter = dev_filter
     self.devs = []
     self.opened_pcidevs = []
     self.opened_pci_resources = {}
@@ -135,6 +136,7 @@ class SMICtx:
     pattern = os.path.join('/tmp', 'am_*.lock')
     for d in [f[8:-5] for f in glob.glob(pattern)]:
       if d.startswith("usb"): continue
+      if self.dev_filter is not None and d != self.dev_filter: continue
       if d not in self.opened_pcidevs:
         self._open_am_device(d)
 
@@ -406,7 +408,7 @@ if __name__ == "__main__":
 
   try:
     if not args.list: os.system('clear')
-    smi_ctx = SMICtx()
+    smi_ctx = SMICtx(args.dev)
     while True:
       smi_ctx.rescan_devs()
       smi_ctx.draw(args.list)

extra/bench_log.py

@@ -35,7 +35,7 @@ class WallTimeEvent:
     return self
   def __exit__(self, *_):
     self.time = time.monotonic() - self.start
-    _events[self.event]["wall"].append(self.time)
+    _events[self.event]["wall"].append((self.time, BENCHMARK_LOG.value))
     return False
 
 class KernelTimeEvent:
@@ -47,19 +47,19 @@ class KernelTimeEvent:
     self.start = GlobalCounters.time_sum_s
     return self
   def __exit__(self, *_):
-    _events[self.event]["kernel"].append(GlobalCounters.time_sum_s - self.start)
+    _events[self.event]["kernel"].append((GlobalCounters.time_sum_s - self.start, BENCHMARK_LOG.value))
     return False
 
 def log_event_instant(event:InstantBenchEvent, value:float):
-  _events[event].append(value)
+  _events[event].append((value, BENCHMARK_LOG.value))
 
 if BENCHMARK_LOG:
   INFLUXDB_HOST = getenv("INFLUXDB_HOST", "")
   INFLUXDB_ORG = getenv("INFLUXDB_ORG", "tiny")
   INFLUXDB_TOKEN = getenv("INFLUXDB_TOKEN", "")
 
-  def _create_point(run_id, i, attempt, ref, commit, name, value, run):
-    point = Point(BENCHMARK_LOG.value).tag("id", run_id).tag("index", i)
+  def _create_point(run_id, i, attempt, ref, commit, name, value, log_name, run):
+    point = Point(log_name.replace(':', '_').replace('.', '_')).tag("id", run_id).tag("index", i)
     point = point.tag("device", Device.DEFAULT)
     point = point.tag("attempt", attempt).tag("ref", ref).tag("commit", commit)
     point = point.field(name, value).field("x", run)
@@ -91,12 +91,12 @@ if BENCHMARK_LOG:
       run_id = str(uuid.uuid4())
       if isinstance(event, BenchEvent):
         for event_type, values in _events[event].items():
-          for i, value in enumerate(values):
-            point = _create_point(run_id, i, attempt, ref, commit, f"{event.value}_{event_type}", value, run)
+          for i, (value, log_name) in enumerate(values):
+            point = _create_point(run_id, i, attempt, ref, commit, f"{event.value}_{event_type}", value, log_name, run)
             points.append(point)
       else:
-        for i, value in enumerate(_events[event]):
-          point = _create_point(run_id, i, attempt, ref, commit, event.value, value, run)
+        for i, (value, log_name) in enumerate(_events[event]):
+          point = _create_point(run_id, i, attempt, ref, commit, event.value, value, log_name, run)
           points.append(point)
 
     write_options = WriteOptions(write_type=WriteType.synchronous, retry_interval=5000, max_retries=5, max_retry_delay=30000, exponential_base=2)

extra/bnxt_driver/bnxtdev.py

@@ -0,0 +1,238 @@
+import ctypes, struct
+from tinygrad.helpers import ceildiv, getenv, wait_cond, DEBUG
+from tinygrad.runtime.autogen import bnxt, pci
+from tinygrad.runtime.support.system import PCIDevice, System, ipv4_to_gid
+
+BNXT_DEBUG = getenv("BNXT_DEBUG", 0)
+BNXT_ACCESS, BNXT_INIT_MASK, BNXT_RTR_MASK, BNXT_RTS_MASK = 3, 0xd, 0x41515ad, 0xae005
+BNXT_CHIMP_COMM, BNXT_CHIMP_COMM_TRIGGER = 0x0, 0x100
+BNXT_BACKING_STORE = ((0, 2), (1, 0), (2, 2), (3, 0), (4, 2), (5, 0), (6, 0), (14, 2), (15, 0))
+
+def db_value(xid, typ, index, epoch):
+  return (xid & bnxt.DBC_DBC_XID_MASK | bnxt.DBC_DBC_PATH_ROCE | typ | bnxt.BNXT_QPLIB_DBR_VALID) << 32 | \
+         index & bnxt.DBC_DBC_INDEX_MASK | epoch << bnxt.BNXT_QPLIB_DBR_EPOCH_SHIFT
+
+def _pbl(dev, paddrs, queue=False):
+  if len(paddrs) == 1: return 0, paddrs[0]
+  values = [p | bnxt.PTU_PTE_VALID for p in paddrs]
+  if queue:
+    values[-1] |= bnxt.PTU_PTE_LAST
+    if len(values) > 1: values[-2] |= bnxt.PTU_PTE_NEXT_TO_LAST
+  table, table_paddrs = dev.pci_dev.alloc_sysmem(ceildiv(len(values), 512) * 0x1000)
+  table[:len(values) * 8] = struct.pack(f"<{len(values)}Q", *values)
+  if len(table_paddrs) == 1: return 1, table_paddrs[0]
+  top, top_paddrs = dev.pci_dev.alloc_sysmem(0x1000)
+  top[:len(table_paddrs) * 8] = struct.pack(f"<{len(table_paddrs)}Q", *(p | bnxt.PTU_PTE_VALID for p in table_paddrs))
+  return 2, top_paddrs[0]
+
+def _queue(dev, stride:int=16, aux=False):
+  mem, paddrs = dev.pci_dev.alloc_sysmem(0x1000 + aux * 0x400)
+  level, base = _pbl(dev, paddrs, queue=True)
+  return {"mem":mem, "paddrs":paddrs, "stride":stride, "prod":0, "cons":0, "level":level, "base":base}
+
+def _qread(q, i):
+  off = (i & 15) * q["stride"]
+  return q["mem"][off:off + q["stride"]]
+
+def _qwrite(q, i, data, aux=False):
+  off = 0x1000 + i % 128 * 8 if aux else (i & 15) * q["stride"]
+  q["mem"][off:off + len(data)] = data
+
+class BNXTDev:
+  def __init__(self, pci_dev:PCIDevice, ip:str=getenv("BNXT_IP", "10.0.0.1")):
+    self.pci_dev, self.devfmt = pci_dev, pci_dev.pcibus
+    self.bar0, self.db = pci_dev.map_bar(0, fmt='I'), pci_dev.map_bar(2, fmt='Q')
+    pci_dev.write_config(pci.PCI_COMMAND, pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)
+    self.resp, self.resp_pa = pci_dev.alloc_sysmem(0x1000)
+    self.seq = 0
+
+    ver = self.hwrm("ver_get")
+    if DEBUG >= 2: print(f"bnxt {self.devfmt}: firmware {ver.hwrm_fw_maj_8b}.{ver.hwrm_fw_min_8b}.{ver.hwrm_fw_bld_8b}")
+    self.hwrm("func_reset", timeout_ms=40000)
+    caps = self.hwrm("func_qcaps", fid=0xffff)
+    self.mac, self.port_id = int.from_bytes(bytes(caps.mac_address), 'big'), caps.port_id
+    self.hwrm("func_drv_rgtr")
+    self.db_off = self.hwrm("func_qcfg", fid=0xffff).legacy_l2_db_size_kb * 1024
+
+    self.setup_backing_store()
+    self._open_rcfw()
+    self._open_l2()
+    self.local_gid = ipv4_to_gid(ip)
+    gids, mac = (ctypes.c_uint32 * 4)(*(int.from_bytes(self.local_gid[i:i + 4], 'big') for i in (12, 8, 4, 0))), self.mac.to_bytes(6, 'big')
+    smac = (ctypes.c_uint16 * 3)(*(int.from_bytes(mac[i:i + 2], 'big') for i in (0, 2, 4)))
+    self.gid_id = self.rcfw("add_gid", gid=gids, src_mac=smac).xid
+
+    if DEBUG >= 2: print(f"bnxt {self.devfmt}: booted mac={self.mac.to_bytes(6, 'big').hex(':')} gid={self.local_gid.hex()}")
+
+  def hwrm(self, name, timeout_ms=10000, **fields):
+    inp, out = getattr(bnxt, f"struct_hwrm_{name}_input"), getattr(bnxt, f"struct_hwrm_{name}_output")
+    opcode = getattr(bnxt, f"HWRM_{name.upper()}")
+    self.seq = (self.seq + 1) & 0xffff
+    data = bytes(inp(req_type=opcode, cmpl_ring=bnxt.BNXT_HWRM_NO_CMPL_RING, seq_id=self.seq, target_id=bnxt.BNXT_HWRM_TARGET,
+                     resp_addr=self.resp_pa[0], **fields))
+    self.resp[:] = bytes(len(self.resp))
+    System.memory_barrier()
+    for i, w in enumerate(memoryview(bytearray(data.ljust(bnxt.HWRM_MAX_REQ_LEN, b'\0'))).cast('I')):
+      self.bar0[BNXT_CHIMP_COMM // 4 + i] = w
+    self.bar0[BNXT_CHIMP_COMM_TRIGGER // 4] = 1
+    def hdr(): return bnxt.struct_hwrm_resp_hdr.from_buffer_copy(bytes(self.resp[:8]))
+    wait_cond(lambda: (n := hdr().resp_len) and hdr().seq_id == self.seq and self.resp[n - 1], timeout_ms=timeout_ms, msg=f"HWRM {name}")
+    ret = out.from_buffer_copy(bytes(self.resp[:ctypes.sizeof(out)]))
+    assert ret.error_code == 0, f"HWRM {name}: {ret.error_code}"
+    return ret
+
+  def setup_backing_store(self):
+    counts: dict[int, int] = {}
+    for typ, extra in BNXT_BACKING_STORE:
+      caps = self.hwrm("func_backing_store_qcaps_v2", type=typ)
+      size, splits = caps.entry_size, tuple(getattr(caps, f"split_entry_{j}") for j in range(caps.subtype_valid_cnt))
+      counts[typ] = n = counts[0] if typ == 15 else max(caps.min_num_entries, sum(splits) + extra)
+      # a zero bitmap means the type has a single instance 0
+      for instance in [i for i in range(8) if caps.instance_bit_map >> i & 1] or [0]:
+        mem, paddrs = self.pci_dev.alloc_sysmem(ceildiv(n * size, 0x1000) * 0x1000)
+        if caps.ctx_init_value:
+          for off in range(caps.ctx_init_offset, len(mem), size): mem[off] = caps.ctx_init_value
+        lvl, base = _pbl(self, paddrs)
+        self.hwrm("func_backing_store_cfg_v2", type=typ, instance=instance, entry_size=size, num_entries=n, page_dir=base,
+          page_size_pbl_level=lvl, subtype_valid_cnt=len(splits),
+          flags=bnxt.FUNC_BACKING_STORE_CFG_V2_REQ_FLAGS_BS_CFG_ALL_DONE if typ == 15 else 0,
+          **{f"split_entry_{j}": v for j, v in enumerate(splits)})
+
+  def _open_rcfw(self):
+    self.rcfw_first = True
+
+    self.creq = _queue(self)
+    self.creq_id = self.hwrm("ring_alloc", ring_type=bnxt.RING_ALLOC_REQ_RING_TYPE_NQ, page_tbl_addr=self.creq["base"],
+      page_size=12, page_tbl_depth=self.creq["level"], length=16, int_mode=bnxt.RING_ALLOC_REQ_INT_MODE_MSIX).ring_id
+
+    self.cmdq = _queue(self)
+    self.doorbell(self.creq_id, bnxt.DBC_DBC_TYPE_NQ_ARM, 0, 0)
+    init = bnxt.struct_cmdq_init(cmdq_pbl=self.cmdq["base"], creq_ring_id=self.creq_id,
+                                 cmdq_size_cmdq_lvl=16 << bnxt.CMDQ_INIT_CMDQ_SIZE_SFT)
+
+    System.memory_barrier()
+    for i, w in enumerate(memoryview(bytearray(bytes(init))).cast('I')): self.bar0[bnxt.RCFW_COMM_BASE_OFFSET // 4 + i] = w
+
+    _, p = self.pci_dev.alloc_sysmem(0x1000)
+    self.rcfw("initialize_fw", stat_ctx_id=self.hwrm("stat_ctx_alloc", stats_dma_addr=p[0], stats_dma_length=176).stat_ctx_id,
+      flags=bnxt.CMDQ_INITIALIZE_FW_FLAGS_HW_REQUESTER_RETX_SUPPORTED)
+
+    # RoCE notification ring: never armed or serviced, but CQ and L2 ring allocation require one
+    nq = _queue(self)
+    self.nq_id = self.hwrm("ring_alloc", ring_type=bnxt.RING_ALLOC_REQ_RING_TYPE_NQ, page_tbl_addr=nq["base"],
+      page_size=12, page_tbl_depth=nq["level"], length=16, logical_id=1, int_mode=bnxt.RING_ALLOC_REQ_INT_MODE_MSIX).ring_id
+
+  def rcfw(self, name, timeout_ms=20000, **fields):
+    req_t, resp_t = getattr(bnxt, f"struct_cmdq_{name}"), getattr(bnxt, f"struct_creq_{name}_resp")
+    op = getattr(bnxt, f"CMDQ_BASE_OPCODE_{name.upper()}")
+    data = bytes(req_t(opcode=op, cmd_size=(slots := ceildiv(ctypes.sizeof(req_t), 16)), **fields)).ljust(slots * 16, b'\0')
+    for i in range(slots): _qwrite(self.cmdq, self.cmdq["prod"] + i, data[i * 16:(i + 1) * 16])
+
+    self.cmdq["prod"] += slots
+    prod = self.cmdq["prod"] & 0xffff
+    if self.rcfw_first: prod, self.rcfw_first = prod | 1 << bnxt.FIRMWARE_FIRST_FLAG, False
+
+    System.memory_barrier()
+
+    self.bar0[(bnxt.RCFW_COMM_BASE_OFFSET + bnxt.RCFW_PF_VF_COMM_PROD_OFFSET) // 4] = prod
+    self.bar0[(bnxt.RCFW_COMM_BASE_OFFSET + bnxt.RCFW_COMM_TRIG_OFFSET) // 4] = bnxt.RCFW_CMDQ_TRIG_VAL
+
+    def poll():
+      h = bnxt.struct_creq_base.from_buffer_copy(bytes(_qread(self.creq, self.creq["cons"])))
+      return bool(h.v & bnxt.CREQ_BASE_V) != bool((self.creq["cons"] // 16) & 1)
+    wait_cond(poll, timeout_ms=timeout_ms, msg=f"RCFW {name}")
+
+    ret = resp_t.from_buffer_copy(bytes(_qread(self.creq, self.creq["cons"])))
+    self.creq["cons"] += 1
+
+    # NQ_ARM also publishes the CREQ consumer index, which is what frees ring space for the next command
+    self.doorbell(self.creq_id, bnxt.DBC_DBC_TYPE_NQ_ARM, self.creq["cons"] & 15, (self.creq["cons"] // 16) & 1)
+    assert ret.status == 0, f"RCFW {name}: {ret.status}"
+
+    if BNXT_DEBUG >= 1: print(f"bnxt {self.devfmt}: rcfw {name} xid={getattr(ret, 'xid', 0):#x}")
+    return ret
+
+  def doorbell(self, xid, typ, index, epoch):
+    System.memory_barrier()
+    self.db[self.db_off // 8] = db_value(xid, typ, index, epoch)
+
+  # L2 receive path, required for RoCE ingress even though no ethernet receive buffers are posted
+  def _open_l2(self):
+    cq = _queue(self)
+    ci = self.hwrm("ring_alloc", enables=bnxt.RING_ALLOC_REQ_ENABLES_NQ_RING_ID_VALID, ring_type=bnxt.RING_ALLOC_REQ_RING_TYPE_L2_CMPL,
+      page_tbl_addr=cq["base"], page_size=12, page_tbl_depth=cq["level"], length=16, nq_ring_id=self.nq_id).ring_id
+    rx = _queue(self)
+    ri = self.hwrm("ring_alloc", enables=bnxt.RING_ALLOC_REQ_ENABLES_NQ_RING_ID_VALID |
+      bnxt.RING_ALLOC_REQ_ENABLES_RX_BUF_SIZE_VALID, ring_type=bnxt.RING_ALLOC_REQ_RING_TYPE_RX, page_tbl_addr=rx["base"],
+      page_size=12, page_tbl_depth=rx["level"], length=16, rx_buf_size=640, nq_ring_id=self.nq_id).ring_id
+    vi = self.hwrm("vnic_alloc").vnic_id
+    self.hwrm("vnic_cfg", enables=bnxt.VNIC_CFG_REQ_ENABLES_MRU | bnxt.VNIC_CFG_REQ_ENABLES_DEFAULT_RX_RING_ID |
+      bnxt.VNIC_CFG_REQ_ENABLES_DEFAULT_CMPL_RING_ID, vnic_id=vi, mru=9018,
+      default_rx_ring_id=ri, default_cmpl_ring_id=ci)
+    self.hwrm("cfa_l2_filter_alloc", flags=bnxt.CFA_L2_FILTER_ALLOC_REQ_FLAGS_PATH_RX,
+      enables=bnxt.CFA_L2_FILTER_ALLOC_REQ_ENABLES_L2_ADDR | bnxt.CFA_L2_FILTER_ALLOC_REQ_ENABLES_L2_ADDR_MASK |
+      bnxt.CFA_L2_FILTER_ALLOC_REQ_ENABLES_DST_ID, l2_addr=tuple(self.mac.to_bytes(6, 'big')), l2_addr_mask=(0xff,) * 6, dst_id=vi)
+
+  def register_mem(self, paddrs:list[int], size:int, log_page_size:int=12) -> int:
+    level, base = _pbl(self, paddrs[:ceildiv(size, 1 << log_page_size)])
+    return self.rcfw("register_mr", flags=bnxt.CMDQ_REGISTER_MR_FLAGS_ALLOC_MR,
+      log2_pg_size_lvl=level << bnxt.CMDQ_REGISTER_MR_LVL_SFT | log_page_size << bnxt.CMDQ_REGISTER_MR_LOG2_PG_SIZE_SFT,
+      access=bnxt.CMDQ_REGISTER_MR_ACCESS_LOCAL_WRITE | bnxt.CMDQ_REGISTER_MR_ACCESS_REMOTE_WRITE,
+      log2_pbl_pg_size=12, pbl=base, va=paddrs[0], mr_size=size).xid
+
+class BNXTQP:
+  def __init__(self, dev:BNXTDev):
+    self.dev, self.sq_psn, self.msn = dev, 0, 0
+
+    self.cqq = _queue(dev, ctypes.sizeof(bnxt.struct_cq_base))
+    self.cq_id = dev.rcfw("create_cq", cq_size=16, pbl=self.cqq["base"],
+                          pg_size_lvl=self.cqq["level"], cq_fco_cnq_id=dev.nq_id).xid
+
+    self.sq = _queue(dev, aux=True)
+    self.qpn = dev.rcfw("create_qp", type=bnxt.CMDQ_CREATE_QP_TYPE_RC,
+      sq_size=16, sq_fwo_sq_sge=1, scq_cid=self.cq_id, rcq_cid=self.cq_id,
+      sq_pbl=self.sq["base"], sq_pg_size_sq_lvl=self.sq["level"]).xid
+    self.qp_op(1, BNXT_INIT_MASK, access=BNXT_ACCESS, pkey=0xffff)
+
+  def qp_op(self, state, mask, network_type=0, **fields):
+    self.dev.rcfw("modify_qp", qp_cid=self.qpn, modify_mask=mask,
+      network_type_en_sqd_async_notify_new_state=state | network_type, **fields)
+
+  def connect(self, qpn:int, gid:bytes, mac:int):
+    network_type = bnxt.CMDQ_MODIFY_QP_NETWORK_TYPE_ROCEV2_IPV4
+    dgid = (ctypes.c_uint32 * 4)(*(int.from_bytes(gid[i:i + 4], 'little') for i in (0, 4, 8, 12)))
+    dmac = (ctypes.c_uint16 * 3)(*(int.from_bytes(mac.to_bytes(6, 'big')[i:i + 2], 'little') for i in (0, 2, 4)))
+
+    self.qp_op(2, BNXT_RTR_MASK, network_type=network_type, qp_type=bnxt.CMDQ_MODIFY_QP_QP_TYPE_RC, access=BNXT_ACCESS,
+      pkey=0xffff, dgid=dgid, sgid_index=self.dev.gid_id, hop_limit=64, dest_mac=dmac,
+      path_mtu_pingpong_push_enable=bnxt.CMDQ_MODIFY_QP_PATH_MTU_MTU_1024, max_dest_rd_atomic=4,
+      dest_qp_id=qpn)
+    self.qp_op(3, BNXT_RTS_MASK, network_type=network_type, qp_type=bnxt.CMDQ_MODIFY_QP_QP_TYPE_RC, access=BNXT_ACCESS,
+      max_rd_atomic=1)
+
+    if BNXT_DEBUG >= 1: print(f"bnxt: QP {self.qpn:#x} connected (remote={qpn:#x})")
+
+  def _poll(self, timeout):
+    def poll():
+      base = bnxt.struct_cq_base.from_buffer_copy(bytes(_qread(self.cqq, self.cqq["cons"])))
+      return bool(base.cqe_type_toggle & bnxt.CQ_BASE_TOGGLE) == (not bool((self.cqq["cons"] // 16) & 1))
+    wait_cond(poll, timeout_ms=timeout, msg="BNXT CQ")
+    raw = bytes(_qread(self.cqq, self.cqq["cons"]))
+    self.cqq["cons"] += 1
+    self.dev.doorbell(self.cq_id, bnxt.DBC_DBC_TYPE_CQ, self.cqq["cons"] & 15, (self.cqq["cons"] // 16) & 1)
+    return raw
+
+  def rdma_write(self, rva, rkey, lva, lkey, size, timeout_ms=20000):
+    start = self.sq["prod"] & 15
+    hdr = bytes(bnxt.struct_sq_rdma_hdr(wqe_type=bnxt.SQ_RDMA_HDR_WQE_TYPE_WRITE_WQE,
+                flags=bnxt.SQ_SEND_FLAGS_SIGNAL_COMP, wqe_size=3, length=size, remote_va=rva, remote_key=rkey))
+    for i, data in enumerate((hdr[:16], hdr[16:32], bytes(bnxt.struct_sq_sge(va_or_pa=lva, l_key=lkey, size=size)))):
+      _qwrite(self.sq, start + i, data)
+    nxt = (self.sq_psn + max(1, ceildiv(size, 1024))) & 0xffffff
+    value = start << bnxt.SQ_MSN_SEARCH_START_IDX_SFT | nxt << bnxt.SQ_MSN_SEARCH_NEXT_PSN_SFT | self.sq_psn
+    _qwrite(self.sq, self.msn, struct.pack("<Q", value), aux=True)
+
+    self.msn, self.sq_psn, self.sq["prod"] = (self.msn + 1) % 128, nxt, self.sq["prod"] + 3
+    self.dev.doorbell(self.qpn, bnxt.DBC_DBC_TYPE_SQ, self.sq["prod"] & 15, (self.sq["prod"] // 16) & 1)
+    cqe = bnxt.struct_cq_req.from_buffer_copy(self._poll(timeout_ms))
+    assert cqe.status == 0

extra/bnxt_driver/connect.py

@@ -0,0 +1,118 @@
+#!/usr/bin/env python3
+"""Send and validate one RDMA WRITE between two Broadcom BNXT hosts.
+
+This follows ``extra/mlx_driver/connect.py``: sync the driver, start the remote
+endpoint over SSH, exchange QP/GID/MAC/MR metadata, move both RC QPs to RTS,
+write bytes into the remote MR, and verify the bytes on the remote host.
+
+Both PCI functions must be unbound from bnxt_en/bnxt_re first.
+"""
+import json
+import os
+import subprocess
+import sys
+from typing import Any, IO
+
+TINYGRAD = os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "../.."))
+sys.path.insert(0, TINYGRAD)
+
+from extra.bnxt_driver.bnxtdev import BNXTDev, BNXTQP
+from tinygrad.runtime.support.system import PCIDevice
+
+REMOTE_HOST = os.getenv("REMOTE_HOST", "192.168.52.213")
+REMOTE_USER = os.getenv("REMOTE_USER", "nimlgen")
+LOCAL_PCI = os.getenv("BNXT_PCI", "0000:41:00.0")
+REMOTE_PCI = os.getenv("REMOTE_PCI", "0000:41:00.0")
+LOCAL_IP = os.getenv("LOCAL_IP", "10.0.200.5")
+REMOTE_IP = os.getenv("REMOTE_IP", "10.0.200.6")
+MESSAGE = os.getenv("RDMA_MESSAGE", "Test message, rdma works!").encode()
+REMOTE = f"{REMOTE_USER}@{REMOTE_HOST}"
+SSH = ["ssh", "-o", "BatchMode=yes", "-o", "ConnectTimeout=10", "-o", "StrictHostKeyChecking=accept-new", REMOTE]
+SYNC_FILES = ("tinygrad/runtime/autogen/bnxt.py", "tinygrad/runtime/support/system.py",
+              "extra/bnxt_driver/bnxtdev.py", "extra/bnxt_driver/connect.py")
+
+def read_json(stream:IO[str], what:str) -> dict[str, Any]:
+  for line in iter(stream.readline, ""):
+    print(f"  [remote] {line}", end="")
+    try: value = json.loads(line)
+    except json.JSONDecodeError: continue
+    if isinstance(value, dict): return value
+  raise RuntimeError(f"remote exited before publishing {what}")
+
+def wait_line(stream:IO[str], text:str) -> str:
+  for line in iter(stream.readline, ""):
+    print(f"  [remote] {line}", end="")
+    if text in line: return line
+  raise RuntimeError(f"remote exited before reporting {text!r}")
+
+def send_line(stream:IO[str], value:str|dict[str, Any]):
+  stream.write((json.dumps(value) if isinstance(value, dict) else value) + "\n")
+  stream.flush()
+
+def qp_info(dev:BNXTDev, qp:BNXTQP) -> dict[str, Any]:
+  return {"qpn":qp.qpn, "mac":dev.mac.to_bytes(6, "big").hex(), "gid":dev.local_gid.hex()}
+
+def server():
+  dev = BNXTDev(PCIDevice("bnxt", os.getenv("BNXT_PCI", "0000:41:00.0")), ip=os.getenv("BNXT_IP", REMOTE_IP))
+  qp = BNXTQP(dev)
+  print(json.dumps(qp_info(dev, qp)), flush=True)
+
+  peer = json.loads(sys.stdin.readline())
+  qp.connect(peer["qpn"], bytes.fromhex(peer["gid"]), int(peer["mac"], 16))
+  print("connected", flush=True)
+
+  target, target_paddrs = dev.pci_dev.alloc_sysmem(0x1000)
+  target[:0x1000] = bytes(0x1000)
+  rkey = dev.register_mem(target_paddrs, 0x1000)
+  print(json.dumps({"target_addr":target_paddrs[0], "rkey":rkey}), flush=True)
+
+  assert sys.stdin.readline().strip() == "done"
+  received = bytes(target).rstrip(b"\0")
+  print(f"AS TEXT: {received.decode(errors='replace')!r}", flush=True)
+  print(json.dumps({"data":received.hex()}), flush=True)
+
+def sync_remote():
+  if os.getenv("SYNC", "1") == "0": return
+  print("syncing BNXT driver to remote")
+  subprocess.run(["rsync", "-azR", *SYNC_FILES, f"{REMOTE}:~/tinygrad/"], cwd=TINYGRAD, check=True)
+
+def start_remote() -> subprocess.Popen[str]:
+  print("booting remote")
+  command = (f"cd ~/tinygrad && sudo env PYTHONPATH=. PYTHONUNBUFFERED=1 BNXT_DEBUG={os.getenv('BNXT_DEBUG', '0')} "
+             f"BNXT_PCI={REMOTE_PCI} BNXT_IP={REMOTE_IP} python3 extra/bnxt_driver/connect.py --server")
+  return subprocess.Popen(SSH + [command], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=sys.stderr, text=True)
+
+def client():
+  assert 0 < len(MESSAGE) <= 0x1000
+  sync_remote()
+  remote = start_remote()
+  assert remote.stdin is not None and remote.stdout is not None
+  remote_info = read_json(remote.stdout, "QP information")
+  print("booting local")
+  dev = BNXTDev(PCIDevice("bnxt", LOCAL_PCI), ip=LOCAL_IP)
+  qp = BNXTQP(dev)
+
+  send_line(remote.stdin, qp_info(dev, qp))
+  wait_line(remote.stdout, "connected")
+  qp.connect(remote_info["qpn"], bytes.fromhex(remote_info["gid"]), int(remote_info["mac"], 16))
+  print("both QPs in RTS")
+
+  remote_target = read_json(remote.stdout, "MR information")
+  source, source_paddrs = dev.pci_dev.alloc_sysmem(0x1000)
+  source[:len(MESSAGE)] = MESSAGE
+  lkey = dev.register_mem(source_paddrs, 0x1000)
+  print(f"RDMA WRITE {len(MESSAGE)}B to remote phys 0x{remote_target['target_addr']:x}")
+  qp.rdma_write(remote_target["target_addr"], remote_target["rkey"], source_paddrs[0], lkey, len(MESSAGE))
+
+  send_line(remote.stdin, "done")
+  wait_line(remote.stdout, "AS TEXT")
+  result = read_json(remote.stdout, "RDMA result")
+  assert bytes.fromhex(result["data"]) == MESSAGE
+  print("RDMA WRITE data verified")
+
+  remote.stdin.close()
+  assert remote.wait() == 0
+  print("RDMA WRITE test complete")
+
+if __name__ == "__main__":
+  server() if "--server" in sys.argv else client()

extra/bnxt_driver/loopback.py

@@ -0,0 +1,46 @@
+#!/usr/bin/env python3
+"""Local BNXT RoCEv2 RDMA WRITE loopback using the firmware's PHY loopback mode.
+
+The kernel bnxt_en/bnxt_re modules must be unloaded first.
+
+  sudo PYTHONPATH=. BNXT_PCI=0000:41:00.0 BNXT_IP=10.0.200.5 python3 extra/bnxt_driver/loopback.py
+"""
+import os
+import sys
+import time
+
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "../.."))
+
+from extra.bnxt_driver.bnxtdev import BNXTDev, BNXTQP
+from tinygrad.runtime.autogen import bnxt
+from tinygrad.runtime.support.system import PCIDevice
+
+BUF_SIZE = 0x1000
+BNXT_PCI = os.getenv("BNXT_PCI", "0000:41:00.0")
+BNXT_IP = os.getenv("BNXT_IP", "10.0.200.5")
+
+if __name__ == "__main__":
+  print(f"[init] BNXT at {BNXT_PCI}")
+  dev = BNXTDev(PCIDevice("bnxt", BNXT_PCI), ip=BNXT_IP)
+  tx_qp, rx_qp = BNXTQP(dev), BNXTQP(dev)
+  print(f"[init] loopback-connect TX QP 0x{tx_qp.qpn:x} <-> RX QP 0x{rx_qp.qpn:x}")
+  tx_qp.connect(rx_qp.qpn, dev.local_gid, dev.mac)
+  rx_qp.connect(tx_qp.qpn, dev.local_gid, dev.mac)
+
+  src, src_paddrs = dev.pci_dev.alloc_sysmem(BUF_SIZE)
+  dst, dst_paddrs = dev.pci_dev.alloc_sysmem(BUF_SIZE)
+  message = b"Hello from BNXT RoCE PHY loopback!"
+  src[:BUF_SIZE], dst[:BUF_SIZE] = bytes(BUF_SIZE), bytes(BUF_SIZE)
+  src[:len(message)] = message
+  lkey = dev.register_mem(src_paddrs, BUF_SIZE)
+  rkey = dev.register_mem(dst_paddrs, BUF_SIZE)
+
+  print("[loopback] enabling local PHY loopback")
+  dev.hwrm("port_phy_cfg", port_id=dev.port_id, enables=bnxt.PORT_PHY_CFG_REQ_ENABLES_LPBK, lpbk=bnxt.PORT_PHY_CFG_REQ_LPBK_LOCAL)
+  time.sleep(1)
+  tx_qp.rdma_write(dst_paddrs[0], rkey, src_paddrs[0], lkey, len(message))
+  got = bytes(dst[:len(message)])
+  print(f"[result] {got!r}")
+  assert got == message
+  print("BNXT RoCE PHY loopback RDMA WRITE passed")
+  dev.hwrm("port_phy_cfg", port_id=dev.port_id, enables=bnxt.PORT_PHY_CFG_REQ_ENABLES_LPBK, lpbk=bnxt.PORT_PHY_CFG_REQ_LPBK_NONE)

extra/export_model.py

@@ -241,8 +241,7 @@ export default {model_name};
 def export_model(model, target:str, *inputs, model_name: Optional[str] = "model", stream_weights=False):
   assert Device.DEFAULT in EXPORT_SUPPORTED_DEVICE, f"only {', '.join(EXPORT_SUPPORTED_DEVICE)} are supported"
 
-  # NOTE: NUM_CPU_THREADS=1, since export does not support threading
-  with Context(JIT=2, NUM_CPU_THREADS=1): linear, output_bufs = jit_model(model, *inputs)
+  with Context(JIT=2): linear, output_bufs = jit_model(model, *inputs)
   functions, statements, bufs, bufs_to_save = compile_net(linear, output_bufs)
   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}

extra/gemm/amd_asm_matmul.py

@@ -462,7 +462,7 @@ def test_matmul():
     lds = UOp.placeholder((lds_size,), dtypes.uint8, 0, AddrSpace.LOCAL)
     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.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts]))))
   c = Tensor.custom_kernel(a, b, c, fxn=asm_kernel)[2]
   linear = c.schedule_linear()
 

extra/gemm/cdna_asm_gemm.py

@@ -125,7 +125,7 @@ def custom_mxfp4_gemm(C:UOp, A:UOp, B:UOp, scale_a:UOp, scale_b:UOp, *extra:UOp,
                   arg=KernelInfo(f"mxfp4_gemm_{M}_{N}_{K}",
                                  estimates=Estimates(ops=2*M*N*K, mem=(M*half_k+N*half_k)*A.dtype.itemsize+M*N*C.dtype.itemsize)))
   insts = build_kernel(M, N, K, tile_m, tile_n)
-  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.LINEAR, src=tuple(UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts))))
 
 def _mxfp4_gemm_quantized(a_q:Tensor, b_q:Tensor, scale_a:Tensor, scale_b:Tensor) -> Tensor:
   M, half_k = a_q.shape
@@ -215,7 +215,7 @@ def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
   k = UOp.range(K, 0, AxisType.REDUCE)
   mul = (A.flatten().index((m*UOp.const(K)+k))*
          B.flatten().index((k*UOp.const(N)+n))).cast(dtypes.float32)
-  red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype)
+  red = mul.reduce(k, arg=Ops.ADD).cast(C.dtype)
   store = C.flatten().index((m*UOp.const(N)+n)).store(red).end(m, n)
   return store.sink(arg=KernelInfo(name=f'uop_gemm_{M}_{N}_{K}'))
 

extra/gemm/max_matmul.py

@@ -113,7 +113,6 @@ if __name__ == "__main__":
       }
     elif GEMM_VARIATION == "hcopt" and M == N == K == 4096 and DTYPE_IN == dtypes.half and DTYPE_OUT == dtypes.half and DTYPE_ACC == dtypes.float:
       print("Using CUDA and generated hcopt")
-      # [Opt(op=OptOps.TC, axis=0, amt=0), Opt(op=OptOps.UPCAST, axis=0, amt=4), Opt(op=OptOps.UPCAST, axis=1, amt=4), Opt(op=OptOps.LOCAL, axis=1, amt=4)]
       prog = CUDAProgram(device, "wmma_example", compiler.compile(open(os.path.join(script_dir, 'max_kernels/nv.fp16_fp32_fp16.hcopt.cu')).read()))
       args = (c, a, b)
       kwargs = {

extra/gemm/moe_routing.py

@@ -1,8 +1,53 @@
+import functools, math, pathlib
 from tinygrad import Tensor, dtypes
+from tinygrad.helpers import getenv
 from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
+from tinygrad.renderer import Estimates
+from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
 
 BLOCK_ROW = 256
 
+@functools.cache
+def _router_mfma_fwd(out:UOp, x:UOp, weight:UOp, bias:UOp, *, dname:str) -> UOp:
+  *lead, K = x.shape
+  M = math.prod(lead)
+  E = weight.shape[0]
+  threads = UOp.special(256, "lidx0")
+  workgroups = UOp.special((M + 63) // 64, "gidx0")
+  sink = UOp.sink(out.base, x.base, weight.base, bias.base, threads, workgroups,
+                  arg=KernelInfo(f"moe_router_mfma_{M}_{K}_{E}", estimates=Estimates(ops=2*M*E*K, mem=(M*K+E*K+E)*2+M*E*4)))
+  amd = pathlib.Path(__file__).parent.parent/"thunder"/"amd"
+  src = (amd/"moe_router_mfma.cpp").read_text()
+  lib = HIPCCCompiler("gfx950", [f"-I{(amd/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-DHIP_ENABLE_WARP_SYNC_BUILTINS",
+                                 f"-DROUTER_M={M}", f"-DROUTER_K={K}", f"-DROUTER_E={E}"]).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)))
+
+def _router_mfma_bwd(gradient:UOp, kernel:UOp) -> tuple:
+  _, x_u, weight_u, bias_u = kernel.src[1:5]
+  x, weight, bias = (Tensor(u, device=u.device) for u in (x_u, weight_u, bias_u))
+  reference = x.float() @ weight.float().T + bias.float()
+  grad_x, grad_weight, grad_bias = reference.gradient(x, weight, bias, gradient=Tensor(gradient, device=x_u.device))
+  return None, grad_x.uop, grad_weight.uop, grad_bias.uop
+
+def router_mfma(x:Tensor, weight:Tensor, bias:Tensor) -> Tensor:
+  assert x.ndim >= 2 and weight.ndim == 2 and bias.ndim == 1
+  K = x.shape[-1]
+  E = weight.shape[0]
+  assert weight.shape == (E, K) and bias.shape == (E,)
+  assert x.dtype == weight.dtype == bias.dtype == dtypes.bfloat16
+  assert E == 32 and K % 64 == 0
+  if isinstance(x.device, tuple):
+    assert x.uop.axis == 0, f"router MFMA requires axis-0 sharding, got axis={x.uop.axis}"
+    local_shape = x.uop.shard_shape
+    assert local_shape[-1] == K and math.prod(local_shape[:-1]) % 64 == 0, f"unsupported local router shape {local_shape}"
+  else:
+    assert math.prod(x.shape[:-1]) % 64 == 0
+  x, weight, bias = x.contiguous(), weight.contiguous(), bias.contiguous()
+  out = _sharded_invalids((*x.shape[:-1], E), dtypes.float32, x.device)
+  out, *_ = Tensor.custom_kernel(out, x, weight, bias,
+    fxn=functools.partial(_router_mfma_fwd, dname=str(x.device)), grad_fxn=_router_mfma_bwd)
+  return out
+
 def _sharded_invalids(shape:tuple[int, ...], dtype, device) -> Tensor:
   if isinstance(device, tuple):
     per = Tensor.invalids(shape[0]//len(device), *shape[1:], dtype=dtype, device=device)
@@ -53,7 +98,7 @@ def _ggather_bwd(gradient:UOp, kernel:UOp) -> tuple:
     g, m, j, jo, ji = _kv_ranges(Gk, M, Dk, _blk_for(Dk))
     row = idx.index(g, m).cast(dtypes.weakint)
     val = gout.index(g, m, j).load().cast(dtypes.float32)
-    atomic = UOp(Ops.CUSTOM, dtypes.void, (gtab.index(g, row, j), val), arg=atomic_str)
+    atomic = UOp(Ops.CUSTOM, src=(gtab.index(g, row, j), val), arg=(atomic_str, dtypes.void))
     return atomic.end(g, m, jo, ji).sink(arg=KernelInfo(name=f"ggather_bwd_{M}_{Dk}", opts_to_apply=()))
   grad_table = Tensor.custom_kernel(gt, go, Tensor(idx_u, device=dev), fxn=_bwd_kernel)[0]
   return (None, grad_table.cast(table_u.dtype).uop, None)

extra/gemm/rdna4_asm_matmul.py

@@ -223,7 +223,7 @@ def test_matmul():
     lds = UOp.placeholder((lds_size,), dtypes.uint8, 0, AddrSpace.LOCAL)
     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*2*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.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts]))))
 
   c = Tensor.custom_kernel(a, b, c, fxn=asm_kernel)[2]
   linear = c.schedule_linear()

extra/gemm/tinygrad_nv_matmul.py

@@ -1,6 +1,7 @@
 from tinygrad import Tensor, dtypes, Context
 from tinygrad.helpers import getenv
 from tinygrad.codegen.opt import Opt, OptOps
+from tinygrad.uop.ops import AxisType
 from tinygrad.engine.realize import run_linear
 from dataclasses import replace
 
@@ -13,17 +14,17 @@ if __name__ == "__main__":
   C = A.matmul(B)
   if getenv("GEMV"):
     opts = [
-      Opt(op=OptOps.UNROLL, axis=0, amt=8),
-      Opt(op=OptOps.GROUP, axis=0, amt=32),
+      Opt(op=OptOps.SPLIT, axis=1, arg=(8, AxisType.UNROLL)),
+      Opt(op=OptOps.SPLIT, axis=1, arg=(32, AxisType.GROUP_REDUCE)),
     ]
   else:
     opts = [
       Opt(op=OptOps.TC, axis=0, amt=0),
-      Opt(op=OptOps.UPCAST, axis=0, amt=4),
-      Opt(op=OptOps.UPCAST, axis=1, amt=8),
-      Opt(op=OptOps.LOCAL, axis=0, amt=2),
-      Opt(op=OptOps.LOCAL, axis=1, amt=2),
-      Opt(op=OptOps.LOCAL, axis=0, amt=2),
+      Opt(op=OptOps.SPLIT, axis=0, arg=(4, AxisType.UPCAST)),
+      Opt(op=OptOps.SPLIT, axis=1, arg=(8, AxisType.UPCAST)),
+      Opt(op=OptOps.SPLIT, axis=0, arg=(2, AxisType.LOCAL)),
+      Opt(op=OptOps.SPLIT, axis=1, arg=(2, AxisType.LOCAL)),
+      Opt(op=OptOps.SPLIT, axis=0, arg=(2, AxisType.LOCAL)),
     ]
   linear = C.schedule_linear()
   call = linear.src[-1]

extra/gptoss_kernels/rmsnorm/__init__.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+import functools, math, pathlib
+from tinygrad import Tensor, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.renderer import Estimates
+from extra.gemm.cdna_asm_gemm import FP8_DTYPE
+from extra.llama_kernels import NUM_WG, THREADS_PER_WG, alloc_like, alloc_local, compile_hip, dname_of
+
+def rmsnorm_mul_fwd(x_in:Tensor, weight:Tensor, eps:float) -> tuple[Tensor, Tensor]:
+  x = x_in.float()
+  rrms = (x.square().mean(-1, keepdim=True) + eps).rsqrt()
+  return ((x * rrms) * weight.float()).cast(x_in.dtype), rrms
+
+@functools.cache
+def _rmsnorm_mul_fwd_fxn(x_in_p, w_p, eps, device):
+  return rmsnorm_mul_fwd(Tensor(x_in_p, device=device), Tensor(w_p, device=device), eps)
+
+def _rmsnorm_mul_bwd(grad:UOp, call:UOp) -> tuple:
+  x = Tensor(call.src[1]).float(); weight = Tensor(call.src[2]).float()
+  rrms = Tensor(call.unbound_outputs[1])
+  x_normed = x * rrms                                  # recompute unweighted normed (x is call.src[1])
+  d_y = Tensor(grad).float()
+  dxn = d_y * weight                                   # d/d(x_normed)
+  d_x = rrms * (dxn - x_normed * (dxn * x_normed).mean(-1, keepdim=True))
+  dw = d_y * x_normed
+  d_weight = dw.sum(axis=tuple(range(dw.ndim - 1)))    # reduce batch/seq -> [dim]
+  return (d_x.cast(call.src[1].dtype).uop, d_weight.cast(call.src[2].dtype).uop)
+
+def rmsnorm_mul(x_in:Tensor, weight:Tensor, eps:float) -> tuple[Tensor, Tensor]:
+  fxn = _rmsnorm_mul_fwd_fxn(x_in.as_param(0).uop, weight.as_param(1).uop, eps, x_in.device)
+  outs = UOp.call_with_outputs((fxn[0].uop, fxn[1].uop), x_in.uop, weight.uop, grad_fxn=_rmsnorm_mul_bwd)
+  return Tensor(outs[0]), Tensor(outs[1])
+
+@functools.cache
+def _custom_rmsnorm_mul_quantize_mxfp8_fwd(q:UOp, e8:UOp, rrms:UOp, x:UOp, weight:UOp, *, dname:str, eps:float) -> UOp:
+  *lead, hidden = x.shape
+  rows, padded = math.prod(lead), q.shape[-1]
+  num_wg = min(NUM_WG, rows)
+  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(num_wg, "gidx0")
+  sink = UOp.sink(q.base, e8.base, rrms.base, x.base, weight.base, threads, workgroups,
+                  arg=KernelInfo(f"rmsnorm_mul_quantize_mxfp8_{rows}_{hidden}_{padded}",
+                                 estimates=Estimates(ops=8*rows*hidden, mem=rows*(hidden*2+padded+padded//32+4)+hidden*2)))
+  src = (pathlib.Path(__file__).parent/"rmsnorm_mul_quantize_mxfp8.cpp").read_text()
+  defines = [f"-DN_ELEMS={rows*hidden}", f"-DHIDDEN={hidden}", f"-DPADDED={padded}",
+             f"-DNUM_WG={num_wg}", f"-DTHREADS_PER_WG={THREADS_PER_WG}", f"-DEPS_LITERAL={eps}f"]
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
+                               UOp(Ops.BINARY, arg=compile_hip(src, defines))))
+
+@functools.cache
+def _custom_rmsnorm_mul_quantize_mxfp8_bwd(grad_x:UOp, grad_weight_partial:UOp, grad_q:UOp, x:UOp, weight:UOp, e8:UOp, rrms:UOp,
+                *, dname:str) -> UOp:
+  *lead, hidden = x.shape
+  rows, padded = math.prod(lead), grad_q.shape[-1]
+  num_wg = min(NUM_WG, rows)
+  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(num_wg, "gidx0")
+  sink = UOp.sink(grad_x.base, grad_weight_partial.base, grad_q.base, x.base, weight.base, e8.base, rrms.base,
+                  threads, workgroups,
+                  arg=KernelInfo(f"rmsnorm_mul_quantize_mxfp8_bwd_{rows}_{hidden}_{padded}",
+                                 estimates=Estimates(ops=10*rows*hidden, mem=rows*(hidden*6+padded*2+padded//32+4)+num_wg*hidden*4)))
+  src = (pathlib.Path(__file__).parent/"rmsnorm_mul_quantize_mxfp8_bwd.cpp").read_text()
+  defines = [f"-DN_ELEMS={rows*hidden}", f"-DHIDDEN={hidden}", f"-DPADDED={padded}", f"-DNUM_WG={num_wg}", f"-DTHREADS_PER_WG={THREADS_PER_WG}"]
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
+                               UOp(Ops.BINARY, arg=compile_hip(src, defines))))
+
+def _rmsnorm_mul_quantize_mxfp8_backward(gradient:UOp, kernel:UOp) -> tuple:
+  _, e8_u, rrms_u, x_u, weight_u = kernel.src[1:]
+  device = x_u.device
+  axis = x_u.axis if isinstance(device, tuple) else None
+  *lead, hidden = x_u.shape
+  num_wg = min(NUM_WG, math.prod(lead))
+  grad_x = alloc_like(x_u.shape, x_u.dtype, device, axis)
+  grad_weight_partial = alloc_local((num_wg, hidden), dtypes.float32, device, axis)
+  grad_q = Tensor(gradient, device=device).cast(dtypes.bfloat16).contiguous()
+  grad_x, grad_weight_partial, *_ = Tensor.custom_kernel(
+    grad_x, grad_weight_partial, grad_q, Tensor(x_u, device=device), Tensor(weight_u, device=device),
+    Tensor(e8_u.after(kernel), device=device), Tensor(rrms_u.after(kernel), device=device),
+    fxn=functools.partial(_custom_rmsnorm_mul_quantize_mxfp8_bwd, dname=dname_of(device)))
+  grad_weight = grad_weight_partial.sum(0).cast(weight_u.dtype)
+  return None, None, None, grad_x.uop, grad_weight.uop
+
+def rmsnorm_mul_quantize_mxfp8(x:Tensor, weight:Tensor, eps:float, padded:int|None=None) -> tuple[Tensor, Tensor, Tensor]:
+  """RMSNorm(x)*weight directly to rowwise MXFP8. Returns (q, e8, rrms), without a BF16 normalized round-trip."""
+  assert x.dtype == weight.dtype == dtypes.bfloat16 and x.shape[-1] == weight.shape[0], f"{x.shape=} {weight.shape=}"
+  hidden = x.shape[-1]
+  padded = math.ceil(hidden / 256) * 256 if padded is None else padded
+  assert padded >= hidden and padded % 256 == 0 and hidden % 32 == 0
+  axis = x.uop.axis if isinstance(x.device, tuple) else None
+  q = alloc_like((*x.shape[:-1], padded), FP8_DTYPE, x.device, axis)
+  e8 = alloc_like((*x.shape[:-1], padded // 32), dtypes.uint8, x.device, axis)
+  rrms = alloc_like((*x.shape[:-1], 1), dtypes.float32, x.device, axis)
+  q, e8, rrms, *_ = Tensor.custom_kernel(q, e8, rrms, x, weight,
+                                         fxn=functools.partial(_custom_rmsnorm_mul_quantize_mxfp8_fwd, dname=dname_of(x.device), eps=eps),
+                                         grad_fxn=_rmsnorm_mul_quantize_mxfp8_backward)
+  return q, e8, rrms

extra/gptoss_kernels/rmsnorm/rmsnorm_mul_quantize_mxfp8.cpp

@@ -0,0 +1,95 @@
+#include <hip/hip_runtime.h>
+#include <hip/hip_bf16.h>
+#include <hip/hip_fp8.h>
+
+#ifndef N_ELEMS
+#define N_ELEMS 47185920
+#endif
+#ifndef HIDDEN
+#define HIDDEN 2880
+#endif
+#ifndef PADDED
+#define PADDED 3072
+#endif
+#ifndef NUM_WG
+#define NUM_WG 1024
+#endif
+#ifndef THREADS_PER_WG
+#define THREADS_PER_WG 256
+#endif
+#ifndef EPS_LITERAL
+#define EPS_LITERAL 1e-5f
+#endif
+
+constexpr int ROWS = N_ELEMS / HIDDEN;
+constexpr int BLOCK = 32;
+constexpr int SCALE_BLOCKS = PADDED / BLOCK;
+constexpr float FP8_MAX = 448.0f;
+
+static_assert(N_ELEMS % HIDDEN == 0, "N_ELEMS must be divisible by HIDDEN");
+static_assert(HIDDEN % BLOCK == 0 && PADDED % BLOCK == 0 && PADDED >= HIDDEN,
+              "HIDDEN and PADDED must be block aligned");
+static_assert(SCALE_BLOCKS <= THREADS_PER_WG, "one thread handles each MXFP8 block");
+
+extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void rmsnorm_mul_quantize_mxfp8(
+    __hip_fp8_storage_t *__restrict__ q_out,
+    uint8_t *__restrict__ e8_out,
+    float *__restrict__ rrms_out,
+    const __hip_bfloat16 *__restrict__ x,
+    const __hip_bfloat16 *__restrict__ weight) {
+  __shared__ float reduce[THREADS_PER_WG];
+  __shared__ __hip_bfloat16 x_row[HIDDEN];
+
+  const int tid = threadIdx.x;
+  for (int row = blockIdx.x; row < ROWS; row += NUM_WG) {
+    const long long xbase = (long long)row * HIDDEN;
+    float sum_sq = 0.0f;
+    for (int col = tid; col < HIDDEN; col += THREADS_PER_WG) {
+      __hip_bfloat16 xb = x[xbase + col];
+      x_row[col] = xb;
+      float xf = (float)xb;
+      sum_sq = fmaf(xf, xf, sum_sq);
+    }
+
+    reduce[tid] = sum_sq;
+    __syncthreads();
+    for (int s = THREADS_PER_WG / 2; s > 0; s >>= 1) {
+      if (tid < s) reduce[tid] += reduce[tid + s];
+      __syncthreads();
+    }
+    const float rrms = rsqrtf(reduce[0] * (1.0f / (float)HIDDEN) + EPS_LITERAL);
+    if (tid == 0) rrms_out[row] = rrms;
+
+    if (tid < SCALE_BLOCKS) {
+      const int col_base = tid * BLOCK;
+      float vals[BLOCK];
+      float amax = 0.0f;
+      #pragma unroll
+      for (int i = 0; i < BLOCK; i++) {
+        const int col = col_base + i;
+        float v = 0.0f;
+        if (col < HIDDEN) {
+          float xn = (float)x_row[col] * rrms;
+          __hip_bfloat16 yb = (__hip_bfloat16)(xn * (float)weight[col]);
+          v = (float)yb;
+        }
+        vals[i] = v;
+        amax = fmaxf(amax, fabsf(v));
+      }
+      int e8 = (int)floorf(log2f(fmaxf(amax, 1e-38f))) + 127;
+      e8 = max(0, min(254, e8));
+      const float qscale = exp2f((float)(127 - e8));
+      __hip_fp8_storage_t packed[BLOCK];
+      #pragma unroll
+      for (int i = 0; i < BLOCK; i++) {
+        float v = fmaxf(-FP8_MAX, fminf(FP8_MAX, vals[i] * qscale));
+        packed[i] = __hip_cvt_float_to_fp8(v, __HIP_SATFINITE, __HIP_E4M3);
+      }
+      const long long qbase = (long long)row * PADDED + col_base;
+      *reinterpret_cast<uint4 *>(&q_out[qbase]) = *reinterpret_cast<uint4 *>(&packed[0]);
+      *reinterpret_cast<uint4 *>(&q_out[qbase + 16]) = *reinterpret_cast<uint4 *>(&packed[16]);
+      e8_out[(long long)row * SCALE_BLOCKS + tid] = (uint8_t)e8;
+    }
+    __syncthreads();
+  }
+}

extra/gptoss_kernels/rmsnorm/rmsnorm_mul_quantize_mxfp8_bwd.cpp

@@ -0,0 +1,99 @@
+#include <hip/hip_runtime.h>
+#include <hip/hip_bf16.h>
+
+#ifndef N_ELEMS
+#define N_ELEMS 47185920
+#endif
+#ifndef HIDDEN
+#define HIDDEN 2880
+#endif
+#ifndef PADDED
+#define PADDED 3072
+#endif
+#ifndef NUM_WG
+#define NUM_WG 1024
+#endif
+#ifndef THREADS_PER_WG
+#define THREADS_PER_WG 256
+#endif
+
+constexpr int ROWS = N_ELEMS / HIDDEN;
+constexpr int BLOCK = 32;
+constexpr int SCALE_BLOCKS = PADDED / BLOCK;
+constexpr int ELEMS_PER_THREAD = (HIDDEN + THREADS_PER_WG - 1) / THREADS_PER_WG;
+
+static_assert(N_ELEMS % HIDDEN == 0, "N_ELEMS must be divisible by HIDDEN");
+static_assert(HIDDEN % BLOCK == 0 && PADDED % BLOCK == 0 && PADDED >= HIDDEN,
+              "HIDDEN and PADDED must be block aligned");
+
+extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void rmsnorm_mul_quantize_mxfp8_bwd(
+    __hip_bfloat16 *__restrict__ grad_x,
+    float *__restrict__ grad_weight_partial,
+    const __hip_bfloat16 *__restrict__ grad_q,
+    const __hip_bfloat16 *__restrict__ x,
+    const __hip_bfloat16 *__restrict__ weight,
+    const uint8_t *__restrict__ e8,
+    const float *__restrict__ rrms) {
+  __shared__ float reduce[THREADS_PER_WG];
+  const int tid = threadIdx.x;
+  const int wg = blockIdx.x;
+
+  float w[ELEMS_PER_THREAD];
+  float gw[ELEMS_PER_THREAD];
+  #pragma unroll
+  for (int i = 0; i < ELEMS_PER_THREAD; i++) {
+    int col = tid + i * THREADS_PER_WG;
+    w[i] = col < HIDDEN ? (float)weight[col] : 0.0f;
+    gw[i] = 0.0f;
+  }
+
+  for (int row = wg; row < ROWS; row += NUM_WG) {
+    const long long xbase = (long long)row * HIDDEN;
+    const long long qbase = (long long)row * PADDED;
+    const long long ebase = (long long)row * SCALE_BLOCKS;
+    const float r = rrms[row];
+    float xn[ELEMS_PER_THREAD];
+    float gxn[ELEMS_PER_THREAD];
+    float local_dot = 0.0f;
+
+    #pragma unroll
+    for (int i = 0; i < ELEMS_PER_THREAD; i++) {
+      const int col = tid + i * THREADS_PER_WG;
+      if (col < HIDDEN) {
+        const float xnf = (float)x[xbase + col] * r;
+        const unsigned se = (unsigned)(254 - (int)e8[ebase + col / BLOCK]) << 23;
+        const float qscale = __builtin_bit_cast(float, se);
+        const float gy = (float)grad_q[qbase + col] * qscale;
+        const float gxnf = gy * w[i];
+        xn[i] = xnf;
+        gxn[i] = gxnf;
+        gw[i] += gy * xnf;
+        local_dot = fmaf(gxnf, xnf, local_dot);
+      } else {
+        xn[i] = gxn[i] = 0.0f;
+      }
+    }
+
+    reduce[tid] = local_dot;
+    __syncthreads();
+    for (int s = THREADS_PER_WG / 2; s > 0; s >>= 1) {
+      if (tid < s) reduce[tid] += reduce[tid + s];
+      __syncthreads();
+    }
+    const float mean_term = reduce[0] * (1.0f / (float)HIDDEN);
+
+    #pragma unroll
+    for (int i = 0; i < ELEMS_PER_THREAD; i++) {
+      const int col = tid + i * THREADS_PER_WG;
+      if (col < HIDDEN) grad_x[xbase + col] = (__hip_bfloat16)(r * (gxn[i] - xn[i] * mean_term));
+    }
+    __syncthreads();
+  }
+
+  const long long gwbase = (long long)wg * HIDDEN;
+  #pragma unroll
+  for (int i = 0; i < ELEMS_PER_THREAD; i++) {
+    const int col = tid + i * THREADS_PER_WG;
+    if (col < HIDDEN) grad_weight_partial[gwbase + col] = gw[i];
+  }
+}

extra/hcq/hcq_smi.py

@@ -16,9 +16,12 @@ def _do_reset_device(pci_bus): os.system(f"sudo sh -c 'echo 1 > /sys/bus/pci/dev
 def _is_module_loaded(name: str) -> bool: return os.path.isdir(f"/sys/module/{name}")
 
 def cmd_remove_module(args):
-  modules = ["nvidia_drm", "nvidia_modeset", "nvidia_uvm", "nvidia", "ast"] if args.backend == "nv" else ["amdgpu"]
+  modules = ["nvidia_drm", "nvidia_modeset", "nvidia_uvm", "nvidia"] if args.backend == "nv" else ["amdgpu"]
   to_unload = [m for m in modules if _is_module_loaded(m)]
   if not to_unload: print("Kernel modules are not loaded")
+  elif getattr(args, "expect", False):
+    print(f"Kernel modules are loaded: {to_unload}")
+    sys.exit(1)
   else:
     print("Removing kernel modules:", ", ".join(to_unload))
     try: subprocess.run(["sudo", "modprobe", "-r", *to_unload], check=True)
@@ -60,17 +63,19 @@ def cmd_show_pids(args):
 
 def cmd_kill_pids(args):
   devs = scan_devs_based_on_lock(prefix:={"amd":"am", "nv":"nv"}[args.backend], args)
+  use_sudo = not getattr(args, "sudoless", False)
 
   for dev in devs:
     for i in range(128):
       if i > 0: time.sleep(0.2)
 
       try:
-        try: pid = subprocess.check_output(['sudo', 'lsof', temp(f'{prefix}_{dev}.lock')]).decode('utf-8').strip().split('\n')[1].split()[1]
+        try: pid = subprocess.check_output((['sudo'] if use_sudo else []) +
+                                           ['lsof', temp(f'{prefix}_{dev}.lock')]).decode('utf-8').strip().split('\n')[1].split()[1]
         except subprocess.CalledProcessError: break
 
         print(f"Killing process {pid} (which uses {dev})")
-        subprocess.run(['sudo', 'kill', '-9', pid], check=True)
+        subprocess.run((['sudo'] if use_sudo else []) + ['kill', '-9', pid], check=True)
       except subprocess.CalledProcessError as e:
         print(f"Failed to kill process for device {dev}: {e}", file=sys.stderr)
 
@@ -79,6 +84,7 @@ def add_common_commands(parent_subparsers):
   p_insmod.set_defaults(func=cmd_insert_module)
 
   p_rmmod = parent_subparsers.add_parser("rmmod", help="Remove a kernel module")
+  p_rmmod.add_argument("--expect", action="store_true", help="Just assert that module is already unloaded")
   p_rmmod.set_defaults(func=cmd_remove_module)
 
   p_reset = parent_subparsers.add_parser("reset", help="Reset a device")
@@ -91,6 +97,7 @@ def add_common_commands(parent_subparsers):
 
   p_reset = parent_subparsers.add_parser("kill_pids", help="Kill pids of processes using the device")
   p_reset.add_argument("--pci_bus", default="", help="PCI bus ID of the device")
+  p_reset.add_argument("--sudoless", action="store_true", help="Do not use sudo when detecting or killing pids")
   p_reset.set_defaults(func=cmd_kill_pids)
 
 if __name__ == "__main__":

extra/hcq2/ops_amd2.py

@@ -1,259 +1,226 @@
 from __future__ import annotations
-from typing import cast, Any, Callable
-import os, ctypes, struct, hashlib, functools, importlib, mmap, errno, array, contextlib, sys, weakref, itertools, collections, atexit
+from typing import cast
+import os, ctypes, struct, functools, importlib, mmap, errno, contextlib, sys, itertools, atexit
 assert sys.platform != 'win32'
 from dataclasses import dataclass
-from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, encode_kernargs_clike, make_cmdbuf
-from tinygrad.runtime.support.hcq2 import make_binary_patch
+from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, HWQueue, encode_submit, to_name
 from tinygrad.uop.ops import sint, UOp
-from tinygrad.device import Compiled, BufferSpec, Buffer, Device
+from tinygrad.device import BufferSpec, Buffer
 from tinygrad.dtype import dtypes
-from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, lo32, hi32, colored, prod, ContextVar, TracingKey
-from tinygrad.helpers import VIZ, ceildiv, unwrap, pluralize, to_tuple
+from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, lo32, hi32
+from tinygrad.helpers import ceildiv, unwrap, pluralize
 from tinygrad.renderer.cstyle import HIPRenderer, HIPCCRenderer
 from tinygrad.renderer.llvmir import AMDLLVMRenderer
-from tinygrad.runtime.autogen import kfd, hsa, sqtt, amdgpu_kd, amdgpu_drm
+from tinygrad.runtime.autogen import kfd, hsa, amdgpu_kd, amdgpu_drm
 from tinygrad.runtime.autogen.am import am
 from tinygrad.runtime.support.elf import elf_loader
 from tinygrad.runtime.support.hcq import FileIOInterface, HCQBuffer, MMIOInterface, hcq_filter_visible_devices
 from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager
 from tinygrad.runtime.support.amd import AMDReg, AMDIP, import_module, import_soc, import_pmc
 from tinygrad.runtime.support.system import PCIIfaceBase, PCIAllocationMeta, USBPCIDevice, MAP_FIXED, MAP_NORESERVE
-from tinygrad.runtime.support.usb import USB3
+from tinygrad.runtime.support.usb import USB3, pm_usb_bufferize
 from tinygrad.runtime.support.memory import AddrSpace, BumpAllocator
-from tinygrad.runtime.ops_amd import SQTT, SQTT_ITRACE_SE_MASK, SQTT_LIMIT_SE, SQTT_SIMD_SEL, SQTT_TOKEN_EXCLUDE, PMC
-from tinygrad.runtime.ops_amd import EVENT_INDEX_PARTIAL_FLUSH, WAIT_REG_MEM_FUNCTION_EQ, WAIT_REG_MEM_FUNCTION_NEQ, WAIT_REG_MEM_FUNCTION_GEQ
+from tinygrad.runtime.ops_amd import SQTT, PMC
+from tinygrad.runtime.ops_amd import EVENT_INDEX_PARTIAL_FLUSH, WAIT_REG_MEM_FUNCTION_GEQ
 if getenv("IOCTL"): import extra.hip_gpu_driver.hip_ioctl  # noqa: F401 # pylint: disable=unused-import
 
-from tinygrad.engine.realize import get_runtime, pm_flatten_linear
-from tinygrad.uop import FastEnum, auto
-from tinygrad.uop.ops import Ops, UPat, PatternMatcher, graph_rewrite
+from tinygrad.engine.realize import get_call_arg_uops, get_call_var_uops
+from tinygrad.uop.ops import Ops, UPat, PatternMatcher
 
 # *****************
 # PM4
 
-class PM4Ops(FastEnum):
-  SET_SH_REG = auto(); SET_UCONFIG_REG = auto(); WAIT_REG_MEM = auto(); ACQUIRE_MEM = auto()  # noqa: E702
-  RELEASE_MEM = auto(); DISPATCH_DIRECT = auto(); EVENT_WRITE = auto()  # noqa: E702
+def _queue_args(hq:HWQueue, q) -> list[UOp]: # the ring and its pointers, tagged {name}_{queue} like the device's bufferize rules
+  shapes = [("ring", (q.ring.size,), q.ring.dtype)] + [(n, (1,), dtypes.uint64) for n in ("write_ptr", "doorbell", "put_value")]
+  return [UOp.placeholder(s, d, 0, device=hq.devs, volatile=True, tag=to_name(n, hq.queue)) for n, s, d in shapes]
 
-def pkt3(ctx, op:PM4Ops, *vals):
-  return UOp(Ops.INS, arg=op, src=tuple(UOp.const(x, dtypes.uint32)
-    for x in (ctx.pm4.PACKET3(getattr(ctx.pm4, f"PACKET3_{op.name}"), len(vals) - 1), *vals)))
+def _dw(vals) -> int: return sum(2 if isinstance(x, UOp) and x.dtype.itemsize == 8 else 1 for x in vals)
 
-def wreg(ctx, reg:AMDReg, *args:sint, **kwargs:int):
-  if bool(args) == bool(kwargs): raise RuntimeError('One (and only one) of *args or **kwargs must be specified')
-  if ctx.pm4.PACKET3_SET_SH_REG_START <= reg.addr[0] < ctx.pm4.PACKET3_SET_SH_REG_END:
-    op, set_packet_start = PM4Ops.SET_SH_REG, ctx.pm4.PACKET3_SET_SH_REG_START
-  elif ctx.pm4.PACKET3_SET_UCONFIG_REG_START <= reg.addr[0] < ctx.pm4.PACKET3_SET_UCONFIG_REG_START + 2**16-1:
-    op, set_packet_start = PM4Ops.SET_UCONFIG_REG, ctx.pm4.PACKET3_SET_UCONFIG_REG_START
-  else: raise RuntimeError(f'Cannot set {reg.name} ({reg.addr[0]}) via pm4 packet')
-  return pkt3(ctx, op, reg.addr[0] - set_packet_start, *(args or (reg.encode(**kwargs),)))
+class AMDComputeQueue(HWQueue):
+  q_rewrite = PatternMatcher([
+    (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), lambda ctx, call, prg: ctx.exec(call, prg)),
+    (UPat(Ops.INS, arg=("barrier", dtypes.void)), lambda ctx: ctx.memory_barrier()),
+    (UPat(Ops.INS, arg=("wait", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))), lambda ctx, dst, val: ctx.wait(dst, val)),
+    (UPat(Ops.INS, arg=("timestamp", dtypes.void), src=(UPat(name="dst"),)), lambda ctx, dst: ctx.timestamp(dst)),
+    (UPat(Ops.INS, arg=("store", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))),
+     lambda ctx, dst, val: ctx.signal(dst, val)),
+  ])
 
-def wait_reg_mem(ctx, value, mask=0xffffffff, mem=None, reg=None, reg_done=0, op=WAIT_REG_MEM_FUNCTION_GEQ):
-  wrm_info_dw = ctx.pm4.WAIT_REG_MEM_MEM_SPACE(int(mem is not None)) | ctx.pm4.WAIT_REG_MEM_OPERATION(int(mem is None and reg_done > 0)) \
-              | ctx.pm4.WAIT_REG_MEM_FUNCTION(op) | ctx.pm4.WAIT_REG_MEM_ENGINE(0)
-  return pkt3(ctx, PM4Ops.WAIT_REG_MEM, wrm_info_dw, *(data64_le(mem) if mem is not None else (reg, reg_done)), value, mask, 4)
+  def __init__(self, ctx, submit):
+    super().__init__(ctx, submit)
+    self.pm4, self.gc, self.soc, self.nbio, self.target = self.dev.pm4, self.dev.gc, self.dev.soc, self.dev.nbio, self.dev.target
 
-def acquire_mem(ctx, addr=0x0, sz=(1 << 64)-1, gli=1, glm=1, glk=1, glv=1, gl1=1, gl2=1):
-  if ctx.target[0] != 9:
-    cache_flags_dw = ctx.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLI_INV(gli) \
-                   | ctx.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLM_INV(glm) | ctx.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLM_WB(glm) \
-                   | ctx.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLK_INV(glk) | ctx.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLK_WB(glk) \
-                   | ctx.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLV_INV(glv) | ctx.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GL1_INV(gl1) \
-                   | ctx.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GL2_INV(gl2) | ctx.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GL2_WB(gl2)
-    return pkt3(ctx, PM4Ops.ACQUIRE_MEM, 0, *data64_le(sz), *data64_le(addr), 0, cache_flags_dw)
-  cp_coher_cntl = ctx.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_SH_ICACHE_ACTION_ENA(gli) | \
-                  ctx.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_SH_KCACHE_ACTION_ENA(glk) | \
-                  ctx.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_TC_ACTION_ENA(gl2) | \
-                  ctx.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_TCL1_ACTION_ENA(gl1) | \
-                  ctx.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_TC_WB_ACTION_ENA(gl2)
-  return pkt3(ctx, PM4Ops.ACQUIRE_MEM, cp_coher_cntl, *data64_le(sz), *data64_le(addr), 0x0000000A)
+  def pkt3(self, cmd, *vals): self.q(self.pm4.PACKET3(cmd, _dw(vals) - 1), *vals)
 
-def release_mem(ctx, address=0x0, value=0, data_sel=0, int_sel=2, ctxid=0, cache_flush=False):
-  if ctx.target[0] != 9:
-    cache_flags_dw = 0 if not cache_flush else (ctx.pm4.PACKET3_RELEASE_MEM_GCR_GLV_INV | ctx.pm4.PACKET3_RELEASE_MEM_GCR_GL1_INV \
-                   | ctx.pm4.PACKET3_RELEASE_MEM_GCR_GL2_INV | ctx.pm4.PACKET3_RELEASE_MEM_GCR_GLM_WB \
-                   | ctx.pm4.PACKET3_RELEASE_MEM_GCR_GLM_INV | ctx.pm4.PACKET3_RELEASE_MEM_GCR_GL2_WB | ctx.pm4.PACKET3_RELEASE_MEM_GCR_SEQ)
-    event_dw = ctx.pm4.PACKET3_RELEASE_MEM_EVENT_TYPE(ctx.pm4.CACHE_FLUSH_AND_INV_TS_EVENT) \
-             | ctx.pm4.PACKET3_RELEASE_MEM_EVENT_INDEX(ctx.pm4.event_index__mec_release_mem__end_of_pipe)
-    memsel_dw = ctx.pm4.PACKET3_RELEASE_MEM_DATA_SEL(data_sel) | ctx.pm4.PACKET3_RELEASE_MEM_INT_SEL(int_sel) \
-              | ctx.pm4.PACKET3_RELEASE_MEM_DST_SEL(0)
-  else:
-    cache_flags_dw = 0 if not cache_flush else (ctx.pm4.EOP_TC_WB_ACTION_EN | ctx.pm4.EOP_TC_NC_ACTION_EN)
-    event_dw = ctx.pm4.EVENT_TYPE(ctx.pm4.CACHE_FLUSH_AND_INV_TS_EVENT) | ctx.pm4.EVENT_INDEX(ctx.pm4.event_index__mec_release_mem__end_of_pipe)
-    memsel_dw = ctx.pm4.DATA_SEL(data_sel) | ctx.pm4.INT_SEL(int_sel)
-    ctxid = 0
-  return pkt3(ctx, PM4Ops.RELEASE_MEM, event_dw | cache_flags_dw, memsel_dw, *data64_le(address), *data64_le(value), ctxid)
+  def wreg(self, reg:AMDReg, *args:sint, **kwargs:int):
+    if bool(args) == bool(kwargs): raise RuntimeError('One (and only one) of *args or **kwargs must be specified')
+    if self.pm4.PACKET3_SET_SH_REG_START <= reg.addr[0] < self.pm4.PACKET3_SET_SH_REG_END:
+      set_packet, set_packet_start = self.pm4.PACKET3_SET_SH_REG, self.pm4.PACKET3_SET_SH_REG_START
+    elif self.pm4.PACKET3_SET_UCONFIG_REG_START <= reg.addr[0] < self.pm4.PACKET3_SET_UCONFIG_REG_START + 2**16-1:
+      set_packet, set_packet_start = self.pm4.PACKET3_SET_UCONFIG_REG, self.pm4.PACKET3_SET_UCONFIG_REG_START
+    else: raise RuntimeError(f'Cannot set {reg.name} ({reg.addr[0]}) via pm4 packet')
+    self.pkt3(set_packet, reg.addr[0] - set_packet_start, *(args or (reg.encode(**kwargs),)))
 
-def memory_barrier(ctx):
-  pf = '' if ctx.nbio.version[0] == 2 else '0' if ctx.nbio.version[:2] != (7, 11) else '1'
-  return UOp(Ops.LINEAR, dtypes.void, (
-    wait_reg_mem(ctx, reg=getattr(ctx.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_REQ').addr[0],
-                 reg_done=getattr(ctx.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_DONE').addr[0], value=0xffffffff),
-    acquire_mem(ctx)))
+  def wait_reg_mem(self, value, mask=0xffffffff, mem=None, reg=None, reg_done=0, op=WAIT_REG_MEM_FUNCTION_GEQ):
+    wrm_info_dw = self.pm4.WAIT_REG_MEM_MEM_SPACE(int(mem is not None)) | self.pm4.WAIT_REG_MEM_OPERATION(int(mem is None and reg_done > 0)) \
+                | self.pm4.WAIT_REG_MEM_FUNCTION(op) | self.pm4.WAIT_REG_MEM_ENGINE(0)
+    self.pkt3(self.pm4.PACKET3_WAIT_REG_MEM, wrm_info_dw, *((mem,) if mem is not None else (reg, reg_done)), value, mask, 4)
 
-def pm4_wait(ctx, dst, val): return wait_reg_mem(ctx, val, mem=dst.getaddr(ctx.devs))
+  def acquire_mem(self, addr=0x0, sz=(1 << 64)-1, gli=1, glm=1, glk=1, glv=1, gl1=1, gl2=1):
+    if self.target[0] != 9:
+      cache_flags_dw = self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLI_INV(gli) \
+                     | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLM_INV(glm) | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLM_WB(glm) \
+                     | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLK_INV(glk) | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLK_WB(glk) \
+                     | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLV_INV(glv) | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GL1_INV(gl1) \
+                     | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GL2_INV(gl2) | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GL2_WB(gl2)
+      return self.pkt3(self.pm4.PACKET3_ACQUIRE_MEM, 0, *data64_le(sz), *data64_le(addr), 0, cache_flags_dw)
+    cp_coher_cntl = self.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_SH_ICACHE_ACTION_ENA(gli) | \
+                    self.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_SH_KCACHE_ACTION_ENA(glk) | \
+                    self.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_TC_ACTION_ENA(gl2) | \
+                    self.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_TCL1_ACTION_ENA(gl1) | \
+                    self.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_TC_WB_ACTION_ENA(gl2)
+    return self.pkt3(self.pm4.PACKET3_ACQUIRE_MEM, cp_coher_cntl, *data64_le(sz), *data64_le(addr), 0x0000000A)
 
-def pm4_barrier(ctx): return memory_barrier(ctx)
+  def release_mem(self, address=0x0, value=0, data_sel=0, int_sel=2, ctxid=0, cache_flush=False):
+    if self.target[0] != 9:
+      cache_flags_dw = 0 if not cache_flush else (self.pm4.PACKET3_RELEASE_MEM_GCR_GLV_INV | self.pm4.PACKET3_RELEASE_MEM_GCR_GL1_INV \
+                     | self.pm4.PACKET3_RELEASE_MEM_GCR_GL2_INV | self.pm4.PACKET3_RELEASE_MEM_GCR_GLM_WB \
+                     | self.pm4.PACKET3_RELEASE_MEM_GCR_GLM_INV | self.pm4.PACKET3_RELEASE_MEM_GCR_GL2_WB | self.pm4.PACKET3_RELEASE_MEM_GCR_SEQ)
+      event_dw = self.pm4.PACKET3_RELEASE_MEM_EVENT_TYPE(self.pm4.CACHE_FLUSH_AND_INV_TS_EVENT) \
+               | self.pm4.PACKET3_RELEASE_MEM_EVENT_INDEX(self.pm4.event_index__mec_release_mem__end_of_pipe)
+      memsel_dw = self.pm4.PACKET3_RELEASE_MEM_DATA_SEL(data_sel) | self.pm4.PACKET3_RELEASE_MEM_INT_SEL(int_sel) \
+                | self.pm4.PACKET3_RELEASE_MEM_DST_SEL(0)
+    else:
+      cache_flags_dw = 0 if not cache_flush else (self.pm4.EOP_TC_WB_ACTION_EN | self.pm4.EOP_TC_NC_ACTION_EN)
+      event_dw = self.pm4.EVENT_TYPE(self.pm4.CACHE_FLUSH_AND_INV_TS_EVENT) | \
+                 self.pm4.EVENT_INDEX(self.pm4.event_index__mec_release_mem__end_of_pipe)
+      memsel_dw = self.pm4.DATA_SEL(data_sel) | self.pm4.INT_SEL(int_sel)
+      ctxid = 0
+    addr_w = address if isinstance(address, UOp) else UOp.const(address, dtypes.uint64)
+    val_w = value.cast(dtypes.uint64) if isinstance(value, UOp) else UOp.const(value, dtypes.uint64)
+    self.pkt3(self.pm4.PACKET3_RELEASE_MEM, event_dw | cache_flags_dw, memsel_dw, addr_w, val_w, ctxid)
 
-def pm4_store(ctx, dst, val):
-  if val.op is Ops.BINARY: return None
-  return release_mem(ctx, dst.getaddr(ctx.devs), val, ctx.pm4.data_sel__mec_release_mem__send_32_bit_low,
-                     ctx.pm4.int_sel__mec_release_mem__send_interrupt_after_write_confirm, cache_flush=True)
+  def memory_barrier(self):
+    pf = '' if self.nbio.version[0] == 2 else '0' if self.nbio.version[:2] != (7, 11) else '1'
+    self.wait_reg_mem(reg=getattr(self.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_REQ').addr[0],
+                      reg_done=getattr(self.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_DONE').addr[0], value=0xffffffff)
+    self.acquire_mem()
 
-def pm4_timestamp(ctx, dst):
-  return release_mem(ctx, dst.getaddr(ctx.devs), 0, ctx.pm4.data_sel__mec_release_mem__send_gpu_clock_counter,
-                     ctx.pm4.int_sel__mec_release_mem__none)
+  def exec(self, call:UOp, prg:UOp):
+    data, lib = amd_build_program(self.dev, prg, self.devs)
+    info = prg.arg
 
-def pm4_program(ctx, call, prg):
-  data, info = prg.arg
-  lib_gpu = prg.src[0]
-  args = encode_kernargs_clike(call, prg, ctx.devs)
-  prog_addr = lib_gpu.getaddr(ctx.devs) + data.entry_point_offset
-  scratch_addr = UOp.placeholder((data.private_segment_size,), dtypes.uint8, 0, device=ctx.devs).rtag("scratch").getaddr(ctx.devs)
-  args_addr = args.getaddr(ctx.devs)
+    # kernargs: a nested blob linear inside a getaddr, packed into the tail of the cmdbuf
+    ka_words = [get_call_arg_uops(call)[gi].getaddr(self.devs) for gi in info.globals] + \
+               [b.ccast(v.dtype) for v, b in zip(info.vars, get_call_var_uops(call, prg))] # a bound value is a bare const, the var has the width
+    pad = data.kernargs_alloc_size - sum(w.dtype.itemsize for w in ka_words)
+    assert pad >= 0 and pad % 4 == 0, f"bad kernargs padding {pad}"
+    ka = UOp(Ops.LINEAR, src=tuple(ka_words) + (UOp.const(0, dtypes.uint32),) * (pad // 4))
 
-  user_regs = []
-  if data.enable_private_segment_sgpr:
-    scratch_hilo = data64_le(scratch_addr)
-    user_regs = [scratch_hilo[0], scratch_hilo[1] | 1 << 31, 0xffffffff, 0x20c14000]
-  if data.enable_dispatch_ptr: user_regs += [*data64_le(args_addr + data.kernargs_segment_size)]
-  user_regs += [*data64_le(args_addr)]
+    prog_addr = lib.getaddr(self.devs) + data.entry_point_offset
+    scratch_addr = UOp.placeholder((data.private_segment_size,), dtypes.uint8, 0, device=self.devs).rtag("scratch").getaddr(self.devs)
+    args_addr = ka.getaddr(self.devs)
 
-  dispatch_init = ctx.gc.regCOMPUTE_DISPATCH_INITIATOR.encode(
-    **({'cs_w32_en': int(data.wave32)} if ctx.target[0] != 9 else {}), force_start_at_000=1, compute_shader_en=1)
-  ins = [acquire_mem(ctx, gli=0, gl2=0),
-    wreg(ctx, ctx.gc.regCOMPUTE_PGM_LO, *data64_le(prog_addr >> 8)),
-    wreg(ctx, ctx.gc.regCOMPUTE_PGM_RSRC1, data.rsrc1, data.rsrc2),
-    wreg(ctx, ctx.gc.regCOMPUTE_PGM_RSRC3, data.rsrc3),
-    wreg(ctx, ctx.gc.regCOMPUTE_TMPRING_SIZE, ctx.tmpring_size(data.private_segment_size))]
-  ins += [wreg(ctx, ctx.gc.regCOMPUTE_DISPATCH_SCRATCH_BASE_LO, *data64_le((scratch_addr + data.private_segment_size // ctx.xccs * xcc_id) >> 8))
-          for xcc_id in range(ctx.xccs)]
-  ins += [wreg(ctx, ctx.gc.regCOMPUTE_RESTART_X, 0, 0, 0),
-    wreg(ctx, ctx.gc.regCOMPUTE_USER_DATA_0, *user_regs),
-    wreg(ctx, ctx.gc.regCOMPUTE_RESOURCE_LIMITS, ctx.gc.regCOMPUTE_RESOURCE_LIMITS.encode(waves_per_sh=getenv("WAVES_PER_SH"))),
-    wreg(ctx, ctx.gc.regCOMPUTE_START_X, 0, 0, 0, *(info.local_size or (1, 1, 1)), 0, 0),
-    pkt3(ctx, PM4Ops.DISPATCH_DIRECT, *info.global_size, dispatch_init),
-    pkt3(ctx, PM4Ops.EVENT_WRITE, ctx.pm4.EVENT_TYPE(ctx.soc.CS_PARTIAL_FLUSH) | ctx.pm4.EVENT_INDEX(EVENT_INDEX_PARTIAL_FLUSH))]
-  return UOp(Ops.LINEAR, dtypes.void, tuple(ins))
+    user_regs:list = []
+    if data.enable_private_segment_sgpr: user_regs = [scratch_addr | (1 << 63), 0xffffffff, 0x20c14000]
+    if data.enable_dispatch_ptr: user_regs += [args_addr + data.kernargs_segment_size]
+    user_regs += [args_addr]
 
-pm_pm4_opsel = PatternMatcher([
-  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), pm4_program),
+    dispatch_init = self.gc.regCOMPUTE_DISPATCH_INITIATOR.encode(
+      **({'cs_w32_en': int(data.wave32)} if self.target[0] != 9 else {}), force_start_at_000=1, compute_shader_en=1)
+    self.acquire_mem(gli=0, gl2=0)
+    self.wreg(self.gc.regCOMPUTE_PGM_LO, prog_addr >> 8)
+    self.wreg(self.gc.regCOMPUTE_PGM_RSRC1, data.rsrc1, data.rsrc2)
+    self.wreg(self.gc.regCOMPUTE_PGM_RSRC3, data.rsrc3)
+    self.wreg(self.gc.regCOMPUTE_TMPRING_SIZE, self.dev.tmpring_size(data.private_segment_size))
+    for xcc_id in range(self.dev.xccs):
+      self.wreg(self.gc.regCOMPUTE_DISPATCH_SCRATCH_BASE_LO, (scratch_addr + data.private_segment_size // self.dev.xccs * xcc_id) >> 8)
+    self.wreg(self.gc.regCOMPUTE_RESTART_X, 0, 0, 0)
+    self.wreg(self.gc.regCOMPUTE_USER_DATA_0, *user_regs)
+    self.wreg(self.gc.regCOMPUTE_RESOURCE_LIMITS, self.gc.regCOMPUTE_RESOURCE_LIMITS.encode(waves_per_sh=getenv("WAVES_PER_SH")))
+    self.wreg(self.gc.regCOMPUTE_START_X, 0, 0, 0, *info.local_size, 0, 0)
+    self.pkt3(self.pm4.PACKET3_DISPATCH_DIRECT, *info.global_size, dispatch_init)
+    self.pkt3(self.pm4.PACKET3_EVENT_WRITE, self.pm4.EVENT_TYPE(self.soc.CS_PARTIAL_FLUSH) | self.pm4.EVENT_INDEX(EVENT_INDEX_PARTIAL_FLUSH))
 
-  (UPat(Ops.INS, arg="wait", src=(UPat(name="dst"), UPat(name="val"))), pm4_wait),
-  (UPat(Ops.INS, arg="barrier"), pm4_barrier),
-  (UPat(Ops.INS, arg="timestamp", src=(UPat(name="dst"),)), pm4_timestamp),
-  (UPat(Ops.INS, arg="store", src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), pm4_store),
-])
+  def wait(self, signal:UOp, value:UOp): self.wait_reg_mem(value.cast(dtypes.uint32), mem=signal.getaddr(self.devs))
 
-def pm4_submit(ctx, lin):
-  # ensure compute queues are allocated
-  for d in (devs:=ctx.devs): q = Device[d].compute_queue
-  ring, wptr, doorbell, put_ptr = (UOp.placeholder((b.size,), b.dtype, 0, device=devs).rtag(f"COMPUTE:0_{name}")
-    for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))
+  def timestamp(self, signal:UOp):
+    self.release_mem(signal.getaddr(self.devs), 0, self.pm4.data_sel__mec_release_mem__send_gpu_clock_counter,
+                     self.pm4.int_sel__mec_release_mem__none)
 
-  # the host fence at the start of the batch guarantees the ib is free to reuse
-  size_dw = sum(len(ins.src) for ins in lin.src)
-  assert size_dw < (1 << 20), f"indirect buffer of {size_dw} dwords doesn't fit one packet"
+  def signal(self, signal:UOp, value:UOp):
+    self.release_mem(signal.getaddr(self.devs), value, self.pm4.data_sel__mec_release_mem__send_32_bit_low,
+                     self.pm4.int_sel__mec_release_mem__send_interrupt_after_write_confirm, cache_flush=True)
 
-  ib = UOp.placeholder((size_dw,), dtypes.uint32, next(UOp.unique_num), device=devs, volatile=True).rtag("cmdbuf")
-  cmdbuf = make_cmdbuf(lin, devs, buf=ib)
+  def submit(self, cmdbuf:UOp) -> UOp:
+    q = self.dev.compute_queue
 
-  # the ring itself only carries a packet pointing at the ib, wrapping the ring
-  put = put_ptr.index(zero:=UOp.const(0, dtypes.int))
-  pkt = (ctx.pm4.PACKET3(ctx.pm4.PACKET3_INDIRECT_BUFFER, 2), *data64_le(cmdbuf.getaddr(devs)), size_dw | ctx.pm4.INDIRECT_BUFFER_VALID)
-  write_pkt = UOp.barrier(*[ring.index(((put + off) % q.ring.size).cast(dtypes.int)).store(UOp.const(x, dtypes.uint32)) for off,x in enumerate(pkt)])
+    ring, wptr, doorbell, put = _queue_args(self, q)
 
-  # advance the put/write pointers past the packet
-  bump_put_ptr = put_ptr.index(zero).store(put + len(pkt))
-  bump_wptr = wptr.index(zero).store(put + len(pkt))
-  flush = UOp.barrier(write_pkt, bump_put_ptr, bump_wptr)
-  return doorbell.after(flush).index(zero).store(put + len(pkt))
-
-pm_pm4_submit = PatternMatcher([(UPat(Ops.LINEAR, name="lin"), pm4_submit)])
+    size_dw = cmdbuf.max_numel() // 4
+    p = put.index(0).load()
+    i = UOp.range(size_dw, 10, dtype=dtypes.int, src=(cmdbuf,))
+    copy = ring.index(((p + i.cast(p.dtype)) % q.ring.size).cast(dtypes.int)).store(cmdbuf.bitcast(dtypes.uint32).index(i).load()).end(i)
+    next_put = p + size_dw
+    flush = UOp.barrier(copy, put.index(0).store(next_put), wptr.index(0).store(next_put))
+    return doorbell.after(flush).index(0).store(next_put)
 
 # *****************
 # SDMA
 
-class SDMAOps(FastEnum): COPY = auto(); POLL_REGMEM = auto(); FENCE = auto(); TRAP = auto(); TIMESTAMP = auto()  # noqa: E702
+class AMDSDMAQueue(HWQueue):
+  q_rewrite = PatternMatcher([
+    (UPat(Ops.CALL, src=(UPat(Ops.COPY),), name="call", allow_any_len=True), lambda ctx, call: ctx.copy(call)),
+    (UPat(Ops.INS, arg=("barrier", dtypes.void)), lambda ctx: ()),
+    (UPat(Ops.INS, arg=("wait", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))), lambda ctx, dst, val: ctx.wait(dst, val)),
+    (UPat(Ops.INS, arg=("timestamp", dtypes.void), src=(UPat(name="dst"),)), lambda ctx, dst: ctx.timestamp(dst)),
+    (UPat(Ops.INS, arg=("store", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))),
+     lambda ctx, dst, val: ctx.signal(dst, val)),
+  ])
 
-def sdma_copy(ctx, call):
-  sz = call.src[2].max_numel() * call.src[2].dtype.itemsize
-  hdr = ctx.sdma.SDMA_OP_COPY | ctx.sdma.SDMA_PKT_COPY_LINEAR_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_COPY_LINEAR)
-  return call.ins(SDMAOps.COPY, src=tuple(x for off in range(0, sz, ctx.max_copy_size) for x in (
-    *(UOp.const(v, dtypes.uint32) for v in (hdr, ctx.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(sz-off, ctx.max_copy_size)-1), 0)),
-    *(a + UOp.const(off, dtypes.uint64) if off else a for a in (call.src[2].getaddr(ctx.devs), call.src[1].getaddr(ctx.devs))))))
+  def __init__(self, ctx, submit):
+    super().__init__(ctx, submit)
+    self.sdma, self.target, self.max_copy_size = self.dev.sdma, self.dev.target, self.dev.max_copy_size
 
-def sdma_wait(ctx, ins, dst, val):
-  op = ctx.sdma.SDMA_OP_POLL_REGMEM | ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_FUNC(WAIT_REG_MEM_FUNCTION_GEQ) \
-     | ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_MEM_POLL(1)
-  return ins.ins(SDMAOps.POLL_REGMEM, src=tuple(UOp.const(x, dtypes.uint32) for x in (
-    op, *data64_le(dst.getaddr(ctx.devs)), val, 0xffffffff,
-    ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_INTERVAL(0x04) | ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_RETRY_COUNT(0xfff))))
+  def copy(self, call:UOp):
+    sz = call.src[2].max_numel() * call.src[2].dtype.itemsize
+    hdr = self.sdma.SDMA_OP_COPY | self.sdma.SDMA_PKT_COPY_LINEAR_HEADER_SUB_OP(self.sdma.SDMA_SUBOP_COPY_LINEAR)
+    for off in range(0, sz, self.max_copy_size):
+      self.q(hdr, self.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(sz-off, self.max_copy_size)-1), 0,
+             *(a + UOp.const(off, dtypes.uint64) if off else a for a in (call.src[2].getaddr(self.devs), call.src[1].getaddr(self.devs))))
 
-def sdma_store(ctx, ins, dst, val):
-  op = ctx.sdma.SDMA_OP_FENCE | (ctx.sdma.SDMA_PKT_FENCE_HEADER_MTYPE(3) if ctx.target[0] != 9 else 0)
-  return UOp(Ops.LINEAR, src=(
-    ins.ins(SDMAOps.FENCE, src=tuple(UOp.const(x, dtypes.uint32) for x in (op, *data64_le(dst.getaddr(ctx.devs)), val))),
-    ins.ins(SDMAOps.TRAP, src=tuple(UOp.const(x, dtypes.uint32) for x in (ctx.sdma.SDMA_OP_TRAP, 0)))))
+  def wait(self, signal:UOp, value:UOp):
+    op = self.sdma.SDMA_OP_POLL_REGMEM | self.sdma.SDMA_PKT_POLL_REGMEM_HEADER_FUNC(WAIT_REG_MEM_FUNCTION_GEQ) \
+       | self.sdma.SDMA_PKT_POLL_REGMEM_HEADER_MEM_POLL(1)
+    self.q(op, signal.getaddr(self.devs), value.cast(dtypes.uint32), 0xffffffff,
+           self.sdma.SDMA_PKT_POLL_REGMEM_DW5_INTERVAL(0x04) | self.sdma.SDMA_PKT_POLL_REGMEM_DW5_RETRY_COUNT(0xfff))
 
-def sdma_timestamp(ctx, ins, dst):
-  op = ctx.sdma.SDMA_OP_TIMESTAMP | ctx.sdma.SDMA_PKT_TIMESTAMP_GET_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_TIMESTAMP_GET_GLOBAL)
-  return ins.ins(SDMAOps.TIMESTAMP, src=tuple(UOp.const(x, dtypes.uint32) for x in (op, *data64_le(dst.getaddr(ctx.devs)))))
+  def timestamp(self, signal:UOp):
+    self.q(self.sdma.SDMA_OP_TIMESTAMP | self.sdma.SDMA_PKT_TIMESTAMP_GET_HEADER_SUB_OP(self.sdma.SDMA_SUBOP_TIMESTAMP_GET_GLOBAL),
+           signal.getaddr(self.devs))
 
-pm_sdma_opsel = PatternMatcher([
-  (UPat(Ops.CALL, src=(UPat(Ops.COPY),), name="call", allow_any_len=True), sdma_copy),
+  def signal(self, signal:UOp, value:UOp): # a fence packet then a trap
+    op = self.sdma.SDMA_OP_FENCE | (self.sdma.SDMA_PKT_FENCE_HEADER_MTYPE(3) if self.target[0] != 9 else 0)
+    self.q(op, signal.getaddr(self.devs), value.cast(dtypes.uint32), self.sdma.SDMA_OP_TRAP, 0)
 
-  (UPat(Ops.INS, arg="barrier"), lambda: UOp(Ops.NOOP, dtypes.void, ())),
-  (UPat(Ops.INS, arg="wait", src=(UPat(name="dst"), UPat(name="val")), name="ins"), sdma_wait),
-  (UPat(Ops.INS, arg="timestamp", src=(UPat(name="dst"),), name="ins"), sdma_timestamp),
-  (UPat(Ops.INS, arg="store", src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val")), name="ins"), sdma_store),
-])
+  def submit(self, cmdbuf:UOp) -> UOp:
+    # sdma needs the cmdbuf contiguous in the ring: if it won't fit before the ring end, restart at 0 and zero the tail
+    q = unwrap(self.dev.sdma_queue(int(self.queue.split(":")[1])))
 
-def sdma_submit(cmdbuf, devs):
-  # the cmdbuf to submit + the patch writes that fill it
-  size_dw, zero = cmdbuf.nbytes() // dtypes.uint32.itemsize, UOp.const(0, dtypes.int)
+    ring, wptr, doorbell, put = _queue_args(self, q)
 
-  # the sdma queue's ring and its host-side ring/write/put pointers
-  for d in devs: q = Device[d].sdma_queue(0)
-  ring, wptr, doorbell, put_ptr = (UOp.placeholder((b.size,), b.dtype, 0, device=devs).rtag(f"COPY:0_{name}")
-    for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))
-
-  # sdma needs the cmdbuf contiguous: if it won't fit before the ring end, restart at 0 and zero the tail
-  put_b = put_ptr.index(zero)
-  tail_off_dw = ((put_b % (q.ring.size * 4)) // 4).cast(dtypes.int)
-  fits = (size_dw <= q.ring.size - tail_off_dw).cast(dtypes.int)
-  start_dw = fits * tail_off_dw
-  zero_amt_dw = (1 - fits) * (q.ring.size - tail_off_dw)
-
-  # zero the wrapped tail, then copy the cmdbuf into the ring
-  zi = UOp.range(zero_amt_dw, 0, dtype=dtypes.int, src=(cmdbuf,))
-  zero_tail = ring.index(tail_off_dw + zi).store(UOp.const(0, dtypes.uint32)).end(zi)
-  i = UOp.range(UOp.const(size_dw, dtypes.int), 0, dtype=dtypes.int, src=(cmdbuf,))
-  copy_to_ring = ring.index(start_dw + i).store(cmdbuf.index(i).load()).end(i)
-
-  # advance the put/write pointers past the zeroed tail and the cmdbuf
-  next_put_b = put_b + ((zero_amt_dw + size_dw) * 4).cast(put_b.dtype)
-  bump_put_ptr = put_ptr.index(zero).store(next_put_b)
-  bump_wptr = wptr.index(zero).store(next_put_b)
-
-  # ring the doorbell once the writes have landed
-  flush = UOp.barrier(zero_tail, copy_to_ring, bump_put_ptr, bump_wptr)
-  return doorbell.after(flush).index(zero).store(next_put_b)
-
-pm_sdma_submit = PatternMatcher([(UPat(Ops.LINEAR, name="lin"),
-  lambda ctx, lin: sdma_submit(make_cmdbuf(lin, ctx.devs), ctx.devs))])
-
-@dataclass(frozen=True)
-class AMDEncodeCtx:  # encode-time constants for one queue: devs (every cmdbuf address resolves into these) + gfx version + packet/ip modules
-  devs: tuple[str, ...]; target: tuple[int, ...]; pm4: Any; sdma: Any; soc: Any  # noqa: E702
-  gc: AMDIP; nbio: AMDIP; xccs: int; max_copy_size: int; tmpring_size: Callable  # noqa: E702
-
-def encode_queue(q:UOp) -> UOp|None:
-  d = Device[(devs:=to_tuple(q.arg[0]))[0]]
-  ctx = AMDEncodeCtx(devs, d.target, d.pm4, d.sdma, d.soc, d.gc, d.nbio, d.xccs, d.max_copy_size, d.tmpring_size)
-  opsel, submit = (pm_pm4_opsel, pm_pm4_submit) if q.arg[1].startswith("COMPUTE") else (pm_sdma_opsel, pm_sdma_submit)
-  return submit.rewrite(graph_rewrite(q, opsel + pm_flatten_linear, walk=True, ctx=ctx, name=f"{q.arg[1]} opsel"), ctx)
+    rs, size_dw = q.ring.size, cmdbuf.max_numel() // 4
+    put_b = put.index(0).load()
+    tail = ((put_b % (rs * 4)) // 4).cast(dtypes.int)
+    fits = (size_dw <= rs - tail).cast(dtypes.int)
+    start_dw, zero_amt = fits * tail, (1 - fits) * (rs - tail)
+    zi = UOp.range(zero_amt, 10, dtype=dtypes.int, src=(cmdbuf,))
+    zero_tail = ring.index(tail + zi).store(UOp.const(0, dtypes.uint32)).end(zi)
+    i = UOp.range(size_dw, 11, dtype=dtypes.int, src=(cmdbuf,))
+    copy = ring.index(start_dw + i).store(cmdbuf.bitcast(dtypes.uint32).index(i).load()).end(i)
+    next_put = put_b + ((zero_amt + size_dw) * 4).cast(put_b.dtype)
+    flush = UOp.barrier(zero_tail, copy, put.index(0).store(next_put), wptr.index(0).store(next_put))
+    return doorbell.after(flush).index(0).store(next_put)
 
 @dataclass(frozen=True)
 class AMDProgramData:
@@ -261,34 +228,38 @@ class AMDProgramData:
   private_segment_size:int; kernargs_segment_size:int; kernargs_alloc_size:int
   enable_dispatch_ptr:int; enable_private_segment_sgpr:int
 
-_amd_program_cache:dict[tuple[bytes, tuple[str, ...]], UOp] = {}
-def amd_build_program(prg:UOp) -> UOp:
-  dev = Device[to_tuple(prg.device)[0]] # TODO: rm this
-  # key on the full device tuple: the same lib can be built for different device sets, each needs its own program buffer
-  if (cached:=_amd_program_cache.get(key:=(lib:=prg.src[3].arg, to_tuple(prg.device)))) is None:
-    image, sections, relocs = elf_loader(lib)
-    rodata = next(sh.header.sh_addr for sh in sections if sh.name == ".rodata")
-    for off, sym, typ, addent in relocs:
-      assert typ == 5, f"unknown AMD reloc {typ}"  # R_AMDGPU_REL64
-      image[off:off+8] = struct.pack('<q', sym - off + addent)
-    desc = amdgpu_kd.llvm_amdhsa_kernel_descriptor_t.from_buffer_copy(bytes(image[rodata:rodata+ctypes.sizeof(amdgpu_kd.llvm_amdhsa_kernel_descriptor_t)]))
-    if (lds:=((desc.group_segment_fixed_size+511)//512)&0x1FF) > (dev.iface.props['lds_size_in_kb']*1024)//512:
-      raise RuntimeError("Too many resources requested: group_segment_size")
-    edp = desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_DISPATCH_PTR
-
-    data = AMDProgramData(entry_point_offset=rodata + desc.kernel_code_entry_byte_offset,
-      rsrc1=desc.compute_pgm_rsrc1 | ((1<<20) if dev.target[0]==11 else 0),  # priv=1 on gfx11 for cwsr
-      rsrc2=desc.compute_pgm_rsrc2 | (lds<<15), rsrc3=desc.compute_pgm_rsrc3,
-      wave32=bool(desc.kernel_code_properties & 0x400), private_segment_size=desc.private_segment_fixed_size, kernargs_segment_size=desc.kernarg_size,
-      kernargs_alloc_size=desc.kernarg_size + (ctypes.sizeof(hsa.hsa_kernel_dispatch_packet_t) if edp else 0), enable_dispatch_ptr=edp,
-      enable_private_segment_sgpr=desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER)
-    buf = UOp.placeholder((len(image),), dtypes.uint8, next(UOp.unique_num), device=prg.device).rtag("program")
-    cached = _amd_program_cache[key] = prg.replace(src=(buf.after(make_binary_patch(buf, bytes(image))),), arg=(data, prg.arg))
+_amd_program_cache:dict[tuple[bytes, tuple[str, ...]], tuple[AMDProgramData, UOp]] = {}
+def amd_build_program(dev, prg:UOp, devs:tuple[str, ...]) -> tuple[AMDProgramData, UOp]:
+  # the image parses once per lib, each device set gets its own program buffer of it
+  if (cached:=_amd_program_cache.get(key:=(lib:=prg.src[3].arg, devs))) is None:
+    data, image = _amd_program_image(dev, lib)
+    buf = UOp.placeholder((len(image),), dtypes.uint8, next(UOp.unique_num), device=devs).rtag("program")
+    cached = _amd_program_cache[key] = (data, buf.after(buf.store(UOp(Ops.BINARY, src=(), arg=image).bitcast(buf.dtype))))
   return cached
 
+@functools.cache
+def _amd_program_image(dev, lib:bytes) -> tuple[AMDProgramData, bytes]:
+  image, sections, relocs = elf_loader(lib)
+  rodata = next(sh.header.sh_addr for sh in sections if sh.name == ".rodata")
+  for off, sym, typ, addent in relocs:
+    assert typ == 5, f"unknown AMD reloc {typ}"  # R_AMDGPU_REL64
+    image[off:off+8] = struct.pack('<q', sym - off + addent)
+  desc = amdgpu_kd.llvm_amdhsa_kernel_descriptor_t.from_buffer_copy(bytes(image[rodata:rodata+ctypes.sizeof(amdgpu_kd.llvm_amdhsa_kernel_descriptor_t)]))
+  if (lds:=((desc.group_segment_fixed_size+511)//512)&0x1FF) > (dev.iface.props['lds_size_in_kb']*1024)//512:
+    raise RuntimeError("Too many resources requested: group_segment_size")
+  edp = desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_DISPATCH_PTR
+
+  data = AMDProgramData(entry_point_offset=rodata + desc.kernel_code_entry_byte_offset,
+    rsrc1=desc.compute_pgm_rsrc1 | ((1<<20) if dev.target[0]==11 else 0),  # priv=1 on gfx11 for cwsr
+    rsrc2=desc.compute_pgm_rsrc2 | (lds<<15), rsrc3=desc.compute_pgm_rsrc3,
+    wave32=bool(desc.kernel_code_properties & 0x400), private_segment_size=desc.private_segment_fixed_size, kernargs_segment_size=desc.kernarg_size,
+    kernargs_alloc_size=desc.kernarg_size + (ctypes.sizeof(hsa.hsa_kernel_dispatch_packet_t) if edp else 0), enable_dispatch_ptr=edp,
+    enable_private_segment_sgpr=desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER)
+  return data, bytes(image).ljust(round_up(len(image), 4), b"\x00") # the program is uploaded as whole dwords
+
 class AMDAllocator(HCQAllocator['AMDDevice']):
   def __init__(self, dev:AMDDevice):
-    super().__init__(dev, supports_copy_from_disk=dev.has_copy_queue, supports_transfer=dev.has_copy_queue and not dev.is_usb())
+    super().__init__(dev, supports_copy_from_disk=dev.has_copy_queue, supports_transfer=dev.has_copy_queue and not dev.is_usb)
 
   def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer:
     return self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access or not self.dev.has_copy_queue)
@@ -524,8 +495,7 @@ class PCIIface(PCIIfaceBase):
       cq = d.compute_queue
       for b in (cq.put_value, cq.read_ptr, cq.write_ptr): b._buf.view.view(fmt='Q')[0] = 0
       d.iface.dev_impl.gfx.setup_ring(*cq.params)
-      d.signal('timeline')._buf.cpu_view().mv.cast('Q')[0] = \
-        d.signal('value', 1).as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] - 1
+      (tl:=d.timeline._buf.cpu_view().view(fmt='Q'))[0] = tl[1]
 
   def sleep(self, timeout):
     if hasattr(self.pci_dev, 'irq_poller') and self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))):
@@ -539,29 +509,52 @@ class PCIIface(PCIIfaceBase):
 
   def device_fini(self): self.dev_impl.fini()
 
+class USBIface(PCIIface):
+  def __init__(self, dev, dev_id): # pylint: disable=super-init-not-called
+    if dev_id >= len(visible:=hcq_filter_visible_devices(USB3.list_devices(0xADD1, 0x0001) + USB3.list_devices(0x3801, 0x0001), "AMD")):
+      raise RuntimeError(f"AMD:{dev_id} does not exist ({pluralize('device', len(visible))} available)")
+    self.dev, self.pci_dev, self.vram_bar, self.count = dev, USBPCIDevice("AM", *visible[dev_id]), 0, len(visible)
+    self.dev_impl = AMDev(self.pci_dev)
+    self._compute_props()
+    self.sram = self._dma_region(ctrl_addr=0xf000, sys_addr=0x200000, size=0x80000)
+    self.cq_buf = self._dma_region(ctrl_addr=0xb800, sys_addr=0x822000, size=0x1000) # +12 is the dword that releases an armed read
+    self.usb_handle = unwrap(ctypes.cast(self.pci_dev.usb.usb.handle, ctypes.c_void_p).value)
+
+  def _dma_region(self, ctrl_addr, sys_addr, size):
+    region = self.dev_impl.mm.map_range(vaddr:=self.dev_impl.mm.alloc_vaddr(size=size), size, [(sys_addr, size)], aspace=AddrSpace.SYS, uncached=True)
+    return HCQBuffer(vaddr, size, meta=PCIAllocationMeta(region, has_cpu_mapping=False), view=self.pci_dev.dma_view(ctrl_addr, size), owner=self.dev)
+
+  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, force_devmem=False, **kwargs) -> HCQBuffer:
+    # everything, even host-style signals, lives in vram: gpu writes into the bridge's own memory collide with an armed 0xF2 read stream
+    return super().alloc(size, host=False, uncached=uncached, cpu_access=cpu_access or host, contiguous=contiguous, force_devmem=True, **kwargs)
+
+  def sleep(self, timeout): pass
+
+  # we don't own the sram region, so the buffer never frees it
+  @functools.cached_property
+  def usb_sram(self) -> Buffer:
+    return Buffer(self.dev.device, (b:=self.sram).size, dtypes.uint8, options=BufferSpec(external_ptr=b.va_addr, nolru=True)).allocate(opaque=b)
+
 def _mock(iface, name=None): return type(name or f"MOCK{iface.__name__}", (iface,), {})
 
 class AMDDevice(HCQ2Compiled):
-  pm_lower = PatternMatcher([
-    # prep program
-    (UPat(Ops.PROGRAM, src=(UPat(), UPat(), UPat(), UPat(Ops.BINARY)), name="prg"), amd_build_program),
-
-    # encoding of cmdbuf
-    (UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="q"),)), encode_queue),
-  ])
-
   timestamp_divider = 100.0  # AMD GPU clock: ticks/us
   max_scratch_psize = 0
+  pm_encode = PatternMatcher([
+    (UPat(Ops.CUSTOM_FUNCTION, arg="submit_amd_compute", name="submit"), lambda ctx, submit: encode_submit(AMDComputeQueue(ctx, submit))),
+    (UPat(Ops.CUSTOM_FUNCTION, arg="submit_amd_copy", name="submit"), lambda ctx, submit: encode_submit(AMDSDMAQueue(ctx, submit))),
+  ])
 
-  ifaces = [KFDIface, PCIIface, _mock(KFDIface, "MOCKIface"), _mock(KFDIface), _mock(PCIIface)]
+  ifaces = [KFDIface, PCIIface, USBIface, _mock(KFDIface, "MOCKIface"), _mock(KFDIface), _mock(PCIIface), _mock(USBIface)]
 
   def device_props(self): return self.iface.props
 
   def is_am(self) -> bool: return isinstance(self.iface, (PCIIface,))
-  def is_usb(self) -> bool: return False
 
   def __init__(self, device:str=""):
     self.iface = self._select_iface(device)
+    self.is_usb = isinstance(self.iface, USBIface)
+    if self.is_usb: self.rt_nbytes = 4 << 20
 
     self.target:tuple[int, ...] = ((trgt:=self.iface.props['gfx_target_version']) // 10000, (trgt // 100) % 100, trgt % 100)
     self.arch = "gfx%d%x%x" % self.target
@@ -586,7 +579,7 @@ class AMDDevice(HCQ2Compiled):
 
     self.is_aql = getenv("AMD_AQL", int(self.xccs > 1))
     if self.is_aql:
-      self.pm4_ibs = self.iface.alloc(0x2000 if self.is_usb() else (16 << 20), uncached=True, cpu_access=True)
+      self.pm4_ibs = self.iface.alloc(0x2000 if self.is_usb else (16 << 20), uncached=True, cpu_access=True)
       self.pm4_ib_alloc = BumpAllocator(self.pm4_ibs.size, wrap=True)
 
     self.max_copy_size = 0x40000000 if self.iface.ip_versions[am.SDMA0_HWIP][0] >= 5 else 0x400000
@@ -597,7 +590,11 @@ class AMDDevice(HCQ2Compiled):
 
     # Scratch setup
     self.max_private_segment_size = 0
-    self.pm_bufferize = PatternMatcher([(UPat(Ops.PARAM, tag="scratch", name="b"), lambda ctx, b: ctx[0].scratch_buffer(b.max_numel()))]) + self.pm_bufferize
+    self.pm_bufferize = PatternMatcher([(UPat(Ops.PARAM, tag="scratch", name="b"), lambda ctx, b: ctx.scratch_buffer(b.max_numel()))]) + self.pm_bufferize
+
+    if self.is_usb:
+      self.pm_bufferize = pm_usb_bufferize + self.pm_bufferize
+      raise NotImplementedError("usb amd is not migrated to sealed submits yet") # a usb pm_lower can override the whole submit graph
 
     self.pmc_enabled:bool = PROFILE > 0 and PMC > 0
     if self.pmc_enabled:
@@ -645,7 +642,7 @@ class AMDDevice(HCQ2Compiled):
 
     qname = f"{'COPY' if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA else 'COMPUTE'}:{idx}"
     self.pm_bufferize = PatternMatcher([
-      (UPat(Ops.PARAM, tag=f"{qname}_{name}"), lambda ctx, b=getattr(queue, name): b) for name in ["ring", "write_ptr", "doorbell", "put_value"]
+      (UPat(Ops.PARAM, tag=to_name(name, qname)), lambda ctx, b=getattr(queue, name): b) for name in ["ring", "write_ptr", "doorbell", "put_value"]
     ]) + self.pm_bufferize
 
     return queue
@@ -659,7 +656,7 @@ class AMDDevice(HCQ2Compiled):
     wg_data_size = round_up((vgpr_size_per_cu + sgrp_size_per_cu + lds_size_per_cu + hwreg_size_per_cu) * self.cu_cnt, mmap.PAGESIZE)
     ctl_stack_size = round_up((12 if self.target[0] != 9 else 8) * self.wave_cnt + 8 + 40, mmap.PAGESIZE)
     return self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL if self.is_aql else kfd.KFD_IOC_QUEUE_TYPE_COMPUTE,
-      0x2000 if self.is_usb() else (16 << 20), eop_buffer_size=0x1000,
+      0x2000 if self.is_usb else (16 << 20), eop_buffer_size=0x1000,
       ctx_save_restore_size=0 if self.is_am() else wg_data_size + ctl_stack_size, ctl_stack_size=ctl_stack_size,
       debug_memory_size=round_up(self.wave_cnt * 32, 64))
 
@@ -667,7 +664,7 @@ class AMDDevice(HCQ2Compiled):
     if getenv("AMD_DISABLE_SDMA"): return None
     if idx in self.sdma_queues: return self.sdma_queues[idx]
     with contextlib.suppress(OSError):
-      self.sdma_queues[idx] = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x200 if self.is_usb() else (16 << 20), idx=idx)
+      self.sdma_queues[idx] = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x2000 if self.is_usb else (16 << 20), idx=idx)
     return self.sdma_queues.get(idx, None)
 
   def tmpring_size(self, private_segment_size):

extra/llama_kernels/__init__.py

@@ -1,8 +1,6 @@
 from __future__ import annotations
 import functools, pathlib
-from dataclasses import replace
-from tinygrad import Tensor, dtypes
-from tinygrad.uop.ops import shape_to_shape_arg
+from tinygrad import Tensor
 from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
 
 FP8_MAX = 448.0
@@ -12,13 +10,13 @@ NUM_WG, THREADS_PER_WG = 1024, 256
 @functools.cache
 def _local_abs_max_fxn(x_p, device):
   x = Tensor(x_p, device=device)
-  inner = Tensor(x.uop.replace(src=(shape_to_shape_arg(x.uop.shard_shape),), arg=replace(x.uop.arg, axis=None))) if x.uop.axis is not None else x
+  inner = Tensor(x.uop.src[0]) if x.uop.axis is not None else x  # the per-shard view of the flat param
   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))
+  return Tensor(fxn[0].uop.call_with_output(x.uop))
 
 def shard_shape(shape:tuple, axis:int, ndev:int) -> list:
   s = list(shape)

extra/llama_kernels/quantize_fp8_delayed/__init__.py

@@ -50,7 +50,7 @@ def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_out:UOp, x:UOp, amax_state:
   else: raise NotImplementedError(f"no atomic max for device {device}")
   amax_idx = amax_out.reshape((1,)).index(UOp.const(0))
   max_val = lds[0].load()
-  atomic = UOp(Ops.CUSTOM, dtypes.void, (amax_idx, max_val.bitcast(dtypes.int32), max_val, amax_idx.load()), arg=atomic_arg)
+  atomic = UOp(Ops.CUSTOM, src=(amax_idx, max_val.bitcast(dtypes.int32), max_val, amax_idx.load()), arg=(atomic_arg, dtypes.void))
   return atomic.end(tid, wg).sink(arg=KernelInfo(f"quantize_fp8_with_amax_{n_elems}", opts_to_apply=()))
 
 @functools.cache

extra/llama_kernels/quantize_mxfp4/__init__.py

@@ -12,7 +12,7 @@ def _custom_quantize_mxfp4(row_fp4:UOp, row_scale:UOp, col_fp4:UOp, col_scale:UO
   mem = M*N*2 + M*N + M*N//16  # read bf16, write row+col fp4 + e8m0
   outputs = (row_fp4, row_scale, col_fp4, col_scale)
   sink = UOp.sink(*(o.base for o in outputs), x.base,
-                  *(UOp(Ops.CUSTOM, dtypes.void, (o.base.index(0),), arg="") for o in outputs),
+                  *(UOp(Ops.CUSTOM, src=(o.base.index(0),), arg=("", dtypes.void)) for o in outputs),
                   UOp.special(256, "lidx0"), UOp.special(M//128, "gidx0"), UOp.special(N//64, "gidx1"),
                   arg=KernelInfo(name, estimates=Estimates(ops=12*M*N, mem=mem)))
   src = (pathlib.Path(__file__).parent/"quantize_mxfp4.cpp").read_text()

extra/llama_kernels/rmsnorm/__init__.py

@@ -13,12 +13,13 @@ def _rmsnorm_fwd_fxn(x_in_p, eps, device):
   return rmsnorm_fwd(Tensor(x_in_p, device=device), eps)
 
 def _rmsnorm_bwd(grad:UOp, call:UOp) -> tuple:
-  x_normed = Tensor(call.gettuple(0)).float()
+  outs = call.unbound_outputs
+  x_normed = Tensor(outs[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))
+  d_x = Tensor(outs[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))
+  outs = UOp.call_with_outputs((fxn[0].uop, fxn[1].uop), x_in.uop, grad_fxn=_rmsnorm_bwd)
+  return Tensor(outs[0]), Tensor(outs[1])

extra/mmapeak/mmapeak.py

@@ -3,7 +3,7 @@ import os
 # TODO: there is a timing bug without this
 os.environ["AMD_AQL"] = "1"
 
-from tinygrad import Tensor, Device, GlobalCounters, Context
+from tinygrad import Tensor, Device, GlobalCounters, Context, dtypes
 from tinygrad.helpers import getenv, DEV
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.renderer import Estimates
@@ -37,7 +37,7 @@ def launchBenchmark(instruction, vgprIndices, dense=True, accum=False, **kwargs)
     gidx = UOp.special(NUM_WORKGROUPS, "gidx0")
     FLOPs = FLOPS_PER_MATMUL * NUM_WAVES * NUM_WORKGROUPS * INTERNAL_LOOP * INSTRUCTIONS_PER_LOOP
     sink = UOp.sink(A.base, threads, gidx, arg=KernelInfo(inst.op.name.lower(), estimates=Estimates(ops=FLOPs, mem=0)))
-    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.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts]))))
   dummy = Tensor.zeros(1).contiguous().realize()
   out = Tensor.custom_kernel(dummy, fxn=fxn)[0]
   linear = out.schedule_linear()

extra/models/llama.py

@@ -5,9 +5,9 @@ from tinygrad.helpers import getenv, DEBUG
 
 # https://github.com/facebookresearch/llama/blob/1076b9c51c77ad06e9d7ba8a4c6df775741732bd/llama/model.py#L47
 def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0) -> Tensor:
-  freqs = 1.0 / (theta ** (Tensor.arange(0, dim, 2)[:(dim // 2)] / dim))
-  freqs = Tensor.arange(end).unsqueeze(dim=1) * freqs.unsqueeze(dim=0)
-  return Tensor.stack(freqs.cos(), freqs.sin(), dim=-1).reshape(1, end, 1, dim//2, 2)
+  freqs = 1.0 / (theta ** (Tensor.arange(0, dim, 2, dtype=dtypes.float32)[:(dim // 2)] / dim))
+  freqs = Tensor.arange(end, dtype=dtypes.float32).unsqueeze(dim=1) * freqs.unsqueeze(dim=0)
+  return Tensor.stack(freqs.cos(), freqs.sin(), dim=-1).cast(dtypes.default_float).reshape(1, end, 1, dim//2, 2)
 
 # matches meta, non hugging face weights
 # (a+i*b) * (c+i*d) = (ac-bd) + i*(ad+bc)

extra/onnx_helpers.py

@@ -1,10 +1,12 @@
 from tinygrad import Tensor
+import os
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.nn.onnx import OnnxRunner, OnnxValue
 import numpy as np
 import onnxruntime as ort
 ort_options = ort.SessionOptions()
 ort_options.log_severity_level = 3
+ort_options.intra_op_num_threads = os.cpu_count() or 1
 
 def get_example_inputs(graph_inputs:dict[str, OnnxValue], config={}):
   """

extra/optimization/test_beam_search.py

@@ -89,7 +89,8 @@ class TestBeamSearch(unittest.TestCase):
     s.apply_opt(Opt(OptOps.TC, 0, (-1, 0, 1)))
     up = prod([x for x, t in zip(s.full_shape, s.axis_types) if t in (AxisType.UPCAST, AxisType.UNROLL)])
     actions = get_kernel_actions(s, include_0=False, max_up=int(up))
-    upcasted = [s for s in actions.values() if any(opt.op in (OptOps.UPCAST, OptOps.UNROLL) for opt in s.applied_opts)]
+    upcasted = [s for s in actions.values() if any(o.op is OptOps.SPLIT and o.arg[1] in (AxisType.UPCAST, AxisType.UNROLL)
+                                                  for o in s.applied_opts)]
     assert len(upcasted) > 0, f"expected upcast/unroll actions after TC with max_up={up}, but got none"
 
   def test_max_up(self):
@@ -98,8 +99,8 @@ class TestBeamSearch(unittest.TestCase):
     s = Scheduler(ast, Device[Device.DEFAULT].renderer)
     for max_up in (2, 4):
       actions = get_kernel_actions(s, include_0=False, max_up=max_up)
-      for up_opts in [s.applied_opts for s in actions.values() if any(opt.op in (OptOps.UPCAST, OptOps.UNROLL) for opt in s.applied_opts)]:
-        assert len([opt for opt in up_opts if opt.arg > max_up]) == 0 and len([op for op in up_opts if op.arg <= max_up]) > 0
+      up_opts = [o for s in actions.values() for o in s.applied_opts if o.op is OptOps.SPLIT and o.arg[1] in (AxisType.UPCAST, AxisType.UNROLL)]
+      assert len([opt for opt in up_opts if opt.arg[0] > max_up]) == 0 and len([op for op in up_opts if op.arg[0] <= max_up]) > 0
 
 if __name__ == '__main__':
   unittest.main()

extra/runbook_digitalocean_mi350x.md

@@ -49,13 +49,13 @@ ldconfig
 curl -sL https://raw.githubusercontent.com/geohot/configuration/master/.tmux.conf -o ~/.tmux.conf
 ```
 
-### 1.6 Reload amdgpu driver
-tinygrad's HCQ backend needs `/dev/kfd` which is created by the amdgpu kernel driver.
-If the driver was unloaded, reload it:
+### 1.6 Verify GPU PCI access
+The AM userspace driver accesses the GPUs directly over PCI. Do not load `amdgpu`. `/dev/kfd` is not required.
 ```bash
-modprobe amdgpu
-ls /dev/kfd  # should exist
+rmmod amdgpu
+lspci -nnk -d 1002:
 ```
+The MI350X devices should not show a `Kernel driver in use: amdgpu`.
 
 ## Phase 2: Clone tinygrad
 ```bash
@@ -76,8 +76,23 @@ rclone config create mlc-training s3 provider=Cloudflare \
   endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
 
 mkdir -p /raid/datasets/c4-8b
-rclone copy mlc-training:mlcommons-training-wg-public/llama3_1/datasets/c4/llama3_1_8b/ /raid/datasets/c4-8b/ -P
+(rclone copy mlc-training:mlcommons-training-wg-public/llama3_1/datasets/c4/llama3_1_8b/ /raid/datasets/c4-8b/ -P && \
+  PYTHONPATH=. python3 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/buid_dataset_cache.py) \
+  > /root/dataset_cache.log 2>&1 &
 ```
+Leave this running and proceed to the beam step while the dataset downloads and its cache builds.
+
+### 3.1 Smoke test (beam search, 2 layers, fake data)
+Always run beam first to validate the pipeline:
+```bash
+tmux new-session -d -s beam 'cd /root/tinygrad && COMGR_PATH=/opt/rocm/lib/libamd_comgr.so COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so CC=/opt/rocm/core-7.14/lib/llvm/bin/clang DEV=PCI+AMD:HIP ROCM_PATH=/opt/rocm bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh 2>&1 | tee /root/beam.log'
+```
+
+The beam test runs 10 training steps with 2 layers. Expected results:
+- ~0.29s per step after warmup
+- ~700K GFLOPS, ~7% MFU (low because only 2 layers)
+- ~380 GB VRAM used
+- Loss stable at ~12.55 with random init
 
 Files downloaded (~85GB total, ~6 minutes):
 - `c4-train.en_6_text_document.bin` (79 GB)
@@ -106,25 +121,13 @@ wandb login <API_KEY>
 
 Run training in tmux so it survives SSH disconnects:
 ```bash
-tmux new-session -d -s train 'cd /root/tinygrad && COMGR_PATH=/opt/rocm/lib/libamd_comgr.so COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so CC=/opt/rocm/core-7.14/lib/llvm/bin/clang DEV=AMD:HIP ROCM_PATH=/opt/rocm WANDB=1 bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh 2>&1 | tee /root/train.log'
+tmux new-session -d -s train 'cd /root/tinygrad && COMGR_PATH=/opt/rocm/lib/libamd_comgr.so COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so CC=/opt/rocm/core-7.14/lib/llvm/bin/clang DEV=PCI+AMD:HIP ROCM_PATH=/opt/rocm WANDB=1 bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh 2>&1 | tee /root/train.log'
 ```
 Attach with `tmux attach -t train`.
 
-### 5.1 Smoke test (beam search, 2 layers, real data)
-Always run beam first to validate the pipeline:
+### 5.1 Full training run
 ```bash
-tmux new-session -d -s beam 'cd /root/tinygrad && COMGR_PATH=/opt/rocm/lib/libamd_comgr.so COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so CC=/opt/rocm/core-7.14/lib/llvm/bin/clang DEV=AMD:HIP ROCM_PATH=/opt/rocm bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh 2>&1 | tee /root/beam.log'
-```
-
-The beam test runs 10 training steps with 2 layers. Expected results:
-- ~0.29s per step after warmup
-- ~700K GFLOPS, ~7% MFU (low because only 2 layers)
-- ~380 GB VRAM used
-- Loss stable at ~12.55 with random init
-
-### 5.2 Full training run
-```bash
-tmux new-session -d -s train 'cd /root/tinygrad && COMGR_PATH=/opt/rocm/lib/libamd_comgr.so COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so CC=/opt/rocm/core-7.14/lib/llvm/bin/clang DEV=AMD:HIP ROCM_PATH=/opt/rocm WANDB=1 bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh 2>&1 | tee /root/train.log'
+tmux new-session -d -s train 'cd /root/tinygrad && COMGR_PATH=/opt/rocm/lib/libamd_comgr.so COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so CC=/opt/rocm/core-7.14/lib/llvm/bin/clang DEV=PCI+AMD:HIP ROCM_PATH=/opt/rocm WANDB=1 bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh 2>&1 | tee /root/train.log'
 ```
 
 ## Environment Variable Reference
@@ -134,7 +137,7 @@ tmux new-session -d -s train 'cd /root/tinygrad && COMGR_PATH=/opt/rocm/lib/liba
 | `COMGR_PATH` | `/opt/rocm/lib/libamd_comgr.so` | tinygrad's DLL loader needs explicit path to find comgr 3.3 |
 | `COMGR_3_PATH` | `/opt/rocm/lib/libamd_comgr.so` | comgr 3.x uses a separate `comgr_3` module with its own path var |
 | `CC` | `/opt/rocm/core-7.14/lib/llvm/bin/clang` | System clang doesn't know gfx950; must use ROCm's bundled clang |
-| `DEV` | `AMD:HIP` | Force HIPRenderer (comgr-based) over HIPCCRenderer (hipcc subprocess) |
+| `DEV` | `PCI+AMD:HIP` | Force HIPRenderer (comgr-based) over HIPCCRenderer (hipcc subprocess) |
 | `ROCM_PATH` | `/opt/rocm` | Script defaults to `/opt/rocm-7.1.1` which doesn't exist |
 | `WANDB` | `1` | Enable wandb logging (off by default) |
 
@@ -150,7 +153,7 @@ tmux new-session -d -s train 'cd /root/tinygrad && COMGR_PATH=/opt/rocm/lib/liba
 | ASM GEMM | `extra/gemm/cdna_asm_gemm.py` — gfx950 MFMA assembly, MXFP4 |
 | Flash attention | `extra/thunder/amd/fa.py` |
 | Fused kernels | `extra/llama_kernels/` — rmsnorm, silu, quantize, fused_ce |
-| GPU driver | `tinygrad/runtime/ops_amd.py` — HCQ, direct KFD ioctl |
+| GPU driver | `tinygrad/runtime/ops_amd.py` — HCQ, using the AM userspace PCI interface |
 | Renderer | `tinygrad/renderer/cstyle.py` — HIPRenderer for gfx950 |
 | comgr compiler | `tinygrad/runtime/support/compiler_amd.py` — HIPCompiler using comgr 3.3 |
 
@@ -190,28 +193,5 @@ $ lspci -nn | grep AMD
 ```
 CPU flags include `hypervisor`. `dmesg` shows `Hypervisor detected: KVM`.
 
-### Working path: amdgpu driver (KFDIface)
-The amdgpu driver loads on boot and binds to all 8 GPUs, creating `/dev/kfd` and 64 renderD nodes (`/dev/dri/renderD128` through `/dev/dri/renderD191`). tinygrad's `KFDIface` enumerates GPUs through `/sys/devices/virtual/kfd/kfd/topology/nodes` and uses `/dev/kfd` for ioctl. No PCI device ID patching is needed — the KFD path does not use `PCIIface` or `AMDev._run_discovery()`.
-
-This is the working configuration. No code changes to tinygrad are required.
-
-### PCIIface path (does not work on this VM)
-For reference, the `PCIIface` path was also explored but does not work in this KVM guest:
-
-- `PCIIface` in `ops_amd.py` does not list device ID `0x75b0`. Adding it allows PCI detection but `AMDev._run_discovery()` fails because the VRAM BAR reads all `0xFF`.
-- This was observed with the GPU unbound from any driver, after PCI reset, and with VFIO bound.
-- VFIO binding (`vfio-pci` with `enable_unsafe_noiommu_mode=1`) succeeded but VRAM BAR still reads all `0xFF`.
-- No IOMMU in guest — `dmesg` has no `AMD-Vi` entries, PCI devices have no `iommu_group` symlink.
-
-### amdgpu driver behavior
-On first boot, amdgpu loaded and bound to all 8 GPUs. On one boot it failed to initialize:
-```
-[  799.780369] amdgpu 0000:83:00.0: Failed to alloc msi vectors
-[  799.781476] amdgpu 0000:83:00.0: sw_init of IP block <vega20_ih> failed -22
-[  799.782724] amdgpu 0000:83:00.0: amdgpu_device_ip_init failed
-[  799.793885] amdgpu 0000:83:00.0: Fatal error during GPU init
-```
-On a subsequent boot, amdgpu initialized successfully (SMU initialized, VRAM ready). After unbinding all 8 GPUs from amdgpu, `rmmod amdgpu` wedged the module (stuck in "Unloading" state in `/proc/modules`), requiring a full VM reboot.
-
 ### No fan control
 No `fan*` or `pwm*` hwmon entries exist. Only `temp*`, `power*`, `freq*` are exposed. GPU temps read 56-63°C, power ~265W per GPU.

extra/sqtt/examples/generate_examples.py

@@ -1,4 +1,4 @@
-import os, subprocess, sys, shlex
+import os, subprocess, sys, shlex, pickle
 from pathlib import Path
 from tinygrad.helpers import temp, getenv
 
@@ -23,5 +23,8 @@ if __name__ == "__main__":
       # AM_RESET=1 gets a clear trace, does not work on mi300 machines
       subprocess.run([sys.executable, *shlex.split(test)], cwd=EXAMPLES_DIR.parent.parent.parent,
                      env={**os.environ, "DEV":"AMD", "AM_RESET":"1" if not arch.startswith("gfx9") else "0", "VIZ":"-2", "PYTHONPATH":"."})
+      with open(PROFILE_PATH, "rb") as f: events = pickle.load(f)
+      with open(PROFILE_PATH, "wb") as f:
+        pickle.dump([e for e in events if type(e).__name__ in {"ProfilePMCEvent", "ProfileSQTTEvent", "ProfileProgramEvent"}], f)
       PROFILE_PATH.rename(dest:=EXAMPLES_DIR/arch/f"profile_{name}_run_{i}.pkl")
       print(f"saved SQTT trace to {dest}")

extra/sqtt/roc.py

@@ -1,13 +1,11 @@
 #!/usr/bin/env python3
 import ctypes, pathlib, argparse, pickle, dataclasses, threading, itertools
-from decimal import Decimal
 from typing import Generator
 from tinygrad.helpers import temp, unwrap, DEBUG
 from tinygrad.runtime.ops_amd import ProfileSQTTEvent
 from tinygrad.runtime.autogen import rocprof
 from tinygrad.renderer.amd.dsl import Inst
-from tinygrad.helpers import ProfileEvent, ProfileRangeEvent, ProfilePointEvent
-from tinygrad.device import ProfileProgramEvent
+from tinygrad.device import ProfileDeviceEvent, ProfileProgramEvent
 from test.amd.disasm import disasm
 
 @dataclasses.dataclass(frozen=True)
@@ -39,17 +37,18 @@ class WaveExec(WaveSlot):
     insts_array = (struct*(len(self.insts)//sz)).from_buffer(self.insts)
     for inst in insts_array:
       inst_typ = rocprof.enum_rocprofiler_thread_trace_decoder_inst_category_t.get(inst.category)
-      yield InstExec(inst_typ, inst.pc.address, inst.stall, inst.duration, inst.time)
+      yield InstExec(inst_typ.replace("ROCPROFILER_THREAD_TRACE_DECODER_", "") if inst_typ else "UNKNOWN",
+                     inst.pc.address, inst.stall, inst.duration, inst.time)
 
 @dataclasses.dataclass(frozen=True)
 class OccEvent(WaveSlot):
   time:int
   start:int
 
-RunKey = tuple[str, int]
+RunKey = tuple[int, int]
 
 class _ROCParseCtx:
-  def __init__(self, sqtt_evs:list[ProfileSQTTEvent], disasms:dict[str, dict[int, Inst]]):
+  def __init__(self, sqtt_evs:list[ProfileSQTTEvent], disasms:dict[int, dict[int, Inst]]):
     self.sqtt_evs, self.disasms = iter(sqtt_evs), {k:{k2:(disasm(v2), v2.size()) for k2,v2 in v.items()} for k,v in disasms.items()}
     self.inst_execs:dict[RunKey, list[WaveExec]] = {}
     self.occ_events:dict[RunKey, list[OccEvent]] = {}
@@ -76,7 +75,7 @@ class _ROCParseCtx:
     self.inst_execs.setdefault(unwrap(self.active_run), []).append(WaveExec(ev.wave_id, ev.cu, ev.simd, unwrap(self.active_se), ev.begin_time,
                                                                              ev.end_time, insts_blob))
 
-def decode(sqtt_evs:list[ProfileSQTTEvent], disasms:dict[str, dict[int, Inst]]) -> _ROCParseCtx:
+def decode(sqtt_evs:list[ProfileSQTTEvent], disasms:dict[int, dict[int, Inst]]) -> _ROCParseCtx:
   ROCParseCtx = _ROCParseCtx(sqtt_evs, disasms)
 
   @rocprof.rocprof_trace_decoder_se_data_callback_t
@@ -129,83 +128,9 @@ def decode(sqtt_evs:list[ProfileSQTTEvent], disasms:dict[str, dict[int, Inst]])
     raise exc
   return ROCParseCtx
 
-def unpack_occ(viz_data, i:int, j:int, key:tuple[str, int], data:list, p:ProfileProgramEvent, target:str) -> dict:
-  from tinygrad.viz.serve import amd_decode, create_step, row_tuple
-  steps = viz_data.ctxs[i]["steps"]
-  if len(steps[j+1:]) > 0: return {"steps":[{k:v for k,v in s.items() if k != "data"} for s in steps[j+1:]]}
-  base = unwrap(p.base)
-  disasm:dict[int, Inst] = {addr+base:inst for addr,inst in amd_decode(unwrap(p.lib), target).items()}
-  rctx = decode(data, {p.tag:disasm})
-  cu_events:dict[str, list[ProfileEvent]] = {}
-  # ** inst traces
-  wave_insts:dict[str, dict[str, dict]] = {}
-  inst_units:dict[str, itertools.count] = {}
-  for w in rctx.inst_execs.get(key, []):
-    if (u:=w.wave_loc) not in inst_units: inst_units[u] = itertools.count(0)
-    n = next(inst_units[u])
-    if (events:=cu_events.get(w.cu_loc)) is None: cu_events[w.cu_loc] = events = []
-    events.append(ProfileRangeEvent(f"SIMD:{w.simd}", loc:=f"INST WAVE:{w.wave_id} N:{n}", Decimal(w.begin_time), Decimal(w.end_time)))
-    wave_insts.setdefault(w.cu_loc, {})[f"{u} N:{n}"] = {"wave":w, "disasm":disasm, "prg":p, "run_number":n, "loc":loc}
-  # ** occ traces (only WAVESTART/WAVEEND)
-  units:dict[str, itertools.count] = {}
-  wave_start:dict[str, int] = {}
-  for occ in rctx.occ_events.get(key, []):
-    if (u:=occ.wave_loc) not in units: units[u] = itertools.count(0)
-    if u in inst_units: continue
-    if occ.start: wave_start[u] = occ.time
-    else:
-      if (events:=cu_events.get(occ.cu_loc)) is None: cu_events[occ.cu_loc] = events = []
-      events.append(ProfileRangeEvent(f"SIMD:{occ.simd}", f"OCC WAVE:{occ.wave_id} N:{next(units[u])}", Decimal(wave_start.pop(u)),Decimal(occ.time)))
-  # ** split graph by CU
-  for cu in sorted(cu_events, key=row_tuple):
-    steps.append(create_step(f"{cu} {len(cu_events[cu])}", ("/cu-sqtt", i, len(steps)), depth=1,
-                             data=[ProfilePointEvent(unit, "start", unit, ts=Decimal(0)) for unit in units]+cu_events[cu]))
-    for k in sorted(wave_insts.get(cu, []), key=row_tuple):
-      wd = wave_insts[cu][k]
-      steps.append(create_step(k.replace(cu, ""), ("/amd-sqtt-insts", i, len(steps)), loc=wd["loc"], depth=2,
-                               data={"fxn":unpack_insts, "args":(wd,)}))
-  return {"steps":[{k:v for k,v in s.items() if k != "data"} for s in steps[j+1:]]}
-
-def unpack_insts(viz_data, i:int, j:int, data:dict) -> dict:
-  columns = ["PC", "Instruction", "Hits", "Cycles", "Stall", "Type"]
-  inst_columns = ["N", "Clk", "Idle", "Dur", "Stall"]
-  # Idle:     The total time gap between the completion of previous instruction and the beginning of the current instruction.
-  #           The idle time can be caused by:
-  #             * Arbiter loss
-  #             * Source or destination register dependency
-  #             * Instruction cache miss
-  # Stall:    The total number of cycles the hardware pipe couldn't issue an instruction.
-  # Duration: Total latency in cycles, defined as "Stall time + Issue time" for gfx9 or "Stall time + Execute time" for gfx10+.
-  prev_instr = (w:=data["wave"]).begin_time
-  pc_to_inst = data["disasm"]
-  start_pc = None
-  rows:dict[int, dict] = {}
-  for pc, inst in pc_to_inst.items():
-    if start_pc is None: start_pc = pc
-    rows[pc] = {"pc":pc-start_pc, "inst":str(inst), "hit_count":0, "dur":0, "stall":0, "type":"", "hits":{"cols":inst_columns, "rows":[]}}
-  for e in w.unpack_insts():
-    if not (inst:=rows[e.pc]).get("type"): inst["type"] = str(e.typ).split("_")[-1]
-    inst["hit_count"] += 1
-    inst["dur"] += e.dur
-    inst["stall"] += e.stall
-    inst["hits"]["rows"].append((inst["hit_count"]-1, e.time, max(0, e.time-prev_instr), e.dur, e.stall))
-    prev_instr = max(prev_instr, e.time + e.dur)
-  summary = [{"label":"Total Cycles", "value":w.end_time-w.begin_time}, {"label":"SE", "value":w.se}, {"label":"CU", "value":w.cu},
-             {"label":"SIMD", "value":w.simd}, {"label":"Wave ID", "value":w.wave_id}, {"label":"Run number", "value":data["run_number"]}]
-  return {"rows":[tuple(v.values()) for v in rows.values()], "cols":columns, "metadata":[summary], "ref":viz_data.ref_map.get(data["prg"].name)}
-
-def print_data(data:dict) -> None:
-  from tabulate import tabulate
-  # plaintext
-  if "src" in data: print(data["src"])
-  # table format
-  elif "cols" in data:
-    print(tabulate([r[:len(data["cols"])] for r in data["rows"]], headers=data["cols"], tablefmt="github"))
-
 def main() -> None:
-  import tinygrad.viz.serve as viz
-  from tinygrad.uop.ops import RewriteTrace
-  data = viz.VizData()
+  from tabulate import tabulate
+  from tinygrad.viz.serve import amd_decode
 
   parser = argparse.ArgumentParser()
   parser.add_argument('--profile', type=pathlib.Path, metavar="PATH", help='Path to profile (optional file, default: latest profile)',
@@ -216,26 +141,28 @@ def main() -> None:
 
   with args.profile.open("rb") as f: profile = pickle.load(f)
 
-  viz.get_profile(profile, data=data)
-
   # List all kernels
   if args.kernel is None:
-    for c in data.ctxs:
-      print(c["name"])
-      for s in c["steps"]: print("  "+s["name"])
+    for p in profile:
+      if isinstance(p, ProfileProgramEvent) and p.device.startswith("AMD"): print(p.name)
     return None
 
-  # Find kernel trace
-  trace = next((c for c in data.ctxs if c["name"] == f"SQTT {args.kernel}"), None)
-  if not trace: raise RuntimeError(f"no matching trace for {args.kernel}")
-  n = 0
-  for s in trace["steps"]:
-    if "PKTS" in s["name"]: continue
-    print(s["name"])
-    ret = viz.get_render(data, s["query"])
-    print_data(ret)
-    n += 1
-    if n > args.n: break
+  prg = next((p for p in profile if isinstance(p, ProfileProgramEvent) and p.name == args.kernel), None)
+  dev = next((p for p in profile if isinstance(p, ProfileDeviceEvent) and p.device == prg.device), None)
+  assert prg is not None and dev is not None, "must have program binary and device props"
+  target = f"gfx{dev.props['gfx_target_version']//1000}"
+  sqtt = [p for p in profile if isinstance(p, ProfileSQTTEvent) and p.kern == prg.tag]
+
+  pc_to_inst = {addr+prg.base:inst for addr,inst in amd_decode(prg.lib, target).items()}
+  rctx = decode(sqtt, {prg.tag:pc_to_inst})
+  waves = sorted(itertools.chain.from_iterable(rctx.inst_execs.values()), key=lambda w:(w.se, w.cu, w.simd, w.wave_id, w.begin_time))
+  if not waves: raise RuntimeError(f"no instruction traces for {args.kernel}")
+  run_numbers:dict[str, itertools.count] = {}
+  for w in itertools.islice(waves, args.n):
+    if w.wave_loc not in run_numbers: run_numbers[w.wave_loc] = itertools.count()
+    print(f"{w.wave_loc} N:{next(run_numbers[w.wave_loc])} Total Cycles:{w.end_time-w.begin_time}")
+    rows = [(e.time, f"0x{e.pc:x}", pc_to_inst[e.pc], e.typ, e.dur, e.stall) for e in w.unpack_insts()]
+    print(tabulate(rows, headers=("Timestamp", "PC", "Instruction", "Type", "Duration", "Stall"), tablefmt="github"))
 
 if __name__ == "__main__":
   main()

extra/thunder/amd/moe_router_mfma.cpp

@@ -0,0 +1,80 @@
+#include "kittens.cuh"
+
+using namespace kittens;
+
+#ifndef ROUTER_M
+#define ROUTER_M 16384
+#endif
+#ifndef ROUTER_K
+#define ROUTER_K 2880
+#endif
+#ifndef ROUTER_E
+#define ROUTER_E 32
+#endif
+
+constexpr int BLOCK_M = 64;
+constexpr int BLOCK_K = 64;
+constexpr int NUM_WARPS = 4;
+constexpr int THREADS = NUM_WARPS * WARP_THREADS;
+
+using G = kittens::group<NUM_WARPS>;
+using XST = st_bf<BLOCK_M, BLOCK_K, st_16x32_s>;
+using WST = st_bf<ROUTER_E, BLOCK_K, st_16x32_s>;
+using XRT = rt_bf<16, BLOCK_K, row_l, rt_16x32_s>;
+using WRT = rt_bf<ROUTER_E, BLOCK_K, row_l, rt_16x32_s>;
+using CRT = rt_fl<16, ROUTER_E, col_l, rt_16x16_s>;
+
+static_assert(ROUTER_M % BLOCK_M == 0, "ROUTER_M must be divisible by 64");
+static_assert(ROUTER_K % BLOCK_K == 0, "ROUTER_K must be divisible by 64");
+static_assert(ROUTER_E == 32, "the small-N tile is specialized for 32 experts");
+
+extern "C" __global__ __launch_bounds__(THREADS, 4) void moe_router_mfma(
+    float *__restrict__ out, bf16 *__restrict__ x_ptr, bf16 *__restrict__ weight_ptr,
+    bf16 *__restrict__ bias) {
+  gl<bf16, 1, 1, ROUTER_M, ROUTER_K> X{x_ptr, nullptr, nullptr, nullptr, nullptr};
+  gl<bf16, 1, 1, ROUTER_E, ROUTER_K> W{weight_ptr, nullptr, nullptr, nullptr, nullptr};
+
+  __shared__ XST Xs;
+  __shared__ WST Ws;
+
+  XRT xr;
+  WRT wr;
+  CRT accum;
+  zero(accum);
+
+  const int block_m = __builtin_amdgcn_workgroup_id_x();
+  const int warp_m = warpid();
+
+  #pragma unroll
+  for (int kk = 0; kk < ROUTER_K / BLOCK_K; kk++) {
+    G::load(Xs, X, {0, 0, block_m, kk});
+    G::load(Ws, W, {0, 0, 0, kk});
+    asm volatile("s_waitcnt vmcnt(0)");
+    asm volatile("s_waitcnt lgkmcnt(0)");
+    __builtin_amdgcn_s_barrier();
+
+    load(xr, subtile_inplace<16, BLOCK_K>(Xs, {warp_m, 0}));
+    load(wr, subtile_inplace<ROUTER_E, BLOCK_K>(Ws, {0, 0}));
+    asm volatile("s_waitcnt lgkmcnt(0)");
+    __builtin_amdgcn_s_setprio(1);
+    mma_ABt(accum, xr, wr, accum);
+    __builtin_amdgcn_s_setprio(0);
+    __builtin_amdgcn_sched_barrier(0);
+    __builtin_amdgcn_s_barrier();
+  }
+
+  // A 16x16 MFMA accumulator is column-layout: each lane owns four consecutive rows
+  // at one column. Store all 64x32 FP32 results directly; no padded or undersized output ABI.
+  const int lane = laneid();
+  const int row0 = block_m * BLOCK_M + warp_m * 16 + 4 * (lane / 16);
+  const int lane_col = lane % 16;
+  #pragma unroll
+  for (int j = 0; j < ROUTER_E / 16; j++) {
+    const int col = j * 16 + lane_col;
+    const float b = (float)bias[col];
+    const float vals[4] = {accum.tiles[0][j].data[0].x, accum.tiles[0][j].data[0].y,
+                           accum.tiles[0][j].data[1].x, accum.tiles[0][j].data[1].y};
+    #pragma unroll
+    for (int r = 0; r < 4; r++) out[(long long)(row0 + r) * ROUTER_E + col] = vals[r] + b;
+  }
+}

extra/tinyfs/fetch_file.py

@@ -1,38 +0,0 @@
-from tinygrad.tensor import Tensor
-from tinygrad.helpers import CHUNK_SIZE
-from tinygrad.nn.state import fs_load
-import argparse, math, hashlib
-
-def _python_hash_1mb(data:bytes|bytearray):
-  chunks = [data[i:i+4096] for i in range(0, len(data), 4096)]
-  chunk_hashes = [hashlib.shake_128(chunk).digest(16) for chunk in chunks]
-  return hashlib.shake_128(b''.join(chunk_hashes)).digest(16)
-
-def hash_file(data: bytes|bytearray):
-  if len(data) % CHUNK_SIZE != 0: data += bytes(CHUNK_SIZE - len(data) % CHUNK_SIZE)
-  base_chunks = math.ceil(len(data) / CHUNK_SIZE)
-  tree_depth = math.ceil(math.log(base_chunks, CHUNK_SIZE // 16))
-
-  for _ in range(tree_depth + 1):
-    data_chunks = [data[i:i+CHUNK_SIZE] for i in range(0, len(data), CHUNK_SIZE)]
-    data_chunk_hashes = [_python_hash_1mb(chunk) for chunk in data_chunks]
-    data = b''.join(data_chunk_hashes)
-    if len(data) % CHUNK_SIZE != 0: data += bytes(CHUNK_SIZE - len(data) % CHUNK_SIZE)
-
-  return data[:16]
-
-if __name__ == "__main__":
-  parser = argparse.ArgumentParser()
-  parser.add_argument("--hash", type=str, required=True, help="file hash to fetch")
-  parser.add_argument("--len", type=int, required=True, help="file length to fetch")
-  parser.add_argument("--dest", type=str, required=True, help="destination path to save the file")
-  parser.add_argument("--check", action="store_true", help="verify the file hash after fetching")
-  args = parser.parse_args()
-
-  fs_load(Tensor(bytes.fromhex(args.hash), device="CPU"), args.len).to(f"disk:{args.dest}").realize()
-
-  if args.check:
-    with open(args.dest, "rb") as f:
-      data = f.read()
-      assert hash_file(data) == bytes.fromhex(args.hash), "Hash mismatch after fetching file"
-    print("File hash verified successfully!")

extra/tinyfs/fetch_raid.py

@@ -1,42 +0,0 @@
-import json, multiprocessing, functools
-from pathlib import Path
-
-from tinygrad.tensor import Tensor
-from tinygrad.helpers import tqdm, getenv
-from tinygrad.nn.state import fs_load
-
-raid_root = Path(getenv("RAID_ROOT", "/raid"))
-
-def fetch_file(item):
-  path, info = item
-  h, size = info["hash"], info["size"]
-
-  path = raid_root / Path(path)
-  path.parent.mkdir(parents=True, exist_ok=True)
-
-  try:
-    pt = fs_load(Tensor(bytes.fromhex(h), device="CPU"), size).to(f"disk:{path.as_posix()}").realize()
-  except Exception as e:
-    print(f"error fetching {path}, {h}, {size}: {e}")
-    raise
-
-  pt.uop.buffer.deallocate()
-
-def fetch_mapping(h, l):
-  mapping_tensor = fs_load(Tensor(bytes.fromhex(h)), l).realize()
-  mapping = mapping_tensor.data().tobytes().decode()
-  mapping = json.loads(mapping)
-  mapped_files = mapping.items()
-  return list(mapped_files)
-
-if __name__ == "__main__":
-  h, l = getenv("HASH", "d734f5e3be9f1e9d863bfaa4fc6c1ef2"), getenv("LENGTH", 175866113)
-
-  with multiprocessing.Pool(processes=1) as pool:
-    mapped_files = pool.apply(functools.partial(fetch_mapping, h, l))
-
-  print(f"fetched mapping for {len(mapped_files)} files")
-
-  with multiprocessing.Pool(processes=multiprocessing.cpu_count()) as pool:
-    for _ in tqdm(pool.imap_unordered(fetch_file, mapped_files), total=len(mapped_files)):
-      pass

extra/tinyfs/upload_raid.py

@@ -1,32 +0,0 @@
-from pathlib import Path
-import multiprocessing, json
-
-from tinygrad.tensor import Tensor
-from tinygrad.helpers import tqdm
-from tinygrad.nn.state import fs_store
-
-raid_root = Path("/raid")
-
-def upload_file(path: Path):
-  pt = Tensor(path).realize()
-  h = fs_store(pt).realize()
-  pt.uop.realized.deallocate()
-  return h.data().hex(), path, pt.nbytes()
-
-if __name__ == "__main__":
-  raid_files = sorted([p for p in raid_root.rglob("*") if p.is_file()])
-  print(f"found {len(raid_files)} files in /raid")
-
-  mapping = {}
-  with multiprocessing.Pool(processes=multiprocessing.cpu_count()) as pool:
-    for h, p, s in tqdm(pool.imap_unordered(upload_file, raid_files), total=len(raid_files)):
-      mapping[p.relative_to(raid_root).as_posix()] = {"hash": h, "size": s}
-
-  # sort the mapping by key
-  mapping = dict(sorted(mapping.items()))
-
-  mapping = json.dumps(mapping).encode()
-  mapping_tensor = Tensor(mapping, device="CPU")
-  h = fs_store(mapping_tensor).realize()
-
-  print(f"final hash: {h.data().hex()}, size: {len(mapping)}")

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "tinygrad"
-version = "0.13.0"
+version = "0.14.0"
 description = "You like pytorch? You like micrograd? You love tinygrad! <3"
 authors = [{ name = "George Hotz" }]
 
@@ -84,7 +84,7 @@ testing = [
   "pillow",
   "onnx==1.19.0",
   "onnx2torch",
-  "onnxruntime",
+  "onnxruntime==1.24.1",
   "opencv-python",
   "transformers",
   "sentencepiece",

spec/tinyspec.tex

@@ -23,7 +23,6 @@
 \definecolor{axblue}{HTML}{1565C0}       % GLOBAL
 \definecolor{axcyan}{HTML}{00838F}       % LOCAL
 \definecolor{axbrcyan}{HTML}{00ACC1}     % WARP
-\definecolor{axbrblue}{HTML}{42A5F5}     % THREAD
 \definecolor{axwhite}{HTML}{616161}      % LOOP (gray on white paper)
 \definecolor{axred}{HTML}{C62828}        % REDUCE
 \definecolor{axbrred}{HTML}{E53935}      % GROUP_REDUCE
@@ -50,10 +49,10 @@ All nodes in the tinygrad graph are \textbf{UOps}. A UOp is a tuple $(\mathrm{op
 \toprule
 \textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
 \midrule
-\op{Param} & $(\mathbf{s})$ & slot, dtype, device?, addrspace? &
-  Placeholder with shape $\mathbf{s}$. Substituted in \op{Function}. \\[4pt]
-\op{Buffer} & $(\mathbf{s})$ & slot, dtype, device, addrspace &
-  Concrete buffer slot with shape $\mathbf{s}$. If device is a tuple, it creates the fully sized buffer across multiple devices. \\
+\op{Param} & () & \texttt{ParamArg} &
+  Placeholder with flat storage of $\mathrm{size}$ elements. Substituted in \op{Call}. \\[4pt]
+\op{Buffer} & () & \texttt{ParamArg} &
+  Flat storage of $\mathrm{size}$ elements. \textbf{Unbound} if not allocated yet. \\
 \op{Const} & () & value, dtype &
   A scalar constant with shape $(\ )$. \\
   & & & Form vector consts with \op{Stack} \\
@@ -62,7 +61,21 @@ All nodes in the tinygrad graph are \textbf{UOps}. A UOp is a tuple $(\mathrm{op
 \end{tabular}
 
 \smallskip
-\textbf{addrspace} is \texttt{GLOBAL}, \texttt{LOCAL}, or \texttt{REG}.
+\texttt{ParamArg} contains slot, dtype, concrete size (or \textsc{null} for a scalar), value bounds, alignment, name, addrspace, device, volatility, optional image shape, and an optional bound device buffer (absent for unbound \op{Buffer}s). \textbf{addrspace} is \texttt{GLOBAL}, \texttt{LOCAL}, or \texttt{REG}.
+
+%% ============================================================
+\subsection*{{\color{callblue}Call Ops} \normalfont\small--- function abstraction, like the lambda calculus}
+
+\begin{tabular}{@{}l l l l@{}}
+\toprule
+\textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
+\midrule
+\op{Call} & (body, $a_0$, $a_1$, \ldots) & --- & Substitute each \op{Param} $k$ in body with $a_k$. \\
+\bottomrule
+\end{tabular}
+
+\smallskip
+A value \op{Call} is void: its \op{Sink} body stores to output \op{Param}s bound positionally to the call's unbound \op{Buffer} arguments; output $a_k$ is \op{After}$(a_k, \op{Call})$. Unbound \op{Buffer}s are scoped to their \op{Call}: they are never implicit inputs of the enclosing graph.
 
 %% ============================================================
 \subsection*{{\color{movgreen}Movement Ops} \normalfont\small--- no arithmetic; view, indexing, and reinterpretation only}
@@ -95,20 +108,6 @@ All nodes in the tinygrad graph are \textbf{UOps}. A UOp is a tuple $(\mathrm{op
 \bottomrule
 \end{tabular}
 
-%% ============================================================
-\subsection*{{\color{callblue}Call Ops} \normalfont\small--- function abstraction, like the lambda calculus}
-
-\begin{tabular}{@{}l l l l@{}}
-\toprule
-\textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
-\midrule
-\op{Function} & (body, $a_0$, $a_1$, \ldots) & --- & Substitute each \op{Param} $k$ in \op{Tuple} body with $a_k$. Gradient-able. \\
-\op{Call} & (body, $a_0$, $a_1$, \ldots) & --- & Opaque invocation of a compiled kernel or custom function. \\
-\op{Tuple} & $(v_0, v_1, \ldots)$ & --- & Pack values; required as \op{Function} body to return a value. \\
-\op{GetTuple} & $(T,)$ & idx & Extract element at idx from a \op{Tuple}. \\
-\bottomrule
-\end{tabular}
-
 %% ============================================================
 \subsection*{{\color{loadred}Load Ops} \normalfont\small--- can change device or addrspace}
 
@@ -256,9 +255,9 @@ Every UOp has a \textbf{dtype}, \textbf{shape}, \textbf{device}, \textbf{addrspa
 \toprule
 \textbf{Op} & \textbf{dtype} & \textbf{shape} & \textbf{device} & \textbf{min\_max} \\
 \midrule
-\op{Buffer} & from arg & from $\mathrm{src}[0]$ & from arg & dtype range \\
+\op{Buffer} & from arg & from arg ($\mathrm{size}$) & from arg & dtype range \\
 \op{Const} & from arg & $()$ & \textsc{null} & $[v, v]$ \\
-\op{Param} & from arg & from $\mathrm{src}[0]$ & from arg & from src or dtype range \\[3pt]
+\op{Param} & from arg & from arg ($\mathrm{size}$) & from arg & from src or dtype range \\[3pt]
 Movement ops & $\mathrm{src}[0].\mathrm{dtype}$ & (see op) & $\mathrm{src}[0].\mathrm{device}$ & $\mathrm{src}[0]$ \\
 \op{Unshard} & $\mathrm{src}[0].\mathrm{dtype}$ & $\mathrm{src}[0]$, each $a_k \times n_k$ & $\mathrm{src}[0].\mathrm{device}$ & $\mathrm{src}[0]$ \\
 \op{Reduce} & $\mathrm{src}[0].\mathrm{dtype}$ & remove first $n$ axes & $\mathrm{src}[0].\mathrm{device}$ & dtype range \\[3pt]
@@ -272,7 +271,7 @@ ALU unary & $\mathrm{src}[0].\mathrm{dtype}$ & $\mathrm{src}[0].\mathrm{shape}$
 Other binary & $\mathrm{src}[0].\mathrm{dtype}$ & broadcast & $\mathrm{src}[0].\mathrm{device}$ & dtype range \\
 \op{CmpLt}, \op{CmpNe} & bool & broadcast & $\mathrm{src}[0].\mathrm{device}$ & from intervals \\
 \op{Where} & $\mathrm{src}[1].\mathrm{dtype}$ & broadcast & $\mathrm{src}[0].\mathrm{device}$ & $[\min(b,c),\, \max(B,C)]$ \\[3pt]
-\op{Function}, \op{Call} & $\mathrm{src}[0].\mathrm{dtype}$ & substitute \op{Param} shapes & $\mathrm{src}[1].\mathrm{device}$ & dtype range \\
+\op{Call} & void & --- & first non-null src device & --- \\
 \op{Range} & index & $()$ & \textsc{null} & $[0,\, n{-}1]$ \\
 \op{Index} & $\mathrm{src}[0].\mathrm{dtype}$ & remaining dims & $\mathrm{src}[0].\mathrm{device}$ & $\mathrm{src}[0]$ \\
 \op{Store} & void & $()$ & $\mathrm{src}[0].\mathrm{device}$ & --- \\
@@ -304,7 +303,6 @@ Each kernel's iteration space is a set of \op{Range} axes. Every range has an \t
 {\color{axblue}\texttt{GLOBAL}} & \texttt{g} & --- & --- & GPU global workgroup dimension. \\
 {\color{axcyan}\texttt{LOCAL}} & \texttt{l} & g, L & inner & Workgroup local dimension (shared memory). \\
 {\color{axbrcyan}\texttt{WARP}} & \texttt{w} & \multicolumn{2}{l}{(created by \op{TC})} & Warp-level lanes for tensor cores. \\
-{\color{axbrblue}\texttt{THREAD}} & \texttt{t} & g & outer & CPU thread parallelism. \\
 {\color{axwhite}\texttt{LOOP}} & \texttt{L} & --- & --- & Generic sequential loop (initial state). \\
 {\color{axred}\texttt{REDUCE}} & \texttt{R} & --- & --- & Reduction axis. \\
 {\color{axbrred}\texttt{GROUP\_REDUCE}} & \texttt{G} & R & inner/outer & Shared-memory group reduction. \\
@@ -327,8 +325,6 @@ An optimization is a triple $(\mathrm{op},\;\mathrm{axis},\;\mathrm{arg})$:
   Pad axis to next multiple of $m$ with validity masks. \\[4pt]
 \op{Swap} & axis$_i$ & axis$_j$ &
   Swap two axes $i \leftrightarrow j$. \\
-\op{Nolocals} & --- & --- &
-  Disable local memory; no workgroup dims emitted. \\
 \op{TC} & reduce idx & (tc, opt, mode) &
   Apply tensor core \op{Wmma}: split reduce/output axes into \texttt{WARP}, \texttt{UPCAST}, and \texttt{UNROLL} dims. \\
 \bottomrule
@@ -421,7 +417,7 @@ def allreduce(T):
 %% ============================================================
 \subsection*{{\color{callblue}The \texttt{@function} Decorator} \normalfont\small--- graph capture via tracing}
 
-The \texttt{@function} decorator transforms a Python function on Tensors into a single \op{Function} node.
+The \texttt{@function} decorator transforms a Python function on Tensors into a single \op{Call} node.
 
 \begin{lstlisting}
 @function
@@ -433,14 +429,15 @@ When \texttt{f(x, y)} is called, the decorator:
 
 \begin{enumerate}[leftmargin=1.5em, itemsep=2pt]
   \item \textbf{Extracts inputs}: walks all arguments to find every Tensor, deduplicates by identity.
-  \item \textbf{Runs the function} lazily (no device execution), building a UOp graph from the result.
-  \item \textbf{Parameterizes}: replaces each input UOp with a \op{Param}$(k)$ placeholder.
-  \item \textbf{Wraps the body} in a \op{Tuple} (even for single returns) and creates\\
-        \op{Function}(\op{Tuple}(body), $x$, $y$).
-  \item \textbf{Returns} the result via \op{GetTuple}$(0)$, or one \op{GetTuple} per element for tuple returns.
+  \item \textbf{Runs the function} lazily (no device execution), building a UOp graph from each returned value.
+  \item \textbf{Parameterizes inputs}: replaces each input UOp with a positional \op{Param}$(k)$ placeholder.
+  \item \textbf{Parameterizes outputs}: for each returned value $v_i$, creates an output \op{Param}$(m+i)$ and a matching unbound \op{Buffer} $b_i$ (unique identity), where $m$ is the number of inputs.
+  \item \textbf{Builds the call}: stores every $v_i$ into its output parameter and creates\\
+        \op{Call}(\op{Sink}(\op{Store}(\op{Param}$(m)$, $v_0$), \ldots), $x$, $y$, $b_0$, \ldots).
+  \item \textbf{Returns values}: exposes each result as \op{After}($b_i$, \op{Call}).
 \end{enumerate}
 
-The result is a reusable graph fragment: the body contains only \op{Param} references, not concrete buffers. At schedule time, the \op{Function} is resolved by substituting each \op{Param}$(k)$ back with its corresponding argument $a_k$, or lowered into an opaque \op{Call} if it is to be compiled as a reusable kernel.
+The result is a reusable graph fragment: the body contains only \op{Param} references, not concrete buffers, and a single call can return any number of values. At schedule time, an ordinary value-producing \op{Call} is inlined by positional \op{Param} substitution and each output \op{After} resolves to the value stored in the body. A precompiled call instead materializes real output buffers in place of the unbound \op{Buffer}s and lowers the body to an opaque call that writes them.
 
 %% ============================================================
 \subsection*{Lowering Pipeline \normalfont\small--- from Tensor graph to machine code}

test/amd/hw/test_sop.py

@@ -1002,6 +1002,39 @@ class TestBarrier(unittest.TestCase):
     for tid in range(64):
       self.assertEqual(st.vgpr[tid][0], tid + 100 + 1000, f"tid={tid}")
 
+class TestSMaxMinSCCRegressions(unittest.TestCase):
+  """Regression test: S_MAX sets SCC only on strict inequality (equal operands -> SCC=0)."""
+
+  def test_s_max_i32_equal_scc(self):
+    st = run_program([s_mov_b32(s[4], 64), s_mov_b32(s[5], 64), s_max_i32(s[6], s[4], s[5])], n_lanes=1)
+    self.assertEqual(st.scc, 0)
+    self.assertEqual(st.sgpr[6], 64)
+    st = run_program([s_mov_b32(s[4], 65), s_mov_b32(s[5], 64), s_max_i32(s[6], s[4], s[5])], n_lanes=1)
+    self.assertEqual(st.scc, 1)  # still set when strictly greater
+
+  def test_s_max_u32_equal_scc(self):
+    st = run_program([s_mov_b32(s[4], 64), s_mov_b32(s[5], 64), s_max_u32(s[6], s[4], s[5])], n_lanes=1)
+    self.assertEqual(st.scc, 0)
+
+class TestAbsdiffOverflowRegressions(unittest.TestCase):
+  """Regression test: S_ABSDIFF_I32 computes abs on the WRAPPED 32-bit difference (found by random difftest vs hardware)."""
+
+  def test_s_absdiff_wrapped(self):
+    # |45 - (-2147483647)| overflows int32; hardware takes abs of the wrapped 32-bit difference
+    instructions = [s_mov_b32(s[4], 45), s_mov_b32(s[5], 0x80000001), s_absdiff_i32(s[6], s[4], s[5])]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.sgpr[6], 0x7FFFFFD4)
+    self.assertEqual(st.scc, 1)
+    # INT_MIN - 1 wraps to +2147483647, already positive
+    instructions = [s_mov_b32(s[4], 0x80000000), s_mov_b32(s[5], 1), s_absdiff_i32(s[6], s[4], s[5])]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.sgpr[6], 0x7FFFFFFF)
+    # equality -> 0 and SCC=0
+    instructions = [s_mov_b32(s[4], 7), s_mov_b32(s[5], 7), s_absdiff_i32(s[6], s[4], s[5])]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.sgpr[6], 0)
+    self.assertEqual(st.scc, 0)
+
 
 if __name__ == '__main__':
   unittest.main()

test/amd/hw/test_vop1.py

@@ -1629,5 +1629,66 @@ class TestSwap(unittest.TestCase):
     self.assertEqual(st.vgpr[0][1], 0x55555555)
 
 
+class TestCvtFrexpRegressions(unittest.TestCase):
+  """Regression tests for float<->int conversion and FREXP corner cases (found by random difftest vs hardware)."""
+
+  def test_cvt_i32_f32_nan_is_zero(self):
+    """v_cvt_i32_f32 of NaN is 0, not INT_MIN (x86 cvttss2si returns INT_MIN)."""
+    for nan in (0x7FC00000, 0xFFC00000, 0x7F800001):
+      st = run_program([v_mov_b32_e32(v[0], nan), v_cvt_i32_f32_e32(v[1], v[0])], n_lanes=1)
+      self.assertEqual(st.vgpr[0][1], 0, f"nan=0x{nan:08x}")
+
+  def test_cvt_i32_f32_positive_overflow(self):
+    """v_cvt_i32_f32 saturates positive overflow/inf to INT_MAX, not INT_MIN."""
+    for bits in (0x7F800000, 0x4F000000, 0x4F800000):  # +inf, 2^31, ~2^32
+      st = run_program([v_mov_b32_e32(v[0], bits), v_cvt_i32_f32_e32(v[1], v[0])], n_lanes=1)
+      self.assertEqual(st.vgpr[0][1], 0x7FFFFFFF, f"bits=0x{bits:08x}")
+
+  def test_cvt_i32_f32_negative_overflow(self):
+    """v_cvt_i32_f32 saturates negative overflow/-inf to INT_MIN."""
+    for bits in (0xFF800000, 0xCF000001):  # -inf, below -2^31
+      st = run_program([v_mov_b32_e32(v[0], bits), v_cvt_i32_f32_e32(v[1], v[0])], n_lanes=1)
+      self.assertEqual(st.vgpr[0][1], 0x80000000, f"bits=0x{bits:08x}")
+
+  def test_cvt_u32_f32_nan_is_zero(self):
+    """v_cvt_u32_f32 of NaN is 0, not UINT_MAX."""
+    for nan in (0x7FC00000, 0xFFC00000, 0x7F800001):
+      st = run_program([v_mov_b32_e32(v[0], nan), v_cvt_u32_f32_e32(v[1], v[0])], n_lanes=1)
+      self.assertEqual(st.vgpr[0][1], 0, f"nan=0x{nan:08x}")
+
+  def test_cvt_i32_f64_nan_and_overflow(self):
+    """v_cvt_i32_f64: NaN -> 0, positive overflow/+inf -> INT_MAX."""
+    st = run_program([v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0x7FF80000), v_cvt_i32_f64_e32(v[2], v[0:1])], n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0)
+    st = run_program([v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0x41F00000), v_cvt_i32_f64_e32(v[2], v[0:1])], n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0x7FFFFFFF)  # 2^32 -> INT_MAX
+
+  def test_frexp_f32_denormal(self):
+    """v_frexp_exp/mant_f32 of denormal/zero inputs is (0, signed zero) on hardware."""
+    for bits in (0x00000001, 0x007FFFFF, 0x00000000):
+      st = run_program([v_mov_b32_e32(v[0], bits), v_frexp_exp_i32_f32_e32(v[1], v[0]), v_frexp_mant_f32_e32(v[2], v[0])], n_lanes=1)
+      self.assertEqual(st.vgpr[0][1] & 0xFFFFFFFF, 0, f"exp bits=0x{bits:08x}")
+      self.assertEqual(st.vgpr[0][2], bits & 0x80000000, f"mant bits=0x{bits:08x}")
+    # negative denormal: mant is -0.0
+    st = run_program([v_mov_b32_e32(v[0], 0x80000001), v_frexp_mant_f32_e32(v[2], v[0])], n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0x80000000)
+
+  def test_frexp_f64_denormal(self):
+    """v_frexp_exp_f64 of a denormal returns the normalized exponent (-1073 for min-denormal); zero -> 0."""
+    st = run_program([v_mov_b32_e32(v[0], 1), v_mov_b32_e32(v[1], 0), v_frexp_exp_i32_f64_e32(v[2], v[0:1])], n_lanes=1)
+    self.assertEqual(st.vgpr[0][2] & 0xFFFFFFFF, 0xFFFFFBCF)  # -1073
+    st = run_program([v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0), v_frexp_exp_i32_f64_e32(v[2], v[0:1])], n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0)
+
+  def test_frexp_exp_inf_nan(self):
+    """v_frexp_exp of +/-inf and NaN is 0 on hardware (host frexp gives 129/1024), for both f32 and f64."""
+    for bits in (0x7F800000, 0xFF800000, 0x7FC00000):
+      st = run_program([v_mov_b32_e32(v[0], bits), v_frexp_exp_i32_f32_e32(v[1], v[0])], n_lanes=1)
+      self.assertEqual(st.vgpr[0][1] & 0xFFFFFFFF, 0, f"f32 bits=0x{bits:08x}")
+    for lo, hi in ((0, 0x7FF00000), (0, 0xFFF00000), (0, 0x7FF80000), (1, 0x7FF00000)):
+      st = run_program([v_mov_b32_e32(v[0], lo), v_mov_b32_e32(v[1], hi), v_frexp_exp_i32_f64_e32(v[2], v[0:1])], n_lanes=1)
+      self.assertEqual(st.vgpr[0][2] & 0xFFFFFFFF, 0, f"f64 bits=0x{hi:08x}{lo:08x}")
+
+
 if __name__ == '__main__':
   unittest.main()

test/amd/hw/test_vop2.py

@@ -989,6 +989,53 @@ class TestCarryOps(unittest.TestCase):
     self.assertEqual(st.vgpr[0][0], 0)  # 0xFFFFFFFF + 1 + 0 = 0 (overflow)
     self.assertEqual(st.vcc, 0xDEADBEEF)  # VCC unchanged - carry was discarded
 
+class TestSelectFlushRegressions(unittest.TestCase):
+  """Regression tests: f32 MIN/MAX flush denormal inputs to signed zero (select-style ops propagate inputs bitwise)."""
+
+  def test_v_min_f32_denormal_flush(self):
+    """min(denormal, 1.0) is +0, min(-denormal, -1.0) is -0."""
+    st = run_program([v_mov_b32_e32(v[0], 0x00000001), v_mov_b32_e32(v[1], 0x3F800000), v_min_f32_e32(v[2], v[0], v[1])], n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0x00000000)
+    # flush(-denormal) = -0.0 > -1.0, so the result is -1.0 (both operand orders)
+    st = run_program([v_mov_b32_e32(v[0], 0x80000001), v_mov_b32_e32(v[1], 0xBF800000), v_min_f32_e32(v[2], v[0], v[1])], n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0xBF800000)
+    st = run_program([v_mov_b32_e32(v[1], 0xBF800000), v_mov_b32_e32(v[2], 0x80000001), v_min_f32_e32(v[3], v[1], v[2])], n_lanes=1)
+    self.assertEqual(st.vgpr[0][3], 0xBF800000)
+
+  def test_v_max_f32_denormal_flush(self):
+    """max(-denormal, -1.0) is -0; max(+denormal, -0) is +0."""
+    st = run_program([v_mov_b32_e32(v[0], 0x80000001), v_mov_b32_e32(v[1], 0xBF800000), v_max_f32_e32(v[2], v[0], v[1])], n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0x80000000)
+    st = run_program([v_mov_b32_e32(v[0], 0x00000001), v_mov_b32_e32(v[1], 0x80000000), v_max_f32_e32(v[2], v[0], v[1])], n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0x00000000)
+
+
+class TestCarryExecRegressions(unittest.TestCase):
+  """Regression tests: per-lane VCC writes (carry ops) zero inactive lane bits - VCC = mask & EXEC, never preserved."""
+
+  def test_co_ci_e32_vcc_masked_by_exec(self):
+    """v_sub_co_ci_u32_e32 with EXEC=0xFFFF0000: hw clears inactive VCC bits instead of preserving them."""
+    instructions = [
+      s_mov_b32(EXEC_LO, 0xFFFF0000),
+      s_mov_b32(VCC_LO, 0xFFFFFFFF),  # preset all bits
+      v_mov_b32_e32(v[0], 0xFFFFFFFE), v_mov_b32_e32(v[1], 0x80000000),
+      v_sub_co_ci_u32_e32(v[2], v[0], v[1]),  # active lanes: no borrow
+    ]
+    st = run_program(instructions, n_lanes=32)
+    self.assertEqual(st.vcc, 0x00000000)
+
+  def test_co_ci_e32_vcc_masked_by_exec_ones(self):
+    """Same with all-ones carry: VCC = borrow_mask & EXEC."""
+    instructions = [
+      s_mov_b32(EXEC_LO, 0x0F0F0F0F),
+      s_mov_b32(VCC_LO, 0),
+      v_mov_b32_e32(v[0], 0xFFFFFFFF), v_mov_b32_e32(v[1], 1),
+      v_add_co_ci_u32_e32(v[2], v[0], v[1]),  # all lanes would carry if active
+    ]
+    st = run_program(instructions, n_lanes=32)
+    self.assertEqual(st.vcc, 0x0F0F0F0F)
+    self.assertEqual(st.vgpr[31][2], 0)  # 0xFFFFFFFF + 1 wraps to 0 in active lanes
+
 
 if __name__ == '__main__':
   unittest.main()

test/amd/hw/test_vop3.py

@@ -4,6 +4,7 @@ Includes: v_fma_f32, v_div_scale_f32, v_div_fmas_f32, v_div_fixup_f32,
           v_alignbit_b32, v_bfe_i32, v_mad_u64_u32, v_readlane_b32, v_writelane_b32
 """
 import unittest
+from tinygrad.helpers import OSX
 from test.amd.hw.helpers import *
 
 class TestFMA(unittest.TestCase):
@@ -3264,6 +3265,23 @@ class TestVOP3ClampMAD(unittest.TestCase):
     # 0xFFFF * 2 = 0x1FFFE, low 16 bits = 0xFFFE
     self.assertEqual(st.vgpr[0][3] & 0xFFFF, 0xFFFE, f"expected 0xFFFE, got 0x{st.vgpr[0][3] & 0xFFFF:04x}")
 
+class TestMadNarrowClampRegressions(unittest.TestCase):
+  """Regression tests: mad i16/i24 with clamp saturate to narrow output range (found by random difftest vs hardware)."""
+
+  def test_mad_i16_clamp_sat_max(self):
+    # neg/src-floggled 16-bit mul operands are sign-extended after toggling bit15; sum > INT_MAX saturates
+    instructions = [s_mov_b32(s[4], 1232348160), v_mov_b32_e32(v[3], 0x80000000),
+                    v_mov_b32_e32(v[1], 0x7F7FFFFF), v_mad_i32_i16(v[0], s[4], v[3], v[1], 0, 3, 5, 1)]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][0], 0x7FFFFFFF)
+
+  def test_mad_i24_clamp_sat_min(self):
+    # sext24(-6344704) * sext24(+4210688) << -2^31 saturates to INT_MIN
+    instructions = [s_mov_b32(s[7], 4290772992), v_mov_b32_e32(v[1], 1077936128),
+                    v_mad_i32_i24(v[0], s[7], v[1], v[1], 1, 0, 0, 1)]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][0], 0x80000000)
+
 
 class TestCvtPkF16(unittest.TestCase):
   """Tests for V_CVT_PK_RTZ_F16_F32 - pack two f32 to f16 with round toward zero."""
@@ -3651,6 +3669,80 @@ class TestPermlane(unittest.TestCase):
     self.assertEqual(st.vgpr[21][1], 5)
     self.assertEqual(st.vgpr[31][1], 15)
 
+class TestClampLdExpRegressions(unittest.TestCase):
+  """Regression tests for f32 clamp (-0 -> +0) and ldexp input passthrough."""
+
+  def test_clamp_negative_zero(self):
+    """clmp=1 maps -0.0 to +0.0 (found by random difftest vs hardware)."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0x80000000), v_mov_b32_e32(v[1], 0x80000000),
+      v_add_f32_e64(v[2], v[0], v[1], clmp=1),  # -0 + -0 = -0, clamp -> +0
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0x00000000)
+    instructions = [
+      v_mov_b32_e32(v[0], 0x3F800000), v_mov_b32_e32(v[1], 0x80000000),
+      v_min_f32_e64(v[2], v[0], v[1], clmp=1),  # min(1.0, -0) = -0, clamp -> +0
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0x00000000)
+
+  def test_ldexp_special_inputs(self):
+    """v_ldexp_f32 of 0/-0/inf/NaN propagates the input instead of computing val * 2**exp (0*inf = NaN on host)."""
+    # -0.0 * 2^INT_MIN = -0.0 (src1 as integer exponent; huge negative)
+    instructions = [v_mov_b32_e32(v[0], 0x80000000), v_mov_b32_e32(v[1], 0x80000000), v_ldexp_f32(v[2], v[0], v[1])]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0x80000000)
+    # inf stays inf even with negative exponent
+    instructions = [v_mov_b32_e32(v[0], 0x7F800000), v_mov_b32_e32(v[1], 0xFFFFFF80), v_ldexp_f32(v[2], v[0], v[1])]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0x7F800000)
+
+  def test_ldexp_denormal_flush(self):
+    """v_ldexp_f32/f64 flush denormal inputs to signed zero (found by random difftest vs hardware)."""
+    # ldexp(+denorm, 1) = +0, ldexp(-denorm, 250) = -0
+    for src, exp_val, want in [(0x00000001, 1, 0x00000000), (0x80000001, 250, 0x80000000)]:
+      st = run_program([v_mov_b32_e32(v[0], src), v_mov_b32_e32(v[1], exp_val), v_ldexp_f32(v[2], v[0], v[1])], n_lanes=1)
+      self.assertEqual(st.vgpr[0][2], want)
+
+  def test_v_mul_neg_modifier_nan_sign(self):
+    """neg modifier is a pure sign-bit toggle on a NaN operand; result keeps that sign (found by random difftest)."""
+    # mul(normal, NEG(ABS(qNaN))): NaN payload negated in the operand stays negative qNaN
+    instructions = [v_mov_b32_e32(v[0], 0xC96CF47F), v_mov_b32_e32(v[1], 0x7FC00000),
+                    v_mul_f32_e64(v[2], v[0], v[1], s[0], 0, 7, 6)]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0xFFC00000)
+    # plain neg modifier still applies to non-NaN values: mul(-1.0, NEG(2.0)) = +2.0
+    st = run_program([v_mov_b32_e32(v[0], 0xBF800000), v_mov_b32_e32(v[1], 0x40000000),
+                      v_mul_f32_e64(v[2], v[0], v[1], s[0], 0, 2, 0)], n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0x40000000)
+
+
+class TestNaNPropagationRegressions(unittest.TestCase):
+  """Regression tests: float arithmetic propagates a NaN from the FIRST NaN operand, quieted with its own sign/payload."""
+
+  @unittest.skipIf(OSX, "broken on mac, TODO: why?")
+  def test_mul_nan_priority(self):
+    # first NaN operand wins (sign+payload), not x86's second-source propagation
+    for a, b, want in [(0x7FC00001, 0x7F800003, 0x7FC00001), (0xFFC00005, 0x7F800003, 0xFFC00005),
+                       (0x7F800001, 0xFFC00005, 0x7FC00001), (0xFF9F1800, 0x7F800001, 0xFFDF1800)]:
+      st = run_program([v_mov_b32_e32(v[0], a), v_mov_b32_e32(v[1], b),
+                        v_mul_f32_e32(v[2], v[0], v[1])], n_lanes=1)
+      self.assertEqual(st.vgpr[0][2], want, f"mul({a:#x}, {b:#x})")
+
+class TestMinMaxFlushE64Regressions(unittest.TestCase):
+  """Regression tests: f32 min/max/median flush denormal inputs to signed zero (e64 forms)."""
+
+  def test_v_min3_f32_denormal_flush(self):
+    st = run_program([v_mov_b32_e32(v[0], 0x00000001), v_mov_b32_e32(v[1], 0x3F800000), v_mov_b32_e32(v[2], 0x40000000),
+                      v_min3_f32(v[3], v[0], v[1], v[2])], n_lanes=1)
+    self.assertEqual(st.vgpr[0][3], 0x00000000)  # min(+denorm, 1, 2) = +0
+
+  def test_v_med3_f32_denormal_flush(self):
+    st = run_program([v_mov_b32_e32(v[0], 0x80000001), v_mov_b32_e32(v[1], 0x3F800000), v_mov_b32_e32(v[2], 0x40000000),
+                      v_med3_f32(v[3], v[0], v[1], v[2])], n_lanes=1)
+    self.assertEqual(st.vgpr[0][3], 0x3F800000)  # med(-0, 1, 2) = 1
+
 
 if __name__ == '__main__':
   unittest.main()

test/amd/hw/test_vop3p.py

@@ -457,6 +457,21 @@ class TestWMMAF16(unittest.TestCase):
         self.assertAlmostEqual(lo, 16.0, places=1, msg=f"v[{reg}] lane {lane}: expected 16.0, got {lo}")
         self.assertEqual(result >> 16, 0, msg=f"v[{reg}] lane {lane}: hi bits should be 0")
 
+  def test_v_wmma_f16_16x16x16_f16_inline_zero_accumulator(self):
+    """V_WMMA_F16_16X16X16_F16 with the inline constant 0 as C: D = A @ B, whatever v[128:135] holds."""
+    instructions: list[Inst] = []
+    instructions.append(s_mov_b32(s[0], 0x3c003c00))  # packed f16 1.0
+    for i in range(16, 32):
+      instructions.append(v_mov_b32_e32(v[i], s[0]))
+    instructions.append(s_mov_b32(s[1], 0x57b057b0))  # packed f16 123.0, poison where a VGPR read of "128" would land
+    for i in range(128, 136):
+      instructions.append(v_mov_b32_e32(v[i], s[1]))
+    instructions.append(v_wmma_f16_16x16x16_f16(v[0:7], v[16:23], v[24:31], 0))
+    st = run_program(instructions, n_lanes=32)
+    for lane in range(32):
+      for reg in range(8):
+        self.assertEqual(st.vgpr[lane][reg], 0x4c00, msg=f"v[{reg}] lane {lane}")
+
   def test_v_wmma_f16_16x16x16_f16_with_accumulator(self):
     """V_WMMA_F16_16X16X16_F16 with non-zero accumulator."""
     instructions: list[Inst] = []

test/amd/hw/test_vopc.py

@@ -973,6 +973,71 @@ class TestCmpxPartialWavefront(unittest.TestCase):
     self.assertEqual(st.sgpr[EXEC_LO.offset] & 0xFFFFFFFF, 0x4,
                      "Only lane 2 should be active after v_cmpx_eq_u32_e64")
 
+class TestClassDenormalRegressions(unittest.TestCase):
+  """Regression tests: V_CMP_CLASS classifies denormals as DENORMAL (raw bits), not as zero class."""
+
+  def test_class_pos_denormal(self):
+    for bits in (0x00000001, 0x007FFFFF):
+      instructions = [v_mov_b32_e32(v[0], bits), v_mov_b32_e32(v[1], 0x80), v_cmp_class_f32_e64(VCC_LO, v[0], v[1])]
+      st = run_program(instructions, n_lanes=1)
+      self.assertEqual(st.vcc, 1, f"bits=0x{bits:08x}")  # n_lanes=1
+      # ...and it is not the zero class
+      instructions = [v_mov_b32_e32(v[0], bits), v_mov_b32_e32(v[1], 0x40), v_cmp_class_f32_e64(VCC_LO, v[0], v[1])]
+      st = run_program(instructions, n_lanes=1)
+      self.assertEqual(st.vcc, 0, f"bits=0x{bits:08x}")
+
+  def test_class_neg_denormal(self):
+    instructions = [v_mov_b32_e32(v[0], 0x80000001), v_mov_b32_e32(v[1], 0x10), v_cmp_class_f32_e64(VCC_LO, v[0], v[1])]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vcc, 1)  # n_lanes=1
+    instructions = [v_mov_b32_e32(v[0], 0x80000001), v_mov_b32_e32(v[1], 0x20), v_cmp_class_f32_e64(VCC_LO, v[0], v[1])]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vcc, 0)  # not the negative-zero class
+
+
+class TestIntCmpModRegressions(unittest.TestCase):
+  """Regression tests: int compares (i32/u32) honor abs/neg as bit-level sign clear/flip (not integer abs/negate)."""
+
+  def test_cmp_i32_abs_neg_bit_level(self):
+    # abs(0x80000001) = 1 -> 1 > 1 is false (integer abs would give 2147483647 > 1)
+    instructions = [v_mov_b32_e32(v[0], 0x80000001), v_mov_b32_e32(v[1], 1), v_cmp_gt_i32_e64(VCC_LO, v[0], v[1], abs=1)]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vcc, 0)
+    # neg(0x80000001) flips the sign bit -> 1 > 2 is false (integer negate would give 2147483647 > 2)
+    instructions = [v_mov_b32_e32(v[0], 0x80000001), v_mov_b32_e32(v[1], 2), v_cmp_gt_i32_e64(VCC_LO, v[0], v[1], neg=1)]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vcc, 0)
+
+  def test_cmp_u32_abs_bit_level(self):
+    # abs(0x80000000) = 0 -> 0 < 1 is true
+    instructions = [v_mov_b32_e32(v[0], 0x80000000), v_mov_b32_e32(v[1], 1), v_cmp_lt_u32_e64(VCC_LO, v[0], v[1], abs=1)]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vcc, 1)  # n_lanes=1
+
+
+class TestCmpxSdstRegressions(unittest.TestCase):
+  """Regression tests: V_CMPX_*_E64 writes EXEC only, never SDST (hardware verified)."""
+
+  def test_cmpx_e64_no_sdst(self):
+    instructions = [
+      s_mov_b32(VCC_LO, 0),  # preset VCC to 0
+      v_mov_b32_e32(v[0], 0x3F800000), v_mov_b32_e32(v[1], 0x40000000),
+      v_cmpx_lt_f32_e64(VCC_LO, v[0], v[1]),  # 1.0 < 2.0
+    ]
+    st = run_program(instructions, n_lanes=32)
+    self.assertEqual(st.sgpr[EXEC_LO.offset], 0xFFFFFFFF)  # EXEC updated
+    self.assertEqual(st.vcc, 0)  # but VCC untouched
+
+  def test_cmpx_e64_partial_exec(self):
+    instructions = [
+      s_mov_b32(EXEC_LO, 0x0F0F0F0F),
+      s_mov_b32(VCC_LO, 0xFFFFFFFF),
+      v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0x3F800000),
+      v_cmpx_lt_f32_e64(VCC_LO, v[0], v[1]),
+    ]
+    st = run_program(instructions, n_lanes=32)
+    self.assertEqual(st.sgpr[EXEC_LO.offset], 0x0F0F0F0F)  # EXEC = computed & old EXEC
+
 
 if __name__ == '__main__':
   unittest.main()

test/amd/test_asm_kernel.py

@@ -30,7 +30,7 @@ def custom_add_one(A:UOp) -> UOp:
     s_endpgm(),
   ]
   sink = UOp.sink(A.base, threads, arg=KernelInfo(f"custom_add_one_{A.numel()}", estimates=Estimates(ops=A.numel(), mem=A.numel()*4*2)))
-  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.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts]))))
 
 def custom_add_var(A:UOp, B:UOp) -> UOp:
   A,B = A.flatten(), B.flatten()
@@ -49,7 +49,7 @@ def custom_add_var(A:UOp, B:UOp) -> UOp:
     s_endpgm(),
   ]
   sink = UOp.sink(A.base, B.base, var, threads, arg=KernelInfo(f"custom_add_var_{A.numel()}"))
-  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.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts]))))
 
 def custom_wave_sync(A:UOp, arch:str) -> UOp:
   # 4 waves across 1024 WG — enough to saturate a SIMD with many concurrent WGs
@@ -63,7 +63,7 @@ def custom_wave_sync(A:UOp, arch:str) -> UOp:
     insts += [s_nop(0)]*4
   insts.append(s_endpgm())
   sink = UOp.sink(A.base, threads, wg, arg=KernelInfo("custom_wave_sync"))
-  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.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts]))))
 
 def custom_lds_sync(A:UOp, arch:str) -> UOp:
   A = A.flatten()
@@ -97,7 +97,7 @@ def custom_lds_sync(A:UOp, arch:str) -> UOp:
     isa.s_endpgm(),
   ]
   sink = UOp.sink(A.base, lds, threads, wg, arg=KernelInfo("custom_lds_sync"))
-  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.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts]))))
 
 def custom_handwritten(A:UOp) -> UOp:
   A = A.flatten()
@@ -143,7 +143,7 @@ def custom_handwritten(A:UOp) -> UOp:
   k.emit(r4.s_endpgm())
   insts = k.finalize()
   sink = UOp.sink(A.base, threads, wg, lds, arg=KernelInfo("custom_handwritten"))
-  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.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts]))))
 
 def custom_data_deps(A:UOp) -> UOp:
   A = A.flatten()
@@ -159,7 +159,7 @@ def custom_data_deps(A:UOp) -> UOp:
   k.emit(s_endpgm())
   insts = k.finalize()
   sink = UOp.sink(A.base, threads, arg=KernelInfo("custom_data_deps"))
-  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.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts]))))
 
 @unittest.skipUnless(Device.DEFAULT == "AMD", "requires AMD device")
 class TestAsmKernel(unittest.TestCase):

test/amd/test_emu2_pcode.py

@@ -152,7 +152,7 @@ class TestDSPcodePatterns(unittest.TestCase):
 
   def test_global_atomic_add_f32_parsing(self):
     """Test GLOBAL_ATOMIC_ADD_F32 keeps memory values in float dtype."""
-    vmem = UOp.param(2, dtypes.uint32, (1024,))
+    vmem = UOp.param(2, dtypes.uint32, 1024)
     srcs = {
       'ADDR': UOp.const(0, dtypes.uint64),
       'DATA': UOp.const(0x3f800000, dtypes.uint32),
@@ -183,7 +183,7 @@ class TestDSPcodePatterns(unittest.TestCase):
   def test_mem_read_parsing(self):
     """Test MEM[addr].type read expression parsing."""
     # Create a mock LDS buffer
-    lds = UOp.param(3, dtypes.uint32, (16384,))
+    lds = UOp.param(3, dtypes.uint32, 16384)
     addr = UOp.const(0, dtypes.uint32)
     vrs = {'_lds': lds, 'ADDR': addr, 'OFFSET': UOp.const(0, dtypes.uint32)}
 
@@ -218,7 +218,7 @@ class TestDSPcodePatterns(unittest.TestCase):
     pcode = PCODE.get(DSOp.DS_LOAD_2ADDR_B32)
     self.assertIsNotNone(pcode)
     assert pcode is not None
-    lds = UOp.param(3, dtypes.uint32, (16384,))
+    lds = UOp.param(3, dtypes.uint32, 16384)
     srcs = {
       'ADDR': UOp.const(0, dtypes.uint32),
       'OFFSET0': UOp.const(0, dtypes.uint32),
@@ -299,7 +299,7 @@ class TestConcatWidthParsing(unittest.TestCase):
       self.assertIs(parsed.simplify(), UOp.const(expected, dtypes.uint32))
 
   def test_permlane64_wave64_pcode_indices(self):
-    vgpr = UOp.param(0, dtypes.uint32, (256,))
+    vgpr = UOp.param(0, dtypes.uint32, 256)
     srcs = {
       'SRC0': UOp.const(0, dtypes.uint32),
       'VDST': UOp.const(1, dtypes.uint32),
@@ -330,7 +330,7 @@ class TestAllPcode(unittest.TestCase):
   def _make_srcs(self):
     """Create dummy source variables for pcode parsing."""
     u32, u64 = lambda v=0: UOp.const(v, dtypes.uint32), lambda v=0: UOp.const(v, dtypes.uint64)
-    lds = UOp.param(3, dtypes.uint32, (16384,))
+    lds = UOp.param(3, dtypes.uint32, 16384)
     return {'laneId': u32(), 'laneID': u32(), 'S0': u32(), 'S1': u32(), 'S2': u32(), 'S3': u32(), 'SRC0': u32(),
             'D0': u32(), 'D1': u32(), 'DST': u32(), 'VDST': u32(), 'SDST': u32(),
             'VCC': u64(), 'VCCZ': u32(), 'EXEC': u64(), 'EXEC_LO': u32(), 'EXECZ': u32(), 'SCC': u32(),

test/amd/test_sqtt_examples.py

@@ -88,7 +88,6 @@ def run_rocprof_decoder(blobs: list[bytes], lib: bytes, base: int, target: str):
   if t.is_alive(): raise RuntimeError("rocprof decoder timeout")
   return occupancy_records, wave_insts
 
-@unittest.skip("TODO: fix to not require unpickling UOps.")
 class SQTTExamplesTestBase(unittest.TestCase):
   target: str
   examples: dict

test/amd/test_sqtt_profiler.py

@@ -1,45 +1,81 @@
 import unittest, contextlib
-from tinygrad import Device, Tensor, Context, TinyJit
-from tinygrad.device import Compiled, ProfileProgramEvent, ProfileDeviceEvent
+from tinygrad import Device, Tensor, Context, TinyJit, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.device import Compiled, ProfileProgramEvent
+from tinygrad.runtime.ops_amd import ProfileSQTTEvent
 from tinygrad.engine.realize import run_linear
 from tinygrad.codegen import to_program
 from tinygrad.viz.serve import load_amd_counters, VizData
+from tinygrad.renderer.amd.sqtt import decode, print_packets
+from tinygrad.renderer.amd.dsl import s
 
 @contextlib.contextmanager
 def save_sqtt():
-  data = VizData()
-  yield data.ctxs
+  Device[Device.DEFAULT].synchronize()
+  profile_start = len(Compiled.profile_events)
+  data = []
+  yield data
   Device[Device.DEFAULT].synchronize()
   Device[Device.DEFAULT]._at_profile_finalize()
-  load_amd_counters(data, Compiled.profile_events)
-  data.ctxs[:] = [r for r in data.ctxs if r["name"].startswith("SQTT")]
+  data[:] = [e for e in Compiled.profile_events[:profile_start] if isinstance(e, ProfileProgramEvent)]+Compiled.profile_events[profile_start:]
+
+def map_sqtt(profile:list) -> list[dict]:
+  load_amd_counters(data:=VizData(), profile)
+  return [r for r in data.ctxs if r["name"].startswith("SQTT")]
+
+def custom_asm_cdna(A:UOp):
+  import tinygrad.runtime.autogen.amd.cdna.ins as cdna
+  WAVE_SIZE = 64
+  insts = [cdna.s_nop(0), cdna.s_mov_b32(s[0], 10)]
+  return custom_asm(A, insts+[cdna.s_endpgm()], WAVE_SIZE*2)
+
+def custom_asm_rdna(A:UOp):
+  import tinygrad.runtime.autogen.amd.rdna3.ins as rdna3
+  WAVE_SIZE = 32
+  insts = [rdna3.s_nop(0), rdna3.s_mov_b32(s[0], 10)]
+  return custom_asm(A, insts+[rdna3.s_endpgm()], WAVE_SIZE*2)
+
+def custom_asm(A, insts, num_threads) -> UOp:
+  return UOp(Ops.PROGRAM, src=(UOp.sink(A, UOp.special(num_threads, "lidx0"), arg=KernelInfo("asm")), \
+      UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS,arg=(x,dtypes.void)) for x in insts]))))
 
 @unittest.skipUnless(Device.DEFAULT == "AMD", "only runs on AMD")
 class TestSQTTProfiler(unittest.TestCase):
-  # TODO: can we enable SQTT profiling in context?
   @classmethod
   def setUpClass(cls):
     if not Device[Device.DEFAULT].sqtt_enabled: raise unittest.SkipTest("device must be in SQTT profiling mode")
-
-  def setUp(self):
-    Device[Device.DEFAULT].synchronize()
-    Compiled.profile_events[:] = [e for e in Compiled.profile_events if isinstance(e, (ProfileProgramEvent, ProfileDeviceEvent))]
+    cls.arch = Device[Device.DEFAULT].arch
 
   def test_simple(self):
     t = Tensor.empty(1) + 1
-    with save_sqtt() as sqtt:
+    with save_sqtt() as data:
       linear = t.schedule_linear()
       run_linear(linear)
     fn_name = to_program(linear.src[0].src[0], renderer=Device[Device.DEFAULT].renderer).arg.function_name
+    sqtt = map_sqtt(data)
     self.assertEqual(len(sqtt), 1)
     self.assertEqual(sqtt[0]["name"], f"SQTT {fn_name}")
 
+  def test_asm(self):
+    t = Tensor.empty(1)
+    with save_sqtt() as data:
+      t.custom_kernel(fxn=custom_asm_cdna if self.arch == "gfx950" else custom_asm_rdna)[0].realize()
+    for event in data:
+      if not isinstance(event, ProfileSQTTEvent) or not event.itrace: continue
+      print(f"\n=== SE {event.se} ===")
+      print_packets(decode(event.blob))
+    from test.null.test_viz import write_files, run_cli
+    with write_files(profile=data) as files:
+      out = run_cli(*files, "-s", "asm SQTT SE:0 PKTS", json_fmt=False)[0]["out"]
+    print(out)
+
   def test_multiple_runs(self):
     t = Tensor.empty(1) + 1
-    with save_sqtt() as sqtt:
+    with save_sqtt() as data:
       linear = t.schedule_linear()
       for _ in range(N:=3): run_linear(linear)
     fn_name = to_program(linear.src[0].src[0], renderer=Device[Device.DEFAULT].renderer).arg.function_name
+    sqtt = map_sqtt(data)
     self.assertEqual(len(sqtt), N)
     for i in range(1, N):
       self.assertEqual(sqtt[i]["name"], f"SQTT {fn_name} n{i+1}")
@@ -47,8 +83,9 @@ class TestSQTTProfiler(unittest.TestCase):
   def test_multiple_kernels(self):
     t = ((Tensor.empty(1) + 1).contiguous() + 2)
     linear = t.schedule_linear()
-    with save_sqtt() as sqtt:
+    with save_sqtt() as data:
       run_linear(linear)
+    sqtt = map_sqtt(data)
     self.assertEqual(len(sqtt), len(linear.src))
     for i,call in enumerate(linear.src):
       fn_name = to_program(call.src[0], renderer=Device[Device.DEFAULT].renderer).arg.function_name
@@ -57,8 +94,9 @@ class TestSQTTProfiler(unittest.TestCase):
   def test_multiple_kernels_lower(self):
     t = ((Tensor.empty(1) + 1).contiguous() + 2)
     linear = t.schedule_linear()
-    with save_sqtt() as sqtt:
+    with save_sqtt() as data:
       run_linear(linear)
+    sqtt = map_sqtt(data)
     self.assertEqual(len(sqtt), len(linear.src))
     for i,call in enumerate(linear.src):
       fn_name = to_program(call.src[0], renderer=Device[Device.DEFAULT].renderer).arg.function_name
@@ -68,21 +106,23 @@ class TestSQTTProfiler(unittest.TestCase):
     @TinyJit
     def f(a): return a + 1
     t = Tensor.empty(1)
-    with save_sqtt() as sqtt:
+    with save_sqtt() as data:
       for _ in range(N:=5):
         f(t).realize()
+    sqtt = map_sqtt(data)
     self.assertEqual(len(sqtt), N)
     kernel_name = sqtt[0]["name"]
-    for i,s in enumerate(sqtt[1:], start=1): self.assertEqual(s["name"], f"{kernel_name} n{i+1}")
+    for i,e in enumerate(sqtt[1:], start=1): self.assertEqual(e["name"], f"{kernel_name} n{i+1}")
 
   # TODO: can we trace SQTT for graphed kernels?
   def test_jit_graph(self, kernel_count=3*1):
     @TinyJit
     def f(a): return ((a + 1).contiguous() + 2).contiguous().sum()
     t = Tensor.empty(32)
-    with save_sqtt() as sqtt:
+    with save_sqtt() as data:
       for _ in range(5):
         f(t).realize()
+    sqtt = map_sqtt(data)
     names = [s["name"] for s in sqtt]
     k0, k1, k2 = names[:3]
     for i in range(3, len(sqtt), 3):

test/backend/test_asm_gemm.py

@@ -188,7 +188,7 @@ class TestMXFP4(unittest.TestCase):
     M, N, K = getenv("M", 16384), getenv("N", 4096), getenv("K", 14336)
     a = Tensor.empty(M, K, dtype=dtypes.bfloat16)
     b = Tensor.empty(N, K, dtype=dtypes.bfloat16)
-    asm_gemm(a, b.T, mxfp4=True).realize()
+    for _ in range(getenv("CNT", 1)): asm_gemm(a, b.T, mxfp4=True).realize()
 
 # test the Asm GEMM with Llama shapes, only run on the real machine for speed
 

test/backend/test_assign.py

@@ -1,10 +1,10 @@
 #!/usr/bin/env python
 import unittest
 import numpy as np
-from tinygrad import dtypes, Tensor, TinyJit, GlobalCounters, Variable
+from tinygrad import Device, dtypes, Tensor, TinyJit, GlobalCounters, Variable
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.helpers import temp, DEV, Context
-from test.helpers import assert_kernel_count
+from test.helpers import assert_kernel_count, needs_second_gpu
 
 N = 200  # has to be bigger than the cache to fail
 
@@ -22,7 +22,7 @@ class TestAssign(unittest.TestCase):
     assert ba1 == ba2 and ba1 != bb1
     np.testing.assert_allclose(a.numpy(), (np.arange(N*N)*2).reshape((N,N)))
 
-  def test_assign_zeros_good(self):
+  def test_assign_keeps_identical_tensor(self):
     a = Tensor.zeros(10,10).contiguous()
     a.assign(Tensor.ones(10,10))
     b = Tensor.zeros(10,10).contiguous()
@@ -30,7 +30,7 @@ class TestAssign(unittest.TestCase):
     np.testing.assert_allclose(b.numpy(), 0)
 
   @unittest.skip("TODO: this often crashes in CI")
-  def test_assign_zeros(self):
+  def test_assign_keeps_earlier_identical_tensor(self):
     a = Tensor.zeros(10,10).contiguous()
     b = Tensor.zeros(10,10).contiguous()
     a.assign(Tensor.ones(10,10))
@@ -114,15 +114,6 @@ class TestAssign(unittest.TestCase):
     x.assign(x + 1)
     assert [y0.item(), y1.item(), y2.item(), x.item()] == [0.0, 1.0, 2.0, 3.0]
 
-  def test_assign_add_jit(self):
-    @TinyJit
-    def f(x):
-      x += 1
-      x.realize()
-    x = Tensor([0])
-    for _ in range(5): f(x)
-    assert x.item() == 5
-
   def test_assign_add_jit_other(self):
     @TinyJit
     def f(x):
@@ -180,21 +171,20 @@ class TestAssign(unittest.TestCase):
     Tensor.realize(a.contiguous().assign(1), b.contiguous().assign(2))
     self.assertEqual((a + b).item(), 3)
 
-  def test_assign_diamond_cycle(self):
-    # NOTE: should *not* raise AssertionError from numpy
-    with self.assertRaisesRegex(RuntimeError, "cycle"):
-      a = Tensor.ones(4).contiguous().realize()
-      times_a = a*3
-      a.assign(Tensor.full((4,), 2.).contiguous())
-      new = a + (times_a-1)
+  def test_assign_diamond(self):
+    a = Tensor.ones(4).contiguous().realize()
+    times_a = a*3
+    a.assign(Tensor.full((4,), 2.).contiguous())
+    new = a + (times_a-1)
+    with self.assertRaisesRegex(RuntimeError, "cycle"):  # TODO: broken now, raises
       np.testing.assert_allclose(new.numpy(), 4)
 
-  def test_assign_diamond_contiguous_cycle(self):
-    with self.assertRaisesRegex(RuntimeError, "cycle"):
-      a = Tensor.ones(4).contiguous().realize()
-      times_a = a*3
-      a.assign(Tensor.full((4,), 2.))
-      new = a.contiguous() + times_a-1
+  def test_assign_diamond_contiguous(self):
+    a = Tensor.ones(4).contiguous().realize()
+    times_a = a*3
+    a.assign(Tensor.full((4,), 2.))
+    new = a.contiguous() + times_a-1
+    with self.assertRaisesRegex(RuntimeError, "cycle"):  # TODO: broken now, raises
       np.testing.assert_allclose(new.numpy(), 4)
 
   def test_assign_diamond_possible(self):
@@ -267,13 +257,12 @@ class TestAssign(unittest.TestCase):
     np.testing.assert_equal(b1.numpy(), 608)
 
   def test_crossunder_assign(self):
-    # NOTE: should *not* raise AssertionError from numpy
-    with self.assertRaisesRegex(RuntimeError, "cycle"):
-      a = Tensor.full((4,), 2).contiguous().realize()
-      b = Tensor.full((4,), 3).contiguous().realize()
-      c = a+9
-      a += b
-      b += c
+    a = Tensor.full((4,), 2).contiguous().realize()
+    b = Tensor.full((4,), 3).contiguous().realize()
+    c = a+9
+    a += b
+    b += c
+    with self.assertRaisesRegex(RuntimeError, "cycle"):  # TODO: broken now, raises
       Tensor.realize(a,b)
       np.testing.assert_allclose(a.numpy(), 2+3)
       np.testing.assert_allclose(b.numpy(), 3+2+9)
@@ -356,49 +345,17 @@ class TestAssign(unittest.TestCase):
     # permute and base are the same buffer
     assert ba1 == ba2 and ba1 != bb1
 
-  def test_post_permuted_assignment(self):
-    a = Tensor(np.arange(N*N, dtype=np.float32)).reshape(N,N)
-    b = Tensor(np.arange(N*N, dtype=np.float32)).reshape(N,N)
-    a.realize()
-    b.realize()
-    #GlobalCounters.cache = []
-    ba1 = a.uop.base.realized # noqa: F841
-    bb1 = b.uop.base.realized # noqa: F841
-    a.assign(a.permute(1,0) + b)   # this should not work!
-    a.realize()
-    ba2 = a.uop.base.realized # noqa: F841
-    # NOTE: don't test that it's assigned
-    #assert ba1 == ba2 and ba1 != bb1
-    np.testing.assert_allclose(a.numpy(), np.arange(N*N).reshape((N,N)) + np.arange(N*N).reshape((N,N)).transpose(1,0))
-
-  def test_post_permuted_assignment_alt(self):
+  def _assign_view_of_self(self, view):
     a = Tensor.arange(N*N).reshape(N,N).clone().realize()
     b = Tensor.arange(N*N).reshape(N,N).clone().realize()
-    new_a = (a.T+b).numpy()
-    a.assign(a.T+b)
+    new_a = (view(a)+b).numpy()
+    a.assign(view(a)+b)
     np.testing.assert_allclose(a.numpy(), new_a)
 
-  def test_post_flipped_assignment(self):
-    a = Tensor.arange(N*N).reshape(N,N).clone().realize()
-    b = Tensor.arange(N*N).reshape(N,N).clone().realize()
-    new_a = (a.flip(0)+b).numpy()
-    a.assign(a.flip(0)+b)
-    np.testing.assert_allclose(a.numpy(), new_a)
-
-  def test_post_flipped_assignment_axis1(self):
-    a = Tensor.arange(N*N).reshape(N,N).clone().realize()
-    b = Tensor.arange(N*N).reshape(N,N).clone().realize()
-    new_a = (a.flip(1)+b).numpy()
-    a.assign(a.flip(1)+b)
-    np.testing.assert_allclose(a.numpy(), new_a)
-
-  def test_post_reshape_assignment_fine(self):
-    a = Tensor.arange(N*N).reshape(N, N).clone().realize()
-    b = Tensor.arange(N*N).reshape(N, N).clone().realize()
-    rhs = a.reshape(-1).reshape(N, N)
-    new_a = (rhs+b).numpy()
-    a.assign(rhs+b)  # self-assign with reshape view is fine
-    np.testing.assert_allclose(a.numpy(), new_a)
+  def test_post_permuted_assignment(self): self._assign_view_of_self(lambda a: a.T)
+  def test_post_flipped_assignment(self): self._assign_view_of_self(lambda a: a.flip(0))
+  def test_post_flipped_assignment_axis1(self): self._assign_view_of_self(lambda a: a.flip(1))
+  def test_post_reshape_assignment(self): self._assign_view_of_self(lambda a: a.reshape(-1).reshape(N,N))
 
   @unittest.skip("multi output not supported anymore")
   def test_simple_assignment_multioutput(self):
@@ -421,14 +378,6 @@ class TestAssign(unittest.TestCase):
 
   # NOTE: if the assign target is read/write in a single kernel, it should be contiguous
 
-  def test_permuted_assignment_correct(self):
-    a = Tensor.arange(4 * 4).reshape(4, 4).clone().realize()
-    b = Tensor.arange(4 * 4).reshape(4, 4).clone().realize()
-    a = a.permute(1, 0)
-    new_val = a + b
-    a.assign(new_val)
-    np.testing.assert_equal(a.numpy(), np.arange(4 * 4).reshape(4, 4).transpose(1, 0) + np.arange(4 * 4).reshape(4, 4))
-
   def test_permuted_reduceop_child_dual_use(self):
     a = Tensor.arange(32*32*32).reshape(32, 32, 32).clone().realize()
     b = Tensor.ones(32, 32, dtype=dtypes.int).contiguous().realize()
@@ -526,24 +475,28 @@ class TestAssign(unittest.TestCase):
     a[2:5] = [1, 2, 3]
     np.testing.assert_allclose(a.numpy(), [0., 0., 1., 2., 3., 0., 0., 0.])
 
+  # IEEE 754: 1.0f = 0x3f800000, 2.0f = 0x40000000, 3.0f = 0x40400000, 4.0f = 0x40800000
+  REVERSED = [0x40800000, 0x40400000, 0x40000000, 0x3f800000]
+
   def test_assign_bitcast(self):
-    # assign to a bitcast view should modify the underlying buffer
     a = Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32).realize()
-    # IEEE 754: 1.0f = 0x3f800000, 2.0f = 0x40000000, 3.0f = 0x40400000, 4.0f = 0x40800000
-    a.bitcast(dtypes.uint32).assign(Tensor([0x40800000, 0x40400000, 0x40000000, 0x3f800000], dtype=dtypes.uint32)).realize()
+    a.bitcast(dtypes.uint32).assign(Tensor(self.REVERSED, dtype=dtypes.uint32)).realize()
     np.testing.assert_allclose(a.numpy(), [4.0, 3.0, 2.0, 1.0])
-    # double bitcast
+
+  def test_assign_bitcast_unrealized(self):
+    a = Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32).realize()
+    a.bitcast(dtypes.uint32).assign(Tensor(self.REVERSED, dtype=dtypes.uint32))
+    np.testing.assert_allclose(a.numpy(), [4.0, 3.0, 2.0, 1.0])
+
+  def test_assign_double_bitcast(self):
     b = Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32).realize()
-    b.bitcast(dtypes.uint32).bitcast(dtypes.int32).assign(Tensor([0x40800000, 0x40400000, 0x40000000, 0x3f800000], dtype=dtypes.int32)).realize()
+    b.bitcast(dtypes.uint32).bitcast(dtypes.int32).assign(Tensor(self.REVERSED, dtype=dtypes.int32)).realize()
     np.testing.assert_allclose(b.numpy(), [4.0, 3.0, 2.0, 1.0])
-    # shrink then bitcast
+
+  def test_assign_shrink_then_bitcast(self):
     c = Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32).realize()
-    c[0:2].bitcast(dtypes.uint32).assign(Tensor([0x40800000, 0x40400000], dtype=dtypes.uint32)).realize()
+    c[0:2].bitcast(dtypes.uint32).assign(Tensor(self.REVERSED[:2], dtype=dtypes.uint32)).realize()
     np.testing.assert_allclose(c.numpy(), [4.0, 3.0, 3.0, 4.0])
-    # without .realize()
-    a = Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32).realize()
-    a.bitcast(dtypes.uint32).assign(Tensor([0x40800000, 0x40400000, 0x40000000, 0x3f800000], dtype=dtypes.uint32))
-    np.testing.assert_allclose(a.numpy(), [4.0, 3.0, 2.0, 1.0])
 
   def test_assign_bitcast_different_size(self):
     # assign to a shape-changing bitcast view (only works on DISK currently)
@@ -562,13 +515,6 @@ class TestAssign(unittest.TestCase):
     with self.assertRaisesRegex(RuntimeError, "assign dtype mismatch"):
       a.assign(b)
 
-  def test_assign_dtype_mismatch_int64_to_float32(self):
-    # int64 -> float32 loses precision for large values, should not be implicit
-    a = Tensor.zeros(1, dtype=dtypes.float32).contiguous().realize()
-    b = Tensor([16777217], dtype=dtypes.int64)  # 2^24 + 1, not exactly representable in float32
-    with self.assertRaisesRegex(RuntimeError, "assign dtype mismatch"):
-      a.assign(b)
-
   def test_assign_shape_broadcast(self):
     # shape broadcasting should work when dtypes match
     a = Tensor.zeros(3, 5, dtype=dtypes.float32).contiguous().realize()
@@ -881,14 +827,16 @@ class TestAssignOrdering(unittest.TestCase):
   def test_war_reader_already_depends_on_write(self):
     x = Tensor([1.0]).contiguous().realize()
     y = Tensor([2.0]).contiguous().realize()
-    x_expr = x + 10
+    x_expr = x + 10          # 11, x is read here, before the assign
     x.assign(x * 2)
     y.assign(y + x)
     z = y + x_expr
     Tensor.realize(x, y, z)
-    # TODO: z should be 15: x_expr means 11 (x captured at build time), but the read is fused past the assign and
-    # sees the new bytes. once stale readers are scheduled before the overwrite, update this to 15
-    np.testing.assert_allclose([x.item(), y.item(), z.item()], [2.0, 4.0, 16.0])
+    try:
+      np.testing.assert_allclose([x.item(), y.item(), z.item()], [2.0, 4.0, 15.0])
+    except AssertionError:
+      # TODO: broken now, x_expr reads x after the assign
+      np.testing.assert_allclose([x.item(), y.item(), z.item()], [2.0, 4.0, 16.0])
 
   def test_war_multi_read_then_assign(self):
     devices = ("CPU:0", "CPU:1")
@@ -909,6 +857,152 @@ class TestAssignOrdering(unittest.TestCase):
     self.assertEqual(buf.sum().realize().item(), 6.0)
 
 # TODO: assigns into views of unrealized non-BUFFER bases are silently dropped
+  def test_read_before_two_assigns(self):
+    g = Tensor.full((2,), 4.0).realize()
+    before = g + 1                                     # 5
+    g.assign(0.0)
+    g.assign(g + 4)
+    with self.assertRaisesRegex(RuntimeError, "cycle"):  # TODO: broken now, raises
+      np.testing.assert_allclose((before + g).numpy(), 9)
+
+  def test_read_between_two_assigns(self):
+    a = Tensor.ones(4).realize()
+    b = Tensor.full((4,), 10.).realize()
+    a.assign(b + 1)                                    # a == 11
+    v1 = a * 3                                         # reads 11 -> 33
+    a.assign(b + 100)                                  # a == 110
+    out = (a + v1).numpy()
+    try:
+      np.testing.assert_allclose(out, 143)
+    except AssertionError:
+      # TODO: broken now, v1 reads a after the second assign
+      np.testing.assert_allclose(out, 440)
+
+  def test_two_reads_between_three_assigns(self):
+    a = Tensor.zeros(4).realize()
+    first = a + 100
+    a.assign(Tensor([1., 2., 0., 0.]))
+    second = a + 0
+    a.assign(a + 10)
+    with self.assertRaisesRegex(RuntimeError, "cycle"):  # TODO: broken now, raises
+      np.testing.assert_allclose((first + second + a).numpy(), [112, 114, 110, 110])
+
+  def test_read_before_slice_assign(self):
+    a = Tensor.ones(4).realize()
+    before = a * 3
+    a[0:2].assign(Tensor.full((2,), 2.))
+    out = (a + (before - 1)).numpy()
+    try:
+      np.testing.assert_allclose(out, [4, 4, 3, 3])
+    except AssertionError:
+      # TODO: broken now, before reads the two assigned elements after the assign
+      np.testing.assert_allclose(out, [7, 7, 3, 3])
+
+  def test_read_before_assign_survives_a_realize(self):
+    a = Tensor.ones(4).realize()
+    before = a * 3
+    a.assign(Tensor.full((4,), 5.))
+    a.realize()
+    out = before.numpy()
+    try:
+      np.testing.assert_allclose(out, 3)
+    except AssertionError:
+      # TODO: broken now, before is computed again from the assigned value
+      np.testing.assert_allclose(out, 15)
+
+  def test_loss_read_after_step_is_the_pre_step_loss(self):
+    from tinygrad import nn
+    w = Tensor([2.]).contiguous().realize()
+    x = Tensor([3.]).realize()
+    opt = nn.optim.SGD([w], lr=0.1)
+    with Context(TRAINING=1):
+      loss = (w*x).sum()                               # 6.0
+      loss.backward()
+      opt.step()                                       # w becomes 1.7
+    out = loss.item()
+    try:
+      self.assertAlmostEqual(out, 6.0, places=5)
+    except AssertionError:
+      # TODO: broken now, loss is computed again from the updated weight
+      self.assertAlmostEqual(out, 5.1, places=5)
+
+  def test_rand_realized_out_of_order(self):
+    Tensor.manual_seed(1)
+    r = [Tensor.rand(4) for _ in range(4)]
+    r[3].realize()
+    out_of_order = r[0].numpy()
+    Tensor.manual_seed(1)
+    in_order = [Tensor.rand(4).numpy() for _ in range(4)]
+    try:
+      np.testing.assert_equal(out_of_order, in_order[0])
+    except AssertionError:
+      # TODO: broken now, r[0] returns the fourth set of numbers
+      np.testing.assert_equal(out_of_order, in_order[3])
+
+  def test_batchnorm_stats_are_realized(self):
+    from tinygrad import nn
+    bn, x = nn.BatchNorm(4), Tensor.randn(2, 4, 3, 3).realize()
+    with Context(TRAINING=1): bn(x).realize()
+    try:
+      self.assertTrue(bn.running_mean.uop.base.is_realized)
+    except AssertionError:
+      # TODO: broken now, the stat update is never run because nothing reads it
+      self.assertFalse(bn.running_mean.uop.base.is_realized)
+
+  def test_batchnorm_under_jit_counts_every_call(self):
+    from tinygrad import nn
+    bn, x = nn.BatchNorm(4), Tensor.randn(8, 4, 2, 2).realize()
+    @TinyJit
+    def step(t):
+      with Context(TRAINING=1): return bn(t).sum().realize()
+    for _ in range(4): step(x)
+    out = bn.num_batches_tracked.item()
+    try:
+      self.assertEqual(out, 4)
+    except AssertionError:
+      # TODO: broken now, only the calls whose stat update happened to be captured are counted
+      self.assertEqual(out, 2)
+
+  def test_assign_from_unrealized_tensor_does_not_alias(self):
+    a = Tensor.full((4,), 7.).realize()
+    b = Tensor.ones(4) * 1
+    b.assign(a)
+    b.assign(Tensor.zeros(4))
+    b.realize()
+    try:
+      self.assertListEqual(a.tolist(), [7., 7., 7., 7.])
+    except AssertionError:
+      # TODO: broken now, b shares a's buffer, so the second assign to b overwrites a
+      self.assertListEqual(a.tolist(), [0., 0., 0., 0.])
+
+  def test_assign_to_function_output(self):
+    from tinygrad import function
+    @function
+    def f(x:Tensor) -> Tensor: return x*2
+    out = f(Tensor.ones(4).realize())
+    with self.assertRaisesRegex(RuntimeError, "UOp verification failed"):  # TODO: broken now, raises
+      out.assign(Tensor.full((4,), 9.).realize())
+      self.assertListEqual(out.tolist(), [9., 9., 9., 9.])
+
+  def test_nested_function_assign(self):
+    from tinygrad import function
+    @function
+    def inner(x:Tensor) -> Tensor:
+      x.assign(x+1)
+      return x*2
+    @function
+    def outer(x:Tensor) -> Tensor:
+      y = inner(x)
+      x.assign(x+1)
+      return y+x
+    a = Tensor([1.]).realize()
+    out = outer(a).item()
+    try:
+      self.assertEqual([out, a.item()], [7., 3.])
+    except AssertionError:
+      # TODO: broken now, the inner assign is run twice
+      self.assertEqual([out, a.item()], [6., 4.])
+
 class TestAssignToUnrealizedView(unittest.TestCase):
   def test_copy(self):
     t = Tensor.zeros(2,2, dtype=dtypes.int).to("CPU:0").contiguous().realize()
@@ -1040,7 +1134,6 @@ class TestPartialAssignToSharedBuffer(unittest.TestCase):
     for v, s in zip(views, shapes):
       np.testing.assert_allclose(v.numpy(), np.ones(s))
 
-
 class TestAfterCachePatterns(unittest.TestCase):
   def test_double_store_after(self):
     a = Tensor.zeros(10).contiguous()
@@ -1071,13 +1164,84 @@ class TestAfterCachePatterns(unittest.TestCase):
     np.testing.assert_array_equal(head.numpy(), [3])
     np.testing.assert_array_equal(full.numpy(), [1, 2])
 
-class TestBatchNormRunningStats(unittest.TestCase):
-  @unittest.expectedFailure  # TODO: nothing reads the stat update so it is never scheduled, and the chain grows every step
-  def test_running_stats_are_realized(self):
-    from tinygrad import nn
-    bn, x = nn.BatchNorm(4), Tensor.randn(2, 4, 3, 3).contiguous().realize()
-    with Context(TRAINING=1): bn(x).realize()
-    self.assertTrue(bn.running_mean.uop.base.is_realized)
+class TestMultiAssign(unittest.TestCase):
+  device = tuple(f"{Device.DEFAULT}:{i}" for i in range(2))
+
+  @needs_second_gpu
+  def setUp(self): pass
+
+  def test_multi_assign_realized(self):
+    out = Tensor.zeros(4).shard(self.device, 0).contiguous().realize()
+    ones = Tensor.ones(4).shard(self.device, 0).contiguous().realize()
+    out.assign(ones).realize()
+    self.assertListEqual(out.tolist(), [1,1,1,1])
+
+  def test_multi_assign_unrealized(self):
+    out = Tensor.zeros(4).contiguous().realize().shard(self.device, 0)
+    ones = Tensor.ones(4).shard(self.device, 0).contiguous().realize()
+    out.assign(ones).realize()
+    self.assertListEqual(out.tolist(), [1,1,1,1])
+
+  def test_multi_assign_both_unrealized(self):
+    out = Tensor.zeros(4).contiguous().realize().shard(self.device, 0)
+    ones = Tensor.ones(4).contiguous().realize().shard(self.device, 0)
+    out.assign(ones).realize()
+    self.assertListEqual(out.tolist(), [1,1,1,1])
+
+  def test_multi_assign_scalar(self):
+    out = Tensor.ones(4).shard(self.device, 0).contiguous().realize()
+    out.assign(0).realize()
+    self.assertListEqual(out.tolist(), [0,0,0,0])
+
+  def test_multi_assign_const_like(self):
+    out = Tensor.ones(4).shard(self.device, 0).contiguous().realize()
+    out.assign(out.const_like(7)).realize()
+    self.assertListEqual(out.tolist(), [7,7,7,7])
+
+  def test_multi_assign_piece(self):
+    out = Tensor.zeros(4,4).shard(self.device, 0).contiguous().realize()
+    ones = Tensor.ones(4,1).shard(self.device, 0).contiguous().realize()
+    out[:, 2:3].assign(ones).realize()
+    self.assertListEqual(out.tolist(), [[0,0,1,0], [0,0,1,0], [0,0,1,0], [0,0,1,0]])
+
+  def test_multi_assign_piece_noncontig(self):
+    out = Tensor.zeros(4,4).contiguous().realize().shard(self.device, 0).realize()
+    ones = Tensor.ones(4,1).shard(self.device, 0).contiguous().realize()
+    out[:, 2:3].assign(ones).realize()
+    self.assertListEqual(out.tolist(), [[0,0,1,0], [0,0,1,0], [0,0,1,0], [0,0,1,0]])
+
+  def test_multi_assign_piece_unrealized(self):
+    out = Tensor.zeros(4,4).contiguous().realize().shard(self.device, 0)
+    ones = Tensor.ones(4,1).shard(self.device, 0).contiguous().realize()
+    out[:, 2:3].assign(ones).realize()
+    try:
+      self.assertListEqual(out.tolist(), [[0,0,1,0], [0,0,1,0], [0,0,1,0], [0,0,1,0]])
+    except AssertionError:
+      # TODO: broken now, the write is dropped
+      self.assertListEqual(out.tolist(), [[0,0,0,0], [0,0,0,0], [0,0,0,0], [0,0,0,0]])
+
+  def test_multi_assign_var_offset(self):
+    out = Tensor.zeros(4,4).contiguous().realize().shard(self.device, 0).realize()
+    ones = Tensor.ones(4,1).shard(self.device, 0).contiguous().realize()
+    vi = Variable("i", 0, 3).bind(2)
+    out[:, vi:vi+1].assign(ones).realize()
+    self.assertListEqual(out.tolist(), [[0,0,1,0], [0,0,1,0], [0,0,1,0], [0,0,1,0]])
+
+  def test_multi_assign_var_offset_jit_none(self): self.test_multi_assign_var_offset_jit(None)
+  def test_multi_assign_var_offset_jit(self, shard_axis=0):
+    out = Tensor.zeros(4,6).contiguous().realize().shard(self.device, shard_axis).realize()
+    ones = Tensor.ones(4,1).shard(self.device, shard_axis).contiguous().realize()
+
+    @TinyJit
+    def f(out:Tensor, vi):
+      out[:, vi:vi+1].assign(ones).realize()
+      ones.assign(ones+1).realize()
+
+    vi = Variable("i", 0, 5)
+    for i in range(1,5):
+      GlobalCounters.reset()
+      f(out, vi.bind(i))
+    self.assertListEqual(out.tolist(), [[0,1,2,3,4,0]]*4)
 
 if __name__ == "__main__":
   unittest.main()

test/backend/test_call.py

@@ -5,12 +5,13 @@ from tinygrad.dtype import dtypes
 from tinygrad.renderer.cstyle import CStyleLanguage
 from tinygrad.uop.ops import KernelInfo
 
+# an external call is a CALL on a CUSTOM_FUNCTION body holding the callee (the loaded function pointer)
 def call_out_kernel(F:UOp, C:UOp) -> UOp:
-  call = F[0].load().call(UOp.const(3).cast(dtypes.int), C[0], ret_dtype=dtypes.void)
+  call = UOp.custom_function("callback", F[0].load()).call(UOp.const(3).cast(dtypes.int), C[0], ret_dtype=dtypes.void)
   return C.after(call)[1].store(C.after(call)[0].load() + 1).sink(arg=KernelInfo(name="call_out"))
 
 def call_ret_kernel(F:UOp, C:UOp) -> UOp:
-  val = F[0].load().call(UOp.const(21).cast(dtypes.int), ret_dtype=dtypes.int)
+  val = UOp.custom_function("callback", F[0].load()).call(UOp.const(21).cast(dtypes.int), ret_dtype=dtypes.int)
   return C[0].store(val * 2).sink(arg=KernelInfo(name="call_ret"))
 
 @unittest.skipUnless(isinstance(Device["CPU"].renderer, CStyleLanguage), "TODO: CALL is rendered in C style only")

test/backend/test_custom_kernel.py

@@ -347,7 +347,7 @@ class TestCustomKernel(unittest.TestCase):
     self.assertTrue((c == 2).all().item())
 
   def test_partial_invalid_store_keeps_uncovered_reads(self):
-    x = Tensor([10., 20., 30., 40.])
+    x = Tensor([10., 20., 30., 40.]).realize()
     after = x.uop.after(x.uop.shrink(((0, 2),)).store(Invalid))
     self.assertEqual(Tensor(after).contiguous().tolist(), [10., 20., 30., 40.])
 

test/backend/test_dtype.py

@@ -258,6 +258,11 @@ class TestDoubleDType(TestDType):
       a = [2, 3, 4]
       np.testing.assert_allclose(func(Tensor(a, dtype=self.DTYPE)).numpy(), func(torch.tensor(a, dtype=torch.float64)), rtol=1e-12, atol=1e-12)
 
+  def test_float32_compare_selecting_float64(self):
+    a = Tensor([1.0, 2.0, 5.0, 9.0], dtype=dtypes.float32)
+    p, q = Tensor([10., 20., 30., 40.], dtype=self.DTYPE), Tensor([50., 60., 70., 80.], dtype=self.DTYPE)
+    _test_op(lambda: (a < 3.0).where(p, q), self.DTYPE, [10., 20., 70., 80.])
+
   def test_float64_to_float32_cast_inf(self):
     _test_op(lambda: Tensor([3.4e40, 3.4e38, 1, 0], dtype=dtypes.float64).cast(dtypes.float32),
              dtypes.float32, [float('inf'), 3.4e38, 1, 0])
@@ -323,7 +328,10 @@ class TestUint16DType(TestDType):
 class TestInt32DType(TestDType): DTYPE = dtypes.int32
 class TestUint32DType(TestDType): DTYPE = dtypes.uint32
 
-class TestInt64DType(TestDType): DTYPE = dtypes.int64
+class TestInt64DType(TestDType):
+  DTYPE = dtypes.int64
+  def test_int64_to_uint32_to_int64(self):
+    _test_op(lambda: Tensor([0x12345678ABCDEF01], dtype=dtypes.int64).cast(dtypes.uint32).cast(dtypes.int64), dtypes.int64, [2882400001])
 
 @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
 class TestEmulatedInt64DType(TestInt64DType):
@@ -423,6 +431,11 @@ class TestDtypeUsage(unittest.TestCase):
         t = Tensor([[1, 2], [3, 4]], dtype=d)
         (t*t).max().item()
 
+  def test_where_float16_compare_to_const(self):
+    # t > 0 is CMPLT(0, t): the float16 operand is on the right
+    t = Tensor([-1.0, 1.0], dtype=dtypes.float16)
+    np.testing.assert_equal((t > 0).where(Tensor.ones(2, dtype=dtypes.float16), Tensor.zeros(2, dtype=dtypes.float16)).numpy(), [0.0, 1.0])
+
 @unittest.skipUnless(dtypes.bfloat16 in supported_dtypes, f"no bfloat16 on {Device.DEFAULT}")
 class TestOpsBFloat16(unittest.TestCase):
   def test_cast(self):

test/backend/test_encodings.py

@@ -4,7 +4,7 @@ from tinygrad.uop.ops import UOp, Ops
 from tinygrad.dtype import dtypes
 from tinygrad.renderer.isa.x86 import X86Ops, X86Renderer, RBP, RDI, RSP, RSI, RAX, RDX, XMM, GPR, imm, def_reg
 
-def ins(op, dt, src, tag=None): return UOp(Ops.INS, arg=op, dtype=dt, src=src, tag=tag)
+def ins(op, dt, src, tag=None): return UOp(Ops.INS, arg=(op, dt), src=src, tag=tag)
 
 @unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "only on x86")
 class TestEncodingsX86(unittest.TestCase):
@@ -100,13 +100,6 @@ class TestEncodingsX86(unittest.TestCase):
     # vaddss xmm0, xmm0, xmm8
     self.assertEqual(bytes.fromhex(self.encode(add)), bytes.fromhex("C4 C1 7A 58 C0"))
 
-  # test ymm encoding
-  def test_ymm_encoding(self):
-    xmm0, xmm1 = def_reg(dtypes._uint256, XMM[0]), def_reg(dtypes._uint256, XMM[1])
-    add = ins(X86Ops.VADDPS, dtypes._uint256, (xmm0, xmm1), XMM[0])
-    # vaddps ymm0, ymm0, ymm1
-    self.assertEqual(bytes.fromhex(self.encode(add)), bytes.fromhex("C5 FC 58 C1"))
-
   # test encoding where register is in the immediate field
   def test_reg_in_imm_field(self):
     xmm0, xmm1, xmm2 = def_reg(dtypes.float32, XMM[0]), def_reg(dtypes.float32, XMM[1]), def_reg(dtypes.float32, XMM[2])
@@ -143,7 +136,7 @@ class TestEncodingsX86(unittest.TestCase):
 
   # cmoves have the cmp as the last src even though it is not explicitly used, the cmp doesn't define a reg and is ignored in the encoding
   def test_cmove_ignore_cmp(self):
-    cmove = ins(X86Ops.CMOVE, dtypes.int32, (def_reg(dtypes.int32, RAX), UOp(Ops.INS, arg=X86Ops.CMP)), RDX)
+    cmove = ins(X86Ops.CMOVE, dtypes.int32, (def_reg(dtypes.int32, RAX), UOp(Ops.INS, arg=(X86Ops.CMP, dtypes.void))), RDX)
     # cmove edx, eax
     self.assertEqual(bytes.fromhex(self.encode(cmove)), bytes.fromhex("0F 44 D0"))
 

test/backend/test_graph.py

@@ -7,7 +7,7 @@ from tinygrad.helpers import Context
 from tinygrad.dtype import dtypes
 from tinygrad.engine.jit import MultiGraphRunner
 from tinygrad.engine.realize import run_linear, compile_linear
-from tinygrad.uop.ops import UOp, Ops, buffers
+from tinygrad.uop.ops import UOp, Ops
 
 from test.helpers import needs_second_gpu
 
@@ -39,8 +39,7 @@ def make_view(base, offset_elems, size_elems):
 
 def get_buf_uop(buf:Buffer, cache:dict[Buffer,UOp]) -> UOp:
   if buf not in cache:
-    cache[buf] = u = UOp.new_buffer(buf.device, buf.size, buf.dtype)
-    buffers[u] = buf
+    cache[buf] = UOp.from_buffer(buf)
   return cache[buf]
 
 def copy_call(dst:Buffer, src:Buffer, c:dict[Buffer,UOp]) -> UOp:

test/backend/test_isel.py

@@ -7,7 +7,7 @@ from tinygrad.renderer.isa.x86 import X86Renderer, X86Ops
 from tinygrad.renderer.isa import IselContext
 
 # INDEX on a register value with a constant index extracts a single element (the old GEP)
-def lane(y:UOp, i:int) -> UOp: return y.index(UOp.cconst(i, dtypes.int), dtype=y.dtype)
+def lane(y:UOp, i:int) -> UOp: return y.index(UOp.cconst(i, dtypes.int))
 
 @unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "only x86")
 class TestIselX86(unittest.TestCase):
@@ -20,7 +20,7 @@ class TestIselX86(unittest.TestCase):
       with self.subTest(dtype=dt):
         v = [UOp.variable(str(i), 0, 0, dt) for i in range(nargs)]
         n = self.isel_rewrite(expr(*v))
-        self.assertIs(n.arg, op)
+        self.assertIs(n.arg[0], op)
 
   def test_cmove(self):
     a = UOp.variable("a", 0, 0, dtypes.int32)
@@ -29,9 +29,9 @@ class TestIselX86(unittest.TestCase):
     d = (a != b).where(a, b)
     f = c + d
     n = self.isel_rewrite(f)
-    self.assertTrue(n.src[0].arg is X86Ops.CMOVL and n.src[1].arg is X86Ops.CMOVNE)
+    self.assertTrue(n.src[0].arg[0] is X86Ops.CMOVL and n.src[1].arg[0] is X86Ops.CMOVNE)
     # both comparisons become the same instruction
-    self.assertTrue(n.src[0].src[2] == n.src[1].src[2] and n.src[0].src[2].arg is X86Ops.CMP)
+    self.assertTrue(n.src[0].src[2] == n.src[1].src[2] and n.src[0].src[2].arg[0] is X86Ops.CMP)
 
   def test_vinsertps(self):
     a = UOp.variable("a", 0, 0, dtypes.float32)
@@ -41,12 +41,12 @@ class TestIselX86(unittest.TestCase):
 
     valid = [UOp.stack(lane(a, 0), lane(b, 1), lane(a, 2), lane(b, 3)),
              UOp.stack(lane(a, 3), lane(b, 2), lane(c, 1), d)]
-    for shuf in valid: self.assertIs(self.isel_rewrite(shuf).arg, X86Ops.VINSERTPS)
+    for shuf in valid: self.assertIs(self.isel_rewrite(shuf).arg[0], X86Ops.VINSERTPS)
 
   # complex address is [base + index*scale + displacement]
   def test_complex_address(self):
     a = UOp.variable("a", 0, 0, dtypes.int32)
-    load = UOp.param(0, dtypes.int32, (16,)).index(a + UOp.cconst(1, dtypes.int32)).load()
+    load = UOp.param(0, dtypes.int32, 16).index(a + UOp.cconst(1, dtypes.int32)).load()
     n = self.isel_rewrite(load)
     # displacement is the constant in "a" scaled to the buffer element size, dtype is int8 when the value fits otherwise int32
     self.assertTrue(n.src[2].dtype is dtypes.int8 and n.src[2].src[0].op is Ops.CONST and n.src[2].src[0].val == 4)

test/backend/test_jit.py

@@ -4,9 +4,10 @@ import numpy as np
 
 from test.helpers import assert_jit_cache_len, call_is_graph, not_support_multi_device, needs_second_gpu, KernelCountException
 from test.unit.test_jit import _simple_test
-from tinygrad import Tensor, Variable, TinyJit, Device, dtypes
+from tinygrad import Tensor, TinyJit, Device, dtypes
 from tinygrad.engine.jit import graph_class
 from tinygrad.helpers import JIT, DEV, GlobalCounters
+from tinygrad.runtime.support.hcq2 import HCQ_DEVS
 from tinygrad.uop.ops import Ops
 from tinygrad.renderer.isa.x86 import X86Renderer
 
@@ -16,19 +17,6 @@ class TestJit(unittest.TestCase):
     def add(a, b): return (a+b).realize()
     _simple_test(add)
 
-  @unittest.skipUnless(Device.DEFAULT == "CPU", "core_id is a CPU runtimevar")
-  def test_hcq_core_id_runtimevar_merge(self):
-    N = 262144
-    @TinyJit
-    def f(x, st):
-      y = (x + 1).contiguous().realize()
-      z = x.shrink(((st, st + N),)).contiguous().realize()
-      return y, z
-    x = Tensor.arange(2*N).clone().realize()
-    for _ in range(3): y, z = f(x, Variable("a", 0, N).bind(0))
-    self.assertEqual(y.shape, (2*N,))
-    self.assertEqual(z.shape, (N,))
-
   def test_jit_input_view(self):
     @TinyJit
     def f(x): return (x[2:5].contiguous() + 1).realize()
@@ -235,6 +223,7 @@ class TestJitPrune(unittest.TestCase):
     assert_jit_cache_len(w2_prune, 1)
 
 class TestJitFree(unittest.TestCase):
+  @unittest.skipIf(Device.DEFAULT.split(":")[0] in HCQ_DEVS - {"CPU"}, "hcq2 keeps refs to intermediate buffers")
   def test_free_intermediates(self):
     ext_tensor = Tensor([1,24,23,45,1])
     @TinyJit

test/backend/test_linearizer.py

@@ -2,12 +2,12 @@ import numpy as np
 import unittest
 
 from tinygrad.codegen.opt import Opt, OptOps
-from tinygrad.uop.ops import UOp, Ops, GroupOp, AxisType, buffers
+from tinygrad.uop.ops import UOp, Ops, GroupOp, AxisType
 from tinygrad.device import Device, Buffer
 from tinygrad.tensor import Tensor, _to_np_dtype
 from tinygrad.engine.realize import run_linear
 from tinygrad.codegen import to_program
-from tinygrad.helpers import Context, flatten, dedup, TC_SELECT, TC_OPT, DEV
+from tinygrad.helpers import Context, dedup, TC_SELECT, TC_OPT, DEV
 from tinygrad.dtype import DType, dtypes, AddrSpace
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.cstyle import CUDARenderer
@@ -73,14 +73,6 @@ class TestLinearizer(unittest.TestCase):
       # assert that there is a global load after the reduce ends
       assert any(u.addrspace == AddrSpace.GLOBAL for u in load_idxs)
 
-  def _test_no_nested_ranges(self, lins, skip=None):
-    for l in lins:
-      range_in_acc = flatten([[x for x in u.src if x.op is Ops.RANGE] for u in l.uops if u.op is Ops.BUFFER and u.addrspace is AddrSpace.REG])
-      ranges = [u.op for u in l.uops if (u.op is Ops.RANGE and u in range_in_acc) or (u.op is Ops.END and u.src[0] in range_in_acc)]
-      for i,u in enumerate(ranges):
-        if skip and i in skip: continue
-        assert ranges[i-1] != u, f"multireduce nested the ranges! {ranges[i-1], {u}}"
-
   def test_two_nested_range(self):
     a = Tensor.randn(2, ).realize()
     out = a.reshape(2, 1).expand(2, 3).sum()
@@ -135,7 +127,7 @@ class TestLinearizer(unittest.TestCase):
     # these are of size 3 to avoid float4 coalesce
     r = a[:-1] + a[1:]
 
-    uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=0)]),
+    uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.SPLIT, axis=0, arg=(0, AxisType.UPCAST))]),
                        renderer=Device[Device.DEFAULT].renderer).src[1].src)
     num_loads = len([uop for uop in uops if uop.op is Ops.LOAD])
     assert num_loads <= 4, "more load uops than needed"
@@ -148,7 +140,7 @@ class TestLinearizer(unittest.TestCase):
     a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize()
     r = a.expand([2]) + b.expand([2])
 
-    uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=0)]),
+    uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.SPLIT, axis=0, arg=(0, AxisType.UPCAST))]),
                        renderer=Device[Device.DEFAULT].renderer).src[1].src)
     num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU])
     assert num_ops <= 1, "more alu uops than needed"
@@ -159,7 +151,8 @@ class TestLinearizer(unittest.TestCase):
     r = Tensor.conv2d(x,w,padding=1).relu()
 
     uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0],
-      [Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]), renderer=Device[Device.DEFAULT].renderer).src[1].src)
+      [Opt(op=OptOps.SPLIT, axis=0, arg=(0, AxisType.UPCAST)),
+       Opt(op=OptOps.SPLIT, axis=1, arg=(0, AxisType.UNROLL))]), renderer=Device[Device.DEFAULT].renderer).src[1].src)
     accs = [u for u in uops if u.op is Ops.BUFFER and u.addrspace is AddrSpace.REG]
     stores = [u for u in uops if u.op is Ops.STORE]
     assert len(accs) == 0  # it's removed now
@@ -170,7 +163,7 @@ class TestLinearizer(unittest.TestCase):
   @unittest.skipUnless(Device.DEFAULT == "CPU", "test only for CPU")
   def test_upcast_with_locals_cpu(self):
     out = Tensor.ones(64,64).contiguous() @ Tensor.ones(64,64).contiguous()
-    prg = to_program(replace_opts(out.schedule_linear().src[-1].src[0], [Opt(OptOps.LOCAL, axis=0, arg=4)]),
+    prg = to_program(replace_opts(out.schedule_linear().src[-1].src[0], [Opt(OptOps.SPLIT, axis=0, arg=(4, AxisType.LOCAL))]),
                       renderer=Device[Device.DEFAULT].renderer)
     self.assertEqual(len(prg.src[2].arg.split("for")), 5)
 
@@ -181,7 +174,8 @@ class TestLinearizer(unittest.TestCase):
   def test_upcast_with_locals(self):
     x, y = Tensor.rand(1,128), Tensor.rand(128, 128)
     r = (x@y).relu()
-    opts_to_apply = [Opt(op=OptOps.GROUP, axis=0, arg=8), Opt(op=OptOps.LOCAL, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=4)]
+    opts_to_apply = [Opt(op=OptOps.SPLIT, axis=1, arg=(8, AxisType.GROUP_REDUCE)), Opt(op=OptOps.SPLIT, axis=0, arg=(4, AxisType.LOCAL)),
+                     Opt(op=OptOps.SPLIT, axis=0, arg=(4, AxisType.UPCAST))]
     program = to_program(replace_opts(r.schedule_linear().src[-1].src[0], opts_to_apply), renderer=Device[Device.DEFAULT].renderer)
 
     stores = [u for u in tuple(program.src[1].src) if u.op is Ops.STORE and u.src[0].addrspace != AddrSpace.REG]
@@ -197,7 +191,7 @@ class TestLinearizer(unittest.TestCase):
   def test_zero_fold(self):
     a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize()
     r = Tensor.stack(a, b)
-    uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=0)]),
+    uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.SPLIT, axis=0, arg=(0, AxisType.UPCAST))]),
                        renderer=Device[Device.DEFAULT].renderer).src[1].src)
     num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU])
     assert num_ops == 0, "more alu uops than needed"
@@ -228,7 +222,7 @@ class TestLinearizer(unittest.TestCase):
       (dtypes.float, dtypes.float16, dtypes.float16),
     )
     for tensor_dtype, acc_dtype, expected_dtype in tests:
-      if tensor_dtype in (dts:=Device[Device.DEFAULT].renderer.supported_dtypes()) and acc_dtype in dts and expected_dtype in dts:
+      if tensor_dtype in (dts:=Device[Device.DEFAULT].renderer.supported_dtypes()) and acc_dtype in dts|{None} and expected_dtype in dts:
         a, b = Tensor.rand(8, 8, dtype=tensor_dtype), Tensor.rand(8, 8, dtype=tensor_dtype)
         helper_arg_acc_dtype(a.sum(dtype=acc_dtype), expected_dtype)
         helper_arg_acc_dtype(a.matmul(b, dtype=acc_dtype), expected_dtype)
@@ -240,7 +234,7 @@ class TestLinearizer(unittest.TestCase):
   def test_simple_unroll_no_between_phi_dependencies(self):
     x, y = Tensor.empty(64, 64), Tensor.empty(64, 64)
     r = (x@y).relu()
-    opt = [Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 4)]
+    opt = [Opt(OptOps.SPLIT, 2, (4, AxisType.UNROLL)), Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST))]
     ast = helper_linearizer_opt(r, [opt])
     # the uops graph is reg BUFFER -> 4x STORE 0.0 -> RANGE -> 4x ALU -> 4x STORE -> ENDRANGE
     uops = tuple(to_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).src[1].src)
@@ -253,9 +247,6 @@ class TestLinearizer(unittest.TestCase):
         else:
           assert u.src[1].op in GroupOp.ALU
           assert begin_range < uops.index(u) < end_range
-      # children of END are placed after ENDRANGE
-      if any(x.op is Ops.END and x.src[1].op in GroupOp.ALU for x in u.src):
-        assert end_range < uops.index(u)
 
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
   def test_default_global_reversed(self):
@@ -353,8 +344,9 @@ class TestLinearizer(unittest.TestCase):
   def test_grouped_store_locals_and_globals(self):
     x, y = Tensor.empty(64, 64), Tensor.empty(64, 64)
     out = x@y
-    opt = [Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.GROUPTOP, 0, 8),
-            Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 2)] # upcast accs in both reduces
+    opt = [Opt(OptOps.SPLIT, 0, (4, AxisType.LOCAL)), Opt(OptOps.SPLIT, 3, (8, AxisType.GROUP_REDUCE, True)),
+            Opt(OptOps.SPLIT, 3, (4, AxisType.UNROLL)), Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST)),
+            Opt(OptOps.SPLIT, 1, (2, AxisType.UPCAST))] # upcast accs in both reduces
     ast = helper_linearizer_opt(out, opts=[opt])
     def get_recursive(uop): return set.union(set(uop.src), [uop], *[get_recursive(v) for v in uop.src])
     uops = tuple(to_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).src[1].src)
@@ -392,9 +384,9 @@ class TestLinearizer(unittest.TestCase):
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
   def test_two_grouped_stores_local(self):
-    # GROUP on both reduces puts two LOCAL buffers in one kernel, and the store to each needs its own barrier
+    # GROUP_REDUCE on both reduces puts two LOCAL buffers in one kernel, and the store to each needs its own barrier
     a = Tensor.rand(32, 32).realize()
-    opts = [Opt(OptOps.GROUP, 1, 4), Opt(OptOps.GROUP, 2, 4)]
+    opts = [Opt(OptOps.SPLIT, 3, (4, AxisType.GROUP_REDUCE)), Opt(OptOps.SPLIT, 5, (4, AxisType.GROUP_REDUCE))]
     ast = helper_linearizer_opt(single_kernel_softmax(a), [opts])
     uops = to_program(replace_opts(ast, opts), renderer=Device[Device.DEFAULT].renderer).src[1].src
     self.assertEqual(len([u for u in uops if u.op is Ops.BARRIER]), 2)
@@ -416,12 +408,6 @@ def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]:
   for b in bufs: b.ensure_allocated()
   return ast, bufs
 
-def helper_linearizer_ast(ast:UOp, inputs:list[Tensor], *args, **kwargs):
-  assert isinstance(ast, UOp), "ast must be UOp"
-  inbufs = [x.uop.base.buffer for x in inputs]
-  outbufs = [Buffer(inbufs[-1].device if inbufs else Device.DEFAULT, out.size, out.src[1].dtype).allocate() for out in ast.src]
-  _helper_linearizer_opt_ast(ast, outbufs+inbufs, *args, **kwargs)
-
 def helper_linearizer_opt(r:Tensor|list[Tensor], *args, **kwargs):
   realized_ast, real_bufs = helper_realized_ast(r)
   _helper_linearizer_opt_ast(realized_ast, real_bufs, *args, **kwargs)
@@ -437,8 +423,7 @@ def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[]
                                apply_tc=False, atol=1e-4, rtol=1e-4, color_sizes=[], wanna_output=[], check_default_opt=True):
   outbufs = real_bufs[:len(realized_ast.src)]
   wanna_output = [np.array(x).flatten() for x in wanna_output]
-  buf_uops = [UOp.new_buffer(b.device, b.size, b.dtype) for b in real_bufs]
-  for u,b in zip(buf_uops, real_bufs): buffers[u] = b
+  buf_uops = [UOp.from_buffer(b) for b in real_bufs]
 
   def run_prg(opts):
     ast = realized_ast if opts is None else replace_opts(realized_ast, list(opts))