Repository navigation
Conversation
Add an SME2 sbgemm (bf16 x bf16 -> fp32) kernel in the style of
gemm_kernel_sme2.c, for a 512-bit streaming vector length:
- BFMOPA ZA32 with .H operands accumulates 2-wide dot products, so the
kernel zips rows l and l+1 of the packed [k][m] and [k][n] panels
into the lane layout 2*i + (k & 1); a missing second row of an odd k
is loaded under a false predicate, which zero-fills. C accumulates
in four ZA32 tiles of 16x16 fp32 via the helpers of
sme2_gemm_tile.h (SBGEMM_UNROLL_M 64, SBGEMM_UNROLL_N 16).
- Custom bfloat16 A and B copies (sbgemm_{n,t}copy_{64,16}_sme2.c)
emit [k][w] panels for both operand layouts, in 64-row chunks on
the M side and power-of-two groups on the N side, like the A copies
of the fp32 SME2 kernels (the generic copies emit column-major
groups on the T side).
- Without SME2 at run time, or in the compiler, plain loops run on
the same packed layout.
Also fix a crash of the existing SME2 kernels on cores without
FEAT_SME2 in s2_usable(): clang if-converts the conjunction and
speculates svcntsw() (RDSVL, undefined without FEAT_SME) above the
sysctl test. Reading the vector length through a noinline helper with
inline asm prevents that.
Only the UNROLL values are specific to the SME2 kernel (the packed panel layout); P, Q and R stay at the shared defaults, which are compatible with the 64-row M chunks.
The edit that added the SBGEMM P/Q/R overrides to the VORTEXM4 block matched on the SBGEMM UNROLL undef lines, which occur in several other cpu blocks as well, so the undefs landed in four of them and broke their setparam compilation under DYNAMIC_ARCH (SBGEMM_DEFAULT_P undeclared, e.g. for NEOVERSEV1 with GCC). This is what the review comment about unrelated cpus was about; the overrides themselves were already dropped in 0bc31e8.
sbgemm on the Apple targets so far ran the generic C fallback; with FEAT_BF16 (present at least since M2, probed at run time so older cores keep the fallback) the new kernel accumulates C in sixteen quad registers of eight columns by two row groups. One BFDOT does a two-wide dot per lane, so rows l and l+1 of the packed panels are zipped into the lane layout (i, k & 1); the B pair of a column is duplicated across the lanes and the A pairs cover four rows each. An odd k is a zero vector, blocks narrower than eight run in scalar code. With SBGEMM_UNROLL_M = SBGEMM_UNROLL_N the driver routes the A side through the O copies as well, so both panels come in the power-of-two grouping of OpenMathLib#6142's copies, which this renames to role-based names (oncopy/otcopy for both sides, ncopy/tcopy kept for the M-side-only use of the M4 SME2 target). On an M3 Pro: 39116-case randomized sbgemm sweep against a double-precision reference passes unchanged, and against the generic fallback of develop the packed path is 34x at 1024^3 (57 vs 1.7 GFLOPS) and 63x at 2048x2048x1024 (114 vs 1.8 GFLOPS), about a quarter of the fp32 NEON sgemm of the same build.
0tomhoyt
force-pushed
the
sbgemm-bfdot
branch
from
October 21, 2026 06:58
fe87291 to
0120425
Compare
This branch has not been deployed
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
1 participant
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.
Fills the last big sbgemm gap on the Apple targets: VORTEX (M1-M3) so far ran
sbgemmon the generic C fallback, ~1.7 GFLOPS on an M3 Pro. The new kernel uses NEONBFDOT(FEAT_BF16), probed at run time viahw.optional.arm.FEAT_BF16so cores without it (early M1) keep the plain-loop fallback — no regression path, and GCC builds keep the generic kernels as before.Stacked on #6142 (it reuses and renames that PR's bfloat16 packing copies to role-based names); merge #6142 first and this applies cleanly on top.
Design
SBGEMM_UNROLL_M = SBGEMM_UNROLL_N = 8: with the two unrolls equal the driver routes the A side through the O copies as well (common_sb.h), so both packed panels come in the same power-of-two grouping and the kernel walks them symmetrically.BFDOTdoes a two-wide dot per lane, so rows l and l+1 of the[k][8]panels are zipped into the lane layout(i, k&1): the B pair of a column is duplicated across the lanes, the A pairs cover four rows each, andacc[n][h]holds four consecutive rows of C column n — one quad load/store per (column, row group). Sixteen quad accumulators, i.e. an 8×8 C block per call.Measurements (M3 Pro, 11 threads, best of 5)
For context, fp32
sgemmof the same build reaches 305/403 GFLOPS at the last two shapes, so bf16 lands at ~28% of fp32 — the 8×8 block and the scalar packing copies leave tuning headroom (wider blocks, vectorized transpose in the copies) that can follow once the layout is in place.Verification
cblas_sbgemmsweep (all transpose combinations, m/n/k 1–257 plus 1000–1200 sized cases, α/β variants, odd strides) against a double-precision reference: 0 failures, same as develop — and unlike arm64: SME2 bf16 sbgemm kernel for VORTEXM4 (Apple M4) #6142 this was validated with the vector path actually executing (the M3 Pro has FEAT_BF16).TARGET=VORTEX,DYNAMIC_ARCH=1(correctness re-run through the runtime-detected VORTEX table entry: 0 failures), and CMake all build and pass.