Skip to content

HMMA with an f32 accumulator truncates as the tensor core does: bit-exact with a Jetson Orin (chain change) - #2061

Open
Know-It-All-Brother wants to merge 3 commits into
seedds:mainfrom
Know-It-All-Brother:ampere-hmma-f32-accum
Open

Know-It-All-Brother wants to merge 3 commits into
seedds:mainfrom
Know-It-All-Brother:ampere-hmma-f32-accum

Conversation

@Know-It-All-Brother

@Know-It-All-Brother Know-It-All-Brother commented Nov 9, 2026 •

Copy link
Copy Markdown
Contributor

Part of the Ampere accuracy series (#2040). It depends on #2041 for the test: the commits before this PR's own are #2041's and drop out once it lands, so this PR's own change is the last commit.

What and why

HMMA.16816 with f16 sources and an f32 accumulator added its products with 64 ordinary FMAs. The tensor core adds them its own way, so two thirds of the outputs differed from the Orin (2,693 of 4,096, median 1 ULP, up to 3,328 ULP where the sum nearly cancels).

The tensor core's arithmetic is the A100 model in Khattak and Mikaitis's MATLAB-tensor-core toolbox (models/A100TC.m; the same source #1897 used for the f16 accumulator):

  1. Every product is exact.
  2. Each block of 8 products, plus C, is aligned to the block's largest exponent on a grid with 24 fraction bits (f32's 23 plus one), and every bit below that is truncated.
  3. The aligned terms are summed exactly, and the sum is truncated to f32.
  4. In an m16n8k16 HMMA, the second block takes the first block's result as its C.

I swept the alternatives against the Orin's outputs with an exact model:

Model Outputs matching (of 4,096)
Blocks of 8, 1 extra bit, round toward zero 4,096
Blocks of 8, 0 / 2 / 3 extra bits 50 % / 75 % / 69 %
Blocks of 4 or 16 at most 69 %
Blocks of 8 with a final round to nearest 61 %

AHmmaF32 does this per lane with integer operations, which Metal's fast math can't change. It decodes each half's significand, exponent and sign, multiplies exactly, takes each block's largest exponent, shifts and sums the terms in a 32-bit integer, then normalises with truncation. An HMMA with an Inf or NaN operand keeps the old FMA result.

The cost is large, but only on this instruction. One HMMA.16816.F32 goes from about 230 MSL statements to about 1,870, almost all integer operations. In Ocarina of Time's shader code there are 13,267 f16-accumulator HMMAs and only 2 with an f32 accumulator, so I expect no measurable effect there. The f16-accumulator path, including its simdgroup_matrix route in compute shaders, is unchanged.

The same model, with blocks of 8 for bf16 and 4 for tf32, makes HMMA with bf16 and tf32 sources match the Orin exactly too; #2062 builds on this PR for that.

Reference entry: "HMMA with an f32 accumulator truncates as the tensor core does" in docs/reference/shaders.md.

How to verify

DMSH_TEST_XLATE=dm/gfx/tools/dmsh_xlate_bin/dmsh_xlate AMPERE_ORIN_ONLY=hmma_16816_f32 python3 dm/gfx/tests/test_dmsh_ampere_orin.py

hmma_16816_f32 went from 2,693 of 4,096 differing lanes to 0, and it comes out of KNOWN.

The same model with a final round to nearest f16 also explains every lane of the f16-accumulator ops, including the last 0.03 % that #1897's "sum in f32, round once" leaves. Because f16-accumulator HMMAs are common and performance-sensitive (agl's upscaler), I haven't changed that path here; it's worth its own measured PR. Not checked against the GPU: Inf / NaN inputs, overflow and the sign of a zero result, which aren't in the fixture.

Checked

  • Tests: test_dmsh_ampere_orin.py, test_dmsh_hmma.py; every other test_dmsh_* gives the same result as on main (test_dmsh_ampere_mem.py, test_dmsh_meshstage.py and test_dmsh_pairs.py fail on main the same way); test_nsp_chain.py, test_xlate_anycpu.py, test_docs.py. lint_python.py shows no new findings.
  • Translator: the dlls are rebuilt from this tree with .NET 10; tool_lock.py --update-lock and dmsh_chain.py --header were run (chain 44897f07bf75).
  • Devices and OS: MacBook Pro M5 Max, macOS 27.
  • 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.
  • Games run: none.

Docs

  • docs/reference/shaders.md has the entry, and python3 dm/tools/docs_symptoms.py was run
  • A chain tool changed: dmsh_chain.py --header was run, and dmsh_chain.h, the translator dlls and the lock are part of the change
  • No game files, keys or data extracted from a game

…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
@github-actions github-actions Bot added nvn NVN on Metal and its tests (dm/gfx/nvn, dm/gfx/tests) docs Documentation shader chain The shader translation chain (DMSH_CHAIN): a change rebuilds every game's shaders labels Nov 9, 2026
Know-It-All-Brother added a commit to Know-It-All-Brother/nxnative that referenced this pull request Nov 9, 2026
…edds#2061), with this branch's translator rebuilt from its own source; chain 28f3f9f31e59
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

docs Documentation nvn NVN on Metal and its tests (dm/gfx/nvn, dm/gfx/tests) shader chain The shader translation chain (DMSH_CHAIN): a change rebuilds every game's shaders

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant