Skip to content

Commit 6468bcb

Browse files
martin-frbgmartin-frbg
authored
Merge pull request #5889 from moluopro/develop
ARM64 SME: fix zero-scalar handling and clean up direct kernels
2 parents c61f95c + 130102e commit 6468bcb

6 files changed

Lines changed: 122 additions & 76 deletions

‎kernel/arm64/sgemm_direct_alpha_beta_arm64_sme1.c‎

Lines changed: 29 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@
66
#include "common.h"
77
#include <stdlib.h>
88
#include <inttypes.h>
9-
#include <math.h>
109

1110
#if defined(DYNAMIC_ARCH)
1211
#define COMBINE(a,b) a ## b
@@ -64,27 +63,30 @@ kernel_2x2(const float *A, const float *B, float *C, size_t shared_dim,
6463
#define pg_c_1 pg_b_1
6564

6665
svzero_za();
67-
svfloat32_t beta_vec = svdup_f32(beta);
66+
// beta == 0 must not read C; ZA is already initialized to zero.
67+
if (beta != 0.0f) {
68+
svfloat32_t beta_vec = svdup_f32(beta);
6869
// Load C to ZA
69-
for (size_t i = 0; i < MIN(svl, block_rows); i++) {
70-
svfloat32_t row_c_0 = svld1(pg_c_0, &C[i * ldc]);
71-
row_c_0 = svmul_x(pg, beta_vec, row_c_0);
72-
svwrite_hor_za32_f32_m(/*tile*/0, /*slice*/i, pg_c_0, row_c_0);
73-
74-
svfloat32_t row_c_1 = svld1(pg_c_1, &C[i * ldc + svl]);
75-
row_c_1 = svmul_x(pg, beta_vec, row_c_1);
76-
svwrite_hor_za32_f32_m(/*tile*/1, /*slice*/i, pg_c_1, row_c_1);
70+
for (size_t i = 0; i < MIN(svl, block_rows); i++) {
71+
svfloat32_t row_c_0 = svld1(pg_c_0, &C[i * ldc]);
72+
row_c_0 = svmul_x(pg, beta_vec, row_c_0);
73+
svwrite_hor_za32_f32_m(/*tile*/0, /*slice*/i, pg_c_0, row_c_0);
74+
75+
svfloat32_t row_c_1 = svld1(pg_c_1, &C[i * ldc + svl]);
76+
row_c_1 = svmul_x(pg, beta_vec, row_c_1);
77+
svwrite_hor_za32_f32_m(/*tile*/1, /*slice*/i, pg_c_1, row_c_1);
78+
}
79+
for (size_t i = svl; i < block_rows; i++) {
80+
svfloat32_t row_c_0 = svld1(pg_c_0, &C[i * ldc]);
81+
row_c_0 = svmul_x(pg, beta_vec, row_c_0);
82+
svwrite_hor_za32_f32_m(/*tile*/2, /*slice*/i, pg_c_0, row_c_0);
83+
84+
svfloat32_t row_c_1 = svld1(pg_c_1, &C[i * ldc + svl]);
85+
row_c_1 = svmul_x(pg, beta_vec, row_c_1);
86+
svwrite_hor_za32_f32_m(/*tile*/3, /*slice*/i, pg_c_1, row_c_1);
87+
}
7788
}
78-
for (size_t i = svl; i < block_rows; i++) {
79-
svfloat32_t row_c_0 = svld1(pg_c_0, &C[i * ldc]);
80-
row_c_0 = svmul_x(pg, beta_vec, row_c_0);
81-
svwrite_hor_za32_f32_m(/*tile*/2, /*slice*/i, pg_c_0, row_c_0);
8289

83-
svfloat32_t row_c_1 = svld1(pg_c_1, &C[i * ldc + svl]);
84-
row_c_1 = svmul_x(pg, beta_vec, row_c_1);
85-
svwrite_hor_za32_f32_m(/*tile*/3, /*slice*/i, pg_c_1, row_c_1);
86-
}
87-
8890
svfloat32_t alpha_vec = svdup_f32(alpha);
8991
// Iterate through shared dimension (K)
9092
for (size_t k = 0; k < shared_dim; k++) {
@@ -168,12 +170,18 @@ void SME1_KERNEL2X2(uint64_t m, uint64_t k, uint64_t n, const float* alpha,\
168170
void CNAME (BLASLONG M, BLASLONG N, BLASLONG K, float alpha, float * __restrict A,\
169171
BLASLONG strideA, float * __restrict B, BLASLONG strideB ,\
170172
float beta, float * __restrict R, BLASLONG strideR){
171-
173+
if (alpha == 0.0f || K == 0) {
174+
if (beta == 1.0f)
175+
return;
176+
SME1_KERNEL2X2(M, 0, N, &alpha, A, B, &beta, R);
177+
return;
178+
}
179+
172180
uint64_t m_mod, vl_elms;
173181

174182
vl_elms = sve_cntw();
175183

176-
m_mod = ceil((double)M/(double)vl_elms) * vl_elms;
184+
m_mod = (((uint64_t)M + vl_elms - 1) / vl_elms) * vl_elms;
177185

178186
float *A_mod = (float *) malloc(m_mod*K*sizeof(float));
179187

‎kernel/arm64/sgemm_direct_arm64_sme1.c‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@
66
#include "common.h"
77
#include <stdlib.h>
88
#include <inttypes.h>
9-
#include <math.h>
109
#if defined(DYNAMIC_ARCH)
1110
#define COMBINE(a,b) a ## b
1211
#define COMBINE2(a,b) COMBINE(a,b)
@@ -50,7 +49,7 @@ void CNAME (BLASLONG M, BLASLONG N, BLASLONG K, float * __restrict A,\
5049
uint64_t m_mod, vl_elms;
5150

5251
vl_elms = sve_cntw();
53-
m_mod = ceil((double)M/(double)vl_elms) * vl_elms;
52+
m_mod = (((uint64_t)M + vl_elms - 1) / vl_elms) * vl_elms;
5453

5554
float *A_mod = (float *) malloc(m_mod*K*sizeof(float));
5655

‎kernel/arm64/ssymm_direct_alpha_beta_arm64_sme1.c‎

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@
66
#include "common.h"
77
#include <stdlib.h>
88
#include <inttypes.h>
9-
#include <math.h>
109
// #include "sme_abi.h"
1110
#if defined(HAVE_SME)
1211

@@ -46,6 +45,7 @@ static uint64_t sve_cntw() {
4645

4746
#if defined(__ARM_FEATURE_SME) && defined(__ARM_FEATURE_LOCALLY_STREAMING) && defined(__clang__) && __clang_major__ >= 16
4847

48+
#if defined(UPPER)
4949
__arm_new("za") __arm_locally_streaming
5050
static void ssymm_direct_sme1_preprocessLU(uint64_t nbr, uint64_t nbc,
5151
const float *restrict a, float *restrict a_mod)
@@ -55,7 +55,7 @@ static void ssymm_direct_sme1_preprocessLU(uint64_t nbr, uint64_t nbc,
5555
const uint64_t svl = svcntw();
5656
uint64_t row_batch = svl;
5757

58-
float *restrict pSrc;
58+
const float *restrict pSrc;
5959
float *restrict pDst;
6060
for (uint64_t row_idx = 0; row_idx < nbr; row_idx += row_batch)
6161
{
@@ -123,8 +123,10 @@ static void ssymm_direct_sme1_preprocessLU(uint64_t nbr, uint64_t nbc,
123123
}
124124
}
125125
}
126+
#endif
126127

127128
//
129+
#if defined(LOWER)
128130
__arm_new("za") __arm_locally_streaming
129131
static void ssymm_direct_sme1_preprocessLL(uint64_t nbr, uint64_t nbc,
130132
const float *restrict a, float *restrict a_mod)
@@ -133,7 +135,7 @@ static void ssymm_direct_sme1_preprocessLL(uint64_t nbr, uint64_t nbc,
133135
const uint64_t svl = svcntw();
134136
uint64_t row_batch = svl;
135137

136-
float *restrict pSrc;
138+
const float *restrict pSrc;
137139
float *restrict pDst;
138140
for (uint64_t row_idx = 0; row_idx < nbr; row_idx += row_batch)
139141
{
@@ -201,20 +203,32 @@ static void ssymm_direct_sme1_preprocessLL(uint64_t nbr, uint64_t nbc,
201203
}
202204
}
203205
}
206+
#endif
204207
#else
208+
#if defined(UPPER)
205209
static void ssymm_direct_sme1_preprocessLU(uint64_t nbr, uint64_t nbc,
206210
const float *restrict a, float *restrict a_mod){}
211+
#endif
212+
#if defined(LOWER)
207213
static void ssymm_direct_sme1_preprocessLL(uint64_t nbr, uint64_t nbc,
208214
const float *restrict a, float *restrict a_mod){}
209215
#endif
216+
#endif
210217

211218
//
212219
void CNAME(BLASLONG M, BLASLONG N, float alpha, float *__restrict A,
213220
BLASLONG strideA, float *__restrict B, BLASLONG strideB,
214221
float beta, float *__restrict R, BLASLONG strideR)
215222
{
223+
if (alpha == 0.0f) {
224+
if (beta == 1.0f)
225+
return;
226+
SGEMM_DIRECT2X2(M, 0, N, &alpha, A, B, &beta, R);
227+
return;
228+
}
229+
216230
uint64_t vl_elms = sve_cntw(); // vl_elem = 16
217-
uint64_t m_mod = ceil((double)M / (double)vl_elms) * vl_elms;
231+
uint64_t m_mod = (((uint64_t)M + vl_elms - 1) / vl_elms) * vl_elms;
218232

219233
/* Pre-process the left matrix to make it suitable for
220234
matrix sum of outer-product calculation

‎kernel/arm64/ssyr2k_direct_alpha_beta_arm64_sme1.c‎

Lines changed: 36 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@
66
#include "common.h"
77
#include <stdlib.h>
88
#include <inttypes.h>
9-
#include <math.h>
109

1110
#if defined(DYNAMIC_ARCH)
1211
#define COMBINE(a,b) a ## b
@@ -30,6 +29,7 @@ extern void SGEMM_PREPROCESS(uint64_t nbr, uint64_t nbc,\
3029
const float * restrict a, float * a_mod) ;
3130

3231
/* Function Definitions */
32+
#if !defined(TRANSA)
3333
static uint64_t sve_cntw() {
3434
uint64_t cnt;
3535
asm volatile(
@@ -39,18 +39,21 @@ static uint64_t sve_cntw() {
3939
);
4040
return cnt;
4141
}
42+
#endif
4243

4344
#if defined(__ARM_FEATURE_SME) && defined(__ARM_FEATURE_LOCALLY_STREAMING) && defined(__clang__) && __clang_major__ >= 16
4445
// Outer product kernel.
4546
// Computes a 2SVL x 2SVL block of C, utilizing all four FP32 tiles of ZA.
4647
__attribute__((always_inline)) inline void
47-
kernel_2x2(const float *A, float *B_T, const float *B, float *A_T, float *C, size_t shared_dim,
48+
kernel_2x2(const float *A, const float *B_T, const float *B, const float *A_T, float *C, size_t shared_dim,
4849
size_t ldc, size_t block_rows, size_t block_cols, float alpha,
4950
float beta, uint64_t row_idx, uint64_t col_idx)
5051
__arm_out("za") __arm_streaming {
5152

5253
const uint64_t svl = svcntw();
54+
#if defined(TRANSA)
5355
size_t ldb = ldc;
56+
#endif
5457
// Predicate set-up
5558
svbool_t pg = svptrue_b32();
5659
svbool_t pg_a_0 = svwhilelt_b32_u64(0, block_rows);
@@ -63,26 +66,29 @@ kernel_2x2(const float *A, float *B_T, const float *B, float *A_T, float *C, siz
6366
#define pg_c_1 pg_b_1
6467

6568
svzero_za();
66-
svfloat32_t beta_vec = svdup_f32(beta);
67-
68-
// Load C to ZA
69-
for (size_t i = 0; i < MIN(svl, block_rows); i++) {
70-
svfloat32_t row_c_0 = svld1(pg_c_0, &C[i * ldc]);
71-
row_c_0 = svmul_x(pg, beta_vec, row_c_0);
72-
svwrite_hor_za32_f32_m(/*tile*/0, /*slice*/i, pg_c_0, row_c_0);
73-
74-
svfloat32_t row_c_1 = svld1(pg_c_1, &C[i * ldc + svl]);
75-
row_c_1 = svmul_x(pg, beta_vec, row_c_1);
76-
svwrite_hor_za32_f32_m(/*tile*/1, /*slice*/i, pg_c_1, row_c_1);
77-
}
78-
for (size_t i = svl; i < block_rows; i++) {
79-
svfloat32_t row_c_0 = svld1(pg_c_0, &C[i * ldc]);
80-
row_c_0 = svmul_x(pg, beta_vec, row_c_0);
81-
svwrite_hor_za32_f32_m(/*tile*/2, /*slice*/i, pg_c_0, row_c_0);
82-
83-
svfloat32_t row_c_1 = svld1(pg_c_1, &C[i * ldc + svl]);
84-
row_c_1 = svmul_x(pg, beta_vec, row_c_1);
85-
svwrite_hor_za32_f32_m(/*tile*/3, /*slice*/i, pg_c_1, row_c_1);
69+
// beta == 0 must not read C; ZA is already initialized to zero.
70+
if (beta != 0.0f) {
71+
svfloat32_t beta_vec = svdup_f32(beta);
72+
73+
// Load C to ZA
74+
for (size_t i = 0; i < MIN(svl, block_rows); i++) {
75+
svfloat32_t row_c_0 = svld1(pg_c_0, &C[i * ldc]);
76+
row_c_0 = svmul_x(pg, beta_vec, row_c_0);
77+
svwrite_hor_za32_f32_m(/*tile*/0, /*slice*/i, pg_c_0, row_c_0);
78+
79+
svfloat32_t row_c_1 = svld1(pg_c_1, &C[i * ldc + svl]);
80+
row_c_1 = svmul_x(pg, beta_vec, row_c_1);
81+
svwrite_hor_za32_f32_m(/*tile*/1, /*slice*/i, pg_c_1, row_c_1);
82+
}
83+
for (size_t i = svl; i < block_rows; i++) {
84+
svfloat32_t row_c_0 = svld1(pg_c_0, &C[i * ldc]);
85+
row_c_0 = svmul_x(pg, beta_vec, row_c_0);
86+
svwrite_hor_za32_f32_m(/*tile*/2, /*slice*/i, pg_c_0, row_c_0);
87+
88+
svfloat32_t row_c_1 = svld1(pg_c_1, &C[i * ldc + svl]);
89+
row_c_1 = svmul_x(pg, beta_vec, row_c_1);
90+
svwrite_hor_za32_f32_m(/*tile*/3, /*slice*/i, pg_c_1, row_c_1);
91+
}
8692
}
8793

8894
svfloat32_t alpha_vec = svdup_f32(alpha);
@@ -250,12 +256,19 @@ void CNAME (BLASLONG N, BLASLONG K, float alpha, float * __restrict A, \
250256
BLASLONG strideA, float * __restrict B, BLASLONG strideB, \
251257
float beta, float * __restrict R, BLASLONG strideR)
252258
{
259+
if (alpha == 0.0f || K == 0) {
260+
if (beta == 1.0f)
261+
return;
262+
ssyr2k_direct_sme1_2VLx2VL(N, 0, &alpha, A, B, &beta, R);
263+
return;
264+
}
265+
253266
#if !defined(TRANSA)
254267
uint64_t n_mod, vl_elms;
255268

256269
vl_elms = sve_cntw();
257270

258-
n_mod = ceil((double)N/(double)vl_elms) * vl_elms;
271+
n_mod = (((uint64_t)N + vl_elms - 1) / vl_elms) * vl_elms;
259272

260273
float *A_mod = (float *) malloc(n_mod*K*sizeof(float));
261274
float *B_mod = (float *) malloc(n_mod*K*sizeof(float));

‎kernel/arm64/ssyrk_direct_alpha_beta_arm64_sme1.c‎

Lines changed: 37 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@
66
#include "common.h"
77
#include <stdlib.h>
88
#include <inttypes.h>
9-
#include <math.h>
109
#if defined(HAVE_SME)
1110

1211
#if defined(DYNAMIC_ARCH)
@@ -31,6 +30,7 @@ extern void SGEMM_PREPROCESS (uint64_t nbr, uint64_t nbc,\
3130
const float * restrict a, float * a_mod) ;
3231

3332
/* Function Definitions */
33+
#if !defined(TRANSA)
3434
static uint64_t sve_cntw() {
3535
uint64_t cnt;
3636
asm volatile(
@@ -40,18 +40,21 @@ static uint64_t sve_cntw() {
4040
);
4141
return cnt;
4242
}
43+
#endif
4344

4445
#if defined(__ARM_FEATURE_SME) && defined(__ARM_FEATURE_LOCALLY_STREAMING) && defined(__clang__) && __clang_major__ >= 16
4546
// Outer product kernel.
4647
// Computes a 2SVL x 2SVL block of C, utilizing all four FP32 tiles of ZA.
4748
__attribute__((always_inline)) inline void
48-
kernel_2x2(const float *A, float *B, float *C, size_t shared_dim,
49+
kernel_2x2(const float *A, const float *B, float *C, size_t shared_dim,
4950
size_t ldc, size_t block_rows, size_t block_cols, float alpha,
5051
float beta, uint64_t row_idx, uint64_t col_idx)
5152
__arm_out("za") __arm_streaming {
5253

5354
const uint64_t svl = svcntw();
55+
#if defined(TRANSA)
5456
size_t ldb = ldc;
57+
#endif
5558
// Predicate set-up
5659
svbool_t pg = svptrue_b32();
5760
svbool_t pg_a_0 = svwhilelt_b32_u64(0, block_rows);
@@ -64,28 +67,31 @@ kernel_2x2(const float *A, float *B, float *C, size_t shared_dim,
6467
#define pg_c_1 pg_b_1
6568

6669
svzero_za();
67-
svfloat32_t beta_vec = svdup_f32(beta);
70+
// beta == 0 must not read C; ZA is already initialized to zero.
71+
if (beta != 0.0f) {
72+
svfloat32_t beta_vec = svdup_f32(beta);
6873

69-
// Load C to ZA
70-
for (size_t i = 0; i < MIN(svl, block_rows); i++) {
71-
svfloat32_t row_c_0 = svld1(pg_c_0, &C[i * ldc]);
72-
row_c_0 = svmul_x(pg, beta_vec, row_c_0);
73-
svwrite_hor_za32_f32_m(/*tile*/0, /*slice*/i, pg_c_0, row_c_0);
74-
75-
svfloat32_t row_c_1 = svld1(pg_c_1, &C[i * ldc + svl]);
76-
row_c_1 = svmul_x(pg, beta_vec, row_c_1);
77-
svwrite_hor_za32_f32_m(/*tile*/1, /*slice*/i, pg_c_1, row_c_1);
78-
}
79-
for (size_t i = svl; i < block_rows; i++) {
80-
svfloat32_t row_c_0 = svld1(pg_c_0, &C[i * ldc]);
81-
row_c_0 = svmul_x(pg, beta_vec, row_c_0);
82-
svwrite_hor_za32_f32_m(/*tile*/2, /*slice*/i - svl, pg_c_0, row_c_0);
83-
84-
svfloat32_t row_c_1 = svld1(pg_c_1, &C[i * ldc + svl]);
85-
row_c_1 = svmul_x(pg, beta_vec, row_c_1);
86-
svwrite_hor_za32_f32_m(/*tile*/3, /*slice*/i - svl, pg_c_1, row_c_1);
74+
// Load C to ZA
75+
for (size_t i = 0; i < MIN(svl, block_rows); i++) {
76+
svfloat32_t row_c_0 = svld1(pg_c_0, &C[i * ldc]);
77+
row_c_0 = svmul_x(pg, beta_vec, row_c_0);
78+
svwrite_hor_za32_f32_m(/*tile*/0, /*slice*/i, pg_c_0, row_c_0);
79+
80+
svfloat32_t row_c_1 = svld1(pg_c_1, &C[i * ldc + svl]);
81+
row_c_1 = svmul_x(pg, beta_vec, row_c_1);
82+
svwrite_hor_za32_f32_m(/*tile*/1, /*slice*/i, pg_c_1, row_c_1);
83+
}
84+
for (size_t i = svl; i < block_rows; i++) {
85+
svfloat32_t row_c_0 = svld1(pg_c_0, &C[i * ldc]);
86+
row_c_0 = svmul_x(pg, beta_vec, row_c_0);
87+
svwrite_hor_za32_f32_m(/*tile*/2, /*slice*/i - svl, pg_c_0, row_c_0);
88+
89+
svfloat32_t row_c_1 = svld1(pg_c_1, &C[i * ldc + svl]);
90+
row_c_1 = svmul_x(pg, beta_vec, row_c_1);
91+
svwrite_hor_za32_f32_m(/*tile*/3, /*slice*/i - svl, pg_c_1, row_c_1);
92+
}
8793
}
88-
94+
8995
svfloat32_t alpha_vec = svdup_f32(alpha);
9096
// Iterate through shared dimension (K)
9197
for (size_t k = 0; k < shared_dim; k++) {
@@ -218,12 +224,19 @@ static void ssyrk_direct_sme1_2VLx2VL(uint64_t n, uint64_t k, const float* alpha
218224

219225
void CNAME (BLASLONG N, BLASLONG K, float alpha, float * __restrict A,\
220226
BLASLONG strideA, float beta, float * __restrict C, BLASLONG strideC){
221-
#if !defined(TRANSA)
227+
if (alpha == 0.0f || K == 0) {
228+
if (beta == 1.0f)
229+
return;
230+
ssyrk_direct_sme1_2VLx2VL(N, 0, &alpha, A, &beta, C);
231+
return;
232+
}
233+
234+
#if !defined(TRANSA)
222235
uint64_t n_mod, vl_elms;
223236

224237
vl_elms = sve_cntw();
225238

226-
n_mod = ceil((double)N/(double)vl_elms) * vl_elms;
239+
n_mod = (((uint64_t)N + vl_elms - 1) / vl_elms) * vl_elms;
227240

228241
float *A_mod = (float *) malloc(n_mod*K*sizeof(float));
229242

0 commit comments

Comments
 (0)