Skip to content

Commit e21b7b5

Browse files
akote123akote123
committed
Post sgemm to sve
1 parent f7113bd commit e21b7b5

5 files changed

Lines changed: 1420 additions & 15 deletions

File tree

‎cmake/onnxruntime_mlas.cmake‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -511,6 +511,8 @@ else()
511511
list(APPEND mlas_platform_srcs ${MLAS_SRC_DIR}/sve/mlasi_sve.h)
512512
list(APPEND mlas_platform_srcs ${MLAS_SRC_DIR}/sve/elementwise_sve.cpp)
513513
set_source_files_properties(${MLAS_SRC_DIR}/sve/elementwise_sve.cpp PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+sve+fp16 ")
514+
list(APPEND mlas_platform_srcs ${MLAS_SRC_DIR}/sve/sgemm_sve.cpp)
515+
set_source_files_properties(${MLAS_SRC_DIR}/sve/sgemm_sve.cpp PROPERTIES COMPILE_FLAGS "-march=armv8.2-a+sve -O3 -ffast-math -funroll-loops")
514516
list(APPEND mlas_private_compile_definitions MLAS_USE_SVE)
515517
endif()
516518

‎onnxruntime/core/mlas/lib/mlasi.h‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,7 @@ Module Name:
8181
#endif
8282
#endif
8383

84+
8485
//
8586
// Macro to place variables at a specified alignment.
8687
//
@@ -180,6 +181,7 @@ class MLASCPUIDInfo
180181

181182
// ARM
182183
bool HasArmNeonDot() const { return has_arm_neon_dot_; }
184+
bool HasArmSVE() const { return has_arm_sve_; }
183185

184186
bool HasFp16VectorAcceleration() const { return has_fp16_; }
185187

@@ -206,6 +208,7 @@ class MLASCPUIDInfo
206208

207209
bool has_arm_neon_dot_{false};
208210
bool has_fp16_{false};
211+
bool has_arm_sve_{false};
209212
bool has_arm_neon_i8mm_{false};
210213
bool has_arm_sve_{false};
211214
bool has_arm_sve_i8mm_{false};
@@ -3019,4 +3022,4 @@ MlasPackInt4Elements(uint8_t* Output, UnpackedType ValueLow, UnpackedType ValueH
30193022
{
30203023
static_assert(std::is_same_v<UnpackedType, uint8_t> || std::is_same_v<UnpackedType, int8_t>);
30213024
*Output = static_cast<uint8_t>(((ValueHigh & 0xF) << 4) | (ValueLow & 0xF));
3022-
}
3025+
}

‎onnxruntime/core/mlas/lib/sgemm.cpp‎

Lines changed: 93 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,9 @@ Module Name:
1414
operation (SGEMM).
1515
1616
--*/
17+
#ifdef MLAS_USE_SVE
18+
#include "sve/mlasi_sve.h"
19+
#endif
1720

1821
#include "mlasi.h"
1922

@@ -259,8 +262,15 @@ Return Value:
259262

260263
do {
261264

262-
#if defined(MLAS_NEON_INTRINSICS)
263-
vst4q_f32(D, vld4q_f32(b));
265+
#if defined(MLAS_USE_SVE) && defined(MLAS_NEON_INTRINSICS)
266+
if (MLAS_CPUIDINFO::GetCPUIDInfo().HasArmSve()) {
267+
SVE_LOAD_STORE(D, b);
268+
}
269+
else{
270+
vst4q_f32(D, vld4q_f32(b));
271+
}
272+
#elif defined(MLAS_NEON_INTRINSICS)
273+
vst4q_f32(D, vld4q_f32(b));
264274
#else
265275
MLAS_FLOAT32X4 t0 = MlasLoadFloat32x4(&b[0]);
266276
MLAS_FLOAT32X4 t1 = MlasLoadFloat32x4(&b[4]);
@@ -303,8 +313,14 @@ Return Value:
303313
float* d = D;
304314
const float* b = B;
305315

306-
#if defined(MLAS_NEON_INTRINSICS)
307-
vst4q_f32(d, ZeroFloat32x4x4);
316+
#if defined(MLAS_USE_SVE) && defined(MLAS_NEON_INTRINSICS)
317+
if (MLAS_CPUIDINFO::GetCPUIDInfo().HasArmSve()) {
318+
SVE_ZERO_INITIALIZE(d);
319+
} else {
320+
vst4q_f32(d, ZeroFloat32x4x4);
321+
}
322+
#elif defined(MLAS_NEON_INTRINSICS)
323+
vst4q_f32(d, ZeroFloat32x4x4);
308324
#else
309325
MlasStoreAlignedFloat32x4(d, ZeroFloat32x4);
310326
MlasStoreAlignedFloat32x4(d + 4, ZeroFloat32x4);
@@ -486,6 +502,21 @@ Return Value:
486502
x -= 4;
487503
}
488504

505+
#elif defined(MLAS_USE_SVE)
506+
if(MLAS_CPUIDINFO::GetCPUIDInfo().HasArmSve()) {
507+
SVE_TRANSPOSE(D,b,ldb,x);
508+
}
509+
else
510+
{
511+
while (x >= 4) {
512+
513+
MlasSgemmTransposePackBNx4<16>(&D[0], &b[0], ldb);
514+
515+
D += 16 * 4;
516+
b += 4;
517+
x -= 4;
518+
}
519+
}
489520
#else
490521

491522
while (x >= 4) {
@@ -564,8 +595,15 @@ Return Value:
564595
const float* b = B;
565596

566597
if ((CountY & 8) != 0) {
567-
568-
MlasSgemmTransposePackBNx4<8>(&d[0], &b[0], ldb);
598+
#if defined(MLAS_USE_SVE)
599+
if (MLAS_CPUIDINFO::GetCPUIDInfo().HasArmSve()){
600+
MlasSveTransposePackBNx4<8>(&d[0], &b[0], ldb);}
601+
else{
602+
MlasSgemmTransposePackBNx4<8>(&d[0], &b[0], ldb);
603+
}
604+
#else
605+
MlasSgemmTransposePackBNx4<8>(&d[0], &b[0], ldb);
606+
#endif
569607

570608
d += 8;
571609
b += ldb * 8;
@@ -584,7 +622,15 @@ Return Value:
584622

585623
if ((CountY & 4) != 0) {
586624

587-
MlasSgemmTransposePackBNx4<4>(&d[0], &b[0], ldb);
625+
#if defined(MLAS_USE_SVE)
626+
if (MLAS_CPUIDINFO::GetCPUIDInfo().HasArmSve()){
627+
MlasSveTransposePackBNx4<4>(&d[0], &b[0], ldb);}
628+
else{
629+
MlasSgemmTransposePackBNx4<4>(&d[0], &b[0], ldb);
630+
}
631+
#else
632+
MlasSgemmTransposePackBNx4<4>(&d[0], &b[0], ldb);
633+
#endif
588634

589635
d += 4;
590636
b += ldb * 4;
@@ -631,7 +677,19 @@ Return Value:
631677

632678
if ((CountY & 1) != 0) {
633679

634-
#if defined(MLAS_NEON_INTRINSICS)
680+
#if defined(MLAS_USE_SVE) && defined(MLAS_NEON_INTRINSICS)
681+
if(MLAS_CPUIDINFO::GetCPUIDInfo().HasArmSve())
682+
{
683+
SCATTER_STORE(&d[0],&b[0]);
684+
}
685+
else{
686+
MLAS_FLOAT32X4 t0 = MlasLoadFloat32x4(&b[0]);
687+
688+
MlasStoreLaneFloat32x4<0>(&d[0], t0);
689+
MlasStoreLaneFloat32x4<1>(&d[16], t0);
690+
MlasStoreLaneFloat32x4<2>(&d[32], t0);
691+
MlasStoreLaneFloat32x4<3>(&d[48], t0);}
692+
#elif defined(MLAS_NEON_INTRINSICS)
635693
MLAS_FLOAT32X4 t0 = MlasLoadFloat32x4(&b[0]);
636694

637695
MlasStoreLaneFloat32x4<0>(&d[0], t0);
@@ -945,8 +1003,7 @@ Return Value:
9451003
#endif
9461004

9471005
MLAS_FORCEINLINE
948-
float*
949-
MlasSgemmKernelLoop(
1006+
float* MlasSgemmKernelLoop(
9501007
const float* A,
9511008
const float* B,
9521009
float* C,
@@ -1000,18 +1057,41 @@ Return Value:
10001057
{
10011058
while (CountM > 0) {
10021059

1003-
size_t RowsHandled;
1060+
size_t RowsHandled = 0;
10041061

10051062
#if (defined(MLAS_TARGET_AMD64_IX86) || defined(MLAS_TARGET_POWER) || defined(MLAS_TARGET_S390X) || defined(MLAS_TARGET_LARCH64)) && !defined(FORCE_GENERIC_ALGORITHMS)
10061063
RowsHandled = GetMlasPlatform().GemmFloatKernel(A, B, C, CountK, CountM, CountN, lda, ldc, alpha, ZeroMode);
1064+
10071065
#else
1066+
10081067
if (ZeroMode) {
1068+
1069+
#if defined(MLAS_USE_SVE)
1070+
if (MLAS_CPUIDINFO::GetCPUIDInfo().HasArmSve()) {
1071+
RowsHandled = MlasSgemmKernelZero_sve(A, B, C, CountK, CountM, CountN, lda, ldc, alpha);
1072+
} else {
1073+
RowsHandled = MlasSgemmKernelZero(A, B, C, CountK, CountM, CountN, lda, ldc, alpha);
1074+
}
1075+
#else
10091076
RowsHandled = MlasSgemmKernelZero(A, B, C, CountK, CountM, CountN, lda, ldc, alpha);
1077+
#endif
1078+
10101079
} else {
1080+
1081+
#if defined(MLAS_USE_SVE)
1082+
if (MLAS_CPUIDINFO::GetCPUIDInfo().HasArmSve()) {
1083+
RowsHandled = MlasSgemmKernelAdd_sve(A, B, C, CountK, CountM, CountN, lda, ldc, alpha);
1084+
} else {
1085+
RowsHandled = MlasSgemmKernelAdd(A, B, C, CountK, CountM, CountN, lda, ldc, alpha);
1086+
}
1087+
#else
10111088
RowsHandled = MlasSgemmKernelAdd(A, B, C, CountK, CountM, CountN, lda, ldc, alpha);
1012-
}
10131089
#endif
10141090

1091+
}
1092+
1093+
#endif // platform check
1094+
10151095
C += ldc * RowsHandled;
10161096
A += lda * RowsHandled;
10171097
CountM -= RowsHandled;
@@ -1020,6 +1100,7 @@ Return Value:
10201100
return C;
10211101
}
10221102

1103+
10231104
void
10241105
MlasSgemmOperation(
10251106
CBLAS_TRANSPOSE TransA,

‎onnxruntime/core/mlas/lib/sve/mlasi_sve.h‎

Lines changed: 70 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,14 +20,17 @@ Module Name:
2020
#ifndef __clang__
2121
#pragma GCC push_options
2222
#pragma GCC target("arch=armv8.2-a+sve")
23-
23+
#endif
2424
// Use Clang-specific per-function attribute
2525
#ifdef __clang__
2626
#define MLAS_SVE_TARGET __attribute__((target("arch=armv8.2-a+sve")))
2727
#else
2828
#define MLAS_SVE_TARGET
2929
#endif
3030

31+
32+
#define PACKED_B_BLOCK_WIDTH 16
33+
3134
typedef svfloat32_t MLAS_SVFLOAT32;
3235
typedef svint32_t MLAS_SVINT32;
3336
typedef svuint32_t MLAS_SVUINT32;
@@ -114,6 +117,65 @@ MlasSveLogisticKernel(
114117
size_t N
115118
);
116119

120+
121+
MLAS_SVE_TARGET
122+
MLASCALL
123+
size_t
124+
MlasSgemmKernelAdd_sve(
125+
const float* A,
126+
const float* B,
127+
float* C,
128+
size_t CountK,
129+
size_t CountM,
130+
size_t CountN,
131+
size_t lda,
132+
size_t ldc,
133+
float alpha
134+
);
135+
136+
MLAS_SVE_TARGET
137+
MLASCALL
138+
size_t
139+
MlasSgemmKernelZero_sve(
140+
const float* A,
141+
const float* B,
142+
float* C,
143+
size_t CountK,
144+
size_t CountM,
145+
size_t CountN,
146+
size_t lda,
147+
size_t ldc,
148+
float alpha
149+
);
150+
151+
void MLAS_SVE_TARGET MLASCALL
152+
SCATTER_STORE(float* d, const float* b);
153+
154+
void MLAS_SVE_TARGET MLASCALL
155+
SVE_ZERO_INITIALIZE(float* d);
156+
157+
void MLAS_SVE_TARGET MLASCALL
158+
SVE_LOAD_STORE(float* D, const float* b);
159+
160+
void MLAS_SVE_TARGET MLASCALL
161+
SVE_TRANSPOSE(float*& D, const float*& b, size_t ldb, size_t& x);
162+
163+
template <unsigned N>
164+
void
165+
MlasSveTransposePackBNx4(
166+
float* D,
167+
const float* B,
168+
size_t ldb
169+
);
170+
171+
template <unsigned N>
172+
void
173+
TransposePackBNx8(
174+
float* D,
175+
const float* B,
176+
size_t ldb
177+
);
178+
117179
//MLAS API for SVE intrinsics
118180

119181
MLAS_SVE_TARGET
@@ -323,6 +385,13 @@ MlasSveBroadcastFloat32(const float* Value)
323385
return svld1_f32(svptrue_b32(), Value);
324386
}
325387

388+
MLAS_SVBOOL
389+
MLAS_FORCEINLINE
390+
MlasSveptrue(void)
391+
{
392+
return svptrue_b32();
393+
}
394+
326395
MLAS_SVE_TARGET
327396
MLAS_FORCEINLINE
328397
MLAS_SVFLOAT32
@@ -649,5 +718,4 @@ MlasSveCompareGreaterThan(svbool_t Pred, MLAS_SVFLOAT32 A, MLAS_SVFLOAT32 B)
649718
#pragma GCC pop_options
650719
#endif
651720

652-
#endif
653721

0 commit comments

Comments
 (0)