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/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/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..5656bb8a7b --- /dev/null +++ b/kernel/arm64/sbgemm_kernel_sme2.c @@ -0,0 +1,179 @@ +/*************************************************************************** +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_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); + 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. */ +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); + 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..46514da0e7 100644 --- a/param.h +++ b/param.h @@ -3579,6 +3579,12 @@ 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 +#define SBGEMM_DEFAULT_UNROLL_M 64 +#define SBGEMM_DEFAULT_UNROLL_N 16 + #define SGEMM_DEFAULT_P 1024 #define DGEMM_DEFAULT_P 512 #define CGEMM_DEFAULT_P 256