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.
fe87291 to
0120425
Compare
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.
0120425 to
18be288
Compare
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. 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.
PANEL FORMAT CHANGE: this rewrites the bfloat16 packing copies shared by both sbgemm kernels and renames them to role-based names, so every kernel that consumes these panels migrates in the same commit. Per Michael Tesch's review of the M4 numbers, the per-k-pair zips and predicated loads cost about a third of the single-thread speed. The copies now pack the panels in the (row, k & 1) lane layout that BFMOPA/BFDOT consume directly - element (row r, k) of a block of w rows at [pp * 2w + 2r + (k & 1)] - and SBGEMM_ALIGN_K 2 matches the driver's B-panel stride to the padded pair extent of the panels. The odd k half of the last pair is zero everywhere; the A-side copies additionally zero-pad rows beyond the ragged chunk (full chunk width), while the B-side copies pack their power-of-two groups at the actual width - the SME2 kernel predicates its B load to the group's lanes and the BFDOT kernel runs its vector path on full groups only, with C stores predicated to the valid rows and columns regardless. The SME2 kernel then needs one plain 512-bit load per operand tile and k pair with no zipping (an SME2 multi-vector load could halve that further on the M4); the BFDOT kernel takes two quad loads per operand and k pair and drops its odd-k zero vector. Fallback loops read the same interleaved panels.
18be288 to
e73ab68
Compare
0tomhoyt
changed the title
|
Restructured per review: the branch is now two commits on top of #6142 - the BFDOT kernel alone (consumes #6142's original packing), then a single panel-format-change commit that renames the shared copies to role-based names, moves the k-pair interleave into them, and migrates both kernels together (no layout-mismatched intermediate state; each commit builds standalone). Also picked up from review: predicated B loads (no overread past narrow groups), unused parameter removed, duplicate static dropped, honest packing-contract comment. @tesch1 the SME2 side of the second commit is what needs your M4 harness again when convenient - single plain load per operand tile and k pair now, no zips; panel geometry identical to what you measured. |
|
Thanks for the quick restructure. I retested e73ab68 on the M4 Pro (Apple clang 21):
One thing I noticed, in case it helps: on the SME2 path this head measures slower than bce51f1 here. Square NN, GFLOPS, both heads measured alternately in one session (the machine was busy with other builds, so please read only the comparison; sgemm stayed steady):
A |
The interleaved copies assembled one u32 per iteration with a strided read and two branches, which on the M4 moved the packing to the top of the profile (Michael Tesch's rerun: sbgemm_itcopy 44 -> 985 samples, single-thread sbgemm at n = 1024 down from 1048 to 716 GFLOPS even though the kernel itself got faster, 1610 -> 1111 samples). s2bc_pairs_c now reads the two halves of a k pair as contiguous 16-lane streams along the rows and zips them with one vzip per quad; s2bc_pairs_l loads four pairs of four source rows and writes them in panel order through a 4x4 u32 transpose (vtrn). Odd k tails, ragged rows and the zero padding stay in scalar code. Packing bandwidth on an M3 Pro, k = 1024, 64 rows: pairs_c 12 -> 41 GB/s, pairs_l 14 -> 61 GB/s (memcpy class). Correctness sweeps unchanged (39116 cases, both targets).
|
Thanks - that profile pinpointed it. Vectorized the two packers in 5dea71f: Packing microbenchmark, M3 Pro, k = 1024, 64 rows, read+write bytes (best of 50x100): pairs_c 12 -> 41 GB/s, pairs_l 14 -> 61 GB/s, i.e. most of the 22x profile blowup should be recovered. The full 39116-case sweeps stay green on both targets. Happy to take another spin on your M4 harness whenever convenient - with the kernel already at 1111 samples and packing back near memcpy, the single-thread n = 1024 number should land above the bce51f1 baseline rather than below it. |
This branch has not been deployed
2 participants
Two commits, stacked on #6142 (merge that first; this PR rewrites its branch head, so the diff collapses to just these two commits once it lands):
NEON BFDOT sbgemm kernel for VORTEX.
sbgemmon the Apple targets (M1-M3) so far ran the generic C fallback (~1.7 GFLOPS on an M3 Pro). With FEAT_BF16 (present at least since M2, probed at run time via sysctl; older cores keep the fallback, GCC builds keep the generic kernels) the kernel accumulates C in sixteen quad registers of eight columns by two row groups. WithSBGEMM_UNROLL_M = SBGEMM_UNROLL_Nthe driver routes the A side through the O copies as well, so both panels come in arm64: SME2 bf16 sbgemm kernel for VORTEXM4 (Apple M4) #6142's power-of-two grouping.k-pair interleaved packing (panel format change for both sbgemm kernels, following tesch1's review of the M4 numbers: the per-k-pair zips and predicated loads cost about a third of the single-thread speed). The copies pack the panels in the
(row, k&1)lane layout that BFMOPA/BFDOT consume directly; the odd k half of the last pair is zeroed, the A-side ragged chunk is zero-padded to full width, B groups pack at their actual width with the SME2 kernel predicating its B load (and the BFDOT vector path running on full groups only).SBGEMM_ALIGN_K 2matches the driver's B-panel stride. The SME2 kernel then loads one plain 512-bit vector per operand tile and k pair; the BFDOT kernel two quads per operand and k pair.Measurements (M3 Pro)
sbgemm: 46.7 GFLOPS vs 1.7 on develop (~27x), ~45% of the fp32sgemmof the same build (the scalar fallback would show ~1.5, so this is the vector path executing).Verification
TARGET=VORTEX(vector path),TARGET=VORTEXM4(fallback loops) andDYNAMIC_ARCH(runtime table).test_sbgemmpasses at 1 and 4 threads; the changed files build warning-free.