fix and spec

remove FUNCTION from tinyspec (regen pdf) and viz
match the TUPLE body in the UPat, not in the rewrite
2026-08-28 22:36:07 +00:00 · 2026-08-28 15:19:06 -07:00 · 2026-08-28 15:06:21 -07:00 · 2026-08-28 14:34:13 -07:00 · 2026-08-28 14:30:38 -07:00 · 2026-08-28 14:16:33 -07:00
176 changed files with 8637 additions and 2075 deletions
@@ -194,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
@@ -240,10 +248,10 @@ runs:
      if: inputs.amd == 'true' && runner.os == 'macOS'
      shell: bash
      run: |
-        sudo mkdir -p /usr/local/lib
-        curl -s -H "Authorization: token $GH_TOKEN" curl -s https://api.github.com/repos/tinygrad/amdcomgr_dylib/releases/latest | \
-          jq -r '.assets[] | select(.name == "libamd_comgr.dylib").browser_download_url' | \
-          sudo xargs curl -fL -o /usr/local/lib/libamd_comgr.dylib
+        sudo "$VIRTUAL_ENV/bin/python" -c "
+          from tinygrad.helpers import fetch
+          fetch('https://github.com/tinygrad/amdcomgr_dylib/releases/download/v7.2.0/libamd_comgr.dylib', name='/usr/local/lib/libamd_comgr.dylib',
+                sha256='7712fbe4fcb9fcdea49aeac989876448df975ce0a8ce7c9b15b55c15e7a05935').chmod(0o644)"

    # **** CUDA ****
    - name: Install CUDA
@@ -261,8 +269,11 @@ runs:
      if: inputs.ocelot == 'true'
      shell: bash
      run: |
-        sudo mkdir -p /usr/local/lib
-        sudo curl --output-dir /usr/local/lib -fLO https://github.com/tinygrad/gpuocelot/releases/download/v0.1.0/libgpuocelot.${{ runner.os == 'Linux' && 'so' || 'dylib' }}
+        sudo "$VIRTUAL_ENV/bin/python" -c "
+          from tinygrad.helpers import fetch
+          fetch('https://github.com/tinygrad/gpuocelot/releases/download/v0.1.0/libgpuocelot.${{ runner.os == 'Linux' && 'so' || 'dylib' }}',
+                name='/usr/local/lib/libgpuocelot.${{ runner.os == 'Linux' && 'so' || 'dylib' }}',
+                sha256='${{ runner.os == 'Linux' && 'a24705276a9a187111371465987b3258f8836ef512a34266e3075bc4714e125a' || '5106c998c795a36dec79eb7b2aae324a93d1338236d36eeaae232649ec457663' }}').chmod(0o644)"

    # **** WebGPU ****

@@ -270,8 +281,11 @@ runs:
      if: inputs.webgpu == 'true'
      shell: bash
      run: |
-        sudo mkdir -p /usr/local/lib
-        sudo curl --output-dir /usr/local/lib -fLO https://github.com/wpmed92/pydawn/releases/download/v0.1.6/libwebgpu_dawn.${{ runner.os == 'Linux' && 'so' || 'dylib' }}
+        sudo "$VIRTUAL_ENV/bin/python" -c "
+          from tinygrad.helpers import fetch
+          fetch('https://github.com/wpmed92/pydawn/releases/download/v0.1.6/libwebgpu_dawn.${{ runner.os == 'Linux' && 'so' || 'dylib' }}',
+                name='/usr/local/lib/libwebgpu_dawn.${{ runner.os == 'Linux' && 'so' || 'dylib' }}',
+                sha256='${{ runner.os == 'Linux' && 'cf36091d266a32c9d5080f14662de44cece241987939713282ea0ff558db81c6' || '7e87c7acefda8b6af1a1c5debfedcf62958311284b8fd8d9bcf93e312e6636e3' }}').chmod(0o644)"

    # **** LLVM ****

@@ -285,7 +299,10 @@ runs:
      if: inputs.opencl == 'true'
      shell: bash
      run: |
-        sudo curl -fL https://github.com/sirhcm/tinymesa/releases/download/rusticl-v1/libRusticlOpenCL.so.1.0.0 -o /usr/lib/libRusticlOpenCL.so
+        sudo "$VIRTUAL_ENV/bin/python" -c "
+          from tinygrad.helpers import fetch
+          fetch('https://github.com/sirhcm/tinymesa/releases/download/rusticl-v1/libRusticlOpenCL.so.1.0.0', name='/usr/lib/libRusticlOpenCL.so',
+                sha256='d4f48566d8fd33f6cdd8ef6de35a71966e8a8517e6f68ff3c52dbb43765a2513').chmod(0o644)"
        sudo mkdir -p /etc/OpenCL/vendors
        echo "/usr/lib/libRusticlOpenCL.so" | sudo tee /etc/OpenCL/vendors/rusticl.icd
        echo "RUSTICL_ENABLE=llvmpipe" >> "$GITHUB_ENV"
@@ -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
@@ -94,7 +94,7 @@ jobs:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
-      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
+      HCQ2: '0'
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -102,12 +102,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: |
@@ -141,7 +140,7 @@ jobs:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
-      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
+      HCQ2: '0'
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -149,12 +148,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: |
@@ -190,7 +188,7 @@ jobs:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
-      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
+      HCQ2: '0'
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -198,12 +196,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
@@ -233,7 +230,7 @@ jobs:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
-      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
+      HCQ2: '0'
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -241,12 +238,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: |
@@ -279,7 +275,7 @@ jobs:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
-      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
+      HCQ2: '0'
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -287,12 +283,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
@@ -340,12 +335,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: |
@@ -437,32 +431,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)
+  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: 17
+          - version: '0.11.0'
+            model: policy
+            url: https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_policy.onnx
+            timing: 3.2
+          - version: '0.11.0'
+            model: dmonitoring
+            url: https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/dmonitoring_model.onnx
+            timing: 11
+          - version: '0.11.2'
+            model: supercombo
+            url: https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/433f85f956837606ad1f1cbee4aa7e2158ad23c768dea914b20436c97232741b
+            timing: 26
+          - dev: QCOM:IR3
+            version: '0.11.2'
+            model: supercombo
+            timing: 41
+          - version: '0.11.2'
+            model: dmonitoring
+            url: https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/3e7b31dfbc0a5234f1baf196513b77fc6af12204b8a8ffe8ee0417e48352f316
+            timing: 11
+          # IR3 dmonitoring is slightly slower
+          - dev: QCOM:IR3
+            model: dmonitoring
+            timing: 12
+      fail-fast: false
+    name: openpilot ${{ matrix.version }} compile3 ${{ matrix.model }} (DEV=${{ matrix.dev }})
    runs-on: [self-hosted, Linux, comma]
-    timeout-minutes: 12
+    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 +505,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 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py ${{ matrix.url }}
+    - name: run pickle
+      run: BENCHMARK_LOG="${BENCHMARK_LOG}_run_pickle" RUN_PICKLE=1 taskset -c 4-7 python3 examples/openpilot/compile3.py
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

@@ -588,9 +585,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
@@ -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 ******
@@ -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
@@ -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:
@@ -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
@@ -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,7 +1769,7 @@ 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)

@@ -1784,22 +1783,18 @@ def train_gptoss():
    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)
@@ -1848,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
@@ -1881,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())))

@@ -1891,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,
@@ -282,14 +282,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
@@ -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)
@@ -44,7 +44,7 @@ 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:-1} BENCHMARK=${BENCHMARK:-10}
 if [ -z "$FULL_LAYERS" ]; then
@@ -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
@@ -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

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

@@ -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}'))

@@ -20,34 +20,39 @@ def v_mfma_fp4(dst, a, b, opsel, opsel_hi, scale_a, scale_b):
 def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
  k = Kernel()
  scale_k = K // 32
+  k.emit(s_and_b32(s[1], s[1], LIT, 65535))
  if (tile_m, tile_n) == (128, 512):
-    k.emit(s_and_b32(s[1], s[1], LIT, 65535))
    k.emit(s_mov_b32(s[47], s[2]))
    k.emit(s_mov_b32(s[48], s[3]))
-    k.emit(s_load_dwordx2(s[4:5], s[0:1], s[0], 0, 0, 0, 0, 1))
-    k.emit(s_mov_b32(s[8], 0))
-    k.emit(s_mov_b32(s[9], 0))
-    k.emit(s_load_dwordx2(s[12:13], s[0:1], s[0], 8, 0, 0, 0, 1))
-    k.emit(s_load_dwordx2(s[16:17], s[0:1], s[0], 16, 0, 0, 0, 1))
-    k.emit(s_mov_b32(s[36], N))
-    k.emit(s_mov_b32(s[37], K))
-    k.emit(s_mov_b32(s[38], K))
-    k.emit(s_mov_b32(s[43], M))
-    k.emit(s_mov_b32(s[44], N))
-    k.emit(s_mov_b32(s[45], K))
-    k.emit(s_load_dwordx2(s[20:21], s[0:1], s[0], 24, 0, 0, 0, 1))
-    k.emit(s_load_dwordx2(s[24:25], s[0:1], s[0], 32, 0, 0, 0, 1))
-    k.emit(s_mov_b32(s[39], scale_k))
-    k.emit(s_mov_b32(s[40], scale_k))
-    k.emit(v_lshrrev_b32_e32(v[1], 10))
-    k.emit(v_lshrrev_b32_e32(v[2], 10, v[1]))
-    k.emit(v_and_b32_e32(v[2], LIT, v[2], 1023))
-    k.emit(v_and_b32_e32(v[1], LIT, v[1], 1023))
-    k.emit(v_and_b32_e32(v[0], LIT, v[0], 1023))
-    k.emit(v_lshrrev_b32_e32(v[3], 6))
-    k.emit(v_and_b32_e32(v[0], 63))
-    k.emit(v_readfirstlane_b32_e32(v[46], v[3]))
-    k.emit(s_waitcnt(49279))
+  k.emit(s_load_dwordx2(s[4:5], s[0:1], s[0], 0, 0, 0, 0, 1))
+  k.emit(s_mov_b32(s[8], 0))
+  k.emit(s_mov_b32(s[9], 0))
+  k.emit(s_load_dwordx2(s[12:13], s[0:1], s[0], 8, 0, 0, 0, 1))
+  k.emit(s_load_dwordx2(s[16:17], s[0:1], s[0], 16, 0, 0, 0, 1))
+  k.emit(s_mov_b32(s[36], N))
+  k.emit(s_mov_b32(s[37], K))
+  k.emit(s_mov_b32(s[38], K))
+  k.emit(s_mov_b32(s[43], M))
+  k.emit(s_mov_b32(s[44], N))
+  k.emit(s_mov_b32(s[45], K))
+  k.emit(s_load_dwordx2(s[20:21], s[0:1], s[0], 24, 0, 0, 0, 1))
+  k.emit(s_load_dwordx2(s[24:25], s[0:1], s[0], 32, 0, 0, 0, 1))
+  k.emit(s_mov_b32(s[39], scale_k))
+  k.emit(s_mov_b32(s[40], scale_k))
+  k.emit(v_lshrrev_b32_e32(v[1], 10))
+  k.emit(v_lshrrev_b32_e32(v[2], 10, v[1]))
+  k.emit(v_and_b32_e32(v[2], LIT, v[2], 1023))
+  k.emit(v_and_b32_e32(v[1], LIT, v[1], 1023))
+  k.emit(v_and_b32_e32(v[0], LIT, v[0], 1023))
+  k.emit(v_lshrrev_b32_e32(v[3], 6))
+  k.emit(v_and_b32_e32(v[0], 63))
+  if (tile_m, tile_n) == (256, 256):
+    k.emit(s_mov_b32(s[49], s[2]))
+    k.emit(s_mov_b32(s[47], s[3]))
+  k.emit(v_readfirstlane_b32_e32(v[46], v[3]))
+  k.emit(s_waitcnt(49279))
+
+  if (tile_m, tile_n) == (128, 512):
    for i in range(2):
      k.emit(s_mov_b32(s[6 + i * 8], -16))
      k.emit(s_mov_b32(s[10 + i * 12], -16))
@@ -1213,31 +1218,6 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
    k.emit(s_waitcnt())
    k.emit(s_endpgm())
  elif (tile_m, tile_n) == (192, 256):
-    k.emit(s_and_b32(s[1], s[1], LIT, 65535))
-    k.emit(s_load_dwordx2(s[4:5], s[0:1], s[0], 0, 0, 0, 0, 1))
-    k.emit(s_mov_b32(s[8], 0))
-    k.emit(s_mov_b32(s[9], 0))
-    k.emit(s_load_dwordx2(s[12:13], s[0:1], s[0], 8, 0, 0, 0, 1))
-    k.emit(s_load_dwordx2(s[16:17], s[0:1], s[0], 16, 0, 0, 0, 1))
-    k.emit(s_mov_b32(s[36], N))
-    k.emit(s_mov_b32(s[37], K))
-    k.emit(s_mov_b32(s[38], K))
-    k.emit(s_mov_b32(s[43], M))
-    k.emit(s_mov_b32(s[44], N))
-    k.emit(s_mov_b32(s[45], K))
-    k.emit(s_load_dwordx2(s[20:21], s[0:1], s[0], 24, 0, 0, 0, 1))
-    k.emit(s_load_dwordx2(s[24:25], s[0:1], s[0], 32, 0, 0, 0, 1))
-    k.emit(s_mov_b32(s[39], scale_k))
-    k.emit(s_mov_b32(s[40], scale_k))
-    k.emit(v_lshrrev_b32_e32(v[1], 10))
-    k.emit(v_lshrrev_b32_e32(v[2], 10, v[1]))
-    k.emit(v_and_b32_e32(v[2], LIT, v[2], 1023))
-    k.emit(v_and_b32_e32(v[1], LIT, v[1], 1023))
-    k.emit(v_and_b32_e32(v[0], LIT, v[0], 1023))
-    k.emit(v_lshrrev_b32_e32(v[3], 6))
-    k.emit(v_and_b32_e32(v[0], 63))
-    k.emit(v_readfirstlane_b32_e32(v[46], v[3]))
-    k.emit(s_waitcnt(49279))
    k.emit(s_mul_i32(s[63], LIT, 8, 192))
    k.emit(v_cvt_f32_u32_e32(v[4], s[63]))
    k.emit(s_sub_i32(s[62], 0, s[63]))
@@ -2234,49 +2214,22 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
    k.emit(s_waitcnt())
    k.emit(s_endpgm())
  elif (tile_m, tile_n) == (256, 256):
-    k.emit(s_and_b32(s[1], s[1], LIT, 65535))
-    k.emit(s_load_dwordx2(s[4:5], s[0:1], s[0], 0, 0, 0, 0, 1))
-    k.emit(s_mov_b32(s[8], 0))
-    k.emit(s_mov_b32(s[9], 0))
-    k.emit(s_load_dwordx2(s[12:13], s[0:1], s[0], 8, 0, 0, 0, 1))
-    k.emit(s_load_dwordx2(s[16:17], s[0:1], s[0], 16, 0, 0, 0, 1))
-    k.emit(s_mov_b32(s[40], N))
-    k.emit(s_mov_b32(s[41], K))
-    k.emit(s_mov_b32(s[42], K))
-    k.emit(s_mov_b32(s[43], M))
-    k.emit(s_mov_b32(s[44], N))
-    k.emit(s_mov_b32(s[45], K))
-    k.emit(s_load_dwordx2(s[20:21], s[0:1], s[0], 24, 0, 0, 0, 1))
-    k.emit(s_load_dwordx2(s[24:25], s[0:1], s[0], 32, 0, 0, 0, 1))
-    k.emit(s_mov_b32(s[36], scale_k))
-    k.emit(s_mov_b32(s[37], scale_k))
-    k.emit(v_lshrrev_b32_e32(v[1], 10))
-    k.emit(v_lshrrev_b32_e32(v[2], 10, v[1]))
-    k.emit(v_and_b32_e32(v[2], LIT, v[2], 1023))
-    k.emit(v_and_b32_e32(v[1], LIT, v[1], 1023))
-    k.emit(v_and_b32_e32(v[0], LIT, v[0], 1023))
-    k.emit(v_lshrrev_b32_e32(v[3], 6))
-    k.emit(v_and_b32_e32(v[0], 63))
-    k.emit(s_mov_b32(s[46], s[2]))
-    k.emit(s_mov_b32(s[47], s[3]))
-    k.emit(v_readfirstlane_b32_e32(v[49], v[3]))
-    k.emit(s_waitcnt(49279))
    k.emit(s_add_u32(s[55], s[44], LIT, 255))
    k.emit(s_lshr_b32(s[54], s[55], 8))
    k.emit(s_mul_i32(s[48], s[54], s[47]))
-    k.emit(s_add_i32(s[48], s[48], s[46]))
+    k.emit(s_add_i32(s[48], s[48], s[49]))
    k.emit(s_add_u32(s[55], s[43], LIT, 255))
    k.emit(s_lshr_b32(s[52], s[55], 8))
    k.emit(s_lshl_b32(s[52], s[52], 5))
-    k.emit(s_mov_b32(s[46], 0))
+    k.emit(s_mov_b32(s[49], 0))
    k.label('L2_00E8')
    k.emit(s_cmp_lt_i32(s[48], s[52]))
    k.emit(s_cbranch_scc1(3), target='L2_00FC')
    k.emit(s_sub_i32(s[48], s[48], s[52]))
-    k.emit(s_add_i32(s[46], s[46], 32))
+    k.emit(s_add_i32(s[49], s[49], 32))
    k.emit(s_branch(65531), target='L2_00E8')
    k.label('L2_00FC')
-    k.emit(s_sub_i32(s[54], s[54], s[46]))
+    k.emit(s_sub_i32(s[54], s[54], s[49]))
    k.emit(s_cmp_lt_i32(s[54], 32))
    k.emit(s_cbranch_scc1(3), target='L2_0114')
    k.emit(s_lshr_b32(s[47], s[48], 5))
@@ -2311,7 +2264,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
    k.emit(s_mul_i32(s[52], s[54], s[47]))
    k.emit(s_sub_i32(s[52], s[48], s[52]))
    k.label('L2_0194')
-    k.emit(s_add_i32(s[46], s[52], s[46]))
+    k.emit(s_add_i32(s[49], s[52], s[49]))
    k.emit(s_mov_b32(s[6], -16))
    k.emit(s_mov_b32(s[10], -16))
    k.emit(s_mov_b32(s[18], -16))
@@ -2328,18 +2281,18 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
    k.emit(s_or_b32(s[9], s[9], LIT, 262144))
    k.emit(s_or_b32(s[17], s[17], LIT, 262144))
    k.emit(s_or_b32(s[13], s[13], LIT, 262144))
-    k.emit(s_lshr_b32(s[41], s[41], 1))
-    k.emit(s_mul_i32(s[52], s[41], s[43]))
+    k.emit(s_lshr_b32(s[37], s[37], 1))
+    k.emit(s_mul_i32(s[52], s[37], s[43]))
    k.emit(s_mov_b32(s[14], s[52]))
-    k.emit(s_lshr_b32(s[42], s[42], 1))
-    k.emit(s_mul_i32(s[52], s[42], s[44]))
+    k.emit(s_lshr_b32(s[38], s[38], 1))
+    k.emit(s_mul_i32(s[52], s[38], s[44]))
    k.emit(s_mov_b32(s[18], s[52]))
    k.emit(s_add_u32(s[52], s[43], 31))
    k.emit(s_lshr_b32(s[52], s[52], 5))
    k.emit(s_lshl_b32(s[52], s[52], 5))
-    k.emit(s_mul_i32(s[53], s[52], s[36]))
+    k.emit(s_mul_i32(s[53], s[52], s[39]))
    k.emit(s_mov_b32(s[22], s[53]))
-    k.emit(s_mul_i32(s[53], s[44], s[37]))
+    k.emit(s_mul_i32(s[53], s[44], s[40]))
    k.emit(s_mov_b32(s[26], s[53]))
    k.emit(s_mov_b32(s[23], LIT, 131072))
    k.emit(s_mov_b32(s[27], LIT, 131072))
@@ -2356,23 +2309,23 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
    k.emit(v_add_u32_e32(v[5], v[5], v[6]))
    k.emit(v_and_b32_e32(v[4], 1, v[4]))
    k.emit(v_add_u32_e32(v[5], v[5], v[4]))
-    k.emit(v_mul_lo_u32(v[212], s[41], v[5]))
+    k.emit(v_mul_lo_u32(v[212], s[37], v[5]))
    k.emit(v_and_b32_e32(v[4], 7))
    k.emit(v_lshlrev_b32_e32(v[4], 4, v[4]))
    k.emit(v_add_u32_e32(v[212], v[212], v[4]))
-    k.emit(s_lshr_b32(s[52], s[49], 1))
+    k.emit(s_lshr_b32(s[52], s[46], 1))
    k.emit(s_mul_i32(s[52], s[52], 8))
-    k.emit(s_and_b32(s[53], s[49], 1))
+    k.emit(s_and_b32(s[53], s[46], 1))
    k.emit(s_mul_i32(s[53], s[53], 2))
    k.emit(s_add_u32(s[52], s[52], s[53]))
    k.emit(s_mul_i32(s[53], s[47], LIT, 256))
    k.emit(s_add_u32(s[52], s[52], s[53]))
-    k.emit(s_mul_i32(s[52], s[41], s[52]))
+    k.emit(s_mul_i32(s[52], s[37], s[52]))
    k.emit(v_add_u32_e32(v[212], s[52], v[212]))
-    k.emit(s_mul_i32(s[52], s[41], 32))
+    k.emit(s_mul_i32(s[52], s[37], 32))
    for i in range(7):
      k.emit(v_add_u32_e32(v[213 + i * 1], s[52], v[212 + i * 1]))
-    k.emit(s_mul_i32(s[59], LIT, s[49], 1056))
+    k.emit(s_mul_i32(s[59], LIT, s[46], 1056))
    k.emit(s_add_u32(s[59], LIT, s[59], 4096))
    k.emit(v_and_b32_e32(v[4], 15))
    k.emit(v_lshrrev_b32_e32(v[5], 3, v[4]))
@@ -2396,35 +2349,35 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
    k.emit(v_add_u32_e32(v[221], LIT, v[220], 33792))
    k.emit(v_lshlrev_b32_e32(v[222], 2))
    k.emit(s_mul_i32(s[52], s[47], LIT, 256))
-    k.emit(s_mul_i32(s[53], s[49], 32))
+    k.emit(s_mul_i32(s[53], s[46], 32))
    k.emit(s_add_i32(s[52], s[53], s[52]))
-    k.emit(s_mul_i32(s[53], s[52], s[36]))
+    k.emit(s_mul_i32(s[53], s[52], s[39]))
    k.emit(v_add_u32_e32(v[222], s[53], v[222]))
-    k.emit(s_mul_i32(s[53], LIT, s[36], 128))
+    k.emit(s_mul_i32(s[53], LIT, s[39], 128))
    k.emit(v_add_u32_e32(v[223], s[53], v[222]))
-    k.emit(s_mul_i32(s[60], s[49], LIT, 256))
+    k.emit(s_mul_i32(s[60], s[46], LIT, 256))
    k.emit(s_add_i32(s[60], s[60], 0))
    k.emit(v_lshlrev_b32_e32(v[224], 2))
    k.emit(v_add_u32_e32(v[224], 0, v[224]))
    k.emit(v_lshlrev_b32_e32(v[225], 4))
-    k.emit(s_mul_i32(s[52], s[46], LIT, 256))
-    k.emit(s_mul_i32(s[53], s[49], 64))
+    k.emit(s_mul_i32(s[52], s[49], LIT, 256))
+    k.emit(s_mul_i32(s[53], s[46], 64))
    k.emit(s_add_u32(s[52], s[52], s[53]))
-    k.emit(s_mul_i32(s[52], s[52], s[42]))
+    k.emit(s_mul_i32(s[52], s[52], s[38]))
    k.emit(v_add_u32_e32(v[225], s[52], v[225]))
-    k.emit(s_mul_i32(s[52], 16, s[42]))
+    k.emit(s_mul_i32(s[52], 16, s[38]))
    k.emit(v_add_u32_e32(v[226], s[52], v[225]))
    k.emit(v_add_u32_e32(v[227], s[52], v[226]))
    k.emit(v_add_u32_e32(v[228], s[52], v[227]))
    for i in range(4):
      k.emit(v_add_u32_e32(v[229 + i * 1], LIT, v[225 + i * 1], 1024))
    k.emit(v_lshlrev_b32_e32(v[233], 2))
-    k.emit(s_mul_i32(s[52], s[46], LIT, 256))
-    k.emit(s_mul_i32(s[53], s[49], 64))
+    k.emit(s_mul_i32(s[52], s[49], LIT, 256))
+    k.emit(s_mul_i32(s[53], s[46], 64))
    k.emit(s_add_i32(s[52], s[53], s[52]))
-    k.emit(s_mul_i32(s[53], s[52], s[37]))
+    k.emit(s_mul_i32(s[53], s[52], s[40]))
    k.emit(v_add_u32_e32(v[233], s[53], v[233]))
-    k.emit(s_mul_i32(s[52], 32, s[37]))
+    k.emit(s_mul_i32(s[52], 32, s[40]))
    k.emit(v_add_u32_e32(v[234], s[52], v[233]))
    k.emit(s_mov_b32(s[61], LIT, 128))
    k.emit(s_mov_b32(s[62], LIT, 2048))
@@ -2510,18 +2463,18 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
    k.emit(ds_read_b32(v[201], v[224], v[0], v[0], 0, 0, 1))
    k.emit(ds_read_b32(v[202], v[224], v[0], v[0], 0, 0, 2))
    k.emit(ds_read_b32(v[203], v[224], v[0], v[0], 0, 0, 3))
-    k.emit(s_lshl_b32(s[40], s[40], 1))
+    k.emit(s_lshl_b32(s[36], s[36], 1))
    k.emit(s_mul_i32(s[52], s[47], LIT, 256))
-    k.emit(s_mul_hi_u32(s[53], s[52], s[40]))
+    k.emit(s_mul_hi_u32(s[53], s[52], s[36]))
    k.emit(s_add_u32(s[5], s[5], s[53]))
-    k.emit(s_mul_i32(s[53], s[52], s[40]))
+    k.emit(s_mul_i32(s[53], s[52], s[36]))
    k.emit(s_add_u32(s[4], s[4], s[53]))
    k.emit(s_addc_u32(s[5], 0, s[5]))
    k.emit(s_sub_i32(s[52], s[43], s[52]))
-    k.emit(s_mul_i32(s[52], s[52], s[40]))
+    k.emit(s_mul_i32(s[52], s[52], s[36]))
    k.emit(s_mov_b32(s[6], s[52]))
    k.emit(v_and_b32_e64(v[235], v[0], 15))
-    k.emit(v_mul_lo_u32(v[235], v[235], s[40]))
+    k.emit(v_mul_lo_u32(v[235], v[235], s[36]))
    k.emit(v_lshrrev_b32_e32(v[4], 5))
    k.emit(v_mul_i32_i24_e32(v[4], 16, v[4]))
    k.emit(v_add_u32_e32(v[235], v[4], v[235]))
@@ -2529,12 +2482,12 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
    k.emit(v_and_b32_e32(v[4], 1, v[4]))
    k.emit(v_mul_i32_i24_e32(v[4], 32, v[4]))
    k.emit(v_add_u32_e32(v[235], v[4], v[235]))
-    k.emit(s_mul_i32(s[52], s[46], LIT, 256))
-    k.emit(s_mul_i32(s[53], s[49], 64))
+    k.emit(s_mul_i32(s[52], s[49], LIT, 256))
+    k.emit(s_mul_i32(s[53], s[46], 64))
    k.emit(s_add_i32(s[52], s[52], s[53]))
    k.emit(s_lshl_b32(s[52], s[52], 1))
    k.emit(v_add_u32_e32(v[235], s[52], v[235]))
-    k.emit(s_mul_i32(s[53], s[40], 16))
+    k.emit(s_mul_i32(s[53], s[36], 16))
    for i in range(15):
      k.emit(v_add_u32_e64(v[236 + i * 1], v[235 + i * 1], s[53]))
    k.emit(s_mov_b32(s[50], 0))
@@ -2543,7 +2496,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
      k.emit(s_cmp_lt_u32(LIT, s[51], 512 + i * -256))
      k.emit(s_cselect_b32(s[61 + i * 1], s[61 + i * 1], 0))
      k.emit(s_cselect_b32(s[63 + i * 1], s[63 + i * 1], 0))
-    k.emit(s_cmp_lt_i32(s[49], 2))
+    k.emit(s_cmp_lt_i32(s[46], 2))
    k.emit(s_cbranch_scc0(1367), target='L2_25B8')
    k.label('L2_105C')
    k.emit(s_waitcnt(122))
@@ -53,7 +53,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)
@@ -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()
@@ -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__":
@@ -37,7 +37,7 @@ class PM4Ops(FastEnum):
  RELEASE_MEM = auto(); DISPATCH_DIRECT = auto(); EVENT_WRITE = auto()  # noqa: E702

 def pkt3(ctx, op:PM4Ops, *vals):
-  return UOp(Ops.INS, arg=op, src=tuple(UOp.const(x, dtypes.uint32)
+  return UOp(Ops.INS, arg=(op, dtypes.void), 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 wreg(ctx, reg:AMDReg, *args:sint, **kwargs:int):
@@ -87,7 +87,7 @@ def release_mem(ctx, address=0x0, value=0, data_sel=0, int_sel=2, ctxid=0, cache

 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, (
+  return UOp(Ops.LINEAR, src=(
    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)))
@@ -135,15 +135,15 @@ def pm4_program(ctx, call, prg):
    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))
+  return UOp(Ops.LINEAR, src=tuple(ins))

 pm_pm4_opsel = PatternMatcher([
  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), pm4_program),

-  (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),
+  (UPat(Ops.INS, arg=("wait", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))), pm4_wait),
+  (UPat(Ops.INS, arg=("barrier", dtypes.void)), pm4_barrier),
+  (UPat(Ops.INS, arg=("timestamp", dtypes.void), src=(UPat(name="dst"),)), pm4_timestamp),
+  (UPat(Ops.INS, arg=("store", dtypes.void), src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), pm4_store),
 ])

 def queue_ptrs(devs, qname:str, q:AMDQueueDesc) -> tuple[UOp, ...]:
@@ -207,10 +207,10 @@ def sdma_timestamp(ctx, ins, dst):
 pm_sdma_opsel = PatternMatcher([
  (UPat(Ops.CALL, src=(UPat(Ops.COPY),), name="call", allow_any_len=True), sdma_copy),

-  (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),
+  (UPat(Ops.INS, arg=("barrier", dtypes.void)), lambda: UOp(Ops.NOOP)),
+  (UPat(Ops.INS, arg=("wait", dtypes.void), src=(UPat(name="dst"), UPat(name="val")), name="ins"), sdma_wait),
+  (UPat(Ops.INS, arg=("timestamp", dtypes.void), src=(UPat(name="dst"),), name="ins"), sdma_timestamp),
+  (UPat(Ops.INS, arg=("store", dtypes.void), src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val")), name="ins"), sdma_store),
 ])

 def sdma_submit(cmdbuf, devs):
@@ -254,7 +254,7 @@ def amd_usb_submit(ctx, lin):

  if nb:=usb_arm_bytes(ctx.pre, Device[ctx.devs[0]].iface.usb_sram):
    poke = (ctx.sdma.SDMA_OP_WRITE, *data64_le(Device[ctx.devs[0]].iface.cq_buf.va_addr + 12), 0, 0)
-    lin = lin.replace(src=lin.src + (UOp(Ops.INS, arg="poke", src=tuple(UOp.const(x, dtypes.uint32) for x in poke)),))
+    lin = lin.replace(src=lin.src + (UOp(Ops.INS, arg=("poke", dtypes.void), src=tuple(UOp.const(x, dtypes.uint32) for x in poke)),))

  ib_host, ib_gpu, pkt_dw = usb_ib(ctx.devs, lin, 32 if comp else 0x100, nb)
  pkt = (ctx.pm4.PACKET3(ctx.pm4.PACKET3_INDIRECT_BUFFER,2),*data64_le(ib_gpu.getaddr(ctx.devs)),pkt_dw|ctx.pm4.INDIRECT_BUFFER_VALID) if comp else ()
@@ -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,7 +10,7 @@ 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:
@@ -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
@@ -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()
@@ -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()
@@ -1,10 +1,12 @@
 from tinygrad import Tensor
+from tinygrad.helpers import NUM_CPU_THREADS
 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 = NUM_CPU_THREADS.value

 def get_example_inputs(graph_inputs:dict[str, OnnxValue], config={}):
  """
@@ -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}")
@@ -84,7 +84,7 @@ testing = [
  "pillow",
  "onnx==1.19.0",
  "onnx2torch",
-  "onnxruntime",
+  "onnxruntime==1.24.1",
  "opencv-python",
  "transformers",
  "sentencepiece",
@@ -50,10 +50,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} & () & slot, dtype, size?, device?, addrspace? &
+  Placeholder with flat storage of $\mathrm{size}$ elements. Substituted in \op{Call}. \\[4pt]
+\op{Buffer} & () & slot, dtype, size, device, addrspace &
+  Concrete buffer slot with flat storage of $\mathrm{size}$ elements. \\
 \op{Const} & () & value, dtype &
  A scalar constant with shape $(\ )$. \\
  & & & Form vector consts with \op{Stack} \\
@@ -102,9 +102,8 @@ 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{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{Call} & (body, $a_0$, $a_1$, \ldots) & --- & Substitute each \op{Param} $k$ in body with $a_k$. \\
+\op{Tuple} & $(v_0, v_1, \ldots)$ & --- & Pack values; required as a value-producing \op{Call} body. \\
 \op{GetTuple} & $(T,)$ & idx & Extract element at idx from a \op{Tuple}. \\
 \bottomrule
 \end{tabular}
@@ -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} & $\mathrm{src}[0].\mathrm{dtype}$ & substitute \op{Param} shapes & $\mathrm{src}[1].\mathrm{device}$ & dtype range \\
 \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}$ & --- \\
@@ -421,7 +420,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
@@ -436,11 +435,11 @@ When \texttt{f(x, y)} is called, the decorator:
  \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$).
+        \op{Call}(\op{Tuple}(body), $x$, $y$).
  \item \textbf{Returns} the result via \op{GetTuple}$(0)$, or one \op{GetTuple} per element for tuple returns.
 \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. At schedule time, the value-producing \op{Call} is resolved by substituting each \op{Param}$(k)$ back with its corresponding argument $a_k$, or its body is lowered into a \op{Sink} (opaque) if it is to be compiled as a reusable kernel.

 %% ============================================================
 \subsection*{Lowering Pipeline \normalfont\small--- from Tensor graph to machine code}
@@ -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):
@@ -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(),
@@ -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
@@ -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

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

@@ -1079,5 +1079,80 @@ class TestBatchNormRunningStats(unittest.TestCase):
    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]])
+
+  @unittest.expectedFailure
+  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()
+    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(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()
@@ -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):
@@ -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"))

@@ -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)
@@ -6,7 +6,7 @@ from test.helpers import assert_jit_cache_len, call_is_graph, not_support_multi_
 from test.unit.test_jit import _simple_test
 from tinygrad import Tensor, Variable, TinyJit, Device, dtypes
 from tinygrad.engine.jit import graph_class
-from tinygrad.helpers import JIT, DEV, GlobalCounters
+from tinygrad.helpers import JIT, DEV, GlobalCounters, HCQ2
 from tinygrad.uop.ops import Ops
 from tinygrad.renderer.isa.x86 import X86Renderer

@@ -235,6 +235,7 @@ class TestJitPrune(unittest.TestCase):
    assert_jit_cache_len(w2_prune, 1)

 class TestJitFree(unittest.TestCase):
+  @unittest.skipIf(HCQ2, "hcq2 keeps refs to intermediate buffers")
  def test_free_intermediates(self):
    ext_tensor = Tensor([1,24,23,45,1])
    @TinyJit
@@ -11,16 +11,16 @@ from tinygrad.codegen import to_program
 class TestLinearizerFailure(unittest.TestCase):
  @unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL")
  def test_failure_beam_mnist(self):
-    c0 = UOp.param(0, dtypes.uchar, (4014080,))
+    c0 = UOp.param(0, dtypes.uchar, 4014080)
    c1 = UOp.range(UOp.const(512), 0, AxisType.GLOBAL)
    c2 = UOp.range(UOp.const(784), 1, AxisType.GLOBAL)
    c3 = UOp.range(UOp.const(10), 3, AxisType.GLOBAL)
-    c4 = UOp.param(1, dtypes.int, (512,))
+    c4 = UOp.param(1, dtypes.int, 512)
    c5 = c4.index(c1.valid(UOp.const(True)))
    c6 = UOp.range(UOp.const(6000), 1004, AxisType.REDUCE)
    c7 = UOp.range(UOp.const(3750), 2006, AxisType.REDUCE)
    c8 = UOp.range(UOp.const(16), 2007, AxisType.GROUP_REDUCE)
-    c9 = UOp.param(2, dtypes.uchar, (47040000,))
+    c9 = UOp.param(2, dtypes.uchar, 47040000)
    c10 = c9.index((((c3*UOp.const(4704000))+c2)+(c6*UOp.const(784))).valid(UOp.const(True)))
    c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(6000))+c6)+((c7*UOp.const(16))+c8)).alu(Ops.CMPLT, UOp.const(59999)).where(UOp.const(0).cast(dtypes.int), UOp.const(1).cast(dtypes.int)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(-1).cast(dtypes.int))).where(UOp.const(0).cast(dtypes.uchar), c10).reduce(c6, arg=Ops.ADD)
    c12 = c0.index((((c1*UOp.const(7840))+(c2*UOp.const(10)))+c3).valid(UOp.const(True))).store(c11).end(c1, c2, c3)
@@ -58,6 +58,11 @@ class TestMultiTensor(unittest.TestCase):
    assert X.uop.ended_ranges == X.uop.src[1:]
    (X + X).realize()

+  def test_shard_invalids_contiguous(self):
+    # every store is Invalid, so none of them should become a (empty) kernel
+    t = Tensor.invalids(8).shard(devices_2, axis=0).contiguous()
+    self.assertEqual(len([c for c in t.schedule_linear().src if c.src[0].op is Ops.SINK]), 1)
+
  @unittest.expectedFailure # TODO: fix
  def test_shard_empty(self):
    GlobalCounters.reset()
@@ -182,6 +187,13 @@ class TestMultiTensor(unittest.TestCase):
        a,b = jit_allreduce(Tensor.rand(256, 256))
        np.testing.assert_almost_equal(a.numpy(), b.numpy(), decimal=5)

+  def test_allreduce_all2all_jit(self):
+    with Context(ALL2ALL=2):
+      jit_allreduce = TinyJit(_test_allreduce)
+      for _ in range(5):
+        a,b = jit_allreduce(Tensor.rand(256, 256))
+        np.testing.assert_almost_equal(a.numpy(), b.numpy(), decimal=5)
+
  def test_multitensor_jit_input(self):
    @TinyJit
    def f(x): return (x+1).contiguous().sum()
@@ -135,7 +135,7 @@ class TestNN(unittest.TestCase):
  def test_conv2d_same_padding_large_kernel(self):
    self._test_conv(Conv2d, torch.nn.Conv2d, BS=16, C1=16, DIMS=[28, 33], C2=32, K=9, S=1, P='same')
  def test_conv2d_same_padding_with_dilation(self):
-    self._test_conv(Conv2d, torch.nn.Conv2d, BS=16, C1=3, DIMS=[28, 28], C2=32, K=3, S=1, P='same', D=3)
+    self._test_conv(Conv2d, torch.nn.Conv2d, BS=16, C1=3, DIMS=[28, 31], C2=32, K=(3,5), S=1, P='same', D=(2,3))

  def test_conv2d_same_padding_invalid_stride(self):
    self.assertRaises(ValueError, Conv2d, in_channels=16, out_channels=32, kernel_size=2, stride=2, padding='same')
@@ -6,6 +6,7 @@ from tinygrad.helpers import getenv, DEBUG, DEV, IMAGE, Context
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.renderer.nir import NIRRenderer
+from tinygrad.renderer.isa.x86 import X86Renderer

 TINY_BACKEND = getenv("TINY_BACKEND")
 if TINY_BACKEND:
@@ -358,6 +359,13 @@ class TestOps(unittest.TestCase):
      lambda x: torch.where(x > 0.5, 4, 2).type(torch.int32).permute((1, 0)),
      lambda x: (x > 0.5).where(4, 2).clone().permute((1, 0)), forward_only=True)

+  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "software vulkan evaluates a NaN != x as false")
+  def test_where_nan_cond(self):
+    # a NaN compares false against everything except !=.
+    for fxn in (lambda x: x<1, lambda x: x>1, lambda x: x!=1, lambda x: x==1):
+      helper_test_op(None, lambda x,a,b: torch.where(fxn(x), a, b), lambda x,a,b: fxn(x).where(a, b), forward_only=True,
+                     vals=[[math.nan, 1.0, 2.0, -1.0], [10, 20, 30, 40], [-1, -2, -3, -4]])
+
  def _test_cmp(self, fxn, reverse=True):
    # test different dtypes
    helper_test_op(None, fxn, fxn, forward_only=True, vals=[[0.,1,2], [2.,1,0]])
@@ -808,6 +816,8 @@ class TestOps(unittest.TestCase):
    helper_test_op([], lambda: tor^0x1337, lambda: ten^0x1337, forward_only=True)
    helper_test_op([], lambda: 0x1337^tor, lambda: 0x1337^ten, forward_only=True)

+  # TODO: x86 PARAM dtype fails SPEC=2
+  @Context(SPEC=1 if isinstance(Device[Device.DEFAULT].renderer, X86Renderer) else 2)
  def test_and(self):
    data = [[1,-8,1],[32,1,6]]
    tor = torch.tensor(data, dtype=torch.int)
@@ -865,9 +875,9 @@ class TestOps(unittest.TestCase):
                   lambda: (ten << Tensor([0,2,4], dtype=dtypes.uint32)).cast(dtypes.int32), forward_only=True)
    helper_test_op([], lambda: tor.__lshift__(2), lambda: ten.__lshift__(2).cast(dtypes.int32), forward_only=True)
    helper_test_op([], lambda: tor.bitwise_left_shift(2), lambda: ten.lshift(2).cast(dtypes.int32), forward_only=True)
-    self.helper_test_exception([], lambda: torch.tensor([1.0]) << 2, lambda: Tensor([1.0]) << 2, expected=RuntimeError)
-    self.helper_test_exception([], lambda: tor << torch.tensor([1.0]), lambda: ten << Tensor([1.0]), expected=RuntimeError)
-    self.helper_test_exception([], lambda: tor << 1.0, lambda: ten << 1.0, expected=RuntimeError)
+    self.helper_test_exception([], lambda: torch.tensor([1.0]) << 2, lambda: (Tensor([1.0]) << 2).realize(), expected=RuntimeError)
+    self.helper_test_exception([], lambda: tor << torch.tensor([1.0]), lambda: (ten << Tensor([1.0])).realize(), expected=RuntimeError)
+    self.helper_test_exception([], lambda: tor << 1.0, lambda: (ten << 1.0).realize(), expected=RuntimeError)

  def test_rshift(self):
    data = [[0,1,2],[1<<8,1<<16,1<<31-1]]
@@ -881,8 +891,8 @@ class TestOps(unittest.TestCase):
                   lambda: (ten >> Tensor([0,2,4], dtype=dtypes.uint32)).cast(dtypes.int32), forward_only=True)
    helper_test_op([], lambda: tor.__rshift__(2), lambda: ten.__rshift__(2).cast(dtypes.int32), forward_only=True)
    helper_test_op([], lambda: tor.bitwise_right_shift(2), lambda: ten.rshift(2).cast(dtypes.int32), forward_only=True)
-    self.helper_test_exception([], lambda: torch.tensor([4.0]) >> 1, lambda: Tensor([4.0]) >> 1, expected=RuntimeError)
-    self.helper_test_exception([], lambda: tor >> torch.tensor([1.0]), lambda: ten >> Tensor([1.0]), expected=RuntimeError)
+    self.helper_test_exception([], lambda: torch.tensor([4.0]) >> 1, lambda: (Tensor([4.0]) >> 1).realize(), expected=RuntimeError)
+    self.helper_test_exception([], lambda: tor >> torch.tensor([1.0]), lambda: (ten >> Tensor([1.0])).realize(), expected=RuntimeError)

  def test_lshift_signed(self):
    data = [[-1, -3, 1, 7], [0, -2147483648, 2147483647, -1]]
@@ -1807,9 +1817,10 @@ class TestOps(unittest.TestCase):
      helper_test_op([()], lambda x: torch.nn.functional.hardtanh(x, -val, val), lambda x: x.hardtanh(-val, val), grad_atol=1e-6)
  def test_asinh(self):
    helper_test_op([(45,65)], lambda x: x.asinh(), grad_atol=1e-6)
-    # TODO: this one has larger tol?
-    helper_test_op([(45,65)], lambda x: x.asinh(), atol=1e-2, rtol=2e-2, grad_rtol=2e-2, low=-300, high=-297)
+    helper_test_op([(45,65)], lambda x: x.asinh(), grad_atol=1e-6, low=-300, high=-297)
    helper_test_op([(45,65)], lambda x: x.asinh(), grad_atol=1e-6, low=300, high=303)
+    helper_test_op([(45,65)], lambda x: x.asinh(), grad_atol=1e-6, low=-1e10, high=-1e9)
+    helper_test_op(None, lambda x: x.asinh(), grad_atol=1e-6, vals=[[-1.0, 0.0, 1.0]])
  def test_acosh(self):
    helper_test_op([(45,65)], lambda x: x.acosh(), grad_atol=1e-6)
    helper_test_op([(45,65)], lambda x: x.acosh(), grad_atol=1e-3, grad_rtol=1e-2, low=-300, high=-297)
@@ -2806,7 +2817,7 @@ class TestOps(unittest.TestCase):
        lambda x: Tensor.interpolate(x, size=out_sz, mode="linear"))

  def test_interpolate_linear_corners_aligned(self):
-    for in_sz, out_sz in [((52,),(29,)), ((29,),(52,))]:
+    for in_sz, out_sz in [((52,),(29,)), ((29,),(52,)), ((29,),(1,))]:
      helper_test_op([(2,3)+in_sz],
        lambda x: torch.nn.functional.interpolate(x, size=out_sz, mode="linear", align_corners=True),
        lambda x: Tensor.interpolate(x, size=out_sz, mode="linear", align_corners=True))
@@ -2959,6 +2970,10 @@ class TestOps(unittest.TestCase):
    data = [math.inf, -math.inf, math.nan]
    helper_test_op((), lambda: torch.tensor(data)[torch.tensor([0, 1, 2])], lambda: Tensor(data)[Tensor([0, 1, 2])])

+  def test_fancy_indexing_index_dtypes(self):
+    helper_test_op((), lambda: torch.tensor([10., 20., 30., 40.])[torch.tensor([1, 2, 3, 0])],
+                       lambda: Tensor([10., 20., 30., 40.])[Tensor([1, 2, 3, 0], dtype=dtypes.uint8)])
+
  @slow_test
  def test_slice_fancy_indexing_no_dim_collapse(self):
    a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
@@ -1,6 +1,6 @@
 import unittest
-from tinygrad import Tensor, nn, Device, dtypes, Variable
-from tinygrad.helpers import Context, GlobalCounters, getenv, PCONTIG, DEBUG
+from tinygrad import Tensor, Device, dtypes, Variable
+from tinygrad.helpers import Context, GlobalCounters, getenv, DEBUG
 from tinygrad.uop.ops import graph_rewrite, PatternMatcher, UPat, Ops, UOp
 from tinygrad.codegen.opt import OptOps, Opt
 from tinygrad.renderer.ptx import PTXRenderer
@@ -14,7 +14,7 @@ class TestDoubleMatmul(unittest.TestCase):
      self.ref = (self.a @ self.b @ self.c).realize()

  def _test(self, opts):
-    with Context(PCONTIG=2, DEBUG=max(2, DEBUG.value)):
+    with Context(DEBUG=max(2, DEBUG.value)):
      out = (self.a @ self.b @ self.c).contiguous(arg=opts).realize()

    with Context(DEBUG=0):
@@ -88,16 +88,15 @@ class TestRangeifyEdgeCase(unittest.TestCase):
    res = Tensor.cat(a, c, dim=0)
    self.assertEqual(res.numpy()[-1, :16].tolist(), [512] * 16)

-  def test_pcontig_multi_gather(self):
+  def test_multi_gather(self):
    # regression test: local bufferize must have device set for const_like to work
-    with Context(PCONTIG=2):
-      # NOTE: with uint type, this will become a long and fail on WEBGPU
-      forest = Tensor(list(range(8)), dtype='int')
-      idx = Tensor([0, 0], dtype='int')
-      node_val = forest.gather(0, idx)
-      idx2 = idx * 2 + 1
-      node_val2 = forest.gather(0, idx2)
-      result = (node_val + node_val2).numpy()
+    # NOTE: with uint type, this will become a long and fail on WEBGPU
+    forest = Tensor(list(range(8)), dtype='int')
+    idx = Tensor([0, 0], dtype='int')
+    node_val = forest.gather(0, idx)
+    idx2 = idx * 2 + 1
+    node_val2 = forest.gather(0, idx2)
+    result = (node_val + node_val2).numpy()
    self.assertEqual(result.tolist(), [1, 1])

 if getenv("BIG") > 2:
@@ -118,65 +117,6 @@ def fa():
  GlobalCounters.reset()
  return q.scaled_dot_product_attention(k, v)

-def fa_bw():
-  Tensor.manual_seed(1337)
-  with Context(DEBUG=0):
-    q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize() for _ in range(3)]
-    attn_output = nn.Linear(HEADS*EMB, HEADS*EMB, bias=False)
-    attn_output.weight.realize()
-    target = Tensor.rand(BS, SEQLEN, HEADS*EMB).contiguous().realize()
-
-  GlobalCounters.reset()
-  attn = q.scaled_dot_product_attention(k, v).contiguous().contiguous_backward()
-  attn = attn.transpose(1, 2).reshape(BS, SEQLEN, -1)
-  out = attn_output(attn)
-  loss = (out - target).square().mean()
-  loss.backward()
-  #ret = [out, Tensor.stack(q.grad, k.grad, v.grad, dim=-1)]
-  #ret = [out, Tensor.stack(q.grad, k.grad, dim=-1), v.grad]
-  ret = [out, q.grad, k.grad, v.grad]
-  Tensor.realize(*ret)
-  return ret
-
-@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (NIRRenderer, PTXRenderer)), "broken in LVP and PTX")
-class TestPcontig(unittest.TestCase):
-  def test_flash_attention_bw(self):
-    with Context(PCONTIG=max(2, PCONTIG.value), DEBUG=2):
-      grads = fa_bw()
-      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
-
-    with Context(PCONTIG=0, DEBUG=2):
-      cmp_grads = fa_bw()
-      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
-
-    with Context(DEBUG=0):
-      mses = [((x-y)**2).sum().item() for x,y in zip(grads, cmp_grads)]
-    mse = sum(mses)
-    print(f"mse: {mse}")
-    self.assertLessEqual(mse, 1e-6)
-
-  def test_flash_attention(self, opts=None):
-    with Context(PCONTIG=2, DEBUG=max(2, DEBUG.value)):
-      ret = fa().realize() if opts is None else fa().contiguous(arg=opts).realize()
-      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
-    with Context(DEBUG=2):
-      cmp = fa().realize()
-      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
-    with Context(DEBUG=0):
-      mse = ((cmp-ret)**2).sum().item()
-    print(f"mse: {mse}")
-    self.assertLessEqual(mse, 1e-6)
-
-  def test_flash_attention_opt(self):
-    opts = ()
-    # columns in top matrix
-    opts += (Opt(OptOps.UPCAST, 0, 4),)
-    # columns in bottom matrix
-    opts += (Opt(OptOps.UPCAST, 3, 4),)
-    # rows in all the matrix
-    opts += (Opt(OptOps.UPCAST, 4, 4),)
-    self.test_flash_attention(opts)
-
 # contiguous + reduce can support ranges?

@unittest.skip("pm_rangeify no longer exists. test this in a different way")
@@ -23,8 +23,8 @@ def _test_uop_result(inputs:list[Tensor], sink:UOp, local_size=None):

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

  @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
  def test_gated_store_with_alu_2d(self):
-    a = UOp.param(0, dtypes.int, (8,))
+    a = UOp.param(0, dtypes.int, 8)
    gate_alu_0 = (lidx0:=UOp.special(4, 'lidx0')).ne(0)
    gate_alu_1 = (lidx1:=UOp.special(2, 'lidx1')).ne(0)
    gated_alu_store = UOp(Ops.STORE, src=(a.index((lidx0+lidx1*4).valid(gate_alu_0&gate_alu_1)), UOp.const(1).cast(dtypes.int)))
@@ -78,9 +78,9 @@ class TestCStyleFailures(unittest.TestCase):
@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, WGSLRenderer), "tests for wgsl renderer")
 class TestWGSLFailures(unittest.TestCase):
  def test_folded_packed_store(self):
-    b = UOp.param(0, dtypes.char, (4,))
+    b = UOp.param(0, dtypes.char, 4)
    idx = b.index(UOp.const(0).cast(dtypes.int))
-    store = UOp.store(idx, UOp.load(idx, dtype=dtypes.uint32) & UOp.const(0xffffff00).cast(dtypes.uint32))
+    store = UOp.store(idx, idx.cast(dtypes.uint32).load() & UOp.const(0xffffff00).cast(dtypes.uint32))
    src = Device[Device.DEFAULT].renderer.render(UOp.sink(store, arg=KernelInfo()).toposort())
    self.assertIn("atomicAnd(&data0_4[0],4294967040u);", src)
    self.assertNotIn("atomicAdd", src)
@@ -93,9 +93,9 @@ class TestWGSLFailures(unittest.TestCase):

  # WGSL has a specific select(alt, val, gate) ternary operator instead of gate?val:alt
  def test_gated_load(self):
-    a = UOp.param(0, dtypes.int, (4,))
-    b = UOp.param(1, dtypes.int, (4,))
-    c = UOp.param(2, dtypes.int, (4,))
+    a = UOp.param(0, dtypes.int, 4)
+    b = UOp.param(1, dtypes.int, 4)
+    c = UOp.param(2, dtypes.int, 4)
    lidx0 = UOp.special(4, "lidx0")
    gate = lidx0.ne(0)
    alt = c.index(lidx0).load()
@@ -110,7 +110,7 @@ class TestWGSLFailures(unittest.TestCase):
 class TestPTXFailures(unittest.TestCase):
  @unittest.skip("INDEX can only have a gate ALU parent, not an IF")
  def test_gated_store_with_if(self):
-    a = UOp.param(0, dtypes.int, (4,))
+    a = UOp.param(0, dtypes.int, 4)
    gate_alu = (lidx0:=UOp.special(4, 'lidx0')).ne(0)
    val = UOp.const(1).cast(dtypes.int)
    if_uop = UOp(Ops.IF, src=(gate_alu,))
@@ -147,7 +147,7 @@ class TestSchedule(unittest.TestCase):
    devs = ("CPU:0", "CPU:1")
    x = Tensor.ones(2, device="CPU").shard(devs, axis=0).realize()
    out = (x.sum()*2).reshape(1).to("CPU")
-    run_linear(*check_schedule(out, 5))
+    run_linear(*check_schedule(out, 3))
    np.testing.assert_equal(out.numpy(), [4.])

 class TestLimitBufs(unittest.TestCase):
@@ -365,16 +365,6 @@ class TestCopyFolding(unittest.TestCase):
    b = a.to("CPU")
    self.assertListEqual(b.tolist(), [2.])

-  def test_copy_to_same_device(self):
-    a = Tensor.empty(4).uop
-    b = a.copy_to_device(a.device)
-    check_schedule(b, 1, filter_sink=False) # TODO: 0?
-
-  def test_copy_to_same_device_alt(self):
-    a = Tensor.empty(4, 4).uop
-    b = a.copy_to_device(a.device)
-    check_schedule(b, 1, filter_sink=False) # TODO: 0?
-
  def test_copy_to_same_device_sched(self):
    a = Tensor.ones(4).contiguous().realize().uop.buf_uop
    t = Tensor(a.copy_to_device(a.device))
@@ -384,40 +374,35 @@ class TestCopyFolding(unittest.TestCase):
    assert t.uop.is_realized, f"didn't realize Tensor {t}"
    self.assertListEqual(t.tolist(), [1.,1.,1.,1.])

-  @unittest.skip("same-device copies are no-ops")
-  def test_self_assign_same_device_copy(self):
-    a = Tensor.ones(4, 4).contiguous().realize()
-    # use copy_to_device to bypass Tensor.to() shortcircuit and force a real same-device COPY in the graph
-    a.assign(Tensor(a.uop.copy_to_device(a.device), a.device))
-    run_linear(*check_schedule(a, 2, filter_sink=False))
-    self.assertListEqual(a.tolist(), [[1.]*4]*4)
-
  def test_clone(self):
    a = Tensor.empty(4)
    check_schedule(a.clone(), 1, filter_sink=False)

  def test_shrink_copy(self):
-    a = Tensor.arange(4)
-    view = a.shrink(((0, 2),))
-    b = view.clone()
-    run_linear(*check_schedule(b, 1, filter_sink=False))
-    self.assertEqual(b.uop.base.buffer.size, 2)
-    self.assertEqual(b.uop.numel(), 2)
-    self.assertListEqual(b.tolist(), [0, 1])
+    a = Tensor.arange(4).clone("CPU:1").realize()
+    b = a.to("CPU:2").shrink(((1, 3),)).to("CPU:3")
+    GlobalCounters.reset()
+    run_linear(*check_schedule(b, 3, filter_sink=False))
+    # extra E kernel, copy exactly 4 bytes
+    self.assertEqual(GlobalCounters.global_mem, 4*4 + 2*4*2 + 2*4)
+    self.assertListEqual(b.tolist(), [1, 2])

  def test_expanded_copy(self):
-    a = Tensor.arange(2)
-    view = a.reshape(2, 1).expand(2, 2)
-    b = view.clone()
-    run_linear(*check_schedule(b, 1, filter_sink=False))
-    self.assertEqual(b.uop.base.buffer.size, 4)
-    self.assertEqual(b.uop.numel(), 4)
-    self.assertListEqual(b.tolist(), [[0, 0], [1, 1]])
+    a = Tensor.arange(4).clone("CPU:1").realize()
+    b = a.to("CPU:2").reshape(4, 1).expand(4, 2).to("CPU:3")
+    GlobalCounters.reset()
+    run_linear(*check_schedule(b, 3, filter_sink=False))
+    # TODO: expands before copy
+    self.assertEqual(GlobalCounters.global_mem, 4*4 + (4*4 + 8*4) + 8*4)
+    self.assertListEqual(b.tolist(), [[0, 0], [1, 1], [2, 2], [3, 3]])

  def test_permuted_copy(self):
-    a = Tensor.arange(4)
-    b = a.reshape(2, 2).permute(1, 0)
-    b.realize()
+    a = Tensor.arange(4).clone("CPU:1").realize()
+    b = a.to("CPU:2").reshape(2, 2).permute(1, 0).to("CPU:3")
+    GlobalCounters.reset()
+    run_linear(*check_schedule(b, 3, filter_sink=False))
+    # permutes before copy
+    self.assertEqual(GlobalCounters.global_mem, 4*4 + (4*4 + 4*4) + 4*4)
    self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])

  def test_permute_on_disk(self):
@@ -301,6 +301,14 @@ class TestSetitem(unittest.TestCase):
    self.assertListEqual(z[2:5].tolist(), [2, 2, 2])
    self.assertListEqual(z[6:7].tolist(), [3])

+class TestAssignBitcast(unittest.TestCase):
+  def test_assign_through_bitcast(self):
+    # the dest is unrealized, so callify cannot fold the BITCAST into a buffer view and the STORE keeps a
+    # BITCAST dest; the bitcast has to move to the value side or the store never reaches the buffer
+    a = Tensor.full((4,), 1.0, dtype=dtypes.float32).contiguous()
+    a.bitcast(dtypes.uint32).assign(Tensor([0x40800000, 0x40400000, 0x40000000, 0x3f800000], dtype=dtypes.uint32)).realize()
+    np.testing.assert_allclose(a.numpy(), [4.0, 3.0, 2.0, 1.0])
+
 class TestWithGrad(unittest.TestCase):
  def test_basic_setitem_works(self):
    z = Tensor.rand(8, 8)
@@ -20,8 +20,8 @@ def run_uops(uops_list:list[UOp], bufs:list[Buffer]):

 def uop(uops:list[UOp], op:Ops, dtype:Optional[DType], src:tuple[UOp, ...], arg:Any=None) -> UOp:
  if op is Ops.CONST: uops.append(UOp.const(arg).cast(dtype))
-  elif op is Ops.PARAM: uops.append(UOp.param(arg, dtype, shape=(1,)))
-  else: uops.append(UOp(op, dtype, tuple(src), arg))
+  elif op is Ops.PARAM: uops.append(UOp.param(arg, dtype, 1))
+  else: uops.append(UOp(op, tuple(src), arg))
  return uops[-1]

 def _test_single_value(vals, op, dts):
@@ -62,7 +62,7 @@ def _test_uops_result(output_dtype, uops, res):
 class TestBitcastBufferView(unittest.TestCase):
  @Context(SPEC=2)
  def test_render(self):
-    buf = UOp.param(0, dtypes.uint32, (4,))
+    buf = UOp.param(0, dtypes.uint32, 4)
    uops = to_uops_list([buf.shrink(((1, 3),)).bitcast(dtypes.uint64).index(0).store(1)], ren=Device[Device.DEFAULT].renderer)
    idx = next(u for u in uops if u.op is Ops.INDEX and u.src[0].op is Ops.BITCAST)
    self.assertEqual(idx.src[0].src[0].op, Ops.SHRINK)
@@ -71,7 +71,7 @@ class TestBitcastBufferView(unittest.TestCase):
  @Context(SPEC=2)
  def test_load(self):
    val = 0x1122334455667788
-    src, out = UOp.param(0, dtypes.uint32, (4,)), UOp.param(1, dtypes.uint64, (1,))
+    src, out = UOp.param(0, dtypes.uint32, 4), UOp.param(1, dtypes.uint64, 1)
    ibuf = Buffer(Device.DEFAULT, 4, dtypes.uint32, initial_value=np.array([0, 0x55667788, 0x11223344, 0], dtype=np.uint32).tobytes())
    obuf = Buffer(Device.DEFAULT, 1, dtypes.uint64).allocate()
    run_uops([out.index(0).store(src.shrink(((1, 3),)).bitcast(dtypes.uint64).index(0))], [ibuf, obuf])
@@ -80,7 +80,7 @@ class TestBitcastBufferView(unittest.TestCase):
  @Context(SPEC=2)
  def test_store(self):
    val = 0x1122334455667788
-    dst = UOp.param(0, dtypes.uint32, (6,))
+    dst = UOp.param(0, dtypes.uint32, 6)
    buf = Buffer(Device.DEFAULT, 6, dtypes.uint32, initial_value=bytes(24))
    view = dst.shrink(((1, 5),)).bitcast(dtypes.uint64)  # two stores through one view: it must inline, not get a declared vector-pointer
    run_uops([view.index(0).store(val ^ 0xff), view.index(1).store(val)], [buf])
@@ -249,8 +249,8 @@ class TestLocalAccess(unittest.TestCase):
@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "This only tests assembly backends")
 class TestAssembly(unittest.TestCase):
  def test_bitshift_left(self):
-    g1 = UOp.param(0, dtypes.int32, shape=(3,))
-    out = UOp.param(1, dtypes.int32, shape=(2,))
+    g1 = UOp.param(0, dtypes.int32, 3)
+    out = UOp.param(1, dtypes.int32, 2)
    c1 = UOp.const(2)
    c2 = UOp.const(3)
    l1 = g1.index(c1)
@@ -278,7 +278,7 @@ class TestAssembly(unittest.TestCase):
    self.assertGreaterEqual(len([x.op for x in uops if x.op is Ops.MULACC]), 4)

  def test_mulacc_shl(self):
-    g1 = UOp.param(0, dtypes.int32, shape=(2,))
+    g1 = UOp.param(0, dtypes.int32, 2)
    c1 = UOp.const(0)
    c2 = UOp.const(1)
    expr = g1.index(c1) * UOp.const(4096) + g1.index(c2)
@@ -287,7 +287,7 @@ class TestAssembly(unittest.TestCase):
    self.assertIn(Ops.MULACC, [x.op for x in uops])

  def test_use_cmpeq(self):
-    g = UOp.param(0, dtypes.uint32, shape=(8,))
+    g = UOp.param(0, dtypes.uint32, 8)
    c = UOp.const(7)
    comp = g.index(c).ne(c).ne(True)
    uops = to_uops_list([comp], ren=Device[Device.DEFAULT].renderer)
@@ -43,7 +43,7 @@ def nested_loop_kernel(C:UOp) -> UOp:
  return C[0].store(i[0].load()).sink(arg=KernelInfo(name="nested_loop", opts_to_apply=()))

 def wait_ext_kernel() -> UOp:
-  sig = UOp.param(0, dtypes.int, (1,), volatile=True)
+  sig = UOp.param(0, dtypes.int, 1, volatile=True)
  l = UOp.loop(0)
  v = sig.after(l)[0].load()
  e = v.end(l, v < 1)
@@ -3,10 +3,10 @@ from unittest.mock import patch
 from tinygrad import Device, Tensor
 from tinygrad.device import Buffer
 from tinygrad.dtype import dtypes
-from tinygrad.helpers import getenv
+from tinygrad.helpers import HCQ2
 from tinygrad.runtime.support.hcq2 import HCQ_DEVS, all_devices_in

-@unittest.skipUnless(getenv("HCQ2") and all_devices_in(Device.DEFAULT, HCQ_DEVS), "hcq2 device required")
+@unittest.skipUnless(HCQ2 and all_devices_in(Device.DEFAULT, HCQ_DEVS), "hcq2 device required")
 class TestHCQ2(unittest.TestCase):
  def test_copy_without_copy_queue(self):
    with patch.object(Device[Device.DEFAULT], "has_copy_queue", False):
@@ -28,7 +28,7 @@ def vision_conv_143():
  c48 = (c24&c32).where(c34.index(c45), UOp.const(0.0, dtypes.float))
  c49 = UOp.param(2, dtypes.half, shape=(64, 49, 4))
  c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
-  c63 = UOp.param(3, dtypes.float, (128,))
+  c63 = UOp.param(3, dtypes.float, 128)
  c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5)
  c67 = c0.index((c2*128+c5+c8*4096)).store(c65).end(c8, c2, c5)

@@ -54,7 +54,7 @@ def vision_conv_153():
  c48 = (c24&c32).where(c34.index(c45), UOp.const(0.0, dtypes.float))
  c49 = UOp.param(2, dtypes.half, shape=(128, 49, 4))
  c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
-  c63 = UOp.param(3, dtypes.float, (256,))
+  c63 = UOp.param(3, dtypes.float, 256)
  c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5)
  c67 = c0.index((c2*256+c5+c8*4096)).store(c65).end(c8, c2, c5)

@@ -73,7 +73,7 @@ def dm_conv_172():
  c18 = UOp.range(8, 2, AxisType.REDUCE)
  c23 = UOp.param(2, dtypes.half, shape=(240, 128, 4))
  c35 = c5.index((c7*4+c10+c13*128+c18*1536))*c23.index((c10*4+c2%4+c7*16+c2//4*512))
-  c37 = UOp.param(3, dtypes.float, (960,))
+  c37 = UOp.param(3, dtypes.float, 960)
  c39 = c35.reduce(c7, c10, arg=Ops.ADD)+c37.index(c2)
  c50 = (1.0+((c39+0.044708251953125*(c39*(c39*c39)))*-2.3021129851685216).exp2()).reciprocal()*c39
  c53 = c50.reduce(c18, c13, arg=Ops.ADD)*0.010416666666666666
@@ -7,8 +7,8 @@ BENCHMARK_OPS = {Ops.INDEX, Ops.STAGE}

@functools.cache
 def create_uop(a:int) -> UOp:
-  op, dtype, src, arg, *rest = trace.uop_fields[a]
-  return UOp(op, dtype, tuple(create_uop(s) for s in src), arg, *rest)
+  op, src, arg, *rest = trace.uop_fields[a]
+  return UOp(op, tuple(create_uop(s) for s in src), arg, *rest)

 if __name__ == "__main__":
  # load rewrite trace
@@ -23,7 +23,7 @@ class TestGPUCrash(unittest.TestCase):
    cls.is_cdna = cls.dev.target[0] < 10
    ins = importlib.import_module('tinygrad.runtime.autogen.amd.' + ('cdna' if cls.is_cdna else 'rdna3') + '.ins')
    for rdna3_name, cdna3_name in RDNA3_CDNA3_MAP.items():
-      setattr(cls, rdna3_name, getattr(ins, cdna3_name if cls.is_cdna else rdna3_name))
+      setattr(cls, rdna3_name, staticmethod(getattr(ins, cdna3_name if cls.is_cdna else rdna3_name)))

  def setUp(self):
    # Verify device works before each test
@@ -36,8 +36,8 @@ class TestGPUCrash(unittest.TestCase):

  def _run_insts(self, insts: list[Inst]):
    buf = UOp.new_buffer("AMD", 64, dtypes.uint8)
-    sink = UOp.sink(UOp.param(0, dtypes.uint8, (64,), device="AMD"), UOp.special(1, "lidx0"), arg=KernelInfo("test"))
-    prg = UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=i) for i in insts))))
+    sink = UOp.sink(UOp.param(0, dtypes.uint8, 64, device="AMD"), UOp.special(1, "lidx0"), arg=KernelInfo("test"))
+    prg = UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=(i, dtypes.void)) for i in insts))))
    run_linear(UOp(Ops.LINEAR, src=(prg.call(buf),)), wait=True)

  def _assert_gpu_fault(self, func):
@@ -54,6 +54,12 @@ class TestMainOnnxOps(TestOnnxOps):
    outputs = ["squeezed"]
    self.helper_test_single_op("Squeeze", inputs, attributes, outputs)

+  def test_mean_variance_normalization_axes(self):
+    inputs = {"x": np.random.randn(2, 3, 4, 5).astype(np.float32)}
+    attributes = {"axes": [2, 3]}
+    outputs = ["out"]
+    self.helper_test_single_op("MeanVarianceNormalization", inputs, attributes, outputs)
+
  def test_conv(self):
    # test VALID auto_pad
    inputs = {
@@ -235,6 +241,15 @@ class TestMainOnnxOps(TestOnnxOps):
    outputs = ["y"]
    self.helper_test_single_op("MaxUnpool", inputs, attributes, outputs)

+  def test_maxunpool_pads(self):
+    # per-axis pads shrink the output: spatial dim is (i-1)*stride + kernel - pad_begin - pad_end -> (2, 4), and indices index into that output
+    # NOTE: indices must be in bounds of that output; ORT aborts the process on out-of-bounds indices
+    xT = np.array([[[[5, 6], [7, 8]]]], dtype=np.float32)
+    xI = np.array([[[[0, 3], [4, 7]]]], dtype=np.int64)
+    inputs = {"x": xT, "indices": xI}
+    attributes = {"kernel_shape": [2, 2], "strides": [2, 2], "pads": [1, 0, 1, 0]}
+    self.helper_test_single_op("MaxUnpool", inputs, attributes, ["y"])
+
  def test_averagepool_3d_dilations_large_count_include_pad_is_1_ceil_mode_is_True(self):
    # https://github.com/onnx/onnx/blob/main/docs/Operators.md#examples-13
    inputs = {"x": np.random.randn(1, 1, 32, 32, 32).astype(np.float32)}
@@ -1,5 +1,5 @@
 import unittest, time, itertools
-from tinygrad import Tensor, Context
+from tinygrad import Tensor, Context, dtypes

 class TestScheduleScaling(unittest.TestCase):
  """Test that .schedule() scales linearly with graph size (no O(n^2) behavior)."""
@@ -138,7 +138,7 @@ class TestScheduleScaling(unittest.TestCase):
    def custom_kernel_assign(n):
      def custom_asm(out):
        return UOp(Ops.PROGRAM, src=(UOp.sink(out, arg=KernelInfo(f"fxn_{next(count)}")),
-                                     UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=s_nop(i)) for i in range(n*8)))))
+                                     UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=(s_nop(i), dtypes.void)) for i in range(n*8)))))
      call = Tensor.custom_kernel(Tensor.empty(1), fxn=custom_asm)[0]
      return Tensor.cat(*[Tensor.empty(1).assign(call+i) for i in range(n)])
    self._assert_linear(custom_kernel_assign, n_small=50, n_large=500)
@@ -122,7 +122,7 @@ def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None, vals:tuple
  for buf_dt, data in inputs or []:
    bufs.append(buf:=allocator.alloc(len(data) * buf_dt.itemsize))
    allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + (buf_dt.fmt or ""), *data)))
-  g = UOp.param(0, uop.dtype, (1,))
+  g = UOp.param(0, uop.dtype, 1)
  prg = to_program(UOp.store(g.index(UOp.const(0)), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON")))
  prog = dev.runtime(prg.to_elf())
  prog(out_buf:=allocator.alloc(uop.dtype.itemsize), *bufs, vals=vals)
@@ -1,20 +1,15 @@
-import pathlib, re, ctypes, mmap, collections, functools, copy, os
+import ctypes, mmap, collections, functools, copy, os
 from tinygrad.runtime.autogen import kfd, amdgpu_drm, libc
 import tinygrad.runtime.autogen.am.am as am
 from tinygrad.helpers import from_mv
 from test.mockgpu.driver import VirtDriver, VirtFileDesc, TextFileDesc, DirFileDesc, VirtFile
 from test.mockgpu.amd.amdgpu import AMDGPU, gpu_props, GFX_TARGET_VERSION, MOCKGPU_ARCH

-def ioctls_from_header():
-  # hdrpy = (pathlib.Path(__file__).parent.parent.parent.parent / "tinygrad" / "runtime" / "autogen" / "kfd.py").read_text()
-  # pattern = r'# (AMDKFD_IOC_[A-Z0-9_]+)\s=\s_(IOW?R?).*\(( 0x[0-9a-fA-F]+) ,\s+struct\s([A-Za-z0-9_]+)\s+\)'
-  # matches = re.findall(pattern, hdrpy, re.MULTILINE)
-  hdr = (pathlib.Path(__file__).parent.parent.parent.parent / "extra" / "hip_gpu_driver" / "kfd_ioctl.h").read_text().replace("\\\n", "")
-  pattern = r'#define\s+(AMDKFD_IOC_[A-Z0-9_]+)\s+AMDKFD_(IOW?R?)\((0x[0-9a-fA-F]+),\s+struct\s([A-Za-z0-9_]+)\)'
-  matches = re.findall(pattern, hdr, re.MULTILINE)
-  return type("KFD_IOCTLS", (object, ), {name: int(nr, 0x10) for name, _, nr, _ in matches}), \
-         {int(nr, 0x10): getattr(kfd, "struct_"+sname, None) for name, idir, nr, sname in matches}
-kfd_ioctls, kfd_headers = ioctls_from_header()
+def _ioctl_nr(ioctl: functools.partial) -> int: return ioctl.args[2]
+
+kfd_ioctl_info = {
+  _ioctl_nr(ioctl): (name, ioctl.args[3]) for name, ioctl in vars(kfd).items()
+  if name.startswith("AMDKFD_IOC_") and isinstance(ioctl, functools.partial)}

 class KFDFileDesc(VirtFileDesc):
  def __init__(self, fd, driver):
@@ -116,42 +111,44 @@ class AMDDriver(VirtDriver):

  def kfd_ioctl(self, req, argp):
    nr = req & 0xFF
-    struct = kfd_headers[nr].from_address(argp)
+    if nr not in kfd_ioctl_info: raise RuntimeError(f"unknown kfd ioctl, {nr} unknown")
+    name, struct_type = kfd_ioctl_info[nr]
+    struct = struct_type.from_address(argp)

-    if nr == kfd_ioctls.AMDKFD_IOC_ACQUIRE_VM: pass
-    elif nr == kfd_ioctls.AMDKFD_IOC_RUNTIME_ENABLE: pass
-    elif nr == kfd_ioctls.AMDKFD_IOC_GET_VERSION:
+    if nr == _ioctl_nr(kfd.AMDKFD_IOC_ACQUIRE_VM): pass
+    elif nr == _ioctl_nr(kfd.AMDKFD_IOC_RUNTIME_ENABLE): pass
+    elif nr == _ioctl_nr(kfd.AMDKFD_IOC_GET_VERSION):
      struct.major_version = 1
      struct.minor_version = 14
-    elif nr == kfd_ioctls.AMDKFD_IOC_ALLOC_MEMORY_OF_GPU:
+    elif nr == _ioctl_nr(kfd.AMDKFD_IOC_ALLOC_MEMORY_OF_GPU):
      if struct.gpu_id not in self.gpus: return -1
      struct.handle = self._alloc_handle()
      self.object_by_handle[struct.handle] = copy.deepcopy(struct) # save memory struct to know what mem it is
      # Track signal memory (uncached + coherent) - progress queues when written to
      if struct.flags & kfd.KFD_IOC_ALLOC_MEM_FLAGS_UNCACHED:
        self.track_address(struct.va_addr, struct.va_addr + struct.size, lambda mv,off: None, lambda mv, off: self._emulate_execute())
-    elif nr == kfd_ioctls.AMDKFD_IOC_FREE_MEMORY_OF_GPU:
+    elif nr == _ioctl_nr(kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU):
      self.object_by_handle.pop(struct.handle)
-    elif nr == kfd_ioctls.AMDKFD_IOC_MAP_MEMORY_TO_GPU:
+    elif nr == _ioctl_nr(kfd.AMDKFD_IOC_MAP_MEMORY_TO_GPU):
      dev_ids = (ctypes.c_int32 * struct.n_devices).from_address(struct.device_ids_array_ptr)
      for i in range(struct.n_devices):
        gpu = self.gpus[dev_ids[i]]
        mem_obj = self.object_by_handle[struct.handle]
        gpu.map_range(mem_obj.va_addr, mem_obj.size)
        struct.n_success = i + 1
-    elif nr == kfd_ioctls.AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU:
+    elif nr == _ioctl_nr(kfd.AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU):
      dev_ids = (ctypes.c_int32 * struct.n_devices).from_address(struct.device_ids_array_ptr)
      for i in range(struct.n_devices):
        gpu = self.gpus[dev_ids[i]]
        mem_obj = self.object_by_handle[struct.handle]
        gpu.unmap_range(mem_obj.va_addr, mem_obj.size)
        struct.n_success = i + 1
-    elif nr == kfd_ioctls.AMDKFD_IOC_CREATE_EVENT:
+    elif nr == _ioctl_nr(kfd.AMDKFD_IOC_CREATE_EVENT):
      struct.event_slot_index = self._alloc_next_event_slot()
      struct.event_id = struct.event_slot_index

      if struct.event_type == kfd.KFD_IOC_EVENT_MEMORY: self.mmu_event_ids.append(struct.event_id)
-    elif nr == kfd_ioctls.AMDKFD_IOC_CREATE_QUEUE:
+    elif nr == _ioctl_nr(kfd.AMDKFD_IOC_CREATE_QUEUE):
      gpu = self.gpus[struct.gpu_id]
      if struct.queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA:
        gpu.add_sdma_queue(struct.ring_base_address, struct.ring_size, struct.read_pointer_address, struct.write_pointer_address)
@@ -162,7 +159,7 @@ class AMDDriver(VirtDriver):
      # Track writes to doorbell, calling callback
      struct.doorbell_offset = self._alloc_doorbell(struct.gpu_id)
      self.track_address(struct.doorbell_offset, struct.doorbell_offset + 8, lambda mv,off: None, lambda mv, off: self._emulate_execute())
-    elif nr == kfd_ioctls.AMDKFD_IOC_WAIT_EVENTS:
+    elif nr == _ioctl_nr(kfd.AMDKFD_IOC_WAIT_EVENTS):
      evs = (kfd.struct_kfd_event_data * struct.num_events).from_address(struct.events_ptr)
      for ev in evs:
        if ev.event_id in self.mmu_event_ids and "MOCKGPU_EMU_FAULTADDR" in os.environ:
@@ -170,11 +167,7 @@ class AMDDriver(VirtDriver):
          ev.memory_exception_data.va = int(os.environ["MOCKGPU_EMU_FAULTADDR"], 16)
          ev.memory_exception_data.failure.NotPresent = 1
    else:
-      name = "unknown"
-      for k,v in kfd_ioctls.__dict__.items():
-        if nr == v: name = k
-      assert False, f"unknown kfd ioctl, {nr} {name}"
-      exit(1)
+      raise RuntimeError(f"unsupported kfd ioctl, {nr} {name}")
    return 0

  def _emulate_execute(self):
@@ -321,10 +321,10 @@ def _int_clamp(op_name: str, srcs: dict) -> UOp | None:
 class _Ctx:
  """Context for instruction compilation - holds buffers and helpers."""
  __slots__ = ('inst_size', 'dyn_fields', '_axis_id', 'wave_size', 'vgpr', 'accvgpr')
-  sgpr = UOp.param(0, dtypes.uint32, (SGPR_COUNT,))
-  vmem = UOp.param(2, dtypes.uint32, (1 << 46,))
-  lds = UOp.param(3, dtypes.uint32, (16384,))
-  scratch = UOp.param(4, dtypes.uint8, (1 << 30,))
+  sgpr = UOp.param(0, dtypes.uint32, SGPR_COUNT)
+  vmem = UOp.param(2, dtypes.uint32, 1 << 46)
+  lds = UOp.param(3, dtypes.uint32, 16384)
+  scratch = UOp.param(4, dtypes.uint8, 1 << 30)
  # Cache PARAM UOps by wave_size so all _Ctx instances with same wave_size share identical UOp references
  _vgpr_cache: dict[int, UOp] = {}
  _accvgpr_cache: dict[int, UOp] = {}
@@ -332,10 +332,10 @@ class _Ctx:
  def __init__(self, inst_size: int, wave_size: int = 32):
    self.inst_size, self._axis_id, self.wave_size = inst_size, 0, wave_size
    self.dyn_fields: list[tuple[int, int]] = []  # (lo, hi) of fields read dynamically
-    if wave_size not in _Ctx._vgpr_cache: _Ctx._vgpr_cache[wave_size] = UOp.param(1, dtypes.uint32, (256 * wave_size,))
+    if wave_size not in _Ctx._vgpr_cache: _Ctx._vgpr_cache[wave_size] = UOp.param(1, dtypes.uint32, 256 * wave_size)
    self.vgpr = _Ctx._vgpr_cache[wave_size]
    if wave_size == 64:
-      if wave_size not in _Ctx._accvgpr_cache: _Ctx._accvgpr_cache[wave_size] = UOp.param(5, dtypes.uint32, (256 * wave_size,))
+      if wave_size not in _Ctx._accvgpr_cache: _Ctx._accvgpr_cache[wave_size] = UOp.param(5, dtypes.uint32, 256 * wave_size)
      self.accvgpr = _Ctx._accvgpr_cache[wave_size]
    else:
      self.accvgpr = self.vgpr
@@ -23,6 +23,7 @@ from tinygrad.codegen.decomp.dtype import f2f
 VarVal = UOp | tuple[str, list[str], str]

 def _const(dt, v): return UOp.const(v, dt)
+def _single_value(v: UOp): return v.vmin if v.vmin == v.vmax else None
 def _u32(v): return _const(dtypes.uint32, v)
 def _u64(v): return _const(dtypes.uint64, v)
 def _to_u32(v): return v if v.dtype == dtypes.uint32 else v.bitcast(dtypes.uint32) if v.dtype.itemsize == 4 else v.cast(dtypes.uint32)
@@ -70,8 +71,8 @@ def _expr_bits(v: UOp) -> int:
  if v.op in (Ops.AND, Ops.XOR):
    widths: list[int] = []
    for src in v.src:
-      if src.op == Ops.CONST and isinstance(src.val, int) and src.val > 0 and (src.val & (src.val + 1)) == 0:
-        widths.append(src.val.bit_length())
+      if isinstance(sv:=_single_value(src), int) and sv > 0 and (sv & (sv + 1)) == 0:
+        widths.append(sv.bit_length())
    if widths: return max(widths)
  return v.dtype.bitsize

@@ -159,9 +160,9 @@ def _minmax_reduce(is_max: bool, dt, *args: UOp) -> UOp:
 def _find_two_pi_mul(x):
  if x.op != Ops.MUL or len(x.src) != 2: return None
  for i, s in enumerate(x.src):
-    if s.op == Ops.CONST and abs(s.val - 6.283185307179586) < 1e-5: return (x.src[1-i], 6.283185307179586)
+    if (sv:=_single_value(s)) is not None and abs(sv - 6.283185307179586) < 1e-5: return (x.src[1-i], 6.283185307179586)
    if s.op == Ops.MUL and len(s.src) == 2:
-      vals = [ss.val for ss in s.src if ss.op == Ops.CONST] + [ss.src[0].val for ss in s.src if ss.op == Ops.CAST and ss.src[0].op == Ops.CONST]
+      vals = [sv for ss in s.src if (sv:=_single_value(ss)) is not None]
      if len(vals) == 2 and abs(vals[0] * vals[1] - 6.283185307179586) < 1e-5: return (x.src[1-i], vals[0] * vals[1])
  return None

@@ -178,7 +179,7 @@ def _trig_reduce(x, phase=0.0):

 def _signext(val: UOp) -> UOp:
  for bits, mask, ext in [(4, 0xF, 0xFFFFFFF0), (8, 0xFF, 0xFFFFFF00), (16, 0xFFFF, 0xFFFF0000)]:
-    if (val.op == Ops.AND and len(val.src) == 2 and val.src[1].op == Ops.CONST and val.src[1].val == mask) or val.dtype.itemsize == bits // 8:
+    if (val.op == Ops.AND and len(val.src) == 2 and _single_value(val.src[1]) == mask) or val.dtype.itemsize == bits // 8:
      v32 = val.cast(dtypes.uint32) if val.dtype != dtypes.uint32 else val
      sb = (v32 >> _u32(bits - 1)) & _u32(1)
      return sb.ne(_u32(0)).where(v32 | _u32(ext), v32).cast(dtypes.int)
@@ -549,7 +550,7 @@ class Parser:
        if not dtypes.is_int(right.dtype): right = right.cast(dtypes.uint32)
        return (left >> right) if op == '>>' else (left << right)
      case '+' | '-':
-        if op == '-' and left.op == Ops.CONST and right.op == Ops.CONST: return _const(left.dtype, left.val - right.val)
+        if op == '-' and (lv:=_single_value(left)) is not None and (rv:=_single_value(right)) is not None: return _const(left.dtype, lv - rv)
        return (left + right) if op == '+' else (left - right)
      case '*' | '/':
        # Integer promotion: promote 16-bit integers to 32-bit before multiply to avoid overflow
@@ -559,7 +560,7 @@ class Parser:
          left, right = left.cast(pdt), right.cast(pdt)
        if op == '*': return left * right
        return (left // right) if dtypes.is_int(left.dtype) else (left / right)
-      case '**': return UOp(Ops.EXP2, src=(right.cast(left.dtype),)) if left.op == Ops.CONST and left.val == 2.0 else left
+      case '**': return UOp(Ops.EXP2, src=(right.cast(left.dtype),)) if _single_value(left) == 2.0 else left

  _PREC = [('||',), ('&&',), ('|',), ('^',), ('&',), ('==', '!=', '<>'), ('>=', '<=', '>', '<'), ('>>', '<<'), ('+', '-'), ('*', '/'), ('**',)]

@@ -581,8 +582,8 @@ class Parser:
      return inner.eq(_const(inner.dtype, 0))
    if self.try_eat_val('-', 'OP'):
      inner = self.unary()
-      if inner.op == Ops.CONST:
-        return _const(dtypes.int if inner.dtype == dtypes.uint32 else inner.dtype, -inner.val)
+      if (v:=_single_value(inner)) is not None:
+        return _const(dtypes.int if inner.dtype == dtypes.uint32 else inner.dtype, -v)
      return inner.neg()
    if self.try_eat_val('+', 'OP'): return self.unary()
    return self.postfix()
@@ -721,15 +722,13 @@ class Parser:
      self.eat('OP')
      width = self.parse()
      self.eat('RBRACKET')
-      if width.op == Ops.CONST:
-        w = int(width.val)
+      if isinstance(w:=_single_value(width), int):
        return (base >> _to_u32(first)) & _const(base.dtype, (1 << w) - 1)
      return base
    if self.try_eat('COLON'):
      second = self.parse()
      self.eat('RBRACKET')
-      if first.op == Ops.CONST and second.op == Ops.CONST:
-        a, b = int(first.val), int(second.val)
+      if isinstance(a:=_single_value(first), int) and isinstance(b:=_single_value(second), int):
        if a < b: return _bitreverse(base, b - a + 1)
        hi, lo = a, b
        if lo >= base.dtype.itemsize * 8:
@@ -750,8 +749,7 @@ class Parser:
      dt_suffix = DTYPES.get(self.eat('IDENT').val, dtypes.uint32)
    if var_name is None:
      var_name = self._find_var_name(base)
-    if first.op == Ops.CONST:
-      idx = int(first.val)
+    if isinstance(idx:=_single_value(first), int):
      # Check for array element (var@idx)
      if var_name and f'{var_name}@{idx}' in self.vars:
        v = self.vars[f'{var_name}@{idx}']
@@ -924,7 +922,7 @@ class Parser:

  def _coerce_cmp(self, l: UOp, r: UOp) -> tuple[UOp, UOp]:
    if l.dtype != r.dtype:
-      if r.dtype == dtypes.int and r.op == Ops.CONST and r.val < 0: l = l.cast(dtypes.int)
+      if r.dtype == dtypes.int and isinstance(rv:=_single_value(r), int) and rv < 0: l = l.cast(dtypes.int)
      else: r = r.cast(l.dtype)
    return l, r

@@ -419,7 +419,7 @@ class TestAutoCastType(unittest.TestCase):
    self.check_where_alternate_input_other(3, True, dtypes.weakint)

  def test_where_non_bool_cond_raises(self):
-    with self.assertRaises(RuntimeError): Tensor([1, 0, 2]).where(1, 0)
+    with self.assertRaises(RuntimeError): Tensor([1, 0, 2]).where(1, 0).dtype
    self.check_where_alternate_input_other(False, True, dtypes.bool)

  @given(strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes))
@@ -107,7 +107,7 @@ class TestGroupedDims(unittest.TestCase):

  def test_global_prod_max(self):
    g, l = UOp.range(256, 0, AxisType.GLOBAL), UOp.range(256, 1, AxisType.LOCAL)
-    sink = UOp.param(0, dtypes.float, (512,)).index(g + l).store(UOp.const(1.0)).end(g, l).sink(arg=KernelInfo())
+    sink = UOp.param(0, dtypes.float, 512).index(g + l).store(UOp.const(1.0)).end(g, l).sink(arg=KernelInfo())
    class R(Renderer): global_max, local_max, global_prod_max = (256, 256, 256), (128, 128, 128), (128, 128, 128)
    specials = [u for u in add_gpudims(R(Target()), sink).toposort() if u.op is Ops.SPECIAL]
    self.assertGreater(len([s for s in specials if "lidx" in s.arg]), 1)
@@ -7,14 +7,14 @@ from tinygrad.codegen import to_program

 class TestLinearizerFailures(unittest.TestCase):
  def test_fail_1(self):
-    c0 = UOp.param(0, dtypes.float, (64,))
+    c0 = UOp.param(0, dtypes.float, 64)
    c1 = UOp.range(UOp.const(2), 1, AxisType.WEAK)
    c2 = UOp.range(UOp.const(32), 2, AxisType.WEAK)
    c3 = ((c1*UOp.const(32))+c2)
-    c4 = UOp.param(1, dtypes.float, (163840,))
+    c4 = UOp.param(1, dtypes.float, 163840)
    c5 = UOp.range(UOp.const(2560), 0, AxisType.REDUCE)
    c6 = c4.index(((((((c5//UOp.const(8))%UOp.const(8))*UOp.const(8))+(c5%UOp.const(8)))+(((c2*UOp.const(40))+(c5//UOp.const(64)))*UOp.const(64)))+(c1*UOp.const(81920))))
-    c7 = UOp.param(2, dtypes.float, (64,))
+    c7 = UOp.param(2, dtypes.float, 64)
    c8 = c7.index(c3)
    c9 = ((((c6+(c8*UOp.const(-1.0)))*(c6+(c8*UOp.const(-1.0)))).reduce(c5, arg=Ops.ADD)*UOp.const(0.000390625))+UOp.const(1e-05)).sqrt().reciprocal()
    c10 = c0.index(c3).store(c9).end(c1, c2)
@@ -95,6 +95,26 @@ class TestLLMTokenizer(unittest.TestCase):
    self.assertEqual(template.end_turn(), "[/INST]")
    self.assertEqual(template.role("assistant"), "")

+  def test_tekken_gpt4o_split(self):
+    split = {p: SimpleTokenizer({}, {}, p)._split_to_word.findall for p in ("tekken", "gpt-4o")}
+    shared = {
+      "HelloWorld": ["Hello", "World"],
+      " ÜNICODE": [" ÜNICODE"],  # Ü: non-ascii upper joins the run
+      "é café": ["é", " café"],  # first é is e + U+0301 combining acute (NFD)
+      "เพื่อน วิ": ["เพื่อน", " วิ"],  # thai vowel marks stay in the word
+      "a/b\r\n x": ["a", "/b", "\r\n", " x"],  # punct tail eats /
+    }
+    for s, want in shared.items():
+      self.assertEqual(split["tekken"](s), want, f"tekken {s!r}")
+      self.assertEqual(split["gpt-4o"](s), want, f"gpt-4o {s!r}")
+    differ = [
+      ("12345", list("12345"), ["123", "45"]),  # digits: tekken single, o200k groups {1,3}
+      ("it's I'M don'T", ["it", "'s", " I", "'M", " don", "'T"], ["it's", " I'M", " don'T"]),  # contraction: o200k inline suffix
+    ]
+    for s, tk, go in differ:
+      self.assertEqual(split["tekken"](s), tk, f"tekken {s!r}")
+      self.assertEqual(split["gpt-4o"](s), go, f"gpt-4o {s!r}")
+
  def test_stream_decoder(self):
    """stream_decoder buffers incomplete UTF-8: token 25677 has 3/4 of emoji, token 138 completes it."""
    bs = [*range(33, 127), *range(161, 173), *range(174, 256)]
@@ -40,7 +40,7 @@ class TestPatternMatcher(unittest.TestCase):
      assert len(x.src) == 0
      return UOp(Ops.CONST, src=(UOp(Ops.CONST),))
    matcher = PatternMatcher([(UPat(Ops.CONST, src=(), name="x"), fxn)])
-    c1 = UOp(Ops.CONST, dtypes.float, arg=1.0)
+    c1 = UOp(Ops.CONST, arg=1.0)
    # second rewrite shouldn't match anything
    c1 = matcher.rewrite(c1)
    c1 = matcher.rewrite(c1)
@@ -15,7 +15,7 @@ def simplify_valid_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move
 def simplify_image_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move_where_on_load+indexing_simplify, name="simplify_image_idx")

 def get_gated_load_uop(valid:UOp, idx:UOp):
-  return UOp.param(0, dtypes.float, (1024,)).index(idx.valid(valid)).load()
+  return UOp.param(0, dtypes.float, 1024).index(idx.valid(valid)).load()

 def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UOp]):
  return UOp.param(0, dtypes.float, image_shape).index(idx[1].valid(valid), idx[0].valid(valid)).load()
@@ -505,7 +505,7 @@ class TestDropTrueGate(unittest.TestCase):
    from tinygrad.codegen.late.coalesce import indexing_simplify
    from tinygrad.uop.ops import graph_rewrite
    from tinygrad.uop.symbolic import sym
-    buf = UOp.param(0, dtypes.int, (1,))
+    buf = UOp.param(0, dtypes.int, 1)
    idx = UOp.const(0)
    true_gate = UOp.const(True)
    index_with_gate = buf.index(idx.valid(true_gate))
@@ -551,7 +551,7 @@ class TestRangeShrink(unittest.TestCase):
    # one load guards r < 4, but another load uses r without a gate -> no shrink
    r = Range(0, 204)
    load1 = get_gated_load_uop(r < UOp.const(4), r)
-    load2 = UOp.param(1, dtypes.float, (204,)).index(r).load()
+    load2 = UOp.param(1, dtypes.float, 204).index(r).load()
    ranges = self.get_ranges(UOp.sink(load1, load2))
    self.assert_range_end(ranges, 204)

@@ -575,7 +575,7 @@ class TestRangeShrink(unittest.TestCase):
    from tinygrad.dtype import Invalid
    r = Range(0, 204)
    x = (r < 4).where(UOp.const(1.0), Invalid)
-    ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r < 4).where(x, Invalid)).sink())
+    ranges = self.get_ranges(UOp.param(0, dtypes.float, 204).index(r).store((r < 4).where(x, Invalid)).sink())
    self.assert_range_end(ranges, 4)

  def test_range_shrink_store_where_invalid_flipped(self):
@@ -583,7 +583,7 @@ class TestRangeShrink(unittest.TestCase):
    from tinygrad.dtype import Invalid
    r = Range(0, 204)
    x = (r < 4).where(UOp.const(1.0), Invalid)
-    ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r >= 4).where(Invalid, x)).sink())
+    ranges = self.get_ranges(UOp.param(0, dtypes.float, 204).index(r).store((r >= 4).where(Invalid, x)).sink())
    self.assert_range_end(ranges, 4)

 if __name__ == '__main__':
@@ -9,7 +9,7 @@ class TestTranscendentalFunctions(unittest.TestCase):
  def test_payne_hanek_reduction(self):
    # TODO: Test constant input when constant folding is fixed (or maybe test both variants)
    # Load input value from a buffer to prevent constant folding
-    input_buf = UOp.param(1, dtypes.double, (1,))
+    input_buf = UOp.param(1, dtypes.double, 1)
    loaded_value = input_buf.index(UOp.const(0)).load()
    def eval_payne_hanek_reduction(v:float) -> tuple[float, int]:
      return tuple(eval_uop(u, [(dtypes.float64, [v])]) for u in payne_hanek_reduction(loaded_value))
@@ -202,6 +202,11 @@ class TestUOpGraph(unittest.TestCase):
    invalid_lane_mul = next(u for u in out.src[0].toposort() if u.op is Ops.MUL)
    self.assertIs(invalid_lane_mul.dtype, dtypes.bool)

+  def test_devectorize_zero_sized_scalar_expand(self):
+    from tinygrad.codegen import devectorizer2
+    expanded = UOp.const(1.0).reshape(1, 1).expand(0, 3)
+    self.assertEqual(graph_rewrite(expanded, devectorizer2).shape, (0, 3))
+
  def test_gep_vec_const_fold(self):
    for vec_size in [2, 4, 8]:
      consts = [UOp.const(float(i), dtypes.float) for i in range(vec_size)]
@@ -209,8 +214,8 @@ class TestUOpGraph(unittest.TestCase):
      for i, const in enumerate(consts): self.assertIs(vec.index(i), const)

  def test_cast_alu_fold(self):
-    d0 = UOp.param(0, dtypes.bool, (1,))
-    d1 = UOp.param(1, dtypes.int, (1,))
+    d0 = UOp.param(0, dtypes.bool, 1)
+    d1 = UOp.param(1, dtypes.int, 1)
    idx = UOp.const(0)
    ld = d1.index(idx)
    alu = (ld<1).cast(dtypes.bool)
@@ -219,8 +224,8 @@ class TestUOpGraph(unittest.TestCase):
    self.assertEqual(len([x for x in uops if x.op is Ops.CAST and x.src[0].op is not Ops.CONST]), 0)

  def test_double_cast_fold(self):
-    d0 = UOp.param(0, dtypes.float, (1,))
-    d1 = UOp.param(1, dtypes.int, (1,))
+    d0 = UOp.param(0, dtypes.float, 1)
+    d1 = UOp.param(1, dtypes.int, 1)
    idx = UOp.const(0, dtypes.int)
    ld = d1.index(idx)
    alu = ld.cast(dtypes.float).cast(dtypes.float)
@@ -238,7 +243,7 @@ class TestUOpGraph(unittest.TestCase):

  def test_bitcast_to_same_dtype_fold(self):
    for dt in dtypes.ints + dtypes.floats + (dtypes.bool,):
-      d0 = UOp.param(0, dt, (1,))
+      d0 = UOp.param(0, dt, 1)
      v = d0.index(UOp.const(0))
      uops = to_uops_list([v.bitcast(dt)])
      self.assertEqual(len([x for x in uops if x.op is Ops.BITCAST and x.dtype is dt]), 0, f"dtype = {dt}")
@@ -250,10 +255,10 @@ class TestUOpGraph(unittest.TestCase):

  def test_where_on_gated_load_fold(self):
    ridx0 = UOp.range(100, 0)
-    d0 = UOp.param(0, dtypes.long, (100,))
+    d0 = UOp.param(0, dtypes.long, 100)
    ld = d0.index(ridx0.valid(ridx0<50))
    w = (ridx0<50).where(ld, 5)
-    out = UOp.param(1, dtypes.long, (100,))
+    out = UOp.param(1, dtypes.long, 100)
    uops = to_uops_list([out.index(ridx0).store(w)])
    expected = full_rewrite(UOp.const(5, dtypes.long).sink()).src[0]
    for u in uops:
@@ -262,7 +267,7 @@ class TestUOpGraph(unittest.TestCase):

  def test_where_on_gated_load_folds_swapped_branches(self):
    ridx0 = UOp.range(100, 0)
-    d0 = UOp.param(0, dtypes.long, (100,))
+    d0 = UOp.param(0, dtypes.long, 100)
    ld = d0.index(ridx0.valid((ridx0<50).logical_not()))
    w = (ridx0<50).where(5, ld)
    uops = to_uops_list([w])
@@ -273,11 +278,11 @@ class TestUOpGraph(unittest.TestCase):

  def test_where_on_gated_load_with_cast(self):
    ridx0 = UOp.range(100, 0)
-    d0 = UOp.param(0, dtypes.int, (100,))
+    d0 = UOp.param(0, dtypes.int, 100)
    gate_idx = ridx0.valid((ridx0<50))
    ld = d0.index(gate_idx).cast(dtypes.float)
    w = (ridx0<50).where(ld, 5.0)
-    out = UOp.param(1, dtypes.float, (100,))
+    out = UOp.param(1, dtypes.float, 100)
    uops = to_uops_list([out.index(ridx0).store(w)])
    expected = full_rewrite(UOp.const(5, dtypes.int).sink()).src[0]
    for u in uops:
@@ -286,27 +291,27 @@ class TestUOpGraph(unittest.TestCase):

  def test_where_on_casted_gated_load_extra_cond(self):
    ridx0 = UOp.range(100, 0)
-    d0 = UOp.param(0, dtypes.float, (100,))
+    d0 = UOp.param(0, dtypes.float, 100)
    ld = d0.index(ridx0.valid(ridx0<50))
    w = ((ridx0<50) & (ridx0>30)).where(ld, UOp.const(0.0)).cast(dtypes.half)
-    out = UOp.param(1, dtypes.half, (100,))
+    out = UOp.param(1, dtypes.half, 100)
    uops = to_uops_list([out.index(ridx0).store(w)])
    for u in uops:
      assert u.op is not Ops.WHERE

  def test_where_on_casted_gated_load_extra_cond_swapped(self):
    ridx0 = UOp.range(100, 0)
-    d0 = UOp.param(0, dtypes.float, (100,))
+    d0 = UOp.param(0, dtypes.float, 100)
    ld = d0.index(ridx0.valid(ridx0<50))
    w = ((ridx0<50) & (ridx0>30)).where(UOp.const(0.0), ld).cast(dtypes.half)
-    out = UOp.param(1, dtypes.half, (100,))
+    out = UOp.param(1, dtypes.half, 100)
    uops = to_uops_list([out.index(ridx0).store(w)])
    for u in uops:
      assert u.op is not Ops.WHERE

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

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

  def test_fold_gated_load(self):
-    glbl0 = UOp.param(0, dtypes.int, (1,))
-    glbl1 = UOp.param(1, dtypes.int, (1,))
-    glbl2 = UOp.param(2, dtypes.int, (1,))
+    glbl0 = UOp.param(0, dtypes.int, 1)
+    glbl1 = UOp.param(1, dtypes.int, 1)
+    glbl2 = UOp.param(2, dtypes.int, 1)
    idx = UOp.const(0)
    ld0 = glbl1.index(UOp.invalid())
    ld1 = glbl2.index(idx.valid(UOp.const(True)))
@@ -367,7 +372,7 @@ class TestUOpGraph(unittest.TestCase):
    self.assertEqual(len([u for u in uops if u.op is Ops.LOAD]), 1)

  def test_fold_gated_load_local(self):
-    glbl0 = UOp.param(0, dtypes.int, (16,))
+    glbl0 = UOp.param(0, dtypes.int, 16)
    smem = UOp.placeholder((18,), dtypes.int, slot=0, addrspace=AddrSpace.LOCAL)
    lidx = UOp.special(16, "lidx0")
    st = smem.index(lidx).store(glbl0.index(lidx).load())
@@ -380,7 +385,7 @@ class TestUOpGraph(unittest.TestCase):
    self.assertEqual(len([u for u in uops if u.op is Ops.LOAD]), 2)

  def test_fold_gated_store(self):
-    glbl = UOp.param(0, dtypes.int, (1,))
+    glbl = UOp.param(0, dtypes.int, 1)
    idx0 = UOp.const(0)
    val = UOp.const(42)
    st0 = glbl.index(UOp.invalid()).store(val)
@@ -420,7 +425,7 @@ class TestReduceCollapse(unittest.TestCase):

  def test_reduce_shapeless_const_unroll(self):
    """a REDUCE over a shapeless CONST (e.g. x*0 folded late in codegen) must collapse before the expander"""
-    out = UOp.param(0, dtypes.float, (1,))
+    out = UOp.param(0, dtypes.float, 1)
    red = UOp.const(3.0).cast(dtypes.float).reduce(UOp.range(4, 0, AxisType.UNROLL), arg=(Ops.ADD, 0))
    ast = UOp.sink(out.index(UOp.const(0)).store(red)).replace(arg=KernelInfo())
    uops = full_rewrite_to_sink(ast, Device["CPU"].renderer, optimize=False).toposort()
@@ -429,17 +434,18 @@ class TestReduceCollapse(unittest.TestCase):

 class TestMovementOps(unittest.TestCase):
  def test_pm_mops_partial_reshape_index_removes_reshape(self):
-    from tinygrad.schedule.rangeify import pm_mops
+    from tinygrad.schedule.prepare import pm_mops
    src = UOp.param(0, dtypes.float, shape=(32, 4))
    r0, r1 = UOp.range(4, 0), UOp.range(8, 1)
    result = graph_rewrite(src.reshape((4, 8, 4)).index(r0, r1), pm_mops, name="test")
    self.assertEqual(result.op, Ops.INDEX)
    self.assertIs(result.src[0], src)
    self.assertEqual(result.shape, (4,))
-    self.assertNotIn(Ops.RESHAPE, [u.op for u in result.toposort()])
+    # the only RESHAPE is src itself: the view of the flat param, the extra reshape was folded into the INDEX
+    self.assertEqual([u for u in result.toposort() if u.op is Ops.RESHAPE], [src])

  def test_pm_mops_partial_reshape_index_suffix_mismatch_does_nothing(self):
-    from tinygrad.schedule.rangeify import pm_mops
+    from tinygrad.schedule.prepare import pm_mops
    src = UOp.param(0, dtypes.float, shape=(2, 6))
    result = graph_rewrite(src.reshape((2, 3, 2)).index(UOp.range(2, 0)), pm_mops, name="test")
    self.assertEqual(result.op, Ops.INDEX)
@@ -4,21 +4,21 @@ from tinygrad import UOp
 class TestUOpRepr(unittest.TestCase):
  def test_simple_const(self):
    a = UOp.const(42)
-    self.assertEqual(repr(a), "UOp(Ops.CONST, dtypes.weakint, arg=42, src=())")
+    self.assertEqual(repr(a), "UOp(Ops.CONST, arg=42, src=())")
  def test_different_consts(self):
    a, b = UOp.const(42), UOp.const(3)
    expected = (
-      "UOp(Ops.ADD, dtypes.weakint, arg=None, src=(\n" +
-      "  UOp(Ops.CONST, dtypes.weakint, arg=42, src=()),\n" +
-      "  UOp(Ops.CONST, dtypes.weakint, arg=3, src=()),))"
+      "UOp(Ops.ADD, arg=None, src=(\n" +
+      "  UOp(Ops.CONST, arg=42, src=()),\n" +
+      "  UOp(Ops.CONST, arg=3, src=()),))"
    )
    self.assertEqual(repr(a+b), expected)
  def test_walrus_operator_indentation(self):
    # The reference should have the same indentation as the definition
    a = UOp.const(42)
    expected = (
-      "UOp(Ops.ADD, dtypes.weakint, arg=None, src=(\n" +
-      "  x0:=UOp(Ops.CONST, dtypes.weakint, arg=42, src=()),\n" +
+      "UOp(Ops.ADD, arg=None, src=(\n" +
+      "  x0:=UOp(Ops.CONST, arg=42, src=()),\n" +
      "  x0,))"
    )
    self.assertEqual(repr(a+a), expected)
@@ -26,9 +26,9 @@ class TestUOpRepr(unittest.TestCase):
    # Ensure indentation is consistent at multiple levels
    b = (a:=UOp.const(1)) + a
    expected = (
-      "UOp(Ops.MUL, dtypes.weakint, arg=None, src=(\n" +
-      "  x0:=UOp(Ops.ADD, dtypes.weakint, arg=None, src=(\n" +
-      "    x1:=UOp(Ops.CONST, dtypes.weakint, arg=1, src=()),\n" +
+      "UOp(Ops.MUL, arg=None, src=(\n" +
+      "  x0:=UOp(Ops.ADD, arg=None, src=(\n" +
+      "    x1:=UOp(Ops.CONST, arg=1, src=()),\n" +
      "    x1,)),\n" +
      "  x0,))"
    )
@@ -5,7 +5,7 @@ import z3
 from tinygrad.dtype import dtypes, ConstType, DType, Invalid
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer
 from tinygrad.uop.spec import spec_shared, type_verify
-from tinygrad.uop.symbolic import sym, commutative, pm_simplify_valid, pm_move_where_on_load
+from tinygrad.uop.symbolic import sym, commutative, pm_simplify_valid, pm_move_where_on_load, symbolic_simple
 from tinygrad.uop.validate import uops_to_z3

 def check_uop_against_string(self, v:UOp, s:str):
@@ -448,10 +448,20 @@ class TestSymbolic(unittest.TestCase):
  def test_and_remove(self):
    self.helper_test_variable(uand([uconst(1), Variable("a", 0, 1)]), 0, 1, "a")

+  def test_zero_div_zero_bottom_up(self):
+    # codegen runs symbolic_simple bottom_up, so the 0/0 is rewritten before its consts fold.
+    # without the guard the unsound x/x -> 1 below it claims this one.
+    z = UOp.const(0.0)
+    self.assertTrue(math.isnan(graph_rewrite(z/z, symbolic_simple, bottom_up=True).arg))
+
  def test_masked_shr_fold(self):
    x = UOp.variable('x', 0, 255, dtype=dtypes.uint32, param=True)
    self.helper_test_variable((x & -4) >> 2, 0, 63, "(x>>2)")

+  def test_masked_idiv_fold(self):
+    x = UOp.variable('x', 0, 255, dtype=dtypes.uint32, param=True)
+    self.helper_test_variable((x & -4) // 4, 0, 63, "(x//4)")
+
  def test_bool_or_not_tautology(self):
    a = Variable("a", 0, 10)
    c = a<10
@@ -1023,7 +1033,7 @@ class TestSymbolic(unittest.TestCase):
    cond = Variable("s", 0, 3, dtypes.int) < 2
    a = Variable("a", 0, 3, dtypes.int)
    self.assertIs(graph_rewrite(cond.where(a, a+1).cast(dtypes.half), sym), cond.where(a.cast(dtypes.half), (a+1).cast(dtypes.half)))
-    self.assertIs(graph_rewrite(cond.where(a, uconst(2)).cast(dtypes.half), sym), cond.where(a.cast(dtypes.half), UOp.const(2, dtypes.half)))
+    self.assertIs(graph_rewrite(cond.where(a, uconst(2)).cast(dtypes.half), sym), cond.where(a.cast(dtypes.half), uconst(2.0)))
    self.assertIs(graph_rewrite(cond.where(a, UOp.invalid()).cast(dtypes.half), sym), cond.where(a.cast(dtypes.half), UOp.invalid()))

  def test_where_const_gate_keeps_stated_width(self):
@@ -1382,7 +1392,7 @@ class TestInvalidIndex(unittest.TestCase):

  def test_gated_load_keeps_index_valid(self):
    # the load executes even on gated-off iterations: gated_given_valid must not erase its mask (PADTO OOB shape)
-    buf = UOp.param(0, dtypes.bool, (17,))
+    buf = UOp.param(0, dtypes.bool, 17)
    ridx = Variable("ridx", 0, 31)
    cond = ridx < 17
    load = buf.index(ridx.valid(cond))
@@ -1394,7 +1404,7 @@ class TestStoreLoadFolding(unittest.TestCase):
  """Tests for store(index, load(index)) -> NOOP rule. This rule matches patterns that EMERGE during simplification."""
  def test_store_load_folding(self):
    # store(idx, load(idx)) -> NOOP, including emergent patterns like store(idx, load(idx) + 0)
-    buf = UOp.param(0, dtypes.int, (1,))
+    buf = UOp.param(0, dtypes.int, 1)
    index = buf.index(UOp.const(0))
    # Direct: store(idx, load(idx)) -> NOOP
    self.assertEqual(graph_rewrite(index.store(index.load()), sym).op, Ops.NOOP)
@@ -1407,7 +1417,7 @@ class TestStoreLoadFolding(unittest.TestCase):

 class TestMoveWhereOnLoad(unittest.TestCase):
  def test_bool_index_preserves_dtype(self):
-    buf = UOp.param(0, dtypes.bool, (8,))
+    buf = UOp.param(0, dtypes.bool, 8)
    a = Variable("a", 0, 7)
    r = UOp.range(8, 0)
    # cond has a range that the rewrite can move into the valid: gate (a<4) goes into load valid
@@ -1455,12 +1465,17 @@ class TestGatedUopGivenValid(unittest.TestCase):
    self.assertEqual(idx, (r0 < 3).where(expected_vec, UOp.invalid()))

 class TestRangeSplitting(unittest.TestCase):
+  def test_end_preserves_constant_backedge(self):
+    loop, backedge = UOp.loop(0), UOp.const(False)
+    end = graph_rewrite(UOp(Ops.NOOP).end(loop, backedge), sym)
+    self.assertEqual(end.src, (UOp(Ops.NOOP), loop, backedge))
+
  def test_range_split_on_mod(self):
    # test that mark_range_mod splits RANGE(8) into RANGE(4)*2 + RANGE(2) when used with %2
    from tinygrad.codegen.simplify import pm_split_ranges, pm_flatten_range
    r0 = UOp.range(uconst(8), 0)
    # create a simple expression using the range with mod: store range%2 to a buffer
-    buf = UOp.param(0, dtypes.int, (1,))
+    buf = UOp.param(0, dtypes.int, 1)
    val = (r0 % uconst(2)).cast(dtypes.int)
    store = UOp(Ops.STORE, src=(buf.index(uconst(0)), val))
    sink = UOp(Ops.SINK, src=(UOp(Ops.END, src=(store, r0)),))
@@ -82,7 +82,7 @@ class TestVminVmaxProperties(unittest.TestCase):
  def test_vmin_vmax_multiplication_0_inf(self):
    # vmin and vmax for multiplication with a variable
    x = UOp.const(0.0)
-    y = UOp.load(UOp.param(0, dtypes.float, (1,)), UOp.const(0), dtype=dtypes.float)
+    y = UOp.load(UOp.param(0, dtypes.float, 1), UOp.const(0))
    uop = x * y
    # TODO: these should be 0, but definitely should not be nan
    self.assertEqual(uop.vmin, -math.inf)
@@ -332,7 +332,7 @@ class TestVminVmaxVConst(unittest.TestCase):

  def test_vmin_vmax_vector_with_gep(self):
    # vmin and vmax for a vector constant of bool values
-    d1 = UOp.param(1, dtypes.int, (1,))
+    d1 = UOp.param(1, dtypes.int, 1)
    idx = UOp.const(0)
    val = UOp(Ops.LOAD, src=(d1.index(idx),))
    uop = (val // 32)
@@ -6,7 +6,7 @@ from tinygrad.helpers import Timing, Context, cdiv
 from tinygrad.dtype import dtypes, AddrSpace, ConstFloat, Invalid  # noqa: F401
 from tinygrad.device import Device
 from tinygrad.uop.ops import Ops, AxisType, ParamArg, PatternMatcher, UOp, UPat, dtype_from_uop, exec_alu, graph_rewrite  # noqa: F401  # ParamArg used by eval(str(uop)) roundtrip tests
-from tinygrad.uop.weak import pm_lower_index_dtype
+from tinygrad.uop.weak import pm_lower_weak
 from tinygrad.uop.spec import spec_program, spec_shared, type_verify
 from tinygrad.uop.symbolic import sym, pm_remove_invalid
 from test.helpers import eval_uop, to_uops_list
@@ -38,15 +38,9 @@ class TestDTypeFromUOp(unittest.TestCase):
    self.assertEqual(UOp(Ops.CONST, arg=ConstFloat(3.0)).dtype, dtypes.weakfloat)
    self.assertEqual(UOp(Ops.CONST, arg=True).dtype, dtypes.bool)
    self.assertEqual(UOp(Ops.CONST, arg=Invalid).dtype, dtypes.bool)
-    # an explicit (strong) const dtype is legal until the field is removed
+    # UOp.const at a strong dtype builds the CAST that carries it
    self.assertEqual(UOp.const(3, dtypes.int32).dtype, dtypes.int32)

-  def test_invalid_stated_dtype(self):
-    # UOp.const normalizes a stated dtype away (const_like/full pass their position's); the core constructor does not,
-    # and the spec is what rejects a non-bool Invalid
-    self.assertIs(UOp.const(Invalid, dtypes.float32), UOp.invalid())
-    with self.assertRaises(RuntimeError): type_verify(UOp(Ops.CONST, dtypes.float32, arg=Invalid), spec_shared)
-
  def test_invalid_dtype_and_consumers(self):
    invalid = UOp.invalid()
    self.assertIs(invalid.dtype, dtypes.bool)
@@ -56,15 +50,15 @@ class TestDTypeFromUOp(unittest.TestCase):
      if u.is_invalid)), (dtypes.float32, dtypes.float32, dtypes.bool))
    invalid, value = UOp.invalid(), UOp.const(1, dtypes.float32)
    for u in (UOp.param(0, dtypes.bool, ()).where(value, invalid), value+invalid, UOp.stack(value, invalid)): self.assertIs(u.src[-1], invalid)
-    for u in (UOp(Ops.STACK, dtypes.float32, src=(value, invalid)), UOp(Ops.ADD, dtypes.float32, src=(value, invalid)),
+    for u in (UOp(Ops.STACK, src=(value, invalid)), UOp(Ops.ADD, src=(value, invalid)),
              UOp.const(True).where(value, invalid), UOp(Ops.CMPLT, src=(invalid, value)), UOp(Ops.CMPLT, src=(value, invalid)),
-              UOp.param(0, dtypes.float32, (4,)).index(invalid)): type_verify(u, spec_shared)
+              UOp.param(0, dtypes.float32, 4).index(invalid)): type_verify(u, spec_shared)
    gate, value = UOp.param(0, dtypes.bool, ()), UOp.param(1, dtypes.float, ())
    self.assertIs((out:=graph_rewrite(gate.where(value, UOp.invalid()), pm_remove_invalid)).src[2], UOp.const(0, dtypes.float))
    type_verify(out.sink(), spec_program)

  def test_remove_invalid_stack_lanes(self):
-    stack = UOp(Ops.STACK, dtypes.half, (UOp.const(1, dtypes.half), UOp.invalid()))
+    stack = UOp(Ops.STACK, src=(UOp.const(1, dtypes.half), UOp.invalid()))
    out = graph_rewrite(stack, pm_remove_invalid)
    self.assertEqual(out.src, (UOp.const(1, dtypes.half), UOp.const(0, dtypes.half)))
    type_verify(out.sink(), spec_program)
@@ -73,20 +67,21 @@ class TestLowerIndexDtype(unittest.TestCase):
  def test_gated_shrink_lowers_to_selected_width(self):
    # coalesce builds gated SHRINKs for masked vectorized loads; lowering must resolve them at the
    # width the offset bounds select (this one needs long)
-    buf = UOp.param(0, dtypes.float, (2**31+64,))
+    buf = UOp.param(0, dtypes.float, 2**31+64)
    i = UOp.variable("i", 0, 2**28)
    shrink = UOp(Ops.SHRINK, src=(buf, (i*24).valid(i < 2**28), UOp.const(4)))
-    lowered = graph_rewrite(shrink.sink(), pm_lower_index_dtype)
-    self.assertTrue(all(u.dtype != dtypes.weakint for u in lowered.backward_slice_with_self), "lowering must resolve all weakint")
+    lowered = graph_rewrite(shrink.sink(), pm_lower_weak)
+    self.assertTrue(all(u.op is Ops.CONST for u in lowered.backward_slice_with_self if u.dtype in dtypes.weaks),
+                    "lowering must resolve every weak width, except a typed literal's value half")
    sh = next(u for u in lowered.backward_slice_with_self if u.op is Ops.SHRINK)
    self.assertEqual(sh.src[1].dtype, dtypes.long)

  def test_reg_buffer_size_lowers(self):
    reg = UOp.placeholder((4,), dtypes.float, 0, addrspace=AddrSpace.REG)
-    self.assertEqual(reg.src[0].dtype, dtypes.weakint)
-    lowered = graph_rewrite(reg.sink(), pm_lower_index_dtype)
-    self.assertTrue(all(u.dtype != dtypes.weakint for u in lowered.backward_slice_with_self), "lowering must resolve all weakint")
-    self.assertEqual(next(u for u in lowered.backward_slice_with_self if u.op is Ops.BUFFER).src[0].dtype, dtypes.int)
+    self.assertEqual(reg.arg.size, 4)
+    lowered = graph_rewrite(reg.sink(), pm_lower_weak)
+    self.assertTrue(all(u.op is Ops.CONST for u in lowered.backward_slice_with_self if u.dtype in dtypes.weaks),
+                    "lowering must resolve every weak width, except a typed literal's value half")

 class TestSafeCast(unittest.TestCase):
  def test_cast_folds(self):
@@ -219,7 +214,7 @@ class TestExecALU(unittest.TestCase):

 class TestGatedStoreRewrite(unittest.TestCase):
  def test_tiny_gate_store(self):
-    gmem = UOp.param(0, dtypes.float, (8,))
+    gmem = UOp.param(0, dtypes.float, 8)
    gidx0 = UOp.special(4, 'gidx0')
    gate = gidx0<UOp.const(1)
    idx = UOp(Ops.INDEX, src=(gmem, (gidx0 * UOp.const(2)).valid(gate)))
@@ -235,8 +230,8 @@ class TestGatedStoreRewrite(unittest.TestCase):
    self.assertEqual(len(gated_uops[-1].src), 2)

  def test_gate_some_stores(self):
-    gmem0 = UOp.param(0, dtypes.float, (8,))
-    gmem1 = UOp.param(1, dtypes.float, (8,))
+    gmem0 = UOp.param(0, dtypes.float, 8)
+    gmem1 = UOp.param(1, dtypes.float, 8)
    gidx0 = UOp.special(4, 'gidx0')
    idx = gidx0 * UOp.const(2)
    idx0 = UOp(Ops.INDEX, src=(gmem0, idx.valid(gidx0<UOp.const(1))))
@@ -255,8 +250,8 @@ class TestGatedStoreRewrite(unittest.TestCase):
  # scaled down version of TestLinearizerDumb.test_unmerged_ifs
  @unittest.skip("we don't merge ifs anymore")
  def test_merge_ifs_alt(self):
-    gmem0 = UOp.param(0, dtypes.float, (8,))
-    gmem1 = UOp.param(1, dtypes.float, (8,))
+    gmem0 = UOp.param(0, dtypes.float, 8)
+    gmem1 = UOp.param(1, dtypes.float, 8)
    gidx0 = UOp.special(4, 'gidx0')
    idx = gidx0*UOp.const(2)
    gate = gidx0<UOp.const(1)
@@ -279,10 +274,10 @@ class TestGatedStoreRewrite(unittest.TestCase):
 class TestFastIdiv(unittest.TestCase):
  def test_division_power_of_two(self):
    for dt in (dtypes.int32, dtypes.uint32):
-      g = UOp.param(0, dt, (3,))
+      g = UOp.param(0, dt, 3)
      c = UOp.const(2)
      l = g.index(c)
-      a = UOp(Ops.CDIV, dt, (l, c))
+      a = UOp(Ops.CDIV, src=(l, c))
      uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
      Device[Device.DEFAULT].renderer.render(uops)
      ops = [x.op for x in uops]
@@ -292,9 +287,9 @@ class TestFastIdiv(unittest.TestCase):
  def test_floormod_power_of_two(self):
    # FLOORMOD by a power of two lowers to AND (correct floor mod for any sign in two's complement)
    for dt in (dtypes.int32, dtypes.uint32):
-      g = UOp.param(0, dt, (9,))
+      g = UOp.param(0, dt, 9)
      c = UOp.const(8)
-      a = UOp(Ops.FLOORMOD, dt, (g.index(c), c))
+      a = UOp(Ops.FLOORMOD, src=(g.index(c), c))
      uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
      ops = [x.op for x in uops]
      self.assertIn(Ops.AND, ops, f"For dtype={dt} FLOORMOD by pow2 did not simplify to AND")
@@ -304,9 +299,9 @@ class TestFastIdiv(unittest.TestCase):
  def test_floordiv_power_of_two(self):
    # FLOORDIV by a power of two lowers to a shift, with no round toward zero correction (a shift is exactly floor division)
    for dt in (dtypes.int32, dtypes.uint32, dtypes.int64, dtypes.uint64):
-      g = UOp.param(0, dt, (3,))
+      g = UOp.param(0, dt, 3)
      c = UOp.const(2)
-      a = UOp(Ops.FLOORDIV, dt, (g.index(c), c))
+      a = UOp(Ops.FLOORDIV, src=(g.index(c), c))
      uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
      ops = [x.op for x in uops]
      self.assertIn(Ops.SHR, ops, f"For dtype={dt} FLOORDIV by power of two did not simplify to shift")
@@ -317,7 +312,7 @@ class TestFastIdiv(unittest.TestCase):
  @Context(DISABLE_FAST_IDIV=0)
  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU doesn't support long")
  def test_fast_idiv_and_mod(self):
-    g = UOp.param(0, dtypes.uint32, (4,))
+    g = UOp.param(0, dtypes.uint32, 4)
    c = UOp.const(3)
    l = g.index(c)
    a = UOp(Ops.CDIV, src=(l, c))
@@ -351,7 +346,7 @@ class TestFastIdiv(unittest.TestCase):
  @unittest.expectedFailure
  def test_fast_idiv_overflow(self):
    # This will be possible with a slightly different method for fast_idiv
-    g = UOp.param(0, dtypes.uint32, (8,))
+    g = UOp.param(0, dtypes.uint32, 8)
    c = UOp.const(7).cast(dtypes.uint)
    l = UOp(Ops.LOAD, src=(g.index(c),))
    a = UOp(Ops.CDIV, src=(l, c))
@@ -362,7 +357,7 @@ class TestFastIdiv(unittest.TestCase):
    self.assertNotIn(Ops.CDIV, ops)

  def test_disable_fast_idiv(self):
-    g = UOp.param(0, dtypes.uint32, (4,))
+    g = UOp.param(0, dtypes.uint32, 4)
    c = UOp.const(3)
    l = g.index(c)
    a = UOp(Ops.CDIV, src=(l, c))
@@ -405,8 +400,8 @@ class TestUOpMethod(unittest.TestCase):
      self.assertIs((x != x).simplify(), x.const_like(False, dtypes.bool))

  def test_replace(self):
-    x = UOp.param(0, dtypes.int, (1,))
-    self.assertEqual(x.replace(arg=UOp.param(1, dtypes.int, (1,)).arg).arg.slot, 1)
+    x = UOp.param(0, dtypes.int, 1)
+    self.assertEqual(x.replace(arg=UOp.param(1, dtypes.int, 1).arg).arg.slot, 1)
    with self.assertRaises(AssertionError): x.replace(field="a")

  def test_const_zero_neg_zero_different(self):
@@ -461,16 +456,16 @@ class TestUOpRender(unittest.TestCase):
  def test_render_ssimplified_marg_outside_toposort(self):
    r = UOp.range(UOp.const(16, dtypes.int), 2, AxisType.WEAK, dtype=dtypes.int)
    offset = (r * 2) + (r * 2)
-    shrink = UOp(Ops.SHRINK, src=(UOp.param(0, dtypes.uint, (32,)), offset, UOp.const(2, dtypes.int)))
+    shrink = UOp(Ops.SHRINK, src=(UOp.param(0, dtypes.uint, 32), offset, UOp.const(2, dtypes.int)))
    self.assertIsNot(shrink.src[1], shrink.marg[0][0])
    self.assertEqual(shrink.render(simplify=False), "p0.shrink((((r2*4), 2),))")
    self.assertEqual(UOp.range(1, 0, src=(shrink,), dtype=dtypes.int).render(simplify=False), "r0")

  def test_render_vectorize_empty(self):
-    u = UOp(Ops.STACK, dtype=dtypes.void, src=())
+    u = UOp(Ops.STACK, src=())
    self.assertEqual(u.render(simplify=False), "{}")
  def test_render_vectorize_empty_simplified(self):
-    u = UOp(Ops.STACK, dtype=dtypes.void, src=())
+    u = UOp(Ops.STACK, src=())
    self.assertEqual(u.render(), "{}")
  def test_render_vectorize_same(self):
    u = UOp(Ops.STACK, src=(UOp.const(0),)*3)
@@ -146,7 +146,7 @@ class TestUOpsStats(unittest.TestCase):

  #MULACC should have the same stats as MUL + ADD
  def test_mulacc(self):
-    globl = UOp.param(0, dtypes.int, (3,))
+    globl = UOp.param(0, dtypes.int, 3)
    o1 = UOp.const(1, dtypes.int)
    o2 = UOp.const(2, dtypes.int)
    u1 = globl.index(o1)
@@ -156,7 +156,7 @@ class TestUOpsStats(unittest.TestCase):
    u5 = UOp(Ops.ADD, src=(u4,u3))
    uops = tuple(u5.toposort())

-    globl = UOp.param(0, dtypes.int, (3,))
+    globl = UOp.param(0, dtypes.int, 3)
    o1 = UOp.const(1, dtypes.int)
    o2 = UOp.const(2, dtypes.int)
    u1 = globl.index(o1)
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
geohot	a1125fe213	fix and spec	2026-08-28 15:19:06 -07:00
geohot	5acfb3c71b	remove FUNCTION from tinyspec (regen pdf) and viz	2026-08-28 15:06:21 -07:00
geohot	de08acbc57	match the TUPLE body in the UPat, not in the rewrite	2026-08-28 14:34:13 -07:00
George HotzandGitHub	17339d17cd	Merge branch 'master' into delete-function	2026-08-28 14:30:38 -07:00
geohot	782a83d8c8	delete Ops.FUNCTION: value-producing calls are just CALLs with a TUPLE body a CALL with a TUPLE body (always void dtype) is value-producing and gradient-able, outputs are extracted with GETTUPLE like before; all other CALLs are opaque. TUPLE/GETTUPLE are untouched	2026-08-28 14:16:33 -07:00
chenyuandGitHub	0640cfce35	remove spec dtype check that's trivially true [PR] (#17823 )	2026-08-28 17:00:24 -04:00
chenyuandGitHub	086f45cb81	dtype is not a UOp field anymore (#17821 )	2026-08-28 16:17:29 -04:00
chenyuandGitHub	cb0db2d042	dtype_from_uop is never None (#17819 )	2026-08-28 15:28:09 -04:00
chenyuandGitHub	793c1664f3	don't set UOp dtype in x86 (#17818 )	2026-08-28 15:18:41 -04:00
chenyuandGitHub	bd2711b7c1	put dtype to Ops.INS arg (#17817 )	2026-08-28 15:05:17 -04:00
chenyuandGitHub	b3dad4fa0c	update X86 cast related stuff (#17816 )	2026-08-28 14:30:34 -04:00
sirhcmandGitHub	49627aecb5	ci: use fetch to download libraries (#17814 )	2026-08-28 13:53:19 -04:00
nimlgenandGitHub	959958135d	am: vf (#17802 ) * am: vf * x * Dx * Dx * Dx * x * x	2026-08-28 20:48:59 +03:00
chenyuandGitHub	f196af2327	fix onnx MeanVarianceNormalization arg (#17812 ) axes, not axis	2026-08-28 13:47:47 -04:00
chenyuandGitHub	112e50ec3a	fix onnx.py type annotations (#17811 )	2026-08-28 13:34:43 -04:00
chenyuandGitHub	ae63f23557	fix onnx MaxUnpool (#17809 ) pads can be tuple	2026-08-28 13:13:47 -04:00
chenyuandGitHub	287679a88a	delete X86 not used codes [PR] (#17807 )	2026-08-28 12:50:50 -04:00
Teddy TennantandGitHub	2eacd4fa68	fix Conv2d same padding with per axis dilation (#17805 )	2026-08-28 12:44:43 -04:00
Joel LisenbyandGitHub	7fdc58b1cc	add support for UD quants (#17806 )	2026-08-28 09:38:34 -07:00
chenyuandGitHub	cb857a5a40	remove X86 ymm thingy (#17803 )	2026-08-28 10:18:10 -04:00
chenyuandGitHub	6083de5cff	fix nan compares (#17801 ) * fix nan compares * fix	2026-08-28 09:17:18 -04:00
nimlgenandGitHub	55c8e18930	hcq2 spec2 (#17798 )	2026-08-28 12:51:55 +03:00
nimlgenandGitHub	76dce1eb8d	tiny hcq2 changes (#17797 ) * tiny hcq2 changes * x * x	2026-08-28 11:29:00 +03:00
qazalandGitHub	6242b09066	cleaner mxfp4 gemm prelude (#17796 ) * cleaner mxfp4 prelude * rename sgprs * min diff	2026-08-28 15:59:35 +09:00
wozeparrotandGitHub	92d92c70a9	gptoss: faster grad handling (#17795 )	2026-08-27 22:06:33 -07:00
Teddy TennantandGitHub	26c4b6319e	fix interpolate with align_corners and an output size of 1 (#17790 )	2026-08-27 19:17:05 -07:00
qazalandGitHub	53cad325c7	copy movement op tests construct copies (#17792 ) * copy movement op tests construct copies * hop 3 gpus * assert the amount of bytes moved * work * arange 4, everyone copies 16B first	2026-08-28 11:09:11 +09:00
sirhcmandGitHub	4456e62343	ci: skip apt steps if packages are already installed (#17794 )	2026-08-27 22:06:08 -04:00
George HotzandGitHub	39d7fd0981	axis arg is dead (#17793 )	2026-08-27 18:41:09 -07:00
geohot	4e6bdac412	hotfix: disable HCQ2	2026-08-27 18:23:52 -07:00
chenyuandGitHub	bc485d271f	fix CMPLT between float16 and const 0 (#17791 ) * fix CMPLT between float16 and const 0 * fix	2026-08-27 21:14:34 -04:00
George HotzandGitHub	6b361e4de7	move shape into arg for param/buffer (#17783 ) * move shape into arg for param/buffer * no param_from_shape * drop gratuitous syntax changes * image is a in-graph view, folded into the param arg at render; drop dead multi param sharding * view_as helper, simpler resolve_function, spec update * spec: param/buffer are flat storage, no shape input * image dims live in the param arg from transform_to_image; tighten kernel graph spec * kernel graph spec: only RESHAPE/SHRINK over storage values, not all movement * kernel graph: call args are storage, not views (pm_no_view_args); assert in spec * strip views at the kernel graph level (pm_no_views), move into rangeify * touchups	2026-08-27 18:13:24 -07:00
chenyuandGitHub	38e0fe103b	fix where with float64 branches (#17789 ) * fix where with float64 branches * fix	2026-08-27 21:03:43 -04:00
sirhcmandGitHub	e6ad2907a8	onnxruntime respects NUM_CPU_THREADS (#17788 )	2026-08-27 20:11:27 -04:00
chenyuandGitHub	fede358811	fix fancy indexing with uint8 index (#17787 ) * fix fancy indexing with uint8 index * fix	2026-08-27 19:59:27 -04:00
sirhcmandGitHub	f06832bf6f	test llm with --no_chat_template (#17785 )	2026-08-27 19:01:55 -04:00
geohot	ee3161e924	hotfix: decrease dims in test_attention to get below the 90s limit	2026-08-27 14:44:06 -07:00
geohot	38fa0643cf	hotfix: update max test time from 300 -> 90	2026-08-27 14:27:22 -07:00
George HotzandGitHub	f267a49639	UOp.param accepts single sint as shape (int,) (#17784 ) * UOp.param accepts single sint as shape (int,) Change UOp.param signature from shape:tuple[sint, ...]\|None to shape:tuple[sint, ...]\|sint\|None. A single sint means 1D shape (sint,). Update all callers from (n,) to n syntax. No param_from_shape — multi-dim shapes stay as tuples. * use single int syntax in hcq2 copy_with_kernel	2026-08-27 14:07:33 -07:00
sirhcmandGitHub	14cbecef8f	autogen: use deb for comgr 2 (#17782 )	2026-08-27 16:32:16 -04:00
chenyuandGitHub	54d69f439f	fix test_assign for no subbuffer backends (#17757 ) * test_assign to test/backend * fix	2026-08-27 10:47:21 -04:00
nimlgenandGitHub	980c42d611	hcq2 tiny speedups (#17773 )	2026-08-27 17:42:37 +03:00
nimlgenandGitHub	7eb763a3c2	bnxt to extra (#17772 ) * bnxt to extra * x * x * les	2026-08-27 17:29:17 +03:00
Teddy TennantandGitHub	5c3d044465	fix asinh gradient at zero (#17758 )	2026-08-27 09:51:17 -04:00
nimlgenandGitHub	f19d89e29a	hcq2: a bit cleaner deps (#17771 )	2026-08-27 14:24:59 +03:00
tonyfloatersuandGitHub	2d4e7afcbf	mockgpu/amd: use generated KFD ioctl metadata (#17761 ) * init commit * init commit * minimize to single file scope	2026-08-27 12:24:37 +03:00
qazalandGitHub	0d66815849	set BEAM_PADTO=0 in llama (#17769 )	2026-08-27 16:45:33 +09:00
qazalandGitHub	d16a2510ad	remove some dead scheduler tests (#17768 ) * remove tests that use .uop * that test does nothing	2026-08-27 16:01:28 +09:00
qazalandGitHub	118e812b31	simplify more same device copies (#17766 ) * simplify more same device copies * stack of contigs or copies * rename	2026-08-27 12:54:12 +09:00
sirhcmandGitHub	0820c7a424	benchmarks: openpilot matrix (#17714 )	2026-08-26 18:53:55 -04:00
wozeparrotandGitHub	22ffff3258	gptoss: single jit (#17525 )	2026-08-26 15:32:21 -07:00
George HotzandGitHub	a0a901c8e4	faster qwen 3.8 (#17720 ) * faster qwen * test fix * dead code * fix gguf issue * pretty nt loads * lil * use warp	2026-08-26 15:23:25 -07:00
sirhcmandGitHub	c015351ac5	fix _get_cpu_count for docker --cpus=N in python 3.13+ (#17760 )	2026-08-26 16:53:30 -04:00
nimlgenandGitHub	6074c002e1	hcq2: fix jit (#17747 ) * hcq2: fix reduce * Dx * inputs table * emoty commit	2026-08-26 22:36:36 +03:00
George HotzandGitHub	6042b87272	delete PCONTIG [PR] (#17756 ) * delete PCONTIG * cleanups	2026-08-26 12:07:28 -07:00
sirhcmandGitHub	cc72b9f7be	cleanup BENCHMARK_LOG (#17754 )	2026-08-26 14:55:54 -04:00
chenyuandGitHub	6a3b297548	fix PTX NIR SPEC=2 for bool [pr] (#17753 ) * fix PTX NIR SPEC=2 for bool [pr] storing bool with uint8 needs to pass SPEC * the fix	2026-08-26 14:07:29 -04:00
George HotzandGitHub	ea6c82f3be	small changes from new rangeify (#17752 )	2026-08-26 10:33:31 -07:00
chenyuandGitHub	0abcf09b74	never bufferize_to_store weak input [PR] (#17751 ) github github	2026-08-26 13:25:52 -04:00
chenyuandGitHub	4bdc865131	delete unused rewrite rules [PR] (#17748 )	2026-08-26 10:50:42 -04:00
Teddy TennantandGitHub	4c20f1d357	fix asinh precision loss on negative inputs (#17749 )	2026-08-26 10:08:04 -04:00
qazalandGitHub	ecf79e260d	better all2all schedule test (#17746 ) * better all2all schedule test * deconstruct those numbers * reorder	2026-08-26 15:46:10 +09:00
b1tgandGitHub	9860e5d285	llm tokenizer: fix tekken, add gpt4o (#17733 )	2026-08-25 23:30:35 -07:00
George HotzandGitHub	625c05df1e	fix am_smi to respect dev (#17742 )	2026-08-25 19:21:55 -07:00
sirhcmandGitHub	b49c03fb1c	benchmarks: don't use sudo on mac (#17740 )	2026-08-25 21:06:04 -04:00
George HotzandGitHub	dc04c7820e	lil fixes from new_rangeify (#17741 ) * lil fixes from new_rangeify * gpt sol review	2026-08-25 17:15:28 -07:00
chenyuandGitHub	6ece327cf3	CUSTOM arg is (str, dtype) [PR] (#17737 )	2026-08-25 20:15:06 -04:00
wozeparrotandGitHub	07268b724f	fix: external_test_gpu_crash on python 3.14 (#17739 )	2026-08-25 15:43:30 -07:00
sirhcmandGitHub	55032514ce	pin onnxruntime==1.24.1 (#17738 )	2026-08-25 17:53:15 -04:00
YassineYousfiandGitHub	2824504f90	usb amd: yield between signal polls (#17712 ) * AMD_USB_POLL_US * com * its micro	2026-08-25 14:11:53 -07:00
sirhcmandGitHub	7fef98c86e	benchmarks: remove usage of sudo from linux runners (#17735 )	2026-08-25 16:21:08 -04:00
geohot	b831ca62d9	Reapply "disk cache: thread-local db conn (#17694 )" This reverts commit `df528499ce`.	2026-08-25 13:20:15 -07:00
nimlgenandGitHub	d851aca9ae	hcq2: fix copy/call in usb (#17732 )	2026-08-25 19:03:19 +03:00
chenyuandGitHub	7dc8b666e7	wgsl cast before load for packed [pr] (#17731 ) instead of explicit dtype on load	2026-08-25 11:46:31 -04:00
nimlgenandGitHub	023bfdb380	fix viz for hcq2 (#17729 )	2026-08-25 18:23:22 +03:00
chenyuandGitHub	9f01775cf4	dtype_from_uop for THREEFRY and FDIV (#17727 )	2026-08-25 08:43:59 -04:00
chenyuandGitHub	9607787ce1	delete more explicitly set dtype to UOp [PR] (#17726 ) * delete more explicitly set dtype to UOp [PR] * not that	2026-08-25 08:32:45 -04:00
qazalandGitHub	ab79879613	bring back sqtt examples tests (#17725 )	2026-08-25 17:00:18 +09:00
George HotzandGitHub	a5678317c2	split rangeify to prepare.py (#17722 )	2026-08-24 19:12:13 -07:00
chenyuandGitHub	1d694dd700	remove more explicitly set dtype [PR] (#17721 )	2026-08-24 22:03:22 -04:00
chenyuandGitHub	65ca68567e	const cleanups [PR] (#17719 )	2026-08-24 21:37:14 -04:00
chenyuandGitHub	a7df1a1ace	const are weak 3 [pr] (#17695 )	2026-08-24 20:32:57 -04:00
RaineandGitHub	d9004cff22	use native shifts in Payne Hayek instead of pow2 mul/divs (#17717 ) * use native shifts instead of pow2 mul/divs * wow	2026-08-24 17:11:46 -07:00