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.VORTEX b/kernel/arm64/KERNEL.VORTEX index 46a34469c3..46f4989e03 100644 --- a/kernel/arm64/KERNEL.VORTEX +++ b/kernel/arm64/KERNEL.VORTEX @@ -1 +1,12 @@ include $(KERNELDIR)/KERNEL.NEOVERSEN1 + +# NEON BFDOT sbgemm kernels (SBGEMM_UNROLL_M = SBGEMM_UNROLL_N = 8); GCC builds keep the generic kernels +ifneq ($(C_COMPILER), GCC) +SBGEMM_BETA = ../generic/gemm_beta.c +SBGEMMKERNEL = sbgemm_kernel_bfdot.c +# UNROLL_M = UNROLL_N: the driver reuses the O copies for the A side as well +SBGEMMONCOPY = sbgemm_oncopy_sme2.c +SBGEMMOTCOPY = sbgemm_otcopy_sme2.c +SBGEMMONCOPYOBJ = sbgemm_oncopy$(TSUFFIX).$(SUFFIX) +SBGEMMOTCOPYOBJ = sbgemm_otcopy$(TSUFFIX).$(SUFFIX) +endif diff --git a/kernel/arm64/KERNEL.VORTEXM4 b/kernel/arm64/KERNEL.VORTEXM4 index 13a0a38c59..b972ff20e4 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_sme2.c +SBGEMMITCOPY = sbgemm_tcopy_sme2.c +SBGEMMONCOPY = sbgemm_oncopy_sme2.c +SBGEMMOTCOPY = sbgemm_otcopy_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_bfdot.c b/kernel/arm64/sbgemm_kernel_bfdot.c new file mode 100644 index 0000000000..7a26f54685 --- /dev/null +++ b/kernel/arm64/sbgemm_kernel_bfdot.c @@ -0,0 +1,177 @@ +/*************************************************************************** +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. */ +/* + * NEON BFDOT sbgemm kernel for Apple cores (VORTEX): C += alpha * A * B on the packed panels + * of the level-3 driver, the bfloat16 counterpart of the NEON sgemm kernels. + * + * One BFDOT does a two-wide dot product per lane, so the sbgemm copies pack the panels + * k-pair interleaved (element (row r, k) at [pp * 2w + 2r + (k & 1)]) with the padding + * zeroed. Two quad loads per operand and k pair (32 bytes = 8 rows or columns): the pairs of + * rows 0-3 | 4-7 and of columns 0-3 | 4-7. The B pair of a column is duplicated across the + * lanes per BFDOT, and acc[n][h] holds four consecutive rows of C column n - one quad + * load/store per (column, row group). Blocks narrower than the full eight run in scalar code. + * + * Without FEAT_BF16 at run time, or in the compiler, plain loops run on the same layout. + */ + +#include "common.h" +#include +#include +#if defined(__APPLE__) +#include +#endif + +#define S2BD_UM GEMM_DEFAULT_UNROLL_M +#define S2BD_UN GEMM_DEFAULT_UNROLL_N + +#if S2BD_UM != 8 || S2BD_UN != 8 +#error "the BFDOT sbgemm kernel needs SBGEMM_UNROLL_M 8 and SBGEMM_UNROLL_N 8" +#endif + +/* Group width at offset j0 of a count of n: full groups, then the largest power of two left. + With SBGEMM_UNROLL_M = SBGEMM_UNROLL_N the driver routes the A side through the O copies + as well, so both panels come in this grouping. */ +static inline int s2bd_w(BLASLONG n, BLASLONG j0) { + BLASLONG r = n - j0; + if (r >= S2BD_UN) return S2BD_UN; + int w = 1; + while (2 * w <= r) w *= 2; + return w; +} + +static inline FLOAT s2bd_cvt(IFLOAT h) { + unsigned int u = (unsigned int)h << 16; + FLOAT f; + memcpy(&f, &u, sizeof(f)); + return f; +} + +/* FEAT_BF16 at run time; VORTEX is the Apple target, so sysctl decides. */ +static inline int s2bd_usable(void) { + static int ok = -1; + if (ok < 0) { +#if defined(__APPLE__) + int v = 0; + size_t len = sizeof(v); + ok = sysctlbyname("hw.optional.arm.FEAT_BF16", &v, &len, NULL, 0) == 0 && v; +#else + ok = 0; +#endif + } + return ok; +} + +/* Scalar block of wm rows and wn columns, for tails and for the plain-loop fallback. */ +static void s2bd_block(BLASLONG k, const IFLOAT *a, BLASLONG wm, const IFLOAT *b, BLASLONG wn, FLOAT *C, + BLASLONG ldc, FLOAT alpha) { + const BLASLONG pairs = (k + 1) >> 1; + for (BLASLONG j = 0; j < wn; ++j) + for (BLASLONG i = 0; i < wm; ++i) { + FLOAT s = 0; + for (BLASLONG p = 0; p < pairs; ++p) + s += s2bd_cvt(a[p * 2 * wm + 2 * i]) * s2bd_cvt(b[p * 2 * wn + 2 * j]) + + s2bd_cvt(a[p * 2 * wm + 2 * i + 1]) * s2bd_cvt(b[p * 2 * wn + 2 * j + 1]); + C[i + j * ldc] += alpha * s; + } +} + +#if defined(__clang__) +/* The file is built with the NEON flags of its target; the functions here add FEAT_BF16. */ +#pragma clang attribute push(__attribute__((target("bf16"))), apply_to = function) + +/* Full 8x8 block of C in sixteen quad accumulators, one quad per (column, row group). */ +static void s2bd_mac(BLASLONG k, const IFLOAT *a, const IFLOAT *b, FLOAT *C, BLASLONG ldc, FLOAT alpha) { + float32x4_t acc[8][2]; + for (int n = 0; n < 8; ++n) acc[n][0] = acc[n][1] = vdupq_n_f32(0); + const BLASLONG pairs = (k + 1) >> 1; + for (BLASLONG l = 0; l < pairs; ++l) { + const uint16_t *ap = (const uint16_t *)a + l * 2 * S2BD_UM, *bp = (const uint16_t *)b + l * 2 * S2BD_UN; + const uint16x8_t az0 = vld1q_u16(ap); /* rows 0-3 pairs */ + const uint16x8_t az1 = vld1q_u16(ap + 8); /* rows 4-7 pairs */ + const uint32x4_t bp0 = vreinterpretq_u32_u16(vld1q_u16(bp)); /* columns 0-3 pairs */ + const uint32x4_t bp1 = vreinterpretq_u32_u16(vld1q_u16(bp + 8)); /* columns 4-7 pairs */ +#define S2BD_PAIR(N, BP, L) \ + do { \ + const uint16x8_t bn = vreinterpretq_u16_u32(vdupq_n_u32(vgetq_lane_u32(BP, L))); \ + acc[N][0] = vbfdotq_f32(acc[N][0], az0, bn); \ + acc[N][1] = vbfdotq_f32(acc[N][1], az1, bn); \ + } while (0); + S2BD_PAIR(0, bp0, 0) S2BD_PAIR(1, bp0, 1) S2BD_PAIR(2, bp0, 2) S2BD_PAIR(3, bp0, 3) + S2BD_PAIR(4, bp1, 0) S2BD_PAIR(5, bp1, 1) S2BD_PAIR(6, bp1, 2) S2BD_PAIR(7, bp1, 3) +#undef S2BD_PAIR + } + for (int n = 0; n < 8; ++n) + for (int h = 0; h < 2; ++h) { + FLOAT *p = C + n * ldc + 4 * h; + vst1q_f32(p, vfmaq_n_f32(vld1q_f32(p), acc[n][h], alpha)); + } +} + +static void s2bd_drive(BLASLONG m, BLASLONG n, BLASLONG k, FLOAT alpha, const IFLOAT *sa, const IFLOAT *sb, + FLOAT *C, BLASLONG ldc) { + const BLASLONG pairs = (k + 1) >> 1; + for (BLASLONG i0 = 0; i0 < m;) { + const BLASLONG wm = s2bd_w(m, i0); + const IFLOAT *a = sa + i0 * pairs * 2; + for (BLASLONG j0 = 0; j0 < n;) { + const BLASLONG wn = s2bd_w(n, j0); + FLOAT *cb = C + i0 + j0 * ldc; + if (wm == S2BD_UM && wn == S2BD_UN) + s2bd_mac(k, a, sb + j0 * pairs * 2, cb, ldc, alpha); + else + s2bd_block(k, a, wm, sb + j0 * pairs * 2, wn, cb, ldc, alpha); + j0 += wn; + } + i0 += wm; + } +} + +#pragma clang attribute pop +#endif + +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; +#if defined(__clang__) + if (s2bd_usable()) { + s2bd_drive(m, n, k, alpha, ba, bb, C, ldc); + return 0; + } +#endif + const BLASLONG pairs = (k + 1) >> 1; + for (BLASLONG i0 = 0; i0 < m;) { + const BLASLONG wm = s2bd_w(m, i0); + for (BLASLONG j0 = 0; j0 < n;) { + const BLASLONG wn = s2bd_w(n, j0); + s2bd_block(k, ba + i0 * pairs * 2, wm, bb + j0 * pairs * 2, wn, C + i0 + j0 * ldc, ldc, alpha); + j0 += wn; + } + i0 += wm; + } + return 0; +} diff --git a/kernel/arm64/sbgemm_kernel_sme2.c b/kernel/arm64/sbgemm_kernel_sme2.c new file mode 100644 index 0000000000..b7b28308bc --- /dev/null +++ b/kernel/arm64/sbgemm_kernel_sme2.c @@ -0,0 +1,174 @@ +/*************************************************************************** +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. The sbgemm copies pack the panels k-pair interleaved (element (row r, k) of + * a block of w rows at [pp * 2w + 2r + (k & 1)]) with the padding zeroed, and the kernel + * needs one plain load per operand tile and k pair, no zipping. 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 padded to the full width, B panels 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 times the B panel (wn columns) over the k-pair interleaved panels. + The A loads read the zero padding of the full-width chunk; the B load is predicated to the + group's actual 2*wn lanes, so lanes of a narrower group read as zero. */ +S2_NOINL void s2b_mac(const int SUB, BLASLONG k, const IFLOAT *a, const IFLOAT *b, BLASLONG wn) S2_S S2_ZA { + const svbool_t pt = svptrue_b16(); /* BFMOPA governs per 16-bit element */ + const svbool_t pb = svwhilelt_b16((int64_t)0, (int64_t)2 * wn); + const BLASLONG pairs = (k + 1) >> 1; + for (BLASLONG l = 0; l < pairs; ++l) { + const svbfloat16_t bv = svreinterpret_bf16_u16(svld1_u16(pb, (const uint16_t *)b + l * 2 * wn)); + s2b_op(SUB, 0, pt, bv, svreinterpret_bf16_u16(svld1_u16(pt, (const uint16_t *)a + l * 2 * S2B_UM))); + s2b_op(SUB, 1, pt, bv, svreinterpret_bf16_u16(svld1_u16(pt, (const uint16_t *)a + l * 2 * S2B_UM + 32))); + s2b_op(SUB, 2, pt, bv, svreinterpret_bf16_u16(svld1_u16(pt, (const uint16_t *)a + l * 2 * S2B_UM + 64))); + s2b_op(SUB, 3, pt, bv, svreinterpret_bf16_u16(svld1_u16(pt, (const uint16_t *)a + l * 2 * S2B_UM + 96))); + } +} + +/* 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 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, b, wn); + s2_c_store(1, S2_NT, C, ldc, (int)wn, (int)wm); + } else { + svzero_za(); + s2b_mac(0, k, a, 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 BLASLONG pairs = (k + 1) >> 1; + const int sgn = alpha == (FLOAT)1 || alpha == (FLOAT)-1; + 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 * pairs * 2; + for (BLASLONG j0 = 0; j0 < n;) { + const BLASLONG wn = n - j0 >= S2B_UN ? S2B_UN : s2b_wn(n, j0); + s2b_blk(k, a, wm, sb + j0 * pairs * 2, 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) { + const BLASLONG pairs = (k + 1) >> 1; + 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 * pairs * 2; + for (BLASLONG j0 = 0; j0 < n;) { + const BLASLONG wn = n - j0 >= S2B_UN ? S2B_UN : s2b_wn(n, j0); + const IFLOAT *b = sb + j0 * pairs * 2; + for (BLASLONG j = 0; j < wn; ++j) + for (BLASLONG i = 0; i < wm; ++i) { + FLOAT s = 0; + for (BLASLONG p = 0; p < pairs; ++p) + s += s2b_cvt(a[p * 2 * S2B_UM + 2 * i]) * s2b_cvt(b[p * 2 * wn + 2 * j]) + + s2b_cvt(a[p * 2 * S2B_UM + 2 * i + 1]) * s2b_cvt(b[p * 2 * wn + 2 * j + 1]); + 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_sme2.c b/kernel/arm64/sbgemm_ncopy_sme2.c new file mode 100644 index 0000000000..c7f2d7e052 --- /dev/null +++ b/kernel/arm64/sbgemm_ncopy_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_pairs_l(m, S2BC_W, w, a + j0 * lda, lda, b); /* ragged chunk padded to S2BC_W */ + b += 2 * ((m + 1) / 2) * S2BC_W; + } + return 0; +} diff --git a/kernel/arm64/sbgemm_oncopy_sme2.c b/kernel/arm64/sbgemm_oncopy_sme2.c new file mode 100644 index 0000000000..ba90959501 --- /dev/null +++ b/kernel/arm64/sbgemm_oncopy_sme2.c @@ -0,0 +1,42 @@ +/*************************************************************************** +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. */ + +#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_pairs_l(m, w, w, a + j0 * lda, lda, b); + b += 2 * ((m + 1) / 2) * w; + j0 += w; + } + return 0; +} diff --git a/kernel/arm64/sbgemm_otcopy_sme2.c b/kernel/arm64/sbgemm_otcopy_sme2.c new file mode 100644 index 0000000000..37b121b0e2 --- /dev/null +++ b/kernel/arm64/sbgemm_otcopy_sme2.c @@ -0,0 +1,42 @@ +/*************************************************************************** +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. */ + +#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_pairs_c(m, w, w, a + j0, lda, b); + b += 2 * ((m + 1) / 2) * w; + j0 += w; + } + return 0; +} diff --git a/kernel/arm64/sbgemm_tcopy_sme2.c b/kernel/arm64/sbgemm_tcopy_sme2.c new file mode 100644 index 0000000000..ab567fbb8b --- /dev/null +++ b/kernel/arm64/sbgemm_tcopy_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_pairs_c(m, S2BC_W, w, a + j0, lda, b); /* ragged chunk padded to S2BC_W */ + b += 2 * ((m + 1) / 2) * S2BC_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..89df1cc5ce --- /dev/null +++ b/kernel/arm64/sme2_sbgemm_copy.h @@ -0,0 +1,118 @@ +/*************************************************************************** +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. +*****************************************************************************/ + +/* bfloat16 packing of the SME2 and BFDOT sbgemm kernels: k-pair interleaved panels. + BFMOPA/BFDOT consume two k steps per instruction, so the panels hold element + (row r, k) of a block of w rows at [pp * 2w + 2r + (k & 1)], one k pair per + load. The odd k half of the last pair is zeroed by the copies; the A-side + copies additionally zero-pad the ragged chunk to the full chunk width. The + B-side copies pack their power-of-two groups at the actual width: consumers + must not read beyond 2*w lanes of a w-column group (the SME2 kernel + predicates its B load to them, the BFDOT kernel runs its vector path on full + groups only), and C stores are predicated to the valid rows and columns. */ + +#ifndef SME2_SBGEMM_COPY_H +#define SME2_SBGEMM_COPY_H + +#include +#include + +#define S2BC_W GEMM_DEFAULT_UNROLL_M +#define S2BC_WN GEMM_DEFAULT_UNROLL_N + +/* Width of the group at offset j0 of a count of n: full groups, then the largest power of two left. */ +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; +} + +/* Block of m k-steps by w rows into b (w u32 per k pair); element (k, r) at + a[k + r * lda], so pairs of k are contiguous in the source. Rows beyond rows are zero. + The main loop is a 4x4 u32 transpose: each quad load takes four pairs of one source + row, each store writes one pair of four rows, so the pairs land in the panel order. */ +static inline void s2bc_pairs_l(BLASLONG m, BLASLONG w, BLASLONG rows, const IFLOAT *a, BLASLONG lda, IFLOAT *b) { + const unsigned pairs = (unsigned)((m + 1) >> 1), full = (unsigned)(m >> 1); + unsigned *bo = (unsigned *)b; + memset(bo, 0, (size_t)pairs * w * sizeof(unsigned)); + BLASLONG r = 0; + for (; r + 4 <= rows; r += 4) { + const uint16_t *s0 = (const uint16_t *)a + r * lda, *s1 = s0 + lda, *s2 = s1 + lda, *s3 = s2 + lda; + unsigned p = 0; + for (; p + 4 <= full; p += 4) { + const uint32x4_t v0 = vreinterpretq_u32_u16(vld1q_u16(s0 + 2 * p)); + const uint32x4_t v1 = vreinterpretq_u32_u16(vld1q_u16(s1 + 2 * p)); + const uint32x4_t v2 = vreinterpretq_u32_u16(vld1q_u16(s2 + 2 * p)); + const uint32x4_t v3 = vreinterpretq_u32_u16(vld1q_u16(s3 + 2 * p)); + const uint32x4_t t0 = vtrn1q_u32(v0, v1), t1 = vtrn2q_u32(v0, v1); + const uint32x4_t t2 = vtrn1q_u32(v2, v3), t3 = vtrn2q_u32(v2, v3); + vst1q_u32(bo + (size_t)(p + 0) * w + r, vcombine_u32(vget_low_u32(t0), vget_low_u32(t2))); + vst1q_u32(bo + (size_t)(p + 1) * w + r, vcombine_u32(vget_low_u32(t1), vget_low_u32(t3))); + vst1q_u32(bo + (size_t)(p + 2) * w + r, vcombine_u32(vget_high_u32(t0), vget_high_u32(t2))); + vst1q_u32(bo + (size_t)(p + 3) * w + r, vcombine_u32(vget_high_u32(t1), vget_high_u32(t3))); + } + for (; p < full; ++p) + for (BLASLONG q = 0; q < 4; ++q) { + const IFLOAT *s = a + (r + q) * lda + 2 * p; + bo[(size_t)p * w + r + q] = (unsigned)s[0] | ((unsigned)s[1] << 16); + } + } + for (; r < rows; ++r) + for (unsigned p = 0; p < full; ++p) { + const IFLOAT *s = a + r * lda + 2 * p; + bo[(size_t)p * w + r] = (unsigned)s[0] | ((unsigned)s[1] << 16); + } + if (m & 1) + for (BLASLONG q = 0; q < rows; ++q) bo[(size_t)(pairs - 1) * w + q] = (unsigned)a[q * lda + m - 1]; +} + +/* Same, but element (k, r) at a[r + k * lda]: along the rows both halves of a k pair + form contiguous streams, so the main loop zips one quad of each into the panel order. */ +static inline void s2bc_pairs_c(BLASLONG m, BLASLONG w, BLASLONG rows, const IFLOAT *a, BLASLONG lda, IFLOAT *b) { + const unsigned pairs = (unsigned)((m + 1) >> 1), full = (unsigned)(m >> 1); + unsigned *bo = (unsigned *)b; + memset(bo, 0, (size_t)pairs * w * sizeof(unsigned)); + for (unsigned p = 0; p < full; ++p) { + const uint16_t *se = (const uint16_t *)a + (BLASLONG)(2 * p) * lda; + const uint16_t *so = se + lda; + BLASLONG r = 0; + for (; r + 8 <= rows; r += 8) { + const uint16x8_t e = vld1q_u16(se + r), o = vld1q_u16(so + r); + vst1q_u32(bo + (size_t)p * w + r, vreinterpretq_u32_u16(vzip1q_u16(e, o))); + vst1q_u32(bo + (size_t)p * w + r + 4, vreinterpretq_u32_u16(vzip2q_u16(e, o))); + } + for (; r < rows; ++r) bo[(size_t)p * w + r] = (unsigned)se[r] | ((unsigned)so[r] << 16); + } + if (m & 1) { + const uint16_t *se = (const uint16_t *)a + (BLASLONG)(m - 1) * lda; + for (BLASLONG r = 0; r < rows; ++r) bo[(size_t)(pairs - 1) * w + r] = (unsigned)se[r]; + } +} + +#endif diff --git a/param.h b/param.h index 26d7bdc279..576e9de802 100644 --- a/param.h +++ b/param.h @@ -3579,6 +3579,14 @@ 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 +#undef SBGEMM_ALIGN_K +#define SBGEMM_ALIGN_K 2 /* k-pair interleaved panels: the sbgemm copies pad k to even */ + #define SGEMM_DEFAULT_P 1024 #define DGEMM_DEFAULT_P 512 #define CGEMM_DEFAULT_P 256 @@ -4107,6 +4115,16 @@ Until then, just keep it different than DGEMM_DEFAULT_UNROLL_N to keep copy rout #define USE_SSYR2K_KERNEL_DIRECT 1 #endif /* ARMv9 SME */ +/* NEON BFDOT sbgemm kernel (kernel/arm64/sbgemm_kernel_bfdot.c): 8x8 C blocks */ +#ifdef VORTEX +#undef SBGEMM_DEFAULT_UNROLL_M +#undef SBGEMM_DEFAULT_UNROLL_N +#define SBGEMM_DEFAULT_UNROLL_M 8 +#define SBGEMM_DEFAULT_UNROLL_N 8 +#undef SBGEMM_ALIGN_K +#define SBGEMM_ALIGN_K 2 /* k-pair interleaved panels: the sbgemm copies pad k to even */ +#endif + #if defined(ARMV5) #define SNUMOPT 2 #define DNUMOPT 2