|
| 1 | +/* |
| 2 | + Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. |
| 3 | + SPDX-License-Identifier: BSD-3-Clause-Clear |
| 4 | +*/ |
| 5 | + |
| 6 | +#include "common.h" |
| 7 | +#include <stdlib.h> |
| 8 | +#include <inttypes.h> |
| 9 | +#include <math.h> |
| 10 | +#include "sme_abi.h" |
| 11 | +#if defined(HAVE_SME) |
| 12 | + |
| 13 | +#if defined(__ARM_FEATURE_SME) && defined(__clang__) && __clang_major__ >= 16 |
| 14 | +#include <arm_sme.h> |
| 15 | +#endif |
| 16 | + |
| 17 | +/* Function prototypes */ |
| 18 | +// For Upper and NonUnit Triangular preprocess |
| 19 | +extern void strmm_direct_sme1_preprocess_UN(uint64_t nbr, uint64_t nbc, const float * restrict a, float * a_mod); |
| 20 | +// For Lower and NonUnit Triangular preprocess |
| 21 | +extern void strmm_direct_sme1_preprocess_LN(uint64_t nbr, uint64_t nbc, const float * restrict a, float * a_mod); |
| 22 | + |
| 23 | +/* Function Definitions */ |
| 24 | +static uint64_t sve_cntw() { |
| 25 | + uint64_t cnt; |
| 26 | + asm volatile( |
| 27 | + "rdsvl %[res], #1\n" |
| 28 | + "lsr %[res], %[res], #2\n" |
| 29 | + : [res] "=r" (cnt) :: |
| 30 | + ); |
| 31 | + return cnt; |
| 32 | +} |
| 33 | + |
| 34 | +#if defined(__ARM_FEATURE_SME) && defined(__ARM_FEATURE_LOCALLY_STREAMING) && defined(__clang__) && __clang_major__ >= 16 |
| 35 | +// Outer product kernel. |
| 36 | +// Computes a 2SVL x 2SVL block of C, utilizing all four FP32 tiles of ZA. |
| 37 | +__attribute__((always_inline)) inline void |
| 38 | +kernel_2x2(const float *A, const float *B, float *C, size_t shared_dim, |
| 39 | + size_t ldc, size_t block_rows, size_t block_cols, float alpha, uint64_t row_idx) |
| 40 | + __arm_out("za") __arm_streaming { |
| 41 | + const uint64_t svl = svcntw(); |
| 42 | + size_t ldb = ldc; |
| 43 | + // Predicate set-up |
| 44 | + svbool_t pg = svptrue_b32(); |
| 45 | + svbool_t pg_a_0 = svwhilelt_b32_u64(0, block_rows); |
| 46 | + svbool_t pg_a_1 = svwhilelt_b32_u64(svl, block_rows); |
| 47 | + |
| 48 | +#if (!defined(TRANSA) && !defined(UPPER)) || (defined(TRANSA) && defined(UPPER)) |
| 49 | +#define pg_a_0_full pg_a_0 |
| 50 | +#define pg_a_1_full pg_a_1 |
| 51 | +#endif |
| 52 | + svbool_t pg_b_0 = svwhilelt_b32_u64(0, block_cols); |
| 53 | + svbool_t pg_b_1 = svwhilelt_b32_u64(svl, block_cols); |
| 54 | + |
| 55 | +#define pg_c_0 pg_b_0 |
| 56 | +#define pg_c_1 pg_b_1 |
| 57 | + |
| 58 | + svzero_za(); |
| 59 | + svfloat32_t alpha_vec = svdup_f32(alpha); |
| 60 | + // Iterate through shared dimension (K) |
| 61 | +#if (!defined(TRANSA) && defined(UPPER)) || (defined(TRANSA) && !defined(UPPER)) |
| 62 | + for (size_t k = row_idx, valid_index = 1; k < shared_dim; k++,valid_index++) { |
| 63 | + pg_a_0 = svwhilelt_b32_u64(0, MIN(valid_index, block_rows)); |
| 64 | + pg_a_1 = svwhilelt_b32_u64(svl, MIN(valid_index, block_rows)); |
| 65 | +#else |
| 66 | + for (size_t k = 0; k < MIN(row_idx + block_rows, shared_dim); k++) { |
| 67 | + // If k exceeds row_idx, mask out rows before (k - row_idx) |
| 68 | + // This ensures only valid rows are included for lower triangular logic. |
| 69 | + if (k > row_idx) { |
| 70 | + pg_a_0 = svnot_b_z(pg_a_0_full, svwhilelt_b32_u64(0, k - row_idx)); |
| 71 | + pg_a_1 = svnot_b_z(pg_a_1_full, svwhilelt_b32_u64(svl, k - row_idx)); |
| 72 | + } |
| 73 | +#endif |
| 74 | + |
| 75 | +#if !defined(TRANSA) |
| 76 | + // Load column of A |
| 77 | + svfloat32_t col_a_0 = svld1(pg_a_0, &A[k * svl]); |
| 78 | + svfloat32_t col_a_1 = svld1(pg_a_1, &A[(k + shared_dim) * svl]); |
| 79 | +#else |
| 80 | + svfloat32_t col_a_0 = svld1(pg_a_0, &A[k * shared_dim]); |
| 81 | + svfloat32_t col_a_1 = svld1(pg_a_1, &A[k * shared_dim + svl]); |
| 82 | +#endif |
| 83 | + col_a_0 = svmul_x(pg_a_0, alpha_vec, col_a_0); |
| 84 | + col_a_1 = svmul_x(pg_a_1, alpha_vec, col_a_1); |
| 85 | + // Load row of B |
| 86 | + svfloat32_t row_b_0 = svld1(pg_b_0, &B[k * ldb]); |
| 87 | + svfloat32_t row_b_1 = svld1(pg_b_1, &B[k * ldb + svl]); |
| 88 | + // Perform outer product |
| 89 | + svmopa_za32_m(/*tile*/0, pg_a_0, pg, col_a_0, row_b_0); |
| 90 | + svmopa_za32_m(/*tile*/1, pg_a_0, pg, col_a_0, row_b_1); |
| 91 | + svmopa_za32_m(/*tile*/2, pg_a_1, pg, col_a_1, row_b_0); |
| 92 | + svmopa_za32_m(/*tile*/3, pg_a_1, pg, col_a_1, row_b_1); |
| 93 | + } |
| 94 | + |
| 95 | + // Store to C from ZA |
| 96 | + for (size_t i = 0; i < MIN(svl, block_rows); i++) { |
| 97 | + svst1_hor_za32(/*tile*/0, /*slice*/i, pg_c_0, &C[i * ldc]); |
| 98 | + svst1_hor_za32(/*tile*/1, /*slice*/i, pg_c_1, &C[i * ldc + svl]); |
| 99 | + } |
| 100 | + for (size_t i = svl; i < block_rows; i++) { |
| 101 | + svst1_hor_za32(/*tile*/2, /*slice*/i, pg_c_0, &C[i * ldc]); |
| 102 | + svst1_hor_za32(/*tile*/3, /*slice*/i, pg_c_1, &C[i * ldc + svl]); |
| 103 | + } |
| 104 | + |
| 105 | +} |
| 106 | + |
| 107 | +__arm_new("za") __arm_locally_streaming |
| 108 | +static inline void strmm_direct_alpha_sme1_2VLx2VL(uint64_t m, uint64_t k, uint64_t n, const float* alpha,\ |
| 109 | + const float *ba, float *restrict bb) { |
| 110 | + const uint64_t num_rows = m; |
| 111 | + const uint64_t num_cols = n; |
| 112 | + |
| 113 | + const float *restrict a_ptr = ba; |
| 114 | + const float *restrict b_ptr = bb; |
| 115 | + float *restrict c_ptr = bb; |
| 116 | + |
| 117 | + const uint64_t svl = svcntw(); |
| 118 | + const uint64_t svl_x2 = 2*svl; |
| 119 | + const uint64_t ldc = n; |
| 120 | + |
| 121 | + |
| 122 | + uint64_t row_idx = 0; |
| 123 | +#if (!defined(TRANSA) && defined(UPPER)) || (defined(TRANSA) && !defined(UPPER)) |
| 124 | + // 2x2 loop |
| 125 | + uint64_t row_batch = svl_x2; |
| 126 | + // Block over rows of C (panels of A) |
| 127 | + for (; row_idx < num_rows; row_idx += row_batch) { |
| 128 | + row_batch = MIN(row_batch, num_rows - row_idx); |
| 129 | +#else |
| 130 | + // Calculate the remainder of num_rows divided by 2VL to determine tail tile size |
| 131 | + uint64_t row_batch = num_rows % svl_x2; |
| 132 | + // If there's no remainder, use full tile size (2VL) for initial batch |
| 133 | + if (row_batch == 0) row_batch = svl_x2; |
| 134 | + // Loop from bottom to top, processing rows in batches |
| 135 | + for (uint64_t index = num_rows; index > 0; index -= row_batch, row_batch = svl_x2) { |
| 136 | + // Compute the starting row index for the current batch |
| 137 | + row_idx = index - row_batch; |
| 138 | +#endif |
| 139 | + uint64_t col_idx = 0; |
| 140 | + uint64_t col_batch = svl_x2; |
| 141 | + // Block over column dimension of C |
| 142 | + for (; col_idx < num_cols; col_idx += col_batch) { |
| 143 | + col_batch = MIN(col_batch, num_cols - col_idx); |
| 144 | +#if !defined(TRANSA) |
| 145 | + kernel_2x2(&a_ptr[row_idx * k], &b_ptr[col_idx], |
| 146 | + &c_ptr[row_idx * ldc + col_idx], k, |
| 147 | + ldc, row_batch, col_batch, *alpha, row_idx); |
| 148 | +#else |
| 149 | + kernel_2x2(&a_ptr[row_idx], &b_ptr[col_idx], |
| 150 | + &c_ptr[row_idx * ldc + col_idx], k, |
| 151 | + ldc, row_batch, col_batch, *alpha, row_idx); |
| 152 | +#endif |
| 153 | + } |
| 154 | + } |
| 155 | + |
| 156 | + return; |
| 157 | +} |
| 158 | + |
| 159 | +#else |
| 160 | +void strmm_direct_alpha_sme1_2VLx2VL(uint64_t m, uint64_t k, uint64_t n, const float* alpha,\ |
| 161 | + const float *ba, float *restrict bb){} |
| 162 | +#endif |
| 163 | + |
| 164 | +void CNAME (BLASLONG M, BLASLONG N, BLASLONG K, float alpha, float * __restrict A,\ |
| 165 | + BLASLONG strideA, float * __restrict B, BLASLONG strideB){ |
| 166 | +#if !defined(TRANSA) |
| 167 | + uint64_t m_mod, vl_elms; |
| 168 | + |
| 169 | + vl_elms = sve_cntw(); |
| 170 | + |
| 171 | + m_mod = ceil((double)M/(double)vl_elms) * vl_elms; |
| 172 | + |
| 173 | + float *A_mod = (float *) malloc(m_mod*K*sizeof(float)); |
| 174 | +#if defined(UPPER) |
| 175 | + strmm_direct_sme1_preprocess_UN(M, K, A, A_mod); |
| 176 | +#else |
| 177 | + strmm_direct_sme1_preprocess_LN(M, K, A, A_mod); |
| 178 | +#endif |
| 179 | + /* Calculate B = alpha*A*B*/ |
| 180 | + strmm_direct_alpha_sme1_2VLx2VL(M, K, N, &alpha, A_mod, B); |
| 181 | + free(A_mod); |
| 182 | +#else |
| 183 | + strmm_direct_alpha_sme1_2VLx2VL(M, K, N, &alpha, A, B); |
| 184 | +#endif |
| 185 | +} |
| 186 | + |
| 187 | +#else |
| 188 | +void CNAME (BLASLONG M, BLASLONG N, BLASLONG K, float alpha, float * __restrict A,\ |
| 189 | + BLASLONG strideA, float * __restrict B, BLASLONG strideB){} |
| 190 | + |
| 191 | +#endif |
0 commit comments