From 8dcf13a87e2e0fc8b759ddb531a5f96e58b609aa Mon Sep 17 00:00:00 2001 From: moluopro Date: Thu, 9 Jul 2026 13:40:22 +0800 Subject: [PATCH 1/6] cmake: clean up ARMV9SME Windows Clang flags --- cmake/cc.cmake | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/cmake/cc.cmake b/cmake/cc.cmake index bacdbef5af..6f597692a1 100644 --- a/cmake/cc.cmake +++ b/cmake/cc.cmake @@ -313,13 +313,13 @@ endif () if (${CORE} STREQUAL ARMV9SME) if (NOT DYNAMIC_ARCH) if (${CMAKE_C_COMPILER_ID} STREQUAL "NVHPC" AND NOT NO_SVE) - set (CCOMMON_OPT "${CCOMMON_OPT} -tp=host") + set (CCOMMON_OPT "${CCOMMON_OPT} -tp=host") else () - set (CCOMMON_OPT "${CCOMMON_OPT} -march=armv9-a+sme") - if (${OSNAME} STREQUAL Windows AND ${CMAKE_C_COMPILER_ID} MATCHES "Clang" ) - set (CCOMMON_OPT "${CCOMMON_OPT} --aarch64-stack-hazard-size=0") + set (CCOMMON_OPT "${CCOMMON_OPT} -march=armv9-a+sme") + if (CMAKE_SYSTEM_NAME STREQUAL "Windows" AND CMAKE_C_COMPILER_ID MATCHES "Clang") + set (CCOMMON_OPT "${CCOMMON_OPT} --aarch64-stack-hazard-size=0") + endif () endif () - endif () endif () endif () From a33e238b354543f73ac91192c4dbd3a26e297312 Mon Sep 17 00:00:00 2001 From: moluopro Date: Thu, 9 Jul 2026 13:41:10 +0800 Subject: [PATCH 2/6] cmake: define ARM64 SME direct kernel variants --- kernel/CMakeLists.txt | 23 +++++++++++------------ 1 file changed, 11 insertions(+), 12 deletions(-) diff --git a/kernel/CMakeLists.txt b/kernel/CMakeLists.txt index 94c02dc9dd..38f484a5f2 100644 --- a/kernel/CMakeLists.txt +++ b/kernel/CMakeLists.txt @@ -292,38 +292,38 @@ function (build_core TARGET_CORE KDIR TSUFFIX KERNEL_DEFINITIONS) if (USE_DIRECT_SSYMM) if (ARM64) set (SSYMMDIRECTKERNEL_ALPHA_BETA ssymm_direct_alpha_beta_arm64_sme1.c) - GenerateNamedObjects("${KERNELDIR}/${SSYMMDIRECTKERNEL_ALPHA_BETA}" "" "symm_direct_alpha_betaLU" false "" "" false SINGLE) - GenerateNamedObjects("${KERNELDIR}/${SSYMMDIRECTKERNEL_ALPHA_BETA}" "" "symm_direct_alpha_betaLL" false "" "" false SINGLE) + GenerateNamedObjects("${KERNELDIR}/${SSYMMDIRECTKERNEL_ALPHA_BETA}" "LEFT;UPPER" "symm_direct_alpha_betaLU" false "" "" false SINGLE) + GenerateNamedObjects("${KERNELDIR}/${SSYMMDIRECTKERNEL_ALPHA_BETA}" "LEFT;LOWER" "symm_direct_alpha_betaLL" false "" "" false SINGLE) endif () endif() if (USE_DIRECT_STRMM) if (ARM64) set (STRMMDIRECTKERNEL strmm_direct_arm64_sme1.c) - GenerateNamedObjects("${KERNELDIR}/${STRMMDIRECTKERNEL}" "" "trmm_direct_LNUN" false "" "" false SINGLE) + GenerateNamedObjects("${KERNELDIR}/${STRMMDIRECTKERNEL}" "UPPER" "trmm_direct_LNUN" false "" "" false SINGLE) GenerateNamedObjects("${KERNELDIR}/${STRMMDIRECTKERNEL}" "" "trmm_direct_LNLN" false "" "" false SINGLE) - GenerateNamedObjects("${KERNELDIR}/${STRMMDIRECTKERNEL}" "" "trmm_direct_LTUN" false "" "" false SINGLE) - GenerateNamedObjects("${KERNELDIR}/${STRMMDIRECTKERNEL}" "" "trmm_direct_LTLN" false "" "" false SINGLE) + GenerateNamedObjects("${KERNELDIR}/${STRMMDIRECTKERNEL}" "TRANSA;UPPER" "trmm_direct_LTUN" false "" "" false SINGLE) + GenerateNamedObjects("${KERNELDIR}/${STRMMDIRECTKERNEL}" "TRANSA" "trmm_direct_LTLN" false "" "" false SINGLE) endif () endif () if (USE_DIRECT_SSYRK) if (ARM64) set (SSYRKDIRECTKERNEL_ALPHA_BETA ssyrk_direct_alpha_beta_arm64_sme1.c) - GenerateNamedObjects("${KERNELDIR}/${SSYRKDIRECTKERNEL_ALPHA_BETA}" "" "syrk_direct_alpha_betaUN" false "" "" false SINGLE) - GenerateNamedObjects("${KERNELDIR}/${SSYRKDIRECTKERNEL_ALPHA_BETA}" "" "syrk_direct_alpha_betaUT" false "" "" false SINGLE) + GenerateNamedObjects("${KERNELDIR}/${SSYRKDIRECTKERNEL_ALPHA_BETA}" "UPPER" "syrk_direct_alpha_betaUN" false "" "" false SINGLE) + GenerateNamedObjects("${KERNELDIR}/${SSYRKDIRECTKERNEL_ALPHA_BETA}" "TRANSA;UPPER" "syrk_direct_alpha_betaUT" false "" "" false SINGLE) GenerateNamedObjects("${KERNELDIR}/${SSYRKDIRECTKERNEL_ALPHA_BETA}" "" "syrk_direct_alpha_betaLN" false "" "" false SINGLE) - GenerateNamedObjects("${KERNELDIR}/${SSYRKDIRECTKERNEL_ALPHA_BETA}" "" "syrk_direct_alpha_betaLT" false "" "" false SINGLE) + GenerateNamedObjects("${KERNELDIR}/${SSYRKDIRECTKERNEL_ALPHA_BETA}" "TRANSA" "syrk_direct_alpha_betaLT" false "" "" false SINGLE) endif () endif() if (USE_DIRECT_SSYR2K) if (ARM64) set (SSYR2KDIRECTKERNEL_ALPHA_BETA ssyr2k_direct_alpha_beta_arm64_sme1.c) - GenerateNamedObjects("${KERNELDIR}/${SSYR2KDIRECTKERNEL_ALPHA_BETA}" "" "syr2k_direct_alpha_betaUN" false "" "" false SINGLE) - GenerateNamedObjects("${KERNELDIR}/${SSYR2KDIRECTKERNEL_ALPHA_BETA}" "" "syr2k_direct_alpha_betaUT" false "" "" false SINGLE) + GenerateNamedObjects("${KERNELDIR}/${SSYR2KDIRECTKERNEL_ALPHA_BETA}" "UPPER" "syr2k_direct_alpha_betaUN" false "" "" false SINGLE) + GenerateNamedObjects("${KERNELDIR}/${SSYR2KDIRECTKERNEL_ALPHA_BETA}" "TRANSA;UPPER" "syr2k_direct_alpha_betaUT" false "" "" false SINGLE) GenerateNamedObjects("${KERNELDIR}/${SSYR2KDIRECTKERNEL_ALPHA_BETA}" "" "syr2k_direct_alpha_betaLN" false "" "" false SINGLE) - GenerateNamedObjects("${KERNELDIR}/${SSYR2KDIRECTKERNEL_ALPHA_BETA}" "" "syr2k_direct_alpha_betaLT" false "" "" false SINGLE) + GenerateNamedObjects("${KERNELDIR}/${SSYR2KDIRECTKERNEL_ALPHA_BETA}" "TRANSA" "syr2k_direct_alpha_betaLT" false "" "" false SINGLE) endif () endif() @@ -1555,4 +1555,3 @@ else () set(KERNEL_DEFINITIONS "") build_core("${TARGET_CORE}" "${KDIR}" "${TSUFFIX}" "${KERNEL_DEFINITIONS}") endif () - From 18b41f3c4daa90092906f03291f08bcb8f66f57b Mon Sep 17 00:00:00 2001 From: moluopro Date: Thu, 9 Jul 2026 13:42:12 +0800 Subject: [PATCH 3/6] interface: gate ARM64 STRMM direct fast path --- interface/trsm.c | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/interface/trsm.c b/interface/trsm.c index 069df1aa5c..988b59f6fd 100644 --- a/interface/trsm.c +++ b/interface/trsm.c @@ -361,6 +361,7 @@ void CNAME(enum CBLAS_ORDER order, if (args.m == 0 || args.n == 0) return; +#if defined(TRMM) #if !defined(COMPLEX) && !defined(DOUBLE) && !defined(BFLOAT16) && !defined(HFLOAT16) #if defined(ARCH_ARM64) && (defined(USE_STRMM_KERNEL_DIRECT)||defined(DYNAMIC_ARCH)) #if defined(DYNAMIC_ARCH) @@ -371,15 +372,17 @@ if (strcmp(gotoblas_corename(), "armv9sme") == 0 ) #endif if (order == CblasRowMajor && Diag == CblasNonUnit && Side == CblasLeft && m == lda && n == ldb) { - if (Trans == CblasNoTrans) { + if (Trans == CblasNoTrans || Trans == CblasConjNoTrans) { (Uplo == CblasUpper ? STRMM_DIRECT_LNUN : STRMM_DIRECT_LNLN)(m, n, alpha, a, lda, b, ldb); - } else if (Trans == CblasTrans) { + return; + } else if (Trans == CblasTrans || Trans == CblasConjTrans) { (Uplo == CblasUpper ? STRMM_DIRECT_LTUN : STRMM_DIRECT_LTLN)(m, n, alpha, a, lda, b, ldb); + return; } - return; } #endif #endif +#endif From 95cae64181490e9642d5302154f4a984da3094f7 Mon Sep 17 00:00:00 2001 From: moluopro Date: Thu, 9 Jul 2026 13:42:48 +0800 Subject: [PATCH 4/6] arm64: harden SME SGEMM direct heuristic --- kernel/arm64/sgemm_direct_performant.c | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/kernel/arm64/sgemm_direct_performant.c b/kernel/arm64/sgemm_direct_performant.c index a20670e8b7..1e9dbf0c3b 100644 --- a/kernel/arm64/sgemm_direct_performant.c +++ b/kernel/arm64/sgemm_direct_performant.c @@ -3,13 +3,13 @@ int CNAME(BLASLONG M, BLASLONG N, BLASLONG K) { -if (M<3) return 0; - unsigned long long mnk = M * N * K; + if (M < 3 || N <= 0 || K <= 0) + return 0; + + unsigned long long mnk = (unsigned long long)M * (unsigned long long)N * (unsigned long long)K; /* benchmark performance on M4 peaks around 512 and crosses the graph of the NEON SGEMM at about 3100 */ - if (mnk >= 3100L * 3100L * 3100L) + if (mnk >= 3100ULL * 3100ULL * 3100ULL) return 0; - + return 1; } - - From 6fdc8db640f98d5099546f72abb5033251360c2f Mon Sep 17 00:00:00 2001 From: moluopro Date: Thu, 9 Jul 2026 13:43:51 +0800 Subject: [PATCH 5/6] arm64: tighten SME STRMM direct variants --- kernel/arm64/strmm_direct_arm64_sme1.c | 534 +++++++++++++------------ 1 file changed, 275 insertions(+), 259 deletions(-) diff --git a/kernel/arm64/strmm_direct_arm64_sme1.c b/kernel/arm64/strmm_direct_arm64_sme1.c index 8ecc1c2191..926d20299d 100644 --- a/kernel/arm64/strmm_direct_arm64_sme1.c +++ b/kernel/arm64/strmm_direct_arm64_sme1.c @@ -1,259 +1,275 @@ -/* - Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. - SPDX-License-Identifier: BSD-3-Clause-Clear -*/ - -#include "common.h" -#include -#include -//#include "sme_abi.h" -#if defined(HAVE_SME) - -#if defined(__ARM_FEATURE_SME) && defined(__clang__) && __clang_major__ >= 16 -#include -#endif - -/* Function Definitions */ -static uint64_t sve_cntw() { - uint64_t cnt; - asm volatile( - "rdsvl %[res], #1\n" - "lsr %[res], %[res], #2\n" - : [res] "=r" (cnt) :: - ); - return cnt; -} - -#if defined(__ARM_FEATURE_SME) && defined(__ARM_FEATURE_LOCALLY_STREAMING) && defined(__clang__) && __clang_major__ >= 16 - -// Transpose 1SVL x N panel of A -__attribute__((always_inline)) -inline static void transpose_panel_lower(const float *restrict a, float *restrict b, - uint64_t rows, uint64_t cols, - uint64_t a_step, uint64_t rows_index) -__arm_out("za") __arm_streaming { - // for Lower Trangular Matrix - uint64_t svl = svcntw(); - uint64_t col_batch = svl; - - svzero_za(); - uint64_t last_rows_index = rows_index + rows - 1; - for (uint64_t k = 0; k < cols; k += col_batch) { - if (last_rows_index < k) { - // Early exit: if all rows are above the diagonal, no valid elements remain - break; - } - // Load to horizontal slices - for (uint64_t row = 0; row < rows; row++) { - svbool_t pg_row = svwhilelt_b32_u64(k, MIN(rows_index + row + 1, cols)); - svld1_hor_za32(0, row, pg_row, &a[row * a_step + k]); - } - - // Save from vertical slices - col_batch = MIN(col_batch, cols - k); - for (uint64_t col = 0; col < col_batch; col++) { - svst1_ver_za32(0, col, svptrue_b32(), &b[(col + k) * svl]); - } - } -} - -__attribute__((always_inline)) -inline static void transpose_panel_upper(const float *restrict a, float *restrict b, - uint64_t rows, uint64_t cols, - uint64_t a_step, uint64_t rows_index) -__arm_out("za") __arm_streaming { - // for Upper Trangular Matrix - uint64_t svl = svcntw(); - uint64_t col_batch = svl; - - svzero_za(); - // Start from column k = rows_index to ensure we only process the upper triangle (k >= rows_index) - for (uint64_t k = rows_index; k < cols; k += col_batch) { - // Load to horizontal slices - for (uint64_t row = 0; row < rows; row++) { - svbool_t pg_row = svwhilelt_b32_u64(k, cols); - svld1_hor_za32(0, row, pg_row, &a[row * a_step + k]); - } - - // Save from vertical slices - col_batch = MIN(col_batch, cols - k); - for (uint64_t col = 0, real_col = k; col < col_batch; col++, real_col++) { - // Only the upper triangular part of the matrix is stored. - svbool_t pg_col = svwhilelt_b32_u64(rows_index, real_col + 1); - svst1_ver_za32(0, col, pg_col, &b[(col + k) * svl]); - } - } -} - -__arm_new("za") __arm_locally_streaming -static void strmm_direct_sme1_preprocess(uint64_t nbr, uint64_t nbc, - const float *restrict a, float *restrict a_mod) { - const uint64_t num_rows = nbr; - uint64_t row_batch = svcntsw(); - for (uint64_t row_idx = 0; row_idx < num_rows; row_idx += row_batch) { - // Transpose 1SVL x N panel of A - row_batch = MIN(row_batch, num_rows - row_idx); -#if !defined(UPPER) - transpose_panel_lower(&a[row_idx * nbc], &a_mod[row_idx * nbc], row_batch, nbc, nbc, row_idx); -#else - transpose_panel_upper(&a[row_idx * nbc], &a_mod[row_idx * nbc], row_batch, nbc, nbc, row_idx); -#endif - } -} - -// Outer product kernel. -// Computes a 2SVL x 2SVL block of C, utilizing all four FP32 tiles of ZA. -__attribute__((always_inline)) inline void -kernel_2x2(const float *A, const float *B, float *C, size_t shared_dim, - size_t ldc, size_t block_rows, size_t block_cols, float alpha, uint64_t row_idx) - __arm_out("za") __arm_streaming { - const uint64_t svl = svcntw(); - size_t ldb = ldc; - // Predicate set-up - svbool_t pg = svptrue_b32(); - svbool_t pg_a_0 = svwhilelt_b32_u64(0, block_rows); - svbool_t pg_a_1 = svwhilelt_b32_u64(svl, block_rows); - -#if (!defined(TRANSA) && !defined(UPPER)) || (defined(TRANSA) && defined(UPPER)) -#define pg_a_0_full pg_a_0 -#define pg_a_1_full pg_a_1 -#endif - svbool_t pg_b_0 = svwhilelt_b32_u64(0, block_cols); - svbool_t pg_b_1 = svwhilelt_b32_u64(svl, block_cols); - -#define pg_c_0 pg_b_0 -#define pg_c_1 pg_b_1 - - svzero_za(); - svfloat32_t alpha_vec = svdup_f32(alpha); - // Iterate through shared dimension (K) -#if (!defined(TRANSA) && defined(UPPER)) || (defined(TRANSA) && !defined(UPPER)) - for (size_t k = row_idx, valid_index = 1; k < shared_dim; k++,valid_index++) { - pg_a_0 = svwhilelt_b32_u64(0, MIN(valid_index, block_rows)); - pg_a_1 = svwhilelt_b32_u64(svl, MIN(valid_index, block_rows)); -#else - for (size_t k = 0; k < MIN(row_idx + block_rows, shared_dim); k++) { - // If k exceeds row_idx, mask out rows before (k - row_idx) - // This ensures only valid rows are included for lower triangular logic. - if (k > row_idx) { - pg_a_0 = svnot_b_z(pg_a_0_full, svwhilelt_b32_u64(0, k - row_idx)); - pg_a_1 = svnot_b_z(pg_a_1_full, svwhilelt_b32_u64(svl, k - row_idx)); - } -#endif - -#if !defined(TRANSA) - // Load column of A - svfloat32_t col_a_0 = svld1(pg_a_0, &A[k * svl]); - svfloat32_t col_a_1 = svld1(pg_a_1, &A[(k + shared_dim) * svl]); -#else - svfloat32_t col_a_0 = svld1(pg_a_0, &A[k * shared_dim]); - svfloat32_t col_a_1 = svld1(pg_a_1, &A[k * shared_dim + svl]); -#endif - col_a_0 = svmul_x(pg_a_0, alpha_vec, col_a_0); - col_a_1 = svmul_x(pg_a_1, alpha_vec, col_a_1); - // Load row of B - svfloat32_t row_b_0 = svld1(pg_b_0, &B[k * ldb]); - svfloat32_t row_b_1 = svld1(pg_b_1, &B[k * ldb + svl]); - // Perform outer product - svmopa_za32_m(/*tile*/0, pg_a_0, pg, col_a_0, row_b_0); - svmopa_za32_m(/*tile*/1, pg_a_0, pg, col_a_0, row_b_1); - svmopa_za32_m(/*tile*/2, pg_a_1, pg, col_a_1, row_b_0); - svmopa_za32_m(/*tile*/3, pg_a_1, pg, col_a_1, row_b_1); - } - - // Store to C from ZA - for (size_t i = 0; i < MIN(svl, block_rows); i++) { - svst1_hor_za32(/*tile*/0, /*slice*/i, pg_c_0, &C[i * ldc]); - svst1_hor_za32(/*tile*/1, /*slice*/i, pg_c_1, &C[i * ldc + svl]); - } - for (size_t i = svl; i < block_rows; i++) { - svst1_hor_za32(/*tile*/2, /*slice*/i, pg_c_0, &C[i * ldc]); - svst1_hor_za32(/*tile*/3, /*slice*/i, pg_c_1, &C[i * ldc + svl]); - } - -} - -__arm_new("za") __arm_locally_streaming -static inline void strmm_direct_alpha_sme1_2VLx2VL(uint64_t m, uint64_t k, uint64_t n, const float* alpha,\ - const float *ba, float *restrict bb) { - const uint64_t num_rows = m; - const uint64_t num_cols = n; - - const float *restrict a_ptr = ba; - const float *restrict b_ptr = bb; - float *restrict c_ptr = bb; - - const uint64_t svl = svcntw(); - const uint64_t svl_x2 = 2*svl; - const uint64_t ldc = n; - - - uint64_t row_idx = 0; -#if (!defined(TRANSA) && defined(UPPER)) || (defined(TRANSA) && !defined(UPPER)) - // 2x2 loop - uint64_t row_batch = svl_x2; - // Block over rows of C (panels of A) - for (; row_idx < num_rows; row_idx += row_batch) { - row_batch = MIN(row_batch, num_rows - row_idx); -#else - // Calculate the remainder of num_rows divided by 2VL to determine tail tile size - uint64_t row_batch = num_rows % svl_x2; - // If there's no remainder, use full tile size (2VL) for initial batch - if (row_batch == 0) row_batch = svl_x2; - // Loop from bottom to top, processing rows in batches - for (uint64_t index = num_rows; index > 0; index -= row_batch, row_batch = svl_x2) { - // Compute the starting row index for the current batch - row_idx = index - row_batch; -#endif - uint64_t col_idx = 0; - uint64_t col_batch = svl_x2; - // Block over column dimension of C - for (; col_idx < num_cols; col_idx += col_batch) { - col_batch = MIN(col_batch, num_cols - col_idx); -#if !defined(TRANSA) - kernel_2x2(&a_ptr[row_idx * k], &b_ptr[col_idx], - &c_ptr[row_idx * ldc + col_idx], k, - ldc, row_batch, col_batch, *alpha, row_idx); -#else - kernel_2x2(&a_ptr[row_idx], &b_ptr[col_idx], - &c_ptr[row_idx * ldc + col_idx], k, - ldc, row_batch, col_batch, *alpha, row_idx); -#endif - } - } - - return; -} - -#else -static void strmm_direct_sme1_preprocess(uint64_t nbr, uint64_t nbc, - const float *restrict a, float *restrict a_mod) {} -static void strmm_direct_alpha_sme1_2VLx2VL(uint64_t m, uint64_t k, uint64_t n, const float* alpha,\ - const float *ba, float *restrict bb){} -#endif - -void CNAME (BLASLONG M, BLASLONG N, float alpha, float * __restrict A,\ - BLASLONG strideA, float * __restrict B, BLASLONG strideB){ -#if !defined(TRANSA) - uint64_t m_mod, vl_elms; - - vl_elms = sve_cntw(); - - m_mod = (((uint64_t)M + vl_elms - 1) / vl_elms) * vl_elms; - - float *A_mod = (float *) malloc(m_mod*M*sizeof(float)); - strmm_direct_sme1_preprocess(M, M, A, A_mod); - /* Calculate B = alpha*A*B*/ - strmm_direct_alpha_sme1_2VLx2VL(M, M, N, &alpha, A_mod, B); - free(A_mod); -#else - strmm_direct_alpha_sme1_2VLx2VL(M, M, N, &alpha, A, B); -#endif -} - -#else -void CNAME (BLASLONG M, BLASLONG N, float alpha, float * __restrict A,\ - BLASLONG strideA, float * __restrict B, BLASLONG strideB){ - } - -#endif +/* + Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. + SPDX-License-Identifier: BSD-3-Clause-Clear +*/ + +#include "common.h" +#include +#include +//#include "sme_abi.h" +#if defined(HAVE_SME) + +#if defined(__ARM_FEATURE_SME) && defined(__clang__) && __clang_major__ >= 16 +#include +#endif + +#if !defined(TRANSA) +/* Function Definitions */ +static uint64_t sve_cntw() { + uint64_t cnt; + asm volatile( + "rdsvl %[res], #1\n" + "lsr %[res], %[res], #2\n" + : [res] "=r" (cnt) :: + ); + return cnt; +} +#endif + +#if defined(__ARM_FEATURE_SME) && defined(__ARM_FEATURE_LOCALLY_STREAMING) && defined(__clang__) && __clang_major__ >= 16 + +#if !defined(TRANSA) +// Transpose 1SVL x N panel of A +__attribute__((always_inline)) +inline static void transpose_panel_lower(const float *restrict a, float *restrict b, + uint64_t rows, uint64_t cols, + uint64_t a_step, uint64_t rows_index) +__arm_out("za") __arm_streaming { + // for Lower Trangular Matrix + uint64_t svl = svcntw(); + uint64_t col_batch = svl; + + svzero_za(); + uint64_t last_rows_index = rows_index + rows - 1; + for (uint64_t k = 0; k < cols; k += col_batch) { + if (last_rows_index < k) { + // Early exit: if all rows are above the diagonal, no valid elements remain + break; + } + // Load to horizontal slices + for (uint64_t row = 0; row < rows; row++) { + svbool_t pg_row = svwhilelt_b32_u64(k, MIN(rows_index + row + 1, cols)); + svld1_hor_za32(0, row, pg_row, &a[row * a_step + k]); + } + + // Save from vertical slices + col_batch = MIN(col_batch, cols - k); + for (uint64_t col = 0; col < col_batch; col++) { + svst1_ver_za32(0, col, svptrue_b32(), &b[(col + k) * svl]); + } + } +} + +__attribute__((always_inline)) +inline static void transpose_panel_upper(const float *restrict a, float *restrict b, + uint64_t rows, uint64_t cols, + uint64_t a_step, uint64_t rows_index) +__arm_out("za") __arm_streaming { + // for Upper Trangular Matrix + uint64_t svl = svcntw(); + uint64_t col_batch = svl; + + svzero_za(); + // Start from column k = rows_index to ensure we only process the upper triangle (k >= rows_index) + for (uint64_t k = rows_index; k < cols; k += col_batch) { + // Load to horizontal slices + for (uint64_t row = 0; row < rows; row++) { + svbool_t pg_row = svwhilelt_b32_u64(k, cols); + svld1_hor_za32(0, row, pg_row, &a[row * a_step + k]); + } + + // Save from vertical slices + col_batch = MIN(col_batch, cols - k); + for (uint64_t col = 0, real_col = k; col < col_batch; col++, real_col++) { + // Only the upper triangular part of the matrix is stored. + svbool_t pg_col = svwhilelt_b32_u64(rows_index, real_col + 1); + svst1_ver_za32(0, col, pg_col, &b[(col + k) * svl]); + } + } +} + +__arm_new("za") __arm_locally_streaming +static void strmm_direct_sme1_preprocess(uint64_t nbr, uint64_t nbc, + const float *restrict a, float *restrict a_mod) { + const uint64_t num_rows = nbr; + uint64_t row_batch = svcntsw(); + for (uint64_t row_idx = 0; row_idx < num_rows; row_idx += row_batch) { + // Transpose 1SVL x N panel of A + row_batch = MIN(row_batch, num_rows - row_idx); +#if !defined(UPPER) + transpose_panel_lower(&a[row_idx * nbc], &a_mod[row_idx * nbc], row_batch, nbc, nbc, row_idx); +#else + transpose_panel_upper(&a[row_idx * nbc], &a_mod[row_idx * nbc], row_batch, nbc, nbc, row_idx); +#endif + } +} +#endif + +// Outer product kernel. +// Computes a 2SVL x 2SVL block of C, utilizing all four FP32 tiles of ZA. +__attribute__((always_inline)) inline void +kernel_2x2(const float *A, const float *B, float *C, size_t shared_dim, + size_t ldc, size_t block_rows, size_t block_cols, float alpha, uint64_t row_idx) + __arm_out("za") __arm_streaming { + const uint64_t svl = svcntw(); + size_t ldb = ldc; + // Predicate set-up + svbool_t pg = svptrue_b32(); + svbool_t pg_a_0 = svwhilelt_b32_u64(0, block_rows); + svbool_t pg_a_1 = svwhilelt_b32_u64(svl, block_rows); + +#if (!defined(TRANSA) && !defined(UPPER)) || (defined(TRANSA) && defined(UPPER)) +#define pg_a_0_full pg_a_0 +#define pg_a_1_full pg_a_1 +#endif + svbool_t pg_b_0 = svwhilelt_b32_u64(0, block_cols); + svbool_t pg_b_1 = svwhilelt_b32_u64(svl, block_cols); + +#define pg_c_0 pg_b_0 +#define pg_c_1 pg_b_1 + + svzero_za(); + svfloat32_t alpha_vec = svdup_f32(alpha); + // Iterate through shared dimension (K) +#if (!defined(TRANSA) && defined(UPPER)) || (defined(TRANSA) && !defined(UPPER)) + for (size_t k = row_idx, valid_index = 1; k < shared_dim; k++,valid_index++) { + pg_a_0 = svwhilelt_b32_u64(0, MIN(valid_index, block_rows)); + pg_a_1 = svwhilelt_b32_u64(svl, MIN(valid_index, block_rows)); +#else + for (size_t k = 0; k < MIN(row_idx + block_rows, shared_dim); k++) { + // If k exceeds row_idx, mask out rows before (k - row_idx) + // This ensures only valid rows are included for lower triangular logic. + if (k > row_idx) { + pg_a_0 = svnot_b_z(pg_a_0_full, svwhilelt_b32_u64(0, k - row_idx)); + pg_a_1 = svnot_b_z(pg_a_1_full, svwhilelt_b32_u64(svl, k - row_idx)); + } +#endif + +#if !defined(TRANSA) + // Load column of A + svfloat32_t col_a_0 = svld1(pg_a_0, &A[k * svl]); + svfloat32_t col_a_1 = svld1(pg_a_1, &A[(k + shared_dim) * svl]); +#else + svfloat32_t col_a_0 = svld1(pg_a_0, &A[k * shared_dim]); + svfloat32_t col_a_1 = svld1(pg_a_1, &A[k * shared_dim + svl]); +#endif + col_a_0 = svmul_x(pg_a_0, alpha_vec, col_a_0); + col_a_1 = svmul_x(pg_a_1, alpha_vec, col_a_1); + // Load row of B + svfloat32_t row_b_0 = svld1(pg_b_0, &B[k * ldb]); + svfloat32_t row_b_1 = svld1(pg_b_1, &B[k * ldb + svl]); + // Perform outer product + svmopa_za32_m(/*tile*/0, pg_a_0, pg, col_a_0, row_b_0); + svmopa_za32_m(/*tile*/1, pg_a_0, pg, col_a_0, row_b_1); + svmopa_za32_m(/*tile*/2, pg_a_1, pg, col_a_1, row_b_0); + svmopa_za32_m(/*tile*/3, pg_a_1, pg, col_a_1, row_b_1); + } + + // Store to C from ZA + for (size_t i = 0; i < MIN(svl, block_rows); i++) { + svst1_hor_za32(/*tile*/0, /*slice*/i, pg_c_0, &C[i * ldc]); + svst1_hor_za32(/*tile*/1, /*slice*/i, pg_c_1, &C[i * ldc + svl]); + } + for (size_t i = svl; i < block_rows; i++) { + svst1_hor_za32(/*tile*/2, /*slice*/i, pg_c_0, &C[i * ldc]); + svst1_hor_za32(/*tile*/3, /*slice*/i, pg_c_1, &C[i * ldc + svl]); + } + +} + +__arm_new("za") __arm_locally_streaming +static inline void strmm_direct_alpha_sme1_2VLx2VL(uint64_t m, uint64_t k, uint64_t n, const float* alpha,\ + const float *ba, float *restrict bb) { + const uint64_t num_rows = m; + const uint64_t num_cols = n; + + const float *restrict a_ptr = ba; + const float *restrict b_ptr = bb; + float *restrict c_ptr = bb; + + const uint64_t svl = svcntw(); + const uint64_t svl_x2 = 2*svl; + const uint64_t ldc = n; + + + uint64_t row_idx = 0; +#if (!defined(TRANSA) && defined(UPPER)) || (defined(TRANSA) && !defined(UPPER)) + // 2x2 loop + uint64_t row_batch = svl_x2; + // Block over rows of C (panels of A) + for (; row_idx < num_rows; row_idx += row_batch) { + row_batch = MIN(row_batch, num_rows - row_idx); +#else + // Calculate the remainder of num_rows divided by 2VL to determine tail tile size + uint64_t row_batch = num_rows % svl_x2; + // If there's no remainder, use full tile size (2VL) for initial batch + if (row_batch == 0) row_batch = svl_x2; + // Loop from bottom to top, processing rows in batches + for (uint64_t index = num_rows; index > 0; index -= row_batch, row_batch = svl_x2) { + // Compute the starting row index for the current batch + row_idx = index - row_batch; +#endif + uint64_t col_idx = 0; + uint64_t col_batch = svl_x2; + // Block over column dimension of C + for (; col_idx < num_cols; col_idx += col_batch) { + col_batch = MIN(col_batch, num_cols - col_idx); +#if !defined(TRANSA) + kernel_2x2(&a_ptr[row_idx * k], &b_ptr[col_idx], + &c_ptr[row_idx * ldc + col_idx], k, + ldc, row_batch, col_batch, *alpha, row_idx); +#else + kernel_2x2(&a_ptr[row_idx], &b_ptr[col_idx], + &c_ptr[row_idx * ldc + col_idx], k, + ldc, row_batch, col_batch, *alpha, row_idx); +#endif + } + } + + return; +} + +#else +#if !defined(TRANSA) +static void strmm_direct_sme1_preprocess(uint64_t nbr, uint64_t nbc, + const float *restrict a, float *restrict a_mod) {} +#endif +static void strmm_direct_alpha_sme1_2VLx2VL(uint64_t m, uint64_t k, uint64_t n, const float* alpha,\ + const float *ba, float *restrict bb){} +#endif + +void CNAME (BLASLONG M, BLASLONG N, float alpha, float * __restrict A,\ + BLASLONG strideA, float * __restrict B, BLASLONG strideB){ + if (alpha == 0.0f) { + for (BLASLONG i = 0; i < M; i++) { + float *row = B + i * strideB; + for (BLASLONG j = 0; j < N; j++) { + row[j] = 0.0f; + } + } + return; + } + +#if !defined(TRANSA) + uint64_t m_mod, vl_elms; + + vl_elms = sve_cntw(); + + m_mod = (((uint64_t)M + vl_elms - 1) / vl_elms) * vl_elms; + + float *A_mod = (float *) malloc(m_mod*M*sizeof(float)); + strmm_direct_sme1_preprocess(M, M, A, A_mod); + /* Calculate B = alpha*A*B*/ + strmm_direct_alpha_sme1_2VLx2VL(M, M, N, &alpha, A_mod, B); + free(A_mod); +#else + strmm_direct_alpha_sme1_2VLx2VL(M, M, N, &alpha, A, B); +#endif +} + +#else +void CNAME (BLASLONG M, BLASLONG N, float alpha, float * __restrict A,\ + BLASLONG strideA, float * __restrict B, BLASLONG strideB){ + } + +#endif From 06d553d0b1ff44d5f3419710ba69d077297dfe06 Mon Sep 17 00:00:00 2001 From: moluopro Date: Thu, 9 Jul 2026 13:44:29 +0800 Subject: [PATCH 6/6] arm64: normalize SME direct source line endings --- .../sgemm_direct_alpha_beta_arm64_sme1.c | 450 +++++++------- kernel/arm64/sgemm_direct_arm64_sme1.c | 198 +++---- .../ssymm_direct_alpha_beta_arm64_sme1.c | 510 ++++++++-------- .../ssyrk_direct_alpha_beta_arm64_sme1.c | 556 +++++++++--------- 4 files changed, 857 insertions(+), 857 deletions(-) diff --git a/kernel/arm64/sgemm_direct_alpha_beta_arm64_sme1.c b/kernel/arm64/sgemm_direct_alpha_beta_arm64_sme1.c index 82ae798f05..97c92fa352 100644 --- a/kernel/arm64/sgemm_direct_alpha_beta_arm64_sme1.c +++ b/kernel/arm64/sgemm_direct_alpha_beta_arm64_sme1.c @@ -1,225 +1,225 @@ -/* - Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. - SPDX-License-Identifier: BSD-3-Clause-Clear -*/ - -#include "common.h" -#include -#include - -#if defined(DYNAMIC_ARCH) -#define COMBINE(a,b) a ## b -#define COMBINE2(a,b) COMBINE(a,b) -#define SME1_PREPROCESS_BASE sgemm_direct_sme1_preprocess -#define SME1_PREPROCESS COMBINE2(SME1_PREPROCESS_BASE,TS) -#define SME1_KERNEL2X2_BASE sgemm_direct_alpha_beta_sme1_2VLx2VL -#define SME1_KERNEL2X2 COMBINE2(SME1_KERNEL2X2_BASE,TS) -#else -#define SME1_PREPROCESS sgemm_direct_sme1_preprocess -#define SME1_KERNEL2X2 sgemm_direct_alpha_beta_sme1_2VLx2VL -#endif - -/* Function prototypes */ -extern void SME1_PREPROCESS(uint64_t nbr, uint64_t nbc,\ - const float * restrict a, float * a_mod); - -#if defined(HAVE_SME) -#include "sme_abi.h" - -#if defined(__ARM_FEATURE_SME) && defined(__clang__) && __clang_major__ >= 16 -#include -#endif - -/* Function Definitions */ -static uint64_t sve_cntw() { - uint64_t cnt; - asm volatile( - "rdsvl %[res], #1\n" - "lsr %[res], %[res], #2\n" - : [res] "=r" (cnt) :: - ); - return cnt; -} - -#if defined(__ARM_FEATURE_SME) && defined(__ARM_FEATURE_LOCALLY_STREAMING) && defined(__clang__) && __clang_major__ >= 16 -// Outer product kernel. -// Computes a 2SVL x 2SVL block of C, utilizing all four FP32 tiles of ZA. -__attribute__((always_inline)) inline void -kernel_2x2(const float *A, const float *B, float *C, size_t shared_dim, - size_t ldc, size_t block_rows, size_t block_cols, float alpha, float beta) - __arm_out("za") __arm_streaming { - - const uint64_t svl = svcntw(); - size_t ldb = ldc; - // Predicate set-up - svbool_t pg = svptrue_b32(); - svbool_t pg_a_0 = svwhilelt_b32_u64(0, block_rows); - svbool_t pg_a_1 = svwhilelt_b32_u64(svl, block_rows); - - svbool_t pg_b_0 = svwhilelt_b32_u64(0, block_cols); - svbool_t pg_b_1 = svwhilelt_b32_u64(svl, block_cols); - -#define pg_c_0 pg_b_0 -#define pg_c_1 pg_b_1 - - svzero_za(); - // beta == 0 must not read C; ZA is already initialized to zero. - if (beta != 0.0f) { - svfloat32_t beta_vec = svdup_f32(beta); - // Load C to ZA - for (size_t i = 0; i < MIN(svl, block_rows); i++) { - svfloat32_t row_c_0 = svld1(pg_c_0, &C[i * ldc]); - row_c_0 = svmul_x(pg, beta_vec, row_c_0); - svwrite_hor_za32_f32_m(/*tile*/0, /*slice*/i, pg_c_0, row_c_0); - - svfloat32_t row_c_1 = svld1(pg_c_1, &C[i * ldc + svl]); - row_c_1 = svmul_x(pg, beta_vec, row_c_1); - svwrite_hor_za32_f32_m(/*tile*/1, /*slice*/i, pg_c_1, row_c_1); - } - for (size_t i = svl; i < block_rows; i++) { - svfloat32_t row_c_0 = svld1(pg_c_0, &C[i * ldc]); - row_c_0 = svmul_x(pg, beta_vec, row_c_0); - svwrite_hor_za32_f32_m(/*tile*/2, /*slice*/i, pg_c_0, row_c_0); - - svfloat32_t row_c_1 = svld1(pg_c_1, &C[i * ldc + svl]); - row_c_1 = svmul_x(pg, beta_vec, row_c_1); - svwrite_hor_za32_f32_m(/*tile*/3, /*slice*/i, pg_c_1, row_c_1); - } - } - - svfloat32_t alpha_vec = svdup_f32(alpha); - // Iterate through shared dimension (K) - for (size_t k = 0; k < shared_dim; k++) { - // Load column of A - svfloat32_t col_a_0 = svld1(pg_a_0, &A[k * svl]); - col_a_0 = svmul_x(pg, alpha_vec, col_a_0); - svfloat32_t col_a_1 = svld1(pg_a_1, &A[(k + shared_dim) * svl]); - col_a_1 = svmul_x(pg, alpha_vec, col_a_1); - // Load row of B - svfloat32_t row_b_0 = svld1(pg_b_0, &B[k * ldb]); - svfloat32_t row_b_1 = svld1(pg_b_1, &B[k * ldb + svl]); - // Perform outer product - svmopa_za32_m(/*tile*/0, pg, pg, col_a_0, row_b_0); - svmopa_za32_m(/*tile*/1, pg, pg, col_a_0, row_b_1); - svmopa_za32_m(/*tile*/2, pg, pg, col_a_1, row_b_0); - svmopa_za32_m(/*tile*/3, pg, pg, col_a_1, row_b_1); - } - - // Store to C from ZA - for (size_t i = 0; i < MIN(svl, block_rows); i++) { - svst1_hor_za32(/*tile*/0, /*slice*/i, pg_c_0, &C[i * ldc]); - svst1_hor_za32(/*tile*/1, /*slice*/i, pg_c_1, &C[i * ldc + svl]); - } - for (size_t i = svl; i < block_rows; i++) { - svst1_hor_za32(/*tile*/2, /*slice*/i, pg_c_0, &C[i * ldc]); - svst1_hor_za32(/*tile*/3, /*slice*/i, pg_c_1, &C[i * ldc + svl]); - } -return; -} - -__arm_new("za") __arm_locally_streaming -void SME1_KERNEL2X2(uint64_t m, uint64_t k, uint64_t n, const float* alpha,\ - const float *ba, const float *restrict bb, const float* beta,\ - float *restrict C) { - - const uint64_t num_rows = m; - const uint64_t num_cols = n; - - const float *restrict a_ptr = ba; - const float *restrict b_ptr = bb; - float *restrict c_ptr = C; - - const uint64_t svl = svcntw(); - const uint64_t ldc = n; - - // Block over rows of C (panels of A) - uint64_t row_idx = 0; - - // 2x2 loop - uint64_t row_batch = 2*svl; - - // Block over row dimension of C - for (; row_idx < num_rows; row_idx += row_batch) { - row_batch = MIN(row_batch, num_rows - row_idx); - - uint64_t col_idx = 0; - uint64_t col_batch = 2*svl; - - // Block over column dimension of C - for (; col_idx < num_cols; col_idx += col_batch) { - col_batch = MIN(col_batch, num_cols - col_idx); - - kernel_2x2(&a_ptr[row_idx * k], &b_ptr[col_idx], - &c_ptr[row_idx * ldc + col_idx], k, - ldc, row_batch, col_batch, *alpha, *beta); - } - } - return; -} - -#else -void SME1_KERNEL2X2(uint64_t m, uint64_t k, uint64_t n, const float* alpha,\ - const float *ba, const float *restrict bb, const float* beta,\ - float *restrict C){fprintf(stderr,"empty sgemm_alpha_beta2x2 should never get called!!!\n");} -#endif - -/*void sgemm_kernel_direct (BLASLONG M, BLASLONG N, BLASLONG K,\ - float * __restrict A, BLASLONG strideA, float * __restrict B,\ - BLASLONG strideB , float * __restrict R, BLASLONG strideR) -*/ -void CNAME (BLASLONG M, BLASLONG N, BLASLONG K, float alpha, float * __restrict A,\ - BLASLONG strideA, float * __restrict B, BLASLONG strideB ,\ - float beta, float * __restrict R, BLASLONG strideR){ - if (alpha == 0.0f || K == 0) { - if (beta == 1.0f) - return; - SME1_KERNEL2X2(M, 0, N, &alpha, A, B, &beta, R); - return; - } - - uint64_t m_mod, vl_elms; - - vl_elms = sve_cntw(); - - m_mod = (((uint64_t)M + vl_elms - 1) / vl_elms) * vl_elms; - - float *A_mod = (float *) malloc(m_mod*K*sizeof(float)); - - /* Prevent compiler optimization by reading from memory instead - * of reading directly from vector (z) registers. - * */ - asm volatile("" : : :"p0", "p1", "p2", "p3", "p4", "p5", "p6", "p7", - "p8", "p9", "p10", "p11", "p12", "p13", "p14", "p15", "d8", "d9", "d10", "d11", "d12", "d13", "d14", "d15", - "z0", "z1", "z2", "z3", "z4", "z5", "z6", "z7", - "z8", "z9", "z10", "z11", "z12", "z13", "z14", "z15", - "z16", "z17", "z18", "z19", "z20", "z21", "z22", "z23", - "z24", "z25", "z26", "z27", "z28", "z29", "z30", "z31"); - - /* Pre-process the left matrix to make it suitable for - matrix sum of outer-product calculation - */ - - SME1_PREPROCESS(M, K, A, A_mod); - - asm volatile("" : : :"p0", "p1", "p2", "p3", "p4", "p5", "p6", "p7", - "p8", "p9", "p10", "p11", "p12", "p13", "p14", "p15","d8", "d9", "d10", "d11", "d12", "d13", "d14", "d15", - "z0", "z1", "z2", "z3", "z4", "z5", "z6", "z7", - "z8", "z9", "z10", "z11", "z12", "z13", "z14", "z15", - "z16", "z17", "z18", "z19", "z20", "z21", "z22", "z23", - "z24", "z25", "z26", "z27", "z28", "z29", "z30", "z31"); - - /* Calculate C = alpha*A*B + beta*C */ - - SME1_KERNEL2X2(M, K, N, &alpha, A_mod, B, &beta, R); - - free(A_mod); -} - -#else - -void CNAME (BLASLONG M, BLASLONG N, BLASLONG K, float alpha, float * __restrict A,\ - BLASLONG strideA, float * __restrict B, BLASLONG strideB ,\ - float beta, float * __restrict R, BLASLONG strideR){fprintf(stderr,"empty sgemm_direct_alpha_beta should not be called!!!\n");} -#endif - - +/* + Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. + SPDX-License-Identifier: BSD-3-Clause-Clear +*/ + +#include "common.h" +#include +#include + +#if defined(DYNAMIC_ARCH) +#define COMBINE(a,b) a ## b +#define COMBINE2(a,b) COMBINE(a,b) +#define SME1_PREPROCESS_BASE sgemm_direct_sme1_preprocess +#define SME1_PREPROCESS COMBINE2(SME1_PREPROCESS_BASE,TS) +#define SME1_KERNEL2X2_BASE sgemm_direct_alpha_beta_sme1_2VLx2VL +#define SME1_KERNEL2X2 COMBINE2(SME1_KERNEL2X2_BASE,TS) +#else +#define SME1_PREPROCESS sgemm_direct_sme1_preprocess +#define SME1_KERNEL2X2 sgemm_direct_alpha_beta_sme1_2VLx2VL +#endif + +/* Function prototypes */ +extern void SME1_PREPROCESS(uint64_t nbr, uint64_t nbc,\ + const float * restrict a, float * a_mod); + +#if defined(HAVE_SME) +#include "sme_abi.h" + +#if defined(__ARM_FEATURE_SME) && defined(__clang__) && __clang_major__ >= 16 +#include +#endif + +/* Function Definitions */ +static uint64_t sve_cntw() { + uint64_t cnt; + asm volatile( + "rdsvl %[res], #1\n" + "lsr %[res], %[res], #2\n" + : [res] "=r" (cnt) :: + ); + return cnt; +} + +#if defined(__ARM_FEATURE_SME) && defined(__ARM_FEATURE_LOCALLY_STREAMING) && defined(__clang__) && __clang_major__ >= 16 +// Outer product kernel. +// Computes a 2SVL x 2SVL block of C, utilizing all four FP32 tiles of ZA. +__attribute__((always_inline)) inline void +kernel_2x2(const float *A, const float *B, float *C, size_t shared_dim, + size_t ldc, size_t block_rows, size_t block_cols, float alpha, float beta) + __arm_out("za") __arm_streaming { + + const uint64_t svl = svcntw(); + size_t ldb = ldc; + // Predicate set-up + svbool_t pg = svptrue_b32(); + svbool_t pg_a_0 = svwhilelt_b32_u64(0, block_rows); + svbool_t pg_a_1 = svwhilelt_b32_u64(svl, block_rows); + + svbool_t pg_b_0 = svwhilelt_b32_u64(0, block_cols); + svbool_t pg_b_1 = svwhilelt_b32_u64(svl, block_cols); + +#define pg_c_0 pg_b_0 +#define pg_c_1 pg_b_1 + + svzero_za(); + // beta == 0 must not read C; ZA is already initialized to zero. + if (beta != 0.0f) { + svfloat32_t beta_vec = svdup_f32(beta); + // Load C to ZA + for (size_t i = 0; i < MIN(svl, block_rows); i++) { + svfloat32_t row_c_0 = svld1(pg_c_0, &C[i * ldc]); + row_c_0 = svmul_x(pg, beta_vec, row_c_0); + svwrite_hor_za32_f32_m(/*tile*/0, /*slice*/i, pg_c_0, row_c_0); + + svfloat32_t row_c_1 = svld1(pg_c_1, &C[i * ldc + svl]); + row_c_1 = svmul_x(pg, beta_vec, row_c_1); + svwrite_hor_za32_f32_m(/*tile*/1, /*slice*/i, pg_c_1, row_c_1); + } + for (size_t i = svl; i < block_rows; i++) { + svfloat32_t row_c_0 = svld1(pg_c_0, &C[i * ldc]); + row_c_0 = svmul_x(pg, beta_vec, row_c_0); + svwrite_hor_za32_f32_m(/*tile*/2, /*slice*/i, pg_c_0, row_c_0); + + svfloat32_t row_c_1 = svld1(pg_c_1, &C[i * ldc + svl]); + row_c_1 = svmul_x(pg, beta_vec, row_c_1); + svwrite_hor_za32_f32_m(/*tile*/3, /*slice*/i, pg_c_1, row_c_1); + } + } + + svfloat32_t alpha_vec = svdup_f32(alpha); + // Iterate through shared dimension (K) + for (size_t k = 0; k < shared_dim; k++) { + // Load column of A + svfloat32_t col_a_0 = svld1(pg_a_0, &A[k * svl]); + col_a_0 = svmul_x(pg, alpha_vec, col_a_0); + svfloat32_t col_a_1 = svld1(pg_a_1, &A[(k + shared_dim) * svl]); + col_a_1 = svmul_x(pg, alpha_vec, col_a_1); + // Load row of B + svfloat32_t row_b_0 = svld1(pg_b_0, &B[k * ldb]); + svfloat32_t row_b_1 = svld1(pg_b_1, &B[k * ldb + svl]); + // Perform outer product + svmopa_za32_m(/*tile*/0, pg, pg, col_a_0, row_b_0); + svmopa_za32_m(/*tile*/1, pg, pg, col_a_0, row_b_1); + svmopa_za32_m(/*tile*/2, pg, pg, col_a_1, row_b_0); + svmopa_za32_m(/*tile*/3, pg, pg, col_a_1, row_b_1); + } + + // Store to C from ZA + for (size_t i = 0; i < MIN(svl, block_rows); i++) { + svst1_hor_za32(/*tile*/0, /*slice*/i, pg_c_0, &C[i * ldc]); + svst1_hor_za32(/*tile*/1, /*slice*/i, pg_c_1, &C[i * ldc + svl]); + } + for (size_t i = svl; i < block_rows; i++) { + svst1_hor_za32(/*tile*/2, /*slice*/i, pg_c_0, &C[i * ldc]); + svst1_hor_za32(/*tile*/3, /*slice*/i, pg_c_1, &C[i * ldc + svl]); + } +return; +} + +__arm_new("za") __arm_locally_streaming +void SME1_KERNEL2X2(uint64_t m, uint64_t k, uint64_t n, const float* alpha,\ + const float *ba, const float *restrict bb, const float* beta,\ + float *restrict C) { + + const uint64_t num_rows = m; + const uint64_t num_cols = n; + + const float *restrict a_ptr = ba; + const float *restrict b_ptr = bb; + float *restrict c_ptr = C; + + const uint64_t svl = svcntw(); + const uint64_t ldc = n; + + // Block over rows of C (panels of A) + uint64_t row_idx = 0; + + // 2x2 loop + uint64_t row_batch = 2*svl; + + // Block over row dimension of C + for (; row_idx < num_rows; row_idx += row_batch) { + row_batch = MIN(row_batch, num_rows - row_idx); + + uint64_t col_idx = 0; + uint64_t col_batch = 2*svl; + + // Block over column dimension of C + for (; col_idx < num_cols; col_idx += col_batch) { + col_batch = MIN(col_batch, num_cols - col_idx); + + kernel_2x2(&a_ptr[row_idx * k], &b_ptr[col_idx], + &c_ptr[row_idx * ldc + col_idx], k, + ldc, row_batch, col_batch, *alpha, *beta); + } + } + return; +} + +#else +void SME1_KERNEL2X2(uint64_t m, uint64_t k, uint64_t n, const float* alpha,\ + const float *ba, const float *restrict bb, const float* beta,\ + float *restrict C){fprintf(stderr,"empty sgemm_alpha_beta2x2 should never get called!!!\n");} +#endif + +/*void sgemm_kernel_direct (BLASLONG M, BLASLONG N, BLASLONG K,\ + float * __restrict A, BLASLONG strideA, float * __restrict B,\ + BLASLONG strideB , float * __restrict R, BLASLONG strideR) +*/ +void CNAME (BLASLONG M, BLASLONG N, BLASLONG K, float alpha, float * __restrict A,\ + BLASLONG strideA, float * __restrict B, BLASLONG strideB ,\ + float beta, float * __restrict R, BLASLONG strideR){ + if (alpha == 0.0f || K == 0) { + if (beta == 1.0f) + return; + SME1_KERNEL2X2(M, 0, N, &alpha, A, B, &beta, R); + return; + } + + uint64_t m_mod, vl_elms; + + vl_elms = sve_cntw(); + + m_mod = (((uint64_t)M + vl_elms - 1) / vl_elms) * vl_elms; + + float *A_mod = (float *) malloc(m_mod*K*sizeof(float)); + + /* Prevent compiler optimization by reading from memory instead + * of reading directly from vector (z) registers. + * */ + asm volatile("" : : :"p0", "p1", "p2", "p3", "p4", "p5", "p6", "p7", + "p8", "p9", "p10", "p11", "p12", "p13", "p14", "p15", "d8", "d9", "d10", "d11", "d12", "d13", "d14", "d15", + "z0", "z1", "z2", "z3", "z4", "z5", "z6", "z7", + "z8", "z9", "z10", "z11", "z12", "z13", "z14", "z15", + "z16", "z17", "z18", "z19", "z20", "z21", "z22", "z23", + "z24", "z25", "z26", "z27", "z28", "z29", "z30", "z31"); + + /* Pre-process the left matrix to make it suitable for + matrix sum of outer-product calculation + */ + + SME1_PREPROCESS(M, K, A, A_mod); + + asm volatile("" : : :"p0", "p1", "p2", "p3", "p4", "p5", "p6", "p7", + "p8", "p9", "p10", "p11", "p12", "p13", "p14", "p15","d8", "d9", "d10", "d11", "d12", "d13", "d14", "d15", + "z0", "z1", "z2", "z3", "z4", "z5", "z6", "z7", + "z8", "z9", "z10", "z11", "z12", "z13", "z14", "z15", + "z16", "z17", "z18", "z19", "z20", "z21", "z22", "z23", + "z24", "z25", "z26", "z27", "z28", "z29", "z30", "z31"); + + /* Calculate C = alpha*A*B + beta*C */ + + SME1_KERNEL2X2(M, K, N, &alpha, A_mod, B, &beta, R); + + free(A_mod); +} + +#else + +void CNAME (BLASLONG M, BLASLONG N, BLASLONG K, float alpha, float * __restrict A,\ + BLASLONG strideA, float * __restrict B, BLASLONG strideB ,\ + float beta, float * __restrict R, BLASLONG strideR){fprintf(stderr,"empty sgemm_direct_alpha_beta should not be called!!!\n");} +#endif + + diff --git a/kernel/arm64/sgemm_direct_arm64_sme1.c b/kernel/arm64/sgemm_direct_arm64_sme1.c index 340245b105..b8837ea958 100644 --- a/kernel/arm64/sgemm_direct_arm64_sme1.c +++ b/kernel/arm64/sgemm_direct_arm64_sme1.c @@ -1,99 +1,99 @@ -/* - Copyright (c) 2025 Qualcomm Innovation Center, Inc. All rights reserved. - SPDX-License-Identifier: BSD-3-Clause-Clear -*/ - -#include "common.h" -#include -#include -#if defined(DYNAMIC_ARCH) -#define COMBINE(a,b) a ## b -#define COMBINE2(a,b) COMBINE(a,b) -#define SME1_PREPROCESS_BASE sgemm_direct_sme1_preprocess -#define SME1_PREPROCESS COMBINE2(SME1_PREPROCESS_BASE,TS) -#define SME1_DIRECT2X2_BASE sgemm_direct_sme1_2VLx2VL -#define SME1_DIRECT2X2 COMBINE2(SME1_DIRECT2X2_BASE,TS) -#else -#define SME1_PREPROCESS sgemm_direct_sme1_preprocess -#define SME1_DIRECT2X2 sgemm_direct_sme1_2VLx2VL -#endif -#if defined(HAVE_SME) -/* Function prototypes */ -extern void SME1_PREPROCESS(uint64_t nbr, uint64_t nbc,\ - const float * restrict a, float * a_mod) ; - -extern void SME1_DIRECT2X2(uint64_t m, uint64_t k, uint64_t n,\ - const float * matLeft,\ - const float * restrict matRight,\ - const float * restrict matResult) ; - -/* Function Definitions */ -static uint64_t sve_cntw() { - uint64_t cnt; - asm volatile( - "rdsvl %[res], #1\n" - "lsr %[res], %[res], #2\n" - : [res] "=r" (cnt) :: - ); - return cnt; -} - -/*void sgemm_kernel_direct (BLASLONG M, BLASLONG N, BLASLONG K,\ - float * __restrict A, BLASLONG strideA, float * __restrict B,\ - BLASLONG strideB , float * __restrict R, BLASLONG strideR) -*/ -void CNAME (BLASLONG M, BLASLONG N, BLASLONG K, float * __restrict A,\ - BLASLONG strideA, float * __restrict B, BLASLONG strideB ,\ - float * __restrict R, BLASLONG strideR){ - - uint64_t m_mod, vl_elms; - - vl_elms = sve_cntw(); - m_mod = (((uint64_t)M + vl_elms - 1) / vl_elms) * vl_elms; - - float *A_mod = (float *) malloc(m_mod*K*sizeof(float)); - - /* Prevent compiler optimization by reading from memory instead - * of reading directly from vector (z) registers. - * */ - asm volatile("" : : :"p0", "p1", "p2", "p3", "p4", "p5", "p6", "p7", - "p8", "p9", "p10", "p11", "p12", "p13", "p14", "p15", "d8", "d9", "d10", "d11", "d12", "d13", "d14", "d15", - "z0", "z1", "z2", "z3", "z4", "z5", "z6", "z7", - "z8", "z9", "z10", "z11", "z12", "z13", "z14", "z15", - "z16", "z17", "z18", "z19", "z20", "z21", "z22", "z23", - "z24", "z25", "z26", "z27", "z28", "z29", "z30", "z31"); - - /* Pre-process the left matrix to make it suitable for - matrix sum of outer-product calculation - */ - SME1_PREPROCESS(M, K, A, A_mod); - - /* Calculate C = A*B */ - SME1_DIRECT2X2(M, K, N, A_mod, B, R); - - asm volatile("" : : :"p0", "p1", "p2", "p3", "p4", "p5", "p6", "p7", - "p8", "p9", "p10", "p11", "p12", "p13", "p14", "p15", "d8", "d9", "d10", "d11", "d12", "d13", "d14", "d15", - "z0", "z1", "z2", "z3", "z4", "z5", "z6", "z7", - "z8", "z9", "z10", "z11", "z12", "z13", "z14", "z15", - "z16", "z17", "z18", "z19", "z20", "z21", "z22", "z23", - "z24", "z25", "z26", "z27", "z28", "z29", "z30", "z31"); - free(A_mod); -} - -#else - -void CNAME (BLASLONG M, BLASLONG N, BLASLONG K, float * __restrict A,\ - BLASLONG strideA, float * __restrict B, BLASLONG strideB ,\ - float * __restrict R, BLASLONG strideR){ -fprintf(stderr,"EMPTY sgemm_kernel_direct should never be called \n"); -} -void SME1_DIRECT2X2( uint64_t M , uint64_t K, uint64_t N,\ - const float * restrict A_base,\ - const float * restrict B_base,\ - const float * restrict C_base){}; -void SME1_PREPROCESS(uint64_t nbr, uint64_t nbc,\ - const float * restrict a, float * a_mod){}; - - -#endif - +/* + Copyright (c) 2025 Qualcomm Innovation Center, Inc. All rights reserved. + SPDX-License-Identifier: BSD-3-Clause-Clear +*/ + +#include "common.h" +#include +#include +#if defined(DYNAMIC_ARCH) +#define COMBINE(a,b) a ## b +#define COMBINE2(a,b) COMBINE(a,b) +#define SME1_PREPROCESS_BASE sgemm_direct_sme1_preprocess +#define SME1_PREPROCESS COMBINE2(SME1_PREPROCESS_BASE,TS) +#define SME1_DIRECT2X2_BASE sgemm_direct_sme1_2VLx2VL +#define SME1_DIRECT2X2 COMBINE2(SME1_DIRECT2X2_BASE,TS) +#else +#define SME1_PREPROCESS sgemm_direct_sme1_preprocess +#define SME1_DIRECT2X2 sgemm_direct_sme1_2VLx2VL +#endif +#if defined(HAVE_SME) +/* Function prototypes */ +extern void SME1_PREPROCESS(uint64_t nbr, uint64_t nbc,\ + const float * restrict a, float * a_mod) ; + +extern void SME1_DIRECT2X2(uint64_t m, uint64_t k, uint64_t n,\ + const float * matLeft,\ + const float * restrict matRight,\ + const float * restrict matResult) ; + +/* Function Definitions */ +static uint64_t sve_cntw() { + uint64_t cnt; + asm volatile( + "rdsvl %[res], #1\n" + "lsr %[res], %[res], #2\n" + : [res] "=r" (cnt) :: + ); + return cnt; +} + +/*void sgemm_kernel_direct (BLASLONG M, BLASLONG N, BLASLONG K,\ + float * __restrict A, BLASLONG strideA, float * __restrict B,\ + BLASLONG strideB , float * __restrict R, BLASLONG strideR) +*/ +void CNAME (BLASLONG M, BLASLONG N, BLASLONG K, float * __restrict A,\ + BLASLONG strideA, float * __restrict B, BLASLONG strideB ,\ + float * __restrict R, BLASLONG strideR){ + + uint64_t m_mod, vl_elms; + + vl_elms = sve_cntw(); + m_mod = (((uint64_t)M + vl_elms - 1) / vl_elms) * vl_elms; + + float *A_mod = (float *) malloc(m_mod*K*sizeof(float)); + + /* Prevent compiler optimization by reading from memory instead + * of reading directly from vector (z) registers. + * */ + asm volatile("" : : :"p0", "p1", "p2", "p3", "p4", "p5", "p6", "p7", + "p8", "p9", "p10", "p11", "p12", "p13", "p14", "p15", "d8", "d9", "d10", "d11", "d12", "d13", "d14", "d15", + "z0", "z1", "z2", "z3", "z4", "z5", "z6", "z7", + "z8", "z9", "z10", "z11", "z12", "z13", "z14", "z15", + "z16", "z17", "z18", "z19", "z20", "z21", "z22", "z23", + "z24", "z25", "z26", "z27", "z28", "z29", "z30", "z31"); + + /* Pre-process the left matrix to make it suitable for + matrix sum of outer-product calculation + */ + SME1_PREPROCESS(M, K, A, A_mod); + + /* Calculate C = A*B */ + SME1_DIRECT2X2(M, K, N, A_mod, B, R); + + asm volatile("" : : :"p0", "p1", "p2", "p3", "p4", "p5", "p6", "p7", + "p8", "p9", "p10", "p11", "p12", "p13", "p14", "p15", "d8", "d9", "d10", "d11", "d12", "d13", "d14", "d15", + "z0", "z1", "z2", "z3", "z4", "z5", "z6", "z7", + "z8", "z9", "z10", "z11", "z12", "z13", "z14", "z15", + "z16", "z17", "z18", "z19", "z20", "z21", "z22", "z23", + "z24", "z25", "z26", "z27", "z28", "z29", "z30", "z31"); + free(A_mod); +} + +#else + +void CNAME (BLASLONG M, BLASLONG N, BLASLONG K, float * __restrict A,\ + BLASLONG strideA, float * __restrict B, BLASLONG strideB ,\ + float * __restrict R, BLASLONG strideR){ +fprintf(stderr,"EMPTY sgemm_kernel_direct should never be called \n"); +} +void SME1_DIRECT2X2( uint64_t M , uint64_t K, uint64_t N,\ + const float * restrict A_base,\ + const float * restrict B_base,\ + const float * restrict C_base){}; +void SME1_PREPROCESS(uint64_t nbr, uint64_t nbc,\ + const float * restrict a, float * a_mod){}; + + +#endif + diff --git a/kernel/arm64/ssymm_direct_alpha_beta_arm64_sme1.c b/kernel/arm64/ssymm_direct_alpha_beta_arm64_sme1.c index fd38c225ad..51c753abd5 100644 --- a/kernel/arm64/ssymm_direct_alpha_beta_arm64_sme1.c +++ b/kernel/arm64/ssymm_direct_alpha_beta_arm64_sme1.c @@ -1,255 +1,255 @@ -/* - Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. - SPDX-License-Identifier: BSD-3-Clause-Clear -*/ - -#include "common.h" -#include -#include -// #include "sme_abi.h" -#if defined(HAVE_SME) - -#if defined(__ARM_FEATURE_SME) && defined(__clang__) && __clang_major__ >= 16 -#include -#endif - -#if defined(DYNAMIC_ARCH) -#define COMBINE(a,b) a ## b -#define COMBINE2(a,b) COMBINE(a,b) -#define SGEMM_PREPROCESS_BASE sgemm_direct_sme1_preprocess -#define SGEMM_PREPROCESS COMBINE2(SGEMM_PREPROCESS_BASE,TS) -#define SGEMM_DIRECT2X2_BASE sgemm_direct_alpha_beta_sme1_2VLx2VL -#define SGEMM_DIRECT2X2 COMBINE2(SGEMM_DIRECT2X2_BASE,TS) -#else -#define SGEMM_PREPROCESS sgemm_direct_sme1_preprocess -#define SGEMM_DIRECT2X2 sgemm_direct_alpha_beta_sme1_2VLx2VL -#endif - -/* Function prototypes */ -extern void SGEMM_PREPROCESS(uint64_t nbr, uint64_t nbc,\ - const float * restrict a, float * a_mod); - -extern void SGEMM_DIRECT2X2(uint64_t m, uint64_t k, uint64_t n, const float* alpha,\ - const float *ba, const float *restrict bb, const float* beta,\ - float *restrict C); -/* Function Definitions */ -static uint64_t sve_cntw() { - uint64_t cnt; - asm volatile( - "rdsvl %[res], #1\n" - "lsr %[res], %[res], #2\n" - : [res] "=r" (cnt) :: - ); - return cnt; -} - -#if defined(__ARM_FEATURE_SME) && defined(__ARM_FEATURE_LOCALLY_STREAMING) && defined(__clang__) && __clang_major__ >= 16 - -#if defined(UPPER) -__arm_new("za") __arm_locally_streaming -static void ssymm_direct_sme1_preprocessLU(uint64_t nbr, uint64_t nbc, - const float *restrict a, float *restrict a_mod) -{ - // const uint64_t num_rows = nbr; - // const uint64_t num_cols = nbc; - const uint64_t svl = svcntw(); - uint64_t row_batch = svl; - - const float *restrict pSrc; - float *restrict pDst; - for (uint64_t row_idx = 0; row_idx < nbr; row_idx += row_batch) - { - row_batch = MIN(row_batch, nbr - row_idx); - - // Fill in the lower triangle and Transpose 1SVL x N panel of A - uint64_t col_batch = svl; - - for (uint64_t col_idx = 0; col_idx < nbc; col_idx += col_batch) - { - svzero_za(); - - if (col_idx == row_idx) - { - pSrc = &a[(row_idx)*nbc + col_idx]; - pDst = &a_mod[(col_idx)*svl + row_idx * nbc]; - // Load horizontal slices, filling lower elements - const svbool_t pg_row = svwhilelt_b32_u64(col_idx, nbc); - for (int64_t row = row_batch - 1; row >= 0; row--) - { - svld1_hor_za32(0, row, pg_row, &pSrc[row * nbc]); - svld1_ver_za32(0, row, pg_row, &pSrc[row * nbc]); - } - // Save vertical slices - col_batch = MIN(col_batch, nbc - col_idx); - for (uint64_t col = 0; col < col_batch; col++) - { - svst1_ver_za32(0, col, svptrue_b32(), &pDst[col * svl]); - } - } - else if (col_idx > row_idx) - { - pSrc = &a[(row_idx)*nbc + col_idx]; - pDst = &a_mod[(col_idx)*svl + row_idx * nbc]; - // Load horizontal slices - const svbool_t pg_row = svwhilelt_b32_u64(col_idx, nbc); - for (uint64_t row = 0; row < row_batch; row++) - { - svld1_hor_za32(0, row, pg_row, &pSrc[row * nbc]); - } - // Save vertical slices - col_batch = MIN(col_batch, nbc - col_idx); - for (uint64_t col = 0; col < col_batch; col++) - { - svst1_ver_za32(0, col, svptrue_b32(), &pDst[col * svl]); - } - } - else if (col_idx < row_idx) - { - pSrc = &a[row_idx + col_idx * nbc]; - pDst = &a_mod[(col_idx)*svl + row_idx * nbc]; - // Load horizontal slices - const svbool_t pg_row = svwhilelt_b32_u64(row_idx, nbc); - for (uint64_t row = 0; row < svl; row++) - { - svld1_hor_za32(0, row, pg_row, &pSrc[row * nbc]); - } - // Save vertical slices - col_batch = MIN(col_batch, nbc - col_idx); - for (uint64_t col = 0; col < svl; col++) - { - svst1_hor_za32(0, col, svptrue_b32(), &pDst[col * svl]); - } - } - } - } -} -#endif - -// -#if defined(LOWER) -__arm_new("za") __arm_locally_streaming -static void ssymm_direct_sme1_preprocessLL(uint64_t nbr, uint64_t nbc, - const float *restrict a, float *restrict a_mod) -{ - // const uint64_t num_rows = nbr; - const uint64_t svl = svcntw(); - uint64_t row_batch = svl; - - const float *restrict pSrc; - float *restrict pDst; - for (uint64_t row_idx = 0; row_idx < nbr; row_idx += row_batch) - { - row_batch = MIN(row_batch, nbr - row_idx); - - // Fill in the upper triangle and Transpose 1SVL x N panel of A - uint64_t col_batch = svl; - - for (uint64_t col_idx = 0; col_idx < nbc; col_idx += col_batch) - { - svzero_za(); - - if (col_idx == row_idx) - { - pSrc = &a[(row_idx)*nbc + col_idx]; - pDst = &a_mod[(col_idx)*svl + row_idx * nbc]; - // Load horizontal slices, filling upper elements - const svbool_t pg_row = svwhilelt_b32_u64(col_idx, nbc); - for (uint64_t row = 0; row < row_batch; row++) - { - svld1_hor_za32(0, row, pg_row, &pSrc[row * nbc]); - svld1_ver_za32(0, row, pg_row, &pSrc[row * nbc]); - } - // Save vertical slices - col_batch = MIN(col_batch, nbc - col_idx); - for (uint64_t col = 0; col < col_batch; col++) - { - svst1_ver_za32(0, col, svptrue_b32(), &pDst[col * svl]); - } - } - else if (col_idx > row_idx) - { - pSrc = &a[row_idx + col_idx * nbc]; - pDst = &a_mod[(col_idx)*svl + row_idx * nbc]; - // Load horizontal slices - const svbool_t pg_row = svptrue_b32(); - for (uint64_t row = 0; row < row_batch; row++) - { - svld1_hor_za32(0, row, pg_row, &pSrc[row * nbc]); - } - // Save vertical slices - col_batch = MIN(col_batch, nbc - col_idx); - for (uint64_t col = 0; col < col_batch; col++) - { - svst1_hor_za32(0, col, svptrue_b32(), &pDst[col * svl]); - } - } - else if (col_idx < row_idx) - { - pSrc = &a[(row_idx)*nbc + col_idx]; - pDst = &a_mod[(col_idx)*svl + row_idx * nbc]; - // Load horizontal slices - const svbool_t pg_row = svwhilelt_b32_u64(col_idx, nbc); - for (uint64_t row = 0; row < row_batch; row++) - { - svld1_hor_za32(0, row, pg_row, &pSrc[row * nbc]); - } - // Save vertical slices - col_batch = MIN(col_batch, nbc - col_idx); - for (uint64_t col = 0; col < col_batch; col++) - { - svst1_ver_za32(0, col, svptrue_b32(), &pDst[col * svl]); - } - } - } - } -} -#endif -#else -#if defined(UPPER) -static void ssymm_direct_sme1_preprocessLU(uint64_t nbr, uint64_t nbc, - const float *restrict a, float *restrict a_mod){} -#endif -#if defined(LOWER) -static void ssymm_direct_sme1_preprocessLL(uint64_t nbr, uint64_t nbc, - const float *restrict a, float *restrict a_mod){} -#endif -#endif - -// -void CNAME(BLASLONG M, BLASLONG N, float alpha, float *__restrict A, - BLASLONG strideA, float *__restrict B, BLASLONG strideB, - float beta, float *__restrict R, BLASLONG strideR) -{ - if (alpha == 0.0f) { - if (beta == 1.0f) - return; - SGEMM_DIRECT2X2(M, 0, N, &alpha, A, B, &beta, R); - return; - } - - uint64_t vl_elms = sve_cntw(); // vl_elem = 16 - uint64_t m_mod = (((uint64_t)M + vl_elms - 1) / vl_elms) * vl_elms; - - /* Pre-process the left matrix to make it suitable for - matrix sum of outer-product calculation - */ - float *A_mod = (float *)malloc(m_mod * M * sizeof(float)); - -#if defined(UPPER) - ssymm_direct_sme1_preprocessLU(M, M, A, A_mod); -#elif defined(LOWER) - ssymm_direct_sme1_preprocessLL(M, M, A, A_mod); -#endif - - /* Calculate C = alpha*A*B + beta*C */ - SGEMM_DIRECT2X2(M, M, N, &alpha, A_mod, B, &beta, R); - free(A_mod); -} - -#else - -void CNAME (BLASLONG M, BLASLONG N, float alpha, float * __restrict A,\ - BLASLONG strideA, float * __restrict B, BLASLONG strideB ,\ - float beta, float * __restrict R, BLASLONG strideR){} - -#endif +/* + Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. + SPDX-License-Identifier: BSD-3-Clause-Clear +*/ + +#include "common.h" +#include +#include +// #include "sme_abi.h" +#if defined(HAVE_SME) + +#if defined(__ARM_FEATURE_SME) && defined(__clang__) && __clang_major__ >= 16 +#include +#endif + +#if defined(DYNAMIC_ARCH) +#define COMBINE(a,b) a ## b +#define COMBINE2(a,b) COMBINE(a,b) +#define SGEMM_PREPROCESS_BASE sgemm_direct_sme1_preprocess +#define SGEMM_PREPROCESS COMBINE2(SGEMM_PREPROCESS_BASE,TS) +#define SGEMM_DIRECT2X2_BASE sgemm_direct_alpha_beta_sme1_2VLx2VL +#define SGEMM_DIRECT2X2 COMBINE2(SGEMM_DIRECT2X2_BASE,TS) +#else +#define SGEMM_PREPROCESS sgemm_direct_sme1_preprocess +#define SGEMM_DIRECT2X2 sgemm_direct_alpha_beta_sme1_2VLx2VL +#endif + +/* Function prototypes */ +extern void SGEMM_PREPROCESS(uint64_t nbr, uint64_t nbc,\ + const float * restrict a, float * a_mod); + +extern void SGEMM_DIRECT2X2(uint64_t m, uint64_t k, uint64_t n, const float* alpha,\ + const float *ba, const float *restrict bb, const float* beta,\ + float *restrict C); +/* Function Definitions */ +static uint64_t sve_cntw() { + uint64_t cnt; + asm volatile( + "rdsvl %[res], #1\n" + "lsr %[res], %[res], #2\n" + : [res] "=r" (cnt) :: + ); + return cnt; +} + +#if defined(__ARM_FEATURE_SME) && defined(__ARM_FEATURE_LOCALLY_STREAMING) && defined(__clang__) && __clang_major__ >= 16 + +#if defined(UPPER) +__arm_new("za") __arm_locally_streaming +static void ssymm_direct_sme1_preprocessLU(uint64_t nbr, uint64_t nbc, + const float *restrict a, float *restrict a_mod) +{ + // const uint64_t num_rows = nbr; + // const uint64_t num_cols = nbc; + const uint64_t svl = svcntw(); + uint64_t row_batch = svl; + + const float *restrict pSrc; + float *restrict pDst; + for (uint64_t row_idx = 0; row_idx < nbr; row_idx += row_batch) + { + row_batch = MIN(row_batch, nbr - row_idx); + + // Fill in the lower triangle and Transpose 1SVL x N panel of A + uint64_t col_batch = svl; + + for (uint64_t col_idx = 0; col_idx < nbc; col_idx += col_batch) + { + svzero_za(); + + if (col_idx == row_idx) + { + pSrc = &a[(row_idx)*nbc + col_idx]; + pDst = &a_mod[(col_idx)*svl + row_idx * nbc]; + // Load horizontal slices, filling lower elements + const svbool_t pg_row = svwhilelt_b32_u64(col_idx, nbc); + for (int64_t row = row_batch - 1; row >= 0; row--) + { + svld1_hor_za32(0, row, pg_row, &pSrc[row * nbc]); + svld1_ver_za32(0, row, pg_row, &pSrc[row * nbc]); + } + // Save vertical slices + col_batch = MIN(col_batch, nbc - col_idx); + for (uint64_t col = 0; col < col_batch; col++) + { + svst1_ver_za32(0, col, svptrue_b32(), &pDst[col * svl]); + } + } + else if (col_idx > row_idx) + { + pSrc = &a[(row_idx)*nbc + col_idx]; + pDst = &a_mod[(col_idx)*svl + row_idx * nbc]; + // Load horizontal slices + const svbool_t pg_row = svwhilelt_b32_u64(col_idx, nbc); + for (uint64_t row = 0; row < row_batch; row++) + { + svld1_hor_za32(0, row, pg_row, &pSrc[row * nbc]); + } + // Save vertical slices + col_batch = MIN(col_batch, nbc - col_idx); + for (uint64_t col = 0; col < col_batch; col++) + { + svst1_ver_za32(0, col, svptrue_b32(), &pDst[col * svl]); + } + } + else if (col_idx < row_idx) + { + pSrc = &a[row_idx + col_idx * nbc]; + pDst = &a_mod[(col_idx)*svl + row_idx * nbc]; + // Load horizontal slices + const svbool_t pg_row = svwhilelt_b32_u64(row_idx, nbc); + for (uint64_t row = 0; row < svl; row++) + { + svld1_hor_za32(0, row, pg_row, &pSrc[row * nbc]); + } + // Save vertical slices + col_batch = MIN(col_batch, nbc - col_idx); + for (uint64_t col = 0; col < svl; col++) + { + svst1_hor_za32(0, col, svptrue_b32(), &pDst[col * svl]); + } + } + } + } +} +#endif + +// +#if defined(LOWER) +__arm_new("za") __arm_locally_streaming +static void ssymm_direct_sme1_preprocessLL(uint64_t nbr, uint64_t nbc, + const float *restrict a, float *restrict a_mod) +{ + // const uint64_t num_rows = nbr; + const uint64_t svl = svcntw(); + uint64_t row_batch = svl; + + const float *restrict pSrc; + float *restrict pDst; + for (uint64_t row_idx = 0; row_idx < nbr; row_idx += row_batch) + { + row_batch = MIN(row_batch, nbr - row_idx); + + // Fill in the upper triangle and Transpose 1SVL x N panel of A + uint64_t col_batch = svl; + + for (uint64_t col_idx = 0; col_idx < nbc; col_idx += col_batch) + { + svzero_za(); + + if (col_idx == row_idx) + { + pSrc = &a[(row_idx)*nbc + col_idx]; + pDst = &a_mod[(col_idx)*svl + row_idx * nbc]; + // Load horizontal slices, filling upper elements + const svbool_t pg_row = svwhilelt_b32_u64(col_idx, nbc); + for (uint64_t row = 0; row < row_batch; row++) + { + svld1_hor_za32(0, row, pg_row, &pSrc[row * nbc]); + svld1_ver_za32(0, row, pg_row, &pSrc[row * nbc]); + } + // Save vertical slices + col_batch = MIN(col_batch, nbc - col_idx); + for (uint64_t col = 0; col < col_batch; col++) + { + svst1_ver_za32(0, col, svptrue_b32(), &pDst[col * svl]); + } + } + else if (col_idx > row_idx) + { + pSrc = &a[row_idx + col_idx * nbc]; + pDst = &a_mod[(col_idx)*svl + row_idx * nbc]; + // Load horizontal slices + const svbool_t pg_row = svptrue_b32(); + for (uint64_t row = 0; row < row_batch; row++) + { + svld1_hor_za32(0, row, pg_row, &pSrc[row * nbc]); + } + // Save vertical slices + col_batch = MIN(col_batch, nbc - col_idx); + for (uint64_t col = 0; col < col_batch; col++) + { + svst1_hor_za32(0, col, svptrue_b32(), &pDst[col * svl]); + } + } + else if (col_idx < row_idx) + { + pSrc = &a[(row_idx)*nbc + col_idx]; + pDst = &a_mod[(col_idx)*svl + row_idx * nbc]; + // Load horizontal slices + const svbool_t pg_row = svwhilelt_b32_u64(col_idx, nbc); + for (uint64_t row = 0; row < row_batch; row++) + { + svld1_hor_za32(0, row, pg_row, &pSrc[row * nbc]); + } + // Save vertical slices + col_batch = MIN(col_batch, nbc - col_idx); + for (uint64_t col = 0; col < col_batch; col++) + { + svst1_ver_za32(0, col, svptrue_b32(), &pDst[col * svl]); + } + } + } + } +} +#endif +#else +#if defined(UPPER) +static void ssymm_direct_sme1_preprocessLU(uint64_t nbr, uint64_t nbc, + const float *restrict a, float *restrict a_mod){} +#endif +#if defined(LOWER) +static void ssymm_direct_sme1_preprocessLL(uint64_t nbr, uint64_t nbc, + const float *restrict a, float *restrict a_mod){} +#endif +#endif + +// +void CNAME(BLASLONG M, BLASLONG N, float alpha, float *__restrict A, + BLASLONG strideA, float *__restrict B, BLASLONG strideB, + float beta, float *__restrict R, BLASLONG strideR) +{ + if (alpha == 0.0f) { + if (beta == 1.0f) + return; + SGEMM_DIRECT2X2(M, 0, N, &alpha, A, B, &beta, R); + return; + } + + uint64_t vl_elms = sve_cntw(); // vl_elem = 16 + uint64_t m_mod = (((uint64_t)M + vl_elms - 1) / vl_elms) * vl_elms; + + /* Pre-process the left matrix to make it suitable for + matrix sum of outer-product calculation + */ + float *A_mod = (float *)malloc(m_mod * M * sizeof(float)); + +#if defined(UPPER) + ssymm_direct_sme1_preprocessLU(M, M, A, A_mod); +#elif defined(LOWER) + ssymm_direct_sme1_preprocessLL(M, M, A, A_mod); +#endif + + /* Calculate C = alpha*A*B + beta*C */ + SGEMM_DIRECT2X2(M, M, N, &alpha, A_mod, B, &beta, R); + free(A_mod); +} + +#else + +void CNAME (BLASLONG M, BLASLONG N, float alpha, float * __restrict A,\ + BLASLONG strideA, float * __restrict B, BLASLONG strideB ,\ + float beta, float * __restrict R, BLASLONG strideR){} + +#endif diff --git a/kernel/arm64/ssyrk_direct_alpha_beta_arm64_sme1.c b/kernel/arm64/ssyrk_direct_alpha_beta_arm64_sme1.c index 7142e31cd2..44c180f9ef 100644 --- a/kernel/arm64/ssyrk_direct_alpha_beta_arm64_sme1.c +++ b/kernel/arm64/ssyrk_direct_alpha_beta_arm64_sme1.c @@ -1,278 +1,278 @@ -/* - Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. - SPDX-License-Identifier: BSD-3-Clause-Clear -*/ - -#include "common.h" -#include -#include -#if defined(HAVE_SME) - -#if defined(DYNAMIC_ARCH) -#define COMBINE(a,b) a ## b -#define COMBINE2(a,b) COMBINE(a,b) -#define SGEMM_PREPROCESS_BASE sgemm_direct_sme1_preprocess -#define SGEMM_PREPROCESS COMBINE2(SGEMM_PREPROCESS_BASE,TS) -#define SGEMM_DIRECT2X2_BASE sgemm_direct_alpha_beta_sme1_2VLx2VL -#define SGEMM_DIRECT2X2 COMBINE2(SGEMM_DIRECT2X2_BASE,TS) -#else -#define SGEMM_PREPROCESS sgemm_direct_sme1_preprocess -#define SGEMM_DIRECT2X2 sgemm_direct_alpha_beta_sme1_2VLx2VL -#endif - -#if defined(__ARM_FEATURE_SME) && defined(__clang__) && __clang_major__ >= 16 -#include -#endif - -/* Function prototypes */ -extern void SGEMM_PREPROCESS (uint64_t nbr, uint64_t nbc,\ - - const float * restrict a, float * a_mod) ; - -/* Function Definitions */ -#if !defined(TRANSA) -static uint64_t sve_cntw() { - uint64_t cnt; - asm volatile( - "rdsvl %[res], #1\n" - "lsr %[res], %[res], #2\n" - : [res] "=r" (cnt) :: - ); - return cnt; -} -#endif - -#if defined(__ARM_FEATURE_SME) && defined(__ARM_FEATURE_LOCALLY_STREAMING) && defined(__clang__) && __clang_major__ >= 16 -// Outer product kernel. -// Computes a 2SVL x 2SVL block of C, utilizing all four FP32 tiles of ZA. -__attribute__((always_inline)) inline void -kernel_2x2(const float *A, const float *B, float *C, size_t shared_dim, - size_t ldc, size_t block_rows, size_t block_cols, float alpha, - float beta, uint64_t row_idx, uint64_t col_idx) - __arm_out("za") __arm_streaming { - - const uint64_t svl = svcntw(); -#if defined(TRANSA) - size_t ldb = ldc; -#endif - // Predicate set-up - svbool_t pg = svptrue_b32(); - svbool_t pg_a_0 = svwhilelt_b32_u64(0, block_rows); - svbool_t pg_a_1 = svwhilelt_b32_u64(svl, block_rows); - - svbool_t pg_b_0 = svwhilelt_b32_u64(0, block_cols); - svbool_t pg_b_1 = svwhilelt_b32_u64(svl, block_cols); - -#define pg_c_0 pg_b_0 -#define pg_c_1 pg_b_1 - - svzero_za(); - // beta == 0 must not read C; ZA is already initialized to zero. - if (beta != 0.0f) { - svfloat32_t beta_vec = svdup_f32(beta); - - // Load C to ZA - for (size_t i = 0; i < MIN(svl, block_rows); i++) { - svfloat32_t row_c_0 = svld1(pg_c_0, &C[i * ldc]); - row_c_0 = svmul_x(pg, beta_vec, row_c_0); - svwrite_hor_za32_f32_m(/*tile*/0, /*slice*/i, pg_c_0, row_c_0); - - svfloat32_t row_c_1 = svld1(pg_c_1, &C[i * ldc + svl]); - row_c_1 = svmul_x(pg, beta_vec, row_c_1); - svwrite_hor_za32_f32_m(/*tile*/1, /*slice*/i, pg_c_1, row_c_1); - } - for (size_t i = svl; i < block_rows; i++) { - svfloat32_t row_c_0 = svld1(pg_c_0, &C[i * ldc]); - row_c_0 = svmul_x(pg, beta_vec, row_c_0); - svwrite_hor_za32_f32_m(/*tile*/2, /*slice*/i - svl, pg_c_0, row_c_0); - - svfloat32_t row_c_1 = svld1(pg_c_1, &C[i * ldc + svl]); - row_c_1 = svmul_x(pg, beta_vec, row_c_1); - svwrite_hor_za32_f32_m(/*tile*/3, /*slice*/i - svl, pg_c_1, row_c_1); - } - } - - svfloat32_t alpha_vec = svdup_f32(alpha); - // Iterate through shared dimension (K) - for (size_t k = 0; k < shared_dim; k++) { -#if !defined(TRANSA) - // Load column of A - svfloat32_t col_a_0 = svld1(pg_a_0, &A[k * svl]); - col_a_0 = svmul_x(pg, alpha_vec, col_a_0); - svfloat32_t col_a_1 = svld1(pg_a_1, &A[(k + shared_dim) * svl]); - col_a_1 = svmul_x(pg, alpha_vec, col_a_1); - - // Load row of A**T - svfloat32_t row_b_0 = svld1(pg_b_0, &B[k * svl]); - svfloat32_t row_b_1 = svld1(pg_b_1, &B[(k + shared_dim) * svl]); -#else - // Load column of A**T - svfloat32_t col_a_0 = svld1(pg_a_0, &A[k * ldb]); - col_a_0 = svmul_x(pg, alpha_vec, col_a_0); - - svfloat32_t col_a_1 = svld1(pg_a_1, &A[k * ldb + svl]); - col_a_1 = svmul_x(pg, alpha_vec, col_a_1); - - // Load row of A - svfloat32_t row_b_0 = svld1(pg_b_0, &B[k * ldb]); - svfloat32_t row_b_1 = svld1(pg_b_1, &B[k * ldb + svl]); -#endif - // Perform outer product - svmopa_za32_m(/*tile*/0, pg, pg, col_a_0, row_b_0); - svmopa_za32_m(/*tile*/1, pg, pg, col_a_0, row_b_1); - svmopa_za32_m(/*tile*/2, pg, pg, col_a_1, row_b_0); - svmopa_za32_m(/*tile*/3, pg, pg, col_a_1, row_b_1); - } - -#if defined(UPPER) -#define pg_c_0_full pg_c_0 -#define pg_c_1_full pg_c_1 - - bool need_update_pg_b = true; - size_t last_invalid_index = col_idx - row_idx; - // For Upper, If col_idx - row_idx >= 2*svl, we don't need to update the predicate due to all elements above the digonal - if (col_idx - row_idx >= 2*svl) { - need_update_pg_b = false; - } - // Store to C from ZA - for (size_t i = 0; i < MIN(svl, block_rows); i++, last_invalid_index++) { - if (need_update_pg_b) { - pg_c_0 = svnot_b_z(pg_c_0_full, svwhilelt_b32_u64(0, last_invalid_index)); - pg_c_1 = svnot_b_z(pg_c_1_full, svwhilelt_b32_u64(svl, last_invalid_index)); - } - - svst1_hor_za32(/*tile*/0, /*slice*/i, pg_c_0, &C[i * ldc]); - svst1_hor_za32(/*tile*/1, /*slice*/i, pg_c_1, &C[i * ldc + svl]); - } - for (size_t i = svl; i < block_rows; i++,last_invalid_index++) { - if (need_update_pg_b) { - pg_c_0 = svnot_b_z(pg_c_0_full, svwhilelt_b32_u64(0, last_invalid_index)); - pg_c_1 = svnot_b_z(pg_c_1_full, svwhilelt_b32_u64(svl, last_invalid_index)); - } - svst1_hor_za32(/*tile*/2, /*slice*/i - svl, pg_c_0, &C[i * ldc]); - svst1_hor_za32(/*tile*/3, /*slice*/i - svl, pg_c_1, &C[i * ldc + svl]); - } -#else - // Store to C from ZA - size_t valid_index = row_idx - col_idx + 1; - for (size_t i = 0; i < MIN(svl, block_rows); i++, valid_index++) { - pg_c_0 = svwhilelt_b32_u64(0, MIN(valid_index, block_cols)); - pg_c_1 = svwhilelt_b32_u64(svl, MIN(valid_index, block_cols)); - svst1_hor_za32(/*tile*/0, /*slice*/i, pg_c_0, &C[i * ldc]); - svst1_hor_za32(/*tile*/1, /*slice*/i, pg_c_1, &C[i * ldc + svl]); - } - for (size_t i = svl; i < block_rows; i++, valid_index++) { - pg_c_0 = svwhilelt_b32_u64(0, MIN(valid_index, block_cols)); - pg_c_1 = svwhilelt_b32_u64(svl, MIN(valid_index, block_cols)); - svst1_hor_za32(/*tile*/2, /*slice*/i - svl, pg_c_0, &C[i * ldc]); - svst1_hor_za32(/*tile*/3, /*slice*/i - svl, pg_c_1, &C[i * ldc + svl]); - } -#endif -} - -__arm_new("za") __arm_locally_streaming -static void ssyrk_direct_sme1_2VLx2VL(uint64_t n, uint64_t k, const float* alpha,\ - const float *ba, const float* beta, float *restrict bc) { - const uint64_t num_rows = n; - const uint64_t num_cols = n; - - const float *restrict a_ptr = ba; - const float *restrict b_ptr = ba; - float *restrict c_ptr = bc; - - const uint64_t svl = svcntw(); - const uint64_t ldc = n; - - // Block over rows of C (panels of A) - uint64_t row_idx = 0; - - // 2x2 loop - uint64_t row_batch = 2*svl; - - // Block over row dimension of C - for (; row_idx < num_rows; row_idx += row_batch) { - row_batch = MIN(row_batch, num_rows - row_idx); - uint64_t col_batch = 2*svl; -#if defined(UPPER) - // for UPLO is upper, Start from column col_idx = rows_index to ensure we only process the upper triangle (col_idx >= rows_index) - for (uint64_t col_idx = row_idx; col_idx < num_cols; col_idx += col_batch) { - col_batch = MIN(col_batch, num_cols - col_idx); -#else - // for UPLO is lower, we only process the lower triangle part (col_idx <= row_idxx) - for (uint64_t col_idx = 0; col_idx < num_cols && col_idx <= row_idx; col_idx += col_batch) { -#endif - col_batch = MIN(col_batch, num_cols - col_idx); -#if !defined(TRANSA) - kernel_2x2(&a_ptr[row_idx * k], &b_ptr[col_idx * k], - &c_ptr[row_idx * ldc + col_idx], k, - ldc, row_batch, col_batch, *alpha, *beta, row_idx, col_idx); -#else - kernel_2x2(&a_ptr[row_idx], &b_ptr[col_idx], - &c_ptr[row_idx * ldc + col_idx], k, - ldc, row_batch, col_batch, *alpha, *beta, row_idx, col_idx); -#endif - - } - } - return; -} - -#else -static void ssyrk_direct_sme1_2VLx2VL(uint64_t n, uint64_t k, const float* alpha,\ - const float *ba, const float* beta, float *restrict bc){} -#endif - -void CNAME (BLASLONG N, BLASLONG K, float alpha, float * __restrict A,\ - BLASLONG strideA, float beta, float * __restrict C, BLASLONG strideC){ - if (alpha == 0.0f || K == 0) { - if (beta == 1.0f) - return; - ssyrk_direct_sme1_2VLx2VL(N, 0, &alpha, A, &beta, C); - return; - } - -#if !defined(TRANSA) - uint64_t n_mod, vl_elms; - - vl_elms = sve_cntw(); - - n_mod = (((uint64_t)N + vl_elms - 1) / vl_elms) * vl_elms; - - float *A_mod = (float *) malloc(n_mod*K*sizeof(float)); - - /* Prevent compiler optimization by reading from memory instead - * of reading directly from vector (z) registers. - * */ - asm volatile("" : : :"p0", "p1", "p2", "p3", "p4", "p5", "p6", "p7", - "p8", "p9", "p10", "p11", "p12", "p13", "p14", "p15", - "z0", "z1", "z2", "z3", "z4", "z5", "z6", "z7", - "z8", "z9", "z10", "z11", "z12", "z13", "z14", "z15", - "z16", "z17", "z18", "z19", "z20", "z21", "z22", "z23", - "z24", "z25", "z26", "z27", "z28", "z29", "z30", "z31"); - - /* Pre-process the left matrix to make it suitable for - matrix sum of outer-product calculation - */ - SGEMM_PREPROCESS (N, K, A, A_mod); - asm volatile("" : : :"p0", "p1", "p2", "p3", "p4", "p5", "p6", "p7", - "p8", "p9", "p10", "p11", "p12", "p13", "p14", "p15", - "z0", "z1", "z2", "z3", "z4", "z5", "z6", "z7", - "z8", "z9", "z10", "z11", "z12", "z13", "z14", "z15", - "z16", "z17", "z18", "z19", "z20", "z21", "z22", "z23", - "z24", "z25", "z26", "z27", "z28", "z29", "z30", "z31"); - ssyrk_direct_sme1_2VLx2VL(N, K, &alpha, A_mod, &beta, C); - free(A_mod); -#else - ssyrk_direct_sme1_2VLx2VL(N, K, &alpha, A, &beta, C); -#endif - -} - -#else - -void CNAME (BLASLONG N, BLASLONG K, float alpha, float * __restrict A,\ - BLASLONG strideA, float beta, float * __restrict C, BLASLONG strideC){ -fprintf(stderr,"empty ssyrk_direct kernel should never be called\n"); -} - -#endif +/* + Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. + SPDX-License-Identifier: BSD-3-Clause-Clear +*/ + +#include "common.h" +#include +#include +#if defined(HAVE_SME) + +#if defined(DYNAMIC_ARCH) +#define COMBINE(a,b) a ## b +#define COMBINE2(a,b) COMBINE(a,b) +#define SGEMM_PREPROCESS_BASE sgemm_direct_sme1_preprocess +#define SGEMM_PREPROCESS COMBINE2(SGEMM_PREPROCESS_BASE,TS) +#define SGEMM_DIRECT2X2_BASE sgemm_direct_alpha_beta_sme1_2VLx2VL +#define SGEMM_DIRECT2X2 COMBINE2(SGEMM_DIRECT2X2_BASE,TS) +#else +#define SGEMM_PREPROCESS sgemm_direct_sme1_preprocess +#define SGEMM_DIRECT2X2 sgemm_direct_alpha_beta_sme1_2VLx2VL +#endif + +#if defined(__ARM_FEATURE_SME) && defined(__clang__) && __clang_major__ >= 16 +#include +#endif + +/* Function prototypes */ +extern void SGEMM_PREPROCESS (uint64_t nbr, uint64_t nbc,\ + + const float * restrict a, float * a_mod) ; + +/* Function Definitions */ +#if !defined(TRANSA) +static uint64_t sve_cntw() { + uint64_t cnt; + asm volatile( + "rdsvl %[res], #1\n" + "lsr %[res], %[res], #2\n" + : [res] "=r" (cnt) :: + ); + return cnt; +} +#endif + +#if defined(__ARM_FEATURE_SME) && defined(__ARM_FEATURE_LOCALLY_STREAMING) && defined(__clang__) && __clang_major__ >= 16 +// Outer product kernel. +// Computes a 2SVL x 2SVL block of C, utilizing all four FP32 tiles of ZA. +__attribute__((always_inline)) inline void +kernel_2x2(const float *A, const float *B, float *C, size_t shared_dim, + size_t ldc, size_t block_rows, size_t block_cols, float alpha, + float beta, uint64_t row_idx, uint64_t col_idx) + __arm_out("za") __arm_streaming { + + const uint64_t svl = svcntw(); +#if defined(TRANSA) + size_t ldb = ldc; +#endif + // Predicate set-up + svbool_t pg = svptrue_b32(); + svbool_t pg_a_0 = svwhilelt_b32_u64(0, block_rows); + svbool_t pg_a_1 = svwhilelt_b32_u64(svl, block_rows); + + svbool_t pg_b_0 = svwhilelt_b32_u64(0, block_cols); + svbool_t pg_b_1 = svwhilelt_b32_u64(svl, block_cols); + +#define pg_c_0 pg_b_0 +#define pg_c_1 pg_b_1 + + svzero_za(); + // beta == 0 must not read C; ZA is already initialized to zero. + if (beta != 0.0f) { + svfloat32_t beta_vec = svdup_f32(beta); + + // Load C to ZA + for (size_t i = 0; i < MIN(svl, block_rows); i++) { + svfloat32_t row_c_0 = svld1(pg_c_0, &C[i * ldc]); + row_c_0 = svmul_x(pg, beta_vec, row_c_0); + svwrite_hor_za32_f32_m(/*tile*/0, /*slice*/i, pg_c_0, row_c_0); + + svfloat32_t row_c_1 = svld1(pg_c_1, &C[i * ldc + svl]); + row_c_1 = svmul_x(pg, beta_vec, row_c_1); + svwrite_hor_za32_f32_m(/*tile*/1, /*slice*/i, pg_c_1, row_c_1); + } + for (size_t i = svl; i < block_rows; i++) { + svfloat32_t row_c_0 = svld1(pg_c_0, &C[i * ldc]); + row_c_0 = svmul_x(pg, beta_vec, row_c_0); + svwrite_hor_za32_f32_m(/*tile*/2, /*slice*/i - svl, pg_c_0, row_c_0); + + svfloat32_t row_c_1 = svld1(pg_c_1, &C[i * ldc + svl]); + row_c_1 = svmul_x(pg, beta_vec, row_c_1); + svwrite_hor_za32_f32_m(/*tile*/3, /*slice*/i - svl, pg_c_1, row_c_1); + } + } + + svfloat32_t alpha_vec = svdup_f32(alpha); + // Iterate through shared dimension (K) + for (size_t k = 0; k < shared_dim; k++) { +#if !defined(TRANSA) + // Load column of A + svfloat32_t col_a_0 = svld1(pg_a_0, &A[k * svl]); + col_a_0 = svmul_x(pg, alpha_vec, col_a_0); + svfloat32_t col_a_1 = svld1(pg_a_1, &A[(k + shared_dim) * svl]); + col_a_1 = svmul_x(pg, alpha_vec, col_a_1); + + // Load row of A**T + svfloat32_t row_b_0 = svld1(pg_b_0, &B[k * svl]); + svfloat32_t row_b_1 = svld1(pg_b_1, &B[(k + shared_dim) * svl]); +#else + // Load column of A**T + svfloat32_t col_a_0 = svld1(pg_a_0, &A[k * ldb]); + col_a_0 = svmul_x(pg, alpha_vec, col_a_0); + + svfloat32_t col_a_1 = svld1(pg_a_1, &A[k * ldb + svl]); + col_a_1 = svmul_x(pg, alpha_vec, col_a_1); + + // Load row of A + svfloat32_t row_b_0 = svld1(pg_b_0, &B[k * ldb]); + svfloat32_t row_b_1 = svld1(pg_b_1, &B[k * ldb + svl]); +#endif + // Perform outer product + svmopa_za32_m(/*tile*/0, pg, pg, col_a_0, row_b_0); + svmopa_za32_m(/*tile*/1, pg, pg, col_a_0, row_b_1); + svmopa_za32_m(/*tile*/2, pg, pg, col_a_1, row_b_0); + svmopa_za32_m(/*tile*/3, pg, pg, col_a_1, row_b_1); + } + +#if defined(UPPER) +#define pg_c_0_full pg_c_0 +#define pg_c_1_full pg_c_1 + + bool need_update_pg_b = true; + size_t last_invalid_index = col_idx - row_idx; + // For Upper, If col_idx - row_idx >= 2*svl, we don't need to update the predicate due to all elements above the digonal + if (col_idx - row_idx >= 2*svl) { + need_update_pg_b = false; + } + // Store to C from ZA + for (size_t i = 0; i < MIN(svl, block_rows); i++, last_invalid_index++) { + if (need_update_pg_b) { + pg_c_0 = svnot_b_z(pg_c_0_full, svwhilelt_b32_u64(0, last_invalid_index)); + pg_c_1 = svnot_b_z(pg_c_1_full, svwhilelt_b32_u64(svl, last_invalid_index)); + } + + svst1_hor_za32(/*tile*/0, /*slice*/i, pg_c_0, &C[i * ldc]); + svst1_hor_za32(/*tile*/1, /*slice*/i, pg_c_1, &C[i * ldc + svl]); + } + for (size_t i = svl; i < block_rows; i++,last_invalid_index++) { + if (need_update_pg_b) { + pg_c_0 = svnot_b_z(pg_c_0_full, svwhilelt_b32_u64(0, last_invalid_index)); + pg_c_1 = svnot_b_z(pg_c_1_full, svwhilelt_b32_u64(svl, last_invalid_index)); + } + svst1_hor_za32(/*tile*/2, /*slice*/i - svl, pg_c_0, &C[i * ldc]); + svst1_hor_za32(/*tile*/3, /*slice*/i - svl, pg_c_1, &C[i * ldc + svl]); + } +#else + // Store to C from ZA + size_t valid_index = row_idx - col_idx + 1; + for (size_t i = 0; i < MIN(svl, block_rows); i++, valid_index++) { + pg_c_0 = svwhilelt_b32_u64(0, MIN(valid_index, block_cols)); + pg_c_1 = svwhilelt_b32_u64(svl, MIN(valid_index, block_cols)); + svst1_hor_za32(/*tile*/0, /*slice*/i, pg_c_0, &C[i * ldc]); + svst1_hor_za32(/*tile*/1, /*slice*/i, pg_c_1, &C[i * ldc + svl]); + } + for (size_t i = svl; i < block_rows; i++, valid_index++) { + pg_c_0 = svwhilelt_b32_u64(0, MIN(valid_index, block_cols)); + pg_c_1 = svwhilelt_b32_u64(svl, MIN(valid_index, block_cols)); + svst1_hor_za32(/*tile*/2, /*slice*/i - svl, pg_c_0, &C[i * ldc]); + svst1_hor_za32(/*tile*/3, /*slice*/i - svl, pg_c_1, &C[i * ldc + svl]); + } +#endif +} + +__arm_new("za") __arm_locally_streaming +static void ssyrk_direct_sme1_2VLx2VL(uint64_t n, uint64_t k, const float* alpha,\ + const float *ba, const float* beta, float *restrict bc) { + const uint64_t num_rows = n; + const uint64_t num_cols = n; + + const float *restrict a_ptr = ba; + const float *restrict b_ptr = ba; + float *restrict c_ptr = bc; + + const uint64_t svl = svcntw(); + const uint64_t ldc = n; + + // Block over rows of C (panels of A) + uint64_t row_idx = 0; + + // 2x2 loop + uint64_t row_batch = 2*svl; + + // Block over row dimension of C + for (; row_idx < num_rows; row_idx += row_batch) { + row_batch = MIN(row_batch, num_rows - row_idx); + uint64_t col_batch = 2*svl; +#if defined(UPPER) + // for UPLO is upper, Start from column col_idx = rows_index to ensure we only process the upper triangle (col_idx >= rows_index) + for (uint64_t col_idx = row_idx; col_idx < num_cols; col_idx += col_batch) { + col_batch = MIN(col_batch, num_cols - col_idx); +#else + // for UPLO is lower, we only process the lower triangle part (col_idx <= row_idxx) + for (uint64_t col_idx = 0; col_idx < num_cols && col_idx <= row_idx; col_idx += col_batch) { +#endif + col_batch = MIN(col_batch, num_cols - col_idx); +#if !defined(TRANSA) + kernel_2x2(&a_ptr[row_idx * k], &b_ptr[col_idx * k], + &c_ptr[row_idx * ldc + col_idx], k, + ldc, row_batch, col_batch, *alpha, *beta, row_idx, col_idx); +#else + kernel_2x2(&a_ptr[row_idx], &b_ptr[col_idx], + &c_ptr[row_idx * ldc + col_idx], k, + ldc, row_batch, col_batch, *alpha, *beta, row_idx, col_idx); +#endif + + } + } + return; +} + +#else +static void ssyrk_direct_sme1_2VLx2VL(uint64_t n, uint64_t k, const float* alpha,\ + const float *ba, const float* beta, float *restrict bc){} +#endif + +void CNAME (BLASLONG N, BLASLONG K, float alpha, float * __restrict A,\ + BLASLONG strideA, float beta, float * __restrict C, BLASLONG strideC){ + if (alpha == 0.0f || K == 0) { + if (beta == 1.0f) + return; + ssyrk_direct_sme1_2VLx2VL(N, 0, &alpha, A, &beta, C); + return; + } + +#if !defined(TRANSA) + uint64_t n_mod, vl_elms; + + vl_elms = sve_cntw(); + + n_mod = (((uint64_t)N + vl_elms - 1) / vl_elms) * vl_elms; + + float *A_mod = (float *) malloc(n_mod*K*sizeof(float)); + + /* Prevent compiler optimization by reading from memory instead + * of reading directly from vector (z) registers. + * */ + asm volatile("" : : :"p0", "p1", "p2", "p3", "p4", "p5", "p6", "p7", + "p8", "p9", "p10", "p11", "p12", "p13", "p14", "p15", + "z0", "z1", "z2", "z3", "z4", "z5", "z6", "z7", + "z8", "z9", "z10", "z11", "z12", "z13", "z14", "z15", + "z16", "z17", "z18", "z19", "z20", "z21", "z22", "z23", + "z24", "z25", "z26", "z27", "z28", "z29", "z30", "z31"); + + /* Pre-process the left matrix to make it suitable for + matrix sum of outer-product calculation + */ + SGEMM_PREPROCESS (N, K, A, A_mod); + asm volatile("" : : :"p0", "p1", "p2", "p3", "p4", "p5", "p6", "p7", + "p8", "p9", "p10", "p11", "p12", "p13", "p14", "p15", + "z0", "z1", "z2", "z3", "z4", "z5", "z6", "z7", + "z8", "z9", "z10", "z11", "z12", "z13", "z14", "z15", + "z16", "z17", "z18", "z19", "z20", "z21", "z22", "z23", + "z24", "z25", "z26", "z27", "z28", "z29", "z30", "z31"); + ssyrk_direct_sme1_2VLx2VL(N, K, &alpha, A_mod, &beta, C); + free(A_mod); +#else + ssyrk_direct_sme1_2VLx2VL(N, K, &alpha, A, &beta, C); +#endif + +} + +#else + +void CNAME (BLASLONG N, BLASLONG K, float alpha, float * __restrict A,\ + BLASLONG strideA, float beta, float * __restrict C, BLASLONG strideC){ +fprintf(stderr,"empty ssyrk_direct kernel should never be called\n"); +} + +#endif