diff --git a/kernel/arm64/KERNEL.VORTEXM4 b/kernel/arm64/KERNEL.VORTEXM4 index 13a0a38c59..a266fcbb29 100644 --- a/kernel/arm64/KERNEL.VORTEXM4 +++ b/kernel/arm64/KERNEL.VORTEXM4 @@ -6,8 +6,8 @@ SGEMMKERNEL = gemm_kernel_sme2.c STRMMKERNEL = gemm_kernel_sme2.c SGEMMINCOPY = gemm_ncopy_sme2.c SGEMMITCOPY = gemm_tcopy_sme2.c -SGEMMONCOPY = ../generic/gemm_ncopy_$(SGEMM_UNROLL_N).c -SGEMMOTCOPY = ../generic/gemm_tcopy_$(SGEMM_UNROLL_N).c +SGEMMONCOPY = gemm_ncopy_b_sme2.c +SGEMMOTCOPY = gemm_tcopy_b_sme2.c SGEMMINCOPYOBJ = sgemm_incopy$(TSUFFIX).$(SUFFIX) SGEMMITCOPYOBJ = sgemm_itcopy$(TSUFFIX).$(SUFFIX) SGEMMONCOPYOBJ = sgemm_oncopy$(TSUFFIX).$(SUFFIX) @@ -17,8 +17,8 @@ DGEMMKERNEL = gemm_kernel_sme2.c DTRMMKERNEL = gemm_kernel_sme2.c 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 +DGEMMONCOPY = gemm_ncopy_b_sme2.c +DGEMMOTCOPY = gemm_tcopy_b_sme2.c DGEMMINCOPYOBJ = dgemm_incopy$(TSUFFIX).$(SUFFIX) DGEMMITCOPYOBJ = dgemm_itcopy$(TSUFFIX).$(SUFFIX) DGEMMONCOPYOBJ = dgemm_oncopy$(TSUFFIX).$(SUFFIX) diff --git a/kernel/arm64/gemm_kernel_sme2.c b/kernel/arm64/gemm_kernel_sme2.c index 9c060a0cd8..7ff6bc1d29 100644 --- a/kernel/arm64/gemm_kernel_sme2.c +++ b/kernel/arm64/gemm_kernel_sme2.c @@ -62,15 +62,49 @@ static inline void s2p_krange(BLASLONG i0, BLASLONG j0, int wm, int wn, BLASLONG #ifdef HAVE_SME2_GEMM #pragma clang attribute push(__attribute__((target("sme2"))), apply_to = function) -typedef void (*s2p_bfn)(BLASLONG, const FLOAT *, int, const FLOAT *, int, FLOAT *, BLASLONG, FLOAT, int) S2_S S2_ZA; +typedef void (*s2p_bfn)(BLASLONG, const FLOAT *, int, const FLOAT *, int, FLOAT *, BLASLONG, FLOAT, int, + BLASLONG) S2_S S2_ZA; + +#if defined(TRMMKERNEL) && defined(LEFT) +/* Steps [l0, l1) at the diagonal: tile Q only from step g + Q * VL on (TRANSA: before step g + (Q + 1) * VL). */ +S2_INL void s2p_mac_diag(const int TN, BLASLONG l0, BLASLONG l1, BLASLONG g, const FLOAT *a, int wm, const FLOAT *b, + int wn) S2_S S2_ZA { + const svbool_t pt = S2_PT(), pb = S2_PW(0, wn); + for (BLASLONG l = l0; l < l1; ++l) { + const s2_v bv = S2_LD1(pb, b + l * wn); + const s2_v4 a0 = s2_ld4n(a + l * wm, wm); + const s2_v4 a1 = TN > 4 ? s2_ld4n(a + l * wm + 4 * S2_VL, wm - 4 * S2_VL) : a0; + const BLASLONG q = (l - g) / S2_VL; +#ifndef TRANSA + S2_UNROLL for (int Q = 0; Q < TN; ++Q) if (Q <= q) s2_mopa(Q, pt, bv, s2_pick(a0, a1, Q)); +#else + S2_UNROLL for (int Q = 0; Q < TN; ++Q) if (Q >= q) s2_mopa(Q, pt, bv, s2_pick(a0, a1, Q)); +#endif + } +} +/* Left TRMM: the diagonal steps through s2p_mac_diag (g: the step of the diagonal element in row 0). */ +#ifndef TRANSA +#define S2P_MAC_BLK(TN, BF, SUB) \ + const BLASLONG t = s2_clip(g + (TN - 1) * S2_VL, kb); \ + s2p_mac_diag(TN, 0, t, g, a, wm, b, wn); \ + s2p_mac(TN, BF, SUB, kb - t, a + t * wm, wm, b + t * wn, wn); +#else +#define S2P_MAC_BLK(TN, BF, SUB) \ + const BLASLONG t = s2_clip(g + S2_VL, kb); \ + s2p_mac(TN, BF, SUB, t, a, wm, b, wn); \ + s2p_mac_diag(TN, t, kb, g, a, wm, b, wn); +#endif +#else +#define S2P_MAC_BLK(TN, BF, SUB) s2p_mac(TN, BF, SUB, kb, a, wm, b, wn); +#endif /* Block of wm rows and wn columns of C; cm 0: C += A B, 1: C -= A B, 2: C += alpha A B, 3: C = alpha A B. */ #define S2P_DEF(TN, BF, SUB) \ S2_NOINL void s2p_blk_##TN##_##BF##_##SUB(BLASLONG kb, const FLOAT *a, int wm, const FLOAT *b, int wn, FLOAT *C, \ - BLASLONG ldc, FLOAT alpha, int cm) S2_S S2_ZA { \ + BLASLONG ldc, FLOAT alpha, int cm, BLASLONG g) S2_S S2_ZA { \ if (cm < 2) s2_c_load(1, TN, C, ldc, wn, wm, S2_LOAD, (FLOAT)1); \ else svzero_za(); \ - s2p_mac(TN, BF, SUB, kb, a, wm, b, wn); \ + S2P_MAC_BLK(TN, BF, SUB) \ if (cm < 2) s2_c_store(1, TN, C, ldc, wn, wm); \ else s2_c_store_alpha(TN, C, ldc, wn, wm, alpha, cm == 2); \ } @@ -105,9 +139,9 @@ __arm_new("za") __arm_locally_streaming static void s2p_drive(BLASLONG m, BLASLO const int wm = m - i0 < S2_NR ? (int)(m - i0) : S2_NR; BLASLONG ks, kl; s2p_krange(i0, j0, wm, wn, k, offset, &ks, &kl); - s2p_blk[(wm + S2_VL - 1) / S2_VL - 1][wn == S2_VL][cm == 1](kl, sa + i0 * k + ks * wm, wm, - sb + j0 * k + ks * wn, wn, C + i0 + j0 * ldc, - ldc, alpha, cm); + s2p_blk[(wm + S2_VL - 1) / S2_VL - 1][wn == S2_VL][cm == 1](kl, sa + i0 * k + ks * wm, wm, sb + j0 * k + ks * wn, + wn, C + i0 + j0 * ldc, ldc, alpha, cm, + offset + i0 - ks); } j0 += wn; } diff --git a/kernel/arm64/gemm_ncopy_b_sme2.c b/kernel/arm64/gemm_ncopy_b_sme2.c new file mode 100644 index 0000000000..2e37806b79 --- /dev/null +++ b/kernel/arm64/gemm_ncopy_b_sme2.c @@ -0,0 +1,36 @@ +/*************************************************************************** +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. +*****************************************************************************/ + +/* B panels of the SME2 kernels (GEMM_DEFAULT_UNROLL_N columns, then 8/4/2/1) from op(B) with contiguous columns. */ + +#include "common.h" +#include "sme2_gemm_copy.h" + +int CNAME(BLASLONG m, BLASLONG n, FLOAT *a, BLASLONG lda, FLOAT *b) { + s2c_b_panels(0, m, n, a, lda, b); + return 0; +} diff --git a/kernel/arm64/gemm_tcopy_b_sme2.c b/kernel/arm64/gemm_tcopy_b_sme2.c new file mode 100644 index 0000000000..51c055c322 --- /dev/null +++ b/kernel/arm64/gemm_tcopy_b_sme2.c @@ -0,0 +1,36 @@ +/*************************************************************************** +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. +*****************************************************************************/ + +/* B panels of the SME2 kernels (GEMM_DEFAULT_UNROLL_N columns, then 8/4/2/1) from op(B) with contiguous rows. */ + +#include "common.h" +#include "sme2_gemm_copy.h" + +int CNAME(BLASLONG m, BLASLONG n, FLOAT *a, BLASLONG lda, FLOAT *b) { + s2c_b_panels(1, m, n, a, lda, b); + return 0; +} diff --git a/kernel/arm64/sme2_gemm_copy.h b/kernel/arm64/sme2_gemm_copy.h index e242932c15..49ce80ba0a 100644 --- a/kernel/arm64/sme2_gemm_copy.h +++ b/kernel/arm64/sme2_gemm_copy.h @@ -25,25 +25,162 @@ 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. *****************************************************************************/ -/* Copies into the A panels of the SME2 kernels: m rows of w <= 64 columns, element (i, c) to b[i * w + c]. */ +/* Copies into the A and B panels of the SME2 kernels: m rows of w <= 64 columns, element (i, c) to b[i * w + c]. */ #ifndef SME2_GEMM_COPY_H #define SME2_GEMM_COPY_H #include #include +#include "sme2_gemm_detect.h" #define S2C_W GEMM_DEFAULT_UNROLL_M static inline BLASLONG s2c_clip(BLASLONG x, BLASLONG lo, BLASLONG hi) { return x < lo ? lo : (x > hi ? hi : x); } +#ifdef HAVE_SME2_GEMM +/* Streaming-mode copies: the SME unit moves data faster than NEON. Below S2C_SME elements NEON is faster. */ +#define S2C_SME 512 +#ifdef DOUBLE +#define S2C_VL 8 +#define S2C_WL(a, b) svwhilelt_b64_s64(a, b) +#define S2C_LD svld1_f64 +#define S2C_ST svst1_f64 +#define S2C_DUP svdup_n_f64 +#else +#define S2C_VL 16 +#define S2C_WL(a, b) svwhilelt_b32_s64(a, b) +#define S2C_LD svld1_f32 +#define S2C_ST svst1_f32 +#define S2C_DUP svdup_n_f32 +#endif + +static inline __attribute__((always_inline)) void s2c_rows_s(BLASLONG m, BLASLONG w, const FLOAT *a, BLASLONG lda, + FLOAT *b) __arm_streaming { + for (BLASLONG c = 0; c < w; c += S2C_VL) { + const svbool_t p = S2C_WL(c, w); + BLASLONG i = 0; + for (; i + 4 <= m; i += 4) { + const FLOAT *q = a + i * lda + c; + FLOAT *d = b + i * w + c; + S2C_ST(p, d, S2C_LD(p, q)); + S2C_ST(p, d + w, S2C_LD(p, q + lda)); + S2C_ST(p, d + 2 * w, S2C_LD(p, q + 2 * lda)); + S2C_ST(p, d + 3 * w, S2C_LD(p, q + 3 * lda)); + } + for (; i < m; ++i) S2C_ST(p, b + i * w + c, S2C_LD(p, a + i * lda + c)); + } +} + +/* Tile numbers must be literals; with a constant t the switches fold away. */ +static inline __attribute__((always_inline)) void s2c_ldh(int t, uint32_t s, svbool_t p, const FLOAT *q) __arm_streaming + __arm_inout("za") { +#ifdef DOUBLE + switch (t) { + case 0: svld1_hor_za64(0, s, p, q); break; + case 1: svld1_hor_za64(1, s, p, q); break; + case 2: svld1_hor_za64(2, s, p, q); break; + case 3: svld1_hor_za64(3, s, p, q); break; + case 4: svld1_hor_za64(4, s, p, q); break; + case 5: svld1_hor_za64(5, s, p, q); break; + case 6: svld1_hor_za64(6, s, p, q); break; + default: svld1_hor_za64(7, s, p, q); break; + } +#else + switch (t) { + case 0: svld1_hor_za32(0, s, p, q); break; + case 1: svld1_hor_za32(1, s, p, q); break; + case 2: svld1_hor_za32(2, s, p, q); break; + default: svld1_hor_za32(3, s, p, q); break; + } +#endif +} + +static inline __attribute__((always_inline)) void s2c_stv(int t, uint32_t s, svbool_t p, FLOAT *q) __arm_streaming + __arm_inout("za") { +#ifdef DOUBLE + switch (t) { + case 0: svst1_ver_za64(0, s, p, q); break; + case 1: svst1_ver_za64(1, s, p, q); break; + case 2: svst1_ver_za64(2, s, p, q); break; + case 3: svst1_ver_za64(3, s, p, q); break; + case 4: svst1_ver_za64(4, s, p, q); break; + case 5: svst1_ver_za64(5, s, p, q); break; + case 6: svst1_ver_za64(6, s, p, q); break; + default: svst1_ver_za64(7, s, p, q); break; + } +#else + switch (t) { + case 0: svst1_ver_za32(0, s, p, q); break; + case 1: svst1_ver_za32(1, s, p, q); break; + case 2: svst1_ver_za32(2, s, p, q); break; + default: svst1_ver_za32(3, s, p, q); break; + } +#endif +} + +/* Columns c go into horizontal slice c % VL of tile c / VL and leave transposed through vertical slices. */ +static inline __attribute__((always_inline)) void s2c_cols_s(BLASLONG m, BLASLONG w, const FLOAT *a, BLASLONG lda, + FLOAT *b) __arm_streaming __arm_inout("za") { + BLASLONG i = 0; + if (w <= S2C_VL) /* B panels: two tiles in turn, so that loads into one overlap stores from the other */ + for (; i + 2 * S2C_VL <= m; i += 2 * S2C_VL) { + const svbool_t pt = S2C_WL(0, S2C_VL), pw = S2C_WL(0, w); + for (int c = 0; c < w; ++c) { + s2c_ldh(0, c, pt, a + i + c * lda); + s2c_ldh(1, c, pt, a + i + S2C_VL + c * lda); + } + for (int r = 0; r < S2C_VL; ++r) s2c_stv(0, r, pw, b + (i + r) * w); + for (int r = 0; r < S2C_VL; ++r) s2c_stv(1, r, pw, b + (i + S2C_VL + r) * w); + } + for (; i < m; i += S2C_VL) { + const svbool_t pm = S2C_WL(i, m); + const int rows = m - i < S2C_VL ? (int)(m - i) : S2C_VL; + _Pragma("GCC unroll 8") for (int t = 0; t < S2C_W / S2C_VL; ++t) { + if (t * S2C_VL >= w) break; + for (int s = 0; s < S2C_VL && t * S2C_VL + s < w; ++s) s2c_ldh(t, s, pm, a + i + (t * S2C_VL + s) * lda); + } + for (int r = 0; r < rows; ++r) { + FLOAT *d = b + (i + r) * w; + _Pragma("GCC unroll 8") for (int t = 0; t < S2C_W / S2C_VL; ++t) { + if (t * S2C_VL >= w) break; + s2c_stv(t, r, S2C_WL(t * S2C_VL, w), d + t * S2C_VL); + } + } + } +} + +__arm_locally_streaming static __attribute__((noinline)) void s2c_rows_sme(BLASLONG m, BLASLONG w, const FLOAT *a, + BLASLONG lda, FLOAT *b) { + s2c_rows_s(m, w, a, lda, b); +} + +__arm_new("za") __arm_locally_streaming static __attribute__((noinline)) void s2c_cols_sme(BLASLONG m, BLASLONG w, + const FLOAT *a, BLASLONG lda, + FLOAT *b) { + s2c_cols_s(m, w, a, lda, b); +} +#endif + /* Element (i, c) at a[c + i * lda]. */ static inline void s2c_rows(BLASLONG m, BLASLONG w, const FLOAT *a, BLASLONG lda, FLOAT *b) { +#ifdef HAVE_SME2_GEMM + if (m * w >= S2C_SME && s2_usable()) { + s2c_rows_sme(m, w, a, lda, b); + return; + } +#endif for (BLASLONG i = 0; i < m; ++i) memcpy(b + i * w, a + i * lda, w * sizeof(FLOAT)); } /* Element (i, c) at a[i + c * lda], transposed in registers 4 x 4 (fp32) or 2 x 2 (fp64) at a time. */ static inline void s2c_cols(BLASLONG m, BLASLONG w, const FLOAT *a, BLASLONG lda, FLOAT *b) { +#ifdef HAVE_SME2_GEMM + if (m * w >= S2C_SME && s2_usable()) { + s2c_cols_sme(m, w, a, lda, b); + return; + } +#endif BLASLONG c = 0; #ifdef DOUBLE for (; c + 2 <= w; c += 2) { @@ -85,8 +222,64 @@ static inline void s2c_cols(BLASLONG m, BLASLONG w, const FLOAT *a, BLASLONG lda for (BLASLONG i = 0; i < m; ++i) b[i * w + c] = a[i + c * lda]; } +/* B panels of GEMM_DEFAULT_UNROLL_N columns, then 8/4/2/1 (the layout of generic/gemm_[nt]copy_16.c). */ +#ifdef HAVE_SME2_GEMM +__arm_new("za") __arm_locally_streaming static __attribute__((noinline)) void s2c_b_panels_sme(int rows, BLASLONG m, + BLASLONG n, + const FLOAT *a, + BLASLONG lda, FLOAT *b) { + for (BLASLONG j = 0, w = GEMM_DEFAULT_UNROLL_N; j < n; j += w, b += m * w) { + while (w > n - j) w >>= 1; + if (rows) s2c_rows_s(m, w, a + j, lda, b); + else s2c_cols_s(m, w, a + j * lda, lda, b); + } +} +#endif + +static inline void s2c_b_panels(int rows, BLASLONG m, BLASLONG n, const FLOAT *a, BLASLONG lda, FLOAT *b) { +#ifdef HAVE_SME2_GEMM + if (m * n >= S2C_SME && s2_usable()) { + s2c_b_panels_sme(rows, m, n, a, lda, b); + return; + } +#endif + for (BLASLONG j = 0, w = GEMM_DEFAULT_UNROLL_N; j < n; j += w, b += m * w) { + while (w > n - j) w >>= 1; + if (rows) s2c_rows(m, w, a + j, lda, b); + else s2c_cols(m, w, a + j * lda, lda, b); + } +} + static inline void s2c_zero(BLASLONG m, BLASLONG w, FLOAT *b) { if (m > 0) memset(b, 0, (size_t)(m * w) * sizeof(FLOAT)); } +/* Rows i1..i2 of a diagonal block, diagonal in column d0 + i: zeros left (zl) or right of it, 1 there with unit. */ +static inline void s2c_tri(BLASLONG i1, BLASLONG i2, BLASLONG w, BLASLONG d0, int zl, int unit, FLOAT *b) { + for (BLASLONG i = i1; i < i2; ++i) { + const BLASLONG d = d0 + i; + FLOAT *bi = b + i * w; + for (BLASLONG c = zl ? 0 : d + 1; c < (zl ? d : w); ++c) bi[c] = ZERO; + if (unit) bi[d] = ONE; + } +} + +#ifdef HAVE_SME2_GEMM +static inline __attribute__((always_inline)) void s2c_zero_s(BLASLONG m, BLASLONG w, FLOAT *b) __arm_streaming { + for (BLASLONG k = 0; k < m * w; k += S2C_VL) S2C_ST(S2C_WL(k, m * w), b + k, S2C_DUP(0)); +} + +static inline __attribute__((always_inline)) void s2c_tri_s(BLASLONG i1, BLASLONG i2, BLASLONG w, BLASLONG d0, int zl, + int unit, FLOAT *b) __arm_streaming { + for (BLASLONG i = i1; i < i2; ++i) { + const BLASLONG d = d0 + i; + for (BLASLONG c = 0; c < w; c += S2C_VL) { + const svbool_t pw = S2C_WL(c, w), pd = S2C_WL(c, d), pd1 = S2C_WL(c, d + 1); + S2C_ST(zl ? pd : svbic_b_z(pw, pw, pd1), b + i * w + c, S2C_DUP(0)); + if (unit) S2C_ST(svbic_b_z(pw, pd1, pd), b + i * w + c, S2C_DUP(1)); + } + } +} +#endif + #endif diff --git a/kernel/arm64/sme2_gemm_tile.h b/kernel/arm64/sme2_gemm_tile.h index 4200e5bc25..02db94afe3 100644 --- a/kernel/arm64/sme2_gemm_tile.h +++ b/kernel/arm64/sme2_gemm_tile.h @@ -508,15 +508,29 @@ S2_INL void s2_c_store(const int TM, const int TN, FLOAT *C, BLASLONG ldc, int m } } +/* alpha z, plus the C row at p (add). */ +S2_INL s2_v s2_alpha_v(s2_v z, FLOAT *p, svbool_t pg, FLOAT alpha, int add) S2_S { + return add ? S2_MLA(S2_LD1(pg, p), z, alpha) : S2_MUL(z, alpha); +} /* C = C + alpha ZA (add) or C = alpha ZA from a row of TN tiles: rows < mrows, columns < ncols. */ S2_INL void s2_c_store_alpha(const int TN, FLOAT *C, BLASLONG ldc, int mrows, int ncols, FLOAT alpha, int add) S2_S S2_ZA { - for (int s = 0; s < mrows; ++s) { - FLOAT *row = C + (BLASLONG)s * ldc; + int s = 0; + for (; s + 4 <= mrows; s += 4) { /* four rows per ZA move (MOVA vg4) */ + S2_UNROLL for (int Q = 0; Q < TN; ++Q) { + const svbool_t pg = S2_PW(Q * S2_VL, ncols); + const s2_v4 q = s2_rdh4(Q, s); + S2_UNROLL for (int u = 0; u < 4; ++u) { + FLOAT *p = C + (BLASLONG)(s + u) * ldc + Q * S2_VL; + S2_ST1(pg, p, s2_alpha_v(s2_get4(q, u), p, pg, alpha, add)); + } + } + } + for (; s < mrows; ++s) { S2_UNROLL for (int Q = 0; Q < TN; ++Q) { const svbool_t pg = S2_PW(Q * S2_VL, ncols); - const s2_v z = s2_rdh(Q, s); - S2_ST1(pg, row + Q * S2_VL, add ? S2_MLA(S2_LD1(pg, row + Q * S2_VL), z, alpha) : S2_MUL(z, alpha)); + FLOAT *p = C + (BLASLONG)s * ldc + Q * S2_VL; + S2_ST1(pg, p, s2_alpha_v(s2_rdh(Q, s), p, pg, alpha, add)); } } } diff --git a/kernel/arm64/trmm_ncopy_sme2.c b/kernel/arm64/trmm_ncopy_sme2.c index 20afbdf99a..6ae3b8b66b 100644 --- a/kernel/arm64/trmm_ncopy_sme2.c +++ b/kernel/arm64/trmm_ncopy_sme2.c @@ -31,34 +31,50 @@ THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. #include "sme2_gemm_copy.h" #ifndef LOWER -#define KEEP(r, c) ((r) < (c)) +#define ZL 1 #else -#define KEEP(r, c) ((r) > (c)) +#define ZL 0 #endif #ifdef UNIT -#define DIAG(v) ONE +#define UNITD 1 #else -#define DIAG(v) (v) +#define UNITD 0 #endif -int CNAME(BLASLONG m, BLASLONG n, FLOAT *a, BLASLONG lda, BLASLONG posX, BLASLONG posY, FLOAT *b) { - for (BLASLONG j0 = 0; j0 < n; j0 += S2C_W) { - const BLASLONG w = n - j0 < S2C_W ? n - j0 : S2C_W, lo = posY + j0; - const FLOAT *ap = a + posX + lo * lda; - const BLASLONG i1 = s2c_clip(lo - posX, 0, m), i2 = s2c_clip(lo + w - posX, i1, m); -#ifndef LOWER - s2c_cols(i1, w, ap, lda, b); - s2c_zero(m - i2, w, b + i2 * w); +#if ZL +#define PART(COPY, ZERO) \ + COPY(i2, w, ap, lda, b); \ + ZERO(m - i2, w, b + i2 * w) #else - s2c_zero(i1, w, b); - s2c_cols(m - i2, w, ap + i2, lda, b + i2 * w); +#define PART(COPY, ZERO) \ + ZERO(i1, w, b); \ + COPY(m - i1, w, ap + i1, lda, b + i1 * w) #endif - for (BLASLONG i = i1; i < i2; ++i) - for (BLASLONG c = 0; c < w; ++c) { - const BLASLONG r = posX + i, col = lo + c; - b[i * w + c] = r == col ? DIAG(ap[i + c * lda]) : (KEEP(r, col) ? ap[i + c * lda] : ZERO); - } - b += m * w; +/* Panels of w columns: the rows on the kept side of the diagonal copied, the others zeroed, then the diagonal block. */ +#define PANELS(COPY, ZERO, TRI) \ + for (BLASLONG j0 = 0, w; j0 < n; j0 += w, b += m * w) { \ + w = n - j0 < S2C_W ? n - j0 : S2C_W; \ + const BLASLONG lo = posY + j0, i1 = s2c_clip(lo - posX, 0, m), i2 = s2c_clip(lo + w - posX, i1, m); \ + const FLOAT *ap = a + posX + lo * lda; \ + PART(COPY, ZERO); \ + TRI(i1, i2, w, posX - lo, ZL, UNITD, b); \ } + +#ifdef HAVE_SME2_GEMM +__arm_new("za") __arm_locally_streaming static + __attribute__((noinline)) void panels_sme(BLASLONG m, BLASLONG n, FLOAT *a, BLASLONG lda, BLASLONG posX, + BLASLONG posY, FLOAT *b) { + PANELS(s2c_cols_s, s2c_zero_s, s2c_tri_s) +} +#endif + +int CNAME(BLASLONG m, BLASLONG n, FLOAT *a, BLASLONG lda, BLASLONG posX, BLASLONG posY, FLOAT *b) { +#ifdef HAVE_SME2_GEMM + if (m * n >= S2C_SME && s2_usable()) { + panels_sme(m, n, a, lda, posX, posY, b); + return 0; + } +#endif + PANELS(s2c_cols, s2c_zero, s2c_tri) return 0; } diff --git a/kernel/arm64/trmm_tcopy_sme2.c b/kernel/arm64/trmm_tcopy_sme2.c index 93e61a267c..cbf32fc45a 100644 --- a/kernel/arm64/trmm_tcopy_sme2.c +++ b/kernel/arm64/trmm_tcopy_sme2.c @@ -31,34 +31,49 @@ THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. #include "sme2_gemm_copy.h" #ifndef LOWER -#define KEEP(r, c) ((r) > (c)) +#define ZL 0 #else -#define KEEP(r, c) ((r) < (c)) +#define ZL 1 #endif #ifdef UNIT -#define DIAG(v) ONE +#define UNITD 1 #else -#define DIAG(v) (v) +#define UNITD 0 #endif -int CNAME(BLASLONG m, BLASLONG n, FLOAT *a, BLASLONG lda, BLASLONG posX, BLASLONG posY, FLOAT *b) { - for (BLASLONG j0 = 0; j0 < n; j0 += S2C_W) { - const BLASLONG w = n - j0 < S2C_W ? n - j0 : S2C_W, lo = posY + j0; - const FLOAT *ap = a + lo + posX * lda; - const BLASLONG i1 = s2c_clip(lo - posX, 0, m), i2 = s2c_clip(lo + w - posX, i1, m); -#ifndef LOWER - s2c_zero(i1, w, b); - s2c_rows(m - i2, w, ap + i2 * lda, lda, b + i2 * w); +#if ZL +#define PART(COPY, ZERO) \ + COPY(i2, w, ap, lda, b); \ + ZERO(m - i2, w, b + i2 * w) #else - s2c_rows(i1, w, ap, lda, b); - s2c_zero(m - i2, w, b + i2 * w); +#define PART(COPY, ZERO) \ + ZERO(i1, w, b); \ + COPY(m - i1, w, ap + i1 * lda, lda, b + i1 * w) #endif - for (BLASLONG i = i1; i < i2; ++i) - for (BLASLONG c = 0; c < w; ++c) { - const BLASLONG r = posX + i, col = lo + c; - b[i * w + c] = r == col ? DIAG(ap[c + i * lda]) : (KEEP(r, col) ? ap[c + i * lda] : ZERO); - } - b += m * w; +/* Panels of w columns: the rows on the kept side of the diagonal copied, the others zeroed, then the diagonal block. */ +#define PANELS(COPY, ZERO, TRI) \ + for (BLASLONG j0 = 0, w; j0 < n; j0 += w, b += m * w) { \ + w = n - j0 < S2C_W ? n - j0 : S2C_W; \ + const BLASLONG lo = posY + j0, i1 = s2c_clip(lo - posX, 0, m), i2 = s2c_clip(lo + w - posX, i1, m); \ + const FLOAT *ap = a + lo + posX * lda; \ + PART(COPY, ZERO); \ + TRI(i1, i2, w, posX - lo, ZL, UNITD, b); \ } + +#ifdef HAVE_SME2_GEMM +__arm_locally_streaming static __attribute__((noinline)) void panels_sme(BLASLONG m, BLASLONG n, FLOAT *a, BLASLONG lda, + BLASLONG posX, BLASLONG posY, FLOAT *b) { + PANELS(s2c_rows_s, s2c_zero_s, s2c_tri_s) +} +#endif + +int CNAME(BLASLONG m, BLASLONG n, FLOAT *a, BLASLONG lda, BLASLONG posX, BLASLONG posY, FLOAT *b) { +#ifdef HAVE_SME2_GEMM + if (m * n >= S2C_SME && s2_usable()) { + panels_sme(m, n, a, lda, posX, posY, b); + return 0; + } +#endif + PANELS(s2c_rows, s2c_zero, s2c_tri) return 0; } diff --git a/kernel/arm64/trsm_ncopy_sme2.c b/kernel/arm64/trsm_ncopy_sme2.c index a4632ab2b0..c35c92f67e 100644 --- a/kernel/arm64/trsm_ncopy_sme2.c +++ b/kernel/arm64/trsm_ncopy_sme2.c @@ -53,11 +53,18 @@ int CNAME(BLASLONG m, BLASLONG n, FLOAT *a, BLASLONG lda, BLASLONG offset, FLOAT s2c_zero(i1, w, b); s2c_cols(m - i2, w, ap + i2, lda, b + i2 * w); #endif - for (BLASLONG i = i1; i < i2; ++i) - for (BLASLONG c = 0; c < w; ++c) { - const BLASLONG r = i, col = lo + c; - b[i * w + c] = r == col ? INV(ap[i + c * lda]) : (KEEP(r, col) ? ap[i + c * lda] : ZERO); + /* Diagonal block: a copy, then the zero side of the diagonal (column d of row i) and the inverted diagonal. */ + s2c_cols(i2 - i1, w, ap + i1, lda, b + i1 * w); + for (BLASLONG i = i1; i < i2; ++i) { + const BLASLONG d = i - lo; + FLOAT *bi = b + i * w; + if (KEEP(0, 1)) { + for (BLASLONG c = 0; c < d; ++c) bi[c] = ZERO; + } else { + for (BLASLONG c = d + 1; c < w; ++c) bi[c] = ZERO; } + bi[d] = INV(bi[d]); + } b += m * w; } return 0; diff --git a/kernel/arm64/trsm_tcopy_sme2.c b/kernel/arm64/trsm_tcopy_sme2.c index 016c3a0c66..f010b1d985 100644 --- a/kernel/arm64/trsm_tcopy_sme2.c +++ b/kernel/arm64/trsm_tcopy_sme2.c @@ -53,11 +53,18 @@ int CNAME(BLASLONG m, BLASLONG n, FLOAT *a, BLASLONG lda, BLASLONG offset, FLOAT s2c_rows(i1, w, ap, lda, b); s2c_zero(m - i2, w, b + i2 * w); #endif - for (BLASLONG i = i1; i < i2; ++i) - for (BLASLONG c = 0; c < w; ++c) { - const BLASLONG r = i, col = lo + c; - b[i * w + c] = r == col ? INV(ap[c + i * lda]) : (KEEP(r, col) ? ap[c + i * lda] : ZERO); + /* Diagonal block: a copy, then the zero side of the diagonal (column d of row i) and the inverted diagonal. */ + s2c_rows(i2 - i1, w, ap + i1 * lda, lda, b + i1 * w); + for (BLASLONG i = i1; i < i2; ++i) { + const BLASLONG d = i - lo; + FLOAT *bi = b + i * w; + if (KEEP(0, 1)) { + for (BLASLONG c = 0; c < d; ++c) bi[c] = ZERO; + } else { + for (BLASLONG c = d + 1; c < w; ++c) bi[c] = ZERO; } + bi[d] = INV(bi[d]); + } b += m * w; } return 0;