remove user contiguous in prepare

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

@@ -194,29 +194,21 @@ 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' && steps.apt-pkgs.outputs.installed == 'false'
+      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'
       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' && steps.apt-pkgs.outputs.installed == 'false'
+      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'
       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') && steps.apt-pkgs.outputs.installed == 'false'
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true')
       shell: bash
       run: |
         sudo apt -qq update || true

.github/workflows/autogen.yml

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

.github/workflows/benchmark.yml

@@ -82,22 +82,19 @@ jobs:
   #      pytest -nauto --durations=20
 
   llmbenchmark:
-    name: Benchmark ${{ matrix.model }} (DEV=${{ matrix.dev }})
+    name: LLM (DEV=${{ matrix.dev }})
     runs-on: [self-hosted, "${{ matrix.dev == 'METAL' && 'macOS' || matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
     strategy:
       fail-fast: false
       matrix:
         dev: ['METAL', 'AMD', 'NV']
-        model: ['llama3.2:3b-f16', 'qwen3.8:27b', 'olmoe']
-        # qwen3.8:27b doesn't fit on mac
-        exclude: [{ dev: 'METAL', model: 'qwen3.8:27b' }, { dev: 'AMD', model: 'olmoe' }, { dev: 'NV', model: 'olmoe' }]
-    timeout-minutes: 15
+    timeout-minutes: 30
     defaults:
       run:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: '0'
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -117,10 +114,16 @@ jobs:
         rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
     - name: reset process replay
       run: python3 test/external/process_replay/reset.py
-    - name: Run ${{ matrix.model }}
-      run: |
-        MODEL=${{ matrix.model }}
-        BENCHMARK_LOG=${MODEL//./} JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m $MODEL --benchmark --warmup
+    - name: Run llama3.2
+      run: BENCHMARK_LOG=llama32_3b-f16 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m llama3.2:3b-f16 --benchmark --warmup
+    - name: Run qwen3.8
+      # qwen3.8:27b doesn't fit on mac
+      if: ${{ matrix.dev != 'METAL' }}
+      run: BENCHMARK_LOG=qwen38_27b JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m qwen3.8:27b --benchmark --warmup
+    - name: Run olmoe
+      # just metal for now
+      if: ${{ matrix.dev == 'METAL' }}
+      run: BENCHMARK_LOG=olmoe JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m olmoe --benchmark --warmup
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
 
@@ -137,7 +140,7 @@ jobs:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: '0'
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -185,7 +188,7 @@ jobs:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: '0'
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -227,7 +230,7 @@ jobs:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: '0'
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -272,7 +275,7 @@ jobs:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: '0'
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -319,7 +322,7 @@ jobs:
       fail-fast: false
       matrix:
         dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 11
+    timeout-minutes: 10
     defaults:
       run:
         shell: bash -e -o pipefail {0}
@@ -398,10 +401,9 @@ jobs:
     - name: Test benchmark allreduce
       if: ${{ matrix.dev == 'NV' }}
       run: python test/external/external_benchmark_multitensor_allreduce.py
-    # TODO: HEVC decode timing test
-    # - name: HEVC Decode Benchmark
-    #   if: ${{ matrix.dev == 'NV' }}
-    #   run: IGNORE_BEAM_CACHE=1 VALIDATE=1 MAX_FRAMES=100 ASSERT_FPS=1400 JITBEAM=1 PYTHONPATH=. python3 extra/hevc/decode.py
+    - name: HEVC Decode Benchmark
+      if: ${{ matrix.dev == 'NV' }}
+      run: VALIDATE=1 MAX_FRAMES=100 ASSERT_FPS=1400 JITBEAM=1 PYTHONPATH=. python3 extra/hevc/decode.py
     - uses: actions/upload-artifact@v7
       if: ${{ matrix.dev != 'AMD' }}
       with:
@@ -436,7 +438,13 @@ jobs:
     - name: UsbGPU tiny tests
       run: GMMU=0 DEV=USB+AMD python3.11 test/test_tiny.py
     - name: UsbGPU copy speeds
-      run: SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/external/external_test_usb_asm24.py
+      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: sh extra/setup_tinygpu_osx.sh
+    - name: UsbGPU (USB4/TB) boot time
+      run: DEBUG=3 DEV=PCI+NV:NAK time python3.11 test/test_tiny.py TestTiny.test_plus
+    - name: UsbGPU (USB4/TB) tiny tests
+      run: DEV=PCI+NV:NAK python3.11 test/test_tiny.py
 
   testcomma:
     strategy:
@@ -450,34 +458,34 @@ jobs:
           - version: '0.11.0'
             model: vision
             url: https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
-            timing: 18
+            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.4
+            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: 13
+            timing: 11
           - version: '0.11.2'
             model: supercombo
             url: https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/433f85f956837606ad1f1cbee4aa7e2158ad23c768dea914b20436c97232741b
-            timing: 28
+            timing: 26
           - dev: QCOM:IR3
             version: '0.11.2'
             model: supercombo
-            timing: 29
+            timing: 41
           - version: '0.11.2'
             model: dmonitoring
             url: https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/3e7b31dfbc0a5234f1baf196513b77fc6af12204b8a8ffe8ee0417e48352f316
-            timing: 12.5
+            timing: 11
           # IR3 dmonitoring is slightly slower
           - dev: QCOM:IR3
             model: dmonitoring
-            timing: 13.0
+            timing: 12
       fail-fast: false
     name: openpilot ${{ matrix.version }} compile3 ${{ matrix.model }} (DEV=${{ matrix.dev }})
-    runs-on: [self-hosted, Linux, comma4]
+    runs-on: [self-hosted, Linux, comma]
     timeout-minutes: 5
     defaults:
       run:
@@ -498,9 +506,9 @@ jobs:
     - name: reset process replay
       run: test/external/process_replay/reset.py
     - name: compile
-      run: FLOAT16=1 IMAGE=1 taskset -c 4-7 python3 examples/openpilot/compile3.py ${{ matrix.url }} openpilot.pkl
+      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 - openpilot.pkl
+      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
 
@@ -525,8 +533,8 @@ jobs:
     - name: benchmark MobileNetV2 on DSP
       run: |
         # generate quantized weights
-        ln -s ~/tinygrad/extra/datasets/imagenet extra/datasets/imagenet
-        ln -s ~/tinygrad/testsig-*.so .
+        ln -s /data/home/tiny/tinygrad/extra/datasets/imagenet extra/datasets/imagenet
+        ln -s /data/home/tiny/tinygrad/testsig-*.so .
         PYTHONPATH=. DEV=CPU QUANT=1 CNT=0 python3 examples/test_onnx_imagenet.py https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx /tmp/model.quant.onnx
         # benchmark on DSP with NOOPT=1, the devectorizer has issues
         PYTHONPATH=. DEV=DSP NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
@@ -556,7 +564,7 @@ jobs:
     - name: openpilot run_pickle big_driving_supercombo
       run: BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo_run_pickle RUN_PICKLE=1 PICKLE_OOB=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py - openpilot.pkl
     - name: Test copy speeds
-      run: SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py
+      run: SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py TestDevCopySpeeds
 
   driverbenchmarks:
     name: PCI Driver Benchmark (DEV=${{ matrix.dev }})
@@ -613,28 +621,37 @@ jobs:
       run: |
         GRAPH_ONE_KERNEL=1 NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
         GRAPH_ONE_KERNEL=1 NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
-    # TODO: HEVC decode timing test
-    # - name: HEVC Decode Benchmark
-    #   if: ${{ matrix.dev == 'NV' }}
-    #   run: IGNORE_BEAM_CACHE=1 VALIDATE=1 MAX_FRAMES=100 ASSERT_FPS=1400 JITBEAM=1 PYTHONPATH=. python3 extra/hevc/decode.py
+    - name: HEVC Decode Benchmark
+      if: ${{ matrix.dev == 'NV' }}
+      run: VALIDATE=1 MAX_FRAMES=100 ASSERT_FPS=1400 JITBEAM=1 PYTHONPATH=. python3 extra/hevc/decode.py
     - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
       if: ${{ matrix.dev == 'NV' }}
       run: BENCHMARK_LOG=resnet_10steps MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
     - name: Run 10 MLPerf Bert training steps (1 gpu)
       # TODO: remove BERT_LAYERS once scheduler is fast
       run: BENCHMARK_LOG=bert_10steps CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+    - name: Remote
+      run: |
+        pkill -f 'extra/remote/serve.py' || true
+        PYTHONPATH=. python3 extra/remote/serve.py 6482 &
+        sleep 1
+        DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6482 AM_RESET=1 python3 test/test_tiny.py
+        if [[ "${{ matrix.dev }}" == "AMD" ]]; then
+          DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6482 AM_RESET=1 AMD_AQL=1 python3 test/test_tiny.py
+        fi
+        pkill -f 'extra/remote/serve.py' || true
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
 
   llvmspeed:
     name: LLVM Speed
     runs-on: [self-hosted, Linux, tinyboxrandom]
-    timeout-minutes: 10
+    timeout-minutes: 5
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
       uses: actions/checkout@v6
     - name: Speed Test
-      run: DEV=CPU:LLVM python3 test/speed/external_test_speed_v_torch.py
+      run: DEV=CPU:LLVM THREADS=0 python3 test/speed/external_test_speed_v_torch.py
     - name: Speed Test (BEAM=2)
-      run: IGNORE_BEAM_CACHE=1 BEAM=2 DEV=CPU:LLVM python3 test/speed/external_test_speed_v_torch.py
+      run: BEAM=2 DEV=CPU:LLVM THREADS=0 python3 test/speed/external_test_speed_v_torch.py

.github/workflows/platform.yml

@@ -36,8 +36,6 @@ jobs:
         deps: testing_unit
     - name: Run unit tests
       run: DEV=METAL python -m pytest -n=auto test/unit/ --durations=20
-    - name: Run opt tests
-      run: DEV=METAL python -m pytest -n=auto test/opt --durations=20
     - name: Test tensor core ops (fake)
       run: DEV=METAL DEBUG=3 TC=2 python test/backend/test_ops.py TestOps.test_gemm
     - name: Test tensor core ops (real)
@@ -76,12 +74,11 @@ jobs:
     - name: Run pytest (ptx)
       env:
         DEV: "MOCK+NV:PTX"
-        HCQ_RUNTIME_DEV: PYTHON
         FORWARD_ONLY: 1
         # TODO: failing due to library loading error
         CAPTURE_PROCESS_REPLAY: 0
       run: |
-        python3 -m pytest -n=auto test/device/test_hcq2.py test/test_tiny.py \
+        python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py \
           test/testextra/test_hevc.py::TestHevc::test_hevc_decode_compile --durations=20
     - name: Run process replay tests
       uses: ./.github/actions/process-replay

.github/workflows/test.yml

@@ -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 --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'
+            ::: $'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'
                 # NOTE: qwen is dumb and only knows about female chickens
 
 # ****** Models Tests ******
@@ -518,15 +518,34 @@ jobs:
       - name: Run LLVM test
         run: DEV=MOCKKFD+AMD:LLVM python test/device/test_amd_llvm.py
 
+  hcq2:
+    name: hcq2
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    timeout-minutes: 5
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v6
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: hcq2
+          deps: testing_unit
+          amd: 'true'
+      - name: Run HCQ2 tests
+        run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/test_tiny.py
+      - name: Run HCQ2 multi-device tests
+        run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python -m pytest -n=auto test/backend/test_multitensor.py
+      - name: Run HCQ2 JIT tests
+        run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/unit/test_jit.py
+      - name: Run HCQ2 unit tests
+        run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python -m pytest test/device/test_hcq2.py
+
   testmockam:
     name: Linux (am)
     runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 15
     env:
       DEV: MOCKPCI+AMD
-      HCQ2: 1
-      HCQ_RUNTIME_DEV: PYTHON
-      PYTHONPATH: .
     steps:
       - name: Checkout Code
         uses: actions/checkout@v6
@@ -536,14 +555,20 @@ jobs:
           key: mockam
           deps: testing_unit
           amd: 'true'
-      - name: Run tests on MOCKAM
-        run: python -m pytest test/test_tiny.py test/unit/test_jit.py
+      - name: Run test_tiny on MOCKAM
+        run: python test/test_tiny.py
       - name: Run test_tiny on MOCKUSB
-        run: HCQ2=0 GMMU=0 DEV=MOCKUSB+AMD python test/test_tiny.py
-      - name: Run test_hcq2 on MOCKPCI
-        run: python -m pytest test/device/test_hcq2.py
+        run: GMMU=0 DEV=MOCKUSB+AMD python test/test_tiny.py
+      - name: Run test_hcq on MOCKPCI
+        run: python -m pytest test/device/test_hcq.py
       - name: Run disk copy tests on MOCKPCI
         run: python -m pytest test/unit/test_disk_tensor.py -k test_copy_from_disk
+      - name: Run test_tiny on MOCKPCI Remote
+        run: |
+          python extra/remote/serve.py 6667 &
+          sleep 2
+          REMOTE=127.0.0.1:6667 python test/test_tiny.py
+          REMOTE=127.0.0.1:6667 python -m pytest test/unit/test_disk_tensor.py -k test_copy_from_disk; kill %1
 
   testamd:
     strategy:
@@ -558,9 +583,6 @@ jobs:
     env:
       DEV: MOCKKFD+AMD:${{ matrix.backend == 'amdllvm' && 'LLVM' || '' }}:${{ matrix.arch }}
       SKIP_SLOW_TEST: 1
-      HCQ2: 1
-      HCQ_RUNTIME_DEV: PYTHON
-      PYTHONPATH: .
     steps:
       - name: Checkout Code
         uses: actions/checkout@v6
@@ -577,11 +599,9 @@ jobs:
           DEBUG=5 FORWARD_ONLY=1 python3 test/test_tiny.py TestTiny.test_plus
       - name: Run MXFP4 Llama training on NULL backend
         if: ${{ matrix.backend == 'amd' && matrix.arch == 'gfx950' }}
-        run: HCQ2=0 PYTHONPATH=. DEV=NULL:HIP:gfx950 MXFP4=1 LLAMA_LAYERS=2 BENCHMARK=3 NULL_ALLOW_COPYOUT=1 NO_HIPCC=1 ROCM_PATH=/opt/rocm JITBEAM=0 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/profile.sh
+        run: PYTHONPATH=. DEV=NULL:HIP:gfx950 MXFP4=1 LLAMA_LAYERS=2 BENCHMARK=3 NULL_ALLOW_COPYOUT=1 NO_HIPCC=1 ROCM_PATH=/opt/rocm JITBEAM=0 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/profile.sh
       - name: Run pytest (amd)
-        run: python -m pytest -n=auto test/backend/test_ops.py test/backend/test_dtype.py test/backend/test_dtype_alu.py test/backend/test_linearizer.py test/backend/test_randomness.py test/backend/test_jit.py test/backend/test_graph.py test/backend/test_multitensor.py test/device/test_hcq2.py test/external/external_test_am.py test/backend/test_asm_gemm.py::TestAsmGEMM --durations=20
-      - name: Run opt tests
-        run: python -m pytest -n=auto test/opt --durations=20
+        run: python -m pytest -n=auto test/backend/test_ops.py test/backend/test_dtype.py test/backend/test_dtype_alu.py test/backend/test_linearizer.py test/backend/test_randomness.py test/backend/test_jit.py test/backend/test_graph.py test/backend/test_multitensor.py test/device/test_hcq.py test/external/external_test_am.py test/backend/test_asm_gemm.py::TestAsmGEMM --durations=20
       - name: Run disk copy tests
         run: python -m pytest test/unit/test_disk_tensor.py -k test_copy_from_disk
       - name: Run TRANSCENDENTAL math
@@ -611,7 +631,7 @@ jobs:
           cuda: 'true'
           ocelot: 'true'
       - name: Set env
-        run: printf "${{ matrix.backend == 'ptx' && 'DEV=MOCK+CUDA:PTX' || matrix.backend == 'nv' && 'DEV=MOCK+NV\nSKIP_SLOW_TEST=1\nHCQ_RUNTIME_DEV=PYTHON' }}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'ptx' && 'DEV=MOCK+CUDA:PTX' || matrix.backend == 'nv' && 'DEV=MOCK+NV\nSKIP_SLOW_TEST=1' }}" >> $GITHUB_ENV
       - name: Check Device.DEFAULT and print some source
         run: |
           python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CUDA','NV'], Device.DEFAULT"

AGENTS.md

@@ -5,4 +5,3 @@
 - Run `python -m ruff check .` to lint
 - Read `./tinygrad/viz/README.md` for profiling and debugging rewrite rules
 - Do not do amend commits. Always do a new commit if a force push to origin would be required.
-- tinygrad has user space PCI drivers for AMD and NVIDIA GPUs. Do not insert the unneeded kernel modules.

conftest.py

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

docs/abstractions4.py

@@ -122,7 +122,7 @@ def example_5_custom_assembly(a:Tensor, correct):
         offset_dwords = (self.labels[inst._target] - inst._pos - inst.size()) // 4
         if not -32768 <= offset_dwords <= 32767: raise ValueError(f"branch to '{inst._target}' offset {offset_dwords} exceeds simm16 range")
         inst.simm16 = offset_dwords
-      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in self.instructions]))))
+      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in self.instructions]))))
 
   CU_COUNT = 32
   LANES = 64

docs/developer/speed.md

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

examples/gpt2.py

@@ -57,7 +57,7 @@ class TransformerBlock:
 
   def __call__(self, x:Tensor, start_pos:Variable, mask:Optional[Tensor]):
     h = x + self.attn(self.ln_1(x), start_pos, mask).float()
-    return (h + self.mlp(self.ln_2(h))).clone()
+    return (h + self.mlp(self.ln_2(h))).contiguous()
 
 class Transformer:
   def __init__(self, dim, n_heads, n_layers, norm_eps, vocab_size, max_seq_len=1024):

examples/mlperf/dataloader.py

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

examples/mlperf/model_train.py

@@ -1667,7 +1667,7 @@ 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, fclip_grads
+  from examples.mlperf.optim import GradAccClipAdamW, GradAccClipAdamWGroup, clip_grads
 
   BENCHMARK = getenv("BENCHMARK")
 
@@ -1785,10 +1785,12 @@ def train_gptoss():
 
     Tensor.realize(loss, *grads)
 
-    clipped_grads, grad_norm = fclip_grads(grads, 1.0)
-    optim.fstep(clipped_grads, grad_norm)
+    grad_norm = clip_grads(grads, 1, 1.0)
+    optim.fstep(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")

examples/mlperf/models/gpt_oss.py

@@ -12,9 +12,9 @@ from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
 from tinygrad.uop.ops import Ops, UOp
 from extra.models.llama import apply_rotary_emb
 from extra.llama_kernels.rmsnorm import rmsnorm
-from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8, asm_gemm, can_use_asm_gemm, mx_pack
+from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8, asm_gemm, can_use_asm_gemm
 from extra.gemm.moe_gemm import grouped_mx_gemm
-from extra.gemm.moe_routing import route, dispatch, combine, router_mfma
+from extra.gemm.moe_routing import route, dispatch, combine
 
 FP8_DTYPE = dtypes.fp8e4m3
 FP8_MAX = 448.0
@@ -41,7 +41,7 @@ def _quant_dequant_bwd(grad:UOp, call:UOp) -> tuple:
 
 def quant_dequant_mx(x:Tensor) -> Tensor:
   fxn = _quant_dequant_fwd_fxn(x.as_param(0).uop, x.device)
-  return Tensor(fxn.uop.call_with_output(x.uop, grad_fxn=_quant_dequant_bwd))
+  return Tensor(UOp.maketuple(fxn.uop).call(x.uop, grad_fxn=_quant_dequant_bwd).gettuple(0))
 
 def _mx_scale(e8:Tensor) -> Tensor:
   return _mx_block_scale(e8) if e8.ndim == 2 else _mx_block_scale_3d(e8)
@@ -58,27 +58,10 @@ def _dequant_bwd(grad:UOp, call:UOp) -> tuple:
 
 def dequant_weight(w_q:Tensor, w_scale:Tensor) -> Tensor:
   fxn = _dequant_fwd_fxn(w_q.as_param(0).uop, w_scale.as_param(1).uop, w_q.device)
-  return Tensor(fxn.uop.call_with_output(w_q.uop, w_scale.uop, grad_fxn=_dequant_bwd))
+  call = UOp.maketuple(fxn.uop).call(w_q.uop, w_scale.uop, grad_fxn=_dequant_bwd)
+  return Tensor(call.gettuple(0))
 
-def matmul_mx(x:Tensor|tuple[Tensor, Tensor], w_q:Tensor, w_scale:Tensor) -> Tensor:
-  if isinstance(x, tuple):
-    assert ASM_GEMM, "pre-quantized MXFP8 input requires ASM_GEMM"
-    from extra.gemm.cdna_asm_gemm import asm_gemm, can_use_asm_gemm, mx_pack
-    x_q, x_e8 = x
-    l_shape, padded = x_q.shape[:-1], x_q.shape[-1]
-    x_q, x_e8 = x_q.reshape(-1, padded), x_e8.reshape(-1, padded // 32)
-    K, N = w_q.shape[1], w_q.shape[0]
-    assert padded >= K and (padded - K) % 32 == 0 and x_e8.shape[-1] == padded // 32
-    wq, ws = w_q, w_scale
-    if (pad := padded - K):
-      wq = wq.pad(((0, 0), (0, pad)))
-      ws = ws.pad(((0, 0), (0, pad // 32)), value=127).cast(dtypes.uint8)
-    if (npad := (-N) % 256):
-      wq = wq.pad(((0, npad), (0, 0)))
-      ws = ws.pad(((0, npad), (0, 0)), value=127).cast(dtypes.uint8)
-    assert can_use_asm_gemm(x_q, wq.T)
-    out = asm_gemm(x_q, wq.T, mx=True, mx_scales=(mx_pack(x_e8), x_e8, mx_pack(ws), ws), mx_w_stored=True)
-    return (out[:, :N] if npad else out).reshape(*l_shape, N).cast(dtypes.bfloat16)
+def matmul_mx(x:Tensor, w_q:Tensor, w_scale:Tensor) -> Tensor:
   l_shape = x.shape[:-1]
   if ASM_GEMM:
     from extra.gemm.cdna_asm_gemm import asm_gemm, can_use_asm_gemm, mx_pack
@@ -192,22 +175,8 @@ class GPTOSS:
   def attention(self, x:Tensor, freqs_cis:Tensor, mask:Tensor, sliding:bool, *, attention_norm:Tensor, wqkv:Tensor,
                 wqkv_scale:Tensor, wqkv_bias:Tensor, wo:Tensor, wo_scale:Tensor, wo_bias:Tensor, sinks:Tensor):
     bsz, seqlen, _ = x.shape
-
-    if getenv("FUSED_RMSNORM_MX", 0):
-      from extra.gptoss_kernels.rmsnorm import rmsnorm_mul_quantize_mxfp8
-      x_q, x_e8, rrms = rmsnorm_mul_quantize_mxfp8(x, attention_norm, self.norm_eps)
-      qkv = matmul_mx((x_q, x_e8), wqkv, wqkv_scale) + wqkv_bias
-      norm_saves = [x_q, x_e8, rrms]
-    if getenv("FUSED_RMSNORM_MUL", 0):
-      from extra.gptoss_kernels.rmsnorm import rmsnorm_mul
-      x_normed, rrms = rmsnorm_mul(x, attention_norm, self.norm_eps)
-      qkv = matmul_mx(x_normed, wqkv, wqkv_scale) + wqkv_bias
-      norm_saves = [x_normed, rrms]
-    else:
-      x_normed, rrms = rmsnorm(x, self.norm_eps)
-      qkv = matmul_mx(x_normed * attention_norm, wqkv, wqkv_scale) + wqkv_bias
-      norm_saves = [x_normed, rrms]
-
+    x_normed, rrms = rmsnorm(x, self.norm_eps)
+    qkv = matmul_mx(x_normed * attention_norm, wqkv, wqkv_scale) + wqkv_bias
     qkv = qkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
     xq = qkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
     xk, xv = qkv[:, :, :, self.n_rep], qkv[:, :, :, self.n_rep + 1]
@@ -233,20 +202,14 @@ class GPTOSS:
       attn = (w @ xvm).permute(0, 3, 1, 2, 4).reshape(bsz, seqlen, self.n_heads * self.head_dim)
 
     out = matmul_mx(attn, wo, wo_scale) + wo_bias
-    return out, [attn] + norm_saves + fa_saves
+    return out, [x_normed, rrms, attn] + fa_saves
 
   def feed_forward(self, x:Tensor, *, ffn_norm:Tensor, gate:Tensor, gate_bias:Tensor,
                    w_gate_up:Tensor, w_gate_up_scale:Tensor, w_gate_up_bias:Tensor,
                    w_down:Tensor, w_down_scale:Tensor, w_down_bias:Tensor):
-    if getenv("FUSED_RMSNORM_MUL", 0):
-      from extra.gptoss_kernels.rmsnorm import rmsnorm_mul
-      x_normed, rrms = rmsnorm_mul(x, ffn_norm, self.norm_eps)
-      inp = x_normed
-    else:
-      x_normed, rrms = rmsnorm(x, self.norm_eps)
-      inp = x_normed * ffn_norm
-
-    logits = router_mfma(inp, gate, gate_bias) if getenv("ROUTER_MFMA", 0) else inp.float() @ gate.float().T + gate_bias.float()
+    x_normed, rrms = rmsnorm(x, self.norm_eps)
+    inp = x_normed * ffn_norm
+    logits = inp.float() @ gate.float().T + gate_bias.float()
     dim, inter = self.dim, self.intermediate_size
 
     if getenv("GROUPED_MOE", 0):
@@ -305,25 +268,10 @@ class GPTOSS:
       h, *_ = self.run_layer(h, freqs_cis, mask_full, i % 2 == 0, attn_kwargs, ffn_kwargs, save=save)
 
     h_normed = self.norm(h)
-
-    if getenv("FP8_LMHEAD", 0) and ASM_GEMM:
-      pad = (-self.dim) % 256
-      h2 = h_normed.reshape(-1, self.dim).pad(((0, 0), (0, pad)))
-      w2 = self.output.pad(((0, 0), (0, pad)))
-      hq, he8, hsi = quantize_mxfp8(h2)
-      oq, oe8, _ = quantize_mxfp8(w2)
-      if hsi is not None and can_use_asm_gemm(hq, oq.T):
-        logits = asm_gemm(hq, oq.T, mx=True, mx_scales=(hsi, he8, mx_pack(oe8), oe8), mx_w_stored=False)
-        logits = logits.reshape(bsz, seqlen, self.vocab_size).cast(dtypes.bfloat16)
-      else:
-        logits = h_normed @ self.output.T
-    elif ASM_GEMM:
-      pad = (-self.dim) % 256
-      h_padded, w_padded = h_normed.pad((None, None, (0, pad))), self.output.pad(((0, 0), (0, pad)))
-      logits = asm_gemm(h_padded, w_padded.T) if can_use_asm_gemm(h_padded, w_padded.T) and getenv("VOCAB_ASM", 1) else h_normed @ self.output.T
-    else:
-      logits = h_normed @ self.output.T
-
+    pad = (-self.dim) % 256
+    h_padded, w_padded = h_normed.pad((None, None, (0, pad))), self.output.pad(((0, 0), (0, pad)))
+    if ASM_GEMM and can_use_asm_gemm(h_padded, w_padded.T): logits = asm_gemm(h_padded, w_padded.T)
+    else: logits = h_normed @ self.output.T
     return logits
 
 def _get_pads(uop:UOp) -> list[UOp]:
@@ -334,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(new_grad))
+    grad_buf.uop = grad_buf.uop.after(grad_buf.uop.store(grad_buf.uop + new_grad))
     return
   cur = grad_buf.uop
   for pad in sorted(pads, key=lambda p: p.marg[0][0] if p.op == Ops.PAD else 0, reverse=True):
     if pad.op == Ops.PAD:
-      grad_shrink = tuple((p[0], s+p[0]) for s,p in zip(pad.src[0].shape, pad.marg))
+      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(pad.src[0].cast(cur.dtype)))
+      cur = cur.after(buf_slice.store(buf_slice + pad.src[0].cast(cur.dtype)))
     else:
       cur = cur.after(cur.store(cur + pad.cast(cur.dtype)))
   grad_buf.uop = cur

examples/mlperf/optim.py

@@ -27,11 +27,6 @@ def clip_grads(grads:list[Tensor], grad_acc, clip_norm) -> Tensor:
   for g in grads: g.assign((g * (clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(g.dtype))
   return total_norm
 
-def fclip_grads(grads:list[Tensor], clip_norm) -> Tensor:
-  total_norm = Tensor.stack(*[g.float().square().sum() for g in grads]).sum().sqrt().contiguous()
-  scale = (clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)
-  return [(g * scale).cast(g.dtype) for g in grads], total_norm
-
 class GradAccClipAdamW(Optimizer):
   def __init__(self, params:list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-6, weight_decay=0.0, grad_acc=1, clip_norm=1.0, device=None, fused=FUSE_OPTIM):
     super().__init__(params, lr, device, fused)

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

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

extra/amdflash/README.md

@@ -1,116 +0,0 @@
-# Navi31 flash tools
-
-Utilities for reading and recovering the 2 MiB SPI flash on Navi31 boards.
-Run them from the tinygrad repository root. No image is bundled; keep a verified
-full-ROM backup before performing any write.
-
-`fw_live.py` accesses BAR5 through tinygrad's `PCIDevice.map_bar()` abstraction
-and supports either the custom ASM24 USB-PCIe bridge or native PCIe. Select the
-transport before the subcommand:
-
-```sh
-python3 extra/amdflash/fw_live.py --transport usb probe
-python3 extra/amdflash/fw_live.py --transport pci probe
-```
-
-The default, `--transport auto`, considers USB devices first and then native
-PCI devices. Native PCI access requires the usual tinygrad PCI permissions and
-an unbound kernel driver.
-
-## Access paths and hardware state
-
-The paths are state-dependent and are not interchangeable:
-
-* **`romless.py`** drives SMUIO `ROM_SW_*` directly through the ASM24 bridge.
-  Use it only when an empty or corrupt flash has stalled the PSP PBL. Healthy
-  autonomous boot gates this engine; the usual gated status is
-  `ROM_SW_STATUS=0x04000800`.
-* **`fw_live.py probe`** queries the early PSP boot-firmware mailbox.
-* Firmware-mediated write commands are retained for protocol documentation but
-  are disabled because an exact stock reflash did not validate safely.
-* **`fw_live.py dump`** reads an exact 2 MiB raw image through
-  `ROM_INDEX/ROM_DATA`. It refuses devices where the raw SMUIO controller is
-  unavailable; the NBIO SOC15 function-ROM aperture is not a physical SPI
-  mapping and is deliberately not used as a fallback.
-
-The tools do not reset or power-cycle the board.
-
-## Raw ROM_SW recovery
-
-Identification and read-only operations:
-
-```sh
-python3 extra/amdflash/romless.py info
-python3 extra/amdflash/romless.py read 0 0x40
-python3 extra/amdflash/romless.py dump spi.bin
-python3 extra/amdflash/romless.py verify known-good.bin
-```
-
-Restore an exact 2 MiB image:
-
-```sh
-python3 extra/amdflash/romless.py flash known-good.bin --yes
-```
-
-If GD25 status-register bit `SR2.CMP` protects the complete array, clearing it
-requires separate authorization:
-
-```sh
-python3 extra/amdflash/romless.py flash known-good.bin --clear-cmp --yes
-```
-
-Programming is sector-granular. Every written 4 KiB sector is immediately read
-back and compared with the input. A range can be resumed independently:
-
-```sh
-python3 extra/amdflash/romless.py flash known-good.bin \
-  --start-sector 128 --sector-count 64 --yes
-```
-
-Navi31 ROM_SW details used by the implementation:
-
-* `ROM_SW_COMMAND = (address << 8) | opcode`
-* TX data uses big-endian stream dwords
-* `RETURN_DATA_EN` (bit 19) is clear for TX and set for RX
-* the RX window exposes the preceding transaction, so reads are primed once
-
-## Firmware-mediated access
-
-The read-only commands are:
-
-```sh
-python3 extra/amdflash/fw_live.py probe
-python3 extra/amdflash/fw_live.py dump current-spi.bin
-```
-
-`dump` produces exactly `0x200000` bytes, requires the raw IFWI magic at offset
-zero, rejects mirrored 1 MiB apertures, and restores the ROM controller/index
-state before writing output.
-
-The validated early-firmware sequence is available as:
-
-```sh
-python3 extra/amdflash/fw_live.py --transport usb ifwi-all full-ifwi.bin --yes
-```
-
-It resolves at most Navi31's configured 19 items, streams the item associated
-with terminal phase `0x2xx`, and then stops. PSP selects the destination
-partition; item `0x08` always comes from the payload referenced by the first
-ISH descriptor, matching AMDVBFlash. A hard power cycle is required afterward.
-
-A successful PSP update is not a byte-identical raw rewrite. On the validated
-stock test, both A/B payloads matched the source exactly, PSP selected and
-booted the updated B partition, and firmware changed only its update cookie,
-B descriptor counter/checksum, and generated metadata near `0x1ef000`.
-
-The `stream`, `ifwi-step`, and `live-flash` commands remain disabled. Testing
-showed that the PSP live path parses a raw stock IFWI but fails with status
-`0xC` (`PSP Write To SPI Error`) after writing an `$AMDVBFL` cookie. Use the
-verified ROM_SW path for recovery.
-
-## Safety
-
-ROM_SW erase/program and `ifwi-all` commands require `--yes`; other
-firmware-streaming commands are disabled. Read-only commands still touch controller and mailbox registers but
-do not issue SPI program/erase or PSP transfer-start commands. Preserve a
-known-good full dump outside the repository.

extra/amdflash/common.py

@@ -1,53 +0,0 @@
-from __future__ import annotations
-import struct, sys, time
-from pathlib import Path
-
-ROOT = Path(__file__).resolve().parents[2]
-if str(ROOT) not in sys.path: sys.path.insert(0, str(ROOT))
-from tinygrad.runtime.support.usb import USB3
-from tinygrad.runtime.support.system import PCIDevice, System, USBPCIDevice
-
-USB_IDS = ((0x3801, 0x0001), (0xADD1, 0x0001))
-NAVI31_DEVICES = ((0xffff, (0x744c,)),)
-
-
-def open_gpu(index: int = 0, transport: str = 'auto') -> PCIDevice:
-  """Open an AMD GPU through tinygrad's transport-independent PCI interface."""
-  if transport not in ('auto', 'usb', 'pci'): raise ValueError(f"unsupported transport {transport!r}")
-  candidates = []
-  if transport in ('auto', 'usb'):
-    for vendor, product in USB_IDS:
-      candidates += [(USBPCIDevice, dev) for dev in USB3.list_devices(vendor, product)]
-  if transport in ('auto', 'pci'):
-    candidates += System.list_devices(0x1002, NAVI31_DEVICES)
-  if not candidates: raise RuntimeError(f"no supported {transport} AMD GPU found")
-  if not 0 <= index < len(candidates): raise RuntimeError(f"device index {index} out of range (found {len(candidates)})")
-  cls, descriptor = candidates[index]
-  return cls("AM", *descriptor) if cls is USBPCIDevice else cls("AM", descriptor)
-
-
-class MMIO:
-  """Transport-independent byte view of BAR5."""
-  def __init__(self, pci_dev: PCIDevice): self.bar = pci_dev.map_bar(5, fmt='B')
-
-  def read32(self, offset: int) -> int:
-    return struct.unpack('<I', bytes(self.bar[offset:offset+4]))[0]
-
-  def write32(self, offset: int, value: int):
-    self.write(offset, struct.pack('<I', value & 0xffffffff))
-
-  def read(self, offset: int, size: int) -> bytes:
-    return bytes(self.bar[offset:offset+size])
-
-  def write(self, offset: int, data: bytes):
-    self.bar[offset:offset+len(data)] = data
-
-
-def wait_until(fn, timeout: float, message: str, interval: float = 0.001):
-  if timeout <= 0 or timeout > 60: raise ValueError("timeout must be in (0, 60] seconds")
-  end = time.monotonic() + timeout
-  while True:
-    value = fn()
-    if value: return value
-    if time.monotonic() >= end: raise TimeoutError(message)
-    time.sleep(interval)

extra/amdflash/fw_live.py

@@ -1,292 +0,0 @@
-#!/usr/bin/env python3
-"""Navi31 firmware-mediated flash access and ROM aperture dumping.
-
-Early item streaming must run after autonomous PSP boot but before a host
-driver or AMDev loads SOS. A fully initialized SOS rejects those commands.
-"""
-from __future__ import annotations
-import argparse, struct, sys, time
-from pathlib import Path
-from common import MMIO, open_gpu, wait_until
-
-ROM_CNTL, ROM_INDEX, ROM_DATA = 0x5A380, 0x5A390, 0x5A394
-FLASH_SIZE, INDEX_PAGE = 0x200000, 0x10000
-
-def bswap32(value: int) -> int: return int.from_bytes(value.to_bytes(4, 'little'), 'big')
-
-COMMAND_DATA, COMMAND, DOORBELL = 0x582D0, 0x582CC, 0x58224
-GET_BOOT_PARTITION, GET_FB_STATE, GET_TRANSFER_TYPE = 0x01, 0x06, 0x07
-START_TRANSFER, DATA_TRANSFER, END_TRANSFER = 0x08, 0x09, 0x0A
-SPI_GET_MODEL_ID = 0x0B
-LIVE_ADDR_LO, LIVE_ADDR_HI, LIVE_UPDATE = 0x02, 0x03, 0x04
-PSP_ERRORS = {
-  0x01: "generic error", 0x02: "out of bounds", 0x03: "invalid parameter",
-  0x04: "off-chip boot error", 0x05: "address not set", 0x06: "parse off-chip error",
-  0x07: "address map error", 0x08: "parse on-chip error", 0x09: "full update error",
-  0x0A: "partition update error", 0x0B: "map on-chip error", 0x0C: "write to SPI error",
-  0x0D: "signature validation error", 0x0E: "invalid command", 0x0F: "signature not found",
-  0x10: "state machine not initialized", 0x11: "state machine transfer error",
-  0x12: "initialization error",
-}
-
-
-class PSPFlashMailbox:
-  def __init__(self, pci_dev): self.mmio = MMIO(pci_dev)
-
-  def command(self, command: int, data: int | None = None, *, timeout: float = 10.0) -> tuple[int, int]:
-    status = self.mmio.read32(COMMAND)
-    if not status & 0x80000000:
-      raise RuntimeError(f"PSP mailbox is not ready before command {command:#x}: status={status:#010x}")
-    if data is not None: self.mmio.write32(COMMAND_DATA, data)
-    self.mmio.write32(COMMAND, command << 16)
-    self.mmio.write32(DOORBELL, 1)
-    wait_until(lambda: self.mmio.read32(COMMAND) & 0x80000000, timeout,
-               f"PSP mailbox command {command:#x} timed out")
-    value = self.mmio.read32(COMMAND)
-    return value & 0xffff, self.mmio.read32(COMMAND_DATA)
-
-  def require(self, command: int, data: int | None = None, *, timeout: float = 10.0, name: str = '') -> int:
-    error, response = self.command(command, data, timeout=timeout)
-    if error:
-      detail = PSP_ERRORS.get(error, "unknown error")
-      raise RuntimeError(f"PSP {name or hex(command)} failed: error={error:#x} ({detail})")
-    return response
-
-  def probe(self) -> dict[str, tuple[int, int]]:
-    result = {}
-    for name, command in (("boot_partition", GET_BOOT_PARTITION), ("fb_state", GET_FB_STATE),
-                          ("model_id", SPI_GET_MODEL_ID), ("transfer_type", GET_TRANSFER_TYPE)):
-      result[name] = self.command(command)
-    return result
-
-  def stream(self, payload: bytes, item_type: int, transfer_type: int | None = None):
-    if not payload: raise ValueError("payload is empty")
-    if len(payload) > 0xFFFFFF: raise ValueError("payload exceeds the mailbox's 24-bit size field")
-    if len(payload) & 3: raise ValueError("payload size must be divisible by four")
-    if not 0 <= item_type <= 0xff: raise ValueError("item type must fit in eight bits")
-    if transfer_type is None: transfer_type = self.require(GET_TRANSFER_TYPE, name="GET_TRANSFER_TYPE")
-    requested = transfer_type & 0xff
-    print(f"firmware transfer_type={transfer_type:#x}", flush=True)
-    if requested != item_type:
-      raise RuntimeError(f"firmware requests item {requested:#x}, not {item_type:#x}")
-    self.require(START_TRANSFER, (len(payload) << 8) | item_type, name="START_TRANSFER")
-    sent, started = 0, time.monotonic()
-    try:
-      for offset in range(0, len(payload), 4):
-        word = struct.unpack_from('<I', payload, offset)[0]
-        self.require(DATA_TRANSFER, word, name=f"DATA_TRANSFER@{offset:#x}")
-        sent = offset + 4
-        if sent % 0x1000 == 0:
-          print(f"{sent:#x}/{len(payload):#x} ({sent/(time.monotonic()-started)/1024:.1f} KiB/s)", flush=True)
-      self.require(END_TRANSFER, (sent << 8) | item_type, timeout=60.0, name="END_TRANSFER")
-    except BaseException:
-      # Give firmware a chance to terminate an interrupted partial session. Do
-      # not submit END_TRANSFER twice if firmware rejected the original END.
-      if sent != len(payload):
-        try: self.command(END_TRANSFER, (sent << 8) | item_type, timeout=10.0)
-        except Exception: pass
-      raise
-    print(f"stream complete: type={item_type:#x} size={sent:#x} elapsed={time.monotonic()-started:.1f}s")
-
-
-def resolve_ifwi_item(image: bytes, item_type: int) -> tuple[int, bytes]:
-  """Resolve AMDVBFlash recovery-layout item types to exact IFWI bytes."""
-  if item_type == 0x01: offset, size = 0, 0x54
-  elif item_type in (0x02, 0x03):
-    offset = 0x2000 if item_type == 0x02 else 0x3000
-    if image[offset:offset+4] != b'$PSP': raise ValueError(f"invalid PSP directory at {offset:#x}")
-    size = (struct.unpack_from('<I', image, offset + 8)[0] + 1) * 0x10
-  elif item_type == 0x04: offset, size = 0x10000, 0x1000
-  elif item_type == 0x05: offset, size = 0x11000, 0x1000
-  elif item_type == 0x06: offset, size = 0x12000, 0x20
-  elif item_type == 0x07: offset, size = 0x13000, 0x20
-  elif item_type == 0x80: offset, size = 0x1000, 4
-  elif item_type == 0x81:
-    offset = struct.unpack_from('<I', image, 0x1000)[0]
-    if image[offset:offset+4] != b'$SGN': raise ValueError("invalid $SGN table pointer")
-    size = (struct.unpack_from('<I', image, offset + 8)[0] + 1) * 0x10
-  elif 0x82 <= item_type <= 0x88:
-    table = struct.unpack_from('<I', image, 0x1000)[0]
-    if image[table:table+4] != b'$SGN': raise ValueError("invalid $SGN table pointer")
-    wanted = item_type - 0x81  # 82h..88h map to SIGN_TYPE 1..7
-    count = struct.unpack_from('<I', image, table + 8)[0]
-    entries = [struct.unpack_from('<IIII', image, table + 0x10 + i*0x10) for i in range(count)]
-    match = [entry for entry in entries if entry[0] == wanted]
-    if len(match) != 1: raise ValueError(f"missing $SGN type {wanted}")
-    _, _, size, offset = match[0]
-  elif item_type == 0x89: offset, size = 0x1f0000, 0x100
-  elif item_type == 0x08:
-    # AMDVBFlash's GetPartitionDetails follows the first ISH entry (firmware ID
-    # 0x13c) and streams its payload. PSP, not the host resolver, selects the
-    # destination partition.
-    offset = struct.unpack_from('<I', image, 0x12000 + 0x10)[0]
-    size = struct.unpack_from('<I', image, 0x12000 + 0x18)[0]
-  else:
-    raise ValueError(f"IFWI resolver does not yet support requested item {item_type:#x}")
-  payload = image[offset:offset+size]
-  if len(payload) != size: raise ValueError(f"item {item_type:#x} extends beyond IFWI")
-  print(f"resolved requested item {item_type:#x}: offset={offset:#x} size={size:#x}")
-  return offset, payload
-
-
-class LivePSPFlash:
-  """Linux psp_v13_0_update_spirom protocol, used with SOS and trained VRAM."""
-  def __init__(self, pci_dev): self.mailbox = PSPFlashMailbox(pci_dev)
-
-  def command(self, command: int, data: int | None = None, timeout: float = 10.0):
-    # Same C2PMSG registers, but the live PSP command set uses IDs 2/3/4.
-    return self.mailbox.require(command, data, timeout=timeout, name=f"LIVE_SPI_{command:#x}")
-
-  def update(self, mc_address: int):
-    status = self.mailbox.mmio.read32(COMMAND)
-    if not status & 0x80000000: raise RuntimeError(f"live PSP mailbox is not ready: {status:#x}")
-    self.command(LIVE_ADDR_LO, mc_address & 0xffffffff)
-    self.command(LIVE_ADDR_HI, mc_address >> 32)
-    self.command(LIVE_UPDATE, timeout=60.0)
-
-
-def open_mailbox(args): return PSPFlashMailbox(open_gpu(args.device, args.transport))
-
-
-def reject_unvalidated_firmware_write():
-  raise RuntimeError("firmware writes are disabled: stock reflash validation failed; use romless.py for recovery")
-
-
-def cmd_probe(args):
-  result = open_mailbox(args).probe()
-  for name, (error, response) in result.items(): print(f"{name}: error={error:#x} response={response:#x}")
-  if result['transfer_type'][0] == 0xA: print("update commands gated: reset card and do not initialize AMDev/SOS", file=sys.stderr)
-
-
-def cmd_stream(args):
-  if not args.yes: raise RuntimeError("refusing to stream without --yes")
-  reject_unvalidated_firmware_write()
-  payload = Path(args.image).read_bytes()
-  open_mailbox(args).stream(payload, args.item_type)
-
-
-def cmd_ifwi_step(args):
-  if not args.yes: raise RuntimeError("refusing to stream without --yes")
-  reject_unvalidated_firmware_write()
-  image = Path(args.ifwi).read_bytes()
-  if len(image) != 0x200000: raise ValueError("Navi31 IFWI image must be exactly 2 MiB")
-  mailbox = open_mailbox(args)
-  state = mailbox.require(GET_TRANSFER_TYPE, name="GET_TRANSFER_TYPE")
-  request = state & 0xff
-  _, payload = resolve_ifwi_item(image, request)
-  mailbox.stream(payload, request, transfer_type=state)
-  next_request = mailbox.require(GET_TRANSFER_TYPE, name="GET_TRANSFER_TYPE")
-  print(f"next firmware transfer_type={next_request:#x}")
-
-
-def cmd_ifwi_all(args):
-  if not args.yes: raise RuntimeError("refusing to stream without --yes")
-  image = Path(args.ifwi).read_bytes()
-  if len(image) != 0x200000: raise ValueError("Navi31 IFWI image must be exactly 2 MiB")
-  mailbox = open_mailbox(args)
-  current = mailbox.require(GET_TRANSFER_TYPE, name="GET_TRANSFER_TYPE")
-  for step in range(19):  # Navi31 ROMItemCount from AMDVBFlash ASICDetails.xml
-    request, phase = current & 0xff, current >> 8
-    print(f"IFWI step {step}: state={current:#x} item={request:#x} phase={phase}", flush=True)
-    _, payload = resolve_ifwi_item(image, request)
-    mailbox.stream(payload, request, transfer_type=current)
-    # AMDVBFlash tests the high byte belonging to the item just streamed. Phase
-    # 2 terminates the loop only after that item has completed successfully.
-    if phase == 2:
-      print(f"IFWI stream complete after terminal state {current:#x}; hard power cycle required")
-      return
-    current = mailbox.require(GET_TRANSFER_TYPE, name="GET_TRANSFER_TYPE")
-  raise RuntimeError(f"IFWI stream did not reach terminal phase after 19 items (state={current:#x})")
-
-
-def cmd_live_flash(args):
-  if not args.yes: raise RuntimeError("refusing to flash without --yes")
-  reject_unvalidated_firmware_write()
-  image = Path(args.ifwi).read_bytes()
-  if not image or len(image) > 16 * 1024 * 1024 or len(image) & 3:
-    raise ValueError("live PSP image must be non-empty, 4-byte aligned, and at most 16 MiB")
-  pci_dev = open_gpu(args.device, args.transport)
-  from tinygrad.runtime.support.am.amdev import AMDev
-  started = time.monotonic()
-  adev = AMDev(pci_dev)
-  print(f"AMDev booted, SOS alive={adev.psp.is_sos_alive()}", flush=True)
-  paddr = adev.mm.palloc(len(image), align=0x1000, zero=False)
-  try:
-    adev.vram.view(paddr, len(image), 'B')[:] = image
-    adev.gmc.flush_hdp()
-    mc_address = adev.paddr2mc(paddr)
-    print(f"staged IFWI at VRAM paddr={paddr:#x} mc={mc_address:#x}", flush=True)
-    LivePSPFlash(pci_dev).update(mc_address)
-    print(f"live PSP flash update complete in {time.monotonic()-started:.1f}s")
-  finally:
-    adev.mm.pfree(paddr)
-
-
-def cmd_dump(args):
-  import hashlib
-  pci_dev = open_gpu(args.device, args.transport)
-  mmio, output, started = MMIO(pci_dev), bytearray(), time.monotonic()
-  original_cntl, original_index = mmio.read32(ROM_CNTL), mmio.read32(ROM_INDEX)
-  if original_cntl == 0xFFFFFFFF:
-    raise RuntimeError("raw SMUIO ROM controller is unavailable; the SOC15 function-ROM aperture is not a raw SPI dump")
-  try:
-    # ROM_DATA must be read one dword at a time; a block read increments MMIO
-    # addresses rather than repeatedly reading the flash aperture register.
-    mmio.write32(ROM_CNTL, bswap32(original_cntl | (1 << 29)))
-    for page in range(0, FLASH_SIZE, INDEX_PAGE):
-      mmio.write32(ROM_INDEX, bswap32(page >> 8))
-      for _ in range(INDEX_PAGE // 4): output += struct.pack('<I', mmio.read32(ROM_DATA))
-      print(f"{page+INDEX_PAGE:#08x}/{FLASH_SIZE:#08x}", flush=True)
-  finally:
-    mmio.write32(ROM_INDEX, bswap32(original_index))
-    mmio.write32(ROM_CNTL, bswap32(original_cntl))
-  if len(output) != FLASH_SIZE or output[:4] != b'\xaa\x55\xaa\x55':
-    raise RuntimeError(f"invalid raw flash dump: size={len(output):#x} magic={output[:4].hex()}")
-  if output[:FLASH_SIZE//2] == output[FLASH_SIZE//2:]:
-    raise RuntimeError("ROM aperture contains mirrored 1 MiB halves; refusing to write a non-raw 2 MiB dump")
-  Path(args.output).write_bytes(output)
-  print(f"dumped {len(output):#x} bytes in {time.monotonic()-started:.1f}s sha256={hashlib.sha256(output).hexdigest()}")
-
-
-def parser():
-  p = argparse.ArgumentParser(description=__doc__)
-  p.add_argument('--device', type=int, default=0, help='device index for the selected transport')
-  p.add_argument('--transport', choices=('auto', 'usb', 'pci'), default='auto', help='PCIe transport (default: USB first, then native PCI)')
-  sub = p.add_subparsers(dest='command', required=True)
-  sub.add_parser('probe', help='query firmware mailbox state without writing').set_defaults(func=cmd_probe)
-
-  s = sub.add_parser('stream', help='stream one exact PSP ROM-item payload')
-  s.add_argument('item_type', type=lambda x:int(x, 0))
-  s.add_argument('image')
-  s.add_argument('--yes', action='store_true')
-  s.set_defaults(func=cmd_stream)
-
-  v = sub.add_parser('ifwi-step', help='resolve and stream the next early-firmware-requested item from a 2 MiB IFWI')
-  v.add_argument('ifwi')
-  v.add_argument('--yes', action='store_true')
-  v.set_defaults(func=cmd_ifwi_step)
-
-  a = sub.add_parser('ifwi-all', help='stream requested IFWI items until firmware reports completion')
-  a.add_argument('ifwi')
-  a.add_argument('--yes', action='store_true')
-  a.set_defaults(func=cmd_ifwi_all)
-
-  l = sub.add_parser('live-flash', help='stage an image in VRAM and invoke the PSP v13 live-update command')
-  l.add_argument('ifwi')
-  l.add_argument('--yes', action='store_true')
-  l.set_defaults(func=cmd_live_flash)
-
-  d = sub.add_parser('dump', help='dump the exact 2 MiB flash through ROM_INDEX/ROM_DATA')
-  d.add_argument('output')
-  d.set_defaults(func=cmd_dump)
-  return p
-
-
-def main():
-  args = parser().parse_args()
-  try: args.func(args)
-  except (RuntimeError, TimeoutError, ValueError, OSError) as error:
-    print(f"error: {error}", file=sys.stderr)
-    raise SystemExit(1)
-
-if __name__ == '__main__': main()

extra/amdflash/romless.py

@@ -1,249 +0,0 @@
-#!/usr/bin/env python3
-"""Direct Navi31 ROM_SW access for GD25LQ16E-class 2 MiB SPI flash."""
-from __future__ import annotations
-import argparse, hashlib, sys, time
-from pathlib import Path
-from common import MMIO, open_gpu, wait_until
-
-FLASH_SIZE, SECTOR_SIZE, PAGE_SIZE, MAX_DATA = 0x200000, 0x1000, 0x100, 0x100
-ROM_CNTL, PAGE_MIRROR_CNTL = 0x5A380, 0x5A384
-ROM_SW_CNTL, ROM_SW_STATUS, ROM_SW_COMMAND, ROM_SW_DATA = 0x5A3A0, 0x5A3A4, 0x5A3A8, 0x5A3B0
-GPIO_PAD_MASK, GPIO_PAD_A, GPIO_PAD_EN = 0x5A504, 0x5A508, 0x5A510
-SPI_GPIO_BITS, RETURN_DATA_EN = 0x780, 0x80000
-EXPECTED_JEDEC = b'\xc8\x60\x15'
-
-
-class Navi31SPI:
-  def __init__(self, pci_dev, prescale: int = 8):
-    if not 0 <= prescale <= 15: raise ValueError("prescale must be 0..15")
-    self.mmio = MMIO(pci_dev)
-    rc = self.mmio.read32(ROM_CNTL)
-    # Select the prescaler instead of inheriting a potentially unusable BL value.
-    self.mmio.write32(ROM_CNTL, (rc & 0xE0FFFFFF) | (1 << 28) | (prescale << 24) | 1)
-
-  def transfer(self, opcode: int, *, address: int = 0, address_len: int = 0,
-               data_out: bytes = b'', data_in: int = 0, timeout: float = 2.0) -> bytes:
-    if data_out and data_in: raise ValueError("simultaneous TX and RX is unsupported")
-    if not 0 <= address_len <= 3: raise ValueError("address_len must be 0..3")
-    count = len(data_out) if data_out else data_in
-    if not 0 <= count <= MAX_DATA: raise ValueError(f"transfer data must be <= {MAX_DATA} bytes")
-    ncmd = 1 + address_len
-    m = self.mmio
-    gpio_mask, gpio_a, gpio_en = m.read32(GPIO_PAD_MASK), m.read32(GPIO_PAD_A), m.read32(GPIO_PAD_EN)
-    page_mirror, rom_cntl = m.read32(PAGE_MIRROR_CNTL), m.read32(ROM_CNTL)
-    try:
-      m.write32(GPIO_PAD_MASK, gpio_mask & ~SPI_GPIO_BITS)
-      m.write32(GPIO_PAD_A, gpio_a & ~SPI_GPIO_BITS)
-      m.write32(GPIO_PAD_EN, gpio_en & ~SPI_GPIO_BITS)
-      m.write32(PAGE_MIRROR_CNTL, (page_mirror & 0xF1FFFFFF) | 0x06000000)
-      m.write32(ROM_CNTL, (rom_cntl & ~0xF) | 8)
-      m.write32(ROM_SW_CNTL, 0)
-      m.write32(ROM_SW_STATUS, 0)
-      if m.read32(ROM_SW_STATUS) != 0: raise RuntimeError("ROM_SW_STATUS did not clear")
-
-      # Navi31 serializes the low instruction byte first, followed by ADDRESS[23:0].
-      m.write32(ROM_SW_COMMAND, ((address & 0xFFFFFF) << 8) | (opcode & 0xFF))
-      for offset in range(0, len(data_out), 4):
-        word = data_out[offset:offset+4].ljust(4, b'\0')
-        m.write32(ROM_SW_DATA + offset, int.from_bytes(word, 'big'))
-
-      control = ((ncmd - 1) << 16) | (RETURN_DATA_EN if data_in else 0) | count
-      m.write32(ROM_SW_CNTL, control)
-      m.read32(ROM_SW_CNTL)  # posted-write flush
-      wait_until(lambda: m.read32(ROM_SW_STATUS) & 1, timeout,
-                 f"ROM_SW transaction timeout (status={m.read32(ROM_SW_STATUS):#x}); engine may be gated after SOS boot")
-      return m.read(ROM_SW_DATA, (data_in + 3) & ~3)[:data_in] if data_in else b''
-    finally:
-      m.write32(ROM_SW_CNTL, 0)
-      m.write32(ROM_SW_STATUS, 0)
-      m.write32(ROM_CNTL, rom_cntl)
-      m.write32(PAGE_MIRROR_CNTL, page_mirror)
-      m.write32(GPIO_PAD_A, gpio_a)
-      m.write32(GPIO_PAD_EN, gpio_en)
-      m.write32(GPIO_PAD_MASK, gpio_mask)
-
-
-class GD25LQ16E:
-  def __init__(self, spi: Navi31SPI): self.spi = spi
-
-  def read_register(self, opcode: int, count: int = 1) -> bytes:
-    # Navi31 exposes the preceding transaction's RX capture. Prime identically.
-    self.spi.transfer(opcode, data_in=max(2, count))
-    return self.spi.transfer(opcode, data_in=count)
-
-  def status(self, opcode: int = 0x05) -> int: return self.read_register(opcode)[0]
-  def rdid(self) -> bytes: return self.read_register(0x9F, 4)
-
-  def sfdp(self, count: int = 20) -> bytes:
-    # 5Ah has one dummy byte after its 24-bit address; retain it for diagnostics.
-    self.spi.transfer(0x5A, address_len=3, data_in=count)
-    return self.spi.transfer(0x5A, address_len=3, data_in=count)
-
-  def wait_idle(self, timeout: float = 2.0) -> int:
-    end = time.monotonic() + timeout
-    while time.monotonic() < end:
-      sr1 = self.status()
-      if not sr1 & 1: return sr1
-      time.sleep(0.002)
-    raise TimeoutError(f"flash remained busy for {timeout}s")
-
-  def write_enable(self):
-    self.spi.transfer(0x06)
-    sr1 = self.status()
-    if not sr1 & 2: raise RuntimeError(f"WREN failed (SR1={sr1:#04x})")
-
-  def clear_cmp(self):
-    sr1, sr2 = self.status(), self.status(0x35)
-    if not sr2 & 0x40: return False
-    self.write_enable()
-    # BUSY/WEL are not writable; preserve all protection/QE fields except CMP.
-    self.spi.transfer(0x01, data_out=bytes((sr1 & 0xFC, sr2 & ~0x40)))
-    self.wait_idle(1.0)
-    new_sr2 = self.status(0x35)
-    if new_sr2 & 0x40: raise RuntimeError(f"failed to clear CMP (SR2={new_sr2:#04x})")
-    return True
-
-  def erase_sector(self, address: int):
-    if address & (SECTOR_SIZE - 1): raise ValueError("sector address is not 4 KiB aligned")
-    self.write_enable()
-    self.spi.transfer(0x20, address=address, address_len=3)
-    self.wait_idle(2.0)
-
-  def program_page(self, address: int, data: bytes):
-    if not data or len(data) > PAGE_SIZE or (address & 0xFF) + len(data) > PAGE_SIZE:
-      raise ValueError("page program crosses a 256-byte boundary")
-    self.write_enable()
-    self.spi.transfer(0x02, address=address, address_len=3, data_out=data)
-    self.wait_idle(1.0)
-
-  def read(self, address: int, count: int) -> bytes:
-    if address < 0 or count < 0 or address + count > FLASH_SIZE: raise ValueError("read outside 2 MiB flash")
-    output = bytearray()
-    while count:
-      size = min(count, MAX_DATA)
-      self.spi.transfer(0x03, address=address, address_len=3, data_in=size)
-      output += self.spi.transfer(0x03, address=address, address_len=3, data_in=size)
-      address, count = address + size, count - size
-    return bytes(output)
-
-
-def has_jedec(raw: bytes) -> bool:
-  return EXPECTED_JEDEC in raw + raw[:2]
-
-
-def open_flash(args) -> GD25LQ16E:
-  flash = GD25LQ16E(Navi31SPI(open_gpu(args.device, 'usb'), args.prescale))
-  raw = flash.rdid()
-  if not has_jedec(raw): raise RuntimeError(f"unexpected GD25LQ16E JEDEC capture: {raw.hex()}")
-  return flash
-
-
-def cmd_info(args):
-  f = open_flash(args)
-  sr1, sr2, sr3 = f.status(), f.status(0x35), f.status(0x15)
-  sfdp = f.sfdp(24)
-  pos = sfdp.find(b'SFDP')
-  print(f"JEDEC capture: {f.rdid().hex()} (C8 60 15 detected)")
-  print(f"SR1/SR2/SR3: {sr1:02x}/{sr2:02x}/{sr3:02x}  CMP={'set' if sr2 & 0x40 else 'clear'}")
-  print(f"SFDP capture: {sfdp.hex()}  signature_offset={pos}")
-
-
-def cmd_read(args):
-  data = open_flash(args).read(args.address, args.size)
-  if args.output: Path(args.output).write_bytes(data)
-  else: print(data.hex())
-
-
-def cmd_dump(args):
-  f = open_flash(args)
-  out = Path(args.output)
-  digest = hashlib.sha256()
-  with out.open('wb') as file:
-    for address in range(0, FLASH_SIZE, SECTOR_SIZE):
-      data = f.read(address, SECTOR_SIZE)
-      file.write(data)
-      digest.update(data)
-      if not (address & 0xFFFF): print(f"{address + SECTOR_SIZE:#08x}/{FLASH_SIZE:#08x}", flush=True)
-  print(f"wrote {out} sha256={digest.hexdigest()}")
-
-
-def cmd_verify(args):
-  expected = Path(args.image).read_bytes()
-  if len(expected) != FLASH_SIZE: raise ValueError(f"image must be exactly {FLASH_SIZE:#x} bytes")
-  f = open_flash(args)
-  digest = hashlib.sha256()
-  for address in range(0, FLASH_SIZE, SECTOR_SIZE):
-    got, wanted = f.read(address, SECTOR_SIZE), expected[address:address+SECTOR_SIZE]
-    digest.update(got)
-    if got != wanted:
-      index = next(i for i, (a, b) in enumerate(zip(got, wanted)) if a != b)
-      raise RuntimeError(f"verify mismatch at {address+index:#x}: flash={got[index]:02x} image={wanted[index]:02x}")
-  print(f"verified {FLASH_SIZE:#x} bytes sha256={digest.hexdigest()}")
-
-
-def cmd_flash(args):
-  if not args.yes: raise RuntimeError("refusing to write without --yes")
-  image = Path(args.image).read_bytes()
-  if len(image) != FLASH_SIZE: raise ValueError(f"image must be exactly {FLASH_SIZE:#x} bytes")
-  total_sectors = FLASH_SIZE // SECTOR_SIZE
-  start, count = args.start_sector, args.sector_count if args.sector_count is not None else total_sectors - args.start_sector
-  if not 0 <= start < total_sectors or not 1 <= count <= total_sectors - start: raise ValueError("invalid sector range")
-  f = open_flash(args)
-  if f.status(0x35) & 0x40:
-    if not args.clear_cmp: raise RuntimeError("CMP protects the full array; rerun with --clear-cmp")
-    f.clear_cmp()
-    print("cleared SR2.CMP", flush=True)
-  begin = time.monotonic()
-  for sector in range(start, start + count):
-    address = sector * SECTOR_SIZE
-    wanted = image[address:address+SECTOR_SIZE]
-    f.erase_sector(address)
-    for offset in range(0, SECTOR_SIZE, PAGE_SIZE):
-      page = wanted[offset:offset+PAGE_SIZE]
-      if page != b'\xff' * PAGE_SIZE: f.program_page(address + offset, page)
-    got = f.read(address, SECTOR_SIZE)
-    if got != wanted:
-      index = next(i for i, (a, b) in enumerate(zip(got, wanted)) if a != b)
-      raise RuntimeError(f"verify mismatch at {address+index:#x}: flash={got[index]:02x} image={wanted[index]:02x}")
-    print(f"OK sector {sector:03d}/{total_sectors-1} @{address:#07x} elapsed={time.monotonic()-begin:.1f}s", flush=True)
-
-
-def parser():
-  p = argparse.ArgumentParser(description=__doc__)
-  p.add_argument('--device', type=int, default=0, help='USB bridge device index')
-  p.add_argument('--prescale', type=int, default=8, help='SCK prescaler 0..15 (default: 8)')
-  sub = p.add_subparsers(dest='command', required=True)
-  sub.add_parser('info', help='read JEDEC, status and SFDP').set_defaults(func=cmd_info)
-
-  r = sub.add_parser('read', help='read a flash range')
-  r.add_argument('address', type=lambda x:int(x, 0))
-  r.add_argument('size', type=lambda x:int(x, 0))
-  r.add_argument('-o', '--output')
-  r.set_defaults(func=cmd_read)
-
-  d = sub.add_parser('dump', help='dump the complete 2 MiB flash')
-  d.add_argument('output')
-  d.set_defaults(func=cmd_dump)
-
-  v = sub.add_parser('verify', help='compare the complete flash with an image')
-  v.add_argument('image')
-  v.set_defaults(func=cmd_verify)
-
-  w = sub.add_parser('flash', help='erase, program, and verify one or more sectors')
-  w.add_argument('image')
-  w.add_argument('--start-sector', type=lambda x:int(x, 0), default=0)
-  w.add_argument('--sector-count', type=lambda x:int(x, 0))
-  w.add_argument('--clear-cmp', action='store_true')
-  w.add_argument('--yes', action='store_true')
-  w.set_defaults(func=cmd_flash)
-  return p
-
-
-def main():
-  args = parser().parse_args()
-  try: args.func(args)
-  except (RuntimeError, TimeoutError, ValueError, OSError) as error:
-    print(f"error: {error}", file=sys.stderr)
-    raise SystemExit(1)
-
-if __name__ == '__main__': main()

extra/amdpci/am_smi.py

@@ -66,8 +66,6 @@ class AMSMI(AMDev):
   def __init__(self, pcibus, vram_bar:MMIOInterface, doorbell_bar:MMIOInterface, mmio_bar:MMIOInterface):
     self.pcibus, self.devfmt = pcibus, pcibus
     self.vram, self.doorbell64, self.mmio = vram_bar, doorbell_bar, mmio_bar
-    self.is_vf = bool(self.mmio[am.mmRCC_IOV_FUNC_IDENTIFIER] & 1)
-    self.vf_rlc_gated:list[tuple[int, int]] = []
     self.pci_state = self.read_pci_state()
     if self.pci_state == "D0": self._init_from_d0()
 

extra/export_model.py

@@ -241,7 +241,8 @@ export default {model_name};
 def export_model(model, target:str, *inputs, model_name: Optional[str] = "model", stream_weights=False):
   assert Device.DEFAULT in EXPORT_SUPPORTED_DEVICE, f"only {', '.join(EXPORT_SUPPORTED_DEVICE)} are supported"
 
-  with Context(JIT=2): linear, output_bufs = jit_model(model, *inputs)
+  # NOTE: NUM_CPU_THREADS=1, since export does not support threading
+  with Context(JIT=2, NUM_CPU_THREADS=1): linear, output_bufs = jit_model(model, *inputs)
   functions, statements, bufs, bufs_to_save = compile_net(linear, output_bufs)
   state = get_state_dict(model)
   weight_names = {(id(b), b.offset, b.size, b.dtype): name for name, x in state.items() if (b:=x.uop.base.realized) is not None}

extra/gemm/amd_asm_matmul.py

@@ -462,7 +462,7 @@ def test_matmul():
     lds = UOp.placeholder((lds_size,), dtypes.uint8, 0, AddrSpace.LOCAL)
     sink = UOp.sink(A.base, B.base, C.base, lds, *gidxs, *lidxs, arg=KernelInfo(name=colored("kernel", "cyan"),
                                                                                   estimates=Estimates(ops=N*N*N*2, mem=N*N*4*3)))
-    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts]))))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
   c = Tensor.custom_kernel(a, b, c, fxn=asm_kernel)[2]
   linear = c.schedule_linear()
 

extra/gemm/cdna_asm_gemm.py

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

extra/gemm/gemm_mxfp4.py

@@ -20,39 +20,34 @@ 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))
-  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):
+    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))
     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))
@@ -1218,6 +1213,31 @@ 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]))
@@ -2214,22 +2234,49 @@ 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[49]))
+    k.emit(s_add_i32(s[48], s[48], s[46]))
     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[49], 0))
+    k.emit(s_mov_b32(s[46], 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[49], s[49], 32))
+    k.emit(s_add_i32(s[46], s[46], 32))
     k.emit(s_branch(65531), target='L2_00E8')
     k.label('L2_00FC')
-    k.emit(s_sub_i32(s[54], s[54], s[49]))
+    k.emit(s_sub_i32(s[54], s[54], s[46]))
     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))
@@ -2264,7 +2311,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[49], s[52], s[49]))
+    k.emit(s_add_i32(s[46], s[52], s[46]))
     k.emit(s_mov_b32(s[6], -16))
     k.emit(s_mov_b32(s[10], -16))
     k.emit(s_mov_b32(s[18], -16))
@@ -2281,18 +2328,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[37], s[37], 1))
-    k.emit(s_mul_i32(s[52], s[37], s[43]))
+    k.emit(s_lshr_b32(s[41], s[41], 1))
+    k.emit(s_mul_i32(s[52], s[41], s[43]))
     k.emit(s_mov_b32(s[14], s[52]))
-    k.emit(s_lshr_b32(s[38], s[38], 1))
-    k.emit(s_mul_i32(s[52], s[38], s[44]))
+    k.emit(s_lshr_b32(s[42], s[42], 1))
+    k.emit(s_mul_i32(s[52], s[42], 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[39]))
+    k.emit(s_mul_i32(s[53], s[52], s[36]))
     k.emit(s_mov_b32(s[22], s[53]))
-    k.emit(s_mul_i32(s[53], s[44], s[40]))
+    k.emit(s_mul_i32(s[53], s[44], s[37]))
     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))
@@ -2309,23 +2356,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[37], v[5]))
+    k.emit(v_mul_lo_u32(v[212], s[41], 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[46], 1))
+    k.emit(s_lshr_b32(s[52], s[49], 1))
     k.emit(s_mul_i32(s[52], s[52], 8))
-    k.emit(s_and_b32(s[53], s[46], 1))
+    k.emit(s_and_b32(s[53], s[49], 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[37], s[52]))
+    k.emit(s_mul_i32(s[52], s[41], s[52]))
     k.emit(v_add_u32_e32(v[212], s[52], v[212]))
-    k.emit(s_mul_i32(s[52], s[37], 32))
+    k.emit(s_mul_i32(s[52], s[41], 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[46], 1056))
+    k.emit(s_mul_i32(s[59], LIT, s[49], 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]))
@@ -2349,35 +2396,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[46], 32))
+    k.emit(s_mul_i32(s[53], s[49], 32))
     k.emit(s_add_i32(s[52], s[53], s[52]))
-    k.emit(s_mul_i32(s[53], s[52], s[39]))
+    k.emit(s_mul_i32(s[53], s[52], s[36]))
     k.emit(v_add_u32_e32(v[222], s[53], v[222]))
-    k.emit(s_mul_i32(s[53], LIT, s[39], 128))
+    k.emit(s_mul_i32(s[53], LIT, s[36], 128))
     k.emit(v_add_u32_e32(v[223], s[53], v[222]))
-    k.emit(s_mul_i32(s[60], s[46], LIT, 256))
+    k.emit(s_mul_i32(s[60], s[49], 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[49], LIT, 256))
-    k.emit(s_mul_i32(s[53], s[46], 64))
+    k.emit(s_mul_i32(s[52], s[46], LIT, 256))
+    k.emit(s_mul_i32(s[53], s[49], 64))
     k.emit(s_add_u32(s[52], s[52], s[53]))
-    k.emit(s_mul_i32(s[52], s[52], s[38]))
+    k.emit(s_mul_i32(s[52], s[52], s[42]))
     k.emit(v_add_u32_e32(v[225], s[52], v[225]))
-    k.emit(s_mul_i32(s[52], 16, s[38]))
+    k.emit(s_mul_i32(s[52], 16, s[42]))
     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[49], LIT, 256))
-    k.emit(s_mul_i32(s[53], s[46], 64))
+    k.emit(s_mul_i32(s[52], s[46], LIT, 256))
+    k.emit(s_mul_i32(s[53], s[49], 64))
     k.emit(s_add_i32(s[52], s[53], s[52]))
-    k.emit(s_mul_i32(s[53], s[52], s[40]))
+    k.emit(s_mul_i32(s[53], s[52], s[37]))
     k.emit(v_add_u32_e32(v[233], s[53], v[233]))
-    k.emit(s_mul_i32(s[52], 32, s[40]))
+    k.emit(s_mul_i32(s[52], 32, s[37]))
     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))
@@ -2463,18 +2510,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[36], s[36], 1))
+    k.emit(s_lshl_b32(s[40], s[40], 1))
     k.emit(s_mul_i32(s[52], s[47], LIT, 256))
-    k.emit(s_mul_hi_u32(s[53], s[52], s[36]))
+    k.emit(s_mul_hi_u32(s[53], s[52], s[40]))
     k.emit(s_add_u32(s[5], s[5], s[53]))
-    k.emit(s_mul_i32(s[53], s[52], s[36]))
+    k.emit(s_mul_i32(s[53], s[52], s[40]))
     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[36]))
+    k.emit(s_mul_i32(s[52], s[52], s[40]))
     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[36]))
+    k.emit(v_mul_lo_u32(v[235], v[235], s[40]))
     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]))
@@ -2482,12 +2529,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[49], LIT, 256))
-    k.emit(s_mul_i32(s[53], s[46], 64))
+    k.emit(s_mul_i32(s[52], s[46], LIT, 256))
+    k.emit(s_mul_i32(s[53], s[49], 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[36], 16))
+    k.emit(s_mul_i32(s[53], s[40], 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))
@@ -2496,7 +2543,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[46], 2))
+    k.emit(s_cmp_lt_i32(s[49], 2))
     k.emit(s_cbranch_scc0(1367), target='L2_25B8')
     k.label('L2_105C')
     k.emit(s_waitcnt(122))

extra/gemm/max_matmul.py

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

extra/gemm/moe_routing.py

@@ -1,53 +1,8 @@
-import functools, math, pathlib
 from tinygrad import Tensor, dtypes
-from tinygrad.helpers import getenv
 from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
-from tinygrad.renderer import Estimates
-from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
 
 BLOCK_ROW = 256
 
-@functools.cache
-def _router_mfma_fwd(out:UOp, x:UOp, weight:UOp, bias:UOp, *, dname:str) -> UOp:
-  *lead, K = x.shape
-  M = math.prod(lead)
-  E = weight.shape[0]
-  threads = UOp.special(256, "lidx0")
-  workgroups = UOp.special((M + 63) // 64, "gidx0")
-  sink = UOp.sink(out.base, x.base, weight.base, bias.base, threads, workgroups,
-                  arg=KernelInfo(f"moe_router_mfma_{M}_{K}_{E}", estimates=Estimates(ops=2*M*E*K, mem=(M*K+E*K+E)*2+M*E*4)))
-  amd = pathlib.Path(__file__).parent.parent/"thunder"/"amd"
-  src = (amd/"moe_router_mfma.cpp").read_text()
-  lib = HIPCCCompiler("gfx950", [f"-I{(amd/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-DHIP_ENABLE_WARP_SYNC_BUILTINS",
-                                 f"-DROUTER_M={M}", f"-DROUTER_K={K}", f"-DROUTER_E={E}"]).compile_cached(src)
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
-
-def _router_mfma_bwd(gradient:UOp, kernel:UOp) -> tuple:
-  _, x_u, weight_u, bias_u = kernel.src[1:5]
-  x, weight, bias = (Tensor(u, device=u.device) for u in (x_u, weight_u, bias_u))
-  reference = x.float() @ weight.float().T + bias.float()
-  grad_x, grad_weight, grad_bias = reference.gradient(x, weight, bias, gradient=Tensor(gradient, device=x_u.device))
-  return None, grad_x.uop, grad_weight.uop, grad_bias.uop
-
-def router_mfma(x:Tensor, weight:Tensor, bias:Tensor) -> Tensor:
-  assert x.ndim >= 2 and weight.ndim == 2 and bias.ndim == 1
-  K = x.shape[-1]
-  E = weight.shape[0]
-  assert weight.shape == (E, K) and bias.shape == (E,)
-  assert x.dtype == weight.dtype == bias.dtype == dtypes.bfloat16
-  assert E == 32 and K % 64 == 0
-  if isinstance(x.device, tuple):
-    assert x.uop.axis == 0, f"router MFMA requires axis-0 sharding, got axis={x.uop.axis}"
-    local_shape = x.uop.shard_shape
-    assert local_shape[-1] == K and math.prod(local_shape[:-1]) % 64 == 0, f"unsupported local router shape {local_shape}"
-  else:
-    assert math.prod(x.shape[:-1]) % 64 == 0
-  x, weight, bias = x.contiguous(), weight.contiguous(), bias.contiguous()
-  out = _sharded_invalids((*x.shape[:-1], E), dtypes.float32, x.device)
-  out, *_ = Tensor.custom_kernel(out, x, weight, bias,
-    fxn=functools.partial(_router_mfma_fwd, dname=str(x.device)), grad_fxn=_router_mfma_bwd)
-  return out
-
 def _sharded_invalids(shape:tuple[int, ...], dtype, device) -> Tensor:
   if isinstance(device, tuple):
     per = Tensor.invalids(shape[0]//len(device), *shape[1:], dtype=dtype, device=device)

extra/gemm/rdna4_asm_matmul.py

@@ -223,7 +223,7 @@ def test_matmul():
     lds = UOp.placeholder((lds_size,), dtypes.uint8, 0, AddrSpace.LOCAL)
     sink = UOp.sink(A.base, B.base, C.base, lds, *gidxs, *lidxs,
                     arg=KernelInfo(name=colored("kernel","cyan"), estimates=Estimates(ops=N*N*N*2, mem=N*N*2*3)))
-    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts]))))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
 
   c = Tensor.custom_kernel(a, b, c, fxn=asm_kernel)[2]
   linear = c.schedule_linear()

extra/gemm/simple_matmul.py

@@ -2,7 +2,7 @@ import numpy as np
 from tinygrad import dtypes, Tensor
 from tinygrad.helpers import getenv, get_single_element
 from tinygrad.dtype import _to_np_dtype
-from tinygrad.engine.realize import lower_and_compile
+from tinygrad.engine.realize import compile_linear
 from tinygrad.codegen.opt import OptOps
 
 dtype_in = (dtypes.half if getenv("HALF") else dtypes.bfloat16 if getenv("BFLOAT16") else
@@ -39,7 +39,7 @@ if __name__ == "__main__":
     c = a.matmul(b, dtype=acc_dtype).realize()
 
   if getenv("SHOULD_USE_TC"):
-    linear = lower_and_compile(a.matmul(b, dtype=acc_dtype).schedule_linear())
+    linear = compile_linear(a.matmul(b, dtype=acc_dtype).schedule_linear())
     call = get_single_element(list(linear.src))
     applied_opts = call.src[0].src[0].arg.applied_opts
     assert any(opt.op is OptOps.TC for opt in applied_opts), f"TC not triggered, {applied_opts}"

extra/gemm/tinygrad_nv_matmul.py

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

extra/gptoss_kernels/rmsnorm/__init__.py

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

extra/gptoss_kernels/rmsnorm/rmsnorm_mul_quantize_mxfp8.cpp

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

extra/gptoss_kernels/rmsnorm/rmsnorm_mul_quantize_mxfp8_bwd.cpp

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

extra/hcq2/ops_amd2.py

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

extra/llama_kernels/__init__.py

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

extra/llama_kernels/rmsnorm/__init__.py

@@ -13,13 +13,12 @@ def _rmsnorm_fwd_fxn(x_in_p, eps, device):
   return rmsnorm_fwd(Tensor(x_in_p, device=device), eps)
 
 def _rmsnorm_bwd(grad:UOp, call:UOp) -> tuple:
-  outs = call.unbound_outputs
-  x_normed = Tensor(outs[0]).float()
+  x_normed = Tensor(call.gettuple(0)).float()
   do_float = Tensor(grad).float()
-  d_x = Tensor(outs[1]) * (do_float - x_normed * (do_float * x_normed).mean(-1, keepdim=True))
+  d_x = Tensor(call.gettuple(1)) * (do_float - x_normed * (do_float * x_normed).mean(-1, keepdim=True))
   return (d_x.cast(call.src[1].dtype).uop,)
 
 def rmsnorm(x_in:Tensor, eps:float) -> tuple[Tensor, Tensor]:
   fxn = _rmsnorm_fwd_fxn(x_in.as_param(0).uop, eps, x_in.device)
-  outs = UOp.call_with_outputs((fxn[0].uop, fxn[1].uop), x_in.uop, grad_fxn=_rmsnorm_bwd)
-  return Tensor(outs[0]), Tensor(outs[1])
+  call = UOp.maketuple(fxn[0].uop, fxn[1].uop).call(x_in.uop, grad_fxn=_rmsnorm_bwd)
+  return Tensor(call.gettuple(0)), Tensor(call.gettuple(1))

extra/mmapeak/mmapeak.py

@@ -3,7 +3,7 @@ import os
 # TODO: there is a timing bug without this
 os.environ["AMD_AQL"] = "1"
 
-from tinygrad import Tensor, Device, GlobalCounters, Context, dtypes
+from tinygrad import Tensor, Device, GlobalCounters, Context
 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, dtypes.void)) for x in insts]))))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
   dummy = Tensor.zeros(1).contiguous().realize()
   out = Tensor.custom_kernel(dummy, fxn=fxn)[0]
   linear = out.schedule_linear()

extra/models/llama.py

@@ -139,7 +139,7 @@ class TransformerBlock:
 
   def __call__(self, x:Tensor, start_pos:Union[Variable,int], freqs_cis:Tensor, mask:Optional[Tensor]):
     h = x + self.attention(self.attention_norm(x), start_pos, freqs_cis, mask)
-    return (h + self.feed_forward(self.ffn_norm(h))).clone().contiguous_backward()
+    return (h + self.feed_forward(self.ffn_norm(h))).contiguous().contiguous_backward()
 
 # standard openai sampling
 def sample(logits: Tensor, temp: float, k: int, p: float, af: float, ap: float):
@@ -201,7 +201,7 @@ class Transformer:
     self.tok_embeddings = embedding(vocab_size, dim)
     self.output = nn.Linear(dim, vocab_size, bias=False) if embedding == nn.Embedding else linear(dim, vocab_size, bias=False)
     self.max_context = max_context
-    self.freqs_cis = precompute_freqs_cis(dim // n_heads, self.max_context * 2, rope_theta).clone().is_param_(False)
+    self.freqs_cis = precompute_freqs_cis(dim // n_heads, self.max_context * 2, rope_theta).contiguous().is_param_(False)
     self.forward_jit = TinyJit(self.forward) if jit else None
 
   def forward(self, tokens:Tensor, start_pos:Union[Variable,int], temperature:float, top_k:int, top_p:float, alpha_f:float, alpha_p:float):

extra/onnx_helpers.py

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

extra/optimization/test_beam_search.py

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

extra/remote/bench.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 import os, sys, time
-from extra.remote.hcq1_remote import RemotePCIDevice
+from tinygrad.runtime.support.system import RemotePCIDevice
 
 LAT_N_RUNS = 500
 THROUGHPUT_N_RUNS = 8
@@ -18,7 +18,7 @@ if __name__ == "__main__":
   print(f"connected to {os.environ['REMOTE']}, device: {name}\n")
 
   # ping (minimal server round-trip, no device I/O)
-  from extra.remote.hcq1_remote import RemoteCmd
+  from tinygrad.runtime.support.system import RemoteCmd
   sock = pci.sock
   for _ in range(10): RemotePCIDevice._rpc(sock, 0, RemoteCmd.PING)
   st = time.perf_counter()

extra/remote/hcq1_remote.py

@@ -1,143 +0,0 @@
-from __future__ import annotations
-import os, mmap, array, functools, contextlib, itertools, struct, socket, subprocess, time, enum, atexit
-from tinygrad.helpers import getenv, temp, ceildiv, unwrap, fetch, system, _ensure_downloads_dir, DEBUG, flatten
-from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface
-from tinygrad.runtime.support.system import PCIDevice, System
-
-class RemoteCmd(enum.IntEnum):
-  PROBE,MAP_BAR,MAP_SYSMEM_FD,CFG_READ,CFG_WRITE,RESET,MMIO_READ,MMIO_WRITE,MAP_SYSMEM,SYSMEM_READ,SYSMEM_WRITE,RESIZE_BAR,PING = range(13)
-
-class RemoteMMIOInterface(MMIOInterface):
-  def __init__(self, dev:RemotePCIDevice, residx:int, nbytes:int, fmt='B', off=0, rd_cmd=RemoteCmd.MMIO_READ, wr_cmd=RemoteCmd.MMIO_WRITE):
-    self.dev, self.residx, self.nbytes, self.fmt, self.off, self.el_sz = dev, residx, nbytes, fmt, off, struct.calcsize(fmt)
-    self.rd_cmd, self.wr_cmd = rd_cmd, wr_cmd
-
-  def __getitem__(self, index):
-    sl = index if isinstance(index, slice) else slice(index, index + 1)
-    start, stop = (sl.start or 0) * self.el_sz, (sl.stop or len(self)) * self.el_sz
-    data = self.dev._bulk_read(self.rd_cmd, self.residx, self.off + start, stop - start)
-    result = data if self.fmt == 'B' else list(struct.unpack(f'<{(stop - start) // self.el_sz}{self.fmt}', data))
-    return result if isinstance(index, slice) else result[0]
-
-  def __setitem__(self, index, val):
-    start = (index.start or 0) * self.el_sz if isinstance(index, slice) else index * self.el_sz
-    data = (val if self.fmt == 'B' else struct.pack(f'<{len(val)}{self.fmt}', *val)) if isinstance(index, slice) else struct.pack(f'<{self.fmt}', val)
-    self.dev._bulk_write(self.wr_cmd, self.residx, self.off + start, data)
-
-  def view(self, offset:int=0, size:int|None=None, fmt=None):
-    return RemoteMMIOInterface(self.dev, self.residx, size or (self.nbytes - offset), fmt or self.fmt, self.off + offset, self.rd_cmd, self.wr_cmd)
-
-class RemotePCIDevice(PCIDevice):
-  _bulk_sent:int = 0
-  _bulk_recv:int = 0
-  _rpc_count:int = 0
-  _start_time:float = 0.0
-
-  @staticmethod
-  @functools.cache
-  def remote_sock(host:str, port:int) -> socket.socket:
-    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-    sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
-    sock.settimeout(getenv("REMOTE_TIMEOUT", 3))
-    sock.connect((host, port))
-    sock.settimeout(None)
-    if DEBUG >= 1 and RemotePCIDevice._start_time == 0.0:
-      RemotePCIDevice._start_time = time.perf_counter()
-      def _print_stats():
-        dt = time.perf_counter() - RemotePCIDevice._start_time
-        sent_mb, recv_mb = RemotePCIDevice._bulk_sent / 1e6, RemotePCIDevice._bulk_recv / 1e6
-        print(f"remote: sent {sent_mb:,.2f} MB ({sent_mb/dt:,.2f} MB/s), recv {recv_mb:,.2f} MB ({recv_mb/dt:,.2f} MB/s), "
-              f"{RemotePCIDevice._rpc_count:,} roundtrips in {dt:.2f}s")
-      atexit.register(_print_stats)
-    return sock
-
-  @staticmethod
-  @functools.cache
-  def remote_list(vendor:int, devices:tuple[tuple[int, tuple[int, ...]], ...], base_class:int|None) -> list[tuple[socket.socket, str]]:
-    payload = array.array('I', itertools.chain.from_iterable((m, d) for m, ds in devices for d in ds)).tobytes()
-    def q(r:str) -> list[tuple[socket.socket, str]]:
-      sock = RemotePCIDevice.remote_sock((host:=r.strip().split(":")[0]), (port:=int(r.strip().split(":")[1]) if ":" in r else 6667))
-      data_len, _, _, _ = RemotePCIDevice._rpc(sock, 0, RemoteCmd.PROBE, base_class or 0, len(payload), vendor, payload=payload)
-      return [(sock, f"remote:{host}:{port}:{d}") for d in RemotePCIDevice._recvall(sock, data_len).decode().split('\n')]
-    return flatten([q(r) for r in getenv("REMOTE", "").split(",") if r.strip()])
-
-  @staticmethod
-  def _recvall(sock:socket.socket, n:int) -> bytes:
-    data = b''
-    while len(data) < n and (chunk:=sock.recv(n - len(data))): data += chunk
-    if len(data) < n: raise RuntimeError("Connection closed")
-    return data
-
-  @staticmethod
-  def _rpc(sock:socket.socket, dev_id:int, cmd:int, *args:int, bar:int=0, readout_size:int=0, payload:bytes=b'', has_fd=False):
-    sock.sendall(struct.pack('<BIIQQQ', cmd, dev_id, bar, *(*args, 0, 0, 0)[:3]) + payload)
-    if has_fd:
-      msg, anc, _, _ = sock.recvmsg(17, socket.CMSG_LEN(4))
-      fd = struct.unpack('<i', anc[0][2][:4])[0]
-    else: msg, fd = RemotePCIDevice._recvall(sock, 17), None
-    if (resp:=struct.unpack('<BQQ', msg))[0] != 0:
-      raise RuntimeError(f"RPC failed: {RemotePCIDevice._recvall(sock, resp[1]).decode('utf-8') if resp[1] > 0 else 'unknown error'}")
-    RemotePCIDevice._rpc_count += 1
-    return (resp[1], resp[2]) + ((RemotePCIDevice._recvall(sock, readout_size) if readout_size > 0 else None),) + (fd,)
-
-  def __init__(self, devpref:str, pcibus:str, sock:socket.socket):
-    self.sock, self.pcibus, self.dev_id = sock, pcibus, int(pcibus.split(':')[-1]) if ':' in pcibus else 0
-    self.peer_group = sock.getpeername()[0]
-    for buft in [socket.SO_SNDBUF, socket.SO_RCVBUF]: self.sock.setsockopt(socket.SOL_SOCKET, buft, 64 << 20)
-
-    self.lock_fd = System.flock_acquire(f"{devpref.lower()}_{pcibus.lower()}.lock")
-
-  def _bulk_read(self, cmd:int, idx:int, offset:int, size:int) -> bytes:
-    RemotePCIDevice._bulk_recv += size
-    return unwrap(self._rpc(self.sock, self.dev_id, cmd, offset, size, bar=idx, readout_size=size)[2])
-  def _bulk_write(self, cmd:int, idx:int, offset:int, data:bytes):
-    RemotePCIDevice._bulk_sent += len(data)
-    self.sock.sendall(struct.pack('<BIIQQQ', cmd, self.dev_id, idx, offset, len(data), 0) + data)
-
-  def alloc_sysmem(self, size:int, vaddr:int=0, contiguous:bool=False) -> tuple[MMIOInterface, list[int]]:
-    paddrs_len, handle, _, _ = self._rpc(self.sock, self.dev_id, RemoteCmd.MAP_SYSMEM, size, int(contiguous))
-    paddrs = list(struct.unpack(f'<{paddrs_len // 8}Q', self._recvall(self.sock, paddrs_len)))
-    return RemoteMMIOInterface(self, handle, size, fmt='B', rd_cmd=RemoteCmd.SYSMEM_READ, wr_cmd=RemoteCmd.SYSMEM_WRITE), paddrs
-
-  def reset(self): self._rpc(self.sock, self.dev_id, RemoteCmd.RESET)
-  def read_config(self, offset:int, size:int): return self._rpc(self.sock, self.dev_id, RemoteCmd.CFG_READ, offset, size)[0]
-  def write_config(self, offset:int, value:int, size:int): self._rpc(self.sock, self.dev_id, RemoteCmd.CFG_WRITE, offset, size, value)
-
-  @functools.cache
-  def bar_info(self, bar_idx:int) -> tuple[int, int]: return self._rpc(self.sock, self.dev_id, RemoteCmd.MAP_BAR, bar=bar_idx)[:2]
-  def map_bar(self, bar:int, off:int=0, addr:int=0, size:int|None=None, fmt='B') -> MMIOInterface:
-    return RemoteMMIOInterface(self, bar, size or self.bar_info(bar)[1], fmt).view(off, size, fmt)
-  def resize_bar(self, bar_idx:int): self._rpc(self.sock, self.dev_id, RemoteCmd.RESIZE_BAR, bar=bar_idx)
-
-class APLRemotePCIDevice(RemotePCIDevice):
-  APP_PATH = "/Applications/TinyGPU.app/Contents/MacOS/TinyGPU"
-
-  @classmethod
-  def ensure_app(cls):
-    commit = "c0d024f9ff0e1dc8fdf217f255da7101d91e8323"
-    app_name = f"TinyGPU_{commit}.zip"
-    if (_ensure_downloads_dir() / app_name).is_file() and os.path.exists(cls.APP_PATH): return
-    print("Downloading TinyGPU.app...")
-    with contextlib.suppress(RuntimeError): system("pkill -f TinyGPU")
-    system(f"ditto -xk {fetch(f'https://github.com/tinygrad/tinygpu_releases/raw/{commit}/TinyGPU.zip', name=app_name)} /Applications")
-    print(system(f"{cls.APP_PATH} install"))
-
-  def __init__(self, devpref:str, pcibus:str):
-    self.ensure_app()
-    sock_path, sock = getenv("APL_REMOTE_SOCK", temp("tinygpu.sock")), socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
-    for i in range(100):
-      with contextlib.suppress(ConnectionRefusedError, FileNotFoundError):
-        sock.connect(sock_path)
-        break
-      if i == 0: subprocess.Popen([self.APP_PATH, "server", sock_path], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
-      time.sleep(0.05)
-    else: raise RuntimeError(f"Failed to connect to TinyGPU server at {sock_path}.")
-    super().__init__(devpref, "usb4", sock=sock)
-
-  def alloc_sysmem(self, size:int, vaddr:int=0, contiguous:bool=False) -> tuple[MMIOInterface, list[int]]:
-    mapped_size, _, _, fd = self._rpc(self.sock, self.dev_id, RemoteCmd.MAP_SYSMEM_FD, size, int(contiguous), has_fd=True)
-    memview = MMIOInterface(FileIOInterface(fd=fd).mmap(0, mapped_size, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED, 0), mapped_size, fmt='B')
-
-    # paddrs are returned as (paddr, size) pairs until a (paddr=0, size=0) terminator in the beginning of the mapping.
-    paddrs_raw = list(itertools.takewhile(lambda p: p[1] != 0, zip(memview.view(fmt='Q')[0::2], memview.view(fmt='Q')[1::2])))
-    return memview, [p + i for p, sz in paddrs_raw for i in range(0, sz, 0x1000)][:ceildiv(size, 0x1000)]

extra/remote/serve.py

@@ -1,7 +1,6 @@
 #!/usr/bin/env python3
 import socket, struct, sys
-from tinygrad.runtime.support.system import PCIDevice, System
-from extra.remote.hcq1_remote import RemoteCmd
+from tinygrad.runtime.support.system import PCIDevice, RemoteCmd, System
 from tinygrad.helpers import DEBUG, OSX
 
 def resp(resp0=0, resp1=0, status=0): return struct.pack('<BQQ', status, resp0, resp1)

extra/runbook_digitalocean_mi350x.md

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

extra/sqtt/roc.py

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

extra/thunder/amd/moe_router_mfma.cpp

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

extra/tinyfs/fetch_file.py

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

extra/tinyfs/fetch_raid.py

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

extra/tinyfs/upload_raid.py

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

spec/tinyspec.tex

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

test/amd/hw/test_ds.py

@@ -7,29 +7,6 @@ Includes: ds_store_b32, ds_load_b32, ds_store_2addr_*, ds_load_2addr_*,
 import unittest
 from test.amd.hw.helpers import *
 
-class TestDSSwizzle(unittest.TestCase):
-  def test_modes_and_overlapping_registers(self):
-    for offset in (0x041f, 0x401f, 0x7c1f, 0x00a0, 0x801b, 0xc020, 0xc420, 0xc021, 0xe000, 0xe010, 0xe01f):
-      for dst in (0, 1):
-        with self.subTest(offset=hex(offset), dst=dst):
-          st = run_program([
-            v_add_nc_u32_e32(v[0], 1, v[255]),
-            ds_swizzle_b32(vdst=v[dst], addr=v[0], offset0=offset & 255, offset1=offset >> 8),
-            s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
-          ], n_lanes=32)
-          self.assertEqual(sorted(st.vgpr[i][dst] for i in range(32)), [6]*32 if offset == 0x00a0 else list(range(1, 33)))
-
-  def test_inactive_sources_and_destinations(self):
-    st = run_program([
-      v_add_nc_u32_e32(v[0], 1, v[255]),
-      v_mov_b32_e32(v[1], 99),
-      s_mov_b32(EXEC_LO, 0x55555555),
-      ds_swizzle_b32(vdst=v[1], addr=v[0], offset0=0x1f, offset1=4),
-      s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
-      s_mov_b32(EXEC_LO, 0xffffffff),
-    ], n_lanes=32)
-    self.assertEqual([st.vgpr[i][1] for i in range(32)], [0, 99]*16)
-
 class TestDS2Addr(unittest.TestCase):
   """Tests for DS_*_2ADDR instructions."""
 

test/amd/hw/test_vop3p.py

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

test/amd/test_asm_kernel.py

@@ -30,7 +30,7 @@ def custom_add_one(A:UOp) -> UOp:
     s_endpgm(),
   ]
   sink = UOp.sink(A.base, threads, arg=KernelInfo(f"custom_add_one_{A.numel()}", estimates=Estimates(ops=A.numel(), mem=A.numel()*4*2)))
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts]))))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) 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, dtypes.void)) for x in insts]))))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) 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, dtypes.void)) for x in insts]))))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) 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, dtypes.void)) for x in insts]))))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) 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, dtypes.void)) for x in insts]))))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) 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, dtypes.void)) for x in insts]))))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
 
 @unittest.skipUnless(Device.DEFAULT == "AMD", "requires AMD device")
 class TestAsmKernel(unittest.TestCase):

test/amd/test_emu2_pcode.py

@@ -17,34 +17,6 @@ def _srcs():
 class TestBasicParsing(unittest.TestCase):
   """Test basic pcode parsing for common instruction patterns."""
 
-  def test_c_style_blocks_and_array_access(self):
-    code = """
-      for (i = 0; i < 4; i+=2) {
-        if (mode == 0) {
-          out[i+0] = input[i+1];
-          out[i+1] = input[i+0];
-        } elsif (mode == 1) {
-          out[i+0] = 7;
-          out[i+1] = 8;
-        } else { // identity
-          out[i+0] = input[i+0];
-          out[i+1] = input[i+1];
-        }
-      }
-    """
-    for mode, expected in enumerate(([11, 10, 13, 12], [7, 8, 7, 8], [10, 11, 12, 13])):
-      with self.subTest(mode=mode):
-        result, _ = parse_pcode(code, {'mode': UOp.const(mode, dtypes.uint32)}, {'input': lambda i: i + 10})
-        self.assertEqual([result[f'out@{i}'].simplify().val for i in range(4)], expected)
-
-  def test_colon_concatenation(self):
-    result, _ = parse_pcode('offset = hi:lo;', {'hi': UOp.const(0x12, dtypes.uint8), 'lo': UOp.const(0x34, dtypes.uint8)})
-    self.assertEqual(result['offset'].simplify().val, 0x1234)
-
-  def test_unclosed_c_block(self):
-    with self.assertRaisesRegex(AssertionError, 'unclosed pcode block'):
-      parse_pcode('if (1) {\nvalue = 2;')
-
   def test_v_add_f32(self):
     """Test parsing V_ADD_F32 pcode."""
     _, assigns = parse_pcode(PCODE[VOP2Op.V_ADD_F32_E32], _srcs())
@@ -180,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),
@@ -211,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)}
 
@@ -246,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),
@@ -327,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),
@@ -358,7 +330,7 @@ class TestAllPcode(unittest.TestCase):
   def _make_srcs(self):
     """Create dummy source variables for pcode parsing."""
     u32, u64 = lambda v=0: UOp.const(v, dtypes.uint32), lambda v=0: UOp.const(v, dtypes.uint64)
-    lds = UOp.param(3, dtypes.uint32, 16384)
+    lds = UOp.param(3, dtypes.uint32, (16384,))
     return {'laneId': u32(), 'laneID': u32(), 'S0': u32(), 'S1': u32(), 'S2': u32(), 'S3': u32(), 'SRC0': u32(),
             'D0': u32(), 'D1': u32(), 'DST': u32(), 'VDST': u32(), 'SDST': u32(),
             'VCC': u64(), 'VCCZ': u32(), 'EXEC': u64(), 'EXEC_LO': u32(), 'EXECZ': u32(), 'SCC': u32(),

test/amd/test_pdf.py

@@ -43,16 +43,6 @@ class TestPcodePDF(unittest.TestCase):
     self.assertEqual(pcode[('S_CMOVK_I32', 2)],
       "if SCC then\nD0.i32 = 32'I(signext(SIMM16.i16))\nendif")
 
-  def test_swizzle_spans_blocks_and_pages(self):
-    for arch in ('rdna3', 'rdna4'):
-      with self.subTest(arch=arch):
-        code = self.pcode[arch][('DS_SWIZZLE_B32', 53)]
-        self.assertIn('} elsif (offset >= 0xc000) {', code)
-        self.assertIn('thread_out[i+3]', code)
-        self.assertIn('xor_mask = offset[14:10];', code)
-        self.assertEqual(code.count('{'), code.count('}'))
-        self.assertTrue(code.endswith('\n}'))
-
   def test_pcode_no_examples(self):
     """Pseudocode should not contain example lines with '=>'."""
     for name in ARCHS:

test/amd/test_sqtt_profiler.py

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

test/backend/test_arange.py

@@ -4,7 +4,7 @@ from tinygrad import Tensor, GlobalCounters, dtypes, nn, Device, Variable
 from tinygrad.helpers import Context, getenv, DEV
 from tinygrad.engine.realize import run_linear, estimate_uop, compile_linear
 from tinygrad.renderer.ptx import PTXRenderer
-from test.helpers import needs_second_gpu, check_schedule, assert_kernel_count, KernelCountException, is_hcq2_device
+from test.helpers import needs_second_gpu, check_schedule, assert_kernel_count, KernelCountException
 
 class TestArange(unittest.TestCase):
   def _get_flops(self, tensor, desired):
@@ -153,7 +153,7 @@ class TestIndexing(unittest.TestCase):
       GlobalCounters.reset()
       z = emb(x).realize()
       self.assertLessEqual(GlobalCounters.global_ops, op_limit)
-      assert_kernel_count(3 if is_hcq2_device() else 2)
+      assert_kernel_count(2)
     if getenv("CHECK", 1):
       import torch
       with torch.no_grad():

test/backend/test_assign.py

@@ -4,7 +4,7 @@ import numpy as np
 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, needs_second_gpu, is_hcq2_device
+from test.helpers import assert_kernel_count, needs_second_gpu
 
 N = 200  # has to be bigger than the cache to fail
 
@@ -22,7 +22,7 @@ class TestAssign(unittest.TestCase):
     assert ba1 == ba2 and ba1 != bb1
     np.testing.assert_allclose(a.numpy(), (np.arange(N*N)*2).reshape((N,N)))
 
-  def test_assign_keeps_identical_tensor(self):
+  def test_assign_zeros_good(self):
     a = Tensor.zeros(10,10).contiguous()
     a.assign(Tensor.ones(10,10))
     b = Tensor.zeros(10,10).contiguous()
@@ -30,7 +30,7 @@ class TestAssign(unittest.TestCase):
     np.testing.assert_allclose(b.numpy(), 0)
 
   @unittest.skip("TODO: this often crashes in CI")
-  def test_assign_keeps_earlier_identical_tensor(self):
+  def test_assign_zeros(self):
     a = Tensor.zeros(10,10).contiguous()
     b = Tensor.zeros(10,10).contiguous()
     a.assign(Tensor.ones(10,10))
@@ -43,17 +43,7 @@ class TestAssign(unittest.TestCase):
     # it should copy into the empty buffer
     GlobalCounters.reset()
     c.realize()
-    assert_kernel_count(2 if is_hcq2_device() else 1)
-
-  def test_assign_copy_retained_uses(self):
-    for use in (lambda x: x.reshape(1, 3), lambda x: x + 1):
-      with self.subTest(use=use):
-        x = Tensor([1., 2, 3], device="PYTHON").to(None)
-        retained = use(x)
-        dest = Tensor.empty(3).assign(x)
-        del x
-        dest.realize().assign(0).realize()
-        self.assertEqual(retained.tolist(), [[1., 2, 3]] if retained.ndim == 2 else [2., 3, 4])
+    assert_kernel_count(1)
 
   def test_assign_slice(self):
     X = Tensor([1,2,3,4]).realize()
@@ -124,6 +114,15 @@ class TestAssign(unittest.TestCase):
     x.assign(x + 1)
     assert [y0.item(), y1.item(), y2.item(), x.item()] == [0.0, 1.0, 2.0, 3.0]
 
+  def test_assign_add_jit(self):
+    @TinyJit
+    def f(x):
+      x += 1
+      x.realize()
+    x = Tensor([0])
+    for _ in range(5): f(x)
+    assert x.item() == 5
+
   def test_assign_add_jit_other(self):
     @TinyJit
     def f(x):
@@ -181,20 +180,21 @@ class TestAssign(unittest.TestCase):
     Tensor.realize(a.contiguous().assign(1), b.contiguous().assign(2))
     self.assertEqual((a + b).item(), 3)
 
-  def test_assign_diamond(self):
-    a = Tensor.ones(4).contiguous().realize()
-    times_a = a*3
-    a.assign(Tensor.full((4,), 2.).contiguous())
-    new = a + (times_a-1)
-    with self.assertRaisesRegex(RuntimeError, "cycle"):  # TODO: broken now, raises
+  def test_assign_diamond_cycle(self):
+    # NOTE: should *not* raise AssertionError from numpy
+    with self.assertRaisesRegex(RuntimeError, "cycle"):
+      a = Tensor.ones(4).contiguous().realize()
+      times_a = a*3
+      a.assign(Tensor.full((4,), 2.).contiguous())
+      new = a + (times_a-1)
       np.testing.assert_allclose(new.numpy(), 4)
 
-  def test_assign_diamond_contiguous(self):
-    a = Tensor.ones(4).contiguous().realize()
-    times_a = a*3
-    a.assign(Tensor.full((4,), 2.))
-    new = a.contiguous() + times_a-1
-    with self.assertRaisesRegex(RuntimeError, "cycle"):  # TODO: broken now, raises
+  def test_assign_diamond_contiguous_cycle(self):
+    with self.assertRaisesRegex(RuntimeError, "cycle"):
+      a = Tensor.ones(4).contiguous().realize()
+      times_a = a*3
+      a.assign(Tensor.full((4,), 2.))
+      new = a.contiguous() + times_a-1
       np.testing.assert_allclose(new.numpy(), 4)
 
   def test_assign_diamond_possible(self):
@@ -267,12 +267,13 @@ class TestAssign(unittest.TestCase):
     np.testing.assert_equal(b1.numpy(), 608)
 
   def test_crossunder_assign(self):
-    a = Tensor.full((4,), 2).contiguous().realize()
-    b = Tensor.full((4,), 3).contiguous().realize()
-    c = a+9
-    a += b
-    b += c
-    with self.assertRaisesRegex(RuntimeError, "cycle"):  # TODO: broken now, raises
+    # NOTE: should *not* raise AssertionError from numpy
+    with self.assertRaisesRegex(RuntimeError, "cycle"):
+      a = Tensor.full((4,), 2).contiguous().realize()
+      b = Tensor.full((4,), 3).contiguous().realize()
+      c = a+9
+      a += b
+      b += c
       Tensor.realize(a,b)
       np.testing.assert_allclose(a.numpy(), 2+3)
       np.testing.assert_allclose(b.numpy(), 3+2+9)
@@ -355,17 +356,49 @@ class TestAssign(unittest.TestCase):
     # permute and base are the same buffer
     assert ba1 == ba2 and ba1 != bb1
 
-  def _assign_view_of_self(self, view):
+  def test_post_permuted_assignment(self):
+    a = Tensor(np.arange(N*N, dtype=np.float32)).reshape(N,N)
+    b = Tensor(np.arange(N*N, dtype=np.float32)).reshape(N,N)
+    a.realize()
+    b.realize()
+    #GlobalCounters.cache = []
+    ba1 = a.uop.base.realized # noqa: F841
+    bb1 = b.uop.base.realized # noqa: F841
+    a.assign(a.permute(1,0) + b)   # this should not work!
+    a.realize()
+    ba2 = a.uop.base.realized # noqa: F841
+    # NOTE: don't test that it's assigned
+    #assert ba1 == ba2 and ba1 != bb1
+    np.testing.assert_allclose(a.numpy(), np.arange(N*N).reshape((N,N)) + np.arange(N*N).reshape((N,N)).transpose(1,0))
+
+  def test_post_permuted_assignment_alt(self):
     a = Tensor.arange(N*N).reshape(N,N).clone().realize()
     b = Tensor.arange(N*N).reshape(N,N).clone().realize()
-    new_a = (view(a)+b).numpy()
-    a.assign(view(a)+b)
+    new_a = (a.T+b).numpy()
+    a.assign(a.T+b)
     np.testing.assert_allclose(a.numpy(), new_a)
 
-  def test_post_permuted_assignment(self): self._assign_view_of_self(lambda a: a.T)
-  def test_post_flipped_assignment(self): self._assign_view_of_self(lambda a: a.flip(0))
-  def test_post_flipped_assignment_axis1(self): self._assign_view_of_self(lambda a: a.flip(1))
-  def test_post_reshape_assignment(self): self._assign_view_of_self(lambda a: a.reshape(-1).reshape(N,N))
+  def test_post_flipped_assignment(self):
+    a = Tensor.arange(N*N).reshape(N,N).clone().realize()
+    b = Tensor.arange(N*N).reshape(N,N).clone().realize()
+    new_a = (a.flip(0)+b).numpy()
+    a.assign(a.flip(0)+b)
+    np.testing.assert_allclose(a.numpy(), new_a)
+
+  def test_post_flipped_assignment_axis1(self):
+    a = Tensor.arange(N*N).reshape(N,N).clone().realize()
+    b = Tensor.arange(N*N).reshape(N,N).clone().realize()
+    new_a = (a.flip(1)+b).numpy()
+    a.assign(a.flip(1)+b)
+    np.testing.assert_allclose(a.numpy(), new_a)
+
+  def test_post_reshape_assignment_fine(self):
+    a = Tensor.arange(N*N).reshape(N, N).clone().realize()
+    b = Tensor.arange(N*N).reshape(N, N).clone().realize()
+    rhs = a.reshape(-1).reshape(N, N)
+    new_a = (rhs+b).numpy()
+    a.assign(rhs+b)  # self-assign with reshape view is fine
+    np.testing.assert_allclose(a.numpy(), new_a)
 
   @unittest.skip("multi output not supported anymore")
   def test_simple_assignment_multioutput(self):
@@ -388,6 +421,14 @@ class TestAssign(unittest.TestCase):
 
   # NOTE: if the assign target is read/write in a single kernel, it should be contiguous
 
+  def test_permuted_assignment_correct(self):
+    a = Tensor.arange(4 * 4).reshape(4, 4).clone().realize()
+    b = Tensor.arange(4 * 4).reshape(4, 4).clone().realize()
+    a = a.permute(1, 0)
+    new_val = a + b
+    a.assign(new_val)
+    np.testing.assert_equal(a.numpy(), np.arange(4 * 4).reshape(4, 4).transpose(1, 0) + np.arange(4 * 4).reshape(4, 4))
+
   def test_permuted_reduceop_child_dual_use(self):
     a = Tensor.arange(32*32*32).reshape(32, 32, 32).clone().realize()
     b = Tensor.ones(32, 32, dtype=dtypes.int).contiguous().realize()
@@ -485,34 +526,34 @@ class TestAssign(unittest.TestCase):
     a[2:5] = [1, 2, 3]
     np.testing.assert_allclose(a.numpy(), [0., 0., 1., 2., 3., 0., 0., 0.])
 
-  # IEEE 754: 1.0f = 0x3f800000, 2.0f = 0x40000000, 3.0f = 0x40400000, 4.0f = 0x40800000
-  REVERSED = [0x40800000, 0x40400000, 0x40000000, 0x3f800000]
-
   def test_assign_bitcast(self):
+    # assign to a bitcast view should modify the underlying buffer
     a = Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32).realize()
-    a.bitcast(dtypes.uint32).assign(Tensor(self.REVERSED, dtype=dtypes.uint32)).realize()
+    # IEEE 754: 1.0f = 0x3f800000, 2.0f = 0x40000000, 3.0f = 0x40400000, 4.0f = 0x40800000
+    a.bitcast(dtypes.uint32).assign(Tensor([0x40800000, 0x40400000, 0x40000000, 0x3f800000], dtype=dtypes.uint32)).realize()
     np.testing.assert_allclose(a.numpy(), [4.0, 3.0, 2.0, 1.0])
-
-  def test_assign_bitcast_unrealized(self):
-    a = Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32).realize()
-    a.bitcast(dtypes.uint32).assign(Tensor(self.REVERSED, dtype=dtypes.uint32))
-    np.testing.assert_allclose(a.numpy(), [4.0, 3.0, 2.0, 1.0])
-
-  def test_assign_double_bitcast(self):
+    # double bitcast
     b = Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32).realize()
-    b.bitcast(dtypes.uint32).bitcast(dtypes.int32).assign(Tensor(self.REVERSED, dtype=dtypes.int32)).realize()
+    b.bitcast(dtypes.uint32).bitcast(dtypes.int32).assign(Tensor([0x40800000, 0x40400000, 0x40000000, 0x3f800000], dtype=dtypes.int32)).realize()
     np.testing.assert_allclose(b.numpy(), [4.0, 3.0, 2.0, 1.0])
-
-  def test_assign_shrink_then_bitcast(self):
+    # shrink then bitcast
     c = Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32).realize()
-    c[0:2].bitcast(dtypes.uint32).assign(Tensor(self.REVERSED[:2], dtype=dtypes.uint32)).realize()
+    c[0:2].bitcast(dtypes.uint32).assign(Tensor([0x40800000, 0x40400000], dtype=dtypes.uint32)).realize()
     np.testing.assert_allclose(c.numpy(), [4.0, 3.0, 3.0, 4.0])
+    # without .realize()
+    a = Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32).realize()
+    a.bitcast(dtypes.uint32).assign(Tensor([0x40800000, 0x40400000, 0x40000000, 0x3f800000], dtype=dtypes.uint32))
+    np.testing.assert_allclose(a.numpy(), [4.0, 3.0, 2.0, 1.0])
 
   def test_assign_bitcast_different_size(self):
-    # assign to a shape-changing bitcast view
+    # assign to a shape-changing bitcast view (only works on DISK currently)
     a = Tensor([0]*8, dtype=dtypes.uint8).realize()
     a.bitcast(dtypes.int64).assign(Tensor([12345], dtype=dtypes.int64)).realize()
-    np.testing.assert_equal(a.numpy(), [57, 48, 0, 0, 0, 0, 0, 0])
+    try:
+      np.testing.assert_equal(a.numpy(), [57, 48, 0, 0, 0, 0, 0, 0])
+    except AssertionError:
+      # TODO: broken now
+      np.testing.assert_equal(a.numpy(), [0]*8)
 
   def test_assign_dtype_mismatch(self):
     # assign should not implicitly cast dtypes - this can lose precision
@@ -521,6 +562,13 @@ class TestAssign(unittest.TestCase):
     with self.assertRaisesRegex(RuntimeError, "assign dtype mismatch"):
       a.assign(b)
 
+  def test_assign_dtype_mismatch_int64_to_float32(self):
+    # int64 -> float32 loses precision for large values, should not be implicit
+    a = Tensor.zeros(1, dtype=dtypes.float32).contiguous().realize()
+    b = Tensor([16777217], dtype=dtypes.int64)  # 2^24 + 1, not exactly representable in float32
+    with self.assertRaisesRegex(RuntimeError, "assign dtype mismatch"):
+      a.assign(b)
+
   def test_assign_shape_broadcast(self):
     # shape broadcasting should work when dtypes match
     a = Tensor.zeros(3, 5, dtype=dtypes.float32).contiguous().realize()
@@ -629,7 +677,7 @@ class TestAssign(unittest.TestCase):
     contig.assign(Tensor([1, 4, 3], dtype=dtypes.int64))
     GlobalCounters.reset()
     base.assign(contig).realize()
-    assert_kernel_count(4 if is_hcq2_device() else 2)  # TODO: first copy is dead, could be 1
+    assert_kernel_count(2)  # TODO: first copy is dead, could be 1
     self.assertEqual(base.tolist(), [1,4,3])
 
   def test_nested_after_contiguous_store_no_init(self):
@@ -639,17 +687,9 @@ class TestAssign(unittest.TestCase):
     contig.assign(Tensor([1, 4, 3], dtype=dtypes.int64))
     GlobalCounters.reset()
     base.assign(contig).realize()
-    assert_kernel_count(2 if is_hcq2_device() else 1)
+    assert_kernel_count(1)
     self.assertEqual(base.tolist(), [1,4,3])
 
-  def test_assign_temporary_copy_reshape(self):
-    a = Tensor([[1., 2], [3, 4]], device="PYTHON")
-    c = Tensor.empty(2, 2).assign(a.to(None))
-    GlobalCounters.reset()
-    c.realize()
-    assert_kernel_count(2 if is_hcq2_device() else 1)
-    self.assertEqual(c.tolist(), [[1., 2], [3, 4]])
-
 class TestAssignOrdering(unittest.TestCase):
   """Tests for complex assign orderings that could differ between lazy and eager execution.
 
@@ -841,16 +881,14 @@ class TestAssignOrdering(unittest.TestCase):
   def test_war_reader_already_depends_on_write(self):
     x = Tensor([1.0]).contiguous().realize()
     y = Tensor([2.0]).contiguous().realize()
-    x_expr = x + 10          # 11, x is read here, before the assign
+    x_expr = x + 10
     x.assign(x * 2)
     y.assign(y + x)
     z = y + x_expr
     Tensor.realize(x, y, z)
-    try:
-      np.testing.assert_allclose([x.item(), y.item(), z.item()], [2.0, 4.0, 15.0])
-    except AssertionError:
-      # TODO: broken now, x_expr reads x after the assign
-      np.testing.assert_allclose([x.item(), y.item(), z.item()], [2.0, 4.0, 16.0])
+    # TODO: z should be 15: x_expr means 11 (x captured at build time), but the read is fused past the assign and
+    # sees the new bytes. once stale readers are scheduled before the overwrite, update this to 15
+    np.testing.assert_allclose([x.item(), y.item(), z.item()], [2.0, 4.0, 16.0])
 
   def test_war_multi_read_then_assign(self):
     devices = ("CPU:0", "CPU:1")
@@ -871,147 +909,6 @@ class TestAssignOrdering(unittest.TestCase):
     self.assertEqual(buf.sum().realize().item(), 6.0)
 
 # TODO: assigns into views of unrealized non-BUFFER bases are silently dropped
-  def test_read_before_two_assigns(self):
-    g = Tensor.full((2,), 4.0).realize()
-    before = g + 1                                     # 5
-    g.assign(0.0)
-    g.assign(g + 4)
-    with self.assertRaisesRegex(RuntimeError, "cycle"):  # TODO: broken now, raises
-      np.testing.assert_allclose((before + g).numpy(), 9)
-
-  def test_read_between_two_assigns(self):
-    a = Tensor.ones(4).realize()
-    b = Tensor.full((4,), 10.).realize()
-    a.assign(b + 1)                                    # a == 11
-    v1 = a * 3                                         # reads 11 -> 33
-    a.assign(b + 100)                                  # a == 110
-    out = (a + v1).numpy()
-    try:
-      np.testing.assert_allclose(out, 143)
-    except AssertionError:
-      # TODO: broken now, v1 reads a after the second assign
-      np.testing.assert_allclose(out, 440)
-
-  def test_two_reads_between_three_assigns(self):
-    a = Tensor.zeros(4).realize()
-    first = a + 100
-    a.assign(Tensor([1., 2., 0., 0.]))
-    second = a + 0
-    a.assign(a + 10)
-    with self.assertRaisesRegex(RuntimeError, "cycle"):  # TODO: broken now, raises
-      np.testing.assert_allclose((first + second + a).numpy(), [112, 114, 110, 110])
-
-  def test_read_before_slice_assign(self):
-    a = Tensor.ones(4).realize()
-    before = a * 3
-    a[0:2].assign(Tensor.full((2,), 2.))
-    out = (a + (before - 1)).numpy()
-    try:
-      np.testing.assert_allclose(out, [4, 4, 3, 3])
-    except AssertionError:
-      # TODO: broken now, before reads the two assigned elements after the assign
-      np.testing.assert_allclose(out, [7, 7, 3, 3])
-
-  def test_read_before_assign_survives_a_realize(self):
-    a = Tensor.ones(4).realize()
-    before = a * 3
-    a.assign(Tensor.full((4,), 5.))
-    a.realize()
-    out = before.numpy()
-    try:
-      np.testing.assert_allclose(out, 3)
-    except AssertionError:
-      # TODO: broken now, before is computed again from the assigned value
-      np.testing.assert_allclose(out, 15)
-
-  def test_loss_read_after_step_is_the_pre_step_loss(self):
-    from tinygrad import nn
-    w = Tensor([2.]).contiguous().realize()
-    x = Tensor([3.]).realize()
-    opt = nn.optim.SGD([w], lr=0.1)
-    with Context(TRAINING=1):
-      loss = (w*x).sum()                               # 6.0
-      loss.backward()
-      opt.step()                                       # w becomes 1.7
-    out = loss.item()
-    try:
-      self.assertAlmostEqual(out, 6.0, places=5)
-    except AssertionError:
-      # TODO: broken now, loss is computed again from the updated weight
-      self.assertAlmostEqual(out, 5.1, places=5)
-
-  def test_rand_realized_out_of_order(self):
-    Tensor.manual_seed(1)
-    r = [Tensor.rand(4) for _ in range(4)]
-    r[3].realize()
-    out_of_order = r[0].numpy()
-    Tensor.manual_seed(1)
-    in_order = [Tensor.rand(4).numpy() for _ in range(4)]
-    try:
-      np.testing.assert_equal(out_of_order, in_order[0])
-    except AssertionError:
-      # TODO: broken now, r[0] returns the fourth set of numbers
-      np.testing.assert_equal(out_of_order, in_order[3])
-
-  def test_batchnorm_stats_are_realized(self):
-    from tinygrad import nn
-    bn, x = nn.BatchNorm(4), Tensor.randn(2, 4, 3, 3).realize()
-    with Context(TRAINING=1): bn(x).realize()
-    try:
-      self.assertTrue(bn.running_mean.uop.base.is_realized)
-    except AssertionError:
-      # TODO: broken now, the stat update is never run because nothing reads it
-      self.assertFalse(bn.running_mean.uop.base.is_realized)
-
-  def test_batchnorm_under_jit_counts_every_call(self):
-    from tinygrad import nn
-    bn, x = nn.BatchNorm(4), Tensor.randn(8, 4, 2, 2).realize()
-    @TinyJit
-    def step(t):
-      with Context(TRAINING=1): return bn(t).sum().realize()
-    for _ in range(4): step(x)
-    out = bn.num_batches_tracked.item()
-    try:
-      self.assertEqual(out, 4)
-    except AssertionError:
-      # TODO: broken now, only the calls whose stat update happened to be captured are counted
-      self.assertEqual(out, 2)
-
-  def test_assign_from_unrealized_tensor_does_not_alias(self):
-    a = Tensor.full((4,), 7.).realize()
-    b = Tensor.ones(4) * 1
-    b.assign(a)
-    b.assign(Tensor.zeros(4))
-    b.realize()
-    self.assertListEqual(a.tolist(), [7., 7., 7., 7.])
-
-  def test_assign_to_function_output(self):
-    from tinygrad import function
-    @function
-    def f(x:Tensor) -> Tensor: return x*2
-    out = f(Tensor.ones(4).realize())
-    out.assign(Tensor.full((4,), 9.).realize())
-    self.assertListEqual(out.tolist(), [9., 9., 9., 9.])
-
-  def test_nested_function_assign(self):
-    from tinygrad import function
-    @function
-    def inner(x:Tensor) -> Tensor:
-      x.assign(x+1)
-      return x*2
-    @function
-    def outer(x:Tensor) -> Tensor:
-      y = inner(x)
-      x.assign(x+1)
-      return y+x
-    a = Tensor([1.]).realize()
-    out = outer(a).item()
-    try:
-      self.assertEqual([out, a.item()], [7., 3.])
-    except AssertionError:
-      # TODO: broken now, the inner assign is run twice
-      self.assertEqual([out, a.item()], [6., 4.])
-
 class TestAssignToUnrealizedView(unittest.TestCase):
   def test_copy(self):
     t = Tensor.zeros(2,2, dtype=dtypes.int).to("CPU:0").contiguous().realize()
@@ -1024,12 +921,16 @@ class TestAssignToUnrealizedView(unittest.TestCase):
       # TODO: broken now
       self.assertEqual(c.tolist(), [[0,0],[0,0]])
 
-  def test_clone(self):
+  def test_contiguous(self):
     t = Tensor([[1,2],[3,4]]).contiguous().realize()
-    c = t.permute(1,0).clone()
-    self.assertIs(c.uop.base.op, Ops.AFTER)
+    c = t.permute(1,0).contiguous()  # unrealized CONTIGUOUS
+    self.assertIs(c.uop.base.op, Ops.CONTIGUOUS)
     c[:, 1:2].assign(Tensor.ones(2,1, dtype=dtypes.int).contiguous().realize())
-    self.assertEqual(c.tolist(), [[1,1],[2,1]])
+    try:
+      self.assertEqual(c.tolist(), [[1,1],[2,1]])
+    except AssertionError:
+      # TODO: broken now
+      self.assertEqual(c.tolist(), [[1,3],[2,4]])
 
   def test_contiguous_backward(self):
     t = Tensor([[1,2],[3,4]]).contiguous().realize()
@@ -1042,16 +943,6 @@ class TestAssignToUnrealizedView(unittest.TestCase):
       # TODO: broken now
       self.assertEqual(cb.tolist(), [[1,2],[3,4]])
 
-  def test_detach_buffer_assignment(self):
-    for realized in (False, True):
-      with self.subTest(realized=realized):
-        base = Tensor([1., 2., 3.])
-        if realized: base.realize()
-        detached = base.detach()
-        detached.assign(detached + 1).realize()
-        self.assertEqual(detached.tolist(), [2., 3., 4.])
-        self.assertEqual(base.tolist(), [2., 3., 4.])
-
   def test_detach_copy(self):
     t = Tensor.zeros(2,2, dtype=dtypes.int).to("CPU:0").contiguous().realize()
     d = t.to("CPU:1").detach()  # DETACH(unrealized COPY)
@@ -1063,12 +954,16 @@ class TestAssignToUnrealizedView(unittest.TestCase):
       # TODO: broken now
       self.assertEqual(d.tolist(), [[0,0],[0,0]])
 
-  def test_detach_clone(self):
+  def test_detach_contiguous(self):
     t = Tensor([[1,2],[3,4]]).contiguous().realize()
-    d = t.permute(1,0).clone().detach()
-    self.assertIs(d.uop.base.op, Ops.AFTER)
+    d = t.permute(1,0).contiguous().detach()  # DETACH(unrealized CONTIGUOUS)
+    self.assertIs(d.uop.base.op, Ops.CONTIGUOUS)
     d[:, 1:2].assign(Tensor.ones(2,1, dtype=dtypes.int).contiguous().realize())
-    self.assertEqual(d.tolist(), [[1,1],[2,1]])
+    try:
+      self.assertEqual(d.tolist(), [[1,1],[2,1]])
+    except AssertionError:
+      # TODO: broken now
+      self.assertEqual(d.tolist(), [[1,3],[2,4]])
 
   def test_alu(self):
     a = Tensor([1,2,3,4]).contiguous().realize()
@@ -1114,16 +1009,6 @@ class TestAssignToUnrealizedView(unittest.TestCase):
       # TODO: broken now, silently dropped
       self.assertEqual(c.tolist(), [[5,5],[5,5]])
 
-  def test_detach_assignment_preserves_earlier_update(self):
-    x = Tensor([1., 2.]).detach()
-    state = Tensor([0., 0.]).detach()
-    state.assign(state + x * 2)
-    result = state + 1
-    x.assign(x + 1).realize(state, result)
-    self.assertEqual(x.tolist(), [2., 3.])
-    self.assertEqual(state.tolist(), [2., 4.])
-    self.assertEqual(result.tolist(), [3., 5.])
-
 class TestPartialAssignToSharedBuffer(unittest.TestCase):
   def test_five_slices(self):
     big = Tensor.zeros(50).contiguous().realize()
@@ -1155,6 +1040,7 @@ class TestPartialAssignToSharedBuffer(unittest.TestCase):
     for v, s in zip(views, shapes):
       np.testing.assert_allclose(v.numpy(), np.ones(s))
 
+
 class TestAfterCachePatterns(unittest.TestCase):
   def test_double_store_after(self):
     a = Tensor.zeros(10).contiguous()
@@ -1185,6 +1071,14 @@ class TestAfterCachePatterns(unittest.TestCase):
     np.testing.assert_array_equal(head.numpy(), [3])
     np.testing.assert_array_equal(full.numpy(), [1, 2])
 
+class TestBatchNormRunningStats(unittest.TestCase):
+  @unittest.expectedFailure  # TODO: nothing reads the stat update so it is never scheduled, and the chain grows every step
+  def test_running_stats_are_realized(self):
+    from tinygrad import nn
+    bn, x = nn.BatchNorm(4), Tensor.randn(2, 4, 3, 3).contiguous().realize()
+    with Context(TRAINING=1): bn(x).realize()
+    self.assertTrue(bn.running_mean.uop.base.is_realized)
+
 class TestMultiAssign(unittest.TestCase):
   device = tuple(f"{Device.DEFAULT}:{i}" for i in range(2))
 
@@ -1231,15 +1125,12 @@ class TestMultiAssign(unittest.TestCase):
     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()
-    try:
-      self.assertListEqual(out.tolist(), [[0,0,1,0], [0,0,1,0], [0,0,1,0], [0,0,1,0]])
-    except AssertionError:
-      # TODO: broken now, the write is dropped
-      self.assertListEqual(out.tolist(), [[0,0,0,0], [0,0,0,0], [0,0,0,0], [0,0,0,0]])
+    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()
@@ -1263,6 +1154,5 @@ class TestMultiAssign(unittest.TestCase):
       GlobalCounters.reset()
       f(out, vi.bind(i))
     self.assertListEqual(out.tolist(), [[0,1,2,3,4,0]]*4)
-
 if __name__ == "__main__":
   unittest.main()

test/backend/test_call.py

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

test/backend/test_const_folding.py

@@ -86,8 +86,8 @@ class TestReduceOpsConstFolding(unittest.TestCase):
   def test_zero_size_realize_folded(self):
     # non contiguous folded output doesn't realize
     _check_ast_count(0, Tensor.empty(1, 0).sum())
-    # An explicitly cloned folded constant still owns persistent storage.
-    a = Tensor.empty(1, 0).sum().clone()
+    # contiguous folded const can still schedule
+    a = Tensor.empty(1, 0).sum().contiguous()
     _check_ast_count(2, a+2)
     self.assertIs(a.uop.base.op, Ops.BUFFER)
     np.testing.assert_equal((Tensor.empty(1, 0).sum().contiguous()+2).numpy(), 2)

test/backend/test_custom_kernel.py

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

test/backend/test_dtype.py

@@ -9,7 +9,7 @@ from tinygrad.renderer.nir import NIRRenderer
 from tinygrad import Context, Device, Tensor, dtypes
 from hypothesis import given, settings, strategies as strat
 from test.helpers import rand_for_dtype, min_normal
-from test.unit.test_dtype_spec import _assert_eq, core_dtypes, FP8E4M3_MAX, FP8E5M2_MAX, FP8E4M3FNUZ_MAX, FP8E5M2FNUZ_MAX
+from test.unit.test_dtype_spec import _assert_eq, core_dtypes, dtype_ints, dtype_floats, FP8E4M3_MAX, FP8E5M2_MAX, FP8E4M3FNUZ_MAX, FP8E5M2FNUZ_MAX
 import pytest
 pytestmark = pytest.mark.filterwarnings("ignore")
 
@@ -19,8 +19,7 @@ settings.load_profile("my_profile")
 supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes()
 
 def get_available_cast_dtypes(dtype: DType) -> List[DType]:
-  emulatable = dtypes.fp8s+(dtypes.half,dtypes.bfloat16,dtypes.long)
-  dts = [v for v in dict.fromkeys(DTYPES_DICT.values()) if v != dtype and (v in supported_dtypes or v in emulatable)]
+  dts = [v for k, v in DTYPES_DICT.items() if v != dtype and v in supported_dtypes or v in dtypes.fp8s+(dtypes.half,dtypes.bfloat16,dtypes.long)]
   if dtype in (dtypes.long, dtypes.ulong) and (dtype not in supported_dtypes or dtypes.long in EMULATED_DTYPES.tolist(dtypes)):
     return [dt for dt in dts if dt != dtypes.double] # can't bitcast with no 64-bit support
   if dtype not in supported_dtypes and dtype not in dtypes.fp8s+(dtypes.half,dtypes.bfloat16): return []
@@ -72,14 +71,14 @@ class TestDType(unittest.TestCase):
     self.assertEqual(a.dtype, self.DTYPE)
     _test_to_np(a, _to_np_dtype(self.DTYPE), np.array(self.DATA, dtype=_to_np_dtype(self.DTYPE)))
 
+  def test_casts_to(self):
+    for dtype in get_available_cast_dtypes(self.DTYPE):
+      _test_cast(Tensor(self.DATA, dtype=dtype), self.DTYPE)
+
   def test_casts_from(self):
     for dtype in get_available_cast_dtypes(self.DTYPE):
       _test_cast(Tensor(self.DATA, dtype=self.DTYPE), dtype)
 
-  def test_const_kernel(self):
-    if not get_available_cast_dtypes(self.DTYPE): raise unittest.SkipTest("dtype does not run here")
-    _assert_eq(Tensor.ones((4,4), dtype=self.DTYPE).clone(), self.DTYPE, np.ones((4,4)))
-
   def test_same_size_ops(self):
     for dtype in get_available_cast_dtypes(self.DTYPE):
       if dtype.itemsize == self.DTYPE.itemsize:
@@ -90,10 +89,10 @@ class TestDType(unittest.TestCase):
       if dtype.itemsize > self.DTYPE.itemsize:
         _test_ops(a_dtype=self.DTYPE, b_dtype=dtype)
 
-  def test_downcast_ops(self):
+  def test_upcast_to_ops(self):
     for dtype in get_available_cast_dtypes(self.DTYPE):
       if dtype.itemsize < self.DTYPE.itemsize:
-        _test_ops(a_dtype=self.DTYPE, b_dtype=dtype)
+        _test_ops(a_dtype=dtype, b_dtype=self.DTYPE)
 
   def test_bitcast(self):
     if self.DTYPE == dtypes.bool: raise unittest.SkipTest("no bools in bitcast")
@@ -113,7 +112,12 @@ def _test_ops(a_dtype:DType, b_dtype:DType, target_dtype=None):
   target_dtype = target_dtype or least_upper_dtype(a_dtype, b_dtype)
   if a_dtype == dtypes.bool or b_dtype == dtypes.bool: return
   _assert_eq(Tensor([1,2,3,4], dtype=a_dtype)+Tensor([1,2,3,4], dtype=b_dtype), target_dtype, [2,4,6,8])
+  _assert_eq((Tensor([1], dtype=a_dtype).cast(b_dtype)+Tensor([1], dtype=a_dtype).cast(b_dtype)).cast(a_dtype), a_dtype, [2])
+  _assert_eq(Tensor([1,2,3,4], dtype=a_dtype)*Tensor([1,2,3,4], dtype=b_dtype), target_dtype, [1,4,9,16])
   _assert_eq(Tensor([[1,2],[3,4]], dtype=a_dtype)@Tensor.eye(2, dtype=b_dtype), target_dtype, [[1,2],[3,4]])
+  _assert_eq(Tensor([1,1,1,1], dtype=a_dtype)+Tensor.ones((4,4), dtype=b_dtype), target_dtype, 2*np.ones((4,4)))
+  _assert_eq(Tensor([1,1,1,1], dtype=a_dtype)+Tensor.ones((4,4), dtype=b_dtype).clone(), target_dtype, 2*np.ones((4,4)))
+  _assert_eq(Tensor.ones((4,4), dtype=b_dtype).clone(), b_dtype, np.ones((4,4)))
 
 class TestFp8sConversions(unittest.TestCase):
   @given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E4M3_MAX, max_value=FP8E4M3_MAX))
@@ -254,11 +258,6 @@ 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])
@@ -284,10 +283,14 @@ class TestUint8DType(TestDType):
     _test_op(lambda: Tensor([255, 254, 253, 252], dtype=dtypes.uint8).cast(dtypes.int8), dtypes.int8, [-1, -2, -3, -4])
 
 class TestBitCast(unittest.TestCase):
-  def test_shape_change_bitcast(self):
-    for dt1, dt2 in [(dtypes.uint8, dtypes.int64), (dtypes.int64, dtypes.uint8)]:
-      a = Tensor(rand_for_dtype(dt1, 32).reshape(2, 2, 8), dtype=dt1)
-      _test_op(lambda: a.bitcast(dt2), dt2, _to_torch_storage(a).view(_to_torch_dtype(dt2)).tolist())
+  @given(strat.sampled_from(dtype_ints + dtype_floats), strat.sampled_from(dtype_ints + dtype_floats))
+  def test_shape_change_bitcast(self, dt1, dt2):
+    data = rand_for_dtype(dt1, 32).reshape(2, 2, 8)
+    a = Tensor(data, dtype=dt1)
+    expected = _to_torch_storage(a).view(_to_torch_dtype(dt2))
+    if dt2 in dtypes.fp8s:
+      expected = torch.tensor([fp8_to_float(x, dt2) for x in expected.view(-1).tolist()]).view_as(expected)
+    _test_op(lambda: a.bitcast(dt2), dt2, expected.tolist())
 
   def test_shape_change_bitcast_exceptions(self):
     with self.assertRaises(RuntimeError):
@@ -320,10 +323,7 @@ class TestUint16DType(TestDType):
 class TestInt32DType(TestDType): DTYPE = dtypes.int32
 class TestUint32DType(TestDType): DTYPE = dtypes.uint32
 
-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])
+class TestInt64DType(TestDType): DTYPE = dtypes.int64
 
 @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
 class TestEmulatedInt64DType(TestInt64DType):
@@ -393,9 +393,6 @@ class TestEmulatedFp8e5m2(TestFp8e5m2):
   @classmethod
   def tearDownClass(cls): cls.stack.close()
 
-class TestFp8e4m3fnuz(TestDType): DTYPE = dtypes.fp8e4m3fnuz
-class TestFp8e5m2fnuz(TestDType): DTYPE = dtypes.fp8e5m2fnuz
-
 class TestImplicitFunctionTypeChange(unittest.TestCase):
   def test_functions(self):
     result = []
@@ -426,11 +423,6 @@ class TestDtypeUsage(unittest.TestCase):
         t = Tensor([[1, 2], [3, 4]], dtype=d)
         (t*t).max().item()
 
-  def test_where_float16_compare_to_const(self):
-    # t > 0 is CMPLT(0, t): the float16 operand is on the right
-    t = Tensor([-1.0, 1.0], dtype=dtypes.float16)
-    np.testing.assert_equal((t > 0).where(Tensor.ones(2, dtype=dtypes.float16), Tensor.zeros(2, dtype=dtypes.float16)).numpy(), [0.0, 1.0])
-
 @unittest.skipUnless(dtypes.bfloat16 in supported_dtypes, f"no bfloat16 on {Device.DEFAULT}")
 class TestOpsBFloat16(unittest.TestCase):
   def test_cast(self):

test/backend/test_encodings.py

@@ -4,7 +4,7 @@ from tinygrad.uop.ops import UOp, Ops
 from tinygrad.dtype import dtypes
 from tinygrad.renderer.isa.x86 import X86Ops, X86Renderer, RBP, RDI, RSP, RSI, RAX, RDX, XMM, GPR, imm, def_reg
 
-def ins(op, dt, src, tag=None): return UOp(Ops.INS, arg=(op, dt), src=src, tag=tag)
+def ins(op, dt, src, tag=None): return UOp(Ops.INS, arg=op, dtype=dt, src=src, tag=tag)
 
 @unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "only on x86")
 class TestEncodingsX86(unittest.TestCase):
@@ -100,6 +100,13 @@ 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])
@@ -136,7 +143,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, dtypes.void))), RDX)
+    cmove = ins(X86Ops.CMOVE, dtypes.int32, (def_reg(dtypes.int32, RAX), UOp(Ops.INS, arg=X86Ops.CMP)), RDX)
     # cmove edx, eax
     self.assertEqual(bytes.fromhex(self.encode(cmove)), bytes.fromhex("0F 44 D0"))
 

test/backend/test_graph.py

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

test/backend/test_isel.py

@@ -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[0], op)
+        self.assertIs(n.arg, 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[0] is X86Ops.CMOVL and n.src[1].arg[0] is X86Ops.CMOVNE)
+    self.assertTrue(n.src[0].arg is X86Ops.CMOVL and n.src[1].arg 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[0] is X86Ops.CMP)
+    self.assertTrue(n.src[0].src[2] == n.src[1].src[2] and n.src[0].src[2].arg 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[0], X86Ops.VINSERTPS)
+    for shuf in valid: self.assertIs(self.isel_rewrite(shuf).arg, X86Ops.VINSERTPS)
 
   # complex address is [base + index*scale + displacement]
   def test_complex_address(self):
     a = UOp.variable("a", 0, 0, dtypes.int32)
-    load = UOp.param(0, dtypes.int32, 16).index(a + UOp.cconst(1, dtypes.int32)).load()
+    load = UOp.param(0, dtypes.int32, (16,)).index(a + UOp.cconst(1, dtypes.int32)).load()
     n = self.isel_rewrite(load)
     # displacement is the constant in "a" scaled to the buffer element size, dtype is int8 when the value fits otherwise int32
     self.assertTrue(n.src[2].dtype is dtypes.int8 and n.src[2].src[0].op is Ops.CONST and n.src[2].src[0].val == 4)

test/backend/test_jit.py

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

test/backend/test_linearizer.py

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

test/backend/test_linearizer_dumb.py

@@ -11,20 +11,20 @@ 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)
-    ast = c12.sink(arg=KernelInfo(name='test', applied_opts=(Opt(op=OptOps.SPLIT, axis=4, arg=(16, AxisType.GROUP_REDUCE)),), opts_to_apply=None))
+    ast = c12.sink(arg=KernelInfo(name='test', axis_types=(), dont_use_locals=False, applied_opts=(Opt(op=OptOps.GROUP, axis=1, arg=16),), opts_to_apply=None))
     _ = to_program(ast, Device["METAL"].renderer)
 
 if __name__ == '__main__':

test/backend/test_multitensor.py

@@ -3,7 +3,7 @@ from tinygrad import Tensor, Device, nn, GlobalCounters, TinyJit, dtypes, Variab
 from tinygrad.uop.ops import Ops, UOp, AxisType, graph_rewrite
 from tinygrad.helpers import getenv, prod, Context
 from tinygrad.nn.state import get_parameters
-from tinygrad.engine.realize import run_linear, lower_and_compile, pm_beam
+from tinygrad.engine.realize import run_linear, compile_linear, lower_and_compile, pm_beam
 import numpy as np
 from hypothesis import given, strategies as strat, settings
 from test.helpers import not_support_multi_device, needs_second_gpu, slow, call_is_graph, check_schedule, assert_kernel_count, KernelCountException
@@ -76,7 +76,7 @@ class TestMultiTensor(unittest.TestCase):
     X = Tensor.ones(256).contiguous().realize()
     X.shard_(devices_2, 0)
     out = (X + X)
-    linear = lower_and_compile(out.schedule_linear())
+    linear = compile_linear(out.schedule_linear())
     uops = [call.src[0].src[0] for call in linear.src if call.src[0].op is Ops.PROGRAM]
     run_linear(linear)
     self.assertEqual(len(set(uops)), 1, "function was relinearized")
@@ -187,13 +187,6 @@ 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()

test/backend/test_nn.py

@@ -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, 31], C2=32, K=(3,5), S=1, P='same', D=(2,3))
+    self._test_conv(Conv2d, torch.nn.Conv2d, BS=16, C1=3, DIMS=[28, 28], C2=32, K=3, S=1, P='same', D=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')

test/backend/test_ops.py

@@ -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,13 +359,6 @@ 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]])
@@ -815,6 +809,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)
@@ -822,7 +818,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([], lambda: tor&tor, lambda: ten&ten, forward_only=True)
     helper_test_op([], lambda: tor&0x1337, lambda: ten&0x1337, forward_only=True)
     helper_test_op([], lambda: 0x1337&tor, lambda: 0x1337&ten, forward_only=True)
-    helper_test_op([], lambda: (tor&12)&tor, lambda: (ten&12)&ten, forward_only=True)
 
     data = [[True, True, False, False], [True, False, True, False]]
     tor0, tor1 = torch.tensor(data[0], dtype=torch.bool),  torch.tensor(data[1], dtype=torch.bool)
@@ -953,18 +948,15 @@ class TestOps(unittest.TestCase):
     helper_test_op([(45,65)], lambda x: x.asin(), low=-1, high=1)
     helper_test_op([(45,65)], lambda x: x.asin(), low=-300, high=-297)
     helper_test_op([(45,65)], lambda x: x.asin(), low=300, high=303)
-    helper_test_op(None, lambda x: x.asin(), vals=[[-0.5, 0., 0.5]])
   def test_acos(self):
     # high grad atol
     helper_test_op([(45,65)], lambda x: x.acos(), low=-1, high=1)
     helper_test_op([(45,65)], lambda x: x.acos(), low=-300, high=-297)
     helper_test_op([(45,65)], lambda x: x.acos(), low=300, high=303)
-    helper_test_op(None, lambda x: x.acos(), vals=[[-0.5, 0., 0.5]])
   def test_atan(self):
     helper_test_op([(45,65)], lambda x: x.atan())
     helper_test_op([(45,65)], lambda x: x.atan(), low=-300, high=-297)
     helper_test_op([(45,65)], lambda x: x.atan(), low=300, high=303)
-    helper_test_op(None, lambda x: x.atan(), vals=[[-0.5, 0., 0.5]])
 
   def test_relu(self):
     helper_test_op([(64,64)], lambda x: x.relu())
@@ -979,12 +971,9 @@ class TestOps(unittest.TestCase):
   def test_celu(self):
     for val in range(1, 5):
       helper_test_op([(45,65)], lambda x: torch.nn.functional.celu(x,val), lambda x: x.celu(val))
-      helper_test_op([(3,3)], lambda x: torch.nn.functional.celu(x,val), lambda x: x.celu(val), low=300, high=400)
       helper_test_op([()], lambda x: torch.nn.functional.celu(x,val), lambda x: x.celu(val))
   def test_selu(self):
     helper_test_op([(45,65)], torch.nn.functional.selu, Tensor.selu)
-    helper_test_op([(3,3)], torch.nn.functional.selu, Tensor.selu, low=300, high=400)
-    helper_test_op(None, torch.nn.functional.selu, Tensor.selu, vals=[[-1.,0.,1.]])
     helper_test_op([()], torch.nn.functional.selu, Tensor.selu)
   def test_silu(self):
     helper_test_op([(45,65)], torch.nn.functional.silu, Tensor.silu)
@@ -1051,7 +1040,6 @@ class TestOps(unittest.TestCase):
     helper_test_op(None, torch.logaddexp, Tensor.logaddexp, vals=[[-1.], [-1.0, 2, 3]])
     helper_test_op(None, torch.logaddexp, Tensor.logaddexp, vals=[[-100.0, -200, -300], [-1.0, 2, 3]])
     helper_test_op(None, torch.logaddexp, Tensor.logaddexp, vals=[[1.0, 2000, 30000], [-1.0, 2, 3]])
-    helper_test_op(None, torch.logaddexp, Tensor.logaddexp, vals=[[-math.inf, math.inf, 1.0, -math.inf], [-math.inf, math.inf, -math.inf, 1.0]])
 
   def test_softsign(self):
     helper_test_op([(45,65)], torch.nn.functional.softsign, Tensor.softsign)
@@ -1093,13 +1081,11 @@ class TestOps(unittest.TestCase):
     helper_test_op([(45,65)], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6, low=300, high=400)
     helper_test_op([(45,65)], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6, low=-400, high=-300)
     helper_test_op([()], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6)
-    helper_test_op(None, torch.nn.functional.softplus, Tensor.softplus, vals=[[-math.inf, math.inf, 0.0]], forward_only=True)
 
   def test_erf(self):
     helper_test_op([(45,65)], torch.erf, Tensor.erf)
     helper_test_op([(45,65)], torch.erf, Tensor.erf, low=300, high=400)
     helper_test_op([(45,65)], torch.erf, Tensor.erf, low=-400, high=-300)
-    helper_test_op(None, torch.erf, Tensor.erf, vals=[[-1., 0., 1.]])
     helper_test_op([()], torch.erf, Tensor.erf)
 
   def test_gelu(self):
@@ -1124,7 +1110,6 @@ class TestOps(unittest.TestCase):
   def test_elu(self):
     helper_test_op([(45,65)], torch.nn.functional.elu, Tensor.elu)
     helper_test_op([(45,65)], lambda x: torch.nn.functional.elu(x, alpha=0.1), lambda x: Tensor.elu(x, alpha=0.1))
-    helper_test_op([(3,3)], torch.nn.functional.elu, Tensor.elu, low=300, high=400)
     helper_test_op([()], torch.nn.functional.elu, Tensor.elu)
   def test_relu6(self):
     helper_test_op([(45,65)], torch.nn.functional.relu6, Tensor.relu6)
@@ -1776,9 +1761,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([(45,65)], lambda x: torch.nn.functional.normalize(x, p=3, dim=0), lambda x: x.normalize(p=3, dim=0), atol=1e-7, grad_atol=1e-7)
     helper_test_op([(45,65)], lambda x: torch.nn.functional.normalize(x, p=0), lambda x: x.normalize(p=0), atol=1e-7, grad_atol=1e-7)
     helper_test_op([(45,65)], lambda x: torch.nn.functional.normalize(x, p=-1), lambda x: x.normalize(p=-1), atol=1e-7, grad_atol=1e-7)
-  def test_normalize_int(self):
-    helper_test_op(None, lambda x: torch.nn.functional.normalize(x.float(), p=2), lambda x: x.normalize(p=2), forward_only=True,
-                   vals=[[[3, 4], [6, 8]]])
 
   def test_logsumexp(self):
     helper_test_op([(45,65)], lambda x: torch.logsumexp(x, dim=0), lambda x: x.logsumexp(0), atol=1e-7, grad_atol=1e-7)
@@ -1791,7 +1773,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([(45)], lambda x: torch.logsumexp(x, dim=0), lambda x: x.logsumexp(0), atol=1e-7, grad_atol=1e-7)
     helper_test_op([()], lambda x: torch.logsumexp(x, dim=0), lambda x: x.logsumexp(0), atol=1e-7, grad_atol=1e-7)
     helper_test_op([()], lambda x: torch.logsumexp(x, dim=-1), lambda x: x.logsumexp(-1), atol=1e-7, grad_atol=1e-7)
-    helper_test_op(None, lambda x: torch.logsumexp(x, dim=0), lambda x: x.logsumexp(0), vals=[[-math.inf, -math.inf]], forward_only=True)
 
   @slow_test
   def test_logcumsumexp(self):
@@ -1807,7 +1788,6 @@ class TestOps(unittest.TestCase):
 
   def test_logcumsumexp_numerical(self):
     helper_test_op(None, lambda x: torch.logcumsumexp(x, dim=0), lambda x: x.logcumsumexp(), atol=1e-7, grad_atol=1e-7, vals=[[0.0, 100.0]])
-    helper_test_op(None, lambda x: torch.logcumsumexp(x, dim=0), lambda x: x.logcumsumexp(), vals=[[-math.inf, 0.0, 1.0]], forward_only=True)
 
   def test_sinh(self):
     helper_test_op([(45,65)], lambda x: x.sinh(), grad_atol=1e-6)
@@ -2830,7 +2810,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,)), ((29,),(1,))]:
+    for in_sz, out_sz in [((52,),(29,)), ((29,),(52,))]:
       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))
@@ -2983,10 +2963,6 @@ 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()

test/backend/test_opt_gemm.py

@@ -4,7 +4,7 @@ from tinygrad import Tensor
 from tinygrad.helpers import get_single_element
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.engine.realize import run_linear
-from tinygrad.uop.ops import Ops, UOp, AxisType
+from tinygrad.uop.ops import Ops, UOp
 from test.helpers import replace_opts
 
 class TestOptGemm(unittest.TestCase):
@@ -26,21 +26,20 @@ class TestOptGemm(unittest.TestCase):
     np.testing.assert_allclose(self.res, test, atol=1e-4)
 
   def test_gemm_unrolled_permute_l_44(self):
-    opts = [Opt(op=OptOps.SPLIT, axis=0, arg=(4, AxisType.UPCAST)), Opt(op=OptOps.SPLIT, axis=1, arg=(4, AxisType.UPCAST))]
+    opts = [Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=4)]
     self._test_gemm_unrolled_permute_l(opts)
 
   def test_gemm_unrolled_permute_l_424(self):
     # was failing with LLVM
-    opts = [Opt(op=OptOps.SPLIT, axis=0, arg=(4, AxisType.UPCAST)), Opt(op=OptOps.SPLIT, axis=1, arg=(2, AxisType.UPCAST)),
-            Opt(op=OptOps.SPLIT, axis=0, arg=(4, AxisType.UPCAST))]
+    opts = [Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=2), Opt(op=OptOps.UPCAST, axis=0, arg=4)]
     self._test_gemm_unrolled_permute_l(opts)
 
   def test_gemm_unrolled_permute_l_42(self):
-    opts = [Opt(op=OptOps.SPLIT, axis=0, arg=(4, AxisType.UPCAST)), Opt(op=OptOps.SPLIT, axis=1, arg=(2, AxisType.UPCAST))]
+    opts = [Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=2)]
     self._test_gemm_unrolled_permute_l(opts)
 
   def test_gemm_unrolled_permute_l_22(self):
-    opts = [Opt(op=OptOps.SPLIT, axis=0, arg=(2, AxisType.UPCAST)), Opt(op=OptOps.SPLIT, axis=1, arg=(2, AxisType.UPCAST))]
+    opts = [Opt(op=OptOps.UPCAST, axis=0, arg=2), Opt(op=OptOps.UPCAST, axis=1, arg=2)]
     self._test_gemm_unrolled_permute_l(opts)
 
 if __name__ == '__main__':

test/backend/test_pickle.py

@@ -133,26 +133,6 @@ class TestPickle(unittest.TestCase):
     t2:Tensor = pickle.loads(st)
     np.testing.assert_equal(t.numpy(), t2.numpy())
 
-  def test_pickle_no_storage_aliasing(self):
-    # loading the same pickle twice gives fully independent storage: the buffers (and their BUFFER uops) are never shared
-    t = Tensor([1,2,3,4]).realize()
-    st = pickle.dumps(t)
-    t1, t2 = pickle.loads(st), pickle.loads(st)
-    self.assertIsNot(t1.uop, t2.uop)
-    self.assertIsNot(t1.uop.base.buffer, t2.uop.base.buffer)
-    t1.assign(Tensor([9,9,9,9])).realize()
-    self.assertListEqual(t1.tolist(), [9,9,9,9])
-    self.assertListEqual(t2.tolist(), [1,2,3,4])
-
-  def test_pickle_view_is_self_contained(self):
-    # a pickled graph carries its own buffer: data from earlier loads of related graphs must not leak into it
-    t = Tensor([1,2,3,4]).realize()
-    t1 = pickle.loads(pickle.dumps(t))
-    t1.assign(Tensor([9,9,9,9])).realize()
-    # loading a view of the original tensor must give the pickled values ([2,3]), not the mutated values from the other load
-    v2 = pickle.loads(pickle.dumps(t[1:3]))
-    self.assertListEqual(v2.realize().tolist(), [2,3])
-
   def test_pickle_jit(self):
     @TinyJit
     def add(a, b): return a.sum()+b+1

test/backend/test_quantize_onnx.py

@@ -2,7 +2,7 @@
 import numpy as np
 import tempfile, unittest
 from tinygrad import Tensor, Context, Device, dtypes, UOp
-from tinygrad.uop.ops import Ops, AxisType
+from tinygrad.uop.ops import Ops
 from tinygrad.dtype import AddrSpace
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.engine.realize import run_linear
@@ -98,7 +98,7 @@ class TestQuantizeOnnx(unittest.TestCase):
     X = Tensor(np.random.uniform(0, 255, size=(1, 32, 128, 128)).astype(np.uint8))
     W = Tensor(np.random.uniform(0, 255, size=(64, 32, 1, 1)).astype(np.uint8))
     out = X.conv2d(W, dtype=X.dtype)
-    opts = [Opt(op=OptOps.SPLIT, axis=1, arg=(128, AxisType.UPCAST)), Opt(op=OptOps.SPLIT, axis=3, arg=(4, AxisType.UNROLL))]
+    opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)]
     sexec(out, opts)
 
   def test_prequant_gemm(self):
@@ -106,7 +106,7 @@ class TestQuantizeOnnx(unittest.TestCase):
     X = Tensor(np.random.uniform(0, 255, size=(N,N)).astype(np.uint8))
     W = Tensor(np.random.uniform(0, 255, size=(N,N)).astype(np.uint8))
     out = X.matmul(W, dtype=X.dtype)
-    opts = [Opt(op=OptOps.SPLIT, axis=1, arg=(128, AxisType.UPCAST)), Opt(op=OptOps.SPLIT, axis=3, arg=(4, AxisType.UNROLL))]
+    opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)]
     sexec(out, opts)
 
   # TODO: this has to work
@@ -116,7 +116,7 @@ class TestQuantizeOnnx(unittest.TestCase):
     W = Tensor(np.random.uniform(0, 255, size=(N,N)).astype(wi))
     # this divide is interesting and forces the accumulator to actually be an int
     out = (X.cast("int").matmul(W.cast("int"))//1000).cast("int8")
-    opts = [Opt(op=OptOps.SPLIT, axis=1, arg=(128, AxisType.UPCAST)), Opt(op=OptOps.SPLIT, axis=3, arg=(4, AxisType.UNROLL))]
+    opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)]
     sexec(out, opts)
 
   def test_prequant_gemm_handcode(self):
@@ -200,11 +200,9 @@ class TestQuantizeOnnx(unittest.TestCase):
     self.test_prequant_gemm_intacc(np.uint8, np.int8, src)
 
   def test_prequant_gemm_intacc_32(self):
-    opts = [Opt(op=OptOps.SPLIT, axis=1, arg=(0, AxisType.UPCAST)), Opt(op=OptOps.SPLIT, axis=0, arg=(4, AxisType.UPCAST)),
-            Opt(op=OptOps.SPLIT, axis=3, arg=(0, AxisType.UNROLL))]
+    opts = [Opt(op=OptOps.UPCAST, axis=1, arg=0), Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UNROLL, axis=0, arg=0)]
     self.test_prequant_gemm_intacc(np.uint8, np.int8, N=32, opts=opts)
-  def test_prequant_gemm_intacc_128(self): self.test_prequant_gemm_intacc(np.uint8, np.int8, N=128,
-    opts=[Opt(op=OptOps.SPLIT, axis=1, arg=(128, AxisType.UPCAST)), Opt(op=OptOps.SPLIT, axis=2, arg=(4, AxisType.UNROLL))])
+  def test_prequant_gemm_intacc_128(self): self.test_prequant_gemm_intacc(np.uint8, np.int8, N=128)
   def test_prequant_gemm_intacc_256(self): self.test_prequant_gemm_intacc(np.uint8, np.int8, N=256)
   def test_prequant_gemm_intacc(self, xi=np.uint8, wi=np.uint8, replace_src=None, N=512, clip=True, opts=None):
     X = Tensor(m1:=(np.random.uniform(0, 255, size=(N,N)).astype(xi))).realize()
@@ -213,8 +211,7 @@ class TestQuantizeOnnx(unittest.TestCase):
     out = (X.int().matmul(W.int())//1000)
     if clip: out = out.clip(tg_dtype.min, tg_dtype.max)
     out = out.cast(tg_dtype)
-    opts = [Opt(op=OptOps.SPLIT, axis=1, arg=(128, AxisType.UPCAST)),
-            Opt(op=OptOps.SPLIT, axis=3, arg=(4, AxisType.UNROLL))] if opts is None else opts
+    opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)] if opts is None else opts
     sexec(out, opts, replace_src, run_count=1)
     tout = out.numpy()
     mout = ((m1.astype(np.int32) @ m2.astype(np.int32)) // 1000)
@@ -235,7 +232,7 @@ class TestQuantizeOnnx(unittest.TestCase):
     #out = X.cast(dtypes.int) @ W.cast(dtypes.int)
     #out = X @ W
     out = X.matmul(W, dtype=X.dtype)
-    opts = [Opt(op=OptOps.SPLIT, axis=0, arg=(128, AxisType.UPCAST)), Opt(op=OptOps.SPLIT, axis=2, arg=(4, AxisType.UNROLL))]
+    opts = [Opt(op=OptOps.UPCAST, axis=0, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)]
     sexec(out, opts)
 
 if __name__ == "__main__":

test/backend/test_rangeify.py

@@ -1,16 +1,55 @@
 import unittest
-from tinygrad import Tensor, dtypes, Variable
+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
+from tinygrad.renderer.nir import NIRRenderer
 
+@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (NIRRenderer, PTXRenderer)), "broken in LVP and PTX")
 class TestDoubleMatmul(unittest.TestCase):
-  def test_double_matmul(self):
+  def setUp(self):
     with Context(DEBUG=0):
-      a, b, c = [Tensor.randn(16, 16).contiguous().realize() for _ in range(3)]
-      ref = a.numpy() @ b.numpy() @ c.numpy()
+      self.a, self.b, self.c = [Tensor.randn(16, 16).contiguous().realize() for _ in range(3)]
+      self.ref = (self.a @ self.b @ self.c).realize()
+
+  def _test(self, opts):
     with Context(DEBUG=max(2, DEBUG.value)):
-      out = (a @ b @ c).numpy()
-    self.assertLess(abs(out-ref).max(), 1e-3)
+      out = (self.a @ self.b @ self.c).contiguous(arg=opts).realize()
+
+    with Context(DEBUG=0):
+      err = (out-self.ref).square()
+      self.assertLess(err.max().item(), 1e-4)
+      self.assertLess(err.mean().item(), 1e-6)
+
+  def test_baseline(self): self._test(())
+  def test_upcast_0(self): self._test((Opt(OptOps.UPCAST, 0, 4),))
+  def test_upcast_1(self): self._test((Opt(OptOps.UPCAST, 1, 4),))
+  def test_upcast_2(self): self._test((Opt(OptOps.UPCAST, 2, 4),))
+  def test_upcast_01(self): self._test((Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4)))
+  def test_upcast_01_mismatch(self): self._test((Opt(OptOps.UPCAST, 0, 2), Opt(OptOps.UPCAST, 1, 4)))
+  def test_upcast_02(self): self._test((Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 2, 4)))
+  def test_upcast_12(self): self._test((Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UPCAST, 2, 4)))
+
+  def test_unroll_0(self): self._test((Opt(OptOps.UNROLL, 0, 4),))
+  def test_unroll_1(self): self._test((Opt(OptOps.UNROLL, 1, 4),))
+  def test_unroll_01(self): self._test((Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UNROLL, 1, 4)))
+
+  def test_upcast_0_unroll_0(self): self._test((Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UNROLL, 0, 4)))
+  def test_upcast_1_unroll_0(self): self._test((Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 4)))
+  def test_upcast_2_unroll_0(self): self._test((Opt(OptOps.UPCAST, 2, 4), Opt(OptOps.UNROLL, 0, 4)))
+
+  def test_upcast_0_unroll_1(self): self._test((Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UNROLL, 1, 4)))
+  def test_upcast_1_unroll_1(self): self._test((Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 1, 4)))
+  def test_upcast_2_unroll_1(self): self._test((Opt(OptOps.UPCAST, 2, 4), Opt(OptOps.UNROLL, 1, 4)))
+
+  def test_upcast_1_unroll_1_small(self): self._test((Opt(OptOps.UPCAST, 1, 2), Opt(OptOps.UNROLL, 1, 2)))
+  def test_upcast_1_unroll_1_rev(self): self._test((Opt(OptOps.UNROLL, 1, 2), Opt(OptOps.UPCAST, 1, 2)))
+
+  def test_upcast_01_unroll_01(self):
+    self._test((Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UNROLL, 1, 4)))
+  def test_upcast_12_unroll_01(self):
+    self._test((Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UPCAST, 2, 4), Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UNROLL, 1, 4)))
 
 class TestRangeifyAssign(unittest.TestCase):
   def test_assign_permuted(self):

test/backend/test_renderer_failures.py

@@ -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,7 +78,7 @@ 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, idx.cast(dtypes.uint32).load() & UOp.const(0xffffff00).cast(dtypes.uint32))
     src = Device[Device.DEFAULT].renderer.render(UOp.sink(store, arg=KernelInfo()).toposort())
@@ -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,))

test/backend/test_schedule.py

@@ -115,8 +115,7 @@ class TestSchedule(unittest.TestCase):
     idx = Tensor([1,2,5,6], dtype=dtypes.int32)
     flat_base[idx] = Tensor([99,99,99,99])
     base.assign(flat_base.reshape(4, 4))
-    # The pending clone is already contiguous, so assign-back needs no separate contiguous buffer.
-    sched = check_schedule(base, 2)
+    sched = check_schedule(base, 4)
     run_linear(*sched)
     expected = list(range(16))
     for i, v in zip([1,2,5,6], [99,99,99,99]): expected[i] = v
@@ -380,30 +379,27 @@ class TestCopyFolding(unittest.TestCase):
     check_schedule(a.clone(), 1, filter_sink=False)
 
   def test_shrink_copy(self):
-    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])
+    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])
 
   def test_expanded_copy(self):
-    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]])
+    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]])
 
   def test_permuted_copy(self):
-    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)
+    a = Tensor.arange(4)
+    b = a.reshape(2, 2).permute(1, 0)
+    b.realize()
     self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])
 
   def test_permute_on_disk(self):

test/backend/test_setitem.py

@@ -1,4 +1,4 @@
-import unittest, operator
+import unittest
 from tinygrad import Tensor, TinyJit, Variable, dtypes, Device
 from tinygrad.helpers import Context
 import numpy as np
@@ -18,16 +18,13 @@ class TestSetitem(unittest.TestCase):
       ((4,4,4,4), (slice(1,3), slice(None), slice(None), slice(0,3)), 4),
       ((6,6), (slice(1,5,2), slice(0,5,3)), 1.0),
       ((6,6), (slice(5,1,-2), slice(5,0,-3)), 1.0),
-      ((6,6), (slice(None), slice(0,6,2)), 1.0),
     )
     for shp, slc, val in cases:
-      for realize in (False, True):
-        t = Tensor.zeros(shp).contiguous()
-        if realize: t.realize()
-        t[slc] = val
-        n = np.zeros(shp)
-        n[slc] = val.numpy() if isinstance(val, Tensor) else val
-        np.testing.assert_allclose(t.numpy(), n)
+      t = Tensor.zeros(shp).contiguous()
+      t[slc] = val
+      n = np.zeros(shp)
+      n[slc] = val.numpy() if isinstance(val, Tensor) else val
+      np.testing.assert_allclose(t.numpy(), n)
 
   def test_padded_setitem(self):
     t = Tensor.arange(10)
@@ -75,11 +72,6 @@ class TestSetitem(unittest.TestCase):
     t.detach()[1, 2] = 5
     self.assertEqual(t[1, 2].item(), 5.0)
 
-  def test_setitem_detach_whole(self):
-    t = Tensor.zeros((3, 3)).realize()
-    t.detach()[:] = 5
-    np.testing.assert_equal(t.numpy(), np.full((3, 3), 5.))
-
   def test_setitem_permute(self):
     # setitem on permuted tensor should modify original
     t = Tensor.zeros((2, 3)).contiguous().realize()
@@ -383,32 +375,6 @@ class TestWithGrad(unittest.TestCase):
     with self.assertRaises(RuntimeError):
       y[0] = 99.0
 
-  def test_unrealized_inplace_keeps_storage(self):
-    x = Tensor([1., 2.]).clone()
-    view = x[:1]
-    x += 3
-    x.realize()
-    self.assertEqual(x.tolist(), [4., 5.])
-    self.assertEqual(view.tolist(), [4.])
-
-  def test_unrealized_view_inplace_keeps_storage(self):
-    x = Tensor([1., 2.]).clone()
-    view = x[:1]
-    view += 3
-    view.realize()
-    self.assertEqual(x.tolist(), [4., 2.])
-    self.assertEqual(view.tolist(), [4.])
-
-  def test_set_augmented_backward(self):
-    for op, expected in ((operator.isub, [-1., -1.]), (operator.imul, [1., 2.]), (operator.itruediv, [-0.01, -0.005])):
-      with self.subTest(op=op.__name__):
-        z = Tensor([1.0, 2.0, 3.0, 4.0])
-        x = Tensor([10.0, 20.0])
-        z[:2] = op(z[:2], x)
-        z.sum().backward()
-        np.testing.assert_allclose(z.grad.numpy(), np.ones(4))
-        np.testing.assert_allclose(x.grad.numpy(), expected)
-
 class TestSetitemLoop(unittest.TestCase):
   def test_arange(self):
     N = 10

test/backend/test_tensor_variable.py

@@ -1,7 +1,7 @@
 import unittest
 import numpy as np
 from tinygrad import Device, Tensor, Variable, TinyJit, dtypes
-from tinygrad.helpers import CHECK_OOB, Context
+from tinygrad.helpers import CHECK_OOB
 
 class TestTensorVariable(unittest.TestCase):
   def test_add_tvar(self):
@@ -35,14 +35,7 @@ class TestTensorVariable(unittest.TestCase):
     vv = Variable("a", 1, 10).bind(2)
     self.assertEqual(Tensor(vv).dtype, dtypes.weakint)
     self.assertEqual((Tensor(vv) + Tensor([1], dtype=dtypes.int8)).dtype, dtypes.int8)  # takes the concrete side, no widening
-    self.assertEqual(Tensor(vv).item(), 2)                                              # a read commits by bounds, like a kernel
-
-  def test_weak_read_widens_by_bounds(self):
-    self.assertEqual(Tensor(2**40).item(), 2**40)
-    self.assertEqual(Tensor(Variable("b", 0, 2**40).bind(2**35+3)).item(), 2**35+3)
-
-  def test_long_variable_emulated_raises(self):
-    with Context(EMULATED_DTYPES="long"), self.assertRaises(RuntimeError): Tensor(Variable("c", 0, 2**40).bind(2**35+3)).item()
+    self.assertEqual(Tensor(vv).item(), 2)                                              # a read commits at default_int
 
   def test_variable_tensor_dtype_arg(self):
     vv = Variable("a", 1, 10).bind(2)
@@ -57,8 +50,6 @@ class TestTensorVariable(unittest.TestCase):
     # bound variables in an expression are fine
     self.assertEqual(Tensor(Variable("u", 1, 10).bind(2) + 1).item(), 3)
 
-  def test_negative_variable_on_device(self): self.assertEqual(Tensor(Variable("n", -10, 10).bind(-3)).clone().item(), -3)
-
   def test_shrink_beyond_buffer_variable(self):
     # TODO: shrink by a variable whose vmax exceeds the dim should fail at build, today only CHECK_OOB=1 rejects it
     t = Tensor.ones(3).contiguous()[:Variable("a", 1, 10).bind(5)]

test/backend/test_uops.py

@@ -5,23 +5,23 @@ from tinygrad.tensor import Tensor, _to_np_dtype
 from tinygrad.helpers import Context, ceildiv
 from tinygrad.dtype import dtypes, DType, AddrSpace, ConstFloat  # noqa: F401
 from tinygrad.device import Buffer, Device
-from tinygrad.uop.ops import Ops, UOp, KernelInfo, AxisType
+from tinygrad.uop.ops import Ops, UOp, KernelInfo, AxisType, buffers
 from tinygrad.renderer.cstyle import CStyleLanguage
 from tinygrad.engine.realize import run_linear
 from tinygrad.codegen import to_program
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.renderer.ptx import PTXRenderer
-from tinygrad.runtime.ops_python import PythonRenderer
 from test.helpers import to_uops_list
 
 def run_uops(uops_list:list[UOp], bufs:list[Buffer]):
-  buf_uops = [UOp.from_buffer(b) for b in bufs]
+  buf_uops = [UOp.new_buffer(b.device, b.size, b.dtype) for b in bufs]
+  for u,b in zip(buf_uops, bufs): buffers[u] = b
   run_linear(UOp(Ops.LINEAR, src=(UOp.sink(*uops_list, arg=KernelInfo()).call(*buf_uops),)))
 
 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, 1))
-  else: uops.append(UOp(op, tuple(src), arg))
+  elif op is Ops.PARAM: uops.append(UOp.param(arg, dtype, shape=(1,)))
+  else: uops.append(UOp(op, dtype, tuple(src), arg))
   return uops[-1]
 
 def _test_single_value(vals, op, dts):
@@ -57,12 +57,12 @@ def _test_uops_result(output_dtype, uops, res):
   run_uops([out], [buf])
   return np.frombuffer(buf.as_memoryview(), _to_np_dtype(output_dtype))[0]
 
-@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, (CStyleLanguage, PythonRenderer)) and
-                     dtypes.uint64 in Device[Device.DEFAULT].renderer.supported_dtypes(), "requires buffer bitcast and 64-bit ints")
+@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, CStyleLanguage) and
+                     dtypes.uint64 in Device[Device.DEFAULT].renderer.supported_dtypes(), "requires C-style pointer bitcast and 64-bit ints")
 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,22 +80,12 @@ 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])
     self.assertEqual(np.frombuffer(buf.as_memoryview(), dtype=np.uint64, count=2, offset=4).tolist(), [val ^ 0xff, val])
 
-  def test_vector_load_store(self):
-    for src_dt, dst_dt in [(dtypes.uint8, dtypes.uint32), (dtypes.uint32, dtypes.uint8)]:
-      with self.subTest(src=src_dt, dst=dst_dt):
-        src, dst = [UOp.param(i, dt, 16 // dt.itemsize) for i, dt in enumerate((src_dt, dst_dt))]
-        src, dst = [b.bitcast(dtypes.uint32).index(UOp.stack(*[UOp.const(i) for i in range(4)])) for b in (src, dst)]
-        bufs = [Buffer(Device.DEFAULT, 16 // dt.itemsize, dt, initial_value=bytes(range(16)) if i == 0 else bytes(16))
-                for i, dt in enumerate((src_dt, dst_dt))]
-        run_uops([dst.store(src.load())], bufs)
-        self.assertEqual(bytes(bufs[1].as_memoryview()), bytes(range(16)))
-
 class TestUOps(unittest.TestCase):
   def _equal(self, v1, v2):
     assert isinstance(v2, (float, int, bool))
@@ -259,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, 3)
-    out = UOp.param(1, dtypes.int32, 2)
+    g1 = UOp.param(0, dtypes.int32, shape=(3,))
+    out = UOp.param(1, dtypes.int32, shape=(2,))
     c1 = UOp.const(2)
     c2 = UOp.const(3)
     l1 = g1.index(c1)
@@ -281,14 +271,14 @@ class TestAssembly(unittest.TestCase):
     b = Tensor.empty(1024)
     c = (a*b).sum()
     ast = c.schedule_linear().src[-1].src[0]
-    opts_to_apply = [Opt(OptOps.SPLIT, 0, (4, AxisType.UNROLL))]
+    opts_to_apply = [Opt(OptOps.UNROLL, 0, 4)]
     ast = ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply)))
     program = to_program(ast, Device[Device.DEFAULT].renderer)
     uops = tuple(program.src[1].src)
     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, 2)
+    g1 = UOp.param(0, dtypes.int32, shape=(2,))
     c1 = UOp.const(0)
     c2 = UOp.const(1)
     expr = g1.index(c1) * UOp.const(4096) + g1.index(c2)
@@ -297,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, 8)
+    g = UOp.param(0, dtypes.uint32, shape=(8,))
     c = UOp.const(7)
     comp = g.index(c).ne(c).ne(True)
     uops = to_uops_list([comp], ren=Device[Device.DEFAULT].renderer)

test/backend/test_wait_loop.py

@@ -1,12 +1,13 @@
-import unittest, threading, functools
-from tinygrad import Tensor, UOp, Context
+import unittest, threading
+from tinygrad import Tensor, UOp
 from tinygrad.device import Device, Buffer, BufferSpec
 from tinygrad.dtype import AddrSpace, dtypes
 from tinygrad.engine.realize import run_linear
 from tinygrad.uop.ops import Ops, KernelInfo
-from tinygrad.renderer.isa.x86 import X86Renderer
 
-def wait_loop_kernel(C:UOp, N=10) -> UOp:
+def wait_loop_kernel(C:UOp) -> UOp:
+  N = 10
+
   # a RANGE with no src is a bound-less loop header: a jump target with no induction variable.
   # the compare and conditional backedge are expanded by the renderers from the loop RANGE/END
   l = UOp.loop(0)
@@ -41,21 +42,8 @@ def nested_loop_kernel(C:UOp) -> UOp:
 
   return C[0].store(i[0].load()).sink(arg=KernelInfo(name="nested_loop", opts_to_apply=()))
 
-def pressure_loop_kernel(C:UOp, n=13) -> UOp:
-  vs = [C[j+1].load() for j in range(n)]
-  l = UOp.loop(0)
-
-  i = UOp.placeholder((1,), dtypes.int, 0, addrspace=AddrSpace.REG)
-  i = i.after(i[0].store(0))
-
-  inc = i.after(l)[0].load() + 1
-  st = i[0].store(inc)
-  i = i.after(st.end(l, inc < sum(v & inc for v in vs)))
-
-  return C[0].store(i[0].load()).sink(arg=KernelInfo(name="pressure_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)
@@ -112,25 +100,6 @@ class TestWaitLoop(unittest.TestCase):
     c.realize()
     self.assertEqual(c.item(), 25)
 
-  # TODO: x86's lower_loop builds an Ops.IF node after regalloc, which fails spec_full
-  @(unittest.expectedFailure if isinstance(Device[Device.DEFAULT].renderer, X86Renderer) else lambda f: f)
-  def test_wait_loop_spec(self):
-    c = Tensor.custom_kernel(Tensor.empty(1, dtype=dtypes.int), fxn=functools.partial(wait_loop_kernel, N=7))[0]
-    with Context(SPEC=2): c.realize()
-    self.assertEqual(c.item(), 7)
-
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "TODO: do-while loop under register pressure segfaults on x86")
-  def test_loop_carried_registers(self):
-    # more loads live across the backedge than any register file (x86 15 gprs, arm64 31, sass 255, rdna3 256 vgprs)
-    c = Tensor.custom_kernel(Tensor.ones(301, dtype=dtypes.int), fxn=functools.partial(pressure_loop_kernel, n=300))[0]
-    self.assertEqual(c[0].item(), 2)
-
-  def test_register_pressure_loop(self):
-    c = Tensor.zeros(16, dtype=dtypes.int).contiguous()
-    c = Tensor.custom_kernel(c, fxn=pressure_loop_kernel)[0]
-    c.realize()
-    self.assertEqual(c[0].item(), 1)
-
   def test_loop_in_loop(self):
     c = Tensor.empty(1, dtype=dtypes.int)
     c = Tensor.custom_kernel(c, fxn=loop_in_loop_kernel)[0]

test/device/test_hcq.py

@@ -9,7 +9,6 @@ from tinygrad.runtime.support.system import PCIIfaceBase
 from tinygrad.engine.realize import get_runtime
 from tinygrad.codegen import to_program
 from tinygrad.codegen.opt import Opt, OptOps
-from tinygrad.uop.ops import AxisType
 from tinygrad import Variable
 
 MOCKGPU = DEV.interface.startswith("MOCK")
@@ -168,7 +167,7 @@ class TestHCQ(unittest.TestCase):
     b = a + 1
     si = b.schedule_linear().src[-1]
 
-    prg = to_program(replace_opts(si.src[0], [Opt(op=OptOps.SPLIT, axis=0, arg=(3, AxisType.LOCAL)) for _ in range(3)]), TestHCQ.d0.renderer)
+    prg = to_program(replace_opts(si.src[0], [Opt(op=OptOps.LOCAL, axis=0, arg=3) for _ in range(3)]), TestHCQ.d0.renderer)
     runtime = get_runtime(Device.DEFAULT, prg)
 
     zb = Buffer(Device.DEFAULT, 3 * 3 * 3, dtypes.int, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()

test/device/test_hcq2.py

@@ -1,231 +1,33 @@
-import unittest, contextlib, ctypes, gc, numpy as np
+import unittest, numpy as np
 from unittest.mock import patch
-from tinygrad import Device, Tensor, TinyJit, Variable, dtypes, GlobalCounters
+from tinygrad import Device, Tensor
 from tinygrad.device import Buffer
-from tinygrad.dtype import AddrSpace
-from tinygrad.helpers import Context, dedup, partition
-from tinygrad.uop.ops import Ops, UOp, KernelInfo
-from tinygrad.engine.realize import compile_linear, link_linear, lower_and_compile, run_linear
-from tinygrad.renderer.cstyle import CStyleLanguage
-from tinygrad.runtime.autogen import libc
-from tinygrad.runtime.support.c import init_c_struct_t
-import tinygrad.runtime.support.hcq2 as hcq2
-from tinygrad.runtime.support.hcq2 import HCQ_DEVS, HCQ2Compiled, all_devices_in, hcq_compile_cache, link_linear_cache
-from test.helpers import call_is_hcq
+from tinygrad.dtype import dtypes
+from tinygrad.helpers import getenv
+from tinygrad.runtime.support.hcq2 import HCQ_DEVS, all_devices_in
 
-@contextlib.contextmanager
-def rt_views():
-  calls, orig = [], HCQ2Compiled.rt_view
-  def track(dev, *args, **kwargs):
-    calls.append(dev)
-    return orig(dev, *args, **kwargs)
-  with patch.object(HCQ2Compiled, "rt_view", track): yield calls
+@unittest.skipUnless(getenv("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):
+      np.testing.assert_equal(Tensor(np.arange(61, dtype=np.float32)).to(Device.DEFAULT).contiguous().realize().numpy(), np.arange(61))
 
-def chain(x:Tensor, n:int) -> Tensor:
-  for _ in range(n): x = (x + 1).contiguous()
-  return x
-
-@contextlib.contextmanager
-def encoded_batches():
-  batches, orig = [], hcq2.lower_and_compile
-  with patch.object(hcq2, "lower_and_compile", lambda l, *a, **kw: (batches.extend(c for c in l.src if call_is_hcq(c)), orig(l, *a, **kw))[1]):
-    yield batches
-
-def eager_chain(x:Tensor, n:int=64) -> Tensor: # at hcq_compile's use_rt bound: an eager linear this big bakes its inputs and borrows ring slots
-  for _ in range(n): x = (x + 1).contiguous()
-  return x.realize()
-
-def patch_words(batch:UOp) -> list[UOp]:
-  return [w for s in batch.src[0].toposort() if s.op is Ops.STORE and s.src[0].op is Ops.INDEX and s.src[0].src[1].op is Ops.STACK
-          and s.src[1].op is Ops.STACK for w in s.src[1].src]
-
-def rt_params(batch:UOp) -> list[str]:
-  return dedup([u.arg.name for w in patch_words(batch) for u in w.toposort() if u.op is Ops.PARAM and u.arg.addrspace is AddrSpace.GLOBAL])
-
-@unittest.skipUnless(all_devices_in(Device.DEFAULT, HCQ_DEVS - {"CPU"}), "non-CPU hcq2 device required")
-class TestHCQ2Core(unittest.TestCase):
-  @staticmethod
-  def input(value:int=2) -> Tensor: return Tensor.full((4,), value, dtype=dtypes.int32).contiguous().realize()
-
-  def compiled(self, n:int, jit=False):
-    x, inputs = self.input(), []
-    if jit:
-      f = TinyJit(lambda a: chain(a, n).realize())
-      f(x)
-      return f(x), f.captured._linear, [x.uop.base]
-    out = chain(x, n)
-    return out, compile_linear(out.schedule_linear(), input_uops=inputs), inputs
-
-  def test_jit_has_no_rt_buffers(self):
-    dev = Device[Device.DEFAULT]
-    rings = [dev.rt_buffer(True, host) for host in (False, True)]
-    ranges = [(b._buf.va_addr, b._buf.va_addr + b.nbytes) for b in rings]
-    for n in (1, 65):
-      with self.subTest(kernels=n):
-        x, f = self.input(), TinyJit(lambda a: chain(a, n).realize())
-        for _ in range(2): f(x)
-        for u in f.captured.linear.toposort():
-          if u.op is Ops.BUFFER and (buf:=u.buffer).device == dev.device:
-            addr = buf._buf.va_addr
-            self.assertFalse(any(addr < end and start < addr + buf.nbytes for start, end in ranges))
-
-  def test_small_eager_cached(self):
-    _, compiled, inputs = self.compiled(1)
-    linked = link_linear(compiled, input_uops=inputs)
-    self.assertIs(link_linear(compiled, input_uops=inputs), linked)
-
-  def test_large_eager_not_cached(self):
-    _, compiled, inputs = self.compiled(65)
-    linked = link_linear(compiled, input_uops=inputs)
-    self.assertIsNot(link_linear(compiled, input_uops=inputs), linked)
-    self.assertNotIn(compiled, link_linear_cache)
-
-  def test_double_compile(self):
-    for n in (1, 65):
-      for jit in (False, True):
-        with self.subTest(kernels=n, jit=jit):
-          out, compiled, inputs = self.compiled(n, jit=jit)
-          linked = link_linear(compiled, input_uops=inputs, allow_cache=not jit)
-          before = tuple(inputs)
-          with rt_views() as borrowed:
-            for linear in (compiled, linked):
-              self.assertIs(compile_linear(linear, input_uops=None if jit else inputs), linear)
-          self.assertEqual(tuple(inputs), before)
-          self.assertFalse(borrowed)
-          run_linear(linked, input_uops=inputs, jit=True, wait=True)
-          self.assertEqual(out.tolist(), [2 + n] * 4)
-
-  def test_double_link(self):
-    for n in (1, 65):
-      for jit in (False, True):
-        with self.subTest(kernels=n, jit=jit):
-          out, compiled, inputs = self.compiled(n, jit=jit)
-          linked = link_linear(compiled, input_uops=inputs, allow_cache=not jit)
-          with rt_views() as borrowed:
-            again = link_linear(linked, input_uops=inputs, allow_cache=not jit)
-          self.assertIs(again, linked)
-          self.assertFalse(borrowed)
-          run_linear(again, input_uops=inputs, jit=True, wait=True)
-          self.assertEqual(out.tolist(), [2 + n] * 4)
-
-  def test_jit_new_inputs_each_call(self):
-    @TinyJit
-    def f(a, b): return (a * b + a).contiguous().realize()
-    ins = [(Tensor.full((23,), float(i)).contiguous().realize(), Tensor.full((23,), 2.0).contiguous().realize()) for i in range(6)]
-    for a, b in ins[:3]: f(a, b).tolist() # warm the jit and the copyout
-
-    before = len(hcq_compile_cache)
-    self.assertEqual([f(a, b).tolist() for a, b in ins[3:]], [[i * 3.0] * 23 for i in range(3, 6)])
-    self.assertEqual(len(hcq_compile_cache), before)
-
-  def test_jit_symbolic(self):
-    @TinyJit
-    def f(a): return (a + 1).sum().contiguous().realize()
-    a = Tensor.rand(3, 10).contiguous().realize()
-    for i in range(1, 5):
-      vi = Variable("i", 1, 10).bind(i)
-      np.testing.assert_allclose(f(a[:, :vi]).item(), (a[:, :i] + 1).sum().item(), atol=1e-5, rtol=1e-5)
-
-  def test_map_cpu_buffer_preserves_contents(self):
-    src = Buffer("CPU", 16, dtypes.uint8, preallocate=True)
-    data = bytes(range(16))
-    src.as_memoryview(force_zero_copy=True)[:] = data
-    src.get_buf(Device.DEFAULT)
-    self.assertEqual(bytes(src.as_memoryview(force_zero_copy=True)), data)
-
-  def test_staged_copy_roundtrip(self):
-    # a host buffer the device cannot read copies in chunks through a small ring of staging slots: every rotation must land bit-exact
-    stage = Buffer("CPU", size:=1 << 16, dtypes.uint8, preallocate=True)
-    for npdt in (np.uint8, np.float32):
-      with self.subTest(dtype=npdt.__name__):
-        n = (size // 2 // np.dtype(npdt).itemsize) * 9 + 7 # nine rotations of a two slot ring, plus a short tail
-        data = np.arange(n, dtype=np.int64).astype(npdt)
-        with patch.object(hcq2, "STAGING_SIZE", size), patch.object(hcq2, "STAGING_SLOTS", 2), patch.object(hcq2, "_staging", lambda: stage):
-          out = Tensor(data).to(Device.DEFAULT).contiguous().realize()
-          np.testing.assert_equal(out.numpy(), data)
-
-  def test_rt_patches_are_inputs_and_vars_only(self):
-    x = Tensor.rand(17, 33).contiguous().realize()
-    with encoded_batches() as batches:
-      @TinyJit
-      def f(a): return (a.sin() * 3).contiguous().realize()
-      for _ in range(3): f(x)
-      eager_chain(x)
-
-    jit, eager = partition(batches, lambda c: c.arg.aux.table >= 0)
-    self.assertTrue(jit and eager, f"want both kinds of batch, got {len(jit)} jit and {len(eager)} eager")
-    for c in batches:
-      self.assertTrue(all(n.startswith(("inputs_", "timeline_")) for n in rt_params(c)), f"runtime patch reads {rt_params(c)}")
-      self.assertFalse([u for w in patch_words(c) for u in w.toposort() if u.op is Ops.GETADDR], "addresses bake at link time")
-    self.assertTrue(any(n.startswith("inputs_") for c in jit for n in rt_params(c)), "the jit patches its input addresses in")
-    self.assertFalse(any(n.startswith("inputs_") for c in eager for n in rt_params(c)), "eager bakes its input addresses")
-
-  def test_programs_are_not_call_args(self):
-    # a program is a link-time patch a cmdbuf word addresses: it rides inside that word, no arg or param of its own
-    def nargs(n):
-      x = Tensor.ones(16).contiguous().realize()
-      with encoded_batches() as batches:
-        @TinyJit
-        def f(a):
-          for i in range(n): a = (a * (i + 1.5)).contiguous()
-          return a.realize()
-        for _ in range(3): f(x)
-      return max(c.arg.aux.nargs for c in batches)
-    self.assertEqual(nargs(2), nargs(12))
-
-  def test_caches_hold_no_buffers(self):
-    # an eager template caches without its buffers and the jit's linear compiles once uncached: freeing the tensors frees the device memory
-    def step(i):
-      x = Tensor(np.full(1024, i, np.float32)).to(Device.DEFAULT).realize()
-      @TinyJit
-      def f(a): return (a * 2 + 1).contiguous().realize()
-      for _ in range(3): out = f(x)
-      self.assertEqual(out.tolist(), [2.0 * i + 1] * 1024)
-    step(1) # warms the programs, templates and rings
-    gc.collect()
-    used = GlobalCounters.mem_used
-    for i in range(2, 5): step(i)
-    gc.collect()
-    self.assertEqual(GlobalCounters.mem_used, used)
-
-  def test_device_state_survives_as_link_refs(self):
-    # a buffer the commands only address, never a param of the body, is kept by the linked call as a ref of what its getaddr resolved into
-    dev, names = Device[Device.DEFAULT], {"AMD": ("scratch",), "NV": ("timeline",), "QCOM": ("_stack", "dummy")}[Device.DEFAULT.split(":")[0]]
-    @TinyJit
-    def f(a): return (a * 2 + 1).contiguous().realize()
-    x = Tensor.ones(16).contiguous().realize()
-    for _ in range(3): f(x)
-    call = f.captured.linear.src[0]
-    self.assertIs(call.op, Ops.AFTER, "the linked call sits after its refs")
-    refs = [u.buffer for u in call.src[1:] if u.op is Ops.BUFFER]
-    for n in names: self.assertTrue(any(r is getattr(dev, n) for r in refs), f"{n} is not a ref of the call")
-
-@unittest.skipUnless(isinstance(Device["CPU"].renderer, CStyleLanguage), "CALL is rendered in C style only")
-class TestHCQ2FFI(unittest.TestCase):
-  @staticmethod
-  def _run(body:UOp) -> list[Buffer]:
-    call = hcq2.lower_call(UOp.sink(body, arg=KernelInfo("test_ffi")).call(aux=hcq2.HCQInfo(("CPU",))))
-    assert call is not None
-    linear = hcq2.hcq_link(lower_and_compile(UOp(Ops.LINEAR, src=(call,))), allow_cache=False)
-    run_linear(linear, jit=True)
-    return [u.buffer for u in linear.src[0].without_after.src[1:] if u.op is Ops.BUFFER]
-
-  def test_ffi_ccall(self):
-    with Context(HCQ_RUNTIME_DEV="CPU"):
-      out = UOp.placeholder((1,), dtypes.int32, slot=1, device="CPU", volatile=True, tag="ffi_result")
-      bufs = self._run(out.index(0).store(hcq2.ccall(libc.dll.ffs, 0x10)))
-    self.assertEqual(next(b for b in bufs if b.dtype is dtypes.int)._buf.cpu_view().view(fmt='i')[0], 5)
-
-  def test_ffi_cstruct(self):
-    struct_t = init_c_struct_t(16, (("u8", ctypes.c_uint8, 0), ("u16", ctypes.c_uint16, 2),
-                                  ("u32", ctypes.c_uint32, 4), ("u64", ctypes.c_uint64, 8)))
-    UOp.placeholder((1,), dtypes.uint8, device="CPU") # reserve slot zero for device-owned placeholders
-    with Context(HCQ_RUNTIME_DEV="CPU"):
-      s = hcq2.cstruct(struct_t, u8=0x12, u16=UOp.const(0x3456, dtypes.uint16), u32=0x789ABCDE, u64=0xFEDCBA9876543210)
-      bufs = self._run(s.index(0).load())
-    got = struct_t.from_buffer_copy(bytes(next(b for b in bufs if b.nbytes == ctypes.sizeof(struct_t))._buf.cpu_view()))
-    self.assertEqual((got.u8, got.u16, got.u32, got.u64), (0x12, 0x3456, 0x789ABCDE, 0xFEDCBA9876543210))
+  @unittest.skipIf(Device.DEFAULT == "CPU", "staged copies need a non-CPU hcq2 device")
+  def test_staged_copy_slot_reuse(self):
+    # chunks of a staged copy rotate through the staging buffer slots, many rotations must stay bit-exact in both directions
+    import tinygrad.runtime.support.hcq2 as hcq2
+    buf = Buffer("CPU", 1 << 20, dtypes.uint8, preallocate=True)
+    data = np.random.default_rng(42).integers(0, 256, (5 << 20) + 123, dtype=np.uint8)
+    with patch.object(hcq2, "STAGING_SIZE", 1 << 20), patch.object(hcq2, "STAGING_SLOTS", 4), patch.object(hcq2, "_staging", lambda: buf):
+      np.testing.assert_equal(Tensor(data).to(Device.DEFAULT).realize().numpy(), data)
 
+  def test_overlapping_device_tuples(self):
+    # an op on a wide device tuple followed by an op on an overlapping smaller tuple used to MMU-fault the smaller one
+    d4, d2 = tuple(f"{Device.DEFAULT}:{i}" for i in range(4)), tuple(f"{Device.DEFAULT}:{i}" for i in range(2))
+    ref = Tensor.arange(16).contiguous().realize()
+    Tensor(ref.uop.copy_to_device(d4)).realize()
+    out = Tensor.ones(8).shard(d2, axis=0).contiguous().realize()
+    np.testing.assert_equal(out.numpy(), np.ones(8))
 
 if __name__ == "__main__":
   unittest.main()

test/external/external_benchmark_op_conv.py

@@ -9,7 +9,7 @@ from tinygrad.helpers import dedup, getenv
 from tinygrad.device import Buffer
 from tinygrad.dtype import Invalid
 
-# PYTHONPATH="." DEV=QCOM FLOAT16=1 IMAGE=2 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
+# PYTHONPATH="." DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
 
 def vision_conv_143():
   c0 = UOp.param(0, dtypes.half, shape=(16, 1024, 4))
@@ -28,12 +28,13 @@ 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)
 
   opts = None
   # JITBEAM=2
+  # (Opt(op=OptOps.UPCAST, axis=2, arg=4), Opt(op=OptOps.NOLOCALS, axis=None, arg=None), Opt(op=OptOps.UPCAST, axis=2, arg=2), Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.SWAP, axis=1, arg=2))
   return c67.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))
 
 def vision_conv_153():
@@ -53,12 +54,13 @@ 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)
 
   opts = None
   # JITBEAM=2
+  # (Opt(op=OptOps.UPCAST, axis=2, arg=4), Opt(op=OptOps.NOLOCALS, axis=None, arg=None), Opt(op=OptOps.UPCAST, axis=2, arg=2), Opt(op=OptOps.SWAP, axis=1, arg=2))
   return c67.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))
 
 def dm_conv_172():
@@ -71,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
@@ -79,6 +81,7 @@ def dm_conv_172():
 
   opts = None
   # JITBEAM=2
+  # (Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.GROUPTOP, axis=1, arg=32), Opt(op=OptOps.UNROLL, axis=1, arg=4), Opt(op=OptOps.LOCAL, axis=0, arg=8), Opt(op=OptOps.UNROLL, axis=0, arg=4), Opt(op=OptOps.GROUP, axis=1, arg=0))
   return c55.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))
 
 ast = {143: vision_conv_143, 153: vision_conv_153, 172: dm_conv_172}[getenv("NUM", 143)]()

test/external/external_benchmark_pyrender.py

@@ -7,8 +7,8 @@ BENCHMARK_OPS = {Ops.INDEX, Ops.STAGE}
 
 @functools.cache
 def create_uop(a:int) -> UOp:
-  op, src, arg, *rest = trace.uop_fields[a]
-  return UOp(op, tuple(create_uop(s) for s in src), arg, *rest)
+  op, dtype, src, arg, *rest = trace.uop_fields[a]
+  return UOp(op, dtype, tuple(create_uop(s) for s in src), arg, *rest)
 
 if __name__ == "__main__":
   # load rewrite trace

test/external/external_test_gpu_crash.py

@@ -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, dtypes.void)) 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) for i in insts))))
     run_linear(UOp(Ops.LINEAR, src=(prg.call(buf),)), wait=True)
 
   def _assert_gpu_fault(self, func):

test/external/external_test_onnx_ops.py

@@ -54,12 +54,6 @@ 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 = {
@@ -241,15 +235,6 @@ 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)}

test/external/external_test_onnx_runner.py

@@ -13,10 +13,10 @@ def _check_ast_count(desired_count:int, t:Tensor):
   asts = [call for call in linear.src if call.src[0].op is Ops.SINK]
   assert len(asts) == desired_count, f"{len(asts)} != {desired_count}"
 
-def build_onnx(nodes, from_disk:bool=True, opset_imports=None, **kwargs):
+def build_onnx(nodes, from_disk:bool=True, **kwargs):
   """Helper to build and return an OnnxRunner from ONNX nodes."""
   graph = onnx.helper.make_graph(nodes, 'test', kwargs.get('inputs', []), kwargs.get('outputs', []), kwargs.get('initializers', []))
-  model = onnx.helper.make_model(graph) if opset_imports is None else onnx.helper.make_model(graph, opset_imports=opset_imports)
+  model = onnx.helper.make_model(graph)
   if from_disk:
     with tempfile.TemporaryDirectory() as tmpdir:
       tmp_path = pathlib.Path(tmpdir)
@@ -29,23 +29,6 @@ def build_onnx(nodes, from_disk:bool=True, opset_imports=None, **kwargs):
   return runner
 
 class TestOnnxRunner(unittest.TestCase):
-  def test_tinygrad_contiguous(self):
-    runner = build_onnx(
-        nodes=[
-          onnx.helper.make_node('Add', ['inp', 'one'], ['added']),
-          onnx.helper.make_node('Contiguous', ['added'], ['materialized'], domain='org.tinygrad'),
-          onnx.helper.make_node('Mul', ['materialized', 'two'], ['output'])
-        ],
-        inputs=[onnx.helper.make_tensor_value_info('inp', onnx.TensorProto.FLOAT, (4,))],
-        outputs=[onnx.helper.make_tensor_value_info('output', onnx.TensorProto.FLOAT, (4,))],
-        initializers=[
-          onnx.helper.make_tensor('one', onnx.TensorProto.FLOAT, (), [1.0]),
-          onnx.helper.make_tensor('two', onnx.TensorProto.FLOAT, (), [2.0])
-        ],
-        opset_imports=[onnx.helper.make_opsetid('', 13), onnx.helper.make_opsetid('org.tinygrad', 1)],
-        from_disk=False).to('PYTHON')
-    _check_ast_count(2, runner({'inp': Tensor.empty(4, device='PYTHON')})['output'])
-
   def _test_const_fold_unary_op(self, from_disk:bool):
     runner = build_onnx(
         nodes=[
@@ -179,4 +162,4 @@ class TestOnnxMetadata(unittest.TestCase):
     self.assertEqual(parsed["metadata_props"][1]["value"], "dGVzdA==")
 
 if __name__ == '__main__':
-  unittest.main()
+  unittest.main()

test/external/external_test_schedule_scaling.py

@@ -1,5 +1,5 @@
 import unittest, time, itertools
-from tinygrad import Tensor, Context, dtypes
+from tinygrad import Tensor, Context
 
 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), dtypes.void)) for i in range(n*8)))))
+                                     UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=s_nop(i)) 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)

test/external/external_test_usb_asm24.py

@@ -3,14 +3,13 @@ from tinygrad.helpers import Timing, getenv
 from tinygrad import Tensor, Device
 import numpy as np
 
-class USBTestCase(unittest.TestCase):
+class TestDevCopySpeeds(unittest.TestCase):
   @classmethod
   def setUpClass(cls):
     cls.sz = getenv("SIZE", 2000000)
     cls.dev = Device["AMD"]
     if not cls.dev.is_usb(): raise unittest.SkipTest("only test this on USB devices")
 
-class TestDevCopySpeeds(USBTestCase):
   def testCopyCPUtoDefault(self):
     for _ in range(10):
       t = Tensor.ones(self.sz, device="CPU", dtype='uchar').contiguous().realize()
@@ -25,7 +24,6 @@ class TestDevCopySpeeds(USBTestCase):
       with Timing(f"copyout of {t.nbytes()/1e6:.2f} MB:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns * 1e3:.2f} MB/s"):
         t.to('CPU').realize()
 
-class TestUSBIntegrity(USBTestCase):
   def testValidateCopies(self):
     t = Tensor.randn(self.sz, device="CPU", dtype='uchar').contiguous().realize()
     x = t.to(Device.DEFAULT).realize()
@@ -36,67 +34,5 @@ class TestUSBIntegrity(USBTestCase):
     np.testing.assert_equal(t.numpy(), y.numpy())
     del x, y, t
 
-  def testCopyinBoundaries(self):
-    rng, chunk = np.random.default_rng(0), 0x40000 - 4
-    for size in (1, 3, 508, 509, 0x3ffc, 0x3ffd, chunk, chunk+1, 2*chunk+31):
-      with self.subTest(size=size):
-        a = rng.integers(0, 256, size, dtype=np.uint8)
-        np.testing.assert_array_equal(a, Tensor(a, device="AMD").numpy())
-
-  def testCopyinFenceWrap(self):
-    a = np.arange(2*(0x40000-4)+31, dtype=np.uint8)
-    np.testing.assert_array_equal(a[:31], Tensor(a[:31], device="AMD").numpy())
-    self.dev.synchronize()
-    alloc, usb = self.dev.allocator, self.dev.iface.pci_dev.usb
-    clear = usb.read(0xA808, 1)
-    # Model a completed 256-chunk copy instead of the one-chunk warmup. The next clear tag must still change.
-    alloc._usb_seq += 255
-    usb.write(0xA800, bytes([alloc._usb_seq & 0xff]))
-    np.testing.assert_array_equal(a[:31], Tensor(a[:31], device="AMD").numpy())
-    self.assertNotEqual(clear, usb.read(0xA808, 1))
-    for bits in (8, 24):
-      with self.subTest(bits=bits):
-        alloc._usb_seq = ((alloc._usb_seq >> bits)+2)*(1 << bits)-2
-        usb.write(0xA800, bytes([alloc._usb_seq & 0xff]))
-        np.testing.assert_array_equal(a, Tensor(a, device="AMD").numpy())
-
-  def testCopyinRingWrap(self):
-    rng = np.random.default_rng(0)
-    a = rng.integers(0, 256, 1 << 20, dtype=np.uint8)
-    np.testing.assert_array_equal(a, Tensor(a, device="AMD").numpy())
-    ring = self.dev.sdma_queue(0)
-    # A 16 MiB copyin needs more than 4 KiB of SDMA packets, forcing the submission to wrap.
-    target = ring.ring.nbytes - 0x1000
-    padding = target - ring.put_value % ring.ring.nbytes - 16  # four-dword timeline fence
-    self.assertGreaterEqual(padding, 0)
-    q = self.dev.hw_copy_queue_t()
-    q.q(*([0] * (padding // 4)))
-    q.signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
-    self.dev.synchronize()
-    before = ring.put_value // ring.ring.nbytes
-    a = rng.integers(0, 256, 16 << 20, dtype=np.uint8)
-    t = Tensor(a, device="AMD").realize()
-    self.assertGreater(ring.put_value // ring.ring.nbytes, before)
-    np.testing.assert_array_equal(a, t.numpy())
-
-  def testCopyinStaleSentinel(self):
-    a = np.arange(16, dtype=np.uint8)
-    np.testing.assert_array_equal(a, Tensor(a, device="AMD").numpy())
-    chunk = 0x40000 - 4
-    for case in ("copyout", "reuse"):
-      with self.subTest(case=case):
-        if case == "copyout":
-          # A 512 KiB copyin takes three chunks. Copyout then fills both SRAM windows with the next expected tag.
-          tag = 0x51000000 | ((self.dev.allocator._usb_seq + 3) & 0xFFFFFF)
-          a = np.full(0x80000 // 4, tag, dtype=np.uint32)
-          np.testing.assert_array_equal(a, Tensor(a, device="AMD").numpy())
-          a = np.arange(31, dtype=np.uint8)
-        else:
-          # The first full chunk contains the tag expected by the short third chunk in the same window.
-          tag = 0x51000000 | ((self.dev.allocator._usb_seq + 2) & 0xFFFFFF)
-          a = np.arange(2 * chunk + 31, dtype=np.uint8)
-          a[:chunk].view(np.uint32)[:] = tag
-        np.testing.assert_array_equal(a, Tensor(a, device="AMD").numpy())
-
 if __name__ == "__main__":
   unittest.main()

test/external/speed_v_theoretical.py

@@ -1,10 +1,6 @@
 import unittest
-from dataclasses import replace
-from itertools import islice
-from tinygrad import Tensor, Device
-from tinygrad.codegen import to_program
-from tinygrad.engine.realize import time_call
-from tinygrad.helpers import Context, DEBUG
+from tinygrad import Tensor, TinyJit, Device
+from tinygrad.helpers import Context, DEBUG, GlobalCounters
 from tinygrad.nn import Conv2d
 from tinygrad.nn.state import get_parameters
 
@@ -14,13 +10,6 @@ class TestKernelSpeed(unittest.TestCase):
       # TODO: randn is 20% faster than rand for gemv
       return Tensor.randn(shape, dtype="half").realize()
 
-  def _time_kernel(self, out:Tensor, beam:int):
-    linear = out.schedule_linear()
-    self.assertEqual(len(linear.src), 1, "expected a single kernel")
-    call = linear.src[0]
-    prg = to_program(call.src[0].replace(arg=replace(call.src[0].arg, beam=beam)), Device[out.device].renderer)
-    return min(islice(time_call(call.replace(src=(prg, *call.src[1:])), clear_l2=True), 3, 10))
-
   def _compare(self, tm, tflops, gbs, nv_tflops=None, nv_gbs=None, amd_tflops=None, amd_gbs=None):
       if DEBUG >= 1:
         print(f"{tm=:.6f}")
@@ -45,39 +34,62 @@ class TestKernelSpeed(unittest.TestCase):
 
   def _test_matmul(self, M, K=None, N=None, nv_tflops=None, nv_gbs=None, amd_tflops=None, amd_gbs=None):
     # (MxK) @ (KxN)
+    @TinyJit
+    def f(a, b) -> Tensor: return (a @ b).realize()
+
     if N is None: N = M
     if K is None: K = M
-    a = self._get_tensor(M, K)
-    b = self._get_tensor(K, N)
-    tm = self._time_kernel(c:=a @ b, beam=3)
+    tms = []
+    with Context(BEAM=3):
+      for i in range(10):
+        a = self._get_tensor(M, K)
+        b = self._get_tensor(K, N)
+        if i >= 3:
+          GlobalCounters.time_sum_s = 0
+          with Context(DEBUG=max(DEBUG.value, 2)): c = f(a, b)
+          tms.append(GlobalCounters.time_sum_s)
+        else:
+          c = f(a, b)
 
     ops = 2 * M * N * K
     mems = a.dtype.itemsize * M * K + b.dtype.itemsize * K * N + c.dtype.itemsize * M * N
+    tm = min(tms)
     tflops = ops / tm / 1e12
     gbs = mems / tm / 1e9
     self._compare(tm, tflops, gbs, nv_tflops, nv_gbs, amd_tflops, amd_gbs)
 
   def _test_conv_3x3(self, BS, CIN, COUT, H, W, nv_tflops=None, nv_gbs=None, amd_tflops=None, amd_gbs=None):
+    @TinyJit
+    def f(conv, x) -> Tensor: return conv(x).realize()
+    tms = []
     K = 3
     with Context(BEAM=0, DEBUG=0):
       conv = Conv2d(CIN, COUT, K, padding=1)
       Tensor.realize(*get_parameters(conv))
 
-    x = self._get_tensor(BS, CIN, H, W)
-    tm = self._time_kernel(_c:=conv(x), beam=2)
+    with Context(BEAM=2):
+      for i in range(10):
+        x = self._get_tensor(BS, CIN, H, W)
+        if i >= 3:
+          GlobalCounters.time_sum_s = 0
+          with Context(DEBUG=max(DEBUG.value, 2)): _c = f(conv, x)
+          tms.append(GlobalCounters.time_sum_s)
+        else:
+          _c = f(conv, x)
 
     # naive algo
     ops = 2 * BS * CIN * COUT * K * K * H * W
     mems = x.nbytes() + conv.weight.nbytes() + conv.bias.nbytes() + _c.nbytes()
+    tm = min(tms)
     tflops = ops / tm / 1e12
     gbs = mems / tm / 1e9
     self._compare(tm, tflops, gbs, nv_tflops, nv_gbs, amd_tflops, amd_gbs)
 
   # TODO: why are convs so slow?!?
-  def test_conv_3x3_256_32_32_256_256(self): self._test_conv_3x3(256, 32, 32, 256, 256, nv_tflops=27, amd_tflops=13)
+  def test_conv_3x3_256_32_32_256_256(self): self._test_conv_3x3(256, 32, 32, 256, 256, nv_tflops=27, amd_tflops=14)
 
   # theoretical is nv_tflops=165, amd_tflops=123
-  def test_gemm_4096(self): self._test_matmul(4096, nv_tflops=109, amd_tflops=65)
+  def test_gemm_4096(self): self._test_matmul(4096, nv_tflops=110, amd_tflops=65)
   def test_gemm_8192(self): self._test_matmul(8192, nv_tflops=115, amd_tflops=60)
 
   # theoretical is nv_gbs=1008, amd_gbs=960

test/helpers.py

@@ -65,17 +65,13 @@ def assert_kernel_count(expected:int):
   got = GlobalCounters.kernel_count
   if got != expected: raise KernelCountException(expected, got)
 
-def is_hcq2_device() -> bool: # an hcq2 device stages every copy from the host through a pinned buffer: such a copy is two calls, not one
-  from tinygrad.runtime.support.hcq2 import HCQ_DEVS
-  return Device.DEFAULT.split(":")[0] in HCQ_DEVS - {"CPU"}
-
 def call_is_graph(call:UOp) -> bool:
   ast = call.src[0]
   return ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph"
 
-def call_is_hcq(call:UOp) -> bool: # an hcq2 batch: a compiled body whose aux lists the kernels it submits
-  from tinygrad.runtime.support.hcq2 import HCQInfo
-  return isinstance(getattr(call.without_after.arg, "aux", None), HCQInfo)
+def call_is_hcq(call:UOp) -> bool:
+  ast = call.src[0]
+  return ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "hcq"
 
 def jit_cache_count(linear:UOp) -> int:
   n = 0
@@ -90,10 +86,9 @@ def assert_jit_cache_len(fxn, expected_len):
   if linear is None or not linear.src:
     if expected_len != 0: raise KernelCountException(expected_len, 0)
     return
-  if expected_len and any(call_is_hcq(call) for call in linear.src): # HCQ2: kernels batch into submits, the finalizers carry the batch's kernels
-    count = sum(len(call.without_after.arg.aux.kernels) if call_is_hcq(call) else 1 for call in linear.src)
-    if count != expected_len: raise KernelCountException(expected_len, count)
-    return
+  if expected_len and all(call_is_hcq(call) for call in linear.src): # HCQ2: one batch submitter, or fence + reset + merged calls + finalizer
+    from tinygrad.runtime.support.hcq2 import HCQ_RUNTIME_DEV
+    expected_len = 1 if HCQ_RUNTIME_DEV.value == "CPU" else 4
   if call_is_graph(linear.src[0]):
     if len(linear.src) != 1: raise KernelCountException(1, len(linear.src))
     inner = linear.src[0].src[0].src[0]  # LINEAR UOp inside CUSTOM_FUNCTION
@@ -127,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)

test/mockgpu/amd/emu.py

@@ -69,7 +69,7 @@ from tinygrad.runtime.autogen.amd.cdna import ins as irc
 from tinygrad.renderer.amd.dsl import VCC_LO, EXEC_LO, SCC, ttmp, Inst
 from tinygrad.runtime.autogen.amd.common import Fmt, OpType
 from test.amd.helpers import decode_dpp16
-from test.mockgpu.amd.pcode import parse_pcode, _FUNCS, _set_bits, _to_bool, _to_u32, _val_to_bits, _ftz_f32, _bitreverse, _countbits
+from test.mockgpu.amd.pcode import parse_pcode, _FUNCS, _set_bits, _to_bool, _to_u32, _val_to_bits, _ftz_f32
 
 MASK32 = 0xFFFFFFFF
 
@@ -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
@@ -537,20 +537,15 @@ class _Ctx:
           stores.extend([self.wsgpr_dyn(_c(EXEC_LO.offset), lo), self.wsgpr_dyn(_c(EXEC_LO.offset + 1), hi)])
         else: stores.append(self.wsgpr_dyn(_c(EXEC_LO.offset), _to_u32(val)))
       elif dest.startswith('VCC'): stores.extend(self.wmask(_c(VCC_LO.offset), val))
-      elif dest.startswith('PC'):  # S_SETPC/S_SWAPPC jump: write PC directly (caller skips inc_pc)
-        lo, hi = _split64(val.cast(dtypes.uint64))
-        stores.extend([self.wsgpr_dyn(_c(PC_LO_IDX), lo), self.wsgpr_dyn(_c(PC_HI_IDX), hi)])
     return stores
 
   def compile_sop_pcode(self, op, srcs: dict[str, UOp | int], sdst_reg: UOp, sdst_size: int) -> UOp:
     """Compile a scalar instruction with dynamic destination register."""
     pcode = get_pcode(op)
-    srcs.update(self.base_srcs(self.rexec()), VCC=self.rmask(_c(VCC_LO.offset)), PC=self.rpc().cast(dtypes.int64))
+    srcs.update(self.base_srcs(self.rexec()), VCC=self.rmask(_c(VCC_LO.offset)))
     if 'D0' not in srcs: srcs['D0'] = self.rsgpr_dyn(sdst_reg)  # D0 is current dest value for read-modify-write ops
     _, assigns = parse_pcode(pcode, srcs)
-    # PC-writing ops (S_SETPC/S_SWAPPC) jump instead of advancing to the next instruction
-    inc = [] if any(dest.startswith('PC') for dest, _ in assigns) else self.inc_pc()
-    return UOp.sink(*self.scalar_stores(assigns, sdst_reg, sdst_size), *inc)
+    return UOp.sink(*self.scalar_stores(assigns, sdst_reg, sdst_size), *self.inc_pc())
 
   def compile_lane_pcode(self, op, inst) -> UOp:
     """Compile cross-lane ops (READLANE/WRITELANE/PERMLANE) using pcode parser."""
@@ -683,7 +678,7 @@ def _compile_sopp(inst: ir3.SOPP | ir4.SOPP, ctx: _Ctx) -> UOp:
             'VCCZ': vcc.eq(UOp.const(0, vcc.dtype)).cast(dtypes.uint32),
             'EXECZ': exec_val.eq(UOp.const(0, exec_val.dtype)).cast(dtypes.uint32)}
     for dest, val in parse_pcode(pcode, srcs)[1]:
-      if dest.startswith('PC'):
+      if dest == 'PC' or dest.startswith('PC.'):
         lo, hi = _split64(val.cast(dtypes.uint64))
         return UOp.sink(ctx.wsgpr_dyn(_c(PC_LO_IDX), lo), ctx.wsgpr_dyn(_c(PC_HI_IDX), hi))
   return UOp.sink(*ctx.inc_pc())
@@ -1328,8 +1323,7 @@ def _compile_wmma(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P, ctx: _Ctx) -> UOp:
   vdst_reg = ctx.inst_field(type(inst).vdst)
   src0_r, src1_r = ctx.inst_field(type(inst).src0) - _c(256), ctx.inst_field(type(inst).src1) - _c(256)
   src2_r = ctx.inst_field(type(inst).src2)
-  is_c_vgpr = src2_r >= _c(256)
-  src2_r = is_c_vgpr.where(src2_r - _c(256), src2_r)  # also keeps the unused VGPR-side index in bounds when src2 is a constant
+  src2_r = (src2_r >= 256).where(src2_r - _c(256), src2_r)
   output_type = op_name.split("WMMA_", 1)[1].split("_", 1)[0]
   is_bf16, is_rdna4 = 'BF16' in op_name, isinstance(inst, ir4.VOP3P)
   cvt = _FUNCS['bf16_to_f32' if is_bf16 else 'f16_to_f32']
@@ -1359,15 +1353,12 @@ def _compile_wmma(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P, ctx: _Ctx) -> UOp:
     return n + lane_bit * 16, vgpr
 
   # Accumulator C. RDNA4 f16/bf16 packs two f32 accumulator VGPRs into one f16 VGPR; RDNA3 uses the lo half of each.
-  # src2 may be a VGPR or an inline/scalar constant (128 = int 0, the usual ", 0" C form); the runner must handle both dynamically
-  out_dt = dtypes.float32 if output_type == "F32" else dtypes.int32
-  cbits = ctx.rsrc_dyn(src2_r, None, 32)
-  cval_const = cvt(cbits & UOp.const(0xFFFF, dtypes.uint32)) if output_type in ("F16", "BF16") else cbits.bitcast(out_dt)
   if output_type in ("F16", "BF16"):
-    mat_c = [is_c_vgpr.where(gval(src2_r, *((lane, vgpr // 2, vgpr % 2) if is_rdna4 else (lane, vgpr, 0))), cval_const)
+    mat_c = [gval(src2_r, *((lane, vgpr // 2, vgpr % 2) if is_rdna4 else (lane, vgpr, 0)))
              for m in range(16) for n in range(16) for lane, vgpr in [d_map(m, n)]]
   else:
-    mat_c = [is_c_vgpr.where(ctx.rvgpr_dyn(src2_r + _c(vgpr), UOp.const(lane, dtypes.int)).bitcast(out_dt), cval_const)
+    out_dt = dtypes.float32 if output_type == "F32" else dtypes.int32
+    mat_c = [ctx.rvgpr_dyn(src2_r + _c(vgpr), UOp.const(lane, dtypes.int)).bitcast(out_dt)
              for m in range(16) for n in range(16) for lane, vgpr in [d_map(m, n)]]
   mat_d = [sum(mat_a[r*16+k] * mat_b[c*16+k] for k in range(16)) + mat_c[r*16+c] for r in range(16) for c in range(16)]
 
@@ -1566,19 +1557,9 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA
   has_data1 = is_lds and hasattr(inst, 'data1') and inst.data1 is not None
   data1_reg = ctx.inst_field(type(inst).data1) if is_lds else _c(0)  # type: ignore[union-attr]
 
-  if is_lds and op_name == 'DS_SWIZZLE_B32':
-    # The manual's reverse_bits operates on five-bit lane indices; thread indices wrap within the wave.
-    funcs = {'reverse_bits': lambda x: _bitreverse(x, 32) >> _c(27), 'count_ones': _countbits,
-             'thread_in': lambda x: ctx.rvgpr_dyn(addr_reg, x & _c(ctx.wave_size - 1)),
-             'thread_valid': lambda x: _lane_active(exec_mask, x & _c(ctx.wave_size - 1))}
-    result, _ = parse_pcode(pcode, {'offset0': offset0.cast(dtypes.uint8), 'offset1': offset1.cast(dtypes.uint8)}, funcs)
-    values = [result[f'thread_out@{i}'] for i in range(ctx.wave_size)]
-    # Snapshot every source before writing: destination and source registers may be identical.
-    reads = UOp(Ops.STACK, src=tuple(values))
-    return UOp.sink(*(ctx.wvgpr_dyn(vdst_reg, _c(i), val, exec_mask, after=reads) for i, val in enumerate(values)), *ctx.inc_pc())
-
   # DS_PERMUTE/DS_BPERMUTE: cross-lane VGPR access via pcode
   if is_lds and 'PERMUTE' in op_name:
+    pcode = get_pcode(inst.op)
     srcs = {'ADDR': addr_reg, 'DATA0': vdata_reg, 'VDST': vdst_reg, 'OFFSET': offset,
             'EXEC': exec_mask.cast(dtypes.uint64), '_vgpr': ctx.vgpr, '_wave_size': ctx.wave_size}
     _, assigns = parse_pcode(pcode, srcs)
@@ -1886,8 +1867,8 @@ class WaveState:
       ctypes.memset(self.accvgpr_buf._buf.va_addr, 0, vgpr_size * 4)
     else:
       self.accvgpr_buf = self.vgpr_buf
-    self._vgpr_mv = self.vgpr_buf.as_memoryview(force_zero_copy=True, no_sync=True).cast('I')
-    self._sgpr_mv = self.sgpr_buf.as_memoryview(force_zero_copy=True, no_sync=True).cast('I')
+    self._vgpr_mv = self.vgpr_buf.as_memoryview(force_zero_copy=True).cast('I')
+    self._sgpr_mv = self.sgpr_buf.as_memoryview(force_zero_copy=True).cast('I')
     # Zero memory using ctypes memset (much faster than Python loops)
     ctypes.memset(self.vgpr_buf._buf.va_addr, 0, vgpr_size * 4)
     ctypes.memset(self.sgpr_buf._buf.va_addr, 0, SGPR_COUNT * 4)
@@ -1966,9 +1947,7 @@ def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int,
   # Use Buffer objects with external_ptr=0 for vmem
   vmem_buf = Buffer('CPU', 1 << 40, dtypes.uint32, options=BufferSpec(external_ptr=0)).ensure_allocated()
   lds_buf = Buffer('CPU', max(lds_size // 4, 1), dtypes.uint32).ensure_allocated()
-  # Scratch is per-lane private memory: each wave needs its own region so data spilled before s_barrier survives other waves' execution.
-  n_waves = -(-total_threads // wave_size)
-  scratch_buf = Buffer('CPU', scratch_size * wave_size * n_waves, dtypes.uint8).ensure_allocated() if scratch_size else None
+  scratch_buf = Buffer('CPU', scratch_size * wave_size, dtypes.uint8).ensure_allocated() if scratch_size else None
 
   # Initialize SQTT encoder — emits packets inline as instructions execute (only when profiling)
   if PROFILE:
@@ -1992,10 +1971,9 @@ def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int,
     waves: list[tuple[WaveState, list]] = []
     for wave_start in range(0, total_threads, wave_size):
       st = _init_wave(lib, wave_start, total_threads, lx, ly, lz, args_ptr, rsrc2, scratch_size, arch, gidx, gidy, gidz, user_data, wave_size)
-      scratch_base = scratch_buf._buf.va_addr + (wave_start // wave_size) * scratch_size * wave_size if scratch_buf else 0
       waves.append((st, [ctypes.c_uint64(st.sgpr_buf._buf.va_addr), ctypes.c_uint64(st.vgpr_buf._buf.va_addr),
                          ctypes.c_uint64(vmem_buf._buf.va_addr), ctypes.c_uint64(lds_buf._buf.va_addr),
-                         ctypes.c_uint64(scratch_base if scratch_buf else 0),
+                         ctypes.c_uint64(scratch_buf._buf.va_addr if scratch_buf else 0),
                          ctypes.c_uint64(st.accvgpr_buf._buf.va_addr)]))
     done = [False] * len(waves)
     for _ in range(10_000_000):

test/mockgpu/amd/pcode.py

@@ -630,10 +630,6 @@ class Parser:
         self.eat('DOT')
         dt_name = self.eat('IDENT').val
         return self._handle_mem_load(addr, DTYPES.get(dt_name, dtypes.uint32))
-      if name in self.funcs and self.try_eat('LBRACKET'):
-        index = self.parse()
-        self.eat('RBRACKET')
-        return self.funcs[name](index)
       if name == 'VGPR' and self.at('LBRACKET'):
         self.eat('LBRACKET')
         lane = self.parse()
@@ -908,16 +904,20 @@ class Parser:
         idx2 = (addr + _const(adt, 4)) >> _const(adt, 2)
         val = val.cast(dtypes.uint64) | (mindex(idx2).cast(dtypes.uint64) << _u64(32))
       elif dt in (dtypes.uint8, dtypes.int8): val = (val >> ((addr & _const(adt, 3)).cast(dtypes.uint32) * _u32(8))) & _u32(0xFF)
+      elif dt in (dtypes.uint16, dtypes.int16):
+        val = (val >> (((addr >> _const(adt, 1)) & _const(adt, 1)).cast(dtypes.uint32) * _u32(16))) & _u32(0xFFFF)
       else:
-        # Handle unaligned 16/32-bit loads: combine two consecutive dwords and shift.
-        # The next dword is only read when the value straddles into it, so a load at the end of a buffer stays in bounds.
+        # Handle unaligned 32-bit loads: combine two consecutive dwords and shift.
+        # To avoid OOB at buffer boundaries for aligned loads, clamp idx_hi to idx (safe).
         # Use int64 for the WHERE to avoid 32-bit int overflow in C pointer arithmetic (addr can be >8GB).
         byte_off = (addr & _const(adt, 3)).cast(dtypes.uint32)
+        is_unaligned = byte_off.ne(_u32(0))
         idx_native = (addr >> _const(adt, 2)).cast(dtypes.int64)
         idx_hi_native = ((addr + _const(adt, 4)) >> _const(adt, 2)).cast(dtypes.int64)
-        hi = mindex((byte_off > _u32(4 - dt.itemsize)).where(idx_hi_native, idx_native))
+        safe_idx_hi = is_unaligned.where(idx_hi_native, idx_native)
+        hi = mindex(safe_idx_hi)
         combined = val.cast(dtypes.uint64) | (hi.cast(dtypes.uint64) << UOp.const(32, dtypes.uint64))
-        val = (combined >> (byte_off.cast(dtypes.uint64) * UOp.const(8, dtypes.uint64))).cast(dtypes.uint32)
+        val = is_unaligned.where((combined >> (byte_off.cast(dtypes.uint64) * UOp.const(8, dtypes.uint64))).cast(dtypes.uint32), val)
     return _cast_to(val, dt)
 
   def _coerce_cmp(self, l: UOp, r: UOp) -> tuple[UOp, UOp]:
@@ -1010,24 +1010,20 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
 
     # for loop
     if first == 'for':
-      # C-style loops use an exclusive bound; for/in loops use an inclusive bound.
-      if m := re.fullmatch(r'for\s*\(\s*(\w+)\s*=\s*(\d+);\s*\1\s*<\s*(\d+);\s*\1\s*(\+\+|\+=\s*\d+)\s*\)', line):
-        loop_var, start_val, end_val = m[1], int(m[2]), int(m[3]) - 1
-        step = 1 if m[4] == '++' else int(m[4][2:])
-      else:
-        p = Parser(toks, env, funcs)
-        p.eat_val('for', 'IDENT')
-        loop_var = p.eat('IDENT').val
-        p.eat_val('in', 'IDENT')
-        def parse_bound():
-          if p.at('NUM') and p.peek(1).type == 'QUOTE':
-            p.eat('NUM')
-            p.eat('QUOTE')
-          if p.at('NUM'): return int(p.eat('NUM').val.rstrip('UuLl'))
-          return int(p.parse())
-        start_val = parse_bound()
-        p.eat('COLON')
-        end_val, step = parse_bound(), 1
+      # Parse: for VAR in [SIZE']START : [SIZE']END do
+      p = Parser(toks, env, funcs)
+      p.eat_val('for', 'IDENT')
+      loop_var = p.eat('IDENT').val
+      p.eat_val('in', 'IDENT')
+      def parse_bound():
+        if p.at('NUM') and p.peek(1).type == 'QUOTE':
+          p.eat('NUM')
+          p.eat('QUOTE')
+        if p.at('NUM'): return int(p.eat('NUM').val.rstrip('UuLl'))
+        return int(p.parse())
+      start_val = parse_bound()
+      p.eat('COLON')
+      end_val = parse_bound()
       # Collect body
       i += 1
       body_lines: list[str] = []
@@ -1043,7 +1039,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
       has_break = any('break' in bl.lower() for bl in body_lines)
       found_var = f'_found_{next(_break_var_ids)}' if has_break else None
       if found_var: env[found_var] = block_assigns[found_var] = _const(dtypes.bool, False)
-      for loop_i in range(start_val, end_val + 1, step):
+      for loop_i in range(start_val, end_val + 1):
         subst_lines = [_subst_loop_var(bl, loop_var, loop_i) for bl in body_lines if not (has_break and bl.strip().lower() == 'break')]
         _, iter_assigns, _ = parse_block(subst_lines, 0, {**env, **block_assigns}, funcs, assigns)
         if has_break:
@@ -1232,9 +1228,9 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
       var = toks[0].val
       j, idx_toks = _match_bracket(toks, 1)
       if j < len(toks) and toks[j].type == 'EQUALS':
-        idx_expr = parse_tokens(idx_toks, env, funcs)
         # Static index: var[NUM] = value
-        if isinstance(idx := _single_value(idx_expr), int):
+        if len(idx_toks) == 1 and idx_toks[0].type == 'NUM':
+          idx = int(idx_toks[0].val.rstrip('UuLl'))
           val = parse_tokens(toks[j+1:], env, funcs)
           existing = block_assigns.get(var, env.get(var))
           if existing is not None and isinstance(existing, UOp):
@@ -1246,6 +1242,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
         # Dynamic index: var[expr] = value where var has @-elements
         elems = [(k.split('@')[1], v) for k, v in {**env, **block_assigns}.items() if k.startswith(f'{var}@') and isinstance(v, UOp)]
         if elems:
+          idx_expr = parse_tokens(idx_toks, env, funcs)
           val = parse_tokens(toks[j+1:], env, funcs)
           for elem_idx_str, old_elem in elems:
             elem_idx = int(elem_idx_str)
@@ -1414,29 +1411,16 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
 def parse_expr(expr: str, env: dict[str, VarVal], funcs: dict | None = None) -> UOp:
   return parse_tokens(tokenize(expr.strip().rstrip(';')), env, funcs)
 
-def parse_pcode(pcode: str, srcs: dict[str, UOp | int] | None = None, funcs: dict | None = None) -> tuple[dict, list]:
+def parse_pcode(pcode: str, srcs: dict[str, UOp | int] | None = None) -> tuple[dict, list]:
   env: dict = srcs.copy() if srcs else {}
   assigns: list[tuple[str, UOp]] = []
+  raw_lines = [l.strip().rstrip(';') for l in pcode.split('\n') if l.strip() and not l.strip().startswith('//')]
+  # TODO: pcode.py should tokenize full pcode string instead of line-by-line, then this hack can be removed
   lines: list[str] = []
-  blocks: list[str] = []
-  for raw in pcode.splitlines():
-    line = raw.split('//')[0].strip().rstrip(';')
-    if not line: continue
-    # Both block syntaxes share the same parser; braces supply the implicit end markers.
-    if line.startswith('}') and blocks:
-      end = blocks.pop()
-      line = line[1:].strip()
-      if not line.startswith(('elsif', 'else')): lines.append(end)
-    if m := re.match(r'(if|elsif|else|for)\b.*\{$', line):
-      blocks.append('endfor' if m[1] == 'for' else 'endif')
-      line = line[:-1].rstrip()
-      if m[1] in ('if', 'elsif'): line += ' then'
-    if not line: continue
-    line = re.sub(r'=\s*(\w+):(\w+)$', r'= {\1, \2}', line)
-    if lines and re.search(r'(&&|\|\||[&|+\-*/^])\s*$', lines[-1]): lines[-1] += ' ' + line
-    else: lines.append(line)
-  assert not blocks, "unclosed pcode block"
-  _, final, _ = parse_block(lines, 0, env, {**_FUNCS, **funcs} if funcs else None, assigns=assigns)
+  for l in raw_lines:
+    if lines and re.search(r'(&&|\|\||[&|+\-*/^])\s*$', lines[-1]): lines[-1] = lines[-1] + ' ' + l
+    else: lines.append(l)
+  _, final, _ = parse_block(lines, 0, env, assigns=assigns)
   sliced = set(d.split('[')[0] for d, _ in assigns if '[' in d)
   for var, val in final.items():
     if var in ['D0', 'S0', 'SCC', 'VCC', 'EXEC', 'PC', 'RETURN_DATA', 'VDATA'] and isinstance(val, UOp):

test/mockgpu/nv/nvdriver.py

@@ -53,7 +53,6 @@ class NVDriver(VirtDriver):
                            VirtFile('/dev/nvidia-uvm', functools.partial(NVUVMFileDesc, driver=self))]
 
     self.root_handle = None
-    self.host_ranges: set[int] = set()
 
     self.gpus = {}
     self.next_fd = (1 << 29)
@@ -252,9 +251,7 @@ class NVDriver(VirtDriver):
     elif nr == nv_gpu.UVM_ENABLE_PEER_ACCESS: pass # uvm and shared spaced are setup already, no emulation for now
     elif nr == nv_gpu.UVM_CREATE_EXTERNAL_RANGE:
       st = nv_gpu.UVM_CREATE_EXTERNAL_RANGE_PARAMS.from_address(argp)
-      # Registered host memory already has a CPU mapping; MAP_FIXED would discard its contents.
-      if st.base not in self.host_ranges:
-        libc.mmap(st.base, st.length, mmap.PROT_READ|mmap.PROT_WRITE, libc.MAP_FIXED|mmap.MAP_SHARED|mmap.MAP_ANONYMOUS, -1, 0)
+      libc.mmap(st.base, st.length, mmap.PROT_READ|mmap.PROT_WRITE, libc.MAP_FIXED|mmap.MAP_SHARED|mmap.MAP_ANONYMOUS, -1, 0)
     elif nr == nv_gpu.UVM_MAP_EXTERNAL_ALLOCATION:
       st = nv_gpu.UVM_MAP_EXTERNAL_ALLOCATION_PARAMS.from_address(argp)
       for gpu_attr_id in range(st.gpuAttributesCount):
@@ -268,7 +265,6 @@ class NVDriver(VirtDriver):
     elif nr == nv_gpu.UVM_REGISTER_CHANNEL: pass
     elif nr == nv_gpu.UVM_FREE:
       st = nv_gpu.UVM_FREE_PARAMS.from_address(argp)
-      self.host_ranges.discard(st.base)
       libc.munmap(st.base, st.length)
     else: raise RuntimeError(f"Unknown {nr} to nvidia-uvm")
     return 0
@@ -280,7 +276,6 @@ class NVDriver(VirtDriver):
       st:Any = nv_gpu.nv_ioctl_nvos02_parameters_with_fd.from_address(argp)
       # Track host memory (signal memory) - progress queues when written to
       if st.params.hClass == nv_gpu.NV01_MEMORY_SYSTEM_OS_DESCRIPTOR:
-        self.host_ranges.add(st.params.pMemory)
         self.track_address(st.params.pMemory, st.params.pMemory + st.params.limit + 1,
                            lambda mv,off: None, lambda mv, off: self._gpu_mmio_write(mv, off, None))
     return 0

test/mockgpu/nv/nvgpu.py

@@ -100,11 +100,11 @@ class GPFIFO:
     if qmd.release0_enable:
       rel0 = to_mv(qmd.release0_address_lower + (qmd.release0_address_upper << 32), 0x10).cast('Q')
       rel0[0] = qmd.release0_payload_lower + (qmd.release0_payload_upper << 32)
-      if qmd.release0_structure_size == 0: rel0[1] = int(time.perf_counter() * 1e9) # four words: the timestamp after the payload
+      rel0[1] = int(time.perf_counter() * 1e9)
     if qmd.release1_enable:
       rel1 = to_mv(qmd.release1_address_lower + (qmd.release1_address_upper << 32), 0x10).cast('Q')
       rel1[0] = qmd.release1_payload_lower + (qmd.release1_payload_upper << 32)
-      if qmd.release1_structure_size == 0: rel1[1] = int(time.perf_counter() * 1e9)
+      rel1[1] = int(time.perf_counter() * 1e9)
     if qmd.dependent_qmd0_enable:
       if qmd.dependent_qmd0_action == 1: self.execute_qmd(qmd.dependent_qmd0_pointer << 8)
       else: raise RuntimeError("unsupported dependent qmd action")
@@ -192,10 +192,11 @@ class GPFIFO:
       sz = self._state(nv_gpu.NVC6B5_LINE_LENGTH_IN)
       assert flags == 0x182, f"unsupported flags in _exec_nvc6b5_dma: {flags}"
       ctypes.memmove(dst, src, sz)
-    elif (semaphore_type:=((flags >> 3) & 0b11)) != 0:
-      to_mv(addr:=self._state64(nv_gpu.NVC6B5_SET_SEMAPHORE_A), 4).cast('I')[0] = self._state(nv_gpu.NVC6B5_SET_SEMAPHORE_PAYLOAD)
-      if semaphore_type == nv_gpu.NVC6B5_LAUNCH_DMA_SEMAPHORE_TYPE_RELEASE_FOUR_WORD_SEMAPHORE:
-        to_mv(addr + 8, 8).cast('Q')[0] = int(time.perf_counter() * 1e9)
+    elif ((flags >> 3) & 0b11) != 0:
+      src = to_mv(self._state64(nv_gpu.NVC6B5_SET_SEMAPHORE_A), 0x10).cast('Q')
+      val = self._state(nv_gpu.NVC6B5_SET_SEMAPHORE_PAYLOAD)
+      src[0] = val
+      src[1] = int(time.perf_counter() * 1e9)
     else: raise RuntimeError("unknown nvc6b5_dma flags")
 
   def _exec_pcas2(self):

test/null/test_gc.py

@@ -21,9 +21,8 @@ def tensors_allocated():
   return _allocations_of_type(Tensor)
 
 def bufs_allocated():
-  # count Buffer objects that own storage: a realized (or to-be-realized) BUFFER UOp owns one, views are transient and excluded
   gc.collect()
-  return sum(1 for x in gc.get_objects() if isinstance(x, Buffer) and x._base is None)
+  return _allocations_of_type(Buffer)
 
 class TestGC(unittest.TestCase):
 
@@ -87,33 +86,35 @@ class TestGC(unittest.TestCase):
       print(inspect.getclosurevars(UOp.toposort().fget))
       raise AssertionError(f"never gced {[x for x in gc.get_objects() if isinstance(x, Buffer)]}")
 
-  def test_buffer_ownership(self):
+  def test_buffer_refcount(self):
     init = bufs_allocated()
     a = Tensor.empty(10)
-    # the Buffer object is owned by the BUFFER UOp 1:1, it exists from creation (device memory is still allocated lazily)
-    self.assertEqual(bufs_allocated()-init, 1)
+    self.assertEqual(bufs_allocated()-init, 0)
     a.realize()
     real_buf = a.uop.buffer
-    self.assertIs(a.uop.arg.buffer, real_buf)
+    # after the Tensor UOp is deleted there shouldn't be any references on the Buffer
+    self.assertEqual(real_buf.uop_refcount, 1)
     self.assertEqual(bufs_allocated()-init, 1)
     del a.uop
-    self.assertEqual(bufs_allocated()-init, 1) # the Buffer object is still held here
+    self.assertEqual(real_buf.uop_refcount, 0)
+    self.assertEqual(bufs_allocated()-init, 1) # keep the buffer alive
     del real_buf
     self.assertEqual(bufs_allocated()-init, 0)
 
-  def test_assign_keeps_buffer(self):
+  def test_assign_refcount(self):
     init = bufs_allocated()
     a = Tensor.full((4,), 1.).contiguous()
     a.realize()
     real_buf = a.uop.buffer
+    self.assertEqual(real_buf.uop_refcount, 1)
     a.assign(Tensor.full((4,), 2.))
-    # assign writes in place: the AFTER still references the same Buffer
     self.assertIs(a.uop.src[0].buffer, real_buf)
+    # NOTE: this is still 1, we don't count the ASSIGN
+    self.assertEqual(real_buf.uop_refcount, 1)
     a.realize()
     del a
-    self.assertEqual(bufs_allocated()-init, 1) # the Buffer object is still held here
-    del real_buf
-    self.assertEqual(bufs_allocated()-init, 0)
+    self.assertEqual(real_buf.uop_refcount, 0) # no UOps for this Buffer
+    self.assertEqual(bufs_allocated()-init, 1) # Buffer is alive
 
 if __name__ == '__main__':
   unittest.main()

test/null/test_gpudims.py

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

test/null/test_linearizer_failures.py

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

test/null/test_linearizer_rewrite.py

@@ -2,7 +2,7 @@ import unittest
 from tinygrad import Tensor, Context, Device
 from tinygrad.codegen import to_program
 from tinygrad.codegen.opt import Opt, OptOps
-from tinygrad.uop.ops import KernelInfo, AxisType
+from tinygrad.uop.ops import KernelInfo
 
 class TestLinearizerRewrite(unittest.TestCase):
   def test_reduction(self):
@@ -11,8 +11,8 @@ class TestLinearizerRewrite(unittest.TestCase):
     with Context(SPLIT_REDUCEOP=0):
       si = out.schedule_linear().src[-1]
       opts_to_apply = []
-      opts_to_apply.append(Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST)))
-      opts_to_apply.append(Opt(OptOps.SPLIT, 2, (4, AxisType.UNROLL)))
+      opts_to_apply.append(Opt(OptOps.UPCAST, 0, 4))
+      opts_to_apply.append(Opt(OptOps.UNROLL, 0, 4))
       ast = si.src[0].replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply)))
       prg = to_program(ast, Device["CPU"].renderer)
       print(prg.src[2].arg)
@@ -22,7 +22,7 @@ class TestLinearizerRewrite(unittest.TestCase):
     with Context(SPLIT_REDUCEOP=0):
       si = out.schedule_linear().src[-1]
       opts_to_apply = []
-      opts_to_apply.append(Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST)))
+      opts_to_apply.append(Opt(OptOps.UPCAST, 0, 4))
       ast = si.src[0].replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply)))
       prg = to_program(ast, Device["CPU"].renderer)
       print(prg.src[2].arg)