faster devectorizer with one line

* external benchmark schedule in 5 sec (codex slop)

* prune

* real?

* delete

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

@@ -137,6 +137,7 @@ runs:
       if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true')
       shell: bash
       run: |
+        sudo mkdir -p /var/cache/apt/archives
         sudo chown -R $USER:$USER /var/cache/apt/archives
 
         echo 'Acquire::GzipIndexes "true";' | sudo tee /etc/apt/apt.conf.d/gzip
@@ -214,6 +215,7 @@ runs:
           sudo apt-get -y --allow-unauthenticated --no-install-recommends install ${{ steps.apt-pkgs.outputs.pkgs }}
         fi
 
+        sudo mkdir -p /var/cache/apt/archives
         sudo chown -R $USER:$USER /var/cache/apt/archives/
 
     - name: Add clang to PATH (Linux)

.github/workflows/benchmark.yml

@@ -117,10 +117,10 @@ jobs:
       run: python3 test/external/process_replay/reset.py
     - name: Run llama3.2
       run: BENCHMARK_LOG=llama32_3b-f16 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m llama3.2:3b-f16 --benchmark --warmup
-    - name: Run qwen3.5
-      # qwen3.5:35b-a3b doesn't fit on mac
+    - name: Run qwen3.6
+      # qwen3.6:35b-a3b doesn't fit on mac
       if: ${{ matrix.dev != 'METAL' }}
-      run: BENCHMARK_LOG=qwen35_35b-a3b JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m qwen3.5:35b-a3b --benchmark --warmup
+      run: BENCHMARK_LOG=qwen36_35b-a3b JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m qwen3.6:35b-a3b --benchmark --warmup
     - name: Run olmoe
       # just metal for now
       if: ${{ matrix.dev == 'METAL' }}
@@ -402,6 +402,35 @@ jobs:
       run: PYTHONPATH=. DEV=PCI+NV:NAK python3.11 test/test_tiny.py
 
   testcommalatest:
+    name: comma Benchmark (0.11.2)
+    runs-on: [self-hosted, Linux, comma]
+    timeout-minutes: 12
+    defaults:
+      run:
+        shell: bash -e -o pipefail {0}
+    if: github.repository_owner == 'tinygrad'
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: setup staging db
+      if: github.ref == 'refs/heads/update_benchmark_staging'
+      run: |
+        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
+        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
+    - name: reset process replay
+      run: test/external/process_replay/reset.py
+    - name: openpilot compile3 0.11.2 supercombo
+      run: BENCHMARK_LOG=openpilot_0_11_2_supercombo PYTHONPATH="." ASSERT_MIN_STEP_TIME=26 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/433f85f956837606ad1f1cbee4aa7e2158ad23c768dea914b20436c97232741b
+    - name: openpilot compile3 0.11.2 supercombo (from pickle)
+      run: BENCHMARK_LOG=openpilot_0_11_2_supercombo_run_pickle RUN_PICKLE=1 PYTHONPATH="." ASSERT_MIN_STEP_TIME=26 DEV=QCOM taskset -c 4-7 python3 examples/openpilot/compile3.py
+    - name: IR3 openpilot compile3 0.11.2 supercombo
+      run: BENCHMARK_LOG=ir3_openpilot_0_11_2_supercombo PYTHONPATH="." ASSERT_MIN_STEP_TIME=41 DEV=QCOM:IR3 FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/433f85f956837606ad1f1cbee4aa7e2158ad23c768dea914b20436c97232741b
+    - name: openpilot compile3 0.11.2 dmonitoring
+      run: BENCHMARK_LOG=openpilot_0_11_2_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/3e7b31dfbc0a5234f1baf196513b77fc6af12204b8a8ffe8ee0417e48352f316
+    - name: Run process replay tests
+      uses: ./.github/actions/process-replay
+
+  testcommaold:
     name: comma Benchmark (0.11.0)
     runs-on: [self-hosted, Linux, comma]
     timeout-minutes: 12
@@ -432,35 +461,6 @@ jobs:
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
 
-  testcommaold:
-    name: comma Benchmark (0.10.1)
-    runs-on: [self-hosted, Linux, comma]
-    timeout-minutes: 12
-    defaults:
-      run:
-        shell: bash -e -o pipefail {0}
-    if: github.repository_owner == 'tinygrad'
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: setup staging db
-      if: github.ref == 'refs/heads/update_benchmark_staging'
-      run: |
-        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
-        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
-    - name: reset process replay
-      run: test/external/process_replay/reset.py
-    - name: DEBUG=2 openpilot compile3 0.10.1 driving_vision
-      run: PYTHONPATH="." DEBUG=2 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
-    - name: openpilot compile3 0.10.1 driving_vision
-      run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
-    - name: openpilot compile3 0.10.1 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3.2 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
-    - name: openpilot compile3 0.10.1 dmonitoring
-      run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
-    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
-
   testqualcommdsp:
     name: DSP Benchmark
     runs-on: [self-hosted, Linux, comma4]
@@ -515,12 +515,14 @@ jobs:
       run: |
         echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
         rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
-    - name: openpilot compile3 0.10.1 driving_vision
-      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision PYTHONPATH="." GMMU=0 DEV=USB+AMD:LLVM ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
-    - name: openpilot load_pickle 0.10.1 driving_vision
-      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_load_pickle PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_LOAD_TIME=15 python3 examples/openpilot/load_pickle.py
-    - name: openpilot run_pickle 0.10.1 driving_vision
-      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_run_pickle RUN_PICKLE=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py
+    - name: openpilot compile3 big_driving_supercombo
+      run: BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo PICKLE_OOB=1 PYTHONPATH="." TC_OPT=2 GMMU=0 DEV=USB+AMD:LLVM ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/10926f2c0911821ca0e72439c1c3bf3ec11f0a08789aa14b7ee8f25379b2afa4 openpilot.pkl
+    - name: openpilot load_pickle big_driving_supercombo
+      run: BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo_load_pickle PICKLE_OOB=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_LOAD_TIME=25 python3 examples/openpilot/load_pickle.py openpilot.pkl
+    - 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=64e6 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py TestDevCopySpeeds
 
   driverbenchmarks:
     name: PCI Driver Benchmark (DEV=${{ matrix.dev }})

.github/workflows/szdiff.yml

@@ -14,12 +14,15 @@ jobs:
     outputs:
       branchstat: ${{ steps.brstat.outputs.stat}}
     steps:
-      - name: Check code from PR branch 
+      - name: Check code from PR branch
         uses: actions/checkout@v6
         with:
           repository: ${{ github.event.pull_request.head.repo.full_name }}
           ref: ${{ github.event.pull_request.head.sha }}
           fetch-depth: 0
+          # PR code is only inspected with git rev-list, never executed
+          allow-unsafe-pr-checkout: true
+          persist-credentials: false
       - name: Check whether branch is up-to-date
         id: brstat
         run: |
@@ -51,6 +54,9 @@ jobs:
           repository: ${{ github.event.pull_request.head.repo.full_name }}
           ref: ${{ github.event.pull_request.head.sha }}
           path: pr
+          # PR code is only line-counted by master's sz.py, never executed
+          allow-unsafe-pr-checkout: true
+          persist-credentials: false
         # the base default to tinygrad master and cannot be other fork branch for security purpose
       - name: Checkout code from tinygrad master
         uses: actions/checkout@v6

.github/workflows/test.yml

@@ -167,22 +167,24 @@ jobs:
       uses: ./.github/actions/setup-tinygrad
       with:
         key: unittest-13
+        python-version: '3.11'
         deps: testing_unit
         llvm: 'true'
         amd: 'true'
     - name: Run NULL backend tests
-      run: DEV=NULL python -m pytest -n=auto test/null/ --durations=20
+      run: SPEC=2 DEV=NULL python -m pytest -n=auto test/null/ --durations=20
     - name: Run targeted tests on NULL backend
       run: |
         DEV=NULL python3 -m unittest test.backend.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step
         DEV=NULL VIZ=1 python3 -m pytest -n=auto test/null/test_viz.py
+        DEBUG=7 python -m tinygrad.viz.cli --json | jq empty
     # TODO: too slow
     # - name: Run SDXL on NULL backend
     #   run: DEV=NULL DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
     - name: Run Clip tests for SD MLPerf on NULL backend
       run: DEV=NULL python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
     - name: Run AMD emulated BERT training on NULL backend
-      run: DEV=NULL::gfx1201 NULL_ALLOW_COPYOUT=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+      run: DEV=NULL::gfx1201 NULL_ALLOW_COPYOUT=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
     # TODO: support fake weights
     #- name: Run LLaMA 7B on 4 fake devices
     #  run: DEV=NULL python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing
@@ -199,6 +201,7 @@ jobs:
       uses: ./.github/actions/setup-tinygrad
       with:
         key: unittest-13
+        python-version: '3.11'
         pydeps: "pre-commit"
         deps: testing_unit
         llvm: 'true'
@@ -291,7 +294,7 @@ jobs:
           llvm: 'true'
       - name: Test openpilot model kernel count and gate usage
         run: |
-          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1391 ALLOWED_GATED_READ_IMAGE=58 FLOAT16=1 DEV="CL::IMAGE_PITCH_ALIGNMENT=64" IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1361 ALLOWED_GATED_READ_IMAGE=54 FLOAT16=1 DEV="CL::IMAGE_PITCH_ALIGNMENT=64" IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
           # IMAGE_PITCH_ALIGNMENT=64 matches adreno 630
       - name: Test openpilot CL compile fp32 (test correctness)
         run: |
@@ -350,6 +353,8 @@ jobs:
         run: DEV=NULL NULL_ALLOW_COPYOUT=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
       - name: Test llama 3 training
         run: DEV=NULL NULL_ALLOW_COPYOUT=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=1 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
+      - name: Test gpt-oss training
+        run: DEV=NULL NULL_ALLOW_COPYOUT=1 SAMPLES=32 BS=2 SEQLEN=128 GRADIENT_ACC_STEPS=1 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 MXFP8=1 VOCAB_SIZE=32000 LAYERS=2 EXPERTS=4 MODEL=gptoss PYTHONPATH=. python3 examples/mlperf/model_train.py
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
@@ -478,7 +483,7 @@ jobs:
           wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | sudo tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc
           echo "deb http://apt.llvm.org/$(lsb_release -cs)/ llvm-toolchain-$(lsb_release -cs)-21 main" | sudo tee /etc/apt/sources.list.d/llvm.list
           sudo apt-get update
-          sudo apt-get install llvm-21 llvm-21-tools cloc
+          sudo apt-get install -y llvm-21 llvm-21-tools cloc
       - name: Install rocprof-trace-decoder
         run: sudo PYTHONPATH="." ./extra/sqtt/install_rocprof_decoder.py
       - name: Run AMD renderer tests
@@ -500,6 +505,29 @@ jobs:
       - name: Run LLVM test
         run: DEV=MOCKKFD+AMD:LLVM python test/device/test_amd_llvm.py
 
+  hcq2:
+    name: hcq2
+    runs-on: *linux
+    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 test/unit/test_multitensor.py \
+            TestMultiTensor.test_simple_add TestMultiTensor.test_shard_reduce \
+            TestMultiTensor.test_backward_sum TestMultiTensor.test_matmul_shard_0_0
+      - 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
+
   testmockam:
     name: Linux (am)
     runs-on: *linux
@@ -605,7 +633,7 @@ jobs:
 
   unittestmacos:
     name: MacOS (unit)
-    runs-on: &macos macos-26
+    runs-on: macos-26
     timeout-minutes: 20
     steps:
     - name: Checkout Code
@@ -615,12 +643,8 @@ jobs:
       with:
         key: unittest-macos
         deps: testing_unit
-        amd: 'true'
-        ocelot: 'true'
     - name: Run unit tests
       run: DEV=METAL python -m pytest -n=auto test/unit/ --durations=20
-    - name: Run NULL backend tests
-      run: DEV=NULL python -m pytest -n=auto test/null/ --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)
@@ -631,6 +655,25 @@ jobs:
       run: DEV=METAL python3 -m pytest test/device/test_metal.py
     #- name: Fuzz Test linearizer
     #  run: DEV=METAL DEPTH=4 FUZZ_N=50 FUZZ_MAX_SIZE=1000000 python test/external/fuzz_linearizer.py
+    - name: Run process replay tests
+      uses: ./.github/actions/process-replay
+
+  unittestmacosmock:
+    name: MacOS (unit, mock)
+    runs-on: macos-26
+    timeout-minutes: 20
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: unittest-macos-mock
+        deps: testing_unit
+        amd: 'true'
+        ocelot: 'true'
+    - name: Run NULL backend tests
+      run: SPEC=2 DEV=NULL python -m pytest -n=auto test/null/ --durations=20
     - name: Run pytest (amd)
       env:
         DEV: MOCKKFD+AMD
@@ -654,7 +697,7 @@ jobs:
       matrix:
         group: [1, 2]
     name: MacOS (DEV=METAL) (${{ matrix.group }})
-    runs-on: *macos
+    runs-on: macos-26
     timeout-minutes: 20
     env:
       DEV: METAL
@@ -686,7 +729,7 @@ jobs:
           - 'WEBGPU'
 
     name: MacOS (DEV=${{ matrix.dev }})
-    runs-on: *macos
+    runs-on: macos-26
     timeout-minutes: 20
     steps:
     - name: Checkout Code

AGENTS.md

@@ -3,3 +3,4 @@
 - Run tests with `-n12` for speed (e.g. `python -m pytest test/null/test_dtype.py -x -q -n12`)
 - Run `python -m mypy tinygrad/` to typecheck
 - Run `python -m ruff check .` to lint
+- Read `./tinygrad/viz/README.md` for profiling and debugging rewrite rules

conftest.py

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

docs/abstractions4.py

@@ -88,7 +88,7 @@ def example_3_custom_uop(a:Tensor, correct):
 
     # store all the per lane accumulators to LOCAL
     local_accs = UOp.placeholder((LCLS,), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL)
-    local_accs = local_accs.after(local_accs[lane].store(acc[0]).barrier())
+    local_accs = local_accs.after(local_accs[lane].store(acc[0]))
 
     # accumulate LOCALs into a single per CU accumulator
     late_reduce_loop = UOp.range(LCLS, 3, AxisType.REDUCE)

docs/dtypes.md

@@ -1,7 +1,8 @@
 ::: tinygrad.dtype.DType
 
-::: tinygrad.dtype.dtypes
+::: tinygrad.dtype.DTypes
     options:
+        heading: dtypes
         members: true
         members_order: source
         show_labels: false

examples/beautiful_cifar.py

@@ -9,8 +9,7 @@ from extra.lr_scheduler import OneCycleLR
 GPUS = [f'{Device.DEFAULT}:{i}' for i in range(getenv("GPUS", 1))]
 
 # override tinygrad defaults
-dtypes.default_float = dtypes.half
-Context(FUSE_OPTIM=1).__enter__()
+Context(DEFAULT_FLOAT=dtypes.half, FUSE_OPTIM=1).__enter__()
 
 # from https://github.com/tysam-code/hlb-CIFAR10/blob/main/main.py
 batchsize = getenv("BS", 1024)

examples/gpt2.py

@@ -22,10 +22,6 @@ class Attention:
     self.head_dim = dim // n_heads
 
   def __call__(self, x:Tensor, start_pos:Variable, mask:Optional[Tensor]) -> Tensor:
-    if mask is not None or start_pos.val == 0:
-      # no symbolic shape qkv when consuming prompts
-      start_pos = start_pos.val
-
     if HALF: x = x.half()
     xqkv = self.c_attn(x).reshape(None, None, 3, self.n_heads, self.head_dim)
     xq, xk, xv = [xqkv[:, :, i, :, :] for i in range(3)]
@@ -38,12 +34,8 @@ class Attention:
     # update the cache
     self.cache_kv[:, :, start_pos:start_pos+seqlen, :, :].assign(Tensor.stack(xk, xv)).realize()
 
-    if start_pos > 0:
-      keys = self.cache_kv[0][:, :start_pos+seqlen, :, :]
-      values = self.cache_kv[1][:, :start_pos+seqlen, :, :]
-    else:
-      keys = xk
-      values = xv
+    keys = self.cache_kv[0][:, :start_pos+seqlen, :, :]
+    values = self.cache_kv[1][:, :start_pos+seqlen, :, :]
 
     xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
     return self.c_proj(xq.scaled_dot_product_attention(keys, values, mask).transpose(1, 2).reshape(bsz, seqlen, self.dim))
@@ -86,15 +78,14 @@ class Transformer:
       seqlen = tokens.shape[1]
       tok_emb = self.wte(tokens)
 
-    # not symbolic when consuming the prompt
-    selected_pos = (0, seqlen) if start_pos.val == 0 else (start_pos, start_pos+1)
-    pos_emb = self.wpe(self.allpos.shrink((None, selected_pos)))
+    # start_pos is a bound Variable, so everything below it stays symbolic
+    pos_emb = self.wpe(self.allpos.shrink((None, (start_pos, start_pos+seqlen))))
 
     h = tok_emb + pos_emb
 
     if HALF: h = h.half()
 
-    mask = Tensor.full((1, 1, seqlen, start_pos.val+seqlen), float("-inf"), dtype=h.dtype).triu(start_pos.val+1) if seqlen > 1 else None
+    mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype).triu(start_pos+1) if seqlen > 1 else None
 
     for hi in self.h: h = hi(h, start_pos, mask)
 

examples/mlperf/dataloader.py

@@ -1,11 +1,11 @@
 import os, random, pickle, queue, struct, math, functools, hashlib, time
 from typing import List
 from pathlib import Path
-from multiprocessing import Queue, Process, shared_memory, connection, Lock, cpu_count
+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
+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(getenv("NUM_CPU_THREADS", min(os.cpu_count(), 8)), 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

@@ -1285,7 +1285,7 @@ def train_llama3():
   from examples.mlperf.models.flat_llama import FlatTransformer, apply_grad, FP8_DTYPE, MXFP8
   from examples.llama3 import MODEL_PARAMS
   from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
-  from examples.mlperf.optim import GradAccClipAdamW
+  from examples.mlperf.optim import GradAccClipAdamW, clip_grads
 
   INITMLPERF = getenv("INITMLPERF")
   RUNMLPERF = getenv("RUNMLPERF")
@@ -1462,6 +1462,8 @@ def train_llama3():
 
   @TinyJit
   def minibatch(tokens:Tensor):
+    for nxt in fp8_next_amax: nxt.assign(0)
+    for nxt in fp8_next_grad_amax: nxt.assign(0)
     if is_dp: tokens = tokens.to(None).shard(device, 0)
     if is_mp: tokens = tokens.shard(device)
     if not is_sharding: tokens = tokens.to(None)
@@ -1480,7 +1482,8 @@ def train_llama3():
 
   @TinyJit
   def optim_step():
-    grad_norm = optim.fstep(grads)
+    grad_norm = clip_grads(grads, grad_acc, 1.0)
+    optim.fstep(grads, grad_norm)
     scheduler.step()
 
     for g in grads: g.assign(0)
@@ -1665,7 +1668,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
+  from examples.mlperf.optim import GradAccClipAdamW, GradAccClipAdamWGroup, clip_grads
 
   BENCHMARK = getenv("BENCHMARK")
 
@@ -1708,9 +1711,10 @@ def train_gptoss():
     wandb.init(config=config, **wandb_args, project="MLPerf-gpt-oss")
 
   model_params = GPT_OSS_20B
-  model_params['vocab_size'] = 128256
+  model_params['vocab_size'] = getenv("VOCAB_SIZE", 128256)
   real_vocab_size = model_params['vocab_size']
   if (layers:=getenv("LAYERS")) != 0: model_params['n_layers'] = layers
+  if (experts:=getenv("EXPERTS")) != 0: model_params['n_experts'] = experts
   print(f"model parameters: {model_params}")
 
   model = GPTOSS(**model_params, max_context=SEQLEN)
@@ -1731,7 +1735,12 @@ def train_gptoss():
   is_offload_optim = bool(getenv("OFFLOAD_OPTIM"))
   is_fake_offload = Device.DEFAULT == "NULL"
   optim_device = ("CPU" if not is_fake_offload else "NULL:99") if is_offload_optim else None
-  optim = GradAccClipAdamW(params, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=grad_acc, device=optim_device)
+  params_wd = [p for p in params if p.ndim >= 3]
+  params_no_wd = [p for p in params if p.ndim < 3]
+  optim = GradAccClipAdamWGroup(
+    GradAccClipAdamW(params_wd, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=grad_acc, device=optim_device),
+    GradAccClipAdamW(params_no_wd, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=0.0, grad_acc=grad_acc, device=optim_device),
+  )
 
   for p in optim.params:
     grad_dtype = dtypes.bfloat16 if p.dtype == FP8_DTYPE else p.dtype
@@ -1740,7 +1749,10 @@ def train_gptoss():
 
   from extra.gemm.cdna_asm_gemm import _mx_block_scale
   model_state = get_state_dict(model)
-  fp8_scale_names = {n: f"{n}_scale" for n, t in model_state.items() if t.dtype == FP8_DTYPE}
+  def _scale_key(n):
+    if "." in n and (c:=f"{(b:=n.rsplit('.',1))[0]}_scale.{b[1]}") in model_state: return c
+    return f"{n}_scale"
+  fp8_scale_names = {n: _scale_key(n) for n, t in model_state.items() if t.dtype == FP8_DTYPE}
   fp8_inv_scales = [model_state[sname] for sname in fp8_scale_names.values()]
   for wname, sname in fp8_scale_names.items():
     w, scale = model_state[wname], model_state[sname]
@@ -1753,11 +1765,12 @@ def train_gptoss():
 
   scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)
 
-  # realize everything here
-  if optim.master_params: Tensor.realize(*optim.master_params)
+  if optim.master_params:
+    for m in optim.master_params: m.realize()
   Tensor.realize(*optim.params, *fp8_inv_scales)
 
   @TinyJit
+  @Context(TRAINING=1)
   def minibatch(tokens:Tensor):
     if is_dp: tokens = tokens.to(None).shard(device, 0)
     if not is_sharding: tokens = tokens.to(None)
@@ -1772,7 +1785,8 @@ def train_gptoss():
 
   @TinyJit
   def optim_step():
-    grad_norm = optim.fstep(grads)
+    grad_norm = clip_grads(grads, grad_acc, 1.0)
+    optim.fstep(grads, grad_norm)
     scheduler.step()
 
     for g in grads: g.assign(0)

examples/mlperf/models/flat_llama.py

@@ -25,6 +25,7 @@ FUSED_SILU_W13 = getenv("FUSED_SILU_W13", 0)
 SPLIT_W13 = getenv("SPLIT_W13", 0)
 COLUMNWISE_WEIGHT_SCALE = getenv("COLUMNWISE_WEIGHT_SCALE", 0)
 MXFP8 = getenv("MXFP8", 0)
+MXFP4 = getenv("MXFP4", 0)
 
 FP8_DTYPE = dtypes.fp8e4m3
 FP8_GRAD_DTYPE = dtypes.fp8e5m2
@@ -37,13 +38,18 @@ def quantize_fp8(x:Tensor, amax_state:Tensor|None=None):
   return x_clamped.cast(FP8_DTYPE), scale.float().reciprocal(), new_amax
 
 def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_scale:Tensor|None=None,
-           x_fp8:Tensor|None=None, x_new_amax:Tensor|None=None,
-           grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None, x_prequant_mx:tuple|None=None) -> tuple[Tensor,...]:
+           x_fp8:Tensor|None=None, grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None, x_prequant_mx:tuple|None=None,
+           next_amax_x:Tensor|None=None) -> tuple[Tensor,...]:
   if not fp8:
     if ASM_GEMM:
       from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
       if can_use_asm_gemm(x, w.T): return (asm_gemm(x, w.T),)
     return (x @ w.T,)
+  if MXFP4:
+    assert x is not None, "MXFP4 matmul requires an unquantized input"
+    from extra.gemm.cdna_asm_gemm import asm_gemm, can_use_asm_gemm
+    if can_use_asm_gemm(x, w.T): return (asm_gemm(x, w.T, mxfp4=True),)
+    return (x @ w.T,)
   assert w_inv_scale is not None, "fp8 matmul requires w_inv_scale (weights must be stored in fp8 with per-tensor scale)"
   if MXFP8:
     from extra.gemm.cdna_asm_gemm import asm_gemm, quantize_mxfp8, mx_pack, can_use_asm_gemm, _mx_block_scale
@@ -56,13 +62,14 @@ def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_sca
     else:
       x_phys = (x_q.cast(dtypes.bfloat16) * _mx_block_scale(x_e8)).reshape(*l_shape, x_q.shape[-1])
       out = x_phys @ (w.cast(dtypes.bfloat16) * _mx_block_scale(w_inv_scale)).T
-    return out, (amax_x.detach() if amax_x is not None else None), x_q
+    return out, x_q
   if x_fp8 is None:
-    if FUSED_INPUT_QUANTIZE and amax_x is not None:
+    if FUSED_INPUT_QUANTIZE:
       from extra.llama_kernels.quantize_fp8_delayed import quantize_fp8_delayed
-      x_fp8, _, x_new_amax, _ = quantize_fp8_delayed(x, amax_x, FP8_DTYPE)
+      x_fp8, _ = quantize_fp8_delayed(x, amax_x, next_amax_x, FP8_DTYPE)
     else:
-      x_fp8, _, x_new_amax = quantize_fp8(x, amax_state=amax_x)
+      x_fp8, _, new_amax_x = quantize_fp8(x, amax_state=amax_x)
+      next_amax_x.assign(new_amax_x)
   if ASM_GEMM:
     from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
     if can_use_asm_gemm(x_fp8, w.T):
@@ -73,51 +80,51 @@ def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_sca
       else:
         out = asm_gemm(x_fp8, w.T, x_scale=amax_x, w_scale=w_inv_scale, grad_amax_state=grad_amax_state,
                        next_grad_amax_state=next_grad_amax_state)
-      return out, x_new_amax, x_fp8
-  return (x_fp8.dot(w.T, dtype=dtypes.float) * ((amax_x.float() + 1e-8) / FP8_MAX) * w_inv_scale).cast(dtypes.bfloat16), x_new_amax, x_fp8
+      return out, x_fp8
+  return (x_fp8.dot(w.T, dtype=dtypes.float) * ((amax_x.float() + 1e-8) / FP8_MAX) * w_inv_scale).cast(dtypes.bfloat16), x_fp8
 
 def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor,
-                         grad_amax_state:Tensor, next_grad_amax_state:Tensor):
-  if FUSED_ADD_NORM_MUL_QUANTIZE:
+                         next_amax_x:Tensor, grad_amax_state:Tensor, next_grad_amax_state:Tensor):
+  if FUSED_ADD_NORM_MUL_QUANTIZE and not MXFP4:
     from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_rmsnorm_mul_quantize_fp8
-    x_fp8, new_amax, x_normed, rrms = fused_rmsnorm_mul_quantize_fp8(x, norm, amax_x, eps, FP8_DTYPE)
-    out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x, x_new_amax=new_amax,
+    x_fp8, x_normed, rrms = fused_rmsnorm_mul_quantize_fp8(x, norm, amax_x, eps, FP8_DTYPE, next_amax_x)
+    out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x,
                        grad_amax_state=grad_amax_state, next_grad_amax_state=next_grad_amax_state)
     return out, x_normed, rrms, ret
   x_normed, rrms = rmsnorm(x, eps)
   out, *ret = matmul(x_normed * norm, w, amax_x=amax_x, w_inv_scale=w_inv_scale, grad_amax_state=grad_amax_state,
-                     next_grad_amax_state=next_grad_amax_state)
+                     next_grad_amax_state=next_grad_amax_state, next_amax_x=next_amax_x)
   return out, x_normed, rrms, ret
 
 def add_norm_quantize_matmul(x:Tensor, residual:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor,
-                             grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None):
-  if FUSED_ADD_NORM_MUL_QUANTIZE:
+                             next_amax_x:Tensor, grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None):
+  if FUSED_ADD_NORM_MUL_QUANTIZE and not MXFP4:
     from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_add_rmsnorm_mul_quantize_fp8
-    x_fp8, new_amax, h, x_normed, rrms = fused_add_rmsnorm_mul_quantize_fp8(x, residual, norm, amax_x, eps, FP8_DTYPE)
-    out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x, x_new_amax=new_amax,
+    x_fp8, h, x_normed, rrms = fused_add_rmsnorm_mul_quantize_fp8(x, residual, norm, amax_x, eps, FP8_DTYPE, next_amax_x)
+    out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x,
                        grad_amax_state=grad_amax_state, next_grad_amax_state=next_grad_amax_state)
     return out, h, x_normed, rrms, ret
   h = x + residual
   x_normed, rrms = rmsnorm(h, eps)
   out, *ret = matmul(x_normed * norm, w, amax_x=amax_x, w_inv_scale=w_inv_scale, grad_amax_state=grad_amax_state,
-                     next_grad_amax_state=next_grad_amax_state)
+                     next_grad_amax_state=next_grad_amax_state, next_amax_x=next_amax_x)
   return out, h, x_normed, rrms, ret
 
 def silu_w13_quantize_matmul(x_w13:Tensor, w2:Tensor, s_2:Tensor,
-                             amax_x2:Tensor,
+                             amax_x2:Tensor, next_amax_x2:Tensor,
                              grad_amax_xw13:Tensor, next_grad_amax_xw13:Tensor,
                              grad_amax_xout:Tensor, next_grad_amax_xout:Tensor):
-  if FUSED_SILU_W13:
+  if FUSED_SILU_W13 and not MXFP4:
     from extra.llama_kernels.cast_amax import fused_quantize_fp8_w13
-    x2_fp8, new_amax_x2 = fused_quantize_fp8_w13(x_w13, amax_x2, FP8_DTYPE, grad_amax_state=grad_amax_xw13,
-                                                 next_grad_amax_state=next_grad_amax_xw13)
-    out, *ret = matmul(None, w2, w_inv_scale=s_2, x_fp8=x2_fp8, amax_x=amax_x2, x_new_amax=new_amax_x2,
+    x2_fp8 = fused_quantize_fp8_w13(x_w13, amax_x2, FP8_DTYPE, grad_amax_state=grad_amax_xw13,
+                                                 next_grad_amax_state=next_grad_amax_xw13, amax_out=next_amax_x2)
+    out, *ret = matmul(None, w2, w_inv_scale=s_2, x_fp8=x2_fp8, amax_x=amax_x2,
                        grad_amax_state=grad_amax_xout, next_grad_amax_state=next_grad_amax_xout)
     return out, ret
   hidden = x_w13.shape[-1] // 2
   x_w1, x_w3 = x_w13[..., :hidden], x_w13[..., hidden:]
   out, *ret = matmul(x_w1.silu() * x_w3, w2, amax_x=amax_x2, w_inv_scale=s_2, grad_amax_state=grad_amax_xout,
-                     next_grad_amax_state=next_grad_amax_xout)
+                     next_grad_amax_state=next_grad_amax_xout, next_amax_x=next_amax_x2)
   return out, ret
 
 class FlatTransformer:
@@ -154,7 +161,7 @@ class FlatTransformer:
     self.tok_embeddings = nn.Embedding(vocab_size, dim)
     self.tok_embeddings.weight = Tensor.normal(vocab_size, dim, mean=0.0, std=0.02, dtype=dtypes.bfloat16)
     self.output = Tensor.normal(1, vocab_size, dim, mean=0.0, std=0.02, dtype=dtypes.bfloat16)
-    self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).contiguous().is_param_(False)
+    self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).clone().is_param_(False)
 
     def _amax(): return Tensor.full((), FP8_MAX, dtype=dtypes.float32).contiguous().is_param_(False)
     names = ["xqkv", "xo", "x2"]
@@ -178,6 +185,9 @@ class FlatTransformer:
       from extra.gemm.cdna_asm_gemm import quantize_mxfp8
       w_q, w_e8, _ = quantize_mxfp8(w.reshape(self.n_layers * out_features, in_features))
       return w_q.reshape(self.n_layers, out_features, in_features), w_e8.reshape(self.n_layers, out_features, in_features // 32)
+    if MXFP4:
+      # FP4 is produced dynamically so optimizer updates always start from the current BF16 weight.
+      return w.cast(dtypes.bfloat16), Tensor.ones(self.n_layers)
     amax = (w.abs().max(axis=2) if COLUMNWISE_WEIGHT_SCALE else w.abs().flatten(1).max(1)).detach()
     scale = FP8_MAX / (amax + 1e-8)
     inv_scale = (amax + 1e-8) / FP8_MAX
@@ -186,89 +196,87 @@ class FlatTransformer:
 
   def attention(self, x:Tensor, freqs_cis:Tensor, *, attention_norm:Tensor, wqkv:Tensor, wo:Tensor,
                 amax_xqkv:Tensor, amax_xo:Tensor, s_qkv:Tensor, s_o:Tensor,
+                next_amax_xqkv:Tensor, next_amax_xo:Tensor,
                 grad_amax_xqkv:Tensor, grad_amax_xo:Tensor, next_grad_amax_xqkv:Tensor, next_grad_amax_xo:Tensor):
     bsz, seqlen, _ = x.shape
-    amaxs, saves = [], []
+    saves = []
 
-    xqkv, x_normed, rrms, (new_amax, *s) = norm_quantize_matmul(x, attention_norm, wqkv, s_qkv, self.norm_eps,
+    xqkv, x_normed, rrms, s = norm_quantize_matmul(x, attention_norm, wqkv, s_qkv, self.norm_eps,
                                                                   amax_x=amax_xqkv, grad_amax_state=grad_amax_xqkv,
-                                                                  next_grad_amax_state=next_grad_amax_xqkv)
-    amaxs.append(new_amax)
+                                                                  next_grad_amax_state=next_grad_amax_xqkv, next_amax_x=next_amax_xqkv)
     saves.extend([x_normed, rrms, *s, xqkv])
-    xqkv = xqkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
-    xq = xqkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
-    xk = xqkv[:, :, :, self.n_rep].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
-    xv = xqkv[:, :, :, self.n_rep+1].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
-
-    xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
-    xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16)
     if getenv("HK_FLASH_ATTENTION"):
-      from extra.thunder.amd.fa import flash_attention
+      from extra.thunder.amd.fa import flash_attention, fused_qkv_rope
+      xq, xk, xv = fused_qkv_rope(xqkv, freqs_cis, self.n_heads, self.n_kv_heads, self.head_dim)
       attn, *save = flash_attention(xq, xk, xv, is_causal=True, write_flat=True)
       saves.extend(save)
     else:
+      xqkv = xqkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
+      xq = xqkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
+      xk = xqkv[:, :, :, self.n_rep].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
+      xv = xqkv[:, :, :, self.n_rep+1].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
+      xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
+      xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16)
       xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
       attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True).transpose(1, 2)
     attn = attn.reshape(bsz, seqlen, -1)
 
-    out, new_amax, *s = matmul(attn, wo, amax_x=amax_xo, w_inv_scale=s_o, grad_amax_state=grad_amax_xo,
-                               next_grad_amax_state=next_grad_amax_xo)
-    amaxs.append(new_amax)
+    out, *s = matmul(attn, wo, amax_x=amax_xo, w_inv_scale=s_o, grad_amax_state=grad_amax_xo,
+                               next_grad_amax_state=next_grad_amax_xo, next_amax_x=next_amax_xo)
     saves.extend([*s, out])
-    return out, amaxs, saves
+    return out, saves
 
   def feed_forward(self, x:Tensor, residual:Tensor, **kwargs):
-    amaxs, saves = [], []
+    saves = []
 
     if SPLIT_W13:
       h = x + residual
       x_normed, rrms = rmsnorm(h, self.norm_eps)
       saves.extend([x_normed, rrms])
       inp = x_normed * kwargs["ffn_norm"]
-      x_w1, new_amax, *s = matmul(inp, kwargs["w1"], amax_x=kwargs["amax_x1"], w_inv_scale=kwargs["s_1"],
-                                  grad_amax_state=kwargs["grad_amax_xw1"], next_grad_amax_state=kwargs["next_grad_amax_xw1"])
-      amaxs.append(new_amax)
+      x_w1, *s = matmul(inp, kwargs["w1"], amax_x=kwargs["amax_x1"], w_inv_scale=kwargs["s_1"],
+                                  grad_amax_state=kwargs["grad_amax_xw1"], next_grad_amax_state=kwargs["next_grad_amax_xw1"],
+                                  next_amax_x=kwargs["next_amax_x1"])
       saves.extend([*s, x_w1])
-      x_w3, new_amax, *s = matmul(inp, kwargs["w3"], amax_x=kwargs["amax_x3"], w_inv_scale=kwargs["s_3"],
-                                  grad_amax_state=kwargs["grad_amax_xw3"], next_grad_amax_state=kwargs["next_grad_amax_xw3"])
-      amaxs.append(new_amax)
+      x_w3, *s = matmul(inp, kwargs["w3"], amax_x=kwargs["amax_x3"], w_inv_scale=kwargs["s_3"],
+                                  grad_amax_state=kwargs["grad_amax_xw3"], next_grad_amax_state=kwargs["next_grad_amax_xw3"],
+                                  next_amax_x=kwargs["next_amax_x3"])
       saves.extend([*s, x_w3])
       if FUSED_SILU_W13 and MXFP8:
         from extra.llama_kernels.fused_silu_mul_quantize_mxfp8 import fused_silu_mul_quantize_mxfp8
         aq, ae8, asi = fused_silu_mul_quantize_mxfp8(x_w1.reshape(-1, x_w1.shape[-1]), x_w3.reshape(-1, x_w3.shape[-1]))
-        out, new_amax, *s = matmul(None, kwargs["w2"], x_prequant_mx=(aq, ae8, asi), amax_x=kwargs["amax_x2"],
-                                   w_inv_scale=kwargs["s_2"], grad_amax_state=kwargs["grad_amax_xout"],
-                                   next_grad_amax_state=kwargs["next_grad_amax_xout"])
+        out, *s = matmul(None, kwargs["w2"], x_prequant_mx=(aq, ae8, asi), amax_x=kwargs["amax_x2"],
+                         w_inv_scale=kwargs["s_2"], grad_amax_state=kwargs["grad_amax_xout"],
+                         next_grad_amax_state=kwargs["next_grad_amax_xout"], next_amax_x=kwargs["next_amax_x2"])
         out = out.reshape(*x_w1.shape[:-1], kwargs["w2"].shape[0])
       else:
-        out, new_amax, *s = matmul(x_w1.silu() * x_w3, kwargs["w2"], amax_x=kwargs["amax_x2"], w_inv_scale=kwargs["s_2"],
-                                   grad_amax_state=kwargs["grad_amax_xout"], next_grad_amax_state=kwargs["next_grad_amax_xout"])
-      amaxs.append(new_amax)
+        out, *s = matmul(x_w1.silu() * x_w3, kwargs["w2"], amax_x=kwargs["amax_x2"], w_inv_scale=kwargs["s_2"],
+                         grad_amax_state=kwargs["grad_amax_xout"], next_grad_amax_state=kwargs["next_grad_amax_xout"],
+                         next_amax_x=kwargs["next_amax_x2"])
       saves.extend([*s, out])
     else:
-      x_w13, h, x_normed, rrms, (new_amax, *s) = add_norm_quantize_matmul(x, residual, kwargs["ffn_norm"], kwargs["w13"], kwargs["s_13"],
+      x_w13, h, x_normed, rrms, s = add_norm_quantize_matmul(x, residual, kwargs["ffn_norm"], kwargs["w13"], kwargs["s_13"],
                                                                           self.norm_eps, amax_x=kwargs["amax_x13"],
+                                                                          next_amax_x=kwargs["next_amax_x13"],
                                                                           grad_amax_state=kwargs["grad_amax_xw13"],
                                                                           next_grad_amax_state=kwargs["next_grad_amax_xw13"])
-      amaxs.append(new_amax)
       saves.extend([x_normed, rrms, *s, x_w13])
-      out, (new_amax, *s) = silu_w13_quantize_matmul(x_w13, kwargs["w2"], kwargs["s_2"], amax_x2=kwargs["amax_x2"],
+      out, s = silu_w13_quantize_matmul(x_w13, kwargs["w2"], kwargs["s_2"], amax_x2=kwargs["amax_x2"],
+                                                     next_amax_x2=kwargs["next_amax_x2"],
                                                      grad_amax_xw13=kwargs["grad_amax_xw13"],
                                                      next_grad_amax_xw13=kwargs["next_grad_amax_xw13"],
                                                      grad_amax_xout=kwargs["grad_amax_xout"],
                                                      next_grad_amax_xout=kwargs["next_grad_amax_xout"])
-      amaxs.append(new_amax)
       saves.extend([*s, out])
-    return out, h, amaxs, saves
+    return out, h, saves
 
   @function(precompile=True, precompile_backward=True)
   def run_layer(self, x:Tensor, freqs_cis:Tensor, attn_kwargs:dict, ffn_kwargs:dict, save:bool=True):
-    attn, attn_amaxs, attn_saves = self.attention(x, freqs_cis, **attn_kwargs)
-    ffn, h, ffn_amaxs, ffn_saves = self.feed_forward(x, attn, **ffn_kwargs)
+    attn, attn_saves = self.attention(x, freqs_cis, **attn_kwargs)
+    ffn, h, ffn_saves = self.feed_forward(x, attn, **ffn_kwargs)
     h = h + ffn
-    amaxs = tuple(a.detach() for a in (*attn_amaxs, *ffn_amaxs))
-    if save: return (h, *amaxs, *attn_saves, *ffn_saves)
-    else: return (h, *amaxs)
+    if save: return (h, *attn_saves, *ffn_saves)
+    else: return (h,)
 
   def shard(self, device:tuple[str, ...], mp:bool=False):
     from tinygrad.nn.state import get_parameters
@@ -313,26 +321,27 @@ class FlatTransformer:
 
   def __call__(self, tokens:Tensor, save:bool=True):
     h = self.tok_embeddings(tokens)
-    freqs_cis = self.freqs_cis.cast(h.dtype)[:, :tokens.shape[1], :, :, :]
+    freqs_cis = self.freqs_cis.cast(h.dtype)
+    if not getenv("HK_FLASH_ATTENTION"): freqs_cis = freqs_cis[:, :tokens.shape[1], :, :, :]
     a, na, ga, nga, s = self._fp8_amax, self._fp8_next_amax, self._fp8_grad_amax, self._fp8_next_grad_amax, self._fp8_inv_scale
     for i in range(self.n_layers):
       attn_kwargs = dict(attention_norm=self.attention_norm[i], wqkv=self.wqkv[i], wo=self.wo[i],
                          amax_xqkv=a["xqkv"][i], amax_xo=a["xo"][i], s_qkv=s["wqkv"][i], s_o=s["wo"][i],
+                         next_amax_xqkv=na["xqkv"][i], next_amax_xo=na["xo"][i],
                          grad_amax_xqkv=ga["xqkv"][i], grad_amax_xo=ga["xo"][i],
                          next_grad_amax_xqkv=nga["xqkv"][i], next_grad_amax_xo=nga["xo"][i])
       ffn_kwargs = dict(ffn_norm=self.ffn_norm[i], w2=self.w2[i],
-                        amax_x2=a["x2"][i], s_2=s["w2"][i], grad_amax_xout=ga["xout"][i], next_grad_amax_xout=nga["xout"][i])
+                        amax_x2=a["x2"][i], s_2=s["w2"][i], grad_amax_xout=ga["xout"][i], next_grad_amax_xout=nga["xout"][i],
+                        next_amax_x2=na["x2"][i])
       if SPLIT_W13:
         ffn_kwargs.update(w1=self.w1[i], w3=self.w3[i], amax_x1=a["x1"][i], amax_x3=a["x3"][i],
+                          next_amax_x1=na["x1"][i], next_amax_x3=na["x3"][i],
                           s_1=s["w1"][i], s_3=s["w3"][i], grad_amax_xw1=ga["xw1"][i], grad_amax_xw3=ga["xw3"][i],
                           next_grad_amax_xw1=nga["xw1"][i], next_grad_amax_xw3=nga["xw3"][i])
       else:
         ffn_kwargs.update(w13=self.w13[i], amax_x13=a["x13"][i], s_13=s["w13"][i], grad_amax_xw13=ga["xw13"][i],
-                          next_grad_amax_xw13=nga["xw13"][i])
-      h, *ret = self.run_layer(h, freqs_cis, attn_kwargs, ffn_kwargs, save=save)
-      amax_names = ["xqkv", "xo"] + (["x1", "x3"] if SPLIT_W13 else ["x13"]) + ["x2"]
-      for name, new_val in zip(amax_names, ret[:len(amax_names)]):
-        na[name][i].assign(new_val)
+                          next_grad_amax_xw13=nga["xw13"][i], next_amax_x13=na["x13"][i])
+      h, *_ = self.run_layer(h, freqs_cis, attn_kwargs, ffn_kwargs, save=save)
 
     logits = matmul(self.norm(h), self.output[0], fp8=False)[0]
     return logits
@@ -415,6 +424,9 @@ if __name__ == "__main__":
   @TinyJit
   def fwd_bwd(tokens:Tensor):
     with Timing("python forward: "):
+      for amax_dict in (model._fp8_next_amax, model._fp8_next_grad_amax):
+        for ts in amax_dict.values():
+          for nxt in ts: nxt.assign(0)
       logits = model(tokens[:, :-1], save=llama_size=="8B")
       loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
     with Timing("python backward: "):

examples/mlperf/models/gpt_oss.py

@@ -10,13 +10,17 @@ if __name__ == "__main__":
 from tinygrad import Tensor, nn, function, getenv, dtypes, TinyJit
 from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
 from tinygrad.uop.ops import Ops, UOp
-from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
+from extra.models.llama import apply_rotary_emb
 from extra.llama_kernels.rmsnorm import rmsnorm
-from extra.gemm.cdna_asm_gemm import _mx_block_scale, quantize_mxfp8
+from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8
+from extra.gemm.moe_gemm import grouped_mx_gemm
+from extra.gemm.moe_routing import route, dispatch, combine
 
 FP8_DTYPE = dtypes.fp8e4m3
 FP8_MAX = 448.0
-INIT_STD = 0.008
+INIT_STD = 0.02
+ASM_GEMM = getenv("ASM_GEMM", 0)
+
 
 def _quant_dequant_fwd(x:Tensor) -> Tensor:
   # x (2d bf16) -> bf16 value after an mxfp8 round-trip (1x32 block scaling on the last axis)
@@ -39,16 +43,18 @@ def quant_dequant_mx(x:Tensor) -> Tensor:
   fxn = _quant_dequant_fwd_fxn(x.as_param(0).uop, x.device)
   return Tensor(UOp.maketuple(fxn.uop).call(x.uop, grad_fxn=_quant_dequant_bwd).gettuple(0))
 
+def _mx_scale(e8:Tensor) -> Tensor:
+  return _mx_block_scale(e8) if e8.ndim == 2 else _mx_block_scale_3d(e8)
+
 def _dequant_fwd(w_q:Tensor, w_scale:Tensor) -> Tensor:
-  return w_q.cast(dtypes.bfloat16) * _mx_block_scale(w_scale)
+  return w_q.cast(dtypes.bfloat16) * _mx_scale(w_scale)
 
 @functools.cache
 def _dequant_fwd_fxn(wq_p, ws_p, device):
   return _dequant_fwd(Tensor(wq_p, device=device), Tensor(ws_p, device=device))
 
 def _dequant_bwd(grad:UOp, call:UOp) -> tuple:
-  w_scale = Tensor(call.src[2])
-  return ((Tensor(grad).cast(dtypes.bfloat16) * _mx_block_scale(w_scale).cast(dtypes.bfloat16)).uop, None)
+  return (Tensor(grad).cast(dtypes.bfloat16).uop, None)
 
 def dequant_weight(w_q:Tensor, w_scale:Tensor) -> Tensor:
   fxn = _dequant_fwd_fxn(w_q.as_param(0).uop, w_scale.as_param(1).uop, w_q.device)
@@ -57,16 +63,45 @@ def dequant_weight(w_q:Tensor, w_scale:Tensor) -> Tensor:
 
 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
+    x2, K, N = x.reshape(-1, x.shape[-1]), x.shape[-1], w_q.shape[0]
+    wq, ws = w_q, w_scale
+    if (pad := (-K) % 256):
+      x2 = x2.pad(((0, 0), (0, pad)))
+      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)
+    x_q, x_e8, x_si = quantize_mxfp8(x2)
+    if x_si is not None and can_use_asm_gemm(x_q, wq.T):
+      out = asm_gemm(x_q, wq.T, mx=True, mx_scales=(x_si, x_e8, mx_pack(ws), ws), mx_w_stored=True)
+      return (out[:, :N] if npad else out).reshape(*l_shape, N).cast(dtypes.bfloat16)
   x_phys = quant_dequant_mx(x.reshape(-1, x.shape[-1])).reshape(*l_shape, x.shape[-1])
   w_phys = dequant_weight(w_q, w_scale)
   return (x_phys @ w_phys.T).cast(dtypes.bfloat16)
 
+def _pad_to_mult(t:Tensor, axis:int, mult:int=256) -> Tensor:
+  if (r := (-t.shape[axis]) % mult) == 0: return t
+  pads = [(0, 0)] * t.ndim
+  pads[axis] = (0, r)
+  return t.pad(tuple(pads))
+
+def _pad_cols(t:Tensor) -> Tensor: return _pad_to_mult(t, -1)
+def _pad_rows(t:Tensor) -> Tensor: return _pad_to_mult(t, -2)
+
 def swiglu(x:Tensor, limit:float=7.0, alpha:float=1.702) -> Tensor:
   x_glu, x_linear = x[..., ::2], x[..., 1::2]
   x_glu = x_glu.clamp(max_=limit)
   x_linear = x_linear.clamp(-limit, limit)
   return (x_glu * (alpha * x_glu).sigmoid()) * (x_linear + 1)
 
+def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0) -> Tensor:
+  freqs = 1.0 / (theta ** (Tensor.arange(0, dim, 2, dtype=dtypes.float32)[:(dim // 2)] / dim))
+  freqs = Tensor.arange(end, dtype=dtypes.float32).unsqueeze(dim=1) * freqs.unsqueeze(dim=0)
+  return Tensor.stack(freqs.cos(), freqs.sin(), dim=-1).cast(dtypes.default_float).reshape(1, end, 1, dim//2, 2)
+
 class GPTOSS:
   def __init__(self, dim:int, n_layers:int, n_heads:int, n_kv_heads:int, head_dim:int, n_experts:int, experts_per_tok:int,
                intermediate_size:int, vocab_size:int, norm_eps:float=1e-5, rope_theta:int=150000, sliding_window:int=128,
@@ -92,9 +127,9 @@ class GPTOSS:
     self.ffn_norm = Tensor.ones(n_layers, dim).contiguous()
     self.gate = Tensor.normal(n_layers, n_experts, dim, mean=0.0, std=INIT_STD, dtype=dtypes.bfloat16)
     self.gate_bias = Tensor.zeros(n_layers, n_experts, dtype=dtypes.bfloat16).contiguous()
-    self.w_gate_up, self.w_gate_up_scale = self._quant_weight(n_layers, n_experts, intermediate_size * 2, dim)
+    self.w_gate_up, self.w_gate_up_scale = self._quant_weight(n_layers, n_experts, intermediate_size * 2, dim, moe=True)
     self.w_gate_up_bias = Tensor.zeros(n_layers, n_experts, intermediate_size * 2, dtype=dtypes.bfloat16).contiguous()
-    self.w_down, self.w_down_scale = self._quant_weight(n_layers, n_experts, dim, intermediate_size, std=scaled_std)
+    self.w_down, self.w_down_scale = self._quant_weight(n_layers, n_experts, dim, intermediate_size, std=scaled_std, moe=True)
     self.w_down_bias = Tensor.zeros(n_layers, n_experts, dim, dtype=dtypes.bfloat16).contiguous()
 
     # output
@@ -104,19 +139,40 @@ class GPTOSS:
     self.output = Tensor.normal(vocab_size, dim, mean=0.0, std=INIT_STD, dtype=dtypes.bfloat16)
     self.freqs_cis = precompute_freqs_cis(head_dim, max_context * 2, rope_theta).contiguous().is_param_(False)
 
-  def _quant_weight(self, *shape:int, std:float=INIT_STD):
-    w = Tensor.zeros(*shape) if getenv("ZEROS") else Tensor.normal(*shape, mean=0.0, std=std)
-    w_q, w_e8, _ = quantize_mxfp8(w)
-    return w_q, w_e8.is_param_(False)
+  def _quant_weight(self, *shape:int, std:float=INIT_STD, moe:bool=False):
+    def _one(*s:int):
+      w = Tensor.zeros(*s) if getenv("ZEROS") else Tensor.normal(*s, mean=0.0, std=std)
+      w_q, w_e8, _ = quantize_mxfp8(_pad_cols(_pad_rows(w)) if moe else w)
+      return w_q, w_e8.is_param_(False)
+    if moe:
+      qs = [_one(*shape[1:]) for _ in range(shape[0])]
+      return [q[0] for q in qs], [q[1] for q in qs]
+    return _one(*shape)
 
-  def _attn_mask(self, seqlen:int, sliding:bool, dtype) -> Tensor:
+  def _attn_mask(self, seqlen:int, dtype) -> Tensor:
     i, j = Tensor.arange(seqlen).reshape(seqlen, 1), Tensor.arange(seqlen).reshape(1, seqlen)
-    allowed = j <= i
-    if sliding: allowed = allowed & (i - j < self.sliding_window)
-    return allowed.where(0.0, -1e30).cast(dtype).contiguous()
+    return (j <= i).where(0.0, -1e30).cast(dtype).contiguous()
 
-  def attention(self, x:Tensor, freqs_cis:Tensor, mask:Tensor, *, attention_norm:Tensor, wqkv:Tensor, wqkv_scale:Tensor,
-                wqkv_bias:Tensor, wo:Tensor, wo_scale:Tensor, wo_bias:Tensor, sinks:Tensor):
+  def _sliding_attention(self, xq:Tensor, xk:Tensor, xv:Tensor, sinks:Tensor) -> Tensor:
+    bsz, seqlen, H, hd = xq.shape
+    KV, R, W = self.n_kv_heads, self.n_rep, self.sliding_window
+    assert seqlen % W == 0, f"seqlen {seqlen} must be a multiple of sliding_window {W} for banded attention"
+    nb = seqlen // W
+    q = xq.reshape(bsz, seqlen, KV, R, hd).permute(0, 2, 3, 1, 4).reshape(bsz, KV, R, nb, W, hd).float()
+    k, v = (x.permute(0, 2, 1, 3).reshape(bsz, KV, 1, nb, W, hd).float() for x in (xk, xv))
+    kk, vv = (x.pad((None, None, None, (1, 0), None, None))[:, :, :, :nb].cat(x, dim=-2) for x in (k, v))
+    sc = (q @ kk.transpose(-1, -2)) * self.sm_scale   # (B,KV,R,nb,W,2W)
+    i, j, pv = Tensor.arange(W).reshape(W, 1), Tensor.arange(2 * W).reshape(1, 2 * W), Tensor.arange(nb).reshape(nb, 1, 1) >= 1
+    sc = ((j > i) & (j <= i + W) & (pv | (j >= W))).where(sc, -float("inf"))
+    sink = sinks.reshape(1, KV, R, 1, 1, 1).float()
+    m = sc.max(-1, keepdim=True).maximum(sink)
+    e = (sc - m).exp()
+    p = (e / (e.sum(-1, keepdim=True) + (sink - m).exp())).cast(dtypes.bfloat16)
+    attn = p @ vv.cast(dtypes.bfloat16)
+    return attn.reshape(bsz, KV, R, seqlen, hd).permute(0, 3, 1, 2, 4).reshape(bsz, seqlen, H * hd)
+
+  def attention(self, x:Tensor, freqs_cis:Tensor, mask:Tensor, sliding:bool, *, attention_norm:Tensor, wqkv:Tensor,
+                wqkv_scale:Tensor, wqkv_bias:Tensor, wo:Tensor, wo_scale:Tensor, wo_bias:Tensor, sinks:Tensor):
     bsz, seqlen, _ = x.shape
     x_normed, rrms = rmsnorm(x, self.norm_eps)
     qkv = matmul_mx(x_normed * attention_norm, wqkv, wqkv_scale) + wqkv_bias
@@ -124,16 +180,23 @@ class GPTOSS:
     xq = qkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
     xk, xv = qkv[:, :, :, self.n_rep], qkv[:, :, :, self.n_rep + 1]
     xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
+    xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16)  # (B,N,H,D)/(B,N,KV,D)
 
-    xq = xq.cast(dtypes.bfloat16).reshape(bsz, seqlen, self.n_kv_heads, self.n_rep, self.head_dim).permute(0, 2, 3, 1, 4)
-    xk = xk.cast(dtypes.bfloat16).permute(0, 2, 1, 3).unsqueeze(2)
-    xv = xv.cast(dtypes.bfloat16).permute(0, 2, 1, 3).unsqueeze(2)
-    scores = (xq @ xk.transpose(-2, -1)).float() * self.sm_scale + mask
-    sink = sinks.reshape(1, self.n_kv_heads, self.n_rep, 1, 1).float()
-    m = scores.max(-1, keepdim=True).maximum(sink)
-    e = (scores - m).exp()
-    w = (e / (e.sum(-1, keepdim=True) + (sink - m).exp())).cast(dtypes.bfloat16)
-    attn = (w @ xv).permute(0, 3, 1, 2, 4).reshape(bsz, seqlen, self.n_heads * self.head_dim)
+    if sliding:
+      attn = self._sliding_attention(xq, xk, xv, sinks)
+    elif getenv("HK_FLASH_ATTENTION"):
+      from extra.thunder.amd.fa import flash_attention
+      attn, *_ = flash_attention(xq, xk, xv, is_causal=True, write_flat=True, sinks=sinks)
+      attn = attn.reshape(bsz, seqlen, self.n_heads * self.head_dim)
+    else:
+      xqm = xq.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep, self.head_dim).permute(0, 2, 3, 1, 4)
+      xkm, xvm = xk.permute(0, 2, 1, 3).unsqueeze(2), xv.permute(0, 2, 1, 3).unsqueeze(2)
+      scores = (xqm @ xkm.transpose(-2, -1)).float() * self.sm_scale + mask
+      sink = sinks.reshape(1, self.n_kv_heads, self.n_rep, 1, 1).float()
+      m = scores.max(-1, keepdim=True).maximum(sink)
+      e = (scores - m).exp()
+      w = (e / (e.sum(-1, keepdim=True) + (sink - m).exp())).cast(dtypes.bfloat16)
+      attn = (w @ xvm).permute(0, 3, 1, 2, 4).reshape(bsz, seqlen, self.n_heads * self.head_dim)
 
     out = matmul_mx(attn, wo, wo_scale) + wo_bias
     return out, [x_normed, rrms, attn]
@@ -143,22 +206,37 @@ class GPTOSS:
                    w_down:Tensor, w_down_scale:Tensor, w_down_bias:Tensor):
     x_normed, rrms = rmsnorm(x, self.norm_eps)
     inp = x_normed * ffn_norm
-
     logits = inp.float() @ gate.float().T + gate_bias.float()
-    thresh = logits.topk(self.experts_per_tok)[0][..., -1:]
-    weights = (logits >= thresh).where(logits, -float("inf")).softmax(-1)
+    dim, inter = self.dim, self.intermediate_size
 
-    out = None
-    for e in range(self.n_experts):
-      gate_up = matmul_mx(inp, w_gate_up[e], w_gate_up_scale[e]) + w_gate_up_bias[e]
-      y = (matmul_mx(swiglu(gate_up, self.swiglu_limit), w_down[e], w_down_scale[e]) + w_down_bias[e]).contiguous()
-      contrib = weights[..., e:e+1].cast(y.dtype) * y
-      out = contrib if out is None else out + contrib
+    if getenv("GROUPED_MOE", 0):
+      bsz, seqlen = x.shape[:2]
+      inp, logits = inp.reshape(-1, dim), logits.reshape(-1, self.n_experts)
+      r = route(logits, self.experts_per_tok, self.n_experts)
+      onehot = r.rows_e.one_hot(self.n_experts).float()
+      xg = dispatch(_pad_cols(inp.cast(dtypes.bfloat16)), r)
+      h = grouped_mx_gemm(xg, (w_gate_up, w_gate_up_scale), r.off)[:, :2*inter] + (onehot @ w_gate_up_bias.float()).cast(dtypes.bfloat16)
+      y = swiglu(h, self.swiglu_limit)
+      z = grouped_mx_gemm(_pad_cols(y.cast(dtypes.bfloat16)), (w_down, w_down_scale), r.off)[:, :dim] \
+          + (onehot @ w_down_bias.float()).cast(dtypes.bfloat16)
+      out = combine(z, r, inp.shape[0], self.experts_per_tok).reshape(bsz, seqlen, dim)
+    else:
+      thresh = logits.topk(self.experts_per_tok)[0][..., -1:]
+      weights = (logits >= thresh).where(logits, -float("inf")).softmax(-1)
+
+      out = None
+      for e in range(self.n_experts):
+        gu_q, gu_s = w_gate_up[e][:2*inter, :dim].contiguous(), w_gate_up_scale[e][:2*inter, :dim//32].contiguous()
+        dn_q, dn_s = w_down[e][:dim, :inter].contiguous(), w_down_scale[e][:dim, :inter//32].contiguous()
+        gate_up = matmul_mx(inp, gu_q, gu_s) + w_gate_up_bias[e]
+        y = (matmul_mx(swiglu(gate_up, self.swiglu_limit), dn_q, dn_s) + w_down_bias[e]).contiguous()
+        contrib = weights[..., e:e+1].cast(y.dtype) * y
+        out = contrib if out is None else out + contrib
     return out, [x_normed, rrms]
 
   @function(precompile=True, precompile_backward=True)
-  def run_layer(self, x:Tensor, freqs_cis:Tensor, mask:Tensor, attn_kwargs:dict, ffn_kwargs:dict, save:bool=True):
-    attn, attn_saves = self.attention(x, freqs_cis, mask, **attn_kwargs)
+  def run_layer(self, x:Tensor, freqs_cis:Tensor, mask:Tensor, sliding:bool, attn_kwargs:dict, ffn_kwargs:dict, save:bool=True):
+    attn, attn_saves = self.attention(x, freqs_cis, mask, sliding, **attn_kwargs)
     h = x + attn
     ffn, ffn_saves = self.feed_forward(h, **ffn_kwargs)
     h = h + ffn
@@ -175,8 +253,7 @@ class GPTOSS:
     h = self.tok_embeddings(tokens)
     bsz, seqlen = tokens.shape
     freqs_cis = self.freqs_cis.cast(h.dtype)[:, :seqlen, :, :, :]
-    mask_full = self._attn_mask(seqlen, False, dtypes.float32)
-    mask_sliding = self._attn_mask(seqlen, True, dtypes.float32)
+    mask_full = None if getenv("HK_FLASH_ATTENTION") else self._attn_mask(seqlen, dtypes.float32)
     for i in range(self.n_layers):
       attn_kwargs = dict(attention_norm=self.attention_norm[i], wqkv=self.wqkv[i], wqkv_scale=self.wqkv_scale[i],
                          wqkv_bias=self.wqkv_bias[i], wo=self.wo[i], wo_scale=self.wo_scale[i], wo_bias=self.wo_bias[i],
@@ -184,8 +261,7 @@ class GPTOSS:
       ffn_kwargs = dict(ffn_norm=self.ffn_norm[i], gate=self.gate[i], gate_bias=self.gate_bias[i],
                         w_gate_up=self.w_gate_up[i], w_gate_up_scale=self.w_gate_up_scale[i], w_gate_up_bias=self.w_gate_up_bias[i],
                         w_down=self.w_down[i], w_down_scale=self.w_down_scale[i], w_down_bias=self.w_down_bias[i])
-      mask = mask_sliding if i % 2 == 0 else mask_full
-      h, *_ = self.run_layer(h, freqs_cis, mask, attn_kwargs, ffn_kwargs, save=save)
+      h, *_ = self.run_layer(h, freqs_cis, mask_full, i % 2 == 0, attn_kwargs, ffn_kwargs, save=save)
 
     logits = self.norm(h) @ self.output.T
     return logits

examples/mlperf/optim.py

@@ -1,8 +1,8 @@
 from tinygrad.tensor import Tensor
 from tinygrad.dtype import dtypes
-from tinygrad.nn.optim import Optimizer
+from tinygrad.nn.optim import Optimizer, OptimizerGroup
 from tinygrad.helpers import FUSE_OPTIM, getenv
-from tinygrad.uop.ops import UOp, Ops
+from tinygrad.uop.ops import UOp, Ops, AxisType
 
 STOCHASTIC_ROUND = getenv("STOCHASTIC_ROUND", 0)
 MASTER_WEIGHTS = getenv("MASTER_WEIGHTS", 0)
@@ -21,6 +21,12 @@ def stochastic_round_bf16(x:Tensor) -> Tensor:
   noise = (noise * 0xFFFF).cast(dtypes.uint32)
   return ((bits + noise) & 0xFFFF0000).bitcast(dtypes.float32).cast(dtypes.bfloat16)
 
+def clip_grads(grads:list[Tensor], grad_acc, clip_norm) -> Tensor:
+  for g in grads: g.assign(g / grad_acc)
+  total_norm = Tensor.stack(*[g.float().square().sum() for g in grads]).sum().sqrt().contiguous()
+  for g in grads: g.assign((g * (clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(g.dtype))
+  return 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)
@@ -36,46 +42,29 @@ class GradAccClipAdamW(Optimizer):
       self.master_params = None
 
   def _zero_shard(self, t:Tensor) -> Tensor:
-    if not self.zero or (t.shape[0] % len(self.device)) != 0: return t
-    return Tensor(t.uop._shard(0, len(self.device)).multi(0)).clone()
+    if not self.zero or t.ndim < 2 or (t.shape[0] % len(self.device)) != 0: return t
+    return Tensor(t.uop._shard(0, UOp.range(len(self.device), -1, AxisType.DEVICE)).unshard(0)).clone()
 
   def _zero_gather(self, t:Tensor) -> Tensor:
     if not isinstance(t.device, tuple) or t.uop.axis != 0: return t
     n, sz = len(t.device), t.shape[0] // len(t.device)
     return Tensor.cat(*[t[p*sz:(p+1)*sz] for p in range(n)], dim=0)
 
-  def fstep(self, grads:list[Tensor]):
-    if self.fused:
-      out, extra = self._step([], grads)
-      updates = [out[0][self.pos_params[i]:self.pos_params[i+1]].reshape(tt.shape) for i, tt in enumerate(self.params)]
-    else:
-      updates, extra = self._step([], grads)
+  def fschedule_step(self, grads:list[Tensor]) -> list[Tensor]:
+    updates, extra = self._step([], grads)
     for i, tt in enumerate(self.params): tt.assign(self._apply_update(tt, updates[i], self.master_params[i] if self.master_params else None))
-    # collect inv_scale tensors attached to fp8 params (set by _apply_update)
     fp8_inv_scales = [tt._inv_scale for tt in self.params if hasattr(tt, '_inv_scale')]
     fp8_next_inv_scales = [tt._next_inv_scale for tt in self.params if hasattr(tt, '_next_inv_scale')]
-    to_realize = extra+self.params+self.buffers+(self.master_params or [])+fp8_inv_scales+fp8_next_inv_scales
+    return extra + self.params + self.buffers + (self.master_params or []) + fp8_inv_scales + fp8_next_inv_scales
 
-    Tensor.realize(*to_realize)
-    return extra[-1]
+  def fstep(self, grads:list[Tensor], grad_norm:Tensor|None=None):
+    Tensor.realize(*([grad_norm] if grad_norm is not None else []), *self.fschedule_step(grads))
 
   def _step(self, params:list[Tensor], grads:list[Tensor]) -> tuple[list[Tensor], list[Tensor]]:
     grads = list(grads)
 
     for i in range(len(grads)):
       if grads[i].device != self.m[i].device: grads[i] = grads[i].to(self.m[i].device)
-
-    if self.fused:
-      grads[0].assign(grads[0] / self.grad_acc)
-      total_norm = grads[0].float().square().sum().sqrt()
-      grads[0].assign((grads[0] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[0].dtype))
-    else:
-      for i in range(len(grads)):
-        grads[i].assign(grads[i] / self.grad_acc)
-      total_norm = Tensor.stack(*[g.float().square().sum() for g in grads]).sum().sqrt().contiguous()
-      for i in range(len(grads)):
-        grads[i].assign((grads[i] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[i].dtype))
-
     ret = []
     self.b1_t *= self.b1
     self.b2_t *= self.b2
@@ -88,7 +77,7 @@ class GradAccClipAdamW(Optimizer):
       v_hat = v_new / (1.0 - self.b2_t)
       up = m_hat / (v_hat.sqrt() + self.eps)
       ret.append(self.lr * up)
-    return ret, [self.b1_t, self.b2_t] + self.m + self.v + [total_norm]
+    return ret, [self.b1_t, self.b2_t] + self.m + self.v
 
   def _apply_update(self, t:Tensor, up:Tensor, master:Tensor|None=None) -> Tensor:
     w = master if master is not None else t
@@ -96,7 +85,7 @@ class GradAccClipAdamW(Optimizer):
     up = up.float().shard_like(w) + self.lr.to(w.device) * wd * w.detach()
     new_w = w.detach() - up
     if master is not None: master.assign(new_w)
-    if self.zero: new_w = self._zero_gather(new_w)
+    if self.zero and not (MXFP8 and t.dtype in dtypes.fp8s): new_w = self._zero_gather(new_w)
     # when master is offloaded to a different device than the param, results are resharded back onto the param's (sharded) device
     offloaded = master is not None and master.device != t.device
     if STOCHASTIC_ROUND and t.dtype == dtypes.bfloat16:
@@ -106,9 +95,10 @@ class GradAccClipAdamW(Optimizer):
       if MXFP8:
         from extra.gemm.cdna_asm_gemm import quantize_mxfp8
         w_q, w_e8, _ = quantize_mxfp8(new_w.reshape(-1, new_w.shape[-1]))
+        if self.zero: w_q, w_e8 = self._zero_gather(w_q), self._zero_gather(w_e8)
         new_e8 = w_e8.reshape(t._inv_scale.shape)
         t._inv_scale.assign(new_e8.shard_like(t._inv_scale) if offloaded else new_e8)
-        ret = w_q.reshape(new_w.shape)
+        ret = w_q.reshape(t.shape)
         return ret.shard_like(t) if offloaded else ret
       from examples.mlperf.models.flat_llama import FP8_MAX
       if IMMEDIATE_SCALE:
@@ -131,3 +121,24 @@ class GradAccClipAdamW(Optimizer):
       return ret.shard_like(t) if offloaded else ret
     out = new_w.cast(t.dtype)
     return out.shard_like(t) if offloaded else out
+
+class GradAccClipAdamWGroup(OptimizerGroup):
+  def __init__(self, *optimizers:GradAccClipAdamW):
+    super().__init__(*optimizers)
+    for o in self.optimizers[1:]: o.lr = self.optimizers[0].lr
+  def fstep(self, grads:list[Tensor], grad_norm:Tensor|None=None):
+    offset = 0
+    to_realize = []
+    for o in self.optimizers:
+      n = len(o.params)
+      to_realize += o.fschedule_step(grads[offset:offset+n])
+      offset += n
+    Tensor.realize(*to_realize, *([grad_norm] if grad_norm is not None else []))
+  @property
+  def lr(self): return self.optimizers[0].lr
+  @property
+  def device(self): return self.optimizers[0].device
+  @property
+  def master_params(self):
+    mp = [mp for o in self.optimizers for mp in (o.master_params or [])]
+    return mp if mp else None

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

@@ -20,7 +20,7 @@ export ZERO_OPTIM=${ZERO_OPTIM:-1}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=${DP:-8} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
+export DP=${DP:-8} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
 export GBS=$((BS * GRADIENT_ACC_STEPS))
 
 export MODEL="gptoss"
@@ -34,7 +34,7 @@ export SEED=${SEED:-5760}
 export DATA_SEED=${DATA_SEED:-5760}
 
 export JITBEAM=${JITBEAM:-3}
-export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
 
 export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
 if [ -z "$FULL_LAYERS" ]; then

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

@@ -20,7 +20,7 @@ export ZERO_OPTIM=${ZERO_OPTIM:-1}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=${DP:-8} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
+export DP=${DP:-8} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
 export GBS=$((BS * GRADIENT_ACC_STEPS))
 
 export MODEL="gptoss"
@@ -34,6 +34,6 @@ export SEED=${SEED:-$RANDOM}
 export DATA_SEED=${DATA_SEED:-5760}
 
 export JITBEAM=${JITBEAM:-3}
-export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
 
 python3 examples/mlperf/model_train.py

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

@@ -2,5 +2,4 @@
 export BENCHMARK=${BENCHMARK:-5}
 export EVAL_BS=0
 VIZ=${VIZ:--1} FULL_LAYERS=1 DEBUG=${DEBUG:--0} examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh
-SRC="AMD"; [[ $DEV == NULL* ]] && SRC="NULL"
-[ "$BENCHMARK" -le 3 ] || python -m tinygrad.viz.cli -s "$SRC" -t --interval "train @ 2" "train @ 3"
+[ "$BENCHMARK" -le 3 ] || [[ $DEV == NULL* ]] || python -m tinygrad.viz.cli -s AMD -t --interval "train @ 2" "train @ 3"

examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/gpt_oss/implementations/tinybox_8xMI350X/dev_beam.sh

@@ -0,0 +1,44 @@
+#!/usr/bin/env bash
+
+export PYTHONPATH="."
+export PATH="/opt/rocm-7.1.1/bin:$PATH"
+export ROCM_PATH="/opt/rocm-7.1.1"
+export DEV=${DEV:-AMD}
+export CHECK_OOB=0
+export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+export DEVICE_IN_FUNCTION_BUG=1
+
+export DEBUG=${DEBUG:-2}
+export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
+export ALL2ALL=${ALL2ALL:-1}
+export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
+export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
+export USE_ATOMICS=${USE_ATOMICS:-1}
+export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
+export MXFP8=${MXFP8:-1}
+export ZERO_OPTIM=${ZERO_OPTIM:-1}
+export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=${DP:-8} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
+export GBS=$((BS * GRADIENT_ACC_STEPS))
+
+export MODEL="gptoss"
+export BASEDIR="/raid/datasets/c4-8b/"
+export EVAL_TARGET=3.34 EVAL_FREQ=12288
+export END_LR="4e-5" WARMUP_STEPS=128 MAX_STEPS=1200000
+export SAMPLES=$((MAX_STEPS * GBS))
+export SEQLEN=${SEQLEN:-8192}
+
+export SEED=${SEED:-5760}
+export DATA_SEED=${DATA_SEED:-5760}
+
+export JITBEAM=${JITBEAM:-3}
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+
+export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
+if [ -z "$FULL_LAYERS" ]; then
+  export LAYERS=${LAYERS:-2}
+fi
+
+python3 examples/mlperf/model_train.py

examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/gpt_oss/implementations/tinybox_8xMI350X/dev_run.sh

@@ -0,0 +1,39 @@
+#!/usr/bin/env bash
+
+export PYTHONPATH="."
+export PATH="/opt/rocm-7.1.1/bin:$PATH"
+export ROCM_PATH="/opt/rocm-7.1.1"
+export DEV=${DEV:-AMD}
+export CHECK_OOB=0
+export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+export DEVICE_IN_FUNCTION_BUG=1
+
+export DEBUG=${DEBUG:-0}
+export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
+export ALL2ALL=${ALL2ALL:-1}
+export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
+export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
+export USE_ATOMICS=${USE_ATOMICS:-1}
+export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
+export MXFP8=${MXFP8:-1}
+export ZERO_OPTIM=${ZERO_OPTIM:-1}
+export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=${DP:-8} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
+export GBS=$((BS * GRADIENT_ACC_STEPS))
+
+export MODEL="gptoss"
+export BASEDIR="/raid/datasets/c4-8b/"
+export EVAL_TARGET=3.34 EVAL_FREQ=12288
+export END_LR="4e-5" WARMUP_STEPS=128 MAX_STEPS=1200000
+export SAMPLES=$((MAX_STEPS * GBS))
+export SEQLEN=${SEQLEN:-8192}
+
+export SEED=${SEED:-$RANDOM}
+export DATA_SEED=${DATA_SEED:-5760}
+
+export JITBEAM=${JITBEAM:-3}
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+
+python3 examples/mlperf/model_train.py

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

@@ -0,0 +1,54 @@
+#!/usr/bin/env bash
+
+export PYTHONPATH="."
+export PATH="/opt/rocm-7.1.1/bin:$PATH"
+export ROCM_PATH="/opt/rocm-7.1.1"
+export DEV=${DEV:-AMD}
+export CHECK_OOB=0
+export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+export DEVICE_IN_FUNCTION_BUG=1
+
+export DEBUG=${DEBUG:-2}
+export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
+export ALL2ALL=${ALL2ALL:-1}
+export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
+export USE_ATOMICS=${USE_ATOMICS:-1}
+export ASM_GEMM=${ASM_GEMM:-1}
+export WQKV=${WQKV:-1}
+export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
+export FP8=${FP8:-1}
+export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
+export FAST_CE=${FAST_CE:-1}
+export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
+export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-1}
+export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
+export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
+export SPLIT_W13=${SPLIT_W13:-0}
+export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
+export GBS=$((BS * GRADIENT_ACC_STEPS))
+
+export MODEL="llama3"
+export BASEDIR="/raid/datasets/c4-8b/"
+export SMALL=1
+export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
+export EVAL_TARGET=3.3 EVAL_FREQ=12288
+export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
+export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
+export SAMPLES=$((MAX_STEPS * GBS))
+export SEQLEN=${SEQLEN:-8192}
+
+export SEED=${SEED:-5760}
+export DATA_SEED=${DATA_SEED:-5760}
+
+export JITBEAM=${JITBEAM:-3}
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
+
+export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
+if [ -z "$FULL_LAYERS" ]; then
+  export LLAMA_LAYERS=${LLAMA_LAYERS:-2}
+fi
+
+python3 examples/mlperf/model_train.py

examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam_mp.sh

@@ -0,0 +1,54 @@
+#!/usr/bin/env bash
+
+export PYTHONPATH="."
+export PATH="/opt/rocm-7.1.1/bin:$PATH"
+export ROCM_PATH="/opt/rocm-7.1.1"
+export DEV=${DEV:-AMD}
+export CHECK_OOB=0
+export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+export DEVICE_IN_FUNCTION_BUG=1
+
+export DEBUG=${DEBUG:-2}
+export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
+export ALL2ALL=${ALL2ALL:-1}
+export LATE_ALLREDUCE=${LATE_ALLREDUCE:-1}
+export USE_ATOMICS=${USE_ATOMICS:-1}
+export ASM_GEMM=${ASM_GEMM:-1}
+export WQKV=${WQKV:-1}
+export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
+export FP8=${FP8:-1}
+export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
+export FAST_CE=${FAST_CE:-0}
+export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0}
+export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0}
+export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0}
+export FUSED_SILU_W13=${FUSED_SILU_W13:-0}
+export SPLIT_W13=${SPLIT_W13:-1}
+export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=${DP:-1} MP=${MP:-8} BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
+export GBS=$((BS * GRADIENT_ACC_STEPS))
+
+export MODEL="llama3"
+export BASEDIR="/raid/datasets/c4-8b/"
+export SMALL=1
+export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
+export EVAL_TARGET=3.3 EVAL_FREQ=12288
+export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
+export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
+export SAMPLES=$((MAX_STEPS * GBS))
+export SEQLEN=${SEQLEN:-8192}
+
+export SEED=${SEED:-5760}
+export DATA_SEED=${DATA_SEED:-5760}
+
+export JITBEAM=${JITBEAM:-3}
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
+
+export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
+if [ -z "$FULL_LAYERS" ]; then
+  export LLAMA_LAYERS=${LLAMA_LAYERS:-2}
+fi
+
+python3 examples/mlperf/model_train.py

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

@@ -0,0 +1,49 @@
+#!/usr/bin/env bash
+
+export PYTHONPATH="."
+export PATH="/opt/rocm-7.1.1/bin:$PATH"
+export ROCM_PATH="/opt/rocm-7.1.1"
+export DEV=${DEV:-AMD}
+export CHECK_OOB=0
+export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+export DEVICE_IN_FUNCTION_BUG=1
+
+export DEBUG=${DEBUG:-0}
+export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
+export ALL2ALL=${ALL2ALL:-1}
+export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
+export USE_ATOMICS=${USE_ATOMICS:-1}
+export ASM_GEMM=${ASM_GEMM:-1}
+export WQKV=${WQKV:-1}
+export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
+export FP8=${FP8:-1}
+export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
+export FAST_CE=${FAST_CE:-1}
+export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
+export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-1}
+export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
+export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
+export SPLIT_W13=${SPLIT_W13:-0}
+export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
+export GBS=$((BS * GRADIENT_ACC_STEPS))
+
+export MODEL="llama3"
+export BASEDIR="/raid/datasets/c4-8b/"
+export SMALL=1
+export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
+export EVAL_TARGET=3.3 EVAL_FREQ=12288
+export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
+export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
+export SAMPLES=$((MAX_STEPS * GBS))
+export SEQLEN=${SEQLEN:-8192}
+
+export SEED=${SEED:-$RANDOM}
+export DATA_SEED=${DATA_SEED:-5760}
+
+export JITBEAM=${JITBEAM:-3}
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
+
+python3 examples/mlperf/model_train.py

examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run_mp.sh

@@ -0,0 +1,49 @@
+#!/usr/bin/env bash
+
+export PYTHONPATH="."
+export PATH="/opt/rocm-7.1.1/bin:$PATH"
+export ROCM_PATH="/opt/rocm-7.1.1"
+export DEV=${DEV:-AMD}
+export CHECK_OOB=0
+export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+export DEVICE_IN_FUNCTION_BUG=1
+
+export DEBUG=${DEBUG:-0}
+export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
+export ALL2ALL=${ALL2ALL:-1}
+export LATE_ALLREDUCE=${LATE_ALLREDUCE:-1}
+export USE_ATOMICS=${USE_ATOMICS:-1}
+export ASM_GEMM=${ASM_GEMM:-1}
+export WQKV=${WQKV:-1}
+export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
+export FP8=${FP8:-1}
+export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
+export FAST_CE=${FAST_CE:-0}
+export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0}
+export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0}
+export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0}
+export FUSED_SILU_W13=${FUSED_SILU_W13:-0}
+export SPLIT_W13=${SPLIT_W13:-1}
+export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=${DP:-1} MP=${MP:-8} BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-32}
+export GBS=$((BS * GRADIENT_ACC_STEPS))
+
+export MODEL="llama3"
+export BASEDIR="/raid/datasets/c4-8b/"
+export SMALL=1
+export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
+export EVAL_TARGET=3.3 EVAL_FREQ=12288
+export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
+export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
+export SAMPLES=$((MAX_STEPS * GBS))
+export SEQLEN=${SEQLEN:-8192}
+
+export SEED=${SEED:-$RANDOM}
+export DATA_SEED=${DATA_SEED:-5760}
+
+export JITBEAM=${JITBEAM:-3}
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
+
+python3 examples/mlperf/model_train.py

examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/profile.sh

@@ -0,0 +1,5 @@
+#!/bin/bash
+export BENCHMARK=${BENCHMARK:-5}
+export EVAL_BS=0
+VIZ=${VIZ:--1} FULL_LAYERS=1 DEBUG=${DEBUG:--0} examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh
+[ "$BENCHMARK" -le 3 ] || [[ $DEV == NULL* ]] || python -m tinygrad.viz.cli -s AMD -t --interval "train @ 2" "train @ 3"

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

@@ -0,0 +1,58 @@
+#!/usr/bin/env bash
+set -e  # Exit on any error
+set -o pipefail  # Make pipeline fail if any command fails
+
+export PYTHONPATH="."
+export PATH="/opt/rocm-7.1.1/bin:$PATH"
+export ROCM_PATH="/opt/rocm-7.1.1"
+export DEV=AMD
+export CHECK_OOB=0
+export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+export DEVICE_IN_FUNCTION_BUG=1
+
+export HK_FLASH_ATTENTION=1
+export ALL2ALL=1
+export LATE_ALLREDUCE=0
+export USE_ATOMICS=1
+export ASM_GEMM=1
+export WQKV=1
+export MASTER_WEIGHTS=1
+export FP8=1
+export ALLREDUCE_CAST=1
+export FAST_CE=1
+export FUSED_INPUT_QUANTIZE=1
+export FUSED_GRAD_QUANTIZE=1
+export FUSED_ADD_NORM_MUL_QUANTIZE=1
+export FUSED_SILU_W13=1
+export SPLIT_W13=0
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=8 MP=1 BS=16 EVAL_BS=8 GRADIENT_ACC_STEPS=2
+export GBS=$((BS * GRADIENT_ACC_STEPS))
+
+export MODEL="llama3"
+export BASEDIR="/raid/datasets/c4-8b/"
+export SMALL=1
+export LLAMA3_SIZE=8B
+export EVAL_TARGET=3.3 EVAL_FREQ=12288
+export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
+export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
+export SAMPLES=$((MAX_STEPS * GBS))
+export SEQLEN=8192
+
+export SEED=$RANDOM
+export DATA_SEED=$SEED
+
+export JITBEAM=3
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
+
+export LOGMLPERF=1
+
+DATETIME=$(date "+%m%d%H%M")
+LOGFILE="llama31_8b_8xMI350x_${DATETIME}_${SEED}.log"
+
+# beam
+FAKEDATA=1 BENCHMARK=10 INITMLPERF=1 LLAMA_LAYERS=2 python3 examples/mlperf/model_train.py | tee "$LOGFILE"
+
+# run
+RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a "$LOGFILE"

examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/startup_walltime.sh

@@ -0,0 +1,10 @@
+#!/bin/bash
+export BENCHMARK=5
+export EVAL_BS=0
+export FAKEDATA=1
+export NULL_ALLOW_COPYOUT=1
+export HIP_VISIBLE_DEVICES=""
+export DEV=NULL:HIP:gfx950
+export JITBEAM=0
+export LLAMA_LAYERS=${LLAMA_LAYERS:-"2"}
+time examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh

examples/mlperf/training_submission_v6.1/tinycorp/systems/tinybox_8xMI300X.json

@@ -0,0 +1,38 @@
+{
+    "submitter": "tinycorp",
+    "division": "closed",
+    "status": "Available on-premise",
+    "system_name": "tinybox 8xMI300X",
+    "number_of_nodes": "1",
+    "host_processors_per_node": "2",
+    "host_processor_model_name": "AMD EPYC 9354",
+    "host_processor_core_count": "32",
+    "host_processor_vcpu_count": "64",
+    "host_processor_frequency": "",
+    "host_processor_caches": "",
+    "host_processor_interconnect": "",
+    "host_memory_capacity": "2304GB",
+    "host_storage_type": "NVMe SSD",
+    "host_storage_capacity": "3x 4TB raid array",
+    "host_networking": "",
+    "host_networking_topology": "",
+    "host_memory_configuration": "24x 96GB DDR5",
+    "accelerators_per_node": "8",
+    "accelerator_model_name": "AMD Instinct MI300X 192GB HBM3",
+    "accelerator_host_interconnect": "PCIe 5.0 x16",
+    "accelerator_frequency": "",
+    "accelerator_on-chip_memories": "",
+    "accelerator_memory_configuration": "HBM3",
+    "accelerator_memory_capacity": "192GB",
+    "accelerator_interconnect": "",
+    "accelerator_interconnect_topology": "",
+    "cooling": "air",
+    "hw_notes": "",
+    "framework": "tinygrad, branch mlperf_training_v5.0",
+    "other_software_stack": {
+        "python": "3.10.16",
+        "ROCm": "3.0.0+94441cb"
+    },
+    "operating_system": "Ubuntu 24.04.1 LTS",
+    "sw_notes": ""
+  }

examples/mlperf/training_submission_v6.1/tinycorp/systems/tinybox_8xMI350X.json

@@ -0,0 +1,38 @@
+{
+    "submitter": "tinycorp",
+    "division": "closed",
+    "status": "Available on-premise",
+    "system_name": "tinybox 8xMI350X",
+    "number_of_nodes": "1",
+    "host_processors_per_node": "2",
+    "host_processor_model_name": "AMD EPYC 9575F",
+    "host_processor_core_count": "32",
+    "host_processor_vcpu_count": "64",
+    "host_processor_frequency": "",
+    "host_processor_caches": "",
+    "host_processor_interconnect": "",
+    "host_memory_capacity": "3072 GiB",
+    "host_storage_type": "NVMe SSD",
+    "host_storage_capacity": "4TB",
+    "host_networking": "",
+    "host_networking_topology": "",
+    "host_memory_configuration": "24x 128GB DDR5",
+    "accelerators_per_node": "8",
+    "accelerator_model_name": "AMD Instinct MI350X 288GB HBM3e",
+    "accelerator_host_interconnect": "PCIe 5.0 x16",
+    "accelerator_frequency": "",
+    "accelerator_on-chip_memories": "",
+    "accelerator_memory_configuration": "HBM3",
+    "accelerator_memory_capacity": "288GB",
+    "accelerator_interconnect": "",
+    "accelerator_interconnect_topology": "",
+    "cooling": "air",
+    "hw_notes": "",
+    "framework": "tinygrad, branch mlperf_training_v6.0",
+    "other_software_stack": {
+        "python": "3.12.3",
+        "ROCm": "7.1.1"
+    },
+    "operating_system": "Ubuntu 24.04.3 LTS",
+    "sw_notes": ""
+  }

examples/mlperf/training_submission_v6.1/tinycorp/systems/tinybox_green.json

@@ -0,0 +1,38 @@
+{
+  "submitter": "tinycorp",
+  "division": "closed",
+  "status": "Available on-premise",
+  "system_name": "tinybox green",
+  "number_of_nodes": "1",
+  "host_processors_per_node": "1",
+  "host_processor_model_name": "AMD EPYC 7532",
+  "host_processor_core_count": "32",
+  "host_processor_vcpu_count": "64",
+  "host_processor_frequency": "",
+  "host_processor_caches": "",
+  "host_processor_interconnect": "",
+  "host_memory_capacity": "128GB",
+  "host_storage_type": "NVMe SSD",
+  "host_storage_capacity": "4 TB raid array + 1 TB boot",
+  "host_networking": "",
+  "host_networking_topology": "",
+  "host_memory_configuration": "8x 16GB DDR4",
+  "accelerators_per_node": "6",
+  "accelerator_model_name": "NVIDIA GeForce RTX 4090",
+  "accelerator_host_interconnect": "PCIe 4.0 x16",
+  "accelerator_frequency": "",
+  "accelerator_on-chip_memories": "",
+  "accelerator_memory_configuration": "GDDR6X",
+  "accelerator_memory_capacity": "24GB",
+  "accelerator_interconnect": "",
+  "accelerator_interconnect_topology": "",
+  "cooling": "air",
+  "hw_notes": "",
+  "framework": "tinygrad, branch mlperf_training_v5.0",
+  "other_software_stack": {
+    "python": "3.10.12",
+    "CUDA": "12.4"
+  },
+  "operating_system": "Ubuntu 22.04.4",
+  "sw_notes": ""
+}

examples/mlperf/training_submission_v6.1/tinycorp/systems/tinybox_red.json

@@ -0,0 +1,37 @@
+{
+  "submitter": "tinycorp",
+  "division": "closed",
+  "status": "Available on-premise",
+  "system_name": "tinybox red",
+  "number_of_nodes": "1",
+  "host_processors_per_node": "1",
+  "host_processor_model_name": "AMD EPYC 7532",
+  "host_processor_core_count": "32",
+  "host_processor_vcpu_count": "64",
+  "host_processor_frequency": "",
+  "host_processor_caches": "",
+  "host_processor_interconnect": "",
+  "host_memory_capacity": "128GB",
+  "host_storage_type": "NVMe SSD",
+  "host_storage_capacity": "4 TB raid array + 1 TB boot",
+  "host_networking": "",
+  "host_networking_topology": "",
+  "host_memory_configuration": "8x 16GB DDR4",
+  "accelerators_per_node": "6",
+  "accelerator_model_name": "AMD Radeon RX 7900 XTX",
+  "accelerator_host_interconnect": "PCIe 4.0 x16",
+  "accelerator_frequency": "",
+  "accelerator_on-chip_memories": "",
+  "accelerator_memory_configuration": "GDDR6",
+  "accelerator_memory_capacity": "24GB",
+  "accelerator_interconnect": "",
+  "accelerator_interconnect_topology": "",
+  "cooling": "air",
+  "hw_notes": "",
+  "framework": "tinygrad, branch mlperf_training_v5.0",
+  "other_software_stack": {
+    "python": "3.10.12"
+  },
+  "operating_system": "Ubuntu 22.04.4",
+  "sw_notes": ""
+}

examples/openpilot/compile3.py

@@ -1,4 +1,4 @@
-import os, sys, pickle, time, re
+import os, sys, pickle, time, re, tempfile, struct, shutil, io
 import numpy as np
 if "JIT_BATCH_SIZE" not in os.environ: os.environ["JIT_BATCH_SIZE"] = "0"
 
@@ -9,6 +9,39 @@ from tinygrad.nn.onnx import OnnxRunner
 
 OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx"
 OUTPUT = sys.argv[2] if len(sys.argv) > 2 else "/tmp/openpilot.pkl"
+PICKLE_OOB = getenv("PICKLE_OOB")
+
+def dump_pickle(obj, f):
+  if PICKLE_OOB:
+    # allows pickling when buffers don't fit in (CPU) RAM
+    # from openpilot/selfdrive/modeld/helpers.py
+    with tempfile.TemporaryFile(dir=".") as tmp:
+      def buffer_callback(pb: pickle.PickleBuffer):
+        m = pb.raw()
+        tmp.write(struct.pack('<q', m.nbytes))
+        tmp.write(m)
+        pb.release() # keep peak ram at ~1 buffer
+      stream = io.BytesIO()
+      pickle.Pickler(stream, protocol=5, buffer_callback=buffer_callback).dump(obj)
+      opcodes = stream.getvalue()
+      f.write(struct.pack('<q', len(opcodes)))
+      f.write(opcodes)
+      tmp.seek(0)
+      shutil.copyfileobj(tmp, f)
+  else: pickle.dump(obj, f)
+
+def load_pickle(f):
+  if PICKLE_OOB:
+    # allows unpickling when buffers don't fit in (CPU) RAM
+    # from openpilot/selfdrive/modeld/helpers.py
+    opcodes = f.read(struct.unpack('<q', f.read(8))[0])
+    def buffers():
+      while (h := f.read(8)):
+        pb = pickle.PickleBuffer(bytearray(struct.unpack('<q', h)[0]))
+        f.readinto(pb)
+        yield pb
+    return pickle.load(io.BytesIO(opcodes), buffers=buffers())
+  else: return pickle.load(f)
 
 def compile(onnx_file):
   run_onnx = OnnxRunner(onnx_file)
@@ -28,8 +61,8 @@ def compile(onnx_file):
     inputs = {k:Tensor(v.numpy(), device=Device.DEFAULT).realize() if 'img' in k else v for k,v in inputs.items()}
   print("created tensors")
 
-  run_onnx_jit = TinyJit(lambda **kwargs:
-                         next(iter(run_onnx({k:v.to(Device.DEFAULT) for k,v in kwargs.items()}).values())).cast('float32'), prune=True)
+  @TinyJit(prune=True)
+  def run_onnx_jit(**kwargs): return next(iter(run_onnx({k:v.to(Device.DEFAULT) for k,v in kwargs.items()}).values())).cast('float32')
   for i in range(3):
     GlobalCounters.reset()
     print(f"run {i}")
@@ -65,38 +98,13 @@ def compile(onnx_file):
   if (allowed_gated_read_image:=getenv("ALLOWED_GATED_READ_IMAGE", -1)) != -1:
     assert gated_read_image_count == allowed_gated_read_image, f"different gated read_image! {gated_read_image_count=}, {allowed_gated_read_image=}"
 
-  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_RPT", 1):
-    from extra.gemm.qcom_openpilot_vision_fp16 import patch_fp32_rpt
-    if (patched:=patch_fp32_rpt(run_onnx_jit)): print(f"repeat-packed {patched} QCOM vision kernels")
-  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_SCHEDULE", 1):
-    from extra.gemm.qcom_openpilot_schedule_projection import patch_projection
-    if (patched:=patch_projection(run_onnx_jit)): print(f"rescheduled {patched} QCOM vision kernels")
-  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_FULL_RPT", 1):
-    from extra.gemm.qcom_openpilot_inverse_full_rpt import patch_model as patch_full_rpt
-    if (patched:=patch_full_rpt(run_onnx_jit)): print(f"fully repeat-packed {patched} QCOM vision kernels")
-  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_DEDUPE", 1):
-    from extra.gemm.qcom_openpilot_dedupe_head import dedupe_identical_calls
-    if (removed:=dedupe_identical_calls(run_onnx_jit)): print(f"deduplicated {len(removed)} QCOM kernels")
-  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_PACK_CONV", 1):
-    from extra.gemm.qcom_openpilot_pack_conv_weights import patch_conv
-    if (patched:=patch_conv(run_onnx_jit)): print(f"packed weights for {patched} QCOM convolution kernels")
-  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_LEVEL_SCHEDULE", 1):
-    from extra.gemm.qcom_openpilot_level_schedule import schedule_levels
-    if (moved:=schedule_levels(run_onnx_jit)): print(f"rescheduled {moved} QCOM kernels by dependency level")
-  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_BATCH_HEAD", 1):
-    from extra.gemm.qcom_openpilot_batch_head import batch_head
-    if (combined:=batch_head(run_onnx_jit)): print(f"batched {combined} groups of QCOM head kernels")
-  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_INPUT_PACK", 1):
-    from extra.gemm.qcom_openpilot_input_pack import patch_input_pack
-    if (patched:=patch_input_pack(run_onnx_jit)): print(f"vectorized {patched} QCOM input kernel")
-  with open(OUTPUT, "wb") as f:
-    pickle.dump(run_onnx_jit, f)
+  with open(OUTPUT, "wb") as f: dump_pickle(run_onnx_jit, f)
   mdl_sz = os.path.getsize(onnx_file)
   pkl_sz = os.path.getsize(OUTPUT)
   print(f"mdl size is {mdl_sz/1e6:.2f}M")
   print(f"pkl size is {pkl_sz/1e6:.2f}M")
   print("**** compile done ****")
-  return run_onnx_jit, inputs, test_val
+  return inputs, test_val
 
 def test_vs_compile(run, inputs, test_val=None):
 
@@ -160,16 +168,15 @@ def bench(run, inputs):
 
 if __name__ == "__main__":
   if getenv("RUN_PICKLE"):
-    with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f)
+    with open(OUTPUT, "rb") as f: pickle_loaded = load_pickle(f)
     inputs = {name: Tensor(Tensor.randn(*view.shape, dtype=dtype).numpy(), device=device)
               for name, (view, _vars, dtype, device) in zip(pickle_loaded.captured.expected_names, pickle_loaded.captured.expected_input_info)}
     test_vs_compile(pickle_loaded, inputs)
   else:
     onnx_file = fetch(OPENPILOT_MODEL)
-    pickle_loaded, inputs, outputs = compile(onnx_file)
+    inputs, outputs = compile(onnx_file)
 
-    if OUTPUT != os.devnull:
-      with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f)
+    with open(OUTPUT, "rb") as f: pickle_loaded = load_pickle(f)
 
     test_vs_compile(pickle_loaded, inputs, outputs)
     if getenv("SELFTEST"):

examples/openpilot/load_pickle.py

@@ -1,5 +1,6 @@
-import sys, pickle
+import sys
 from extra.bench_log import WallTimeEvent, BenchEvent
+from examples.openpilot.compile3 import load_pickle
 from tinygrad.helpers import getenv
 
 PKL = sys.argv[1] if len(sys.argv) > 1 else "/tmp/openpilot.pkl"
@@ -7,7 +8,7 @@ PKL = sys.argv[1] if len(sys.argv) > 1 else "/tmp/openpilot.pkl"
 load_times = []
 
 for _ in range(10):
-  with WallTimeEvent(BenchEvent.STEP) as wte: pickle.load(open(PKL, 'rb'))
+  with WallTimeEvent(BenchEvent.STEP) as wte: load_pickle(open(PKL, 'rb'))
   load_times.append(wte.time)
   print(f"pickle load: {wte.time:6.2f} s")
 

extra/export_model.py

@@ -241,8 +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"
 
-  # NOTE: CPU_COUNT=1, since export does not support threading
-  with Context(JIT=2, CPU_COUNT=1): 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}
@@ -264,7 +264,7 @@ def export_model(model, target:str, *inputs, model_name: Optional[str] = "model"
         if getattr(dim, "op", None) is Ops.ADD and len(dim.src) == 2 and \
            any(s.op is Ops.PARAM and s.addrspace is AddrSpace.ALU for s in dim.src) and any(s.op is Ops.CONST for s in dim.src):
           name, val = dim.src if dim.src[1].op is Ops.CONST else reversed(dim.src)
-          global_size[j] = f"_{name.expr}[0] + {val.arg}"
+          global_size[j] = f"_{name.expr}[0] + {val.val}"
 
   prg = ""
   if target == "clang":

extra/f16_decompress.py

@@ -5,10 +5,10 @@ def bit_extract(x: Tensor, e: int, s: int) -> Tensor:
   return (x >> s) & mask
 
 def u16_to_f16(x: Tensor) -> Tensor:
-  sign = bit_extract(x, 15, 15).float()
+  sign = bit_extract(x, 15, 15).bool()
   exponent = bit_extract(x, 14, 10).float()
   fraction = bit_extract(x, 9, 0).float()
-  return sign.where(-1, 1) * exponent.where((exponent - 15.0).exp2() * (1 + fraction / 1024.0), 6.103515625e-5 * (fraction / 1024.0))
+  return sign.where(-1, 1) * exponent.bool().where((exponent - 15.0).exp2() * (1 + fraction / 1024.0), 6.103515625e-5 * (fraction / 1024.0))
 
 def u32_to_f16(oo: Tensor) -> Tensor:
   f1 = u16_to_f16(oo>>16)

extra/fp8/fp8_linear.py

@@ -18,9 +18,9 @@ def custom_matmul(output: UOp, inp: UOp, weight: UOp) -> UOp:
   SEQ = inp.shape[1]
   OUT = weight.shape[0]
   IN = weight.shape[-1]
-  seq_idx = UOp.range(SEQ, 2, AxisType.LOOP)
-  out_idx = UOp.range(OUT, 3, AxisType.LOOP)
-  batch_idx = UOp.range(output.size//SEQ//OUT, 1, AxisType.LOOP)
+  seq_idx = UOp.range(SEQ, 2)
+  out_idx = UOp.range(OUT, 3)
+  batch_idx = UOp.range(output.size//SEQ//OUT, 1)
   reduce_idx = UOp.range(IN, 0, AxisType.REDUCE)
   product = (inp.index((seq_idx*IN+reduce_idx+batch_idx*IN*SEQ)) * weight.index((out_idx*IN+reduce_idx))).cast(dtypes.float)
   reduced = product.reduce(reduce_idx, arg=Ops.ADD)
@@ -53,7 +53,7 @@ class FP8Linear:
     x_fp8, x_scale = quantize_to_fp8(x)
     GPUS = self.weight.device
     if isinstance(GPUS, tuple) and len(GPUS) > 1:
-      y = Tensor(Tensor.empty((batch//len(GPUS), seq, self.weight.shape[0]), dtype=dtypes.float, device=GPUS).uop.multi(0), device=GPUS)
+      y = Tensor(Tensor.empty((batch//len(GPUS), seq, self.weight.shape[0]), dtype=dtypes.float, device=GPUS).uop.unshard(0), device=GPUS)
     else:
       y = Tensor.empty((batch, seq, self.weight.shape[0]), dtype=dtypes.float)
     y = Tensor.custom_kernel(y, x_fp8, w_fp8, fxn=custom_matmul, grad_fxn=custom_matmul_backward)[0]

extra/gemm/QUALCOMM.md

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

extra/gemm/amd_copy_matmul.py

@@ -58,8 +58,8 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
   B_copy = B_local.permute((1,0)) if use_wmma else B_local
   A_store = A_copy.reshape(-1, THREADS_PER_BLOCK)[:, tid].store(a[k_tile].reshape(-1, THREADS_PER_BLOCK)[:, tid])
   B_store = B_copy.reshape(-1, THREADS_PER_BLOCK)[:, tid].store(b[k_tile].reshape(-1, THREADS_PER_BLOCK)[:, tid])
-  barrier = UOp.barrier(A_store, B_store)
-  A_local, B_local = A_local.after(barrier), B_local.after(barrier)
+  # NOTE: no explicit barrier needed, the AFTER on the LOCAL buffers implies it in late codegen
+  A_local, B_local = A_local.after(A_store, B_store), B_local.after(A_store, B_store)
 
   # -- COMPUTE --
   lane_m, lane_n = lane // LANES_PER_WAVE_N, lane % LANES_PER_WAVE_N
@@ -70,8 +70,8 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
 
   if use_wmma:
     k = UOp.range(BLOCK_K // WMMA_K, 101, AxisType.REDUCE)
-    tile_m = UOp.range(TM // WMMA_ACC, 200, AxisType.LOOP)
-    tile_n = UOp.range(TN, 201, AxisType.LOOP)
+    tile_m = UOp.range(TM // WMMA_ACC, 200)
+    tile_n = UOp.range(TN, 201)
 
     acc_frag = acc.reshape(TM // WMMA_ACC, WMMA_ACC, TN).permute(0,2,1)[tile_m, tile_n]
     a_frag = A_local.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, BLOCK_K // WMMA_K, WMMA_K)[wave_m, tile_m, lane_n, k]
@@ -80,7 +80,7 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
       # NOTE: since this is part of K, these 2 can be anywhere in the frags and long as a and b match
       a_frag = a_frag.reshape(2, 8)[lane_m, :]
       b_frag = b_frag.reshape(2, 8)[lane_m, :]
-    wmma = UOp.wmma(a_frag, b_frag, acc_frag.after(k), ((16, 16, 16), 'AMD', 32))
+    wmma = UOp.wmma(a_frag, b_frag, acc_frag.after(k), (16, 16, 16), 'AMD', 32)
     acc_store = acc_frag.store(wmma).end(tile_m, tile_n)
   else:
     # registers for LOCAL -> REG
@@ -96,8 +96,8 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
     b_frag = b_frag.reshape(1, TN).expand(TM, TN)
     acc_store = acc.store(acc.after(k) + (a_frag * b_frag))
 
-  # store accumulator and loop
-  acc = acc.after(acc_store.end(k).barrier().end(k_tile))
+  # store accumulator and loop (the barrier at the end of the loop is implied by the LOCAL buffers stored and loaded in the loop)
+  acc = acc.after(acc_store.end(k).end(k_tile))
 
   # store accumulator to output (unified)
   c = c.reshape(WAVES_M, TM//UNROLL_M, LANES_PER_WAVE_M, UNROLL_M,

extra/gemm/amd_flash_attention.py

@@ -13,7 +13,7 @@ WMMA_ACC = WMMA_M // LANES_PER_WAVE_M
 THREADS_PER_BLOCK = WARP_SIZE * WAVES_M * WAVES_N
 LDS_PAD = 4  # pad LDS rows to reduce bank conflicts
 
-WMMA_ARG = ((WMMA_M, WMMA_N, WMMA_K), 'AMD', 32)
+WMMA_ARG = (WMMA_M, WMMA_N, WMMA_K), 'AMD', 32
 LOG2E = math.log2(math.e)
 
 def warp_shfl_xor(val, offset, lane):
@@ -84,20 +84,20 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
     q.reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid])
   K_store = KV_lds.reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid].store(
     k[n_tile].reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid])
-  qk_load_barrier = UOp.barrier(UOp.group(Q_store, K_store))
-  Q_lds = Q_lds.after(qk_load_barrier)
-  KV_lds_k = KV_lds.after(qk_load_barrier)
+  # NOTE: no explicit barrier needed, the AFTER on the LOCAL buffers implies it in late codegen
+  Q_lds = Q_lds.after(UOp.group(Q_store, K_store))
+  KV_lds_k = KV_lds.after(UOp.group(Q_store, K_store))
 
   # -- S = Q @ K^T via WMMA (re-init each n_tile) --
   S_reg = UOp.placeholder((TM, TN), dtypes.float, slot=6, addrspace=AddrSpace.REG)
   S_reg = S_reg.after(S_reg.after(n_tile).store(S_reg.const_like(0)))
   k_qk = UOp.range(D // WMMA_K, 101, AxisType.REDUCE)
-  tm1 = UOp.range(TM // WMMA_ACC, 200, AxisType.LOOP)
-  tn1 = UOp.range(TN, 201, AxisType.LOOP)
+  tm1 = UOp.range(TM // WMMA_ACC, 200)
+  tn1 = UOp.range(TN, 201)
   S_frag = S_reg.reshape(TM // WMMA_ACC, WMMA_ACC, TN).permute(0, 2, 1)[tm1, tn1]
   q_frag = Q_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, D // WMMA_K, WMMA_K)[wave_m, tm1, lane_n, k_qk]
   k_frag = KV_lds_k.reshape(WAVES_N, TN, WMMA_N, D // WMMA_K, WMMA_K)[wave_n, tn1, lane_n, k_qk]
-  qk = UOp.wmma(q_frag, k_frag, S_frag.after(k_qk), WMMA_ARG)
+  qk = UOp.wmma(q_frag, k_frag, S_frag.after(k_qk), *WMMA_ARG)
   qk_done = S_frag.store(qk).end(tm1, tn1).end(k_qk)
   S_reg = S_reg.after(qk_done)
 
@@ -110,7 +110,7 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
   rm2 = UOp.range(TN, 261, AxisType.REDUCE)
   m_ij = m_ij.after(m_ij.store(m_ij.after(rm2).maximum(S_reg[:, rm2])).end(rm2))
   # warp reduce max (in-place)
-  ri_w = UOp.range(TM, 270, AxisType.LOOP)
+  ri_w = UOp.range(TM, 270)
   m_ij = m_ij.after(m_ij[ri_w].store(warp_reduce_max(m_ij[ri_w], lane)).end(ri_w))
 
   # compute P = exp(S - m_ij) in S_reg
@@ -120,7 +120,7 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
   p_local = p_local.after(p_local.after(n_tile).store(p_local.const_like(0)))
   rp2 = UOp.range(TN, 291, AxisType.REDUCE)
   p_local = p_local.after(p_local.store(p_local.after(rp2) + S_reg[:, rp2]).end(rp2))
-  ri_ws = UOp.range(TM, 295, AxisType.LOOP)
+  ri_ws = UOp.range(TM, 295)
   p_sum = p_local.after(p_local[ri_ws].store(warp_reduce_sum(p_local[ri_ws], lane)).end(ri_ws))
 
   # write P = exp(S - m_ij) to P_lds (reuses slot 0, Q no longer needed)
@@ -130,11 +130,11 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
   P_store = P_write[tid].store(S_reg.cast(dtypes.half))
 
   # -- online softmax correction --
-  ri4 = UOp.range(TM, 330, AxisType.LOOP)
+  ri4 = UOp.range(TM, 330)
   m_new_val = m_i[ri4].maximum(m_ij[ri4])
   alpha_val = ((m_i[ri4] - m_new_val) * LOG2E).exp2()
   beta_val = ((m_ij[ri4] - m_new_val) * LOG2E).exp2()
-  rj4 = UOp.range(TD, 331, AxisType.LOOP)
+  rj4 = UOp.range(TD, 331)
   correction = UOp.group(
     acc[ri4, rj4].store(alpha_val * acc[ri4, rj4]).end(rj4),
     l_i[ri4].store(alpha_val * l_i[ri4] + beta_val * p_sum[ri4]),
@@ -147,21 +147,21 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
   # load V into KV_lds (must wait for QK WMMA to finish reading K from KV_lds)
   V_store = KV_lds.after(qk_done).reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid].store(
     v[n_tile].reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid])
-  pv_barrier = UOp.barrier(UOp.group(P_store, V_store))
-  P_lds = P_lds.after(pv_barrier)
-  KV_lds_v = KV_lds.after(pv_barrier)
+  # NOTE: no explicit barrier needed, the AFTER on the LOCAL buffers implies it in late codegen
+  P_lds = P_lds.after(UOp.group(P_store, V_store))
+  KV_lds_v = KV_lds.after(UOp.group(P_store, V_store))
 
   # -- acc += P @ V via WMMA --
   k_pv = UOp.range(BLOCK_N // WMMA_K, 400, AxisType.REDUCE)
-  tm2 = UOp.range(TM // WMMA_ACC, 401, AxisType.LOOP)
-  tn2 = UOp.range(TD, 402, AxisType.LOOP)
+  tm2 = UOp.range(TM // WMMA_ACC, 401)
+  tn2 = UOp.range(TD, 402)
   acc_frag = acc.reshape(TM // WMMA_ACC, WMMA_ACC, TD).permute(0, 2, 1)[tm2, tn2]
   p_frag = P_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, BLOCK_N // WMMA_K, WMMA_K)[wave_m, tm2, lane_n, k_pv]
   v_frag = KV_lds_v.reshape(WAVES_N, TD, WMMA_N, BLOCK_N // WMMA_K, WMMA_K)[wave_n, tn2, lane_n, k_pv]
-  pv = UOp.wmma(p_frag, v_frag, acc_frag.after(k_pv), WMMA_ARG)
+  pv = UOp.wmma(p_frag, v_frag, acc_frag.after(k_pv), *WMMA_ARG)
 
   # end KV tile loop
-  n_tile_end = acc_frag.store(pv).end(tm2, tn2).end(k_pv).barrier().end(n_tile)
+  n_tile_end = acc_frag.store(pv).end(tm2, tn2).end(k_pv).end(n_tile)
   acc = acc.after(n_tile_end)
   l_i = l_i.after(n_tile_end)
   m_i = m_i.after(n_tile_end)

extra/gemm/amd_uop_matmul.py

@@ -28,10 +28,10 @@ REG_TILES_PER_WAVE_M = BLOCK_M // (WAVES_PER_BLOCK_M * LANES_PER_WAVE_M * TM)
 assert WAVES_PER_BLOCK_M*REG_TILES_PER_WAVE_M*LANES_PER_WAVE_M*TM == BLOCK_M, "M reshape is wrong"
 assert WAVES_PER_BLOCK_N*REG_TILES_PER_WAVE_N*LANES_PER_WAVE_N*TN == BLOCK_N, "N reshape is wrong"
 
-def rngs_for_shape(shape:tuple[sint, ...], rng:int, axis_type=AxisType.LOOP): return [UOp.range(s, rng+i, axis_type) for i,s in enumerate(shape)]
+def rngs_for_shape(shape:tuple[sint, ...], rng:int, axis_type=AxisType.WEAK): return [UOp.range(s, rng+i, axis_type) for i,s in enumerate(shape)]
 def copy(dest:UOp, src:UOp, rng:int, upcast=False):
   assert dest.shape == src.shape
-  rngs = rngs_for_shape(src.shape, rng, AxisType.UPCAST if upcast else AxisType.LOOP)
+  rngs = rngs_for_shape(src.shape, rng, AxisType.UPCAST if upcast else AxisType.WEAK)
   return dest[*rngs].store(src[*rngs]).end(*rngs)
 
 def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
@@ -66,9 +66,8 @@ def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
   B_local = UOp.placeholder((BLOCK_K, BLOCK_N), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
   B_local_store = copy(B_local.reshape(-1, THREADS_PER_BLOCK)[:, tid], b.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=200)
 
-  # TODO: can we automate barrier?
-  barrier = UOp.barrier(A_local_store, B_local_store)
-  A_local, B_local = A_local.after(barrier), B_local.after(barrier)
+  # NOTE: no explicit barrier needed, the AFTER on the LOCAL buffers implies it in late codegen
+  A_local, B_local = A_local.after(A_local_store, B_local_store), B_local.after(A_local_store, B_local_store)
 
   # open inner k range
   k = UOp.range(BLOCK_K, 3, AxisType.REDUCE)
@@ -102,7 +101,7 @@ def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
   sink = c_regs[*rngs].store(c_regs.after(k)[*rngs] + A_col[iter_m, t_m] * B_row[iter_n, t_n]).end(iter_m, iter_n, t_m, t_n)
 
   # Close k, sync, and close K tiles
-  sink = sink.end(k).barrier().end(k_tile_range)
+  sink = sink.end(k).end(k_tile_range)
 
   # ---------------------------
   # REG -> GLOBAL (epilogue)

extra/gemm/cdna_asm_gemm.py

@@ -1,8 +1,9 @@
-import atexit, functools, pathlib
+import atexit, functools, math, pathlib
 from tinygrad import Tensor, Device, dtypes
+from tinygrad.dtype import AddrSpace
 from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
 from tinygrad.renderer import Estimates
-from tinygrad.helpers import getenv, all_same, DEBUG
+from tinygrad.helpers import getenv, all_same, DEBUG, ceildiv
 from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
 from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8
 
@@ -72,7 +73,7 @@ def hk_fp8_atb_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, g_amax:Tensor|
     elif b.uop.axis == 2: inv, out_axis = Tensor.invalids(1, M, N // ndev, dtype=dtypes.bfloat16, device=a.device), 2
     elif a.uop.axis == 2: inv, out_axis = Tensor.invalids(1, M // ndev, N, dtype=dtypes.bfloat16, device=a.device), 1
     else: inv, out_axis, reduce_out = Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a.device), 0, True
-    out = Tensor(inv.uop.multi(out_axis), device=a.device)
+    out = Tensor(inv.uop.unshard(out_axis), device=a.device)
     dname = a.device[0]
   else:
     out = Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a.device)
@@ -107,6 +108,23 @@ def custom_hk_mxfp8_gemm(C:UOp, A:UOp, B:UOp, scale_A:UOp, scale_B:UOp, *extra:U
   return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
                                UOp(Ops.BINARY, arg=lib)))
 
+# ** MXFP4 GEMM custom kernel
+
+@functools.cache
+def custom_mxfp4_gemm(C:UOp, A:UOp, B:UOp, scale_a:UOp, scale_b:UOp, *extra:UOp, tile_m:int, tile_n:int) -> UOp:
+  from extra.gemm.gemm_mxfp4 import build_kernel
+  M, half_k = math.prod(A.shape[:-1]), A.shape[-1]
+  N, half_k_b = math.prod(B.shape[:-1]), B.shape[-1]
+  K = half_k * 2
+  assert half_k == half_k_b and math.prod(C.shape[:-1]) == M and C.shape[-1] == N
+  threads = UOp.special(256, "lidx0")
+  groups_x, groups_y = UOp.special(ceildiv(N, tile_n), "gidx0"), UOp.special(ceildiv(M, tile_m), "gidx1")
+  lds = UOp.placeholder((163840,), dtypes.uint8, 0, AddrSpace.LOCAL)
+  sink = UOp.sink(C.base, A.base, B.base, scale_a.base, scale_b.base, *(x.base for x in extra), lds, threads, groups_x, groups_y,
+                  arg=KernelInfo(f"custom_mxfp4_gemm_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K)))
+  insts = build_kernel(M, N, K, tile_m, tile_n)
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=x) for x in insts))))
+
 def quantize_mxfp8(x:Tensor) -> tuple[Tensor, Tensor, Tensor]:
   # 1x32 block scaling along the last axis
   *batch, K = x.shape
@@ -119,6 +137,50 @@ def quantize_mxfp8(x:Tensor) -> tuple[Tensor, Tensor, Tensor]:
   packed = mx_pack(e8) if len(batch) == 1 and scale_K % 4 == 0 else None
   return x_clamped.cast(FP8_DTYPE), e8, packed
 
+def _mxfp4_shuffle_weight(x:Tensor) -> Tensor:
+  # shuffle_weight(x, layout=(16, 16)) on the packed uint8 buffer.
+  if x.ndim == 3:
+    ndev, rows, half_k = x.shape
+    return x.reshape(ndev, rows//16, 16, half_k//32, 2, 16).permute(0, 1, 3, 4, 2, 5).reshape(ndev, rows, half_k).contiguous()
+  rows, half_k = x.shape
+  return x.reshape(rows//16, 16, half_k//32, 2, 16).permute(0, 2, 3, 1, 4).reshape(rows, half_k).contiguous()
+
+def _mxfp4_shuffle_scales(x:Tensor) -> Tensor:
+  # e8m0_shuffle: each 256x8 scale tile is arranged for the raw MFMA scale loads.
+  if x.ndim == 3:
+    ndev, rows, scale_k = x.shape
+    return x.reshape(ndev, rows//32, 2, 16, scale_k//8, 2, 4).permute(0, 1, 4, 6, 3, 5, 2).reshape(ndev, rows, scale_k).contiguous()
+  rows, scale_k = x.shape
+  return x.reshape(rows//32, 2, 16, scale_k//8, 2, 4).permute(0, 3, 5, 2, 4, 1).reshape(rows, scale_k).contiguous()
+
+def quantize_mxfp4(x:Tensor) -> tuple[Tensor, Tensor, Tensor]:
+  # OCP MXFP4: 1x32 blocks, E2M1 values packed low-nibble first, and E8M0 scales.
+  *batch, K = x.shape
+  rows = math.prod(batch)
+  assert x.ndim >= 2 and K % 256 == 0 and rows % 32 == 0, \
+    f"mxfp4 quantization needs rows%32 and K%256, got {x.shape}"
+  xb = x.float().reshape(*batch, K//32, 32)
+  amax = xb.abs().max(axis=-1)
+
+  # even scale rounding: round the fp32 significand before choosing 2^(floor(log2)-2).
+  amax_rounded = ((amax.bitcast(dtypes.uint32) + 0x200000) & 0xFF800000).bitcast(dtypes.float32)
+  scale_exp = (amax_rounded.maximum(2**-126).log2().floor() - 2).clamp(-127, 127)
+  e8 = (scale_exp + 127).cast(dtypes.uint8)
+  scaled = xb * (-scale_exp).exp2().reshape(*batch, K//32, 1)
+
+  mag = scaled.abs()
+  code = sum(x.cast(dtypes.uint8) for x in
+             (mag > .25, mag >= .75, mag > 1.25, mag >= 1.75, mag > 2.5, mag >= 3.5, mag > 5.0))
+  code = code | ((scaled < 0).cast(dtypes.uint8) << 3)
+  code = code.reshape(*batch, K)
+  packed = code[..., 0::2] | (code[..., 1::2] << 4)
+  if isinstance(x.device, tuple) and x.uop.axis == x.ndim-2 and x.shape[x.uop.axis] == len(x.device):
+    axis = x.uop.axis
+    order = (axis, *range(axis), *range(axis+1, e8.ndim))
+    e8_local = e8.permute(order)
+    return packed, e8, _mxfp4_shuffle_scales(e8_local.reshape(e8_local.shape[0], -1, K//32))
+  return packed, e8, _mxfp4_shuffle_scales(e8.reshape(rows, K//32))
+
 def mx_pack(e8:Tensor) -> Tensor:
   rows, scale_K = e8.shape
   return e8.reshape(rows, scale_K // 4, 4).bitcast(dtypes.uint32).reshape(rows, scale_K // 4).permute(1, 0).contiguous()
@@ -128,6 +190,11 @@ def _mx_block_scale(e8:Tensor) -> Tensor:
   rows, scale_K = e8.shape
   return (e8.cast(dtypes.float32) - 127.0).exp2().reshape(rows, scale_K, 1).expand(rows, scale_K, 32).reshape(rows, scale_K*32)
 
+def _mx_block_scale_3d(e8:Tensor) -> Tensor:
+  # batched (E, rows, scale_K) dequant scale 2^(e8-127) broadcast to (E, rows, scale_K*32)
+  E, rows, scale_K = e8.shape
+  return (e8.cast(dtypes.float32) - 127.0).exp2().reshape(E, rows, scale_K, 1).expand(E, rows, scale_K, 32).reshape(E, rows, scale_K*32)
+
 counters = {"used":0, "todos":[]}
 def todo(msg:str) -> bool: counters["todos"].append(msg); return False
 def _asm_gemm_report():
@@ -166,13 +233,13 @@ def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
   M, K = A.shape[0]*A.shape[1], A.shape[2]
   K2, N = B.shape[(1 if B.ndim == 3 else 0):]
   assert K == K2
-  m = UOp.range(M, 1, AxisType.LOOP)
-  n = UOp.range(N, 2, AxisType.LOOP)
+  m = UOp.range(M, 1)
+  n = UOp.range(N, 2)
   k = UOp.range(K, 0, AxisType.REDUCE)
-  mul = (A.flatten().index((m*UOp.const(dtypes.index, K)+k))*
-         B.flatten().index((k*UOp.const(dtypes.index, N)+n))).cast(dtypes.float32)
+  mul = (A.flatten().index((m*UOp.const(K)+k))*
+         B.flatten().index((k*UOp.const(N)+n))).cast(dtypes.float32)
   red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype)
-  store = C.flatten().index((m*UOp.const(dtypes.index, N)+n)).store(red).end(m, n)
+  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}'))
 
 # ** bf16 A @ B.T kernel in C
@@ -229,7 +296,7 @@ def hk_bf16_atb_gemm(a:Tensor, b:Tensor) -> Tensor:
     elif b.uop.axis == 2: inv, out_axis = Tensor.invalids(1, M, N // ndev, dtype=a.dtype, device=a.device), 2
     elif a.uop.axis == 2: inv, out_axis = Tensor.invalids(1, M // ndev, N, dtype=a.dtype, device=a.device), 1
     else: inv, out_axis, reduce_out = Tensor.invalids(1, M, N, dtype=a.dtype, device=a.device), 0, True
-    out = Tensor(inv.uop.multi(out_axis), device=a.device)
+    out = Tensor(inv.uop.unshard(out_axis), device=a.device)
     dname = a.device[0]
   else:
     out = Tensor.invalids(1, M, N, dtype=a.dtype, device=a.device)
@@ -239,7 +306,6 @@ def hk_bf16_atb_gemm(a:Tensor, b:Tensor) -> Tensor:
   if reduce_out: out = out.sum(0)
   return out.squeeze(0) if out.ndim == 3 else out
 
-
 # ** backward gemm, might use the asm gemm
 
 def custom_gemm_bw(gradient:UOp, kernel:UOp, n_scales:int=2, has_grad_amax:bool=False, has_w_post:bool=False):
@@ -275,10 +341,7 @@ def custom_gemm_bw(gradient:UOp, kernel:UOp, n_scales:int=2, has_grad_amax:bool=
       elif getenv("FUSED_GRAD_QUANTIZE", 0):
         grad_amax_t = Tensor(grad_amax_state, device=a.device)
         g_amax = grad_amax_t
-        g_fp8, _, new_grad_amax, _ = quantize_fp8_delayed(g_t, g_amax)
-        store_effect = next_grad_amax_state.store(new_grad_amax.uop)
-        assert g_fp8.uop.op is Ops.AFTER, f"expected AFTER, got {g_fp8.uop.op}"
-        g_fp8 = Tensor(g_fp8.uop.replace(src=g_fp8.uop.src + (store_effect,)), device=a.device)
+        g_fp8, _ = quantize_fp8_delayed(g_t, g_amax, Tensor(next_grad_amax_state, device=a.device))
       else:
         grad_amax_t = Tensor(grad_amax_state, device=a.device)
         g_amax = grad_amax_t
@@ -339,13 +402,29 @@ def custom_mx_gemm_bw(gradient:UOp, kernel:UOp, has_w_post:bool, w_stored:bool=F
   if wp is not None: grad_b = grad_b / wp.reshape(-1, 1)
   return (None, grad_a.uop, grad_b.uop) + tuple(None for _ in inputs[3:])
 
+# ** mxfp4 gemm backward
+
+def custom_mxfp4_gemm_bw(gradient:UOp, kernel:UOp):
+  # The raw kernel consumes quantized buffers, while the final two inputs retain the BF16 operands for STE gradients.
+  inputs = kernel.src[1:]  # (out, a_q, b_q, scale_a, scale_b, a, w)
+  assert len(inputs) == 7
+  a, w = Tensor(inputs[5], device=inputs[5].device), Tensor(inputs[6], device=inputs[6].device)
+  g = Tensor(gradient, device=a.device)[:a.shape[0]].cast(dtypes.bfloat16)
+  grad_a = asm_gemm(g, w, mxfp4=True)
+  a_flat, g_flat = a.reshape(-1, a.shape[-1]), g.reshape(-1, g.shape[-1])
+  grad_w = asm_gemm(g_flat.T, a_flat, mxfp4=True)
+  return (None, None, None, None, None, grad_a.uop, grad_w.uop)
+
 # ** main gemm function
 
 def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=None, grad_amax_state:Tensor|None=None,
              next_grad_amax_state:Tensor|None=None,
              w_post_scale:Tensor|None=None, mx:bool=False, mx_scales:tuple|None=None, mx_w_stored:bool=False, g_amax:Tensor|None=None,
-             a_pretranspose:Tensor|None=None) -> Tensor:
+             a_pretranspose:Tensor|None=None, mxfp4:bool=False) -> Tensor:
   assert can_use_asm_gemm(a, b), f"{counters['todos'][-1]}"
+  if mxfp4:
+    assert not mx and mx_scales is None, "mxfp4 owns quantization; mx/mx_scales are for mxfp8"
+    assert a.dtype == dtypes.bfloat16, f"cannot quantize {a.dtype} to mxfp4"
   counters["used"] += 1
   unfold_batch = a.ndim == 3 and isinstance(a.device, tuple) and a.uop.axis == 2 and b.uop.axis == 0
   if unfold_batch:
@@ -353,7 +432,7 @@ def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=N
     a = a.reshape(a.shape[0]*a.shape[1], a.shape[2])
   squeeze = a.ndim == 2
   if squeeze: a = a.unsqueeze(0)
-  out_dtype = dtypes.bfloat16 if a.dtype == FP8_DTYPE else a.dtype
+  out_dtype = dtypes.bfloat16 if a.dtype == FP8_DTYPE or mxfp4 else a.dtype
 
   batch, M, K = a.shape
   N = b.shape[1]
@@ -364,11 +443,11 @@ def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=N
 
   if is_multi:
     if n_sharded:
-      out = Tensor(Tensor.invalids(batch, M, N//len(a.device), dtype=out_dtype, device=a.device).uop.multi(2), device=a.device)
+      out = Tensor(Tensor.invalids(batch, M, N//len(a.device), dtype=out_dtype, device=a.device).uop.unshard(2), device=a.device)
     elif m_sharded:
-      out = Tensor(Tensor.invalids(batch, M, N, dtype=out_dtype, device=a.device).uop.multi(1), device=a.device)
+      out = Tensor(Tensor.invalids(batch, M, N, dtype=out_dtype, device=a.device).uop.unshard(1), device=a.device)
     else:
-      out = Tensor(Tensor.invalids(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=out_dtype, device=a.device).uop.multi(0),
+      out = Tensor(Tensor.invalids(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=out_dtype, device=a.device).uop.unshard(0),
                    device=a.device)
   else:
     out = Tensor.invalids(batch, M, N, dtype=out_dtype, device=a.device)
@@ -376,7 +455,21 @@ def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=N
   renderer = Device[dname:=(a.device[0] if is_multi else a.device)].renderer
   dname, arch = dname.split(":")[0], renderer.target.arch
   if arch.startswith("gfx950") and getenv("USE_ASM", 1):
-    if mx:
+    if mxfp4:
+      tile_m, tile_n = next((tm, tn) for tm, tn in ((256, 256), (192, 256), (128, 512)) if (batch*M) % tm == N % tn == 0)
+      fxn = functools.partial(custom_mxfp4_gemm, tile_m=tile_m, tile_n=tile_n)
+      w = b.T
+      if k_sharded:
+        ndev = len(a.device)
+        a_q, _, scale_a = quantize_mxfp4(a.reshape(batch, M, ndev, K))
+        b_q, _, scale_b = quantize_mxfp4(w.reshape(w.shape[0], ndev, K))
+        b_q = _mxfp4_shuffle_weight(b_q.permute(1, 0, 2))
+      else:
+        a_q, _, scale_a = quantize_mxfp4(a.reshape(batch*M, K))
+        b_q, _, scale_b = quantize_mxfp4(w)
+        a_q, b_q = a_q.reshape(batch, M, K//2).contiguous(), _mxfp4_shuffle_weight(b_q)
+      out = Tensor.custom_kernel(out, a_q, b_q, scale_a, scale_b, a, w, fxn=fxn, grad_fxn=custom_mxfp4_gemm_bw)[0]
+    elif mx:
       # mxfp8 1x32 block scaling
       if mx_scales is not None:
         a_si, a_e8, b_si, b_e8 = mx_scales

extra/gemm/dsp_int8_gemm.py

@@ -1,141 +0,0 @@
-#!/usr/bin/env python3
-"""Standalone correctness/throughput harness for the Hexagon HVX int8 GEMM."""
-import argparse
-
-import numpy as np
-
-from tinygrad import Device, dtypes
-from tinygrad.device import Buffer
-
-
-KERNEL = r"""
-typedef int int32x32 __attribute__((aligned(128),vector_size(128)));
-typedef unsigned char uchar4 __attribute__((aligned(4),vector_size(4)));
-typedef signed char char128 __attribute__((aligned(128),vector_size(128)));
-typedef unsigned char uchar128 __attribute__((aligned(128),vector_size(128)));
-typedef unsigned char uchar256 __attribute__((aligned(256),vector_size(256)));
-union V256 { uchar256 vec256; struct { uchar128 lo128, hi128; }; };
-
-__attribute__((noinline)) void gemm(unsigned char * restrict __attribute__((align_value(128))) out,
-                                    unsigned char * restrict __attribute__((align_value(128))) weight,
-                                    signed char * restrict __attribute__((align_value(128))) activation) {
-  for (int n = 0; n < 512; n++) {
-    int noff = n << 9;
-    for (int mb = 0; mb < 4; mb++) {
-      int moff = mb << 7;
-      int32x32 acc0 = __builtin_HEXAGON_V6_vd0_128B();
-      int32x32 acc1 = __builtin_HEXAGON_V6_vd0_128B();
-      int32x32 acc2 = __builtin_HEXAGON_V6_vd0_128B();
-      int32x32 acc3 = __builtin_HEXAGON_V6_vd0_128B();
-      for (int k4 = 0; k4 < 128; k4++) {
-        uchar4 w4 = *((uchar4 *)(weight + noff + (k4 << 2)));
-        int aoff = moff + (k4 << 11);
-        char128 x0 = *((char128 *)(activation + aoff));
-        char128 x1 = *((char128 *)(activation + aoff + 512));
-        char128 x2 = *((char128 *)(activation + aoff + 1024));
-        char128 x3 = *((char128 *)(activation + aoff + 1536));
-        union V256 s01, s23, slo, shi;
-        s01.vec256 = __builtin_HEXAGON_V6_vshufoeb_128B(x1, x0);
-        s23.vec256 = __builtin_HEXAGON_V6_vshufoeb_128B(x3, x2);
-        slo.vec256 = __builtin_HEXAGON_V6_vdealvdd_128B(s23.lo128, s01.lo128, 2);
-        shi.vec256 = __builtin_HEXAGON_V6_vdealvdd_128B(s23.hi128, s01.hi128, 2);
-        uchar128 w = __builtin_HEXAGON_V6_lvsplatw_128B(*((unsigned int *)&w4));
-        acc0 = __builtin_HEXAGON_V6_vrmpybusv_acc_128B(acc0, w, slo.lo128);
-        acc1 = __builtin_HEXAGON_V6_vrmpybusv_acc_128B(acc1, w, shi.lo128);
-        acc2 = __builtin_HEXAGON_V6_vrmpybusv_acc_128B(acc2, w, slo.hi128);
-        acc3 = __builtin_HEXAGON_V6_vrmpybusv_acc_128B(acc3, w, shi.hi128);
-      }
-      acc0 /= 1000; acc1 /= 1000; acc2 /= 1000; acc3 /= 1000;
-      uchar128 packed = __builtin_HEXAGON_V6_vpackhub_sat_128B(
-        __builtin_HEXAGON_V6_vpackwh_sat_128B(acc3, acc2),
-        __builtin_HEXAGON_V6_vpackwh_sat_128B(acc1, acc0));
-      packed = __builtin_HEXAGON_V6_vshuffb_128B(packed);
-      packed = __builtin_HEXAGON_V6_vshuffb_128B(packed);
-      *((uchar128 *)(out + noff + moff)) = packed;
-    }
-  }
-}
-
-struct dcvs_v2_req { int type; int pad; _Bool dcvs_enable; char dcvs_option; _Bool set_latency; int latency;
-  _Bool set_dcvs_params; short pad2; char target_corner; char min_corner; char max_corner; int pad3[3]; };
-typedef union { struct { void *pv; unsigned int len; } buf; struct { int fd; unsigned int offset; } dma; } remote_arg;
-int HAP_power_set(void *, void *);
-void *HAP_mmap(void *, int, int, int, int, long);
-int HAP_munmap(void *, int);
-unsigned long long HAP_perf_get_time_us(void);
-
-int entry(unsigned long long handle, unsigned int sc, remote_arg *pra) {
-  struct dcvs_v2_req req = {.type=7, .dcvs_enable=0, .set_latency=1, .latency=100,
-    .set_dcvs_params=1, .target_corner=6};
-  HAP_power_set((void *)handle, (void *)&req);
-  if ((sc >> 24) != 2) return 0;
-  int *sizes = (int *)pra[0].buf.pv, *offs = (int *)pra[1].buf.pv;
-  void *out = HAP_mmap(0, sizes[0], 3, 0, pra[3].dma.fd, 0) + offs[0];
-  void *weight = HAP_mmap(0, sizes[1], 3, 0, pra[4].dma.fd, 0) + offs[1];
-  void *activation = HAP_mmap(0, sizes[2], 3, 0, pra[5].dma.fd, 0) + offs[2];
-  unsigned long long start = HAP_perf_get_time_us();
-  gemm(out, weight, activation);
-  *(unsigned long long *)pra[2].buf.pv = HAP_perf_get_time_us() - start;
-  HAP_munmap(out-offs[0], sizes[0]); HAP_munmap(weight-offs[1], sizes[1]); HAP_munmap(activation-offs[2], sizes[2]);
-  return 0;
-}
-"""
-
-
-def main() -> None:
-  parser = argparse.ArgumentParser()
-  parser.add_argument("--iters", type=int, default=5)
-  parser.add_argument("--check", action="store_true")
-  parser.add_argument("--raw", action="store_true", help="store raw int32 accumulators without requantization")
-  args = parser.parse_args()
-  dev = Device["DSP"]
-  source = KERNEL
-  if args.raw:
-    source = source.replace(
-      "unsigned char * restrict __attribute__((align_value(128))) out,\n                                    unsigned char * restrict",
-      "int * restrict __attribute__((align_value(128))) out,\n                                    unsigned char * restrict", 1)
-    old = """      acc0 /= 1000; acc1 /= 1000; acc2 /= 1000; acc3 /= 1000;
-      uchar128 packed = __builtin_HEXAGON_V6_vpackhub_sat_128B(
-        __builtin_HEXAGON_V6_vpackwh_sat_128B(acc3, acc2),
-        __builtin_HEXAGON_V6_vpackwh_sat_128B(acc1, acc0));
-      packed = __builtin_HEXAGON_V6_vshuffb_128B(packed);
-      packed = __builtin_HEXAGON_V6_vshuffb_128B(packed);
-      *((uchar128 *)(out + noff + moff)) = packed;"""
-    new = """      int base = noff + moff;
-      *((int32x32 *)(out + base + 0)) = acc0;
-      *((int32x32 *)(out + base + 32)) = acc1;
-      *((int32x32 *)(out + base + 64)) = acc2;
-      *((int32x32 *)(out + base + 96)) = acc3;"""
-    if old not in source: raise RuntimeError("raw-kernel source pattern not found")
-    source = source.replace(old, new)
-  lib = dev.compiler.compile(source)
-  prg = dev.runtime("entry", lib)
-  rng = np.random.default_rng(0)
-  # Kernel contract is weight[N,K] and activation[K,M], both contiguous.
-  weight_np = rng.integers(0, 16, (512, 512), dtype=np.uint8)
-  activation_np = rng.integers(-8, 8, (512, 512), dtype=np.int8)
-  out_dtype = dtypes.int if args.raw else dtypes.uint8
-  bufs = [Buffer("DSP", 512*512, dt, preallocate=True) for dt in (out_dtype, dtypes.uint8, dtypes.int8)]
-  bufs[1].copyin(memoryview(weight_np).cast("B"))
-  bufs[2].copyin(memoryview(activation_np).cast("B"))
-  for _ in range(2): prg(*(x._buf for x in bufs), wait=True)
-  times = [prg(*(x._buf for x in bufs), wait=True) for _ in range(args.iters)]
-  best = min(times)
-  print(f"{2*512**3/best/1e9:.1f} GOPS ({best*1e3:.3f} ms)")
-  if args.check:
-    raw = bytearray(bufs[0].nbytes)
-    bufs[0].copyout(memoryview(raw))
-    expected_dot = weight_np.astype(np.int32) @ activation_np.astype(np.int32)
-    if args.raw:
-      got = np.frombuffer(raw, dtype=np.int32).reshape(512, 4, 4, 32).transpose(0, 1, 3, 2).reshape(512, 512)
-      expected = expected_dot
-      delta = np.abs(got.astype(np.int64)-expected.astype(np.int64))
-    else:
-      got = np.frombuffer(raw, dtype=np.uint8).reshape(512, 512)
-      expected = (expected_dot // 1000).clip(0, 255).astype(np.uint8)
-      delta = np.abs(got.astype(np.int16)-expected.astype(np.int16))
-    print(f"check={np.array_equal(got, expected)} max_abs={delta.max()} mismatches={np.count_nonzero(delta)}")
-    if not np.array_equal(got, expected): raise SystemExit(1)
-
-
-if __name__ == "__main__": main()

extra/gemm/ir3asm.py

@@ -1,440 +0,0 @@
-"""ir3 assembler for Adreno a6xx (A630).
-
-Constructs complete QCOM shader binaries from instruction listings.
-Uses a compiled "donor" kernel for the binary envelope (header, metadata,
-buffer descriptors, sampler info) and replaces the shader instructions
-and register counts.
-
-Encoding reference: derived from Mesa ir3 disassembly of known-good shaders.
-Instruction format: 64 bits (8 bytes) stored as two little-endian 32-bit words.
-"""
-import struct
-
-# ============================================================
-# HELPERS
-# ============================================================
-
-def _hreg(name):
-    """Parse 'hr3.z' -> half-register number 14."""
-    if isinstance(name, int): return name
-    r, c = name.replace('hr','').replace('r','').split('.')
-    return int(r) * 4 + 'xyzw'.index(c)
-
-def _freg(name):
-    """Parse 'r3.z' -> full-register number 14."""
-    if isinstance(name, int): return name
-    r, c = name.replace('r','').split('.')
-    return int(r) * 4 + 'xyzw'.index(c)
-
-def _pack(lo, hi):
-    return struct.pack('<II', lo & 0xFFFFFFFF, hi & 0xFFFFFFFF)
-
-# ============================================================
-# CAT0: FLOW CONTROL
-# ============================================================
-
-def NOP(rpt=0):
-    """(rptN)nop"""
-    return _pack(0, (rpt & 0x7F) << 8)
-
-def NOP_SS(rpt=0):
-    """(ss)(rptN)nop -- wait until prior instructions have consumed their sources."""
-    return _pack(0, 0x1000 | ((rpt & 0x7F) << 8))
-
-def END():
-    """end"""
-    return _pack(0, 0x03000000)
-
-def BR(offset, inv=True):
-    """br !p0.x, #offset  (inv=True means branch when predicate is FALSE)
-    offset is signed, relative to the branch instruction."""
-    return struct.pack('<iI', offset, 0x00900000 if inv else 0x00800000)
-
-def JUMP(offset):
-    """jump #offset. Offset is signed, relative to the jump instruction."""
-    return struct.pack('<iI', offset, 0x01000000)
-
-# ============================================================
-# CAT1: MOVE / CONVERT
-# ============================================================
-
-def MOV_S32(dst, imm, sy=False):
-    """(sy?)mov.s32s32 rDST, #imm"""
-    return _pack(imm, ((0x30 if sy else 0x20) << 24) | (0x55 << 16) | (0x40 << 8) | (_freg(dst) & 0xFF))
-
-def MOV_F32(dst, src, rpt=0, sy=False, ss=False, r=False):
-    """(sy?)(ss?)(rptN?)mov.f32f32 rDST, (r?)rSRC"""
-    return _pack(_freg(src), (0x30044000 if sy else 0x20044000) | (0x1000 if ss else 0) |
-                 (0x800 if r else 0) | ((rpt & 0x7F) << 8) | (_freg(dst) & 0xFF))
-
-def MOV_H(dst, src, rpt=0, r=False):
-    """(rptN?)mov.f16f16 hrDST, (r?)hrSRC."""
-    return _pack(_hreg(src), 0x20000000 | (0x800 if r else 0) | ((rpt & 0x7F) << 8) | (_hreg(dst) & 0xFF))
-
-def MOV_H_IMM(dst, imm_u16=0, rpt=0):
-    """(rptN?)mov.f16f16 hrDST, h(imm) -- imm is raw fp16 bits (0=zero, 0x3c00=1.0)."""
-    return _pack(imm_u16, 0x20400000 | ((rpt & 0x7F) << 8) | (_hreg(dst) & 0xFF))
-
-def COV_F16F32(dst, src, sy=False, rpt=0, r=False):
-    """(sy?)(rptN?)cov.f16f32 rDST, (r?)hrSRC"""
-    return _pack(_hreg(src), ((0x30 if sy else 0x20) << 24) | 0x004000 | (0x800 if r else 0) |
-                 ((rpt & 0x7f) << 8) | (_freg(dst) & 0xFF))
-
-# ============================================================
-# CAT2: INTEGER / FLOAT ALU (2 operands)
-# ============================================================
-
-def ADD_S(dst, src1, imm, nop=0, ss=False):
-    """(ss?)(nopN?)add.s rDST, rSRC1, #imm  (signed immediate add)"""
-    d, s = _freg(dst), _freg(src1)
-    hi_base = 0x42300000 | (d & 0xFF)
-    if nop > 0:
-        hi_base = (hi_base & 0xFF00FFFF) | (0x38 << 16) | ((nop & 0x7) << 11)
-    if ss: hi_base |= 0x1000
-    lo = ((0x27 if imm < 0 else 0x20) << 24) | ((imm & 0xFF) << 16) | (s & 0xFF)
-    return _pack(lo, hi_base)
-
-def ADD_S_REG(dst, src1, src2, nop=0):
-    """(nopN?)add.s rDST, rSRC1, rSRC2"""
-    d, s1, s2 = _freg(dst), _freg(src1), _freg(src2)
-    hi_base = 0x42300000 | (d & 0xFF)
-    if nop > 0:
-        hi_base = (hi_base & 0xFF00FFFF) | (0x38 << 16) | ((nop & 0x7) << 11)
-    return _pack(((s2 & 0xFF) << 16) | (s1 & 0xFF), hi_base)
-
-def ADD_S_CONST_REG(dst, const_src, src2, nop=0):
-    """(nopN?)add.s rDST, cSRC1, rSRC2"""
-    d, c1, s2 = _freg(dst), _freg(const_src.replace('c', 'r', 1)), _freg(src2)
-    hi_base = 0x42300000 | (d & 0xFF)
-    if nop > 0:
-        hi_base = (hi_base & 0xFF00FFFF) | (0x38 << 16) | ((nop & 0x7) << 11)
-    return _pack(((s2 & 0xFF) << 16) | 0x1000 | (c1 & 0xFF), hi_base)
-
-def ADD_F(dst, src1, src2, rpt=0, r1=False, r2=False, sy=False):
-    """Vector-capable add.f; full registers use the same scalar indices."""
-    hi = (0x50100000 if sy else 0x40100000) | (0x800 if r1 else 0) | (0x80000 if r2 else 0)
-    return _pack(((_hreg(src2) & 0xFF) << 16) | (_hreg(src1) & 0xFF),
-                 hi | ((rpt & 0x7f) << 8) | (_hreg(dst) & 0xFF))
-
-def SUB_F(dst, src1, src2, rpt=0, r1=False, r2=False, sy=False):
-    """Vector-capable add.f with a negated second source."""
-    hi = (0x50100000 if sy else 0x40100000) | (0x800 if r1 else 0) | (0x80000 if r2 else 0)
-    return _pack(0x40000000 | ((_hreg(src2) & 0xFF) << 16) | (_hreg(src1) & 0xFF),
-                 hi | ((rpt & 0x7f) << 8) | (_hreg(dst) & 0xFF))
-
-def ADD_U(dst, src1_const, src2):
-    """add.u rDST, cSRC1, rSRC2  -- src1 is constant register"""
-    # From: 42100008_00031050 = add.u r2.x, c20.x, r0.w
-    return _pack((_freg(src2) << 16) | 0x1050, 0x42100000 | (_freg(dst) & 0xFF))
-
-def CMPS_S_EQ(src1, imm, nop=0):
-    """(nopN?)cmps.s.eq p0.x, rSRC1, #imm"""
-    hi = 0x42b400f8
-    if nop > 0:
-        hi = (hi & 0xFF00FFFF) | (0xb4 << 16) | ((nop & 0x7) << 11)
-    # Integer immediates use the low bits of the source descriptor for bits 8+.
-    # Keeping this fixed at 0x20 silently truncated loop bounds above 255.
-    lo = ((0x20 | (imm >> 8)) << 24) | ((imm & 0xFF) << 16) | (_freg(src1) & 0xFF)
-    return _pack(lo, hi)
-
-def CMPS_S_LT_REG(src1, src2, nop=0):
-    """(nopN?)cmps.s.lt p0.x, rSRC1, rSRC2"""
-    hi = 0x42b000f8
-    if nop > 0: hi = (hi & 0xFF00FFFF) | (0xb0 << 16) | ((nop & 0x7) << 11)
-    return _pack(((_freg(src2) & 0xff) << 16) | (_freg(src1) & 0xff), hi)
-
-def SHL_B(dst, src, imm, jp=False, ss=False, nop=0):
-    """(ss?)(jp?)(nopN?)shl.b rDST, rSRC, #imm"""
-    hi = (0x4ed00000 if jp else 0x46d00000) | (_freg(dst) & 0xFF)
-    if ss: hi |= 1 << 12
-    if nop & 1: hi |= 1 << 11
-    if nop & 2: hi |= 1 << 19
-    return _pack((0x20 << 24) | ((imm & 0xFF) << 16) | (_freg(src) & 0xFF), hi)
-
-def SHR_B(dst, src, imm):
-    """shr.b rDST, rSRC, #imm"""
-    return _pack((0x20 << 24) | ((imm & 0xFF) << 16) | (_freg(src) & 0xFF), 0x46f00000 | (_freg(dst) & 0xFF))
-
-def AND_B(dst, src, imm, nop=0):
-    """(nopN?)and.b rDST, rSRC, #imm"""
-    hi = 0x43900000 | (_freg(dst) & 0xFF)
-    if nop & 1: hi |= 1 << 11
-    if nop & 2: hi |= 1 << 19
-    return _pack((0x20 << 24) | ((imm & 0xFF) << 16) | (_freg(src) & 0xFF), hi)
-
-def AND_B_CONST(dst, src, const_src, nop=0):
-    """(nopN?)and.b rDST, rSRC, cSRC2"""
-    d, s, c = _freg(dst), _freg(const_src.replace('c', 'r', 1)), _freg(src)
-    hi = 0x43900000 | (d & 0xFF)
-    if nop & 1: hi |= 1 << 11
-    if nop & 2: hi |= 1 << 19
-    return _pack((0x10 << 24) | ((c & 0xFF) << 16) | (s & 0xFF), hi)
-
-def OR_B(dst, src, imm, ss=False):
-    """(ss?)or.b rDST, rSRC, #imm"""
-    return _pack((0x20 << 24) | ((imm & 0xFF) << 16) | (_freg(src) & 0xFF),
-                 0x43b00000 | (0x1000 if ss else 0) | (_freg(dst) & 0xFF))
-
-def CMPS_U_LT(dst, src1, src2_const):
-    """cmps.u.lt rDST, rSRC1, cSRC2"""
-    # From: 42900010_10500008 = cmps.u.lt r4.x, r2.x, c20.x
-    return _pack(0x10500000 | (_freg(src1) & 0xFF), 0x42900000 | (_freg(dst) & 0xFF))
-
-def CMPS_U_LT_REG(dst, src1, src2, sy=False):
-    """(sy?)cmps.u.lt rDST, rSRC1, rSRC2"""
-    hi = (0x52900000 if sy else 0x42900000) | (_freg(dst) & 0xff)
-    return _pack(((_freg(src2) & 0xff) << 16) | (_freg(src1) & 0xff), hi)
-
-# ============================================================
-# CAT3: MAD (3 operands)
-# ============================================================
-
-def MAD_F16(dst, src1, src2, src3, rpt=0, sy=False, r=False, r1=False, r3=False):
-    """(sy?)(rptN?)mad.f16 hrDST, (r1?)hrSRC1, (r?)hrSRC2, (r?)hrSRC3
-    When rpt>0, r1 auto-increments src1 and r auto-increments src2/src3/dst."""
-    d, s1, s2, s3 = _hreg(dst), _hreg(src1), _hreg(src2), _hreg(src3)
-    hi = ((0x73 if sy else 0x63) << 24) | ((s2 >> 1) << 16) | ((((s2 & 1) << 7) | (0x08 if r1 else 0) | (rpt & 0x7F)) << 8) | (d & 0xFF)
-    lo = (0x20000000 if (r or r3) else 0) | ((s3 & 0xFF) << 16) | (0x8000 if r else 0) | (s1 & 0xFF)
-    return _pack(lo, hi)
-
-def MAD_F32(dst, src1, src2, src3, rpt=0, sy=False, r=False, r1=False):
-    """(sy?)(rptN?)mad.f32 rDST, rSRC1, (r?)rSRC2, (r?)rSRC3"""
-    d, s1, s2, s3 = _freg(dst), _freg(src1), _freg(src2), _freg(src3)
-    hi = ((0x73 if sy else 0x63) << 24) | (0x80 << 16) | ((s2 >> 1) << 16) | \
-         ((((s2 & 1) << 7) | (0x08 if r1 else 0) | (rpt & 0x7F)) << 8) | (d & 0xFF)
-    lo = (0x20000000 if r else 0) | ((s3 & 0xFF) << 16) | (0x8000 if r else 0) | (s1 & 0xFF)
-    return _pack(lo, hi)
-
-def DP4ACC(dst, src1, src2, src3, sy=False, mixed=False, signed=None):
-    """A6xx packed 4x int8 dot product accumulated into a full int32 register.
-
-    ``mixed=False`` selects unsigned*unsigned. ``mixed=True`` selects the
-    pre-A7xx mixed signedness mode used by A630 (signed lhs, unsigned rhs).
-    The instruction has no repeat form on this generation.
-    """
-    if signed is not None: mixed = signed
-    d, s1, s2, s3 = _freg(dst), _freg(src1), _freg(src2), _freg(src3)
-    hi = ((0x76 if sy else 0x66) << 24) | (0x80 << 16) | ((s2 >> 1) << 16)
-    hi |= (((s2 & 1) << 7) | 0x40) << 8
-    # AL-OP is bit 13 and the pre-A7 signed/unsigned selector is bit 14.
-    lo = ((s3 & 0xff) << 16) | 0x2000 | (0x4000 if mixed else 0) | (s1 & 0xff)
-    return _pack(lo, hi | (d & 0xff))
-
-# ============================================================
-# CAT3: SHLG / SHRM (shift with merge)
-# ============================================================
-
-def SHLG(dst, imm, src1, src2, nop=0):
-    """(nopN?)shlg rDST, #imm, rSRC1, rSRC2.
-
-    This covers the packed image-coordinate forms emitted by the a6xx compiler
-    for GEMM kernels. The low byte encodes the shift immediate and bits 23:16
-    encode src2; the remaining source mode bits are pattern-specific.
-    """
-    d, s1, s2 = _freg(dst), _freg(src1), _freg(src2)
-    if (s1, s2) in ((_freg('r0.y'), _freg('r0.z')), (_freg('r0.z'), _freg('r0.x'))):
-        hi_mid, lo_mid = 0x80, 0xb0
-        if (s1, s2) == (_freg('r0.z'), _freg('r0.x')): hi_mid = 0x81
-    elif (s1, s2) in ((_freg('r0.w'), _freg('r0.x')), (_freg('r0.w'), _freg('r0.y'))):
-        hi_mid, lo_mid = 0x81, 0x30
-    else:
-        raise ValueError('unsupported SHLG source pattern %s, %s' % (src1, src2))
-    hi = (0x65 << 24) | (hi_mid << 16) | (0x84 << 8) | (d & 0xFF)
-    lo = ((s2 & 0xFF) << 16) | (lo_mid << 8) | (imm & 0xFF)
-    return _pack(lo, hi)
-
-def SHLG_IMM(dst, imm, src, merge):
-    """shlg rDST, #imm, rSRC, #merge.
-
-    Observed in compiler address generation for widened column stores, e.g.
-    65b08402_10803002 = shlg r0.z, 2, r24.y, 128.
-    """
-    d, s = _freg(dst), _freg(src)
-    hi = (0x65 << 24) | ((0x80 | ((s >> 1) & 0x7f)) << 16) | (0x84 << 8) | (d & 0xff)
-    lo = (0x10 << 24) | ((merge & 0xffff) << 16) | 0x3000 | (imm & 0xff)
-    return _pack(lo, hi)
-
-def SHRM(dst, shift, src1, merge):
-    """shrm rDST, #shift, rSRC1, #merge.
-
-    Observed compiler form for subgroup row offsets, e.g.
-    64000402_100c3003 = shrm r0.z, 3, r0.x, 12.
-    """
-    d, s1 = _freg(dst), _freg(src1)
-    if s1 != _freg('r0.x'):
-        raise ValueError('unsupported SHRM source %s' % src1)
-    hi = 0x64000400 | (d & 0xFF)
-    lo = (0x10 << 24) | ((merge & 0xFF) << 16) | 0x3000 | (shift & 0xFF)
-    return _pack(lo, hi)
-
-# ============================================================
-# CAT5: TEXTURE (ISAM)
-# ============================================================
-
-def ISAM_F16(dst, coord, tex=0, samp=0, sy=False, wrmask=0xf):
-    """isam.1d (f16)(xyzw) hrDST, rCOORD, s#SAMP, t#TEX
-    dst: first half-register of the xyzw quad
-    coord: full-register containing the (int2) coordinate pair"""
-    return _pack((tex * 2) << 24 | ((samp & 0x7) << 21) | (_freg(coord) * 2 + 1),
-                 (0xb0000000 if sy else 0xa0000000) | ((wrmask & 0xf) << 8) | (_hreg(dst) & 0xFF))
-
-def ISAM_F32(dst, coord, tex=0, samp=0):
-    """isam.1d (f32)(xyzw) rDST, rCOORD, s#SAMP, t#TEX"""
-    return _pack((tex * 2) << 24 | ((samp & 0x7) << 21) | (_freg(coord) * 2 + 1), 0xa0001f00 | (_freg(dst) & 0xFF))
-
-def ISAM_U32(dst, coord, tex=0, samp=0):
-    """isam.1d (u32)(xyzw) rDST, rCOORD, s#SAMP, t#TEX"""
-    return _pack((tex * 2) << 24 | ((samp & 0x7) << 21) | (_freg(coord) * 2 + 1), 0xa0003f00 | (_freg(dst) & 0xFF))
-
-def COV_S32S16(dst, src, rpt=0, r=False, sy=False):
-    """cov.s32s16 hDST, rSRC, optionally repeating over four packed lanes."""
-    hi = (0x30150000 if sy else 0x20150000) | ((rpt & 0x7) << 8) | (0x800 if r else 0) | (_hreg(dst) & 0xff)
-    return _pack(_freg(src) & 0xff, hi)
-
-def SHRG_H(dst, src, shift=16, rpt=0, r=False):
-    """shrg hDST, #shift, rSRC, #0 for extracting packed high half lanes."""
-    s = _freg(src)
-    hi = 0x65004400 | (((s >> 1) & 0x7f) << 16) | ((rpt & 0x7) << 8) | (_hreg(dst) & 0xff)
-    lo = 0x10003000 | (0x8000 if r else 0) | (shift & 0xff)
-    return _pack(lo, hi)
-
-def QUAD_BRCST(dst, src, idx, typ=3, wrmask=1, sy=False, jp=False):
-    """quad_shuffle.brcst.{typ} DST, SRC, IDX"""
-    half = typ in (0, 2, 4, 6)
-    d = _hreg(dst) if half else _freg(dst)
-    s = _hreg(src) if half else _freg(src)
-    i = _hreg(idx) if half else _freg(idx)
-    lo = (0 if half else 1) | ((s & 0xff) << 1) | ((i & 0xff) << 9)
-    hi = 0xa7e00000 | ((typ & 7) << 12) | ((wrmask & 0xf) << 8) | (d & 0xff)
-    if jp: hi |= 1 << 27
-    if sy: hi |= 1 << 28
-    return _pack(lo, hi)
-
-# ============================================================
-# CAT6: LOAD / STORE
-# ============================================================
-
-def STG_F16(addr, data_hreg, count=4, sy=False):
-    """(sy?)stg.f16 g[rADDR], hrDATA, count"""
-    # Encoding from compiled kernels:
-    # c0c01100_04800000 = stg.f16 g[r2.x], hr0.x, 4
-    # c0c01500_04800008 = stg.f16 g[r2.z], hr1.x, 4
-    # c0c01900_04800010 = stg.f16 g[r3.x], hr2.x, 4
-    # c0c01d00_04800018 = stg.f16 g[r3.z], hr3.x, 4
-    # hi pattern: c0c0XX00 where XX encodes the address register
-    # lo pattern: 048000YY where YY encodes the data register
-    a, d = _freg(addr), _hreg(data_hreg)
-    # addr encoding: r2.x=8 -> 0x11, r2.z=10 -> 0x15, r3.x=12 -> 0x19, r3.z=14 -> 0x1d
-    # Pattern: (addr * 2 + 1) = 17,21,25,29 = 0x11,0x15,0x19,0x1d
-    addr_enc = a * 2 + 1
-    hi = (0xd0c00000 if sy else 0xc0c00000) | (addr_enc << 8)
-    lo = 0x04800000 | ((d * 2) & 0xFF)
-    return _pack(lo, hi)
-
-def STG_U32(addr, data_reg, count=1, sy=False):
-    """(sy?)stg.u32 g[rADDR], rDATA, count"""
-    a, d = _freg(addr), _freg(data_reg)
-    hi = (0xd0c00000 if sy else 0xc0c00000) | (3 << 17) | ((a * 2 + 1) << 8)
-    lo = ((count & 0x7) << 24) | 0x00800000 | ((d << 1) & 0x1FE)
-    return _pack(lo, hi)
-
-def STG_F32(addr, data_reg, count=4, sy=False):
-    """(sy?)stg.f32 g[rADDR], rDATA, count"""
-    a, d = _freg(addr), _freg(data_reg)
-    hi = (0xd0c00000 if sy else 0xc0c00000) | (1 << 17) | ((a * 2 + 1) << 8)
-    lo = ((count & 0x7) << 24) | 0x00800000 | ((d << 1) & 0x1FE)
-    return _pack(lo, hi)
-
-def STIB_F32(data_reg, coord_reg, sy=False):
-    """Typed 2D image store of float4 data to integer (x,y) coordinates."""
-    hi = (0xd0220000 if sy else 0xc0220000) | (_freg(data_reg) & 0xff)
-    lo = ((_freg(coord_reg) & 0xff) << 24) | 0x00677a00
-    return _pack(lo, hi)
-
-def GETFIBERID(dst):
-    """getfiberid.u32 rDST"""
-    return _pack(0x00c98000, 0xc0260000 | (_freg(dst) & 0xff))
-
-def SHFL(dst, src, idx, mode=7, typ=2, sy=False, jp=False):
-    """shfl.{mode}.{typ} DST, SRC, IDX
-
-    mode: xor=1, up=2, down=3, rup=6, rdown=7.
-    typ: f16=0, f32=1, u16=2, u32=3, s16=4, s32=5.
-    idx can be an immediate int or a full register. For half types, dst/src are
-    half-register indices; SRC2 is always a full register/immediate per Mesa.
-    """
-    d = _hreg(dst) if typ in (0, 2, 4, 6) else _freg(dst)
-    s = _hreg(src) if typ in (0, 2, 4, 6) else _freg(src)
-    if isinstance(idx, int):
-        idx_im, idx_bits = 1, idx & 0xff
-    else:
-        idx_im, idx_bits = 0, _freg(idx) & 0xff
-    lo = ((s & 0xff) << 1) | (idx_im << 23) | (idx_bits << 24)
-    hi = (0xc0000000 | (0x1b << 22) | (2 << 20) | ((typ & 7) << 17) |
-          ((mode & 7) << 13) | (d & 0xff))
-    if jp: hi |= 1 << 27
-    if sy: hi |= 1 << 28
-    return _pack(lo, hi)
-
-# ============================================================
-# CAT3 SPECIAL: SAD.S32
-# ============================================================
-
-def SAD_S32(dst, src1_const, src2, src3, nop=0):
-    """(nopN?)sad.s32 rDST, cSRC1, (neg)rSRC2, rSRC3"""
-    # From: 67888009_40101051 = sad.s32 r2.y, c20.y, (neg)r4.y, r4.x
-    d, s2, s3 = _freg(dst), _freg(src2), _freg(src3)
-    hi_src2 = 0x80 | ((s2 >> 1) & 0xF)
-    # Observed nop3 form uses 0x88 in the third byte; plain sad.s32 uses 0x80.
-    hi_nop = 0x88 if nop > 0 else 0x80
-    hi = (0x67 << 24) | (hi_src2 << 16) | (hi_nop << 8) | (d & 0xFF)
-    lo = 0x40000000 | (s3 << 16) | 0x1051
-    return _pack(lo, hi)
-
-# ============================================================
-# BINARY ENVELOPE
-# ============================================================
-
-def get_envelope(dev, src):
-    """Compile an OpenCL kernel and return the binary as a mutable envelope."""
-    lib = bytearray(dev.compiler.compile_cached(src))
-    img_off = struct.unpack_from('<I', lib, 0xc0)[0]
-    img_sz = struct.unpack_from('<I', lib, 0x100)[0]
-    reg_off = struct.unpack_from('<I', lib, 0x34)[0]
-    return lib, img_off, img_sz, reg_off
-
-def inject(lib, img_off, img_sz, reg_off, shader_bytes, fregs, hregs, mergedregs=None):
-    """Replace shader binary and register counts in the envelope."""
-    lib = bytearray(lib)
-    shader = bytearray(shader_bytes)
-    if len(shader) > img_sz:
-        raise ValueError(f"shader is {len(shader)} bytes but donor image is only {img_sz} bytes")
-    # Pad to original size
-    while len(shader) < img_sz:
-        shader += NOP()
-    lib[img_off:img_off+img_sz] = shader[:img_sz]
-    if mergedregs is True: fregs |= 1 << 31
-    if mergedregs is False: hregs |= 1 << 31
-    struct.pack_into('<I', lib, reg_off + 0x14, fregs)
-    struct.pack_into('<I', lib, reg_off + 0x18, hregs)
-    return bytes(lib)
-
-def disasm(shader_bytes, gpu_id=630):
-    """Disassemble shader binary using Mesa's ir3_isa_disasm."""
-    import ctypes, tempfile
-    from tinygrad.runtime.autogen import mesa
-    from tinygrad.helpers import data64
-    with tempfile.TemporaryFile('w+', buffering=1) as tf:
-        @ctypes.CFUNCTYPE(None, ctypes.c_void_p, ctypes.c_uint32, ctypes.c_void_p)
-        def hd(data, n, instr):
-            fst, snd = data64(ctypes.cast(instr, ctypes.POINTER(ctypes.c_uint64)).contents.value)
-            print(f"{n:04} [{fst:08x}_{snd:08x}] ", end="", flush=True, file=tf)
-        libc = ctypes.CDLL(None)
-        libc.setlinebuf(fp:=ctypes.cast(libc.fdopen(tf.fileno(), b"w"), ctypes.POINTER(mesa.struct__IO_FILE)))
-        mesa.ir3_isa_disasm(bytes(shader_bytes), len(shader_bytes), fp, mesa.struct_isa_decode_options(gpu_id, True, 0, True, pre_instr_cb=hd))
-        tf.seek(0)
-        return tf.read()
-
-def assemble(instr_list):
-    """Assemble a list of instruction bytes into a shader binary."""
-    return b''.join(instr_list)

extra/gemm/metal_uop_matmul.py

@@ -1,5 +1,5 @@
 from tinygrad import UOp, dtypes
-from tinygrad.uop.ops import AxisType, Ops, KernelInfo, AddrSpace
+from tinygrad.uop.ops import AxisType, KernelInfo, AddrSpace
 from extra.gemm.amd_uop_matmul import test_matmul
 
 N = 2048
@@ -27,11 +27,8 @@ def hand_spec_tc_cores():
   acc = acc[0].set(0.0)
   acc = acc[1].set(0.0)
 
-  # TODO: make this simple
-  wmma_arg = ('WMMA_8_8_8_float_float', (8, 8, 8), dtypes.float, dtypes.float, 'METAL', 32, (((3, 2),), ((3, 2),), ((3, 2),)), ())
-
   acc_load = UOp.stack(acc.after(gk)[0], acc.after(gk)[1])
-  out = UOp(Ops.WMMA, dtypes.float, (a_tc, b_tc, acc_load), arg=wmma_arg)
+  out = UOp.wmma(a_tc, b_tc, acc_load, (8, 8, 8), 'METAL', 32)
 
   end_loop = UOp.group(*[acc[i].store(out.index(i)) for i in range(2)]).end(gk)
 

extra/gemm/mi350x_uop_matmul.py

@@ -6,7 +6,7 @@ os.environ["AMD_LLVM"] = "0"
 from tinygrad import Tensor, Context, dtypes, UOp, GlobalCounters
 from tinygrad.helpers import DEBUG, getenv
 from tinygrad.dtype import AddrSpace
-from tinygrad.uop.ops import AxisType, KernelInfo, Ops
+from tinygrad.uop.ops import AxisType, KernelInfo
 
 WARP_SIZE = 64
 
@@ -79,7 +79,7 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
   # this is the big accumulator
   acc = UOp.placeholder((BLOCK_N//TC_N, BLOCK_M//TC_M//WARPGROUP_SIZE), dtypes.float, 0, AddrSpace.REG)
   assert acc.size*WARP_SIZE*WARPGROUP_SIZE*4 == BLOCK_M*BLOCK_N
-  acc = acc[init_l:=UOp.range(acc.size, 500)].set(UOp.const(dtypes.float, (0.0,)*4), end=init_l)
+  acc = acc[init_l:=UOp.range(acc.size, 500)].set(UOp.const((0.0,)*4, dtypes.float), end=init_l)
 
   # create locals (note A is permuted, and the stride is changed to avoid bank conflicts)
   def make_locals(slot) -> tuple[UOp, UOp]:
@@ -137,8 +137,7 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
     acc_load = acc_after[N_inner_loop, M_inner_loop]
 
     # do WMMA
-    wmma_arg = ('WMMA_16_16_32_half_float', (16, 16, 32), dtypes.half, dtypes.float, 'AMD', 64, ((), (), ((3, 2), (2, 2))), ())
-    out = UOp(Ops.WMMA, dtypes.float, (Ar[M_inner_loop], Br[N_inner_loop], acc_load), arg=wmma_arg)
+    out = UOp.wmma(Ar[M_inner_loop], Br[N_inner_loop], acc_load, (16, 16, 32), 'AMD', 64)
 
     # store back the acc
     acc_store = acc[N_inner_loop, M_inner_loop].store(out)
@@ -193,8 +192,7 @@ acc = acc[init_l:=UOp.range(4, 1)].set(0.0, end=init_l)
 
 # do the wmma
 acc_load = UOp.stack(*[acc.after(K_loop)[i] for i in range(4)])
-wmma_arg = ('WMMA_16_16_32_half_float', (16, 16, 32), dtypes.half, dtypes.float, 'AMD', 64, ((), (), ((3, 2), (2, 2))), ())
-out = UOp(Ops.WMMA, dtypes.float, (A_in, B_in, acc_load), arg=wmma_arg)
+out = UOp.wmma(A_in, B_in, acc_load, (16, 16, 32), 'AMD', 64)
 
 # store back the acc
 acc = acc.after(UOp.group(*[acc[i].store(out.index(i)) for i in range(4)]).end(K_loop))

extra/gemm/mi350x_uop_matmul_2.py

@@ -6,7 +6,7 @@ os.environ["AMD_LLVM"] = "0"
 from tinygrad import Tensor, Context, dtypes, UOp, GlobalCounters
 from tinygrad.helpers import DEBUG, getenv
 from tinygrad.dtype import AddrSpace
-from tinygrad.uop.ops import sint, AxisType, KernelInfo, Ops
+from tinygrad.uop.ops import AxisType, KernelInfo
 
 WARP_SIZE = 64
 
@@ -29,7 +29,7 @@ TID_SIZE = WARPGROUP_SIZE*WARP_SIZE
 
 def copy(dest:UOp, src:UOp, rng:int, set=False, upcast=()):
   assert dest.shape == src.shape
-  rngs = [UOp.range(s, rng+i, AxisType.UPCAST if i in upcast else AxisType.LOOP) for i,s in enumerate(src.shape)]
+  rngs = [UOp.range(s, rng+i, AxisType.UPCAST if i in upcast else AxisType.WEAK) for i,s in enumerate(src.shape)]
   copy = dest[*rngs].store(src[*rngs]).end(*rngs)
   return dest.after(copy) if set else copy
 
@@ -60,8 +60,7 @@ def compute_on_locals(acc:UOp, Asl:UOp, Bsl:UOp, rng:int, afters:tuple[UOp, ...]
   acc_load = acc_after[N_inner_loop, M_inner_loop]
 
   # do WMMA
-  wmma_arg = ('WMMA_16_16_32_half_float', (16, 16, 32), dtypes.half, dtypes.float, 'AMD', 64, ((), (), ((3, 2), (2, 2))), ())
-  out = UOp(Ops.WMMA, dtypes.float, (Ar[M_inner_loop], Br[N_inner_loop], acc_load), arg=wmma_arg)
+  out = UOp.wmma(Ar[M_inner_loop], Br[N_inner_loop], acc_load, (16, 16, 32), 'AMD', 64)
 
   # store back the acc
   acc_store = acc[N_inner_loop, M_inner_loop].store(out)

extra/gemm/moe_gemm.py

@@ -0,0 +1,111 @@
+import functools, pathlib
+from tinygrad import Tensor, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.renderer import Estimates
+from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
+from extra.gemm.cdna_asm_gemm import quantize_mxfp8, _mx_block_scale, _mx_block_scale_3d
+
+@functools.cache
+def custom_hk_grouped_mxfp8_gemm(C:UOp, A:UOp, B:UOp, scale_A:UOp, scale_B:UOp, *extra:UOp, dname:str, n_experts:int) -> UOp:
+  M, K = A.shape
+  E, N, K2 = B.shape
+  assert K == K2, f"{A.shape} {B.shape}"
+  assert E == n_experts, f"{E} != {n_experts}"
+  threads = UOp.special(64 * 8, "lidx0")
+  workgroups = UOp.special((M // 256) * (N // 256), "gidx0")
+  sink_inputs = (C.base, A.base, B.base, scale_A.base, scale_B.base, extra[0].base, extra[1].base, extra[2].base, threads, workgroups)
+  sink = UOp.sink(*sink_inputs,
+                  arg=KernelInfo(f"hk_grouped_mxfp8_gemm_{E}_{M}_{N}_{K}",
+                                 estimates=Estimates(ops=2*M*N*K, mem=(M*K+E*N*K)*A.dtype.itemsize+M*N*C.dtype.itemsize)))
+  kittens_path = pathlib.Path(__file__).parent.parent/"thunder"/"amd"
+  src = (kittens_path/"grouped_mxfp8_gemm.cpp").read_text()
+  lib = HIPCCCompiler("gfx950", [f"-I{(kittens_path/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-ffast-math",
+                                 "-DHIP_ENABLE_WARP_SYNC_BUILTINS", f"-DGEMM_M={M}", f"-DGEMM_N={N}", f"-DGEMM_K={K}",
+                                 f"-DGEMM_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)))
+
+@functools.cache
+def custom_hk_grouped_mxfp8_wgrad(C:UOp, A:UOp, B:UOp, scale_A:UOp, scale_B:UOp, expert_off:UOp, *, dname:str, n_experts:int) -> UOp:
+  N, M = A.shape
+  K, M2 = B.shape
+  assert M == M2, f"{A.shape} {B.shape}"
+  E = n_experts
+  threads = UOp.special(64 * 8, "lidx0")
+  workgroups = UOp.special(E * (N // 256) * (K // 256), "gidx0")
+  sink = UOp.sink(C.base, A.base, B.base, scale_A.base, scale_B.base, expert_off.base, threads, workgroups,
+                  arg=KernelInfo(f"hk_grouped_mxfp8_wgrad_{E}_{M}_{N}_{K}",
+                                 estimates=Estimates(ops=2*M*N*K, mem=(N*M+K*M)*A.dtype.itemsize+E*N*K*C.dtype.itemsize)))
+  kittens_path = pathlib.Path(__file__).parent.parent/"thunder"/"amd"
+  src = (kittens_path/"grouped_mxfp8_wgrad.cpp").read_text()
+  lib = HIPCCCompiler("gfx950", [f"-I{(kittens_path/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-ffast-math",
+                                 "-DHIP_ENABLE_WARP_SYNC_BUILTINS", f"-DWGRAD_M={M}", f"-DWGRAD_N={N}", f"-DWGRAD_K={K}",
+                                 f"-DWGRAD_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 grouped_mx_wgrad(g:Tensor, xg:Tensor, expert_off:Tensor, n_experts:int) -> Tensor:
+  from extra.llama_kernels.transpose_quantize_mxfp8 import transpose_quantize_mxfp8
+  M, N = g.shape
+  M2, K = xg.shape
+  assert M == M2, f"{g.shape} {xg.shape}"
+  assert M % 128 == 0 and N % 256 == 0 and K % 256 == 0, f"wgrad needs M%128,N%256,K%256, got {g.shape} {xg.shape}"
+  gT, _, g_si = transpose_quantize_mxfp8(g.contiguous())
+  xT, _, x_si = transpose_quantize_mxfp8(xg.contiguous())
+  dname = (g.device[0] if isinstance(g.device, tuple) else g.device).split(":")[0]
+  is_multi = isinstance(g.device, tuple)
+  inv = Tensor.invalids(1, n_experts * N, K, dtype=dtypes.bfloat16, device=g.device)
+  out = Tensor(inv.uop.unshard(0), device=g.device) if is_multi else inv
+  out = Tensor.custom_kernel(out, gT, xT, g_si, x_si, expert_off,
+                             fxn=functools.partial(custom_hk_grouped_mxfp8_wgrad, dname=dname, n_experts=n_experts))[0]
+  out = out.sum(0) if is_multi else out.squeeze(0)
+  return out.reshape(n_experts, N, K)
+
+def mx_pack_3d(e8:Tensor) -> Tensor:
+  E, rows, scale_K = e8.shape
+  return e8.reshape(E, rows, scale_K // 4, 4).bitcast(dtypes.uint32).reshape(E, rows, scale_K // 4).permute(0, 2, 1).contiguous()
+
+@functools.cache
+def custom_grouped_mx_gemm_bw(gradient:UOp, kernel:UOp, w_stored:bool=False) -> tuple:
+  inputs = kernel.src[1:]
+  aq = Tensor(inputs[1], device=inputs[1].device)
+  bq = Tensor(inputs[2], device=inputs[2].device)
+  ae8 = Tensor(inputs[5], device=inputs[5].device)
+  be8 = Tensor(inputs[6], device=inputs[6].device)
+  E, N = bq.shape[0], bq.shape[1]
+  M, K = aq.shape
+  g = Tensor(gradient, device=aq.device).reshape(M, N).cast(dtypes.bfloat16)
+  x_phys = (aq.cast(dtypes.bfloat16) * _mx_block_scale(ae8).cast(dtypes.bfloat16))
+  w_phys = (bq.cast(dtypes.bfloat16) * _mx_block_scale_3d(be8).cast(dtypes.bfloat16))
+  expert_off = Tensor(inputs[7], device=inputs[7].device)
+  grad_x = grouped_mx_gemm(g, w_phys.transpose(1, 2), expert_off)
+  grad_w = grouped_mx_wgrad(g, x_phys, expert_off, E)
+  grad_xq = grad_x * _mx_block_scale(ae8).cast(dtypes.bfloat16)
+  grad_wq = grad_w.contiguous() if w_stored else (grad_w * _mx_block_scale_3d(be8).cast(dtypes.bfloat16)).contiguous()
+  return (None, grad_xq.uop, grad_wq.uop) + tuple(None for _ in inputs[3:])
+
+_grouped_bw_stored = functools.partial(custom_grouped_mx_gemm_bw, w_stored=True)
+
+def grouped_mx_gemm(x:Tensor, w:Tensor|tuple[Tensor, Tensor], expert_off:Tensor) -> Tensor:
+  if (pre_quantized := isinstance(w, tuple)):
+    w_q, w_e8 = w
+    E, N, K2 = w_q.shape
+  else:
+    E, N, K2 = w.shape
+  M, K = x.shape
+  assert K == K2, f"shape mismatch {x.shape} {w.shape}"
+  assert M % 256 == 0 and N % 256 == 0 and K % 128 == 0, f"grouped mxfp8 needs M%256,N%256,K%128, got {x.shape} {w.shape}"
+  dname = (x.device[0] if isinstance(x.device, tuple) else x.device).split(":")[0]
+  x_q, x_e8, x_si = quantize_mxfp8(x)
+  if not pre_quantized: w_q, w_e8, _ = quantize_mxfp8(w)
+  w_si = mx_pack_3d(w_e8)
+  xe_in, out_shape = x_e8.reshape(M, K // 32), (M, N)
+  if isinstance(x.device, tuple) and (row_axis := x.uop.axis) is not None:
+    ndev = len(x.device)
+    out = Tensor(Tensor.invalids(*(s // ndev if i == row_axis else s for i, s in enumerate(out_shape)),
+                                 dtype=dtypes.bfloat16, device=x.device).uop.unshard(row_axis), device=x.device)
+  else:
+    out = Tensor.invalids(*out_shape, dtype=dtypes.bfloat16, device=x.device)
+  return Tensor.custom_kernel(out, x_q, w_q, x_si, w_si, xe_in, w_e8, expert_off,
+                              fxn=functools.partial(custom_hk_grouped_mxfp8_gemm, dname=dname, n_experts=E),
+                              grad_fxn=(_grouped_bw_stored if pre_quantized else custom_grouped_mx_gemm_bw))[0]

extra/gemm/moe_routing.py

@@ -0,0 +1,130 @@
+from tinygrad import Tensor, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
+
+BLOCK_ROW = 256
+
+def _sharded_invalids(shape:tuple[int, ...], dtype, device) -> Tensor:
+  if isinstance(device, tuple):
+    return Tensor(Tensor.invalids(shape[0] // len(device), *shape[1:], dtype=dtype, device=device).uop.multi(0), device=device)
+  return Tensor.invalids(*shape, dtype=dtype, device=device)
+
+def _atomic_add(device:str) -> str:
+  return "__hip_atomic_fetch_add({0}, {1}, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);" if device == "AMD" \
+    else "__atomic_fetch_add({0}, {1}, __ATOMIC_RELAXED);"
+
+def _blk_for(D:int) -> int:
+  blk = 64
+  while D % blk: blk //= 2
+  return blk
+
+def _kv_ranges(G, N, D, BLK):
+  g = UOp.range(G, 0)
+  m = UOp.range(N, 1)
+  jo = UOp.range(D // BLK, 2)
+  ji = UOp.range(BLK, 3, AxisType.LOCAL)
+  return g, m, jo * BLK + ji, jo, ji
+
+def _ggather_fwd_kernel(out:UOp, table:UOp, idx:UOp) -> UOp:
+  G, M, D = out.shape
+  g, m, j, jo, ji = _kv_ranges(G, M, D, _blk_for(D))
+  row = idx.index(g, m).cast(dtypes.weakint)
+  val = table.index(g, row, j).load()
+  return out.index(g, m, j).store(val).end(g, m, jo, ji).sink(
+    arg=KernelInfo(name=f"ggather_fwd_{M}_{D}", opts_to_apply=()))
+
+def _ggather_zero_kernel(out:UOp) -> UOp:
+  i = UOp.range(out.numel(), 0)
+  return out.flatten().index(i).store(UOp.const(0.0, out.dtype)).end(i).sink(arg=KernelInfo(name="ggather_zero"))
+
+def _sharded_zeros(shape:tuple[int, ...], dtype, device) -> Tensor:
+  return Tensor.custom_kernel(_sharded_invalids(shape, dtype, device), fxn=_ggather_zero_kernel)[0]
+
+def _ggather_bwd(gradient:UOp, kernel:UOp) -> tuple:
+  _, table_u, idx_u = kernel.src[1:4]
+  dev = table_u.device
+  device = (dev[0] if isinstance(dev, tuple) else dev).split(":")[0]
+  G, R, D = table_u.shape
+  gt = _sharded_zeros((G, R, D), dtypes.float32, dev)
+  go = Tensor(gradient, device=dev)
+  atomic_str = _atomic_add(device)
+  def _bwd_kernel(gtab:UOp, gout:UOp, idx:UOp) -> UOp:
+    Gk, M, Dk = gout.shape
+    g, m, j, jo, ji = _kv_ranges(Gk, M, Dk, _blk_for(Dk))
+    row = idx.index(g, m).cast(dtypes.weakint)
+    val = gout.index(g, m, j).load().cast(dtypes.float32)
+    atomic = UOp(Ops.CUSTOM, dtypes.void, (gtab.index(g, row, j), val), arg=atomic_str)
+    return atomic.end(g, m, jo, ji).sink(arg=KernelInfo(name=f"ggather_bwd_{M}_{Dk}", opts_to_apply=()))
+  grad_table = Tensor.custom_kernel(gt, go, Tensor(idx_u, device=dev), fxn=_bwd_kernel)[0]
+  return (None, grad_table.cast(table_u.dtype).uop, None)
+
+def grouped_gather_rows(table:Tensor, idx:Tensor, n_groups:int) -> Tensor:
+  G, R, D = table.shape
+  M = idx.shape[1]
+  out = _sharded_invalids((G, M, D), table.dtype, table.device)
+  return Tensor.custom_kernel(out, table, idx, fxn=_ggather_fwd_kernel, grad_fxn=_ggather_bwd)[0]
+
+def _gscatter_fwd_kernel(out:UOp, src:UOp, idx:UOp) -> UOp:
+  G, M, D = out.shape
+  k = idx.shape[1] // src.shape[1]
+  g, m, j, jo, ji = _kv_ranges(G, idx.shape[1], D, _blk_for(D))
+  row = idx.index(g, m).cast(dtypes.weakint)
+  val = src.index(g, (m // k).cast(dtypes.weakint), j).load()
+  return out.index(g, row, j).store(val).end(g, m, jo, ji).sink(
+    arg=KernelInfo(name=f"gscatter_fwd_{idx.shape[1]}_{D}", opts_to_apply=()))
+
+def _gscatter_bwd(gradient:UOp, kernel:UOp) -> tuple:
+  _, src_u, idx_u = kernel.src[1:4]
+  dev = src_u.device
+  G, T_l, D = src_u.shape
+  k = idx_u.shape[1] // T_l
+  sel = grouped_gather_rows(Tensor(gradient, device=dev), Tensor(idx_u, device=dev), G)
+  return (None, sel.reshape(G, T_l, k, D).sum(2).cast(src_u.dtype).uop, None)
+
+def grouped_scatter_rows(src:Tensor, idx:Tensor, m_l:int) -> Tensor:
+  G, T_l, D = src.shape
+  zero = _sharded_zeros((G, m_l, D), src.dtype, src.device)
+  return Tensor.custom_kernel(zero, src, idx, fxn=_gscatter_fwd_kernel, grad_fxn=_gscatter_bwd)[0]
+
+def m_max_for(t_local:int, experts_per_tok:int, n_experts:int) -> int:
+  return (-(-t_local * experts_per_tok // BLOCK_ROW) + n_experts) * BLOCK_ROW
+
+class Routing:
+  def __init__(self, weights:Tensor, dest_row:Tensor, off:Tensor, m_l:int, n_groups:int, t_local:int):
+    self.weights, self.dest_row = weights, dest_row
+    self.off = off
+    self.m_l, self.n_groups, self.t_local = m_l, n_groups, t_local
+
+  @property
+  def rows_e(self) -> Tensor:
+    G, E = self.off.shape[0], self.off.shape[1] - 1
+    tr = Tensor.arange(self.m_l // BLOCK_ROW, dtype=dtypes.int32).reshape(1, -1, 1) * BLOCK_ROW
+    tr = tr.shard(self.off.device) if isinstance(self.off.device, tuple) else tr.to(self.off.device)
+    tile_e = ((tr >= self.off[:, :E].reshape(G, 1, E)).sum(-1) - 1).cast(dtypes.int32)
+    return tile_e.reshape(-1, 1).expand(-1, BLOCK_ROW).reshape(-1)
+
+def n_groups_of(t:Tensor) -> int:
+  return len(t.device) if isinstance(t.device, tuple) else 1
+
+def route(logits:Tensor, experts_per_tok:int, n_experts:int) -> Routing:
+  T, E = logits.shape
+  k, G = experts_per_tok, n_groups_of(logits)
+  assert T % G == 0, f"tokens {T} must split across {G} devices"
+  T_l, m_l = T // G, m_max_for(T // G, k, n_experts)
+
+  topv, topi = logits.reshape(G, T_l, E).topk(k)
+  weights = topv.softmax(-1)
+  m = topi.reshape(G, T_l * k).cast(dtypes.int32).one_hot(E).cast(dtypes.int32)
+
+  pad = ((m.sum(1) + (BLOCK_ROW - 1)) // BLOCK_ROW) * BLOCK_ROW
+  off = pad.cumsum(1).pad(((0, 0), (1, 0)))
+  dest_row = ((m.cumsum(1) + off[:, :E].reshape(G, 1, E)) * m).sum(-1).sub(1).cast(dtypes.int32)
+  return Routing(weights, dest_row, off, m_l, G, T_l)
+
+def dispatch(x:Tensor, r:Routing) -> Tensor:
+  G, D = r.n_groups, x.shape[-1]
+  return grouped_scatter_rows(x.reshape(G, r.t_local, D), r.dest_row, r.m_l).reshape(G * r.m_l, D)
+
+def combine(y:Tensor, r:Routing, n_tokens:int, experts_per_tok:int) -> Tensor:
+  G, D, k = r.n_groups, y.shape[-1], experts_per_tok
+  sel = grouped_gather_rows(y.reshape(G, r.m_l, D), r.dest_row, G).reshape(G, r.t_local, k, D)
+  return (sel * r.weights.reshape(G, r.t_local, k, 1).cast(sel.dtype)).sum(2).reshape(n_tokens, D).cast(y.dtype)

extra/gemm/openpilot_onnx_precision_scan.py

@@ -1,58 +0,0 @@
-#!/usr/bin/env python3
-"""Rapidly measure ONNX output sensitivity to FP16-rounded initializers."""
-import argparse, copy
-
-import numpy as np
-import onnx
-import onnxruntime as ort
-from onnx import numpy_helper
-
-
-def main() -> None:
-  ap = argparse.ArgumentParser()
-  ap.add_argument("model")
-  ap.add_argument("corpus")
-  ap.add_argument("--case", type=int, default=9)
-  ap.add_argument("--chunks", type=int, default=8)
-  ap.add_argument("--start", type=int, default=0)
-  ap.add_argument("--stop", type=int)
-  ap.add_argument("--list", action="store_true")
-  args = ap.parse_args()
-
-  model = onnx.load(args.model)
-  consumers: dict[str, list[str]] = {}
-  for node in model.graph.node:
-    for name in node.input: consumers.setdefault(name, []).append(node.op_type)
-  initializers = [(init, numpy_helper.to_array(init)) for init in model.graph.initializer]
-  selected = [(init, arr) for init, arr in initializers if arr.dtype == np.float32 and arr.ndim >= 2 and
-              any(op in {"Conv", "Gemm", "MatMul"} for op in consumers.get(init.name, []))]
-  if args.list:
-    for i, (init, arr) in enumerate(selected): print(i, init.name, arr.shape, consumers.get(init.name))
-
-  # Listing an initializer as a graph input lets one ORT session override it at run time.
-  known_inputs = {x.name for x in model.graph.input}
-  for init, _ in selected:
-    if init.name not in known_inputs:
-      model.graph.input.append(copy.deepcopy(onnx.helper.make_tensor_value_info(init.name, init.data_type, init.dims)))
-  session_options = ort.SessionOptions()
-  session_options.log_severity_level = 3
-  session = ort.InferenceSession(model.SerializeToString(), session_options, providers=["CPUExecutionProvider"])
-  corpus = np.load(args.corpus)
-  feeds = {spec.name: corpus[f"case{args.case}:input:{spec.name}"] for spec in session.get_inputs()
-           if f"case{args.case}:input:{spec.name}" in corpus}
-  expected = corpus[f"case{args.case}:output"].astype(np.float32)
-
-  def check(indices: list[int]) -> tuple[float, float]:
-    overrides = {selected[i][0].name: selected[i][1].astype(np.float16).astype(np.float32) for i in indices}
-    got = session.run(None, feeds | overrides)[0].astype(np.float32)
-    delta = np.abs(expected.reshape(got.shape)-got)
-    return float(delta.max()), float(delta.mean())
-
-  scan = list(range(args.start, len(selected) if args.stop is None else args.stop))
-  print(f"selected={len(selected)} scan={scan[0]}..{scan[-1]} baseline={check([])} scan_error={check(scan)}")
-  for chunk in np.array_split(np.asarray(scan), args.chunks):
-    ids = [int(x) for x in chunk]
-    print(f"range={ids[0]}..{ids[-1]} count={len(ids)} error={check(ids)}")
-
-
-if __name__ == "__main__": main()

extra/gemm/qcom_4x16_check.py

@@ -1,266 +0,0 @@
-#!/usr/bin/env python3
-"""Random-matrix oracle for the hand IR3 4x16 FP16 GEMM."""
-import os, struct
-
-import numpy as np
-
-from tinygrad import Device, dtypes
-from tinygrad.device import Buffer
-from extra.gemm import qcom_intensity_gemm as q
-from extra.gemm.ir3asm import disasm, get_envelope, inject
-
-
-def upload(x: np.ndarray) -> Buffer:
-  ret = Buffer("QCOM", x.size, dtypes.half).allocate()
-  ret.copyin(memoryview(np.ascontiguousarray(x)).cast("B"))
-  return ret
-
-
-def strip_redundant_mad_sy(lib: bytes) -> bytes:
-  ret = bytearray(lib)
-  io, sz = struct.unpack_from('<I', ret, 0xc0)[0], struct.unpack_from('<I', ret, 0x100)[0]
-  seen = False
-  for off in range(io, io+sz, 8):
-    hi = struct.unpack_from('<I', ret, off+4)[0]
-    if (hi >> 24) == 0x73:
-      if seen: struct.pack_into('<I', ret, off+4, (hi & 0x0fffffff) | 0x60000000)
-      else: seen = True
-  return bytes(ret)
-
-
-def restore_all_mad_sy(lib: bytes) -> bytes:
-  ret = bytearray(lib)
-  io, sz = struct.unpack_from('<I', ret, 0xc0)[0], struct.unpack_from('<I', ret, 0x100)[0]
-  for off in range(io, io+sz, 8):
-    hi = struct.unpack_from('<I', ret, off+4)[0]
-    if (hi >> 24) == 0x63: struct.pack_into('<I', ret, off+4, (hi & 0x0fffffff) | 0x70000000)
-  return bytes(ret)
-
-
-def restore_original_mad_sy(lib: bytes, original: bytes) -> bytes:
-  ret = bytearray(lib)
-  io, sz = struct.unpack_from('<I', ret, 0xc0)[0], struct.unpack_from('<I', ret, 0x100)[0]
-  for off in range(io, io+sz, 8):
-    hi = struct.unpack_from('<I', ret, off+4)[0]
-    old_hi = struct.unpack_from('<I', original, off+4)[0]
-    if (hi >> 24) == 0x63 and (old_hi >> 24) == 0x73:
-      struct.pack_into('<I', ret, off+4, (hi & 0x0fffffff) | 0x70000000)
-  return bytes(ret)
-
-
-def main() -> None:
-  m, n, k = int(os.getenv("M", "128")), int(os.getenv("N", "1024")), int(os.getenv("K", "384"))
-  stride = int(os.getenv("STRIDE", str(n)))
-  ncols = int(os.getenv("NCOLS", "4"))
-  threads = int(os.getenv("THREADS", "128"))
-  rng = np.random.default_rng(int(os.getenv("SEED", "4")))
-  a = (rng.standard_normal((m, k))*0.05).astype(np.float16)
-  b = (rng.standard_normal((k, n))*0.05).astype(np.float16)
-  if pattern := os.getenv("PATTERN", ""):
-    a.fill(0)
-    b.fill(0)
-    if pattern == "row":
-      a[:, 0] = np.arange(1, m+1)
-      b[0, :] = 1
-    elif pattern == "col":
-      a[:, 0] = 1
-      b[0, :] = (np.arange(n) % 251) + 1
-    elif pattern.startswith("k"):
-      kk = int(pattern[1:])
-      a[:, kk] = np.arange(1, m+1)
-      b[kk, :] = 1
-    else: raise ValueError(f"unknown PATTERN={pattern!r}")
-  q.M, q.N, q.K, q.K4 = m, stride, k, k//4
-  dev = Device["QCOM"]
-  compiler = bool(int(os.getenv("COMPILER", "0")))
-  env_ncols = int(os.getenv("ENV_NCOLS", str(ncols)))
-  direct_env = compiler or bool(int(os.getenv("ENV_DIRECT", "0")))
-  image_store = bool(int(os.getenv("IMAGE_STORE", "0")))
-  output_float = bool(int(os.getenv("OUTPUT_FLOAT", "0")))
-  dynamic_splits = int(os.getenv("DYNAMIC_SPLIT", "0"))
-  env_src = q.make_direct_image_donor_src(env_ncols, threads) if image_store else \
-            q.make_direct_donor_src(env_ncols if direct_env else ncols, threads) if direct_env else q.make_donor_src(env_ncols, threads)
-  env, io, sz, ro = get_envelope(dev, env_src)
-  fast = bool(int(os.getenv("FAST", "1")))
-  preserve_coords = bool(int(os.getenv("PRESERVE_COORDS", "0")))
-  high_inputs = bool(int(os.getenv("HIGH_INPUTS", "0")))
-  high_store = bool(int(os.getenv("HIGH_STORE", "0")))
-  low_a = bool(int(os.getenv("LOW_A", "0")))
-  inc = bool(int(os.getenv("INC", str(int(fast and not preserve_coords)))))
-  persistent = bool(int(os.getenv("PERSISTENT", str(int(fast and inc and not preserve_coords)))))
-  unroll = int(os.getenv("K_UNROLL", str(4 if k % 16 == 0 else 1)))
-  k_count = int(os.getenv("K_COUNT", str(k//4)))
-  if compiler:
-    patch_mode = os.getenv("PATCH_COMPILER", "0")
-    if patch_mode == "sync": lib = strip_redundant_mad_sy(env)
-    elif patch_mode != "0":
-      from extra.gemm.qcom_gemm import patch_kernel
-      lib = patch_kernel(env)
-      if patch_mode == "rpt": lib = restore_all_mad_sy(lib)
-      elif patch_mode == "original": lib = restore_original_mad_sy(lib, env)
-    else: lib = env
-    shader = bytes(env[io:io+sz])
-  else:
-    safe_store = bool(int(os.getenv("SAFE_STORE", "0")))
-    compact = bool(int(os.getenv("COMPACT", str(int(not safe_store)))))
-    isolated = bool(int(os.getenv("ISOLATED", "0")))
-    shader, _ = q.build_4x16_isolated_shader(dev, threads, k_unroll=unroll) if isolated else q.build_4xn_shader(
-      dev, threads, ncols=ncols, direct=True, compact_acc=compact,
-      store_constant=bool(int(os.getenv("STORE_CONSTANT", "0"))),
-      donor_store=bool(int(os.getenv("DONOR_STORE", "0"))),
-      native_store=bool(int(os.getenv("NATIVE_STORE", "0"))),
-      safe_store=safe_store,
-      linear_store=bool(int(os.getenv("LINEAR_STORE", "0"))),
-      image_store=image_store,
-      preserve_coords=preserve_coords,
-      preload_b=bool(int(os.getenv("PRELOAD_B", "0"))),
-      preload_b_safe_coords=bool(int(os.getenv("PRELOAD_B_SAFE_COORDS", "0"))),
-      high_inputs=high_inputs,
-      high_store=high_store,
-      copy_b_probe=bool(int(os.getenv("COPY_B_PROBE", "0"))),
-      thread_store=bool(int(os.getenv("THREAD_STORE", "0"))),
-      repeat_first_store=bool(int(os.getenv("REPEAT_FIRST_STORE", "0"))),
-      repair_row1_store=bool(int(os.getenv("REPAIR_ROW1_STORE", "0"))),
-      repeat_each_store=bool(int(os.getenv("REPEAT_EACH_STORE", "0"))),
-      post_constant=bool(int(os.getenv("POST", "0"))),
-      stable_bx=fast and not preserve_coords, stable_ay=fast, low_a_coords=low_a,
-      inc_coords=inc, persistent_coords=persistent,
-      serial_b_cols=bool(int(os.getenv("SERIAL", "0"))),
-      single_cols_all=bool(int(os.getenv("SINGLE_COLS_ALL", "0"))),
-      first_sync_only=bool(int(os.getenv("FIRST_SYNC_ONLY", str(int(fast))))),
-      no_store=bool(int(os.getenv("NO_STORE", "0"))),
-      skip_a_loads=bool(int(os.getenv("SKIP_A_LOADS", "0"))),
-      skip_b_loads=bool(int(os.getenv("SKIP_B_LOADS", "0"))),
-      k_unroll=unroll, b_first=fast and ncols == 4 and not preserve_coords,
-      k_count=None if dynamic_splits else k_count,
-      coord_delay=int(os.getenv("COORD_DELAY", "-1" if fast else "4")),
-      stable_settle_delay=int(os.getenv("STABLE_SETTLE_DELAY", "5")),
-      row_sync=bool(int(os.getenv("ROW_SYNC", "0"))), store_row_shift=int(os.getenv("STORE_ROW_SHIFT", "10")),
-      store_gap=int(os.getenv("STORE_GAP", "-1")),
-      safe_b_y=bool(int(os.getenv("SAFE_B_Y", "0"))), sync_b_y=bool(int(os.getenv("SYNC_B_Y", "0"))),
-      separate_b_coords=bool(int(os.getenv("SEPARATE_B_COORDS", "0"))),
-      high_b_coords=bool(int(os.getenv("HIGH_B_COORDS", "0"))),
-      reuse_separate_b_y=bool(int(os.getenv("REUSE_SEPARATE_B_Y", "0"))),
-      persistent_b_coords=bool(int(os.getenv("PERSISTENT_B_COORDS", "0"))),
-      interleave_second_pair=bool(int(os.getenv("INTERLEAVE_SECOND_PAIR", "0"))),
-      pipeline=bool(int(os.getenv("PIPELINE", "0"))),
-      acc_hr=int(os.getenv("ACC_HR")) if os.getenv("ACC_HR") else None,
-      high_a_only=bool(int(os.getenv("HIGH_A_ONLY", "0"))),
-      save_output_coords=bool(int(os.getenv("SAVE_OUTPUT_COORDS", "0"))),
-      vector_init=bool(int(os.getenv("VECTOR_INIT", "0"))),
-      dynamic_split_k=dynamic_splits,
-      alu_order=os.getenv("ALU_ORDER", "auto"),
-      first_cols_only=bool(int(os.getenv("FIRST_COLS_ONLY", "0"))), first_cols_offset=int(os.getenv("FIRST_COLS_OFFSET", "0")))
-    merged_opt = os.getenv("MERGEDREGS")
-    mergedregs = None if merged_opt is None else bool(int(merged_opt))
-    native = bool(int(os.getenv("NATIVE_STORE", "0")))
-    save_output = bool(int(os.getenv("SAVE_OUTPUT_COORDS", "0")))
-    persistent_b = bool(int(os.getenv("PERSISTENT_B_COORDS", "0")))
-    default_fregs = (23 if isolated else 30 if high_store else 28 if native else 19 if save_output and persistent_b else
-                     18 if bool(int(os.getenv("HIGH_B_COORDS", "0"))) else
-                     16 if bool(int(os.getenv("SAFE_B_Y", "0"))) else 11 if save_output or bool(int(os.getenv("THREAD_STORE", "0"))) else
-                     8 if high_inputs and low_a else 10)
-    acc_hr = int(os.getenv("ACC_HR", "0"))
-    default_hregs = (28 if isolated else max(acc_hr + 4*ncols, 36 if bool(int(os.getenv("HIGH_A_ONLY", "0"))) else
-                     44 if high_inputs and low_a else 48 if high_inputs else 32 if not compact else 12 + 4*ncols))
-    lib = inject(env, io, sz, ro, shader, fregs=int(os.getenv("FREGS", str(default_fregs))),
-                 hregs=int(os.getenv("HREGS", str(default_hregs))), mergedregs=mergedregs)
-  if int(os.getenv("PRINT_META", "0")):
-    asm = disasm(shader)
-    print(f"shader_instrs={len(shader)//8} mad_f16={asm.count('mad.f16')} isam={asm.count('isam')} sy={asm.count('(sy)')}")
-  if int(os.getenv("DUMP", "0")):
-    print(disasm(shader))
-    return
-  ab, bb = upload(a), upload(b.reshape(k, n//4, 4))
-  cb = Buffer("QCOM", max(1, dynamic_splits)*m*stride, dtypes.float if output_float else dtypes.half).allocate()
-  cb.copyin(memoryview(np.zeros((max(1, dynamic_splits)*m, stride), np.float32 if output_float else np.float16)).cast("B"))
-  specs = ([((0, dtypes.half, (m, stride//4, 4)),), ((0, dtypes.half, (m, k//4, 4)),),
-            ((1, dtypes.half, (k, n//4, 4)),)] if image_store and not output_float else
-           [((0, dtypes.float, (m, stride//4, 4)),), ((0, dtypes.half, (m, k//4, 4)),),
-            ((1, dtypes.half, (k, n//4, 4)),)] if image_store else
-           [((0, dtypes.half, (m, k//4, 4)),), ((1, dtypes.half, (k, n//4, 4)),),
-            ((2, dtypes.half, (max(1, dynamic_splits)*m*stride,)),)])
-  prg = dev.runtime("gemm_h", lib, buf_dtypes=specs)
-  call_bufs = (cb._buf, ab._buf, bb._buf) if image_store else (ab._buf, bb._buf, cb._buf)
-  tile_m = (threads//32)*4
-  # Each workgroup covers 32 lanes * ncols half4 vectors = 128*ncols scalar columns;
-  # thread count changes only the number of 4-row subtiles in Y.
-  times = [prg(*call_bufs, global_size=(n//(128*ncols), (m//tile_m)*max(1, dynamic_splits), 1), local_size=(threads, 1, 1), wait=True)*1e3
-           for _ in range(10)]
-  if int(os.getenv("NO_STORE", "0")):
-    print(f"K={k} ncols={ncols} compute_only_ms={min(times):.4f}")
-    return
-  got = np.empty((max(1, dynamic_splits)*m, stride), np.float32 if output_float else np.float16)
-  cb.copyout(memoryview(got).cast("B"))
-  if dynamic_splits:
-    split_got = got.reshape(dynamic_splits, m, stride).astype(np.float32)
-    if int(os.getenv("SPLIT_STATS", "0")):
-      chunk = k//dynamic_splits
-      split_expected = np.stack([a[:, s*chunk:(s+1)*chunk].astype(np.float32) @
-                                 b[s*chunk:(s+1)*chunk].astype(np.float32) for s in range(dynamic_splits)])
-      print("split_err", [[float(np.abs(split_got[x, :, :n]-split_expected[y]).mean())
-                           for y in range(dynamic_splits)] for x in range(dynamic_splits)])
-      print("split_norm", [float(np.abs(split_got[x, :, :n]).mean()) for x in range(dynamic_splits)])
-    got = split_got.sum(axis=0)
-  if int(os.getenv("RAW_STATS", "0")):
-    nz = np.flatnonzero(got.reshape(-1))
-    print("c_va", hex(cb._buf.va_addr), "raw_nonzero", len(nz), "head", nz[:128].tolist(), "tail", nz[-32:].tolist())
-  if int(os.getenv("THREAD_STORE", "0")):
-    raw, got = got.reshape(-1, 4, ncols, 4), np.empty_like(got)
-    nz = np.flatnonzero(raw.reshape(-1))
-    print("thread_nonzero_head", nz[:64].tolist(), "threads", np.unique(nz//(16*ncols))[:64].tolist(),
-          "thread_count", len(np.unique(nz//(16*ncols))))
-    # The thread-major kernel reserves tile slots using the physical output
-    # stride, even when only a logical prefix of columns is launched.
-    storage_gx_count = stride//(128*ncols)
-    launched_gx_count = n//(128*ncols)
-    for gy in range(m//16):
-      for gx in range(launched_gx_count):
-        for lid in range(128):
-          tm, tid = lid//32, lid%32
-          thread = (gy*storage_gx_count+gx)*128+lid
-          col_base = gx*32*ncols+tid
-          for row in range(4):
-            for col in range(ncols): got[gy*16+tm*4+row, (col_base+col*32)*4:(col_base+col*32+1)*4] = raw[thread, row, col]
-  got = got[:, :n]
-  expected = (np.full((m, n), 1024.0, np.float32) if int(os.getenv("POST", "0")) else
-              np.broadcast_to(b[0].astype(np.float32), (m, n)) if int(os.getenv("COPY_B_PROBE", "0")) else
-              np.full((m, n), float(k), np.float32) if int(os.getenv("SKIP_A_LOADS", "0")) and int(os.getenv("SKIP_B_LOADS", "0")) else
-              a[:, :k_count*4].astype(np.float32) @ b[:k_count*4].astype(np.float32))
-  delta = np.abs(got.astype(np.float32)-expected)
-  checked = np.ones((m, n), dtype=bool)
-  if int(os.getenv("FIRST_COLS_ONLY", "0")):
-    selected_parity = int(os.getenv("FIRST_COLS_OFFSET", "0")) & 1
-    for block in range(n//128):
-      if (block % ncols) % 2 != selected_parity:
-        delta[:, block*128:(block+1)*128] = 0
-        checked[:, block*128:(block+1)*128] = False
-  correct = np.allclose(got[checked], expected[checked], rtol=2e-2, atol=2e-2)
-  print(f"K={k} fast={fast} ms={min(times):.4f} max={delta.max():.8g} mean={delta.mean():.8g} "
-        f"finite={np.isfinite(got[checked]).all()} allclose={correct}")
-  print("samples", got[0, :16].tolist(), expected[0, :16].tolist())
-  if pattern:
-    print("pattern_blocks", [(x, got[0, x:x+8].tolist()) for x in range(0, n, 32)])
-  print("worst", np.unravel_index(int(np.nanargmax(delta)), delta.shape),
-        "col_means", [float(delta[:, x:x+128].mean()) for x in range(0, n, 128)],
-        "row_means", [float(delta[x:x+16].mean()) for x in range(0, m, 16)])
-  for out_block in (1, 3):
-    x = out_block * 128
-    print("block_match", out_block,
-          [float(np.abs(got[:, x:x+128].astype(np.float32)-expected[:, y:y+128]).mean())
-           for y in range(0, n, 128)])
-  bad = np.argwhere(delta > 0.02)
-  print("bad_count", len(bad), "bad_head", bad[:32].tolist())
-  print("bad_rows", [(int(r), int((bad[:, 0] == r).sum())) for r in np.unique(bad[:, 0])],
-        "bad_col_range", (int(bad[:, 1].min()), int(bad[:, 1].max())) if len(bad) else None)
-  print("row1_match", [float(np.abs(got[1].astype(np.float32)-expected[r]).mean()) for r in range(16)])
-  for probe_row in (127, 128, m-1):
-    if probe_row >= m: continue
-    probe_cols = checked[probe_row]
-    row_delta = np.abs(expected[:, probe_cols] - got[probe_row, probe_cols].astype(np.float32)).mean(axis=1)
-    nearest = np.argsort(row_delta)[:4]
-    print("row_match", probe_row, [(int(r), float(row_delta[r])) for r in nearest])
-  if not correct: raise SystemExit(1)
-
-
-if __name__ == "__main__": main()

extra/gemm/qcom_8x8_check.py

@@ -1,277 +0,0 @@
-#!/usr/bin/env python3
-"""Random-matrix oracle for the high-throughput 8x8 IR3 GEMM."""
-import os, hashlib
-import numpy as np
-
-from tinygrad import Device, dtypes
-from tinygrad.device import Buffer
-from extra.gemm import qcom_8x4_gemm as q8
-from extra.gemm import qcom_intensity_gemm as q4
-from extra.gemm.ir3asm import disasm, get_envelope, inject
-
-
-def main():
-  m, n, k = int(os.getenv("M", "128")), int(os.getenv("N", "512")), int(os.getenv("K", "192"))
-  batch = int(os.getenv("BATCH", "1"))
-  threads = int(os.getenv("THREADS", "128"))
-  stride = int(os.getenv("STRIDE", str(max(1024, n))))
-  k_start, k_count = int(os.getenv("K_START", "0")), int(os.getenv("K_COUNT", str(k//4)))
-  rng = np.random.default_rng(int(os.getenv("SEED", "0")))
-  a_np = (rng.standard_normal((batch*m, k))*0.05).astype(np.float16)
-  b_np = (rng.standard_normal((batch*k, n))*0.05).astype(np.float16)
-  batch_horizontal = batch > 1 and bool(int(os.getenv("BATCH_HORIZONTAL", "1")))
-  batch_repeat_b = batch > 1 and not batch_horizontal and bool(int(os.getenv("BATCH_REPEAT_B", "0")))
-  batch_repeat_b_x = batch > 1 and not batch_horizontal and bool(int(os.getenv("BATCH_REPEAT_B_X", "0")))
-  b_storage = (np.concatenate([b_np[x*k:(x+1)*k] for x in range(batch) for _ in range(m//8)], axis=1)
-               if batch_repeat_b_x else
-               np.concatenate([b_np[x*k:(x+1)*k] for x in range(batch)], axis=1) if batch_horizontal else
-               np.concatenate([b_np[x*k:(x+1)*k] for x in range(batch) for _ in range(m//8)])
-               if batch_repeat_b else b_np)
-  pattern = os.getenv("PATTERN", "")
-  if pattern:
-    a_np.fill(0)
-    b_np.fill(0)
-    if pattern == "row":
-      a_np[:, 0] = np.arange(1, m+1, dtype=np.float16)
-      b_np[0, :] = 1
-    elif pattern == "col":
-      a_np[:, 0] = 1
-      b_np[0, :] = (np.arange(n, dtype=np.float16) % 251) + 1
-    elif pattern.startswith("k"):
-      kk = int(pattern[1:])
-      a_np[:, kk] = np.arange(1, m+1, dtype=np.float16)
-      b_np[kk, :] = 1
-    elif pattern.startswith("cross"):
-      ak, bk = map(int, pattern[5:].split("_"))
-      a_np[:, ak] = np.arange(1, m+1, dtype=np.float16)
-      b_np[bk, :] = 1
-    elif pattern == "ones":
-      a_np.fill(1)
-      b_np.fill(1)
-    else: raise ValueError(f"unknown PATTERN={pattern!r}")
-  q8.M, q8.N, q8.K, q8.K4 = batch*m, stride, k, k//4
-  dev = Device["QCOM"]
-  image_store = bool(int(os.getenv("IMAGE_STORE", "0")))
-  batch_const_mask = batch > 1 and bool(int(os.getenv("BATCH_CONST_MASK", "0")))
-  batch_z = batch > 1 and bool(int(os.getenv("BATCH_Z", "0")))
-  loop_instrs = -1
-  if bool(int(os.getenv("COMPILER", "0"))):
-    lib, _, _, _ = get_envelope(dev, q8.make_donor_src8(2, 128))
-  else:
-    wide = bool(int(os.getenv("WIDE", "0")))
-    tri = bool(int(os.getenv("TRI", "0")))
-    env_src = q4.make_direct_image_donor_src(4, threads) if image_store else q8.make_donor_src8(4, threads)
-    if batch_const_mask:
-      if not image_store: raise ValueError("BATCH_CONST_MASK currently requires IMAGE_STORE=1")
-      groups_per_batch = m // ((threads//32)*8)
-      env_src = env_src.replace("for(int k4=0", f"int batch=get_group_id(1)/{groups_per_batch};for(int k4=0")
-      env_src = env_src.replace("k4*4+", f"batch*{k}+k4*4+")
-    env, io, sz, ro = get_envelope(dev, env_src)
-    if int(os.getenv("PERSISTENT8", "0")):
-      shader, hregs, fregs, loop_instrs = q8.build_8x8_persistent_shader(dev, threads, k_count=k_count,
-        store_row_shift=int(os.getenv("STORE_ROW_SHIFT", "10")), pipeline_b=bool(int(os.getenv("PERSISTENT_PIPELINE", "0"))),
-        b_reuse_gap=int(os.getenv("B_REUSE_GAP", "0")), double_b=bool(int(os.getenv("PERSISTENT_DOUBLE_B", "0"))),
-        rotate_b=bool(int(os.getenv("PERSISTENT_ROTATE_B", "0"))), pipeline_a=bool(int(os.getenv("PERSISTENT_PIPELINE_A", "0"))),
-        one_sync=bool(int(os.getenv("PERSISTENT_ONE_SYNC", "0"))), one_sync_wait=int(os.getenv("PERSISTENT_ONE_SYNC_WAIT", "0")),
-        stagger_b=bool(int(os.getenv("PERSISTENT_STAGGER_B", "0"))),
-        stagger_rows=int(os.getenv("STAGGER_ROWS", "2")), masked_prefetch_a4=bool(int(os.getenv("MASKED_PREFETCH_A4", "0"))),
-        lagged_a4=bool(int(os.getenv("LAGGED_A4", "0"))), dual_a_tile=bool(int(os.getenv("DUAL_A_TILE", "0"))),
-        stream_a4_gap=int(os.getenv("STREAM_A4_GAP", "-1")),
-        dynamic_a4_dual=bool(int(os.getenv("DYNAMIC_A4_DUAL", "0"))),
-        dynamic_a4_wait=int(os.getenv("DYNAMIC_A4_WAIT", "0")),
-        dynamic_b_prefetch=bool(int(os.getenv("DYNAMIC_B_PREFETCH", "0"))),
-        dynamic_b_rows=int(os.getenv("DYNAMIC_B_ROWS", "1")),
-        dynamic_b_gap=int(os.getenv("DYNAMIC_B_GAP", "0")),
-        rotate_low_banks=bool(int(os.getenv("ROTATE_LOW_BANKS", "0"))),
-        rotate_no_prefetch=bool(int(os.getenv("ROTATE_NO_PREFETCH", "0"))),
-        batch_m=m if batch > 1 and bool(int(os.getenv("BATCH_SHADER", "1"))) else 0,
-        batch_n=n if batch > 1 and bool(int(os.getenv("BATCH_SHADER", "1"))) else 0,
-        batch_k=k if batch > 1 and bool(int(os.getenv("BATCH_SHADER", "1"))) else 0,
-        batch_b_offset=bool(int(os.getenv("BATCH_B_OFFSET", "1"))),
-        batch_row_offset=bool(int(os.getenv("BATCH_ROW_OFFSET", "1"))),
-        batch_horizontal=batch_horizontal,
-        batch_repeat_b=batch_repeat_b,
-        batch_repeat_b_x=batch_repeat_b_x,
-        batch_fixed_b=-2 if batch_z else int(os.getenv("BATCH_FIXED_B", "-1")),
-        batch_const_mask=batch_const_mask,
-        image_store_gap=int(os.getenv("IMAGE_STORE_GAP", "16")),
-        image_store=image_store)
-    elif int(os.getenv("PACKED_B8", "0")):
-      shader, hregs, fregs, loop_instrs = q8.build_8x8_bpacked_shader(dev, threads, coord_delay=int(os.getenv("ADELAY", "5")),
-        merged_alias=bool(int(os.getenv("PACKED_B_ALIAS", "0"))))
-    elif int(os.getenv("PACKED8", "0")):
-      shader, hregs, fregs, loop_instrs = q8.build_8x8_packed8_shader(dev, threads, coord_delay=int(os.getenv("ADELAY", "2")))
-    elif tri:
-      shader, hregs, fregs, loop_instrs = q8.build_8x4_shader(dev, 128, os.getenv("TRI_VARIANT", "serial"), 3,
-        a_coord_delay=int(os.getenv("ADELAY", "-1")), b_coord_delay=int(os.getenv("BDELAY", "-1")),
-        post_constant=bool(int(os.getenv("POST_CONSTANT", "0"))), image_store=image_store)
-    elif wide:
-      if image_store: raise ValueError("WIDE image store is not implemented")
-      shader, hregs, fregs, loop_instrs = q8.build_8x16_split_a_unroll_shader(dev, 128, k_unroll=int(os.getenv("KUNROLL", "4")),
-        b_coord_delay=int(os.getenv("BDELAY", "0")), fast_coords=True, safe_coords=bool(int(os.getenv("SAFE_COORDS", "1"))),
-        add256_store_mode=os.getenv("STORE_MODE", "tight"), alu_order=os.getenv("ALU_ORDER", "row_col_kk"),
-        post_constant=bool(int(os.getenv("POST_CONSTANT", "0"))),
-        skip_a_loads=bool(int(os.getenv("SKIP_A_LOADS", "0"))), skip_b_loads=bool(int(os.getenv("SKIP_B_LOADS", "0"))),
-        store_row_shift=int(os.getenv("STORE_ROW_SHIFT", "10")))
-    elif int(os.getenv("LIFETIME", "0")):
-      shader, hregs, fregs, loop_instrs = q8.build_8x8_lifetime_shader(dev, 128, k_unroll=int(os.getenv("KUNROLL", "4")),
-        b_coord_delay=int(os.getenv("BDELAY", "0")), a_coord_delay=int(os.getenv("ADELAY", "0")),
-        k_start=k_start, k_count=k_count, post_sequence=bool(int(os.getenv("POST_SEQUENCE", "0"))))
-    elif int(os.getenv("SELF_COORDS", "0")):
-      shader, hregs, fregs, loop_instrs = q8.build_8x8_selfcoord_shader(
-        dev, 128, coord_delay=int(os.getenv("ADELAY", "0")), post_sequence=bool(int(os.getenv("POST_SEQUENCE", "0"))))
-    elif int(os.getenv("BASE", "0")):
-      shader, hregs, fregs, loop_instrs = q8.build_8x8_split_a_shader(dev, 128,
-        a_coord_delay=int(os.getenv("ADELAY", "4")), b_coord_delay=int(os.getenv("BDELAY", "4")),
-        post_constant=bool(int(os.getenv("POST_CONSTANT", "0"))),
-        thread_store_gx=n//256 if int(os.getenv("THREAD_STORE", "0")) else 0,
-        thread_store_lid_reg=None if os.getenv("SAVE_REG", "r28.x") == "none" else os.getenv("SAVE_REG", "r28.x"),
-        thread_store_group_regs=("r36.y", "r36.z") if int(os.getenv("SAVE_GROUPS", "0")) else None,
-        row_sync=bool(int(os.getenv("ROW_SYNC", "0"))), reserved_out=int(os.getenv("RESERVED_OUT", "-1")))
-    else:
-      hist = bool(int(os.getenv("HIST", "0")))
-      common = dict(k_unroll=int(os.getenv("KUNROLL", "8")), b_coord_delay=int(os.getenv("BDELAY", "0")),
-        fast_coords=True, prefetch_next_b=bool(int(os.getenv("PREFETCH", "0"))), add256_store_mode=os.getenv("STORE_MODE", "tight"),
-        prefetch_next_a=bool(int(os.getenv("PREFETCH_A", "0"))),
-        grouped_b=bool(int(os.getenv("GROUPED_B", "0"))), grouped_b_cols=bool(int(os.getenv("GROUPED_B_COLS", "0"))),
-        stream_col1=bool(int(os.getenv("STREAM_COL1", "0"))), stream_col1_sync=bool(int(os.getenv("STREAM_COL1_SYNC", "0"))),
-        add256_gap=int(os.getenv("ADD256_GAP", "16")),
-        add256_offset_before_gap=bool(int(os.getenv("ADD256_OFFSET_BEFORE_GAP", "0"))),
-        alu_order=os.getenv("ALU_ORDER", "row_col_kk"),
-        post_constant=bool(int(os.getenv("POST_CONSTANT", "0"))))
-      if hist:
-        shader, hregs, fregs, loop_instrs = q8.build_8x8_split_a_unroll_shader(dev, 128, **common)
-      else:
-        shader, hregs, fregs, loop_instrs = q8.build_8x8_split_a_unroll_shader(dev, threads, **common, k_start=k_start, k_count=k_count,
-        thread_store_gx=n//256 if int(os.getenv("THREAD_STORE", "0")) else 0,
-        post_sequence=bool(int(os.getenv("POST_SEQUENCE", "0"))),
-        a_coord_delay=int(os.getenv("ADELAY", "4")), unroll_gap=int(os.getenv("GAP", "0")),
-        relaxed_sync=bool(int(os.getenv("RELAXED_SYNC", "0"))), sync_mask=int(os.getenv("SYNC_MASK", "7"), 0),
-        sync_wait=int(os.getenv("SYNC_WAIT", "0")), high_inputs=bool(int(os.getenv("HIGH_INPUTS", "0"))), image_store=image_store,
-        mid_acc=bool(int(os.getenv("MID_ACC", "0"))),
-        safe_coords=bool(int(os.getenv("SAFE_COORDS", "0"))), low_stable_coords=bool(int(os.getenv("LOW_STABLE_COORDS", "0"))),
-        triple_coords=bool(int(os.getenv("TRIPLE_COORDS", "0"))),
-        dual_a_coords=bool(int(os.getenv("DUAL_A_COORDS", "0"))),
-        high_pair_coords=bool(int(os.getenv("HIGH_PAIR_COORDS", "0"))),
-        high_a=bool(int(os.getenv("HIGH_A", "0"))),
-        low_a=bool(int(os.getenv("LOW_A", "0"))),
-        high_pair_b=bool(int(os.getenv("HIGH_PAIR_B", "0"))), high_pair_a=bool(int(os.getenv("HIGH_PAIR_A", "0"))),
-        serial_safe_coords=bool(int(os.getenv("SERIAL_SAFE_COORDS", "0"))),
-        separate_coords=bool(int(os.getenv("SEPARATE_COORDS", "0"))), buffer_a=bool(int(os.getenv("BUFFER_A", "0"))),
-        prefetch_loop_b=bool(int(os.getenv("PREFETCH_LOOP_B", "0"))), preload_a8=bool(int(os.getenv("PRELOAD_A8", "0"))),
-        reuse_b=bool(int(os.getenv("REUSE_B", "0"))), row_stream=bool(int(os.getenv("ROW_STREAM", "0"))),
-        phase_stream=bool(int(os.getenv("PHASE_STREAM", "0"))), split_low_pairs=bool(int(os.getenv("SPLIT_LOW_PAIRS", "0"))),
-        quad_a=bool(int(os.getenv("QUAD_A", "0"))), quad_map=os.getenv("QUAD_MAP", "0123"),
-        sampler_source_sync=bool(int(os.getenv("SOURCE_SYNC", "0"))),
-        stream_b_a8=bool(int(os.getenv("STREAM_B_A8", "0"))), store_row_shift=int(os.getenv("STORE_ROW_SHIFT", "10")),
-        source_hold_delay=int(os.getenv("SOURCE_HOLD_DELAY", "-1")), one_sync_tile=bool(int(os.getenv("ONE_SYNC_TILE", "0"))),
-        interleave_a4=bool(int(os.getenv("INTERLEAVE_A4", "0"))), interleave_a_reuse_gap=int(os.getenv("A_REUSE_GAP", "0")),
-        single_high_coord=bool(int(os.getenv("SINGLE_HIGH_COORD", "0"))))
-    assert len(shader) <= sz
-    if int(os.getenv("DISASM", "0")): print(disasm(shader))
-    lib = inject(env, io, sz, ro, shader, fregs=int(os.getenv("FREGS", str(fregs))), hregs=int(os.getenv("HREGS", str(hregs))),
-                 mergedregs=False if bool(int(os.getenv("SEPARATE_REGS", "0"))) else None)
-    if int(os.getenv("PRINT_META", "0")): print("shader_meta", fregs, hregs, len(shader), loop_instrs, hashlib.sha1(lib).hexdigest()[:8])
-  a, b = Buffer("QCOM", a_np.size, dtypes.half).allocate(), Buffer("QCOM", b_storage.size, dtypes.half).allocate()
-  c = Buffer("QCOM", batch*m*stride, dtypes.half).allocate()
-  q8.buf_copyin(a, memoryview(a_np).cast("B"))
-  q8.buf_copyin(b, memoryview(b_storage).cast("B"))
-  if not int(os.getenv("NO_INIT", "0")):
-    q8.buf_copyin(c, memoryview(np.zeros(batch*m*stride, dtype=np.float16)).cast("B"))
-  packed8 = bool(int(os.getenv("PACKED8", "0")))
-  packed_b8 = bool(int(os.getenv("PACKED_B8", "0")))
-  specs = ([((0, dtypes.half, (batch*m, stride//4, 4)),), ((0, dtypes.half, (batch*m, k//4, 4)),),
-            ((1, dtypes.half, (k, batch*(m//8)*n//4, 4)),) if batch_repeat_b_x else
-            ((1, dtypes.half, (k, batch*n//4, 4)),) if batch_horizontal else
-            ((1, dtypes.half, ((batch*k*(m//8) if batch_repeat_b else batch*k), n//4, 4)),)] if image_store else
-           [((0, dtypes.uint32, (m, k//8, 4)),), ((0, dtypes.uint32, (k, n//8, 4)),), ((0, dtypes.half, None),)] if packed8 else
-           [((0, dtypes.half, (m, k//4, 4)),), ((0, dtypes.uint32, (k, n//8, 4)),), ((0, dtypes.half, None),)] if packed_b8 else
-           [((0, dtypes.half, (batch*m, k//4, 4)),),
-            ((0, dtypes.half, (k, batch*n//4, 4)),) if batch_horizontal else ((0, dtypes.half, (batch*k, n//4, 4)),),
-            ((0, dtypes.half, None),)])
-  prg = dev.runtime("gemm_h", lib, buf_dtypes=specs)
-  if int(os.getenv("PRINT_META", "0")):
-    print("buffer_specs", specs)
-    print("runtime_meta", {k:v for k,v in vars(prg).items() if k in ("wgid", "lid", "max_threads", "prg")})
-  call_bufs = (c._buf, a._buf, b._buf) if image_store else (a._buf, b._buf, c._buf)
-  tile_n = 384 if bool(int(os.getenv("TRI", "0"))) else 512 if bool(int(os.getenv("WIDE", "0"))) else 256
-  tile_m = (threads//32)*8
-  global_size = ((n//tile_n, m//tile_m, batch) if batch_z else
-                 (batch*n//tile_n, m//tile_m, 1) if batch_horizontal else
-                 (n//tile_n, batch*m//tile_m, 1))
-  times = [prg(*call_bufs, global_size=global_size,
-               local_size=(threads, 1, 1), wait=True) for _ in range(int(os.getenv("BENCH_RUNS", "10")))]
-  elapsed = min(times)
-  got = np.empty(batch*m*stride, dtype=np.float16)
-  q8.buf_copyout(c, memoryview(got).cast("B"))
-  if int(os.getenv("RAW_STATS", "0")):
-    nz = np.flatnonzero(got)
-    print("raw_nonzero", nz.size, "first", nz[:64].tolist(), "last", nz[-64:].tolist(),
-          "values", np.unique(got[nz])[:16].tolist())
-  if int(os.getenv("THREAD_STORE", "0")) or int(os.getenv("DECODE_THREAD", "0")):
-    raw, matrix = got[:m*n].reshape(-1, 8, 2, 4), np.empty((m, n), np.float16)
-    if pattern:
-      print("raw_lids=", [[float(raw[lid, row, 0, 0]) for row in range(8)] for lid in range(0, 128, 8)])
-    gx_count = n//256
-    for gy in range(m//tile_m):
-      for gx in range(gx_count):
-        for lid in range(threads):
-          tm, tid = lid//32, lid%32
-          thread = (gy*gx_count+gx)*threads+lid
-          for row in range(8):
-            for col in range(2):
-              x = (gx*64+tid+col*32)*4
-              matrix[gy*tile_m+tm*8+row, x:x+4] = raw[thread, row, col]
-    got = matrix.astype(np.float32)
-  else: got = got.reshape(batch*m, stride)[:, :n].astype(np.float32)
-  if int(os.getenv("POST_SEQUENCE", "0")):
-    tile = np.empty((8, 256), np.float32)
-    for row in range(8):
-      for col in range(2): tile[row, col*128:(col+1)*128] = row*2+col+1
-    expected = np.tile(tile, (m//8, n//256))
-  else: expected = (np.full((batch*m, n), 1024, np.float32) if int(os.getenv("POST_CONSTANT", "0")) else
-              np.concatenate([a_np[x*m:(x+1)*m, k_start*4:(k_start+k_count)*4].astype(np.float32) @
-                              b_np[x*k+k_start*4:x*k+(k_start+k_count)*4].astype(np.float32)
-                              for x in range(batch)]))
-  delta = np.abs(expected-got)
-  if (reserved_out := int(os.getenv("RESERVED_OUT", "-1"))) >= 0:
-    row, col = divmod(reserved_out, 2)
-    delta[row::8, col*128:(col+1)*128] = 0
-  correct = np.allclose(expected, got, rtol=2e-2, atol=2e-2)
-  gflops = batch*2*m*n*(k_count*4)/elapsed/1e9
-  print(f"shape={batch}x{m}x{n}x{k_count*4} accumulate=fp16 elapsed_ms={elapsed*1e3:.3f} gflops={gflops:.1f} "
-        f"max_abs={delta.max():.9g} mean_abs={delta.mean():.9g} allclose={correct}")
-  bad = ~np.isfinite(got) | (delta > .02)
-  bad_idx = np.argwhere(bad)
-  print(f"bad_count={bad_idx.shape[0]}")
-  if int(os.getenv("VERBOSE", "0")):
-    for r in range(8): print(f"row{r} expected={expected[r,:8].tolist()} got={got[r,:8].tolist()}")
-    print("block_max=", [[float(delta[r:r+8, c:c+128].max()) for c in range(0, n, 128)] for r in range(0, m, 8)])
-    print("local_rows=", [(lr, float(delta[lr::8].max()), float(delta[lr::8].mean())) for lr in range(8)])
-    print("bad_by_row=", [(int(r), int(bad[r].sum())) for r in np.flatnonzero(bad.any(axis=1))])
-    print("bad_first=", [(int(r), int(c), float(expected[r, c]), float(got[r, c])) for r, c in bad_idx[:64]])
-  if int(os.getenv("POST_SEQUENCE", "0")):
-    print("sequence_blocks=", [[np.unique(got[row, col:col+128], return_counts=True) for col in range(0, n, 128)] for row in range(8)])
-  if int(os.getenv("VERBOSE", "0")):
-    row0_matches = np.abs(expected-got[0]).mean(axis=1)
-    print("row0_matches=", [(int(i), float(row0_matches[i])) for i in np.argsort(row0_matches)[:8]])
-    if batch > 1:
-      for row in range(0, batch*m, (threads//32)*8):
-        candidates = [a_np[row].astype(np.float32) @ b_np[x*k:(x+1)*k].astype(np.float32) for x in range(batch)]
-        print("batch_map=", row, [(x, float(np.abs(c-got[row]).mean())) for x, c in enumerate(candidates)])
-  if int(os.getenv("VERBOSE", "0")) and not int(os.getenv("POST_SEQUENCE", "0")) and not int(os.getenv("POST_CONSTANT", "0")):
-    contrib = np.stack([a_np[0, kk*4:kk*4+4].astype(np.float32) @
-                        b_np[kk*4:kk*4+4].astype(np.float32) for kk in range(k_start, k_start+k_count)])
-    excluded = np.abs((expected[0][None, :]-contrib)-got[0]).mean(axis=1)
-    prefixes = np.abs(np.cumsum(contrib, axis=0)-got[0]).mean(axis=1)
-    print("row0_k=", "exclude", [(k_start+int(i), float(excluded[i])) for i in np.argsort(excluded)[:4]],
-          "prefix", [(k_start+int(i)+1, float(prefixes[i])) for i in np.argsort(prefixes)[:4]])
-    if k_count == 1:
-      cs = np.stack([a_np[:, k_start*4+j:k_start*4+j+1].astype(np.float32) @
-                     b_np[k_start*4+j:k_start*4+j+1].astype(np.float32) for j in range(4)])
-      subset = [(mask, float(np.abs(sum((cs[j] for j in range(4) if mask & (1<<j)), np.zeros_like(got))-got).mean())) for mask in range(16)]
-      print("component_subsets=", sorted(subset, key=lambda x:x[1])[:8])
-  if not correct: raise SystemExit(1)
-
-
-if __name__ == "__main__": main()

extra/gemm/qcom_8x8_probe.py

@@ -1,111 +0,0 @@
-#!/usr/bin/env python3
-"""Dependency-free lane-mapping probe for the thread-major 8x8 shader."""
-import ctypes, os, random, struct
-
-from tinygrad import Device, dtypes
-from tinygrad.device import Buffer
-from extra.gemm import qcom_8x4_gemm as q8
-from extra.gemm.ir3asm import get_envelope, inject
-
-
-def half_bytes(values): return bytearray(struct.pack(f"<{len(values)}e", *values))
-
-
-def main():
-  m, n, k = int(os.getenv("M", "32")), int(os.getenv("N", "256")), int(os.getenv("K", "192"))
-  pattern = os.getenv("PATTERN", "row")
-  int8_b = bool(int(os.getenv("INT8_B", "0")))
-  a = [0.0] * (m*k)
-  b = [0.0] * (k*n)
-  if pattern == "row":
-    for row in range(m): a[row*k] = row+1
-    for col in range(n): b[col] = 1
-  elif pattern == "col":
-    for row in range(m): a[row*k] = 1
-    for col in range(n): b[col] = col % 251 + 1
-  elif pattern == "random":
-    rng = random.Random(int(os.getenv("SEED", "0")))
-    a = [rng.uniform(-0.05, 0.05) for _ in a]
-    b = [rng.uniform(-0.05, 0.05) for _ in b]
-  else: raise ValueError(pattern)
-  # The oracle must use the exact FP16 values consumed by the images.
-  a = list(struct.unpack(f"<{len(a)}e", half_bytes(a)))
-  if int8_b:
-    bq = [max(-127, min(127, round(x*127))) for x in b]
-    b = [x/127.0 for x in bq]
-    b_bytes = bytearray((x & 0xff) for x in bq)
-  else:
-    b = list(struct.unpack(f"<{len(b)}e", half_bytes(b)))
-    b_bytes = half_bytes(b)
-
-  q8.M, q8.N, q8.K, q8.K4 = m, n, k, k//4
-  dev = Device["QCOM"]
-  compiler = bool(int(os.getenv("COMPILER", "0")))
-  tight_store = bool(int(os.getenv("TIGHT_STORE", "0")))
-  mode = os.getenv("MODE", "")
-  if compiler:
-    lib, _, _, _ = get_envelope(dev, q8.make_donor_src8(2, 128))
-  elif mode:
-    env, io, sz, ro = get_envelope(dev, q8.make_donor_src8(4, 128))
-    if mode == "pipeline": shader, hregs, fregs, _, _ = q8.build_8x8_pipelined_shader(dev, 128, 4, 4, thread_store_gx=n//256)
-    elif mode == "pipeline4": shader, hregs, fregs, _, _ = q8.build_8x8_pipeline4_shader(dev, 128, 4, 4)
-    elif mode == "batch2": shader, hregs, fregs, _, _ = q8.build_8x8_batch2_shader(dev, 128, 4, 4)
-    else: raise ValueError(mode)
-    assert len(shader) <= sz
-    lib = inject(env, io, sz, ro, shader, fregs=fregs, hregs=hregs)
-  elif int(os.getenv("BASE", "0")):
-    env, io, sz, ro = get_envelope(dev, q8.make_donor_src8(4, 128))
-    shader, hregs, fregs, _ = q8.build_8x8_split_a_shader(dev, 128,
-      a_coord_delay=int(os.getenv("ADELAY", "3")), b_coord_delay=int(os.getenv("BDELAY", "3")),
-      pre_mad_nops=int(os.getenv("PMAD", "-1")), grouped_b=bool(int(os.getenv("GROUPED_B", "0"))),
-      grouped_b_cols=bool(int(os.getenv("GROUPED_COLS", "0"))), thread_store_gx=0 if tight_store else 1,
-      add256_store_mode="tight" if tight_store else "donor")
-    lib = inject(env, io, sz, ro, shader, fregs=fregs, hregs=hregs)
-  else:
-    env, io, sz, ro = get_envelope(dev, q8.make_donor_src8(4, 128))
-    shader, hregs, fregs, _ = q8.build_8x8_split_a_unroll_shader(dev, 128,
-      k_unroll=int(os.getenv("KUNROLL", "8")), b_coord_delay=int(os.getenv("BDELAY", "0")),
-      fast_coords=bool(int(os.getenv("FAST", "1"))), prefetch_next_b=bool(int(os.getenv("PREFETCH", "0"))),
-      thread_store_gx=0 if tight_store else 1, add256_store_mode="tight" if tight_store else "donor",
-      post_sequence=bool(int(os.getenv("POST_SEQUENCE", "0"))), a_coord_delay=int(os.getenv("ADELAY", "4")),
-      unroll_gap=int(os.getenv("GAP", "0")), relaxed_sync=bool(int(os.getenv("RELAXED_SYNC", "0"))),
-      sync_mask=int(os.getenv("SYNC_MASK", "7"), 0), sync_wait=int(os.getenv("SYNC_WAIT", "0")))
-    lib = inject(env, io, sz, ro, shader, fregs=fregs, hregs=hregs)
-  ab = Buffer("QCOM", len(a), dtypes.half).allocate()
-  bb = Buffer("QCOM", len(b), dtypes.int8 if int8_b else dtypes.half).allocate()
-  cb = Buffer("QCOM", m*n, dtypes.half).allocate()
-  for buf, raw in ((ab, half_bytes(a)), (bb, b_bytes), (cb, bytearray(m*n*2))):
-    src = (ctypes.c_ubyte * len(raw)).from_buffer(raw)
-    ctypes.memmove(int(buf._buf.va_addr), ctypes.addressof(src), len(raw))
-  specs = [((0, dtypes.half, (m, k//4, 4)),), ((0, dtypes.int8 if int8_b else dtypes.half, (k, n//4, 4)),),
-           ((0, dtypes.half, None),)]
-  prg = dev.runtime("gemm_h", lib, buf_dtypes=specs)
-  times = [prg(ab._buf, bb._buf, cb._buf, global_size=(n//256, m//32, 1), local_size=(128, 1, 1), wait=True)*1e3 for _ in range(5)]
-  print("elapsed_ms=", min(times))
-  out = bytearray(m*n*2)
-  ctypes.memmove(ctypes.addressof((ctypes.c_ubyte * len(out)).from_buffer(out)), int(cb._buf.va_addr), len(out))
-  raw = struct.unpack(f"<{m*n}e", out)
-  if int(os.getenv("DUMP_RAW", "0")):
-    for row in range(min(m, 16)): print("raw", row, list(raw[row*n:row*n+min(n, 64)]))
-    return
-  if pattern == "random":
-    worst = total = 0.0
-    worst_at = None
-    for row in range(m):
-      tm, rr = row//8, row%8
-      for col in range(n):
-        tid, cc, lane = (col//4)%32, col//128, col%4
-        got = raw[row*n+col] if compiler or tight_store or mode in ("pipeline4", "batch2") else raw[(tm*32+tid)*64 + rr*8 + cc*4 + lane]
-        expected = sum(a[row*k+kk] * b[kk*n+col] for kk in range(k))
-        delta = abs(got-expected)
-        if delta > worst: worst, worst_at = delta, (row, col, got, expected)
-        total += delta
-    print("max_abs=", worst, "mean_abs=", total/(m*n), "worst_at=", worst_at)
-    if worst > 0.02: raise SystemExit(1)
-    return
-  for lid in range(min(128, m*n//64)):
-    vals = [raw[lid*64+row*8] for row in range(8)]
-    print(lid, vals)
-
-
-if __name__ == "__main__": main()

extra/gemm/qcom_alu_peak.py

@@ -1,277 +0,0 @@
-#!/usr/bin/env python3
-"""Adreno 630 FP16 MAD throughput benchmark."""
-import argparse, ctypes, struct
-
-from tinygrad import Device, dtypes
-from tinygrad.device import Buffer
-from extra.gemm.ir3asm import *
-from extra.gemm.ir3asm import _hreg
-from extra.gemm.qcom_intensity_gemm import M, N, K4, make_donor_src, prologue_4x2, store_output
-
-
-def make_bufs(dev):
-    a = Buffer(dev.device, (K4)*M*4, dtypes.half, preallocate=True)
-    b = Buffer(dev.device, (N//4)*(K4*4)*4, dtypes.half, preallocate=True)
-    c = Buffer(dev.device, M*N, dtypes.half, preallocate=True)
-    ctypes.memset(int(a._buf.va_addr), 0, a.nbytes)
-    ctypes.memset(int(b._buf.va_addr), 0, b.nbytes)
-    ctypes.memset(int(c._buf.va_addr), 0, c.nbytes)
-    return a, b, c
-
-
-def emit_mov_h_block(instrs, start, end, src):
-    pos = start
-    while pos < end:
-        rpt = min(3, end - pos - 1)
-        instrs.append(MOV_H(pos, src, rpt=rpt))
-        pos += rpt + 1
-
-
-def build_compiler_pattern_shader(dev, threads, loops, pairs, store):
-    if pairs < 2: raise ValueError('compiler-pattern needs at least two x/y MAD pairs; smaller shaders have caused QCOM hangs')
-    instrs = prologue_4x2(dev, threads)
-    instrs += [MOV_S32('r8.x', 0, sy=True), MOV_H_IMM('hr0.x', 0x3c00)]
-    emit_mov_h_block(instrs, 1, _hreg('hr8.x'), 0)
-
-    loop_start = len(instrs)
-    for _ in range(pairs):
-        # This mirrors the vec16 OpenCL MAD peak lowering: one vector MAD into x,
-        # then one vector MAD into y. The split scalar lane avoids clobbering hr0.y.
-        instrs += [
-            MAD_F16('hr0.z', 'hr0.z', 'hr4.y', 'hr4.y', rpt=3, r=True, r1=True),
-            MAD_F16('hr1.z', 'hr1.z', 'hr5.y', 'hr5.y', rpt=3, r=True, r1=True),
-            MAD_F16('hr2.z', 'hr2.z', 'hr6.y', 'hr6.y', rpt=3, r=True, r1=True),
-            MAD_F16('hr3.z', 'hr3.z', 'hr7.y', 'hr7.y', rpt=2, r=True, r1=True),
-            MAD_F16('hr0.x', 'hr0.x', 'hr0.y', 'hr0.y'),
-            MAD_F16('hr4.y', 'hr0.z', 'hr4.y', 'hr0.z', rpt=3, r=True, r1=True),
-            MAD_F16('hr5.y', 'hr1.z', 'hr5.y', 'hr1.z', rpt=3, r=True, r1=True),
-            MAD_F16('hr6.y', 'hr2.z', 'hr6.y', 'hr2.z', rpt=3, r=True, r1=True),
-            MAD_F16('hr7.y', 'hr3.z', 'hr7.y', 'hr3.z', rpt=2, r=True, r1=True),
-            MAD_F16('hr0.y', 'hr0.x', 'hr0.y', 'hr0.x'),
-        ]
-    instrs += [
-        ADD_S('r8.y', 'r8.x', 1),
-        CMPS_S_EQ('r8.x', loops - 1, nop=1),
-        MOV_F32('r8.x', 'r8.y'),
-        NOP(rpt=3),
-    ]
-    loop_end = len(instrs)
-    instrs.append(BR(loop_start - loop_end))
-    if store: store_output(instrs, 'r7.x', 'r7.y', 0)
-    instrs.append(END())
-    return assemble(instrs), loop_end - loop_start, 8, 9, pairs * 32 * 2
-
-
-def build_alu_shader(dev, threads, groups, rpt, loops, unroll, independent, r1):
-    if rpt > 3: raise ValueError('mad.f16 repeat counts above rpt3 encode other flags on A630, not more FP16 lanes')
-    if not (1 <= loops <= 256): raise ValueError('loops must be in 1..256; current immediate compare encodes only 8 bits')
-    width = rpt + 1
-    instrs = prologue_4x2(dev, threads)
-    instrs += [
-        MOV_S32('r6.z', 0, sy=True),
-        MOV_H_IMM('hr0.x', 0x3c00),
-        MOV_H_IMM('hr16.x', 0), MOV_H('hr16.y', 'hr16.x', rpt=2),
-    ]
-    emit_mov_h_block(instrs, 1, max(width, 4), 0)
-    emit_mov_h_block(instrs, _hreg('hr4.x'), _hreg('hr4.x') + max(width, 4), 0)
-    acc0 = _hreg('hr16.x')
-    hregs = (acc0 + groups * width + 3) // 4
-    emit_mov_h_block(instrs, acc0 + 4, acc0 + groups * width, acc0)
-
-    loop_start = len(instrs)
-    for _ in range(unroll):
-        for g in range(groups):
-            src1 = (g * width) % max(width, 4)
-            src2 = _hreg('hr4.x') + ((g * width) % max(width, 4))
-            src3 = src1 if independent else acc0 + g * width
-            instrs.append(MAD_F16(acc0 + g * width, src1, src2, src3, rpt=rpt, r=True, r1=r1))
-    instrs += [
-        ADD_S('r0.x', 'r6.z', 1),
-        CMPS_S_EQ('r6.z', loops - 1, nop=1),
-        MOV_F32('r6.z', 'r0.x'),
-        NOP(rpt=3),
-    ]
-    loop_end = len(instrs)
-    instrs.append(BR(loop_start - loop_end))
-    store_output(instrs, 'r7.x', 'r7.y', acc0)
-    instrs.append(END())
-    return assemble(instrs), loop_end - loop_start, hregs, None, unroll * groups * width * 2
-
-
-def gemm_check_inputs(rows, ncols):
-    a = [[((row * 3 + kk) % 4 + 1) / 8 for kk in range(4)] for row in range(rows)]
-    b = [[[[((col * 7 + kk * 3 + lane) % 4 + 1) / 8 for lane in range(4)] for kk in range(4)] for col in range(ncols)]][0]
-    return a, b
-
-
-def half_raw(value):
-    return struct.unpack('<H', struct.pack('<e', value))[0]
-
-
-def build_gemm_pattern_shader(dev, threads, loops, rows, ncols, unroll, order, bmode, r1, check_pattern=False, store_group=0):
-    if loops != 1: raise ValueError('gemm-pattern is a one-shot ALU body benchmark; use --loops 1 so loop-control regs do not clobber A/B sources')
-    if rows not in (4, 8): raise ValueError('rows must be 4 or 8')
-    if ncols < 1: raise ValueError('ncols must be positive')
-    instrs = prologue_4x2(dev, threads)
-    instrs += [MOV_S32('r6.z', 0, sy=True)]
-
-    # A lives in hr0..hr(rows-1). B either reuses one 4-texel column group or
-    # allocates one 4-texel group per output col4. Accumulators start at hr16 to
-    # match the working GEMM kernels and avoid low full-register aliases.
-    a_base = 0
-    b_base = rows * 4
-    b_groups = ncols if bmode == 'percol' else 1
-    b_end = b_base + b_groups * 16
-    acc0 = max(_hreg('hr16.x'), ((b_end + 3) // 4) * 4)
-    if check_pattern:
-        check_a, check_b = gemm_check_inputs(rows, ncols)
-        for row in range(rows):
-            for kk in range(4): instrs.append(MOV_H_IMM(a_base + row * 4 + kk, half_raw(check_a[row][kk])))
-        for col in range(b_groups):
-            for kk in range(4):
-                for lane in range(4): instrs.append(MOV_H_IMM(b_base + col * 16 + kk * 4 + lane, half_raw(check_b[col][kk][lane])))
-    else:
-        instrs.append(MOV_H_IMM('hr0.x', 0x3c00))
-        emit_mov_h_block(instrs, 1, rows * 4, 0)
-        emit_mov_h_block(instrs, b_base, b_end, 0)
-    for lane in range(acc0, acc0 + rows * ncols * 4): instrs.append(MOV_H_IMM(lane, 0))
-    hregs = (max(b_end, acc0 + rows * ncols * 4) + 3) // 4
-
-    loop_start = len(instrs)
-    def emit(row, kk, col):
-        b_col = col if bmode == 'percol' else 0
-        instrs.append(MAD_F16(acc0 + (row * ncols + col) * 4, a_base + row * 4 + kk, b_base + b_col * 16 + kk * 4,
-                              acc0 + (row * ncols + col) * 4, rpt=3, r=True, r1=r1))
-    for _ in range(unroll):
-        if order == 'kk_row_col':
-            for kk in range(4):
-                for row in range(rows):
-                    for col in range(ncols): emit(row, kk, col)
-        elif order == 'kk_col_row':
-            for kk in range(4):
-                for col in range(ncols):
-                    for row in range(rows): emit(row, kk, col)
-        elif order == 'col_kk_row':
-            for col in range(ncols):
-                for kk in range(4):
-                    for row in range(rows): emit(row, kk, col)
-        elif order == 'row_kk_col':
-            for row in range(rows):
-                for kk in range(4):
-                    for col in range(ncols): emit(row, kk, col)
-        elif order == 'row_col_kk':
-            for row in range(rows):
-                for col in range(ncols):
-                    for kk in range(4): emit(row, kk, col)
-        else: raise ValueError('unknown order %s' % order)
-    instrs += [
-        ADD_S('r0.x', 'r6.z', 1),
-        CMPS_S_EQ('r6.z', loops - 1, nop=1),
-        MOV_F32('r6.z', 'r0.x'),
-        NOP(rpt=3),
-    ]
-    loop_end = len(instrs)
-    instrs.append(BR(loop_start - loop_end))
-    if check_pattern:
-        # The per-column B register bank aliases the donor prologue's r7 output
-        # coordinates. Every lane computes the same diagnostic tile, so use one
-        # common output address and bit-check the selected accumulator vector.
-        instrs += [MOV_S32('r7.x', 0), MOV_S32('r7.y', 0), NOP(rpt=2)]
-    store_output(instrs, 'r7.x', 'r7.y', acc0 + store_group * 4)
-    instrs.append(END())
-    return assemble(instrs), loop_end - loop_start, hregs, None, unroll * rows * ncols * 4 * 4 * 2
-
-
-def run(args):
-    dev = Device['QCOM']
-    env_ncols = max(4 if args.gemm_pattern else 2, args.ncols if args.gemm_pattern else 2)
-    envelope, img_off, img_sz, reg_off = get_envelope(dev, make_donor_src(env_ncols, args.threads))
-    if args.compiler_pattern:
-        shader, loop_instrs, hregs, fregs, flops_per_thread_loop = build_compiler_pattern_shader(dev, args.threads, args.loops, args.pairs, args.store)
-    elif args.gemm_pattern:
-        shader, loop_instrs, hregs, fregs, flops_per_thread_loop = build_gemm_pattern_shader(
-            dev, args.threads, args.loops, args.rows, args.ncols, args.unroll, args.order, args.bmode, args.r1, args.check_gemm)
-    else:
-        shader, loop_instrs, hregs, fregs, flops_per_thread_loop = build_alu_shader(dev, args.threads, args.groups, args.rpt, args.loops, args.unroll, args.independent, args.r1)
-        width = args.rpt + 1
-    if fregs is None: fregs = args.fregs
-    if hregs > 48 and not args.allow_invalid_regs:
-        print('skipped: groups=%d needs hregs=%d, but A630 addressable GPR half registers stop at hr47 (hregs=48).' % (args.groups, hregs))
-        return
-    if len(shader) > img_sz:
-        print('skipped: shader is %d bytes but envelope has only %d bytes.' % (len(shader), img_sz))
-        return
-    lib = inject(envelope, img_off, img_sz, reg_off, shader, fregs=fregs, hregs=hregs)
-    asm = disasm(shader)
-    reg_count = fregs + (hregs + 1) // 2
-    wave_pairs = 96 // reg_count
-    mode = 'compiler-pattern' if args.compiler_pattern else ('gemm-pattern' if args.gemm_pattern else ('independent' if args.independent else 'accumulate'))
-    print('mode=%s r1=%d rows=%d ncols=%d bmode=%s order=%s groups=%d rpt=%d width=%d unroll=%d pairs=%d fregs=%d hregs=%d reg_count=%d wave_pairs=%d loop_instrs=%d shader_instrs=%d mad=%d rpt3=%d' % (
-        mode, args.r1, args.rows, args.ncols, args.bmode, args.order, args.groups, args.rpt, args.rpt + 1, args.unroll, args.pairs, fregs, hregs, reg_count, wave_pairs, loop_instrs, len(shader)//8, asm.count('mad.f16'), asm.count('(rpt3)mad.f16')))
-    if args.disasm: print(asm)
-
-    a, b, c = make_bufs(dev)
-    # Runtime buffer metadata now carries image shape separately from the scalar dtype.
-    buf_dtypes = [((0, dtypes.half, (M, K4, 4)),), ((0, dtypes.half, (K4*4, N//4, 4)),), ((0, dtypes.half, None),)]
-    prg = dev.runtime('gemm_h', lib, buf_dtypes=buf_dtypes)
-    tile_m = (args.threads // 32) * 4
-    gs, ls = (8, M // tile_m, 1), (args.threads, 1, 1)
-    for _ in range(5): prg(a._buf, b._buf, c._buf, global_size=gs, local_size=ls, wait=True)
-    times = []
-    for _ in range(args.iters):
-        t = prg(a._buf, b._buf, c._buf, global_size=gs, local_size=ls, wait=True)
-        if t: times.append(t)
-    best = min(times)
-    median = sorted(times)[len(times) // 2]
-    total_threads = gs[0] * gs[1] * args.threads
-    flops = total_threads * args.loops * flops_per_thread_loop
-    print('%.1f GFLOPS best (%.3f ms), %.1f GFLOPS median (%.3f ms), flops=%d runs=%d' %
-          (flops / best / 1e9, best * 1e3, flops / median / 1e9, median * 1e3, flops, len(times)))
-    if args.check_gemm:
-        if not args.gemm_pattern or args.bmode != 'percol' or args.loops != 1:
-            raise ValueError('--check-gemm requires --gemm-pattern --bmode percol --loops 1')
-        check_a, check_b = gemm_check_inputs(args.rows, args.ncols)
-        checked = 0
-        for group in range(args.rows * args.ncols):
-            check_shader, _, check_hregs, _, _ = build_gemm_pattern_shader(
-                dev, args.threads, args.loops, args.rows, args.ncols, args.unroll, args.order, args.bmode, args.r1,
-                check_pattern=True, store_group=group)
-            check_lib = inject(envelope, img_off, img_sz, reg_off, check_shader, fregs=fregs, hregs=check_hregs)
-            check_prg = dev.runtime('gemm_h', check_lib, buf_dtypes=buf_dtypes)
-            check_prg(a._buf, b._buf, c._buf, global_size=gs, local_size=ls, wait=True)
-            raw = c.copyout(memoryview(bytearray(c.nbytes))).cast('H')
-            row, col = divmod(group, args.ncols)
-            expected = [half_raw(args.unroll * sum(check_a[row][kk] * check_b[col][kk][lane] for kk in range(4))) for lane in range(4)]
-            bad = next((i for i, value in enumerate(raw[:4]) if value != expected[i]), None)
-            if bad is not None:
-                got = struct.unpack('<e', struct.pack('<H', raw[bad]))[0]
-                want = struct.unpack('<e', struct.pack('<H', expected[bad]))[0]
-                raise RuntimeError('GEMM CHECK FAIL group=%d index=%d got=%r expected=%r' % (group, bad, got, want))
-            checked += 4
-        print('GEMM CHECK PASS groups=%d scalar_outputs=%d bit_exact=true shape_per_thread=%dx%dx%d' %
-              (args.rows * args.ncols, checked, args.rows, args.ncols * 4, args.unroll * 4))
-
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--groups', type=int, default=16)
-    parser.add_argument('--rpt', type=int, choices=(0, 1, 3), default=3)
-    parser.add_argument('--loops', type=int, default=K4)
-    parser.add_argument('--unroll', type=int, default=1)
-    parser.add_argument('--pairs', type=int, default=8, help='compiler-pattern vector MAD pairs per loop')
-    parser.add_argument('--independent', action='store_true', help='remove loop-carried accumulator dependency for raw FMA issue peak')
-    parser.add_argument('--r1', action='store_true', help='auto-increment mad.f16 source1 across repeat lanes')
-    parser.add_argument('--compiler-pattern', action='store_true', help='use the vec16 OpenCL peak MAD source/destination pattern')
-    parser.add_argument('--gemm-pattern', action='store_true', help='use true GEMM-style acc=A_scalar*B_half4+acc MADs')
-    parser.add_argument('--check-gemm', action='store_true', help='use nonuniform exact inputs and bit-check every GEMM accumulator')
-    parser.add_argument('--rows', type=int, choices=(4, 8), default=4)
-    parser.add_argument('--ncols', type=int, default=4)
-    parser.add_argument('--bmode', choices=('reuse', 'percol'), default='reuse')
-    parser.add_argument('--order', choices=('kk_row_col', 'kk_col_row', 'col_kk_row', 'row_kk_col', 'row_col_kk'), default='kk_row_col')
-    parser.add_argument('--store', action='store_true', help='store one result after the ALU loop')
-    parser.add_argument('--threads', type=int, choices=(64, 128, 256), default=128)
-    parser.add_argument('--fregs', type=int, default=8)
-    parser.add_argument('--iters', type=int, default=20)
-    parser.add_argument('--allow-invalid-regs', action='store_true')
-    parser.add_argument('--disasm', action='store_true')
-    run(parser.parse_args())

extra/gemm/qcom_asm_gemm.py

@@ -1,440 +0,0 @@
-#!/usr/bin/env python3
-"""Hand-assembled GEMM kernels for Adreno 630.
-
-Tests:
-1. Pure ALU kernel (MAD throughput ceiling)
-2. Pure LOAD kernel (texture throughput ceiling)
-3. Full GEMM with optimal isam/mad interleaving
-"""
-import struct, ctypes, math
-from tinygrad import Device, dtypes
-from tinygrad.device import Buffer
-from extra.gemm.ir3asm import *
-
-dev = Device['QCOM']
-
-# ============================================================
-# DONOR KERNEL: compile the 4-row GEMM for the binary envelope
-# ============================================================
-DONOR_SRC = (
-    '#pragma OPENCL EXTENSION cl_khr_fp16 : enable\n'
-    'const sampler_t smp=CLK_NORMALIZED_COORDS_FALSE|CLK_ADDRESS_CLAMP|CLK_FILTER_NEAREST;\n'
-    '__attribute__((reqd_work_group_size(128,1,1)))\n'
-    '__kernel void gemm_h(read_only image2d_t A, read_only image2d_t B, __global half *C) {\n'
-    '  int lid=get_local_id(0); int tm=lid>>5; int tn=lid&31;\n'
-    '  int row=get_group_id(1)*16+tm*4; int col4=get_group_id(0)*32+tn;\n'
-    '  half4 r0c0=(half4)(0); for(int k4=0;k4<256;k4++){\n'
-    '    half4 a=read_imageh(A,smp,(int2)(k4,row));\n'
-    '    half4 b0=read_imageh(B,smp,(int2)(col4,k4*4));\n'
-    '    r0c0+=a.xxxx*b0;\n'
-    '  }\n'
-    '  vstore4(r0c0, 0, C+row*1024+col4*4);\n'
-    '}\n'
-)
-
-# Use the 4-row GEMM as donor since it has the right metadata for image textures
-_DONOR4 = (
-    '#pragma OPENCL EXTENSION cl_khr_fp16 : enable\n'
-    'const sampler_t smp=CLK_NORMALIZED_COORDS_FALSE|CLK_ADDRESS_CLAMP|CLK_FILTER_NEAREST;\n'
-    '__attribute__((reqd_work_group_size(128,1,1)))\n'
-    '__kernel void gemm_h(read_only image2d_t A, read_only image2d_t B, __global half *C) {\n'
-    '  int lid=get_local_id(0); int tm=lid>>5; int tn=lid&31;\n'
-    '  int row=get_group_id(1)*16+tm*4; int col4=get_group_id(0)*32+tn;\n'
-    '  half4 r0c0=(half4)(0),r0c1=(half4)(0),r0c2=(half4)(0),r0c3=(half4)(0);\n'
-    '  half4 r1c0=(half4)(0),r1c1=(half4)(0),r1c2=(half4)(0),r1c3=(half4)(0);\n'
-    '  half4 r2c0=(half4)(0),r2c1=(half4)(0),r2c2=(half4)(0),r2c3=(half4)(0);\n'
-    '  half4 r3c0=(half4)(0),r3c1=(half4)(0),r3c2=(half4)(0),r3c3=(half4)(0);\n'
-    '  for (int k4=0;k4<256;k4++) {\n'
-    '    half4 ar0=read_imageh(A,smp,(int2)(k4,row));\n'
-    '    half4 ar1=read_imageh(A,smp,(int2)(k4,row+1));\n'
-    '    half4 ar2=read_imageh(A,smp,(int2)(k4,row+2));\n'
-    '    half4 ar3=read_imageh(A,smp,(int2)(k4,row+3));\n'
-    '    half4 b0=read_imageh(B,smp,(int2)(col4,k4*4));\n'
-    '    half4 b1=read_imageh(B,smp,(int2)(col4,k4*4+1));\n'
-    '    half4 b2=read_imageh(B,smp,(int2)(col4,k4*4+2));\n'
-    '    half4 b3=read_imageh(B,smp,(int2)(col4,k4*4+3));\n'
-    '    r0c0+=ar0.xxxx*b0; r0c1+=ar0.yyyy*b1; r0c2+=ar0.zzzz*b2; r0c3+=ar0.wwww*b3;\n'
-    '    r1c0+=ar1.xxxx*b0; r1c1+=ar1.yyyy*b1; r1c2+=ar1.zzzz*b2; r1c3+=ar1.wwww*b3;\n'
-    '    r2c0+=ar2.xxxx*b0; r2c1+=ar2.yyyy*b1; r2c2+=ar2.zzzz*b2; r2c3+=ar2.wwww*b3;\n'
-    '    r3c0+=ar3.xxxx*b0; r3c1+=ar3.yyyy*b1; r3c2+=ar3.zzzz*b2; r3c3+=ar3.wwww*b3;\n'
-    '  }\n'
-    '  vstore4(r0c0+r0c1+r0c2+r0c3, 0, C+row*1024+col4*4);\n'
-    '  vstore4(r1c0+r1c1+r1c2+r1c3, 0, C+(row+1)*1024+col4*4);\n'
-    '  vstore4(r2c0+r2c1+r2c2+r2c3, 0, C+(row+2)*1024+col4*4);\n'
-    '  vstore4(r3c0+r3c1+r3c2+r3c3, 0, C+(row+3)*1024+col4*4);\n'
-    '}\n'
-)
-envelope, img_off, img_sz, reg_off = get_envelope(dev, _DONOR4)
-
-M, N, K = 1024, 1024, 1024
-K4 = K // 4  # 256
-
-def make_bufs():
-    a = Buffer(dev.device, (K//4)*M*4, dtypes.half, preallocate=True)
-    b = Buffer(dev.device, (N//4)*K*4, dtypes.half, preallocate=True)
-    c = Buffer(dev.device, M*N, dtypes.half, preallocate=True)
-    ctypes.memset(int(a._buf.va_addr), 0, a.nbytes)
-    ctypes.memset(int(b._buf.va_addr), 0, b.nbytes)
-    return a, b, c
-
-def bench(lib, gs, ls, label, flops=2*1024*1024*1024, iters=20):
-    a, b, c = make_bufs()
-    try:
-        prg = dev.runtime('gemm_h', lib, buf_dtypes=[((0, dtypes.half, (M, K//4, 4)),),
-                                                     ((1, dtypes.half, (K, N//4, 4)),),
-                                                     ((2, dtypes.half, None),)])
-        for _ in range(5):
-            prg(a._buf, b._buf, c._buf, global_size=gs, local_size=ls, wait=True)
-        times = []
-        for _ in range(iters):
-            t = prg(a._buf, b._buf, c._buf, global_size=gs, local_size=ls, wait=True)
-            if t: times.append(t)
-        if times:
-            best = min(times)
-            gflops = flops / best / 1e9
-            print("  %s: %.1f GFLOPS (%.0fus)" % (label, gflops, best*1e6))
-            return gflops
-    except Exception as e:
-        print("  %s: ERROR %s" % (label, str(e)[:80]))
-    return 0
-
-# ============================================================
-# Register plan for 4-row GEMM (matching the compiled kernel):
-#
-# Address/coordinate registers (full):
-#   r0.x(0)   = lid (hardware input)
-#   r0.y(1)   = group_id(1) + lid_row_offset
-#   r0.z(2)   = tm = lid >> 5
-#   r0.w(3)   = group_id(0) + lid_col_offset
-#   r2.y(9)   = A coord (k4 value for isam)
-#   r2.z(10)  = A row0 coord
-#   r2.w(11)  = A coord duplicate
-#   r3.x(12)  = A row0+1 coord
-#   r3.y(13)  = A coord dup
-#   r3.z(14)  = A row0+2 coord
-#   r3.w(15)  = A coord dup
-#   r4.x(16)  = A row0+3 coord
-#   r4.y(17)  = B col coord
-#   r4.z(18)  = B K offset
-#   r4.w(19)  = B K offset
-#   r5.y(21)  = B col coord dup
-#   r5.w(23)  = B col coord dup
-#   r6.x(24)  = temp
-#   r6.y(25)  = k4*4 base
-#   r6.z(26)  = k4 counter
-#   r7.x(28)  = row base addr
-#   r7.y(29)  = col4 base addr
-#
-# Texture result registers (half):
-#   hr0(0-3)   = A row3 texel (or temp)
-#   hr1(4-7)   = A row2 texel
-#   hr2(8-11)  = A row1 texel
-#   hr3(12-15) = A row0 texel
-#   hr4(16-19) = B texel (shared across all rows)
-#
-# Accumulator registers (half): 64 values = 16 groups of 4
-#   Row0: hr13.z(54)-hr16.y(65) = 4 groups: K0-K3
-#   Row1: hr17.z(70)-hr20.y(81) = 4 groups  [WRONG, let me read the actual mapping]
-#
-# Actually, from the disasm the accumulator mapping is:
-#   Row0 K0: hr20.z(82),hr20.w(83),hr21.x(84),hr21.y(85)
-#   Row0 K1: hr21.z(86),hr21.w(87),hr22.x(88),hr22.y(89) 
-#   Row0 K2: hr22.z(90),hr22.w(91),hr23.x(92),hr23.y(93)
-#   Row0 K3: hr23.z(94),hr23.w(95),hr24.x(96),hr24.y(97)
-#   Row1 K0: hr24.z(98),hr24.w(99),hr25.x(100),hr25.y(101)
-#   Row1 K1: hr25.z(102),hr25.w(103),hr26.x(104),hr26.y(105)
-#   Row1 K2: hr26.z(106),hr26.w(107),hr27.x(108),hr27.y(109)
-#   Row1 K3: hr27.z(110),hr27.w(111),hr28.x(112),hr28.y(113)
-#   Row2 K0: hr28.z(114),hr28.w(115),hr29.x(116),hr29.y(117)
-#   Row2 K1: hr29.z(118),hr29.w(119),hr30.x(120),hr30.y(121)
-#   Row2 K2: (from rpt1+rpt1, noncontiguous)
-#   Row2 K3: (from rpt3)
-#   Row3 K0: hr17.z(70),hr17.w(71),hr18.x(72),hr18.y(73)
-#   ... etc
-# This is messy. Let me use a CLEAN register plan instead.
-# ============================================================
-
-# ============================================================
-# TEST 1: PURE ALU - 16 (rpt3)mad.f16 in a loop, no texture loads
-# ============================================================
-
-print("=== TEST 1: Pure ALU (MAD throughput ceiling) ===")
-
-# Accumulator regs: hr20.x(80) through hr35.w(143) = 64 half-regs = 16 groups of 4
-# Source A: hr0.x(0) - hr0.w(3)
-# Source B: hr4.x(16) - hr7.w(31) (unused, just for mad operands)
-
-alu_instrs = [
-    MOV_S32('r6.z', 0, sy=True),     # counter = 0
-    MOV_H_IMM('hr0.x', 0x3c00),      # hr0.x = 1.0 (fp16)
-    MOV_H('hr0.y', 'hr0.x', rpt=2),  # hr0.y,z,w = 1.0
-    MOV_H_IMM('hr20.x', 0),           # zero first acc
-]
-# Zero all 64 accumulator regs (hr20.x=80 through hr35.w=143)
-for base in range(84, 144, 4):
-    alu_instrs.append(MOV_H(base, 80, rpt=3))
-# Set source B regs to 1.0
-for base in range(16, 32, 4):
-    alu_instrs.append(MOV_H(base, 0, rpt=3))
-
-# Loop label will be here
-loop_start = len(alu_instrs)
-
-# 16x (rpt3)mad.f16 = 64 MADs per iteration
-for g in range(16):
-    acc = 80 + g * 4           # accumulator base: hr20.x + g*4
-    src1 = g % 4               # hr0.x, hr0.y, hr0.z, hr0.w (cycling)
-    src2 = 16 + (g % 4) * 4   # hr4.x, hr5.x, hr6.x, hr7.x
-    alu_instrs.append(MAD_F16(acc, src1, src2, acc, rpt=3, r=True))
-
-# Loop control
-alu_instrs.append(ADD_S('r6.z', 'r6.z', 1))
-alu_instrs.append(CMPS_S_EQ('r6.z', K4 - 1))
-
-loop_end = len(alu_instrs)
-alu_instrs.append(BR(loop_start - loop_end))
-
-# Epilogue: sum and store (minimal - just write something)
-alu_instrs.append(ADD_F('hr0.x', 80, 84))
-alu_instrs.append(ADD_F('hr0.y', 88, 92))
-alu_instrs.append(ADD_F('hr0.z', 96, 100))
-alu_instrs.append(ADD_F('hr0.w', 104, 108))
-alu_instrs.append(NOP(rpt=5))
-alu_instrs.append(STG_F16('r0.z', 'hr0.x'))
-alu_instrs.append(END())
-
-shader_alu = assemble(alu_instrs)
-lib_alu = inject(envelope, img_off, img_sz, reg_off, shader_alu, fregs=8, hregs=64)
-
-print("  Shader: %d instrs (loop body: %d)" % (len(alu_instrs), loop_end - loop_start))
-print("  Disasm loop body:")
-asm = disasm(shader_alu)
-lines = asm.strip().split('\n')
-for line in lines[loop_start:loop_end+2]:
-    print("    " + line[:120])
-
-total_mads = 64 * K4  # 64 MADs per iter * 256 iters
-total_threads = 128 * (M // 128) * (M // 16)  # same grid as GEMM
-total_flops = total_mads * 2 * total_threads
-bench(lib_alu, (M//128, M//16, 1), (128, 1, 1), "PURE ALU", flops=total_flops)
-
-# ============================================================
-# TEST 2: PURE LOAD - 8 isam per iteration, accumulate results
-# ============================================================
-
-print("\n=== TEST 2: Pure LOAD (texture throughput ceiling) ===")
-
-# Same coordinate setup as the real GEMM but no MAD - just isam + add
-# We reuse the donor kernel's prologue for coordinate setup.
-# Actually let's just build it from scratch with minimal coord math.
-
-load_instrs = [
-    MOV_S32('r6.y', 3, sy=True),     # k4*4 base = 3 (initial)
-    MOV_S32('r6.z', 0),               # k4 counter = 0
-    MOV_H_IMM('hr20.x', 0),           # zero accumulator
-    MOV_H('hr20.y', 'hr20.x', rpt=2), # hr20.y,z,w = 0
-    # Compute row and col4 from lid
-    MOV_F32('r0.y', 'r52.x'),         # gid1 (from hardware constant)
-    NOP(rpt=2),
-    ADD_S('r0.y', 'r0.y', 0),         # r0.y = gid1 (simplified; real kernel adds c7.y)
-]
-# Copy the coordinate setup from the compiled kernel (lines 0-20)
-# Actually this is getting complex. Let me just build a simple version:
-# Use the compiled kernel verbatim but NOP out all the MADs.
-
-# Load the full 4-row donor kernel
-donor4 = (
-    '#pragma OPENCL EXTENSION cl_khr_fp16 : enable\n'
-    'const sampler_t smp=CLK_NORMALIZED_COORDS_FALSE|CLK_ADDRESS_CLAMP|CLK_FILTER_NEAREST;\n'
-    '__attribute__((reqd_work_group_size(128,1,1)))\n'
-    '__kernel void gemm_h(read_only image2d_t A, read_only image2d_t B, __global half *C) {\n'
-    '  int lid=get_local_id(0); int tm=lid>>5; int tn=lid&31;\n'
-    '  int row=get_group_id(1)*16+tm*4; int col4=get_group_id(0)*32+tn;\n'
-    '  half4 r0c0=(half4)(0),r0c1=(half4)(0),r0c2=(half4)(0),r0c3=(half4)(0);\n'
-    '  half4 r1c0=(half4)(0),r1c1=(half4)(0),r1c2=(half4)(0),r1c3=(half4)(0);\n'
-    '  half4 r2c0=(half4)(0),r2c1=(half4)(0),r2c2=(half4)(0),r2c3=(half4)(0);\n'
-    '  half4 r3c0=(half4)(0),r3c1=(half4)(0),r3c2=(half4)(0),r3c3=(half4)(0);\n'
-    '  for (int k4=0;k4<256;k4++) {\n'
-    '    half4 ar0=read_imageh(A,smp,(int2)(k4,row));\n'
-    '    half4 ar1=read_imageh(A,smp,(int2)(k4,row+1));\n'
-    '    half4 ar2=read_imageh(A,smp,(int2)(k4,row+2));\n'
-    '    half4 ar3=read_imageh(A,smp,(int2)(k4,row+3));\n'
-    '    half4 b0=read_imageh(B,smp,(int2)(col4,k4*4));\n'
-    '    half4 b1=read_imageh(B,smp,(int2)(col4,k4*4+1));\n'
-    '    half4 b2=read_imageh(B,smp,(int2)(col4,k4*4+2));\n'
-    '    half4 b3=read_imageh(B,smp,(int2)(col4,k4*4+3));\n'
-    '    r0c0+=ar0.xxxx*b0; r0c1+=ar0.yyyy*b1; r0c2+=ar0.zzzz*b2; r0c3+=ar0.wwww*b3;\n'
-    '    r1c0+=ar1.xxxx*b0; r1c1+=ar1.yyyy*b1; r1c2+=ar1.zzzz*b2; r1c3+=ar1.wwww*b3;\n'
-    '    r2c0+=ar2.xxxx*b0; r2c1+=ar2.yyyy*b1; r2c2+=ar2.zzzz*b2; r2c3+=ar2.wwww*b3;\n'
-    '    r3c0+=ar3.xxxx*b0; r3c1+=ar3.yyyy*b1; r3c2+=ar3.zzzz*b2; r3c3+=ar3.wwww*b3;\n'
-    '  }\n'
-    '  vstore4(r0c0+r0c1+r0c2+r0c3, 0, C+row*1024+col4*4);\n'
-    '  vstore4(r1c0+r1c1+r1c2+r1c3, 0, C+(row+1)*1024+col4*4);\n'
-    '  vstore4(r2c0+r2c1+r2c2+r2c3, 0, C+(row+2)*1024+col4*4);\n'
-    '  vstore4(r3c0+r3c1+r3c2+r3c3, 0, C+(row+3)*1024+col4*4);\n'
-    '}\n'
-)
-
-lib4, io4, isz4, ro4 = get_envelope(dev, donor4)
-shader4 = bytearray(lib4[io4:io4+isz4])
-total4 = isz4 // 8
-
-# NOP out all MAD instructions
-for i in range(total4):
-    lo, hi = struct.unpack_from('<II', shader4, i*8)
-    if (hi >> 24) in (0x63, 0x73) and ((hi >> 24) & 0xF) == 3:
-        struct.pack_into('<Q', shader4, i*8, 0)
-
-lib_load = inject(lib4, io4, isz4, ro4, shader4, fregs=8, hregs=31)
-bench(lib_load, (M//128, M//16, 1), (128, 1, 1), "PURE LOAD")
-
-# ============================================================
-# TEST 3: FULL GEMM - patched 4-row kernel (sy-stripped + rpt3)
-# ============================================================
-
-print("\n=== TEST 3: Patched GEMM (sy-stripped + rpt3) ===")
-
-# Take the compiled 4-row kernel, strip extra (sy), convert to rpt3
-shader_gemm = bytearray(lib4[io4:io4+isz4])
-
-# Strip extra (sy) flags - keep only the first one
-first_sy = False
-for i in range(total4):
-    lo, hi = struct.unpack_from('<II', shader_gemm, i*8)
-    if (hi >> 24) in (0x63, 0x73) and ((hi >> 24) & 0xF) == 3 and (hi >> 28) == 7:
-        if first_sy:
-            struct.pack_into('<I', shader_gemm, i*8+4, (hi & 0x0FFFFFFF) | 0x60000000)
-        else:
-            first_sy = True
-
-# Convert eligible 4-scalar MAD groups to (rpt3)
-i = 0
-while i < total4 - 3:
-    lo0, hi0 = struct.unpack_from('<II', shader_gemm, i*8)
-    if not ((hi0 >> 24) in (0x63, 0x73) and ((hi0 >> 24) & 0xF) == 3) or (hi0 >> 8) & 0x7F > 0 or (hi0 & 0xFF) != ((lo0 >> 16) & 0xFF):
-        i += 1; continue
-    d0, s1_0 = hi0 & 0xFF, lo0 & 0xFF
-    s2_0 = ((hi0 >> 16) & 0xFF) * 2 + (((hi0 >> 8) & 0xFF) >> 7)
-    ok = True
-    for j in range(1, 4):
-        lj, hj = struct.unpack_from('<II', shader_gemm, (i+j)*8)
-        if not ((hj >> 24) in (0x63, 0x73) and ((hj >> 24) & 0xF) == 3): ok = False; break
-        dj, rpj = hj & 0xFF, (hj >> 8) & 0x7F
-        s1j, s3j = lj & 0xFF, (lj >> 16) & 0xFF
-        s2j = ((hj >> 16) & 0xFF) * 2 + (((hj >> 8) & 0xFF) >> 7)
-        if rpj != 0 or s1j != s1_0 or dj != d0+j or s2j != s2_0+j or s3j != d0+j: ok = False; break
-    if ok:
-        rb = ((hi0 >> 8) & 0x80) | 3
-        struct.pack_into('<I', shader_gemm, i*8+4, (hi0 & 0xFFFF00FF) | (rb << 8))
-        struct.pack_into('<I', shader_gemm, i*8, lo0 | 0x20000000)
-        for j in range(1, 4): struct.pack_into('<Q', shader_gemm, (i+j)*8, 0)
-        i += 4
-    else:
-        i += 1
-
-# Merge (rpt1)+(rpt1) -> (rpt3)
-for i in range(total4 - 1):
-    lo0, hi0 = struct.unpack_from('<II', shader_gemm, i*8)
-    lo1, hi1 = struct.unpack_from('<II', shader_gemm, (i+1)*8)
-    if hi0 == 0 or hi1 == 0: continue
-    if not ((hi0 >> 24) in (0x63, 0x73) and ((hi0 >> 24) & 0xF) == 3): continue
-    if not ((hi1 >> 24) in (0x63, 0x73) and ((hi1 >> 24) & 0xF) == 3): continue
-    if (hi0 >> 8) & 0x7F != 1 or (hi1 >> 8) & 0x7F != 1: continue
-    d0, d1 = hi0 & 0xFF, hi1 & 0xFF
-    s10, s11 = lo0 & 0xFF, lo1 & 0xFF
-    s20 = ((hi0 >> 16) & 0xFF) * 2 + (((hi0 >> 8) & 0xFF) >> 7)
-    s21 = ((hi1 >> 16) & 0xFF) * 2 + (((hi1 >> 8) & 0xFF) >> 7)
-    if s10 != s11 or d1 != d0 + 2 or s21 != s20 + 2: continue
-    rb = ((hi0 >> 8) & 0x80) | 3
-    struct.pack_into('<I', shader_gemm, i*8+4, (hi0 & 0xFFFF00FF) | (rb << 8))
-    struct.pack_into('<Q', shader_gemm, (i+1)*8, 0)
-
-lib_gemm = inject(lib4, io4, isz4, ro4, shader_gemm, fregs=8, hregs=31)
-
-# Count stats
-asm_gemm = disasm(shader_gemm)
-print("  mad.f16: %d, (rpt3): %d, isam: %d, (sy): %d" % (
-    asm_gemm.count('mad.f16'), asm_gemm.count('(rpt3)mad.f16'),
-    asm_gemm.count('isam'), asm_gemm.count('(sy)')))
-
-bench(lib_gemm, (M//128, M//16, 1), (128, 1, 1), "PATCHED GEMM")
-
-# ============================================================
-# TEST 4: FULL GEMM at different sizes
-# ============================================================
-
-print("\n=== TEST 4: Patched GEMM at various sizes ===")
-for dim in [512, 768, 1024, 2048]:
-    if dim % 128 != 0 or dim % 16 != 0: continue
-    K4d = dim // 4
-    src_d = donor4.replace('k4<256', 'k4<%d' % K4d)
-    for s in ['row*1024', '(row+1)*1024', '(row+2)*1024', '(row+3)*1024']:
-        src_d = src_d.replace(s, s.replace('1024', str(dim)))
-    lib_d, io_d, isz_d, ro_d = get_envelope(dev, src_d)
-    s_d = bytearray(lib_d[io_d:io_d+isz_d])
-    t_d = isz_d // 8
-    # Apply same patches
-    fsy = False
-    for i in range(t_d):
-        lo, hi = struct.unpack_from('<II', s_d, i*8)
-        if (hi >> 24) in (0x63, 0x73) and ((hi >> 24) & 0xF) == 3 and (hi >> 28) == 7:
-            if fsy: struct.pack_into('<I', s_d, i*8+4, (hi & 0x0FFFFFFF) | 0x60000000)
-            else: fsy = True
-    i = 0
-    while i < t_d - 3:
-        lo0, hi0 = struct.unpack_from('<II', s_d, i*8)
-        if not ((hi0 >> 24) in (0x63, 0x73) and ((hi0 >> 24) & 0xF) == 3) or (hi0 >> 8) & 0x7F > 0 or (hi0 & 0xFF) != ((lo0 >> 16) & 0xFF):
-            i += 1; continue
-        d0, s1_0 = hi0 & 0xFF, lo0 & 0xFF
-        s2_0 = ((hi0 >> 16) & 0xFF) * 2 + (((hi0 >> 8) & 0xFF) >> 7)
-        ok = True
-        for j in range(1, 4):
-            lj, hj = struct.unpack_from('<II', s_d, (i+j)*8)
-            if not ((hj >> 24) in (0x63, 0x73) and ((hj >> 24) & 0xF) == 3): ok = False; break
-            dj, rpj = hj & 0xFF, (hj >> 8) & 0x7F
-            s1j, s3j = lj & 0xFF, (lj >> 16) & 0xFF
-            s2j = ((hj >> 16) & 0xFF) * 2 + (((hj >> 8) & 0xFF) >> 7)
-            if rpj != 0 or s1j != s1_0 or dj != d0+j or s2j != s2_0+j or s3j != d0+j: ok = False; break
-        if ok:
-            rb = ((hi0 >> 8) & 0x80) | 3
-            struct.pack_into('<I', s_d, i*8+4, (hi0 & 0xFFFF00FF) | (rb << 8))
-            struct.pack_into('<I', s_d, i*8, lo0 | 0x20000000)
-            for j in range(1, 4): struct.pack_into('<Q', s_d, (i+j)*8, 0)
-            i += 4
-        else: i += 1
-    for i in range(t_d - 1):
-        lo0, hi0 = struct.unpack_from('<II', s_d, i*8)
-        lo1, hi1 = struct.unpack_from('<II', s_d, (i+1)*8)
-        if hi0 == 0 or hi1 == 0: continue
-        if not ((hi0 >> 24) in (0x63, 0x73) and ((hi0 >> 24) & 0xF) == 3): continue
-        if not ((hi1 >> 24) in (0x63, 0x73) and ((hi1 >> 24) & 0xF) == 3): continue
-        if (hi0 >> 8) & 0x7F != 1 or (hi1 >> 8) & 0x7F != 1: continue
-        d0v, d1v = hi0 & 0xFF, hi1 & 0xFF
-        s10, s11 = lo0 & 0xFF, lo1 & 0xFF
-        s20 = ((hi0 >> 16) & 0xFF) * 2 + (((hi0 >> 8) & 0xFF) >> 7)
-        s21 = ((hi1 >> 16) & 0xFF) * 2 + (((hi1 >> 8) & 0xFF) >> 7)
-        if s10 != s11 or d1v != d0v + 2 or s21 != s20 + 2: continue
-        rb = ((hi0 >> 8) & 0x80) | 3
-        struct.pack_into('<I', s_d, i*8+4, (hi0 & 0xFFFF00FF) | (rb << 8))
-        struct.pack_into('<Q', s_d, (i+1)*8, 0)
-    ld = inject(lib_d, io_d, isz_d, ro_d, s_d, fregs=8, hregs=31)
-    M2 = N2 = K2 = dim
-    a2 = Buffer(dev.device, (K2//4)*M2*4, dtypes.half, preallocate=True)
-    b2 = Buffer(dev.device, (N2//4)*K2*4, dtypes.half, preallocate=True)
-    c2 = Buffer(dev.device, M2*N2, dtypes.half, preallocate=True)
-    ctypes.memset(int(a2._buf.va_addr), 0, a2.nbytes)
-    ctypes.memset(int(b2._buf.va_addr), 0, b2.nbytes)
-    try:
-        prg_d = dev.runtime('gemm_h', ld, [[(0, dtypes.imageh((M2, K2//4)))], [(1, dtypes.imageh((K2, N2//4)))], [(2, dtypes.half.ptr())]])
-        gs_d = (dim//128, dim//16, 1)
-        for _ in range(5): prg_d(a2._buf, b2._buf, c2._buf, global_size=gs_d, local_size=(128,1,1), wait=True)
-        ts = []
-        for _ in range(20):
-            t = prg_d(a2._buf, b2._buf, c2._buf, global_size=gs_d, local_size=(128,1,1), wait=True)
-            if t: ts.append(t)
-        if ts:
-            best = min(ts)
-            gf = 2*dim*dim*dim / best / 1e9
-            print("  %dx%d: %.1f GFLOPS (%.1fms)" % (dim, dim, gf, best*1e3))
-    except Exception as e:
-        print("  %d: ERROR %s" % (dim, str(e)[:60]))

extra/gemm/qcom_blocked_accum_model.py

@@ -1,106 +0,0 @@
-#!/usr/bin/env python3
-"""Patch openpilot's 4x16 FP32 GEMM with FP16 K4 partials and FP32 totals."""
-import argparse, itertools, pickle, struct
-
-from tinygrad.engine.jit import create_graph_call
-from tinygrad.uop.ops import Ops, UOp
-from extra.gemm.ir3asm import BR, CMPS_S_EQ, COV_F16F32, ISAM_F16, JUMP, MAD_F16, MAD_F32, MOV_F32, MOV_H_IMM, MOV_S32, NOP, inject
-from extra.gemm.qcom_ir3_matmul_patch import plain_name
-
-
-def blocked_image(image:bytes, block:int=1, direct_branch:bool=False, outer_iters:int|None=None, no_back_edge:bool=False) -> bytes:
-  instrs = [image[i:i+8] for i in range(0, len(image), 8)]
-  if len(instrs) != 349: raise ValueError(f"expected 349 instructions, got {len(instrs)}")
-  if block not in (1, 2, 4): raise ValueError(f"block must be 1, 2, or 4, got {block}")
-  # The compiler's loop is 47..100. Preserve its coordinate arithmetic and
-  # loop control, but sample native half vectors into a disjoint register bank.
-  # Each partial vector contains four output columns. Accumulate four scalar K
-  # terms per substep. Several substeps can share one partial before promotion.
-  body = [MOV_H_IMM(f"hr{34+row}.x", 0, rpt=3) for row in range(4)]
-  for substep in range(block):
-    # Drain the preceding half MADs before reusing their texture-source
-    # registers. Reissuing ISAM into a still-live half register can deadlock.
-    if substep: body += [MOV_F32("r0.x", "r0.x", sy=True), NOP(rpt=2)]
-    body += instrs[47:55]
-    for dst, coord in zip(("hr26.x", "hr27.x", "hr28.x", "hr29.x"), ("r0.x", "r1.x", "r2.x", "r3.x")):
-      body.append(ISAM_F16(dst, coord, 1, 1))
-    body += instrs[63:71]
-    for dst, coord in zip(("hr30.x", "hr31.x", "hr32.x", "hr33.x"), ("r4.x", "r5.x", "r6.x", "r7.x")):
-      body.append(ISAM_F16(dst, coord, 0, 0))
-    first = True
-    for kk in range(4):
-      for row in range(4):
-        body.append(MAD_F16(f"hr{34+row}.x", 4*(30+row)+kk, f"hr{26+kk}.x", f"hr{34+row}.x",
-                              rpt=3, sy=first, r=True))
-        first = False
-    # Keep the compare even between substeps: besides setting p0 it provides
-    # the latency slot needed by add r0.x -> mov r12.w. The final compare below
-    # overwrites p0 before loop control.
-    if substep != block-1: body += instrs[95:100]
-  # r4 is dead after all texture operations and supplies scalar 1.0 to vector
-  # MADs, giving FP32 total += promoted_partial without a separate add opcode.
-  body.append(MOV_S32("r4.x", 0x3f800000))
-  for row in range(4): body.append(COV_F16F32(f"r{row}.x", f"hr{34+row}.x", sy=(row == 0), rpt=3, r=True))
-  for row in range(4): body.append(MAD_F32(f"r{8+row}.x", "r4.x", f"r{row}.x", f"r{8+row}.x", rpt=3, r=True))
-  loop_limit = 95 if outer_iters is None else outer_iters*block-1
-  body += instrs[95:97] + [CMPS_S_EQ("r12.w", loop_limit, nop=1)] + instrs[98:100]
-
-  out = instrs[:47] + body
-  if no_back_edge:
-    pass
-  elif block == 1 or direct_branch:
-    out.append(BR(47-len(out), inv=True))
-  else:
-    # A6xx conditional branches have a much shorter reliable backward range
-    # than unconditional jumps. Branch past a long-range jump when complete.
-    branch_index = len(out)
-    out += [BR(2, inv=False), JUMP(47-(branch_index+1))]
-  out += instrs[101:]
-  while len(out) > len(instrs) and out[-1] == NOP(): out.pop()
-  if len(out) > len(instrs): raise ValueError(f"patched shader grew beyond envelope: {len(out)} > {len(instrs)}")
-  out += [NOP()] * (len(instrs)-len(out))
-  return b"".join(out)
-
-
-def patch_lib(lib:bytes, block:int, direct_branch:bool=False, outer_iters:int|None=None, no_back_edge:bool=False) -> bytes:
-  image_off = struct.unpack_from("<I", lib, 0xc0)[0]
-  image_size = struct.unpack_from("<I", lib, 0x100)[0]
-  reg_off = struct.unpack_from("<I", lib, 0x34)[0]
-  image = blocked_image(lib[image_off:image_off+image_size], block, direct_branch, outer_iters, no_back_edge)
-  return inject(lib, image_off, image_size, reg_off, image, fregs=13, hregs=38)
-
-
-def main() -> None:
-  parser = argparse.ArgumentParser()
-  parser.add_argument("input")
-  parser.add_argument("output")
-  parser.add_argument("--global-size", default="12,8,1")
-  parser.add_argument("--block", type=int, default=1)
-  parser.add_argument("--direct-branch", action="store_true")
-  parser.add_argument("--outer-iters", type=int, help="diagnostic loop limit; normal model execution requires 96/block iterations")
-  parser.add_argument("--no-back-edge", action="store_true", help="diagnostic: execute one outer body with no loop branch")
-  args = parser.parse_args()
-  target_global = tuple(int(x) for x in args.global_size.split(","))
-  with open(args.input, "rb") as f: jit = pickle.load(f)
-  slots = [x.arg.slot for x in jit.captured.linear.toposort()
-           if x.op is Ops.BUFFER and hasattr(x.arg, "slot") and x.arg.slot >= 0]
-  UOp.unique_num = itertools.count(max(slots, default=-1)+1)
-  outer = jit.captured.linear.src[0]
-  batch = outer.src[0].src[0].src
-  cache, replacements = {}, {}
-  for call in batch:
-    if call.op is not Ops.CALL or call.src[0].op is not Ops.PROGRAM: continue
-    program = call.src[0]
-    if plain_name(program.arg.name) != "gemm_h" or tuple(program.arg.global_size) != target_global: continue
-    old_lib = program.src[3].arg
-    new_lib = cache.setdefault(old_lib, patch_lib(old_lib, args.block, args.direct_branch, args.outer_iters, args.no_back_edge))
-    replacements[call] = call.replace(src=(program.replace(src=program.src[:3]+(program.src[3].replace(arg=new_lib),)), *call.src[1:]))
-  if not replacements: raise ValueError(f"no gemm_h calls with global size {target_global}")
-  new_outer = create_graph_call([replacements.get(call, call) for call in batch])
-  jit.captured._linear = jit.captured.linear.substitute({outer:new_outer}, walk=True)
-  jit.captured.__dict__.pop("linear", None)
-  with open(args.output, "wb") as f: pickle.dump(jit, f)
-  print(f"patched {len(replacements)} calls across {len(cache)} binaries with block={args.block} direct_branch={args.direct_branch}")
-
-
-if __name__ == "__main__": main()

extra/gemm/qcom_blocked_source_check.py

@@ -1,90 +0,0 @@
-#!/usr/bin/env python3
-"""Validate and time an OpenCL blocked-half/FP32 GEMM on QCOM."""
-import argparse
-
-import numpy as np
-
-from tinygrad import Device, dtypes
-from tinygrad.device import Buffer
-
-
-def upload(values:np.ndarray, dtype) -> Buffer:
-  return Buffer("QCOM", values.size, dtype, initial_value=np.ascontiguousarray(values).tobytes())
-
-
-def source(m:int, n:int, k:int, stride:int, block4:int, linear:bool=False, ldib:bool=False) -> str:
-  assert m % 16 == 0 and n % 128 == 0 and k % (block4*4) == 0
-  image_type = "read_write image2d_t" if ldib else "read_only image1d_buffer_t" if linear else "read_only image2d_t"
-  def coord(index:str) -> str: return f"(int2)(({index})&16383,({index})>>14)"
-  def a_load(row:str) -> str: return coord(f"({row})*{k//4}+k4") if ldib else f"{row}*{k//4}+k4" if linear else f"(int2)(k4,{row})"
-  def b_load(krow:str) -> str: return coord(f"({krow})*{n//4}+col4") if ldib else f"{krow}*{n//4}+col4" if linear else f"(int2)(col4,{krow})"
-  image_args = "," if (linear or ldib) else ",smp,"
-  return f"""#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-const sampler_t smp=CLK_NORMALIZED_COORDS_FALSE|CLK_ADDRESS_CLAMP|CLK_FILTER_NEAREST;
-__attribute__((reqd_work_group_size(128,1,1)))
-__kernel void gemm_blocked({image_type} A,{image_type} B,__global float *C) {{
-  int lid=get_local_id(0), row=get_group_id(1)*16+(lid>>5)*4;
-  int col4=get_group_id(0)*32+(lid&31);
-  float4 t0=(float4)(0),t1=(float4)(0),t2=(float4)(0),t3=(float4)(0);
-  for(int kb=0;kb<{k//4};kb+={block4}) {{
-    half4 h0=(half4)(0),h1=(half4)(0),h2=(half4)(0),h3=(half4)(0);
-    #pragma unroll
-    for(int q=0;q<{block4};q++) {{
-      int k4=kb+q;
-      half4 a0=read_imageh(A{image_args}{a_load('row+0')});
-      half4 a1=read_imageh(A{image_args}{a_load('row+1')});
-      half4 a2=read_imageh(A{image_args}{a_load('row+2')});
-      half4 a3=read_imageh(A{image_args}{a_load('row+3')});
-      half4 b0=read_imageh(B{image_args}{b_load('k4*4+0')});
-      half4 b1=read_imageh(B{image_args}{b_load('k4*4+1')});
-      half4 b2=read_imageh(B{image_args}{b_load('k4*4+2')});
-      half4 b3=read_imageh(B{image_args}{b_load('k4*4+3')});
-      h0+=a0.xxxx*b0+a0.yyyy*b1+a0.zzzz*b2+a0.wwww*b3;
-      h1+=a1.xxxx*b0+a1.yyyy*b1+a1.zzzz*b2+a1.wwww*b3;
-      h2+=a2.xxxx*b0+a2.yyyy*b1+a2.zzzz*b2+a2.wwww*b3;
-      h3+=a3.xxxx*b0+a3.yyyy*b1+a3.zzzz*b2+a3.wwww*b3;
-    }}
-    t0+=convert_float4(h0);t1+=convert_float4(h1);t2+=convert_float4(h2);t3+=convert_float4(h3);
-  }}
-  vstore4(t0,0,C+(row+0)*{stride}+col4*4);vstore4(t1,0,C+(row+1)*{stride}+col4*4);
-  vstore4(t2,0,C+(row+2)*{stride}+col4*4);vstore4(t3,0,C+(row+3)*{stride}+col4*4);
-}}"""
-
-
-def main() -> None:
-  ap = argparse.ArgumentParser()
-  ap.add_argument("--m", type=int, default=128)
-  ap.add_argument("--n", type=int, default=1536)
-  ap.add_argument("--k", type=int, default=384)
-  ap.add_argument("--stride", type=int, default=2048)
-  ap.add_argument("--block4", type=int, default=4)
-  ap.add_argument("--seed", type=int, default=0)
-  ap.add_argument("--float-a", action="store_true", help="sample an FP32 activation image with read_imageh")
-  ap.add_argument("--linear", action="store_true", help="use image1d_buffer_t with explicit flattened indices")
-  ap.add_argument("--ldib", action="store_true", help="use read-write image2d_t and LDIB with flattened 2D indices")
-  args = ap.parse_args()
-  rng = np.random.default_rng(args.seed)
-  av = (rng.standard_normal((args.m, args.k))*0.05).astype(np.float32 if args.float_a else np.float16)
-  bv = (rng.standard_normal((args.k, args.n))*0.05).astype(np.float16)
-  a, b = upload(av, dtypes.float if args.float_a else dtypes.half), upload(bv, dtypes.half)
-  c = upload(np.zeros(args.m*args.stride, np.float32), dtypes.float)
-  src = source(args.m, args.n, args.k, args.stride, args.block4, args.linear, args.ldib)
-  if args.ldib:
-    ashape = ((args.m*(args.k//4)+16383)//16384, 16384, 4)
-    bshape = ((args.k*(args.n//4)+16383)//16384, 16384, 4)
-  else:
-    ashape = (1, args.m*(args.k//4), 4) if args.linear else (args.m, args.k//4, 4)
-    bshape = (1, args.k*(args.n//4), 4) if args.linear else (args.k, args.n//4, 4)
-  specs = [((0, dtypes.float if args.float_a else dtypes.half, ashape),),
-           ((1, dtypes.half, bshape),), ((2, dtypes.float, (args.m*args.stride,)),)]
-  program = Device["QCOM"].runtime("gemm_blocked", Device["QCOM"].compiler.compile(src), buf_dtypes=specs)
-  times = [program(a._buf, b._buf, c._buf, global_size=(args.n//128, args.m//16, 1),
-                   local_size=(128, 1, 1), wait=True)*1e3 for _ in range(8)]
-  storage = c.numpy().reshape(args.m, args.stride)
-  got, expected = storage[:, :args.n], av.astype(np.float32) @ bv.astype(np.float32)
-  delta = np.abs(got-expected)
-  print(f"block4={args.block4} ms={min(times):.4f} max_abs={float(delta.max()):.9g} "
-        f"mean_abs={float(delta.mean()):.9g} allclose={np.allclose(got, expected, rtol=1e-2, atol=1e-2)}")
-
-
-if __name__ == "__main__": main()

extra/gemm/qcom_cache_sweep.py

@@ -1,39 +0,0 @@
-#!/usr/bin/env python3
-"""Sweep QCOM compute texture/UAV partition registers on one captured model."""
-import argparse, os, pickle, time
-
-import numpy as np
-
-from tinygrad import Tensor
-from tinygrad.engine.realize import graph_cache
-
-
-def main() -> None:
-  parser = argparse.ArgumentParser()
-  parser.add_argument("model")
-  parser.add_argument("corpus")
-  parser.add_argument("--case", type=int, default=9)
-  parser.add_argument("--pairs", default="128:64,1:1,1:64,64:1,32:32,64:32,128:32,64:64")
-  parser.add_argument("--runs", type=int, default=5)
-  args = parser.parse_args()
-  with open(args.model, "rb") as f: model = pickle.load(f)
-  corpus = np.load(args.corpus)
-  inputs = {}
-  for name, (view, _vars, dtype, device) in zip(model.captured.expected_names, model.captured.expected_input_info):
-    key=f"case{args.case}:input:{name}"
-    inputs[name] = Tensor(corpus[key if key in corpus else name].astype(np.dtype(dtype.fmt), copy=False), device=device).realize()
-  output_key=f"case{args.case}:output"
-  expected = corpus[output_key if output_key in corpus else "out"]
-  for pair in args.pairs.split(","):
-    tsize, usize = pair.split(":")
-    os.environ["QCOM_TSIZE"], os.environ["QCOM_USIZE"] = tsize, usize
-    graph_cache.clear()
-    for _ in range(2): got = model(**inputs).numpy()
-    start = time.perf_counter()
-    for _ in range(args.runs): got = model(**inputs).numpy()
-    elapsed = (time.perf_counter()-start)*1000/args.runs
-    delta = np.abs(got.astype(np.float32)-expected.reshape(got.shape).astype(np.float32))
-    print(f"tsize={tsize} usize={usize} ms={elapsed:.3f} max_abs={float(delta.max()):.9g}")
-
-
-if __name__ == "__main__": main()

extra/gemm/qcom_cached_target3_boundary.py

@@ -1,64 +0,0 @@
-#!/usr/bin/env python3
-"""Compare the exact cached target-3 GEMM with its graph replacement."""
-import argparse, pickle
-
-import numpy as np
-
-from tinygrad import Device, dtypes
-from tinygrad.device import Buffer
-from tinygrad.uop.ops import Ops
-from extra.gemm.qcom_ir3_matmul_patch import plain_name
-
-
-def upload(x:np.ndarray, dtype) -> Buffer:
-  ret = Buffer("QCOM", x.size, dtype).allocate()
-  ret.copyin(memoryview(np.ascontiguousarray(x)).cast("B"))
-  return ret
-
-
-def read(buf:Buffer, count:int, dtype) -> np.ndarray:
-  ret = np.empty(count, dtype=dtype)
-  buf.copyout(memoryview(ret).cast("B"))
-  return ret
-
-
-def batch(model): return model.captured.linear.src[0].src[0].src[0].src
-
-
-def main() -> None:
-  parser = argparse.ArgumentParser()
-  parser.add_argument("reference")
-  parser.add_argument("candidate")
-  args = parser.parse_args()
-  with open(args.reference, "rb") as f: reference = pickle.load(f)
-  with open(args.candidate, "rb") as f: candidate = pickle.load(f)
-  ref_call = next(c for c in batch(reference) if c.op is Ops.CALL and c.src[0].op is Ops.PROGRAM and
-                  plain_name(c.src[0].arg.name) == "gemm_h" and tuple(c.src[0].arg.global_size) == (12, 8, 1))
-  cbatch = batch(candidate)
-  cand_call = next(cbatch[i] for i in range(len(cbatch)-1) if cbatch[i].op is Ops.CALL and
-                   cbatch[i].src[0].op is Ops.PROGRAM and plain_name(cbatch[i+1].src[0].arg.name) == "cached_epi3")
-  rng = np.random.default_rng(7)
-  a_np = (rng.standard_normal((128, 384))*0.05).astype(np.float16)
-  a = upload(a_np, dtypes.half)
-  ref_out = upload(np.zeros(128*2048, np.float32), dtypes.float)
-  cand_out = upload(np.zeros(128*2048, np.float16), dtypes.half)
-  dev = Device["QCOM"]
-  ref_runtime = dev.runtime("ref", ref_call.src[0].src[3].arg, buf_dtypes=ref_call.src[0].arg.aux[0])
-  cand_runtime = dev.runtime("cand", cand_call.src[0].src[3].arg, buf_dtypes=cand_call.src[0].arg.aux[0])
-  ref_runtime(a._buf, ref_call.src[2].buffer._buf, ref_out._buf,
-              global_size=ref_call.src[0].arg.global_size, local_size=ref_call.src[0].arg.local_size, wait=True)
-  cand_runtime(a._buf, cand_call.src[2].buffer._buf, cand_out._buf,
-               global_size=cand_call.src[0].arg.global_size, local_size=cand_call.src[0].arg.local_size, wait=True)
-  ref = read(ref_out, 128*2048, np.float32).reshape(128, 2048)[:, :1536]
-  got = read(cand_out, 128*2048, np.float16).reshape(128, 2048)[:, :1536].astype(np.float32)
-  weight = np.asarray(ref_call.src[2].buffer.numpy()).reshape(384, 1536)
-  cpu = a_np.astype(np.float32) @ weight.astype(np.float32)
-  delta = np.abs(got-ref)
-  at = np.unravel_index(int(delta.argmax()), delta.shape)
-  print(f"max_abs={float(delta[at]):.9g} mean_abs={float(delta.mean()):.9g} at={at} "
-        f"got={float(got[at]):.9g} reference={float(ref[at]):.9g}")
-  print(f"weight_max={float(np.max(np.abs(weight))):.9g} cpu_max={float(np.max(np.abs(cpu))):.9g} "
-        f"reference_max={float(np.max(np.abs(ref))):.9g} candidate_max={float(np.max(np.abs(got))):.9g}")
-
-
-if __name__ == "__main__": main()

extra/gemm/qcom_cached_target3_check.py

@@ -1,67 +0,0 @@
-#!/usr/bin/env python3
-"""Compare the cached exact target-3 GEMM against the THREAD128 FP16 hand kernel."""
-import pickle
-import numpy as np
-
-from tinygrad import Device, dtypes
-from tinygrad.device import Buffer
-from tinygrad.uop.ops import Ops
-from extra.gemm import qcom_intensity_gemm as q
-from extra.gemm.ir3asm import get_envelope, inject
-from extra.gemm.qcom_ir3_matmul_patch import plain_name
-
-
-def upload(x:np.ndarray, dtype) -> Buffer:
-  ret = Buffer("QCOM", x.size, dtype).allocate()
-  ret.copyin(memoryview(np.ascontiguousarray(x)).cast("B"))
-  return ret
-
-
-def read(buf:Buffer, count:int, dtype) -> np.ndarray:
-  ret = np.empty(count, dtype=dtype)
-  buf.copyout(memoryview(ret).cast("B"))
-  return ret
-
-
-with open("/data/openpilot_p3_rpt245679.pkl", "rb") as f: model = pickle.load(f)
-batch = model.captured.linear.src[0].src[0].src[0].src
-call = next(x for x in batch if x.op is Ops.CALL and x.src[0].op is Ops.PROGRAM and
-            plain_name(x.src[0].arg.name) == "gemm_h" and tuple(x.src[0].arg.global_size) == (12, 8, 1))
-dev, rng = Device["QCOM"], np.random.default_rng(7)
-a = upload((rng.standard_normal(128*384)*0.05).astype(np.float16), dtypes.half)
-a_np = read(a, 128*384, np.float16).reshape(128, 384)
-w_np = np.array(call.src[2].buffer.numpy(), copy=True).reshape(384, 384, 4).reshape(384, 1536)
-exact_out = upload(np.zeros(128*2048, np.float32), dtypes.float)
-hand_out = upload(np.zeros(128*2048, np.float16), dtypes.half)
-
-exact = dev.runtime("gemm_h", call.src[0].src[3].arg, buf_dtypes=call.src[0].arg.aux[0])
-exact(a._buf, call.src[2].buffer._buf, exact_out._buf,
-      global_size=call.src[0].arg.global_size, local_size=call.src[0].arg.local_size, wait=True)
-
-q.M, q.N, q.K, q.K4 = 128, 1536, 384, 96
-env, io, sz, ro = get_envelope(dev, q.make_direct_image_donor_src(4, 128))
-shader, _ = q.build_4xn_shader(dev, 128, ncols=4, direct=True, compact_acc=True,
-  stable_bx=True, stable_ay=True, inc_coords=True, persistent_coords=True,
-  first_sync_only=True, k_unroll=4, b_first=True, coord_delay=-1, stable_settle_delay=0,
-  store_row_shift=11, image_store=True, high_inputs=True)
-lib = inject(env, io, sz, ro, shader, fregs=10, hregs=48)
-hand = dev.runtime("gemm_h", lib, buf_dtypes=[((0, dtypes.half, (128, 512, 4)),),
-  ((0, dtypes.half, (128, 96, 4)),), ((1, dtypes.half, (384, 384, 4)),)])
-hand(hand_out._buf, a._buf, call.src[2].buffer._buf,
-     global_size=(3, 8, 1), local_size=(128, 1, 1), wait=True)
-
-expected = read(exact_out, 128*2048, np.float32).reshape(128, 2048)[:, :1536]
-got = read(hand_out, 128*2048, np.float16).reshape(128, 2048)[:, :1536].astype(np.float32)
-delta = np.abs(got-expected)
-cpu0 = a_np[0].astype(np.float32) @ w_np.astype(np.float32)
-for name, value in (("exact", expected[0]), ("hand", got[0])):
-  d_cpu = np.abs(value-cpu0)
-  print(name+"_cpu0", "max_abs", float(d_cpu.max()), "mean_abs", float(d_cpu.mean()))
-at = np.unravel_index(int(np.argmax(delta)), delta.shape)
-print("exact_hand", "max_abs", float(delta[at]), "mean_abs", float(delta.mean()), "at", at,
-      "got", float(got[at]), "expected", float(expected[at]))
-for tile in range(3):
-  d = np.abs(got[:, tile*512:(tile+1)*512]-expected[:, tile*512:(tile+1)*512])
-  print("tile", tile, "max_abs", float(d.max()), "mean_abs", float(d.mean()))
-print("timing_ms", min(hand(hand_out._buf, a._buf, call.src[2].buffer._buf,
-  global_size=(3,8,1), local_size=(128,1,1), wait=True) for _ in range(20))*1e3)

extra/gemm/qcom_compact_gemm.py

@@ -1,234 +0,0 @@
-#!/usr/bin/env python3
-"""Compact the GEMM loop by removing NOP instructions and adjusting branch offsets."""
-import struct, ctypes, tempfile
-from tinygrad import Device, dtypes
-from tinygrad.device import Buffer
-from tinygrad.runtime.autogen import mesa
-from tinygrad.helpers import data64
-
-dev = Device['QCOM']
-
-src = (
-    '#pragma OPENCL EXTENSION cl_khr_fp16 : enable\n'
-    'const sampler_t smp = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_CLAMP | CLK_FILTER_NEAREST;\n'
-    '__attribute__((reqd_work_group_size(128, 1, 1)))\n'
-    '__kernel void gemm_h(read_only image2d_t A, read_only image2d_t B, __global half *C) {\n'
-    '  int lid = get_local_id(0);\n'
-    '  int row = get_group_id(1) * 4 + (lid >> 5);\n'
-    '  int col4 = get_group_id(0) * 32 + (lid & 31);\n'
-    '  half4 acc0=(half4)(0), acc1=(half4)(0), acc2=(half4)(0), acc3=(half4)(0);\n'
-    '  for (int k4 = 0; k4 < 256; k4++) {\n'
-    '    half4 a = read_imageh(A, smp, (int2)(k4, row));\n'
-    '    half4 b0 = read_imageh(B, smp, (int2)(col4, k4*4));\n'
-    '    half4 b1 = read_imageh(B, smp, (int2)(col4, k4*4+1));\n'
-    '    half4 b2 = read_imageh(B, smp, (int2)(col4, k4*4+2));\n'
-    '    half4 b3 = read_imageh(B, smp, (int2)(col4, k4*4+3));\n'
-    '    acc0 += a.xxxx * b0;\n'
-    '    acc1 += a.yyyy * b1;\n'
-    '    acc2 += a.zzzz * b2;\n'
-    '    acc3 += a.wwww * b3;\n'
-    '  }\n'
-    '  half4 r = acc0 + acc1 + acc2 + acc3;\n'
-    '  vstore4(r, 0, C + row*1024 + col4*4);\n'
-    '}\n'
-)
-
-lib = bytearray(dev.compiler.compile_cached(src))
-image_offset = struct.unpack_from('<I', lib, 0xc0)[0]
-image_size_orig = struct.unpack_from('<I', lib, 0x100)[0]
-shader = bytearray(lib[image_offset:image_offset+image_size_orig])
-total = image_size_orig // 8
-
-def ri(buf, line):
-    off = line * 8
-    return struct.unpack_from('<I', buf, off+4)[0], struct.unpack_from('<I', buf, off)[0]
-
-def wi(buf, line, hi, lo):
-    off = line * 8
-    struct.pack_into('<I', buf, off, lo)
-    struct.pack_into('<I', buf, off+4, hi)
-
-def rn(r):
-    return "hr%d.%s" % (r // 4, "xyzw"[r % 4])
-
-def get_disasm(binary):
-    with tempfile.TemporaryFile('w+', buffering=1) as tf:
-        @ctypes.CFUNCTYPE(None, ctypes.c_void_p, ctypes.c_uint32, ctypes.c_void_p)
-        def hd(data, n, instr):
-            fst, snd = data64(ctypes.cast(instr, ctypes.POINTER(ctypes.c_uint64)).contents.value)
-            print(f"{n:04} [{fst:08x}_{snd:08x}] ", end="", flush=True, file=tf)
-        libc = ctypes.CDLL(None)
-        libc.setlinebuf(fp:=ctypes.cast(libc.fdopen(tf.fileno(), b"w"), ctypes.POINTER(mesa.struct__IO_FILE)))
-        mesa.ir3_isa_disasm(bytes(binary), len(binary), fp, mesa.struct_isa_decode_options(630, True, 0, True, pre_instr_cb=hd))
-        tf.seek(0)
-        return tf.read()
-
-# Step 1: Apply register remap (48->44, 49->45)
-for old_r, new_r in [(48, 44), (49, 45)]:
-    for i in range(total):
-        hi, lo = ri(shader, i)
-        if hi == 0 and lo == 0: continue
-        changed = False
-        if (hi & 0xFF) == old_r: hi = (hi & 0xFFFFFF00) | new_r; changed = True
-        if (lo & 0xFF) == old_r: lo = (lo & 0xFFFFFF00) | new_r; changed = True
-        if ((lo >> 16) & 0xFF) == old_r: lo = (lo & 0xFF00FFFF) | (new_r << 16); changed = True
-        if changed: wi(shader, i, hi, lo)
-
-# Step 2: Convert all eligible MAD groups to (rpt3)
-# First convert 4x scalar -> rpt3
-i = 0
-while i < total - 3:
-    hi0, lo0 = ri(shader, i)
-    if not ((hi0 >> 24) in (0x63, 0x73) and ((hi0 >> 24) & 0x0F) == 0x3): i += 1; continue
-    dst0, rpt0 = hi0 & 0xFF, (hi0 >> 8) & 0x7F
-    src1_0, src3_0 = lo0 & 0xFF, (lo0 >> 16) & 0xFF
-    src2_0 = ((hi0 >> 16) & 0xFF) * 2 + (((hi0 >> 8) & 0xFF) >> 7)
-    if rpt0 > 0 or dst0 != src3_0: i += 1; continue
-    ok = True
-    for j in range(1, 4):
-        hj, lj = ri(shader, i+j)
-        if not ((hj >> 24) in (0x63, 0x73) and ((hj >> 24) & 0x0F) == 0x3): ok = False; break
-        dj, rpj = hj & 0xFF, (hj >> 8) & 0x7F
-        s1j, s3j = lj & 0xFF, (lj >> 16) & 0xFF
-        s2j = ((hj >> 16) & 0xFF) * 2 + (((hj >> 8) & 0xFF) >> 7)
-        if rpj != 0 or s1j != src1_0 or dj != dst0+j or s2j != src2_0+j or s3j != dst0+j: ok = False; break
-    if ok:
-        rpt_byte_new = ((hi0 >> 8) & 0x80) | 3
-        hi_new = (hi0 & 0xFFFF00FF) | (rpt_byte_new << 8)
-        wi(shader, i, hi_new, lo0 | 0x20000000)
-        for j in range(1, 4): wi(shader, i+j, 0, 0)
-        i += 4
-    else:
-        i += 1
-
-# Merge (rpt1)+(rpt1) -> (rpt3)
-for i in range(total - 1):
-    hi0, lo0 = ri(shader, i)
-    hi1, lo1 = ri(shader, i+1)
-    if hi0 == 0 or hi1 == 0: continue
-    if not ((hi0 >> 24) in (0x63, 0x73) and ((hi0 >> 24) & 0x0F) == 0x3): continue
-    if not ((hi1 >> 24) in (0x63, 0x73) and ((hi1 >> 24) & 0x0F) == 0x3): continue
-    rpt0 = (hi0 >> 8) & 0x7F
-    rpt1v = (hi1 >> 8) & 0x7F
-    if rpt0 != 1 or rpt1v != 1: continue
-    dst0, dst1 = hi0 & 0xFF, hi1 & 0xFF
-    src1_0, src1_1 = lo0 & 0xFF, lo1 & 0xFF
-    src2_0 = ((hi0 >> 16) & 0xFF) * 2 + (((hi0 >> 8) & 0xFF) >> 7)
-    src2_1 = ((hi1 >> 16) & 0xFF) * 2 + (((hi1 >> 8) & 0xFF) >> 7)
-    if src1_0 != src1_1 or dst1 != dst0 + 2 or src2_1 != src2_0 + 2: continue
-    rpt_byte_new = ((hi0 >> 8) & 0x80) | 3
-    wi(shader, i, (hi0 & 0xFFFF00FF) | (rpt_byte_new << 8), lo0)
-    wi(shader, i+1, 0, 0)
-
-# Step 3: COMPACT - remove NOP instructions from the loop body
-# Find the branch and loop target
-branch_line = None
-for i in range(total):
-    hi, lo = ri(shader, i)
-    if (hi >> 20) == 0x009:
-        branch_line = i
-        br_offset_raw = lo
-        br_offset = struct.unpack('<i', struct.pack('<I', lo))[0]
-        target_line = i + 1 + br_offset
-
-if branch_line is None:
-    print("ERROR: no branch found")
-    exit(1)
-
-print("Branch at line %d, target line %d (offset %d)" % (branch_line, target_line, br_offset))
-
-# Count NOPs in the LOOP (between target_line and branch_line inclusive)
-loop_nops = []
-for i in range(target_line, branch_line + 1):
-    hi, lo = ri(shader, i)
-    if hi == 0 and lo == 0:
-        loop_nops.append(i)
-
-print("Loop body: lines %d-%d (%d instrs), %d NOPs to remove" % (
-    target_line, branch_line, branch_line - target_line + 1, len(loop_nops)))
-
-# Build new instruction stream: remove NOPs from the loop body
-# Also need to handle: some "NOPs" are actually (nop2), (nop3) etc which are
-# instruction modifiers, not standalone NOPs. Only remove pure 00000000_00000000 NOPs.
-new_instrs = []
-old_to_new = {}  # map old line numbers to new line numbers
-
-for i in range(total):
-    hi, lo = ri(shader, i)
-    # Remove pure NOPs that are inside the loop
-    if hi == 0 and lo == 0 and target_line <= i <= branch_line:
-        continue  # skip this NOP
-    old_to_new[i] = len(new_instrs)
-    new_instrs.append((hi, lo))
-
-new_total = len(new_instrs)
-print("Compacted: %d -> %d instructions (removed %d)" % (total, new_total, total - new_total))
-
-# Fix the branch offset
-if branch_line in old_to_new and target_line in old_to_new:
-    new_branch = old_to_new[branch_line]
-    new_target = old_to_new[target_line]
-    new_br_offset = new_target - new_branch - 1
-    # Update the branch instruction
-    br_hi, br_lo = new_instrs[new_branch]
-    new_instrs[new_branch] = (br_hi, struct.unpack('<I', struct.pack('<i', new_br_offset))[0])
-    print("Branch: old offset %d -> new offset %d" % (br_offset, new_br_offset))
-
-# Build new shader binary - KEEP SAME SIZE by padding with NOPs at the end
-new_shader = bytearray()
-for hi, lo in new_instrs:
-    new_shader += struct.pack('<II', lo, hi)
-
-# Pad to original size with end + nop instructions
-while len(new_shader) < image_size_orig:
-    new_shader += struct.pack('<II', 0x00000000, 0x00000000)  # nop padding
-
-new_image_size = image_size_orig  # keep same size!
-print("New shader: %d bytes = %d real instrs + %d padding" % (new_image_size, new_total, (image_size_orig - new_total*8)//8))
-
-# Don't resize - just replace shader in-place
-lib_new = bytearray(lib)
-lib_new[image_offset:image_offset+image_size_orig] = new_shader
-# image_size stays the same - no need to update
-
-# Verify disassembly
-print("\n=== COMPACTED KERNEL ===")
-asm = get_disasm(bytes(new_shader))
-mad_count = asm.count('mad.f16')
-rpt3_count = asm.count('(rpt3)mad.f16')
-isam_count = asm.count('isam')
-nop_count = asm.count('nop')
-print("instrs=%d mad=%d rpt3=%d isam=%d nop=%d" % (new_total, mad_count, rpt3_count, isam_count, nop_count))
-
-for line in asm.strip().split('\n'):
-    if line.strip():
-        print(line[:120])
-
-# Benchmark
-a_imgdt = dtypes.imageh((1024, 256))
-b_imgdt = dtypes.imageh((1024, 256))
-a_buf = Buffer(dev.device, 256*1024*4, dtypes.half, preallocate=True)
-b_buf = Buffer(dev.device, 256*1024*4, dtypes.half, preallocate=True)
-c_buf = Buffer(dev.device, 1024*1024, dtypes.half, preallocate=True)
-ctypes.memset(int(a_buf._buf.va_addr), 0, a_buf.nbytes)
-ctypes.memset(int(b_buf._buf.va_addr), 0, b_buf.nbytes)
-
-try:
-    prg = dev.runtime('gemm_h', bytes(lib_new), [[(0, a_imgdt)], [(1, b_imgdt)], [(2, dtypes.half.ptr())]])
-    gs = (1024 // 128, 1024 // 4, 1)
-    ls = (128, 1, 1)
-
-    for _ in range(5):
-        prg(a_buf._buf, b_buf._buf, c_buf._buf, global_size=gs, local_size=ls, wait=True)
-
-    times = []
-    for _ in range(30):
-        t = prg(a_buf._buf, b_buf._buf, c_buf._buf, global_size=gs, local_size=ls, wait=True)
-        if t: times.append(t)
-
-    if times:
-        best = min(times)
-        gflops = 2 * 1024 * 1024 * 1024 / best / 1e9
-        print("\n*** COMPACTED: %.1f GFLOPS (%.0fus) ***" % (gflops, best * 1e6))
-except Exception as e:
-    print("ERROR: %s" % str(e)[:200])

extra/gemm/qcom_compare_call_outputs.py

@@ -1,59 +0,0 @@
-#!/usr/bin/env python3
-"""Compare captured buffers after selected calls in two OpenPilot pickles."""
-import argparse
-import pickle
-
-import numpy as np
-
-from tinygrad import Tensor
-from tinygrad.engine.jit import _prepare_jit_inputs, create_graph_call
-from tinygrad.engine.realize import resolve_params, run_linear
-from tinygrad.uop.ops import Ops, UOp
-from extra.gemm.qcom_ir3_matmul_patch import plain_name
-
-
-def capture(model, corpus, name: str, arg_index: int) -> np.ndarray:
-  inputs = {key: Tensor(corpus[key], device=device).realize()
-            for key, (_view, _vars, _dtype, device) in zip(model.captured.expected_names, model.captured.expected_input_info)}
-  input_uops, var_vals, _names, _info = _prepare_jit_inputs((), inputs)
-  batch = model.captured.linear.src[0].src[0].src[0].src
-  index, call = next((i, call) for i, call in enumerate(batch) if call.op is Ops.CALL and
-                     call.src[0].op is Ops.PROGRAM and plain_name(call.src[0].arg.name) == name)
-  run_linear(UOp(Ops.LINEAR, src=(create_graph_call(list(batch[:index+1])),)), var_vals,
-             input_uops=input_uops, jit=True, wait=True)
-  resolved = resolve_params(call, tuple(input_uops))
-  output = resolved[call.src[0].arg.outs[0] if arg_index < 0 else arg_index]
-  return output.buffer.numpy().copy()
-
-
-def main() -> None:
-  parser = argparse.ArgumentParser()
-  parser.add_argument("left")
-  parser.add_argument("left_name")
-  parser.add_argument("right")
-  parser.add_argument("right_name")
-  parser.add_argument("corpus")
-  parser.add_argument("--left-arg", type=int, default=-1)
-  parser.add_argument("--right-arg", type=int, default=-1)
-  args = parser.parse_args()
-  with open(args.left, "rb") as f:
-    left = pickle.load(f)
-  with open(args.right, "rb") as f:
-    right = pickle.load(f)
-  corpus = np.load(args.corpus)
-  a = capture(left, corpus, args.left_name, args.left_arg)
-  b = capture(right, corpus, args.right_name, args.right_arg)
-  if a.size == 32*1088*4 and b.size == 2048*16*4:
-    image, expected = a.reshape(32, 1088, 4), np.empty((2048, 16, 4), dtype=a.dtype)
-    for row in range(2048):
-      idx1, block = row >> 2, row & 3
-      expected[row] = image[idx1 >> 4, (idx1 & 15)*68+block*17:(idx1 & 15)*68+block*17+16]
-    a = expected.reshape(-1)
-  delta = np.abs(a.astype(np.float32)-b.astype(np.float32))
-  print("shape", a.shape, b.shape, "max", float(delta.max()), "mean", float(delta.mean()))
-  print("left", a[:32])
-  print("right", b[:32])
-
-
-if __name__ == "__main__":
-  main()

extra/gemm/qcom_compare_pickles.py

@@ -1,71 +0,0 @@
-#!/usr/bin/env python3
-"""Compare two compiled-model pickles on identical deterministic inputs."""
-import argparse, os, pickle, time
-
-import numpy as np
-
-from tinygrad import Tensor
-from tinygrad.engine.realize import graph_cache
-
-
-def main() -> None:
-  parser = argparse.ArgumentParser()
-  parser.add_argument("reference")
-  parser.add_argument("candidates", nargs="+")
-  parser.add_argument("--seeds", default="123", help="comma-separated deterministic input seeds")
-  parser.add_argument("--scale", type=float, default=1.0, help="scale applied to generated normal inputs")
-  parser.add_argument("--rtol", type=float, default=1e-2)
-  parser.add_argument("--atol", type=float, default=1e-2)
-  parser.add_argument("--runs", type=int, default=5)
-  parser.add_argument("--corpus", help="NPZ corpus with caseN:input:name arrays; seed values select case indices")
-  parser.add_argument("--candidate-constlen", type=int)
-  parser.add_argument("--corpus-output", action="store_true", help="compare with corpus out/caseN:output instead of rerun reference")
-  args = parser.parse_args()
-  with open(args.reference, "rb") as f: reference = pickle.load(f)
-  seeds = [int(x) for x in args.seeds.split(",")]
-  corpus = np.load(args.corpus) if args.corpus else None
-
-  def make_inputs(seed):
-    rng = np.random.default_rng(seed)
-    inputs = {}
-    for name, (view, _vars, dtype, device) in zip(reference.captured.expected_names, reference.captured.expected_input_info):
-      corpus_key = f"case{seed}:input:{name}"
-      arr = (corpus[corpus_key if corpus_key in corpus else name].astype(np.dtype(dtype.fmt), copy=False) if corpus is not None else
-             (rng.standard_normal(view.shape)*args.scale).astype(np.dtype(dtype.fmt)))
-      inputs[name] = Tensor(arr, device=device).realize()
-    return inputs
-
-  def run(model, inputs):
-    for _ in range(2): out = model(**inputs).numpy()
-    start = time.perf_counter()
-    for _ in range(args.runs): out = model(**inputs).numpy()
-    return np.array(out, copy=True), (time.perf_counter()-start)*1000.0/args.runs
-
-  inputs_by_seed = [make_inputs(seed) for seed in seeds]
-  refs, ref_times = zip(*(run(reference, inputs) for inputs in inputs_by_seed))
-  if args.corpus_output:
-    if corpus is None: raise ValueError("--corpus-output requires --corpus")
-    refs = tuple(np.asarray(corpus[f"case{seed}:output" if f"case{seed}:output" in corpus else "out"]) for seed in seeds)
-  print(f"reference_ms={np.mean(ref_times):.3f} seeds={seeds} scale={args.scale:g}")
-  if args.candidate_constlen is not None:
-    os.environ["QCOM_CONSTLEN"] = str(args.candidate_constlen)
-    graph_cache.clear()
-  failed = False
-  for candidate_path in args.candidates:
-    with open(candidate_path, "rb") as f: candidate = pickle.load(f)
-    results = [run(candidate, inputs) for inputs in inputs_by_seed]
-    got_times = [x[1] for x in results]
-    deltas = [np.abs(ref-got) for ref, (got, _) in zip(refs, results)]
-    closes = [np.allclose(ref, got, rtol=args.rtol, atol=args.atol) for ref, (got, _) in zip(refs, results)]
-    worst_seed = int(np.argmax([x.max() for x in deltas]))
-    worst = np.unravel_index(np.argmax(deltas[worst_seed]), deltas[worst_seed].shape)
-    print(f"candidate={candidate_path} candidate_ms={np.mean(got_times):.3f}")
-    print(f"max_abs={max(x.max() for x in deltas):.9g} mean_abs={np.mean([x.mean() for x in deltas]):.9g} "
-          f"allclose={all(closes)} per_seed={closes}")
-    print(f"worst_seed={seeds[worst_seed]} worst={worst} reference={refs[worst_seed][worst]!r} "
-          f"candidate={results[worst_seed][0][worst]!r}")
-    failed |= not all(closes)
-  if failed: raise SystemExit(1)
-
-
-if __name__ == "__main__": main()

extra/gemm/qcom_compiler_isam_schedule.py

@@ -1,28 +0,0 @@
-#!/usr/bin/env python3
-"""Disassemble the proprietary compiler's consecutive image-read schedule."""
-import struct
-
-from tinygrad import Device
-from extra.gemm.ir3asm import disasm
-
-
-SRC = r"""#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-const sampler_t smp=CLK_NORMALIZED_COORDS_FALSE|CLK_ADDRESS_CLAMP|CLK_FILTER_NEAREST;
-__attribute__((reqd_work_group_size(128,1,1)))
-__kernel void reads(read_only image2d_t X,__global half *O) {
-  int x=get_global_id(0), y=get_group_id(1)*4;
-  half4 a=read_imageh(X,smp,(int2)(x,y+0));
-  half4 b=read_imageh(X,smp,(int2)(x,y+1));
-  half4 c=read_imageh(X,smp,(int2)(x,y+2));
-  half4 d=read_imageh(X,smp,(int2)(x,y+3));
-  vstore4(a+b+c+d,0,O+x*4+y*4096);
-}"""
-
-
-def main() -> None:
-  lib=Device["QCOM"].compiler.compile(SRC)
-  off,size=struct.unpack_from("<I",lib,0xc0)[0],struct.unpack_from("<I",lib,0x100)[0]
-  print(disasm(lib[off:off+size]))
-
-
-if __name__ == "__main__": main()

extra/gemm/qcom_compiler_shapes.py

@@ -1,48 +0,0 @@
-#!/usr/bin/env python3
-import argparse, ctypes
-
-from tinygrad import Device, dtypes
-from tinygrad.device import Buffer
-from extra.gemm.qcom_8x4_gemm import M, N, K, check_all_ones, fill_half, make_donor_src8
-
-
-def make_bufs(dev):
-    a = Buffer(dev.device, (K//4)*M*4, dtypes.half, preallocate=True)
-    b = Buffer(dev.device, (N//4)*K*4, dtypes.half, preallocate=True)
-    c = Buffer(dev.device, M*N, dtypes.half, preallocate=True)
-    if hasattr(a._buf, 'va_addr'):
-        ctypes.memset(int(a._buf.va_addr), 0, a.nbytes)
-        ctypes.memset(int(b._buf.va_addr), 0, b.nbytes)
-        ctypes.memset(int(c._buf.va_addr), 0, c.nbytes)
-    return a, b, c
-
-
-def run(args):
-    dev = Device[Device.DEFAULT]
-    src = make_donor_src8(args.ncols, args.threads)
-    lib = dev.compiler.compile_cached(src)
-    a_img, b_img = dtypes.imageh((M, K//4)), dtypes.imageh((K, N//4))
-    a, b, c = make_bufs(dev)
-    fill_half(a, 0x3c00)
-    fill_half(b, 0x3c00)
-    prg = dev.runtime('gemm_h', lib, [[(0, a_img)], [(1, b_img)], [(2, dtypes.half.ptr())]])
-    tile_m = (args.threads // 32) * 8
-    gs, ls = (N // (128 * args.ncols), M // tile_m, 1), (args.threads, 1, 1)
-    prg(a._buf, b._buf, c._buf, global_size=gs, local_size=ls, wait=True)
-    if not check_all_ones(c): return
-    for _ in range(args.warmup): prg(a._buf, b._buf, c._buf, global_size=gs, local_size=ls, wait=True)
-    times = []
-    for _ in range(args.iters):
-        t = prg(a._buf, b._buf, c._buf, global_size=gs, local_size=ls, wait=True)
-        if t: times.append(t)
-    best = min(times)
-    print('compiler8 ncols=%d scalar_tile=8x%d threads=%d %.1f GFLOPS (%.3f ms)' % (args.ncols, args.ncols * 4, args.threads, 2*M*N*K / best / 1e9, best * 1e3))
-
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--ncols', type=int, choices=(1, 2, 4), default=2)
-    parser.add_argument('--threads', type=int, choices=(128, 256), default=128)
-    parser.add_argument('--warmup', type=int, default=5)
-    parser.add_argument('--iters', type=int, default=20)
-    run(parser.parse_args())

extra/gemm/qcom_coord_delay_model.py

@@ -1,55 +0,0 @@
-#!/usr/bin/env python3
-"""Remove redundant coordinate-settle repeats from cached openpilot GEMMs."""
-import argparse, itertools, pickle, struct
-
-from tinygrad.engine.jit import create_graph_call
-from tinygrad.uop.ops import Ops, UOp
-from extra.gemm.ir3asm import NOP
-from extra.gemm.qcom_ir3_matmul_patch import plain_name
-
-
-def patch_lib(lib:bytes, keep_first:bool) -> bytes:
-  image_off = struct.unpack_from("<I", lib, 0xc0)[0]
-  image_size = struct.unpack_from("<I", lib, 0x100)[0]
-  instrs = [lib[image_off+i:image_off+i+8] for i in range(0, image_size, 8)]
-  if len(instrs) != 349: raise ValueError(f"expected 349 instructions, got {len(instrs)}")
-  delay_indices = (55, 57, 59, 61, 71, 73, 75, 77)
-  for position, index in enumerate(delay_indices):
-    if instrs[index] != NOP(rpt=4): raise ValueError(f"unexpected instruction at delay {index}: {instrs[index].hex()}")
-    if not (keep_first and position in (0, 4)): instrs[index] = NOP()
-  ret = bytearray(lib)
-  ret[image_off:image_off+image_size] = b"".join(instrs)
-  return bytes(ret)
-
-
-def main() -> None:
-  parser = argparse.ArgumentParser()
-  parser.add_argument("input")
-  parser.add_argument("output")
-  parser.add_argument("--global-size", default="12,8,1")
-  parser.add_argument("--keep-first", action="store_true")
-  args = parser.parse_args()
-  target_global = tuple(int(x) for x in args.global_size.split(","))
-  with open(args.input, "rb") as f: jit = pickle.load(f)
-  slots = [x.arg.slot for x in jit.captured.linear.toposort()
-           if x.op is Ops.BUFFER and hasattr(x.arg, "slot") and x.arg.slot >= 0]
-  UOp.unique_num = itertools.count(max(slots, default=-1)+1)
-  outer, cache, replacements = jit.captured.linear.src[0], {}, {}
-  batch = outer.src[0].src[0].src
-  for call in batch:
-    if call.op is not Ops.CALL or call.src[0].op is not Ops.PROGRAM: continue
-    program = call.src[0]
-    if plain_name(program.arg.name) != "gemm_h" or tuple(program.arg.global_size) != target_global: continue
-    old_lib = program.src[3].arg
-    new_lib = cache.setdefault(old_lib, patch_lib(old_lib, args.keep_first))
-    new_program = program.replace(src=program.src[:3]+(program.src[3].replace(arg=new_lib),))
-    replacements[call] = call.replace(src=(new_program, *call.src[1:]))
-  if not replacements: raise ValueError(f"no gemm_h calls with global size {target_global}")
-  new_outer = create_graph_call([replacements.get(call, call) for call in batch])
-  jit.captured._linear = jit.captured.linear.substitute({outer:new_outer}, walk=True)
-  jit.captured.__dict__.pop("linear", None)
-  with open(args.output, "wb") as f: pickle.dump(jit, f)
-  print(f"patched {len(replacements)} calls across {len(cache)} binaries keep_first={args.keep_first}")
-
-
-if __name__ == "__main__": main()

extra/gemm/qcom_distributed_fp32_check.py

@@ -1,111 +0,0 @@
-#!/usr/bin/env python3
-"""Two-device 1024 GEMM: row-partitioned FP16 inputs, true FP32 MAD accumulation, full oracle."""
-import json, os, subprocess, tempfile, time
-from concurrent.futures import ThreadPoolExecutor
-from pathlib import Path
-
-import numpy as np
-
-
-N = 1024
-REMOTE_REPO = os.getenv("REMOTE_REPO", "/data/openpilot/tinygrad_repo")
-
-
-def worker() -> None:
-  from tinygrad import Device, dtypes
-  from tinygrad.device import Buffer
-  from extra.gemm import qcom_intensity_gemm as q
-  from extra.gemm.ir3asm import get_envelope, inject
-
-  first, rows, seed = int(os.environ["ROW_START"]), int(os.environ["ROWS"]), int(os.getenv("SEED", "1001"))
-  if rows != N//2 or first not in (0, N//2): raise ValueError("worker partition must be one half of N=1024")
-  rng = np.random.default_rng(seed)
-  a_full = (rng.standard_normal((N, N), dtype=np.float32)*np.float32(1/32)).astype(np.float16)
-  b_np = (rng.standard_normal((N, N), dtype=np.float32)*np.float32(1/32)).astype(np.float16)
-  a_np = np.ascontiguousarray(a_full[first:first+rows])
-  q.M, q.N, q.K, q.K4 = rows, N, N, N//4
-  dev = Device["QCOM"]
-  env, io, sz, ro = get_envelope(dev, q.make_direct_image_donor_src(2, 64))
-  shader, hregs, fregs, loop_instrs = q.build_4x8_fp32_rotate_shader(
-    dev, 64, k_count=N//4, k_unroll=3)
-  lib = inject(env, io, sz, ro, shader, fregs=fregs, hregs=hregs, mergedregs=False)
-
-  def upload(x: np.ndarray, dtype):
-    ret = Buffer("QCOM", x.size, dtype).allocate()
-    ret.copyin(memoryview(x).cast("B"))
-    return ret
-
-  a, b = upload(a_np, dtypes.half), upload(b_np, dtypes.half)
-  c = Buffer("QCOM", rows*N, dtypes.float).allocate()
-  prg = dev.runtime("gemm_h", lib, buf_dtypes=[
-    ((0, dtypes.float, (rows, N//4, 4)),), ((0, dtypes.half, (rows, N//4, 4)),),
-    ((1, dtypes.half, (N, N//4, 4)),)])
-  for _ in range(2):
-    prg(c._buf, a._buf, b._buf, global_size=(N//256, rows//8, 1), local_size=(64, 1, 1), wait=True)
-  if start_ns := int(os.getenv("START_TIME_NS", "0")):
-    delay = (start_ns-time.time_ns())/1e9
-    if delay > 0: time.sleep(delay)
-  times = [prg(c._buf, a._buf, b._buf, global_size=(N//256, rows//8, 1),
-               local_size=(64, 1, 1), wait=True) for _ in range(int(os.getenv("BENCH_RUNS", "20")))]
-  got = np.empty((rows, N), np.float32); c.copyout(memoryview(got).cast("B"))
-  expected = a_np.astype(np.float32) @ b_np.astype(np.float32)
-  delta = np.abs(got-expected)
-  bad = ~np.isclose(got, expected, rtol=1e-3, atol=8e-3)
-  output = Path(os.getenv("PART_OUTPUT", f"/tmp/qcom_fp32_rows_{first}.npy"))
-  np.save(output, got)
-  result = {"first": first, "rows": rows, "elapsed": min(times), "times": times, "bad_count": int(bad.sum()),
-            "max_abs": float(delta.max()), "mean_abs": float(delta.mean()), "fregs": fregs,
-            "loop_instrs": loop_instrs, "output": str(output)}
-  print("RESULT_JSON="+json.dumps(result, sort_keys=True))
-  if bad.any(): raise SystemExit(1)
-
-
-def coordinator() -> None:
-  hosts = os.getenv("QCOM_HOSTS", "tc3,tc4").split(",")
-  if len(hosts) != 2: raise ValueError("QCOM_HOSTS must name exactly two devices")
-  ssh_opts = ["-o", "ConnectTimeout=8", "-o", "BatchMode=yes"]
-  here = Path(__file__).resolve()
-  kernel = here.with_name("qcom_intensity_gemm.py")
-  for host in hosts:
-    subprocess.run(["scp", "-q", *ssh_opts, str(here), str(kernel), f"{host}:{REMOTE_REPO}/extra/gemm/"], check=True,
-                   timeout=15)
-  start_ns = time.time_ns()+15_000_000_000
-
-  def launch(item: tuple[str, int]) -> tuple[str, dict]:
-    host, first = item
-    remote_output = f"/tmp/qcom_fp32_rows_{first}.npy"
-    cmd = (f"cd {REMOTE_REPO} && PYTHONPATH=. DEV=QCOM IMAGE=1 FLOAT16=1 HCQ2=1 MODE=worker "
-           f"ROW_START={first} ROWS={N//2} SEED={int(os.getenv('SEED', '1001'))} PART_OUTPUT={remote_output} "
-           f"START_TIME_NS={start_ns} BENCH_RUNS={int(os.getenv('BENCH_RUNS', '20'))} "
-           f".venv/bin/python extra/gemm/{here.name}")
-    done = subprocess.run(["ssh", *ssh_opts, host, cmd], check=True, text=True, capture_output=True, timeout=60)
-    line = next(x for x in done.stdout.splitlines() if x.startswith("RESULT_JSON="))
-    result = json.loads(line.removeprefix("RESULT_JSON="))
-    return host, result
-
-  with ThreadPoolExecutor(max_workers=2) as pool:
-    results = list(pool.map(launch, zip(hosts, (0, N//2))))
-  with tempfile.TemporaryDirectory() as tmp:
-    parts = []
-    for host, result in results:
-      local = Path(tmp)/f"part_{result['first']}.npy"
-      subprocess.run(["scp", "-q", *ssh_opts, f"{host}:{result['output']}", str(local)], check=True, timeout=15)
-      parts.append((result["first"], np.load(local), result))
-    parts.sort()
-    got = np.concatenate([x[1] for x in parts])
-  rng = np.random.default_rng(int(os.getenv("SEED", "1001")))
-  a_np = (rng.standard_normal((N, N), dtype=np.float32)*np.float32(1/32)).astype(np.float16)
-  b_np = (rng.standard_normal((N, N), dtype=np.float32)*np.float32(1/32)).astype(np.float16)
-  expected = a_np.astype(np.float32) @ b_np.astype(np.float32)
-  delta = np.abs(got-expected); bad = ~np.isclose(got, expected, rtol=1e-3, atol=8e-3)
-  paired = [max(parts[0][2]["times"][i], parts[1][2]["times"][i]) for i in range(len(parts[0][2]["times"]))]
-  best_i = int(np.argmin(paired)); elapsed = paired[best_i]
-  print(f"shape={N}x{N}x{N} devices={','.join(hosts)} inputs=fp16 accumulate=fp32 elapsed_ms={elapsed*1e3:.3f} "
-        f"gflops={2*N**3/elapsed/1e9:.1f} outputs={N*N} bad_count={int(bad.sum())} "
-        f"max_abs={float(delta.max()):.9g} mean_abs={float(delta.mean()):.9g} allclose={not bool(bad.any())} "
-        f"part_ms={[round(x[2]['times'][best_i]*1e3, 3) for x in parts]} paired_iteration={best_i}")
-  if bad.any() or 2*N**3/elapsed/1e9 <= 400: raise SystemExit(1)
-
-
-if __name__ == "__main__":
-  worker() if os.getenv("MODE") == "worker" else coordinator()

extra/gemm/qcom_dp4_check.py

@@ -1,75 +0,0 @@
-#!/usr/bin/env python3
-"""Exact randomized oracle and benchmark for A630 packed UINT8 dp4acc GEMM."""
-import os, random, struct
-
-from tinygrad import Device, dtypes
-from tinygrad.device import Buffer
-from extra.gemm import qcom_intensity_gemm as q
-from extra.gemm.ir3asm import get_envelope, inject
-
-
-def pack4(xs): return sum((int(x) & 0xff) << (8*i) for i, x in enumerate(xs))
-
-
-def main():
-  m, n, k = int(os.getenv("M", "16")), int(os.getenv("N", "128")), int(os.getenv("K", "192"))
-  seed, check = int(os.getenv("SEED", "0")), bool(int(os.getenv("CHECK", "1")))
-  rng = random.Random(seed)
-  ones = bool(int(os.getenv("ONES", "0")))
-  if check:
-    signed_a = bool(int(os.getenv("SIGNED_A", "0")))
-    val_range = int(os.getenv("VAL_RANGE", "8"))
-    av = [1 if ones else rng.randrange(-val_range, val_range) if signed_a else rng.randrange(val_range) for _ in range(m*k)]
-    bv = [1 if ones else rng.randrange(val_range) for _ in range(k*n)]
-    ap = [pack4(av[row*k+ki*16+c*4:row*k+ki*16+c*4+4]) for row in range(m) for ki in range(k//16) for c in range(4)]
-    bp = [pack4([bv[(ki*16+j*4+l)*n+col4*4+c] for l in range(4)])
-          for ki in range(k//16) for j in range(4) for col4 in range(n//4) for c in range(4)]
-  else:
-    # Throughput-only runs do not need to spend O(MNK) time packing Python
-    # integers. The shader executes the same instructions for zero words.
-    av = bv = []
-    ap, bp = [0] * (m*k//4), [0] * (k*n//4)
-  combined = bool(int(os.getenv("COMBINED", "0")))
-  if combined:
-    width, bheight = max(k//16, n//4), k//4
-    packed = [0] * ((bheight+m)*width*4)
-    for y in range(bheight): packed[y*width*4:y*width*4+(n//4)*4] = bp[y*(n//4)*4:(y+1)*(n//4)*4]
-    for row in range(m): packed[(bheight+row)*width*4:(bheight+row)*width*4+(k//16)*4] = ap[row*(k//16)*4:(row+1)*(k//16)*4]
-
-  q.M, q.N, q.K, q.K4 = m, n, k, k//4
-  dev = Device["QCOM"]
-  env, io, sz, ro = get_envelope(dev, q.make_direct_donor_src_u32(1, 128))
-  const_inputs, const_output = bool(int(os.getenv("CONST_INPUTS", "0"))), bool(int(os.getenv("CONST_OUTPUT", "0")))
-  shader, hregs, fregs, _ = q.build_4x4_dp4_shader(dev, 128, k, constant_inputs=const_inputs, constant_output=const_output,
-    constant_a=bool(int(os.getenv("CONST_A", "0"))), constant_b=bool(int(os.getenv("CONST_B", "0"))),
-    combined_b_height=k//4 if combined else 0, mixed=bool(int(os.getenv("MIXED", "0"))),
-    initial_acc=int(os.getenv("INITIAL_ACC", "0")), coord_delay=int(os.getenv("COORD_DELAY", "4")))
-  assert len(shader) <= sz, (len(shader), sz)
-  lib = inject(env, io, sz, ro, shader, fregs=fregs, hregs=hregs)
-  ab, bb, cb = (Buffer("QCOM", size, dtype).allocate() for size, dtype in
-                ((len(packed) if combined else len(ap), dtypes.uint32),
-                 (len(packed) if combined else len(bp), dtypes.uint32), (m*n, dtypes.int32)))
-  ab.copyin(memoryview(bytearray(struct.pack(f"<{ab.size}I", *(packed if combined else ap)))))
-  bb.copyin(memoryview(bytearray(struct.pack(f"<{bb.size}I", *(packed if combined else bp)))))
-  cb.copyin(memoryview(bytearray(m*n*4)))
-  specs = ([((0, dtypes.uint32, (k//4+m, max(k//16, n//4), 4)),),
-            ((1, dtypes.uint32, (k//4+m, max(k//16, n//4), 4)),)] if combined else
-           [((0, dtypes.uint32, (m, k//16, 4)),), ((1, dtypes.uint32, (k//4, n//4, 4)),)]) + [((0, dtypes.int32, None),)]
-  prg = dev.runtime("gemm_h", lib, buf_dtypes=specs)
-  times = [prg(ab._buf, bb._buf, cb._buf, global_size=(n//128, m//16, 1), local_size=(128, 1, 1), wait=True) for _ in range(10)]
-  print(f"elapsed_ms={min(times)*1e3:.4f} gops={2*m*n*k/min(times)/1e9:.1f}")
-  if check:
-    outb = bytearray(m*n*4)
-    cb.copyout(memoryview(outb))
-    got = struct.unpack(f"<{m*n}i", outb)
-    worst = 0
-    for row in range(m):
-      for col in range(n):
-        expected = sum(av[row*k+kk]*bv[kk*n+col] for kk in range(k)) + int(os.getenv("INITIAL_ACC", "0"))
-        worst = max(worst, abs(got[row*n+col]-expected))
-    print("first=", list(got[:16]))
-    print("max_abs=", worst)
-    if worst: raise SystemExit(1)
-
-
-if __name__ == "__main__": main()

extra/gemm/qcom_dp4_compiler_check.py

@@ -1,40 +0,0 @@
-#!/usr/bin/env python3
-"""Check Qualcomm's compiler-generated packed uint8 dot-product instruction."""
-import struct, time
-
-import numpy as np
-
-from tinygrad import Device, dtypes
-from tinygrad.device import Buffer
-from extra.gemm.ir3asm import disasm
-
-
-def main() -> None:
-  source = """__kernel void dp4(__global int *O,__global uint *A,__global uint *B) {
-  int i=get_global_id(0); uchar4 a=as_uchar4(A[i]),b=as_uchar4(B[i]);
-  O[i]=(int)a.x*(int)b.x+(int)a.y*(int)b.y+(int)a.z*(int)b.z+(int)a.w*(int)b.w;
-}"""
-  dev = Device["QCOM"]
-  lib = dev.compiler.compile(source)
-  image_off, image_size = struct.unpack_from("<I", lib, 0xc0)[0], struct.unpack_from("<I", lib, 0x100)[0]
-  print("\n".join(x for x in disasm(lib[image_off:image_off+image_size]).splitlines()
-                  if "dp4" in x or "mad" in x or "mul" in x or "stg" in x))
-  rng = np.random.default_rng(0)
-  n = 131072
-  a8, b8 = rng.integers(0, 16, (n, 4), dtype=np.uint8), rng.integers(0, 16, (n, 4), dtype=np.uint8)
-  a, b = a8.view(np.uint32).reshape(-1), b8.view(np.uint32).reshape(-1)
-  ab, bb, ob = Buffer("QCOM", n, dtypes.uint).allocate(), Buffer("QCOM", n, dtypes.uint).allocate(), Buffer("QCOM", n, dtypes.int).allocate()
-  ab.copyin(memoryview(a).cast("B"))
-  bb.copyin(memoryview(b).cast("B"))
-  prg = dev.runtime("dp4", lib, buf_dtypes=[((0, dtypes.int, None),), ((1, dtypes.uint, None),), ((2, dtypes.uint, None),)])
-  times = [prg(ob._buf, ab._buf, bb._buf, global_size=(n//128, 1, 1), local_size=(128, 1, 1), wait=True) for _ in range(20)]
-  out = np.empty(n, np.int32)
-  ob.copyout(memoryview(out).cast("B"))
-  expected = (a8.astype(np.int32)*b8.astype(np.int32)).sum(axis=1)
-  print(f"min_us={min(times)*1e6:.3f} max_abs={int(np.max(np.abs(out-expected)))} first={out[:8].tolist()}")
-
-
-if __name__ == "__main__":
-  start = time.perf_counter()
-  main()
-  print(f"wall_ms={(time.perf_counter()-start)*1e3:.1f}")

extra/gemm/qcom_dp4_openpilot_graph.py

@@ -1,166 +0,0 @@
-#!/usr/bin/env python3
-"""Replace selected cached OpenPilot FP16 GEMMs with dynamically-scaled A630 DP4 kernels."""
-import argparse, itertools, pickle
-from dataclasses import replace
-
-import numpy as np
-
-from tinygrad import Device, dtypes
-from tinygrad.engine.jit import create_graph_call
-from tinygrad.uop.ops import Ops, UOp
-from extra.gemm import qcom_intensity_gemm as q
-from extra.gemm.ir3asm import get_envelope, inject
-from extra.gemm.qcom_ir3_matmul_patch import plain_name
-
-
-def aux(*specs):
-  return (tuple(((i, dtype, shape),) for i, (dtype, shape) in enumerate(specs)),)
-
-
-def build_program(template:UOp, name:str, source:str, lib:bytes, global_size, local_size, specs, outs, ins):
-  info = replace(template.arg, name=name, global_size=global_size, local_size=local_size,
-                 globals=tuple(range(len(specs))), outs=outs, ins=ins, aux=aux(*specs))
-  return template.replace(arg=info, src=template.src[:2]+(template.src[2].replace(arg=source), template.src[3].replace(arg=lib)))
-
-
-def pack_unsigned_weights(matrix:np.ndarray) -> tuple[np.ndarray, np.ndarray]:
-  """Per-output-channel symmetric int8 quantization, biased to uint8 for A630 signed*unsigned DP4."""
-  k, n = matrix.shape
-  scale = np.max(np.abs(matrix), axis=0).astype(np.float32) / 127.0
-  scale[scale == 0] = 1.0
-  signed = np.clip(np.rint(matrix/scale), -127, 127).astype(np.int16)
-  unsigned = (signed+128).astype(np.uint8).reshape(k//16, 4, 4, n//4, 4)
-  words = np.zeros((k//16, 4, n//4, 4), dtype=np.uint32)
-  for lane in range(4): words |= unsigned[:, :, lane].astype(np.uint32) << (8*lane)
-  return words.reshape(k//4, n//4, 4), scale
-
-
-def main() -> None:
-  ap = argparse.ArgumentParser()
-  ap.add_argument("input")
-  ap.add_argument("output")
-  ap.add_argument("--indices", required=True, help="comma-separated indices in the cached gemm_h call sequence")
-  args = ap.parse_args()
-  selected = {int(x) for x in args.indices.split(",") if x}
-  with open(args.input, "rb") as f: model = pickle.load(f)
-  batch = model.captured.linear.src[0].src[0].src[0].src
-  existing_slots = [x.arg.slot for x in model.captured.linear.toposort()
-                    if x.op is Ops.BUFFER and hasattr(x.arg, "slot") and x.arg.slot >= 0]
-  UOp.unique_num = itertools.count(max(existing_slots, default=-1)+1)
-  dev = Device["QCOM"]
-
-  pack_sources, pack_libs, dp4_libs = {}, {}, {}
-  epi_sources, epi_libs = {}, {}
-  replacements = {}
-  candidates = [(i, call) for i, call in enumerate(batch) if call.op is Ops.CALL and call.src[0].op is Ops.PROGRAM and
-                plain_name(call.src[0].arg.name) == "gemm_h" and int(call.src[0].arg.global_size[0]) in (3, 12)]
-  for occurrence, (index, call) in enumerate(candidates):
-    if occurrence not in selected: continue
-    gsx = int(call.src[0].arg.global_size[0])
-    m, k, n = (128, 384, 1536) if gsx == 12 else (128, 1536, 384)
-    epi_call = batch[index+1]
-    expected_epi = "epi3_fp32" if gsx == 12 else "epi_fp32"
-    if epi_call.op is not Ops.CALL or plain_name(epi_call.src[0].arg.name) != expected_epi:
-      raise ValueError(f"cached GEMM {occurrence} is followed by {plain_name(epi_call.src[0].arg.name)}, expected {expected_epi}")
-
-    if k not in pack_libs:
-      pack_source = f"""#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-const sampler_t smp=CLK_NORMALIZED_COORDS_FALSE|CLK_ADDRESS_CLAMP|CLK_FILTER_NEAREST;
-__attribute__((reqd_work_group_size(128,1,1)))
-__kernel void qpack(__global uint *O,__global float *S,__global int *SUM,read_only image2d_t A) {{
-  int lid=get_local_id(0),row=get_group_id(0); __local float vmax[128]; __local int vsum[128];
-  float mx=0.0f; for(int k4=lid;k4<{k//4};k4+=128) {{
-    float4 v=fabs(convert_float4(read_imageh(A,smp,(int2)(k4,row))));
-    mx=fmax(mx,fmax(fmax(v.x,v.y),fmax(v.z,v.w))); }}
-  vmax[lid]=mx; barrier(CLK_LOCAL_MEM_FENCE);
-  for(int d=64;d;d>>=1) {{ if(lid<d) vmax[lid]=fmax(vmax[lid],vmax[lid+d]); barrier(CLK_LOCAL_MEM_FENCE); }}
-  float sc=vmax[0]==0.0f?1.0f:vmax[0]/127.0f; int sm=0;
-  for(int k4=lid;k4<{k//4};k4+=128) {{
-    float4 v=convert_float4(read_imageh(A,smp,(int2)(k4,row)))/(float4)(sc);
-    char4 z=convert_char4_sat_rte(v); O[row*{k//4}+k4]=as_uint(z);
-    sm+=(int)z.x+(int)z.y+(int)z.z+(int)z.w; }}
-  vsum[lid]=sm; barrier(CLK_LOCAL_MEM_FENCE);
-  for(int d=64;d;d>>=1) {{ if(lid<d) vsum[lid]+=vsum[lid+d]; barrier(CLK_LOCAL_MEM_FENCE); }}
-  if(lid==0) {{ S[row]=sc; SUM[row]=vsum[0]; }}
-}}"""
-      pack_sources[k] = pack_source
-      pack_libs[k] = dev.compiler.compile(pack_source)
-    pack = build_program(call.src[0], "qpack", pack_sources[k], pack_libs[k], (m, 1, 1), (128, 1, 1),
-                         ((dtypes.uint, (m*k//4,)), (dtypes.float, (m,)), (dtypes.int, (m,)),
-                          (dtypes.half, (m, k//4, 4))), (0, 1, 2), (3,))
-
-    if (m, n, k) not in dp4_libs:
-      q.M, q.N, q.K, q.K4 = m, n, k, k//4
-      env, io, sz, ro = get_envelope(dev, q.make_direct_donor_src_u32(1, 128))
-      shader, hregs, fregs, _ = q.build_4x4_dp4_shader(dev, 128, k, mixed=True, coord_delay=4)
-      dp4_libs[(m, n, k)] = inject(env, io, sz, ro, shader, fregs=fregs, hregs=hregs)
-    dp4 = build_program(call.src[0], "gemm_h", "packed signed-u8 DP4 GEMM", dp4_libs[(m, n, k)],
-                        (n//128, m//16, 1), (128, 1, 1),
-                        ((dtypes.uint, (m, k//16, 4)), (dtypes.uint, (k//4, n//4, 4)),
-                         (dtypes.int, (m*n,))), (2,), (0, 1))
-
-    if gsx not in epi_libs:
-      if gsx == 12:
-        epi_source = """#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-__attribute__((reqd_work_group_size(128,1,1)))
-__kernel void epi3_dp4(__global half *O,__global float *S,__global float *B,__global int *C,
-                        __global float *AS,__global int *SUM,__global float *WS) {
-  int t=get_global_id(0),row=t/384,col=t-row*384,y=row>>2,r=row&3,o=(y*1536+r*384+col)*4;
-  int4 d=vload4(0,C+row*1536+col*4)-(int4)(128*SUM[row]);
-  float4 z=convert_float4(d)*(float4)(AS[row])*vload4(0,WS+col*4);
-  z=select((float4)(0),z,isgreater(z,(float4)(0)));
-  vstore4(convert_half4((float4)(*S)*z*z+(float4)(*B)),0,O+o);
-}"""
-      else:
-        epi_source = """#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-const sampler_t smp=CLK_NORMALIZED_COORDS_FALSE|CLK_ADDRESS_CLAMP|CLK_FILTER_NEAREST;
-__attribute__((reqd_work_group_size(128,1,1)))
-__kernel void epi_dp4(write_only image2d_t O,read_only image2d_t X,read_only image2d_t S,__global int *C,
-                       __global float *AS,__global int *SUM,__global float *WS) {
-  int t=get_global_id(0),row=t/96,col=t-row*96;
-  int4 d=vload4(0,C+row*384+col*4)-(int4)(128*SUM[row]);
-  float4 v=convert_float4(d)*(float4)(AS[row])*vload4(0,WS+col*4);
-  write_imagef(O,(int2)(t,0),read_imagef(X,smp,(int2)(t,0))*read_imagef(S,smp,(int2)(col,0))+v);
-}"""
-      epi_sources[gsx], epi_libs[gsx] = epi_source, dev.compiler.compile(epi_source)
-
-    matrix = np.asarray(call.src[2].buffer.numpy(), dtype=np.float32).reshape(k, n)
-    packed_weight_np, weight_scale_np = pack_unsigned_weights(matrix)
-    packed_weight = UOp.new_buffer("QCOM", packed_weight_np.size, dtypes.uint)
-    packed_weight.buffer.ensure_allocated(); packed_weight.buffer.copyin(memoryview(packed_weight_np).cast("B"))
-    weight_scale = UOp.new_buffer("QCOM", n, dtypes.float)
-    weight_scale.buffer.ensure_allocated(); weight_scale.buffer.copyin(memoryview(weight_scale_np).cast("B"))
-    packed_activation = UOp.new_buffer("QCOM", m*k//4, dtypes.uint); packed_activation.buffer.ensure_allocated()
-    activation_scale = UOp.new_buffer("QCOM", m, dtypes.float); activation_scale.buffer.ensure_allocated()
-    activation_sum = UOp.new_buffer("QCOM", m, dtypes.int); activation_sum.buffer.ensure_allocated()
-    scratch = UOp.new_buffer("QCOM", m*n, dtypes.int); scratch.buffer.ensure_allocated()
-
-    pack_call = pack.call(packed_activation, activation_scale, activation_sum, call.src[1])
-    dp4_call = dp4.call(packed_activation, packed_weight, scratch)
-    if gsx == 12:
-      epi = build_program(epi_call.src[0], "epi3_dp4", epi_sources[gsx], epi_libs[gsx], (384, 1, 1), (128, 1, 1),
-                          ((dtypes.half, (128*1536,)), (dtypes.float, (1,)), (dtypes.float, (1,)),
-                           (dtypes.int, (m*n,)), (dtypes.float, (m,)), (dtypes.int, (m,)), (dtypes.float, (n,))),
-                          (0,), (1, 2, 3, 4, 5, 6))
-      epi_new = epi.call(epi_call.src[1], epi_call.src[2], epi_call.src[3], scratch,
-                         activation_scale, activation_sum, weight_scale)
-    else:
-      epi = build_program(epi_call.src[0], "epi_dp4", epi_sources[gsx], epi_libs[gsx], (96, 1, 1), (128, 1, 1),
-                          ((dtypes.float, (1, 12288, 4)), (dtypes.float, (1, 12288, 4)), (dtypes.float, (1, 96, 4)),
-                           (dtypes.int, (m*n,)), (dtypes.float, (m,)), (dtypes.int, (m,)), (dtypes.float, (n,))),
-                          (0,), (1, 2, 3, 4, 5, 6))
-      epi_new = epi.call(epi_call.src[1], epi_call.src[2], epi_call.src[3], scratch,
-                         activation_scale, activation_sum, weight_scale)
-    replacements[index] = (pack_call, dp4_call)
-    replacements[index+1] = (epi_new,)
-    print(f"occurrence={occurrence} geometry={gsx} shape={m}x{n}x{k}")
-
-  outer = model.captured.linear.src[0]
-  new_batch = [new for i, old in enumerate(batch) for new in replacements.get(i, (old,))]
-  model.captured._linear = model.captured.linear.substitute({outer:create_graph_call(new_batch)}, walk=True)
-  model.captured.__dict__.pop("linear", None)
-  with open(args.output, "wb") as f: pickle.dump(model, f)
-  print(f"wrote {args.output} with {len(replacements)//2} DP4 GEMMs")
-
-
-if __name__ == "__main__": main()

extra/gemm/qcom_dump_prefix.py

@@ -1,50 +0,0 @@
-#!/usr/bin/env python3
-"""Run one cached model graph prefix and dump the selected call output."""
-import argparse, pickle
-
-import numpy as np
-
-from tinygrad import Tensor
-from tinygrad.engine.jit import _prepare_jit_inputs, create_graph_call
-from tinygrad.engine.realize import resolve_params, run_linear
-from tinygrad.uop.ops import Ops, UOp
-
-
-def main() -> None:
-  parser = argparse.ArgumentParser()
-  parser.add_argument("model")
-  parser.add_argument("corpus")
-  parser.add_argument("output")
-  parser.add_argument("--case", type=int, default=9)
-  parser.add_argument("--index", type=int)
-  parser.add_argument("--indices", help="comma-separated indices; output must contain a {index} placeholder")
-  parser.add_argument("--individual-last", action="store_true")
-  args = parser.parse_args()
-  if (args.index is None) == (args.indices is None): parser.error("pass exactly one of --index or --indices")
-  indices = [args.index] if args.index is not None else [int(x) for x in args.indices.split(",")]
-  if len(indices) > 1 and "{index}" not in args.output: parser.error("multi-index output must contain {index}")
-  with open(args.model, "rb") as f: model = pickle.load(f)
-  corpus = np.load(args.corpus)
-  inputs = {}
-  for name, (view, _vars, dtype, device) in zip(model.captured.expected_names, model.captured.expected_input_info):
-    arr = corpus[f"case{args.case}:input:{name}"].astype(np.dtype(dtype.fmt), copy=False)
-    inputs[name] = Tensor(arr, device=device).realize()
-  input_uops, var_vals = _prepare_jit_inputs((), inputs)[:2]
-  batch = model.captured.linear.src[0].src[0].src[0].src
-  for index in indices:
-    prefix_end = index if args.individual_last else index+1
-    if prefix_end:
-      run_linear(UOp(Ops.LINEAR, src=(create_graph_call(batch[:prefix_end]),)), var_vals,
-                 input_uops=input_uops, jit=True, wait=True)
-    if args.individual_last:
-      run_linear(UOp(Ops.LINEAR, src=(batch[index],)), var_vals, input_uops=input_uops, jit=True, wait=True)
-    call = batch[index]
-    call_args = resolve_params(call, tuple(input_uops))
-    out_buffer = call_args[call.src[0].arg.outs[0]]
-    out = np.asarray(out_buffer.buffer.numpy()).copy()
-    output = args.output.format(index=index)
-    np.save(output, out)
-    print(f"index={index} shape={out.shape} dtype={out.dtype} min={float(out.min())} max={float(out.max())} mean={float(out.mean())}")
-
-
-if __name__ == "__main__": main()

extra/gemm/qcom_find_divergence.py

@@ -1,86 +0,0 @@
-#!/usr/bin/env python3
-"""Locate the first call whose output differs between two compiled QCOM models."""
-import argparse, pickle
-
-import numpy as np
-
-from tinygrad import Tensor
-from tinygrad.engine.jit import create_graph_call
-from tinygrad.engine.jit import _prepare_jit_inputs
-from tinygrad.engine.realize import resolve_params, run_linear
-from tinygrad.uop.ops import Ops, UOp
-from extra.gemm.qcom_ir3_matmul_patch import plain_name
-
-
-def load(path: str):
-  with open(path, "rb") as f: return pickle.load(f)
-
-
-def batch(model): return model.captured.linear.src[0].src[0].src[0].src
-
-
-def prepare(model, corpus, case: int):
-  inputs = {}
-  for name, (view, _vars, dtype, device) in zip(model.captured.expected_names, model.captured.expected_input_info):
-    arr = corpus[f"case{case}:input:{name}"].astype(np.dtype(dtype.fmt), copy=False)
-    inputs[name] = Tensor(arr, device=device).realize()
-  return _prepare_jit_inputs((), inputs)[:2]
-
-
-def main() -> None:
-  parser = argparse.ArgumentParser()
-  parser.add_argument("reference")
-  parser.add_argument("candidate")
-  parser.add_argument("corpus")
-  parser.add_argument("--case", type=int, default=9)
-  parser.add_argument("--threshold", type=float, default=1e-3)
-  parser.add_argument("--graph-prefix", type=int)
-  args = parser.parse_args()
-  reference, candidate = load(args.reference), load(args.candidate)
-  corpus = np.load(args.corpus)
-  rb, cb = batch(reference), batch(candidate)
-  if len(rb) != len(cb): raise ValueError(f"batch lengths differ: {len(rb)} != {len(cb)}")
-  ri, rv = prepare(reference, corpus, args.case)
-  ci, cv = prepare(candidate, corpus, args.case)
-  if args.graph_prefix is not None:
-    index = args.graph_prefix
-    rc, cc = rb[index], cb[index]
-    run_linear(UOp(Ops.LINEAR, src=(create_graph_call(rb[:index+1]),)), rv, input_uops=ri, jit=True, wait=True)
-    rargs = resolve_params(rc, tuple(ri))
-    snapshots = {out_index: np.asarray(rargs[out_index].buffer.numpy(), dtype=np.float32).copy()
-                 for out_index in rc.src[0].arg.outs}
-    run_linear(UOp(Ops.LINEAR, src=(create_graph_call(cb[:index+1]),)), cv, input_uops=ci, jit=True, wait=True)
-    cargs = resolve_params(cc, tuple(ci))
-    for out_index, candidate_out_index in zip(rc.src[0].arg.outs, cc.src[0].arg.outs):
-      delta = np.abs(snapshots[out_index]-np.asarray(cargs[candidate_out_index].buffer.numpy(), dtype=np.float32))
-      print(f"{index}: graph-prefix output={out_index} max_abs={float(delta.max(initial=0)):.9g} mean_abs={float(delta.mean()):.9g}")
-    return
-  for index, (rc, cc) in enumerate(zip(rb, cb)):
-    run_linear(UOp(Ops.LINEAR, src=(rc,)), rv, input_uops=ri, jit=True, wait=True)
-    reference_outputs = {}
-    if rc.op is Ops.CALL:
-      rargs = resolve_params(rc, tuple(ri))
-      reference_outputs = {out_index: (rargs[out_index].dtype, rargs[out_index].buffer.nbytes,
-        np.asarray(rargs[out_index].buffer.numpy(), dtype=np.float32).copy()) for out_index in rc.src[0].arg.outs}
-    run_linear(UOp(Ops.LINEAR, src=(cc,)), cv, input_uops=ci, jit=True, wait=True)
-    if rc.op is not Ops.CALL or cc.op is not Ops.CALL: continue
-    rn = plain_name(rc.src[0].arg.name) if rc.src[0].op is Ops.PROGRAM else str(rc.op)
-    cn = plain_name(cc.src[0].arg.name) if cc.src[0].op is Ops.PROGRAM else str(cc.op)
-    cargs = resolve_params(cc, tuple(ci))
-    maximum = 0.0
-    for out_index, candidate_out_index in zip(rc.src[0].arg.outs, cc.src[0].arg.outs):
-      reference_dtype, reference_nbytes, ro = reference_outputs[out_index]
-      cout = cargs[candidate_out_index]
-      if reference_dtype != cout.dtype or reference_nbytes != cout.buffer.nbytes:
-        print(f"{index}: {rn} -> {cn}: incompatible output {out_index}")
-        continue
-      co = np.asarray(cout.buffer.numpy(), dtype=np.float32)
-      delta = np.abs(ro-co)
-      out_maximum, mean = float(delta.max(initial=0)), float(delta.mean())
-      maximum = max(maximum, out_maximum)
-      if out_maximum > args.threshold or rn != cn:
-        print(f"{index}: {rn} -> {cn}: output={out_index} max_abs={out_maximum:.9g} mean_abs={mean:.9g}")
-    if maximum > args.threshold: break
-
-
-if __name__ == "__main__": main()

extra/gemm/qcom_fp16_acc_general_check.py

@@ -1,170 +0,0 @@
-#!/usr/bin/env python3
-"""Full random-matrix correctness check for the fast A630 FP16-accumulate GEMM."""
-import ctypes, importlib.util, os, struct
-
-import numpy as np
-
-from tinygrad import Device, dtypes
-from tinygrad.device import Buffer
-from tinygrad.helpers import ceildiv
-from tinygrad.runtime.ops_qcom import dcache_flush
-from extra.gemm.ir3asm import disasm, get_envelope, inject
-from extra.gemm.qcom_gemm import patch_kernel
-
-if module_path := os.getenv("QCOM_INTENSITY_MODULE"):
-  spec = importlib.util.spec_from_file_location("qcom_intensity_snapshot", module_path)
-  if spec is None or spec.loader is None: raise RuntimeError(f"cannot load {module_path}")
-  q = importlib.util.module_from_spec(spec)
-  spec.loader.exec_module(q)
-else:
-  from extra.gemm import qcom_intensity_gemm as q
-
-
-def install_random_safe_store() -> None:
-  """Replace the all-ones-only repeated move in the preserved fast kernel's epilogue."""
-  original_mov_h = q.MOV_H
-  def random_safe_mov_h(dst, src, rpt=0, r=False):
-    # In this path every repeated MOV broadcasts the zero accumulator seed. On A630 the
-    # relative-source repeat also walks the source, so use an immediate vector fill instead.
-    return q.MOV_H_IMM(dst, 0, rpt=rpt) if rpt else original_mov_h(dst, src, r=r)
-  q.MOV_H = random_safe_mov_h
-  hand_addr = [bytes.fromhex(x) for x in (
-    "1c000a200000d04e 1d0002200100d046 0000000000100000 000061100201b843 000060100600b043 0000010000003042 "
-    "0100020002013842 0100060001003042 000001200000d046 020001200200d046 030001200600d046 010001200700d046 "
-    "0000000003401520 0000000000000000 0600000001401520 0700000000401520 0000000000100000 5010030008001042 "
-    "501002000a001042 501001000c001042 501000000e001042 0000000000100000 0800501010009042 03001f201100f046 "
-    "0a00501012009042 02001f201300f046 0c00501014009042 01001f201500f046 0e00501016009042 00001f201700f046 "
-    "0000000000100000 5110104009808867 511012400b808967 511014400d808a67 519016400f888b67").split()]
-  gap = int(os.getenv("STORE_GAP", "0"))
-  hand_stores = []
-  for row, addr in enumerate(("r2.x", "r2.z", "r3.x", "r3.z")):
-    hand_stores.append(q.STG_F16(addr, row*4))
-    if row != 3: hand_stores.append(q.NOP(rpt=gap))
-  def emit(instrs, acc0, ncols, *_args, **_kwargs):
-    for col in range(ncols):
-      if col: instrs.append(q.ADD_S("r7.y", "r7.y", 32))
-      instrs += hand_addr
-      for row in range(4):
-        for lane in range(4): instrs.append(q.MOV_H(row*4+lane, acc0+(row*ncols+col)*4+lane))
-      instrs += hand_stores
-  q.emit_hand4_stores = emit
-
-
-def upload(x: np.ndarray, dtype) -> Buffer:
-  ret = Buffer("QCOM", x.size, dtype).allocate()
-  raw = memoryview(np.ascontiguousarray(x)).cast("B")
-  ret.copyin(raw) if hasattr(ret, "copyin") else Device[ret.device].allocator._copyin(ret._buf, raw)
-  ptr = ret._buf.cpu_view().addr
-  dcache_flush().fxn(ctypes.c_uint64(ptr & ~63), ceildiv(ptr + ret.nbytes - (ptr & ~63), 64))
-  return ret
-
-
-def main() -> None:
-  m, n, k = (int(os.getenv(name, "1024")) for name in ("M", "N", "K"))
-  seed, threads = int(os.getenv("SEED", "901")), int(os.getenv("THREADS", "128"))
-  ncols = int(os.getenv("NCOLS", "4"))
-  if m % 16 or n % (128*ncols) or k % 16: raise ValueError("M, N, K must divide the selected tile")
-  if threads not in (64, 128, 256): raise ValueError("THREADS must be 64, 128, or 256")
-  rng = np.random.default_rng(seed)
-  a_np = (np.eye(m, k, dtype=np.float16) if os.getenv("PATTERN") == "identity" else
-          (rng.standard_normal((m, k), dtype=np.float32)*np.float32(0.05)).astype(np.float16))
-  b_np = (rng.standard_normal((k, n), dtype=np.float32)*np.float32(0.05)).astype(np.float16)
-
-  q.M, q.N, q.K, q.K4 = m, n, k, k//4
-  install_random_safe_store()
-  dev = Device["QCOM"]
-  if os.getenv("COMPILER_DIRECT"):
-    compiler_partial = bool(int(os.getenv("COMPILER_PARTIAL", "0")))
-    compiler_image = bool(int(os.getenv("COMPILER_IMAGE", "0")))
-    compiler_src = (q.make_direct_image_donor_src(ncols, threads) if compiler_image else
-                    q.make_donor_src(ncols, threads) if compiler_partial else q.make_direct_donor_src(ncols, threads))
-    lib = dev.compiler.compile_cached(compiler_src)
-    if os.getenv("PATCH_COMPILER", "1") != "0":
-      lib = patch_kernel(lib, os.getenv("PATCH_SYNC", "1") != "0", os.getenv("MERGE_PAIRS", "1") != "0", int(os.getenv("MAX_GROUPS", "-1")))
-    if os.getenv("PRINT_ASM"):
-      io, sz = struct.unpack_from("<I", lib, 0xc0)[0], struct.unpack_from("<I", lib, 0x100)[0]
-      print(disasm(lib[io:io+sz]))
-    loop_instrs = -1
-  else:
-    image_output = bool(int(os.getenv("IMAGE_OUTPUT", "0")))
-    env_src = q.make_direct_image_donor_src(4, threads) if image_output else q.make_donor_src(4, threads)
-    env, io, sz, ro = get_envelope(dev, env_src)
-    low_a = bool(int(os.getenv("LOW_A", "0")))
-    high_inputs = bool(int(os.getenv("HIGH_INPUTS", "0")))
-    high_a_only = bool(int(os.getenv("HIGH_A_ONLY", "0")))
-    extra = ({"high_inputs": high_inputs, "high_a_only": high_a_only}
-             if "high_inputs" in __import__("inspect").signature(q.build_4xn_shader).parameters else {})
-    extra["serial_b_cols"] = bool(int(os.getenv("SERIAL_B_COLS", "0")))
-    k_unroll = int(os.getenv("K_UNROLL", "4"))
-    if image_output:
-      persistent = bool(int(os.getenv("PERSISTENT", "0")))
-      advanced = ncols == 4
-      shader, loop_instrs = q.build_4xn_shader(dev, threads, ncols=ncols, direct=True, b_first=advanced, compact_acc=True,
-        stable_bx=advanced, stable_ay=advanced, low_a_coords=low_a, inc_coords=persistent and advanced,
-        persistent_coords=persistent and advanced, k_unroll=k_unroll,
-        alu_order="row_col_kk", first_sync_only=bool(int(os.getenv("FIRST_SYNC_ONLY", "1"))),
-        row_sync=bool(int(os.getenv("ROW_SYNC", "0"))), coord_delay=int(os.getenv("COORD_DELAY", "-1")),
-        image_store=True, preserve_coords=bool(int(os.getenv("PRESERVE_COORDS", "1"))),
-        safe_b_y=bool(int(os.getenv("SAFE_B_Y", "0"))), sync_b_y=bool(int(os.getenv("SYNC_B_Y", "0"))),
-        separate_b_coords=bool(int(os.getenv("SEPARATE_B_COORDS", "0"))), high_b_coords=bool(int(os.getenv("HIGH_B_COORDS", "0"))),
-        persistent_b_x=bool(int(os.getenv("PERSISTENT_B_X", "0"))), **extra)
-    else:
-      advanced = ncols == 4
-      shader, loop_instrs = q.build_4xn_shader(dev, threads, ncols=ncols, direct=True, b_first=advanced, compact_acc=True,
-        stable_bx=advanced, stable_ay=advanced, low_a_coords=low_a, inc_coords=advanced, persistent_coords=advanced, k_unroll=k_unroll,
-        alu_order="row_col_kk", first_sync_only=bool(int(os.getenv("FIRST_SYNC_ONLY", "1"))),
-        row_sync=bool(int(os.getenv("ROW_SYNC", "0"))), coord_delay=int(os.getenv("COORD_DELAY", "-1")),
-        safe_b_y=bool(int(os.getenv("SAFE_B_Y", "0"))), sync_b_y=bool(int(os.getenv("SYNC_B_Y", "0"))),
-        separate_b_coords=bool(int(os.getenv("SEPARATE_B_COORDS", "0"))), high_b_coords=bool(int(os.getenv("HIGH_B_COORDS", "0"))),
-        persistent_b_x=bool(int(os.getenv("PERSISTENT_B_X", "0"))), **extra)
-    lib = inject(env, io, sz, ro, shader, fregs=13 if bool(int(os.getenv("PERSISTENT_B_X", "0"))) else 8 if low_a else 10,
-                 hregs=48 if high_inputs else 36 if high_a_only else 28)
-    asm = disasm(shader)
-    if asm.count("mad.f16") != 16*ncols*k_unroll or asm.count("mad.f32") != 0:
-      raise RuntimeError("unexpected accumulator instruction mix")
-
-  a, b = upload(a_np, dtypes.half), upload(b_np, dtypes.half)
-  image_output = (bool(int(os.getenv("IMAGE_OUTPUT", "0"))) and not os.getenv("COMPILER_DIRECT")) or \
-                 bool(int(os.getenv("COMPILER_IMAGE", "0")))
-  c_np, c_dtype = (np.zeros((m, n), np.float32), dtypes.float) if image_output else (np.zeros((m, n), np.float16), dtypes.half)
-  c = upload(c_np, c_dtype)
-  if image_output:
-    # The image envelope declares C first so QCOMArgsState assigns its sole IBO to C,
-    # followed by A/B in texture slots 0/1 as expected by the injected prologue.
-    prg = dev.runtime("gemm_h", lib, buf_dtypes=[
-      ((0, dtypes.float, (m, n//4, 4)),), ((0, dtypes.half, (m, k//4, 4)),), ((0, dtypes.half, (k, n//4, 4)),)])
-    args = (c._buf, a._buf, b._buf)
-  else:
-    prg = dev.runtime("gemm_h", lib, buf_dtypes=[
-      ((0, dtypes.half, (m, k//4, 4)),), ((0, dtypes.half, (k, n//4, 4)),), ((0, dtypes.half, None),)])
-    args = (a._buf, b._buf, c._buf)
-  gs, ls = (n//(128*ncols), m//((threads//32)*4), 1), (threads, 1, 1)
-  for _ in range(int(os.getenv("WARM", "3"))): prg(*args, global_size=gs, local_size=ls, wait=True)
-  times = [prg(*args, global_size=gs, local_size=ls, wait=True) for _ in range(int(os.getenv("RUNS", "10")))]
-
-  got = np.empty((m, n), c_np.dtype)
-  raw = memoryview(got).cast("B")
-  c.copyout(raw) if hasattr(c, "copyout") else Device[c.device].allocator._copyout(raw, c._buf)
-  expected = a_np.astype(np.float32) @ b_np.astype(np.float32)
-  delta = np.abs(got.astype(np.float32)-expected)
-  rtol, atol = float(os.getenv("RTOL", "0.02")), float(os.getenv("ATOL", "0.02"))
-  bad = ~np.isclose(got, expected, rtol=rtol, atol=atol)
-  if os.getenv("SHOW_VALUES"):
-    print("bad_by_row_mod4", [int(bad[r::4].sum()) for r in range(4)])
-    print("bad_by_col_block", [int(bad[:, c:c+128].sum()) for c in range(0, n, 128)])
-    for row in range(4):
-      print(f"row={row} got={got[row, :32].astype(np.float32).tolist()}")
-      print(f"row={row} exp={expected[row, :32].tolist()}")
-    if os.getenv("PATTERN") == "identity":
-      for row in range(16):
-        mse = np.mean((expected[:, :128]-got[row, :128])**2, axis=1)
-        print(f"identity_row={row} nearest_expected_row={int(np.argmin(mse))} mse={float(mse.min()):.9g}")
-  best = min(x for x in times if x is not None)
-  print(f"shape={m}x{n}x{k} inputs=fp16 accumulate=fp16 elapsed_ms={best*1e3:.3f} "
-        f"gflops={2*m*n*k/best/1e9:.1f} outputs={m*n} bad_count={int(bad.sum())} "
-        f"max_abs={float(delta.max()):.9g} mean_abs={float(delta.mean()):.9g} "
-        f"rtol={rtol:g} atol={atol:g} allclose={not bool(bad.any())} loop_instrs={loop_instrs}")
-  if bad.any(): raise SystemExit(1)
-
-
-if __name__ == "__main__": main()

extra/gemm/qcom_fp32_4x4_check.py

@@ -1,79 +0,0 @@
-#!/usr/bin/env python3
-"""Randomized oracle for the compact FP32-accumulating 4x4 QCOM GEMM."""
-import argparse
-
-import numpy as np
-
-from tinygrad import Device, dtypes
-from tinygrad.device import Buffer
-from extra.gemm import qcom_intensity_gemm as q
-from extra.gemm.ir3asm import get_envelope, inject
-
-
-def upload(x:np.ndarray, dtype) -> Buffer:
-  ret = Buffer("QCOM", x.size, dtype).allocate()
-  ret.copyin(memoryview(np.ascontiguousarray(x)).cast("B"))
-  return ret
-
-
-def main() -> None:
-  ap = argparse.ArgumentParser()
-  ap.add_argument("--m", type=int, default=128)
-  ap.add_argument("--n", type=int, default=1536)
-  ap.add_argument("--k", type=int, default=384)
-  ap.add_argument("--stride", type=int, default=0)
-  ap.add_argument("--threads", type=int, default=128, choices=(64, 128, 256))
-  ap.add_argument("--seed", type=int, default=0)
-  ap.add_argument("--coord-delay", type=int, default=4)
-  ap.add_argument("--first-wait-only", action="store_true")
-  ap.add_argument("--batch-coords", action="store_true")
-  ap.add_argument("--quad-map", action="store_true", help="map each quad to one output column")
-  ap.add_argument("--quad-b", action="store_true", help="load B in one quad lane and broadcast it")
-  ap.add_argument("--quad-b-load-all", action="store_true", help="load B in all lanes before broadcasting lane zero")
-  ap.add_argument("--quad-b-shfl-mode", type=int, default=0, help="use scalar relative shuffles instead of vector quad broadcast")
-  ap.add_argument("--post-constant", action="store_true", help="replace accumulators with 1024 before storing")
-  ap.add_argument("--float-inputs", action="store_true", help="store both sampled inputs as float32 images")
-  args = ap.parse_args()
-  rng = np.random.default_rng(args.seed)
-  input_np_dtype = np.float32 if args.float_inputs else np.float16
-  input_dtype = dtypes.float if args.float_inputs else dtypes.half
-  a_np = (rng.standard_normal((args.m, args.k))*0.05).astype(input_np_dtype)
-  b_np = (rng.standard_normal((args.k, args.n))*0.05).astype(input_np_dtype)
-  stride = args.stride or (2048 if args.n > 1024 else 1024)
-  q.M, q.N, q.K, q.K4 = args.m, stride, args.k, args.k//4
-  dev = Device["QCOM"]
-  env, io, sz, ro = get_envelope(dev, q.make_direct_donor_src_fp32(1, args.threads))
-  shader, hregs, fregs, _ = q.build_4x4_fp32_compact_preload_shader(
-    dev, args.threads, coord_delay=args.coord_delay, sampler_per_texture=True, post_constant=args.post_constant,
-    batch_coords=args.batch_coords, first_coord_wait_only=args.first_wait_only,
-    quad_map=args.quad_map, quad_b=args.quad_b, quad_b_load_all=args.quad_b_load_all,
-    quad_b_shfl_mode=args.quad_b_shfl_mode)
-  lib = inject(env, io, sz, ro, shader, fregs=fregs, hregs=hregs)
-  a, b = upload(a_np, input_dtype), upload(b_np, input_dtype)
-  c = upload(np.zeros(args.m*stride, np.float32), dtypes.float)
-  specs = [((0, input_dtype, (args.m, args.k//4, 4)),),
-           ((1, input_dtype, (args.k, args.n//4, 4)),), ((2, dtypes.float, (args.m*stride,)),)]
-  program = dev.runtime("gemm_h", lib, buf_dtypes=specs)
-  tile_m = (args.threads//32)*4
-  times = [program(a._buf, b._buf, c._buf, global_size=(args.n//128, args.m//tile_m, 1),
-                   local_size=(args.threads, 1, 1), wait=True)*1e3 for _ in range(10)]
-  got_storage = np.empty((args.m, stride), np.float32)
-  c.copyout(memoryview(got_storage).cast("B"))
-  got = got_storage[:, :args.n]
-  expected = np.full((args.m,args.n), 1024, np.float32) if args.post_constant else a_np.astype(np.float32) @ b_np.astype(np.float32)
-  delta = np.abs(expected-got)
-  worst = np.unravel_index(int(delta.argmax()), delta.shape)
-  passed = bool(np.allclose(expected, got, rtol=1e-4, atol=1e-4))
-  print(f"ms={min(times):.4f} max_abs={float(delta.max()):.9g} mean_abs={float(delta.mean()):.9g} "
-        f"worst={worst} allclose={passed}")
-  if not passed:
-    nz = np.argwhere(got_storage != 0)
-    print("storage_nonzero=", int(nz.shape[0]), "first=", nz[:32].tolist())
-    row_cost = np.mean(np.abs(expected[:,None,:]-got[None,:,:]), axis=2)
-    print("best_expected_row_for_got=", [(int(j), int(np.argmin(row_cost[:,j])), float(np.min(row_cost[:,j]))) for j in range(min(32,args.m))])
-    print("row0_expected=", expected[0,:16].tolist())
-    print("row0_got=", got[0,:16].tolist())
-  if not passed: raise SystemExit(1)
-
-
-if __name__ == "__main__": main()

extra/gemm/qcom_fp32_check.py

@@ -1,109 +0,0 @@
-#!/usr/bin/env python3
-"""Random-matrix oracle for hand-assembled QCOM FP32-accumulating GEMMs."""
-import argparse
-
-import numpy as np
-
-from tinygrad import Device, dtypes
-from tinygrad.device import Buffer
-from extra.gemm import qcom_intensity_gemm as q
-from extra.gemm import qcom_8x4_gemm as q8
-from extra.gemm.ir3asm import get_envelope, inject
-
-
-def main() -> None:
-  parser = argparse.ArgumentParser()
-  parser.add_argument("--m", type=int, default=192)
-  parser.add_argument("--n", type=int, default=512)
-  parser.add_argument("--k", type=int, default=768)
-  parser.add_argument("--stride", type=int, default=0)
-  parser.add_argument("--seed", type=int, default=0)
-  parser.add_argument("--threads", type=int, default=128, choices=(64, 128))
-  parser.add_argument("--preload-b", action="store_true")
-  parser.add_argument("--batch-coords", action="store_true")
-  parser.add_argument("--hand-store", action="store_true")
-  parser.add_argument("--interleaved-a", action="store_true")
-  parser.add_argument("--no-store", action="store_true", help="profile compute only; skip output validation")
-  parser.add_argument("--compiler", action="store_true", help="run the unmodified compiler-generated 4x8 kernel")
-  parser.add_argument("--pipeline", action="store_true", help="run the double-buffered hand 4x8 kernel")
-  parser.add_argument("--alu-order", default="kk_row_col")
-  parser.add_argument("--coord-delay", type=int, default=-1)
-  parser.add_argument("--identity-b", action="store_true", help="make the output expose the first N activation columns")
-  parser.add_argument("--float-inputs", action="store_true", help="store both sampled inputs as float32 images")
-  parser.add_argument("--eight-row", action="store_true", help="test the scalar FP32 8x4 kernel")
-  parser.add_argument("--compact-preload", action="store_true", help="test the compact FP32 4x4 preload kernel")
-  parser.add_argument("--first-wait-only", action="store_true")
-  parser.add_argument("--quad-map", action="store_true")
-  parser.add_argument("--quad-b", action="store_true")
-  parser.add_argument("--quad-b-load-all", action="store_true")
-  args = parser.parse_args()
-  n_tile = 128 if args.eight_row or args.compact_preload else 256
-  tile_m = (args.threads//32) * (8 if args.eight_row else 4)
-  assert args.m % tile_m == 0 and args.n % n_tile == 0 and args.k % 4 == 0
-
-  rng = np.random.default_rng(args.seed)
-  input_np_dtype = np.float32 if args.float_inputs else np.float16
-  input_dtype = dtypes.float if args.float_inputs else dtypes.half
-  a_np = (rng.standard_normal((args.m, args.k))*0.1).astype(input_np_dtype)
-  b_np = (rng.standard_normal((args.k, args.n))*0.1).astype(input_np_dtype)
-  if args.identity_b:
-    b_np.fill(0)
-    np.fill_diagonal(b_np, np.float16(1))
-  stride = args.stride or (2048 if args.n > 1024 else 1024)
-  q.M, q.N, q.K, q.K4 = args.m, stride, args.k, args.k//4
-  q8.M, q8.N, q8.K, q8.K4 = args.m, stride, args.k, args.k//4
-  dev = Device["QCOM"]
-  donor_src = q8.make_donor_src8_fp32(1, args.threads) if args.eight_row else q.make_direct_donor_src_fp32(2, args.threads)
-  envelope, image_offset, image_size, register_offset = get_envelope(dev, donor_src)
-  if args.compiler:
-    lib = bytes(envelope)
-  else:
-    if args.eight_row:
-      shader, hregs, fregs, _ = q8.build_8x8_fp32_shader(
-        dev, args.threads, ncols=1, b_coord_delay=args.coord_delay, alu_order=args.alu_order)
-    elif args.compact_preload:
-      shader, hregs, fregs, _ = q.build_4x4_fp32_compact_preload_shader(
-        dev, args.threads, coord_delay=args.coord_delay, sampler_per_texture=True,
-        batch_coords=args.batch_coords, quad_map=args.quad_map, quad_b=args.quad_b,
-        quad_b_load_all=args.quad_b_load_all, first_coord_wait_only=args.first_wait_only)
-    elif args.pipeline:
-      shader, hregs, fregs, _ = q.build_4x8_fp32_pipeline_shader(
-        dev, args.threads, coord_delay=args.coord_delay, sampler_per_texture=True, no_store=args.no_store)
-    else:
-      shader, hregs, fregs, _ = q.build_4x8_fp32_low_shader(
-        dev, args.threads, coord_delay=args.coord_delay, sampler_per_texture=True, alu_order=args.alu_order,
-        preload_b=args.preload_b, batch_coords=args.batch_coords, hand_store=args.hand_store,
-        interleaved_a=args.interleaved_a, no_store=args.no_store)
-    lib = inject(envelope, image_offset, image_size, register_offset, shader, fregs=fregs, hregs=hregs)
-
-  a_upload = a_np.reshape(args.m//4, 4, args.k).transpose(0, 2, 1).copy() if args.interleaved_a else a_np
-  a = Buffer("QCOM", a_upload.size, input_dtype, initial_value=memoryview(a_upload).cast("B").tobytes())
-  b = Buffer("QCOM", b_np.size, input_dtype, initial_value=memoryview(b_np).cast("B").tobytes())
-  c = Buffer("QCOM", args.m*stride, dtypes.float,
-             initial_value=memoryview(np.zeros(args.m*stride, dtype=np.float32)).cast("B").tobytes())
-  a_shape = (args.m//4, args.k, 4) if args.interleaved_a else (args.m, args.k//4, 4)
-  specs = [((0, input_dtype, a_shape),), ((0, input_dtype, (args.k, args.n//4, 4)),),
-           ((0, dtypes.float, None),)]
-  program = dev.runtime("gemm_h" if args.eight_row else "gemm_f", lib, buf_dtypes=specs)
-  for _ in range(3):
-    elapsed = program(a._buf, b._buf, c._buf, global_size=(args.n//n_tile, args.m//tile_m, 1),
-                      local_size=(args.threads, 1, 1), wait=True)
-  if args.no_store:
-    print(f"elapsed_ms={elapsed*1e3:.3f} compute_only=True")
-    return
-  got_flat = np.empty(args.m*stride, dtype=np.float32)
-  got_flat[:] = c.numpy()
-  got = got_flat.reshape(args.m, stride)[:, :args.n]
-  expected = a_np.astype(np.float32) @ b_np.astype(np.float32)
-  delta = np.abs(expected-got)
-  worst = np.unravel_index(np.argmax(delta), delta.shape)
-  passed = bool(np.allclose(expected, got, rtol=1e-4, atol=1e-4))
-  print(f"elapsed_ms={elapsed*1e3:.3f} max_abs={float(delta.max()):.9g} mean_abs={float(delta.mean()):.9g} "
-        f"worst={worst} expected={expected[worst]!r} got={got[worst]!r} allclose={passed}")
-  print("max_abs row16 x col256=", [[float(delta[r:r+16, c:c+256].max()) for c in range(0, args.n, 256)]
-                                      for r in range(0, args.m, 16)])
-  print(f"nonzero={np.count_nonzero(got)/got.size:.3f} got_row0={got[0, :8].tolist()} expected_row0={expected[0, :8].tolist()}")
-  if not passed: raise SystemExit(1)
-
-
-if __name__ == "__main__": main()

extra/gemm/qcom_fp32_wide_check.py

@@ -1,135 +0,0 @@
-#!/usr/bin/env python3
-"""Full-random oracle and timer for the streamed wide FP32-accumulate GEMM."""
-import os
-
-import numpy as np
-
-from tinygrad import Device, dtypes
-from tinygrad.device import Buffer
-from extra.gemm import qcom_intensity_gemm as q
-from extra.gemm.ir3asm import get_envelope, inject
-
-
-def upload(x: np.ndarray, dtype) -> Buffer:
-  raw = np.ascontiguousarray(x)
-  return Buffer("QCOM", raw.size, dtype, initial_value=memoryview(raw).cast("B").tobytes())
-
-
-def main() -> None:
-  m, n, k = (int(os.getenv(x, d)) for x, d in (("M", "256"), ("N", "512"), ("K", "1024")))
-  batch = int(os.getenv("BATCH", "1"))
-  batch_y = batch > 1 and bool(int(os.getenv("BATCH_Y", "1")))
-  ncols, threads = int(os.getenv("NCOLS", "4")), int(os.getenv("THREADS", "128"))
-  custom_rows = int(os.getenv("CUSTOM_ROWS", "0"))
-  stride, seed = int(os.getenv("STRIDE", str(n))), int(os.getenv("SEED", "0"))
-  k_start = int(os.getenv("K_START", "0"))
-  rows8, square8 = bool(int(os.getenv("ROWS8", "0"))), bool(int(os.getenv("SQUARE8", "0")))
-  quad_a = bool(int(os.getenv("QUAD_A", "0")))
-  quad_split = bool(int(os.getenv("QUAD_SPLIT", "0")))
-  rows8 = rows8 or square8
-  tile_m = (threads//32)*(custom_rows or (8 if rows8 else 4))
-  tile_n = 32*ncols if quad_split else 128*ncols
-  assert m % tile_m == 0 and n % tile_n == 0 and k % 4 == 0
-  rng = np.random.default_rng(seed)
-  a_np = (rng.standard_normal((batch, m, k))*0.05).astype(np.float16)
-  b_np = (rng.standard_normal((batch, k, n))*0.05).astype(np.float16)
-  if os.getenv("PATTERN") == "ones":
-    a_np.fill(1)
-    b_np.fill(1)
-  q.M, q.N, q.K, q.K4 = m, stride, k, k//4
-  dev = Device["QCOM"]
-  rotate_buffer = bool(int(os.getenv("ROTATE_BUFFER", "0")))
-  swap_groups = bool(int(os.getenv("SWAP_GROUPS", "0")))
-  column_z = bool(int(os.getenv("COLUMN_Z", "0")))
-  if column_z and (batch != 1 or swap_groups): raise ValueError("COLUMN_Z requires BATCH=1 and SWAP_GROUPS=0")
-  output_half = bool(int(os.getenv("OUTPUT_HALF", "0")))
-  int8_b = bool(int(os.getenv("INT8_B", "0")))
-  int8_a = bool(int(os.getenv("INT8_A", "0")))
-  env_src = q.make_direct_donor_src_fp32(4, threads) if rotate_buffer else \
-            q.make_direct_image_donor_src(ncols, threads, swap_groups=swap_groups)
-  env, io, sz, ro = get_envelope(dev, env_src)
-  if quad_split:
-    if ncols != 2: raise ValueError("QUAD_SPLIT=1 requires NCOLS=2")
-    shader, hregs, fregs, loop_instrs = q.build_4x8_fp32_quad_splitk_shader(
-      dev, threads, store_gap=int(os.getenv("STORE_GAP", "16")), no_reduce=bool(int(os.getenv("NO_REDUCE", "0"))),
-      k_count=int(os.getenv("K_COUNT", str(k//4))))
-  elif custom_rows:
-    if ncols != 2: raise ValueError("CUSTOM_ROWS requires NCOLS=2")
-    shader, hregs, fregs, loop_instrs = q.build_rx8_fp32_shader(
-      dev, threads, rows=custom_rows, store_gap=int(os.getenv("STORE_GAP", "16")))
-  elif quad_a:
-    if ncols != 2: raise ValueError("QUAD_A=1 requires NCOLS=2")
-    shader, hregs, fregs, loop_instrs = q.build_4x8_fp32_quad_a_shader(
-      dev, threads, store_gap=int(os.getenv("STORE_GAP", "16")))
-  elif square8:
-    if ncols != 2: raise ValueError("SQUARE8=1 requires NCOLS=2")
-    shader, hregs, fregs, loop_instrs = q.build_8x8_fp32_shader(
-      dev, threads, store_gap=int(os.getenv("STORE_GAP", "16")))
-  elif rows8:
-    if ncols != 1: raise ValueError("ROWS8=1 requires NCOLS=1")
-    shader, hregs, fregs, loop_instrs = q.build_8x4_fp32_shader(
-      dev, threads, store_gap=int(os.getenv("STORE_GAP", "16")))
-  elif bool(int(os.getenv("WAKSMAN", "0"))):
-    if ncols != 2: raise ValueError("WAKSMAN=1 requires NCOLS=2")
-    shader, hregs, fregs, loop_instrs = q.build_4x8_waksman_fp32_shader(
-      dev, threads, store_gap=int(os.getenv("STORE_GAP", "16")), no_q=bool(int(os.getenv("WAKSMAN_NO_Q", "0"))))
-  elif bool(int(os.getenv("ROTATE", "0"))):
-    if ncols != 2: raise ValueError("ROTATE=1 requires NCOLS=2")
-    shader, hregs, fregs, loop_instrs = q.build_4x8_fp32_rotate_shader(
-      dev, threads, store_gap=int(os.getenv("STORE_GAP", "16")),
-      post_constant=bool(int(os.getenv("POST_CONSTANT", "0"))), image_store=not rotate_buffer,
-      k_count=int(os.getenv("K_COUNT", str(k//4))), batch_stride=m if batch > 1 else 0, batch_from_row=batch_y, k_start=k_start,
-      k_unroll=int(os.getenv("K_UNROLL", "3")), swap_groups=swap_groups, col_from_z=column_z)
-  else:
-    shader, hregs, fregs, loop_instrs = q.build_4xn_fp32_stream_shader(
-      dev, threads, ncols=ncols, coord_delay=int(os.getenv("COORD_DELAY", "-1")),
-      sync_each_col=bool(int(os.getenv("SYNC_EACH_COL", "1"))), store_gap=int(os.getenv("STORE_GAP", "16")),
-      post_constant=bool(int(os.getenv("POST_CONSTANT", "0"))), pipeline_b=bool(int(os.getenv("PIPELINE_B", "0"))),
-      component_stream=bool(int(os.getenv("COMPONENT_STREAM", "0"))),
-      component_sync_kk=int(os.getenv("COMPONENT_SYNC_KK", "0")))
-  lib = inject(env, io, sz, ro, shader, fregs=fregs, hregs=hregs, mergedregs=False)
-  a_upload = a_np.reshape(batch*m//4, 4, k).transpose(0, 2, 1).copy() if quad_split else a_np
-  if int8_a:
-    a_upload = np.clip(np.rint(a_upload.astype(np.float32)*127), -127, 127).astype(np.int8)
-    a_np = a_upload.astype(np.float32)/127
-  if int8_b:
-    b_upload = np.clip(np.rint(b_np.astype(np.float32)*127), -127, 127).astype(np.int8)
-    b_np = b_upload.astype(np.float32)/127
-  else: b_upload = b_np
-  a, b = upload(a_upload, dtypes.int8 if int8_a else dtypes.half), upload(b_upload, dtypes.int8 if int8_b else dtypes.half)
-  c = upload(np.zeros((batch*m, stride), np.float16 if output_half else np.float32), dtypes.half if output_half else dtypes.float)
-  specs = ([((0, dtypes.half, (m, k//4, 4)),), ((1, dtypes.half, (k, n//4, 4)),),
-            ((0, dtypes.float, None),)] if rotate_buffer else
-           [((0, dtypes.half if output_half else dtypes.float, (batch*m, stride//4, 4)),),
-            ((0, dtypes.int8 if int8_a else dtypes.half,
-              (batch*m//4, k, 4) if quad_split else (batch*m, k//4, 4)),),
-            ((1, dtypes.int8 if int8_b else dtypes.half, (batch*k, n//4, 4)),)])
-  prg = dev.runtime("gemm_h", lib, buf_dtypes=specs)
-  runs = int(os.getenv("BENCH_RUNS", "10"))
-  call_bufs = (a._buf, b._buf, c._buf) if rotate_buffer else (c._buf, a._buf, b._buf)
-  launch_x, launch_y = ((batch*m//tile_m if batch_y else m//tile_m), n//tile_n) if swap_groups else \
-                       (1 if column_z else n//tile_n, batch*m//tile_m if batch_y else m//tile_m)
-  launch_z = n//tile_n if column_z else 1 if batch_y else batch
-  times = [prg(*call_bufs, global_size=(launch_x, launch_y, launch_z),
-               local_size=(threads, 1, 1), wait=True) for _ in range(runs)]
-  got_storage = np.empty((batch*m, stride), np.float16 if output_half else np.float32)
-  got_storage.reshape(-1)[:] = c.numpy().reshape(-1)
-  got = got_storage[:, :n].reshape(batch, m, n)
-  expected = np.full((batch, m, n), 1024, np.float32) if bool(int(os.getenv("POST_CONSTANT", "0"))) else \
-             a_np[:, :, k_start*4:(k_start+int(os.getenv("K_COUNT", str(k//4))))*4].astype(np.float32) @ \
-             b_np[:, k_start*4:(k_start+int(os.getenv("K_COUNT", str(k//4))))*4].astype(np.float32)
-  delta = np.abs(got-expected)
-  passed = bool(np.allclose(got, expected, rtol=1e-4, atol=1e-4))
-  elapsed = min(times)
-  bad = int(np.count_nonzero(~np.isclose(got, expected, rtol=1e-4, atol=1e-4)))
-  print(f"shape={batch}x{m}x{n}x{k} inputs=fp16 accumulate=fp32 elapsed_ms={elapsed*1e3:.3f} "
-        f"gflops={batch*2*m*n*k/elapsed/1e9:.1f} fregs={fregs} loop_instrs={loop_instrs} "
-        f"max_abs={float(delta.max()):.9g} mean_abs={float(delta.mean()):.9g} allclose={passed} bad_count={bad}")
-  if not passed:
-    where = np.argwhere(~np.isclose(got, expected, rtol=1e-4, atol=1e-4))
-    print("bad_head=", where[:32].tolist(), "got_head=", got.reshape(-1)[:32].tolist(),
-          "expected_head=", expected.reshape(-1)[:32].tolist())
-    raise SystemExit(1)
-
-
-if __name__ == "__main__": main()

extra/gemm/qcom_fuse_attention.py

@@ -1,101 +0,0 @@
-#!/usr/bin/env python3
-"""Fuse OpenPilot's QK, softmax, and AV calls with online-softmax attention."""
-import argparse, pickle
-from dataclasses import replace
-
-from tinygrad import Device, dtypes
-from tinygrad.engine.jit import create_graph_call
-from tinygrad.uop.ops import Ops, UOp
-from extra.gemm.qcom_ir3_matmul_patch import plain_name
-
-QK, SM, AV = "r_12_32_32_4_4_8_4", "softmax512", "r_32_96_4_4_32_4"
-
-
-def aux(*specs): return (tuple(((i, dtype, shape),) for i, (dtype, shape) in enumerate(specs)),)
-
-
-def build_program(template:UOp, source:str, lib:bytes):
-  specs = ((dtypes.float, (1, 12288, 4)), (dtypes.float, (1, 13824, 4)),
-           (dtypes.float, (1, 13824, 4)), (dtypes.float, (1, 12672, 4)))
-  info = replace(template.arg, name="attention_online", global_size=(12, 8, 1), local_size=(8, 16, 1),
-                 globals=(0, 1, 2, 3), outs=(0,), ins=(1, 2, 3), aux=aux(*specs))
-  return template.replace(arg=info, src=template.src[:2]+(template.src[2].replace(arg=source), template.src[3].replace(arg=lib)))
-
-
-def make_source() -> str:
-  qdecls = "\n".join(f"  float4 q{r};" for r in range(8))
-  qloads = "\n".join(f"    q{r}=read_imagef(Q,smp,(int2)(h*9+{r}+qg*432+ql*108,0));" for r in range(8))
-  score = []
-  for r in range(8):
-    score += [f"    int kb{r}=keyb*36+h*1152+{r*4};",
-              f"    float4 k{r}0=read_imagef(K,smp,(int2)(kb{r},0));",
-              f"    float4 k{r}1=read_imagef(K,smp,(int2)(kb{r}+1,0));",
-              f"    float4 k{r}2=read_imagef(K,smp,(int2)(kb{r}+2,0));",
-              f"    float4 k{r}3=read_imagef(K,smp,(int2)(kb{r}+3,0));",
-              f"    s+=q{r}.xxxx*k{r}0+q{r}.yyyy*k{r}1+q{r}.zzzz*k{r}2+q{r}.wwww*k{r}3;"]
-  updates = ["    float4 aa,bb;"]
-  for lane, comp in enumerate("xyzw"):
-    updates += [f"    float nm{lane}=fmax(mx,s.{comp});",
-                f"    aa.{comp}=exp2((mx-nm{lane})*1.4426950408889634f);",
-                f"    bb.{comp}=exp2((s.{comp}-nm{lane})*1.4426950408889634f); mx=nm{lane};"]
-  return f"""const sampler_t smp=CLK_NORMALIZED_COORDS_FALSE|CLK_ADDRESS_CLAMP|CLK_FILTER_NEAREST;
-__attribute__((reqd_work_group_size(8,16,1)))
-__kernel void attention_online(write_only image2d_t O,read_only image2d_t Q,read_only image2d_t K,read_only image2d_t V) {{
-  int x=get_global_id(0),query=get_global_id(1),ox=get_local_id(0),ly=get_local_id(1);
-  int h=x>>3,qg=query>>2,ql=query&3;
-  __local float4 la[16],lb[16];
-{qdecls}
-  if(ox==0) {{
-{qloads}
-  }}
-  float mx=-INFINITY,den=0.0f; float4 acc=(float4)(0.0f);
-  for(int keyb=0;keyb<32;keyb++) {{
-    if(ox==0) {{
-      float4 s=(float4)(0.0f);
-{chr(10).join(score)}
-      s*=0.1767766922712326f;
-{chr(10).join(updates)}
-      la[ly]=aa; lb[ly]=bb;
-    }}
-    barrier(CLK_LOCAL_MEM_FENCE);
-    int vb=x*132+keyb*4;
-    float4 v0=read_imagef(V,smp,(int2)(vb,0)),v1=read_imagef(V,smp,(int2)(vb+1,0));
-    float4 v2=read_imagef(V,smp,(int2)(vb+2,0)),v3=read_imagef(V,smp,(int2)(vb+3,0));
-    float4 aa=la[ly],bb=lb[ly];
-    acc=acc*(float4)(aa.x)+v0*(float4)(bb.x); den=den*aa.x+bb.x;
-    acc=acc*(float4)(aa.y)+v1*(float4)(bb.y); den=den*aa.y+bb.y;
-    acc=acc*(float4)(aa.z)+v2*(float4)(bb.z); den=den*aa.z+bb.z;
-    acc=acc*(float4)(aa.w)+v3*(float4)(bb.w); den=den*aa.w+bb.w;
-    barrier(CLK_LOCAL_MEM_FENCE);
-  }}
-  write_imagef(O,(int2)(x+qg*384+ql*96,0),acc/(float4)(den));
-}}"""
-
-
-def main() -> None:
-  ap = argparse.ArgumentParser()
-  ap.add_argument("input"); ap.add_argument("output")
-  args = ap.parse_args()
-  with open(args.input, "rb") as f: model = pickle.load(f)
-  batch = model.captured.linear.src[0].src[0].src[0].src
-  source = make_source(); lib = Device["QCOM"].compiler.compile(source)
-  replacements, count = {}, 0
-  for i in range(len(batch)-2):
-    calls = batch[i:i+3]
-    if not all(x.op is Ops.CALL and x.src[0].op is Ops.PROGRAM for x in calls): continue
-    if tuple(plain_name(x.src[0].arg.name) for x in calls) != (QK, SM, AV): continue
-    qk, _sm, av = calls
-    program = build_program(qk.src[0], source, lib)
-    replacements[i] = (program.call(av.src[1], qk.src[2], qk.src[3], av.src[3]),)
-    replacements[i+1] = replacements[i+2] = ()
-    count += 1
-  if count != 18: raise ValueError(f"expected 18 attention triples, found {count}")
-  outer = model.captured.linear.src[0]
-  new_batch = [new for i, old in enumerate(batch) for new in replacements.get(i, (old,))]
-  model.captured._linear = model.captured.linear.substitute({outer:create_graph_call(new_batch)}, walk=True)
-  model.captured.__dict__.pop("linear", None)
-  with open(args.output, "wb") as f: pickle.dump(model, f)
-  print(f"wrote {args.output}: fused {count} attention triples, calls {len(batch)} -> {len(new_batch)}")
-
-
-if __name__ == "__main__": main()

extra/gemm/qcom_fuse_mlp.py

@@ -1,91 +0,0 @@
-#!/usr/bin/env python3
-"""Fuse OpenPilot transformer MLP projection pairs through local memory."""
-import argparse, itertools, pickle
-from dataclasses import replace
-
-from tinygrad import Device, dtypes
-from tinygrad.engine.jit import create_graph_call
-from tinygrad.uop.ops import Ops, UOp
-from extra.gemm.qcom_ir3_matmul_patch import plain_name
-
-
-def aux(*specs): return (tuple(((i, dtype, shape),) for i, (dtype, shape) in enumerate(specs)),)
-
-
-SOURCE = r"""#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-const sampler_t smp=CLK_NORMALIZED_COORDS_FALSE|CLK_ADDRESS_CLAMP|CLK_FILTER_NEAREST;
-__attribute__((reqd_work_group_size(128,1,1)))
-__kernel void fused_mlp(write_only image2d_t O,read_only image2d_t A,read_only image2d_t W1,
-                        __global float *MUL,__global float *BIAS,read_only image2d_t W2,
-                        read_only image2d_t X,read_only image2d_t S) {
-  int lid=get_local_id(0),row=get_group_id(1);
-  __local half4 hidden[384];
-  for(int tile=0;tile<3;tile++) {
-    int n4=lid+tile*128;
-    float4 z=(float4)(0.0f);
-    for(int k4=0;k4<96;k4++) {
-      float4 a=convert_float4(read_imageh(A,smp,(int2)(k4,row)));
-      float4 w0=convert_float4(read_imageh(W1,smp,(int2)(n4,k4*4+0)));
-      float4 w1=convert_float4(read_imageh(W1,smp,(int2)(n4,k4*4+1)));
-      float4 w2=convert_float4(read_imageh(W1,smp,(int2)(n4,k4*4+2)));
-      float4 w3=convert_float4(read_imageh(W1,smp,(int2)(n4,k4*4+3)));
-      z+=a.x*w0+a.y*w1+a.z*w2+a.w*w3;
-    }
-    z=select((float4)(0.0f),z,isgreater(z,(float4)(0.0f)));
-    hidden[n4]=convert_half4((float4)(*MUL)*z*z+(float4)(*BIAS));
-  }
-  barrier(CLK_LOCAL_MEM_FENCE);
-  if(lid<96) {
-    float4 z=(float4)(0.0f);
-    for(int k4=0;k4<384;k4++) {
-      float4 a=convert_float4(hidden[k4]);
-      float4 w0=convert_float4(read_imageh(W2,smp,(int2)(lid,k4*4+0)));
-      float4 w1=convert_float4(read_imageh(W2,smp,(int2)(lid,k4*4+1)));
-      float4 w2=convert_float4(read_imageh(W2,smp,(int2)(lid,k4*4+2)));
-      float4 w3=convert_float4(read_imageh(W2,smp,(int2)(lid,k4*4+3)));
-      z+=a.x*w0+a.y*w1+a.z*w2+a.w*w3;
-    }
-    int t=row*96+lid;
-    write_imagef(O,(int2)(t,0),read_imagef(X,smp,(int2)(t,0))*read_imagef(S,smp,(int2)(lid,0))+z);
-  }
-}"""
-
-
-def build_program(template:UOp, lib:bytes):
-  specs = ((dtypes.float, (1, 12288, 4)), (dtypes.half, (128, 96, 4)),
-           (dtypes.half, (384, 384, 4)), (dtypes.float, (1,)), (dtypes.float, (1,)),
-           (dtypes.half, (1536, 96, 4)), (dtypes.float, (1, 12288, 4)),
-           (dtypes.float, (1, 96, 4)))
-  info = replace(template.arg, name="fused_mlp", global_size=(1, 128, 1), local_size=(128, 1, 1),
-                 globals=tuple(range(8)), outs=(0,), ins=(1,2,3,4,5,6,7), aux=aux(*specs))
-  return template.replace(arg=info, src=template.src[:2]+(template.src[2].replace(arg=SOURCE), template.src[3].replace(arg=lib)))
-
-
-def main() -> None:
-  ap=argparse.ArgumentParser(); ap.add_argument("input"); ap.add_argument("output"); args=ap.parse_args()
-  with open(args.input,"rb") as f: model=pickle.load(f)
-  existing=[x.arg.slot for x in model.captured.linear.toposort() if x.op is Ops.BUFFER and hasattr(x.arg,"slot") and x.arg.slot>=0]
-  UOp.unique_num=itertools.count(max(existing,default=-1)+1)
-  batch=model.captured.linear.src[0].src[0].src[0].src
-  lib=Device["QCOM"].compiler.compile(SOURCE)
-  repl, count={},0
-  for i in range(len(batch)-3):
-    calls=batch[i:i+4]
-    if not all(x.op is Ops.CALL and x.src[0].op is Ops.PROGRAM for x in calls): continue
-    names=tuple(plain_name(x.src[0].arg.name) for x in calls)
-    if names != ("gemm_h","epi3_fp32","gemm_h","epi_fp32"): continue
-    if tuple(calls[0].src[0].arg.global_size)!=(12,8,1) or tuple(calls[2].src[0].arg.global_size)!=(3,8,1): continue
-    g1,e1,g2,e2=calls; p=build_program(g1.src[0],lib)
-    repl[i]=(p.call(e2.src[1],g1.src[1],g1.src[2],e1.src[2],e1.src[3],g2.src[2],e2.src[2],e2.src[3]),)
-    repl[i+1]=repl[i+2]=repl[i+3]=()
-    count+=1
-  if count!=17: raise ValueError(f"expected 17 MLPs, found {count}")
-  outer=model.captured.linear.src[0]
-  new_batch=[new for i,old in enumerate(batch) for new in repl.get(i,(old,))]
-  model.captured._linear=model.captured.linear.substitute({outer:create_graph_call(new_batch)},walk=True)
-  model.captured.__dict__.pop("linear",None)
-  with open(args.output,"wb") as f: pickle.dump(model,f)
-  print(f"wrote {args.output}: fused {count} MLPs, calls {len(batch)} -> {len(new_batch)}")
-
-
-if __name__=="__main__": main()

extra/gemm/qcom_fusion_check.py

@@ -1,93 +0,0 @@
-#!/usr/bin/env python3
-"""Compare the first fused OpenPilot graph operation with its original call sequence."""
-import argparse, pickle
-
-import numpy as np
-
-from tinygrad import Tensor
-from tinygrad.engine.jit import _prepare_jit_inputs
-from tinygrad.engine.realize import run_linear
-from tinygrad.uop.ops import Ops, UOp
-from extra.gemm.qcom_ir3_matmul_patch import plain_name
-
-
-def load(path):
-  with open(path, "rb") as f: return pickle.load(f)
-
-
-def batch(model): return model.captured.linear.src[0].src[0].src[0].src
-
-
-def prepared(model, corpus_path, case=0):
-  corpus = np.load(corpus_path)
-  legacy = "names" in corpus and "case0:output" not in corpus
-  names = corpus["names"].tolist() if legacy else []
-  inputs = {}
-  for name, (view, _vars, dtype, device) in zip(model.captured.expected_names, model.captured.expected_input_info):
-    key = f"s{case}_input_{names.index(name)}" if legacy else f"case{case}:input:{name}"
-    arr = corpus[key].astype(np.dtype(dtype.fmt), copy=False)
-    inputs[name] = Tensor(arr, device=device).realize()
-  return _prepare_jit_inputs((), inputs)[:2]
-
-
-def main():
-  ap = argparse.ArgumentParser()
-  ap.add_argument("base")
-  ap.add_argument("candidate")
-  ap.add_argument("--fused", required=True)
-  ap.add_argument("--original", required=True, help="comma-separated original program names")
-  ap.add_argument("--occurrence", type=int, default=0, help="zero-based matching fusion occurrence")
-  ap.add_argument("--corpus", default="/data/openpilot_validation_5seeds.npz")
-  ap.add_argument("--case", type=int, default=0)
-  ap.add_argument("--side", choices=("base", "candidate"), help="run and dump only one model in this process")
-  ap.add_argument("--dump", help=".npy output path for --side")
-  ap.add_argument("--dump-inputs", help="optional .npz path containing selected call arguments")
-  args = ap.parse_args()
-  originals = args.original.split(",")
-  if args.side:
-    if not args.dump: ap.error("--side requires --dump")
-    model = load(args.base if args.side == "base" else args.candidate)
-    calls = batch(model)
-    if args.side == "base":
-      indices = [i for i in range(len(calls)-len(originals)+1)
-                 if [plain_name(x.src[0].arg.name) for x in calls[i:i+len(originals)]] == originals]
-      end = indices[args.occurrence]+len(originals)
-    else:
-      indices = [i for i, x in enumerate(calls) if x.op is Ops.CALL and x.src[0].op is Ops.PROGRAM and
-                 plain_name(x.src[0].arg.name) == args.fused]
-      end = indices[args.occurrence]+1
-    iu, vv = prepared(model, args.corpus, args.case)
-    run_linear(UOp(Ops.LINEAR, src=tuple(calls[:end])), vv, input_uops=iu, jit=True, wait=True)
-    last_call = calls[end-1]
-    out_index = last_call.src[0].arg.outs[0]
-    out = np.array(last_call.src[out_index+1].buffer.numpy(), copy=True)
-    np.save(args.dump, out)
-    if args.dump_inputs:
-      selected = calls[indices[args.occurrence]:end]
-      np.savez(args.dump_inputs, **{f"call{ci}_arg{ai}":np.array(arg.buffer.numpy(), copy=True)
-                                    for ci, call in enumerate(selected) for ai, arg in enumerate(call.src[1:])
-                                    if arg.op in (Ops.BUFFER, Ops.SLICE)})
-    print(args.side, "index", end-1, "shape", out.shape, "min", float(out.min()), "max", float(out.max()))
-    return
-  base, cand = load(args.base), load(args.candidate)
-  bb, cb = batch(base), batch(cand)
-  bis = [i for i in range(len(bb)-len(originals)+1)
-         if [plain_name(x.src[0].arg.name) for x in bb[i:i+len(originals)]] == originals]
-  cis = [i for i, x in enumerate(cb) if x.op is Ops.CALL and x.src[0].op is Ops.PROGRAM and
-         plain_name(x.src[0].arg.name) == args.fused]
-  bi, ci = bis[args.occurrence], cis[args.occurrence]
-  iu, vv = prepared(base, args.corpus, args.case)
-  run_linear(UOp(Ops.LINEAR, src=tuple(bb[:bi+len(originals)])), vv, input_uops=iu, jit=True, wait=True)
-  bo = np.array(bb[bi+len(originals)-1].src[1].buffer.numpy(), copy=True)
-  iu, vv = prepared(cand, args.corpus, args.case)
-  run_linear(UOp(Ops.LINEAR, src=tuple(cb[:ci+1])), vv, input_uops=iu, jit=True, wait=True)
-  co = np.array(cb[ci].src[1].buffer.numpy(), copy=True)
-  d = np.abs(bo.astype(np.float32)-co.astype(np.float32))
-  at = np.unravel_index(int(d.argmax()), d.shape)
-  print("indices", bi, ci, "shape", bo.shape, "max_abs", float(d[at]), "mean_abs", float(d.mean()),
-        "at", at, "base", float(bo[at]), "candidate", float(co[at]))
-  print("base_stats", float(bo.min()), float(bo.max()), float(np.mean(np.abs(bo))))
-  print("candidate_stats", float(co.min()), float(co.max()), float(np.mean(np.abs(co))))
-
-
-if __name__ == "__main__": main()

extra/gemm/qcom_gemm.py

@@ -1,164 +0,0 @@
-#!/usr/bin/env python3
-"""FP16 GEMM benchmark for Adreno 630 with binary patching.
-
-Achieves ~190 GFLOPS via:
-1. 4 rows x 4 cols per thread IMAGE kernel (read_imageh)
-2. Binary patching to strip redundant (sy) sync flags
-3. Binary patching to convert scalar MADs to (rpt3)mad.f16
-
-Usage:
-  DEV=QCOM python3 extra/gemm/qcom_gemm.py
-  DEV=QCOM python3 extra/gemm/qcom_gemm.py --m 512 --n 512 --k 512
-"""
-import struct, ctypes, math, argparse
-from tinygrad import Device, dtypes
-from tinygrad.device import Buffer
-
-def ri(buf, l):
-    o = l * 8
-    return struct.unpack_from('<I', buf, o+4)[0], struct.unpack_from('<I', buf, o)[0]
-
-def wi(buf, l, h, lo):
-    o = l * 8
-    struct.pack_into('<I', buf, o, lo)
-    struct.pack_into('<I', buf, o+4, h)
-
-def patch_kernel(lib, strip_sync=True, merge_pairs=True, max_groups=-1):
-    """Strip redundant (sy) and convert eligible MAD groups to (rpt3)."""
-    lib = bytearray(lib)
-    io = struct.unpack_from('<I', lib, 0xc0)[0]
-    isz = struct.unpack_from('<I', lib, 0x100)[0]
-    s = bytearray(lib[io:io+isz])
-    t = isz // 8
-
-    # Strip all (sy) except the first on mad.f16 instructions
-    if strip_sync:
-        first_sy = False
-        for i in range(t):
-            h, lo = ri(s, i)
-            if (h >> 24) in (0x63, 0x73) and ((h >> 24) & 0xF) == 3 and (h >> 28) == 7:
-                if first_sy:
-                    wi(s, i, (h & 0x0FFFFFFF) | 0x60000000, lo)
-                else:
-                    first_sy = True
-
-    # Convert groups of 4 scalar MADs to (rpt3)
-    i = packed_groups = 0
-    while i < t - 3:
-        h0, l0 = ri(s, i)
-        if not ((h0 >> 24) in (0x63, 0x73) and ((h0 >> 24) & 0xF) == 3):
-            i += 1; continue
-        d0, r0 = h0 & 0xFF, (h0 >> 8) & 0x7F
-        s1 = l0 & 0xFF; s3 = (l0 >> 16) & 0xFF
-        s2 = ((h0 >> 16) & 0xFF) * 2 + (((h0 >> 8) & 0xFF) >> 7)
-        if r0 > 0 or d0 != s3:
-            i += 1; continue
-        ok = True
-        for j in range(1, 4):
-            hj, lj = ri(s, i+j)
-            if not ((hj >> 24) in (0x63, 0x73) and ((hj >> 24) & 0xF) == 3):
-                ok = False; break
-            dj = hj & 0xFF; rj = (hj >> 8) & 0x7F
-            s1j = lj & 0xFF; s3j = (lj >> 16) & 0xFF
-            s2j = ((hj >> 16) & 0xFF) * 2 + (((hj >> 8) & 0xFF) >> 7)
-            if rj != 0 or s1j != s1 or dj != d0+j or s2j != s2+j or s3j != d0+j:
-                ok = False; break
-        if ok and (max_groups < 0 or packed_groups < max_groups):
-            rb = ((h0 >> 8) & 0x80) | 3
-            # Repeat needs relative src2 as well as relative dst/src3. Without bit 15,
-            # every output lane incorrectly reuses the first weight component.
-            wi(s, i, (h0 & 0xFFFF00FF) | (rb << 8), l0 | 0x20008000)
-            for j in range(1, 4):
-                wi(s, i+j, 0, 0)
-            packed_groups += 1
-            i += 4
-        else:
-            i += 1
-
-    # Merge (rpt1)+(rpt1) into (rpt3)
-    for i in range(t - 1) if merge_pairs else ():
-        h0, l0 = ri(s, i); h1, l1 = ri(s, i+1)
-        if h0 == 0 or h1 == 0: continue
-        if not ((h0 >> 24) in (0x63, 0x73) and ((h0 >> 24) & 0xF) == 3): continue
-        if not ((h1 >> 24) in (0x63, 0x73) and ((h1 >> 24) & 0xF) == 3): continue
-        if (h0 >> 8) & 0x7F != 1 or (h1 >> 8) & 0x7F != 1: continue
-        d0, d1 = h0 & 0xFF, h1 & 0xFF
-        s10, s11 = l0 & 0xFF, l1 & 0xFF
-        s20 = ((h0 >> 16) & 0xFF) * 2 + (((h0 >> 8) & 0xFF) >> 7)
-        s21 = ((h1 >> 16) & 0xFF) * 2 + (((h1 >> 8) & 0xFF) >> 7)
-        if s10 != s11 or d1 != d0 + 2 or s21 != s20 + 2: continue
-        rb = ((h0 >> 8) & 0x80) | 3
-        wi(s, i, (h0 & 0xFFFF00FF) | (rb << 8), l0)
-        wi(s, i+1, 0, 0)
-
-    lib[io:io+isz] = s
-    return bytes(lib)
-
-
-def make_gemm_src(M, N, K, nrows=4):
-    """Generate 4-row FP16 IMAGE GEMM kernel source."""
-    K4 = K // 4
-    TM = (128 // 32) * nrows  # 4 * nrows
-    src = '#pragma OPENCL EXTENSION cl_khr_fp16 : enable\n'
-    src += 'const sampler_t smp=CLK_NORMALIZED_COORDS_FALSE|CLK_ADDRESS_CLAMP|CLK_FILTER_NEAREST;\n'
-    src += '__attribute__((reqd_work_group_size(128,1,1)))\n'
-    src += '__kernel void gemm_h(read_only image2d_t A,read_only image2d_t B,__global half *C){\n'
-    src += 'int lid=get_local_id(0);int tm=lid>>5;int tn=lid&31;\n'
-    src += 'int row=get_group_id(1)*%d+tm*%d;int col4=get_group_id(0)*32+tn;\n' % (TM, nrows)
-    for r in range(nrows):
-        src += 'half4 r%dc0=(half4)(0),r%dc1=(half4)(0),r%dc2=(half4)(0),r%dc3=(half4)(0);\n' % (r,r,r,r)
-    src += 'for(int k4=0;k4<%d;k4++){\n' % K4
-    for r in range(nrows):
-        src += 'half4 a%d=read_imageh(A,smp,(int2)(k4,row+%d));\n' % (r, r)
-    for b in range(4):
-        src += 'half4 b%d=read_imageh(B,smp,(int2)(col4,k4*4+%d));\n' % (b, b)
-    for r in range(nrows):
-        src += 'r%dc0+=a%d.xxxx*b0;r%dc1+=a%d.yyyy*b1;r%dc2+=a%d.zzzz*b2;r%dc3+=a%d.wwww*b3;\n' % (r,r,r,r,r,r,r,r)
-    src += '}\n'
-    for r in range(nrows):
-        src += 'vstore4(r%dc0+r%dc1+r%dc2+r%dc3,0,C+(row+%d)*%d+col4*4);\n' % (r,r,r,r,r,N)
-    src += '}\n'
-    return src, TM
-
-
-def run_gemm(args):
-    dev = Device['QCOM']
-    M, N, K = args.m, args.n, args.k
-    print("device=%s M=%d N=%d K=%d" % (dev.device, M, N, K))
-
-    src, TM = make_gemm_src(M, N, K, nrows=4)
-    lib = patch_kernel(dev.compiler.compile_cached(src))
-
-    a_img = dtypes.imageh((M, K//4))
-    b_img = dtypes.imageh((K, N//4))
-    a_buf = Buffer(dev.device, (K//4)*M*4, dtypes.half, preallocate=True)
-    b_buf = Buffer(dev.device, (N//4)*K*4, dtypes.half, preallocate=True)
-    c_buf = Buffer(dev.device, M*N, dtypes.half, preallocate=True)
-    ctypes.memset(int(a_buf._buf.va_addr), 0, a_buf.nbytes)
-    ctypes.memset(int(b_buf._buf.va_addr), 0, b_buf.nbytes)
-
-    prg = dev.runtime('gemm_h', lib, [[(0, a_img)], [(1, b_img)], [(2, dtypes.half.ptr())]])
-    gs = (N // 128, M // TM, 1)
-    ls = (128, 1, 1)
-
-    for _ in range(5):
-        prg(a_buf._buf, b_buf._buf, c_buf._buf, global_size=gs, local_size=ls, wait=True)
-
-    times = []
-    for _ in range(args.iters):
-        t = prg(a_buf._buf, b_buf._buf, c_buf._buf, global_size=gs, local_size=ls, wait=True)
-        if t: times.append(t)
-
-    if times:
-        best = min(times)
-        gflops = 2 * M * N * K / best / 1e9
-        print("%.1f GFLOPS  (%.1f ms)  %.0f%% of 690 peak" % (gflops, best*1e3, gflops/690*100))
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--m", type=int, default=1024)
-    parser.add_argument("--n", type=int, default=1024)
-    parser.add_argument("--k", type=int, default=1024)
-    parser.add_argument("--iters", type=int, default=20)
-    run_gemm(parser.parse_args())

extra/gemm/qcom_globalb_half_check.py

@@ -1,50 +0,0 @@
-#!/usr/bin/env python3
-"""Randomized oracle and benchmark for compiler FP16 GEMM with linear global weights."""
-import os
-import numpy as np
-
-from tinygrad import Device, dtypes
-from tinygrad.device import Buffer
-from extra.gemm.qcom_8x4_gemm import buf_copyin, buf_copyout
-
-
-def main() -> None:
-  m, n, k = (int(os.getenv(x, d)) for x, d in (("M", 128), ("N", 384), ("K", 1536)))
-  source = f"""#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-const sampler_t smp=CLK_NORMALIZED_COORDS_FALSE|CLK_ADDRESS_CLAMP|CLK_FILTER_NEAREST;
-__attribute__((reqd_work_group_size(128,1,1)))
-__kernel void gemm_globalb(read_only image2d_t A,__global half *B,__global half *C) {{
-  int lid=get_local_id(0),tm=lid>>5,tid=lid&31;
-  int row=get_group_id(1)*16+tm*4,col4=get_group_id(0)*32+tid;
-  half4 r0=(half4)(0),r1=(half4)(0),r2=(half4)(0),r3=(half4)(0);
-  for(int k4=0;k4<{k//4};k4++) {{
-    half4 a0=read_imageh(A,smp,(int2)(k4,row)),a1=read_imageh(A,smp,(int2)(k4,row+1));
-    half4 a2=read_imageh(A,smp,(int2)(k4,row+2)),a3=read_imageh(A,smp,(int2)(k4,row+3));
-    int p=(k4*4)*{n}+col4*4;
-    half4 b0=vload4(0,B+p),b1=vload4(0,B+p+{n}),b2=vload4(0,B+p+{2*n}),b3=vload4(0,B+p+{3*n});
-    r0+=a0.xxxx*b0+a0.yyyy*b1+a0.zzzz*b2+a0.wwww*b3;
-    r1+=a1.xxxx*b0+a1.yyyy*b1+a1.zzzz*b2+a1.wwww*b3;
-    r2+=a2.xxxx*b0+a2.yyyy*b1+a2.zzzz*b2+a2.wwww*b3;
-    r3+=a3.xxxx*b0+a3.yyyy*b1+a3.zzzz*b2+a3.wwww*b3;
-  }}
-  vstore4(r0,0,C+row*{n}+col4*4); vstore4(r1,0,C+(row+1)*{n}+col4*4);
-  vstore4(r2,0,C+(row+2)*{n}+col4*4); vstore4(r3,0,C+(row+3)*{n}+col4*4);
-}}"""
-  dev = Device["QCOM"]
-  lib = dev.compiler.compile(source)
-  rng = np.random.default_rng(4)
-  a = (rng.standard_normal((m, k))*0.05).astype(np.float16)
-  b = (rng.standard_normal((k, n))*0.05).astype(np.float16)
-  ab, bb, cb = (Buffer("QCOM", x.size, dtypes.half).allocate() for x in (a, b, np.empty((m, n), np.float16)))
-  buf_copyin(ab, memoryview(a).cast("B")); buf_copyin(bb, memoryview(b).cast("B"))
-  prg = dev.runtime("gemm_globalb", lib, buf_dtypes=[((0, dtypes.half, (m, k//4, 4)),),
-                    ((0, dtypes.half, None),), ((0, dtypes.half, None),)])
-  times = [prg(ab._buf, bb._buf, cb._buf, global_size=(n//128, m//16, 1), local_size=(128, 1, 1), wait=True) for _ in range(10)]
-  got = np.empty((m, n), np.float16); buf_copyout(cb, memoryview(got).cast("B"))
-  expected = a.astype(np.float32) @ b.astype(np.float32)
-  err = np.abs(got.astype(np.float32)-expected)
-  print(f"ms={min(times)*1e3:.4f} gflops={2*m*n*k/min(times)/1e9:.1f} max={err.max():.8g} mean={err.mean():.8g} "
-        f"allclose={np.allclose(got, expected, rtol=.01, atol=.01)} finite={np.isfinite(got).all()}")
-
-
-if __name__ == "__main__": main()

extra/gemm/qcom_graph_critical_path.py

@@ -1,54 +0,0 @@
-#!/usr/bin/env python3
-"""Estimate a captured QCOM graph's data-dependency critical path from a call profile."""
-import argparse, hashlib, pickle, re, sys
-
-from tinygrad.uop.ops import Ops
-from extra.gemm.qcom_ir3_matmul_patch import plain_name
-
-LINE = re.compile(r"^\s*[\d.]+ ms.*?total=\s*[\d.]+ ms (\S+?)(?: global=|$)")
-
-
-def main() -> None:
-  ap = argparse.ArgumentParser()
-  ap.add_argument("model")
-  ap.add_argument("profile")
-  ap.add_argument("--min-duration", type=float, default=.1)
-  args = ap.parse_args()
-  durations = {}
-  profile = sys.stdin if args.profile == "-" else open(args.profile)
-  for line in profile:
-    if not (m := LINE.match(line)): continue
-    key = m.group(1)
-    durations[key] = float(line.split("ms", 1)[0])
-  with open(args.model, "rb") as f: model = pickle.load(f)
-  batch = model.captured.linear.src[0].src[0].src[0].src
-  finish, writer, records = {}, {}, []
-  for index, call in enumerate(batch):
-    if call.op is not Ops.CALL or call.src[0].op is not Ops.PROGRAM: continue
-    program = call.src[0]
-    name = plain_name(program.arg.name)
-    digest = hashlib.sha1(program.src[3].arg).hexdigest()[:8]
-    key = f"{name}#{digest}"
-    duration = durations.get(key, durations.get(name, 0.0))
-    cid = len(records)
-    deps = [(finish.get(writer.get(arg), 0.0), writer.get(arg)) for arg in call.src[1:]]
-    start, pred = max(deps, default=(0.0, None), key=lambda x:x[0])
-    finish[cid] = start+duration
-    for out in program.arg.outs: writer[call.src[out+1]] = cid
-    records.append((cid, index, key, duration, start, start+duration, pred))
-  end = max(records, key=lambda x:x[5])
-  chain, cur = [], end[0]
-  by_call = {x[0]:x for x in records}
-  while cur is not None:
-    rec = by_call[cur]
-    chain.append(rec)
-    cur = rec[6]
-  chain.reverse()
-  print(f"profiled_total_ms={sum(x[3] for x in records):.3f} critical_path_ms={end[5]:.3f} "
-        f"profiled_calls={sum(x[3] > 0 for x in records)}/{len(records)}")
-  print(f"critical path (profiled calls >={args.min_duration} ms):")
-  for _, index, key, duration, start, stop, _ in chain:
-    if duration >= args.min_duration: print(f"{index:4d} {start:8.3f}->{stop:8.3f} {duration:7.3f} {key}")
-
-
-if __name__ == "__main__": main()

extra/gemm/qcom_ibo_probe.py

@@ -1,99 +0,0 @@
-#!/usr/bin/env python3
-"""Compile minimal image-buffer kernels and show the generated A630 ISA."""
-import numpy as np
-import struct
-
-from tinygrad import Device, dtypes
-from tinygrad.device import Buffer
-from extra.gemm.ir3asm import disasm, get_envelope
-
-
-KERNELS = {
-  "half4": r"""#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-__attribute__((reqd_work_group_size(128,1,1)))
-__kernel void probe(read_only image1d_buffer_t src, __global half4 *dst) {
-  int i = get_global_id(0); dst[i] = read_imageh(src, i);
-}""",
-  "float4": r"""__attribute__((reqd_work_group_size(128,1,1)))
-__kernel void probe(read_only image1d_buffer_t src, __global float4 *dst) {
-  int i = get_global_id(0); dst[i] = read_imagef(src, i);
-}""",
-  "uint4": r"""__attribute__((reqd_work_group_size(128,1,1)))
-__kernel void probe(read_only image1d_buffer_t src, __global uint4 *dst) {
-  int i = get_global_id(0); dst[i] = read_imageui(src, i);
-}""",
-  "read_write_half4": r"""#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-__attribute__((reqd_work_group_size(128,1,1)))
-__kernel void probe(read_write image1d_buffer_t src, __global half4 *dst) {
-  int i = get_global_id(0); dst[i] = read_imageh(src, i);
-}""",
-  "read_write_2d": r"""#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-__attribute__((reqd_work_group_size(128,1,1)))
-__kernel void probe(read_write image2d_t src, __global half4 *dst) {
-  int i = get_global_id(0); dst[i] = read_imageh(src, (int2)(i, 0));
-}""",
-}
-
-
-def binfos(lib:bytes, name:str="probe") -> list[tuple[int, int]]:
-  u32 = lambda off: struct.unpack_from("<I", lib, off)[0]
-  image_desc_off = u32(0x110)
-  samp_count = u32(image_desc_off + 0xdc)
-  off = (image_desc_off + 0x158 + len(name) + 3) & -4
-  off += 8 * samp_count
-  ret = []
-  while off + 32 <= len(lib):
-    vals = struct.unpack_from("<8I", lib, off)
-    if vals[0] == 0: break
-    ret.append((vals[3] * 4, vals[7]))
-    off += vals[0]
-  return ret
-
-
-def main() -> None:
-  dev = Device["QCOM"]
-  for name, src in KERNELS.items():
-    try:
-      lib, image_off, image_size, _ = get_envelope(dev, src)
-      prg = dev.runtime("probe", bytes(lib), buf_dtypes=[])
-      print(f"=== {name}: image={image_size} tex={prg.tex_cnt} ibo={prg.ibo_cnt} samp={prg.samp_cnt} binfos={binfos(bytes(lib))} ===")
-      print(disasm(bytes(lib[image_off:image_off+image_size])))
-    except Exception as exc:
-      print(f"=== {name}: ERROR {type(exc).__name__}: {exc} ===")
-
-  count = 4096
-  values = np.random.default_rng(123).standard_normal((count, 4)).astype(np.float16)
-  src_buf = Buffer("QCOM", values.size, dtypes.half).allocate()
-  dst_buf = Buffer("QCOM", values.size, dtypes.half).allocate()
-  src_buf.copyin(memoryview(values).cast("B"))
-  src = KERNELS["half4"]
-  lib = dev.compiler.compile(src)
-  specs = [((0, dtypes.half, (1, count, 4)),), ((1, dtypes.half, None),)]
-  prg = dev.runtime("probe", lib, buf_dtypes=specs)
-  times = [prg(src_buf._buf, dst_buf._buf, global_size=(count//128, 1, 1),
-               local_size=(128, 1, 1), wait=True)*1e3 for _ in range(20)]
-  got = np.empty_like(values)
-  dst_buf.copyout(memoryview(got).cast("B"))
-  print(f"=== half4 runtime: best_ms={min(times):.6f} exact={np.array_equal(got, values)} "
-        f"max_abs={float(np.max(np.abs(got.astype(np.float32)-values.astype(np.float32))))} ===")
-
-  count = 147456
-  values = np.random.default_rng(456).standard_normal((count, 4)).astype(np.float16)
-  src_buf = Buffer("QCOM", values.size, dtypes.half).allocate()
-  dst_buf = Buffer("QCOM", values.size, dtypes.half).allocate()
-  src_buf.copyin(memoryview(values).cast("B"))
-  lib = dev.compiler.compile(KERNELS["read_write_half4"])
-  specs = [((0, dtypes.half, (1, count, 4)),), ((1, dtypes.half, None),)]
-  prg = dev.runtime("probe", lib, buf_dtypes=specs)
-  times = [prg(src_buf._buf, dst_buf._buf, global_size=(count//128, 1, 1),
-               local_size=(128, 1, 1), wait=True)*1e3 for _ in range(20)]
-  got = np.empty_like(values)
-  dst_buf.copyout(memoryview(got).cast("B"))
-  print(f"=== read_write_half4 runtime: best_ms={min(times):.6f} exact={np.array_equal(got, values)} "
-        f"max_abs={float(np.max(np.abs(got.astype(np.float32)-values.astype(np.float32))))} ===")
-  print("expected_head", values[:4].tolist(), "got_head", got[:4].tolist())
-  bad = np.flatnonzero(np.any(got != values, axis=1))
-  print("first_bad", int(bad[0]) if bad.size else None, "bad_vectors", int(bad.size))
-
-
-if __name__ == "__main__": main()

extra/gemm/qcom_int8_cached_weights.py

@@ -1,99 +0,0 @@
-#!/usr/bin/env python3
-"""Quantize selected cached QCOM GEMM weights to normalized int8 textures."""
-import argparse, itertools, pickle
-from dataclasses import replace
-
-import numpy as np
-
-from tinygrad import Device, dtypes
-from tinygrad.engine.jit import create_graph_call
-from tinygrad.uop.ops import Ops, UOp
-from extra.gemm.qcom_ir3_matmul_patch import plain_name
-
-
-def graph_batch(model): return model.captured.linear.src[0].src[0].src[0].src
-
-
-def adapt_aux_dtype(aux, index, dtype):
-  if isinstance(aux, tuple) and len(aux) == 3 and aux[0] == index and isinstance(aux[0], int):
-    return (aux[0], dtype, aux[2])
-  return tuple(adapt_aux_dtype(x, index, dtype) for x in aux) if isinstance(aux, tuple) else aux
-
-
-def main() -> None:
-  parser = argparse.ArgumentParser()
-  parser.add_argument("input")
-  parser.add_argument("output")
-  parser.add_argument("--geometry", type=int, choices=(3, 12), required=True)
-  parser.add_argument("--indices", default="", help="comma-separated occurrence indices; default is all")
-  parser.add_argument("--per-channel", action="store_true", help="scale each output channel independently")
-  args = parser.parse_args()
-  with open(args.input, "rb") as f: model = pickle.load(f)
-  batch = graph_batch(model)
-  existing_slots = [x.arg.slot for x in model.captured.linear.toposort()
-                    if x.op is Ops.BUFFER and hasattr(x.arg, "slot") and x.arg.slot >= 0]
-  UOp.unique_num = itertools.count(max(existing_slots, default=-1) + 1)
-  candidates = [(i, call) for i, call in enumerate(batch) if call.op is Ops.CALL and call.src[0].op is Ops.PROGRAM and
-                plain_name(call.src[0].arg.name) == "gemm_h" and int(call.src[0].arg.global_size[0]) == args.geometry]
-  selected = {int(x) for x in args.indices.split(",") if x} if args.indices else set(range(len(candidates)))
-  replacements = {}
-  for occurrence, (index, call) in enumerate(candidates):
-    if occurrence not in selected: continue
-    if index+1 >= len(batch): raise ValueError(f"GEMM {occurrence} has no epilogue")
-    epi_call = batch[index+1]
-    epi_name = plain_name(epi_call.src[0].arg.name)
-    expected_epi = "epi_fp32" if args.geometry == 3 else "epi3_fp32"
-    if epi_name != expected_epi: raise ValueError(f"GEMM {occurrence} is followed by {epi_name}, expected {expected_epi}")
-    weights = np.asarray(call.src[2].buffer.numpy(), dtype=np.float32)
-    k, n = ((1536, 384) if args.geometry == 3 else (384, 1536))
-    scales = np.max(np.abs(weights.reshape(k, n)), axis=0) if args.per_channel else np.asarray([np.max(np.abs(weights))])
-    if not np.isfinite(scales).all():
-      raise ValueError(f"invalid scale range {scales.min()}..{scales.max()} for GEMM {occurrence}")
-    scales[scales == 0] = 1.0
-    quantized = np.clip(np.rint(weights.reshape(k, n)/scales.reshape(1, -1)*127.0), -127, 127).astype(np.int8)
-    weight = UOp.new_buffer("QCOM", quantized.size, dtypes.int8)
-    weight.buffer.ensure_allocated()
-    weight.buffer.copyin(memoryview(quantized).cast("B"))
-    program = call.src[0].replace(arg=replace(call.src[0].arg, aux=adapt_aux_dtype(call.src[0].arg.aux, 1, dtypes.int8)))
-    replacements[index] = call.replace(src=(program, call.src[1], weight, *call.src[3:]))
-
-    epi_program = epi_call.src[0]
-    source = epi_program.src[2].arg
-    needle = "float4 v=vload4(0,C+row*1024+col*4);" if args.geometry == 3 else \
-             "float4 z=vload4(0,C+row*2048+col*4);"
-    if args.per_channel:
-      source = source.replace("__global float *C)", "__global float *C,__global float *Q)")
-      replacement = needle + (" v*=vload4(0,Q+col*4);" if args.geometry == 3 else " z*=vload4(0,Q+col*4);")
-    else:
-      scale = float(scales[0])
-      replacement = needle + (f" v*=(float4)({scale:.9g}f);" if args.geometry == 3 else f" z*=(float4)({scale:.9g}f);")
-    if needle not in source: raise ValueError(f"epilogue source pattern missing for GEMM {occurrence}")
-    source = source.replace(needle, replacement)
-    lib = Device["QCOM"].compiler.compile(source)
-    if args.per_channel:
-      scale_buf = UOp.new_buffer("QCOM", n, dtypes.float)
-      scale_buf.buffer.ensure_allocated()
-      scale_buf.buffer.copyin(memoryview(np.ascontiguousarray(scales, dtype=np.float32)).cast("B"))
-      info = epi_program.arg
-      old_aux = info.aux[0]
-      info = replace(info, globals=info.globals+(len(epi_call.src)-1,), ins=info.ins+(len(epi_call.src)-1,),
-                     aux=(old_aux+(((len(epi_call.src)-1, dtypes.float, (n,)),),),))
-      epi_program = epi_program.replace(arg=info, src=epi_program.src[:2] +
-                                        (epi_program.src[2].replace(arg=source), epi_program.src[3].replace(arg=lib)))
-      replacements[index+1] = epi_call.replace(src=(epi_program, *epi_call.src[1:], scale_buf))
-      print(f"geometry={args.geometry} occurrence={occurrence} scale={scales.min():.8g}..{scales.max():.8g}")
-    else:
-      epi_program = epi_program.replace(src=epi_program.src[:2] +
-                                        (epi_program.src[2].replace(arg=source), epi_program.src[3].replace(arg=lib)))
-      replacements[index+1] = epi_call.replace(src=(epi_program, *epi_call.src[1:]))
-      print(f"geometry={args.geometry} occurrence={occurrence} scale={scale:.8g}")
-
-  outer = model.captured.linear.src[0]
-  new_outer = create_graph_call([replacements.get(i, call) for i, call in enumerate(batch)])
-  model.captured._linear = model.captured.linear.substitute({outer: new_outer}, walk=True)
-  model.captured.__dict__.pop("linear", None)
-  with open(args.output, "wb") as f: pickle.dump(model, f)
-  print(f"wrote {args.output} with {len(replacements)//2} int8 GEMMs")
-
-
-if __name__ == "__main__": main()

extra/gemm/qcom_local_gemm_check.py

@@ -1,81 +0,0 @@
-#!/usr/bin/env python3
-"""Random-data oracle and benchmark for a cooperative-local FP16 QCOM GEMM."""
-import argparse, statistics
-
-import numpy as np
-
-from tinygrad import Device, dtypes
-from tinygrad.device import Buffer
-
-
-SRC = r"""
-#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-__attribute__((reqd_work_group_size(128,1,1)))
-__kernel void local_gemm(__global const half *A, __global const half *B, __global half *C) {
-  __local half la[32*16];
-  __local half lb[16*64];
-  const int lid=get_local_id(0), lr=lid>>5, lc=lid&31;
-  const int row0=get_group_id(1)*32+lr*8;
-  const int col0=get_group_id(0)*64+lc*2;
-  half2 c0=(half2)(0),c1=(half2)(0),c2=(half2)(0),c3=(half2)(0);
-  half2 c4=(half2)(0),c5=(half2)(0),c6=(half2)(0),c7=(half2)(0);
-  for (int k0=0;k0<@K@;k0+=16) {
-    for (int i=lid;i<32*16;i+=128) la[i]=A[(get_group_id(1)*32+i/16)*@K@+k0+i%16];
-    for (int i=lid;i<16*64;i+=128) lb[i]=B[(k0+i/64)*@N@+get_group_id(0)*64+i%64];
-    barrier(CLK_LOCAL_MEM_FENCE);
-    #pragma unroll
-    for (int kk=0;kk<16;kk++) {
-      half2 b=vload2(0,lb+kk*64+lc*2);
-      c0+=la[(lr*8+0)*16+kk]*b; c1+=la[(lr*8+1)*16+kk]*b;
-      c2+=la[(lr*8+2)*16+kk]*b; c3+=la[(lr*8+3)*16+kk]*b;
-      c4+=la[(lr*8+4)*16+kk]*b; c5+=la[(lr*8+5)*16+kk]*b;
-      c6+=la[(lr*8+6)*16+kk]*b; c7+=la[(lr*8+7)*16+kk]*b;
-    }
-    barrier(CLK_LOCAL_MEM_FENCE);
-  }
-  vstore2(c0,0,C+(row0+0)*@N@+col0); vstore2(c1,0,C+(row0+1)*@N@+col0);
-  vstore2(c2,0,C+(row0+2)*@N@+col0); vstore2(c3,0,C+(row0+3)*@N@+col0);
-  vstore2(c4,0,C+(row0+4)*@N@+col0); vstore2(c5,0,C+(row0+5)*@N@+col0);
-  vstore2(c6,0,C+(row0+6)*@N@+col0); vstore2(c7,0,C+(row0+7)*@N@+col0);
-}
-"""
-
-
-def upload(x:np.ndarray, dtype) -> Buffer:
-  ret = Buffer("QCOM", x.size, dtype).allocate()
-  ret.copyin(memoryview(np.ascontiguousarray(x)).cast("B"))
-  return ret
-
-
-def main() -> None:
-  ap = argparse.ArgumentParser()
-  ap.add_argument("--m", type=int, default=128)
-  ap.add_argument("--n", type=int, default=1536)
-  ap.add_argument("--k", type=int, default=384)
-  ap.add_argument("--seed", type=int, default=0)
-  ap.add_argument("--runs", type=int, default=10)
-  args = ap.parse_args()
-  if args.m % 32 or args.n % 64 or args.k % 16: raise ValueError("M,N,K must divide the 32x64x16 tile")
-  rng = np.random.default_rng(args.seed)
-  a = (rng.standard_normal((args.m,args.k))*.05).astype(np.float16)
-  b = (rng.standard_normal((args.k,args.n))*.05).astype(np.float16)
-  ab, bb = upload(a, dtypes.half), upload(b, dtypes.half)
-  cb = upload(np.zeros((args.m,args.n), np.float16), dtypes.half)
-  dev = Device["QCOM"]
-  src = SRC.replace("@K@", str(args.k)).replace("@N@", str(args.n))
-  lib = dev.compiler.compile_cached(src)
-  prg = dev.runtime("local_gemm", lib, buf_dtypes=[((0,dtypes.half,None),)]*3)
-  gs, ls = (args.n//64,args.m//32,1), (128,1,1)
-  for _ in range(2): prg(ab._buf,bb._buf,cb._buf,global_size=gs,local_size=ls,wait=True)
-  times = [prg(ab._buf,bb._buf,cb._buf,global_size=gs,local_size=ls,wait=True)*1e3 for _ in range(args.runs)]
-  got = np.empty((args.m,args.n),np.float16)
-  cb.copyout(memoryview(got).cast("B"))
-  expected = a.astype(np.float32) @ b.astype(np.float32)
-  delta = np.abs(got.astype(np.float32)-expected)
-  med, best = statistics.median(times), min(times)
-  print(f"best_ms={best:.4f} median_ms={med:.4f} gflops={2*args.m*args.n*args.k/best/1e6:.1f} "
-        f"max_abs={delta.max():.9g} mean_abs={delta.mean():.9g} allclose={np.allclose(got,expected,rtol=.02,atol=.02)}")
-  if not np.allclose(got,expected,rtol=.02,atol=.02): raise SystemExit(1)
-
-
-if __name__ == "__main__": main()

extra/gemm/qcom_local_image_gemm_check.py

@@ -1,152 +0,0 @@
-#!/usr/bin/env python3
-"""Random-data benchmark for cooperative image-to-local FP16 GEMM."""
-import argparse, statistics
-
-import numpy as np
-
-from tinygrad import Device, dtypes
-from tinygrad.device import Buffer
-
-
-def source(n:int, k:int, stride:int, bk4:int, fp32_acc:bool=False) -> str:
-  acc_t, zero, conv = ("float4", "(float4)(0)", "convert_float4") if fp32_acc else ("half4", "(half4)(0)", "")
-  out_t = "float" if fp32_acc else "half"
-  return f"""#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-const sampler_t smp=CLK_NORMALIZED_COORDS_FALSE|CLK_ADDRESS_CLAMP|CLK_FILTER_NEAREST;
-__attribute__((reqd_work_group_size(128,1,1)))
-__kernel void local_image_gemm(read_only image2d_t A, read_only image2d_t B, __global {out_t} *C) {{
-  __local half4 la[{32*bk4}];
-  __local half4 lb[{bk4*4*32}];
-  int lid=get_local_id(0), tm=lid>>5, tid=lid&31;
-  int row0=get_group_id(1)*32+tm*8, col4=get_group_id(0)*32+tid;
-  {acc_t} c0={zero},c1={zero},c2={zero},c3={zero};
-  {acc_t} c4={zero},c5={zero},c6={zero},c7={zero};
-  for(int kb=0;kb<{k//4};kb+={bk4}) {{
-    for(int i=lid;i<{32*bk4};i+=128) {{
-      int r=i/{bk4},q=i-r*{bk4};
-      la[i]=read_imageh(A,smp,(int2)(kb+q,get_group_id(1)*32+r));
-    }}
-    for(int i=lid;i<{bk4*4*32};i+=128) {{
-      int y=i>>5,x=i&31;
-      lb[i]=read_imageh(B,smp,(int2)(get_group_id(0)*32+x,kb*4+y));
-    }}
-    barrier(CLK_LOCAL_MEM_FENCE);
-    #pragma unroll
-    for(int q=0;q<{bk4};q++) {{
-      {acc_t} a0={conv}(la[(tm*8+0)*{bk4}+q]),a1={conv}(la[(tm*8+1)*{bk4}+q]);
-      {acc_t} a2={conv}(la[(tm*8+2)*{bk4}+q]),a3={conv}(la[(tm*8+3)*{bk4}+q]);
-      {acc_t} a4={conv}(la[(tm*8+4)*{bk4}+q]),a5={conv}(la[(tm*8+5)*{bk4}+q]);
-      {acc_t} a6={conv}(la[(tm*8+6)*{bk4}+q]),a7={conv}(la[(tm*8+7)*{bk4}+q]);
-      {acc_t} b0={conv}(lb[(q*4+0)*32+tid]),b1={conv}(lb[(q*4+1)*32+tid]);
-      {acc_t} b2={conv}(lb[(q*4+2)*32+tid]),b3={conv}(lb[(q*4+3)*32+tid]);
-      c0+=a0.xxxx*b0+a0.yyyy*b1+a0.zzzz*b2+a0.wwww*b3;
-      c1+=a1.xxxx*b0+a1.yyyy*b1+a1.zzzz*b2+a1.wwww*b3;
-      c2+=a2.xxxx*b0+a2.yyyy*b1+a2.zzzz*b2+a2.wwww*b3;
-      c3+=a3.xxxx*b0+a3.yyyy*b1+a3.zzzz*b2+a3.wwww*b3;
-      c4+=a4.xxxx*b0+a4.yyyy*b1+a4.zzzz*b2+a4.wwww*b3;
-      c5+=a5.xxxx*b0+a5.yyyy*b1+a5.zzzz*b2+a5.wwww*b3;
-      c6+=a6.xxxx*b0+a6.yyyy*b1+a6.zzzz*b2+a6.wwww*b3;
-      c7+=a7.xxxx*b0+a7.yyyy*b1+a7.zzzz*b2+a7.wwww*b3;
-    }}
-    barrier(CLK_LOCAL_MEM_FENCE);
-  }}
-  vstore4(c0,0,C+(row0+0)*{stride}+col4*4); vstore4(c1,0,C+(row0+1)*{stride}+col4*4);
-  vstore4(c2,0,C+(row0+2)*{stride}+col4*4); vstore4(c3,0,C+(row0+3)*{stride}+col4*4);
-  vstore4(c4,0,C+(row0+4)*{stride}+col4*4); vstore4(c5,0,C+(row0+5)*{stride}+col4*4);
-  vstore4(c6,0,C+(row0+6)*{stride}+col4*4); vstore4(c7,0,C+(row0+7)*{stride}+col4*4);
-}}"""
-
-
-def global_b_source(n:int, k:int, stride:int) -> str:
-  rows = "\n".join(f"  half4 c{r}=(half4)(0);" for r in range(8))
-  aloads = "\n".join(f"    half4 a{r}=read_imageh(A,smp,(int2)(q,row0+{r}));" for r in range(8))
-  mads = "\n".join(f"    c{r}+=a{r}.xxxx*b0+a{r}.yyyy*b1+a{r}.zzzz*b2+a{r}.wwww*b3;" for r in range(8))
-  stores = "\n".join(f"  vstore4(c{r},0,C+(row0+{r})*{stride}+col4*4);" for r in range(8))
-  return f"""#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-const sampler_t smp=CLK_NORMALIZED_COORDS_FALSE|CLK_ADDRESS_CLAMP|CLK_FILTER_NEAREST;
-__attribute__((reqd_work_group_size(128,1,1)))
-__kernel void local_image_gemm(read_only image2d_t A, __global half *B, __global half *C) {{
-  int lid=get_local_id(0),tm=lid>>5,tid=lid&31;
-  int row0=get_group_id(1)*32+tm*8,col4=get_group_id(0)*32+tid;
-{rows}
-  for(int q=0;q<{k//4};q++) {{
-{aloads}
-    int p=q*4*{n}+col4*4;
-    half4 b0=vload4(0,B+p),b1=vload4(0,B+p+{n});
-    half4 b2=vload4(0,B+p+{2*n}),b3=vload4(0,B+p+{3*n});
-{mads}
-  }}
-{stores}
-}}"""
-
-
-def local_b_fp32_source(n:int, k:int, stride:int, bk4:int) -> str:
-  aloads = ",".join(f"a{r}=convert_float4(read_imageh(A,smp,(int2)(kb+q,row0+{r})))" for r in range(8))
-  mads = "\n".join(f"      c{r}+=a{r}.xxxx*b0+a{r}.yyyy*b1+a{r}.zzzz*b2+a{r}.wwww*b3;" for r in range(8))
-  stores = "\n".join(f"  vstore4(c{r},0,C+(row0+{r})*{stride}+col4*4);" for r in range(8))
-  return f"""#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-const sampler_t smp=CLK_NORMALIZED_COORDS_FALSE|CLK_ADDRESS_CLAMP|CLK_FILTER_NEAREST;
-__attribute__((reqd_work_group_size(128,1,1)))
-__kernel void local_image_gemm(read_only image2d_t A,read_only image2d_t B,__global float *C) {{
-  __local half4 lb[{bk4*4*32}];
-  int lid=get_local_id(0),tm=lid>>5,tid=lid&31,row0=get_group_id(1)*32+tm*8,col4=get_group_id(0)*32+tid;
-  float4 c0=(float4)(0),c1=(float4)(0),c2=(float4)(0),c3=(float4)(0);
-  float4 c4=(float4)(0),c5=(float4)(0),c6=(float4)(0),c7=(float4)(0);
-  for(int kb=0;kb<{k//4};kb+={bk4}) {{
-    for(int i=lid;i<{bk4*4*32};i+=128) {{ int y=i>>5,x=i&31;lb[i]=read_imageh(B,smp,(int2)(get_group_id(0)*32+x,kb*4+y)); }}
-    barrier(CLK_LOCAL_MEM_FENCE);
-    #pragma unroll
-    for(int q=0;q<{bk4};q++) {{
-      float4 {aloads};
-      float4 b0=convert_float4(lb[(q*4+0)*32+tid]),b1=convert_float4(lb[(q*4+1)*32+tid]);
-      float4 b2=convert_float4(lb[(q*4+2)*32+tid]),b3=convert_float4(lb[(q*4+3)*32+tid]);
-{mads}
-    }}
-    barrier(CLK_LOCAL_MEM_FENCE);
-  }}
-{stores}
-}}"""
-
-
-def upload(values:np.ndarray, dtype) -> Buffer:
-  ret = Buffer("QCOM", values.size, dtype).allocate()
-  ret.copyin(memoryview(np.ascontiguousarray(values)).cast("B"))
-  return ret
-
-
-def main() -> None:
-  ap = argparse.ArgumentParser()
-  ap.add_argument("--m", type=int, default=128); ap.add_argument("--n", type=int, default=1536)
-  ap.add_argument("--k", type=int, default=384); ap.add_argument("--stride", type=int, default=2048)
-  ap.add_argument("--bk4", type=int, choices=(2, 4, 8, 16), default=8)
-  ap.add_argument("--global-b", action="store_true")
-  ap.add_argument("--fp32-acc", action="store_true")
-  ap.add_argument("--b-only", action="store_true")
-  ap.add_argument("--seed", type=int, default=0); ap.add_argument("--runs", type=int, default=10)
-  args = ap.parse_args()
-  if args.m%32 or args.n%128 or (args.k//4)%args.bk4: raise ValueError("shape does not divide tile")
-  rng=np.random.default_rng(args.seed)
-  av=(rng.standard_normal((args.m,args.k))*.05).astype(np.float16)
-  bv=(rng.standard_normal((args.k,args.n))*.05).astype(np.float16)
-  a,b=upload(av,dtypes.half),upload(bv,dtypes.half)
-  out_np, out_dtype = (np.float32, dtypes.float) if args.fp32_acc else (np.float16, dtypes.half)
-  c=upload(np.zeros((args.m,args.stride),out_np),out_dtype)
-  dev=Device["QCOM"]
-  src=(global_b_source(args.n,args.k,args.stride) if args.global_b else local_b_fp32_source(args.n,args.k,args.stride,args.bk4)
-       if args.b_only else source(args.n,args.k,args.stride,args.bk4,args.fp32_acc))
-  specs=[((0,dtypes.half,(args.m,args.k//4,4)),),
-         ((1,dtypes.half,None),) if args.global_b else ((1,dtypes.half,(args.k,args.n//4,4)),),((2,out_dtype,None),)]
-  prg=dev.runtime("local_image_gemm",dev.compiler.compile(src),buf_dtypes=specs)
-  gs,ls=(args.n//128,args.m//32,1),(128,1,1)
-  for _ in range(2): prg(a._buf,b._buf,c._buf,global_size=gs,local_size=ls,wait=True)
-  times=[prg(a._buf,b._buf,c._buf,global_size=gs,local_size=ls,wait=True)*1e3 for _ in range(args.runs)]
-  storage=np.empty((args.m,args.stride),out_np); c.copyout(memoryview(storage).cast("B"))
-  got=storage[:,:args.n].astype(np.float32); expected=av.astype(np.float32)@bv.astype(np.float32)
-  delta=np.abs(got-expected); best=min(times)
-  print(f"bk4={args.bk4} best_ms={best:.4f} median_ms={statistics.median(times):.4f} "
-        f"gflops={2*args.m*args.n*args.k/best/1e6:.1f} max_abs={delta.max():.9g} "
-        f"mean_abs={delta.mean():.9g} accumulate={'fp32' if args.fp32_acc else 'fp16'} "
-        f"allclose={np.allclose(got,expected,rtol=1e-4 if args.fp32_acc else .02,atol=1e-4 if args.fp32_acc else .02)}")
-
-
-if __name__ == "__main__": main()

extra/gemm/qcom_openpilot_batch_head.py

@@ -1,110 +0,0 @@
-#!/usr/bin/env python3
-"""Batch adjacent independent openpilot head kernels into one QCOM launch."""
-import argparse, pickle, re
-from dataclasses import replace
-
-from tinygrad import Device
-from tinygrad.engine.jit import create_graph_call
-from tinygrad.uop.ops import Ops
-from extra.gemm.qcom_ir3_matmul_patch import plain_name
-
-MAX_BATCH={"r_256_4_128_4":4,"r_128_16_4_16_4":4,"r_128_16_4_32_4":4,
-           "r_8_16_4_8_4":4,"r_8_4_8_4":4,"r_8_4_8_4n1":4}
-MAX_BATCH.update({"r_16_16_4_8_4":4,"r_4_16_4_8_4":4,"r_16_16_4_4":4,
-                  "r_4_4_4_4":4,"r_16_16_4_4n1":4,"r_4_4_4_4n1":4})
-
-
-def batched_source(source:str, name:str, batch_count:int) -> str:
-  match=re.search(r"__kernel void \w+\((.*?)\) \{",source,re.S)
-  if match is None: raise RuntimeError("kernel signature not found")
-  declarations=[x.strip() for x in match.group(1).split(",")]
-  arg_names=[x.rsplit(" ",1)[1] for x in declarations]
-  renamed=[]
-  bodies=[]
-  body=source[match.end():source.rfind("}")]
-  local_decls=re.findall(r"__attribute__\s*\(\(aligned \(\d+\)\)\)\s*__local\s+[^;]+;",body)
-  hoisted=[]
-  for batch in range(batch_count):
-    mapping={arg:f"{arg}_{batch}" for arg in arg_names}
-    renamed.extend(decl.rsplit(" ",1)[0]+" "+mapping[arg] for decl,arg in zip(declarations,arg_names))
-    branch=body
-    for declaration in local_decls:
-      local_match=re.search(r"(\w+)(\[[^;]+;)$",declaration)
-      if local_match is None: raise RuntimeError(f"local declaration not understood: {declaration}")
-      old=local_match.group(1)
-      new=f"{old}_{batch}"
-      hoisted.append(declaration[:local_match.start(1)]+new+local_match.group(2))
-      branch=branch.replace(declaration,"")
-      branch=re.sub(rf"\b{re.escape(old)}\b",new,branch)
-    for old,new in mapping.items(): branch=re.sub(rf"\b{re.escape(old)}\b",new,branch)
-    bodies.append(branch)
-  prefix=source[:match.start()]
-  count=len(declarations)
-  order=tuple(batch*count for batch in range(batch_count))+tuple(
-    batch*count+arg for batch in range(batch_count) for arg in range(1,count))
-  branches=" else ".join((f"if (get_group_id(1)=={batch}) " if batch < batch_count-1 else "")+f"{{{body}}}"
-                         for batch,body in enumerate(bodies))
-  return f"{prefix}__kernel void {name}_batch{batch_count}({','.join(renamed[i] for i in order)}) {{\n" \
-         f"{''.join(hoisted)}\n{branches}\n}}"
-
-
-def independent(calls:list) -> bool:
-  outputs={call.src[out+1] for call in calls for out in call.src[0].arg.outs}
-  return not any(arg in outputs for call in calls for i,arg in enumerate(call.src[1:]) if i not in call.src[0].arg.outs)
-
-
-def batch_head(model) -> int:
-  outer=model.captured.linear.src[0]
-  batch=list(outer.src[0].src[0].src)
-  new_batch=[]
-  combined=0
-  index=0
-  cache={}
-  while index < len(batch):
-    first=batch[index]
-    name=plain_name(first.src[0].arg.name) if first.op is Ops.CALL and first.src[0].op is Ops.PROGRAM else ""
-    if index+1 < len(batch) and name in MAX_BATCH:
-      calls=[first]
-      while index+len(calls) < len(batch) and len(calls) < MAX_BATCH[name]:
-        candidate=batch[index+len(calls)]
-        candidate_name=plain_name(candidate.src[0].arg.name) if candidate.op is Ops.CALL and candidate.src[0].op is Ops.PROGRAM else ""
-        if candidate_name != name or first.src[0].src[3].arg != candidate.src[0].src[3].arg: break
-        calls.append(candidate)
-      if len(calls) > 1 and independent(calls):
-        batch_count=len(calls)
-        program=first.src[0]
-        source=batched_source(program.src[2].arg,name,batch_count)
-        if source not in cache: cache[source]=Device["QCOM"].compiler.compile_cached(source)
-        aux0=program.arg.aux[0]
-        count=len(aux0)
-        ordered_aux=tuple(aux0[0] for _ in calls)+tuple(entry for _ in calls for entry in aux0[1:])
-        combined_aux=tuple(tuple((new_index,dtype,shape) for _old_index,dtype,shape in entry)
-                           for new_index,entry in enumerate(ordered_aux))
-        info=replace(program.arg,name=f"{name}_batch{batch_count}",global_size=(program.arg.global_size[0],batch_count,1),
-                     globals=tuple(range(count*batch_count)),outs=tuple(range(batch_count)),
-                     ins=tuple(range(batch_count,count*batch_count)),aux=(combined_aux,))
-        program=program.replace(arg=info,src=program.src[:2]+
-          (program.src[2].replace(arg=source),program.src[3].replace(arg=cache[source])))
-        new_batch.append(first.replace(src=(program,*[call.src[1] for call in calls],
-                                                *[arg for call in calls for arg in call.src[2:]])))
-        combined+=1
-        index+=batch_count
-        continue
-    new_batch.append(first)
-    index+=1
-  model.captured._linear=model.captured.linear.substitute({outer:create_graph_call(new_batch)},walk=True)
-  model.captured.__dict__.pop("linear",None)
-  return combined
-
-
-def main() -> None:
-  parser=argparse.ArgumentParser()
-  parser.add_argument("input")
-  parser.add_argument("output")
-  args=parser.parse_args()
-  with open(args.input,"rb") as f:model=pickle.load(f)
-  print("combined",batch_head(model))
-  with open(args.output,"wb") as f:pickle.dump(model,f)
-
-
-if __name__ == "__main__":main()

extra/gemm/qcom_openpilot_blocked_forward.py

@@ -1,78 +0,0 @@
-#!/usr/bin/env python3
-"""Experimental 64-term FP16 partial / FP32 total OpenPilot projection."""
-import argparse, pickle
-from dataclasses import replace
-
-from tinygrad import Device
-from tinygrad.engine.jit import create_graph_call
-from tinygrad.uop.ops import Ops
-from extra.gemm.qcom_ir3_matmul_patch import plain_name
-
-TARGET = "r_32_192_4_4_64_4"
-
-SOURCE = r"""#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-const sampler_t smp=CLK_NORMALIZED_COORDS_FALSE|CLK_ADDRESS_CLAMP|CLK_FILTER_NEAREST;
-inline float4 gelu(float4 v) {
-  return ((float4)(1)/(1+exp2((v+(float4)(0.044708251953125f)*v*v*v)*(float4)(-2.3021129851685216f))))*v;
-}
-__kernel void r_32_192_4_4_64_4(write_only image2d_t O, read_only image2d_t A,
-                                read_only image2d_t W, read_only image2d_t B) {
-  int n=get_global_id(0), m=get_global_id(1), abase=m*260;
-  float4 t0=(float4)(0),t1=(float4)(0),t2=(float4)(0),t3=(float4)(0);
-  for (int kb=0;kb<64;kb+=16) {
-    half4 r0=(half4)(0),r1=(half4)(0),r2=(half4)(0),r3=(half4)(0);
-    for (int k=kb;k<kb+16;k++) {
-      half4 a0=read_imageh(A,smp,(int2)(abase+k,0));
-      half4 a1=read_imageh(A,smp,(int2)(abase+k+65,0));
-      half4 a2=read_imageh(A,smp,(int2)(abase+k+130,0));
-      half4 a3=read_imageh(A,smp,(int2)(abase+k+195,0));
-      int x=k*4;
-      half4 w0=read_imageh(W,smp,(int2)(x,n));
-      half4 w1=read_imageh(W,smp,(int2)(x+1,n));
-      half4 w2=read_imageh(W,smp,(int2)(x+2,n));
-      half4 w3=read_imageh(W,smp,(int2)(x+3,n));
-      r0+=(half4)(a0.x)*w0; r0+=(half4)(a0.y)*w1; r0+=(half4)(a0.z)*w2; r0+=(half4)(a0.w)*w3;
-      r1+=(half4)(a1.x)*w0; r1+=(half4)(a1.y)*w1; r1+=(half4)(a1.z)*w2; r1+=(half4)(a1.w)*w3;
-      r2+=(half4)(a2.x)*w0; r2+=(half4)(a2.y)*w1; r2+=(half4)(a2.z)*w2; r2+=(half4)(a2.w)*w3;
-      r3+=(half4)(a3.x)*w0; r3+=(half4)(a3.y)*w1; r3+=(half4)(a3.z)*w2; r3+=(half4)(a3.w)*w3;
-    }
-    t0+=convert_float4(r0); t1+=convert_float4(r1); t2+=convert_float4(r2); t3+=convert_float4(r3);
-  }
-  float4 b=read_imagef(B,smp,(int2)(n,0));
-  write_imagef(O,(int2)(n,m),gelu(t0+b));
-  write_imagef(O,(int2)(n+192,m),gelu(t1+b));
-  write_imagef(O,(int2)(n+384,m),gelu(t2+b));
-  write_imagef(O,(int2)(n+576,m),gelu(t3+b));
-}"""
-
-
-def patch_model(model, block4:int=16) -> int:
-  if 64 % block4: raise ValueError("block4 must divide 64")
-  source = SOURCE.replace("kb<64;kb+=16", f"kb<64;kb+={block4}").replace("k<kb+16", f"k<kb+{block4}")
-  outer = model.captured.linear.src[0]
-  batch, patched = list(outer.src[0].src[0].src), 0
-  lib = Device["QCOM"].compiler.compile_cached(source)
-  for index, call in enumerate(batch):
-    if call.op is not Ops.CALL or call.src[0].op is not Ops.PROGRAM or plain_name(call.src[0].arg.name) != TARGET: continue
-    program = call.src[0]
-    program = program.replace(arg=replace(program.arg, global_size=(24, 1, 1), local_size=(8, 32, 1)),
-                              src=program.src[:2]+(program.src[2].replace(arg=source), program.src[3].replace(arg=lib)))
-    batch[index] = call.replace(src=(program, *call.src[1:]))
-    patched += 1
-  if patched:
-    model.captured._linear = model.captured.linear.substitute({outer:create_graph_call(batch)}, walk=True)
-    model.captured.__dict__.pop("linear", None)
-  return patched
-
-
-def main() -> None:
-  ap = argparse.ArgumentParser()
-  ap.add_argument("input"); ap.add_argument("output")
-  ap.add_argument("--block4", type=int, default=16)
-  args = ap.parse_args()
-  with open(args.input, "rb") as f: model = pickle.load(f)
-  print("patched", patch_model(model, args.block4))
-  with open(args.output, "wb") as f: pickle.dump(model, f)
-
-
-if __name__ == "__main__": main()

extra/gemm/qcom_openpilot_dedupe_head.py

@@ -1,79 +0,0 @@
-#!/usr/bin/env python3
-"""Remove byte-identical duplicate linear chains in the driving-vision head."""
-import argparse, hashlib, pickle
-
-from tinygrad.engine.jit import create_graph_call
-from tinygrad.uop.ops import Ops, UOp
-from extra.gemm.qcom_ir3_matmul_patch import plain_name
-
-TARGETS = {"r_128_16_4_32_4", "r_256_4_128_4", "r_128_16_4_16_4"}
-
-
-def dedupe_identical_calls(model, all_calls:bool=True) -> list[tuple[int, str]]:
-  """Alias calls with identical programs, inputs, and byte-identical constants."""
-  outer = model.captured.linear.src[0]
-  batch = outer.src[0].src[0].src
-  produced:dict[UOp, UOp] = {}
-  static_hash:dict[UOp, str] = {}
-  seen:dict[tuple, tuple[UOp, ...]] = {}
-  new_batch, removed = [], []
-
-  def representative(buf:UOp) -> UOp:
-    while buf in produced and produced[buf] is not buf: buf = produced[buf]
-    return buf
-
-  def content_hash(buf:UOp) -> str:
-    if buf not in static_hash:
-      static_hash[buf] = hashlib.sha256(memoryview(buf.buffer.numpy()).cast("B")).hexdigest()
-    return static_hash[buf]
-
-  for index, original in enumerate(batch):
-    call = original.replace(src=tuple(representative(x) if x in produced else x for x in original.src))
-    if call.op is not Ops.CALL or call.src[0].op is not Ops.PROGRAM or (not all_calls and plain_name(call.src[0].arg.name) not in TARGETS):
-      new_batch.append(call)
-      if call.op is Ops.CALL and call.src[0].op is Ops.PROGRAM:
-        for out_index in call.src[0].arg.outs: produced[original.src[out_index+1]] = call.src[out_index+1]
-      continue
-    program = call.src[0]
-    output_indices = set(program.arg.outs)
-    signature_args = []
-    for arg_index, (before, after) in enumerate(zip(original.src[1:], call.src[1:])):
-      if arg_index in output_indices: continue
-      if before.op is Ops.PARAM:
-        signature_args.append(("param", before.arg))
-      elif before in produced:
-        signature_args.append(("dynamic", representative(before)))
-      else:
-        signature_args.append((str(after.dtype), after.buffer.size, content_hash(after)))
-    signature = (plain_name(program.arg.name), program.src[3].arg, tuple(signature_args))
-    outputs = tuple(original.src[i+1] for i in program.arg.outs)
-    if signature in seen:
-      canonical_outputs = seen[signature]
-      for output, canonical in zip(outputs, canonical_outputs): produced[output] = representative(canonical)
-      removed.append((index, plain_name(program.arg.name)))
-    else:
-      new_batch.append(call)
-      canonical_outputs = tuple(call.src[i+1] for i in program.arg.outs)
-      seen[signature] = canonical_outputs
-      for output, canonical in zip(outputs, canonical_outputs): produced[output] = canonical
-
-  # Apply aliases to consumers which occur after the duplicate chains.
-  new_batch = [call.replace(src=tuple(representative(x) if x in produced else x for x in call.src)) for call in new_batch]
-  model.captured._linear = model.captured.linear.substitute({outer:create_graph_call(new_batch)}, walk=True)
-  model.captured.__dict__.pop("linear", None)
-  return removed
-
-
-def main() -> None:
-  parser = argparse.ArgumentParser()
-  parser.add_argument("input")
-  parser.add_argument("output")
-  parser.add_argument("--all", action="store_true", help="deduplicate every program family, not only the head linears")
-  args = parser.parse_args()
-  with open(args.input, "rb") as f: model = pickle.load(f)
-  removed = dedupe_identical_calls(model, args.all)
-  with open(args.output, "wb") as f: pickle.dump(model, f)
-  print(f"removed {len(removed)} duplicate head calls: {removed}")
-
-
-if __name__ == "__main__": main()

extra/gemm/qcom_openpilot_fast_gelu.py

@@ -1,54 +0,0 @@
-#!/usr/bin/env python3
-"""Replace selected QCOM GELU epilogues with a bounded polynomial approximation."""
-import argparse, pickle, re
-
-from tinygrad import Device
-from tinygrad.engine.jit import create_graph_call
-from tinygrad.uop.ops import Ops
-from extra.gemm.qcom_ir3_matmul_patch import plain_name
-
-
-def fast_gelu_source(source:str) -> tuple[str, int]:
-  variables=set(re.findall(r"float (alu\d+) =",source))
-  replaced=0
-  for var in variables:
-    old=f"((1/(1.0f+exp2((({var}+(0.044708251953125f*{var}*{var}*{var}))*-2.3021129851685216f))))*{var})"
-    # Degree-10 approximation of the model's exact tanh-GELU on |x| < 4.
-    # GELU(-x)=GELU(x)-x lets one polynomial cover both signs; outside this
-    # interval ReLU differs from the source expression by less than 1.3e-4.
-    coeffs=(1.95458887333,2.17220398188,-0.215882554761,-0.00733160096454,0.28997582181,-0.274775761974,
-            0.0167240224759,0.132422938329,-0.0634056438625,-0.022554589963,0.0179724326925)
-    t=f"(fabs({var})*0.5f-1.0f)"
-    poly=f"{coeffs[-1]:.10g}f"
-    for coefficient in reversed(coeffs[:-1]): poly=f"({coefficient:.10g}f+{t}*{poly})"
-    new=f"((fabs({var})>=4.0f)?max({var},0.0f):({poly}+min({var},0.0f)))"
-    if old in source:
-      source=source.replace(old,new)
-      replaced+=1
-  return source,replaced
-
-
-def patch_model(model,names:set[str]) -> int:
-  outer=model.captured.linear.src[0]
-  batch=list(outer.src[0].src[0].src)
-  compiler,cache,patched=Device["QCOM"].compiler,{},0
-  for index,call in enumerate(batch):
-    if call.op is not Ops.CALL or call.src[0].op is not Ops.PROGRAM or plain_name(call.src[0].arg.name) not in names: continue
-    program=call.src[0]
-    source,count=fast_gelu_source(program.src[2].arg)
-    if not count: continue
-    if source not in cache: cache[source]=compiler.compile(source)
-    program=program.replace(src=program.src[:2]+(program.src[2].replace(arg=source),program.src[3].replace(arg=cache[source])))
-    batch[index]=call.replace(src=(program,*call.src[1:]))
-    patched+=1
-  if patched:
-    model.captured._linear=model.captured.linear.substitute({outer:create_graph_call(batch)},walk=True)
-    model.captured.__dict__.pop("linear",None)
-  return patched
-
-
-if __name__ == "__main__":
-  ap=argparse.ArgumentParser();ap.add_argument("input");ap.add_argument("output");ap.add_argument("--names",required=True);args=ap.parse_args()
-  with open(args.input,"rb") as f:model=pickle.load(f)
-  print("patched",patch_model(model,set(args.names.split(","))))
-  with open(args.output,"wb") as f:pickle.dump(model,f)

extra/gemm/qcom_openpilot_firstconv_tile8.py

@@ -1,111 +0,0 @@
-#!/usr/bin/env python3
-"""Replace driving_vision's first convolution with a wider spatial tile."""
-import argparse, os, pickle
-from dataclasses import replace
-
-from tinygrad import Device, dtypes
-from tinygrad.engine.jit import create_graph_call
-from tinygrad.uop.ops import Ops, UOp
-from extra.gemm.qcom_ir3_matmul_patch import plain_name
-
-TARGET = "r_64_32_16_4_4_6_3_3_4"
-
-
-def make_source(spatial:int, output_blocks:int, split:bool=False) -> str:
-  fp32 = bool(int(os.getenv("FP32_TILE", "0")))
-  vec, read, scalar = ("float4", "read_imagef", "float4") if fp32 else ("half4", "read_imageh", "half4")
-  local_x = 16//output_blocks
-  local_y = 128//local_x
-  lines = ["#pragma OPENCL EXTENSION cl_khr_fp16 : enable", """
-const sampler_t smp=CLK_NORMALIZED_COORDS_FALSE|CLK_ADDRESS_CLAMP|CLK_FILTER_NEAREST;
-inline float4 gelu(float4 v) {
-  return ((float4)(1)/(1+exp2((v+(float4)(0.044708251953125f)*v*v*v)*(float4)(-2.3021129851685216f))))*v;
-}
-""", f"__attribute__((reqd_work_group_size({local_x},{local_y},1)))", """
-__kernel void firstconv_tile8(write_only image2d_t O,read_only image2d_t A,
-                              read_only image2d_t W,read_only image2d_t B) {
-  int ob=get_global_id(0), xb=get_global_id(1), y=get_global_id(2);
-"""]
-  lines += [f"  {vec} z{s}_{n}=({vec})(0);" for s in range(spatial) for n in range(output_blocks)]
-  lines.append("  for(int ic=0;ic<6;ic++) for(int ky=0;ky<3;ky++) for(int kx=0;kx<3;kx++) {")
-  lines.append(f"    int ax=xb*{spatial*12}+kx*6+ic, ay=y*2+ky-1;")
-  lines += [f"    {vec} a{s}={read}(A,smp,(int2)(ax+{12*s-6},ay));" for s in range(spatial)]
-  for n in range(output_blocks):
-    lines.append(f"    int wp{n}=ic*12+kx*4+ky*72+(ob*{output_blocks}+{n})*216;")
-    lines += [f"    {vec} w{n}{k}={read}(W,smp,(int2)(wp{n}+{k},0));" for k in (0, 1, 2, 3)]
-  for s in range(spatial):
-    for n in range(output_blocks):
-      lines.append(f"    z{s}_{n}+=({scalar})(a{s}.x)*w{n}0+({scalar})(a{s}.y)*w{n}1+"
-                   f"({scalar})(a{s}.z)*w{n}2+({scalar})(a{s}.w)*w{n}3;")
-  lines.append("  }")
-  if not split:
-    for n in range(output_blocks):
-      lines.append(f"  float4 b{n}=read_imagef(B,smp,(int2)(ob*{output_blocks}+{n},0));")
-  for s in range(spatial):
-    for n in range(output_blocks):
-      raw = f"z{s}_{n}" if fp32 else f"convert_float4(z{s}_{n})"
-      value = raw if split else f"gelu({raw}+b{n})"
-      lines.append(f"  write_imagef(O,(int2)(ob*{output_blocks}+{n}+xb*{spatial*16}+{s*16},y),{value});")
-  lines.append("}")
-  return "\n".join(lines)
-
-
-EPILOGUE = r"""#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-const sampler_t smp=CLK_NORMALIZED_COORDS_FALSE|CLK_ADDRESS_CLAMP|CLK_FILTER_NEAREST;
-inline float4 gelu(float4 v) {
-  return ((float4)(1)/(1+exp2((v+(float4)(0.044708251953125f)*v*v*v)*(float4)(-2.3021129851685216f))))*v;
-}
-__attribute__((reqd_work_group_size(128,1,1)))
-__kernel void firstconv_gelu(write_only image2d_t O,read_only image2d_t B,read_only image2d_t T) {
-  int x=get_global_id(0),y=get_global_id(1);
-  write_imagef(O,(int2)(x,y),gelu(convert_float4(read_imageh(T,smp,(int2)(x,y)))+read_imagef(B,smp,(int2)(x&15,0))));
-}"""
-
-
-def patch_model(model, spatial:int, output_blocks:int, split:bool=False) -> int:
-  if spatial*output_blocks not in (4, 8) or 16%output_blocks: raise ValueError("tile must contain four or eight vectors")
-  outer, source = model.captured.linear.src[0], make_source(spatial, output_blocks, split)
-  batch, lib, patched = list(outer.src[0].src[0].src), Device["QCOM"].compiler.compile(source), 0
-  replacements:dict[int, tuple[UOp, ...]] = {}
-  epi_lib = Device["QCOM"].compiler.compile(EPILOGUE) if split else None
-  for index, call in enumerate(batch):
-    if call.op is not Ops.CALL or call.src[0].op is not Ops.PROGRAM or plain_name(call.src[0].arg.name) != TARGET: continue
-    old = call.src[0]
-    local_x, local_y = 16//output_blocks, 128//(16//output_blocks)
-    global_y = (128//spatial)//local_y
-    info = replace(old.arg, name=f"firstconv_tile{spatial}x{output_blocks*4}", global_size=(1, global_y, 64),
-                   local_size=(local_x, local_y, 1))
-    program = old.replace(arg=info, src=old.src[:2]+(old.src[2].replace(arg=source), old.src[3].replace(arg=lib)))
-    if split:
-      temporary = UOp.new_buffer("QCOM", call.src[1].buffer.size, dtypes.half, num=-3_000_000)
-      temporary.buffer.ensure_allocated()
-      compute = call.replace(src=(program, temporary, *call.src[2:]))
-      epi_aux = ((((0, dtypes.half, (64, 2048, 4)),), ((1, dtypes.half, (1, 16, 4)),),
-                  ((2, dtypes.half, (64, 2048, 4)),)),)
-      epi_info = replace(old.arg, name="firstconv_gelu", global_size=(16, 64, 1), local_size=(128, 1, 1),
-                         globals=(0, 1, 2), outs=(0,), ins=(1, 2), aux=epi_aux)
-      epi_program = old.replace(arg=epi_info, src=old.src[:2]+(old.src[2].replace(arg=EPILOGUE), old.src[3].replace(arg=epi_lib)))
-      replacements[index] = (compute, epi_program.call(call.src[1], call.src[4], temporary))
-    else: replacements[index] = (call.replace(src=(program, *call.src[1:])),)
-    patched += 1
-  if patched:
-    new_batch = [new for index, call in enumerate(batch) for new in replacements.get(index, (call,))]
-    model.captured._linear = model.captured.linear.substitute({outer:create_graph_call(new_batch)}, walk=True)
-    model.captured.__dict__.pop("linear", None)
-  return patched
-
-
-def main() -> None:
-  parser=argparse.ArgumentParser()
-  parser.add_argument("input")
-  parser.add_argument("output")
-  parser.add_argument("--spatial", type=int, default=8)
-  parser.add_argument("--output-blocks", type=int, default=1)
-  parser.add_argument("--split", action="store_true")
-  args=parser.parse_args()
-  with open(args.input, "rb") as f: model=pickle.load(f)
-  print("patched", patch_model(model, args.spatial, args.output_blocks, args.split))
-  with open(args.output, "wb") as f: pickle.dump(model, f)
-
-
-if __name__ == "__main__": main()

extra/gemm/qcom_openpilot_fix_half8_weights.py

@@ -1,103 +0,0 @@
-#!/usr/bin/env python3
-"""Repair the lane order in packed half8 OpenPilot projection weights."""
-import argparse
-import pickle
-import struct
-
-import numpy as np
-
-from tinygrad import Device
-from tinygrad.device import Buffer
-from tinygrad.engine.jit import create_graph_call
-from tinygrad.uop.ops import Ops
-from extra.gemm.ir3asm import BR, COV_S32S16, ISAM_F16, MAD_F16, SHRG_H
-from extra.gemm.qcom_ir3_matmul_patch import plain_name
-
-TARGET = "r_32_192_4_4_64_4"
-
-
-def fix_repeat_mads(lib: bytes) -> bytes:
-  image_offset = struct.unpack_from("<I", lib, 0xC0)[0]
-  image_size = struct.unpack_from("<I", lib, 0x100)[0]
-  instructions = [lib[x:x+8] for x in range(image_offset, image_offset+image_size, 8)]
-  if instructions[71] != BR(25-71):
-    raise RuntimeError("unexpected half8 loop layout")
-  # Keep sampler outputs, accumulators, and activations in disjoint banks.
-  instructions[35] = ISAM_F16("hr23.x", "r5.w", 0)
-  instructions[36] = ISAM_F16("hr22.x", "r6.y", 0)
-  instructions[37] = ISAM_F16("hr21.x", "r6.w", 0)
-  instructions[39] = ISAM_F16("hr20.x", "r7.y", 0)
-  unpack = []
-  for destination, source in ((12, "r3.x"), (14, "r2.x"), (16, "r1.x"), (18, "r0.x")):
-    unpack.append(COV_S32S16(f"hr{destination}.x", source, rpt=3, r=True, sy=not unpack))
-    unpack.append(SHRG_H(f"hr{destination+1}.x", source, rpt=3, r=True))
-  mads = []
-  rows = (("hr10.x", "hr11.x", "hr23"), ("hr8.x", "hr9.x", "hr22"),
-          ("hr6.x", "hr7.x", "hr21"), ("hr4.x", "hr5.x", "hr20"))
-  for component, (weight0, weight1) in zip("xyzw", (("hr12.x", "hr13.x"), ("hr14.x", "hr15.x"),
-                                                         ("hr16.x", "hr17.x"), ("hr18.x", "hr19.x"))):
-    for accumulator0, accumulator1, activation in rows:
-      mads.append(MAD_F16(accumulator0, f"{activation}.{component}", weight0, accumulator0,
-                          rpt=3, r=True))
-      mads.append(MAD_F16(accumulator1, f"{activation}.{component}", weight1, accumulator1, rpt=3, r=True))
-  output = instructions[:48] + unpack + mads + instructions[70:]
-  output[89] = BR(25-89)
-  output = output[:len(instructions)]
-  patched = bytearray(lib[:image_offset] + b"".join(output) + lib[image_offset+image_size:])
-  register_offset = struct.unpack_from("<I", patched, 0x34)[0]
-  old_hregs = struct.unpack_from("<I", patched, register_offset+0x18)[0]
-  struct.pack_into("<I", patched, register_offset+0x18, (old_hregs & 0x80000000) | 24)
-  return bytes(patched)
-
-
-def patch_model(model, fix_weights: bool = True, recompile: bool = False, fix_mads: bool = False) -> int:
-  outer = model.captured.linear.src[0]
-  batch = list(outer.src[0].src[0].src)
-  seen, patched = set(), 0
-  for index, call in enumerate(batch):
-    if call.op is not Ops.CALL or call.src[0].op is not Ops.PROGRAM or plain_name(call.src[0].arg.name) != TARGET:
-      continue
-    if recompile or fix_mads:
-      program = call.src[0]
-      lib = Device["QCOM"].compiler.compile_cached(program.src[2].arg) if recompile else fix_repeat_mads(program.src[3].arg)
-      program = program.replace(src=program.src[:3] + (program.src[3].replace(arg=lib),))
-      batch[index] = call.replace(src=(program, *call.src[1:]))
-    weight = call.src[3].buffer
-    if not fix_weights or id(weight) in seen or weight.dtype.itemsize != 4:
-      patched += 1
-      continue
-    seen.add(id(weight))
-    # The old pack transposed two adjacent float4 output channels before
-    # bitcasting to uint4, producing a0,b0,a1,b1,... in each half8 pixel.
-    # The kernel consumes half8.lo/hi as complete float4 channels.
-    packed = weight.numpy().view(np.float16).reshape(-1, 8)
-    corrected = np.ascontiguousarray(packed[:, (0, 2, 4, 6, 1, 3, 5, 7)])
-    raw = memoryview(corrected).cast("B")
-    if hasattr(weight, "copyin"):
-      weight.copyin(raw)
-    else:
-      weight.copy_from(Buffer("PYTHON", weight.size, weight.dtype, opaque=raw))
-    patched += 1
-  if recompile or fix_mads:
-    model.captured._linear = model.captured.linear.substitute({outer: create_graph_call(batch)}, walk=True)
-    model.captured.__dict__.pop("linear", None)
-  return patched
-
-
-def main() -> None:
-  parser = argparse.ArgumentParser()
-  parser.add_argument("input")
-  parser.add_argument("output")
-  parser.add_argument("--skip-weights", action="store_true")
-  parser.add_argument("--recompile", action="store_true")
-  parser.add_argument("--fix-mads", action="store_true")
-  args = parser.parse_args()
-  with open(args.input, "rb") as f:
-    model = pickle.load(f)
-  print("patched", patch_model(model, not args.skip_weights, args.recompile, args.fix_mads))
-  with open(args.output, "wb") as f:
-    pickle.dump(model, f)
-
-
-if __name__ == "__main__":
-  main()

extra/gemm/qcom_openpilot_forward_raw.py

@@ -1,133 +0,0 @@
-#!/usr/bin/env python3
-"""Raw 8x8 FP16-accumulate projection for the padded OpenPilot vision layout."""
-import argparse, os, pickle, struct
-from dataclasses import replace
-
-from tinygrad import Device
-from tinygrad.engine.jit import create_graph_call
-from tinygrad.uop.ops import Ops
-from extra.gemm.ir3asm import (ADD_S, ADD_S_REG, AND_B, BR, CMPS_S_EQ, COV_F16F32, END, ISAM_F16, MAD_F16, MOV_F32,
-                               MOV_H_IMM, MOV_S32, NOP, NOP_SS, SHL_B, SHR_B, STIB_F32, assemble, inject)
-from extra.gemm import qcom_8x4_gemm as q8
-from extra.gemm.qcom_8x4_gemm import prologue_8x4
-from extra.gemm.qcom_ir3_matmul_patch import plain_name
-from extra.gemm.qcom_openpilot_forward_tile8 import SOURCE
-
-TARGET = "r_32_192_4_4_64_4"
-
-
-def build_raw_shader(dev) -> tuple[bytes, int, int]:
-  instrs = prologue_8x4(dev, 128)
-  # The donor produces row=gid1*32+(lid>>5)*8 and col=gid0*32+(lid&31).
-  # Widen col to a two-col4 tile: gid0*64+tid, with the second column at +32.
-  instrs += [MOV_F32("r12.x", "r51.w"), NOP(rpt=2), SHL_B("r12.x", "r12.x", 5), NOP(rpt=2),
-             ADD_S_REG("r7.y", "r7.y", "r12.x"), NOP(rpt=2)]
-
-  # Precompute the eight padded-A row bases. A is a 1D image laid out as
-  # (row&31)*260 + (row>>5)*65 + k4.
-  instrs += [SHR_B("r12.y", "r7.x", 5), AND_B("r12.z", "r7.x", 31), NOP(rpt=2),
-             SHL_B("r12.w", "r12.y", 6), SHL_B("r13.x", "r12.z", 8), SHL_B("r13.y", "r12.z", 2),
-             ADD_S_REG("r12.w", "r12.w", "r12.y"), ADD_S_REG("r13.x", "r13.x", "r13.y"), NOP(rpt=2),
-             ADD_S_REG("r13.x", "r13.x", "r12.w"), MOV_S32("r13.y", 260), NOP(rpt=2)]
-  row_bases = ("r13.x", "r13.z", "r13.w", "r14.x", "r14.y", "r14.z", "r14.w", "r15.x")
-  for index, dst in enumerate(row_bases[1:], 1):
-    instrs += [ADD_S_REG(dst, row_bases[index-1], "r13.y"), NOP(rpt=2)]
-
-  acc0 = 12 * 4
-  for base in range(acc0, acc0+16*4, 4): instrs.append(MOV_H_IMM(base, 0, rpt=3))
-  instrs += [MOV_S32("r6.z", 0), MOV_S32("r6.y", 3, sy=True)]
-  loop_start = len(instrs)
-
-  b_pairs = tuple((f"r{16+i//2}.{'xz'[i&1]}", f"r{16+i//2}.{'yw'[i&1]}") for i in range(8))
-  for component in range(4):
-    for col in range(2):
-      xreg, yreg = b_pairs[component*2+col]
-      instrs.append(MOV_F32(xreg, "r6.y") if component == 3 else ADD_S(xreg, "r6.y", component-3))
-      instrs.append(MOV_F32(yreg, "r7.y") if col == 0 else ADD_S(yreg, "r7.y", 32))
-  instrs.append(NOP(rpt=3))
-  for index, (xreg, _) in enumerate(b_pairs): instrs.append(ISAM_F16(index*4, xreg, 1))
-
-  a_pairs = (("r20.x", "r20.y"), ("r20.z", "r20.w"), ("r21.x", "r21.y"), ("r21.z", "r21.w"))
-  def load_a(first_row: int) -> None:
-    nonlocal instrs
-    for slot, ((xreg, yreg), base) in enumerate(zip(a_pairs, row_bases[first_row:first_row+4])):
-      instrs += [ADD_S_REG(xreg, base, "r6.z"), MOV_S32(yreg, 0)]
-    instrs.append(NOP(rpt=3))
-    for slot, (xreg, _) in enumerate(a_pairs): instrs.append(ISAM_F16((8+slot)*4, xreg, 0))
-
-  def mads(first_row: int) -> None:
-    first = True
-    for slot, row in enumerate(range(first_row, first_row+4)):
-      for component in range(4):
-        for col in range(2):
-          acc = acc0+(row*2+col)*4
-          instrs.append(MAD_F16(acc, (8+slot)*4+component, (component*2+col)*4, acc, rpt=3, r=True, sy=first))
-          first = False
-
-  load_a(0)
-  mads(0)
-  instrs.append(NOP_SS())
-  load_a(4)
-  mads(4)
-  instrs += [ADD_S("r0.x", "r6.z", 1), ADD_S("r6.y", "r6.y", 4), CMPS_S_EQ("r6.z", 63, nop=1),
-             MOV_F32("r6.z", "r0.x"), NOP(rpt=3)]
-  loop_end = len(instrs)
-  instrs.append(BR(loop_start-loop_end))
-
-  if os.getenv("RAW_NO_STORE"):
-    instrs.append(END())
-    return assemble(instrs), 24, 28
-
-  # Typed image stores. p=row>>5 is constant within a tile; output x is
-  # col+p*192 and output y is row&31.
-  instrs += [SHL_B("r12.w", "r12.y", 7), SHL_B("r13.x", "r12.y", 6),
-             ADD_S_REG("r12.w", "r12.w", "r13.x"), ADD_S_REG("r12.w", "r12.w", "r7.y"), NOP(rpt=2)]
-  for row in range(8):
-    for col in range(2):
-      instrs.append(MOV_F32("r22.x", "r12.w") if col == 0 else ADD_S("r22.x", "r12.w", 32))
-      instrs.append(MOV_F32("r22.y", "r12.z") if row == 0 else ADD_S("r22.y", "r12.z", row))
-      instrs += [COV_F16F32("r23.x", acc0+(row*2+col)*4, sy=True, rpt=3, r=True), NOP(rpt=5),
-                 STIB_F32("r23.x", "r22.x"), NOP(rpt=8)]
-  instrs.append(END())
-  return assemble(instrs), 24, 28
-
-
-def raw_lib(dev) -> bytes:
-  lib = dev.compiler.compile_cached(SOURCE)
-  image_off, image_size = struct.unpack_from("<I", lib, 0xc0)[0], struct.unpack_from("<I", lib, 0x100)[0]
-  reg_off = struct.unpack_from("<I", lib, 0x34)[0]
-  if os.getenv("RAW_GENERAL"):
-    threads = int(os.getenv("RAW_THREADS", "128"))
-    q8.K, q8.K4 = 256, 64
-    shader, hregs, fregs, _ = q8.build_8x8_split_a_unroll_shader(
-      dev, threads, k_unroll=8, b_coord_delay=0, fast_coords=True,
-      prefetch_next_b=True, no_store=True)
-  else:
-    shader, fregs, hregs = build_raw_shader(dev)
-  return inject(lib, image_off, image_size, reg_off, shader, fregs, hregs)
-
-
-def patch_model(model) -> int:
-  outer = model.captured.linear.src[0]
-  batch, patched, lib = list(outer.src[0].src[0].src), 0, raw_lib(Device["QCOM"])
-  threads = int(os.getenv("RAW_THREADS", "128")) if os.getenv("RAW_GENERAL") else 128
-  for index, call in enumerate(batch):
-    if call.op is not Ops.CALL or call.src[0].op is not Ops.PROGRAM or plain_name(call.src[0].arg.name) != TARGET: continue
-    program = call.src[0].replace(arg=replace(call.src[0].arg, global_size=(3, 512//threads, 1), local_size=(threads, 1, 1)),
-                                  src=call.src[0].src[:3]+(call.src[0].src[3].replace(arg=lib),))
-    batch[index] = call.replace(src=(program, *call.src[1:]))
-    patched += 1
-  if patched:
-    model.captured._linear = model.captured.linear.substitute({outer:create_graph_call(batch)}, walk=True)
-    model.captured.__dict__.pop("linear", None)
-  return patched
-
-
-if __name__ == "__main__":
-  parser = argparse.ArgumentParser()
-  parser.add_argument("input")
-  parser.add_argument("output")
-  args = parser.parse_args()
-  with open(args.input, "rb") as f: model = pickle.load(f)
-  print("patched", patch_model(model))
-  with open(args.output, "wb") as f: pickle.dump(model, f)

extra/gemm/qcom_openpilot_forward_raw_probe.py

@@ -1,21 +0,0 @@
-#!/usr/bin/env python3
-"""Compile/disassemble the raw-accumulator variant of the vision forward tile."""
-import struct
-
-from tinygrad import Device
-from extra.gemm.ir3asm import disasm
-from extra.gemm.qcom_openpilot_forward_tile8 import SOURCE
-
-RAW_TAIL = r"""  int r=row0;
-#define STORE(v) { int m=r&31,p=r>>5; write_imageh(O,(int2)(n0+p*192,m),v.lo); write_imageh(O,(int2)(n1+p*192,m),v.hi); r++; }
-  STORE(c0); STORE(c1); STORE(c2); STORE(c3); STORE(c4); STORE(c5); STORE(c6); STORE(c7);
-}
-"""
-RAW_SOURCE = SOURCE[:SOURCE.index("  float4 b0=")] + RAW_TAIL
-
-if __name__ == "__main__":
-  lib = Device["QCOM"].compiler.compile_cached(RAW_SOURCE)
-  image_off, image_size = struct.unpack_from("<I", lib, 0xc0)[0], struct.unpack_from("<I", lib, 0x100)[0]
-  lines = [x for x in disasm(lib[image_off:image_off+image_size]).splitlines() if not x.rstrip().endswith(":")]
-  print("COUNT", len(lines))
-  for index, line in enumerate(lines): print(f"{index}: {line}")

extra/gemm/qcom_openpilot_forward_tile8.py

@@ -1,333 +0,0 @@
-#!/usr/bin/env python3
-"""High-intensity 8x8 FP16-accumulate tile for OpenPilot vision projections."""
-import argparse, os, pickle, struct
-from dataclasses import replace
-
-from tinygrad import Device
-from tinygrad.engine.jit import create_graph_call
-from tinygrad.uop.ops import Ops
-from extra.gemm.ir3asm import ADD_S, ADD_S_REG, BR, ISAM_F16, MAD_F16, MOV_F32, MOV_H_IMM, NOP, inject
-from extra.gemm.qcom_ir3_matmul_patch import plain_name
-
-TARGET = "r_32_192_4_4_64_4"
-SOURCE = r"""#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-const sampler_t smp=CLK_NORMALIZED_COORDS_FALSE|CLK_ADDRESS_CLAMP|CLK_FILTER_NEAREST;
-inline float4 gelu(float4 v) { return ((float4)(1)/(1+exp2((v+(float4)(0.044708251953125f)*v*v*v)*(float4)(-2.3021129851685216f))))*v; }
-__attribute__((reqd_work_group_size(128,1,1)))
-__kernel void r_32_192_4_4_64_4(write_only image2d_t O,read_only image2d_t A,read_only image2d_t W,read_only image2d_t B) {
-  int lid=get_local_id(0),tm=lid>>5,tid=lid&31,row0=get_group_id(1)*32+tm*8,n0=get_group_id(0)*64+tid*2,n1=n0+1;
-  half8 c0=(half8)(0),c1=(half8)(0),c2=(half8)(0),c3=(half8)(0);
-  half8 c4=(half8)(0),c5=(half8)(0),c6=(half8)(0),c7=(half8)(0);
-  for(int k=0;k<64;k++) {
-    int x=k*4;
-    half8 w0=(half8)(read_imageh(W,smp,(int2)(x,n0)),read_imageh(W,smp,(int2)(x,n1)));
-    half8 w1=(half8)(read_imageh(W,smp,(int2)(x+1,n0)),read_imageh(W,smp,(int2)(x+1,n1)));
-    half8 w2=(half8)(read_imageh(W,smp,(int2)(x+2,n0)),read_imageh(W,smp,(int2)(x+2,n1)));
-    half8 w3=(half8)(read_imageh(W,smp,(int2)(x+3,n0)),read_imageh(W,smp,(int2)(x+3,n1)));
-    int r=row0,base=(r&31)*260+(r>>5)*65;
-    half4 a0=read_imageh(A,smp,(int2)(base+k,0)); r++; base=(r&31)*260+(r>>5)*65;
-    half4 a1=read_imageh(A,smp,(int2)(base+k,0)); r++; base=(r&31)*260+(r>>5)*65;
-    half4 a2=read_imageh(A,smp,(int2)(base+k,0)); r++; base=(r&31)*260+(r>>5)*65;
-    half4 a3=read_imageh(A,smp,(int2)(base+k,0)); r++; base=(r&31)*260+(r>>5)*65;
-    half4 a4=read_imageh(A,smp,(int2)(base+k,0)); r++; base=(r&31)*260+(r>>5)*65;
-    half4 a5=read_imageh(A,smp,(int2)(base+k,0)); r++; base=(r&31)*260+(r>>5)*65;
-    half4 a6=read_imageh(A,smp,(int2)(base+k,0)); r++; base=(r&31)*260+(r>>5)*65;
-    half4 a7=read_imageh(A,smp,(int2)(base+k,0));
-    c0+=(half8)(a0.x)*w0+(half8)(a0.y)*w1+(half8)(a0.z)*w2+(half8)(a0.w)*w3;
-    c1+=(half8)(a1.x)*w0+(half8)(a1.y)*w1+(half8)(a1.z)*w2+(half8)(a1.w)*w3;
-    c2+=(half8)(a2.x)*w0+(half8)(a2.y)*w1+(half8)(a2.z)*w2+(half8)(a2.w)*w3;
-    c3+=(half8)(a3.x)*w0+(half8)(a3.y)*w1+(half8)(a3.z)*w2+(half8)(a3.w)*w3;
-    c4+=(half8)(a4.x)*w0+(half8)(a4.y)*w1+(half8)(a4.z)*w2+(half8)(a4.w)*w3;
-    c5+=(half8)(a5.x)*w0+(half8)(a5.y)*w1+(half8)(a5.z)*w2+(half8)(a5.w)*w3;
-    c6+=(half8)(a6.x)*w0+(half8)(a6.y)*w1+(half8)(a6.z)*w2+(half8)(a6.w)*w3;
-    c7+=(half8)(a7.x)*w0+(half8)(a7.y)*w1+(half8)(a7.z)*w2+(half8)(a7.w)*w3;
-  }
-  float4 b0=read_imagef(B,smp,(int2)(n0,0)),b1=read_imagef(B,smp,(int2)(n1,0));
-  int r=row0;
-#define STORE(v) { int m=r&31,p=r>>5; write_imagef(O,(int2)(n0+p*192,m),gelu(convert_float4(v.lo)+b0)); \
- write_imagef(O,(int2)(n1+p*192,m),gelu(convert_float4(v.hi)+b1)); r++; }
-  STORE(c0); STORE(c1); STORE(c2); STORE(c3); STORE(c4); STORE(c5); STORE(c6); STORE(c7);
-}"""
-
-SOURCE_4X2 = r"""#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-const sampler_t smp=CLK_NORMALIZED_COORDS_FALSE|CLK_ADDRESS_CLAMP|CLK_FILTER_NEAREST;
-inline float4 gelu(float4 v) { return ((float4)(1)/(1+exp2((v+(float4)(0.044708251953125f)*v*v*v)*(float4)(-2.3021129851685216f))))*v; }
-__attribute__((reqd_work_group_size(128,1,1)))
-__kernel void r_32_192_4_4_64_4(write_only image2d_t O,read_only image2d_t A,read_only image2d_t W,read_only image2d_t B) {
-  int lid=get_local_id(0),tm=lid>>5,tid=lid&31,row0=get_group_id(1)*16+tm*4,n0=get_group_id(0)*64+tid*2,n1=n0+1;
-  half8 c0=(half8)(0),c1=(half8)(0),c2=(half8)(0),c3=(half8)(0);
-  for(int k=0;k<64;k++) {
-    int x=k*4;
-    half8 w0=(half8)(read_imageh(W,smp,(int2)(x,n0)),read_imageh(W,smp,(int2)(x,n1)));
-    half8 w1=(half8)(read_imageh(W,smp,(int2)(x+1,n0)),read_imageh(W,smp,(int2)(x+1,n1)));
-    half8 w2=(half8)(read_imageh(W,smp,(int2)(x+2,n0)),read_imageh(W,smp,(int2)(x+2,n1)));
-    half8 w3=(half8)(read_imageh(W,smp,(int2)(x+3,n0)),read_imageh(W,smp,(int2)(x+3,n1)));
-    int r=row0,base=(r&31)*260+(r>>5)*65;
-    half4 a0=read_imageh(A,smp,(int2)(base+k,0)); r++; base=(r&31)*260+(r>>5)*65;
-    half4 a1=read_imageh(A,smp,(int2)(base+k,0)); r++; base=(r&31)*260+(r>>5)*65;
-    half4 a2=read_imageh(A,smp,(int2)(base+k,0)); r++; base=(r&31)*260+(r>>5)*65;
-    half4 a3=read_imageh(A,smp,(int2)(base+k,0));
-    c0+=(half8)(a0.x)*w0+(half8)(a0.y)*w1+(half8)(a0.z)*w2+(half8)(a0.w)*w3;
-    c1+=(half8)(a1.x)*w0+(half8)(a1.y)*w1+(half8)(a1.z)*w2+(half8)(a1.w)*w3;
-    c2+=(half8)(a2.x)*w0+(half8)(a2.y)*w1+(half8)(a2.z)*w2+(half8)(a2.w)*w3;
-    c3+=(half8)(a3.x)*w0+(half8)(a3.y)*w1+(half8)(a3.z)*w2+(half8)(a3.w)*w3;
-  }
-  float4 b0=read_imagef(B,smp,(int2)(n0,0)),b1=read_imagef(B,smp,(int2)(n1,0));
-  int r=row0;
-#define STORE(v) { int m=r&31,p=r>>5; write_imagef(O,(int2)(n0+p*192,m),gelu(convert_float4(v.lo)+b0)); \
- write_imagef(O,(int2)(n1+p*192,m),gelu(convert_float4(v.hi)+b1)); r++; }
-  STORE(c0); STORE(c1); STORE(c2); STORE(c3);
-}"""
-
-
-def pack_tile(lib:bytes) -> bytes:
-  image_off,image_size=struct.unpack_from("<I",lib,0xc0)[0],struct.unpack_from("<I",lib,0x100)[0]
-  reg_off=struct.unpack_from("<I",lib,0x34)[0]
-  ins=[lib[i:i+8] for i in range(image_off,image_off+image_size,8)]
-  if len(ins)<500: raise RuntimeError(f"unexpected tile8 shader length {len(ins)}")
-  # Pair each pair of output-channel texture rows in consecutive half registers.
-  for index,dst,coord in ((171,"hr3.x","r3.x"),(174,"hr4.x","r4.z"),
-                          (178,"hr5.x","r3.z"),(181,"hr6.x","r5.x"),
-                          (185,"hr7.x","r4.x"),(188,"hr8.x","r5.z"),
-                          (191,"hr32.x","r6.x"),(194,"hr33.x","r6.z")):
-    ins[index]=ISAM_F16(dst,coord,tex=1,sy=index==194)
-  out=ins[:198]
-  for activation,acc in (("hr15",30),("hr14",28),("hr13",26),("hr12",24),
-                         ("hr11",22),("hr10",20),("hr2",18),("hr0",16)):
-    for component,(weight_lo,weight_hi) in zip("xyzw",((3,4),(5,6),(7,8),(32,33))):
-      for offset,weight in enumerate((weight_lo,weight_hi)):
-        out.append(MAD_F16(f"hr{acc+offset}.x",f"{activation}.{component}",f"hr{weight}.x",f"hr{acc+offset}.x",
-                           rpt=3,r=True,sy=len(out)==198))
-  out.append(ins[366])
-  out.append(BR(143-len(out)))
-  out+=ins[368:]
-  if len(out)>len(ins): raise RuntimeError(f"packed shader grew from {len(ins)} to {len(out)}")
-  out += [NOP()]*(len(ins)-len(out))
-  fregs,hregs=struct.unpack_from("<II",lib,reg_off+0x14)
-  hregs=(hregs&0x80000000)|max(hregs&0x7fffffff,34)
-  return inject(lib,image_off,image_size,reg_off,b"".join(out),fregs,hregs)
-
-
-def pack_4x2_tile(lib: bytes, wg256: bool = False) -> bytes:
-  """Collapse the compiler's scalarized 4-row x 2-output FP16 dot products."""
-  image_off, image_size = struct.unpack_from("<I", lib, 0xc0)[0], struct.unpack_from("<I", lib, 0x100)[0]
-  reg_off = struct.unpack_from("<I", lib, 0x34)[0]
-  ins = [lib[i:i+8] for i in range(image_off, image_off+image_size, 8)]
-  if len(ins) < 450: raise RuntimeError(f"unexpected 4x2 shader length {len(ins)}")
-
-  # The donor loads four A vectors into hr11,hr10,hr5,hr4 and four pairs of
-  # W vectors into the low/high output banks below. Accumulators are already
-  # initialized in hr12..hr19; update them directly instead of materializing
-  # scalar partial sums and adding those afterwards.
-  loop_body, control, epilogue, loop_target = ((104, 228, 235, 65) if wg256 else (108, 232, 239, 69))
-  out = list(ins[:loop_body])
-  weight_pairs = ((6, 0), (9, 3), (8, 2), (7, 1))
-  first = True
-  rows = ((11, (18, 19)), (10, (16, 17)), (5, (14, 15)), (4, (12, 13)))
-  # Component-major order leaves eight instructions between updates of the
-  # same accumulator, hiding the dependent half-MAD latency.
-  for component, (weight_lo, weight_hi) in zip("xyzw", weight_pairs):
-    for column, weight in enumerate((weight_lo, weight_hi)):
-      for activation, accs in rows:
-        acc = accs[column]
-        out.append(MAD_F16(f"hr{acc}.x", f"hr{activation}.{component}", f"hr{weight}.x", f"hr{acc}.x",
-                           rpt=3, r=True, sy=first))
-        first = False
-
-  # Retain the compiler's K/coordinate updates and predicate setup, then
-  # relocate the loop backedge to the unchanged load block at instruction 69.
-  out += ins[control:epilogue-1]
-  out.append(BR(loop_target-len(out)))
-  out += ins[epilogue:]
-  if len(out) > len(ins): raise RuntimeError(f"packed shader grew from {len(ins)} to {len(out)}")
-  out += [NOP()] * (len(ins)-len(out))
-  fregs, hregs = struct.unpack_from("<II", lib, reg_off+0x14)
-  hregs = (hregs & 0x80000000) | 20
-  return inject(lib, image_off, image_size, reg_off, b"".join(out), fregs, hregs)
-
-
-def pack_4x2_wide_tile(lib: bytes) -> bytes:
-  """Use one rpt7 half-MAD for both adjacent output vectors."""
-  image_off, image_size = struct.unpack_from("<I", lib, 0xc0)[0], struct.unpack_from("<I", lib, 0x100)[0]
-  reg_off = struct.unpack_from("<I", lib, 0x34)[0]
-  ins = [lib[i:i+8] for i in range(image_off, image_off+image_size, 8)]
-  if len(ins) < 450: raise RuntimeError(f"unexpected 4x2 shader length {len(ins)}")
-
-  # Pack each low/high output weight pair into adjacent high half registers so
-  # a single rpt7 instruction covers all eight output lanes. High destinations
-  # avoid the merged-register aliases of the still-live coordinate registers.
-  # hr12..19 remain the compiler donor's initialized accumulator bank.
-  for index, dst, coord in ((72, "hr20.x", "r0.x"), (73, "hr21.x", "r0.z"),
-                            (77, "hr22.x", "r1.x"), (80, "hr23.x", "r1.z"),
-                            (84, "hr24.x", "r2.x"), (87, "hr25.x", "r2.z"),
-                            (90, "hr26.x", "r3.x"), (93, "hr27.x", "r3.z")):
-    ins[index] = ISAM_F16(dst, coord, tex=1)
-
-  out = list(ins[:108])
-  first = True
-  for component, weight in zip("xyzw", (20, 22, 24, 26)):
-    for activation, acc in ((11, 18), (10, 16), (5, 14), (4, 12)):
-      out.append(MAD_F16(f"hr{acc}.x", f"hr{activation}.{component}", f"hr{weight}.x", f"hr{acc}.x",
-                         rpt=7, r=True, sy=first))
-      first = False
-  out += ins[232:238]
-  out.append(BR(69-len(out)))
-  out += ins[239:]
-  if len(out) > len(ins): raise RuntimeError(f"packed shader grew from {len(ins)} to {len(out)}")
-  out += [NOP()] * (len(ins)-len(out))
-  fregs, hregs = struct.unpack_from("<II", lib, reg_off+0x14)
-  hregs = (hregs & 0x80000000) | 28
-  return inject(lib, image_off, image_size, reg_off, b"".join(out), fregs, hregs)
-
-
-def pack_split_tile(lib: bytes) -> bytes:
-  """Use the verified split-A register footprint while retaining the fused epilogue."""
-  image_off, image_size = struct.unpack_from("<I", lib, 0xc0)[0], struct.unpack_from("<I", lib, 0x100)[0]
-  reg_off = struct.unpack_from("<I", lib, 0x34)[0]
-  ins = [lib[i:i+8] for i in range(image_off, image_off+image_size, 8)]
-  if len(ins) < 500: raise RuntimeError(f"unexpected tile8 shader length {len(ins)}")
-
-  # The compiler donor already calculated all padded vision-image row bases in
-  # r9.w..r11.z and keeps K4/weight-x in r8.y/r8.x. Keep that coordinate work,
-  # but use the 28-half-register layout of the validated general 8x8 kernel:
-  # weights hr0..7, four reusable A vectors hr8..11, accumulators hr12..27.
-  out = ins[:127]
-  for acc in range(12, 28): out.append(MOV_H_IMM(f"hr{acc}.x", 0, rpt=3))
-  out.append(NOP())
-  if len(out) != 144: raise RuntimeError(f"split prologue ended at {len(out)}")
-
-  weight_coords = tuple((f"r{20+i//2}.{'xz'[i&1]}", f"r{20+i//2}.{'yw'[i&1]}") for i in range(8))
-  for component in range(4):
-    for col in range(2):
-      xreg, yreg = weight_coords[component*2+col]
-      out.append(MOV_F32(xreg, "r8.x") if component == 3 else ADD_S(xreg, "r8.x", component-3))
-      out.append(MOV_F32(yreg, "r8.w" if col == 0 else "r8.z"))
-  out.append(NOP(rpt=5))
-  for weight, (xreg, _) in enumerate(weight_coords): out.append(ISAM_F16(f"hr{weight}.x", xreg, 1))
-
-  activation_bases = ("r11.z", "r11.y", "r11.x", "r10.w", "r10.z", "r10.y", "r10.x", "r9.w")
-  def load_rows(first_row: int) -> None:
-    for slot, row in enumerate(range(first_row, first_row+4)):
-      out.extend((ADD_S_REG("r25.x", "r8.y", activation_bases[row]), NOP(rpt=5),
-                  ISAM_F16(f"hr{8+slot}.x", "r25.x", 0)))
-
-  def mad_rows(first_row: int) -> None:
-    first = True
-    for slot, row in enumerate(range(first_row, first_row+4)):
-      for component, (weight0, weight1) in zip("xyzw", ((0, 1), (2, 3), (4, 5), (6, 7))):
-        for col, weight in enumerate((weight0, weight1)):
-          acc = 12 + row*2 + col
-          out.append(MAD_F16(f"hr{acc}.x", f"hr{8+slot}.{component}", f"hr{weight}.x", f"hr{acc}.x",
-                             rpt=3, r=True, sy=first))
-          first = False
-
-  load_rows(0)
-  mad_rows(0)
-  load_rows(4)
-  mad_rows(4)
-  out += ins[195:198] + [ins[366]]
-  out.append(BR(143-len(out)))
-
-  # Retain the exact compiler-generated bias/GELU/image-store epilogue. Its
-  # first stage converts the old accumulator bank; redirect those sources to
-  # the compact bank without disturbing any later full-register scheduling.
-  lane_map = {}
-  for row in range(8):
-    for col in range(2):
-      old_vec, new_vec = 30-row*2+col, 12+row*2+col
-      for lane in range(4): lane_map[old_vec*4+lane] = new_vec*4+lane
-  epilogue = list(ins[368:])
-  for index in range(4, min(31, len(epilogue))):
-    lo, hi = struct.unpack("<II", epilogue[index])
-    if lo in lane_map and (hi & 0x00F04000) == 0x00004000:
-      epilogue[index] = struct.pack("<II", lane_map[lo], hi)
-  out += epilogue
-  fregs, hregs = struct.unpack_from("<II", lib, reg_off+0x14)
-  hregs = (hregs & 0x80000000) | 28
-  return inject(lib, image_off, image_size, reg_off, b"".join(out), fregs, hregs)
-
-
-def pack_split_safe_tile(lib: bytes) -> bytes:
-  """Split the eight A rows without changing the donor's sampler register assignment."""
-  image_off, image_size = struct.unpack_from("<I", lib, 0xc0)[0], struct.unpack_from("<I", lib, 0x100)[0]
-  reg_off = struct.unpack_from("<I", lib, 0x34)[0]
-  ins = [lib[i:i+8] for i in range(image_off, image_off+image_size, 8)]
-  if len(ins) < 500: raise RuntimeError(f"unexpected tile8 shader length {len(ins)}")
-  out = list(ins[:144])
-
-  # Preserve the compiler's proven sampler destinations. Only hr32/hr33 are
-  # remapped, to the dead activation slots hr10/hr11, reducing hregs 34 -> 32.
-  weight_regs = (6, 1, 9, 5, 8, 4, 7, 3)
-  weight_coords = ("r3.x", "r4.z", "r3.z", "r5.x", "r4.x", "r5.z", "r6.x", "r6.z")
-  weight_loads = ((168, 171), (172, 174), (175, 178), (179, 181),
-                  (182, 185), (186, 188), (189, 191), (192, 194))
-  def load_weights() -> None:
-    for index, (((start, load), coord), weight) in enumerate(zip(zip(weight_loads, weight_coords), weight_regs)):
-      out.extend(ins[start:load])
-      out.append(ISAM_F16(f"hr{weight}.x", coord, 1))
-
-  activation_bases = ("r11.z", "r11.y", "r11.x", "r10.w", "r10.z", "r10.y", "r10.x", "r9.w")
-  activation_regs = (15, 14, 13, 12)
-  def load_rows(first_row: int) -> None:
-    for slot, row in enumerate(range(first_row, first_row+4)):
-      out.extend((ADD_S_REG("r25.x", "r8.y", activation_bases[row]), NOP(rpt=5),
-                  ISAM_F16(f"hr{activation_regs[slot]}.x", "r25.x", 0)))
-
-  def mad_rows(first_row: int) -> None:
-    first = True
-    for slot, row in enumerate(range(first_row, first_row+4)):
-      for component, weights in zip("xyzw", ((6, 1), (9, 5), (8, 4), (7, 3))):
-        for col, weight in enumerate(weights):
-          acc = 30-row*2+col
-          out.append(MAD_F16(f"hr{acc}.x", f"hr{activation_regs[slot]}.{component}", f"hr{weight}.x", f"hr{acc}.x",
-                             rpt=3, r=True, sy=first))
-          first = False
-
-  load_rows(0)
-  load_weights()
-  mad_rows(0)
-  load_rows(4)
-  mad_rows(4)
-  out += ins[195:198] + [ins[366]]
-  out.append(BR(143-len(out)))
-  out += ins[368:]
-  fregs, hregs = struct.unpack_from("<II", lib, reg_off+0x14)
-  hregs = (hregs & 0x80000000) | 32
-  return inject(lib, image_off, image_size, reg_off, b"".join(out), fregs, hregs)
-
-
-def patch_model(model) -> int:
-  outer=model.captured.linear.src[0]; batch=list(outer.src[0].src[0].src)
-  source = SOURCE_4X2 if os.getenv("TILE4X2") else SOURCE
-  if os.getenv("TILE4X2_WG256"):
-    source = source.replace("reqd_work_group_size(128,1,1)", "reqd_work_group_size(256,1,1)") \
-                   .replace("get_group_id(1)*16+tm*4", "get_group_id(1)*32+tm*4")
-  if os.getenv("TILE_LINEAR"):
-    source = "\n".join(line for line in source.splitlines() if not line.startswith("inline float4 gelu"))
-    source = source.replace("gelu(convert", "(convert")
-  raw_lib = Device["QCOM"].compiler.compile_cached(source)
-  lib=((raw_lib if os.getenv("TILE4X2_RAW") else pack_4x2_wide_tile(raw_lib) if os.getenv("TILE4X2_WIDE") else
-        pack_4x2_tile(raw_lib, bool(os.getenv("TILE4X2_WG256")))) if os.getenv("TILE4X2") else
-       pack_split_safe_tile(raw_lib) if os.getenv("TILE8_SAFE") else pack_split_tile(raw_lib)); patched=0
-  for index,call in enumerate(batch):
-    if call.op is not Ops.CALL or call.src[0].op is not Ops.PROGRAM or plain_name(call.src[0].arg.name)!=TARGET: continue
-    program=call.src[0]
-    program=program.replace(arg=replace(program.arg,global_size=((3,4 if os.getenv("TILE4X2_WG256") else 8,1)
-                                                                  if os.getenv("TILE4X2") else (3,4,1)),
-      local_size=((256,1,1) if os.getenv("TILE4X2_WG256") else (128,1,1))),
-      src=program.src[:2]+(program.src[2].replace(arg=source),program.src[3].replace(arg=lib)))
-    batch[index]=call.replace(src=(program,*call.src[1:])); patched+=1
-  if patched:
-    model.captured._linear=model.captured.linear.substitute({outer:create_graph_call(batch)},walk=True)
-    model.captured.__dict__.pop("linear",None)
-  return patched
-
-
-def main() -> None:
-  ap=argparse.ArgumentParser(); ap.add_argument("input"); ap.add_argument("output"); args=ap.parse_args()
-  with open(args.input,"rb") as f:model=pickle.load(f)
-  print("patched",patch_model(model))
-  with open(args.output,"wb") as f:pickle.dump(model,f)
-
-
-if __name__=="__main__":main()

extra/gemm/qcom_openpilot_forward_tile8x4.py

@@ -1,130 +0,0 @@
-#!/usr/bin/env python3
-"""High-intensity 8x4 FP16-accumulate tile for OpenPilot vision projections."""
-import argparse, os, pickle, struct
-from dataclasses import replace
-
-from tinygrad import Device
-from tinygrad.engine.jit import create_graph_call
-from tinygrad.uop.ops import Ops
-from extra.gemm.ir3asm import BR, ISAM_F16, MAD_F16, MOV_H_IMM, NOP, inject
-from extra.gemm.qcom_ir3_matmul_patch import plain_name
-
-TARGET = "r_32_192_4_4_64_4"
-SOURCE = r"""#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-const sampler_t smp=CLK_NORMALIZED_COORDS_FALSE|CLK_ADDRESS_CLAMP|CLK_FILTER_NEAREST;
-inline float4 gelu(float4 v) { return ((float4)(1)/(1+exp2((v+(float4)(0.044708251953125f)*v*v*v)*(float4)(-2.3021129851685216f))))*v; }
-__attribute__((reqd_work_group_size(128,1,1)))
-__kernel void r_32_192_4_4_64_4(write_only image2d_t O,read_only image2d_t A,read_only image2d_t W,read_only image2d_t B) {
-  int lid=get_local_id(0),tm=lid>>5,tid=lid&31,row0=get_group_id(1)*32+tm*8,n=get_group_id(0)*32+tid;
-  half4 c0=(half4)(0),c1=(half4)(0),c2=(half4)(0),c3=(half4)(0);
-  half4 c4=(half4)(0),c5=(half4)(0),c6=(half4)(0),c7=(half4)(0);
-  for(int k=0;k<64;k++) {
-    int x=k*4;
-    half4 w0=read_imageh(W,smp,(int2)(x,n));
-    half4 w1=read_imageh(W,smp,(int2)(x+1,n));
-    half4 w2=read_imageh(W,smp,(int2)(x+2,n));
-    half4 w3=read_imageh(W,smp,(int2)(x+3,n));
-    int r=row0,base=(r&31)*260+(r>>5)*65;
-    half4 a0=read_imageh(A,smp,(int2)(base+k,0)); r++; base=(r&31)*260+(r>>5)*65;
-    half4 a1=read_imageh(A,smp,(int2)(base+k,0)); r++; base=(r&31)*260+(r>>5)*65;
-    half4 a2=read_imageh(A,smp,(int2)(base+k,0)); r++; base=(r&31)*260+(r>>5)*65;
-    half4 a3=read_imageh(A,smp,(int2)(base+k,0)); r++; base=(r&31)*260+(r>>5)*65;
-    half4 a4=read_imageh(A,smp,(int2)(base+k,0)); r++; base=(r&31)*260+(r>>5)*65;
-    half4 a5=read_imageh(A,smp,(int2)(base+k,0)); r++; base=(r&31)*260+(r>>5)*65;
-    half4 a6=read_imageh(A,smp,(int2)(base+k,0)); r++; base=(r&31)*260+(r>>5)*65;
-    half4 a7=read_imageh(A,smp,(int2)(base+k,0));
-    c0+=(half4)(a0.x)*w0+(half4)(a0.y)*w1+(half4)(a0.z)*w2+(half4)(a0.w)*w3;
-    c1+=(half4)(a1.x)*w0+(half4)(a1.y)*w1+(half4)(a1.z)*w2+(half4)(a1.w)*w3;
-    c2+=(half4)(a2.x)*w0+(half4)(a2.y)*w1+(half4)(a2.z)*w2+(half4)(a2.w)*w3;
-    c3+=(half4)(a3.x)*w0+(half4)(a3.y)*w1+(half4)(a3.z)*w2+(half4)(a3.w)*w3;
-    c4+=(half4)(a4.x)*w0+(half4)(a4.y)*w1+(half4)(a4.z)*w2+(half4)(a4.w)*w3;
-    c5+=(half4)(a5.x)*w0+(half4)(a5.y)*w1+(half4)(a5.z)*w2+(half4)(a5.w)*w3;
-    c6+=(half4)(a6.x)*w0+(half4)(a6.y)*w1+(half4)(a6.z)*w2+(half4)(a6.w)*w3;
-    c7+=(half4)(a7.x)*w0+(half4)(a7.y)*w1+(half4)(a7.z)*w2+(half4)(a7.w)*w3;
-  }
-  float4 b=read_imagef(B,smp,(int2)(n,0));
-  int r=row0;
-#define STORE(v) { int m=r&31,p=r>>5; write_imagef(O,(int2)(n+p*192,m),gelu(convert_float4(v)+b)); r++; }
-  STORE(c0); STORE(c1); STORE(c2); STORE(c3); STORE(c4); STORE(c5); STORE(c6); STORE(c7);
-}"""
-
-
-def pack_tile(lib:bytes) -> bytes:
-  image_off,image_size=struct.unpack_from("<I",lib,0xc0)[0],struct.unpack_from("<I",lib,0x100)[0]
-  reg_off=struct.unpack_from("<I",lib,0x34)[0]
-  ins=[lib[i:i+8] for i in range(image_off,image_off+image_size,8)]
-  if len(ins)<500: raise RuntimeError(f"unexpected tile8 shader length {len(ins)}")
-  # Pair each pair of output-channel texture rows in consecutive half registers.
-  for index,dst,coord in ((171,"hr3.x","r3.x"),(174,"hr4.x","r4.z"),
-                          (178,"hr5.x","r3.z"),(181,"hr6.x","r5.x"),
-                          (185,"hr7.x","r4.x"),(188,"hr8.x","r5.z"),
-                          (191,"hr32.x","r6.x"),(194,"hr33.x","r6.z")):
-    ins[index]=ISAM_F16(dst,coord,tex=1,sy=index==194)
-  out=ins[:198]
-  for activation,acc in (("hr15",30),("hr14",28),("hr13",26),("hr12",24),
-                         ("hr11",22),("hr10",20),("hr2",18),("hr0",16)):
-    for component,(weight_lo,weight_hi) in zip("xyzw",((3,4),(5,6),(7,8),(32,33))):
-      for offset,weight in enumerate((weight_lo,weight_hi)):
-        out.append(MAD_F16(f"hr{acc+offset}.x",f"{activation}.{component}",f"hr{weight}.x",f"hr{acc+offset}.x",
-                           rpt=3,r=True,sy=len(out)==198))
-  out.append(ins[366])
-  out.append(BR(143-len(out)))
-  out+=ins[368:]
-  if len(out)>len(ins): raise RuntimeError(f"packed shader grew from {len(ins)} to {len(out)}")
-  out += [NOP()]*(len(ins)-len(out))
-  fregs,hregs=struct.unpack_from("<II",lib,reg_off+0x14)
-  hregs=(hregs&0x80000000)|max(hregs&0x7fffffff,34)
-  return inject(lib,image_off,image_size,reg_off,b"".join(out),fregs,hregs)
-
-
-def pack_compact_tile(lib:bytes) -> bytes:
-  """Replace the compiler's split partial sums with eight direct half4 accumulators."""
-  image_off,image_size=struct.unpack_from("<I",lib,0xc0)[0],struct.unpack_from("<I",lib,0x100)[0]
-  reg_off=struct.unpack_from("<I",lib,0x34)[0]
-  ins=[lib[i:i+8] for i in range(image_off,image_off+image_size,8)]
-  if len(ins) < 500: raise RuntimeError(f"unexpected tile8x4 shader length {len(ins)}")
-  out=ins[:91] + [MOV_H_IMM(f"hr{acc}.x",0,rpt=3) for acc in range(10,18)] + ins[99:123]
-  out += ins[123:126] + ins[126:129] + ins[177:179] + [ISAM_F16("hr18.x","r4.x",1)] \
-       + ins[211:213] + [ISAM_F16("hr19.x","r4.z",1)]
-  first=True
-  for weight,component in (("hr8.x","x"),("hr9.x","y"),("hr18.x","z"),("hr19.x","w")):
-    for activation,acc in zip(("hr7","hr6","hr5","hr4","hr3","hr2","hr1","hr0"),range(17,9,-1)):
-      out.append(MAD_F16(f"hr{acc}.x",f"{activation}.{component}",weight,f"hr{acc}.x",rpt=3,r=True,sy=first))
-      first=False
-  out += ins[256:261]
-  out.append(BR(99-len(out)))
-  out += ins[262:]
-  if len(out)>len(ins): raise RuntimeError(f"compact shader grew from {len(ins)} to {len(out)}")
-  out += [NOP()]*(len(ins)-len(out))
-  fregs,hregs=struct.unpack_from("<II",lib,reg_off+0x14)
-  hregs=(hregs&0x80000000)|20
-  return inject(lib,image_off,image_size,reg_off,b"".join(out),fregs,hregs)
-
-
-def patch_model(model) -> int:
-  outer=model.captured.linear.src[0]; batch=list(outer.src[0].src[0].src)
-  lib=pack_compact_tile(Device["QCOM"].compiler.compile_cached(SOURCE)); patched=0
-  seen=0
-  for index,call in enumerate(batch):
-    if patched >= int(os.getenv("MAX_PATCH", "1000000")): break
-    if call.op is not Ops.CALL or call.src[0].op is not Ops.PROGRAM or plain_name(call.src[0].arg.name)!=TARGET: continue
-    if seen != int(os.getenv("TARGET_INDEX", str(seen))): seen+=1; continue
-    seen+=1
-    program=call.src[0]
-    program=program.replace(arg=replace(program.arg,global_size=(6,4,1),local_size=(128,1,1)),
-      src=program.src[:2]+(program.src[2].replace(arg=SOURCE),program.src[3].replace(arg=lib)))
-    batch[index]=call.replace(src=(program,*call.src[1:])); patched+=1
-  if patched:
-    model.captured._linear=model.captured.linear.substitute({outer:create_graph_call(batch)},walk=True)
-    model.captured.__dict__.pop("linear",None)
-  return patched
-
-
-def main() -> None:
-  ap=argparse.ArgumentParser(); ap.add_argument("input"); ap.add_argument("output"); args=ap.parse_args()
-  with open(args.input,"rb") as f:model=pickle.load(f)
-  print("patched",patch_model(model))
-  with open(args.output,"wb") as f:pickle.dump(model,f)
-
-
-if __name__=="__main__":main()

extra/gemm/qcom_openpilot_fuse_vision_mlp.py

@@ -1,317 +0,0 @@
-#!/usr/bin/env python3
-"""Fuse the six driving_vision MLP pairs through work-group local memory."""
-import argparse, pickle, re, struct
-from dataclasses import replace
-
-from tinygrad import Device, dtypes
-from tinygrad.engine.jit import create_graph_call
-from tinygrad.uop.ops import Ops
-from extra.gemm.ir3asm import BR, ISAM_F16, MAD_F16, NOP, inject
-from extra.gemm.qcom_ir3_matmul_patch import plain_name
-
-FORWARD, INVERSE = "r_32_192_4_4_64_4", "r_32_64_4_4_192_4"
-
-SOURCE = r"""#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-const sampler_t smp=CLK_NORMALIZED_COORDS_FALSE|CLK_ADDRESS_CLAMP|CLK_FILTER_NEAREST;
-inline float4 gelu(float4 v) {
-  return ((float4)(1)/(1+exp2((v+(float4)(0.044708251953125f)*v*v*v)*(float4)(-2.3021129851685216f))))*v;
-}
-__attribute__((reqd_work_group_size(64,1,1)))
-__kernel void fused_vision_mlp(write_only image2d_t O,read_only image2d_t X,read_only image2d_t A,
-                               read_only image2d_t W1,read_only image2d_t B1,read_only image2d_t W2,
-                               read_only image2d_t B2,read_only image2d_t S) {
-  int lid=get_local_id(0),g=get_group_id(1);
-  __local half4 hidden[768];
-  {
-    half4 z0[3]={(half4)(0),(half4)(0),(half4)(0)};
-    half4 z1[3]={(half4)(0),(half4)(0),(half4)(0)};
-    half4 z2[3]={(half4)(0),(half4)(0),(half4)(0)};
-    half4 z3[3]={(half4)(0),(half4)(0),(half4)(0)};
-    for(int k=0;k<64;k++) {
-      int ax=g*260+k,wx=k*4;
-      half4 a0=read_imageh(A,smp,(int2)(ax,0)),a1=read_imageh(A,smp,(int2)(ax+65,0));
-      half4 a2=read_imageh(A,smp,(int2)(ax+130,0)),a3=read_imageh(A,smp,(int2)(ax+195,0));
-      #pragma unroll 3
-      for(int j=0;j<3;j++) {
-        int n=lid+j*64;
-        half4 w0=read_imageh(W1,smp,(int2)(wx,n)),w1=read_imageh(W1,smp,(int2)(wx+1,n));
-        half4 w2=read_imageh(W1,smp,(int2)(wx+2,n)),w3=read_imageh(W1,smp,(int2)(wx+3,n));
-        z0[j]+=(half4)(a0.x)*w0+(half4)(a0.y)*w1+(half4)(a0.z)*w2+(half4)(a0.w)*w3;
-        z1[j]+=(half4)(a1.x)*w0+(half4)(a1.y)*w1+(half4)(a1.z)*w2+(half4)(a1.w)*w3;
-        z2[j]+=(half4)(a2.x)*w0+(half4)(a2.y)*w1+(half4)(a2.z)*w2+(half4)(a2.w)*w3;
-        z3[j]+=(half4)(a3.x)*w0+(half4)(a3.y)*w1+(half4)(a3.z)*w2+(half4)(a3.w)*w3;
-      }
-    }
-    #pragma unroll 3
-    for(int j=0;j<3;j++) {
-      int n=lid+j*64; float4 b=read_imagef(B1,smp,(int2)(n,0));
-      hidden[n]=convert_half4(gelu(convert_float4(z0[j])+b));
-      hidden[192+n]=convert_half4(gelu(convert_float4(z1[j])+b));
-      hidden[384+n]=convert_half4(gelu(convert_float4(z2[j])+b));
-      hidden[576+n]=convert_half4(gelu(convert_float4(z3[j])+b));
-    }
-  }
-  barrier(CLK_LOCAL_MEM_FENCE);
-  {
-    half4 z0=(half4)(0),z1=(half4)(0),z2=(half4)(0),z3=(half4)(0);
-    for(int k=0;k<192;k++) {
-      int wx=k*4;
-      half4 a0=hidden[k],a1=hidden[192+k],a2=hidden[384+k],a3=hidden[576+k];
-      half4 w0=read_imageh(W2,smp,(int2)(wx,lid)),w1=read_imageh(W2,smp,(int2)(wx+1,lid));
-      half4 w2=read_imageh(W2,smp,(int2)(wx+2,lid)),w3=read_imageh(W2,smp,(int2)(wx+3,lid));
-      z0+=(half4)(a0.x)*w0+(half4)(a0.y)*w1+(half4)(a0.z)*w2+(half4)(a0.w)*w3;
-      z1+=(half4)(a1.x)*w0+(half4)(a1.y)*w1+(half4)(a1.z)*w2+(half4)(a1.w)*w3;
-      z2+=(half4)(a2.x)*w0+(half4)(a2.y)*w1+(half4)(a2.z)*w2+(half4)(a2.w)*w3;
-      z3+=(half4)(a3.x)*w0+(half4)(a3.y)*w1+(half4)(a3.z)*w2+(half4)(a3.w)*w3;
-    }
-    int x=lid+g*256; float4 b=read_imagef(B2,smp,(int2)(lid,0)),s=read_imagef(S,smp,(int2)(lid,0));
-    write_imagef(O,(int2)(x,0),read_imagef(X,smp,(int2)(x,0))+(convert_float4(z0)+b)*s);
-    write_imagef(O,(int2)(x+64,0),read_imagef(X,smp,(int2)(x+64,0))+(convert_float4(z1)+b)*s);
-    write_imagef(O,(int2)(x+128,0),read_imagef(X,smp,(int2)(x+128,0))+(convert_float4(z2)+b)*s);
-    write_imagef(O,(int2)(x+192,0),read_imagef(X,smp,(int2)(x+192,0))+(convert_float4(z3)+b)*s);
-  }
-}"""
-
-SOURCE_TILED = r"""#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-const sampler_t smp=CLK_NORMALIZED_COORDS_FALSE|CLK_ADDRESS_CLAMP|CLK_FILTER_NEAREST;
-inline float4 gelu(float4 v) {
-  return ((float4)(1)/(1+exp2((v+(float4)(0.044708251953125f)*v*v*v)*(float4)(-2.3021129851685216f))))*v;
-}
-__attribute__((reqd_work_group_size(64,1,1)))
-__kernel void fused_vision_mlp(write_only image2d_t O,read_only image2d_t X,read_only image2d_t A,
-                               read_only image2d_t W1,read_only image2d_t B1,read_only image2d_t W2,
-                               read_only image2d_t B2,read_only image2d_t S) {
-  int lid=get_local_id(0),g=get_group_id(1);
-  __local half4 hidden[256];
-  half4 o0=(half4)(0),o1=(half4)(0),o2=(half4)(0),o3=(half4)(0);
-  #pragma unroll 3
-  for(int j=0;j<3;j++) {
-    int n=lid+j*64;
-    half4 z0=(half4)(0),z1=(half4)(0),z2=(half4)(0),z3=(half4)(0);
-    for(int k=0;k<64;k++) {
-      int ax=g*260+k,wx=k*4;
-      half4 a0=read_imageh(A,smp,(int2)(ax,0)),a1=read_imageh(A,smp,(int2)(ax+65,0));
-      half4 a2=read_imageh(A,smp,(int2)(ax+130,0)),a3=read_imageh(A,smp,(int2)(ax+195,0));
-      half4 w0=read_imageh(W1,smp,(int2)(wx,n)),w1=read_imageh(W1,smp,(int2)(wx+1,n));
-      half4 w2=read_imageh(W1,smp,(int2)(wx+2,n)),w3=read_imageh(W1,smp,(int2)(wx+3,n));
-      z0+=(half4)(a0.x)*w0+(half4)(a0.y)*w1+(half4)(a0.z)*w2+(half4)(a0.w)*w3;
-      z1+=(half4)(a1.x)*w0+(half4)(a1.y)*w1+(half4)(a1.z)*w2+(half4)(a1.w)*w3;
-      z2+=(half4)(a2.x)*w0+(half4)(a2.y)*w1+(half4)(a2.z)*w2+(half4)(a2.w)*w3;
-      z3+=(half4)(a3.x)*w0+(half4)(a3.y)*w1+(half4)(a3.z)*w2+(half4)(a3.w)*w3;
-    }
-    float4 b=read_imagef(B1,smp,(int2)(n,0));
-    hidden[lid]=convert_half4(gelu(convert_float4(z0)+b));
-    hidden[64+lid]=convert_half4(gelu(convert_float4(z1)+b));
-    hidden[128+lid]=convert_half4(gelu(convert_float4(z2)+b));
-    hidden[192+lid]=convert_half4(gelu(convert_float4(z3)+b));
-    barrier(CLK_LOCAL_MEM_FENCE);
-    for(int k=0;k<64;k++) {
-      int wx=(j*64+k)*4;
-      half4 a0=hidden[k],a1=hidden[64+k],a2=hidden[128+k],a3=hidden[192+k];
-      half4 w0=read_imageh(W2,smp,(int2)(wx,lid)),w1=read_imageh(W2,smp,(int2)(wx+1,lid));
-      half4 w2=read_imageh(W2,smp,(int2)(wx+2,lid)),w3=read_imageh(W2,smp,(int2)(wx+3,lid));
-      o0+=(half4)(a0.x)*w0+(half4)(a0.y)*w1+(half4)(a0.z)*w2+(half4)(a0.w)*w3;
-      o1+=(half4)(a1.x)*w0+(half4)(a1.y)*w1+(half4)(a1.z)*w2+(half4)(a1.w)*w3;
-      o2+=(half4)(a2.x)*w0+(half4)(a2.y)*w1+(half4)(a2.z)*w2+(half4)(a2.w)*w3;
-      o3+=(half4)(a3.x)*w0+(half4)(a3.y)*w1+(half4)(a3.z)*w2+(half4)(a3.w)*w3;
-    }
-    barrier(CLK_LOCAL_MEM_FENCE);
-  }
-  int x=lid+g*256; float4 b=read_imagef(B2,smp,(int2)(lid,0)),s=read_imagef(S,smp,(int2)(lid,0));
-  write_imagef(O,(int2)(x,0),read_imagef(X,smp,(int2)(x,0))+(convert_float4(o0)+b)*s);
-  write_imagef(O,(int2)(x+64,0),read_imagef(X,smp,(int2)(x+64,0))+(convert_float4(o1)+b)*s);
-  write_imagef(O,(int2)(x+128,0),read_imagef(X,smp,(int2)(x+128,0))+(convert_float4(o2)+b)*s);
-  write_imagef(O,(int2)(x+192,0),read_imagef(X,smp,(int2)(x+192,0))+(convert_float4(o3)+b)*s);
-}"""
-
-SOURCE_PARALLEL = r"""#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-const sampler_t smp=CLK_NORMALIZED_COORDS_FALSE|CLK_ADDRESS_CLAMP|CLK_FILTER_NEAREST;
-inline float4 gelu(float4 v) {
-  return ((float4)(1)/(1+exp2((v+(float4)(0.044708251953125f)*v*v*v)*(float4)(-2.3021129851685216f))))*v;
-}
-__attribute__((reqd_work_group_size(192,1,1)))
-__kernel void fused_vision_mlp(write_only image2d_t O,read_only image2d_t X,read_only image2d_t A,
-                               read_only image2d_t W1,read_only image2d_t B1,read_only image2d_t W2,
-                               read_only image2d_t B2,read_only image2d_t S) {
-  int lid=get_local_id(0),g=get_group_id(1),n=lid;
-  __local half4 hidden[768];
-  __local half4 partial[768];
-  half4 z0=(half4)(0),z1=(half4)(0),z2=(half4)(0),z3=(half4)(0);
-  for(int k=0;k<64;k++) {
-    int ax=g*260+k,wx=k*4;
-    half4 a0=read_imageh(A,smp,(int2)(ax,0)),a1=read_imageh(A,smp,(int2)(ax+65,0));
-    half4 a2=read_imageh(A,smp,(int2)(ax+130,0)),a3=read_imageh(A,smp,(int2)(ax+195,0));
-    half4 w0=read_imageh(W1,smp,(int2)(wx,n)),w1=read_imageh(W1,smp,(int2)(wx+1,n));
-    half4 w2=read_imageh(W1,smp,(int2)(wx+2,n)),w3=read_imageh(W1,smp,(int2)(wx+3,n));
-    z0+=(half4)(a0.x)*w0+(half4)(a0.y)*w1+(half4)(a0.z)*w2+(half4)(a0.w)*w3;
-    z1+=(half4)(a1.x)*w0+(half4)(a1.y)*w1+(half4)(a1.z)*w2+(half4)(a1.w)*w3;
-    z2+=(half4)(a2.x)*w0+(half4)(a2.y)*w1+(half4)(a2.z)*w2+(half4)(a2.w)*w3;
-    z3+=(half4)(a3.x)*w0+(half4)(a3.y)*w1+(half4)(a3.z)*w2+(half4)(a3.w)*w3;
-  }
-  float4 b1=read_imagef(B1,smp,(int2)(n,0));
-  hidden[n]=convert_half4(gelu(convert_float4(z0)+b1));
-  hidden[192+n]=convert_half4(gelu(convert_float4(z1)+b1));
-  hidden[384+n]=convert_half4(gelu(convert_float4(z2)+b1));
-  hidden[576+n]=convert_half4(gelu(convert_float4(z3)+b1));
-  barrier(CLK_LOCAL_MEM_FENCE);
-  int part=lid>>6; n=lid&63;
-  half4 o0=(half4)(0),o1=(half4)(0),o2=(half4)(0),o3=(half4)(0);
-  for(int k=part*64;k<(part+1)*64;k++) {
-    int wx=k*4;
-    half4 a0=hidden[k],a1=hidden[192+k],a2=hidden[384+k],a3=hidden[576+k];
-    half4 w0=read_imageh(W2,smp,(int2)(wx,n)),w1=read_imageh(W2,smp,(int2)(wx+1,n));
-    half4 w2=read_imageh(W2,smp,(int2)(wx+2,n)),w3=read_imageh(W2,smp,(int2)(wx+3,n));
-    o0+=(half4)(a0.x)*w0+(half4)(a0.y)*w1+(half4)(a0.z)*w2+(half4)(a0.w)*w3;
-    o1+=(half4)(a1.x)*w0+(half4)(a1.y)*w1+(half4)(a1.z)*w2+(half4)(a1.w)*w3;
-    o2+=(half4)(a2.x)*w0+(half4)(a2.y)*w1+(half4)(a2.z)*w2+(half4)(a2.w)*w3;
-    o3+=(half4)(a3.x)*w0+(half4)(a3.y)*w1+(half4)(a3.z)*w2+(half4)(a3.w)*w3;
-  }
-  int po=part*256+n;
-  partial[po]=o0; partial[po+64]=o1; partial[po+128]=o2; partial[po+192]=o3;
-  barrier(CLK_LOCAL_MEM_FENCE);
-  if(part==0) {
-    o0=partial[n]+partial[256+n]+partial[512+n];
-    o1=partial[64+n]+partial[320+n]+partial[576+n];
-    o2=partial[128+n]+partial[384+n]+partial[640+n];
-    o3=partial[192+n]+partial[448+n]+partial[704+n];
-    int x=n+g*256; float4 b=read_imagef(B2,smp,(int2)(n,0)),s=read_imagef(S,smp,(int2)(n,0));
-    write_imagef(O,(int2)(x,0),read_imagef(X,smp,(int2)(x,0))+(convert_float4(o0)+b)*s);
-    write_imagef(O,(int2)(x+64,0),read_imagef(X,smp,(int2)(x+64,0))+(convert_float4(o1)+b)*s);
-    write_imagef(O,(int2)(x+128,0),read_imagef(X,smp,(int2)(x+128,0))+(convert_float4(o2)+b)*s);
-    write_imagef(O,(int2)(x+192,0),read_imagef(X,smp,(int2)(x+192,0))+(convert_float4(o3)+b)*s);
-  }
-}"""
-
-
-def pack_tiled_inner(lib:bytes) -> bytes:
-  """Pack the compiler's scalar 4x4 half outer product into 16 repeated MADs."""
-  image_off, image_size = struct.unpack_from("<I", lib, 0xc0)[0], struct.unpack_from("<I", lib, 0x100)[0]
-  reg_off = struct.unpack_from("<I", lib, 0x34)[0]
-  ins = [lib[i:i+8] for i in range(image_off, image_off+image_size, 8)]
-  if len(ins) < 500: raise RuntimeError(f"unexpected fused shader length {len(ins)}")
-  # 37:49 loads the four spatial activation vectors. 49:57 loads W0/W1; the
-  # compiler's W2/W3 coordinates are at 79:82 and 98:101. Give all four weights
-  # stable destinations, then accumulate directly into z0..z3.
-  out = list(ins[:57])
-  out += ins[79:82] + [ISAM_F16("hr6.x", "r0.x", 2)]
-  out += ins[98:101] + [ISAM_F16("hr7.x", "r0.x", 2)]
-  loop_start = 37
-  first = True
-  for activation, acc in zip(("hr3", "hr2", "hr1", "hr0"), ("hr13.y", "hr12.y", "hr11.y", "hr10.y")):
-    for component, weight in zip("xyzw", ("hr4.x", "hr5.x", "hr6.x", "hr7.x")):
-      out.append(MAD_F16(acc, f"{activation}.{component}", weight, acc, rpt=3, r=True, sy=first))
-      first = False
-  out += ins[124:130]
-  out.append(BR(loop_start-len(out)))
-  # Compacting the first inner loop relocates the GELU, local-memory inverse,
-  # and outer-j tail. Preserve branches wholly inside that tail, and rebuild
-  # every branch whose target remains in the untouched prologue.
-  tail_start = len(out)
-  shift = tail_start-131
-  for old_index, instruction in enumerate(ins[131:], 131):
-    lo, hi = struct.unpack("<iI", instruction)
-    new_index = old_index+shift
-    if hi in (0x00800000, 0x00900000) and (old_target:=old_index+lo) < 131:
-      instruction = BR(old_target-new_index, inv=hi == 0x00900000)
-    out.append(instruction)
-  if len(out) > len(ins): raise RuntimeError(f"packed fused shader grew from {len(ins)} to {len(out)}")
-  out += [NOP()] * (len(ins)-len(out))
-  fregs, hregs = struct.unpack_from("<II", lib, reg_off+0x14)
-  return inject(lib, image_off, image_size, reg_off, b"".join(out), fregs, hregs)
-
-
-def pack_tiled_inverse(lib:bytes) -> bytes:
-  """Pack the local-hidden x W2 4x4 outer product in the inverse phase."""
-  image_off, image_size = struct.unpack_from("<I", lib, 0xc0)[0], struct.unpack_from("<I", lib, 0x100)[0]
-  reg_off = struct.unpack_from("<I", lib, 0x34)[0]
-  ins = [lib[i:i+8] for i in range(image_off, image_off+image_size, 8)]
-  if len(ins) < 529: raise RuntimeError(f"unexpected fused shader length {len(ins)}")
-  out = list(ins[:396])
-  out += ins[396:407] + [ISAM_F16("hr12.x", "r3.y", 4)]
-  out += ins[408:410] + [ISAM_F16("hr13.x", "r3.w", 4)]
-  out += ins[432:434] + [ISAM_F16("hr14.x", "r4.y", 4)]
-  out += ins[450:452] + [ISAM_F16("hr15.x", "r4.w", 4)]
-  first = True
-  for activation, acc in (("hr2", "hr9.y"), ("hr3", "hr8.y"), ("hr4", "hr7.y"), ("hr5", "hr6.y")):
-    for component, weight in zip("xyzw", ("hr12.x", "hr13.x", "hr14.x", "hr15.x")):
-      out.append(MAD_F16(acc, f"{activation}.{component}", weight, acc, rpt=3, r=True, sy=first))
-      first = False
-  out += ins[470:476]
-  out.append(BR(400-len(out)))
-  tail_start = len(out)
-  shift = tail_start-477
-  for old_index, instruction in enumerate(ins[477:], 477):
-    lo, hi = struct.unpack("<iI", instruction)
-    new_index = old_index+shift
-    if hi in (0x00800000, 0x00900000) and (old_target:=old_index+lo) < 477:
-      instruction = BR(old_target-new_index, inv=hi == 0x00900000)
-    out.append(instruction)
-  if len(out) > len(ins): raise RuntimeError(f"packed fused shader grew from {len(ins)} to {len(out)}")
-  out += [NOP()] * (len(ins)-len(out))
-  fregs, hregs = struct.unpack_from("<II", lib, reg_off+0x14)
-  return inject(lib, image_off, image_size, reg_off, b"".join(out), fregs, hregs)
-
-
-def fp32_acc_source(source:str) -> str:
-  """Keep the hidden local tile in half, but accumulate both projections in float."""
-  source = source.replace("read_imageh(", "read_imagef(")
-  source = source.replace("half4 z", "float4 z").replace("half4 o", "float4 o")
-  source = source.replace("half4 a", "float4 a").replace("half4 w", "float4 w")
-  source = source.replace("(half4)(a", "(float4)(a")
-  source = source.replace("(half4)(0)", "(float4)(0)")
-  source = re.sub(r"=hidden\[([^]]+)\]", r"=convert_float4(hidden[\1])", source)
-  return source
-
-
-def patch_model(model, packed:bool=True, barriers:bool=True, inverse_packed:bool=False, parallel:bool=False,
-                fp32_acc:bool=False) -> int:
-  outer=model.captured.linear.src[0]
-  batch=list(outer.src[0].src[0].src)
-  # Keep the local-memory barriers until the packed shader has been checked for
-  # both timing and changed-input correctness on device.
-  source = SOURCE_PARALLEL if parallel else SOURCE_TILED
-  if fp32_acc: source = fp32_acc_source(source)
-  if not barriers: source=source.replace("barrier(CLK_LOCAL_MEM_FENCE);", "")
-  lib=Device["QCOM"].compiler.compile(source)
-  if inverse_packed and not fp32_acc: lib=pack_tiled_inverse(lib)
-  if packed and not fp32_acc: lib=pack_tiled_inner(lib)
-  specs=((dtypes.half,(1,8192,4)),(dtypes.half,(1,8192,4)),(dtypes.half,(1,8320,4)),
-         (dtypes.half,(192,320,4)),(dtypes.half,(1,192,4)),(dtypes.half,(64,768,4)),
-         (dtypes.half,(1,64,4)),(dtypes.half,(1,64,4)))
-  aux=(tuple(((i,dtype,shape),) for i,(dtype,shape) in enumerate(specs)),)
-  replacements, skip, patched = {}, set(), 0
-  for index in range(len(batch)-1):
-    forward,inverse=batch[index:index+2]
-    if not all(x.op is Ops.CALL and x.src[0].op is Ops.PROGRAM for x in (forward,inverse)): continue
-    if (plain_name(forward.src[0].arg.name),plain_name(inverse.src[0].arg.name)) != (FORWARD,INVERSE): continue
-    info=replace(forward.src[0].arg,name="fused_vision_mlp",global_size=(1,32,1),local_size=((192 if parallel else 64),1,1),
-                 globals=tuple(range(8)),outs=(0,),ins=tuple(range(1,8)),aux=aux)
-    program=forward.src[0].replace(arg=info,src=forward.src[0].src[:2]+
-      (forward.src[0].src[2].replace(arg=source),forward.src[0].src[3].replace(arg=lib)))
-    replacements[index]=program.call(inverse.src[1],inverse.src[2],forward.src[2],forward.src[3],forward.src[4],
-                                     inverse.src[4],inverse.src[5],inverse.src[6])
-    skip.add(index+1)
-    patched+=1
-  new_batch=[replacements.get(i,call) for i,call in enumerate(batch) if i not in skip]
-  if patched:
-    model.captured._linear=model.captured.linear.substitute({outer:create_graph_call(new_batch)},walk=True)
-    model.captured.__dict__.pop("linear",None)
-  return patched
-
-
-def main() -> None:
-  ap=argparse.ArgumentParser()
-  ap.add_argument("input")
-  ap.add_argument("output")
-  ap.add_argument("--no-pack", action="store_true")
-  ap.add_argument("--no-barrier", action="store_true")
-  ap.add_argument("--inverse-pack", action="store_true")
-  ap.add_argument("--parallel", action="store_true")
-  ap.add_argument("--fp32-acc", action="store_true")
-  args=ap.parse_args()
-  with open(args.input,"rb") as f:model=pickle.load(f)
-  print("patched",patch_model(model, not args.no_pack, not args.no_barrier, args.inverse_pack, args.parallel, args.fp32_acc))
-  with open(args.output,"wb") as f:pickle.dump(model,f)
-
-
-if __name__=="__main__": main()