mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-26 19:52:38 +00:00
POWER10: Optimized SGEMM in MLAS (#8121)
* POWER10: Optimized SGEMM in MLAS This patch introduces new optimized version of SGEMM in MLAS using power10 Matrix-Multiply Assist (MMA) feature introduced in POWER ISA v3.1. This patch makes use of new POWER10 compute instructions for matrix multiplication operation. * Adjust tabs in cmake Changing tabs to spaces as per review comment. * Adjust tabs in new sgemm file Changing tabs to spaces in SgemmKernelPOWER10.cpp. * Reusing functions using common header Co-authored-by: Rajalakshmi Srinivasaraghavan <rajis@linux.ibm.com>
This commit is contained in:
RajalakshmiSR
GitHub
parent
9b75be3d3e
commit
32ceaf4532
7 changed files with 874 additions and 405 deletions
|
|
@ -202,6 +202,29 @@ else()
|
|||
set(mlas_platform_srcs
|
||||
${ONNXRUNTIME_ROOT}/core/mlas/lib/power/SgemmKernelPower.cpp
|
||||
)
|
||||
check_cxx_compiler_flag("-mcpu=power10" HAS_POWER10)
|
||||
if(HAS_POWER10)
|
||||
set(CMAKE_REQUIRED_FLAGS "-mcpu=power10")
|
||||
check_cxx_source_compiles("
|
||||
#include <altivec.h>
|
||||
int main() {
|
||||
__vector_quad acc0;
|
||||
__builtin_mma_xxsetaccz (&acc0);
|
||||
return 0;
|
||||
}"
|
||||
COMPILES_P10
|
||||
)
|
||||
if(COMPILES_P10)
|
||||
set(mlas_platform_srcs_power10
|
||||
${ONNXRUNTIME_ROOT}/core/mlas/lib/power/SgemmKernelPOWER10.cpp
|
||||
)
|
||||
set_source_files_properties(${mlas_platform_srcs_power10} PROPERTIES COMPILE_FLAGS "-O2 -mcpu=power10")
|
||||
set(mlas_platform_srcs
|
||||
${mlas_platform_srcs}
|
||||
${mlas_platform_srcs_power10}
|
||||
)
|
||||
endif()
|
||||
endif()
|
||||
elseif(X86)
|
||||
enable_language(ASM)
|
||||
|
||||
|
|
|
|||
|
|
@ -497,6 +497,7 @@ extern "C" {
|
|||
#endif
|
||||
#elif defined(MLAS_TARGET_POWER)
|
||||
MLAS_GEMM_FLOAT_KERNEL MlasSgemmKernel;
|
||||
MLAS_GEMM_FLOAT_KERNEL MlasSgemmKernelPOWER10;
|
||||
#else
|
||||
MLAS_GEMM_FLOAT_KERNEL MlasSgemmKernelZero;
|
||||
MLAS_GEMM_FLOAT_KERNEL MlasSgemmKernelAdd;
|
||||
|
|
@ -702,7 +703,7 @@ struct MLAS_PLATFORM {
|
|||
|
||||
MLAS_PLATFORM(void);
|
||||
|
||||
#if defined(MLAS_TARGET_AMD64_IX86)
|
||||
#if defined(MLAS_TARGET_AMD64_IX86) || defined(MLAS_TARGET_POWER)
|
||||
MLAS_GEMM_FLOAT_KERNEL* GemmFloatKernel;
|
||||
const MLAS_GEMM_U8X8_DISPATCH* GemmU8S8Dispatch;
|
||||
const MLAS_GEMM_U8X8_DISPATCH* GemmU8U8Dispatch;
|
||||
|
|
|
|||
|
|
@ -17,6 +17,10 @@ Abstract:
|
|||
|
||||
#include "mlasi.h"
|
||||
|
||||
#if defined(MLAS_TARGET_POWER) && defined(__linux__)
|
||||
#include <sys/auxv.h>
|
||||
#endif
|
||||
|
||||
#if defined(MLAS_TARGET_ARM64)
|
||||
#if defined(_WIN32)
|
||||
// N.B. Support building with downlevel versions of the Windows SDK.
|
||||
|
|
@ -368,6 +372,19 @@ Return Value:
|
|||
}
|
||||
|
||||
#endif // MLAS_TARGET_ARM64
|
||||
#if defined(MLAS_TARGET_POWER)
|
||||
this->GemmFloatKernel = MlasSgemmKernel;
|
||||
#if defined(__linux__)
|
||||
#if (defined(__GNUC__) && ((__GNUC__ > 10) || (__GNUC__== 10 && __GNUC_MINOR__ >= 2))) || \
|
||||
(defined(__clang__) && (__clang_major__ >= 12))
|
||||
unsigned long hwcap2 = getauxval(AT_HWCAP2);
|
||||
bool HasP10Instructions = ((hwcap2 & PPC_FEATURE2_MMA) && (hwcap2 & PPC_FEATURE2_ARCH_3_1));
|
||||
if (HasP10Instructions) {
|
||||
this->GemmFloatKernel = MlasSgemmKernelPOWER10;
|
||||
}
|
||||
#endif
|
||||
#endif
|
||||
#endif
|
||||
|
||||
}
|
||||
|
||||
|
|
|
|||
412
onnxruntime/core/mlas/lib/power/SgemmKernelPOWER10.cpp
Normal file
412
onnxruntime/core/mlas/lib/power/SgemmKernelPOWER10.cpp
Normal file
|
|
@ -0,0 +1,412 @@
|
|||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
SgemmKernelPower.cpp
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements the kernels for the single precision matrix/matrix
|
||||
multiply operation (SGEMM).
|
||||
|
||||
--*/
|
||||
|
||||
#include "SgemmKernelpower.h"
|
||||
struct MlasSgemmBroadcastAElementsMMA
|
||||
{
|
||||
template<size_t RowCount, size_t Row>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Iteration(
|
||||
MLAS_FLOAT32X4 ABroadcast[RowCount],
|
||||
const float* A,
|
||||
size_t lda
|
||||
)
|
||||
{
|
||||
ABroadcast[0][Row] = A [Row * lda];
|
||||
}
|
||||
};
|
||||
|
||||
template<size_t RowCount>
|
||||
MLAS_FORCEINLINE
|
||||
void
|
||||
MlasSgemmComputeAElements(
|
||||
MLAS_FLOAT32X4 AElements[RowCount],
|
||||
MLAS_FLOAT32X4 ABroadcast[RowCount]
|
||||
)
|
||||
{
|
||||
__vector float a1,a2;
|
||||
a1 = vec_mergee (AElements[0], AElements[1]);
|
||||
a2 = vec_mergee (AElements[2], AElements[3]);
|
||||
ABroadcast[0] =vec_xxpermdi(a1,a2,0);
|
||||
ABroadcast[2] =vec_xxpermdi(a1,a2,3);
|
||||
a1 = vec_mergeo (AElements[0], AElements[1]);
|
||||
a2 = vec_mergeo (AElements[2], AElements[3]);
|
||||
ABroadcast[1] =vec_xxpermdi(a1,a2,0);
|
||||
ABroadcast[3] =vec_xxpermdi(a1,a2,3);
|
||||
}
|
||||
template<size_t RowCount>
|
||||
MLAS_FORCEINLINE
|
||||
void
|
||||
MlasSgemmComputeBlockMMA(
|
||||
__vector_quad acc[8],
|
||||
MLAS_FLOAT32X4 ABroadcast,
|
||||
MLAS_FLOAT32X4 A2Broadcast,
|
||||
const float* B,
|
||||
size_t CountM
|
||||
)
|
||||
{
|
||||
MLAS_FLOAT32X4 BElements[4];
|
||||
typedef __vector unsigned char vec_t;
|
||||
|
||||
BElements[0] = MlasLoadFloat32x4(B);
|
||||
BElements[1] = MlasLoadFloat32x4(B + 4);
|
||||
BElements[2] = MlasLoadFloat32x4(B + 8);
|
||||
BElements[3] = MlasLoadFloat32x4(B + 12);
|
||||
__builtin_mma_xvf32gerpp (&acc[0], (vec_t) ABroadcast, (vec_t )BElements[0]);
|
||||
__builtin_mma_xvf32gerpp (&acc[1], (vec_t) ABroadcast, (vec_t )BElements[1]);
|
||||
__builtin_mma_xvf32gerpp (&acc[2], (vec_t) ABroadcast, (vec_t )BElements[2]);
|
||||
__builtin_mma_xvf32gerpp (&acc[3], (vec_t) ABroadcast, (vec_t )BElements[3]);
|
||||
if (CountM == 8) {
|
||||
__builtin_mma_xvf32gerpp (&acc[4], (vec_t) A2Broadcast, (vec_t )BElements[0]);
|
||||
__builtin_mma_xvf32gerpp (&acc[5], (vec_t) A2Broadcast, (vec_t )BElements[1]);
|
||||
__builtin_mma_xvf32gerpp (&acc[6], (vec_t) A2Broadcast, (vec_t )BElements[2]);
|
||||
__builtin_mma_xvf32gerpp (&acc[7], (vec_t) A2Broadcast, (vec_t )BElements[3]);
|
||||
}
|
||||
}
|
||||
template<size_t VectorCount>
|
||||
struct MlasSgemmStoreVectorMMA
|
||||
{
|
||||
template<size_t RowCount, size_t Row>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Iteration(
|
||||
MLAS_FLOAT32X4 Result[4],
|
||||
float* C,
|
||||
size_t ldc,
|
||||
MLAS_FLOAT32X4 AlphaBroadcast,
|
||||
bool ZeroMode
|
||||
)
|
||||
{
|
||||
MLAS_FLOAT32X4 *rowC;
|
||||
if (ZeroMode) {
|
||||
rowC = (MLAS_FLOAT32X4 *) &C[Row * ldc + VectorCount];
|
||||
rowC[0] = Result[Row] * AlphaBroadcast;
|
||||
} else {
|
||||
rowC = (MLAS_FLOAT32X4 *) &C[Row * ldc + VectorCount];
|
||||
rowC[0] += Result[Row] * AlphaBroadcast;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
struct MlasSgemmMultiplyAlphaTrailingMMA
|
||||
{
|
||||
template<size_t RowCount, size_t Row>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Iteration(
|
||||
MLAS_FLOAT32X4 Accumulators[RowCount],
|
||||
MLAS_FLOAT32X4 AlphaBroadcast
|
||||
)
|
||||
{
|
||||
Accumulators[Row] = MlasMultiplyFloat32x4(Accumulators[Row], AlphaBroadcast);
|
||||
}
|
||||
};
|
||||
template<unsigned Lane>
|
||||
struct MlasSgemmStoreScalarMMA
|
||||
{
|
||||
template<size_t RowCount, size_t Row>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Iteration(
|
||||
MLAS_FLOAT32X4 Accumulators[RowCount],
|
||||
float* C,
|
||||
size_t ldc,
|
||||
bool ZeroMode
|
||||
)
|
||||
{
|
||||
float* c = C + Row * ldc + Lane;
|
||||
float Value = Accumulators[Row][Lane];
|
||||
if (!ZeroMode) {
|
||||
Value += *c;
|
||||
}
|
||||
|
||||
*c = Value;
|
||||
}
|
||||
};
|
||||
|
||||
template<size_t RowCount>
|
||||
MLAS_FORCEINLINE
|
||||
size_t
|
||||
MlasSgemmMMAProcessCount(
|
||||
const float* A,
|
||||
const float* B,
|
||||
float* C,
|
||||
size_t CountM,
|
||||
size_t CountK,
|
||||
size_t CountN,
|
||||
size_t lda,
|
||||
size_t ldc,
|
||||
MLAS_FLOAT32X4 AlphaBroadcast,
|
||||
bool ZeroMode
|
||||
)
|
||||
{
|
||||
do {
|
||||
|
||||
const float* a = A;
|
||||
size_t k = CountK;
|
||||
|
||||
MLAS_FLOAT32X4 Accumulators[2][RowCount] = {{ 0 }};
|
||||
MLAS_FLOAT32X4 Result[RowCount];
|
||||
MLAS_FLOAT32X4 AElements[RowCount];
|
||||
MLAS_FLOAT32X4 ABroadcast[RowCount] = { 0 };
|
||||
MLAS_FLOAT32X4 A2Broadcast[RowCount] = { 0 };
|
||||
__vector_quad acc[8];
|
||||
|
||||
//
|
||||
// Clear the block accumulators.
|
||||
//
|
||||
__builtin_mma_xxsetaccz(&acc[0]);
|
||||
__builtin_mma_xxsetaccz(&acc[1]);
|
||||
__builtin_mma_xxsetaccz(&acc[2]);
|
||||
__builtin_mma_xxsetaccz(&acc[3]);
|
||||
__builtin_mma_xxsetaccz(&acc[4]);
|
||||
__builtin_mma_xxsetaccz(&acc[5]);
|
||||
__builtin_mma_xxsetaccz(&acc[6]);
|
||||
__builtin_mma_xxsetaccz(&acc[7]);
|
||||
|
||||
//
|
||||
// Compute the output block.
|
||||
//
|
||||
while (k >= 4) {
|
||||
|
||||
MlasLoopUnroll<RowCount, MlasSgemmLoadAElements>()(AElements, a, lda);
|
||||
MlasSgemmComputeAElements<RowCount>(AElements, ABroadcast);
|
||||
if (CountM == 8) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmLoadAElements>()(AElements, a + ( lda * 4), lda);
|
||||
MlasSgemmComputeAElements<RowCount>(AElements, A2Broadcast);
|
||||
}
|
||||
MlasSgemmComputeBlockMMA<RowCount>(&acc[0], ABroadcast[0], A2Broadcast[0], B, CountM);
|
||||
MlasSgemmComputeBlockMMA<RowCount>(&acc[0], ABroadcast[1], A2Broadcast[1], B+16, CountM);
|
||||
MlasSgemmComputeBlockMMA<RowCount>(&acc[0], ABroadcast[2], A2Broadcast[2], B+32, CountM);
|
||||
MlasSgemmComputeBlockMMA<RowCount>(&acc[0], ABroadcast[3], A2Broadcast[3], B+48, CountM);
|
||||
B += 16 * 4;
|
||||
a += 4;
|
||||
k -= 4;
|
||||
}
|
||||
|
||||
while (k > 0) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmBroadcastAElementsMMA>()(ABroadcast, a, lda);
|
||||
if (CountM == 8) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmBroadcastAElementsMMA>()(A2Broadcast, a + (lda * 4), lda);
|
||||
}
|
||||
MlasSgemmComputeBlockMMA<RowCount>(&acc[0], ABroadcast[0], A2Broadcast[0], B, CountM);
|
||||
a += 1;
|
||||
B += 16;
|
||||
k -= 1;
|
||||
}
|
||||
if (CountN >= 16) {
|
||||
|
||||
//
|
||||
// Store the entire output block.
|
||||
//
|
||||
__builtin_mma_disassemble_acc ((void *)Result, &acc[0]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<0>>()(Result, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
__builtin_mma_disassemble_acc ((void *)Result, &acc[1]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<4>>()(Result, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
__builtin_mma_disassemble_acc ((void *)Result, &acc[2]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<8>>()(Result, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
__builtin_mma_disassemble_acc ((void *)Result, &acc[3]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<12>>()(Result, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
if (CountM == 8) {
|
||||
__builtin_mma_disassemble_acc ((void *)Result, &acc[4]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<0>>()(Result, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode);
|
||||
__builtin_mma_disassemble_acc ((void *)Result, &acc[5]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<4>>()(Result, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode);
|
||||
__builtin_mma_disassemble_acc ((void *)Result, &acc[6]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<8>>()(Result, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode);
|
||||
__builtin_mma_disassemble_acc ((void *)Result, &acc[7]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<12>>()(Result, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode);
|
||||
}
|
||||
} else {
|
||||
|
||||
//
|
||||
// Store the partial output block.
|
||||
//
|
||||
|
||||
if (CountN >= 12) {
|
||||
__builtin_mma_disassemble_acc ((void *)Result, &acc[0]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<0>>()(Result, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
__builtin_mma_disassemble_acc ((void *)Result, &acc[1]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<4>>()(Result, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
__builtin_mma_disassemble_acc ((void *)Result, &acc[2]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<8>>()(Result, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
if (CountM == 8) {
|
||||
__builtin_mma_disassemble_acc ((void *)Result, &acc[4]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<0>>()(Result, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode);
|
||||
__builtin_mma_disassemble_acc ((void *)Result, &acc[5]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<4>>()(Result, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode);
|
||||
__builtin_mma_disassemble_acc ((void *)Result, &acc[6]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<8>>()(Result, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode);
|
||||
if (CountN - 12 > 0) {
|
||||
__builtin_mma_disassemble_acc ((void *)Accumulators[1], &acc[7]);
|
||||
}
|
||||
}
|
||||
if (CountN - 12 > 0) {
|
||||
__builtin_mma_disassemble_acc ((void *)Accumulators[0], &acc[3]);
|
||||
}
|
||||
} else if (CountN >= 8) {
|
||||
__builtin_mma_disassemble_acc ((void *)Result, &acc[0]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<0>>()(Result, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
__builtin_mma_disassemble_acc ((void *)Result, &acc[1]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<4>>()(Result, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
if (CountM == 8) {
|
||||
__builtin_mma_disassemble_acc ((void *)Result, &acc[4]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<0>>()(Result, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode);
|
||||
__builtin_mma_disassemble_acc ((void *)Result, &acc[5]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<4>>()(Result, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode);
|
||||
if (CountN - 8 > 0) {
|
||||
__builtin_mma_disassemble_acc ((void *)Accumulators[1], &acc[6]);
|
||||
}
|
||||
}
|
||||
if (CountN - 8 > 0) {
|
||||
__builtin_mma_disassemble_acc ((void *)Accumulators[0], &acc[2]);
|
||||
}
|
||||
} else if (CountN >= 4) {
|
||||
__builtin_mma_disassemble_acc ((void *)Result, &acc[0]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<0>>()(Result, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
if (CountM == 8) {
|
||||
__builtin_mma_disassemble_acc ((void *)Result, &acc[4]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<0>>()(Result, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode);
|
||||
if (CountN - 4 > 0) {
|
||||
__builtin_mma_disassemble_acc ((void *)Accumulators[1], &acc[5]);
|
||||
}
|
||||
}
|
||||
if (CountN - 4 > 0) {
|
||||
__builtin_mma_disassemble_acc ((void *)Accumulators[0], &acc[1]);
|
||||
}
|
||||
} else {
|
||||
__builtin_mma_disassemble_acc ((void *)Accumulators[0], &acc[0]);
|
||||
if (CountM == 8) {
|
||||
__builtin_mma_disassemble_acc ((void *)Accumulators[1], &acc[4]);
|
||||
}
|
||||
}
|
||||
|
||||
//
|
||||
// Store the remaining unaligned columns.
|
||||
//
|
||||
|
||||
C += (CountN & ~3);
|
||||
CountN &= 3;
|
||||
|
||||
if (CountN > 0) {
|
||||
|
||||
MlasLoopUnroll<RowCount, MlasSgemmMultiplyAlphaTrailingMMA>()(Accumulators[0], AlphaBroadcast);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreScalarMMA<0>>()(Accumulators[0], C, ldc, ZeroMode);
|
||||
if (CountM == 8) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmMultiplyAlphaTrailingMMA>()(Accumulators[1], AlphaBroadcast);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreScalarMMA<0>>()(Accumulators[1], C + (ldc*4), ldc, ZeroMode);
|
||||
}
|
||||
if (CountN >= 2) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreScalarMMA<1>>()(Accumulators[0], C, ldc, ZeroMode);
|
||||
if (CountM == 8) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreScalarMMA<1>>()(Accumulators[1], C + (ldc*4), ldc, ZeroMode);
|
||||
}
|
||||
}
|
||||
if (CountN >= 3) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreScalarMMA<2>>()(Accumulators[0], C, ldc, ZeroMode);
|
||||
if (CountM == 8) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreScalarMMA<2>>()(Accumulators[1], C + (ldc*4), ldc, ZeroMode);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
break;
|
||||
}
|
||||
|
||||
C += 16;
|
||||
CountN -= 16;
|
||||
|
||||
} while (CountN > 0);
|
||||
|
||||
return CountM;
|
||||
}
|
||||
|
||||
size_t
|
||||
MLASCALL
|
||||
MlasSgemmKernelPOWER10(
|
||||
const float* A,
|
||||
const float* B,
|
||||
float* C,
|
||||
size_t CountK,
|
||||
size_t CountM,
|
||||
size_t CountN,
|
||||
size_t lda,
|
||||
size_t ldc,
|
||||
float alpha,
|
||||
bool ZeroMode
|
||||
)
|
||||
/*++
|
||||
|
||||
Routine Description:
|
||||
|
||||
This routine is an inner kernel to compute matrix multiplication for a
|
||||
set of rows.
|
||||
|
||||
Arguments:
|
||||
|
||||
A - Supplies the address of matrix A.
|
||||
|
||||
B - Supplies the address of matrix B. The matrix data has been packed using
|
||||
MlasSgemmCopyPackB or MlasSgemmTransposePackB.
|
||||
|
||||
C - Supplies the address of matrix C.
|
||||
|
||||
CountK - Supplies the number of columns from matrix A and the number of rows
|
||||
from matrix B to iterate over.
|
||||
|
||||
CountM - Supplies the maximum number of rows that can be processed for
|
||||
matrix A and matrix C. The actual number of rows handled for this
|
||||
invocation depends on the kernel implementation.
|
||||
|
||||
CountN - Supplies the number of columns from matrix B and matrix C to
|
||||
iterate over.
|
||||
|
||||
lda - Supplies the first dimension of matrix A.
|
||||
|
||||
ldc - Supplies the first dimension of matrix C.
|
||||
|
||||
alpha - Supplies the scalar multiplier (see SGEMM definition).
|
||||
|
||||
ZeroMode - Supplies true if the output matrix must be zero initialized,
|
||||
else false if the output matrix is accumulated into.
|
||||
|
||||
Return Value:
|
||||
|
||||
Returns the number of rows handled.
|
||||
|
||||
--*/
|
||||
{
|
||||
size_t RowsHandled;
|
||||
MLAS_FLOAT32X4 AlphaBroadcast = MlasBroadcastFloat32x4(alpha);
|
||||
|
||||
if (CountM >= 8) {
|
||||
RowsHandled = MlasSgemmMMAProcessCount<4>(A, B, C, 8 ,CountK, CountN, lda, ldc, AlphaBroadcast, ZeroMode);
|
||||
} else if (CountM >= 4) {
|
||||
RowsHandled = MlasSgemmMMAProcessCount<4>(A, B, C, 4, CountK, CountN, lda, ldc, AlphaBroadcast, ZeroMode);
|
||||
} else if (CountM >= 2) {
|
||||
RowsHandled = MlasSgemmProcessCount<2>(A, B, C, CountK, CountN, lda, ldc, AlphaBroadcast, ZeroMode);
|
||||
} else {
|
||||
RowsHandled = MlasSgemmProcessCount<1>(A, B, C, CountK, CountN, lda, ldc, AlphaBroadcast, ZeroMode);
|
||||
}
|
||||
|
||||
return RowsHandled;
|
||||
}
|
||||
|
|
@ -14,407 +14,7 @@ Abstract:
|
|||
multiply operation (SGEMM).
|
||||
|
||||
--*/
|
||||
|
||||
#include "mlasi.h"
|
||||
|
||||
//
|
||||
// Templates to ensure that a loop is unrolled.
|
||||
//
|
||||
|
||||
template<size_t Count, size_t Index>
|
||||
struct MlasLoopUnrollStep
|
||||
{
|
||||
template<typename IterationType, typename... IterationArgs>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Step(
|
||||
IterationArgs&&... Arguments
|
||||
)
|
||||
{
|
||||
IterationType::template Iteration<Count, Index>(Arguments...);
|
||||
MlasLoopUnrollStep<Count, Index + 1>::template Step<IterationType>(Arguments...);
|
||||
}
|
||||
};
|
||||
|
||||
template<size_t Count>
|
||||
struct MlasLoopUnrollStep<Count, Count>
|
||||
{
|
||||
template<typename IterationType, typename... IterationArgs>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Step(
|
||||
IterationArgs&&...
|
||||
)
|
||||
{
|
||||
// Terminate the loop.
|
||||
}
|
||||
};
|
||||
|
||||
template<size_t Count, typename IteratorType>
|
||||
struct MlasLoopUnroll
|
||||
{
|
||||
template<typename... IterationArgs>
|
||||
MLAS_FORCEINLINE
|
||||
void
|
||||
operator()(
|
||||
IterationArgs&&... Arguments
|
||||
)
|
||||
{
|
||||
MlasLoopUnrollStep<Count, 0>::template Step<IteratorType>(Arguments...);
|
||||
}
|
||||
};
|
||||
|
||||
//
|
||||
// Templates used with loop unrolling to perform an action on one row of the
|
||||
// output.
|
||||
//
|
||||
|
||||
struct MlasSgemmZeroAccumulators
|
||||
{
|
||||
template<size_t RowCount, size_t Row>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Iteration(
|
||||
MLAS_FLOAT32X4 Accumulators[RowCount][4]
|
||||
)
|
||||
{
|
||||
Accumulators[Row][0] = MlasZeroFloat32x4();
|
||||
Accumulators[Row][1] = MlasZeroFloat32x4();
|
||||
Accumulators[Row][2] = MlasZeroFloat32x4();
|
||||
Accumulators[Row][3] = MlasZeroFloat32x4();
|
||||
}
|
||||
};
|
||||
|
||||
struct MlasSgemmLoadAElements
|
||||
{
|
||||
template<size_t RowCount, size_t Row>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Iteration(
|
||||
MLAS_FLOAT32X4 AElements[RowCount],
|
||||
const float* A,
|
||||
size_t lda
|
||||
)
|
||||
{
|
||||
AElements[Row] = MlasLoadFloat32x4(A + Row * lda);
|
||||
}
|
||||
};
|
||||
|
||||
struct MlasSgemmBroadcastAElements
|
||||
{
|
||||
template<size_t RowCount, size_t Row>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Iteration(
|
||||
MLAS_FLOAT32X4 ABroadcast[RowCount],
|
||||
const float* A,
|
||||
size_t lda
|
||||
)
|
||||
{
|
||||
ABroadcast[Row] = MlasBroadcastFloat32x4(A + Row * lda);
|
||||
}
|
||||
};
|
||||
|
||||
template<unsigned Lane>
|
||||
struct MlasSgemmSplatAElements
|
||||
{
|
||||
template<size_t RowCount, size_t Row>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Iteration(
|
||||
MLAS_FLOAT32X4 AElements[RowCount],
|
||||
MLAS_FLOAT32X4 ABroadcast[RowCount]
|
||||
)
|
||||
{
|
||||
ABroadcast[Row] = vec_splat(AElements[Row], Lane);
|
||||
}
|
||||
};
|
||||
|
||||
struct MlasSgemmMultiplyAddRow
|
||||
{
|
||||
template<size_t RowCount, size_t Row>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Iteration(
|
||||
MLAS_FLOAT32X4 Accumulators[RowCount][4],
|
||||
MLAS_FLOAT32X4 ABroadcast[RowCount],
|
||||
MLAS_FLOAT32X4 BElements[4]
|
||||
)
|
||||
{
|
||||
Accumulators[Row][0] = MlasMultiplyAddFloat32x4(ABroadcast[Row], BElements[0], Accumulators[Row][0]);
|
||||
Accumulators[Row][1] = MlasMultiplyAddFloat32x4(ABroadcast[Row], BElements[1], Accumulators[Row][1]);
|
||||
Accumulators[Row][2] = MlasMultiplyAddFloat32x4(ABroadcast[Row], BElements[2], Accumulators[Row][2]);
|
||||
Accumulators[Row][3] = MlasMultiplyAddFloat32x4(ABroadcast[Row], BElements[3], Accumulators[Row][3]);
|
||||
}
|
||||
};
|
||||
|
||||
template<size_t RowCount>
|
||||
MLAS_FORCEINLINE
|
||||
void
|
||||
MlasSgemmComputeBlock(
|
||||
MLAS_FLOAT32X4 Accumulators[RowCount][4],
|
||||
MLAS_FLOAT32X4 ABroadcast[RowCount],
|
||||
const float* B
|
||||
)
|
||||
{
|
||||
MLAS_FLOAT32X4 BElements[4];
|
||||
|
||||
BElements[0] = MlasLoadFloat32x4(B);
|
||||
BElements[1] = MlasLoadFloat32x4(B + 4);
|
||||
BElements[2] = MlasLoadFloat32x4(B + 8);
|
||||
BElements[3] = MlasLoadFloat32x4(B + 12);
|
||||
|
||||
MlasLoopUnroll<RowCount, MlasSgemmMultiplyAddRow>()(Accumulators, ABroadcast, BElements);
|
||||
}
|
||||
|
||||
struct MlasSgemmMultiplyAlphaRow
|
||||
{
|
||||
template<size_t Count, size_t Index>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Iteration(
|
||||
MLAS_FLOAT32X4 Accumulators[4],
|
||||
MLAS_FLOAT32X4 AlphaBroadcast
|
||||
)
|
||||
{
|
||||
Accumulators[Index] = MlasMultiplyFloat32x4(Accumulators[Index], AlphaBroadcast);
|
||||
}
|
||||
};
|
||||
|
||||
struct MlasSgemmMultiplyAlphaAddRow
|
||||
{
|
||||
template<size_t Count, size_t Index>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Iteration(
|
||||
MLAS_FLOAT32X4 Accumulators[4],
|
||||
MLAS_FLOAT32X4 AlphaBroadcast,
|
||||
const float* C
|
||||
)
|
||||
{
|
||||
Accumulators[Index] = MlasMultiplyAddFloat32x4(Accumulators[Index],
|
||||
AlphaBroadcast, MlasLoadFloat32x4(C + Index * 4));
|
||||
}
|
||||
};
|
||||
|
||||
struct MlasSgemmStoreRow
|
||||
{
|
||||
template<size_t Count, size_t Index>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Iteration(
|
||||
MLAS_FLOAT32X4 Accumulators[4],
|
||||
float* C
|
||||
)
|
||||
{
|
||||
MlasStoreFloat32x4(C + Index * 4, Accumulators[Index]);
|
||||
}
|
||||
};
|
||||
|
||||
template<size_t VectorCount>
|
||||
struct MlasSgemmStoreVector
|
||||
{
|
||||
template<size_t RowCount, size_t Row>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Iteration(
|
||||
MLAS_FLOAT32X4 Accumulators[RowCount][4],
|
||||
float* C,
|
||||
size_t ldc,
|
||||
MLAS_FLOAT32X4 AlphaBroadcast,
|
||||
bool ZeroMode
|
||||
)
|
||||
{
|
||||
float* c = C + Row * ldc;
|
||||
|
||||
if (ZeroMode) {
|
||||
MlasLoopUnroll<VectorCount, MlasSgemmMultiplyAlphaRow>()(Accumulators[Row], AlphaBroadcast);
|
||||
} else {
|
||||
MlasLoopUnroll<VectorCount, MlasSgemmMultiplyAlphaAddRow>()(Accumulators[Row], AlphaBroadcast, c);
|
||||
}
|
||||
|
||||
MlasLoopUnroll<VectorCount, MlasSgemmStoreRow>()(Accumulators[Row], c);
|
||||
|
||||
//
|
||||
// Shift down any unaligned elements to the bottom for further processing.
|
||||
//
|
||||
|
||||
if (VectorCount < 4) {
|
||||
Accumulators[Row][0] = Accumulators[Row][VectorCount];
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
struct MlasSgemmMultiplyAlphaTrailing
|
||||
{
|
||||
template<size_t RowCount, size_t Row>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Iteration(
|
||||
MLAS_FLOAT32X4 Accumulators[RowCount][4],
|
||||
MLAS_FLOAT32X4 AlphaBroadcast
|
||||
)
|
||||
{
|
||||
Accumulators[Row][0] = MlasMultiplyFloat32x4(Accumulators[Row][0], AlphaBroadcast);
|
||||
}
|
||||
};
|
||||
|
||||
template<unsigned Lane>
|
||||
struct MlasSgemmStoreScalar
|
||||
{
|
||||
template<size_t RowCount, size_t Row>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Iteration(
|
||||
MLAS_FLOAT32X4 Accumulators[RowCount][4],
|
||||
float* C,
|
||||
size_t ldc,
|
||||
bool ZeroMode
|
||||
)
|
||||
{
|
||||
float* c = C + Row * ldc + Lane;
|
||||
float Value = MlasExtractLaneFloat32x4<Lane>(Accumulators[Row][0]);
|
||||
|
||||
if (!ZeroMode) {
|
||||
Value += *c;
|
||||
}
|
||||
|
||||
*c = Value;
|
||||
}
|
||||
};
|
||||
|
||||
template<size_t RowCount>
|
||||
MLAS_FORCEINLINE
|
||||
size_t
|
||||
MlasSgemmProcessCount(
|
||||
const float* A,
|
||||
const float* B,
|
||||
float* C,
|
||||
size_t CountK,
|
||||
size_t CountN,
|
||||
size_t lda,
|
||||
size_t ldc,
|
||||
MLAS_FLOAT32X4 AlphaBroadcast,
|
||||
bool ZeroMode
|
||||
)
|
||||
{
|
||||
do {
|
||||
|
||||
const float* a = A;
|
||||
size_t k = CountK;
|
||||
|
||||
MLAS_FLOAT32X4 Accumulators[RowCount][4];
|
||||
MLAS_FLOAT32X4 AElements[RowCount];
|
||||
MLAS_FLOAT32X4 ABroadcast[RowCount];
|
||||
|
||||
//
|
||||
// Clear the block accumulators.
|
||||
//
|
||||
|
||||
MlasLoopUnroll<RowCount, MlasSgemmZeroAccumulators>()(Accumulators);
|
||||
|
||||
//
|
||||
// Compute the output block.
|
||||
//
|
||||
|
||||
while (k >= 4) {
|
||||
|
||||
MlasLoopUnroll<RowCount, MlasSgemmLoadAElements>()(AElements, a, lda);
|
||||
|
||||
MlasLoopUnroll<RowCount, MlasSgemmSplatAElements<0>>()(AElements, ABroadcast);
|
||||
MlasSgemmComputeBlock<RowCount>(Accumulators, ABroadcast, B);
|
||||
|
||||
MlasLoopUnroll<RowCount, MlasSgemmSplatAElements<1>>()(AElements, ABroadcast);
|
||||
MlasSgemmComputeBlock<RowCount>(Accumulators, ABroadcast, B + 16);
|
||||
|
||||
MlasLoopUnroll<RowCount, MlasSgemmSplatAElements<2>>()(AElements, ABroadcast);
|
||||
MlasSgemmComputeBlock<RowCount>(Accumulators, ABroadcast, B + 32);
|
||||
|
||||
MlasLoopUnroll<RowCount, MlasSgemmSplatAElements<3>>()(AElements, ABroadcast);
|
||||
MlasSgemmComputeBlock<RowCount>(Accumulators, ABroadcast, B + 48);
|
||||
|
||||
a += 4;
|
||||
B += 16 * 4;
|
||||
k -= 4;
|
||||
}
|
||||
|
||||
while (k > 0) {
|
||||
|
||||
MlasLoopUnroll<RowCount, MlasSgemmBroadcastAElements>()(ABroadcast, a, lda);
|
||||
MlasSgemmComputeBlock<RowCount>(Accumulators, ABroadcast, B);
|
||||
|
||||
a += 1;
|
||||
B += 16;
|
||||
k -= 1;
|
||||
}
|
||||
|
||||
if (CountN >= 16) {
|
||||
|
||||
//
|
||||
// Store the entire output block.
|
||||
//
|
||||
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVector<4>>()(Accumulators, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
|
||||
} else {
|
||||
|
||||
//
|
||||
// Store the partial output block.
|
||||
//
|
||||
|
||||
if (CountN >= 12) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVector<3>>()(Accumulators, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
} else if (CountN >= 8) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVector<2>>()(Accumulators, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
} else if (CountN >= 4) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVector<1>>()(Accumulators, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
}
|
||||
|
||||
//
|
||||
// Store the remaining unaligned columns.
|
||||
//
|
||||
|
||||
C += (CountN & ~3);
|
||||
CountN &= 3;
|
||||
|
||||
if (CountN > 0) {
|
||||
|
||||
MlasLoopUnroll<RowCount, MlasSgemmMultiplyAlphaTrailing>()(Accumulators, AlphaBroadcast);
|
||||
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreScalar<0>>()(Accumulators, C, ldc, ZeroMode);
|
||||
|
||||
if (CountN >= 2) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreScalar<1>>()(Accumulators, C, ldc, ZeroMode);
|
||||
}
|
||||
|
||||
if (CountN >= 3) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreScalar<2>>()(Accumulators, C, ldc, ZeroMode);
|
||||
}
|
||||
}
|
||||
|
||||
break;
|
||||
}
|
||||
|
||||
C += 16;
|
||||
CountN -= 16;
|
||||
|
||||
} while (CountN > 0);
|
||||
|
||||
return RowCount;
|
||||
}
|
||||
#include "SgemmKernelpower.h"
|
||||
|
||||
size_t
|
||||
MLASCALL
|
||||
|
|
|
|||
418
onnxruntime/core/mlas/lib/power/SgemmKernelpower.h
Normal file
418
onnxruntime/core/mlas/lib/power/SgemmKernelpower.h
Normal file
|
|
@ -0,0 +1,418 @@
|
|||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
SgemmKernelPower.cpp
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements the kernels for the single precision matrix/matrix
|
||||
multiply operation (SGEMM).
|
||||
|
||||
--*/
|
||||
|
||||
#include "mlasi.h"
|
||||
|
||||
//
|
||||
// Templates to ensure that a loop is unrolled.
|
||||
//
|
||||
|
||||
template<size_t Count, size_t Index>
|
||||
struct MlasLoopUnrollStep
|
||||
{
|
||||
template<typename IterationType, typename... IterationArgs>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Step(
|
||||
IterationArgs&&... Arguments
|
||||
)
|
||||
{
|
||||
IterationType::template Iteration<Count, Index>(Arguments...);
|
||||
MlasLoopUnrollStep<Count, Index + 1>::template Step<IterationType>(Arguments...);
|
||||
}
|
||||
};
|
||||
|
||||
template<size_t Count>
|
||||
struct MlasLoopUnrollStep<Count, Count>
|
||||
{
|
||||
template<typename IterationType, typename... IterationArgs>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Step(
|
||||
IterationArgs&&...
|
||||
)
|
||||
{
|
||||
// Terminate the loop.
|
||||
}
|
||||
};
|
||||
|
||||
template<size_t Count, typename IteratorType>
|
||||
struct MlasLoopUnroll
|
||||
{
|
||||
template<typename... IterationArgs>
|
||||
MLAS_FORCEINLINE
|
||||
void
|
||||
operator()(
|
||||
IterationArgs&&... Arguments
|
||||
)
|
||||
{
|
||||
MlasLoopUnrollStep<Count, 0>::template Step<IteratorType>(Arguments...);
|
||||
}
|
||||
};
|
||||
|
||||
//
|
||||
// Templates used with loop unrolling to perform an action on one row of the
|
||||
// output.
|
||||
//
|
||||
|
||||
struct MlasSgemmZeroAccumulators
|
||||
{
|
||||
template<size_t RowCount, size_t Row>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Iteration(
|
||||
MLAS_FLOAT32X4 Accumulators[RowCount][4]
|
||||
)
|
||||
{
|
||||
Accumulators[Row][0] = MlasZeroFloat32x4();
|
||||
Accumulators[Row][1] = MlasZeroFloat32x4();
|
||||
Accumulators[Row][2] = MlasZeroFloat32x4();
|
||||
Accumulators[Row][3] = MlasZeroFloat32x4();
|
||||
}
|
||||
};
|
||||
|
||||
struct MlasSgemmLoadAElements
|
||||
{
|
||||
template<size_t RowCount, size_t Row>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Iteration(
|
||||
MLAS_FLOAT32X4 AElements[RowCount],
|
||||
const float* A,
|
||||
size_t lda
|
||||
)
|
||||
{
|
||||
AElements[Row] = MlasLoadFloat32x4(A + Row * lda);
|
||||
}
|
||||
};
|
||||
|
||||
struct MlasSgemmBroadcastAElements
|
||||
{
|
||||
template<size_t RowCount, size_t Row>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Iteration(
|
||||
MLAS_FLOAT32X4 ABroadcast[RowCount],
|
||||
const float* A,
|
||||
size_t lda
|
||||
)
|
||||
{
|
||||
ABroadcast[Row] = MlasBroadcastFloat32x4(A + Row * lda);
|
||||
}
|
||||
};
|
||||
|
||||
template<unsigned Lane>
|
||||
struct MlasSgemmSplatAElements
|
||||
{
|
||||
template<size_t RowCount, size_t Row>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Iteration(
|
||||
MLAS_FLOAT32X4 AElements[RowCount],
|
||||
MLAS_FLOAT32X4 ABroadcast[RowCount]
|
||||
)
|
||||
{
|
||||
ABroadcast[Row] = vec_splat(AElements[Row], Lane);
|
||||
}
|
||||
};
|
||||
|
||||
struct MlasSgemmMultiplyAddRow
|
||||
{
|
||||
template<size_t RowCount, size_t Row>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Iteration(
|
||||
MLAS_FLOAT32X4 Accumulators[RowCount][4],
|
||||
MLAS_FLOAT32X4 ABroadcast[RowCount],
|
||||
MLAS_FLOAT32X4 BElements[4]
|
||||
)
|
||||
{
|
||||
Accumulators[Row][0] = MlasMultiplyAddFloat32x4(ABroadcast[Row], BElements[0], Accumulators[Row][0]);
|
||||
Accumulators[Row][1] = MlasMultiplyAddFloat32x4(ABroadcast[Row], BElements[1], Accumulators[Row][1]);
|
||||
Accumulators[Row][2] = MlasMultiplyAddFloat32x4(ABroadcast[Row], BElements[2], Accumulators[Row][2]);
|
||||
Accumulators[Row][3] = MlasMultiplyAddFloat32x4(ABroadcast[Row], BElements[3], Accumulators[Row][3]);
|
||||
}
|
||||
};
|
||||
|
||||
template<size_t RowCount>
|
||||
MLAS_FORCEINLINE
|
||||
void
|
||||
MlasSgemmComputeBlock(
|
||||
MLAS_FLOAT32X4 Accumulators[RowCount][4],
|
||||
MLAS_FLOAT32X4 ABroadcast[RowCount],
|
||||
const float* B
|
||||
)
|
||||
{
|
||||
MLAS_FLOAT32X4 BElements[4];
|
||||
|
||||
BElements[0] = MlasLoadFloat32x4(B);
|
||||
BElements[1] = MlasLoadFloat32x4(B + 4);
|
||||
BElements[2] = MlasLoadFloat32x4(B + 8);
|
||||
BElements[3] = MlasLoadFloat32x4(B + 12);
|
||||
|
||||
MlasLoopUnroll<RowCount, MlasSgemmMultiplyAddRow>()(Accumulators, ABroadcast, BElements);
|
||||
}
|
||||
|
||||
struct MlasSgemmMultiplyAlphaRow
|
||||
{
|
||||
template<size_t Count, size_t Index>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Iteration(
|
||||
MLAS_FLOAT32X4 Accumulators[4],
|
||||
MLAS_FLOAT32X4 AlphaBroadcast
|
||||
)
|
||||
{
|
||||
Accumulators[Index] = MlasMultiplyFloat32x4(Accumulators[Index], AlphaBroadcast);
|
||||
}
|
||||
};
|
||||
|
||||
struct MlasSgemmMultiplyAlphaAddRow
|
||||
{
|
||||
template<size_t Count, size_t Index>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Iteration(
|
||||
MLAS_FLOAT32X4 Accumulators[4],
|
||||
MLAS_FLOAT32X4 AlphaBroadcast,
|
||||
const float* C
|
||||
)
|
||||
{
|
||||
Accumulators[Index] = MlasMultiplyAddFloat32x4(Accumulators[Index],
|
||||
AlphaBroadcast, MlasLoadFloat32x4(C + Index * 4));
|
||||
}
|
||||
};
|
||||
|
||||
struct MlasSgemmStoreRow
|
||||
{
|
||||
template<size_t Count, size_t Index>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Iteration(
|
||||
MLAS_FLOAT32X4 Accumulators[4],
|
||||
float* C
|
||||
)
|
||||
{
|
||||
MlasStoreFloat32x4(C + Index * 4, Accumulators[Index]);
|
||||
}
|
||||
};
|
||||
|
||||
template<size_t VectorCount>
|
||||
struct MlasSgemmStoreVector
|
||||
{
|
||||
template<size_t RowCount, size_t Row>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Iteration(
|
||||
MLAS_FLOAT32X4 Accumulators[RowCount][4],
|
||||
float* C,
|
||||
size_t ldc,
|
||||
MLAS_FLOAT32X4 AlphaBroadcast,
|
||||
bool ZeroMode
|
||||
)
|
||||
{
|
||||
float* c = C + Row * ldc;
|
||||
|
||||
if (ZeroMode) {
|
||||
MlasLoopUnroll<VectorCount, MlasSgemmMultiplyAlphaRow>()(Accumulators[Row], AlphaBroadcast);
|
||||
} else {
|
||||
MlasLoopUnroll<VectorCount, MlasSgemmMultiplyAlphaAddRow>()(Accumulators[Row], AlphaBroadcast, c);
|
||||
}
|
||||
|
||||
MlasLoopUnroll<VectorCount, MlasSgemmStoreRow>()(Accumulators[Row], c);
|
||||
|
||||
//
|
||||
// Shift down any unaligned elements to the bottom for further processing.
|
||||
//
|
||||
|
||||
if (VectorCount < 4) {
|
||||
Accumulators[Row][0] = Accumulators[Row][VectorCount];
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
struct MlasSgemmMultiplyAlphaTrailing
|
||||
{
|
||||
template<size_t RowCount, size_t Row>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Iteration(
|
||||
MLAS_FLOAT32X4 Accumulators[RowCount][4],
|
||||
MLAS_FLOAT32X4 AlphaBroadcast
|
||||
)
|
||||
{
|
||||
Accumulators[Row][0] = MlasMultiplyFloat32x4(Accumulators[Row][0], AlphaBroadcast);
|
||||
}
|
||||
};
|
||||
|
||||
template<unsigned Lane>
|
||||
struct MlasSgemmStoreScalar
|
||||
{
|
||||
template<size_t RowCount, size_t Row>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Iteration(
|
||||
MLAS_FLOAT32X4 Accumulators[RowCount][4],
|
||||
float* C,
|
||||
size_t ldc,
|
||||
bool ZeroMode
|
||||
)
|
||||
{
|
||||
float* c = C + Row * ldc + Lane;
|
||||
float Value = MlasExtractLaneFloat32x4<Lane>(Accumulators[Row][0]);
|
||||
|
||||
if (!ZeroMode) {
|
||||
Value += *c;
|
||||
}
|
||||
|
||||
*c = Value;
|
||||
}
|
||||
};
|
||||
|
||||
template<size_t RowCount>
|
||||
MLAS_FORCEINLINE
|
||||
size_t
|
||||
MlasSgemmProcessCount(
|
||||
const float* A,
|
||||
const float* B,
|
||||
float* C,
|
||||
size_t CountK,
|
||||
size_t CountN,
|
||||
size_t lda,
|
||||
size_t ldc,
|
||||
MLAS_FLOAT32X4 AlphaBroadcast,
|
||||
bool ZeroMode
|
||||
)
|
||||
{
|
||||
do {
|
||||
|
||||
const float* a = A;
|
||||
size_t k = CountK;
|
||||
|
||||
MLAS_FLOAT32X4 Accumulators[RowCount][4];
|
||||
MLAS_FLOAT32X4 AElements[RowCount];
|
||||
MLAS_FLOAT32X4 ABroadcast[RowCount];
|
||||
|
||||
//
|
||||
// Clear the block accumulators.
|
||||
//
|
||||
|
||||
MlasLoopUnroll<RowCount, MlasSgemmZeroAccumulators>()(Accumulators);
|
||||
|
||||
//
|
||||
// Compute the output block.
|
||||
//
|
||||
|
||||
while (k >= 4) {
|
||||
|
||||
MlasLoopUnroll<RowCount, MlasSgemmLoadAElements>()(AElements, a, lda);
|
||||
|
||||
MlasLoopUnroll<RowCount, MlasSgemmSplatAElements<0>>()(AElements, ABroadcast);
|
||||
MlasSgemmComputeBlock<RowCount>(Accumulators, ABroadcast, B);
|
||||
|
||||
MlasLoopUnroll<RowCount, MlasSgemmSplatAElements<1>>()(AElements, ABroadcast);
|
||||
MlasSgemmComputeBlock<RowCount>(Accumulators, ABroadcast, B + 16);
|
||||
|
||||
MlasLoopUnroll<RowCount, MlasSgemmSplatAElements<2>>()(AElements, ABroadcast);
|
||||
MlasSgemmComputeBlock<RowCount>(Accumulators, ABroadcast, B + 32);
|
||||
|
||||
MlasLoopUnroll<RowCount, MlasSgemmSplatAElements<3>>()(AElements, ABroadcast);
|
||||
MlasSgemmComputeBlock<RowCount>(Accumulators, ABroadcast, B + 48);
|
||||
|
||||
a += 4;
|
||||
B += 16 * 4;
|
||||
k -= 4;
|
||||
}
|
||||
|
||||
while (k > 0) {
|
||||
|
||||
MlasLoopUnroll<RowCount, MlasSgemmBroadcastAElements>()(ABroadcast, a, lda);
|
||||
MlasSgemmComputeBlock<RowCount>(Accumulators, ABroadcast, B);
|
||||
|
||||
a += 1;
|
||||
B += 16;
|
||||
k -= 1;
|
||||
}
|
||||
|
||||
if (CountN >= 16) {
|
||||
|
||||
//
|
||||
// Store the entire output block.
|
||||
//
|
||||
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVector<4>>()(Accumulators, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
|
||||
} else {
|
||||
|
||||
//
|
||||
// Store the partial output block.
|
||||
//
|
||||
|
||||
if (CountN >= 12) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVector<3>>()(Accumulators, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
} else if (CountN >= 8) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVector<2>>()(Accumulators, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
} else if (CountN >= 4) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVector<1>>()(Accumulators, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
}
|
||||
|
||||
//
|
||||
// Store the remaining unaligned columns.
|
||||
//
|
||||
|
||||
C += (CountN & ~3);
|
||||
CountN &= 3;
|
||||
|
||||
if (CountN > 0) {
|
||||
|
||||
MlasLoopUnroll<RowCount, MlasSgemmMultiplyAlphaTrailing>()(Accumulators, AlphaBroadcast);
|
||||
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreScalar<0>>()(Accumulators, C, ldc, ZeroMode);
|
||||
|
||||
if (CountN >= 2) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreScalar<1>>()(Accumulators, C, ldc, ZeroMode);
|
||||
}
|
||||
|
||||
if (CountN >= 3) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreScalar<2>>()(Accumulators, C, ldc, ZeroMode);
|
||||
}
|
||||
}
|
||||
|
||||
break;
|
||||
}
|
||||
|
||||
C += 16;
|
||||
CountN -= 16;
|
||||
|
||||
} while (CountN > 0);
|
||||
|
||||
return RowCount;
|
||||
}
|
||||
|
||||
|
|
@ -1061,10 +1061,8 @@ Return Value:
|
|||
|
||||
size_t RowsHandled;
|
||||
|
||||
#if defined(MLAS_TARGET_AMD64_IX86)
|
||||
#if defined(MLAS_TARGET_AMD64_IX86) || defined(MLAS_TARGET_POWER)
|
||||
RowsHandled = MlasPlatform.GemmFloatKernel(A, B, C, CountK, CountM, CountN, lda, ldc, alpha, ZeroMode);
|
||||
#elif defined(MLAS_TARGET_POWER)
|
||||
RowsHandled = MlasSgemmKernel(A, B, C, CountK, CountM, CountN, lda, ldc, alpha, ZeroMode);
|
||||
#else
|
||||
if (ZeroMode) {
|
||||
RowsHandled = MlasSgemmKernelZero(A, B, C, CountK, CountM, CountN, lda, ldc, alpha);
|
||||
|
|
|
|||
Loading…
Reference in a new issue