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.
|
Thanks - it doesn't make sense to undefine the SBGEMM_DEFAULT_(P,Q,R) of unrelated cpus though, this will only break builds that were using either dedicated sbgemm kernels or the generic C fallback |
|
Regarding your comment, a dedicated VORTEXM4 build is not expected/intended to work on M1 to M3 hardware - cpu autodetection at build or (in the case of a DYNAMIC_ARCH-built library) runtime will select the non-SME "VORTEX" target for these. |
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.
|
Fair point — dropped in 0bc31e8. The VORTEXM4 block now only overrides the UNROLL values (these define the packed panel layout, same |
|
Understood on the intended usage — I'd misframed that section (this machine is an M3 Pro, so I was testing the dedicated build directly on unsupported hardware). Two clarifications on why the
By the same token I'm fine leaving the SME1 |
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.
|
Now I see what you were pointing at, and I owe you a correction: my reply above fixed only the undefs inside the VORTEXM4 block. The edit that had added them matched on the SBGEMM UNROLL undef lines, which occur verbatim in several other cpu blocks, so the same three undefs had landed in four unrelated blocks as well - and their setparam then failed under DYNAMIC_ARCH (this is also what broke CI on #6145: |
|
Tested on an M4 Pro (SME2, SVL 512), Apple clang 21, static As posted, the SME2 path fails on M4 in two ways:
The fix (both in - const svbool_t pt = svptrue_b32();
+ const svbool_t pt = svptrue_b16();
...
-__arm_new("za") __arm_locally_streaming static void s2b_blk(BLASLONG k, const IFLOAT *a, BLASLONG wm, const IFLOAT *b,
- BLASLONG wn, FLOAT *C, BLASLONG ldc, FLOAT alpha,
- int sgn) {
+S2_NOINL void s2b_blk(BLASLONG k, const IFLOAT *a, BLASLONG wm, const IFLOAT *b, BLASLONG wn, FLOAT *C, BLASLONG ldc,
+ FLOAT alpha, int sgn) S2_S S2_ZA {With the fix:
Speed (GFLOPS, square NN, M4 Pro):
Two notes on speed, not blockers:
|
Two fixes from Michael Tesch's M4 Pro testing:
- s2b_blk had its own __arm_new("za"), so the call from the equally-
annotated s2b_drive set up a lazy ZA save that __arm_tpidr2_save
(OpenBLAS's aborting weak stub) committed. It now shares ZA from the
drive function through S2_S S2_ZA like the fp32/fp64 block functions.
- BFMOPA governs per 16-bit element, but the predicates were
svptrue_b32(), which enables only the even bf16 lanes and drops the
k+1 half of every pair (every k >= 2 wrong, k = 1 right). Now
svptrue_b16().
Also apply the suggested thread cap to SBGEMM: sme_level3_threads()
capped only real precision (!BFLOAT16), so default threads contended
for the SME units - 839 GFLOPS at n = 4096 against about 1190 on one
thread and 2347 with one thread per unit on the M4 Pro. The condition
is now (!BFLOAT16 || !BGEMM): bgemm keeps the NEON kernels and its
threading.
With the fixes his run reports the sbgemm/BLAS/CBLAS/utest suites
passing and 3000 randomized trials clean; 1193 GFLOPS at 2048^2 single
thread against 0.48 on develop.
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.
|
Thank you for the thorough test run - both fixes and the thread cap are in bce51f1 (kernel predicates, S2_S S2_ZA block handoff, and the If you are up for one more measurement: #6145 gives VORTEX (M1-M3) a NEON BFDOT sbgemm (57 GFLOPS at 1024^3 on an M3 Pro, 34x over the generic fallback). It is stacked on this PR and reuses its packing copies, so a quick |
This branch has not been deployed
3 participants
Adds an SME2
sbgemm(bf16 × bf16 → fp32) kernel in the style of the fp32/fp64 SME2 kernels of #6074, plus the packed-layout fixes it needed, and one crash fix in the shared SME2 detection.What it does
BFMOPA ZA32with.Hoperands accumulates 2-wide dot products, so one instruction covers two k steps. The kernel (kernel/arm64/sbgemm_kernel_sme2.c) zips rows l and l+1 of the standard packed[k][m]/[k][n]panels into the lane layout2*i + (k&1)that the outer product needs; a missing second row of an odd k is loaded under a false predicate, which zero-fills. C accumulates in four ZA32 tiles of 16×16 fp32 through the load/store helpers ofsme2_gemm_tile.h(SBGEMM_UNROLL_M 64,SBGEMM_UNROLL_N 16), the bfloat16 counterpart of the fp32 micro-kernel shape.New bfloat16 packing copies (
sbgemm_{n,t}copy_{64,16}_sme2.c) emit[k][w]panels for both operand layouts — 64-row chunks on the M side, power-of-two groups on the N side. These are needed because the generic copies emit column-major groups on the T side, which the kernel does not expect. Without SME2 at run time (or in the compiler), plain loops run on the same packed layout.Fix included: the SME2 run-time check could not return false
s2_usable()computedsme2 && svcntsw() == 16. clang if-converts the conjunction and speculatessvcntsw()(RDSVL — undefined without FEAT_SME) above the sysctl test, so the guard itself trapped instead of returning false. I ran into this by testing a dedicatedTARGET=VORTEXM4build on an M3 Pro (unsupported, but it showed the guard never gets a chance to work); with DYNAMIC_ARCH the VORTEXM4 code is only entered on M4, so this is defensive. The vector length is now read through anoinlinehelper with inline asm, which clang will not speculate.Status
make TARGET=VORTEXM4(Apple clang 21),make DYNAMIC_ARCH=1, and CMake all build cleanly; GCC keeps the generic kernels as before (the KERNEL entries sit inside the existing clang guard).cblas_sbgemmcases (all four transpose combinations, m/n/k from 1 to 257 plus 1000–1200 sized cases, α ∈ {1, −1, 0.5, 3}, β ∈ {0, 1, 0.25}, odd strides) against a double-precision reference: 0 failures on this branch, same as develop.bfmopa× 4 per k-pair on ZA0–ZA3 with predicated zips). The exercised path is the fallback loops plus the full driver/packing pipeline, which shares the packed-layout contract with the SME2 path. Marking as draft until someone can run the BLAS tests and a benchmark on an M4 — happy to iterate on the numbers.