Repository navigation
nir/divergence: treat cmat_muladd_agx as divergent - #43
joshuaswarren wants to merge 2 commits into
Conversation
The flash attention shader of llama.cpp accumulates the P times V product in a half precision matrix, so its cmat_muladd has half A and B and a half accumulator. lower_matrix_instr asserts that the accumulator is float or bfloat, so compiling that shader aborts on a build with assertions. The lowering already widens the accumulator to 32 bits for the hardware multiply-add and converts the result back to the result type, so a half accumulator needs no new code: allow it in the assertion. Add a unit test that lowers a 16x16x16 half multiply-add with a float and with a half accumulator and checks that no cmat_muladd is left and that the hardware multiply-add takes a 32-bit accumulator. The half case aborts in the assertion without the change. Signed-off-by: Joshua Warren <816217+joshuaswarren@users.noreply.github.com>
agx_nir_lower_cmat turns every cooperative matrix multiply-add into cmat_muladd_agx, and the divergence analysis run on the lowered shader has no case for it. With assertions enabled the first shader that uses cooperative matrices dies in visit_intrinsic with "Unhandled intrinsic", for example every matrix multiply pipeline of llama.cpp. Each lane holds its own part of the result, which is how the other cooperative matrix multiply-add intrinsics are classified, so list it with them. The CmatLower tests now run the divergence analysis after the lowering and expect the multiply-add result to be divergent. They abort on the unhandled intrinsic without this change. Signed-off-by: Joshua Warren <816217+joshuaswarren@users.noreply.github.com>
 [Medium impact] The PR affects AGX cooperative-matrix shader compilation and its shared divergence analysis. The PR appears safe to merge, with a non-blocking cleanup issue in the new test helper. Findings
SummaryAdds
Reviews (1) · Last reviewed commit: "nir/divergence: treat cmat_muladd_agx as..." · Reviewed by Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| ralloc_free(b.shader); | ||
| return left == 0 ? agx : 0; |
There was a problem hiding this comment.
lowered_muladd_count calls glsl_type_singleton_init_or_ref() on every run but never releases that reference. Freeing the shader does not release the shared type cache, so these tests keep it alive for the rest of the test run. Add glsl_type_singleton_decref() after ralloc_free(b.shader), or use a fixture that pairs setup and cleanup.
| ralloc_free(b.shader); | |
| return left == 0 ? agx : 0; | |
| ralloc_free(b.shader); | |
| glsl_type_singleton_decref(); | |
| return left == 0 ? agx : 0; |
maralcbr
left a comment
maralcbr
left a comment
There was a problem hiding this comment.
Thanks @joshuaswarren, this is the right fix. cmat_muladd_agx now sits with the AMD, Panfrost and NVIDIA multiply-adds in the always-divergent list (nir_divergence_analysis.c:1097). That matches its SUBGROUP flag, and each lane's result depends on the other lanes' fragments. NDEBUG builds already fell back to divergent at :1161-1164. Assert builds, including CI and any assert build on an M5 where cooperative matrices are on by default (hk_physical_device.c:1953-1955), now compile these shaders instead of aborting. Only the AGX lowering emits this intrinsic, so other GPU drivers aren't affected. No blockers in the code.
One thing before merge:
- CI isn't complete yet. The aarch64 build passed, and it builds with
-Db_ndebug=falseand runsagx_tests(build.sh:15,:34), so the new divergence check inCmatLowerreally ran with assertions on. The x86_64 build was still pending when I looked, and there's noci / resultyet.
The diff still carries a8b0011 from #41, as the body says, so landing #41 first keeps this one to be59f0f. Adding "stacked on #41" to the title would make that order obvious, but it's optional.
Nits, fine as follow-ups:
test-cmat-accumulator.cpp:77: please addglsl_type_singleton_decref()afterralloc_free, as the other asahi tests do. That code comes from #41's commit, so it may be simplest to fix it there.
For release, not merge: the evidence so far comes from drm-shim on x86 with #40 and #41 applied. A llama.cpp MUL_MAT / FLASH_ATTN_EXT run on an M5 at the final head is still needed on edge, with the exact commits recorded (AGENTS.md).
Reviewed at be59f0fc4acb.
— Marcelo's agent
| } | ||
| } | ||
| } | ||
| ralloc_free(b.shader); |
There was a problem hiding this comment.
nit: glsl_type_singleton_init_or_ref() at line 41 has no matching glsl_type_singleton_decref(). Please add one after ralloc_free(b.shader), as the other asahi tests do.
— Marcelo's agent
agx_nir_lower_cmatturns every cooperative matrix multiply-add intocmat_muladd_agx, andnir_divergence_analysishas no case for that intrinsic. With assertions enabled, the first shader that uses cooperative matrices aborts invisit_intrinsicwith "NIR divergence analysis: Unhandled intrinsic". Release builds compile the assert out and fall through to the default, so this only shows on an assert-enabled Mesa.It was found offline, without the M5: Mesa at the M5 prefix with #40 and #41 applied, built for x86 with the AGX Vulkan driver and the noop DRM shim set to the M5 Max (G17C B1) identity, then llama.cpp's own
test-backend-opsforMUL_MATandFLASH_ATTN_EXTthrough its Vulkan backend. Both stop in this assertion on the first matrix multiply pipeline, so it is separate from #40 and #41 and sits behind both.The fix lists
cmat_muladd_agxwith the other cooperative matrix multiply-add intrinsics (cmat_muladd_amd,cmat_muladd_pan,cmat_muladd_nv), whose results are divergent because each lane holds its own part of the result.Test: the
CmatLowertests from #41 now runnir_divergence_analysisafter the lowering and expect the multiply-add result to be divergent. On an assert-enabled build (b_ndebug=false) they abort on the unhandled intrinsic without this change and pass with it. This branch is based on #41 because the test file is added there; the change itself is one line innir_divergence_analysis.c.