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/7] 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/7] 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/7] 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/7] 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); From 37c89676226ed20dae885b55ca6c739e3c977337 Mon Sep 17 00:00:00 2001 From: tomhoyt <0tomhoyt@users.noreply.github.com> Date: Thu, 22 Oct 2026 11:12:41 +0800 Subject: [PATCH 5/7] arm64: NEON BFDOT sbgemm kernel for VORTEX (Apple M1-M3) sbgemm on the Apple targets so far ran the generic C fallback; with FEAT_BF16 (present at least since M2, probed at run time so older cores keep the fallback) the new kernel accumulates C in sixteen quad registers of eight columns by two row groups. One BFDOT does a two-wide dot per lane, so rows l and l+1 of the packed panels are zipped into the lane layout (i, k & 1); the B pair of a column is duplicated across the lanes and the A pairs cover four rows each. An odd k is a zero vector, blocks narrower than eight run in scalar code. With SBGEMM_UNROLL_M = SBGEMM_UNROLL_N the driver routes the A side through the O copies as well, so both panels come in the power-of-two grouping of #6142's copies. On an M3 Pro: 39116-case randomized sbgemm sweep against a double-precision reference passes unchanged, and against the generic fallback of develop the packed path is 34x at 1024^3 (57 vs 1.7 GFLOPS) and 63x at 2048x2048x1024 (114 vs 1.8 GFLOPS), about a quarter of the fp32 NEON sgemm of the same build. --- kernel/arm64/KERNEL.VORTEX | 11 ++ kernel/arm64/sbgemm_kernel_bfdot.c | 176 +++++++++++++++++++++++++++++ param.h | 8 ++ 3 files changed, 195 insertions(+) create mode 100644 kernel/arm64/sbgemm_kernel_bfdot.c diff --git a/kernel/arm64/KERNEL.VORTEX b/kernel/arm64/KERNEL.VORTEX index 46a34469c3..59df87d6cc 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_ncopy_16_sme2.c +SBGEMMOTCOPY = sbgemm_tcopy_16_sme2.c +SBGEMMONCOPYOBJ = sbgemm_oncopy$(TSUFFIX).$(SUFFIX) +SBGEMMOTCOPYOBJ = sbgemm_otcopy$(TSUFFIX).$(SUFFIX) +endif diff --git a/kernel/arm64/sbgemm_kernel_bfdot.c b/kernel/arm64/sbgemm_kernel_bfdot.c new file mode 100644 index 0000000000..dcb4c6a42d --- /dev/null +++ b/kernel/arm64/sbgemm_kernel_bfdot.c @@ -0,0 +1,176 @@ +/*************************************************************************** +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 kernel zips rows l and l+1 of the + * standard [k][m] and [k][n] panels into the lane layout (i, k & 1): the B pair of a column + * is duplicated across the lanes, the A pairs cover four rows each, and acc[n][h] holds four + * consecutive rows of C column n - one quad load/store per (column, row group). A missing + * second row of an odd k is a zero vector, and blocks narrower than eight rows or columns 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) { + 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 += s2bd_cvt(a[l * wm + i]) * s2bd_cvt(b[l * wn + j]); + 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); + static const uint16_t zero[8]; /* the missing row of an odd k */ + for (BLASLONG l = 0; l < k; l += 2) { + const uint16x8_t ak0 = vld1q_u16((const uint16_t *)a + l * 8); + const uint16x8_t ak1 = vld1q_u16(l + 1 < k ? (const uint16_t *)a + (l + 1) * 8 : zero); + const uint16x8_t bk0 = vld1q_u16((const uint16_t *)b + l * 8); + const uint16x8_t bk1 = vld1q_u16(l + 1 < k ? (const uint16_t *)b + (l + 1) * 8 : zero); + const uint16x8_t az0 = vzip1q_u16(ak0, ak1); /* (m0,k0),(m0,k1),...,(m3,k1) */ + const uint16x8_t az1 = vzip2q_u16(ak0, ak1); /* (m4,k0),...,(m7,k1) */ + const uint32x4_t bp0 = vreinterpretq_u32_u16(vzip1q_u16(bk0, bk1)); /* (n0,k0),(n0,k1),... */ + const uint32x4_t bp1 = vreinterpretq_u32_u16(vzip2q_u16(bk0, bk1)); +/* vgetq_lane needs a literal lane, hence one macro expansion per column. */ +#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) { + for (BLASLONG i0 = 0; i0 < m;) { + const BLASLONG wm = s2bd_w(m, i0); + const IFLOAT *a = sa + i0 * k; + 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 * k, cb, ldc, alpha); + else + s2bd_block(k, a, wm, sb + j0 * k, 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 + 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 * k, wm, bb + j0 * k, wn, C + i0 + j0 * ldc, ldc, alpha); + j0 += wn; + } + i0 += wm; + } + return 0; +} diff --git a/param.h b/param.h index 46514da0e7..a4fc7b2f3a 100644 --- a/param.h +++ b/param.h @@ -4113,6 +4113,14 @@ 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 +#endif + #if defined(ARMV5) #define SNUMOPT 2 #define DNUMOPT 2 From e73ab683792519e1833df367d6bdcbe98f8ccc01 Mon Sep 17 00:00:00 2001 From: tomhoyt <0tomhoyt@users.noreply.github.com> Date: Thu, 22 Oct 2026 11:16:55 +0800 Subject: [PATCH 6/7] arm64: pack the sbgemm panels k-pair interleaved PANEL FORMAT CHANGE: this rewrites the bfloat16 packing copies shared by both sbgemm kernels and renames them to role-based names, so every kernel that consumes these panels migrates in the same commit. Per Michael Tesch's review of the M4 numbers, the per-k-pair zips and predicated loads cost about a third of the single-thread speed. The copies now pack the panels in the (row, k & 1) lane layout that BFMOPA/BFDOT consume directly - element (row r, k) of a block of w rows at [pp * 2w + 2r + (k & 1)] - and SBGEMM_ALIGN_K 2 matches the driver's B-panel stride to the padded pair extent of the panels. The odd k half of the last pair is zero everywhere; the A-side copies additionally zero-pad rows beyond the ragged chunk (full chunk width), while the B-side copies pack their power-of-two groups at the actual width - the SME2 kernel predicates its B load to the group's lanes and the BFDOT kernel runs its vector path on full groups only, with C stores predicated to the valid rows and columns regardless. The SME2 kernel then needs one plain 512-bit load per operand tile and k pair with no zipping (an SME2 multi-vector load could halve that further on the M4); the BFDOT kernel takes two quad loads per operand and k pair and drops its odd-k zero vector. Fallback loops read the same interleaved panels. --- kernel/arm64/KERNEL.VORTEX | 4 +- kernel/arm64/KERNEL.VORTEXM4 | 8 +- kernel/arm64/sbgemm_kernel_bfdot.c | 45 +++++----- kernel/arm64/sbgemm_kernel_sme2.c | 85 +++++++++---------- ...mm_ncopy_64_sme2.c => sbgemm_ncopy_sme2.c} | 4 +- ...m_tcopy_16_sme2.c => sbgemm_oncopy_sme2.c} | 7 +- ...m_ncopy_16_sme2.c => sbgemm_otcopy_sme2.c} | 7 +- ...mm_tcopy_64_sme2.c => sbgemm_tcopy_sme2.c} | 4 +- kernel/arm64/sme2_sbgemm_copy.h | 56 ++++++++---- param.h | 4 + 10 files changed, 124 insertions(+), 100 deletions(-) rename kernel/arm64/{sbgemm_ncopy_64_sme2.c => sbgemm_ncopy_sme2.c} (94%) rename kernel/arm64/{sbgemm_tcopy_16_sme2.c => sbgemm_oncopy_sme2.c} (92%) rename kernel/arm64/{sbgemm_ncopy_16_sme2.c => sbgemm_otcopy_sme2.c} (92%) rename kernel/arm64/{sbgemm_tcopy_64_sme2.c => sbgemm_tcopy_sme2.c} (94%) diff --git a/kernel/arm64/KERNEL.VORTEX b/kernel/arm64/KERNEL.VORTEX index 59df87d6cc..46f4989e03 100644 --- a/kernel/arm64/KERNEL.VORTEX +++ b/kernel/arm64/KERNEL.VORTEX @@ -5,8 +5,8 @@ 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_ncopy_16_sme2.c -SBGEMMOTCOPY = sbgemm_tcopy_16_sme2.c +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 d3f5944028..b972ff20e4 100644 --- a/kernel/arm64/KERNEL.VORTEXM4 +++ b/kernel/arm64/KERNEL.VORTEXM4 @@ -22,10 +22,10 @@ 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 +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) diff --git a/kernel/arm64/sbgemm_kernel_bfdot.c b/kernel/arm64/sbgemm_kernel_bfdot.c index dcb4c6a42d..7a26f54685 100644 --- a/kernel/arm64/sbgemm_kernel_bfdot.c +++ b/kernel/arm64/sbgemm_kernel_bfdot.c @@ -30,12 +30,12 @@ THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. * 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 kernel zips rows l and l+1 of the - * standard [k][m] and [k][n] panels into the lane layout (i, k & 1): the B pair of a column - * is duplicated across the lanes, the A pairs cover four rows each, and acc[n][h] holds four - * consecutive rows of C column n - one quad load/store per (column, row group). A missing - * second row of an odd k is a zero vector, and blocks narrower than eight rows or columns run - * in scalar code. + * 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. */ @@ -90,10 +90,13 @@ static inline int s2bd_usable(void) { /* 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 l = 0; l < k; ++l) s += s2bd_cvt(a[l * wm + i]) * s2bd_cvt(b[l * wn + j]); + 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; } } @@ -106,17 +109,13 @@ static void s2bd_block(BLASLONG k, const IFLOAT *a, BLASLONG wm, const IFLOAT *b 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); - static const uint16_t zero[8]; /* the missing row of an odd k */ - for (BLASLONG l = 0; l < k; l += 2) { - const uint16x8_t ak0 = vld1q_u16((const uint16_t *)a + l * 8); - const uint16x8_t ak1 = vld1q_u16(l + 1 < k ? (const uint16_t *)a + (l + 1) * 8 : zero); - const uint16x8_t bk0 = vld1q_u16((const uint16_t *)b + l * 8); - const uint16x8_t bk1 = vld1q_u16(l + 1 < k ? (const uint16_t *)b + (l + 1) * 8 : zero); - const uint16x8_t az0 = vzip1q_u16(ak0, ak1); /* (m0,k0),(m0,k1),...,(m3,k1) */ - const uint16x8_t az1 = vzip2q_u16(ak0, ak1); /* (m4,k0),...,(m7,k1) */ - const uint32x4_t bp0 = vreinterpretq_u32_u16(vzip1q_u16(bk0, bk1)); /* (n0,k0),(n0,k1),... */ - const uint32x4_t bp1 = vreinterpretq_u32_u16(vzip2q_u16(bk0, bk1)); -/* vgetq_lane needs a literal lane, hence one macro expansion per column. */ + 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))); \ @@ -136,16 +135,17 @@ static void s2bd_mac(BLASLONG k, const IFLOAT *a, const IFLOAT *b, FLOAT *C, BLA 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 * k; + 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 * k, cb, ldc, alpha); + s2bd_mac(k, a, sb + j0 * pairs * 2, cb, ldc, alpha); else - s2bd_block(k, a, wm, sb + j0 * k, wn, cb, ldc, alpha); + s2bd_block(k, a, wm, sb + j0 * pairs * 2, wn, cb, ldc, alpha); j0 += wn; } i0 += wm; @@ -163,11 +163,12 @@ int CNAME(BLASLONG m, BLASLONG n, BLASLONG k, FLOAT alpha, IFLOAT *ba, IFLOAT *b 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 * k, wm, bb + j0 * k, wn, C + i0 + j0 * ldc, ldc, alpha); + s2bd_block(k, ba + i0 * pairs * 2, wm, bb + j0 * pairs * 2, wn, C + i0 + j0 * ldc, ldc, alpha); j0 += wn; } i0 += wm; diff --git a/kernel/arm64/sbgemm_kernel_sme2.c b/kernel/arm64/sbgemm_kernel_sme2.c index 5656bb8a7b..b7b28308bc 100644 --- a/kernel/arm64/sbgemm_kernel_sme2.c +++ b/kernel/arm64/sbgemm_kernel_sme2.c @@ -31,15 +31,15 @@ THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. * 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 + * 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 ([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. + * 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" @@ -89,42 +89,33 @@ 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 { +/* 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 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); + 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 static void s2b_blk(BLASLONG k, const IFLOAT *a, BLASLONG wm, const IFLOAT *b, BLASLONG wn, FLOAT *C, +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, wm, b, wn); + 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, wm, b, wn); + s2b_mac(0, k, a, b, wn); s2_c_store_alpha(S2_NT, C, ldc, (int)wn, (int)wm, alpha, 1); } } @@ -132,15 +123,16 @@ S2_NOINL static void s2b_blk(BLASLONG k, const IFLOAT *a, BLASLONG wm, const IFL __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 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); + 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; } - j0 += wn; } } @@ -149,20 +141,23 @@ __arm_new("za") __arm_locally_streaming static void s2b_drive(BLASLONG m, BLASLO 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; + 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 l = 0; l < k; ++l) s += s2b_cvt(a[l * wm + i]) * s2b_cvt(b[l * wn + j]); + 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; } - j0 += wn; } } diff --git a/kernel/arm64/sbgemm_ncopy_64_sme2.c b/kernel/arm64/sbgemm_ncopy_sme2.c similarity index 94% rename from kernel/arm64/sbgemm_ncopy_64_sme2.c rename to kernel/arm64/sbgemm_ncopy_sme2.c index 5a8ca0503e..c7f2d7e052 100644 --- a/kernel/arm64/sbgemm_ncopy_64_sme2.c +++ b/kernel/arm64/sbgemm_ncopy_sme2.c @@ -34,8 +34,8 @@ THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 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; + 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_tcopy_16_sme2.c b/kernel/arm64/sbgemm_oncopy_sme2.c similarity index 92% rename from kernel/arm64/sbgemm_tcopy_16_sme2.c rename to kernel/arm64/sbgemm_oncopy_sme2.c index 27d4db82fa..ba90959501 100644 --- a/kernel/arm64/sbgemm_tcopy_16_sme2.c +++ b/kernel/arm64/sbgemm_oncopy_sme2.c @@ -26,8 +26,7 @@ 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). */ +/* bfloat16 B panels of the SME2 sbgemm kernel from op(B) with contiguous columns. */ #include "common.h" #include "sme2_sbgemm_copy.h" @@ -35,8 +34,8 @@ THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 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; + 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_ncopy_16_sme2.c b/kernel/arm64/sbgemm_otcopy_sme2.c similarity index 92% rename from kernel/arm64/sbgemm_ncopy_16_sme2.c rename to kernel/arm64/sbgemm_otcopy_sme2.c index 84afe06a61..37b121b0e2 100644 --- a/kernel/arm64/sbgemm_ncopy_16_sme2.c +++ b/kernel/arm64/sbgemm_otcopy_sme2.c @@ -26,8 +26,7 @@ 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). */ +/* bfloat16 B panels of the SME2 sbgemm kernel from op(B) with contiguous rows. */ #include "common.h" #include "sme2_sbgemm_copy.h" @@ -35,8 +34,8 @@ THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 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; + 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_64_sme2.c b/kernel/arm64/sbgemm_tcopy_sme2.c similarity index 94% rename from kernel/arm64/sbgemm_tcopy_64_sme2.c rename to kernel/arm64/sbgemm_tcopy_sme2.c index 7a41716b9a..ab567fbb8b 100644 --- a/kernel/arm64/sbgemm_tcopy_64_sme2.c +++ b/kernel/arm64/sbgemm_tcopy_sme2.c @@ -34,8 +34,8 @@ THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 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; + 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_sbgemm_copy.h b/kernel/arm64/sme2_sbgemm_copy.h index b76b31e614..d2f1098e7b 100644 --- a/kernel/arm64/sme2_sbgemm_copy.h +++ b/kernel/arm64/sme2_sbgemm_copy.h @@ -25,9 +25,15 @@ 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. */ +/* 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 @@ -37,18 +43,7 @@ THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. #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. */ +/* 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; @@ -57,4 +52,35 @@ static inline BLASLONG s2bc_wn(BLASLONG n, BLASLONG j0) { 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. */ +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); + unsigned *bo = (unsigned *)b; + for (unsigned p = 0; p < pairs; ++p) + for (BLASLONG r = 0; r < w; ++r) { + unsigned v = 0; + if (r < rows && 2 * p < m) { + const IFLOAT *s = a + r * lda + 2 * p; + v = (unsigned)s[0] | (2 * p + 1 < m ? (unsigned)s[1] << 16 : 0); + } + bo[(size_t)p * w + r] = v; + } +} + +/* Same, but element (k, r) at a[r + k * lda]: pairs of k one lda apart in the source. */ +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); + unsigned *bo = (unsigned *)b; + for (unsigned p = 0; p < pairs; ++p) + for (BLASLONG r = 0; r < w; ++r) { + unsigned v = 0; + if (r < rows) { + if (2 * p < m) v = (unsigned)a[r + (BLASLONG)(2 * p) * lda]; + if (2 * p + 1 < m) v |= (unsigned)a[r + (BLASLONG)(2 * p + 1) * lda] << 16; + } + bo[(size_t)p * w + r] = v; + } +} + #endif diff --git a/param.h b/param.h index a4fc7b2f3a..576e9de802 100644 --- a/param.h +++ b/param.h @@ -3584,6 +3584,8 @@ USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. #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 @@ -4119,6 +4121,8 @@ Until then, just keep it different than DGEMM_DEFAULT_UNROLL_N to keep copy rout #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) From 5dea71f3483b698a9354ca35015d188ecbaf6a33 Mon Sep 17 00:00:00 2001 From: tomhoyt <0tomhoyt@users.noreply.github.com> Date: Thu, 22 Oct 2026 19:22:18 +0800 Subject: [PATCH 7/7] arm64: vectorize the k-pair interleaved packing The interleaved copies assembled one u32 per iteration with a strided read and two branches, which on the M4 moved the packing to the top of the profile (Michael Tesch's rerun: sbgemm_itcopy 44 -> 985 samples, single-thread sbgemm at n = 1024 down from 1048 to 716 GFLOPS even though the kernel itself got faster, 1610 -> 1111 samples). s2bc_pairs_c now reads the two halves of a k pair as contiguous 16-lane streams along the rows and zips them with one vzip per quad; s2bc_pairs_l loads four pairs of four source rows and writes them in panel order through a 4x4 u32 transpose (vtrn). Odd k tails, ragged rows and the zero padding stay in scalar code. Packing bandwidth on an M3 Pro, k = 1024, 64 rows: pairs_c 12 -> 41 GB/s, pairs_l 14 -> 61 GB/s (memcpy class). Correctness sweeps unchanged (39116 cases, both targets). --- kernel/arm64/sme2_sbgemm_copy.h | 70 ++++++++++++++++++++++++--------- 1 file changed, 51 insertions(+), 19 deletions(-) diff --git a/kernel/arm64/sme2_sbgemm_copy.h b/kernel/arm64/sme2_sbgemm_copy.h index d2f1098e7b..89df1cc5ce 100644 --- a/kernel/arm64/sme2_sbgemm_copy.h +++ b/kernel/arm64/sme2_sbgemm_copy.h @@ -38,6 +38,7 @@ THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. #ifndef SME2_SBGEMM_COPY_H #define SME2_SBGEMM_COPY_H +#include #include #define S2BC_W GEMM_DEFAULT_UNROLL_M @@ -53,34 +54,65 @@ static inline BLASLONG s2bc_wn(BLASLONG n, BLASLONG j0) { } /* 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. */ + 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); + const unsigned pairs = (unsigned)((m + 1) >> 1), full = (unsigned)(m >> 1); unsigned *bo = (unsigned *)b; - for (unsigned p = 0; p < pairs; ++p) - for (BLASLONG r = 0; r < w; ++r) { - unsigned v = 0; - if (r < rows && 2 * p < m) { - const IFLOAT *s = a + r * lda + 2 * p; - v = (unsigned)s[0] | (2 * p + 1 < m ? (unsigned)s[1] << 16 : 0); + 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); } - bo[(size_t)p * w + r] = v; + } + 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]: pairs of k one lda apart in the source. */ +/* 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); + const unsigned pairs = (unsigned)((m + 1) >> 1), full = (unsigned)(m >> 1); unsigned *bo = (unsigned *)b; - for (unsigned p = 0; p < pairs; ++p) - for (BLASLONG r = 0; r < w; ++r) { - unsigned v = 0; - if (r < rows) { - if (2 * p < m) v = (unsigned)a[r + (BLASLONG)(2 * p) * lda]; - if (2 * p + 1 < m) v |= (unsigned)a[r + (BLASLONG)(2 * p + 1) * lda] << 16; - } - bo[(size_t)p * w + r] = v; + 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