From cf5090235d6a85a75ce5abdb0fb534ba4ae1cc9a Mon Sep 17 00:00:00 2001 From: tomhoyt <0tomhoyt@users.noreply.github.com> Date: Tue, 20 Oct 2026 22:54:04 +0800 Subject: [PATCH 1/4] arm64: SME2 bf16 sbgemm kernel for VORTEXM4 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. --- cmake/prebuild.cmake | 4 + kernel/arm64/KERNEL.VORTEXM4 | 11 ++ kernel/arm64/sbgemm_kernel_sme2.c | 180 ++++++++++++++++++++++++++++ kernel/arm64/sbgemm_ncopy_16_sme2.c | 43 +++++++ kernel/arm64/sbgemm_ncopy_64_sme2.c | 41 +++++++ kernel/arm64/sbgemm_tcopy_16_sme2.c | 43 +++++++ kernel/arm64/sbgemm_tcopy_64_sme2.c | 41 +++++++ kernel/arm64/sme2_gemm_detect.h | 17 ++- kernel/arm64/sme2_sbgemm_copy.h | 60 ++++++++++ param.h | 24 ++++ 10 files changed, 462 insertions(+), 2 deletions(-) create mode 100644 kernel/arm64/sbgemm_kernel_sme2.c create mode 100644 kernel/arm64/sbgemm_ncopy_16_sme2.c create mode 100644 kernel/arm64/sbgemm_ncopy_64_sme2.c create mode 100644 kernel/arm64/sbgemm_tcopy_16_sme2.c create mode 100644 kernel/arm64/sbgemm_tcopy_64_sme2.c create mode 100644 kernel/arm64/sme2_sbgemm_copy.h diff --git a/cmake/prebuild.cmake b/cmake/prebuild.cmake index 61ed48c000..b587f04575 100644 --- a/cmake/prebuild.cmake +++ b/cmake/prebuild.cmake @@ -1278,6 +1278,10 @@ endif () set(SGEMM_UNROLL_N 16) set(DGEMM_UNROLL_M 64) set(DGEMM_UNROLL_N 8) + if (BUILD_BFLOAT16) + set(SBGEMM_UNROLL_M 64) + set(SBGEMM_UNROLL_N 16) + endif () else () set(SGEMM_UNROLL_M 16) set(SGEMM_UNROLL_N 4) diff --git a/kernel/arm64/KERNEL.VORTEXM4 b/kernel/arm64/KERNEL.VORTEXM4 index 13a0a38c59..d3f5944028 100644 --- a/kernel/arm64/KERNEL.VORTEXM4 +++ b/kernel/arm64/KERNEL.VORTEXM4 @@ -19,6 +19,17 @@ DGEMMINCOPY = gemm_ncopy_sme2.c DGEMMITCOPY = gemm_tcopy_sme2.c DGEMMONCOPY = ../generic/gemm_ncopy_$(DGEMM_UNROLL_N).c DGEMMOTCOPY = ../generic/gemm_tcopy_$(DGEMM_UNROLL_N).c + +SBGEMM_BETA = ../generic/gemm_beta.c +SBGEMMKERNEL = sbgemm_kernel_sme2.c +SBGEMMINCOPY = sbgemm_ncopy_64_sme2.c +SBGEMMITCOPY = sbgemm_tcopy_64_sme2.c +SBGEMMONCOPY = sbgemm_ncopy_16_sme2.c +SBGEMMOTCOPY = sbgemm_tcopy_16_sme2.c +SBGEMMINCOPYOBJ = sbgemm_incopy$(TSUFFIX).$(SUFFIX) +SBGEMMITCOPYOBJ = sbgemm_itcopy$(TSUFFIX).$(SUFFIX) +SBGEMMONCOPYOBJ = sbgemm_oncopy$(TSUFFIX).$(SUFFIX) +SBGEMMOTCOPYOBJ = sbgemm_otcopy$(TSUFFIX).$(SUFFIX) DGEMMINCOPYOBJ = dgemm_incopy$(TSUFFIX).$(SUFFIX) DGEMMITCOPYOBJ = dgemm_itcopy$(TSUFFIX).$(SUFFIX) DGEMMONCOPYOBJ = dgemm_oncopy$(TSUFFIX).$(SUFFIX) diff --git a/kernel/arm64/sbgemm_kernel_sme2.c b/kernel/arm64/sbgemm_kernel_sme2.c new file mode 100644 index 0000000000..e5159985dc --- /dev/null +++ b/kernel/arm64/sbgemm_kernel_sme2.c @@ -0,0 +1,180 @@ +/*************************************************************************** +Copyright (c) 2026, The OpenBLAS Project +All rights reserved. +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are +met: +1. Redistributions of source code must retain the above copyright +notice, this list of conditions and the following disclaimer. +2. Redistributions in binary form must reproduce the above copyright +notice, this list of conditions and the following disclaimer in +the documentation and/or other materials provided with the +distribution. +3. Neither the name of the OpenBLAS project nor the names of +its contributors may be used to endorse or promote products +derived from this software without specific prior written permission. +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +ARE DISCLAIMED. IN NO EVENT SHALL THE OPENBLAS PROJECT OR CONTRIBUTORS BE +LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) +HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT +LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF +THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +*****************************************************************************/ + +/* A panels of the SME2 kernels (sme2_gemm_packed.h) from op(A) with contiguous columns. */ +/* + * SME2 bf16 GEMM kernel for a 512-bit streaming vector length, the bfloat16 counterpart of + * gemm_kernel_sme2.c: C += alpha * A * B on the packed panels of the level-3 driver. + * + * BFMOPA ZA32 with .H operands accumulates 2-wide dot products, so one instruction covers + * two k steps: per k pair the kernel zips rows l and l+1 of the standard [k][m] and [k][n] + * panels into the lane layout 2*i + (k & 1) that the outer product needs (a missing second + * row is loaded under a false predicate, which zero-fills). C accumulates in four ZA32 + * tiles of 16x16 fp32: tile rows are the wn columns of the block, tile columns the wm rows, + * with the C load/store helpers of sme2_gemm_tile.h. + * + * A panels come in chunks of SBGEMM_UNROLL_M rows ([k][64], the last chunk narrower), the + * B panels of the generic copies in power-of-two groups of SBGEMM_UNROLL_N columns. Without + * SME2 at run time, or in the compiler, plain loops run on the same layout. + */ + +#include "common.h" +#include + +#include "sme2_gemm_detect.h" + +#define S2B_UM GEMM_DEFAULT_UNROLL_M +#define S2B_UN GEMM_DEFAULT_UNROLL_N + +#if S2B_UM != 64 || S2B_UN != 16 +#error "the SME2 sbgemm kernel needs SBGEMM_UNROLL_M 64 and SBGEMM_UNROLL_N 16" +#endif + +/* Width of the B panel group at column j0 of n: full groups, then the largest power of two left. */ +static inline int s2b_wn(BLASLONG n, BLASLONG j0) { + BLASLONG r = n - j0; + if (r >= S2B_UN) return S2B_UN; + int w = 1; + while (2 * w <= r) w *= 2; + return w; +} + +static inline FLOAT s2b_cvt(IFLOAT h) { + unsigned int u = (unsigned int)h << 16; + FLOAT f; + memcpy(&f, &u, sizeof(f)); + return f; +} + +#ifdef HAVE_SME2_GEMM +#include "sme2_gemm_tile.h" + +#if S2_VL != 16 || S2_NR != 64 || S2_NT != 4 +#error "the SME2 sbgemm kernel needs a 512-bit streaming vector length" +#endif + +/* The file is built with the SME flags of its target (no SME2); the functions here enable SME2 themselves. */ +#pragma clang attribute push(__attribute__((target("sme2"))), apply_to = function) + +S2_INL void s2b_op(const int SUB, int t, svbool_t p, svbfloat16_t b, svbfloat16_t a) S2_S S2_ZA { + switch (t) { + case 0: SUB ? svmops_za32_bf16_m(0, p, p, b, a) : svmopa_za32_bf16_m(0, p, p, b, a); break; + case 1: SUB ? svmops_za32_bf16_m(1, p, p, b, a) : svmopa_za32_bf16_m(1, p, p, b, a); break; + case 2: SUB ? svmops_za32_bf16_m(2, p, p, b, a) : svmopa_za32_bf16_m(2, p, p, b, a); break; + default: SUB ? svmops_za32_bf16_m(3, p, p, b, a) : svmopa_za32_bf16_m(3, p, p, b, a); break; + } +} + +/* ZA += or -= the A panel (wm rows) times the B panel (wn columns) over k steps of the packed layout. */ +S2_NOINL void s2b_mac(const int SUB, BLASLONG k, const IFLOAT *a, BLASLONG wm, const IFLOAT *b, BLASLONG wn) S2_S + S2_ZA { + const svbool_t pt = svptrue_b32(); + const svbool_t pm0 = svwhilelt_b16_s32(0, (int32_t)wm); + const svbool_t pm1 = svwhilelt_b16_s32(S2B_UN * 2, (int32_t)wm); + const svbool_t pn = svwhilelt_b16_s32(0, (int32_t)wn); + const svbool_t pf = svpfalse_b(); + for (BLASLONG l = 0; l < k; l += 2) { + const int two = l + 1 < k; + const IFLOAT *ar = a + l * wm, *br = b + l * wn; + const svuint16_t e0 = svld1_u16(pm0, ar), e1 = svld1_u16(pm1, ar + S2B_UN * 2); + const svuint16_t o0 = svld1_u16(two ? pm0 : pf, ar + wm), o1 = svld1_u16(two ? pm1 : pf, ar + wm + S2B_UN * 2); + const svbfloat16_t a0 = svreinterpret_bf16_u16(svzip1_u16(e0, o0)); + const svbfloat16_t a1 = svreinterpret_bf16_u16(svzip2_u16(e0, o0)); + const svbfloat16_t a2 = svreinterpret_bf16_u16(svzip1_u16(e1, o1)); + const svbfloat16_t a3 = svreinterpret_bf16_u16(svzip2_u16(e1, o1)); + const svbfloat16_t bv = svreinterpret_bf16_u16(svzip1_u16(svld1_u16(pn, br), svld1_u16(two ? pn : pf, br + wn))); + s2b_op(SUB, 0, pt, bv, a0); + s2b_op(SUB, 1, pt, bv, a1); + s2b_op(SUB, 2, pt, bv, a2); + s2b_op(SUB, 3, pt, bv, a3); + } +} + +/* Block of wm rows and wn columns of C; sgn 1 for C += alpha AB with alpha = +-1 (the AB part + added or subtracted in ZA), sgn 0 for a general alpha scaled on the way out of ZA. */ +__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) { + if (sgn) { + s2_c_load(1, S2_NT, C, ldc, (int)wn, (int)wm, S2_LOAD, (FLOAT)1); + s2b_mac(alpha == (FLOAT)-1, k, a, wm, b, wn); + s2_c_store(1, S2_NT, C, ldc, (int)wn, (int)wm); + } else { + svzero_za(); + s2b_mac(0, k, a, wm, b, wn); + s2_c_store_alpha(S2_NT, C, ldc, (int)wn, (int)wm, alpha, 1); + } +} + +__arm_new("za") __arm_locally_streaming static void s2b_drive(BLASLONG m, BLASLONG n, BLASLONG k, FLOAT alpha, + const IFLOAT *sa, const IFLOAT *sb, FLOAT *C, + BLASLONG ldc) { + const int sgn = alpha == (FLOAT)1 || alpha == (FLOAT)-1; + for (BLASLONG j0 = 0; j0 < n;) { + const BLASLONG wn = n - j0 >= S2B_UN ? S2B_UN : s2b_wn(n, j0); + const IFLOAT *b = sb + j0 * k; + for (BLASLONG i0 = 0; i0 < m; i0 += S2B_UM) { + const BLASLONG wm = m - i0 < S2B_UM ? m - i0 : S2B_UM; + s2b_blk(k, sa + i0 * k, wm, b, wn, C + i0 + j0 * ldc, ldc, alpha, sgn); + } + j0 += wn; + } +} + +#pragma clang attribute pop +#endif + +static void s2b_loops(BLASLONG m, BLASLONG n, BLASLONG k, FLOAT alpha, const IFLOAT *sa, const IFLOAT *sb, FLOAT *C, + BLASLONG ldc) { + for (BLASLONG j0 = 0; j0 < n;) { + const BLASLONG wn = s2b_wn(n, j0); + const IFLOAT *b = sb + j0 * k; + for (BLASLONG i0 = 0; i0 < m; i0 += S2B_UM) { + const BLASLONG wm = m - i0 < S2B_UM ? m - i0 : S2B_UM; + const IFLOAT *a = sa + i0 * k; + for (BLASLONG j = 0; j < wn; ++j) + for (BLASLONG i = 0; i < wm; ++i) { + FLOAT s = 0; + for (BLASLONG l = 0; l < k; ++l) s += s2b_cvt(a[l * wm + i]) * s2b_cvt(b[l * wn + j]); + C[i0 + i + (j0 + j) * ldc] += alpha * s; + } + } + j0 += wn; + } +} + +int CNAME(BLASLONG m, BLASLONG n, BLASLONG k, FLOAT alpha, IFLOAT *ba, IFLOAT *bb, FLOAT *C, BLASLONG ldc) { + if (m <= 0 || n <= 0 || k <= 0) return 0; +#ifdef HAVE_SME2_GEMM + if (s2_usable()) { + s2b_drive(m, n, k, alpha, ba, bb, C, ldc); + return 0; + } +#endif + s2b_loops(m, n, k, alpha, ba, bb, C, ldc); + return 0; +} diff --git a/kernel/arm64/sbgemm_ncopy_16_sme2.c b/kernel/arm64/sbgemm_ncopy_16_sme2.c new file mode 100644 index 0000000000..84afe06a61 --- /dev/null +++ b/kernel/arm64/sbgemm_ncopy_16_sme2.c @@ -0,0 +1,43 @@ +/*************************************************************************** +Copyright (c) 2026, The OpenBLAS Project +All rights reserved. +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are +met: +1. Redistributions of source code must retain the above copyright +notice, this list of conditions and the following disclaimer. +2. Redistributions in binary form must reproduce the above copyright +notice, this list of conditions and the following disclaimer in +the documentation and/or other materials provided with the +distribution. +3. Neither the name of the OpenBLAS project nor the names of +its contributors may be used to endorse or promote products +derived from this software without specific prior written permission. +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +ARE DISCLAIMED. IN NO EVENT SHALL THE OPENBLAS PROJECT OR CONTRIBUTORS BE +LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) +HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT +LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF +THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +*****************************************************************************/ + +/* A panels of the SME2 kernels (sme2_gemm_packed.h) from op(A) with contiguous columns. */ +/* bfloat16 B panels of the SME2 sbgemm kernel from op(B) with contiguous columns, + in [k][w] groups like the A panels (the generic ncopy packs column-major groups). */ + +#include "common.h" +#include "sme2_sbgemm_copy.h" + +int CNAME(BLASLONG m, BLASLONG n, IFLOAT *a, BLASLONG lda, IFLOAT *b) { + for (BLASLONG j0 = 0; j0 < n;) { + const BLASLONG w = s2bc_wn(n, j0); + s2bc_cols(m, w, a + j0 * lda, lda, b); + b += m * w; + j0 += w; + } + return 0; +} diff --git a/kernel/arm64/sbgemm_ncopy_64_sme2.c b/kernel/arm64/sbgemm_ncopy_64_sme2.c new file mode 100644 index 0000000000..5a8ca0503e --- /dev/null +++ b/kernel/arm64/sbgemm_ncopy_64_sme2.c @@ -0,0 +1,41 @@ +/*************************************************************************** +Copyright (c) 2026, The OpenBLAS Project +All rights reserved. +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are +met: +1. Redistributions of source code must retain the above copyright +notice, this list of conditions and the following disclaimer. +2. Redistributions in binary form must reproduce the above copyright +notice, this list of conditions and the following disclaimer in +the documentation and/or other materials provided with the +distribution. +3. Neither the name of the OpenBLAS project nor the names of +its contributors may be used to endorse or promote products +derived from this software without specific prior written permission. +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +ARE DISCLAIMED. IN NO EVENT SHALL THE OPENBLAS PROJECT OR CONTRIBUTORS BE +LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) +HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT +LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF +THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +*****************************************************************************/ + +/* A panels of the SME2 kernels (sme2_gemm_packed.h) from op(A) with contiguous columns. */ +/* bfloat16 A panels of the SME2 sbgemm kernel from op(A) with contiguous columns. */ + +#include "common.h" +#include "sme2_sbgemm_copy.h" + +int CNAME(BLASLONG m, BLASLONG n, IFLOAT *a, BLASLONG lda, IFLOAT *b) { + for (BLASLONG j0 = 0; j0 < n; j0 += S2BC_W) { + const BLASLONG w = n - j0 < S2BC_W ? n - j0 : S2BC_W; + s2bc_cols(m, w, a + j0 * lda, lda, b); + b += m * w; + } + return 0; +} diff --git a/kernel/arm64/sbgemm_tcopy_16_sme2.c b/kernel/arm64/sbgemm_tcopy_16_sme2.c new file mode 100644 index 0000000000..27d4db82fa --- /dev/null +++ b/kernel/arm64/sbgemm_tcopy_16_sme2.c @@ -0,0 +1,43 @@ +/*************************************************************************** +Copyright (c) 2026, The OpenBLAS Project +All rights reserved. +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are +met: +1. Redistributions of source code must retain the above copyright +notice, this list of conditions and the following disclaimer. +2. Redistributions in binary form must reproduce the above copyright +notice, this list of conditions and the following disclaimer in +the documentation and/or other materials provided with the +distribution. +3. Neither the name of the OpenBLAS project nor the names of +its contributors may be used to endorse or promote products +derived from this software without specific prior written permission. +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +ARE DISCLAIMED. IN NO EVENT SHALL THE OPENBLAS PROJECT OR CONTRIBUTORS BE +LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) +HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT +LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF +THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +*****************************************************************************/ + +/* A panels of the SME2 kernels (sme2_gemm_packed.h) from op(A) with contiguous columns. */ +/* bfloat16 B panels of the SME2 sbgemm kernel from op(B) with contiguous rows, + in [k][w] groups like the A panels (the generic tcopy packs column-major groups). */ + +#include "common.h" +#include "sme2_sbgemm_copy.h" + +int CNAME(BLASLONG m, BLASLONG n, IFLOAT *a, BLASLONG lda, IFLOAT *b) { + for (BLASLONG j0 = 0; j0 < n;) { + const BLASLONG w = s2bc_wn(n, j0); + s2bc_rows(m, w, a + j0, lda, b); + b += m * w; + j0 += w; + } + return 0; +} diff --git a/kernel/arm64/sbgemm_tcopy_64_sme2.c b/kernel/arm64/sbgemm_tcopy_64_sme2.c new file mode 100644 index 0000000000..7a41716b9a --- /dev/null +++ b/kernel/arm64/sbgemm_tcopy_64_sme2.c @@ -0,0 +1,41 @@ +/*************************************************************************** +Copyright (c) 2026, The OpenBLAS Project +All rights reserved. +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are +met: +1. Redistributions of source code must retain the above copyright +notice, this list of conditions and the following disclaimer. +2. Redistributions in binary form must reproduce the above copyright +notice, this list of conditions and the following disclaimer in +the documentation and/or other materials provided with the +distribution. +3. Neither the name of the OpenBLAS project nor the names of +its contributors may be used to endorse or promote products +derived from this software without specific prior written permission. +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +ARE DISCLAIMED. IN NO EVENT SHALL THE OPENBLAS PROJECT OR CONTRIBUTORS BE +LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) +HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT +LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF +THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +*****************************************************************************/ + +/* A panels of the SME2 kernels (sme2_gemm_packed.h) from op(A) with contiguous columns. */ +/* bfloat16 A panels of the SME2 sbgemm kernel from op(A) with contiguous rows. */ + +#include "common.h" +#include "sme2_sbgemm_copy.h" + +int CNAME(BLASLONG m, BLASLONG n, IFLOAT *a, BLASLONG lda, IFLOAT *b) { + for (BLASLONG j0 = 0; j0 < n; j0 += S2BC_W) { + const BLASLONG w = n - j0 < S2BC_W ? n - j0 : S2BC_W; + s2bc_rows(m, w, a + j0, lda, b); + b += m * w; + } + return 0; +} diff --git a/kernel/arm64/sme2_gemm_detect.h b/kernel/arm64/sme2_gemm_detect.h index 8830e7dd32..474955ed64 100644 --- a/kernel/arm64/sme2_gemm_detect.h +++ b/kernel/arm64/sme2_gemm_detect.h @@ -41,7 +41,20 @@ THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. #include #endif -/* SME2 with a 512-bit streaming vector length (the kernels assume 16 fp32 lanes). */ +/* SME2 with a 512-bit streaming vector length (the kernels assume 16 fp32 lanes). + The vector length is read through inline asm: as a pure intrinsic svcntsw() can be + speculated above the FEAT_SME2 test (clang if-converts the conjunction), and RDSVL + is undefined on cores without FEAT_SME. */ +/* Noinline, so that clang cannot if-convert the call above the FEAT_SME2 test. */ +#if defined(__GNUC__) || defined(__clang__) +__attribute__((noinline)) +#endif +static int s2_sme2_svl_lanes(void) { + uint64_t svl; + __asm__ volatile("rdsvl %0, #1" : "=r"(svl)); /* streaming vector length in bytes */ + return (int)(svl / 4); +} + static inline int s2_usable(void) { static int ok = -1; if (ok < 0) { @@ -53,7 +66,7 @@ static inline int s2_usable(void) { #elif defined(__linux__) sme2 = (getauxval(AT_HWCAP2) & (1UL << 37)) != 0; /* HWCAP2_SME2 */ #endif - ok = sme2 && svcntsw() == 16; + ok = sme2 && s2_sme2_svl_lanes() == 16; } return ok; } diff --git a/kernel/arm64/sme2_sbgemm_copy.h b/kernel/arm64/sme2_sbgemm_copy.h new file mode 100644 index 0000000000..b76b31e614 --- /dev/null +++ b/kernel/arm64/sme2_sbgemm_copy.h @@ -0,0 +1,60 @@ +/*************************************************************************** +Copyright (c) 2026, The OpenBLAS Project +All rights reserved. +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are +met: +1. Redistributions of source code must retain the above copyright +notice, this list of conditions and the following disclaimer. +2. Redistributions in binary form must reproduce the above copyright +notice, this list of conditions and the following disclaimer in +the documentation and/or other materials provided with the +distribution. +3. Neither the name of the OpenBLAS project nor the names of +its contributors may be used to endorse or promote products +derived from this software without specific prior written permission. +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +ARE DISCLAIMED. IN NO EVENT SHALL THE OPENBLAS PROJECT OR CONTRIBUTORS BE +LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) +HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT +LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF +THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +*****************************************************************************/ + +/* A panels of the SME2 kernels (sme2_gemm_packed.h) from op(A) with contiguous columns. */ +/* bfloat16 A panels of the SME2 sbgemm kernel (sbgemm_kernel_sme2.c), the same + [k][m] layout in chunks of GEMM_UNROLL_M rows as gemm_ncopy_sme2.c. */ + +#ifndef SME2_SBGEMM_COPY_H +#define SME2_SBGEMM_COPY_H + +#include + +#define S2BC_W GEMM_DEFAULT_UNROLL_M +#define S2BC_WN GEMM_DEFAULT_UNROLL_N + +/* Element (i, c) at a[c + i * lda], bfloat16: a plain transpose for now. */ +static inline void s2bc_cols(BLASLONG m, BLASLONG w, const IFLOAT *a, BLASLONG lda, IFLOAT *b) { + for (BLASLONG i = 0; i < m; ++i) + for (BLASLONG c = 0; c < w; ++c) b[i * w + c] = a[i + c * lda]; +} + +/* Element (i, c) at a[i * lda + c], bfloat16: contiguous rows. */ +static inline void s2bc_rows(BLASLONG m, BLASLONG w, const IFLOAT *a, BLASLONG lda, IFLOAT *b) { + for (BLASLONG i = 0; i < m; ++i) memcpy(b + i * w, a + i * lda, (size_t)w * sizeof(IFLOAT)); +} + +/* Width of the group at column j0 of n for the B copies: full groups of S2BC_N, then powers of two. */ +static inline BLASLONG s2bc_wn(BLASLONG n, BLASLONG j0) { + BLASLONG r = n - j0; + if (r >= S2BC_WN) return S2BC_WN; + BLASLONG w = 1; + while (2 * w <= r) w *= 2; + return w; +} + +#endif diff --git a/param.h b/param.h index 26d7bdc279..05145aea33 100644 --- a/param.h +++ b/param.h @@ -3285,6 +3285,9 @@ USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. #undef SBGEMM_DEFAULT_UNROLL_M #undef SBGEMM_DEFAULT_UNROLL_N +#undef SBGEMM_DEFAULT_P +#undef SBGEMM_DEFAULT_Q +#undef SBGEMM_DEFAULT_R #define SBGEMM_DEFAULT_UNROLL_M 8 #define SBGEMM_DEFAULT_UNROLL_N 8 @@ -3346,6 +3349,9 @@ USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. #undef SBGEMM_DEFAULT_UNROLL_M #undef SBGEMM_DEFAULT_UNROLL_N +#undef SBGEMM_DEFAULT_P +#undef SBGEMM_DEFAULT_Q +#undef SBGEMM_DEFAULT_R #define SBGEMM_DEFAULT_UNROLL_M 16 #define SBGEMM_DEFAULT_UNROLL_N 8 @@ -3579,17 +3585,29 @@ USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. #define ZGEMM_DEFAULT_UNROLL_M 4 #define ZGEMM_DEFAULT_UNROLL_N 4 +/* SME2 sbgemm kernel (kernel/arm64/sbgemm_kernel_sme2.c): four ZA32 tiles of 16x16 fp32 */ +#undef SBGEMM_DEFAULT_UNROLL_M +#undef SBGEMM_DEFAULT_UNROLL_N +#undef SBGEMM_DEFAULT_P +#undef SBGEMM_DEFAULT_Q +#undef SBGEMM_DEFAULT_R +#define SBGEMM_DEFAULT_UNROLL_M 64 +#define SBGEMM_DEFAULT_UNROLL_N 16 + #define SGEMM_DEFAULT_P 1024 +#define SBGEMM_DEFAULT_P 1024 #define DGEMM_DEFAULT_P 512 #define CGEMM_DEFAULT_P 256 #define ZGEMM_DEFAULT_P 128 #define SGEMM_DEFAULT_Q 1024 +#define SBGEMM_DEFAULT_Q 1024 #define DGEMM_DEFAULT_Q 512 #define CGEMM_DEFAULT_Q 512 #define ZGEMM_DEFAULT_Q 512 #define SGEMM_DEFAULT_R 4096 +#define SBGEMM_DEFAULT_R 4096 #define DGEMM_DEFAULT_R 4096 #define CGEMM_DEFAULT_R 4096 #define ZGEMM_DEFAULT_R 2048 @@ -3848,6 +3866,9 @@ is a big desktop or server with abundant cache rather than a phone or embedded d #undef SBGEMM_ALIGN_K #undef SBGEMM_DEFAULT_UNROLL_M #undef SBGEMM_DEFAULT_UNROLL_N +#undef SBGEMM_DEFAULT_P +#undef SBGEMM_DEFAULT_Q +#undef SBGEMM_DEFAULT_R #define SBGEMM_ALIGN_K 4 #define SBGEMM_DEFAULT_UNROLL_M 8 #define SBGEMM_DEFAULT_UNROLL_N 4 @@ -3906,6 +3927,9 @@ is a big desktop or server with abundant cache rather than a phone or embedded d #undef SBGEMM_ALIGN_K #undef SBGEMM_DEFAULT_UNROLL_M #undef SBGEMM_DEFAULT_UNROLL_N +#undef SBGEMM_DEFAULT_P +#undef SBGEMM_DEFAULT_Q +#undef SBGEMM_DEFAULT_R #define SBGEMM_ALIGN_K 4 #define SBGEMM_DEFAULT_UNROLL_M 8 #define SBGEMM_DEFAULT_UNROLL_N 8 From 0bc31e8075df540e4fe785a1030d871f4b6cb918 Mon Sep 17 00:00:00 2001 From: tomhoyt <0tomhoyt@users.noreply.github.com> Date: Wed, 21 Oct 2026 09:15:57 +0800 Subject: [PATCH 2/4] arm64: keep the shared SBGEMM P/Q/R defaults for VORTEXM4 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. --- param.h | 6 ------ 1 file changed, 6 deletions(-) diff --git a/param.h b/param.h index 05145aea33..98bddcc6f8 100644 --- a/param.h +++ b/param.h @@ -3588,26 +3588,20 @@ USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. /* SME2 sbgemm kernel (kernel/arm64/sbgemm_kernel_sme2.c): four ZA32 tiles of 16x16 fp32 */ #undef SBGEMM_DEFAULT_UNROLL_M #undef SBGEMM_DEFAULT_UNROLL_N -#undef SBGEMM_DEFAULT_P -#undef SBGEMM_DEFAULT_Q -#undef SBGEMM_DEFAULT_R #define SBGEMM_DEFAULT_UNROLL_M 64 #define SBGEMM_DEFAULT_UNROLL_N 16 #define SGEMM_DEFAULT_P 1024 -#define SBGEMM_DEFAULT_P 1024 #define DGEMM_DEFAULT_P 512 #define CGEMM_DEFAULT_P 256 #define ZGEMM_DEFAULT_P 128 #define SGEMM_DEFAULT_Q 1024 -#define SBGEMM_DEFAULT_Q 1024 #define DGEMM_DEFAULT_Q 512 #define CGEMM_DEFAULT_Q 512 #define ZGEMM_DEFAULT_Q 512 #define SGEMM_DEFAULT_R 4096 -#define SBGEMM_DEFAULT_R 4096 #define DGEMM_DEFAULT_R 4096 #define CGEMM_DEFAULT_R 4096 #define ZGEMM_DEFAULT_R 2048 From c151d85b488b1528cf7ee0f43aefbc0ab4a557fc Mon Sep 17 00:00:00 2001 From: tomhoyt <0tomhoyt@users.noreply.github.com> Date: Wed, 21 Oct 2026 14:54:48 +0800 Subject: [PATCH 3/4] arm64: drop stray SBGEMM P/Q/R undefs picked up by other cpu blocks 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. --- param.h | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/param.h b/param.h index 98bddcc6f8..46514da0e7 100644 --- a/param.h +++ b/param.h @@ -3285,9 +3285,6 @@ USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. #undef SBGEMM_DEFAULT_UNROLL_M #undef SBGEMM_DEFAULT_UNROLL_N -#undef SBGEMM_DEFAULT_P -#undef SBGEMM_DEFAULT_Q -#undef SBGEMM_DEFAULT_R #define SBGEMM_DEFAULT_UNROLL_M 8 #define SBGEMM_DEFAULT_UNROLL_N 8 @@ -3349,9 +3346,6 @@ USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. #undef SBGEMM_DEFAULT_UNROLL_M #undef SBGEMM_DEFAULT_UNROLL_N -#undef SBGEMM_DEFAULT_P -#undef SBGEMM_DEFAULT_Q -#undef SBGEMM_DEFAULT_R #define SBGEMM_DEFAULT_UNROLL_M 16 #define SBGEMM_DEFAULT_UNROLL_N 8 @@ -3860,9 +3854,6 @@ is a big desktop or server with abundant cache rather than a phone or embedded d #undef SBGEMM_ALIGN_K #undef SBGEMM_DEFAULT_UNROLL_M #undef SBGEMM_DEFAULT_UNROLL_N -#undef SBGEMM_DEFAULT_P -#undef SBGEMM_DEFAULT_Q -#undef SBGEMM_DEFAULT_R #define SBGEMM_ALIGN_K 4 #define SBGEMM_DEFAULT_UNROLL_M 8 #define SBGEMM_DEFAULT_UNROLL_N 4 @@ -3921,9 +3912,6 @@ is a big desktop or server with abundant cache rather than a phone or embedded d #undef SBGEMM_ALIGN_K #undef SBGEMM_DEFAULT_UNROLL_M #undef SBGEMM_DEFAULT_UNROLL_N -#undef SBGEMM_DEFAULT_P -#undef SBGEMM_DEFAULT_Q -#undef SBGEMM_DEFAULT_R #define SBGEMM_ALIGN_K 4 #define SBGEMM_DEFAULT_UNROLL_M 8 #define SBGEMM_DEFAULT_UNROLL_N 8 From bce51f1f02c882a2b25ebcdc0326b1585a3efeed Mon Sep 17 00:00:00 2001 From: tomhoyt <0tomhoyt@users.noreply.github.com> Date: Thu, 22 Oct 2026 08:11:54 +0800 Subject: [PATCH 4/4] arm64: fix the SME2 sbgemm kernel's ZA handoff and BFMOPA predicates 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. --- common_thread.h | 2 +- kernel/arm64/sbgemm_kernel_sme2.c | 7 +++---- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/common_thread.h b/common_thread.h index 8317d07941..e9e2776596 100644 --- a/common_thread.h +++ b/common_thread.h @@ -157,7 +157,7 @@ extern int blas_omp_threads_local; extern int blas_is_num_threads_set_explicitly; #if defined(ARCH_ARM64) && defined(OS_DARWIN) && defined(__clang__) && (defined(VORTEXM4) || defined(DYNAMIC_ARCH)) && \ - !defined(COMPLEX) && !defined(BFLOAT16) && !defined(HFLOAT16) + !defined(COMPLEX) && (!defined(BFLOAT16) || !defined(BGEMM)) && !defined(HFLOAT16) #include #ifdef DYNAMIC_ARCH extern char *gotoblas_corename(void); diff --git a/kernel/arm64/sbgemm_kernel_sme2.c b/kernel/arm64/sbgemm_kernel_sme2.c index e5159985dc..5656bb8a7b 100644 --- a/kernel/arm64/sbgemm_kernel_sme2.c +++ b/kernel/arm64/sbgemm_kernel_sme2.c @@ -92,7 +92,7 @@ S2_INL void s2b_op(const int SUB, int t, svbool_t p, svbfloat16_t b, svbfloat16_ /* ZA += or -= the A panel (wm rows) times the B panel (wn columns) over k steps of the packed layout. */ S2_NOINL void s2b_mac(const int SUB, BLASLONG k, const IFLOAT *a, BLASLONG wm, const IFLOAT *b, BLASLONG wn) S2_S S2_ZA { - const svbool_t pt = svptrue_b32(); + const svbool_t pt = svptrue_b16(); /* BFMOPA governs per 16-bit element */ const svbool_t pm0 = svwhilelt_b16_s32(0, (int32_t)wm); const svbool_t pm1 = svwhilelt_b16_s32(S2B_UN * 2, (int32_t)wm); const svbool_t pn = svwhilelt_b16_s32(0, (int32_t)wn); @@ -116,9 +116,8 @@ S2_NOINL void s2b_mac(const int SUB, BLASLONG k, const IFLOAT *a, BLASLONG wm, c /* Block of wm rows and wn columns of C; sgn 1 for C += alpha AB with alpha = +-1 (the AB part added or subtracted in ZA), sgn 0 for a general alpha scaled on the way out of ZA. */ -__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 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_S S2_ZA { if (sgn) { s2_c_load(1, S2_NT, C, ldc, (int)wn, (int)wm, S2_LOAD, (FLOAT)1); s2b_mac(alpha == (FLOAT)-1, k, a, wm, b, wn);