Repository navigation
HMMA with an f16 accumulator outside compute shaders matches the tensor core exactly (was 34 % off) (chain change) - #2067
Open
Know-It-All-Brother wants to merge 5 commits into
Conversation
…re_orin.py runs about 190 operations, each the GPU compiler's own machine code for one PTX operation, through the translator, spirv-cross, the chain's MSL rewrites and Metal on this Mac's GPU, and compares every output word with what a Jetson Orin (sm_87) computed from the same 1,024 input rows. The translator tests checked the MSL a few encodings give, not the values; this one finds the FSETP.NE, SHFL lane, F2FP, rounding-mode and signed-zero differences the tracking issue lists, which KNOWN names until each fix takes its entry out. dm/gfx/tools/ampere_orin/capture.py records the fixture on any Linux machine with an Ampere GPU and NVIDIA's driver, with no toolkit or compiler on either side
…ules with its functions, so a rewrite that uses one (struct in PR seedds#1902's hfma_msl) runs as it does in a build instead of raising NameError
…core does: the translator added the 16 products to C with a chain of FMAs rounded to nearest, while the tensor core aligns each block of 8 exact products and C to the largest exponent with one bit more than f32's fraction, truncates the bits below, adds them exactly and truncates the sum to f32 (Khattak and Mikaitis's A100 model, MATLAB-tensor-core's A100TC). AHmmaF32 computes that model with integer operations on the bits, so Metal's fast math can't change it, and an Inf or NaN operand keeps the FMA chain's result; it costs about 1,640 more MSL statements per HMMA. An exact Python version of the model matches all 4,096 of a Jetson Orin's outputs (blocks of 4 or 16, other extra-bit counts or a final rounding to nearest match at most 75 %), and with blocks of 8 for bf16 and 4 for tf32 it also matches all of hmma_16816_bf16 and hmma_1688_tf32; test_dmsh_ampere_orin.py's hmma_16816_f32 went from 2,693 of 4,096 lanes differing to none and leaves KNOWN; chain 44897f07bf75
…ore does: HMMA.16816.F32.BF16 and HMMA.1688.F32.TF32 stopped the translator (HMMA size 1 source type 1, size 0 source type 2), since AHmma only knew f16. A bf16 is now read in m16n8k16's f16 fragment layout and a tf32 in m16n8k8's own layout (one element a register: A at rows g and g + 8, columns t and t + 4; B at k t and t + 4), with its low 13 bits ignored as the tensor core does (with them kept, no output of the Orin's 4,096 matches). AHmmaF32's model takes both, in blocks of 8 for bf16 and of 4 for tf32, with the 8-bit exponents brought to its scale, and a right shift past 31 that only those exponents reach is clamped. test_dmsh_ampere_orin.py's hmma_16816_bf16 and hmma_1688_tf32, which the translator refused, now match all 4,096 outputs each (an FMA chain left 1,334 and 2,189 differing) and leave KNOWN; chain b0cbe916d032
…16n8k16 outside a compute shader) sums as Ampere's tensor core does: the per-thread HMMA added all of k to C with f32 FMAs and rounded once to f16, while the tensor core aligns each block of 8 exact products and C to the largest exponent with 24 fraction bits, truncates the bits below, adds them exactly and rounds the sum to f16 (to nearest, even), which is the next block's C. AHmmaExact (seedds#2062's generalised model, f16 sources only with an f16 accumulator) now computes it for both accumulators with integer operations (AHmmaF16Block rounds a block's sum to f16 bits), about 900 more MSL statements for an m16n8k8 and 1,760 for an m16n8k16; an Inf or NaN operand keeps the FMA chain. An exact Python version matches every output of hmma_16816_f16 and hmma_1688_f16 and all 4,718,592 of fxd0h's Orin outputs (3 x 4,096 warps, single and chained HMMA.16816), where PR seedds#1897's f32 sum a block matches 99.97 %; through the per-thread path on this Mac's GPU all of them match too (main: 1,603,197 differ). A compute shader's m16n8k16 keeps PR seedds#1897's simdgroup_matrix path: the same model there, each lane summing its outputs in integers, also matches all of them, but isolated on an M5 Max (best of 200 dispatches of 20 x 23 threadgroups, synthetic inputs) Ocarina of Time's NDSR fa3033d05c0a41d6 took 5.71 ms against 0.249 and d99fbeeb53d38847 6.22 against 0.254, and all 16,620 of that game's HMMAs are m16n8k16 in compute shaders, so none of them change. test_dmsh_ampere_orin.py's hmma_1688_f16 went from 1 of 4,096 lanes differing to none and leaves APPROX; hmma_16816_f32, hmma_16816_bf16, hmma_1688_tf32 and the compute shaders' MSL are unchanged; chain f7357f2435ad
github-actions
Bot
added
nvn
Know-It-All-Brother
added a commit
to Know-It-All-Brother/nxnative
that referenced
this pull request
Nov 9, 2026
…exactly (PR seedds#2067), with this branch's translator rebuilt from its own source; chain 997c02b12c13
2 tasks done
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Part of the Ampere accuracy series (#2040). It depends on #2041 for the test and builds on #2061 and #2062 (the tensor-core model): the commits before this PR's own are theirs and drop out once those land, so this PR's own change is the last commit.
What and why
HMMA with an f16 accumulator has two paths. In compute shaders, m16n8k16 goes through Apple's matrix hardware with #1897's rounding (each block of 8 summed in f32 and rounded once to f16), which matches the Orin on 99.97 % of outputs. Everywhere else (other shader stages, and m16n8k8) it runs per thread, and that path summed all 16 products with FMAs and rounded once. Forced through the per-thread path, fxd0h's Orin recordings (gist) differ in 1,603,197 of 4,718,592 outputs (34 %). The Orin suite didn't catch it, because its m16n8k16 op runs as compute.
The fix uses #2061's tensor-core model with a final round to f16. It's the same A100 model, with blocks of 8, alignment truncated to 24 fraction bits and the sums exact, but each block's sum is rounded to the nearest-even f16 and becomes the next block's C. Against the Orin it matches:
hmma_16816_f16/hmma_1688_f16With the per-thread path forced on the GPU, all 4,718,592 outputs now match.
AHmmaF32becomesAHmmaExact, taking the accumulator type, with the shared block sum split out. The f32-accumulator, bf16 and tf32 MSL is byte-identical to #2062's. An Inf or NaN operand keeps the old FMA result.The compute path is unchanged, on purpose. The same exact model in the compute body also matches all 4.7 million outputs, but on Ocarina of Time's two NDSR upscaler shaders, isolated on an M5 Max (best of 200), it costs 23 to 30 times as much: 0.249 → 5.71 ms and 0.254 → 6.22 ms at a 20 × 23 grid. A cheap correction doesn't work either: a safe bound on where the f32 sum and the truncated sum round differently flags 97 % of HMMAs. So compute keeps #1897's method and its 0.03 % gap. That patch isn't part of this PR.
Cost: the per-thread path goes from 163 to 1,062 MSL statements for m16n8k8, and from 252 to 2,015 for m16n8k16. In Ocarina of Time's shader code, every one of its 16,620 HMMAs is m16n8k16 with an f16 accumulator in a compute shader, so none of them takes this path. I haven't timed it on the GPU.
Reference entry: "HMMA with an f16 accumulator outside compute" in
docs/reference/shaders.md.How to verify
hmma_1688_f16(per thread) goes from 1 differing lane of 4,096 to 0 and leavesAPPROX.hmma_16816_f16(compute) is unchanged at 0 and stays inAPPROX, since its path still has #1897's 0.03 %. The 34 % figure needs the per-thread path forced on for m16n8k16, which I did with a scratch translator that turns the compute route off; it isn't a switch in this PR.Not checked against the GPU: the sign of a zero result, and Inf, NaN or overflow outputs (fxd0h's data and the fixture have none).
Checked
test_dmsh_ampere_orin.py,test_dmsh_hmma.py; every othertest_dmsh_*gives the same result as on main (test_dmsh_ampere_mem.py,test_dmsh_meshstage.pyandtest_dmsh_pairs.pyfail on main the same way);test_nsp_chain.py,test_xlate_anycpu.py,test_docs.py.lint_python.pyshows no new findings.tool_lock.py --update-lockanddmsh_chain.py --headerwere run (chain f7357f2435ad).regress.py: not run. My only local Switch 2 build's shader package comes from an older shader split than main's map, so main's chain can't build it.Docs
docs/reference/shaders.mdhas the entry, andpython3 dm/tools/docs_symptoms.pywas rundmsh_chain.py --headerwas run, anddmsh_chain.h, the translator dlls and the lock are part of the change