Repository navigation
honeykrisp: compute set for M1/M1 Max (coopmat, precise math, dependency-tracked CDM barrier, submit poll) - #3
Closed
joshuaswarren wants to merge 16 commits into
Closed
joshuaswarren wants to merge 16 commits into
joshuaswarren wants to merge 16 commits into
Conversation
…add32 Port the GPU-verified HW matrix path from the agx-llm-perf showcase: opcodes+packer (byte-exact vs applegpu), NIR intrinsic bridge, agx_nir_lower_simdmat lowering, validator widths, and gated advertisement of an 8x8x8 fp32 coopmat shape. Everything is behind AGX_SIMDMAT; default Honeykrisp behavior is unchanged.
…le coopmat shaders
VK_KHR_cooperative_matrix on Honeykrisp only advertised 8x8x8 fp32 even
though agx_nir_lower_simdmat.c already tiles a 16x16x16 coopmat into a
2x2x2 grid of simd_matrix_fmadd ops, and the G13 matrix unit takes fp16
operands. Now advertised (subgroup scope, N in {8, 16}, NxNxN):
fp32 A/B, fp32 C/D simd_matrix_fmadd32 with 32-bit pairs
fp16 A/B, fp32 C/D simd_matrix_fmadd32 with 16-bit A/B pairs
fp16 A/B, fp16 C/D simd_matrix_fmadd16
The mixed form is the one Apple's own Metal mul_mm kernel emits
(simd_matrix_fmadd32 r15_r16, r49l_r49h, r46l_r46h, r15_r16); the generic
ALU packer already encodes the per-source size bits, so the lowering
passes A/B at their own width instead of up-converting to fp32, and the
nir intrinsic gains a 16-bit dest form that selects fmadd16.
Bug found while verifying: with fp16 A and fp32 D the register allocator's
early-kill let D (a 32-bit pair) land on top of A (a 16-bit pair inside
that pair). The hardware then produced zeros for that fmadd: 16x16x16
fp16/fp32 failed in every input/layout variant, always sub-block (1,1),
whose last fmadd was the only one with D overlapping A. Apple's compiler
only ever emits D == C, so the RA now keeps A and B live until D is
placed (D may still reuse C).
Verification on Apple M1 (G13G B1), receipts/cmshape.c, one coopMatMulAdd
per run, exact host reference, integer-valued inputs bit-exact and random
floats within tolerance, identity/ones/randint/randf x row/col-major:
stock Mesa 26.1.7: FAIL, extension not exposed
5bb2b28 (before): 8x8x8 fp32 8/8 PASS; 16x16x16 not advertised
this commit, no RA guard: 40/48 (16x16x16 fp16/fp32: 8/8 FAIL)
this commit: 48/48 PASS, plus 48/48 randint over 4 seeds
mlx-omarchy Release tests (primitive 98, reduce_ops 34, runtime 40,
matmul_family 14 cases) pass identically on the stock ICD and on this
driver with AGX_SIMDMAT=1.
Bug: (word >> s) & mask with a non-constant s returns zero on Honeykrisp
whenever s has any bit set above bit 4 (s >= 32). mlx-omarchy hit this as
the "dynamic shift-then-mask byte extraction" miscompile family: bool
scatter dropping writes across word lanes, select picking the wrong
operand under broadcast conditions, reduce_general's load_truthy reading
zero, masked_scatter's scan undercounting (s_scan[255] = 32, not 128).
Root cause: agx_nir_fuse_algebraic_late rewrites
iand(ushr(a, b), (1 << bits) - 1) -> ubitfield_extract(a, b, bits)
NIR's ushr uses only the low 5 bits of its count, and by this point the
main nir_opt_algebraic has already relied on that to strip explicit masks
(ushr(a, iand(b, 31)) -> ushr(a, b); ushr(a, ishl(iand(b, 3), 3)) ->
ushr(a, ishl(b, 3))). ubitfield_extract has no such wrap: an offset >= 32
is undefined in NIR and bfeil returns 0 for it on the hardware.
lower_sm5_shift, which masks the count of every surviving shift, runs
after the fusion and never sees the ushr it consumed.
Fix: fuse directly only when the count is a constant below 32; otherwise
fuse with an explicit iand(b, 31) on the offset. Same instruction count as
the unfused shift+and, and the mask CSEs with one the shader already had.
Verification on Apple M1 (G13G B1), standalone Vulkan compute reproducer
(20 variants, 1024 lanes each, host-compared):
stock Mesa 26.1.7: 9 of 16 base variants FAIL (byte_dyn,
byte_dyn_const, byte_dyn_gid, scan_dyn, scan_dyn_const, scan_dyn_rand,
select_bcast, select_dyn, shift_dyn_mask), plus all 4 boundary
variants FAIL (count exactly 32: 1016/1024 wrong; 63: 516/1024;
uniform 40: 1020/1024; divergent 8/40: 509/1024, only the 40 lanes).
this commit: 20 of 20 PASS, 0 mismatches.
mlx-omarchy suites, stock vs this commit (AGX_SIMDMAT=1), identical:
primitive 98/98 (2699373 assertions), reduce_ops 34/34 (7152),
runtime 40/40 (6307), indexing_ops 53/53 (4093), select_layout 11/11
(1908), scatter_determinism 21/21 (127).
…duction
Three fp32 rounding defects on Honeykrisp (Apple M1 G13G, Mesa 26.1.7)
that mlx-omarchy patches around in its shaders:
a) GLSL "/" was lowered by nir_opt_algebraic (lower_fdiv) to
a * frcp(b), and frcp was only Newton-refined when the instruction
was `precise`. Even with an exact reciprocal a * RN(1/b) is not
RN(a/b): 80/25 and 10/25 come out one ulp low, 6.55859375/0.030016447
one ulp high. Measured on 1008 pairs: 299 mismatches (298 at 1 ulp,
1 at 2 ulp), 1/b: 204 mismatches; 1M random pairs: 305792 mismatches.
fma() itself is IEEE (0/1000 mismatches), the mlx-omarchy note
blaming fma contraction was a misdiagnosis.
b) log(x) was flog2(x) * ln2 on the hardware log2, which is only
absolutely accurate: log(3) = 1.0986121892929077 instead of
1.0986123085021973, 2792/4107 mismatches on [0.5, 64] with errors up
to 357 ulp near 1 (log(1.0000001) = 1.65e-7 vs 1.0e-7); log2 had
2905/4107 mismatches.
c) sin/cos reduced the argument as fract(x / 2pi) in fp32, so the
reduced angle lost ulp(x) of precision: sin(3.0) 14 ulp off,
sin(1e5) 1.06e-3 absolute error, sin(1e8) = 0.0, sin(-1e-3) 571 ulp
off, sin(pi_f32) = 0, and the hardware sin_pt_2 core itself is 2 to
3 ulp off (cos(1e-3) 3 ulp).
The lowering moves to a new the AGX compiler pass file agx_nir_lower_math.c,
built only from the exact fp32 fma and marked no-fast-math so
nir_opt_algebraic cannot reassociate it:
- lower_fdiv is dropped from agx_nir_options; fdiv32 becomes Markstein
division (y = refined rcp, q = a y, r = fma(-b, q, a), q + r y) with
the NaN/zero guard keeping q for b = 0, overflow and signed zero.
frcp32 is always refined. fdiv16 keeps a * frcp(b).
- fmul(flog2(x), ln2) as emitted by nir_flog, and any remaining flog2,
become a double-float evaluation of log(m) = 2s + 2s^3 P(s^2) with
s = (m-1)/(m+1) kept as a pair, m folded into [sqrt(1/2), sqrt(2)),
e ln2 split hi/lo and Fast2Sum accumulation (about 40 bits before the
final rounding). Zero/subnormal/negative/inf/NaN are handled by bit
tests, not the hardware log2.
- fsin/fcos: Cody-Waite in radians with pi/2 in three pieces, the k C2
product and its subtraction as TwoProduct/TwoSum so the reduced
argument is a pair accurate to an ulp, for |x| < 2^22; an integer
Payne-Hanek (24-bit mantissa times a 128-bit window of 2/pi, 64
fraction bits kept) for |x| >= 2^22 up to FLT_MAX. The quadrant then
selects a Taylor polynomial in quadrant units (sin degree 13, cos
degree 14, leading terms as pairs, low part of the argument fed to
the derivative). The hardware sin_pt_1/sin_pt_2 core (nir_op_fsin_agx)
is no longer used by the lowering. fp16 sin/cos go through the fp32
path.
Verification (receipts/2026-09-08-precise-math.json, standalone Vulkan
compute reproducer in receipts/precise-math/):
- division: 1008-pair table incl. 80/25, 10/25, 1000 random: 0
mismatches (stock 299); 1/b 0 (stock 204); 1M random pairs 0
mismatches for both a/b and 1/b; b = 0, inf, NaN, overflow, flushed
subnormal and -0/5 = -0 match host IEEE.
- log: 4107 points on [0.5, 64] plus pinned exact cases: 0 mismatches
(stock 2792), log(3) = 1.0986123085021973; log2 0 (stock 2905);
1M random on [1e-37, 1e38]: log 3 at 1 ulp, log2 14 at 1 ulp,
rest exact.
- sin/cos at 1e5, 5e5, 1e6, 5e6, 1e8: all 0 ulp from the double
reference except cos(5e6) at 1 ulp (stock: 2.8e5 to 1.06e9 ulp);
1310 mixed points max 1 ulp; 1M random over [1e-3, 1e9] and [-10, 10]:
sin 42346 and cos 41184 at 1 ulp, rest exact, max 1 ulp (stock
>= 8 ulp for 627689 / 512816 of them).
- mlx-omarchy suites on the same Release tree, stock vs fixed:
primitive 98/98, eq_math 7/7, fast_ops 30/30, runtime 40/40,
reduce_ops 34/34 both; complex_ops 30/33 both with the same three
pre-existing cases (failed assertions 31 -> 29). The upstream
ops_tests log(3) pin through libmlx: FAIL on stock, PASS on fixed,
as do 80/25, 10/25 and sin/cos up to the 1e5 limit mlx enforces.
- timing: 16M-element elementwise dispatch is memory bound at ~2.45 ms
on both drivers (sin 2.474 -> 2.463 ms median, within spread). A
32-deep in-register chain shows the ALU cost: sin 6.64 -> 39.6 ms,
log 2.57 -> 34.8 ms, div 2.84 -> 7.24 ms, mul unchanged.
…it late Symptom: any shader converting fp32 to a 64-bit integer aborts in the backend with "Unhandled ALU op fdiv" (agx_compile.c agx_emit_alu); the mlx-omarchy suite omarchy_copy_offset_tests crashes on "dtype converting copies cover the integer family and bool". Fragment shaders with perspective-correct varyings hit the same path. Cause: 2022dde dropped lower_fdiv from the compiler options so nir_opt_algebraic no longer splits fdiv into fmul/frcp, and only ran agx_nir_lower_fdiv in agx_preprocess_nir. nir_lower_int64 (f2u64/f2i64: fdiv by 2^32, plus the frem that lower_fmod turns into a second fdiv) and agx_nir_lower_interpolation run after preprocessing, so their fdiv reached the backend, which has no fdiv. Fix: agx_nir_lower_fdiv_late runs in agx_optimize_nir after the int64 cleanup loop and lowers only fdiv, leaving frcp alone so the reciprocals built by the first pass and by the log lowering are not refined twice. Verification: hk/precise-math driver alone reproduced the abort on jwm1; with this commit omarchy_copy_offset_tests passes 25/25 (see the honeykrisp-omarchy integration receipt).
…0 disables)
Replace the five getenv("AGX_SIMDMAT") checks (agx_preprocess_nir lowering,
the extension table, the feature bit, the supported stages, and the
properties query) with agx_simdmat_enabled(), which reads the variable through
debug_get_bool_option with a default of true. Setting AGX_SIMDMAT=0 restores
the previous behaviour.
The 25/25 mlx-omarchy battery and every reproducer in
receipts/2026-09-08-honeykrisp-omarchy-integration.json ran with the gate on
(AGX_SIMDMAT=1), so the default now matches the verified configuration and a
system-installed driver exposes coopmat to unmodified wheels.
… path (G13X)
hk_dispatch_with_usc_launch flushed the full CDM barrier after every launch
unconditionally, ~4.6-5.1 us per dispatch on compute streams of hundreds of
small dependent dispatches. Key the barrier on what the app declared instead:
vkCmdPipelineBarrier2 marks a pending dependency and the next CDM launch
carries the barrier and clears it. Launches the app left unordered have no
Vulkan-defined ordering to enforce, so their drain is skipped. The flag lives
on the command buffer so it survives CS splits and merge_control_streams;
hk_EndCommandBuffer lands an unconsumed barrier on the trailing edge, and
driver-internal ordering (query availability writes, hk_queue_write) keeps its
explicit flush. The compute batch also stays open across a compute-to-compute
vkCmdPipelineBarrier2 when no graphics batch is open.
On G13X the barrier carries the measured designed set {unk_4,5,6,8} + USC
cache invalidate plus bits 0-2 (0x17f). Bit bisection showed every one of the
designed bits is required for bit-exact results: removing any corrupts, and a
USC-invalidate-only barrier is fast but wrong (320/320 greedy token flips).
G13G/G14 keep the previous kitchen-sink barrier.
New perftests: HK_PERFTEST=alwayscdmbarrier (restore the unconditional flush),
nocdmbarrier, usccdmbarrier, designedusccdmbarrier, and HK_CDM_BARRIER_MASK
(exact flag word) for bisecting.
Measured on an M1 Max (G13C), MLX-on-Vulkan Qwen3.5-class 4-bit greedy decode:
the designed-set trim 63.50 -> 68.34 tok/s (10 interleaved paired reps, +7.7%,
identical token digests); with the app declaring only real dependencies
(MLX dependency gating on) the tracked barrier adds +2.9% (81.3 -> 83.7 tok/s),
and 16 interleaved runs of the 10-prompt x 10-pass greedy contract produced
identical token digests in both arms.
…_US) HK_SUBMIT_POLL_US=<us> (default 0 = off) makes vk_drm_syncobj_wait_many poll with a zero timeout for up to that long before falling back to the blocking timeline wait with the original deadline. Short submit->signal round trips (one submit + one wait per step of a stepwise decoder loop) pay a host wake-up premium under the sleeping wait; the poll observes the signal the moment the kernel publishes it. After a budget expires without seeing the signal the poll is disabled for budget << misses (capped at 1024x), so long waits (shader cache compile storms, whole-frame jobs) pay at most one budget per backoff period; any poll hit resets the state. Semantics preserved: same ioctl with timeout 0 during the poll budget, original absolute deadline on the fallback, ETIME still maps to VK_TIMEOUT, non-ETIME errors surface immediately, and drivers that do not set sync_wait_poll_us take the identical single-call path. Measured on an M1 Max (G13C): a Parakeet TDT decoder loop (278 submit+wait round trips) 341 -> 254 ms (-25%) at HK_SUBMIT_POLL_US=2000 (poll budget sweep 0/100/200/500/1000/2000/5000 us: 347/314/292/266/258/248/250 ms), transcript unchanged; LLM decode/prefill throughput and token digests unchanged.
Author
|
This is the bottom of a stack: #3 → #4 → #5. Review in that order. #4 adds the rest of the |
This was referenced Sep 30, 2026
…ction step Dj stress-testing an M2 Pro found a third division failure class the original review missed (Basecamp 10277024959, 2026-09-13): when the exact quotient sits just below the normal boundary, the GPU zeroes the correction step's intermediate value and the Markstein step is silently skipped, giving 1 ulp errors -- 18,382 wrong results over 1M random pairs. The first commit in this series refined a * rcp(b) in place, so the flush landed on an intermediate, not on the result. Divide the normalized mantissas instead and restore the exponent afterwards, so the refinement always runs on operands near one and the flush happens on the exact quotient, before rounding. The rounding decision is integer: the half-ulp point is exact by construction, and the residual of the second Markstein step is exact by the fma contract, so the comparison against it decides the tie without a third refinement. The reciprocal takes the same path, and the log2 series divisor (near two) keeps the unguarded Newton form. This re-lands the division part of the fix first made on 2026-09-13 and pulled with it during the 2026-09-20 A/B isolation windows; it was never re-applied to the omacom/mesa stack. Verified against an exact model over 262,144 pairs covering edge bit patterns and the binade boundaries, and by an independent exact-rational checker; see the receipt in the PR.
nir_builder_opcodes_h.py renders the intrinsic builders from a mako template that references FP_MATH_CTRL, but the index was never imported or passed to render(), so mako resolved the name to UNDEFINED and silently skipped the builder-defaulting block for every FP_MATH_CTRL intrinsic. NoContraction on cooperative-matrix operations was therefore dropped in every driver, not just this one, and the generated header is missing the twelve fp_math_ctrl defaulting blocks. Import the index and hand it to the template. The generated header grows exactly 24 lines, twelve two-line defaulting blocks. (cherry picked from commit 864b4e8)
cmshape2 walks the advertised matrix shapes, layouts, workgroup sizes and pointer widths against a host oracle, and checks that NoContraction survives and that a divergent component write is not lost. run_matrix2 drives the 30 cases that cover the fixes in this series. precise_math compares division, reciprocal, log and log2 against an exact model over 262,144 inputs, including the binade boundaries and the literal-versus-buffer pair for ln2. compute_runner is the small Vulkan dispatch harness it uses. (cherry picked from commit fb62375)
Dj's stress runs found (Basecamp 10277024959, 2026-09-13) that a precise log2(x) * ln2 gave different bits depending on whether ln2 was a literal or read from a buffer, on 26.8% of inputs. The stack carries an amended form of the log fix in which a multiply of a hardware flog2 by the fp32 ln2 constant is rewritten to the correctly rounded soft log. That reopens the finding from the other side: the literal form now gets the single-rounding soft log, while the buffer form still computes soft log2 and multiplies, so the two forms of the same source expression differ by 1 ulp on 25.3% of inputs (66,332 of 262,144 in the literal_vs_buffer reproducer), and the precise multiply check fails the same way. Drop the literal-only rewrite. log2 is lowered to the double-float soft log2 everywhere, and log() keeps the one rounded multiply the source asks for, from a literal or from a buffer: the two forms now agree bit for bit on every input, powers of two still come out exact because the series returns the exponent unchanged, and log stays within 1 ulp on the 262,144-input reproducer. (cherry picked from commit 155e743)
Author
|
Dj, I went back through your September review and fixed what was still open. I also found two regressions of our own in the stack. Fixes are pushed (fast-forward): #3 57809ff, #4 29074c4, #5 e129cd1, #6 6dc1fba. From your review:
Regressions we found and fixed:
CTS 1.4.6 on G13X, #6 before vs after: compute 12,379 pass / 2 fail / 48,457 not supported on both (the same two copy_ssbo_bounds cases), subgroups 12,939 / 0 / 35,766 on both. No case that passed before fails now. float_controls fp32 is identical to stock on G13X and G13G. Could you take another look at #3? |
joshuaswarren
changed the title
joshuaswarren
force-pushed
the
jw16/omacom-pr
branch
from
October 8, 2026 13:51
57809ff to
c56e086
Compare
Author
Author
Author
This was referenced Oct 29, 2026
Author
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.
NOT READY TO MERGE
Head:
c56e08634bba83c00b7345494a4cef9bdf681686.GitHub reports MERGEABLE against
main: no text conflicts. This is not a test pass.Superseded by #6. Do not merge this old layer on its own. Keep it open until the replacement is qualified and the agreed closure step is reached.
Earlier review, source and test notes are kept below with their original bindings. They do not change this status.
Honeykrisp (M1 / M1 Max): compute performance and correctness set
Eleven commits on top of current
main, all from running MLX-on-Vulkan LLM inference (Qwen3.5-class GatedDeltaNet + attention, 4-bit) and Parakeet speech decoding on Apple M1 Max (G13C) hardware under Apple Silicon Linux.Contents
Cooperative matrix / compiler (built and used daily for prefill GEMMs):
honeykrisp: hardware 8x8x8 fp32 cooperative matrix via simd_matrix_fmadd32and the ISA/signature follow-ups,advertise 16x16x16 and fp16 cooperative matrix shapes,enable VK_KHR_cooperative_matrix by default (AGX_SIMDMAT=0 disables).agx: keep the mod-32 shift count when fusing ushr+iand into ubfe.agx: correctly rounded fdiv/frcp, faithful log/log2, exact sin/cos reductionand the late-fdiv lowering.Compute submission cost (the two new performance commits):
honeykrisp: dependency-track the per-launch CDM barrier on the Vulkan path (G13X): the barrier is emitted at the next launch after avkCmdPipelineBarrier2instead of after every launch, carries the designed bit set (every bit required, bisected), compute batch stays open across compute-to-compute barriers.HK_PERFTEST=alwayscdmbarrierrestores the old behaviour; new bisect knobsnocdmbarrier,usccdmbarrier,designedusccdmbarrier,HK_CDM_BARRIER_MASK.hk: opt-in bounded poll before blocking syncobj waits (HK_SUBMIT_POLL_US): default off; polls with a zero timeout before the blocking wait, with backoff after misses so long waits are unaffected.Evidence (M1 Max, Linux 7.1.x, greedy MLX Qwen3.5-class 4-bit)
HK_SUBMIT_POLL_US=2000; sweep 0/100/200/500/1000/2000/5000 us: 347/314/292/266/258/248/250 ms. Transcript, mel and encoder outputs unchanged; LLM decode and prefill unchanged.main(Mesa 26.3-devel) and re-validated with the same gates: decode 88.57 vs 88.21 tok/s and prefill 900.9 vs 895.7 tok/s against the previous build, all digests identical, Parakeet contract green.Scope and caveats
HK_SUBMIT_POLL_USis opt-in in the driver; the consuming runtime (MLX Omarchy backend) sets 2000 by default at device init.simd_matrix_fmaddshaders; hardware verification receipts and reproducers live in the fork history and are not included here.