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 |
This branch has not been deployed
2 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.