work

examples/openpilot/compile3.py

@@ -65,6 +65,30 @@ def compile(onnx_file):
   if (allowed_gated_read_image:=getenv("ALLOWED_GATED_READ_IMAGE", -1)) != -1:
     assert gated_read_image_count == allowed_gated_read_image, f"different gated read_image! {gated_read_image_count=}, {allowed_gated_read_image=}"
 
+  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_RPT", 1):
+    from extra.gemm.qcom_openpilot_vision_fp16 import patch_fp32_rpt
+    if (patched:=patch_fp32_rpt(run_onnx_jit)): print(f"repeat-packed {patched} QCOM vision kernels")
+  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_SCHEDULE", 1):
+    from extra.gemm.qcom_openpilot_schedule_projection import patch_projection
+    if (patched:=patch_projection(run_onnx_jit)): print(f"rescheduled {patched} QCOM vision kernels")
+  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_FULL_RPT", 1):
+    from extra.gemm.qcom_openpilot_inverse_full_rpt import patch_model as patch_full_rpt
+    if (patched:=patch_full_rpt(run_onnx_jit)): print(f"fully repeat-packed {patched} QCOM vision kernels")
+  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_DEDUPE", 1):
+    from extra.gemm.qcom_openpilot_dedupe_head import dedupe_identical_calls
+    if (removed:=dedupe_identical_calls(run_onnx_jit)): print(f"deduplicated {len(removed)} QCOM kernels")
+  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_PACK_CONV", 1):
+    from extra.gemm.qcom_openpilot_pack_conv_weights import patch_conv
+    if (patched:=patch_conv(run_onnx_jit)): print(f"packed weights for {patched} QCOM convolution kernels")
+  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_LEVEL_SCHEDULE", 1):
+    from extra.gemm.qcom_openpilot_level_schedule import schedule_levels
+    if (moved:=schedule_levels(run_onnx_jit)): print(f"rescheduled {moved} QCOM kernels by dependency level")
+  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_BATCH_HEAD", 1):
+    from extra.gemm.qcom_openpilot_batch_head import batch_head
+    if (combined:=batch_head(run_onnx_jit)): print(f"batched {combined} groups of QCOM head kernels")
+  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_INPUT_PACK", 1):
+    from extra.gemm.qcom_openpilot_input_pack import patch_input_pack
+    if (patched:=patch_input_pack(run_onnx_jit)): print(f"vectorized {patched} QCOM input kernel")
   with open(OUTPUT, "wb") as f:
     pickle.dump(run_onnx_jit, f)
   mdl_sz = os.path.getsize(onnx_file)
@@ -72,7 +96,7 @@ def compile(onnx_file):
   print(f"mdl size is {mdl_sz/1e6:.2f}M")
   print(f"pkl size is {pkl_sz/1e6:.2f}M")
   print("**** compile done ****")
-  return inputs, test_val
+  return run_onnx_jit, inputs, test_val
 
 def test_vs_compile(run, inputs, test_val=None):
 
@@ -142,9 +166,10 @@ if __name__ == "__main__":
     test_vs_compile(pickle_loaded, inputs)
   else:
     onnx_file = fetch(OPENPILOT_MODEL)
-    inputs, outputs = compile(onnx_file)
+    pickle_loaded, inputs, outputs = compile(onnx_file)
 
-    with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f)
+    if OUTPUT != os.devnull:
+      with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f)
 
     test_vs_compile(pickle_loaded, inputs, outputs)
     if getenv("SELFTEST"):

extra/gemm/QUALCOMM.md

@@ -0,0 +1,973 @@
+# 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/dsp_int8_gemm.py

@@ -0,0 +1,141 @@
+#!/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

@@ -0,0 +1,440 @@
+"""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/openpilot_onnx_precision_scan.py

@@ -0,0 +1,58 @@
+#!/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

@@ -0,0 +1,266 @@
+#!/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

@@ -0,0 +1,277 @@
+#!/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

@@ -0,0 +1,111 @@
+#!/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

@@ -0,0 +1,277 @@
+#!/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

@@ -0,0 +1,440 @@
+#!/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

@@ -0,0 +1,106 @@
+#!/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

@@ -0,0 +1,90 @@
+#!/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

@@ -0,0 +1,39 @@
+#!/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

@@ -0,0 +1,64 @@
+#!/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

@@ -0,0 +1,67 @@
+#!/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

@@ -0,0 +1,234 @@
+#!/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

@@ -0,0 +1,59 @@
+#!/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

@@ -0,0 +1,71 @@
+#!/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

@@ -0,0 +1,28 @@
+#!/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

@@ -0,0 +1,48 @@
+#!/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

@@ -0,0 +1,55 @@
+#!/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

@@ -0,0 +1,111 @@
+#!/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

@@ -0,0 +1,75 @@
+#!/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

@@ -0,0 +1,40 @@
+#!/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

@@ -0,0 +1,166 @@
+#!/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

@@ -0,0 +1,50 @@
+#!/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

@@ -0,0 +1,86 @@
+#!/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

@@ -0,0 +1,170 @@
+#!/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

@@ -0,0 +1,79 @@
+#!/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

@@ -0,0 +1,109 @@
+#!/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

@@ -0,0 +1,135 @@
+#!/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

@@ -0,0 +1,101 @@
+#!/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

@@ -0,0 +1,91 @@
+#!/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

@@ -0,0 +1,93 @@
+#!/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

@@ -0,0 +1,164 @@
+#!/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

@@ -0,0 +1,50 @@
+#!/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

@@ -0,0 +1,54 @@
+#!/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

@@ -0,0 +1,99 @@
+#!/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

@@ -0,0 +1,99 @@
+#!/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

@@ -0,0 +1,81 @@
+#!/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

@@ -0,0 +1,152 @@
+#!/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

@@ -0,0 +1,110 @@
+#!/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

@@ -0,0 +1,78 @@
+#!/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

@@ -0,0 +1,79 @@
+#!/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

@@ -0,0 +1,54 @@
+#!/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

@@ -0,0 +1,111 @@
+#!/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

@@ -0,0 +1,103 @@
+#!/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

@@ -0,0 +1,133 @@
+#!/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

@@ -0,0 +1,21 @@
+#!/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

@@ -0,0 +1,333 @@
+#!/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

@@ -0,0 +1,130 @@
+#!/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

@@ -0,0 +1,317 @@
+#!/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()

extra/gemm/qcom_openpilot_geometry.py

@@ -0,0 +1,43 @@
+#!/usr/bin/env python3
+"""Apply validated equivalent launch geometries to OpenPilot QCOM kernels."""
+import argparse, pickle
+from dataclasses import replace
+
+from tinygrad.engine.jit import create_graph_call
+from tinygrad.uop.ops import Ops
+from extra.gemm.qcom_ir3_matmul_patch import plain_name
+
+
+GEOMETRIES = {
+  "r_32_192_4_4_64_4": ((16, 2, 1), (12, 16, 1)),
+  "r_8_384_4_4_128_4": ((48, 1, 1), (8, 8, 1)),
+  "r_64_32_16_4_4_6_3_3_4": ((4, 2, 32), (4, 16, 2)),
+}
+
+
+def patch_geometry(model) -> int:
+  outer = model.captured.linear.src[0]
+  batch, patched = list(outer.src[0].src[0].src), 0
+  for index, call in enumerate(batch):
+    if call.op is not Ops.CALL or call.src[0].op is not Ops.PROGRAM: continue
+    if (geometry := GEOMETRIES.get(plain_name(call.src[0].arg.name))) is None: continue
+    program = call.src[0].replace(arg=replace(call.src[0].arg, global_size=geometry[0], local_size=geometry[1]))
+    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:
+  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_geometry(model))
+  with open(args.output, "wb") as f: pickle.dump(model, f)
+
+
+if __name__ == "__main__": main()

extra/gemm/qcom_openpilot_geometry_sweep.py

@@ -0,0 +1,52 @@
+#!/usr/bin/env python3
+"""Sweep equivalent launch geometries for global-ID-only OpenPilot kernels."""
+import argparse, itertools, pickle, statistics
+
+from tinygrad import Tensor
+from tinygrad.engine.jit import _prepare_jit_inputs
+from tinygrad.engine.realize import get_runtime, resolve_params
+from tinygrad.uop.ops import Ops
+from extra.gemm.qcom_ir3_matmul_patch import plain_name
+
+
+def divisors(value:int) -> list[int]: return [x for x in range(1, value+1) if value%x == 0]
+
+
+def main() -> None:
+  parser = argparse.ArgumentParser()
+  parser.add_argument("model")
+  parser.add_argument("name")
+  parser.add_argument("--runs", type=int, default=7)
+  parser.add_argument("--max-threads", type=int, default=256)
+  args = parser.parse_args()
+  with open(args.model, "rb") as f: model = pickle.load(f)
+  inputs = {name:Tensor.zeros(*view.shape, dtype=dtype, device=device).contiguous().realize()
+            for name, (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)
+  model(**inputs).numpy()
+  calls = [call for call in model.captured.linear.src[0].src[0].src[0].src
+           if call.op is Ops.CALL and call.src[0].op is Ops.PROGRAM and plain_name(call.src[0].arg.name) == args.name]
+  call, program = calls[0], calls[0].src[0]
+  source = program.src[2].arg
+  if any(token in source for token in ("get_group_id", "get_local_id", "barrier(", "__local")):
+    raise ValueError("kernel is not global-ID-only")
+  resolved = resolve_params(call, tuple(input_uops))
+  bufs = [u.buffer.ensure_allocated()._buf for u in resolved]
+  runtime = get_runtime(resolved[0].device, program)
+  total = tuple(int(g*l) for g,l in zip(program.arg.global_size, program.arg.local_size))
+  candidates = []
+  for local in itertools.product(*(divisors(x) for x in total)):
+    threads = local[0]*local[1]*local[2]
+    if threads > args.max_threads or threads < 16: continue
+    global_size = tuple(total[i]//local[i] for i in range(3))
+    candidates.append((local, global_size))
+  results = []
+  for local, global_size in candidates:
+    for _ in range(2): runtime(*bufs, global_size=global_size, local_size=local, vals=(), wait=True)
+    times = [runtime(*bufs, global_size=global_size, local_size=local, vals=(), wait=True)*1e3 for _ in range(args.runs)]
+    results.append((statistics.median(times), min(times), local, global_size))
+  for median, best, local, global_size in sorted(results)[:30]:
+    print(f"median_ms={median:.5f} best_ms={best:.5f} local={local} global={global_size}")
+
+
+if __name__ == "__main__": main()

extra/gemm/qcom_openpilot_head_gemv.py

@@ -0,0 +1,122 @@
+#!/usr/bin/env python3
+"""Transpose head-GEMV weight microtiles for vector FP32 accumulation."""
+import argparse, itertools, os, pickle, struct
+from dataclasses import replace
+
+import numpy as np
+
+from tinygrad import Device, dtypes
+from tinygrad.device import Buffer
+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
+from extra.gemm.ir3asm import BR, ISAM_F32, MAD_F32, MOV_F32, NOP, inject
+
+TARGET = "r_128_16_4_32_4_batch2"
+
+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(16,1,1)))
+__kernel void head_gemv_2048_512(write_only image2d_t O, read_only image2d_t A,
+                                 read_only image2d_t W, read_only image2d_t B) {
+  int out4=get_group_id(0), lid=get_local_id(0);
+  float4 z=(float4)(0.0f);
+  for (int r=0;r<32;r++) {
+    float4 a=read_imagef(A,smp,(int2)(lid*32+r,0));
+    int x=lid*128+r*4;
+    float4 w0=read_imagef(W,smp,(int2)(x+0,out4));
+    float4 w1=read_imagef(W,smp,(int2)(x+1,out4));
+    float4 w2=read_imagef(W,smp,(int2)(x+2,out4));
+    float4 w3=read_imagef(W,smp,(int2)(x+3,out4));
+    z+=(float4)(a.x)*w0; z+=(float4)(a.y)*w1;
+    z+=(float4)(a.z)*w2; z+=(float4)(a.w)*w3;
+  }
+  __local float4 partial[16];
+  partial[lid]=z;
+  barrier(CLK_LOCAL_MEM_FENCE);
+  if (lid==0) {
+    z=(float4)(0.0f);
+    for (int i=0;i<16;i++) z+=partial[i];
+    z+=read_imagef(B,smp,(int2)(out4,0));
+    z=select((float4)(0.0f),convert_float4(convert_half4(z)),isgreater(z,(float4)(0.0f)));
+    write_imagef(O,(int2)(out4,0),z);
+  }
+}"""
+
+
+def packed_weight(weight:UOp) -> UOp:
+  original = np.asarray(weight.buffer.numpy()).reshape(128, 2112, 4)
+  packed = np.array(original, copy=True)
+  for out4 in range(128):
+    for lid in range(16):
+      for r in range(32):
+        x = lid*128+r*4
+        packed[out4, x:x+4] = original[out4, x:x+4].T
+  return UOp.from_buffer(Buffer("QCOM", packed.size, weight.dtype, initial_value=packed.tobytes()))
+
+
+def pack_lib(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) != 127: raise RuntimeError(f"expected 127 head GEMV instructions, got {len(ins)}")
+  init_lo,init_hi = struct.unpack("<II",ins[14])
+  init = struct.pack("<II",init_lo,(init_hi&~0xff)|24|0x300)
+  out = list(ins[:14]) + [init,NOP(),NOP(),NOP()]
+  loop_start = len(out)
+  out += ins[18:21]
+  for coord_ins,reg,coord in zip((21,28,35,42),(7,8,9,10),("r2.y","r2.w","r3.y","r3.w")):
+    out += [ins[coord_ins],NOP(rpt=5),ISAM_F32(f"r{reg}.x",coord,1,0)]
+  for component,weight_reg in zip("xyzw",range(7,11)):
+    if os.getenv("HEAD_SCALAR"):
+      for lane in "xyzw":
+        out.append(MAD_F32(f"r6.{lane}",f"r0.{component}",f"r{weight_reg}.{lane}",f"r6.{lane}",
+                           sy=component=="x" and lane=="x"))
+    else:
+      out.append(MAD_F32("r6.x",f"r0.{component}",f"r{weight_reg}.x","r6.x",rpt=3,r=True,sy=component=="x"))
+  out += ins[49:55]
+  out.append(BR(loop_start-len(out)))
+  out += [ins[56],MOV_F32("r0.x","r6.x",rpt=3,r=True),NOP(),NOP(),NOP()] + ins[61:]
+  out += [NOP()]*(len(ins)-len(out))
+  if len(out) != len(ins): raise RuntimeError(f"packed head GEMV overflow: {len(out)}")
+  fregs,hregs = struct.unpack_from("<II",lib,reg_off+0x14)
+  return inject(lib,image_off,image_size,reg_off,b"".join(out),max(fregs&0x7fffffff,11)|(fregs&0x80000000),hregs)
+
+
+def aux(*specs): return (tuple(((i, dtype, shape),) for i, (dtype, shape) in enumerate(specs)),)
+
+
+def patch_model(model) -> int:
+  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)
+  outer = model.captured.linear.src[0]
+  batch, new_batch, patched = outer.src[0].src[0].src, [], 0
+  lib = pack_lib(Device["QCOM"].compiler.compile_cached(SOURCE))
+  for call in 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:
+      new_batch.append(call)
+      continue
+    template = call.src[0]
+    info = replace(template.arg, name="head_gemv_2048_512", global_size=(128,1,1), local_size=(16,1,1),
+                   globals=(0,1,2,3), outs=(0,), ins=(1,2,3),
+                   aux=aux((dtypes.half,(1,136,4)), (dtypes.half,(1,520,4)),
+                           (dtypes.half,(128,2112,4)), (dtypes.half,(1,128,4))))
+    program = template.replace(arg=info, src=template.src[:2]+(template.src[2].replace(arg=SOURCE), template.src[3].replace(arg=lib)))
+    new_batch += [program.call(call.src[1], call.src[3], packed_weight(call.src[4]), call.src[5]),
+                  program.call(call.src[2], call.src[6], packed_weight(call.src[7]), call.src[8])]
+    patched += 1
+  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"); 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_input_pack.py

@@ -0,0 +1,49 @@
+"""Vectorize the driving-vision uint8 input normalization kernel on QCOM."""
+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 = "E_8192_3_4_2_4"
+SOURCE = r"""#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+__kernel void E_8192_3_4_2_4(write_only image2d_t O,__global uchar *A,__global uchar *B,
+                             __global half *MEAN,__global half *STD) {
+  int c=get_global_id(0),i=get_global_id(1),off=(c<<17)+(i<<2),mc=c<<2;
+  uchar4 a0=vload4(0,A+off),a1=vload4(0,A+off+32768);
+  uchar4 a2=vload4(0,A+off+65536),a3=vload4(0,A+off+98304);
+  uchar4 b0=vload4(0,B+off),b1=vload4(0,B+off+32768);
+  uchar4 b2=vload4(0,B+off+65536),b3=vload4(0,B+off+98304);
+  half4 ma=vload4(0,MEAN+mc),mb=vload4(0,MEAN+mc+12);
+  half4 ia=(half4)(1)/vload4(0,STD+mc),ib=(half4)(1)/vload4(0,STD+mc+12);
+  int x=c+(i&7)*24,y=i>>3;
+  write_imagef(O,(int2)(x,y),convert_float4(((half4)(a0.x,a1.x,a2.x,a3.x)-ma)*ia));
+  write_imagef(O,(int2)(x+3,y),convert_float4(((half4)(b0.x,b1.x,b2.x,b3.x)-mb)*ib));
+  write_imagef(O,(int2)(x+6,y),convert_float4(((half4)(a0.y,a1.y,a2.y,a3.y)-ma)*ia));
+  write_imagef(O,(int2)(x+9,y),convert_float4(((half4)(b0.y,b1.y,b2.y,b3.y)-mb)*ib));
+  write_imagef(O,(int2)(x+12,y),convert_float4(((half4)(a0.z,a1.z,a2.z,a3.z)-ma)*ia));
+  write_imagef(O,(int2)(x+15,y),convert_float4(((half4)(b0.z,b1.z,b2.z,b3.z)-mb)*ib));
+  write_imagef(O,(int2)(x+18,y),convert_float4(((half4)(a0.w,a1.w,a2.w,a3.w)-ma)*ia));
+  write_imagef(O,(int2)(x+21,y),convert_float4(((half4)(b0.w,b1.w,b2.w,b3.w)-mb)*ib));
+}"""
+
+
+def patch_input_pack(jit) -> int:
+  outer = jit.captured.linear.src[0]
+  batch = outer.src[0].src[0].src
+  lib, new_batch, replaced = None, [], 0
+  for call in batch:
+    name = plain_name(call.src[0].arg.name) if call.op is Ops.CALL and call.src[0].op is Ops.PROGRAM else ""
+    if name == TARGET:
+      if lib is None: lib = Device["QCOM"].compiler.compile_cached(SOURCE)
+      program = call.src[0]
+      program = program.replace(arg=replace(program.arg, global_size=(1, 64, 1), local_size=(3, 128, 1)),
+                                src=program.src[:2] +
+                                (program.src[2].replace(arg=SOURCE), program.src[3].replace(arg=lib)))
+      call, replaced = call.replace(src=(program, *call.src[1:])), replaced+1
+    new_batch.append(call)
+  if replaced:
+    jit.captured._linear = jit.captured.linear.substitute({outer:create_graph_call(new_batch)}, walk=True)
+    jit.captured.__dict__.pop("linear", None)
+  return replaced

extra/gemm/qcom_openpilot_inverse_full_rpt.py

@@ -0,0 +1,124 @@
+#!/usr/bin/env python3
+"""Pack all four FP32 accumulator vectors in the OpenPilot inverse projection."""
+import argparse, hashlib, pickle, struct
+
+from tinygrad.engine.jit import create_graph_call
+from tinygrad.uop.ops import Ops
+from extra.gemm.ir3asm import BR, MAD_F32, MOV_F32, NOP
+from extra.gemm.qcom_ir3_matmul_patch import plain_name
+
+TARGET="r_32_64_4_4_192_4"
+INVERSE_W_TARGETS={"r_512_16_4_4_48_4","r_128_32_4_4_96_4"}
+SAFE_DONORS={"d4c281a1","1fe26758","e34e7e58"}
+
+
+def replace_src2(ins:bytes, src2:int) -> bytes:
+  lo,hi=struct.unpack("<II",ins)
+  return struct.pack("<II",(lo&0xff00ffff)|(src2<<16),hi)
+
+
+def replace_low_src(ins:bytes, src:int) -> bytes:
+  lo,hi=struct.unpack("<II",ins)
+  return struct.pack("<II",(lo&0xffffff00)|src,hi)
+
+
+def pack_inverse_full(lib:bytes) -> bytes:
+  off,size=struct.unpack_from("<I",lib,0xc0)[0],struct.unpack_from("<I",lib,0x100)[0]
+  instrs=[lib[i:i+8] for i in range(off,off+size,8)]
+  if len(instrs)!=175: raise RuntimeError(f"expected 175 inverse instructions, got {len(instrs)}")
+  # Move loop control from r13.x into the existing r12.w zero register. This
+  # makes r13-r16 four contiguous accumulator vectors without growing the
+  # shader's declared register file.
+  out=instrs[:11]
+  for acc in ("r13.x","r14.x","r15.x","r16.x"): out.append(MOV_F32(acc,"r12.w",rpt=3))
+  loop_start=len(out)
+  body=list(instrs[16:32])
+  body[0]=replace_low_src(body[0],51)   # add r8.z, r12.w, 192
+  body[1]=replace_low_src(body[1],51)   # add r9.x, r12.w, 384
+  body[2]=replace_src2(body[2],51)      # add r9.z, c28.y, r12.w
+  body[3]=replace_low_src(body[3],51)   # mov r10.x, r12.w
+  out+=body
+  for component,weight in zip("xyzw",("r5.x","r2.x","r3.x","r4.x")):
+    out.append(MAD_F32("r13.x",f"r7.{component}",weight,"r13.x",rpt=3,r=True,sy=component=="x"))
+  out+=instrs[48:60]
+  control=list(instrs[60:65])
+  control[0]=replace_low_src(control[0],51)  # increment r12.w
+  control[2]=replace_low_src(control[2],51)  # compare r12.w
+  control[3]=MOV_F32("r12.w","r0.x")
+  out+=control
+  out.append(BR(loop_start-len(out)))
+  tail=list(instrs[66:131])
+  # The first residual moved from r12.w to r13.x; y/z/w were already in r13.
+  tail[90-66]=replace_src2(tail[90-66],52)
+  out+=tail
+  out += [NOP()]*(len(instrs)-len(out))
+  if len(out)!=len(instrs): raise RuntimeError(f"packed image has {len(out)} instructions")
+  return lib[:off]+b"".join(out)+lib[off+size:]
+
+
+def with_fregs(lib:bytes, count:int) -> bytes:
+  out=bytearray(lib)
+  regoff=struct.unpack_from("<I",out,0x34)[0]+0x14
+  regs=struct.unpack_from("<I",out,regoff)[0]
+  struct.pack_into("<I",out,regoff,(regs&0x80000000)|max(regs&0x7fffffff,count))
+  return bytes(out)
+
+
+def pack_inverse_w_full(lib:bytes) -> bytes:
+  off,size=struct.unpack_from("<I",lib,0xc0)[0],struct.unpack_from("<I",lib,0x100)[0]
+  ins=[lib[i:i+8] for i in range(off,off+size,8)]
+  if len(ins) not in (174,178): raise RuntimeError(f"expected 174/178 inverse-W instructions, got {len(ins)}")
+  out=ins[:19]+[MOV_F32("r17.x","r12.z",rpt=3)]
+  loop_start=len(out)
+  out+=ins[19:35]
+  for component,weight in zip("xyzw",("r5.x","r2.x","r3.x","r4.x")):
+    out.append(MAD_F32("r17.x",f"r7.{component}",weight,"r17.x",rpt=3,r=True,sy=component=="x"))
+  out+=ins[51:68]
+  out.append(BR(loop_start-len(out)))
+  tail=list(ins[69:])
+  mapping={50:68,52:69,53:70,54:71}
+  first_store=next(i for i,x in enumerate(tail) if struct.unpack_from("<I",x,4)[0]>>24==0xc0)
+  for i in range(first_store):
+    lo,_=struct.unpack("<II",tail[i])
+    src2=(lo>>16)&0xff
+    if src2 in mapping: tail[i]=replace_src2(tail[i],mapping[src2])
+  out+=tail
+  out += [NOP()]*(len(ins)-len(out))
+  return with_fregs(lib[:off]+b"".join(out)+lib[off+size:],18)
+
+
+def patch_model(model,names:set[str]|None=None) -> int:
+  outer=model.captured.linear.src[0]
+  batch=list(outer.src[0].src[0].src)
+  cache,patched={},0
+  for index,call in enumerate(batch):
+    if call.op is not Ops.CALL or call.src[0].op is not Ops.PROGRAM: continue
+    name=plain_name(call.src[0].arg.name)
+    if name not in (names if names is not None else {TARGET}|INVERSE_W_TARGETS): continue
+    program=call.src[0]
+    old=program.src[3].arg
+    # These transforms relocate fixed compiler registers. A different QCOM compiler allocation can
+    # have the same instruction count but different live values and must not be patched by index.
+    if hashlib.sha1(old).hexdigest()[:8] not in SAFE_DONORS: continue
+    if old not in cache:
+      cache[old]=pack_inverse_full(old) if name==TARGET else pack_inverse_w_full(old)
+    program=program.replace(src=program.src[:3]+(program.src[3].replace(arg=cache[old]),))
+    batch[index]=call.replace(src=(program,*call.src[1:]))
+    patched+=1
+  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("--names",help="comma-separated program families")
+  args=ap.parse_args()
+  with open(args.input,"rb") as f:model=pickle.load(f)
+  print("patched",patch_model(model,set(args.names.split(",")) if args.names else None))
+  with open(args.output,"wb") as f:pickle.dump(model,f)
+
+
+if __name__=="__main__":main()

extra/gemm/qcom_openpilot_large_fp16.py

@@ -0,0 +1,106 @@
+#!/usr/bin/env python3
+"""Replace the large 2048x192x64 vision projection with a checked FP16 GEMM."""
+import argparse
+import pickle
+
+import numpy as np
+
+from tinygrad import Device, dtypes
+from tinygrad.device import Buffer
+from tinygrad.engine.jit import create_graph_call
+from tinygrad.uop.ops import Ops, UOp
+from extra.gemm import qcom_8x4_gemm as q8
+from extra.gemm.ir3asm import get_envelope, inject
+from extra.gemm.qcom_ir3_matmul_patch import plain_name
+from extra.gemm.qcom_openpilot_graph import build_program
+
+TARGET = "r_512_48_4_4_16_4"
+M, N, K, PAD_N, STRIDE = 2048, 192, 64, 256, 1024
+
+PACK = 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 pack_large(__global half *O,read_only image2d_t A) {
+  int t=get_global_id(0),row=t/16,k4=t-row*16,idx1=row>>2,block=row&3;
+  int x=(idx1&15)*68+k4+block*17,y=idx1>>4;
+  vstore4(read_imageh(A,smp,(int2)(x,y)),0,O+t*4);
+}"""
+
+EPI = 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 epi_large(write_only image2d_t O,read_only image2d_t B,__global half *C) {
+  int t=get_global_id(0),row=t/48,col=t-row*48,idx1=row>>2,block=row&3;
+  float4 z=convert_float4(vload4(0,C+row*1024+col*4));
+  write_imagef(O,(int2)(col+block*48,idx1),gelu(z+read_imagef(B,smp,(int2)(col,0))));
+}"""
+
+
+def copyin(buffer, array: np.ndarray) -> None:
+  raw = memoryview(np.ascontiguousarray(array)).cast("B")
+  if hasattr(buffer, "copyin"):
+    buffer.copyin(raw)
+  else:
+    buffer.copy_from(Buffer("PYTHON", buffer.size, buffer.dtype, opaque=raw))
+
+
+def patch_model(model) -> int:
+  outer = model.captured.linear.src[0]
+  batch, replacements = list(outer.src[0].src[0].src), {}
+  dev = Device["QCOM"]
+  pack_lib, epi_lib = dev.compiler.compile_cached(PACK), dev.compiler.compile_cached(EPI)
+  q8.M, q8.N, q8.K, q8.K4 = M, STRIDE, K, K//4
+  envelope, image_offset, image_size, register_offset = get_envelope(dev, q8.make_donor_src8(4, 128))
+  shader, hregs, _fregs, _ = q8.build_8x8_split_a_unroll_shader(
+    dev, 128, k_unroll=8, b_coord_delay=0, fast_coords=True,
+    prefetch_next_b=True, add256_store_mode="tight", high_a=True, split_low_pairs=True)
+  gemm_lib = inject(envelope, image_offset, image_size, register_offset, shader, fregs=10, hregs=hregs)
+  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
+    template = call.src[0]
+    activation = UOp.new_buffer("QCOM", M*K, dtypes.half)
+    activation.buffer.ensure_allocated()
+    weight = UOp.new_buffer("QCOM", K*PAD_N, dtypes.half)
+    weight.buffer.ensure_allocated()
+    original = call.src[3].buffer.numpy().view(np.float16).reshape(48, 72, 4)
+    packed_weight = np.zeros((K, PAD_N//4, 4), dtype=np.float16)
+    packed_weight[:, :48] = original[:, :K].transpose(1, 0, 2)
+    copyin(weight.buffer, packed_weight)
+    temporary = UOp.new_buffer("QCOM", M*STRIDE, dtypes.half)
+    temporary.buffer.ensure_allocated()
+    pack = build_program(template, "pack_large", PACK, pack_lib, (M*(K//4)//128, 1, 1), (128, 1, 1),
+                         ((dtypes.half, (M*K,)), (dtypes.half, (32, 1088, 4))), (0,), (1,))
+    gemm = build_program(template, "gemm_h", "checked FP16 8x8 GEMM", gemm_lib,
+                         (PAD_N//256, M//32, 1), (128, 1, 1),
+                         ((dtypes.half, (M, K//4, 4)), (dtypes.half, (K, PAD_N//4, 4)),
+                          (dtypes.half, (M*STRIDE,))), (2,), (0, 1))
+    epi = build_program(template, "epi_large", EPI, epi_lib, (M*(N//4)//128, 1, 1), (128, 1, 1),
+                        ((dtypes.half, (512, 192, 4)), (dtypes.half, (1, 48, 4)),
+                         (dtypes.half, (M*STRIDE,))), (0,), (1, 2))
+    replacements[index] = (pack.call(activation, call.src[2]), gemm.call(activation, weight, temporary),
+                           epi.call(call.src[1], call.src[4], temporary))
+  if replacements:
+    new_batch = [item for index, call in enumerate(batch) for item 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 len(replacements)
+
+
+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("patched", patch_model(model))
+  with open(args.output, "wb") as f:
+    pickle.dump(model, f)
+
+
+if __name__ == "__main__":
+  main()

extra/gemm/qcom_openpilot_level_schedule.py

@@ -0,0 +1,47 @@
+#!/usr/bin/env python3
+"""Group ready OpenPilot graph calls by program family without crossing dependency levels."""
+import argparse, pickle
+from collections import defaultdict
+
+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 schedule_levels(model) -> int:
+  outer=model.captured.linear.src[0]
+  batch=list(outer.src[0].src[0].src)
+  writer,levels={},{}
+  grouped=defaultdict(list)
+  for sequence,call in enumerate(batch):
+    if call.op is not Ops.CALL or call.src[0].op is not Ops.PROGRAM:
+      grouped[sequence].append((sequence,call))
+      continue
+    deps={writer[arg] for arg in call.src[1:] if arg in writer}
+    level=1+max((levels[dep] for dep in deps),default=-1)
+    levels[sequence]=level
+    grouped[level].append((sequence,call))
+    for output in call.src[0].arg.outs: writer[call.src[output+1]]=sequence
+  scheduled=[]
+  moved=0
+  for entries in grouped.values():
+    ordered=sorted(entries,key=lambda item:(plain_name(item[1].src[0].arg.name),item[0]))
+    scheduled.extend(call for _index,call in ordered)
+    moved+=sum(old_index!=entries[new_index][0] for new_index,(old_index,_call) in enumerate(ordered))
+  if scheduled != batch:
+    model.captured._linear=model.captured.linear.substitute({outer:create_graph_call(scheduled)},walk=True)
+    model.captured.__dict__.pop("linear",None)
+  return moved
+
+
+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("moved",schedule_levels(model))
+  with open(args.output,"wb") as f:pickle.dump(model,f)
+
+
+if __name__=="__main__":main()

extra/gemm/qcom_openpilot_levels.py

@@ -0,0 +1,25 @@
+#!/usr/bin/env python3
+"""Print dependency-level parallelism in a captured OpenPilot graph."""
+import argparse, pickle
+from collections import defaultdict
+
+from tinygrad.uop.ops import Ops
+from extra.gemm.qcom_ir3_matmul_patch import plain_name
+
+
+def main() -> None:
+  ap = argparse.ArgumentParser(); ap.add_argument("model"); args = ap.parse_args()
+  with open(args.model, "rb") as f: model = pickle.load(f)
+  batch = model.captured.linear.src[0].src[0].src[0].src
+  writer, level, groups = {}, {}, defaultdict(list)
+  for index, call in enumerate(batch):
+    if call.op is not Ops.CALL or call.src[0].op is not Ops.PROGRAM: continue
+    deps = {writer[x] for x in call.src[1:] if x in writer}
+    level[index] = 1 + max((level[x] for x in deps), default=-1)
+    groups[level[index]].append((index, plain_name(call.src[0].arg.name), call.src[0].arg.global_size, call.src[0].arg.local_size))
+    for out in call.src[0].arg.outs: writer[call.src[out+1]] = index
+  for lev, calls in groups.items():
+    if len(calls) > 1: print(f"level={lev} calls={len(calls)}", *calls, sep="\n  ")
+
+
+if __name__ == "__main__": main()

extra/gemm/qcom_openpilot_native_fp16.py

@@ -0,0 +1,52 @@
+#!/usr/bin/env python3
+"""Experimental native-FP16 accumulator rewrite for driving_vision kernels."""
+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 half_acc_source(source:str) -> str:
+  # buf0 is the matrix accumulator in these generated reduction kernels.  val0..7
+  # are the two four-vector matrix operands; later values belong to the FP32
+  # bias/residual epilogue and intentionally remain float.
+  source = source.replace("float buf0[16];", "half buf0[16];")
+  source = source.replace("float buf0[4];", "half buf0[4];")
+  for index in range(8):
+    source = re.sub(fr"float4 val{index} = read_imagef\(", fr"half4 val{index} = read_imageh(", source)
+  return source
+
+
+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: continue
+    program = call.src[0]
+    if plain_name(program.arg.name) not in names or not any(x in program.src[2].arg for x in ("float buf0[16];", "float buf0[4];")): continue
+    source = half_acc_source(program.src[2].arg)
+    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
+
+
+def main() -> None:
+  parser = argparse.ArgumentParser()
+  parser.add_argument("input")
+  parser.add_argument("output")
+  parser.add_argument("--names", required=True, help="comma-separated exact display names")
+  args = parser.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)
+
+
+if __name__ == "__main__": main()

extra/gemm/qcom_openpilot_pack_conv_weights.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python3
+"""Prepack static weights for the slow stride-2 openpilot 7x7 convolutions."""
+import argparse, pickle, re
+
+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, UOp
+from extra.gemm.qcom_ir3_matmul_patch import plain_name
+
+TARGETS={"r_16_8_16_2_4_4_7_7", "r_8_4_32_2_4_4_7_7", "r_4_2_64_2_4_4_7_7"}
+
+
+def packed_source(source:str) -> str:
+  start=source.index("    half val0")
+  end=source.index("    int alu20",start)
+  weight_name=re.search(r"__global half\* (data2_\d+)",source).group(1)  # type: ignore[union-attr]
+  block="    int wp=(((alu0*2+alu2)*7+Ridx0)*28);\n"+"\n".join(
+    f"    float4 w{i}=convert_float4(vload4(0,{weight_name}+wp+{i*4}));" for i in range(7))+"\n"
+  source=source[:start]+block+source[end:]
+  accum_start=source.index("    *(buf0+0)",start)
+  loop_prefix=source[start:accum_start]
+  casts=re.findall(r"    float (cast\d+) = \(\(float\)\(val(\d+)\)\);\n",loop_prefix)
+  assert len(casts) == 28
+  source=source[:start]+re.sub(r"    float cast\d+ = \(\(float\)\(val\d+\)\);\n", "", loop_prefix)+source[accum_start:]
+  mapping={cast:("w0.x" if int(val) == 27 else f"w{int(val)%7+1}.x" if int(val) < 6 else
+                       f"w{(int(val)-6)%7}.{'yzw'[(int(val)-6)//7]}") for cast,val in casts}
+  for old,new in sorted(mapping.items(),key=lambda item:-len(item[0])):
+    source=re.sub(rf"\b{old}\b",new,source)
+  return source
+
+
+def pack_weight_buffer(weight:UOp) -> UOp:
+  original=np.asarray(weight.buffer.numpy()).reshape(-1)
+  outputs=original.size//896
+  assert outputs*896 == original.size
+  packed=np.empty((outputs,2,7,7,4),dtype=np.float16)
+  for output in range(outputs):
+    for parity in range(2):
+      for row in range(7):
+        base=output*896+parity+row*28
+        for tap in range(7):
+          for component in range(4): packed[output,parity,row,tap,component]=original[base+component*224+tap*4]
+  buf=Buffer("QCOM",packed.size,weight.dtype,initial_value=bytearray(packed.tobytes()))
+  return UOp.from_buffer(buf)
+
+
+def patch_conv(model) -> int:
+  outer=model.captured.linear.src[0]
+  batch=outer.src[0].src[0].src
+  new_batch=[]
+  replaced=0
+  for call in batch:
+    name=plain_name(call.src[0].arg.name) if call.op is Ops.CALL and call.src[0].op is Ops.PROGRAM else ""
+    if name in TARGETS:
+      program=call.src[0]
+      source=packed_source(program.src[2].arg)
+      lib=Device["QCOM"].compiler.compile_cached(source)
+      program=program.replace(src=program.src[:2]+(program.src[2].replace(arg=source),program.src[3].replace(arg=lib)))
+      call=call.replace(src=(program,call.src[1],call.src[2],pack_weight_buffer(call.src[3]),*call.src[4:]))
+      replaced+=1
+    new_batch.append(call)
+  model.captured._linear=model.captured.linear.substitute({outer:create_graph_call(new_batch)},walk=True)
+  model.captured.__dict__.pop("linear",None)
+  return replaced
+
+
+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("patched",patch_conv(model))
+  with open(args.output,"wb") as f:pickle.dump(model,f)
+
+
+if __name__ == "__main__":main()

extra/gemm/qcom_openpilot_pack_firstconv_weights.py

@@ -0,0 +1,89 @@
+#!/usr/bin/env python3
+"""Repack driving-vision first-convolution weights for output-thread locality."""
+import argparse, itertools, pickle, 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, UOp
+from extra.gemm.qcom_ir3_matmul_patch import plain_name
+from extra.gemm.ir3asm import BR, MAD_F32, NOP
+
+TARGET = "r_64_32_16_4_4_6_3_3_4"
+
+
+def packed_weight(weight:UOp, group4:bool=False) -> UOp:
+  original = np.asarray(weight.buffer.numpy()).reshape(-1)
+  assert original.size == 16*216*4
+  packed = np.empty_like(original)
+  for out4 in range(16):
+    for ky in range(3):
+      for ic in range(6):
+        for kx in range(3):
+          old = (out4*216 + ky*72 + ic*12 + kx*4)*4
+          tap = ky*18 + ic*3 + kx
+          new_pixel = ((out4//4)*864 + tap*16 + (out4%4)*4) if group4 else (tap*16 + out4)*4
+          new = new_pixel*4
+          packed[new:new+16] = original[old:old+16]
+  return UOp.from_buffer(Buffer("QCOM", packed.size, weight.dtype, initial_value=packed.tobytes()))
+
+
+def repeat_pack(lib:bytes) -> bytes:
+  image_offset, image_size = struct.unpack_from("<I", lib, 0xc0)[0], struct.unpack_from("<I", lib, 0x100)[0]
+  instrs = [lib[image_offset+i:image_offset+i+8] for i in range(0, image_size, 8)]
+  if len(instrs) != 261: raise RuntimeError(f"expected 261 repacked first-conv instructions, got {len(instrs)}")
+  out = instrs[:58]
+  first = True
+  for component, weight in zip("xyzw", ("r5", "r2", "r3", "r4")):
+    out.append(MAD_F32("r11.x", f"r7.{component}", f"{weight}.x", "r11.x", sy=first, r=True))
+    first = False
+    out.append(MAD_F32("r12.y", f"r7.{component}", f"{weight}.y", "r12.y", rpt=2, r=True))
+  for component, weight in zip("xyzw", ("r5.x", "r2.x", "r3.x", "r4.x")):
+    for acc, activation in zip(("r13.x", "r14.x", "r15.x"), ("r6", "r1", "r0")):
+      out.append(MAD_F32(acc, f"{activation}.{component}", weight, acc, rpt=3, r=True))
+  out += instrs[122:128]
+  out.append(BR(29-len(out)))
+  out += instrs[129:]
+  out[91] = BR(24-91)
+  out[98] = BR(22-98)
+  if len(out) > len(instrs): raise RuntimeError(f"packed shader grew to {len(out)} instructions")
+  out += [NOP()]*(len(instrs)-len(out))
+  return lib[:image_offset]+b"".join(out)+lib[image_offset+image_size:]
+
+
+def patch_model(model, rpt:bool=False, group4:bool=False) -> int:
+  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)
+  outer = model.captured.linear.src[0]
+  batch, patched = list(outer.src[0].src[0].src), 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, source = call.src[0], call.src[0].src[2].arg
+    old = "int alu18 = ((Ridx0*12)+(Ridx3<<2)+(Ridx2*72)+(idx0*216));"
+    new = ("int alu18 = ((idx0>>2)*864+(Ridx2*18+Ridx0*3+Ridx3)*16+(idx0&3)*4);" if group4 else
+           "int alu18 = (((Ridx2*18+Ridx0*3+Ridx3)*16+idx0)*4);")
+    if old not in source: raise RuntimeError("unexpected first-convolution source")
+    source = source.replace(old, new)
+    lib = Device["QCOM"].compiler.compile(source)
+    if rpt: lib = repeat_pack(lib)
+    program = program.replace(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], call.src[2], packed_weight(call.src[3], group4), *call.src[4:]))
+    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("--rpt", action="store_true")
+  ap.add_argument("--group4", action="store_true"); args=ap.parse_args()
+  with open(args.input,"rb") as f: model=pickle.load(f)
+  print("patched",patch_model(model, args.rpt, args.group4))
+  with open(args.output,"wb") as f: pickle.dump(model,f)
+
+
+if __name__ == "__main__": main()

extra/gemm/qcom_openpilot_profile_calls.py

@@ -0,0 +1,89 @@
+#!/usr/bin/env python3
+"""Time representative compiled openpilot kernels outside the graph."""
+import argparse, hashlib, pickle, statistics
+
+from tinygrad import Tensor
+from tinygrad.device import Buffer
+from tinygrad.engine.jit import _prepare_jit_inputs
+from tinygrad.engine.realize import get_runtime, resolve_params
+from tinygrad.uop.ops import Ops
+from extra.gemm.qcom_ir3_matmul_patch import plain_name
+
+
+def main() -> None:
+  parser = argparse.ArgumentParser()
+  parser.add_argument("model")
+  parser.add_argument("reference", nargs="?")
+  parser.add_argument("--runs", type=int, default=10)
+  parser.add_argument("--unique", action="store_true", help="separate programs with the same display name")
+  parser.add_argument("--name", help="only profile programs with this display name")
+  parser.add_argument("--details", action="store_true", help="show launch geometry and register metadata")
+  parser.add_argument("--list-only", action="store_true", help="list program families without running the model")
+  parser.add_argument("--no-warm", action="store_true", help="profile captured calls without first executing the full graph")
+  parser.add_argument("--output-dtype", help="only profile calls with an output dtype whose name contains this text")
+  parser.add_argument("--dry-bufs", action="store_true", help="print resolved buffer/runtime metadata without executing")
+  parser.add_argument("--program-hash", help="only profile this eight-character library SHA1 prefix")
+  parser.add_argument("--fresh-bufs", action="store_true", help="execute with fresh dedicated buffers instead of captured buffers")
+  parser.add_argument("--local-size", help="override launch local size, for example 8,4,4")
+  parser.add_argument("--global-size", help="override launch work-group counts, for example 1,8,16")
+  parser.add_argument("--zero-bufs", action="store_true", help="zero all execution buffers before profiling")
+  parser.add_argument("--zero-buffer-indices", help="comma-separated execution buffer indices to zero")
+  parser.add_argument("--all", action="store_true", help="profile every program family instead of only the top entries")
+  args = parser.parse_args()
+  local_size_override = tuple(int(x) for x in args.local_size.split(",")) if args.local_size else None
+  global_size_override = tuple(float(x) for x in args.global_size.split(",")) if args.global_size else None
+  with open(args.model, "rb") as f: model = pickle.load(f)
+  batch = model.captured.linear.src[0].src[0].src[0].src
+  if args.list_only:
+    families: dict[tuple[str, tuple, tuple], int] = {}
+    for call in batch:
+      if call.op is Ops.CALL and call.src[0].op is Ops.PROGRAM:
+        program = call.src[0]
+        key = (plain_name(program.arg.name), tuple(program.arg.global_size), tuple(program.arg.local_size))
+        families[key] = families.get(key, 0) + 1
+    for (name, global_size, local_size), count in sorted(families.items()):
+      print(f"{count:3d} {name} global={global_size} local={local_size}")
+    return
+  inputs = {}
+  for name, (view, _vars, dtype, device) in zip(model.captured.expected_names, model.captured.expected_input_info):
+    inputs[name] = Tensor.zeros(*view.shape, dtype=dtype, device=device).contiguous().realize()
+  input_uops, var_vals, _names, _info = _prepare_jit_inputs((), inputs)
+  if not args.no_warm: model(**inputs).numpy()
+  families: dict[object, list] = {}
+  for call in batch:
+    if call.op is Ops.CALL and call.src[0].op is Ops.PROGRAM:
+      name = plain_name(call.src[0].arg.name)
+      if args.name is not None and name != args.name: continue
+      if args.output_dtype is not None and not any(args.output_dtype in str(call.src[i+1].dtype) for i in call.src[0].arg.outs): continue
+      if args.program_hash is not None and hashlib.sha1(call.src[0].src[3].arg).hexdigest()[:8] != args.program_hash: continue
+      key = (name, hashlib.sha1(call.src[0].src[3].arg).hexdigest()[:8]) if args.unique else name
+      families.setdefault(key, []).append(call)
+  ranked = sorted(families.items(), key=lambda x: sum(int(c.src[0].src[0].arg.estimates.ops) for c in x[1]), reverse=True)
+  for key, calls in (ranked if args.all else ranked[:80 if args.unique else 30]):
+    name = f"{key[0]}#{key[1]}" if isinstance(key, tuple) else key
+    call, program = calls[0], calls[0].src[0]
+    resolved = resolve_params(call, tuple(input_uops))
+    bufs = ([Buffer(resolved[0].device, u.buffer.size, u.dtype).allocate()._buf for u in resolved] if args.fresh_bufs else
+            [u.buffer.ensure_allocated()._buf for u in resolved])
+    runtime = get_runtime(resolved[0].device, program)
+    if args.dry_bufs:
+      print(name, "bufs", [(x.dtype, x.buffer.size, hex(int(b.va_addr)), int(b.va_addr)%4096, b.size) for x,b in zip(resolved, bufs)],
+            "runtime_offs", runtime.buf_offs, "tex", runtime.tex_cnt, "ibo", runtime.ibo_cnt)
+      continue
+    zero_indices = set(range(len(bufs))) if args.zero_bufs else {int(x) for x in (args.zero_buffer_indices or "").split(",") if x}
+    for i in zero_indices: bufs[i].cpu_view().mv[:] = bytes(bufs[i].size)
+    launch_local = local_size_override or program.arg.local_size
+    launch_global = global_size_override or (tuple(program.arg.global_size[i]*program.arg.local_size[i]/launch_local[i] for i in range(3))
+                                                if local_size_override else program.arg.global_size)
+    for _ in range(2): runtime(*bufs, global_size=launch_global, local_size=launch_local, vals=(), wait=True)
+    times = [runtime(*bufs, global_size=launch_global, local_size=launch_local,
+                     vals=(), wait=True)*1e3 for _ in range(args.runs)]
+    ops = int(program.src[0].arg.estimates.ops)
+    med = statistics.median(times)
+    detail = ""
+    if args.details:
+      detail = f" global={program.arg.global_size} local={program.arg.local_size}"
+    print(f"{med:8.4f} ms {ops/med/1e6:8.1f} GFLOP/s x{len(calls):2d} total={med*len(calls):8.3f} ms {name}{detail}")
+
+
+if __name__ == "__main__": main()

extra/gemm/qcom_openpilot_quant_forward.py

@@ -0,0 +1,77 @@
+#!/usr/bin/env python3
+"""Quantize the six vision MLP forward weights with per-output-lane scales."""
+import argparse, itertools, pickle, re
+from dataclasses import replace
+
+import numpy as np
+
+from tinygrad import Device, dtypes
+from tinygrad.device import Buffer
+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_32_192_4_4_64_4"
+
+
+def upload(buffer:Buffer, data:np.ndarray) -> None:
+  raw=memoryview(np.ascontiguousarray(data)).cast("B")
+  if hasattr(buffer, "copyin"): buffer.copyin(raw)
+  else: buffer.copy_from(Buffer("PYTHON", buffer.size, buffer.dtype, opaque=raw))
+
+
+def adapt_aux_dtype(aux, index:int, 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 patch_model(model) -> int:
+  outer=model.captured.linear.src[0]
+  batch=list(outer.src[0].src[0].src)
+  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)
+  cache,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,source=call.src[0],call.src[0].src[2].arg
+    match=re.search(r"(__kernel void \w+\(.*?)(\)) \{",source,re.S)
+    if match is None: raise RuntimeError("kernel signature not found")
+    source=source[:match.start(2)]+", read_only image2d_t qscale"+source[match.start(2):]
+    bias=re.search(r"(float4 (val\d+) = read_imagef\(data3_[^;]+;)",source)
+    if bias is None: raise RuntimeError("bias load not found")
+    source=source[:bias.end()]+"\n  float4 qs = read_imagef(qscale, smp, (int2)(idx0,0));"+source[bias.end():]
+    bias_name=bias.group(2)
+    for lane in range(16):
+      component="xyzw"[lane&3]
+      pattern=fr"\(\(\*\(buf0\+{lane}\)\)\+{bias_name}\.{component}\)"
+      replacement=f"(((*(buf0+{lane}))*qs.{component})+{bias_name}.{component})"
+      source,count=re.subn(pattern,replacement,source)
+      if count != 1: raise RuntimeError(f"output lane {lane} matched {count} times")
+
+    weight=np.asarray(call.src[3].buffer.numpy(),dtype=np.float32).reshape(192,320,4)
+    scale=np.max(np.abs(weight),axis=1,keepdims=True).astype(np.float32)
+    scale[scale == 0]=1.0
+    quantized=np.clip(np.rint(weight/scale*127.0),-127,127).astype(np.int8)
+    qweight=UOp.new_buffer("QCOM",quantized.size,dtypes.int8); qweight.buffer.ensure_allocated(); upload(qweight.buffer,quantized)
+    qscale=UOp.new_buffer("QCOM",scale.size,dtypes.half); qscale.buffer.ensure_allocated(); upload(qscale.buffer,scale.astype(np.float16))
+
+    info=program.arg
+    aux=adapt_aux_dtype(info.aux,2,dtypes.int8)
+    aux=(aux[0]+(((4,dtypes.half,(1,192,4)),),),)
+    info=replace(info,globals=info.globals+(4,),ins=info.ins+(4,),aux=aux)
+    if source not in cache: cache[source]=Device["QCOM"].compiler.compile(source)
+    lib=cache[source]
+    program=program.replace(arg=info,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],call.src[2],qweight,call.src[4],qscale))
+    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"); 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)

extra/gemm/qcom_openpilot_quant_head.py

@@ -0,0 +1,87 @@
+#!/usr/bin/env python3
+"""Experimentally quantize the largest OpenPilot head vector projection."""
+import argparse, itertools, pickle
+from dataclasses import replace
+
+import numpy as np
+
+from tinygrad import Device, dtypes
+from tinygrad.device import Buffer
+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_512_4_256_4"
+
+
+def upload(buffer:Buffer, data:np.ndarray) -> None:
+  raw = memoryview(np.ascontiguousarray(data)).cast("B")
+  if hasattr(buffer, "copyin"): buffer.copyin(raw)
+  else: buffer.copy_from(Buffer("PYTHON", buffer.size, buffer.dtype, opaque=raw))
+
+
+def adapt_aux_dtype(aux, index:int, 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 patch_model(model) -> int:
+  outer = model.captured.linear.src[0]
+  batch = list(outer.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)
+  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, source = call.src[0], call.src[0].src[2].arg
+    signature = "read_only image2d_t data3_1_512_4)"
+    if signature not in source: raise RuntimeError("unexpected head kernel signature")
+    source = source.replace(signature, "read_only image2d_t data3_1_512_4, read_only image2d_t qscale)")
+    bias = "float4 val5 = read_imagef(data3_1_512_4, smp, (int2)(idx0,0));"
+    source = source.replace(bias, bias+"\n  float4 qs = read_imagef(qscale, smp, (int2)(idx0,0));")
+    for lane, component in enumerate("xyzw"):
+      old = f"((*(buf0+{lane}))+val5.{component})"
+      new = f"(((*(buf0+{lane}))*qs.{component})+val5.{component})"
+      if old not in source: raise RuntimeError(f"missing output lane {lane}")
+      source = source.replace(old, new)
+
+    weight = np.asarray(call.src[3].buffer.numpy(), dtype=np.float32).reshape(512, 1088, 4)
+    scale = np.max(np.abs(weight), axis=1, keepdims=True).astype(np.float32)
+    scale[scale == 0] = 1.0
+    quantized = np.clip(np.rint(weight/scale*127.0), -127, 127).astype(np.int8)
+    qweight = UOp.new_buffer("QCOM", quantized.size, dtypes.int8)
+    qweight.buffer.ensure_allocated()
+    upload(qweight.buffer, quantized)
+    qscale = UOp.new_buffer("QCOM", scale.size, dtypes.half)
+    qscale.buffer.ensure_allocated()
+    upload(qscale.buffer, scale.astype(np.float16))
+
+    info = program.arg
+    aux = adapt_aux_dtype(info.aux, 2, dtypes.int8)
+    aux = (aux[0] + (((4, dtypes.half, (1, 512, 4)),),),)
+    info = replace(info, globals=info.globals+(4,), ins=info.ins+(4,), aux=aux)
+    lib = Device["QCOM"].compiler.compile(source)
+    program = program.replace(arg=info, 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], call.src[2], qweight, call.src[4], qscale))
+    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:
+  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)
+
+
+if __name__ == "__main__": main()

extra/gemm/qcom_openpilot_schedule_projection.py

@@ -0,0 +1,112 @@
+#!/usr/bin/env python3
+"""Reschedule independent texture addresses in the dominant vision projection."""
+import argparse, pickle, struct
+
+from tinygrad.engine.jit import create_graph_call
+from tinygrad.uop.ops import Ops
+from extra.gemm.ir3asm import ADD_S, BR, ISAM_F32, MOV_F32, NOP
+from extra.gemm.qcom_ir3_matmul_patch import plain_name
+
+TARGET="r_32_192_4_4_64_4"
+FIRST_CONV_TARGET="r_64_32_16_4_4_6_3_3_4"
+FORWARD_STYLE={TARGET,"r_32_64_4_4_64_4","r_8_384_4_4_128_4"}
+GAP_STYLE={"r_512_16_4_4_16_4","r_512_48_4_4_16_4","r_128_32_4_4_32_4","r_128_96_4_4_32_4"}
+INVERSE_W_STYLE={"r_512_16_4_4_48_4","r_128_32_4_4_96_4"}
+INVERSE_STYLE={"r_32_64_4_4_192_4"}
+TARGETS=FORWARD_STYLE|GAP_STYLE|INVERSE_W_STYLE|INVERSE_STYLE|{FIRST_CONV_TARGET}
+
+
+def schedule_first_conv(lib:bytes) -> bytes:
+  image_offset, image_size=struct.unpack_from("<I",lib,0xc0)[0],struct.unpack_from("<I",lib,0x100)[0]
+  instrs=[lib[i:i+8] for i in range(image_offset,image_offset+image_size,8)]
+  if len(instrs) != 262: raise RuntimeError(f"expected 262 first-conv instructions, got {len(instrs)}")
+  # Use registers that the subsequent texture loads overwrite, allowing all
+  # eight independent input/weight coordinates to precede the texture reads.
+  addresses=[]
+  for index,(register,offset) in enumerate(zip(("r0","r2","r3","r4"),(-36,-24,-12,0))):
+    addresses.append(MOV_F32(f"{register}.x","r16.w",ss=index > 0) if offset == 0 else
+                     ADD_S(f"{register}.x","r16.w",offset,ss=index > 0))
+    addresses.append(MOV_F32(f"{register}.y","r16.z"))
+  addresses.extend(instrs[i] for i in (48,51,54,57))
+  loads=[ISAM_F32(dst,f"{coord}.x",tex=0) for dst,coord in zip(("r7.x","r6.x","r1.x","r0.x"),("r0","r2","r3","r4"))]
+  loads.extend(ISAM_F32(dst,coord,tex=1) for dst,coord in zip(("r2.x","r3.x","r4.x","r5.x"),("r8.x","r8.z","r9.x","r9.z")))
+  out=instrs[:32]+addresses+loads+instrs[60:86]
+  out.append(BR(31-len(out)))
+  out.extend(instrs[87:92])
+  out.append(BR(26-len(out)))
+  out.extend(instrs[93:99])
+  out.append(BR(24-len(out)))
+  out.extend(instrs[100:])
+  out.extend([NOP()]*(len(instrs)-len(out)))
+  if len(out) != len(instrs): raise RuntimeError(f"scheduled first conv has {len(out)} instructions")
+  return lib[:image_offset]+b"".join(out)+lib[image_offset+image_size:]
+
+
+def schedule_native_f16(lib:bytes) -> bytes:
+  image_offset, image_size=struct.unpack_from("<I",lib,0xc0)[0],struct.unpack_from("<I",lib,0x100)[0]
+  instrs=[lib[i:i+8] for i in range(image_offset,image_offset+image_size,8)]
+  if len(instrs) != 222: raise RuntimeError(f"expected 222 native-FP16 instructions, got {len(instrs)}")
+  addresses=(21,24,27,30,33,40,47,54)
+  loads=(23,26,29,32,35,42,49,56)
+  mads=tuple(range(36,40))+tuple(range(43,47))+tuple(range(50,54))+tuple(range(57,61))
+  out=instrs[:21]+[instrs[i] for i in addresses]+[instrs[i] for i in loads]+[instrs[i] for i in mads]+instrs[61:67]
+  out.append(BR(21-len(out)))
+  out.extend(instrs[68:])
+  out.extend([NOP()]*(len(instrs)-len(out)))
+  if len(out) != len(instrs): raise RuntimeError(f"scheduled native FP16 has {len(out)} instructions")
+  return lib[:image_offset]+b"".join(out)+lib[image_offset+image_size:]
+
+
+def schedule_loads(lib:bytes, name:str) -> bytes:
+  image_offset, image_size=struct.unpack_from("<I",lib,0xc0)[0],struct.unpack_from("<I",lib,0x100)[0]
+  instrs=[lib[i:i+8] for i in range(image_offset,image_offset+image_size,8)]
+  if name == TARGET and len(instrs) == 222: return schedule_native_f16(lib)
+  if len(instrs) < 160: raise RuntimeError(f"expected projection shader, got {len(instrs)} instructions")
+  # The compiler emits address, rpt5 nop, texture-read eight times. Calculate
+  # every independent address first, then issue the reads as one contiguous run.
+  if name in FORWARD_STYLE: start,body_end=20,66
+  elif name in GAP_STYLE: start,body_end=26,84
+  elif name in INVERSE_W_STYLE: start,body_end=19,76
+  elif name in INVERSE_STYLE: start,body_end=16,73
+  else: raise RuntimeError(f"unsupported projection {name}")
+  address_indices=tuple(start+3*i for i in range(8))
+  load_indices=tuple(start+3*i+2 for i in range(8))
+  out=instrs[:start]+[instrs[i] for i in address_indices]+[instrs[i] for i in load_indices]+instrs[start+24:body_end]
+  branch_index=len(out)
+  out.append(BR(start-branch_index))
+  out.extend(instrs[body_end+1:])
+  out.extend([NOP()]*(len(instrs)-len(out)))
+  if len(out) != len(instrs): raise RuntimeError(f"scheduled image has {len(out)} instructions")
+  return lib[:image_offset]+b"".join(out)+lib[image_offset+image_size:]
+
+
+def patch_projection(model) -> int:
+  outer=model.captured.linear.src[0]
+  batch=outer.src[0].src[0].src
+  new_batch=[]
+  cache={}
+  replaced=0
+  for call in batch:
+    name=plain_name(call.src[0].arg.name) if call.op is Ops.CALL and call.src[0].op is Ops.PROGRAM else ""
+    if name in TARGETS:
+      program=call.src[0]
+      old=program.src[3].arg
+      if old not in cache: cache[old]=schedule_first_conv(old) if name == FIRST_CONV_TARGET else schedule_loads(old,name)
+      program=program.replace(src=program.src[:3]+(program.src[3].replace(arg=cache[old]),))
+      call=call.replace(src=(program,*call.src[1:]))
+      replaced+=1
+    new_batch.append(call)
+  model.captured._linear=model.captured.linear.substitute({outer:create_graph_call(new_batch)},walk=True)
+  model.captured.__dict__.pop("linear",None)
+  return replaced
+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("patched",patch_projection(model))
+  with open(args.output,"wb") as f:pickle.dump(model,f)
+
+
+if __name__ == "__main__":main()

extra/gemm/qcom_openpilot_skip_mlp.py

@@ -0,0 +1,40 @@
+#!/usr/bin/env python3
+"""Sensitivity experiment: alias selected vision MLP residual branches to identity."""
+import argparse, pickle
+
+from tinygrad.engine.jit import create_graph_call
+from tinygrad.uop.ops import Ops
+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"
+
+
+def patch_model(model, selected:set[int]) -> int:
+  outer=model.captured.linear.src[0]
+  batch=list(outer.src[0].src[0].src)
+  aliases, remove, pair = {}, set(), 0
+  for index in range(len(batch)-1):
+    a,b=batch[index:index+2]
+    if not all(x.op is Ops.CALL and x.src[0].op is Ops.PROGRAM for x in (a,b)): continue
+    if (plain_name(a.src[0].arg.name),plain_name(b.src[0].arg.name)) != (FORWARD,INVERSE): continue
+    if pair in selected:
+      aliases[b.src[1]] = b.src[2]
+      remove.update((index,index+1))
+    pair += 1
+  def representative(x):
+    while x in aliases: x=aliases[x]
+    return x
+  new_batch=[call.replace(src=tuple(representative(x) for x in call.src)) for i,call in enumerate(batch) if i not in remove]
+  model.captured._linear=model.captured.linear.substitute({outer:create_graph_call(new_batch)},walk=True)
+  model.captured.__dict__.pop("linear",None)
+  return len(remove)//2
+
+
+def main() -> None:
+  ap=argparse.ArgumentParser(); ap.add_argument("input"); ap.add_argument("output"); ap.add_argument("--indices",required=True); args=ap.parse_args()
+  with open(args.input,"rb") as f:model=pickle.load(f)
+  print("skipped",patch_model(model,{int(x) for x in args.indices.split(",")}))
+  with open(args.output,"wb") as f:pickle.dump(model,f)
+
+
+if __name__ == "__main__":main()

extra/gemm/qcom_openpilot_vision_fp16.py

@@ -0,0 +1,378 @@
+#!/usr/bin/env python3
+"""Replace selected driving_vision 1x1 convolutions with vector FP16-acc kernels."""
+import argparse, 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, MAD_F32, NOP
+from extra.gemm.qcom_ir3_matmul_patch import plain_name
+
+TARGET = "r_32_192_4_4_64_4"
+INVERSE_TARGET = "r_32_64_4_4_192_4"
+FIRST_CONV_TARGET = "r_64_32_16_4_4_6_3_3_4"
+FULL_Y_TARGETS = {TARGET, "r_32_64_4_4_64_4"}
+FULL_Z_TARGETS = {"r_8_384_4_4_128_4"}
+GAP_Y_TARGETS = {"r_512_16_4_4_16_4", "r_512_48_4_4_16_4", "r_128_32_4_4_32_4", "r_128_96_4_4_32_4"}
+INVERSE_W_TARGETS = {"r_512_16_4_4_48_4", "r_128_32_4_4_96_4"}
+OTHER_INVERSE_TARGETS: set[str] = set()
+FP32_TARGETS = FULL_Y_TARGETS | FULL_Z_TARGETS | GAP_Y_TARGETS | INVERSE_W_TARGETS | OTHER_INVERSE_TARGETS | {INVERSE_TARGET, FIRST_CONV_TARGET}
+
+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;
+  half4 r0=(half4)(0),r1=(half4)(0),r2=(half4)(0),r3=(half4)(0);
+  for (int k=0;k<64;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;
+  }
+  float4 b=read_imagef(B,smp,(int2)(n,0));
+  write_imagef(O,(int2)(n,m),gelu(convert_float4(r0)+b));
+  write_imagef(O,(int2)(n+192,m),gelu(convert_float4(r1)+b));
+  write_imagef(O,(int2)(n+384,m),gelu(convert_float4(r2)+b));
+  write_imagef(O,(int2)(n+576,m),gelu(convert_float4(r3)+b));
+}"""
+
+
+def pack_mads(lib:bytes) -> bytes:
+  image_offset, image_size = struct.unpack_from("<I", lib, 0xc0)[0], struct.unpack_from("<I", lib, 0x100)[0]
+  image = lib[image_offset:image_offset+image_size]
+  instrs = [image[i:i+8] for i in range(0, len(image), 8)]
+  if len(instrs) != 222: raise RuntimeError(f"expected 222 instructions, got {len(instrs)}")
+  out = instrs[:36]
+  activations = ("hr3", "hr2", "hr1", "hr0")
+  accumulators = ("hr8.x", "hr7.x", "hr6.x", "hr5.x")
+  load_chunks = (instrs[52:55], instrs[68:71], instrs[87:90], ())
+  for component, load_after in zip("xyzw", load_chunks):
+    for index, (acc, activation) in enumerate(zip(accumulators, activations)):
+      out.append(MAD_F16(acc, f"{activation}.{component}", "hr4.x", acc, rpt=3, sy=index == 0, r=True))
+    out += list(load_after)
+  out += instrs[102:108]
+  branch_index = len(out)
+  out.append(BR(21-branch_index))
+  out += instrs[109:]
+  out += [NOP()] * (len(instrs)-len(out))
+  if len(out) != len(instrs): raise RuntimeError(f"packed image has {len(out)} instructions")
+  return lib[:image_offset] + b"".join(out) + lib[image_offset+image_size:]
+
+
+def pack_preloaded_mads(lib:bytes) -> bytes:
+  """Keep all four weight vectors resident and synchronize the sampler once per K."""
+  lib = pack_mads(lib)
+  image_offset, image_size = struct.unpack_from("<I", lib, 0xc0)[0], struct.unpack_from("<I", lib, 0x100)[0]
+  instrs = [lib[image_offset+i:image_offset+i+8] for i in range(0, image_size, 8)]
+  if len(instrs) != 222: raise RuntimeError(f"expected 222 instructions, got {len(instrs)}")
+  out = instrs[:23] + instrs[23:33]
+  for coord_setup, coord, weight in ((instrs[33:35], "r0.x", "hr9.x"), (instrs[40:42], "r0.z", "hr10.x"),
+                                     (instrs[47:49], "r1.x", "hr11.x"), (instrs[54:56], "r1.z", "hr12.x")):
+    out += coord_setup + [ISAM_F16(weight, coord, 1)]
+  activations, accumulators = ("hr3", "hr2", "hr1", "hr0"), ("hr8.x", "hr7.x", "hr6.x", "hr5.x")
+  first = True
+  for component, weight in zip("xyzw", ("hr9.x", "hr10.x", "hr11.x", "hr12.x")):
+    for acc, activation in zip(accumulators, activations):
+      out.append(MAD_F16(acc, f"{activation}.{component}", weight, acc, rpt=3, sy=first, r=True))
+      first = False
+  out += instrs[61:68] + instrs[68:]
+  if len(out) != len(instrs): raise RuntimeError(f"preloaded image has {len(out)} instructions")
+  patched = bytearray(lib[:image_offset] + b"".join(out) + lib[image_offset+image_size:])
+  reg_offset = struct.unpack_from("<I", patched, 0x34)[0]
+  old_hregs = struct.unpack_from("<I", patched, reg_offset+0x18)[0]
+  struct.pack_into("<I", patched, reg_offset+0x18, (old_hregs & 0x80000000) | 13)
+  return bytes(patched)
+
+
+def pack_fp32_mads(lib:bytes, component:str="y") -> bytes:
+  image_offset, image_size = struct.unpack_from("<I", lib, 0xc0)[0], struct.unpack_from("<I", lib, 0x100)[0]
+  image = lib[image_offset:image_offset+image_size]
+  instrs = [image[i:i+8] for i in range(0, len(image), 8)]
+  if len(instrs) < 116: raise RuntimeError(f"expected FP32 matmul loop through instruction 115, got {len(instrs)}")
+  out = instrs[:44]
+  for k_component, weight in zip("xyzw", ("r5.x", "r2.x", "r3.x", "r4.x")):
+    for acc, activation in zip(tuple(f"r{reg}.{component}" for reg in range(13, 17)), ("r7", "r6", "r1", "r0")):
+      out.append(MAD_F32(acc, f"{activation}.{k_component}", weight, acc, rpt=3,
+                         sy=len(out) == 44, r=True))
+  out += instrs[108:114]
+  branch_index = len(out)
+  out.append(BR(20-branch_index))
+  out += instrs[115:]
+  out += [NOP()] * (len(instrs)-len(out))
+  if len(out) != len(instrs): raise RuntimeError(f"packed FP32 image has {len(out)} instructions")
+  return lib[:image_offset] + b"".join(out) + lib[image_offset+image_size:]
+
+
+def pack_gap_y_fp32_mads(lib:bytes) -> bytes:
+  image_offset, image_size = struct.unpack_from("<I", lib, 0xc0)[0], struct.unpack_from("<I", lib, 0x100)[0]
+  image = lib[image_offset:image_offset+image_size]
+  instrs = [image[i:i+8] for i in range(0, len(image), 8)]
+  if len(instrs) < 237: raise RuntimeError(f"expected at least 237 gap-Y instructions, got {len(instrs)}")
+  out = instrs[:66]
+  for component, weight in zip("xyzw", ("r5.x", "r2.x", "r3.x", "r4.x")):
+    for acc, activation in zip(("r14.y", "r15.y", "r16.y"), ("r6", "r1", "r0")):
+      out.append(MAD_F32(acc, f"{activation}.{component}", weight, acc, rpt=3, r=True))
+  out += instrs[114:120]
+  branch_index = len(out)
+  out.append(BR(26-branch_index))
+  out += instrs[121:]
+  out += [NOP()] * (len(instrs)-len(out))
+  if len(out) != len(instrs): raise RuntimeError(f"packed gap-Y image has {len(out)} instructions")
+  return lib[:image_offset] + b"".join(out) + lib[image_offset+image_size:]
+
+
+def pack_inverse_w_fp32_mads(lib:bytes) -> bytes:
+  image_offset, image_size = struct.unpack_from("<I", lib, 0xc0)[0], struct.unpack_from("<I", lib, 0x100)[0]
+  image = lib[image_offset:image_offset+image_size]
+  instrs = [image[i:i+8] for i in range(0, len(image), 8)]
+  if len(instrs) < 174: raise RuntimeError(f"expected at least 174 inverse-W instructions, got {len(instrs)}")
+  out = instrs[:59]
+  for component, weight in zip("xyzw", ("r5.x", "r2.x", "r3.x", "r4.x")):
+    for acc, activation in zip(("r13.w", "r14.w", "r15.w"), ("r6", "r1", "r0")):
+      out.append(MAD_F32(acc, f"{activation}.{component}", weight, acc, rpt=3, r=True))
+  out += instrs[107:112]
+  branch_index = len(out)
+  out.append(BR(19-branch_index))
+  out += instrs[113:]
+  out += [NOP()] * (len(instrs)-len(out))
+  if len(out) != len(instrs): raise RuntimeError(f"packed inverse-W image has {len(out)} instructions")
+  return lib[:image_offset] + b"".join(out) + lib[image_offset+image_size:]
+
+
+def pack_other_inverse_fp32_mads(lib:bytes) -> bytes:
+  image_offset, image_size = struct.unpack_from("<I", lib, 0xc0)[0], struct.unpack_from("<I", lib, 0x100)[0]
+  image = lib[image_offset:image_offset+image_size]
+  instrs = [image[i:i+8] for i in range(0, len(image), 8)]
+  if len(instrs) != 179: raise RuntimeError(f"expected 179 other-inverse instructions, got {len(instrs)}")
+  # The first output vector is split by loop-control registers. Keep it scalar;
+  # the remaining three vectors are contiguous from r14.y through r17.x.
+  out = instrs[:60]
+  for component, weight in zip("xyzw", ("r5.x", "r2.x", "r3.x", "r4.x")):
+    for acc, activation in zip(("r14.y", "r15.y", "r16.y"), ("r6", "r1", "r0")):
+      out.append(MAD_F32(acc, f"{activation}.{component}", weight, acc, rpt=3, r=True))
+  out += instrs[108:114]
+  branch_index = len(out)
+  out.append(BR(20-branch_index))
+  out += instrs[115:]
+  out += [NOP()] * (len(instrs)-len(out))
+  if len(out) != len(instrs): raise RuntimeError(f"packed other-inverse image has {len(out)} instructions")
+  return lib[:image_offset] + b"".join(out) + lib[image_offset+image_size:]
+
+
+def pack_first_conv_fp32_mads(lib:bytes) -> bytes:
+  image_offset, image_size = struct.unpack_from("<I", lib, 0xc0)[0], struct.unpack_from("<I", lib, 0x100)[0]
+  image = lib[image_offset:image_offset+image_size]
+  instrs = [image[i:i+8] for i in range(0, len(image), 8)]
+  if len(instrs) != 262: raise RuntimeError(f"expected 262 first-conv instructions, got {len(instrs)}")
+  out = instrs[:60]
+  first = True
+  for component, weight in zip("xyzw", ("r5", "r2", "r3", "r4")):
+    out.append(MAD_F32("r11.w", f"r7.{component}", f"{weight}.x", "r11.w", sy=first, r=True))
+    first = False
+    out.append(MAD_F32("r12.y", f"r7.{component}", f"{weight}.y", "r12.y", rpt=2, r=True))
+  for component, weight in zip("xyzw", ("r5.x", "r2.x", "r3.x", "r4.x")):
+    for acc, activation in zip(("r13.x", "r14.x", "r15.x"), ("r6", "r1", "r0")):
+      out.append(MAD_F32(acc, f"{activation}.{component}", weight, acc, rpt=3, r=True))
+  out += instrs[124:130]
+  branch_index = len(out)
+  out.append(BR(31-branch_index))
+  out += instrs[131:]
+  # Compacting the innermost loop also relocates the two enclosing-loop branches.
+  # Their targets remain in the untouched prologue, so rebuild their relative offsets.
+  out[92] = BR(26-92)
+  out[99] = BR(24-99)
+  out += [NOP()] * (len(instrs)-len(out))
+  if len(out) != len(instrs): raise RuntimeError(f"packed first-conv image has {len(out)} instructions")
+  return lib[:image_offset] + b"".join(out) + lib[image_offset+image_size:]
+
+
+def pack_native_f16_forward_mads(lib:bytes) -> bytes:
+  image_offset, image_size = struct.unpack_from("<I", lib, 0xc0)[0], struct.unpack_from("<I", lib, 0x100)[0]
+  image = lib[image_offset:image_offset+image_size]
+  instrs = [image[i:i+8] for i in range(0, len(image), 8)]
+  if len(instrs) == 232:
+    # Current QCOM compiler scalarizes the 4x4 outer product into 64 MADs even
+    # though both the four output lanes and accumulator lanes are contiguous.
+    # Preserve its sampler schedule and replace only that scalar MAD block.
+    # Hoist all independent texture coordinates, then issue the eight samples
+    # contiguously. The original compiler inserted rpt5 after every coordinate.
+    address_indices = (20, 23, 26, 29, 32, 35, 38, 41)
+    load_indices = (22, 25, 28, 31, 34, 37, 40, 43)
+    out = instrs[:20] + [instrs[i] for i in address_indices] + [instrs[i] for i in load_indices]
+    for component, weight in zip("xyzw", ("hr5.x", "hr2.x", "hr3.x", "hr4.x")):
+      for acc, activation in (("hr8.x", "hr7"), ("hr9.x", "hr6"), ("hr10.x", "hr1"), ("hr11.x", "hr0")):
+        out.append(MAD_F16(acc, f"{activation}.{component}", weight, acc, rpt=3, r=True,
+                           sy=len(out) == 36))
+    out += instrs[108:114]
+    out.append(BR(20-len(out)))
+    out += instrs[115:]
+    out += [NOP()] * (len(instrs)-len(out))
+    if len(out) != len(instrs): raise RuntimeError(f"packed native-FP16 image has {len(out)} instructions")
+    return lib[:image_offset] + b"".join(out) + lib[image_offset+image_size:]
+  if len(instrs) != 241: raise RuntimeError(f"expected 232/241 native-FP16 forward instructions, got {len(instrs)}")
+  already_packed = instrs[53] == MAD_F16("hr0.x", "hr7.y", "hr8.x", "hr0.x", rpt=3, r=True)
+  out = instrs[:52]
+  for component, weight in zip("xyzw", ("hr11.x", "hr8.x", "hr9.x", "hr10.x")):
+    for acc, activation in (("hr0.x", "hr7"), ("hr1.x", "hr4"), ("hr2.x", "hr5"), ("hr3.x", "hr6")):
+      out.append(MAD_F16(acc, f"{activation}.{component}", weight, acc, rpt=3, r=True))
+  out += instrs[68:74] if already_packed else instrs[116:122]
+  branch_index = len(out)
+  out.append(BR(20-branch_index))
+  out += instrs[75:] if already_packed else instrs[123:]
+  out += [NOP()] * (len(instrs)-len(out))
+  if len(out) != len(instrs): raise RuntimeError(f"packed native-FP16 image has {len(out)} instructions")
+  return lib[:image_offset] + b"".join(out) + lib[image_offset+image_size:]
+
+
+def pack_inverse_fp32_mads(lib:bytes) -> bytes:
+  image_offset, image_size = struct.unpack_from("<I", lib, 0xc0)[0], struct.unpack_from("<I", lib, 0x100)[0]
+  image = lib[image_offset:image_offset+image_size]
+  instrs = [image[i:i+8] for i in range(0, len(image), 8)]
+  if len(instrs) != 175: raise RuntimeError(f"expected 175 inverse instructions, got {len(instrs)}")
+  # The first output vector straddles loop-control r13.x, so retain its scalar
+  # instructions. The remaining r14/r15/r16 accumulator vectors are contiguous.
+  out = instrs[:56]
+  for component, weight in zip("xyzw", ("r5.x", "r2.x", "r3.x", "r4.x")):
+    for acc, activation in zip(("r14.x", "r15.x", "r16.x"), ("r6", "r1", "r0")):
+      out.append(MAD_F32(acc, f"{activation}.{component}", weight, acc, rpt=3, r=True))
+  out += instrs[104:109]
+  branch_index = len(out)
+  out.append(BR(16-branch_index))
+  out += instrs[110:]
+  out += [NOP()] * (len(instrs)-len(out))
+  if len(out) != len(instrs): raise RuntimeError(f"packed inverse image has {len(out)} instructions")
+  return lib[:image_offset] + b"".join(out) + lib[image_offset+image_size:]
+
+
+def pack_first_conv_f16_mads(lib:bytes) -> bytes:
+  """Vectorize the native-half first convolution's four 4x4 dot products."""
+  image_offset, image_size = struct.unpack_from("<I", lib, 0xc0)[0], struct.unpack_from("<I", lib, 0x100)[0]
+  instrs = [lib[image_offset+i:image_offset+i+8] for i in range(0, image_size, 8)]
+  if len(instrs) < 200: raise RuntimeError(f"expected native-FP16 first-conv shader, got {len(instrs)} instructions")
+  out = instrs[:60]
+  first = True
+  for activation, acc in (("hr7", "hr8.x"), ("hr6", "hr9.x"), ("hr1", "hr10.x"), ("hr0", "hr11.x")):
+    for component, weight in zip("xyzw", ("hr5.x", "hr2.x", "hr3.x", "hr4.x")):
+      out.append(MAD_F16(acc, f"{activation}.{component}", weight, acc, rpt=3, r=True, sy=first))
+      first = False
+  out += instrs[124:130]
+  out.append(BR(31-len(out)))
+  out += instrs[131:]
+  # The compact inner loop relocates both enclosing-loop branches too.
+  out[88] = BR(26-88)
+  out[95] = BR(24-95)
+  if len(out) > len(instrs): raise RuntimeError(f"packed native-FP16 first conv grew from {len(instrs)} to {len(out)}")
+  out += [NOP()] * (len(instrs)-len(out))
+  return lib[:image_offset] + b"".join(out) + lib[image_offset+image_size:]
+
+
+def patch_native_f16_firstconv(jit) -> int:
+  outer = jit.captured.linear.src[0]
+  batch, replaced = list(outer.src[0].src[0].src), 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) != FIRST_CONV_TARGET: continue
+    program = call.src[0]
+    program = program.replace(src=program.src[:3] + (program.src[3].replace(arg=pack_first_conv_f16_mads(program.src[3].arg)),))
+    batch[index], replaced = call.replace(src=(program, *call.src[1:])), replaced+1
+  if replaced:
+    jit.captured._linear = jit.captured.linear.substitute({outer:create_graph_call(batch)}, walk=True)
+    jit.captured.__dict__.pop("linear", None)
+  return replaced
+
+
+def patch_fp32_rpt(jit, names:set[str]|None=None) -> int:
+  """Apply the verified FP32-accumulate repeat packing to a captured vision JIT."""
+  outer = jit.captured.linear.src[0]
+  batch = outer.src[0].src[0].src
+  new_batch, replaced = [], 0
+  for call in batch:
+    name = plain_name(call.src[0].arg.name) if call.op is Ops.CALL and call.src[0].op is Ops.PROGRAM else ""
+    if name in FP32_TARGETS and (names is None or name in names):
+      program = call.src[0]
+      patch = (pack_fp32_mads if name in FULL_Y_TARGETS else
+               (lambda lib:pack_fp32_mads(lib, "z")) if name in FULL_Z_TARGETS else
+               pack_gap_y_fp32_mads if name in GAP_Y_TARGETS else
+               pack_inverse_w_fp32_mads if name in INVERSE_W_TARGETS else
+               pack_other_inverse_fp32_mads if name in OTHER_INVERSE_TARGETS else
+               pack_first_conv_fp32_mads if name == FIRST_CONV_TARGET else pack_inverse_fp32_mads)
+      program = program.replace(src=program.src[:3] + (program.src[3].replace(arg=patch(program.src[3].arg)),))
+      if name == INVERSE_TARGET:
+        program = program.replace(arg=replace(program.arg, global_size=(8, 2, 1), local_size=(8, 16, 1)))
+      call, replaced = call.replace(src=(program, *call.src[1:])), replaced+1
+    new_batch.append(call)
+  if replaced:
+    jit.captured._linear = jit.captured.linear.substitute({outer:create_graph_call(new_batch)}, walk=True)
+    jit.captured.__dict__.pop("linear", None)
+  return replaced
+
+
+def main() -> None:
+  ap = argparse.ArgumentParser()
+  ap.add_argument("input")
+  ap.add_argument("output")
+  ap.add_argument("--fp32-rpt", action="store_true")
+  ap.add_argument("--fp32-rpt-names", help="comma-separated subset for --fp32-rpt")
+  ap.add_argument("--native-f16-rpt", action="store_true")
+  ap.add_argument("--preload-f16", action="store_true")
+  ap.add_argument("--native-f16-firstconv-rpt", action="store_true")
+  args = ap.parse_args()
+  with open(args.input, "rb") as f: jit = pickle.load(f)
+  if args.native_f16_firstconv_rpt:
+    replaced = patch_native_f16_firstconv(jit)
+    if not replaced: raise RuntimeError(f"no {FIRST_CONV_TARGET} calls found")
+    with open(args.output, "wb") as f: pickle.dump(jit, f)
+    print(f"patched {replaced} call(s)")
+    return
+  if args.fp32_rpt:
+    replaced = patch_fp32_rpt(jit, set(args.fp32_rpt_names.split(",")) if args.fp32_rpt_names else None)
+    if not replaced: raise RuntimeError(f"no {TARGET} calls found")
+    with open(args.output, "wb") as f: pickle.dump(jit, f)
+    print(f"patched {replaced} call(s)")
+    return
+  hand_lib = None if args.fp32_rpt else (pack_preloaded_mads if args.preload_f16 else pack_mads)(Device["QCOM"].compiler.compile_cached(SOURCE))
+  outer = jit.captured.linear.src[0]
+  batch = outer.src[0].src[0].src
+  new_batch, replaced = [], 0
+  for call in batch:
+    name = plain_name(call.src[0].arg.name) if call.op is Ops.CALL and call.src[0].op is Ops.PROGRAM else ""
+    if name == TARGET or (args.fp32_rpt and name in FP32_TARGETS):
+      program = call.src[0]
+      if args.native_f16_rpt:
+        program = program.replace(src=program.src[:3] + (program.src[3].replace(arg=pack_native_f16_forward_mads(program.src[3].arg)),))
+      elif args.fp32_rpt:
+        patch = (pack_fp32_mads if name in FULL_Y_TARGETS else
+                 (lambda lib:pack_fp32_mads(lib, "z")) if name in FULL_Z_TARGETS else
+                 pack_gap_y_fp32_mads if name in GAP_Y_TARGETS else
+                 pack_inverse_w_fp32_mads if name in INVERSE_W_TARGETS else
+                 pack_other_inverse_fp32_mads if name in OTHER_INVERSE_TARGETS else
+                 pack_first_conv_fp32_mads if name == FIRST_CONV_TARGET else pack_inverse_fp32_mads)
+        program = program.replace(src=program.src[:3] + (program.src[3].replace(arg=patch(program.src[3].arg)),))
+        if name == INVERSE_TARGET:
+          program = program.replace(arg=replace(program.arg, global_size=(8, 2, 1), local_size=(8, 16, 1)))
+      else:
+        program = program.replace(src=program.src[:2] +
+          (program.src[2].replace(arg=SOURCE), program.src[3].replace(arg=hand_lib)))
+      call, replaced = call.replace(src=(program, *call.src[1:])), replaced+1
+    new_batch.append(call)
+  if not replaced: raise RuntimeError(f"no {TARGET} calls found")
+  jit.captured._linear = jit.captured.linear.substitute({outer:create_graph_call(new_batch)}, walk=True)
+  jit.captured.__dict__.pop("linear", None)
+  with open(args.output, "wb") as f: pickle.dump(jit, f)
+  print(f"patched {replaced} call(s)")
+
+
+if __name__ == "__main__": main()

extra/gemm/qcom_openpilot_weight_stats.py

@@ -0,0 +1,67 @@
+#!/usr/bin/env python3
+"""Inspect dominant cached OpenPilot GEMM weight distributions."""
+import argparse, pickle
+
+import numpy as np
+
+from tinygrad.uop.ops import Ops
+from extra.gemm.qcom_ir3_matmul_patch import plain_name
+
+
+def main() -> None:
+  parser = argparse.ArgumentParser()
+  parser.add_argument("model")
+  parser.add_argument("--svd", default="", help="comma-separated GEMM indices for singular-value tail analysis")
+  parser.add_argument("--compare", help="second pickle whose cached GEMM weights should be compared")
+  args = parser.parse_args()
+  with open(args.model, "rb") as f: model = pickle.load(f)
+  batch = model.captured.linear.src[0].src[0].src[0].src
+  rows = []
+  for call in 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) != "gemm_h": continue
+    weight = np.asarray(call.src[2].buffer.numpy(), dtype=np.float32).reshape(-1)
+    absw = np.abs(weight)
+    scale = float(absw.max())
+    rows.append((tuple(call.src[0].arg.global_size), scale, float(np.mean(weight == 0)),
+                 *(float(np.mean(absw <= scale*x)) for x in (1/1024, 1/512, 1/256, 1/128, 1/64))))
+  for i, row in enumerate(rows):
+    print(f"{i:2d} gs={row[0]} max={row[1]:.7g} zero={row[2]:.5f} "
+          f"near=[{row[3]:.4f},{row[4]:.4f},{row[5]:.4f},{row[6]:.4f},{row[7]:.4f}]")
+  if not rows:
+    produced, seen = set(), set()
+    for call in batch:
+      if call.op is not Ops.CALL or call.src[0].op is not Ops.PROGRAM: continue
+      for arg_index, buf in enumerate(call.src[1:]):
+        if arg_index in call.src[0].arg.outs or buf.op is not Ops.BUFFER or buf in produced or buf in seen or buf.buffer.size < 4096: continue
+        seen.add(buf)
+        values = np.asarray(buf.buffer.numpy())
+        if not np.issubdtype(values.dtype, np.floating): continue
+        absw = np.abs(values.astype(np.float32)); scale = float(absw.max())
+        print(f"{plain_name(call.src[0].arg.name)} arg={arg_index} size={values.size} max={scale:.7g} "
+              f"zero={float(np.mean(values == 0)):.6f} near={float(np.mean(absw <= scale/1024)):.6f}")
+      for out in call.src[0].arg.outs: produced.add(call.src[out+1])
+  if args.svd:
+    selected = {int(x) for x in args.svd.split(",")}
+    gemms = [call for call in 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"]
+    for i in sorted(selected):
+      call = gemms[i]
+      weight = np.asarray(call.src[2].buffer.numpy(), dtype=np.float32)
+      matrix = weight.reshape((384, 1536) if tuple(call.src[0].arg.global_size) == (12, 8, 1) else (1536, 384))
+      singular = np.linalg.svd(matrix, compute_uv=False)
+      energy = np.cumsum(singular[::-1]**2)[::-1]
+      total = energy[0]
+      ranks = (32, 64, 96, 128, 192, 256, 320)
+      tails = [float(np.sqrt(energy[r]/total)) if r < len(singular) else 0.0 for r in ranks]
+      print(f"svd {i:2d} shape={matrix.shape} rel_frob_tail=" + ",".join(f"r{r}:{e:.5f}" for r, e in zip(ranks, tails)))
+  if args.compare:
+    with open(args.compare, "rb") as f: other = pickle.load(f)
+    other_batch = other.captured.linear.src[0].src[0].src[0].src
+    lhs = [call for call in 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"]
+    rhs = [call for call in other_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"]
+    for i, (a, b) in enumerate(zip(lhs, rhs)):
+      av, bv = np.asarray(a.src[2].buffer.numpy()), np.asarray(b.src[2].buffer.numpy())
+      print(f"compare {i:2d} max={float(np.max(np.abs(av.astype(np.float32)-bv.astype(np.float32)))):.9g} equal={np.array_equal(av,bv)}")
+
+
+if __name__ == "__main__": main()

extra/gemm/qcom_pack_banked_check.py

@@ -0,0 +1,45 @@
+#!/usr/bin/env python3
+"""Check FP32 banked-image to row-major FP16 packing on QCOM."""
+import argparse
+
+import numpy as np
+
+from tinygrad import Device, dtypes
+from tinygrad.device import Buffer
+
+
+def main() -> None:
+  ap = argparse.ArgumentParser()
+  ap.add_argument("--m", type=int, default=192)
+  ap.add_argument("--k", type=int, default=768)
+  ap.add_argument("--constant", action="store_true")
+  args = ap.parse_args()
+  rng = np.random.default_rng(0)
+  source = rng.standard_normal((args.m//4, args.k, 4)).astype(np.float32)
+  expected = source.transpose(0, 2, 1).reshape(args.m, args.k).astype(np.float16)
+  inp = Buffer("QCOM", source.size, dtypes.float).allocate()
+  out = Buffer("QCOM", source.size, dtypes.half).allocate()
+  inp.copyin(memoryview(source).cast("B"))
+  out.copyin(memoryview(np.zeros(source.size, np.float16)).cast("B"))
+  kernel = 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 pack_banked(__global uint *O, read_only image2d_t X) {{
+  int t=get_global_id(0),row=t/{args.k//4},k4=t-row*{args.k//4},lane=row&3,y=row>>2,x=k4*4;
+  float4 p0=read_imagef(X,smp,(int2)(x+0,y)),p1=read_imagef(X,smp,(int2)(x+1,y));
+  float4 p2=read_imagef(X,smp,(int2)(x+2,y)),p3=read_imagef(X,smp,(int2)(x+3,y));
+  float4 v=lane==0?(float4)(p0.x,p1.x,p2.x,p3.x):lane==1?(float4)(p0.y,p1.y,p2.y,p3.y):
+           lane==2?(float4)(p0.z,p1.z,p2.z,p3.z):(float4)(p0.w,p1.w,p2.w,p3.w);
+  vstore2(as_uint2({"(half4)(1.0h)" if args.constant else "convert_half4(v)"}),0,O+t*2);
+}}"""
+  specs = [((0, dtypes.half, None),), ((1, dtypes.float, (args.m//4, args.k, 4)),)]
+  program = Device["QCOM"].runtime("pack_banked", Device["QCOM"].compiler.compile(kernel), buf_dtypes=specs)
+  times = [program(out._buf, inp._buf, global_size=(args.m*args.k//(4*128), 1, 1),
+                   local_size=(128, 1, 1), wait=True)*1e3 for _ in range(5)]
+  got = np.empty_like(expected)
+  out.copyout(memoryview(got).cast("B"))
+  delta = np.abs(got.astype(np.float32)-expected.astype(np.float32))
+  print(f"ms={min(times):.4f} max_abs={float(delta.max()):.9g} nonzero={int(np.count_nonzero(got))}/{got.size}")
+
+
+if __name__ == "__main__": main()

extra/gemm/qcom_packed8_gemm_check.py

@@ -0,0 +1,71 @@
+#!/usr/bin/env python3
+"""Validate a zero-copy uint4 view that fetches eight packed FP16 values."""
+import argparse, statistics
+
+import numpy as np
+
+from tinygrad import Device, dtypes
+from tinygrad.device import Buffer
+from extra.gemm.ir3asm import disasm, get_envelope
+
+
+def source(k:int, stride:int) -> str:
+  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 packed8(read_only image2d_t A,read_only image2d_t 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,col8=get_group_id(0)*32+tid;
+  half8 c0=(half8)(0),c1=(half8)(0),c2=(half8)(0),c3=(half8)(0);
+  for(int k8=0;k8<{k//8};k8++) {{
+    half8 a0=as_half8(read_imageui(A,smp,(int2)(k8,row+0)));
+    half8 a1=as_half8(read_imageui(A,smp,(int2)(k8,row+1)));
+    half8 a2=as_half8(read_imageui(A,smp,(int2)(k8,row+2)));
+    half8 a3=as_half8(read_imageui(A,smp,(int2)(k8,row+3)));
+    half8 b0=as_half8(read_imageui(B,smp,(int2)(col8,k8*8+0)));
+    half8 b1=as_half8(read_imageui(B,smp,(int2)(col8,k8*8+1)));
+    half8 b2=as_half8(read_imageui(B,smp,(int2)(col8,k8*8+2)));
+    half8 b3=as_half8(read_imageui(B,smp,(int2)(col8,k8*8+3)));
+    half8 b4=as_half8(read_imageui(B,smp,(int2)(col8,k8*8+4)));
+    half8 b5=as_half8(read_imageui(B,smp,(int2)(col8,k8*8+5)));
+    half8 b6=as_half8(read_imageui(B,smp,(int2)(col8,k8*8+6)));
+    half8 b7=as_half8(read_imageui(B,smp,(int2)(col8,k8*8+7)));
+    c0+=a0.s0*b0+a0.s1*b1+a0.s2*b2+a0.s3*b3+a0.s4*b4+a0.s5*b5+a0.s6*b6+a0.s7*b7;
+    c1+=a1.s0*b0+a1.s1*b1+a1.s2*b2+a1.s3*b3+a1.s4*b4+a1.s5*b5+a1.s6*b6+a1.s7*b7;
+    c2+=a2.s0*b0+a2.s1*b1+a2.s2*b2+a2.s3*b3+a2.s4*b4+a2.s5*b5+a2.s6*b6+a2.s7*b7;
+    c3+=a3.s0*b0+a3.s1*b1+a3.s2*b2+a3.s3*b3+a3.s4*b4+a3.s5*b5+a3.s6*b6+a3.s7*b7;
+  }}
+  vstore8(c0,0,C+(row+0)*{stride}+col8*8); vstore8(c1,0,C+(row+1)*{stride}+col8*8);
+  vstore8(c2,0,C+(row+2)*{stride}+col8*8); vstore8(c3,0,C+(row+3)*{stride}+col8*8);
+}}"""
+
+
+def upload(x:np.ndarray, dtype) -> Buffer:
+  ret=Buffer("QCOM",x.size,dtype).allocate(); raw=memoryview(np.ascontiguousarray(x)).cast("B")
+  Device["QCOM"].allocator._copyin(ret._buf,raw); 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("--seed",type=int,default=0); ap.add_argument("--runs",type=int,default=10)
+  ap.add_argument("--disasm",action="store_true"); args=ap.parse_args()
+  if args.m%16 or args.n%256 or args.k%8: raise ValueError("shape does not divide 16x256x8 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); c=upload(np.zeros((args.m,args.stride),np.float16),dtypes.half)
+  dev=Device["QCOM"]; lib=dev.compiler.compile(source(args.k,args.stride))
+  if args.disasm:
+    env,off,size,_=get_envelope(dev,source(args.k,args.stride)); print(disasm(bytes(env[off:off+size])))
+  specs=[((0,dtypes.uint32,(args.m,args.k//8,4)),),((1,dtypes.uint32,(args.k,args.n//8,4)),),((2,dtypes.half,None),)]
+  prg=dev.runtime("packed8",lib,buf_dtypes=specs); gs,ls=(args.n//256,args.m//16,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),np.float16); raw=memoryview(storage).cast("B")
+  Device["QCOM"].allocator._copyout(raw,c._buf)
+  got=storage[:,:args.n].astype(np.float32); expected=av.astype(np.float32)@bv.astype(np.float32); d=np.abs(got-expected); best=min(times)
+  print(f"best_ms={best:.4f} median_ms={statistics.median(times):.4f} gflops={2*args.m*args.n*args.k/best/1e6:.1f} "
+        f"max_abs={d.max():.9g} mean_abs={d.mean():.9g} allclose={np.allclose(got,expected,rtol=.02,atol=.02)}")
+
+
+if __name__=="__main__": main()

extra/gemm/qcom_packed_alias_probe.py

@@ -0,0 +1,36 @@
+#!/usr/bin/env python3
+"""Inspect uint4 image-load to half8 register layout on A630."""
+import numpy as np
+
+from tinygrad import Device, dtypes
+from tinygrad.device import Buffer
+from extra.gemm.ir3asm import NOP, disasm, get_envelope, inject
+
+
+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 probe(read_only image2d_t X,__global half *O) {
+  int i=get_global_id(0); half8 h=as_half8(read_imageui(X,smp,(int2)(i,0))); vstore8(h,0,O+i*8);
+}"""
+
+
+def main() -> None:
+  dev=Device["QCOM"]; lib,off,size,ro=get_envelope(dev,SRC); print(disasm(bytes(lib[off:off+size])))
+  values=np.arange(64*8,dtype=np.float16).view(np.uint32).reshape(1,64,4)
+  src=Buffer("QCOM",values.size,dtypes.uint32).allocate(); out=Buffer("QCOM",128*8,dtypes.half).allocate()
+  src.copyin(memoryview(values).cast("B"))
+  prg=dev.runtime("probe",bytes(lib),buf_dtypes=[((0,dtypes.uint32,(1,64,4)),),((1,dtypes.half,None),)])
+  prg(src._buf,out._buf,global_size=(1,1,1),local_size=(128,1,1),wait=True)
+  got=np.empty(128*8,np.float16); out.copyout(memoryview(got).cast("B")); print("head",got[:16].tolist())
+  shader=bytearray(lib[off:off+size])
+  for i in range(13,21): shader[i*8:(i+1)*8]=NOP()
+  for mode in (None, True, False):
+    patched=inject(lib,off,size,ro,bytes(shader),fregs=2,hregs=2,mergedregs=mode)
+    out.copyin(memoryview(np.zeros(128*8,np.float16)).cast("B"))
+    hand=dev.runtime("probe",patched,buf_dtypes=[((0,dtypes.uint32,(1,64,4)),),((1,dtypes.half,None),)])
+    hand(src._buf,out._buf,global_size=(1,1,1),local_size=(128,1,1),wait=True)
+    out.copyout(memoryview(got).cast("B")); print("direct_alias",mode,got[:16].tolist())
+
+
+if __name__=="__main__": main()

extra/gemm/qcom_parallel_gemm_check.py

@@ -0,0 +1,109 @@
+#!/usr/bin/env python3
+"""Two-stream, arbitrary-input FP16 GEMM benchmark for Adreno 630.
+
+The output rows are split between two independent KGSL contexts.  Timing is
+the wall-clock union of both dispatches; correctness is checked after joining
+the two row partitions into one dense C matrix.
+"""
+from __future__ import annotations
+
+import multiprocessing as mp
+import os, time
+from typing import Any
+
+import numpy as np
+
+
+def worker(rank: int, a: np.ndarray, b: np.ndarray, barrier: Any, start_at: Any, conn: Any, iters: int) -> None:
+  # Import/open the device after fork so every worker owns an independent KGSL
+  # context and command queue.
+  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 get_envelope, inject
+
+  cpus = sorted(os.sched_getaffinity(0))
+  os.sched_setaffinity(0, {cpus[rank % len(cpus)]})
+
+  rows, k, n = a.shape[0], a.shape[1], b.shape[1]
+  threads, stride = 128, n
+  q8.M, q8.N, q8.K, q8.K4 = rows, stride, k, k//4
+  dev = Device["QCOM"]
+  env, io, sz, ro = get_envelope(dev, q8.make_donor_src8(4, threads))
+  shader, hregs, fregs, _ = q8.build_8x8_split_a_unroll_shader(
+    dev, threads, k_unroll=2, b_coord_delay=-1, fast_coords=True,
+    stream_col1=True, add256_store_mode="tight", a_coord_delay=-1,
+    relaxed_sync=True, sync_mask=14, separate_coords=True)
+  assert len(shader) <= sz
+  lib = inject(env, io, sz, ro, shader, fregs=fregs, hregs=hregs)
+
+  ab = Buffer("QCOM", a.size, dtypes.half).allocate()
+  bb = Buffer("QCOM", b.size, dtypes.half).allocate()
+  cb = Buffer("QCOM", rows*stride, dtypes.half).allocate()
+  ab.copyin(memoryview(np.ascontiguousarray(a)).cast("B"))
+  bb.copyin(memoryview(np.ascontiguousarray(b)).cast("B"))
+  cb.copyin(memoryview(np.zeros(rows*stride, dtype=np.float16)).cast("B"))
+  specs = [((0, dtypes.half, (rows, k//4, 4)),), ((0, dtypes.half, (k, n//4, 4)),),
+           ((0, dtypes.half, None),)]
+  prg = dev.runtime("gemm_h", lib, buf_dtypes=specs)
+  launch = dict(global_size=(n//256, rows//32, 1), local_size=(threads, 1, 1), wait=True)
+
+  for _ in range(5): prg(ab._buf, bb._buf, cb._buf, **launch)
+  samples = []
+  for _ in range(iters):
+    barrier.wait()
+    while time.perf_counter() < start_at.value: pass
+    start = time.perf_counter()
+    event_time = prg(ab._buf, bb._buf, cb._buf, **launch)
+    end = time.perf_counter()
+    samples.append((start, end, event_time))
+
+  got = np.empty((rows, stride), dtype=np.float16)
+  cb.copyout(memoryview(got).cast("B"))
+  conn.send((rank, samples, got))
+  conn.close()
+
+
+def main() -> None:
+  m, n, k = (int(os.getenv(x, "1024")) for x in ("M", "N", "K"))
+  seed, iters = int(os.getenv("SEED", "0")), int(os.getenv("ITERS", "10"))
+  if m % 64 or n % 256 or k % 8:
+    raise ValueError("parallel tiled kernel requires M%64 == 0, N%256 == 0, K%8 == 0")
+
+  # Independent streams make B generation invariant to the chosen M split.
+  a = (np.random.default_rng(seed).standard_normal((m, k))*0.05).astype(np.float16)
+  b = (np.random.default_rng(seed+1).standard_normal((k, n))*0.05).astype(np.float16)
+  halves = (a[:m//2], a[m//2:])
+  barrier = mp.Barrier(3)
+  start_at = mp.Value("d", 0.0, lock=False)
+  pipes = [mp.Pipe(duplex=False) for _ in range(2)]
+  procs = [mp.Process(target=worker, args=(rank, halves[rank], b, barrier, start_at, pipes[rank][1], iters)) for rank in range(2)]
+  for proc in procs: proc.start()
+  for _, send_conn in pipes: send_conn.close()
+  for _ in range(iters):
+    start_at.value = time.perf_counter() + 0.01
+    barrier.wait()
+  results = [pipes[rank][0].recv() for rank in range(2)]
+  for proc in procs:
+    proc.join()
+    if proc.exitcode: raise RuntimeError(f"worker exited with status {proc.exitcode}")
+  results.sort(key=lambda x: x[0])
+
+  overlap_times = [max(results[0][1][i][1], results[1][1][i][1]) -
+                   min(results[0][1][i][0], results[1][1][i][0]) for i in range(iters)]
+  best = min(overlap_times)
+  got = np.concatenate((results[0][2], results[1][2]), axis=0).astype(np.float32)
+  expected = a.astype(np.float32) @ b.astype(np.float32)
+  delta = np.abs(expected-got)
+  correct = np.allclose(expected, got, rtol=2e-2, atol=2e-2)
+  bad = ~np.isfinite(got) | (delta > 0.02)
+  event_ms = [[sample[2]*1e3 for sample in result[1]] for result in results]
+  print(f"shape={m}x{n}x{k} streams=2 accumulate=fp16 elapsed_ms={best*1e3:.3f} "
+        f"gflops={2*m*n*k/best/1e9:.1f} max_abs={delta.max():.9g} mean_abs={delta.mean():.9g} "
+        f"allclose={correct} bad_count={bad.sum()}")
+  print(f"worker_event_ms={[round(min(x), 3) for x in event_ms]}")
+  if not correct: raise SystemExit(1)
+
+
+if __name__ == "__main__": main()

extra/gemm/qcom_patch_gemm.py

@@ -0,0 +1,210 @@
+#!/usr/bin/env python3
+"""Clean register remap: identify all accumulators, remap group 3 to be consecutive, apply all 4 rpt3."""
+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 = struct.unpack_from('<I', lib, 0x100)[0]
+shader = bytearray(lib[image_offset:image_offset+image_size])
+total = image_size // 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])
+
+# Map all accumulator registers from MAD instructions
+print("=== ACCUMULATOR ANALYSIS ===")
+for i in range(total):
+    hi, lo = ri(shader, i)
+    if not ((hi >> 24) in (0x63, 0x73) and ((hi >> 24) & 0x0F) == 0x3): continue
+    dst = hi & 0xFF
+    src1 = lo & 0xFF
+    rpt_byte = (hi >> 8) & 0xFF
+    src2_hi = (hi >> 16) & 0xFF
+    src2 = (src2_hi << 1) | ((rpt_byte >> 7) & 1)
+    src3 = (lo >> 16) & 0xFF
+    rpt = rpt_byte & 0x7F
+    sy = (hi >> 28) == 0x7
+    r_flag = (lo >> 29) & 1
+    flags = ("(sy)" if sy else "") + ("(rpt%d)" % rpt if rpt else "")
+    rf = "(r)" if r_flag else ""
+    print("  L%02d: %smad.f16 %s, %s, %s%s, %s%s" % (i, flags, rn(dst), rn(src1), rf, rn(src2), rf, rn(src3)))
+
+# From the analysis, we know:
+# Group 0 (src1=hr0.x=0): accs at hr14.z..hr15.y = 58,59,60,61 - CONSECUTIVE
+# Group 1 (src1=hr0.y=1): accs at hr13.z..hr14.y = 54,55,56,57 - CONSECUTIVE (already rpt3)
+# Group 2 (src1=hr0.z=2): accs at hr12.z..hr13.y = 50,51,52,53 - CONSECUTIVE
+# Group 3 (src1=hr0.w=3): split as (rpt1) hr10.z,hr10.w + (rpt1) hr12.x,hr12.y = 42,43,48,49 - NOT CONSECUTIVE
+
+# Fix: remap 48->44, 49->45 (making group 3 = 42,43,44,45)
+# Check 44 and 45 are free
+used = set()
+for i in range(total):
+    hi, lo = ri(shader, i)
+    if hi == 0 and lo == 0: continue
+    used.add(hi & 0xFF)
+    used.add(lo & 0xFF)
+    used.add((lo >> 16) & 0xFF)
+
+print("\nRegs 44,45 in use:", 44 in used, 45 in used)
+
+# Do the remap
+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)
+
+print("Remapped 48->44, 49->45")
+
+# Now find groups of 4 consecutive non-rpt MADs and convert to rpt3
+patched = 0
+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 = hi0 & 0xFF
+    rpt0 = (hi0 >> 8) & 0x7F
+    src1_0 = lo0 & 0xFF
+    src3_0 = (lo0 >> 16) & 0xFF
+    src2_hi0 = (hi0 >> 16) & 0xFF
+    rpt_byte0 = (hi0 >> 8) & 0xFF
+    src2_0 = (src2_hi0 << 1) | ((rpt_byte0 >> 7) & 1)
+    
+    if rpt0 > 0 or dst0 != src3_0:
+        i += 1; continue
+    
+    # Check next 3
+    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 = hj & 0xFF
+        s1j = lj & 0xFF
+        rpj = (hj >> 8) & 0x7F
+        s3j = (lj >> 16) & 0xFF
+        s2j_hi = (hj >> 16) & 0xFF
+        rpj_byte = (hj >> 8) & 0xFF
+        s2j = (s2j_hi << 1) | ((rpj_byte >> 7) & 1)
+        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)
+        lo_new = lo0 | 0x20000000
+        wi(shader, i, hi_new, lo_new)
+        for j in range(1, 4):
+            wi(shader, i+j, 0, 0)
+        patched += 1
+        print("  rpt3 at line %d: %s" % (i, rn(dst0)))
+        i += 4
+    else:
+        i += 1
+
+# Also merge (rpt1)+(rpt1) into (rpt3) where registers are consecutive
+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 = hi0 & 0xFF
+    dst1 = hi1 & 0xFF
+    src1_0 = lo0 & 0xFF
+    src1_1 = 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: continue
+    if dst1 != dst0 + 2: continue
+    if src2_1 != src2_0 + 2: continue
+    # Merge: change first to rpt3, NOP second
+    rpt_byte_new = ((hi0 >> 8) & 0x80) | 3
+    hi_new = (hi0 & 0xFFFF00FF) | (rpt_byte_new << 8)
+    wi(shader, i, hi_new, lo0)
+    wi(shader, i+1, 0, 0)
+    patched += 1
+    print("  Merged rpt1+rpt1 -> rpt3 at line %d: %s" % (i, rn(dst0)))
+
+print("Total patched: %d groups" % patched)
+
+# Write back
+lib[image_offset:image_offset+image_size] = shader
+
+# 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)
+
+prg = dev.runtime('gemm_h', bytes(lib), [[(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*** RESULT: %.1f GFLOPS (%.0fus) ***" % (gflops, best*1e6))

extra/gemm/qcom_pickle_disasm.py

@@ -0,0 +1,56 @@
+#!/usr/bin/env python3
+"""Disassemble one representative program family from a compiled-model pickle."""
+import argparse, hashlib, pickle, struct
+
+from tinygrad.uop.ops import Ops
+from extra.gemm.ir3asm import disasm
+from extra.gemm.qcom_ir3_matmul_patch import plain_name
+
+
+def main() -> None:
+  parser = argparse.ArgumentParser()
+  parser.add_argument("model")
+  parser.add_argument("family")
+  parser.add_argument("--start", type=int, default=0)
+  parser.add_argument("--stop", type=int)
+  parser.add_argument("--index", type=int, default=0)
+  parser.add_argument("--source", action="store_true")
+  parser.add_argument("--meta", action="store_true")
+  parser.add_argument("--list", action="store_true")
+  parser.add_argument("--calls", help="list graph calls in a start:stop index range")
+  args = parser.parse_args()
+  with open(args.model, "rb") as f: model = pickle.load(f)
+  batch = model.captured.linear.src[0].src[0].src[0].src
+  if args.calls:
+    start, stop = (int(x) for x in args.calls.split(":"))
+    for i, call in enumerate(batch[start:stop], start):
+      if call.op is not Ops.CALL or call.src[0].op is not Ops.PROGRAM: continue
+      def desc(x):
+        return f"slot={getattr(x.arg, 'slot', '?')} dtype={x.dtype} size={x.buffer.size}" if x.op is Ops.BUFFER else str(x.op)
+      print(i, plain_name(call.src[0].arg.name), "outs", call.src[0].arg.outs,
+            "args", [desc(x) for x in call.src[1:]])
+    return
+  programs = [call.src[0] for call in batch if call.op is Ops.CALL and call.src[0].op is Ops.PROGRAM and
+              plain_name(call.src[0].arg.name) == args.family]
+  if args.list:
+    for i, item in enumerate(programs): print(i, item.arg.global_size, hashlib.sha1(item.src[3].arg).hexdigest()[:8])
+    return
+  program = programs[args.index]
+  if args.meta:
+    from tinygrad.engine.realize import get_runtime
+    runtime = get_runtime("QCOM", program)
+    print("global", program.arg.global_size, "local", program.arg.local_size, "globals", program.arg.globals)
+    print("aux", program.arg.aux)
+    print("buf_offs", runtime.buf_offs, "tex", runtime.tex_cnt, "ibo", runtime.ibo_cnt, "nir", runtime.NIR)
+    return
+  if args.source:
+    print(program.src[2].arg)
+    return
+  lib = program.src[3].arg
+  image_off, image_size = struct.unpack_from("<I", lib, 0xc0)[0], struct.unpack_from("<I", lib, 0x100)[0]
+  lines = [line for line in disasm(lib[image_off:image_off+image_size]).splitlines() if not line.rstrip().endswith(":")]
+  stop = len(lines) if args.stop is None else args.stop
+  for i, line in enumerate(lines[args.start:stop], args.start): print(f"{i:3d}: {line}")
+
+
+if __name__ == "__main__": main()

extra/gemm/qcom_planar_b_check.py

@@ -0,0 +1,148 @@
+#!/usr/bin/env python3
+"""Randomized GEMM using four planar B textures to remove coordinate churn."""
+import os
+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 *
+from extra.gemm.ir3asm import _hreg
+
+
+def make_envelope_src(k4: int, n: int, threads: int = 128) -> str:
+  src = 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({threads},1,1)))
+__kernel void gemm_h(write_only image2d_t C,read_only image2d_t A,read_only image2d_t B0,
+                     read_only image2d_t B1,read_only image2d_t B2,read_only image2d_t B3) {{
+  int lid=get_local_id(0), row=get_group_id(1)*32+(lid>>5)*8;
+  int col4=get_group_id(0)*64+(lid&31);
+  half4 c[16]; for(int i=0;i<16;i++) c[i]=(half4)(0);
+  for(int k4=0;k4<{k4};k4++) {{
+    half4 a0=read_imageh(A,smp,(int2)(k4,row));
+    half4 b0=read_imageh(B0,smp,(int2)(col4,k4));
+    half4 b1=read_imageh(B1,smp,(int2)(col4,k4));
+    half4 b2=read_imageh(B2,smp,(int2)(col4,k4));
+    half4 b3=read_imageh(B3,smp,(int2)(col4,k4));
+    c[0]+=a0.xxxx*b0+a0.yyyy*b1+a0.zzzz*b2+a0.wwww*b3;
+  }}
+  for(int i=0;i<16;i++) write_imageh(C,(int2)(col4+(i&1)*32,row+(i>>1)),c[i]);
+}}'''
+  pad = ''.join('  c[0]=mad(c[0],(half4)(1.0009765625h),(half4)(0.0009765625h));\n' for _ in range(512))
+  return src.replace('  for(int i=0;i<16;i++) write_imageh', pad+'  for(int i=0;i<16;i++) write_imageh')
+
+
+def build_shader(dev, threads: int, k4: int, unroll: int, no_store: bool = False, control_only: bool = False,
+                 post_constant: bool = False) -> tuple[bytes, int, int]:
+  instrs = q8.prologue_8x4(dev, threads)
+  q8.emit_col_stride(instrs, 2)
+  kz, row, col, col1 = 'r6.z', 'r7.x', 'r7.y', 'r6.w'
+  acc0 = _hreg('hr12.x')
+  for base in range(acc0, acc0+64, 4): q8.emit_hvec_imm(instrs, base, 0)
+  instrs += [ADD_S(col1, col, 32),
+             MOV_F32('r14.x', col), MOV_F32('r14.z', col1)]
+  # Four independent sampler coordinate pairs avoid source WAR hazards between
+  # the four outstanding A reads.  They are reused only after the first MAD group.
+  a_coords = (('r15.x', 'r15.y'), ('r15.z', 'r15.w'),
+              ('r16.x', 'r16.y'), ('r16.z', 'r16.w'))
+  b_regs = [[_hreg(f'hr{i}.x') for i in range(4)], [_hreg(f'hr{i}.x') for i in range(4, 8)]]
+  a_regs = [_hreg(f'hr{i}.x') for i in range(8, 12)]
+
+  def mad(r: int, c: int, kk: int, sy: bool = False) -> None:
+    dst = acc0+(r*2+c)*4
+    instrs.append(MAD_F16(dst, a_regs[r&3]+kk, b_regs[c][kk], dst, rpt=3, sy=sy, r=True))
+
+  loop_start = len(instrs)
+  for ku in range(unroll):
+    if ku: instrs += [ADD_S(kz, kz, 1), NOP(rpt=2)]
+    if control_only: continue
+    instrs += [MOV_F32('r14.y', kz), MOV_F32('r14.w', kz)]
+    for c, cr in enumerate(('r14.x', 'r14.z')):
+      for kk in range(4): instrs.append(ISAM_F16(b_regs[c][kk], cr, 1+kk))
+    instrs += [MOV_F32(ac[0], kz) for ac in a_coords]
+    for r in range(4):
+      ac = a_coords[r]
+      instrs += [MOV_F32(ac[1], row) if r == 0 else ADD_S(ac[1], row, r), NOP(rpt=int(os.getenv('ACOORD_DELAY', '5'))),
+                 ISAM_F16(a_regs[r], ac[0], 0)]
+    first = True
+    for r in range(4):
+      for c in range(2):
+        for kk in range(4):
+          mad(r, c, kk, first)
+          first = False
+    for r in range(4):
+      ac = a_coords[r]
+      instrs += [ADD_S(ac[1], row, r+4), NOP(rpt=int(os.getenv('ACOORD_DELAY', '5'))), ISAM_F16(a_regs[r], ac[0], 0)]
+    first = True
+    for r in range(4, 8):
+      for c in range(2):
+        for kk in range(4):
+          mad(r, c, kk, first)
+          first = False
+  instrs += [ADD_S('r0.x', kz, 1), NOP(rpt=2), CMPS_S_EQ(kz, k4-1, nop=1), MOV_F32(kz, 'r0.x'), NOP(rpt=3)]
+  loop_end = len(instrs)
+  instrs.append(BR(loop_start-loop_end))
+  if post_constant:
+    for base in range(acc0, acc0+64, 4): q8.emit_hvec_imm(instrs, base, 0x6400)
+  if not no_store:
+    for r in range(8):
+      for c in range(2):
+        instrs += [MOV_F32('r14.x', col) if c == 0 else ADD_S('r14.x', col, 32),
+                   MOV_F32('r14.y', row) if r == 0 else ADD_S('r14.y', row, r),
+                   COV_F16F32('r19.x', acc0+(r*2+c)*4, sy=True, rpt=3, r=True), NOP(rpt=5),
+                   STIB_F32('r19.x', 'r14.x'), NOP(rpt=16)]
+  instrs.append(END())
+  return assemble(instrs), 28, 17
+
+
+def main() -> None:
+  m, n, k = (int(os.getenv(x, '1024')) for x in ('M','N','K'))
+  seed, threads, unroll = int(os.getenv('SEED','0')), 128, int(os.getenv('KUNROLL','4'))
+  if m%32 or n%256 or k%4 or (k//4)%unroll: raise ValueError('unsupported tile shape')
+  rng = np.random.default_rng(seed)
+  a = (rng.standard_normal((m,k))*0.05).astype(np.float16)
+  b = (rng.standard_normal((k,n))*0.05).astype(np.float16)
+  planes = [np.ascontiguousarray(b[p::4]) for p in range(4)]
+  q8.M=q4.M=m; q8.N=q4.N=n; q8.K=q4.K=k; q8.K4=q4.K4=k//4
+  dev = Device['QCOM']
+  env, io, sz, ro = get_envelope(dev, make_envelope_src(k//4, n, threads))
+  no_store = bool(int(os.getenv('NO_STORE', '0')))
+  build_k4 = int(os.getenv('BUILD_K4', str(k//4)))
+  shader, hregs, fregs = build_shader(dev, threads, build_k4, unroll, no_store,
+                                      bool(int(os.getenv('CONTROL_ONLY', '0'))), bool(int(os.getenv('POST_CONSTANT', '0'))))
+  if len(shader)>sz: raise ValueError(f'shader {len(shader)} > envelope {sz}')
+  if int(os.getenv('DUMP', '0')):
+    print(f'shader_bytes={len(shader)} envelope_bytes={sz} fregs={fregs} hregs={hregs}')
+    print(disasm(shader))
+    return
+  lib = bytes(env) if int(os.getenv('COMPILER', '0')) else \
+        inject(env, io, sz, ro, shader, fregs=int(os.getenv('FREGS',str(fregs))), hregs=int(os.getenv('HREGS',str(hregs))),
+               mergedregs=False if int(os.getenv('SEPARATE_REGS', '0')) else None)
+  ab=Buffer('QCOM',a.size,dtypes.half).allocate(); ab.copyin(memoryview(a).cast('B'))
+  pbs=[]
+  for p in planes:
+    pb=Buffer('QCOM',p.size,dtypes.half).allocate(); pb.copyin(memoryview(p).cast('B')); pbs.append(pb)
+  cb=Buffer('QCOM',m*n,dtypes.half).allocate(); cb.copyin(memoryview(np.zeros((m,n),np.float16)).cast('B'))
+  specs=[((0,dtypes.half,(m,n//4,4)),),((1,dtypes.half,(m,k//4,4)),),((2,dtypes.half,(k//4,n//4,4)),),
+         ((3,dtypes.half,(k//4,n//4,4)),),((4,dtypes.half,(k//4,n//4,4)),),((5,dtypes.half,(k//4,n//4,4)),)]
+  prg=dev.runtime('gemm_h',lib,buf_dtypes=specs)
+  args=(cb._buf,ab._buf,pbs[0]._buf,pbs[1]._buf,pbs[2]._buf,pbs[3]._buf)
+  times=[prg(*args,global_size=(n//256,m//32,1),local_size=(threads,1,1),wait=True) for _ in range(10)]
+  if no_store:
+    print(f'no_store elapsed_ms={min(times)*1e3:.3f}')
+    return
+  got=np.empty((m,n),np.float16); cb.copyout(memoryview(got).cast('B'))
+  expected=(np.full((m,n),1024,np.float32) if int(os.getenv('POST_CONSTANT','0')) else
+            a[:,:build_k4*4].astype(np.float32)@b[:build_k4*4].astype(np.float32)); delta=np.abs(expected-got.astype(np.float32))
+  correct=np.allclose(expected,got,rtol=2e-2,atol=2e-2); best=min(times)
+  print(f'shape={m}x{n}x{k} planar_b=4 accumulate=fp16 elapsed_ms={best*1e3:.3f} gflops={2*m*n*k/best/1e9:.1f} '
+        f'max_abs={delta.max():.9g} mean_abs={delta.mean():.9g} allclose={correct} bad_count={(delta>.02).sum()}')
+  if not correct and int(os.getenv('DEBUG', '0')):
+    print('bad_by_row', np.count_nonzero(delta > .02, axis=1).tolist())
+    print('max_by_row', delta.max(axis=1).tolist())
+  if not correct: raise SystemExit(1)
+
+
+if __name__=='__main__': main()

extra/gemm/qcom_pointer_add_schedule.py

@@ -0,0 +1,17 @@
+#!/usr/bin/env python3
+"""Disassemble a compiler-generated 64-bit global-pointer increment."""
+import struct
+
+from tinygrad import Device
+from extra.gemm.ir3asm import disasm
+
+SRC = r"""#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+__attribute__((reqd_work_group_size(128,1,1)))
+__kernel void stores(__global half *O) {
+  int t=get_global_id(0); __global half *p=O+t*4;
+  vstore4((half4)(1),0,p); vstore4((half4)(2),0,p+128);
+}"""
+
+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]))

extra/gemm/qcom_precision_correction.py

@@ -0,0 +1,55 @@
+#!/usr/bin/env python3
+"""Measure whether a fast model's output error admits a cross-validated correction."""
+import argparse, pickle
+
+import numpy as np
+
+from tinygrad import Tensor
+
+
+def main() -> None:
+  parser = argparse.ArgumentParser()
+  parser.add_argument("model")
+  parser.add_argument("corpus")
+  args = parser.parse_args()
+  with open(args.model, "rb") as f: model = pickle.load(f)
+  corpus = np.load(args.corpus)
+  expected, actual = [], []
+  for case, seed in enumerate(corpus["seeds"].tolist()):
+    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()
+    got = np.asarray(model(**inputs).numpy(), dtype=np.float32).reshape(-1)
+    ref = corpus[f"case{case}:output"].astype(np.float32).reshape(-1)
+    actual.append(got.copy())
+    expected.append(ref)
+    print(f"captured case={case} seed={seed}")
+  x, y = np.stack(actual), np.stack(expected)
+  raw = np.abs(y-x)
+  print(f"raw max={float(raw.max()):.9g} mean={float(raw.mean()):.9g}")
+  for case in np.argsort(raw.max(axis=1))[-5:]:
+    worst = int(raw[case].argmax())
+    print(f"raw_case={int(case)} worst_index={worst} expected={float(y[case, worst]):.9g} "
+          f"actual={float(x[case, worst]):.9g} max={float(raw[case, worst]):.9g} "
+          f"bias_mean={float((y[case]-x[case]).mean()):.9g}")
+
+  corrected_mean, corrected_affine = [], []
+  for held_out in range(len(x)):
+    train = np.arange(len(x)) != held_out
+    bias = (y[train]-x[train]).mean(axis=0)
+    corrected_mean.append(y[held_out]-(x[held_out]+bias))
+    xm, ym = x[train].mean(axis=0), y[train].mean(axis=0)
+    covariance = ((x[train]-xm)*(y[train]-ym)).sum(axis=0)
+    variance = ((x[train]-xm)**2).sum(axis=0)
+    slope = np.divide(covariance, variance, out=np.ones_like(covariance), where=variance > 1e-12)
+    intercept = ym-slope*xm
+    corrected_affine.append(y[held_out]-(slope*x[held_out]+intercept))
+  for name, residual in (("mean", corrected_mean), ("affine", corrected_affine)):
+    error = np.abs(np.stack(residual))
+    maxima = error.max(axis=1)
+    print(f"loo_{name} max={float(error.max()):.9g} mean={float(error.mean()):.9g} "
+          f"passing={int((maxima <= 0.01).sum())}/{len(maxima)} per_case={maxima.tolist()}")
+
+
+if __name__ == "__main__": main()

extra/gemm/qcom_prefix_check.py

@@ -0,0 +1,74 @@
+#!/usr/bin/env python3
+"""Compare baseline and replacement at the first target-5 boundary on a real model input."""
+import 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 inputs_for(model):
+  corpus = np.load("/tmp/openpilot_onnx_scale8_5seed.npz")
+  inputs = {}
+  for name, (view, _vars, dtype, device) in zip(model.captured.expected_names, model.captured.expected_input_info):
+    inputs[name] = Tensor(corpus[f"case0:input:{name}"].astype(np.dtype(dtype.fmt), copy=False), device=device).realize()
+  return _prepare_jit_inputs((), inputs)[:2]
+
+
+def batch(model): return model.captured.linear.src[0].src[0].src[0].src
+
+
+def main() -> None:
+  base = load("/data/openpilot_target1_fp32hand.pkl")
+  cand = load("/data/openpilot_target1_wide5_fp32_v4.pkl")
+  input_uops, var_vals = inputs_for(base)
+  bb = batch(base)
+  bi = next(i for i, x in enumerate(bb) if x.op is Ops.CALL and x.src[0].op is Ops.PROGRAM and
+            plain_name(x.src[0].arg.name) == "r_48_128_4_4_96_4")
+  run_linear(UOp(Ops.LINEAR, src=tuple(bb[:bi+1])), var_vals, input_uops=input_uops, jit=True, wait=True)
+  bo = np.array(bb[bi].src[1].buffer.numpy(), copy=True)
+  ba = np.array(bb[bi].src[4].buffer.numpy(), copy=True)
+  br = np.array(bb[bi].src[2].buffer.numpy(), copy=True)
+
+  cb = batch(cand)
+  ci = next(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) == "gemm_f" and x.src[0].arg.aux[0][0][0][2] == (48, 384, 4))
+  assert plain_name(cb[ci+1].src[0].arg.name) == "epi2_fp32"
+  run_linear(UOp(Ops.LINEAR, src=tuple(cb[:ci+2])), var_vals, input_uops=input_uops, jit=True, wait=True)
+  co = np.array(cb[ci+1].src[1].buffer.numpy(), copy=True)
+  ca = np.array(cb[ci].src[1].buffer.numpy(), copy=True)
+  cr = np.array(cb[ci+1].src[2].buffer.numpy(), copy=True)
+  delta = np.abs(bo-co)
+  print("indices", bi, ci, "max_abs", float(delta.max()), "mean_abs", float(delta.mean()),
+        "allclose", bool(np.allclose(bo, co, rtol=1e-5, atol=1e-5)))
+  print("samples", bo[:8].tolist(), co[:8].tolist())
+  print("activation", float(np.abs(ba-ca).max()), "residual", float(np.abs(br-cr).max()))
+  print("activation_stats", float(ba.min()), float(ba.max()), float(np.mean(np.abs(ba))),
+        "residual_stats", float(br.min()), float(br.max()))
+  for name, u in (("activation", cb[ci].src[1]), ("temporary", cb[ci].src[3]),
+                  ("output", cb[ci+1].src[1]), ("residual", cb[ci+1].src[2])):
+    b = u.buffer._buf
+    print(name, hex(b.va_addr), b.size)
+  tmp = np.array(cb[ci].src[3].buffer.numpy(), copy=True).reshape(192, 1024)[:, :512]
+  w = np.array(cb[ci].src[2].buffer.numpy(), copy=True).reshape(384, 512).astype(np.float32)
+  a = ca.reshape(48, 384, 4).transpose(0, 2, 1).reshape(192, 384).astype(np.float32)
+  td = np.abs(tmp-a@w)
+  print("temporary_cpu", float(td.max()), float(td.mean()), np.unravel_index(td.argmax(), td.shape))
+  for name, u in (("base_output", bb[bi].src[1]), ("base_residual", bb[bi].src[2]),
+                  ("base_weight", bb[bi].src[3]), ("base_activation", bb[bi].src[4])):
+    b = u.buffer._buf
+    print(name, hex(b.va_addr), b.size)
+  np.save("/tmp/openpilot_target5_activation.npy", ba)
+  np.save("/tmp/openpilot_target5_residual.npy", br)
+  np.save("/tmp/openpilot_target5_baseline_output.npy", bo)
+
+
+if __name__ == "__main__": main()

extra/gemm/qcom_recompile_graph.py

@@ -0,0 +1,42 @@
+#!/usr/bin/env python3
+"""Recompile every source-backed PROGRAM in a captured QCOM graph.
+
+This preserves the scheduled graph, static buffers, and mixed-precision
+boundaries while allowing an apples-to-apples comparison of QCOM compilers.
+"""
+import argparse, itertools, pickle
+
+from tinygrad import Device, Context
+from tinygrad.uop.ops import Ops, UOp
+
+
+def main() -> None:
+  parser = argparse.ArgumentParser()
+  parser.add_argument("input")
+  parser.add_argument("output")
+  parser.add_argument("--compiler-device", default="QCOM:IR3")
+  args = parser.parse_args()
+
+  with open(args.input, "rb") as f: jit = pickle.load(f)
+  existing_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(existing_slots, default=-1) + 1)
+
+  # The suffix selects a renderer/compiler target, not a physical device ID.
+  # Open the base device while that target is active.
+  with Context(DEV=args.compiler_device): compiler = Device[args.compiler_device.split(":", 1)[0]].compiler
+  programs = [x for x in jit.captured.linear.toposort() if x.op is Ops.PROGRAM]
+  replacements, binaries = {}, {}
+  for number, program in enumerate(programs, 1):
+    source = program.src[2].arg
+    if source not in binaries:
+      print(f"compiling source {len(binaries)+1}: {program.arg.name} ({len(source)} bytes)", flush=True)
+      binaries[source] = compiler.compile_cached(source)
+      print(f"compiled {len(binaries)} unique sources ({number}/{len(programs)} programs)", flush=True)
+    replacements[program] = program.replace(src=program.src[:3] + (program.src[3].replace(arg=binaries[source]),))
+  jit.captured._linear = jit.captured.linear.substitute(replacements, walk=True)
+  with open(args.output, "wb") as f: pickle.dump(jit, f)
+  print(f"wrote {args.output}: {len(programs)} programs, {len(binaries)} unique sources")
+
+
+if __name__ == "__main__": main()

extra/gemm/qcom_replace_up4x16.py

@@ -0,0 +1,59 @@
+#!/usr/bin/env python3
+"""Replace cached openpilot FFN-up GEMMs with the random-checked 4x16 kernel."""
+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 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 build_program(template:UOp, lib:bytes) -> UOp:
+  specs = ((dtypes.float, (128, 512, 4)), (dtypes.half, (128, 96, 4)), (dtypes.half, (384, 384, 4)))
+  aux = (tuple(((i, dtype, shape),) for i, (dtype, shape) in enumerate(specs)),)
+  # The QCOM ELF argument parser locates descriptors using the embedded symbol name length.
+  # Keep the donor's `gemm_h` name or its image arguments are parsed as empty.
+  info = replace(template.arg, name="gemm_h", global_size=(3, 8, 1), local_size=(128, 1, 1),
+                 globals=(0, 1, 2), outs=(0,), ins=(1, 2), aux=aux)
+  return template.replace(arg=info, src=template.src[:2] +
+                          (template.src[2].replace(arg="random-checked separate-bank 4x16 FP16 GEMM"),
+                           template.src[3].replace(arg=lib)))
+
+
+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: 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
+  targets = [call for call in 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 tuple(call.src[0].arg.global_size) == (12, 8, 1)]
+  if len(targets) != 18: raise RuntimeError(f"expected 18 FFN-up GEMMs, found {len(targets)}")
+
+  q.M, q.N, q.K, q.K4 = 128, 1536, 384, 96
+  dev = Device["QCOM"]
+  envelope, image_off, image_size, reg_off = 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, image_store=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, separate_b_coords=True)
+  lib = inject(envelope, image_off, image_size, reg_off, shader, fregs=10, hregs=28, mergedregs=False)
+  program = build_program(targets[0].src[0], lib)
+  replacements = {call: program.call(call.src[3], call.src[1], call.src[2]) for call in targets}
+  new_batch = [replacements.get(call, call) for call in batch]
+  new_outer = create_graph_call(new_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"replaced {len(targets)} FFN-up GEMMs")
+
+
+if __name__ == "__main__": main()

extra/gemm/qcom_shfl_probe.py

@@ -0,0 +1,404 @@
+#!/usr/bin/env python3
+"""Subgroup shuffle probes for Adreno 630 ir3.
+
+This keeps the same image+buffer donor envelope as the GEMM experiments, but
+only uses the C buffer. The semantic probe stores per-lane u32 values so the
+exact source lane selected by each shfl mode is visible.
+"""
+import argparse, ctypes, struct
+
+from tinygrad import Device, dtypes
+from tinygrad.device import Buffer
+from extra.gemm.ir3asm import *
+
+M = N = K = 1024
+K4 = K // 4
+
+
+def make_donor_src(ncols=1, threads=128):
+    tn = 32 * ncols
+    tm = (threads // 32) * 4
+    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(%d,1,1)))\n' % threads
+    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 tid=lid&31;\n'
+    src += 'int row=get_group_id(1)*%d+tm*4;int col4=get_group_id(0)*%d+tid;\n' % (tm, tn)
+    for r in range(4):
+        for c in range(ncols):
+            src += 'half4 r%dd%dc0=(half4)(0),r%dd%dc1=(half4)(0),r%dd%dc2=(half4)(0),r%dd%dc3=(half4)(0);\n' % (r,c,r,c,r,c,r,c)
+    src += 'for(int k4=0;k4<%d;k4++){\n' % K4
+    for r in range(4): src += 'half4 a%d=read_imageh(A,smp,(int2)(k4,row+%d));\n' % (r, r)
+    for c in range(ncols):
+        for b in range(4): src += 'half4 b%d_%d=read_imageh(B,smp,(int2)(col4+%d,k4*4+%d));\n' % (c, b, c*32, b)
+    for r in range(4):
+        for c in range(ncols):
+            src += 'r%dd%dc0+=a%d.xxxx*b%d_0;r%dd%dc1+=a%d.yyyy*b%d_1;r%dd%dc2+=a%d.zzzz*b%d_2;r%dd%dc3+=a%d.wwww*b%d_3;\n' % (r,c,r,c,r,c,r,c,r,c,r,c,r,c,r,c)
+    src += '}\n'
+    for r in range(4):
+        for c in range(ncols):
+            src += 'vstore4(r%dd%dc0+r%dd%dc1+r%dd%dc2+r%dd%dc3,0,C+(row+%d)*%d+(col4+%d)*4);\n' % (r,c,r,c,r,c,r,c,r,N,c*32)
+    src += '}\n'
+    return src
+
+
+def prologue(dev, threads):
+    lib, io, isz, _ = get_envelope(dev, make_donor_src(1, threads))
+    pro = bytearray(lib[io:io + 21 * 8])
+    return [bytes(pro[i:i+8]) for i in range(0, len(pro), 8)]
+
+
+def emit_addr(instrs, row_reg, col_reg):
+    instrs += [
+        SHL_B('r0.x', row_reg, 10, jp=True),
+        SHL_B('r0.y', col_reg, 2),
+        ADD_S_REG('r0.x', 'r0.x', 'r0.y'),
+        SHL_B('r0.x', 'r0.x', 1),
+        ADD_U('r2.x', 'c20.x', 'r0.x'),
+        CMPS_U_LT('r6.w', 'r2.x', 'c20.x'),
+        SHR_B('r6.y', 'r0.x', 31),
+        SAD_S32('r2.y', 'c20.y', 'r6.y', 'r6.w', nop=3),
+    ]
+
+
+def store_u32(instrs, row_reg, col_reg, data_reg):
+    emit_addr(instrs, row_reg, col_reg)
+    instrs += [STG_U32('r2.x', data_reg), NOP()]
+
+
+def emit_linear_addr(instrs, idx_reg, byte_base=0, join=False):
+    instrs.append(SHL_B('r0.y', idx_reg, 2, jp=join, ss=join, nop=3 if join else 0))
+    if byte_base:
+        instrs += [MOV_S32('r0.z', byte_base), NOP(rpt=2), ADD_S_REG('r0.y', 'r0.y', 'r0.z'), NOP(rpt=16)]
+    instrs += [
+        ADD_U('r2.x', 'c20.x', 'r0.y'),
+        NOP(rpt=16),
+        CMPS_U_LT('r6.w', 'r2.x', 'c20.x'),
+        SHR_B('r6.y', 'r0.y', 31),
+        NOP(rpt=16),
+        SAD_S32('r2.y', 'c20.y', 'r6.y', 'r6.w', nop=3),
+    ]
+
+
+def store_u32_linear(instrs, idx_reg, data_reg, byte_base=0, join=False):
+    emit_linear_addr(instrs, idx_reg, byte_base, join=join)
+    instrs += [NOP(rpt=16), STG_U32('r2.x', data_reg, sy=True), NOP(rpt=16)]
+
+
+def build_semantics_shader(dev, threads):
+    instrs = [
+        AND_B('r8.x', 'r0.x', 31),
+    ]
+    instrs += [
+        SHFL('r8.y', 'r8.x', 0, mode=7, typ=3),
+        SHFL('r8.z', 'r8.x', 1, mode=7, typ=3),
+        SHFL('r8.w', 'r8.x', 2, mode=7, typ=3),
+        SHFL('r9.x', 'r8.x', 4, mode=7, typ=3),
+        SHFL('r9.y', 'r8.x', 8, mode=7, typ=3),
+        SHFL('r9.z', 'r8.x', 16, mode=7, typ=3),
+        MOV_S32('r10.x', 16),
+        SHFL('r9.w', 'r8.x', 'r10.x', mode=1, typ=3),
+        MOV_S32('r10.x', 0),
+        QUAD_BRCST('r11.x', 'r8.x', 'r10.x', typ=3),
+        MOV_S32('r10.y', 1),
+        QUAD_BRCST('r11.y', 'r8.x', 'r10.y', typ=3),
+        MOV_S32('r10.z', 2),
+        QUAD_BRCST('r11.z', 'r8.x', 'r10.z', typ=3),
+        MOV_S32('r10.w', 3),
+        QUAD_BRCST('r11.w', 'r8.x', 'r10.w', typ=3),
+        NOP(rpt=3),
+    ]
+    regs = ['r8.x', 'r8.y', 'r8.z', 'r8.w', 'r9.x', 'r9.y', 'r9.z', 'r9.w', 'r11.x', 'r11.y', 'r11.z', 'r11.w']
+    for case, reg in enumerate(regs):
+        store_u32_linear(instrs, 'r8.x', reg, case * threads * 4)
+    instrs.append(END())
+    return assemble(instrs)
+
+
+def build_quad_semantics_shader(dev, threads):
+    instrs = prologue(dev, threads)
+    instrs += [
+        # Recover lid=(row-tile lane)*32+column lane from persistent coordinates.
+        SHR_B('r12.x', 'r7.x', 2), AND_B('r12.x', 'r12.x', 3), SHL_B('r12.x', 'r12.x', 5),
+        AND_B('r12.y', 'r7.y', 31), ADD_S_REG('r12.x', 'r12.x', 'r12.y'), NOP(rpt=2),
+        MOV_F32('r8.x', 'r12.x'),
+        MOV_S32('r10.x', 0),
+        QUAD_BRCST('r11.x', 'r8.x', 'r10.x', typ=3),
+        QUAD_BRCST('r13.x', 'r8.x', 'r10.x', typ=3, sy=True),
+        MOV_S32('r10.y', 1),
+        QUAD_BRCST('r11.y', 'r8.x', 'r10.y', typ=3),
+        MOV_S32('r10.z', 2),
+        QUAD_BRCST('r11.z', 'r8.x', 'r10.z', typ=3),
+        MOV_S32('r10.w', 3),
+        QUAD_BRCST('r11.w', 'r8.x', 'r10.w', typ=3),
+        NOP(rpt=5),
+    ]
+    regs = ['r8.x', 'r11.x', 'r13.x', 'r11.y', 'r11.z', 'r11.w']
+    for case, reg in enumerate(regs):
+        store_u32_linear(instrs, 'r12.x', reg, case * threads * 4)
+    instrs.append(END())
+    return assemble(instrs)
+
+
+def build_quad_map_shader(dev, threads):
+    instrs = []
+    instrs += [
+        MOV_F32('r14.x', 'r0.x'),
+        MOV_F32('r0.z', 'r0.x'),
+        AND_B('r0.x', 'r0.z', 31, nop=3),
+        MOV_S32('r9.x', 1),
+        NOP(rpt=5),
+    ]
+    store_u32_linear(instrs, 'r14.x', 'r9.x', 0)
+    instrs += [
+        MOV_S32('r10.x', 0),
+        QUAD_BRCST('r1.x', 'r0.x', 'r10.x', typ=3, sy=True),
+        MOV_S32('r10.x', 1),
+        QUAD_BRCST('r1.y', 'r0.x', 'r10.x', typ=3),
+        MOV_S32('r10.x', 2),
+        QUAD_BRCST('r1.z', 'r0.x', 'r10.x', typ=3),
+        MOV_S32('r10.x', 3),
+        QUAD_BRCST('r1.w', 'r0.x', 'r10.x', typ=3),
+        NOP(rpt=5),
+    ]
+    for case, reg in enumerate(['r1.x', 'r1.y', 'r1.z', 'r1.w'], start=1):
+        store_u32_linear(instrs, 'r14.x', reg, case * threads * 4)
+    instrs.append(END())
+    return assemble(instrs)
+
+
+def build_modes_shader(dev, threads):
+    instrs = prologue(dev, threads)
+    instrs += [MOV_F32('r8.x', 'r0.x'), MOV_F32('r15.x', 'r8.x')]
+    cases = []
+    dsts = ['r8.y', 'r8.z', 'r8.w', 'r9.x', 'r9.y', 'r9.z', 'r9.w', 'r10.x', 'r10.y', 'r10.z', 'r10.w', 'r11.x', 'r11.y', 'r11.z', 'r11.w', 'r12.x', 'r12.y', 'r12.z', 'r12.w', 'r13.x']
+    for mode in [1, 2, 3, 6, 7]:
+        for idx in [1, 2, 3, 4]:
+            if not dsts: break
+            dst = dsts.pop(0)
+            instrs.append(SHFL(dst, 'r8.x', idx, mode=mode, typ=3))
+            cases.append((f'm{mode}i{idx}', dst))
+        if not dsts: break
+    instrs.append(NOP(rpt=5))
+    for case, (_, reg) in enumerate(cases): store_u32_linear(instrs, 'r15.x', reg, case * threads * 4)
+    instrs.append(END())
+    return assemble(instrs), [name for name, _ in cases]
+
+
+def build_bench_shader(dev, threads, op, kind, ops_per_iter):
+    instrs = [
+        MOV_S32('r6.z', 0, sy=True),
+        MOV_F32('r14.x', 'r0.x'),
+        AND_B('r8.x', 'r0.x', 31),
+    ]
+    instrs += [
+        MOV_F32('r8.y', 'r8.x'), MOV_F32('r8.z', 'r8.x'), MOV_F32('r8.w', 'r8.x'),
+        MOV_F32('r9.x', 'r8.x'), MOV_F32('r9.y', 'r8.x'),
+        MOV_F32('r9.z', 'r8.x'), MOV_F32('r9.w', 'r8.x'),
+    ]
+    if op == 'quad': instrs.append(MOV_S32('r15.x', 0))
+    loop_start = len(instrs)
+    if kind == 'chain':
+        for _ in range(ops_per_iter):
+            instrs.append(SHFL('r8.x', 'r8.x', 1, mode=7, typ=3) if op == 'shfl' else QUAD_BRCST('r8.x', 'r8.x', 'r15.x', typ=3))
+    elif kind == 'throughput':
+        dsts = ['r10.x', 'r10.y', 'r10.z', 'r10.w', 'r11.x', 'r11.y', 'r11.z', 'r11.w',
+                'r12.x', 'r12.y', 'r12.z', 'r12.w', 'r13.x', 'r13.y', 'r13.z', 'r13.w']
+        srcs = ['r8.x', 'r8.y', 'r8.z', 'r8.w', 'r9.x', 'r9.y', 'r9.z', 'r9.w']
+        for i in range(ops_per_iter):
+            instrs.append(SHFL(dsts[i % len(dsts)], srcs[i % len(srcs)], 1, mode=7, typ=3) if op == 'shfl' else QUAD_BRCST(dsts[i % len(dsts)], srcs[i % len(srcs)], 'r15.x', typ=3))
+    else:
+        raise ValueError(kind)
+    instrs += [
+        ADD_S('r0.x', 'r6.z', 1),
+        CMPS_S_EQ('r6.z', K4 - 1, nop=1),
+        MOV_F32('r6.z', 'r0.x'),
+        NOP(rpt=3),
+    ]
+    loop_end = len(instrs)
+    instrs.append(BR(loop_start - loop_end))
+    store_u32_linear(instrs, 'r14.x', 'r8.x' if kind == 'chain' else 'r10.x')
+    instrs.append(END())
+    return assemble(instrs), loop_end - loop_start
+
+
+def build_branch_shader(dev, threads):
+    instrs = [
+        MOV_F32('r14.x', 'r0.x'),
+        MOV_F32('r0.z', 'r0.x'),
+        AND_B('r8.x', 'r0.z', 31),
+        MOV_S32('r9.x', 0),
+        AND_B('r0.x', 'r0.z', 3, nop=3),
+        CMPS_S_EQ('r0.x', 0),
+        NOP(rpt=5),
+        BR(2),
+        MOV_S32('r9.x', 1),
+    ]
+    store_u32_linear(instrs, 'r14.x', 'r0.x', join=True)
+    store_u32_linear(instrs, 'r14.x', 'r9.x', threads * 4)
+    instrs.append(END())
+    return assemble(instrs)
+
+
+def build_fiber_shader(dev, threads):
+    instrs = [
+        MOV_F32('r14.x', 'r0.x'),
+        GETFIBERID('r9.x'),
+        NOP(rpt=5),
+    ]
+    store_u32_linear(instrs, 'r14.x', 'r9.x')
+    instrs.append(END())
+    return assemble(instrs)
+
+
+def make_runtime(dev, shader, threads, disasm_shader=False):
+    envelope, io, isz, ro = get_envelope(dev, make_donor_src(2, threads))
+    if disasm_shader:
+        print(disasm(shader))
+    print('shader_instrs=%d bytes=%d envelope_bytes=%d' % (len(shader)//8, len(shader), isz))
+    lib = inject(envelope, io, isz, ro, shader, fregs=16, hregs=16, mergedregs=False)
+    return 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),)])
+
+
+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)
+    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 read_u32_linear(c, idx):
+    off = idx * 4
+    return struct.unpack_from('<I', c.as_memoryview(), off)[0]
+
+
+def run_semantics(args):
+    dev = Device['QCOM']
+    if args.quad_map: shader, mode_names = build_quad_map_shader(dev, args.threads), ['src   ', 'qbc0  ', 'qbc1  ', 'qbc2  ', 'qbc3  ']
+    elif args.shfl_modes: shader, mode_names = build_modes_shader(dev, args.threads)
+    else: shader = build_quad_semantics_shader(dev, args.threads) if args.quad_only else build_semantics_shader(dev, args.threads)
+    prg = make_runtime(dev, shader, args.threads, args.disasm)
+    a, b, c = make_bufs(dev)
+    prg(a._buf, b._buf, c._buf, global_size=(1, 1, 1), local_size=(args.threads, 1, 1), wait=True)
+    copied = bytearray(c.nbytes)
+    c.copyout(memoryview(copied))
+    names = mode_names if args.shfl_modes or args.quad_map else ['src   ', 'qbc0  ', 'sqbc0 ', 'qbc1  ', 'qbc2  ', 'qbc3  '] if args.quad_only else ['src   ', 'rdn0  ', 'rdn1  ', 'rdn2  ', 'rdn4  ', 'rdn8  ', 'rdn16 ', 'xor16 ', 'qbc0  ', 'qbc1  ', 'qbc2  ', 'qbc3  ']
+    vals = [[struct.unpack_from('<I', copied, (case * args.threads + lane)*4)[0] for lane in range(32)] for case in range(len(names))]
+    for name, row_vals in zip(names, vals):
+        print('%s: %s' % (name, ' '.join('%02d' % (v & 0xff) for v in row_vals)))
+    if args.quad_map:
+        mv = copied
+        for case, name in enumerate(names):
+            nz = []
+            base = case * args.threads
+            for i in range(args.threads * 2):
+                v = struct.unpack_from('<I', mv, (base + i) * 4)[0]
+                if v: nz.append((i, v & 0xff))
+            print('nonzero %s: %s' % (name.strip(), nz[:64]))
+
+
+def run_bench(args):
+    dev = Device['QCOM']
+    shader, loop_instrs = build_bench_shader(dev, args.threads, args.op, args.bench, args.ops_per_iter)
+    prg = make_runtime(dev, shader, args.threads, args.disasm)
+    print('loop_instrs=%d op=%s ops_per_iter=%d' % (loop_instrs, args.op, args.ops_per_iter))
+    a, b, c = make_bufs(dev)
+    gs = (args.groups, 64, 1)
+    ls = (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)
+    total_threads = args.groups * 64 * args.threads
+    total_ops = total_threads * K4 * args.ops_per_iter
+    print('%.1f G%s/s  (%.3f ms)' % (total_ops / best / 1e9, args.op.upper(), best * 1e3))
+
+
+def run_branch(args):
+    dev = Device['QCOM']
+    shader = build_branch_shader(dev, args.threads)
+    prg = make_runtime(dev, shader, args.threads, args.disasm)
+    a, b, c = make_bufs(dev)
+    prg(a._buf, b._buf, c._buf, global_size=(1, 1, 1), local_size=(args.threads, 1, 1), wait=True)
+    mods = [read_u32_linear(c, lane) for lane in range(args.threads)]
+    vals = [read_u32_linear(c, args.threads + lane) for lane in range(args.threads)]
+    for base in range(0, args.threads, 32):
+        print('mod %03d: %s' % (base, ' '.join('%d' % (v & 0xff) for v in mods[base:base+32])))
+        print('br  %03d: %s' % (base, ' '.join('%d' % (v & 0xff) for v in vals[base:base+32])))
+    mv = c.as_memoryview()
+    nz0, nz1 = [], []
+    for i in range(4096):
+        v0 = struct.unpack_from('<I', mv, i * 4)[0]
+        v1 = struct.unpack_from('<I', mv, args.threads * 4 + i * 4)[0]
+        if v0: nz0.append((i, v0 & 0xff))
+        if v1: nz1.append((i, v1 & 0xff))
+    print('nonzero mod:', nz0[:64])
+    print('nonzero br :', nz1[:64])
+
+
+def run_fiber(args):
+    dev = Device['QCOM']
+    shader = build_fiber_shader(dev, args.threads)
+    prg = make_runtime(dev, shader, args.threads, args.disasm)
+    a, b, c = make_bufs(dev)
+    prg(a._buf, b._buf, c._buf, global_size=(1, 1, 1), local_size=(args.threads, 1, 1), wait=True)
+    vals = [read_u32_linear(c, lane) for lane in range(args.threads)]
+    for base in range(0, args.threads, 32): print('fiber %03d: %s' % (base, ' '.join('%02d' % (v & 0xff) for v in vals[base:base+32])))
+
+
+def run_compiled_branch(args):
+    dev = Device['QCOM']
+    src = '__attribute__((reqd_work_group_size(%d,1,1)))\n' % args.threads
+    src += '__kernel void gemm_h(__global uint *C){int lid=get_local_id(0);'
+    src += 'if((lid&3)==0) C[lid]=1; else C[lid]=0;}\n'
+    lib = bytearray(dev.compiler.compile_cached(src))
+    io = struct.unpack_from('<I', lib, 0xc0)[0]
+    isz = struct.unpack_from('<I', lib, 0x100)[0]
+    print('compiled_branch_bytes=%d' % isz)
+    print(disasm(bytes(lib[io:io+isz])))
+
+
+def run_compiled_image_branch(args):
+    dev = Device['QCOM']
+    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(%d,1,1)))\n' % args.threads
+    src += '__kernel void gemm_h(read_only image2d_t A,__global half *C){int lid=get_local_id(0);half v=(half)0;'
+    src += 'if((lid&3)==0) v=read_imageh(A,smp,(int2)(0,lid)).x; C[lid]=v;}\n'
+    lib = bytearray(dev.compiler.compile_cached(src))
+    io = struct.unpack_from('<I', lib, 0xc0)[0]
+    isz = struct.unpack_from('<I', lib, 0x100)[0]
+    print('compiled_image_branch_bytes=%d' % isz)
+    print(disasm(bytes(lib[io:io+isz])))
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--threads', type=int, choices=(64, 128), default=128)
+    parser.add_argument('--disasm', action='store_true')
+    parser.add_argument('--op', choices=('shfl', 'quad'), default='shfl')
+    parser.add_argument('--bench', choices=('throughput', 'chain'))
+    parser.add_argument('--branch', action='store_true')
+    parser.add_argument('--fiberid', action='store_true')
+    parser.add_argument('--compiled-branch', action='store_true')
+    parser.add_argument('--compiled-image-branch', action='store_true')
+    parser.add_argument('--quad-only', action='store_true')
+    parser.add_argument('--quad-map', action='store_true')
+    parser.add_argument('--shfl-modes', action='store_true')
+    parser.add_argument('--ops-per-iter', type=int, default=16)
+    parser.add_argument('--groups', type=int, default=8)
+    parser.add_argument('--iters', type=int, default=20)
+    args = parser.parse_args()
+    if args.compiled_image_branch: run_compiled_image_branch(args)
+    elif args.compiled_branch: run_compiled_branch(args)
+    elif args.fiberid: run_fiber(args)
+    elif args.branch: run_branch(args)
+    elif args.bench: run_bench(args)
+    else: run_semantics(args)

extra/gemm/qcom_split_check.py

@@ -0,0 +1,101 @@
+#!/usr/bin/env python3
+"""Check a hand-assembled split-K GEMM against a NumPy partial matmul."""
+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 disasm, get_envelope, inject
+
+
+def main() -> None:
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--m", type=int, default=128)
+  parser.add_argument("--n", type=int, default=384)
+  parser.add_argument("--ncols", type=int, default=0, help="128-column blocks computed by each workgroup")
+  parser.add_argument("--k", type=int, default=1536)
+  parser.add_argument("--k-start", type=int, default=0, help="first K/4 iteration")
+  parser.add_argument("--k-count", type=int, default=48, help="number of K/4 iterations")
+  parser.add_argument("--seed", type=int, default=0)
+  parser.add_argument("--scale", type=float, default=0.25)
+  parser.add_argument("--impulse-k", type=int, default=-1,
+                      help="zero both inputs except one K plane (useful for coordinate diagnostics)")
+  parser.add_argument("--compiler", action="store_true", help="check the compiler donor over the full K range")
+  parser.add_argument("--donor-store", action="store_true")
+  parser.add_argument("--store-constant", action="store_true")
+  parser.add_argument("--serial-b-cols", action="store_true")
+  parser.add_argument("--disasm", action="store_true")
+  args = parser.parse_args()
+  assert args.m % 16 == 0 and args.n % 128 == 0 and args.n <= 512
+
+  rng = np.random.default_rng(args.seed)
+  a_np = (rng.standard_normal((args.m, args.k))*args.scale).astype(np.float16)
+  b_np = (rng.standard_normal((args.k, args.n))*args.scale).astype(np.float16)
+  if args.impulse_k >= 0:
+    assert 0 <= args.impulse_k < args.k
+    a_keep, b_keep = a_np[:, args.impulse_k].copy(), b_np[args.impulse_k].copy()
+    a_np.fill(0)
+    b_np.fill(0)
+    a_np[:, args.impulse_k], b_np[args.impulse_k] = a_keep, b_keep
+  dev = Device["QCOM"]
+  q.M, q.N, q.K, q.K4 = args.m, 1024, args.k, args.k//4
+  ncols = args.ncols or args.n//128
+  assert args.n % (128*ncols) == 0
+  if args.compiler:
+    args.k_start, args.k_count = 0, args.k//4
+    env, io, sz, _ = get_envelope(dev, q.make_direct_donor_src(ncols, 128))
+    shader = bytes(env[io:io+sz])
+    lib = bytes(env)
+  else:
+    # Keep enough instruction capacity for the explicitly initialized hand
+    # shader even when it computes only one column block.
+    env, io, sz, ro = get_envelope(dev, q.make_donor_src(max(ncols, 3), 128))
+    unroll = 4 if args.k_count % 4 == 0 else 2 if args.k_count % 2 == 0 else 1
+    shader, _ = q.build_4xn_shader(dev, 128, ncols=ncols, direct=True, compact_acc=True,
+      alu_order="row_col_kk", k_unroll=unroll, first_sync_only=False, coord_delay=4, serial_b_cols=args.serial_b_cols,
+      k_start=args.k_start, k_count=args.k_count, donor_store=args.donor_store,
+      store_constant=args.store_constant)
+    lib = inject(env, io, sz, ro, shader, fregs=10, hregs=20+4*ncols)
+  if args.disasm:
+    asm = disasm(shader)
+    print(asm)
+
+  a = Buffer("QCOM", a_np.size, dtypes.half).allocate()
+  b = Buffer("QCOM", b_np.size, dtypes.half).allocate()
+  c = Buffer("QCOM", args.m*1024, dtypes.half).allocate()
+  a.copyin(memoryview(a_np).cast("B"))
+  b.copyin(memoryview(b_np).cast("B"))
+  c.copyin(memoryview(np.zeros(args.m*1024, dtype=np.float16)).cast("B"))
+  buf_dtypes = [((0, dtypes.half, (args.m, args.k//4, 4)),),
+                ((0, dtypes.half, (args.k, args.n//4, 4)),), ((0, dtypes.half, None),)]
+  prg = dev.runtime("gemm_h", lib, buf_dtypes=buf_dtypes)
+  elapsed = prg(a._buf, b._buf, c._buf, global_size=(args.n//(128*ncols), args.m//16, 1), local_size=(128, 1, 1), wait=True)
+  got_flat = np.empty(args.m*1024, dtype=np.float16)
+  c.copyout(memoryview(got_flat).cast("B"))
+  got = got_flat.reshape(args.m, 1024)[:, :args.n].astype(np.float32)
+  lo, hi = args.k_start*4, (args.k_start+args.k_count)*4
+  expected = a_np[:, lo:hi].astype(np.float32) @ b_np[lo:hi].astype(np.float32)
+  delta = np.abs(expected-got)
+  worst = np.unravel_index(np.argmax(delta), delta.shape)
+  def best_axis(left, right):
+    left = left-left.mean(axis=1, keepdims=True)
+    right = right-right.mean(axis=1, keepdims=True)
+    corr = left@right.T/(np.linalg.norm(left, axis=1, keepdims=True)*np.linalg.norm(right, axis=1)[None, :]+1e-20)
+    match = np.argmax(np.abs(corr), axis=1)
+    return match, corr[np.arange(len(match)), match]
+  row_match, row_corr = best_axis(got, expected)
+  col_match, col_corr = best_axis(got.T, expected.T)
+  print(f"elapsed_ms={elapsed*1e3:.3f} max_abs={delta.max():.9g} mean_abs={delta.mean():.9g} "
+        f"worst={worst} expected={expected[worst]!r} got={got[worst]!r}")
+  print(f"rows match={row_match[:16].tolist()} corr={np.round(row_corr[:16], 3).tolist()} median={np.median(np.abs(row_corr)):.3f}")
+  print(f"cols match={col_match[:16].tolist()} corr={np.round(col_corr[:16], 3).tolist()} median={np.median(np.abs(col_corr)):.3f}")
+  print(f"expected[0,:8]={expected[0, :8].tolist()} got[0,:8]={got[0, :8].tolist()} "
+        f"within4_deltas={[float(np.max(np.abs(got[:, 0]-got[:, i]))) for i in range(1, 4)]}")
+  print("max_abs row4-groups x col128-blocks=", [[float(delta[r:r+4, c:c+128].max())
+        for c in range(0, args.n, 128)] for r in range(0, args.m, 4)])
+  if not np.allclose(expected, got, rtol=2e-2, atol=2e-2): raise SystemExit(1)
+
+
+if __name__ == "__main__": main()

extra/gemm/qcom_strassen2_check.py

@@ -0,0 +1,124 @@
+#!/usr/bin/env python3
+"""Fully checked two-level Strassen FP16 GEMM for Adreno 630."""
+import os
+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 get_envelope, inject
+
+
+def add(*xs: np.ndarray) -> np.ndarray: return sum(xs, np.zeros_like(xs[0]))
+def sub(x: np.ndarray, y: np.ndarray) -> np.ndarray: return x-y
+
+
+def operands(blocks: list[np.ndarray], side: str) -> list[np.ndarray]:
+  x11, x12, x21, x22 = blocks
+  if side == "A": return [add(x11,x22), add(x21,x22), x11, x22, add(x11,x12), sub(x21,x11), sub(x12,x22)]
+  return [add(x11,x22), x11, sub(x12,x22), sub(x21,x11), x22, add(x11,x12), add(x21,x22)]
+
+
+def combine(ms: list[np.ndarray]) -> list[np.ndarray]:
+  m1,m2,m3,m4,m5,m6,m7 = ms
+  return [m1+m4-m5+m7, m3+m5, m2+m4, m1-m2+m3+m6]
+
+
+def operand_coeffs(side: str) -> list[np.ndarray]:
+  eye = np.eye(16, dtype=np.int32).reshape(4,4,16)
+  top = operands([eye[:2,:2], eye[:2,2:], eye[2:,:2], eye[2:,2:]], side)
+  return [x for t in top for x in operands([t[0,0],t[0,1],t[1,0],t[1,1]], side)]
+
+
+def output_coeffs() -> list[np.ndarray]:
+  eye = np.eye(49, dtype=np.int32).reshape(7,7,49)
+  inner = [np.array(combine(list(eye[p]))).reshape(2,2,49) for p in range(7)]
+  outer = combine(inner)
+  grid = np.empty((4,4,49),np.int32)
+  grid[:2,:2],grid[:2,2:],grid[2:,:2],grid[2:,2:] = outer
+  return [grid[r,c] for r in range(4) for c in range(4)]
+
+
+def expr(coeff: np.ndarray, names: list[str]) -> str:
+  terms: list[str] = []
+  for c, name in zip(coeff.tolist(), names):
+    terms += ([name]*c if c > 0 else [f"(-{name})"]*(-c))
+  return "+".join(terms) if terms else "(half4)(0)"
+
+
+def prep_src(side: str) -> str:
+  coeffs, names = operand_coeffs(side), [f"x{i}" for i in range(16)]
+  lines = ["#pragma OPENCL EXTENSION cl_khr_fp16 : enable",
+           "__attribute__((reqd_work_group_size(128,1,1)))",
+           "__kernel void prep(__global const half *X,__global half *O){",
+           "int i=get_global_id(0),r=i>>6,c=(i&63)<<2,o=r*256+c;"]
+  for br in range(4):
+    for bc in range(4):
+      j=br*4+bc
+      lines.append(f"half4 x{j}=vload4(0,X+({br}*256+r)*1024+{bc}*256+c);")
+  for p, c in enumerate(coeffs): lines.append(f"vstore4({expr(c,names)},0,O+{p}*65536+o);")
+  return "\n".join(lines+["}"])
+
+
+def post_src() -> str:
+  coeffs = output_coeffs()
+  lines = ["#pragma OPENCL EXTENSION cl_khr_fp16 : enable",
+           "__attribute__((reqd_work_group_size(128,1,1)))",
+           "__kernel void post(__global const half *M,__global half *C){",
+           "int i=get_global_id(0),block=i>>14,j=i&16383,r=j>>6,c=(j&63)<<2; half4 v=(half4)(0);"]
+  for block, coeff in enumerate(coeffs):
+    nz = np.flatnonzero(coeff)
+    loads = [f"vload4(0,M+{p}*65536+r*256+c)" for p in nz]
+    e = expr(coeff[nz], loads)
+    lines.append(f"{'if' if block == 0 else 'else if'}(block=={block}) v={e};")
+  lines += ["int br=block>>2,bc=block&3;vstore4(v,0,C+(br*256+r)*1024+bc*256+c);", "}"]
+  return "\n".join(lines)
+
+
+def alloc_half(count: int) -> Buffer: return Buffer("QCOM", count, dtypes.half).allocate()
+
+
+def main() -> None:
+  seed, runs = int(os.getenv("SEED", "367")), int(os.getenv("BENCH_RUNS", "5"))
+  rng = np.random.default_rng(seed)
+  a_np=(rng.standard_normal((1024,1024))*0.05).astype(np.float16)
+  b_np=(rng.standard_normal((1024,1024))*0.05).astype(np.float16)
+  dev=Device["QCOM"]
+  a,b,pa,pb,pm,c = alloc_half(1024**2),alloc_half(1024**2),alloc_half(49*256**2),alloc_half(49*256**2),alloc_half(49*256**2),alloc_half(1024**2)
+  a.copyin(memoryview(a_np).cast("B")); b.copyin(memoryview(b_np).cast("B"))
+  spec=((0,dtypes.half,None),)
+  prepa=dev.runtime("prep",dev.compiler.compile(prep_src("A")),buf_dtypes=[spec,spec])
+  prepb=dev.runtime("prep",dev.compiler.compile(prep_src("B")),buf_dtypes=[spec,spec])
+  post=dev.runtime("post",dev.compiler.compile(post_src()),buf_dtypes=[spec,spec])
+
+  q8.M=q8.N=q8.K=256; q8.K4=64
+  env,io,sz,ro=get_envelope(dev,q4.make_direct_image_donor_src(4,128))
+  shader,hregs,fregs,_=q8.build_8x8_persistent_shader(dev,128,batch_m=256,batch_n=256,batch_k=256,
+    batch_fixed_b=-2,dynamic_a4_dual=True,image_store=True)
+  lib=inject(env,io,sz,ro,shader,fregs=fregs,hregs=hregs,mergedregs=False)
+  images=[((0,dtypes.half,(49*256,64,4)),),((0,dtypes.half,(49*256,64,4)),),((1,dtypes.half,(49*256,64,4)),)]
+  gemm=dev.runtime("gemm_h",lib,buf_dtypes=images)
+
+  def iteration() -> tuple[float,float,float,float]:
+    ta=prepa(a._buf,pa._buf,global_size=(128,1,1),local_size=(128,1,1),wait=True)
+    tb=prepb(b._buf,pb._buf,global_size=(128,1,1),local_size=(128,1,1),wait=True)
+    tg=gemm(pm._buf,pa._buf,pb._buf,global_size=(1,8,49),local_size=(128,1,1),wait=True)
+    to=post(pm._buf,c._buf,global_size=(2048,1,1),local_size=(128,1,1),wait=True)
+    return ta,tb,tg,to
+
+  for _ in range(2): iteration()
+  measured=[iteration() for _ in range(runs)]
+  best=min(measured,key=sum); elapsed=sum(best)
+  got=np.empty((1024,1024),np.float16); c.copyout(memoryview(got).cast("B"))
+  expected=a_np.astype(np.float32)@b_np.astype(np.float32)
+  delta=np.abs(got.astype(np.float32)-expected); correct=np.allclose(got,expected,rtol=2e-2,atol=2e-2)
+  bad=~np.isfinite(got)|(delta>.02)
+  print(f"shape=1024x1024x1024 algorithm=strassen2 accumulate=fp16 elapsed_ms={elapsed*1e3:.3f} gflops={2*1024**3/elapsed/1e9:.1f} "
+        f"prepA_ms={best[0]*1e3:.3f} prepB_ms={best[1]*1e3:.3f} gemm_ms={best[2]*1e3:.3f} post_ms={best[3]*1e3:.3f} "
+        f"max_abs={delta.max():.9g} mean_abs={delta.mean():.9g} allclose={correct}")
+  print(f"bad_count={int(bad.sum())}")
+  if not correct: raise SystemExit(1)
+
+
+if __name__ == "__main__": main()

extra/gemm/qcom_strassen2_fused_fp32_check.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python3
+"""Two-level Strassen full GEMM with arbitrary FP16 inputs and FP32 dot accumulation."""
+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
+from extra.gemm.qcom_strassen4_fused2_fp32_check import alloc, combine2_src, transform2_src
+
+
+def main() -> None:
+  n, seed = int(os.getenv("N", "1024")), int(os.getenv("SEED", "981"))
+  if n not in (1024, 2048): raise ValueError("N must be 1024 or 2048")
+  leaf = n//4
+  wg, combine_wg = int(os.getenv("TRANSFORM_WG", "256")), int(os.getenv("COMBINE_WG", "128"))
+  transform_vec, combine_vec = int(os.getenv("TRANSFORM_VEC", "2")), int(os.getenv("COMBINE_VEC", "4"))
+  leaf_unroll = int(os.getenv("LEAF_UNROLL", "3"))
+  rng = np.random.default_rng(seed)
+  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)
+  a, b = alloc(n*n, dtypes.half), alloc(n*n, dtypes.half)
+  a.copyin(memoryview(a_np).cast("B")); b.copyin(memoryview(b_np).cast("B"))
+  pa, pb, pm, c = (alloc(49*leaf*leaf, dtypes.half), alloc(49*leaf*leaf, dtypes.half),
+                   alloc(49*leaf*leaf, dtypes.half), alloc(n*n, dtypes.float))
+  dev = Device["QCOM"]
+  transforms = {side: dev.runtime("transform", dev.compiler.compile(
+    transform2_src(n, side, True, True, wg, transform_vec, False, False, True)), buf_dtypes=[
+      ((0, dtypes.half, (49*leaf*leaf,)),), ((0, dtypes.half, (n*n,)),)]) for side in "AB"}
+  combine = dev.runtime("combine", dev.compiler.compile(
+    combine2_src(n, combine_wg, True, False, combine_vec, False)), buf_dtypes=[
+      ((0, dtypes.float, (n*n,)),), ((0, dtypes.half, (49*leaf*leaf,)),)])
+
+  q.M = q.N = q.K = leaf; q.K4 = leaf//4
+  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=leaf//4, batch_stride=leaf, batch_from_row=True, k_unroll=leaf_unroll)
+  lib = inject(env, io, sz, ro, shader, fregs=fregs, hregs=hregs, mergedregs=False)
+  gemms = {}
+  def gemm(batch: int):
+    if batch not in gemms:
+      gemms[batch] = dev.runtime("gemm_h", lib, buf_dtypes=[
+        ((0, dtypes.half, (batch*leaf, leaf//4, 4)),), ((0, dtypes.half, (batch*leaf, leaf//4, 4)),),
+        ((1, dtypes.half, (batch*leaf, leaf//4, 4)),)])
+    return gemms[batch]
+  def grid(groups: int) -> tuple[int, int, int]:
+    gx = min(groups, 1024)
+    while groups % gx: gx -= 1
+    return gx, groups//gx, 1
+
+  transform_groups = leaf*leaf//transform_vec//wg
+  combine_groups = leaf*leaf//combine_vec//combine_wg
+  times = []
+  times.append(transforms["A"](pa._buf, a._buf, global_size=grid(transform_groups), local_size=(wg, 1, 1), wait=True))
+  times.append(transforms["B"](pb._buf, b._buf, global_size=grid(transform_groups), local_size=(wg, 1, 1), wait=True))
+  input_bytes, output_bytes = leaf*leaf*2, leaf*leaf*2
+  max_batch = 8192//leaf
+  for first in range(0, 49, max_batch):
+    batch = min(max_batch, 49-first)
+    times.append(gemm(batch)(pm._buf.offset(first*output_bytes, batch*output_bytes),
+      pa._buf.offset(first*input_bytes, batch*input_bytes), pb._buf.offset(first*input_bytes, batch*input_bytes),
+      global_size=(max(1, leaf//256), batch*leaf//8, 1), local_size=(64, 1, 1), wait=True))
+  times.append(combine(c._buf, pm._buf, global_size=grid(combine_groups), local_size=(combine_wg, 1, 1), wait=True))
+  elapsed = sum(x for x in times if x is not None)
+  got = np.empty((n, 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)
+  print(f"shape={n}x{n}x{n} algorithm=strassen2 inputs=fp16 accumulate=fp32 elapsed_ms={elapsed*1e3:.3f} "
+        f"gflops={2*n**3/elapsed/1e9:.1f} transform_ms={(times[0]+times[1])*1e3:.3f} "
+        f"gemm_ms={sum(times[2:-1])*1e3:.3f} combine_ms={times[-1]*1e3:.3f} outputs={n*n} "
+        f"bad_count={int(bad.sum())} max_abs={float(delta.max()):.9g} mean_abs={float(delta.mean()):.9g} "
+        f"allclose={not bool(bad.any())} transform_wg={wg} combine_wg={combine_wg} transform_vec={transform_vec} "
+        f"combine_vec={combine_vec} leaf_unroll={leaf_unroll} loop_instrs={loop_instrs}")
+  if bad.any(): raise SystemExit(1)
+
+
+if __name__ == "__main__": main()

extra/gemm/qcom_strassen3_fused_fp32_check.py

@@ -0,0 +1,102 @@
+#!/usr/bin/env python3
+"""Three-level Strassen full GEMM with arbitrary FP16 inputs and FP32 dot accumulation."""
+import os
+
+import numpy as np
+
+from tinygrad import Device, dtypes
+from extra.gemm import qcom_intensity_gemm as q
+from extra.gemm.ir3asm import get_envelope, inject
+from extra.gemm.qcom_strassen4_fused2_fp32_check import (
+  alloc, combine1_batch_src, combine2_src, transform1_batch_src, transform2_src,
+)
+
+
+def main() -> None:
+  n, seed = int(os.getenv("N", "1024")), int(os.getenv("SEED", "991"))
+  if n != 1024: raise ValueError("this performance gate is intentionally fixed at N=1024")
+  parent, leaf = n//4, n//8
+  wg, combine_wg = int(os.getenv("TRANSFORM_WG", "256")), int(os.getenv("COMBINE_WG", "128"))
+  transform_vec, combine_vec = int(os.getenv("TRANSFORM_VEC", "2")), int(os.getenv("COMBINE_VEC", "4"))
+  leaf_unroll = int(os.getenv("LEAF_UNROLL", "3"))
+  leaf_tile = os.getenv("LEAF_TILE", "4x8")
+  rng = np.random.default_rng(seed)
+  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)
+  a, b = alloc(n*n, dtypes.half), alloc(n*n, dtypes.half)
+  a.copyin(memoryview(a_np).cast("B")); b.copyin(memoryview(b_np).cast("B"))
+  pcount = 49*parent*parent
+  lcount = 343*leaf*leaf
+  pa, pb, la, lb, lm, pm = (alloc(pcount, dtypes.half), alloc(pcount, dtypes.half),
+                             alloc(lcount, dtypes.half), alloc(lcount, dtypes.half),
+                             alloc(lcount, dtypes.half), alloc(pcount, dtypes.half))
+  c = alloc(n*n, dtypes.float)
+  dev = Device["QCOM"]
+  transforms2 = {side: dev.runtime("transform2", dev.compiler.compile(
+    transform2_src(n, side, True, True, wg, transform_vec, False, False, True)), buf_dtypes=[
+      ((0, dtypes.half, (pcount,)),), ((0, dtypes.half, (n*n,)),)]) for side in "AB"}
+  transforms1 = {side: dev.runtime("transform1", dev.compiler.compile(
+    transform1_batch_src(parent, 49, side, wg, transform_vec)), buf_dtypes=[
+      ((0, dtypes.half, (lcount,)),), ((0, dtypes.half, (pcount,)),)]) for side in "AB"}
+  combine1 = dev.runtime("combine1", dev.compiler.compile(
+    combine1_batch_src(parent, 49, combine_wg, combine_vec)), buf_dtypes=[
+      ((0, dtypes.half, (pcount,)),), ((0, dtypes.half, (lcount,)),)])
+  combine2 = dev.runtime("combine2", dev.compiler.compile(
+    combine2_src(n, combine_wg, True, False, combine_vec, False)), buf_dtypes=[
+      ((0, dtypes.float, (n*n,)),), ((0, dtypes.half, (pcount,)),)])
+
+  q.M = q.N = q.K = leaf; q.K4 = leaf//4
+  env, io, sz, ro = get_envelope(dev, q.make_direct_image_donor_src(2, 64))
+  if leaf_tile == "4x8":
+    shader, hregs, fregs, loop_instrs = q.build_4x8_fp32_rotate_shader(
+      dev, 64, k_count=leaf//4, batch_stride=leaf, batch_from_row=True, k_unroll=leaf_unroll)
+    rows_per_group = 8
+  elif leaf_tile == "8x8":
+    shader, hregs, fregs, loop_instrs = q.build_8x8_fp32_shader(dev, 64, batch_stride=leaf)
+    rows_per_group = 16
+  else:
+    raise ValueError("LEAF_TILE must be 4x8 or 8x8")
+  lib = inject(env, io, sz, ro, shader, fregs=fregs, hregs=hregs, mergedregs=False)
+  gemms = {}
+  def gemm(batch: int):
+    if batch not in gemms:
+      gemms[batch] = dev.runtime("gemm_h", lib, buf_dtypes=[
+        ((0, dtypes.half, (batch*leaf, leaf//4, 4)),), ((0, dtypes.half, (batch*leaf, leaf//4, 4)),),
+        ((1, dtypes.half, (batch*leaf, leaf//4, 4)),)])
+    return gemms[batch]
+  def grid(groups: int) -> tuple[int, int, int]:
+    gx = min(groups, 1024)
+    while groups % gx: gx -= 1
+    return gx, groups//gx, 1
+
+  times = {"transform": 0.0, "gemm": 0.0, "combine": 0.0}
+  t2groups = parent*parent//transform_vec//wg
+  t1groups = 49*leaf*leaf//transform_vec//wg
+  c1groups = 49*leaf*leaf//combine_vec//combine_wg
+  c2groups = parent*parent//combine_vec//combine_wg
+  for prg, out, inp in ((transforms2["A"], pa, a), (transforms2["B"], pb, b)):
+    times["transform"] += prg(out._buf, inp._buf, global_size=grid(t2groups), local_size=(wg, 1, 1), wait=True)
+  for prg, out, inp in ((transforms1["A"], la, pa), (transforms1["B"], lb, pb)):
+    times["transform"] += prg(out._buf, inp._buf, global_size=grid(t1groups), local_size=(wg, 1, 1), wait=True)
+  matrix_bytes = leaf*leaf*2
+  max_batch = 8192//leaf
+  for first in range(0, 343, max_batch):
+    batch = min(max_batch, 343-first)
+    span = batch*matrix_bytes
+    times["gemm"] += gemm(batch)(lm._buf.offset(first*matrix_bytes, span), la._buf.offset(first*matrix_bytes, span),
+      lb._buf.offset(first*matrix_bytes, span), global_size=(1, batch*leaf//rows_per_group, 1), local_size=(64, 1, 1), wait=True)
+  times["combine"] += combine1(pm._buf, lm._buf, global_size=grid(c1groups), local_size=(combine_wg, 1, 1), wait=True)
+  times["combine"] += combine2(c._buf, pm._buf, global_size=grid(c2groups), local_size=(combine_wg, 1, 1), wait=True)
+  elapsed = sum(times.values())
+  got = np.empty((n, 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)
+  print(f"shape={n}x{n}x{n} algorithm=strassen3 inputs=fp16 accumulate=fp32 elapsed_ms={elapsed*1e3:.3f} "
+        f"gflops={2*n**3/elapsed/1e9:.1f} transform_ms={times['transform']*1e3:.3f} "
+        f"gemm_ms={times['gemm']*1e3:.3f} combine_ms={times['combine']*1e3:.3f} outputs={n*n} "
+        f"bad_count={int(bad.sum())} max_abs={float(delta.max()):.9g} mean_abs={float(delta.mean()):.9g} "
+        f"allclose={not bool(bad.any())} leaf_tile={leaf_tile} leaf_unroll={leaf_unroll} loop_instrs={loop_instrs}")
+  if bad.any(): raise SystemExit(1)
+
+
+if __name__ == "__main__": main()

extra/gemm/qcom_strassen4_fused2_fp32_check.py

@@ -0,0 +1,577 @@
+#!/usr/bin/env python3
+"""Four-level full Strassen GEMM using register-fused pairs of transforms/combines."""
+import ctypes, gc, hashlib, mmap, os, subprocess, time
+from dataclasses import dataclass
+from pathlib import Path
+
+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
+from extra.gemm.qcom_strassen_fused_transform_check import TA, TB
+from extra.gemm.qcom_strassen_fused_combine_check import COMBINE, combine_expr
+
+
+def alloc(count: int, dtype) -> Buffer: return Buffer("QCOM", count, dtype).allocate()
+
+
+def alloc_external(count: int, dtype) -> tuple[Buffer, mmap.mmap]:
+  nbytes = count*dtype.itemsize
+  backing = mmap.mmap(-1, nbytes)
+  ptr = ctypes.addressof(ctypes.c_char.from_buffer(backing))
+  return Buffer("QCOM", count, dtype).allocate(external_ptr=ptr), backing
+
+
+def transform2_src(n: int, side: str, input_half: bool, output_half: bool = False, wg: int = 256, vec: int = 4,
+                   input_image: bool = False, output_image: bool = False, compute_half: bool = False) -> str:
+  if output_image and vec != 4: raise ValueError("image transform output requires vec=4")
+  if compute_half and not (input_half and output_half): raise ValueError("half transform compute requires half input and output")
+  leaf, t = n//4, TA if side == "A" else TB
+  groups = leaf*leaf//vec//wg
+  grid_x = min(groups, 1024)
+  while groups % grid_x: grid_x -= 1
+  sid = "get_global_id(0)" if groups <= 1024 else f"get_global_id(0)+get_group_id(1)*{grid_x*wg}"
+  loads, stores = [], []
+  for q0 in range(4):
+    for q1 in range(4):
+      rb = ((q0>>1)<<1)|(q1>>1)
+      cb = ((q0&1)<<1)|(q1&1)
+      offset = f"({rb}*{leaf}+r)*{n}+{cb}*{leaf}+xv*{vec}"
+      if input_image:
+        value = f"read_image{'h' if input_half else 'f'}(I,smp,(int2)({cb*leaf//4}+xv,{rb*leaf}+r))"
+        if input_half and not compute_half: value = f"convert_float4({value})"
+      else:
+        value = ((f"vload{vec}(0,I+{offset})" if compute_half else f"convert_float{vec}(vload{vec}(0,I+{offset}))")
+                 if input_half else f"vload{vec}(0,I+{offset})") if vec > 1 else \
+                ((f"I[{offset}]" if compute_half else f"convert_float(I[{offset}])") if input_half else f"I[{offset}]")
+      ctype = "half" if compute_half else "float"
+      loads.append(f"{ctype if vec == 1 else f'{ctype}{vec}'} v{q0*4+q1}={value};")
+  # Factor the Kronecker transform one inner path at a time.  The direct
+  # 49-expression form repeats each inner sum for every outer path (95 vector
+  # adds); this schedule computes four inner temporaries, consumes them into
+  # seven outputs, then reuses the registers (55 vector adds total).
+  scalar_type = "half" if compute_half else "float"
+  vtype = scalar_type if vec == 1 else f"{scalar_type}{vec}"
+  for p1 in range(7):
+    stores.append("{")
+    for q0 in range(4):
+      terms = [("+" if t[p1][q1] > 0 else "-")+f"v{q0*4+q1}" for q1 in range(4) if t[p1][q1]]
+      stores.append(f"{vtype} w{q0}={''.join(terms).lstrip('+')};")
+    for p0 in range(7):
+      terms = [("+" if t[p0][q0] > 0 else "-")+f"w{q0}" for q0 in range(4) if t[p0][q0]]
+      p, expr = p0*7+p1, "".join(terms).lstrip("+")
+      oval = ((expr if compute_half else f"convert_half{vec}({expr})") if vec > 1 else
+              (expr if compute_half else f"convert_half({expr})")) if output_half else expr
+      offset = f"{p*leaf*leaf}+r*{leaf}+xv*{vec}"
+      if output_image: stores.append(f"write_image{'h' if output_half else 'f'}(O,(int2)(xv,{p*leaf}+r),{oval});")
+      else: stores.append(f"vstore{vec}({oval},0,O+{offset});" if vec > 1 else f"O[{offset}]={oval};")
+    stores.append("}")
+  itype, otype = ("half" if input_half else "float"), ("half" if output_half else "float")
+  return f"""#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+{'const sampler_t smp=CLK_NORMALIZED_COORDS_FALSE|CLK_ADDRESS_CLAMP|CLK_FILTER_NEAREST;' if input_image else ''}
+__attribute__((reqd_work_group_size({wg},1,1)))
+__kernel void transform({'write_only image2d_t O' if output_image else f'__global {otype} *O'},
+                        {'read_only image2d_t I' if input_image else f'__global const {itype} *I'}) {{
+  uint s={sid},r=s/{leaf//vec},xv=s%{leaf//vec};
+  {''.join(loads)}
+  {''.join(stores)}
+}}"""
+
+
+def transform2_path_src(n: int, side: str, path: int, wg: int = 256, vec: int = 4, compute_half: bool = False) -> str:
+  leaf = n//4
+  groups = leaf*leaf//vec//wg
+  grid_x = min(groups, 1024)
+  while groups % grid_x: grid_x -= 1
+  sid = "get_global_id(0)" if groups <= 1024 else f"get_global_id(0)+get_group_id(1)*{grid_x*wg}"
+  t = TA if side == "A" else TB
+  p0, p1 = divmod(path, 7)
+  loads, terms = [], []
+  for q0 in range(4):
+    for q1 in range(4):
+      coeff = t[p0][q0]*t[p1][q1]
+      if not coeff: continue
+      rb = ((q0>>1)<<1)|(q1>>1)
+      cb = ((q0&1)<<1)|(q1&1)
+      name = f"v{q0}_{q1}"
+      value = f"read_imageh(I,smp,(int2)({cb*leaf//4}+xv,{rb}*{leaf}+r))"
+      loads.append(f"{'half' if compute_half else 'float'}{vec} {name}="
+                   f"{value if compute_half else f'convert_float{vec}({value})'}; ")
+      terms.append(("+" if coeff > 0 else "-")+name)
+  expr = "".join(terms).lstrip("+")
+  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({wg},1,1)))
+__kernel void transform_path(__global half *O,read_only image2d_t I) {{
+  uint s={sid},r=s/{leaf//vec},xv=s%{leaf//vec};
+  {''.join(loads)}
+  vstore{vec}({expr if compute_half else f'convert_half{vec}({expr})'},0,O+r*{leaf}+xv*{vec});
+}}"""
+
+
+def transform1_batch_src(n: int, batch: int, side: str, wg: int = 256, vec: int = 2, input_half: bool = True) -> str:
+  """One additional Strassen level over a dense batch of FP16 matrices."""
+  leaf, t = n//2, TA if side == "A" else TB
+  groups = batch*leaf*leaf//vec//wg
+  grid_x = min(groups, 1024)
+  while groups % grid_x: grid_x -= 1
+  sid = "get_global_id(0)" if groups <= 1024 else f"get_global_id(0)+get_group_id(1)*{grid_x*wg}"
+  scalar, vtype = ("half" if input_half else "float"), ("half" if input_half else "float") if vec == 1 else \
+                  f"{'half' if input_half else 'float'}{vec}"
+  loads, stores = [], []
+  for q in range(4):
+    rb, cb = q//2, q%2
+    offset = f"b*{n*n}+({rb}*{leaf}+r)*{n}+{cb}*{leaf}+xv*{vec}"
+    loads.append(f"{vtype} v{q}={'I['+offset+']' if vec == 1 else f'vload{vec}(0,I+{offset})'};")
+  for p in range(7):
+    terms = [("+" if t[p][q] > 0 else "-")+f"v{q}" for q in range(4) if t[p][q]]
+    value = "".join(terms).lstrip("+")
+    offset = f"(b*7+{p})*{leaf*leaf}+r*{leaf}+xv*{vec}"
+    value = value if input_half else (f"convert_half({value})" if vec == 1 else f"convert_half{vec}({value})")
+    stores.append(f"O[{offset}]={value};" if vec == 1 else f"vstore{vec}({value},0,O+{offset});")
+  return f"""#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+__attribute__((reqd_work_group_size({wg},1,1)))
+__kernel void transform1(__global half *O,__global const {scalar} *I) {{
+  uint s={sid},b=s/{leaf*leaf//vec},z=s%{leaf*leaf//vec},r=z/{leaf//vec},xv=z%{leaf//vec};
+  {''.join(loads)}
+  {''.join(stores)}
+}}"""
+
+
+def combine1_batch_src(n: int, batch: int, wg: int = 256, vec: int = 4) -> str:
+  """Inverse of transform1_batch_src, retaining FP16 product storage semantics."""
+  leaf = n//2
+  groups = batch*leaf*leaf//vec//wg
+  grid_x = min(groups, 1024)
+  while groups % grid_x: grid_x -= 1
+  sid = "get_global_id(0)" if groups <= 1024 else f"get_global_id(0)+get_group_id(1)*{grid_x*wg}"
+  body = []
+  names = [f"convert_float{vec}(vload{vec}(0,M+(b*7+{p})*{leaf*leaf}+r*{leaf}+xv*{vec}))" for p in range(7)]
+  for u in range(4):
+    value = f"convert_half{vec}({combine_expr(names, u)})"
+    body.append(f"vstore{vec}({value},0,C+b*{n*n}+({u//2}*{leaf}+r)*{n}+{u%2}*{leaf}+xv*{vec});")
+  return f"""#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+__attribute__((reqd_work_group_size({wg},1,1)))
+__kernel void combine1(__global half *C,__global const half *M) {{
+  uint s={sid},b=s/{leaf*leaf//vec},z=s%{leaf*leaf//vec},r=z/{leaf//vec},xv=z%{leaf//vec};
+  {''.join(body)}
+}}"""
+
+
+def combine2_src(n: int, wg: int = 128, input_half: bool = False, output_half: bool = False, vec: int = 4,
+                 compute_half: bool = False) -> str:
+  if compute_half and not (input_half and output_half): raise ValueError("half combine compute requires half input and output")
+  leaf = n//4
+  groups = leaf*leaf//vec//wg
+  grid_x = min(groups, 1024)
+  while groups % grid_x: grid_x -= 1
+  sid = "get_global_id(0)" if groups <= 1024 else f"get_global_id(0)+get_group_id(1)*{grid_x*wg}"
+  body = []
+  for u in range(4):
+    for p0 in range(7):
+      names = [f"vload{vec}(0,M+{(p0*7+p1)*leaf*leaf}+r*{leaf}+xv*{vec})" for p1 in range(7)]
+      if input_half and not compute_half: names = [f"convert_float{vec}({name})" for name in names]
+      body.append(f"{'half' if compute_half else 'float'}{vec} d{u}_{p0}={combine_expr(names, u)};")
+    ds = [f"d{u}_{x}" for x in range(7)]
+    for v in range(4):
+      block_r, block_c = (v>>1)*2+(u>>1), (v&1)*2+(u&1)
+      value = combine_expr(ds, v)
+      if output_half and not compute_half: value = f"convert_half{vec}({value})"
+      body.append(f"vstore{vec}({value},0,C+({block_r}*{leaf}+r)*{n}+{block_c}*{leaf}+xv*{vec});")
+  return f"""#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+__attribute__((reqd_work_group_size({wg},1,1)))
+__kernel void combine(__global {'half' if output_half else 'float'} *C,__global const {'half' if input_half else 'float'} *M) {{
+  uint s={sid},r=s/{leaf//vec},xv=s%{leaf//vec};
+  {''.join(body)}
+}}"""
+
+
+def combine2_block_src(n: int, block: int, wg: int = 128, input_half: bool = False,
+                       output_half: bool = False, vec: int = 4) -> str:
+  """One final block with only its mathematically nonzero product reads."""
+  leaf, block_r, block_c = n//4, block//4, block%4
+  u = ((block_r&1)<<1)|(block_c&1)
+  v = ((block_r>>1)<<1)|(block_c>>1)
+  terms = []
+  for p0 in range(7):
+    for p1 in range(7):
+      coeff = COMBINE[v][p0]*COMBINE[u][p1]
+      if not coeff: continue
+      value = f"vload{vec}(0,M+{(p0*7+p1)*leaf*leaf}+r*{leaf}+xv*{vec})"
+      if input_half: value = f"convert_float{vec}({value})"
+      terms.append(("+" if coeff > 0 else "-")+value)
+  value = "".join(terms).lstrip("+")
+  if output_half: value = f"convert_half{vec}({value})"
+  return f"""#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+__attribute__((reqd_work_group_size({wg},1,1)))
+__kernel void combine(__global {'half' if output_half else 'float'} *C,
+                      __global const {'half' if input_half else 'float'} *M) {{
+  uint s=get_global_id(0),r=s/{leaf//vec},xv=s%{leaf//vec};
+  vstore{vec}({value},0,C+({block_r}*{leaf}+r)*{n}+{block_c}*{leaf}+xv*{vec});
+}}"""
+
+
+@dataclass
+class Times:
+  transform: float = 0
+  gemm: float = 0
+  combine: float = 0
+  wall: float | None = None
+  @property
+  def total(self) -> float: return self.wall if self.wall is not None else self.transform+self.gemm+self.combine
+
+
+def cpu_pipeline_lib(threads: int):
+  src = Path(__file__).with_name("qcom_strassen_cpu_pipeline.c")
+  digest = hashlib.sha1(src.read_bytes()).hexdigest()[:12]
+  so = Path(f"/tmp/qcom_strassen_cpu_pipeline_{digest}.so")
+  if not so.exists():
+    subprocess.run(["clang", "-O3", "-march=armv8.2-a+fp16", "-fopenmp", "-shared", "-fPIC", str(src), "-o", str(so)], check=True)
+  lib = ctypes.CDLL(str(so))
+  lib.set_threads.argtypes = [ctypes.c_int]
+  lib.transform2_f16.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int, ctypes.c_int]
+  lib.combine2_f16.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int]
+  lib.cache_clean.argtypes = [ctypes.c_void_p, ctypes.c_int64]
+  lib.cache_invalidate.argtypes = [ctypes.c_void_p, ctypes.c_int64]
+  lib.set_threads(threads)
+  return lib
+
+
+def main() -> None:
+  n, seed = int(os.getenv("N", "4096")), int(os.getenv("SEED", "701"))
+  if n//16 not in (64, 128, 256, 512, 768): raise ValueError("N/16 must be one of 64, 128, 256, 512, or 768")
+  child, mid, wg = n//4, n//16, int(os.getenv("TRANSFORM_WG", "256"))
+  strassen5 = bool(int(os.getenv("STRASSEN5", "0")))
+  leaf = mid//2 if strassen5 else mid
+  combine_wg = int(os.getenv("COMBINE_WG", "256"))
+  leaf_threads = int(os.getenv("LEAF_THREADS", "64"))
+  if leaf_threads not in (64, 128): raise ValueError("LEAF_THREADS must be 64 or 128")
+  leaf_halfwave = bool(int(os.getenv("LEAF_HALFWAVE", "0")))
+  if leaf_halfwave and (strassen5 or leaf != 64 or leaf_threads != 64):
+    raise ValueError("LEAF_HALFWAVE requires a four-level N=1024 run with LEAF_THREADS=64")
+  transform_half = bool(int(os.getenv("TRANSFORM_HALF", "1")))
+  product_half = bool(int(os.getenv("PRODUCT_HALF", "1")))
+  transform_vec = int(os.getenv("TRANSFORM_VEC", os.getenv("MEMORY_VEC", "2")))
+  combine_vec = int(os.getenv("COMBINE_VEC", os.getenv("MEMORY_VEC", "4")))
+  transform_image = bool(int(os.getenv("TRANSFORM_IMAGE", "0")))
+  transform_output_image = bool(int(os.getenv("TRANSFORM_OUTPUT_IMAGE", "0")))
+  transform_compute_half = bool(int(os.getenv("TRANSFORM_COMPUTE_HALF", "1")))
+  combine_compute_half = bool(int(os.getenv("COMBINE_COMPUTE_HALF", "0")))
+  split_top_combine = bool(int(os.getenv("SPLIT_TOP_COMBINE", "0")))
+  cpu_pipeline = bool(int(os.getenv("CPU_PIPELINE", "0")))
+  stream_top = bool(int(os.getenv("STREAM_TOP", "0")))
+  cpu_threads = int(os.getenv("CPU_THREADS", "6"))
+  if cpu_pipeline and stream_top: raise ValueError("CPU_PIPELINE and STREAM_TOP are mutually exclusive")
+  if strassen5 and cpu_pipeline: raise ValueError("STRASSEN5 does not support CPU_PIPELINE")
+  if strassen5 and (not transform_half or not product_half):
+    raise ValueError("STRASSEN5 currently requires FP16 transform and product storage")
+  if transform_image and transform_vec != 4: raise ValueError("TRANSFORM_IMAGE requires TRANSFORM_VEC=4")
+  if transform_output_image and transform_vec != 4: raise ValueError("TRANSFORM_OUTPUT_IMAGE requires TRANSFORM_VEC=4")
+  if cpu_pipeline and transform_output_image: raise ValueError("CPU_PIPELINE does not support image transform output")
+  if transform_vec not in (1, 2, 4, 8, 16) or combine_vec not in (2, 4, 8, 16):
+    raise ValueError("TRANSFORM_VEC must be 1, 2, 4, 8, or 16; COMBINE_VEC must be 2, 4, 8, or 16")
+  rng = np.random.default_rng(seed)
+  def random_half() -> np.ndarray:
+    out = np.empty((n, n), np.float16)
+    for first in range(0, n, 64):
+      out[first:first+64] = rng.standard_normal((min(64, n-first), n), dtype=np.float32)*np.float32(1/32)
+    return out
+  a_np, b_np = random_half(), random_half()
+  dev = Device["QCOM"]
+  a, b = alloc(n*n, dtypes.half), alloc(n*n, dtypes.half)
+  a.copyin(memoryview(a_np).cast("B"))
+  b.copyin(memoryview(b_np).cast("B"))
+  gpu_ref = bool(int(os.getenv("GPU_REF", "1" if n >= 8192 else "0")))
+  if gpu_ref:
+    del a_np, b_np
+    gc.collect()
+
+  def transform_runtime(size: int, side: str, input_half: bool):
+    idt = dtypes.half if input_half else dtypes.float
+    odt = dtypes.half if transform_half else dtypes.float
+    input_image = transform_image or (transform_output_image and size == child and not stream_top)
+    return dev.runtime("transform", dev.compiler.compile(
+      transform2_src(size, side, input_half, transform_half, wg, transform_vec, input_image, transform_output_image,
+                     transform_compute_half)), buf_dtypes=[
+      ((0, odt, (49*size//4, size//16, 4)),) if transform_output_image else ((0, odt, (49*(size//4)**2,)),),
+      ((0, idt, (size, size//4, 4)),) if input_image else ((0, idt, (size*size,)),)])
+
+  path_vec = int(os.getenv("PATH_VEC", "4"))
+  if path_vec != 4: raise ValueError("PATH_VEC must be 4 for image-backed streamed paths")
+  top_paths = ({side: [dev.runtime("transform_path", dev.compiler.compile(
+    transform2_path_src(n, side, p, wg, path_vec, transform_compute_half)), buf_dtypes=[
+    ((0, dtypes.half, (child*child,)),), ((0, dtypes.half, (n, n//4, 4)),)]) for p in range(49)] for side in "AB"}
+    if cpu_pipeline or stream_top else None)
+  top_t = None if cpu_pipeline or stream_top else {side: transform_runtime(n, side, True) for side in "AB"}
+  child_t = None if cpu_pipeline else {side: transform_runtime(child, side, transform_half) for side in "AB"}
+  extra_t = ({side: dev.runtime("transform1", dev.compiler.compile(
+    transform1_batch_src(mid, 49, side, wg, transform_vec)), buf_dtypes=[
+      ((0, dtypes.half, (343*leaf*leaf,)),), ((0, dtypes.half, (49*mid*mid,)),)]) for side in "AB"}
+    if strassen5 and not cpu_pipeline else None)
+  top_c_specs = [((0, dtypes.float, (n*n,)),),
+                 ((0, dtypes.half if product_half else dtypes.float, (49*child*child,)),)]
+  top_c = ([dev.runtime("combine", dev.compiler.compile(
+    combine2_block_src(n, block, combine_wg, product_half, False, combine_vec)), buf_dtypes=top_c_specs) for block in range(16)]
+    if split_top_combine else dev.runtime("combine", dev.compiler.compile(
+      combine2_src(n, combine_wg, product_half, False, combine_vec)), buf_dtypes=top_c_specs))
+  child_c = (None if cpu_pipeline else dev.runtime("combine", dev.compiler.compile(
+    combine2_src(child, combine_wg, product_half, product_half, combine_vec, combine_compute_half)), buf_dtypes=[
+      ((0, dtypes.half if product_half else dtypes.float, (child*child,)),),
+      ((0, dtypes.half if product_half else dtypes.float, (49*mid*mid,)),)]))
+  extra_c = (dev.runtime("combine1", dev.compiler.compile(
+    combine1_batch_src(mid, 49, combine_wg, combine_vec)), buf_dtypes=[
+      ((0, dtypes.half, (49*mid*mid,)),), ((0, dtypes.half, (343*leaf*leaf,)),)])
+    if strassen5 and not cpu_pipeline else None)
+
+  q.M = q.N = q.K = leaf
+  q.K4 = leaf//4
+  leaf_unroll = int(os.getenv("LEAF_UNROLL", "3"))
+  leaf_batch_from_row = leaf > 0 and not (leaf & (leaf-1))
+  env, io, sz, ro = get_envelope(dev, q.make_direct_image_donor_src(2, leaf_threads))
+  if leaf_halfwave:
+    shader, hregs, fregs, loop_instrs = q.build_4x4_fp32_halfwave_batch_shader(dev, leaf_threads, batch_stride=leaf)
+  else:
+    shader, hregs, fregs, loop_instrs = q.build_4x8_fp32_rotate_shader(
+      dev, leaf_threads, k_count=leaf//4, batch_stride=leaf, batch_from_row=leaf_batch_from_row, k_unroll=leaf_unroll)
+  lib = inject(env, io, sz, ro, shader, fregs=fregs, hregs=hregs, mergedregs=False)
+  gemms = {}
+  def gemm_runtime(batch: int):
+    if batch not in gemms:
+      gemms[batch] = dev.runtime("gemm_h", lib, buf_dtypes=[
+        ((0, dtypes.half if product_half else dtypes.float, (batch*leaf, leaf//4, 4)),),
+        ((0, dtypes.half if transform_half else dtypes.float, (batch*leaf, leaf//4, 4)),),
+        ((1, dtypes.half if transform_half else dtypes.float, (batch*leaf, leaf//4, 4)),)])
+    return gemms[batch]
+
+  def grid(groups: int) -> tuple[int, int, int]:
+    gx = min(groups, 1024)
+    while groups % gx: gx -= 1
+    return gx, groups//gx, 1
+
+  transform_dt = dtypes.half if transform_half else dtypes.float
+  product_dt = dtypes.half if product_half else dtypes.float
+  if cpu_pipeline:
+    top_a = top_b = None
+  elif stream_top:
+    top_a, top_b = alloc(child*child, transform_dt), alloc(child*child, transform_dt)
+  else:
+    top_a, top_b = alloc(49*child*child, transform_dt), alloc(49*child*child, transform_dt)
+  top_m = alloc(49*child*child, product_dt)
+  times = Times()
+  candidate_start = time.perf_counter()
+  top_groups = (child*child//transform_vec+wg-1)//wg
+  path_groups = (child*child//path_vec+wg-1)//wg
+  top_combine_groups = (child*child//combine_vec+combine_wg-1)//combine_wg
+  if not cpu_pipeline and not stream_top:
+    times.transform += top_t["A"](top_a._buf, a._buf, global_size=grid(top_groups), local_size=(wg, 1, 1), wait=True)
+    times.transform += top_t["B"](top_b._buf, b._buf, global_size=grid(top_groups), local_size=(wg, 1, 1), wait=True)
+  child_input_bytes = child*child*(2 if transform_half else 4)
+  child_output_bytes = child*child*(2 if product_half else 4)
+  leaf_input_bytes = leaf*leaf*(2 if transform_half else 4)
+  leaf_output_bytes = leaf*leaf*(2 if product_half else 4)
+  child_transform_groups = (mid*mid//transform_vec+wg-1)//wg
+  extra_transform_groups = (49*leaf*leaf//transform_vec+wg-1)//wg
+  extra_combine_groups = (49*leaf*leaf//combine_vec+combine_wg-1)//combine_wg
+  child_combine_groups = (mid*mid//combine_vec+combine_wg-1)//combine_wg
+  def launch_leaf_products(lm, la, lb, wait: bool, products: int = 49) -> None:
+    max_batch = 8192//leaf
+    for first in range(0, products, max_batch):
+      batch = min(max_batch, products-first)
+      ispan, ospan = batch*leaf_input_bytes, batch*leaf_output_bytes
+      elapsed = gemm_runtime(batch)(lm._buf.offset(first*leaf_output_bytes, ospan), la._buf.offset(first*leaf_input_bytes, ispan),
+        lb._buf.offset(first*leaf_input_bytes, ispan),
+        global_size=((1, batch*leaf//16, 1) if leaf_halfwave else
+                     (max(1, leaf//256), batch*leaf//((leaf_threads//32)*4), 1) if leaf_batch_from_row else
+                     (max(1, leaf//256), leaf//((leaf_threads//32)*4), batch)),
+        local_size=(leaf_threads, 1, 1), wait=wait)
+      if elapsed is not None: times.gemm += elapsed
+
+  cpu_transform_work = cpu_combine_work = 0.0
+  child_combine_time = top_combine_time = 0.0
+  if cpu_pipeline:
+    if not transform_half or not product_half or transform_image:
+      raise ValueError("CPU_PIPELINE requires FP16 transform/product buffers and buffer-backed transforms")
+    lib_cpu = cpu_pipeline_lib(cpu_threads)
+    slots = []
+    for _ in range(2):
+      lm, lm_backing = alloc_external(49*leaf*leaf, product_dt)
+      slots.append((alloc(49*leaf*leaf, transform_dt), alloc(49*leaf*leaf, transform_dt), lm, lm_backing))
+    top_parent_slots = [(alloc_external(child*child, transform_dt), alloc_external(child*child, transform_dt)) for _ in range(2)]
+    scratch_a = ctypes.create_string_buffer(49*leaf*leaf*2)
+    scratch_b = ctypes.create_string_buffer(49*leaf*leaf*2)
+    scratch_c = ctypes.create_string_buffer(child*child*2)
+    scratch_a_addr, scratch_b_addr, scratch_c_addr = ctypes.addressof(scratch_a), ctypes.addressof(scratch_b), ctypes.addressof(scratch_c)
+    def prepare_operands(a_parent, b_parent, slot) -> None:
+      nonlocal cpu_transform_work
+      la, lb, _, _ = slot
+      cpu_start = time.perf_counter()
+      lib_cpu.transform2_f16(scratch_a_addr, a_parent._buf.cpu_view().addr, child, 0)
+      lib_cpu.transform2_f16(scratch_b_addr, b_parent._buf.cpu_view().addr, child, 1)
+      ctypes.memmove(la._buf.cpu_view().addr, scratch_a_addr, la._buf.size)
+      ctypes.memmove(lb._buf.cpu_view().addr, scratch_b_addr, lb._buf.size)
+      cpu_transform_work += time.perf_counter()-cpu_start
+
+    (a_parent, _), (b_parent, _) = top_parent_slots[0]
+    times.transform += top_paths["A"][0](a_parent._buf, a._buf, global_size=grid(path_groups), local_size=(wg, 1, 1), wait=True)
+    times.transform += top_paths["B"][0](b_parent._buf, b._buf, global_size=grid(path_groups), local_size=(wg, 1, 1), wait=True)
+    lib_cpu.cache_invalidate(a_parent._buf.cpu_view().addr, a_parent._buf.size)
+    lib_cpu.cache_invalidate(b_parent._buf.cpu_view().addr, b_parent._buf.size)
+    prepare_operands(a_parent, b_parent, slots[0])
+    for p in range(49):
+      la, lb, lm, _ = slots[p&1]
+      path_signal = None
+      if p+1 < 49:
+        (next_a_parent, _), (next_b_parent, _) = top_parent_slots[(p+1)&1]
+        top_paths["A"][p+1](next_a_parent._buf, a._buf, global_size=grid(path_groups), local_size=(wg, 1, 1), wait=False)
+        top_paths["B"][p+1](next_b_parent._buf, b._buf, global_size=grid(path_groups), local_size=(wg, 1, 1), wait=False)
+        path_signal = dev.timeline_value-1
+      launch_leaf_products(lm, la, lb, wait=False)
+      if path_signal is not None:
+        dev.timeline_signal.wait(path_signal)
+        lib_cpu.cache_invalidate(next_a_parent._buf.cpu_view().addr, next_a_parent._buf.size)
+        lib_cpu.cache_invalidate(next_b_parent._buf.cpu_view().addr, next_b_parent._buf.size)
+        prepare_operands(next_a_parent, next_b_parent, slots[(p+1)&1])
+      if p:
+        cpu_start = time.perf_counter()
+        prev_lm = slots[(p-1)&1][2]
+        lib_cpu.cache_invalidate(prev_lm._buf.cpu_view().addr, prev_lm._buf.size)
+        out = top_m._buf.offset((p-1)*child_output_bytes, child_output_bytes)
+        lib_cpu.combine2_f16(scratch_c_addr, prev_lm._buf.cpu_view().addr, leaf)
+        ctypes.memmove(out.cpu_view().addr, scratch_c_addr, out.size)
+        cpu_combine_work += time.perf_counter()-cpu_start
+      dev.synchronize()
+    la, lb, lm, _ = slots[0]
+    lib_cpu.cache_invalidate(lm._buf.cpu_view().addr, lm._buf.size)
+    cpu_start = time.perf_counter()
+    out = top_m._buf.offset(48*child_output_bytes, child_output_bytes)
+    lib_cpu.combine2_f16(scratch_c_addr, lm._buf.cpu_view().addr, leaf)
+    ctypes.memmove(out.cpu_view().addr, scratch_c_addr, out.size)
+    cpu_combine_work += time.perf_counter()-cpu_start
+  else:
+    for p in range(49):
+      if stream_top:
+        times.transform += top_paths["A"][p](top_a._buf, a._buf, global_size=grid(path_groups), local_size=(wg, 1, 1), wait=True)
+        times.transform += top_paths["B"][p](top_b._buf, b._buf, global_size=grid(path_groups), local_size=(wg, 1, 1), wait=True)
+      mid_a, mid_b = alloc(49*mid*mid, transform_dt), alloc(49*mid*mid, transform_dt)
+      pa = top_a._buf if stream_top else top_a._buf.offset(p*child_input_bytes, child_input_bytes)
+      pb = top_b._buf if stream_top else top_b._buf.offset(p*child_input_bytes, child_input_bytes)
+      times.transform += child_t["A"](mid_a._buf, pa, global_size=grid(child_transform_groups), local_size=(wg, 1, 1), wait=True)
+      times.transform += child_t["B"](mid_b._buf, pb, global_size=grid(child_transform_groups), local_size=(wg, 1, 1), wait=True)
+      if strassen5:
+        la, lb = alloc(343*leaf*leaf, transform_dt), alloc(343*leaf*leaf, transform_dt)
+        lm, mid_m = alloc(343*leaf*leaf, product_dt), alloc(49*mid*mid, product_dt)
+        times.transform += extra_t["A"](la._buf, mid_a._buf, global_size=grid(extra_transform_groups), local_size=(wg, 1, 1), wait=True)
+        times.transform += extra_t["B"](lb._buf, mid_b._buf, global_size=grid(extra_transform_groups), local_size=(wg, 1, 1), wait=True)
+        launch_leaf_products(lm, la, lb, wait=True, products=343)
+        elapsed = extra_c(mid_m._buf, lm._buf, global_size=grid(extra_combine_groups), local_size=(combine_wg, 1, 1), wait=True)
+        times.combine += elapsed
+        child_combine_time += elapsed
+      else:
+        la, lb, lm, mid_m = mid_a, mid_b, alloc(49*leaf*leaf, product_dt), None
+        launch_leaf_products(lm, la, lb, wait=True)
+      out = top_m._buf.offset(p*child_output_bytes, child_output_bytes)
+      elapsed = child_c(out, (mid_m if strassen5 else lm)._buf, global_size=grid(child_combine_groups),
+                        local_size=(combine_wg, 1, 1), wait=True)
+      times.combine += elapsed
+      child_combine_time += elapsed
+  allocation_start = time.perf_counter()
+  if cpu_pipeline: del top_parent_slots
+  else: del top_a, top_b
+  gc.collect()
+  c = alloc(n*n, dtypes.float)
+  allocation_overhead = time.perf_counter()-allocation_start
+  if split_top_combine:
+    split_groups = (child*child//combine_vec+combine_wg-1)//combine_wg
+    top_combine_time = sum(prg(c._buf, top_m._buf, global_size=grid(split_groups), local_size=(combine_wg, 1, 1), wait=True)
+                           for prg in top_c)
+  else:
+    top_combine_time = top_c(c._buf, top_m._buf, global_size=grid(top_combine_groups), local_size=(combine_wg, 1, 1), wait=True)
+  times.combine += top_combine_time
+  if cpu_pipeline: times.wall = time.perf_counter()-candidate_start-allocation_overhead
+
+  rtol = float(os.getenv("RTOL", "1e-3"))
+  atol = float(os.getenv("ATOL", "8e-3" if product_half else "5e-3" if transform_half else "2e-3"))
+  ref_block = int(os.getenv("REF_BLOCK", "256"))
+  bad_count, max_abs, sum_abs, err2, ref2 = 0, 0.0, 0.0, 0.0, 0.0
+  reference_ms = 0.0
+  if gpu_ref:
+    # Free recursive intermediates, then compute an independent conventional
+    # FP32-accumulating GEMM in row slabs. It is an oracle only and is excluded
+    # from times.total; slab streaming avoids a second N*N float allocation.
+    del top_m, la, lb, lm
+    if not cpu_pipeline: del mid_a, mid_b
+    if strassen5: del mid_m
+    if cpu_pipeline: del slots, scratch_a, scratch_b, scratch_c
+    gc.collect()
+    if n % ref_block or ref_block % 16: raise ValueError("GPU reference requires REF_BLOCK to divide N and be a multiple of 16")
+    ref = alloc(ref_block*n, dtypes.float)
+    q.M, q.N, q.K = ref_block, n, n
+    q.K4 = n//4
+    renv, rio, rsz, rro = get_envelope(dev, q.make_direct_image_donor_src(2, 128))
+    ref_k4 = int(os.getenv("REF_K4", "256"))
+    references = []
+    for kstart in range(0, n//4, ref_k4):
+      kcount = min(ref_k4, n//4-kstart)
+      rshader, rhregs, rfregs, _ = q.build_4x8_fp32_rotate_shader(dev, 128, k_count=kcount, k_start=kstart)
+      rlib = inject(renv, rio, rsz, rro, rshader, fregs=rfregs, hregs=rhregs, mergedregs=False)
+      references.append(dev.runtime("gemm_h", rlib, buf_dtypes=[
+        ((0, dtypes.float, (ref_block, n//4, 4)),), ((0, dtypes.half, (ref_block, n//4, 4)),),
+        ((1, dtypes.half, (n, n//4, 4)),)]))
+    for first in range(0, n, ref_block):
+      av = a._buf.offset(first*n*2, ref_block*n*2)
+      got = np.empty((ref_block, n), np.float32)
+      expected = np.zeros((ref_block, n), np.float32)
+      partial = np.empty((ref_block, n), np.float32)
+      for reference in references:
+        reference_ms += reference(ref._buf, av, b._buf, global_size=(n//256, ref_block//16, 1),
+                                  local_size=(128, 1, 1), wait=True)*1e3
+        ref.copyout(memoryview(partial).cast("B"))
+        expected += partial
+      c.view(ref_block*n, dtypes.float, first*n*4).ensure_allocated().copyout(memoryview(got).cast("B"))
+      err = got-expected
+      delta = np.abs(err)
+      bad_count += int(np.count_nonzero(~np.isclose(got, expected, rtol=rtol, atol=atol)))
+      max_abs = max(max_abs, float(delta.max()))
+      sum_abs += float(delta.astype(np.float64).sum())
+      err2 += float(np.square(err.astype(np.float64)).sum())
+      ref2 += float(np.square(expected.astype(np.float64)).sum())
+  else:
+    got = np.empty((n, n), np.float32)
+    c.copyout(memoryview(got).cast("B"))
+    del a, b, c, top_m, la, lb, lm
+    gc.collect()
+    ref_k_block = int(os.getenv("REF_K_BLOCK", "256"))
+    for first in range(0, n, ref_block):
+      last = min(n, first+ref_block)
+      expected = np.zeros((last-first, n), np.float32)
+      for kfirst in range(0, n, ref_k_block):
+        klast = min(n, kfirst+ref_k_block)
+        expected += a_np[first:last, kfirst:klast].astype(np.float32) @ b_np[kfirst:klast].astype(np.float32)
+      err = got[first:last]-expected
+      delta = np.abs(err)
+      bad_count += int(np.count_nonzero(~np.isclose(got[first:last], expected, rtol=rtol, atol=atol)))
+      max_abs = max(max_abs, float(delta.max()))
+      sum_abs += float(delta.astype(np.float64).sum())
+      err2 += float(np.square(err.astype(np.float64)).sum())
+      ref2 += float(np.square(expected.astype(np.float64)).sum())
+  mean_abs, rel_l2 = sum_abs/(n*n), (err2/ref2)**0.5
+  print(f"shape={n}x{n}x{n} algorithm={'strassen5_fused2' if strassen5 else 'strassen4_fused2'} "
+        f"inputs=fp16 accumulate=fp32 elapsed_ms={times.total*1e3:.3f} "
+        f"gflops={2*n**3/times.total/1e9:.1f} transform_storage={'fp16' if transform_half else 'fp32'} "
+        f"product_storage={'fp16' if product_half else 'fp32'} "
+        f"transform_vec={transform_vec} combine_vec={combine_vec} leaf_unroll={leaf_unroll} "
+        f"pipeline={'cpu_gpu' if cpu_pipeline else 'stream_gpu' if stream_top else 'serial_gpu'} "
+        f"cpu_threads={cpu_threads if cpu_pipeline else 0} "
+        f"transform_ms={times.transform*1e3:.3f} gemm_ms={times.gemm*1e3:.3f} "
+        f"combine_ms={times.combine*1e3:.3f} child_combine_ms={child_combine_time*1e3:.3f} "
+        f"top_combine_ms={top_combine_time*1e3:.3f} cpu_transform_work_ms={cpu_transform_work*1e3:.3f} "
+        f"cpu_combine_work_ms={cpu_combine_work*1e3:.3f} outputs={n*n} bad_count={bad_count} "
+        f"max_abs={max_abs:.9g} mean_abs={mean_abs:.9g} rel_l2={rel_l2:.9g} "
+        f"rtol={rtol:g} atol={atol:g} allclose={bad_count == 0} oracle={'gpu_direct' if gpu_ref else 'numpy'} "
+        f"reference_ms={reference_ms:.3f} loop_instrs={loop_instrs}")
+  if bad_count: raise SystemExit(1)
+
+
+if __name__ == "__main__": main()

extra/gemm/qcom_strassen5_fp32_check.py

@@ -0,0 +1,228 @@
+#!/usr/bin/env python3
+"""Full Strassen GEMM with dense FP16 inputs, FP32 transforms/accumulation, and output validation."""
+import os
+from dataclasses import dataclass
+
+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 transform_src(n: int, side: str, parallel: bool = False, buffer_output: bool = False,
+                  fp32: bool = False, fp32_input: bool = False, wg: int = 128) -> str:
+  h = n//2
+  vals = (["x0+x3", "x2+x3", "x0", "x3", "x0+x1", "x2-x0", "x1-x3"] if side == "A" else
+          ["x0+x3", "x0", "x1-x3", "x2-x0", "x3", "x0+x1", "x2+x3"])
+  vec, read, write = ("float4", "read_imagef" if fp32_input else "convert_float4(read_imageh", "write_imagef") if fp32 else \
+                     ("half4", "read_imageh", "write_imageh")
+  def rd(coord: str) -> str:
+    call = f"{read}(X,smp,(int2)({coord}))"
+    return call+")" if fp32 and not fp32_input else call
+  if parallel:
+    # One output pixel per invocation avoids seven dependent image stores in a
+    # single thread.  It rereads source quadrants, but substantially increases
+    # the number of memory operations the GPU can keep in flight.
+    cases = "\n".join(f"{'if' if p == 0 else 'else if'}(p=={p}) v={v};" for p, v in enumerate(vals))
+    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({wg},1,1)))
+__kernel void transform(write_only image2d_t O,read_only image2d_t X) {{
+  uint i=get_global_id(0),unit={h*h//4},b=i/(7*unit),j=i%(7*unit),pix=j%unit,p=j/unit,r=pix/{h//4},x=pix%{h//4};
+  uint iy=b*{n};
+  {vec} x0={rd('x,iy+r')},x1={rd('x+'+str(h//4)+',iy+r')};
+  {vec} x2={rd('x,iy+r+'+str(h))},x3={rd('x+'+str(h//4)+',iy+r+'+str(h))},v;
+  {cases}
+  {write}(O,(int2)(x,(b*7+p)*{h}+r),v);
+}}"""
+  stores = "\n".join(f"write_imageh(O,(int2)(x,{p*h}+r),{v});" for p, v in enumerate(vals))
+  if buffer_output:
+    stores = "\n".join(f"O[((b*7+{p})*{h}+r)*{h//4}+x]={v};" for p, v in enumerate(vals))
+  elif fp32:
+    stores = "\n".join(f"write_imagef(O,(int2)(x,{p*h}+r),{v});" for p, v in enumerate(vals))
+  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({wg},1,1)))
+__kernel void transform({'__global '+vec+' *O' if buffer_output else 'write_only image2d_t O'},read_only image2d_t X) {{
+  uint i=get_global_id(0),unit={h*h//4},b=i/unit,pix=i%unit,r=pix/{h//4},x=pix%{h//4},iy=b*{n};
+  {vec} x0={rd('x,iy+r')},x1={rd('x+'+str(h//4)+',iy+r')};
+  {vec} x2={rd('x,iy+r+'+str(h))},x3={rd('x+'+str(h//4)+',iy+r+'+str(h))};
+  {stores if buffer_output else stores.replace('(x,', '(x,b*7*'+str(h)+'+')}
+}}"""
+
+
+def combine_src(n: int, wg: int = 128) -> str:
+  h = n//2
+  def rd(p: int) -> str: return f"read_imagef(M,smp,(int2)(xx,(b*7+{p})*{h}+rr))"
+  return f"""const sampler_t smp=CLK_NORMALIZED_COORDS_FALSE|CLK_ADDRESS_CLAMP|CLK_FILTER_NEAREST;
+__attribute__((reqd_work_group_size({wg},1,1)))
+__kernel void combine(write_only image2d_t C,read_only image2d_t M) {{
+  uint i=get_global_id(0),unit={n*n//4},b=i/unit,pix=i%unit,r=pix/{n//4},x=pix%{n//4},rr=r%{h},xx=x%{h//4}; float4 v;
+  if(r<{h}&&x<{h//4}) v={rd(0)}+{rd(3)}-{rd(4)}+{rd(6)};
+  else if(r<{h}) v={rd(2)}+{rd(4)};
+  else if(x<{h//4}) v={rd(1)}+{rd(3)};
+  else v={rd(0)}-{rd(1)}+{rd(2)}+{rd(5)};
+  write_imagef(C,(int2)(x,b*{n}+r),v);
+}}"""
+
+
+def alloc(count: int, dtype) -> Buffer:
+  return Buffer("QCOM", count, dtype).allocate()
+
+
+@dataclass
+class Times:
+  transform: float = 0.0
+  gemm: float = 0.0
+  combine: float = 0.0
+  leaves: int = 0
+
+  @property
+  def total(self) -> float: return self.transform+self.gemm+self.combine
+
+
+def main() -> None:
+  size, levels = int(os.getenv("N", "8192")), int(os.getenv("LEVELS", "5"))
+  seed = int(os.getenv("SEED", "701"))
+  parallel_transform = bool(int(os.getenv("PARALLEL_TRANSFORM", "0")))
+  buffer_transform = bool(int(os.getenv("BUFFER_TRANSFORM", "0")))
+  fp32_transform = bool(int(os.getenv("FP32_TRANSFORM", "1")))
+  memory_wg = int(os.getenv("MEMORY_WG", "256"))
+  if parallel_transform and buffer_transform: raise ValueError("parallel buffer transform is not implemented")
+  leaf = size >> levels
+  if leaf != 256: raise ValueError("current FP32 leaf requires N >> LEVELS == 256")
+  rng = np.random.default_rng(seed)
+  data = os.getenv("DATA", "gaussian")
+  if data == "exact":
+    # Dense signed powers of two keep all transforms and the FP32 reference
+    # exactly representable.  This is a strict indexing/instruction oracle.
+    a_np = (rng.integers(0, 2, (size, size), dtype=np.int8)*2-1).astype(np.float16)*np.float16(1/256)
+    b_np = (rng.integers(0, 2, (size, size), dtype=np.int8)*2-1).astype(np.float16)*np.float16(1/256)
+  elif data == "gaussian":
+    a_np = rng.normal(0, 1/32, (size, size)).astype(np.float16)
+    b_np = rng.normal(0, 1/32, (size, size)).astype(np.float16)
+  else: raise ValueError("DATA must be exact or gaussian")
+  dev = Device["QCOM"]
+  a, b, c = alloc(size*size, dtypes.half), alloc(size*size, dtypes.half), alloc(size*size, dtypes.float)
+  a.copyin(memoryview(a_np).cast("B"))
+  b.copyin(memoryview(b_np).cast("B"))
+
+  # Rebuild with the selected workgroup size; this only affects the bandwidth
+  # transforms/combine, not the 128-thread hand-written GEMM leaf.
+  transform_libs = {(n, side): dev.compiler.compile(transform_src(n, side, parallel_transform, buffer_transform,
+                                                                  fp32_transform, fp32_transform and n != size, memory_wg))
+                    for n in (size >> x for x in range(levels)) for side in "AB"}
+  combine_libs = {n: dev.compiler.compile(combine_src(n, memory_wg)) for n in (size >> x for x in range(levels))}
+  transforms, combines = {}, {}
+
+  def get_transform(n: int, side: str, count: int):
+    key = (n, side, count)
+    if key not in transforms:
+      odt, idt = (dtypes.float if fp32_transform else dtypes.half), \
+                 (dtypes.float if fp32_transform and n != size else dtypes.half)
+      transforms[key] = dev.runtime("transform", transform_libs[(n, side)],
+        buf_dtypes=[((0, odt, (count*7*(n//2)*(n//2),)),) if buffer_transform else
+                    ((0, odt, (count*7*(n//2), n//8, 4)),),
+                    ((0, idt, (count*n, n//4, 4)),)])
+    return transforms[key]
+
+  def get_combine(n: int, count: int):
+    key = (n, count)
+    if key not in combines:
+      combines[key] = dev.runtime("combine", combine_libs[n],
+        buf_dtypes=[((0, dtypes.float, (count*n, n//4, 4)),),
+                    ((0, dtypes.float, (count*7*(n//2), n//8, 4)),)])
+    return combines[key]
+
+  q.M = q.N = q.K = leaf
+  q.K4 = leaf//4
+  env, io, sz, ro = get_envelope(dev, q.make_direct_image_donor_src(2, 128))
+  shader, hregs, fregs, loop_instrs = q.build_4x8_fp32_rotate_shader(
+    dev, 128, k_count=leaf//4, batch_stride=leaf, batch_from_row=True)
+  lib = inject(env, io, sz, ro, shader, fregs=fregs, hregs=hregs, mergedregs=False)
+  gemms = {}
+  def get_gemm(count: int):
+    if count not in gemms:
+      out_image = ((0, dtypes.float, (count*7*leaf, leaf//4, 4)),)
+      leaf_dt = dtypes.float if fp32_transform else dtypes.half
+      a_image = ((0, leaf_dt, (count*7*leaf, leaf//4, 4)),)
+      b_image = ((1, leaf_dt, (count*7*leaf, leaf//4, 4)),)
+      gemms[count] = dev.runtime("gemm_h", lib, buf_dtypes=[out_image, a_image, b_image])
+    return gemms[count]
+  print(f"shape={size}x{size}x{size} levels={levels} leaf={leaf} leaf_fregs={fregs} leaf_loop={loop_instrs}", flush=True)
+
+  keepalive: list[Buffer] = []
+  times = Times()
+
+  max_batch = int(os.getenv("BATCH_NODES", "4"))
+  if not 1 <= max_batch <= 4: raise ValueError("BATCH_NODES must be 1..4 for the 8192-row image limit")
+
+  def recurse(n: int, count: int, ax, bx, out) -> None:
+    h = n//2
+    transform_dt = dtypes.float if fp32_transform else dtypes.half
+    ac, bc, mm = alloc(count*7*h*h, transform_dt), alloc(count*7*h*h, transform_dt), alloc(count*7*h*h, dtypes.float)
+    keepalive.extend((ac, bc, mm))
+    groups = (count*(7 if parallel_transform else 1)*h*h//4 + memory_wg-1)//memory_wg
+    times.transform += get_transform(n, "A", count)(ac._buf, ax, global_size=(groups, 1, 1), local_size=(memory_wg, 1, 1), wait=True)
+    times.transform += get_transform(n, "B", count)(bc._buf, bx, global_size=(groups, 1, 1), local_size=(memory_wg, 1, 1), wait=True)
+    if n == size and count == 1 and int(os.getenv("DEBUG_STAGE", "0")):
+      tnp = np.float32 if fp32_transform else np.float16
+      ah = np.empty((7, h, h), tnp)
+      bh = np.empty((7, h, h), tnp)
+      ac.copyout(memoryview(ah).cast("B"))
+      bc.copyout(memoryview(bh).cast("B"))
+      aq = [a_np[:h,:h], a_np[:h,h:], a_np[h:,:h], a_np[h:,h:]]
+      bq = [b_np[:h,:h], b_np[:h,h:], b_np[h:,:h], b_np[h:,h:]]
+      ae = [aq[0]+aq[3], aq[2]+aq[3], aq[0], aq[3], aq[0]+aq[1], aq[2]-aq[0], aq[1]-aq[3]]
+      be = [bq[0]+bq[3], bq[0], bq[1]-bq[3], bq[2]-bq[0], bq[3], bq[0]+bq[1], bq[2]+bq[3]]
+      print("stage_transform", max(float(np.max(np.abs(ah[i].astype(np.float32)-ae[i]))) for i in range(7)),
+            max(float(np.max(np.abs(bh[i].astype(np.float32)-be[i]))) for i in range(7)), flush=True)
+      print("stage_a_parts", [(i, int(np.count_nonzero(ah[i] != ae[i])), ah[i,0,:8].tolist(), ae[i][0,:8].tolist())
+                              for i in range(7)], flush=True)
+    if h == leaf:
+      times.gemm += get_gemm(count)(mm._buf, ac._buf, bc._buf, global_size=(1, count*7*leaf//16, 1),
+                                    local_size=(128, 1, 1), wait=True)
+      times.leaves += count*7
+      if n == size and count == 1 and int(os.getenv("DEBUG_STAGE", "0")):
+        mh = np.empty((7, h, h), np.float32)
+        mm.copyout(memoryview(mh).cast("B"))
+        print("stage_gemm", max(float(np.max(np.abs(mh[i]-ae[i].astype(np.float32)@be[i].astype(np.float32)))) for i in range(7)), flush=True)
+    else:
+      hb, fb = h*h*(4 if fp32_transform else 2), h*h*4
+      children = count*7
+      # The child's transform output stacks seven h/2-row matrices per input.
+      # Keep every typed image at or below A630's 8192-row limit.
+      child_batch = min(max_batch, 8192//(7*(h//2)))
+      for first in range(0, children, child_batch):
+        batch = min(child_batch, children-first)
+        recurse(h, batch, ac._buf.offset(first*hb, batch*hb), bc._buf.offset(first*hb, batch*hb),
+                mm._buf.offset(first*fb, batch*fb))
+    groups_out = (count*n*n//4 + memory_wg-1)//memory_wg
+    times.combine += get_combine(n, count)(out, mm._buf, global_size=(groups_out, 1, 1), local_size=(memory_wg, 1, 1), wait=True)
+    # Recursive children have completed before their parents return. Dropping
+    # these references bounds live storage to one seven-way branch per level.
+    keepalive.pop()
+    keepalive.pop()
+    keepalive.pop()
+
+  recurse(size, 1, a._buf, b._buf, c._buf)
+  got = np.empty((size, size), np.float32)
+  c.copyout(memoryview(got).cast("B"))
+  expected = a_np.astype(np.float32) @ b_np.astype(np.float32)
+  delta = np.abs(got-expected)
+  if data == "exact": wrong = np.flatnonzero(got.reshape(-1).view(np.uint32) != expected.reshape(-1).view(np.uint32))
+  else:
+    rtol, atol = float(os.getenv("RTOL", "1e-3")), float(os.getenv("ATOL", "1e-3"))
+    wrong = np.flatnonzero(~np.isclose(got.reshape(-1), expected.reshape(-1), rtol=rtol, atol=atol))
+  rel_l2 = float(np.linalg.norm((got-expected).astype(np.float64))/np.linalg.norm(expected.astype(np.float64)))
+  conventional = 2*size**3
+  print(f"elapsed_ms={times.total*1e3:.3f} gflops={conventional/times.total/1e9:.1f} "
+        f"transform_ms={times.transform*1e3:.3f} gemm_ms={times.gemm*1e3:.3f} combine_ms={times.combine*1e3:.3f} "
+        f"leaves={times.leaves} data={data} outputs={got.size} bad_count={wrong.size} "
+        f"max_abs={float(delta.max()):.9g} mean_abs={float(delta.mean()):.9g} rel_l2={rel_l2:.9g}", flush=True)
+  if wrong.size: raise SystemExit(1)
+
+
+if __name__ == "__main__": main()

extra/gemm/qcom_strassen_check.py

@@ -0,0 +1,120 @@
+#!/usr/bin/env python3
+"""Fully checked one-level Strassen FP16 GEMM for Adreno 630."""
+import hashlib, os
+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 get_envelope, inject
+
+
+PREP_SRC = r"""
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+__attribute__((reqd_work_group_size(128,1,1)))
+__kernel void strassen_prep(__global const half *A, __global const half *B,
+  __global half *A1, __global half *A2, __global half *A5, __global half *A6, __global half *A7,
+  __global half *B1, __global half *B3, __global half *B4, __global half *B6, __global half *B7) {
+  int i=get_global_id(0), r=i>>7, c=(i&127)<<2, o=r*512+c;
+  half4 a11=vload4(0,A+r*1024+c), a12=vload4(0,A+r*1024+c+512);
+  half4 a21=vload4(0,A+(r+512)*1024+c), a22=vload4(0,A+(r+512)*1024+c+512);
+  half4 b11=vload4(0,B+r*1024+c), b12=vload4(0,B+r*1024+c+512);
+  half4 b21=vload4(0,B+(r+512)*1024+c), b22=vload4(0,B+(r+512)*1024+c+512);
+  vstore4(a11+a22,0,A1+o); vstore4(a21+a22,0,A2+o); vstore4(a11+a12,0,A5+o);
+  vstore4(a21-a11,0,A6+o); vstore4(a12-a22,0,A7+o);
+  vstore4(b11+b22,0,B1+o); vstore4(b12-b22,0,B3+o); vstore4(b21-b11,0,B4+o);
+  vstore4(b11+b12,0,B6+o); vstore4(b21+b22,0,B7+o);
+}
+"""
+
+POST_SRC = r"""
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+__attribute__((reqd_work_group_size(128,1,1)))
+__kernel void strassen_post(__global const half *M1, __global const half *M2,
+  __global const half *M3, __global const half *M4, __global const half *M5,
+  __global const half *M6, __global const half *M7, __global half *C) {
+  int i=get_global_id(0), r=i>>7, c=(i&127)<<2, o=r*512+c;
+  half4 m1=vload4(0,M1+o),m2=vload4(0,M2+o),m3=vload4(0,M3+o),m4=vload4(0,M4+o);
+  half4 m5=vload4(0,M5+o),m6=vload4(0,M6+o),m7=vload4(0,M7+o);
+  vstore4(m1+m4-m5+m7,0,C+r*1024+c); vstore4(m3+m5,0,C+r*1024+c+512);
+  vstore4(m2+m4,0,C+(r+512)*1024+c); vstore4(m1-m2+m3+m6,0,C+(r+512)*1024+c+512);
+}
+"""
+
+
+def alloc_half(count: int) -> Buffer:
+  return Buffer("QCOM", count, dtypes.half).allocate()
+
+
+def vectorize(src: str, vec: int) -> str:
+  if vec == 4: return src
+  if vec not in (8, 16): raise ValueError("TRANSFORM_VEC must be 4, 8, or 16")
+  log_cols = {8: 6, 16: 5}[vec]
+  return (src.replace("i>>7", f"i>>{log_cols}").replace("(i&127)<<2", f"(i&{(512//vec)-1})<<{vec.bit_length()-1}")
+          .replace("half4", f"half{vec}").replace("vload4", f"vload{vec}").replace("vstore4", f"vstore{vec}"))
+
+
+def main() -> None:
+  seed, runs = int(os.getenv("SEED", "307")), int(os.getenv("BENCH_RUNS", "10"))
+  transform_vec = int(os.getenv("TRANSFORM_VEC", "4"))
+  rng = np.random.default_rng(seed)
+  a_np = (rng.standard_normal((1024, 1024))*0.05).astype(np.float16)
+  b_np = (rng.standard_normal((1024, 1024))*0.05).astype(np.float16)
+  dev = Device["QCOM"]
+  a, b = alloc_half(a_np.size), alloc_half(b_np.size)
+  a.copyin(memoryview(a_np).cast("B")); b.copyin(memoryview(b_np).cast("B"))
+  tile_elems, tile_bytes = 512*512, 512*512*2
+  pa, pb, pm = alloc_half(7*tile_elems), alloc_half(7*tile_elems), alloc_half(7*tile_elems)
+  aa = [pa._buf.offset(i*tile_bytes, tile_bytes) for i in range(7)]
+  bb = [pb._buf.offset(i*tile_bytes, tile_bytes) for i in range(7)]
+  mm = [pm._buf.offset(i*tile_bytes, tile_bytes) for i in range(7)]
+  c = alloc_half(1024*1024)
+
+  gspec = ((0, dtypes.half, None),)
+  prep = dev.runtime("strassen_prep", dev.compiler.compile(vectorize(PREP_SRC, transform_vec)), buf_dtypes=[gspec]*12)
+  post = dev.runtime("strassen_post", dev.compiler.compile(vectorize(POST_SRC, transform_vec)), buf_dtypes=[gspec]*8)
+
+  q8.M = q8.N = q8.K = 512; q8.K4 = 128
+  env, io, sz, ro = get_envelope(dev, q4.make_direct_image_donor_src(4, 128))
+  shader, hregs, fregs, loop_instrs = q8.build_8x8_persistent_shader(
+    dev, 128, dynamic_a4_dual=True, image_store=True)
+  lib = inject(env, io, sz, ro, shader, fregs=fregs, hregs=hregs, mergedregs=False)
+  out_spec, std_spec, wide_spec = ((0, dtypes.half, (512, 128, 4)),), ((0, dtypes.half, (512, 128, 4)),), ((0, dtypes.half, (512, 256, 4)),)
+  gemm_ss = dev.runtime("gemm_h", lib, buf_dtypes=[out_spec, std_spec, ((1, dtypes.half, (512, 128, 4)),)])
+  gemm_ws = dev.runtime("gemm_h", lib, buf_dtypes=[out_spec, wide_spec, ((1, dtypes.half, (512, 128, 4)),)])
+  gemm_sw = dev.runtime("gemm_h", lib, buf_dtypes=[out_spec, std_spec, ((1, dtypes.half, (512, 256, 4)),)])
+  quadrant_bytes = (512*1024+512)*2
+  a11, a22, b11, b22 = a._buf, a._buf.offset(quadrant_bytes), b._buf, b._buf.offset(quadrant_bytes)
+  if int(os.getenv("PRINT_META", "0")):
+    print("gemm_meta", fregs, hregs, len(shader), loop_instrs, hashlib.sha1(lib).hexdigest()[:8])
+
+  def iteration() -> tuple[float, list[float], float]:
+    transform_groups = 512*(512//transform_vec)//128
+    tp = prep(a._buf, b._buf, aa[0], aa[1], aa[4], aa[5], aa[6], bb[0], bb[2], bb[3], bb[5], bb[6],
+              global_size=(transform_groups, 1, 1), local_size=(128, 1, 1), wait=True)
+    args = [(gemm_ss, aa[0], bb[0]), (gemm_sw, aa[1], b11), (gemm_ws, a11, bb[2]),
+            (gemm_ws, a22, bb[3]), (gemm_sw, aa[4], b22), (gemm_ss, aa[5], bb[5]), (gemm_ss, aa[6], bb[6])]
+    tg = [prg(mm[i], ax, bx, global_size=(2, 16, 1), local_size=(128, 1, 1), wait=True) for i, (prg, ax, bx) in enumerate(args)]
+    to = post(*mm, c._buf, global_size=(transform_groups, 1, 1), local_size=(128, 1, 1), wait=True)
+    return tp, tg, to
+
+  for _ in range(2): iteration()
+  measured = [iteration() for _ in range(runs)]
+  totals = [tp+sum(tg)+to for tp, tg, to in measured]
+  best_i = int(np.argmin(totals)); tp, tg, to = measured[best_i]; elapsed = totals[best_i]
+
+  got = np.empty((1024, 1024), np.float16); c.copyout(memoryview(got).cast("B"))
+  expected = a_np.astype(np.float32) @ b_np.astype(np.float32)
+  delta = np.abs(got.astype(np.float32)-expected)
+  correct = np.allclose(got, expected, rtol=2e-2, atol=2e-2)
+  bad = ~np.isfinite(got) | (delta > .02)
+  gflops = 2*1024**3/elapsed/1e9
+  print(f"shape=1024x1024x1024 algorithm=strassen1 accumulate=fp16 elapsed_ms={elapsed*1e3:.3f} gflops={gflops:.1f} "
+        f"prep_ms={tp*1e3:.3f} gemm_ms={sum(tg)*1e3:.3f} post_ms={to*1e3:.3f} "
+        f"max_abs={delta.max():.9g} mean_abs={delta.mean():.9g} allclose={correct}")
+  print(f"bad_count={int(bad.sum())} gemm_parts_ms={[round(x*1e3, 3) for x in tg]}")
+  if not correct: raise SystemExit(1)
+
+
+if __name__ == "__main__": main()

extra/gemm/qcom_strassen_cpu_pipeline.c

@@ -0,0 +1,88 @@
+#include <arm_neon.h>
+#include <omp.h>
+#include <stdint.h>
+
+void set_threads(int n) { omp_set_num_threads(n); }
+
+static inline void transform1v(float32x4_t out[7], const float32x4_t x[4], int side) {
+  out[0] = vaddq_f32(x[0], x[3]);
+  if (!side) {
+    out[1] = vaddq_f32(x[2], x[3]); out[2] = x[0]; out[3] = x[3];
+    out[4] = vaddq_f32(x[0], x[1]); out[5] = vsubq_f32(x[2], x[0]); out[6] = vsubq_f32(x[1], x[3]);
+  } else {
+    out[1] = x[0]; out[2] = vsubq_f32(x[1], x[3]); out[3] = vsubq_f32(x[2], x[0]);
+    out[4] = x[3]; out[5] = vaddq_f32(x[0], x[1]); out[6] = vaddq_f32(x[2], x[3]);
+  }
+}
+
+void transform2_f16(_Float16 *restrict out, const _Float16 *restrict in, int n, int side) {
+  const int leaf = n / 4;
+  #pragma omp parallel for schedule(dynamic, 1)
+  for (int r = 0; r < leaf; r++) for (int c = 0; c < leaf; c += 4) {
+    float32x4_t v[4][4], inner[4][7], outer[7];
+    for (int q0 = 0; q0 < 4; q0++) for (int q1 = 0; q1 < 4; q1++) {
+      const int rb = ((q0 >> 1) << 1) | (q1 >> 1);
+      const int cb = ((q0 & 1) << 1) | (q1 & 1);
+      const float16x4_t h = vld1_f16((const float16_t *)(in + (int64_t)(rb * leaf + r) * n + cb * leaf + c));
+      v[q0][q1] = vcvt_f32_f16(h);
+    }
+    for (int q0 = 0; q0 < 4; q0++) transform1v(inner[q0], v[q0], side);
+    for (int p1 = 0; p1 < 7; p1++) {
+      float32x4_t column[4];
+      for (int q0 = 0; q0 < 4; q0++) column[q0] = inner[q0][p1];
+      transform1v(outer, column, side);
+      for (int p0 = 0; p0 < 7; p0++) {
+        const float16x4_t h = vcvt_f16_f32(outer[p0]);
+        vst1_f16((float16_t *)(out + ((int64_t)p0 * 7 + p1) * leaf * leaf + (int64_t)r * leaf + c), h);
+      }
+    }
+  }
+}
+
+static inline void combine1v(float32x4_t out[4], const float32x4_t m[7]) {
+  out[0] = vaddq_f32(vsubq_f32(vaddq_f32(m[0], m[3]), m[4]), m[6]);
+  out[1] = vaddq_f32(m[2], m[4]);
+  out[2] = vaddq_f32(m[1], m[3]);
+  out[3] = vaddq_f32(vaddq_f32(vsubq_f32(m[0], m[1]), m[2]), m[5]);
+}
+
+void combine2_f16(_Float16 *restrict out, const _Float16 *restrict in, int leaf) {
+  const int n = leaf * 4;
+  #pragma omp parallel for schedule(dynamic, 1)
+  for (int r = 0; r < leaf; r++) for (int c = 0; c < leaf; c += 4) {
+    float32x4_t inner[4][7];
+    for (int p0 = 0; p0 < 7; p0++) {
+      float32x4_t m[7], d[4];
+      for (int p1 = 0; p1 < 7; p1++) {
+        const float16x4_t h = vld1_f16((const float16_t *)(in + ((int64_t)p0 * 7 + p1) * leaf * leaf + (int64_t)r * leaf + c));
+        m[p1] = vcvt_f32_f16(h);
+      }
+      combine1v(d, m);
+      for (int u = 0; u < 4; u++) inner[u][p0] = d[u];
+    }
+    for (int u = 0; u < 4; u++) {
+      float32x4_t d[4];
+      combine1v(d, inner[u]);
+      for (int v = 0; v < 4; v++) {
+        const int br = (v >> 1) * 2 + (u >> 1);
+        const int bc = (v & 1) * 2 + (u & 1);
+        const float16x4_t h = vcvt_f16_f32(d[v]);
+        vst1_f16((float16_t *)(out + (int64_t)(br * leaf + r) * n + bc * leaf + c), h);
+      }
+    }
+  }
+}
+
+void cache_clean(void *ptr, int64_t size) {
+  uintptr_t p = (uintptr_t)ptr & ~(uintptr_t)63;
+  const uintptr_t end = ((uintptr_t)ptr + size + 63) & ~(uintptr_t)63;
+  for (; p < end; p += 64) __asm__ volatile("dc cvac, %0" :: "r"(p) : "memory");
+  __asm__ volatile("dsb sy" ::: "memory");
+}
+
+void cache_invalidate(void *ptr, int64_t size) {
+  uintptr_t p = (uintptr_t)ptr & ~(uintptr_t)63;
+  const uintptr_t end = ((uintptr_t)ptr + size + 63) & ~(uintptr_t)63;
+  for (; p < end; p += 64) __asm__ volatile("dc civac, %0" :: "r"(p) : "memory");
+  __asm__ volatile("dsb sy" ::: "memory");
+}

extra/gemm/qcom_strassen_fused_combine_check.py

@@ -0,0 +1,207 @@
+#!/usr/bin/env python3
+"""Fuse two recursive Strassen output combines into one streaming kernel."""
+import os
+
+import numpy as np
+
+from tinygrad import Device, dtypes
+from tinygrad.device import Buffer
+
+
+def alloc(count: int, dtype) -> Buffer: return Buffer("QCOM", count, dtype).allocate()
+
+
+def combine_expr(names: list[str], u: int) -> str:
+  ids, signs = (((0, 3, 4, 6), (1, 1, -1, 1)), ((2, 4), (1, 1)),
+                ((1, 3), (1, 1)), ((0, 1, 2, 5), (1, -1, 1, 1)))[u]
+  return "".join(("+" if s > 0 else "-")+names[i] for i, s in zip(ids, signs)).lstrip("+")
+
+
+COMBINE = ((1, 0, 0, 1, -1, 0, 1), (0, 0, 1, 0, 1, 0, 0),
+           (0, 1, 0, 1, 0, 0, 0), (1, -1, 1, 0, 0, 1, 0))
+
+
+def parallel_source(n: int, wg: int, grid_x: int, fixed_block: int | None = None) -> tuple[str, int, int]:
+  """One invocation per final quadrant, avoiding the serial kernel's redundant product reads."""
+  leaf, paths = n//4, 49
+  width, height = grid_x*leaf, ((paths+grid_x-1)//grid_x)*leaf
+  cases = []
+  for block_r in range(4):
+    for block_c in range(4):
+      u = ((block_r & 1) << 1) | (block_c & 1)
+      v = ((block_r >> 1) << 1) | (block_c >> 1)
+      terms = []
+      for p0 in range(7):
+        for p1 in range(7):
+          coeff = COMBINE[v][p0] * COMBINE[u][p1]
+          if not coeff: continue
+          p = p0*7+p1
+          value = f"M[((({p}/{grid_x})*{leaf}+r)*{width//4})+({p}%{grid_x})*{leaf//4}+x4]"
+          terms.append(("+" if coeff > 0 else "-")+value)
+      expr = "".join(terms).lstrip("+")
+      block = block_r*4+block_c
+      if fixed_block is None: cases.append(f"{'if' if block == 0 else 'else if'}(q=={block})v={expr};")
+      elif block == fixed_block: cases.append(f"v={expr};")
+  q_init = f"q={fixed_block},j=s" if fixed_block is not None else f"unit={leaf*leaf//4},q=s/unit,j=s%unit"
+  return f"""__attribute__((reqd_work_group_size({wg},1,1)))
+__kernel void combine(__global float4 *C,__global const float4 *M) {{
+  uint s=get_global_id(0),{q_init},r=j/{leaf//4},x4=j%{leaf//4};float4 v;
+  {''.join(cases)}
+  C[((q/4*{leaf}+r)*{n//4})+(q%4)*{leaf//4}+x4]=v;
+}}""", width, height
+
+
+def reuse_source(n: int, wg: int, grid_x: int) -> tuple[str, int, int]:
+  """Produce eight quadrants per pass and share each product load across four inner combines."""
+  leaf, paths = n//4, 49
+  width, height = grid_x*leaf, ((paths+grid_x-1)//grid_x)*leaf
+  body = []
+  for vs in ((0, 1), (2, 3)):
+    body.append("{\n")
+    for v in vs:
+      for u in range(4): body.append(f"float4 o{v}_{u}=(float4)(0);\n")
+    used_p0 = [p0 for p0 in range(7) if any(COMBINE[v][p0] for v in vs)]
+    for p0 in used_p0:
+      names = []
+      for p1 in range(7):
+        p = p0*7+p1
+        name = f"x{p0}_{p1}"
+        names.append(name)
+        body.append(f"float4 {name}=M[((({p}/{grid_x})*{leaf}+r)*{width//4})+({p}%{grid_x})*{leaf//4}+x4];\n")
+      for u in range(4): body.append(f"float4 d{p0}_{u}={combine_expr(names, u)};\n")
+      for v in vs:
+        coeff = COMBINE[v][p0]
+        if coeff:
+          op = "+=" if coeff > 0 else "-="
+          for u in range(4): body.append(f"o{v}_{u}{op}d{p0}_{u};\n")
+    for v in vs:
+      for u in range(4):
+        block_r, block_c = (v >> 1)*2+(u >> 1), (v & 1)*2+(u & 1)
+        body.append(f"C[({block_r}*{leaf}+r)*{n//4}+{block_c}*{leaf//4}+x4]=o{v}_{u};\n")
+    body.append("}\n")
+  return f"""__attribute__((reqd_work_group_size({wg},1,1)))
+__kernel void combine(__global float4 *C,__global const float4 *M) {{
+  uint s=get_global_id(0),r=s/{leaf//4},x4=s%{leaf//4};
+  {''.join(body)}
+}}""", width, height
+
+
+def vmajor_source(n: int, wg: int, grid_x: int) -> tuple[str, int, int]:
+  """Accumulate four inner quadrants for one outer quadrant, sharing its product loads."""
+  leaf, paths = n//4, 49
+  width, height = grid_x*leaf, ((paths+grid_x-1)//grid_x)*leaf
+  body = []
+  for v in range(4):
+    body.append("{\n")
+    for u in range(4): body.append(f"float4 o{u}=(float4)(0);\n")
+    for p0, coeff in enumerate(COMBINE[v]):
+      if not coeff: continue
+      body.append("{\n")
+      names = []
+      for p1 in range(7):
+        p = p0*7+p1
+        name = f"z{p1}"
+        names.append(name)
+        body.append(f"float4 {name}=M[((({p}/{grid_x})*{leaf}+r)*{width//4})+({p}%{grid_x})*{leaf//4}+x4];\n")
+      op = "+=" if coeff > 0 else "-="
+      for u in range(4): body.append(f"o{u}{op}{combine_expr(names, u)};\n")
+      body.append("}\n")
+    for u in range(4):
+      block_r, block_c = (v >> 1)*2+(u >> 1), (v & 1)*2+(u & 1)
+      body.append(f"C[({block_r}*{leaf}+r)*{n//4}+{block_c}*{leaf//4}+x4]=o{u};\n")
+    body.append("}\n")
+  source_body = "".join(body)
+  return f"""__attribute__((reqd_work_group_size({wg},1,1)))
+__kernel void combine(__global float4 *C,__global const float4 *M) {{
+  uint s=get_global_id(0),r=s/{leaf//4},x4=s%{leaf//4};
+  {source_body}
+}}""", width, height
+
+
+def source(n: int, wg: int, grid_x: int) -> tuple[str, int, int]:
+  leaf, paths = n//4, 49
+  width, height = grid_x*leaf, ((paths+grid_x-1)//grid_x)*leaf
+  body = []
+  for u in range(4):
+    for p0 in range(7):
+      names = [f"M[((({p0*7+p1}/{grid_x})*{leaf}+r)*{width//4})+({p0*7+p1}%{grid_x})*{leaf//4}+x4]" for p1 in range(7)]
+      body.append(f"float4 d{u}_{p0}={combine_expr(names, u)};")
+    ds = [f"d{u}_{x}" for x in range(7)]
+    for v in range(4):
+      block_r, block_c = (v>>1)*2+(u>>1), (v&1)*2+(u&1)
+      body.append(f"C[({block_r}*{leaf}+r)*{n//4}+{block_c}*{leaf//4}+x4]={combine_expr(ds, v)};")
+  return f"""__attribute__((reqd_work_group_size({wg},1,1)))
+__kernel void combine(__global float4 *C,__global const float4 *M) {{
+  uint s=get_global_id(0),r=s/{leaf//4},x4=s%{leaf//4};
+  {''.join(body)}
+}}""", width, height
+
+
+def image_source(n: int, wg: int, grid_x: int) -> tuple[str, int, int]:
+  leaf, paths = n//4, 49
+  width, height = grid_x*leaf, ((paths+grid_x-1)//grid_x)*leaf
+  body = []
+  for u in range(4):
+    for p0 in range(7):
+      names = [f"read_imagef(M,smp,(int2)({(p0*7+p1)%grid_x}*{leaf//4}+x4,{(p0*7+p1)//grid_x}*{leaf}+r))" for p1 in range(7)]
+      body.append(f"float4 d{u}_{p0}={combine_expr(names, u)};")
+    ds = [f"d{u}_{x}" for x in range(7)]
+    for v in range(4):
+      block_r, block_c = (v>>1)*2+(u>>1), (v&1)*2+(u&1)
+      body.append(f"write_imagef(C,(int2)({block_c}*{leaf//4}+x4,{block_r}*{leaf}+r),{combine_expr(ds, v)});")
+  return f"""const sampler_t smp=CLK_NORMALIZED_COORDS_FALSE|CLK_ADDRESS_CLAMP|CLK_FILTER_NEAREST;
+__attribute__((reqd_work_group_size({wg},1,1)))
+__kernel void combine(write_only image2d_t C,read_only image2d_t M) {{
+  uint s=get_global_id(0),r=s/{leaf//4},x4=s%{leaf//4};
+  {''.join(body)}
+}}""", width, height
+
+
+def combine_once(m: np.ndarray) -> np.ndarray:
+  h = m.shape[1]
+  out = np.empty((m.shape[0]//7, h*2, h*2), np.float32)
+  for b in range(out.shape[0]):
+    x = m[b*7:(b+1)*7]
+    out[b, :h, :h] = x[0]+x[3]-x[4]+x[6]
+    out[b, :h, h:] = x[2]+x[4]
+    out[b, h:, :h] = x[1]+x[3]
+    out[b, h:, h:] = x[0]-x[1]+x[2]+x[5]
+  return out
+
+
+def main() -> None:
+  n, wg, grid_x = int(os.getenv("N", "1024")), int(os.getenv("WG", "128")), int(os.getenv("GRID_X", "32"))
+  leaf, paths = n//4, 49
+  rng = np.random.default_rng(int(os.getenv("SEED", "701")))
+  mats = rng.normal(0, 1/32, (paths, leaf, leaf)).astype(np.float32)
+  image = bool(int(os.getenv("IMAGE_COMBINE", "0")))
+  parallel = bool(int(os.getenv("PARALLEL", "0")))
+  split = bool(int(os.getenv("PARALLEL_SPLIT", "0")))
+  reuse = bool(int(os.getenv("REUSE", "0")))
+  vmajor = bool(int(os.getenv("VMAJOR", "0")))
+  src, width, height = (image_source(n, wg, grid_x) if image else
+                        vmajor_source(n, wg, grid_x) if vmajor else
+                        reuse_source(n, wg, grid_x) if reuse else
+                        parallel_source(n, wg, grid_x) if parallel else source(n, wg, grid_x))
+  storage = np.zeros((height, width), np.float32)
+  for p in range(paths):
+    storage[(p//grid_x)*leaf:(p//grid_x+1)*leaf, (p%grid_x)*leaf:(p%grid_x+1)*leaf] = mats[p]
+  dev = Device["QCOM"]
+  mb, cb = alloc(storage.size, dtypes.float), alloc(n*n, dtypes.float)
+  mb.copyin(memoryview(storage).cast("B"))
+  specs = ([((0, dtypes.float, (n, n//4, 4)),), ((0, dtypes.float, (height, width//4, 4)),)] if image else
+           [((0, dtypes.float, (n*n,)),), ((0, dtypes.float, (storage.size,)),)])
+  prgs = ([dev.runtime("combine", dev.compiler.compile(parallel_source(n, wg, grid_x, block)[0]), buf_dtypes=specs)
+           for block in range(16)] if split else [dev.runtime("combine", dev.compiler.compile(src), buf_dtypes=specs)])
+  groups = leaf*leaf//4//wg*(16 if parallel else 1)
+  if split: groups //= 16
+  times = [sum(prg(cb._buf, mb._buf, global_size=(groups, 1, 1), local_size=(wg, 1, 1), wait=True) for prg in prgs) for _ in range(5)]
+  got = np.empty((n, n), np.float32)
+  cb.copyout(memoryview(got).cast("B"))
+  expected = combine_once(combine_once(mats))
+  delta = np.abs(got-expected)
+  print(f"n={n} best_ms={min(times)*1e3:.3f} max_abs={delta.max():.9g} mean_abs={delta.mean():.9g} "
+        f"exact={np.array_equal(got, expected)}")
+
+
+if __name__ == "__main__": main()

extra/gemm/qcom_strassen_fused_transform_check.py

@@ -0,0 +1,209 @@
+#!/usr/bin/env python3
+"""Fuse every recursive Strassen operand transform into one local-memory kernel."""
+import os
+
+import numpy as np
+
+from tinygrad import Device, dtypes
+from tinygrad.device import Buffer
+
+
+def alloc(count: int, dtype) -> Buffer: return Buffer("QCOM", count, dtype).allocate()
+
+
+TA = ((1, 0, 0, 1), (0, 0, 1, 1), (1, 0, 0, 0), (0, 0, 0, 1),
+      (1, 1, 0, 0), (-1, 0, 1, 0), (0, 1, 0, -1))
+TB = ((1, 0, 0, 1), (1, 0, 0, 0), (0, 1, 0, -1), (-1, 0, 1, 0),
+      (0, 0, 0, 1), (1, 1, 0, 0), (0, 0, 1, 1))
+
+
+def direct2_source(n: int, side: str, wg: int, grid_x: int) -> tuple[str, int, int]:
+  leaf, paths = n//4, 49
+  width, height = grid_x*leaf, ((paths+grid_x-1)//grid_x)*leaf
+  t = TA if side == "A" else TB
+  loads = []
+  for q0 in range(4):
+    for q1 in range(4):
+      rb = ((q0>>1)<<1)|(q1>>1)
+      cb = ((q0&1)<<1)|(q1&1)
+      loads.append(f"float4 v{q0*4+q1}=convert_float4(vload4(0,I+({rb}*{leaf}+r)*{n}+{cb}*{leaf}+x4*4));")
+  stores = []
+  for p0 in range(7):
+    for p1 in range(7):
+      terms = []
+      for q0 in range(4):
+        for q1 in range(4):
+          c = t[p0][q0]*t[p1][q1]
+          if c: terms.append(("+" if c > 0 else "-")+f"v{q0*4+q1}")
+      expr = "".join(terms).lstrip("+")
+      p = p0*7+p1
+      stores.append(f"O[((({p}/{grid_x})*{leaf}+r)*{width//4})+({p}%{grid_x})*{leaf//4}+x4]={expr};")
+  return f"""#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+__attribute__((reqd_work_group_size({wg},1,1)))
+__kernel void transform(__global float4 *O,__global const half *I) {{
+  uint s=get_global_id(0),r=s/{leaf//4},x4=s%{leaf//4};
+  {''.join(loads)}
+  {''.join(stores)}
+}}""", width, height
+
+
+def direct3_source(n: int, side: str, wg: int, vec: int = 2) -> tuple[str, int, int]:
+  """Fuse three transforms, processing one innermost path at a time to bound registers."""
+  leaf, paths = n//8, 343
+  t = TA if side == "A" else TB
+  body = []
+  for p2 in range(7):
+    body.append("{")
+    vals = []
+    for q0 in range(4):
+      for q1 in range(4):
+        terms = []
+        for q2 in range(4):
+          coeff = t[p2][q2]
+          if not coeff: continue
+          rb = ((q0 >> 1) << 2) | ((q1 >> 1) << 1) | (q2 >> 1)
+          cb = ((q0 & 1) << 2) | ((q1 & 1) << 1) | (q2 & 1)
+          load = f"convert_float{vec}(vload{vec}(0,I+({rb}*{leaf}+r)*{n}+{cb}*{leaf}+xv*{vec}))"
+          terms.append(("+" if coeff > 0 else "-")+load)
+        name = f"v{q0*4+q1}"
+        vals.append(name)
+        body.append(f"float{vec} {name}={''.join(terms).lstrip('+')};")
+    for p0 in range(7):
+      for p1 in range(7):
+        terms = []
+        for q0 in range(4):
+          for q1 in range(4):
+            coeff = t[p0][q0]*t[p1][q1]
+            if coeff: terms.append(("+" if coeff > 0 else "-")+vals[q0*4+q1])
+        p = (p0*7+p1)*7+p2
+        body.append(f"vstore{vec}(convert_half{vec}({''.join(terms).lstrip('+')}),0,O+{p*leaf*leaf}+r*{leaf}+xv*{vec});")
+    body.append("}")
+  return f"""#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+__attribute__((reqd_work_group_size({wg},1,1)))
+__kernel void transform(__global half *O,__global const half *I) {{
+  uint s=get_global_id(0),r=s/{leaf//vec},xv=s%{leaf//vec};
+  {''.join(body)}
+}}""", leaf, paths*leaf
+
+
+def direct3_full_source(n: int, side: str, wg: int, vec: int = 2) -> tuple[str, int, int]:
+  """Fuse three transforms while retaining all 64 source subtiles in registers."""
+  leaf, paths = n//8, 343
+  t = TA if side == "A" else TB
+  loads, stores = [], []
+  for q0 in range(4):
+    for q1 in range(4):
+      for q2 in range(4):
+        q = (q0*4+q1)*4+q2
+        rb = ((q0 >> 1) << 2) | ((q1 >> 1) << 1) | (q2 >> 1)
+        cb = ((q0 & 1) << 2) | ((q1 & 1) << 1) | (q2 & 1)
+        loads.append(f"float{vec} v{q}=convert_float{vec}(vload{vec}(0,I+({rb}*{leaf}+r)*{n}+{cb}*{leaf}+xv*{vec}));")
+  for p0 in range(7):
+    for p1 in range(7):
+      for p2 in range(7):
+        terms = []
+        for q0 in range(4):
+          for q1 in range(4):
+            for q2 in range(4):
+              coeff = t[p0][q0]*t[p1][q1]*t[p2][q2]
+              q = (q0*4+q1)*4+q2
+              if coeff: terms.append(("+" if coeff > 0 else "-")+f"v{q}")
+        p = (p0*7+p1)*7+p2
+        stores.append(f"vstore{vec}(convert_half{vec}({''.join(terms).lstrip('+')}),0,O+{p*leaf*leaf}+r*{leaf}+xv*{vec});")
+  return f"""#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+__attribute__((reqd_work_group_size({wg},1,1)))
+__kernel void transform(__global half *O,__global const half *I) {{
+  uint s=get_global_id(0),r=s/{leaf//vec},xv=s%{leaf//vec};
+  {''.join(loads)}
+  {''.join(stores)}
+}}""", leaf, paths*leaf
+
+
+def source(n: int, levels: int, side: str, wg: int, grid_x: int) -> tuple[str, int, int]:
+  leaf, paths = n >> levels, 7**levels
+  width, height = grid_x*leaf, ((paths+grid_x-1)//grid_x)*leaf
+  width4 = width//4
+  max_state = paths
+  stages = []
+  for d in range(levels):
+    old_prefixes, rest = 7**d, 4**(levels-d-1)
+    out_len = old_prefixes*7*rest
+    src, dst = ("x", "y") if d % 2 == 0 else ("y", "x")
+    stages.append(f"""
+  barrier(CLK_LOCAL_MEM_FENCE);
+  for(uint i=lid;i<{out_len};i+={wg}) {{
+    uint np=i/{rest},p=np%7,op=np/7,suf=i%{rest},base=op*{4*rest}+suf;
+    float4 q0={src}[base],q1={src}[base+{rest}],q2={src}[base+{2*rest}],q3={src}[base+{3*rest}],v;
+    {('if(p==0)v=q0+q3;else if(p==1)v=q2+q3;else if(p==2)v=q0;else if(p==3)v=q3;'
+      'else if(p==4)v=q0+q1;else if(p==5)v=q2-q0;else v=q1-q3;' if side == 'A' else
+      'if(p==0)v=q0+q3;else if(p==1)v=q0;else if(p==2)v=q1-q3;else if(p==3)v=q2-q0;'
+      'else if(p==4)v=q3;else if(p==5)v=q0+q1;else v=q2+q3;')}
+    {dst}[i]=v;
+  }}""")
+  final = "y" if levels % 2 else "x"
+  return f"""#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+__attribute__((reqd_work_group_size({wg},1,1)))
+__kernel void transform(__global float4 *O,__global const half *I) {{
+  uint lid=get_local_id(0),s=get_group_id(0),ir=s/{leaf//4},ic4=s%{leaf//4};
+  __local float4 x[{max_state}],y[{max_state}];
+  for(uint q=lid;q<{4**levels};q+={wg}) {{
+    uint rb=0,cb=0;
+    for(uint d=0;d<{levels};d++) {{uint qd=(q>>(2*({levels}-1-d)))&3;rb|=(qd>>1)<<({levels}-1-d);cb|=(qd&1)<<({levels}-1-d);}}
+    x[q]=convert_float4(vload4(0,I+(rb*{leaf}+ir)*{n}+cb*{leaf}+ic4*4));
+  }}
+  {''.join(stages)}
+  barrier(CLK_LOCAL_MEM_FENCE);
+  for(uint p=lid;p<{paths};p+={wg}) {{
+    uint orow=(p/{grid_x})*{leaf}+ir,ocol4=(p%{grid_x})*{leaf//4}+ic4;
+    O[orow*{width4}+ocol4]={final}[p];
+  }}
+}}""", width, height
+
+
+def cpu_transform(x: np.ndarray, levels: int, side: str) -> np.ndarray:
+  mats = [x.astype(np.float32)]
+  for _ in range(levels):
+    out = []
+    for m in mats:
+      h = m.shape[0]//2
+      q0, q1, q2, q3 = m[:h, :h], m[:h, h:], m[h:, :h], m[h:, h:]
+      out.extend(([q0+q3, q2+q3, q0, q3, q0+q1, q2-q0, q1-q3] if side == "A" else
+                  [q0+q3, q0, q1-q3, q2-q0, q3, q0+q1, q2+q3]))
+    mats = out
+  return np.stack(mats)
+
+
+def main() -> None:
+  n, levels = int(os.getenv("N", "1024")), int(os.getenv("LEVELS", "2"))
+  wg, grid_x, seed = int(os.getenv("WG", "128")), int(os.getenv("GRID_X", "128")), int(os.getenv("SEED", "701"))
+  leaf, paths = n >> levels, 7**levels
+  rng = np.random.default_rng(seed)
+  inp = rng.normal(0, 1/32, (n, n)).astype(np.float16)
+  dev = Device["QCOM"]
+  ib = alloc(n*n, dtypes.half)
+  ib.copyin(memoryview(inp).cast("B"))
+  direct2 = levels == 2 and bool(int(os.getenv("DIRECT2", "1")))
+  direct3 = levels == 3 and bool(int(os.getenv("DIRECT3", "0")))
+  direct3_full = direct3 and bool(int(os.getenv("DIRECT3_FULL", "0")))
+  for side in "AB":
+    src, width, height = (direct2_source(n, side, wg, grid_x) if direct2 else
+                          direct3_full_source(n, side, wg) if direct3_full else
+                          direct3_source(n, side, wg) if direct3 else source(n, levels, side, wg, grid_x))
+    odt = dtypes.half if direct3 else dtypes.float
+    ob = alloc(width*height, odt)
+    prg = dev.runtime("transform", dev.compiler.compile(src), buf_dtypes=[
+      ((0, odt, (width*height,)),), ((0, dtypes.half, (n*n,)),)])
+    groups = leaf*leaf//(2 if direct3 else 4)//wg if direct2 or direct3 else leaf*leaf//4
+    times = [prg(ob._buf, ib._buf, global_size=(groups, 1, 1), local_size=(wg, 1, 1), wait=True) for _ in range(3)]
+    got_storage = np.empty((height, width), np.float16 if direct3 else np.float32)
+    ob.copyout(memoryview(got_storage).cast("B"))
+    got = (got_storage.reshape(paths, leaf, leaf) if direct3 else
+           np.stack([got_storage[(p//grid_x)*leaf:(p//grid_x+1)*leaf,
+                                 (p%grid_x)*leaf:(p%grid_x+1)*leaf] for p in range(paths)]))
+    expected = cpu_transform(inp, levels, side).astype(np.float16 if direct3 else np.float32)
+    delta = np.abs(got-expected)
+    print(f"side={side} n={n} levels={levels} paths={paths} layout={height}x{width} best_ms={min(times)*1e3:.3f} "
+          f"max_abs={delta.max():.9g} mean_abs={delta.mean():.9g} exact={np.array_equal(got, expected)}")
+
+
+if __name__ == "__main__": main()

extra/gemm/qcom_strassen_leaf_fused_check.py

@@ -0,0 +1,108 @@
+#!/usr/bin/env python3
+"""One-level full Strassen GEMM with operand transforms fused into FP32 leaves."""
+import os
+
+import numpy as np
+
+from tinygrad import Device, dtypes
+from tinygrad.device import Buffer
+
+
+def alloc(count: int, dtype) -> Buffer: return Buffer("QCOM", count, dtype).allocate()
+
+
+def leaf_src(p: int, n: int = 512) -> str:
+  h, h4 = n//2, n//8
+  ae = (("a0+a3"), ("a2+a3"), "a0", "a3", "a0+a1", "a2-a0", "a1-a3")[p]
+  be = (("b0+b3"), "b0", "b1-b3", "b2-b0", "b3", "b0+b1", "b2+b3")[p]
+  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 leaf(write_only image2d_t C,read_only image2d_t A,read_only image2d_t B) {{
+  int lid=get_local_id(0),tm=lid>>5,tid=lid&31;
+  int row=get_group_id(1)*16+tm*4,col=get_group_id(0)*32+tid;
+  float4 r0=(float4)(0),r1=(float4)(0),r2=(float4)(0),r3=(float4)(0);
+  for(int k4=0;k4<{h4};k4++) {{
+    float4 a0=convert_float4(read_imageh(A,smp,(int2)(k4,row+0)));
+    float4 a1=convert_float4(read_imageh(A,smp,(int2)(k4+{h4},row+0)));
+    float4 a2=convert_float4(read_imageh(A,smp,(int2)(k4,row+{h}+0)));
+    float4 a3=convert_float4(read_imageh(A,smp,(int2)(k4+{h4},row+{h}+0))),aa0={ae};
+    a0=convert_float4(read_imageh(A,smp,(int2)(k4,row+1)));a1=convert_float4(read_imageh(A,smp,(int2)(k4+{h4},row+1)));
+    a2=convert_float4(read_imageh(A,smp,(int2)(k4,row+{h}+1)));a3=convert_float4(read_imageh(A,smp,(int2)(k4+{h4},row+{h}+1)));float4 aa1={ae};
+    a0=convert_float4(read_imageh(A,smp,(int2)(k4,row+2)));a1=convert_float4(read_imageh(A,smp,(int2)(k4+{h4},row+2)));
+    a2=convert_float4(read_imageh(A,smp,(int2)(k4,row+{h}+2)));a3=convert_float4(read_imageh(A,smp,(int2)(k4+{h4},row+{h}+2)));float4 aa2={ae};
+    a0=convert_float4(read_imageh(A,smp,(int2)(k4,row+3)));a1=convert_float4(read_imageh(A,smp,(int2)(k4+{h4},row+3)));
+    a2=convert_float4(read_imageh(A,smp,(int2)(k4,row+{h}+3)));a3=convert_float4(read_imageh(A,smp,(int2)(k4+{h4},row+{h}+3)));float4 aa3={ae};
+    float4 b0=convert_float4(read_imageh(B,smp,(int2)(col,k4*4+0)));
+    float4 b1=convert_float4(read_imageh(B,smp,(int2)(col+{h4},k4*4+0)));
+    float4 b2=convert_float4(read_imageh(B,smp,(int2)(col,k4*4+{h})));
+    float4 b3=convert_float4(read_imageh(B,smp,(int2)(col+{h4},k4*4+{h})));
+    float4 bb0={be};
+    b0=convert_float4(read_imageh(B,smp,(int2)(col,k4*4+1)));
+    b1=convert_float4(read_imageh(B,smp,(int2)(col+{h4},k4*4+1)));
+    b2=convert_float4(read_imageh(B,smp,(int2)(col,k4*4+{h}+1)));
+    b3=convert_float4(read_imageh(B,smp,(int2)(col+{h4},k4*4+{h}+1)));
+    float4 bb1={be};
+    b0=convert_float4(read_imageh(B,smp,(int2)(col,k4*4+2)));
+    b1=convert_float4(read_imageh(B,smp,(int2)(col+{h4},k4*4+2)));
+    b2=convert_float4(read_imageh(B,smp,(int2)(col,k4*4+{h}+2)));
+    b3=convert_float4(read_imageh(B,smp,(int2)(col+{h4},k4*4+{h}+2)));
+    float4 bb2={be};
+    b0=convert_float4(read_imageh(B,smp,(int2)(col,k4*4+3)));
+    b1=convert_float4(read_imageh(B,smp,(int2)(col+{h4},k4*4+3)));
+    b2=convert_float4(read_imageh(B,smp,(int2)(col,k4*4+{h}+3)));
+    b3=convert_float4(read_imageh(B,smp,(int2)(col+{h4},k4*4+{h}+3)));
+    float4 bb3={be};
+    r0+=aa0.x*bb0+aa0.y*bb1+aa0.z*bb2+aa0.w*bb3;
+    r1+=aa1.x*bb0+aa1.y*bb1+aa1.z*bb2+aa1.w*bb3;
+    r2+=aa2.x*bb0+aa2.y*bb1+aa2.z*bb2+aa2.w*bb3;
+    r3+=aa3.x*bb0+aa3.y*bb1+aa3.z*bb2+aa3.w*bb3;
+  }}
+  write_imagef(C,(int2)(col,row+0),r0);write_imagef(C,(int2)(col,row+1),r1);
+  write_imagef(C,(int2)(col,row+2),r2);write_imagef(C,(int2)(col,row+3),r3);
+}}"""
+
+
+def combine_src(n: int = 512) -> str:
+  h = n//2
+  rd = lambda p: f"read_imagef(M,smp,(int2)(x,{p*h}+r))"  # noqa: E731
+  return f"""const sampler_t smp=CLK_NORMALIZED_COORDS_FALSE|CLK_ADDRESS_CLAMP|CLK_FILTER_NEAREST;
+__attribute__((reqd_work_group_size(256,1,1)))
+__kernel void combine(write_only image2d_t C,read_only image2d_t M) {{
+  uint i=get_global_id(0),q=i/{h*h//4},j=i%{h*h//4},r=j/{h//4},x=j%{h//4};float4 v;
+  if(q==0)v={rd(0)}+{rd(3)}-{rd(4)}+{rd(6)};
+  else if(q==1)v={rd(2)}+{rd(4)};else if(q==2)v={rd(1)}+{rd(3)};
+  else v={rd(0)}-{rd(1)}+{rd(2)}+{rd(5)};
+  write_imagef(C,(int2)(x+(q&1)*{h//4},r+(q>>1)*{h}),v);
+}}"""
+
+
+def main() -> None:
+  n, seed = 512, int(os.getenv("SEED", "701"))
+  rng = np.random.default_rng(seed)
+  a_np = rng.normal(0, 1/32, (n, n)).astype(np.float16)
+  b_np = rng.normal(0, 1/32, (n, n)).astype(np.float16)
+  dev = Device["QCOM"]
+  a, b, mm, c = alloc(n*n, dtypes.half), alloc(n*n, dtypes.half), alloc(7*n*n//4, dtypes.float), alloc(n*n, dtypes.float)
+  a.copyin(memoryview(a_np).cast("B"))
+  b.copyin(memoryview(b_np).cast("B"))
+  leaves = [dev.runtime("leaf", dev.compiler.compile(leaf_src(p)), buf_dtypes=[
+    ((0, dtypes.float, (n//2, n//8, 4)),), ((0, dtypes.half, (n, n//4, 4)),), ((1, dtypes.half, (n, n//4, 4)),)]) for p in range(7)]
+  combine = dev.runtime("combine", dev.compiler.compile(combine_src()), buf_dtypes=[
+    ((0, dtypes.float, (n, n//4, 4)),), ((0, dtypes.float, (7*n//2, n//8, 4)),)])
+  times = []
+  for p, prg in enumerate(leaves):
+    out = mm._buf.offset(p*n*n, n*n)
+    times.append(prg(out, a._buf, b._buf, global_size=(2, 16, 1), local_size=(128, 1, 1), wait=True))
+  times.append(combine(c._buf, mm._buf, global_size=(n*n//4//256, 1, 1), local_size=(256, 1, 1), wait=True))
+  got = np.empty((n, 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)
+  elapsed = sum(times)
+  print(f"elapsed_ms={elapsed*1e3:.3f} gflops={2*n**3/elapsed/1e9:.1f} leaf_ms={sum(times[:-1])*1e3:.3f} "
+        f"combine_ms={times[-1]*1e3:.3f} max_abs={delta.max():.9g} mean_abs={delta.mean():.9g} "
+        f"allclose={np.allclose(got, expected, rtol=1e-3, atol=1e-3)} leaf_parts_ms={[round(x*1e3, 3) for x in times[:-1]]}")
+
+
+if __name__ == "__main__": main()

extra/gemm/qcom_strassen_terminal_fused_check.py

@@ -0,0 +1,184 @@
+#!/usr/bin/env python3
+"""Full one-level Strassen GEMM with all transforms and combines fused into one kernel."""
+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 (ADD_F, ADD_S, ADD_S_REG, AND_B, BR, CMPS_S_EQ, END, MAD_F32, MOV_F32, MOV_S32,
+                               NOP, SHR_B, STIB_F32, SUB_F, assemble, get_envelope, inject)
+
+
+def alloc(count: int, dtype) -> Buffer: return Buffer("QCOM", count, dtype).allocate()
+
+
+def source(n: int = 512, wg: int = 128) -> str:
+  h, h4 = n//2, n//8
+  ae = ("a0+a3", "a2+a3", "a0", "a3", "a0+a1", "a2-a0", "a1-a3")
+  be = ("b0+b3", "b0", "b1-b3", "b2-b0", "b3", "b0+b1", "b2+b3")
+  contributions = (((0, 1), (3, 1)), ((2, 1), (3, -1)), ((1, 1), (3, 1)), ((0, 1), (2, 1)),
+                   ((0, -1), (1, 1)), ((3, 1),), ((0, 1),))
+  phases = []
+  for p in range(7):
+    updates = []
+    for quadrant, sign in contributions[p]:
+      op = "+=" if sign > 0 else "-="
+      updates.extend((f"c{quadrant}0{op}m0;", f"c{quadrant}1{op}m1;"))
+    phases.append(f"""
+  {{ float4 m0=(float4)(0),m1=(float4)(0);
+    for(int k4=0;k4<{h4};k4++) {{
+      float4 a0=convert_float4(read_imageh(A,smp,(int2)(k4,row+0)));
+      float4 a1=convert_float4(read_imageh(A,smp,(int2)(k4+{h4},row+0)));
+      float4 a2=convert_float4(read_imageh(A,smp,(int2)(k4,row+{h}+0)));
+      float4 a3=convert_float4(read_imageh(A,smp,(int2)(k4+{h4},row+{h}+0))),aa0={ae[p]};
+      a0=convert_float4(read_imageh(A,smp,(int2)(k4,row+1)));a1=convert_float4(read_imageh(A,smp,(int2)(k4+{h4},row+1)));
+      a2=convert_float4(read_imageh(A,smp,(int2)(k4,row+{h}+1)));a3=convert_float4(read_imageh(A,smp,(int2)(k4+{h4},row+{h}+1)));
+      float4 aa1={ae[p]};
+      float4 b0=convert_float4(read_imageh(B,smp,(int2)(col,k4*4+0)));
+      float4 b1=convert_float4(read_imageh(B,smp,(int2)(col+{h4},k4*4+0)));
+      float4 b2=convert_float4(read_imageh(B,smp,(int2)(col,k4*4+{h}+0)));
+      float4 b3=convert_float4(read_imageh(B,smp,(int2)(col+{h4},k4*4+{h}+0))),bb0={be[p]};
+      b0=convert_float4(read_imageh(B,smp,(int2)(col,k4*4+1)));b1=convert_float4(read_imageh(B,smp,(int2)(col+{h4},k4*4+1)));
+      b2=convert_float4(read_imageh(B,smp,(int2)(col,k4*4+{h}+1)));b3=convert_float4(read_imageh(B,smp,(int2)(col+{h4},k4*4+{h}+1)));
+      float4 bb1={be[p]};
+      b0=convert_float4(read_imageh(B,smp,(int2)(col,k4*4+2)));b1=convert_float4(read_imageh(B,smp,(int2)(col+{h4},k4*4+2)));
+      b2=convert_float4(read_imageh(B,smp,(int2)(col,k4*4+{h}+2)));b3=convert_float4(read_imageh(B,smp,(int2)(col+{h4},k4*4+{h}+2)));
+      float4 bb2={be[p]};
+      b0=convert_float4(read_imageh(B,smp,(int2)(col,k4*4+3)));b1=convert_float4(read_imageh(B,smp,(int2)(col+{h4},k4*4+3)));
+      b2=convert_float4(read_imageh(B,smp,(int2)(col,k4*4+{h}+3)));b3=convert_float4(read_imageh(B,smp,(int2)(col+{h4},k4*4+{h}+3)));
+      float4 bb3={be[p]};
+      m0+=aa0.x*bb0+aa0.y*bb1+aa0.z*bb2+aa0.w*bb3;
+      m1+=aa1.x*bb0+aa1.y*bb1+aa1.z*bb2+aa1.w*bb3;
+    }}
+    {''.join(updates)} }}""")
+  stores = []
+  for quadrant in range(4):
+    qr, qc = quadrant>>1, quadrant&1
+    stores.extend((f"write_imagef(C,(int2)(col+{qc*h4},row+{qr*h}+0),c{quadrant}0);",
+                   f"write_imagef(C,(int2)(col+{qc*h4},row+{qr*h}+1),c{quadrant}1);"))
+  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({wg},1,1)))
+__kernel void gemm(write_only image2d_t C,read_only image2d_t A,read_only image2d_t B) {{
+  int lid=get_local_id(0),tm=lid>>5,tid=lid&31;
+  int row=get_group_id(1)*8+tm*2,col=get_group_id(0)*32+tid;
+  float4 c00=(float4)(0),c01=(float4)(0),c10=(float4)(0),c11=(float4)(0);
+  float4 c20=(float4)(0),c21=(float4)(0),c30=(float4)(0),c31=(float4)(0);
+  {''.join(phases)}
+  {''.join(stores)}
+}}"""
+
+
+def hand_shader(dev, n: int = 512) -> tuple[bytes, int, int]:
+  h, h4 = n//2, n//8
+  instrs = q.prologue_4x2(dev, 128)
+  # Convert the donor's 4x8-within-full-output coordinates to a 2x4 tile
+  # within one quadrant: row/=2 and group-column stride/=2, preserving tid.
+  instrs += [SHR_B("r7.x", "r7.x", 1), AND_B("r6.x", "r7.y", 31),
+             AND_B("r7.y", "r7.y", 0xc0), SHR_B("r7.y", "r7.y", 1),
+             ADD_S_REG("r7.y", "r7.y", "r6.x"), NOP(rpt=2)]
+  acc0, mi0, avec0, bvec0, tmp, state, const = 0, 8, 10, 12, 16, 17, 18
+  row_base, col_base, k4, ky = (q.fvec(state, x) for x in range(4))
+  hreg, h4reg = q.fvec(const, 0), q.fvec(const, 1)
+  instrs += [MOV_F32(row_base, "r7.x"), MOV_F32(col_base, "r7.y"), MOV_S32(hreg, h), MOV_S32(h4reg, h4)]
+  for vec in range(acc0, acc0+8): q.emit_f32_vec_imm(instrs, vec, 0)
+
+  ta = (((0, 1), (3, 1)), ((2, 1), (3, 1)), ((0, 1),), ((3, 1),),
+        ((0, 1), (1, 1)), ((2, 1), (0, -1)), ((1, 1), (3, -1)))
+  tb = (((0, 1), (3, 1)), ((0, 1),), ((1, 1), (3, -1)), ((2, 1), (0, -1)),
+        ((3, 1),), ((0, 1), (1, 1)), ((2, 1), (3, 1)))
+  contributions = (((0, 1), (3, 1)), ((2, 1), (3, -1)), ((1, 1), (3, 1)), ((0, 1), (2, 1)),
+                   ((0, -1), (1, 1)), ((3, 1),), ((0, 1),))
+
+  def coords(dst: int, tex: int, quadrant: int, item: int) -> None:
+    qr, qc = quadrant>>1, quadrant&1
+    if tex == 0:
+      instrs.append(MOV_F32(q.fvec(dst, 0), k4) if not qc else ADD_S_REG(q.fvec(dst, 0), k4, h4reg))
+      instrs.append(MOV_F32(q.fvec(dst, 1), row_base) if not qr else ADD_S_REG(q.fvec(dst, 1), row_base, hreg))
+      if item: instrs.append(ADD_S(q.fvec(dst, 1), q.fvec(dst, 1), item))
+    else:
+      instrs.append(MOV_F32(q.fvec(dst, 0), col_base) if not qc else ADD_S_REG(q.fvec(dst, 0), col_base, h4reg))
+      instrs.append(MOV_F32(q.fvec(dst, 1), ky) if not qr else ADD_S_REG(q.fvec(dst, 1), ky, hreg))
+      if item != 3: instrs.append(ADD_S(q.fvec(dst, 1), q.fvec(dst, 1), item-3))
+
+  def load_expr(dst: int, tex: int, terms: tuple[tuple[int, int], ...], item: int) -> None:
+    nonlocal instrs
+    coords(dst, tex, terms[0][0], item)
+    q.emit_isam_f32_vec(instrs, dst, q.fvec(dst, 0), tex, True)
+    if len(terms) == 1: return
+    coords(tmp, tex, terms[1][0], item)
+    q.emit_isam_f32_vec(instrs, tmp, q.fvec(tmp, 0), tex, True)
+    instrs += [MOV_F32(q.fvec(dst), q.fvec(dst), sy=True), NOP(rpt=5)]
+    op = ADD_F if terms[1][1] > 0 else SUB_F
+    for comp in range(4): instrs.append(op(dst*4+comp, dst*4+comp, tmp*4+comp))
+
+  for p in range(7):
+    for vec in range(mi0, mi0+2): q.emit_f32_vec_imm(instrs, vec, 0)
+    instrs += [MOV_S32(k4, 0), MOV_S32(ky, 3)]
+    loop_start = len(instrs)
+    for kk in range(4): load_expr(bvec0+kk, 1, tb[p], kk)
+    for row in range(2): load_expr(avec0+row, 0, ta[p], row)
+    first = True
+    for kk in range(4):
+      for row in range(2):
+        instrs.append(MAD_F32(q.fvec(mi0+row), q.fvec(avec0+row, kk), q.fvec(bvec0+kk),
+                              q.fvec(mi0+row), rpt=3, sy=first, r=True))
+        first = False
+    instrs += [ADD_S("r6.x", k4, 1), ADD_S(ky, ky, 4), CMPS_S_EQ(k4, h4-1, nop=1),
+               MOV_F32(k4, "r6.x"), NOP(rpt=3)]
+    loop_end = len(instrs)
+    instrs.append(BR(loop_start-loop_end))
+    instrs += [MOV_F32(q.fvec(mi0), q.fvec(mi0), sy=True), NOP(rpt=3)]
+    for quadrant, sign in contributions[p]:
+      for row in range(2):
+        dst, src = acc0+quadrant*2+row, mi0+row
+        op = ADD_F if sign > 0 else SUB_F
+        for comp in range(4): instrs.append(op(dst*4+comp, dst*4+comp, src*4+comp))
+
+  instrs += [MOV_F32("r6.x", "r6.x", sy=True), NOP(rpt=8)]
+  for quadrant in range(4):
+    qr, qc = quadrant>>1, quadrant&1
+    for row in range(2):
+      instrs.append(MOV_F32("r4.x", col_base) if not qc else ADD_S_REG("r4.x", col_base, h4reg))
+      instrs.append(MOV_F32("r4.y", row_base) if not qr else ADD_S_REG("r4.y", row_base, hreg))
+      if row: instrs.append(ADD_S("r4.y", "r4.y", row))
+      instrs += [MOV_F32(q.fvec(acc0+quadrant*2+row), q.fvec(acc0+quadrant*2+row), sy=True), NOP(rpt=5),
+                 STIB_F32(q.fvec(acc0+quadrant*2+row), "r4.x"), NOP(rpt=16)]
+  instrs.append(END())
+  return assemble(instrs), 1, 19
+
+
+def main() -> None:
+  n, seed = 512, int(os.getenv("SEED", "701"))
+  rng = np.random.default_rng(seed)
+  a_np = rng.normal(0, 1/32, (n, n)).astype(np.float16)
+  b_np = rng.normal(0, 1/32, (n, n)).astype(np.float16)
+  dev = Device["QCOM"]
+  a, b, c = alloc(n*n, dtypes.half), alloc(n*n, dtypes.half), alloc(n*n, dtypes.float)
+  a.copyin(memoryview(a_np).cast("B"))
+  b.copyin(memoryview(b_np).cast("B"))
+  src = source()
+  if bool(int(os.getenv("HAND", "0"))):
+    q.M = q.N = q.K = n
+    q.K4 = n//4
+    env, io, sz, ro = get_envelope(dev, src)
+    shader, hregs, fregs = hand_shader(dev, n)
+    if len(shader) > sz: raise RuntimeError(f"shader {len(shader)} exceeds envelope {sz}")
+    lib = inject(env, io, sz, ro, shader, hregs=hregs, fregs=fregs, mergedregs=False)
+  else: lib = dev.compiler.compile(src)
+  prg = dev.runtime("gemm", lib, buf_dtypes=[
+    ((0, dtypes.float, (n, n//4, 4)),), ((0, dtypes.half, (n, n//4, 4)),), ((1, dtypes.half, (n, n//4, 4)),)])
+  times = [prg(c._buf, a._buf, b._buf, global_size=(2, 32, 1), local_size=(128, 1, 1), wait=True) for _ in range(5)]
+  got = np.empty((n, 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)
+  elapsed = min(times)
+  print(f"shape={n}x{n}x{n} algorithm=strassen1_fused inputs=fp16 accumulate=fp32 elapsed_ms={elapsed*1e3:.3f} "
+        f"gflops={2*n**3/elapsed/1e9:.1f} max_abs={delta.max():.9g} mean_abs={delta.mean():.9g} "
+        f"allclose={np.allclose(got, expected, rtol=1e-3, atol=1e-3)}")
+
+
+if __name__ == "__main__": main()

extra/gemm/qcom_subgroup_broadcast_probe.py

@@ -0,0 +1,81 @@
+#!/usr/bin/env python3
+"""Compile and validate A630 subgroup broadcast before using it in GEMM."""
+import os, struct
+
+import numpy as np
+
+from tinygrad import Device, dtypes
+from tinygrad.device import Buffer
+from extra.gemm.ir3asm import disasm, get_envelope, inject, QUAD_BRCST
+from extra.gemm.qcom_8x4_gemm import buf_copyin, buf_copyout
+
+
+SRC = r"""#pragma OPENCL EXTENSION cl_qcom_subgroup_shuffle : enable
+__attribute__((reqd_work_group_size(128,1,1)))
+__kernel void broadcast(__global uint *O, __global uint *X) {
+  uint i=get_global_id(0);
+  uint v=X[i];
+  O[i]=qcom_sub_group_shuffle_xor(v,1,CLK_SUB_GROUP_SHUFFLE_WIDTH_WAVE_SIZE_QCOM,999u);
+}"""
+
+
+def main() -> None:
+  dev=Device["QCOM"]
+  src = SRC
+  if (api := os.getenv("API", "xor")) != "xor":
+    call = f"qcom_sub_group_shuffle_{api}(v,{os.getenv('SELECTOR', 'i&31')},CLK_SUB_GROUP_SHUFFLE_WIDTH_WAVE_SIZE_QCOM,999u)"
+    src = src.replace("qcom_sub_group_shuffle_xor(v,1,CLK_SUB_GROUP_SHUFFLE_WIDTH_WAVE_SIZE_QCOM,999u)", call)
+  lib0, off, size, reg_off = get_envelope(dev, src)
+  lib=bytearray(lib0)
+  if int(os.getenv("QBC", "0")):
+    # Preserve the compiler's proven load/store and dependency schedule, but
+    # replace its shuffle source setup and operation with quad broadcast lane 0.
+    # The compiler put 999 in r0.w; QBC uses its low two bits, selecting lane 3.
+    if int(os.getenv("QBC_ZERO", "0")):
+      from extra.gemm.ir3asm import MOV_F32, MOV_S32, NOP
+      lib[off+11*8:off+12*8] = MOV_S32('r0.w', int(os.getenv("QBC_SELECTOR", "0")))
+      lib[off+12*8:off+13*8] = NOP(rpt=5)
+      lib[off+13*8:off+14*8] = QUAD_BRCST('r0.y', 'r0.z', 'r0.w', typ=3, sy=True)
+      lib[off+14*8:off+15*8] = MOV_F32('r0.w', 'r0.y')
+    elif int(os.getenv("QBC_NO_WAW", "0")):
+      from extra.gemm.ir3asm import NOP
+      lib[off+11*8:off+12*8] = NOP()
+      lib[off+13*8:off+14*8] = QUAD_BRCST('r0.w', 'r0.z', 'r0.y', typ=3, sy=True)
+    elif int(os.getenv("QBC_DST1", "0")):
+      from extra.gemm.ir3asm import MOV_F32, NOP
+      lib[off+11*8:off+12*8] = QUAD_BRCST('r1.x', 'r0.z', 'r0.y', typ=3, sy=True)
+      lib[off+12*8:off+13*8] = NOP(rpt=5)
+      lib[off+13*8:off+14*8] = MOV_F32('r0.w', 'r1.x')
+    elif int(os.getenv("QBC_KNOWN", "0")):
+      lib[off+13*8:off+14*8] = bytes.fromhex("010400000731e0b7")
+      # Preserve the store's r0.w source while avoiding an overlapping qbc destination.
+      lib[off+14*8:off+15*8] = bytes.fromhex("0700000003c00c20")
+    elif int(os.getenv("QBC_CONST_IDX", "0")):
+      # r0.y is the compiler's zero-valued high-address carry here.
+      lib[off+13*8:off+14*8] = QUAD_BRCST('r0.w', 'r0.z', 'r0.y', typ=3, sy=True)
+      if int(os.getenv("QBC_WAIT", "0")):
+        from extra.gemm.ir3asm import NOP
+        lib[off+14*8:off+15*8] = NOP(rpt=int(os.getenv("QBC_WAIT")))
+    else:
+      lib[off+13*8:off+14*8] = (QUAD_BRCST('r0.w', 'r0.w', 'r0.z', typ=3, sy=True) if int(os.getenv("QBC_SWAP", "0")) else
+                                    QUAD_BRCST('r0.w', 'r0.z', 'r0.w', typ=3, sy=True))
+  if int(os.getenv("MERGED0", "0")):
+    lib = bytearray(inject(lib, off, size, reg_off, bytes(lib[off:off+size]), fregs=2, hregs=0, mergedregs=False))
+  print(disasm(lib[off:off+size]))
+  n=256
+  x=np.arange(n,dtype=np.uint32)
+  xb,ob=Buffer("QCOM",n,dtypes.uint).allocate(),Buffer("QCOM",n,dtypes.uint).allocate()
+  buf_copyin(xb, memoryview(x).cast("B"))
+  prg=dev.runtime("broadcast",lib,buf_dtypes=[((0,dtypes.uint,None),),((1,dtypes.uint,None),)])
+  times=[prg(ob._buf,xb._buf,global_size=(2,1,1),local_size=(128,1,1),wait=True) for _ in range(10)]
+  got=np.empty(n,np.uint32); buf_copyout(ob, memoryview(got).cast("B"))
+  print("min_us",min(times)*1e6)
+  print("blocks",[np.unique(got[i:i+32]).tolist() for i in range(0,n,32)])
+  expected = ((np.arange(n, dtype=np.uint32)//4)*4+3 if int(os.getenv("QBC", "0")) else
+              (np.arange(n, dtype=np.uint32)//32)*32 if api == "up" else np.arange(n, dtype=np.uint32)^1)
+  print("exact", bool(np.array_equal(got, expected)), "head", got[:32].tolist())
+  if int(os.getenv("MAP", "0")):
+    print("map16", [np.unique(got[i:i+16]).tolist() for i in range(0, n, 16)])
+
+
+if __name__ == "__main__": main()