diff --git a/cmake/onnxruntime_mlas.cmake b/cmake/onnxruntime_mlas.cmake index 58e899b407..8e13790c11 100644 --- a/cmake/onnxruntime_mlas.cmake +++ b/cmake/onnxruntime_mlas.cmake @@ -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 + 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) diff --git a/onnxruntime/core/mlas/lib/mlasi.h b/onnxruntime/core/mlas/lib/mlasi.h index 9bea597933..5f9dc0e3b1 100644 --- a/onnxruntime/core/mlas/lib/mlasi.h +++ b/onnxruntime/core/mlas/lib/mlasi.h @@ -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; diff --git a/onnxruntime/core/mlas/lib/platform.cpp b/onnxruntime/core/mlas/lib/platform.cpp index 1c5fda25ba..a95e3ed7c6 100644 --- a/onnxruntime/core/mlas/lib/platform.cpp +++ b/onnxruntime/core/mlas/lib/platform.cpp @@ -17,6 +17,10 @@ Abstract: #include "mlasi.h" +#if defined(MLAS_TARGET_POWER) && defined(__linux__) +#include +#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 } diff --git a/onnxruntime/core/mlas/lib/power/SgemmKernelPOWER10.cpp b/onnxruntime/core/mlas/lib/power/SgemmKernelPOWER10.cpp new file mode 100644 index 0000000000..9ba4f8062a --- /dev/null +++ b/onnxruntime/core/mlas/lib/power/SgemmKernelPOWER10.cpp @@ -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 + MLAS_FORCEINLINE + static + void + Iteration( + MLAS_FLOAT32X4 ABroadcast[RowCount], + const float* A, + size_t lda + ) + { + ABroadcast[0][Row] = A [Row * lda]; + } +}; + +template +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 +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 +struct MlasSgemmStoreVectorMMA +{ + template + 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 + MLAS_FORCEINLINE + static + void + Iteration( + MLAS_FLOAT32X4 Accumulators[RowCount], + MLAS_FLOAT32X4 AlphaBroadcast + ) + { + Accumulators[Row] = MlasMultiplyFloat32x4(Accumulators[Row], AlphaBroadcast); + } +}; +template +struct MlasSgemmStoreScalarMMA +{ + template + 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 +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()(AElements, a, lda); + MlasSgemmComputeAElements(AElements, ABroadcast); + if (CountM == 8) { + MlasLoopUnroll()(AElements, a + ( lda * 4), lda); + MlasSgemmComputeAElements(AElements, A2Broadcast); + } + MlasSgemmComputeBlockMMA(&acc[0], ABroadcast[0], A2Broadcast[0], B, CountM); + MlasSgemmComputeBlockMMA(&acc[0], ABroadcast[1], A2Broadcast[1], B+16, CountM); + MlasSgemmComputeBlockMMA(&acc[0], ABroadcast[2], A2Broadcast[2], B+32, CountM); + MlasSgemmComputeBlockMMA(&acc[0], ABroadcast[3], A2Broadcast[3], B+48, CountM); + B += 16 * 4; + a += 4; + k -= 4; + } + + while (k > 0) { + MlasLoopUnroll()(ABroadcast, a, lda); + if (CountM == 8) { + MlasLoopUnroll()(A2Broadcast, a + (lda * 4), lda); + } + MlasSgemmComputeBlockMMA(&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>()(Result, C, ldc, AlphaBroadcast, ZeroMode); + __builtin_mma_disassemble_acc ((void *)Result, &acc[1]); + MlasLoopUnroll>()(Result, C, ldc, AlphaBroadcast, ZeroMode); + __builtin_mma_disassemble_acc ((void *)Result, &acc[2]); + MlasLoopUnroll>()(Result, C, ldc, AlphaBroadcast, ZeroMode); + __builtin_mma_disassemble_acc ((void *)Result, &acc[3]); + MlasLoopUnroll>()(Result, C, ldc, AlphaBroadcast, ZeroMode); + if (CountM == 8) { + __builtin_mma_disassemble_acc ((void *)Result, &acc[4]); + MlasLoopUnroll>()(Result, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode); + __builtin_mma_disassemble_acc ((void *)Result, &acc[5]); + MlasLoopUnroll>()(Result, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode); + __builtin_mma_disassemble_acc ((void *)Result, &acc[6]); + MlasLoopUnroll>()(Result, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode); + __builtin_mma_disassemble_acc ((void *)Result, &acc[7]); + MlasLoopUnroll>()(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>()(Result, C, ldc, AlphaBroadcast, ZeroMode); + __builtin_mma_disassemble_acc ((void *)Result, &acc[1]); + MlasLoopUnroll>()(Result, C, ldc, AlphaBroadcast, ZeroMode); + __builtin_mma_disassemble_acc ((void *)Result, &acc[2]); + MlasLoopUnroll>()(Result, C, ldc, AlphaBroadcast, ZeroMode); + if (CountM == 8) { + __builtin_mma_disassemble_acc ((void *)Result, &acc[4]); + MlasLoopUnroll>()(Result, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode); + __builtin_mma_disassemble_acc ((void *)Result, &acc[5]); + MlasLoopUnroll>()(Result, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode); + __builtin_mma_disassemble_acc ((void *)Result, &acc[6]); + MlasLoopUnroll>()(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>()(Result, C, ldc, AlphaBroadcast, ZeroMode); + __builtin_mma_disassemble_acc ((void *)Result, &acc[1]); + MlasLoopUnroll>()(Result, C, ldc, AlphaBroadcast, ZeroMode); + if (CountM == 8) { + __builtin_mma_disassemble_acc ((void *)Result, &acc[4]); + MlasLoopUnroll>()(Result, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode); + __builtin_mma_disassemble_acc ((void *)Result, &acc[5]); + MlasLoopUnroll>()(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>()(Result, C, ldc, AlphaBroadcast, ZeroMode); + if (CountM == 8) { + __builtin_mma_disassemble_acc ((void *)Result, &acc[4]); + MlasLoopUnroll>()(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()(Accumulators[0], AlphaBroadcast); + MlasLoopUnroll>()(Accumulators[0], C, ldc, ZeroMode); + if (CountM == 8) { + MlasLoopUnroll()(Accumulators[1], AlphaBroadcast); + MlasLoopUnroll>()(Accumulators[1], C + (ldc*4), ldc, ZeroMode); + } + if (CountN >= 2) { + MlasLoopUnroll>()(Accumulators[0], C, ldc, ZeroMode); + if (CountM == 8) { + MlasLoopUnroll>()(Accumulators[1], C + (ldc*4), ldc, ZeroMode); + } + } + if (CountN >= 3) { + MlasLoopUnroll>()(Accumulators[0], C, ldc, ZeroMode); + if (CountM == 8) { + MlasLoopUnroll>()(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; +} diff --git a/onnxruntime/core/mlas/lib/power/SgemmKernelPower.cpp b/onnxruntime/core/mlas/lib/power/SgemmKernelPower.cpp index b56986e7a5..8c42348338 100644 --- a/onnxruntime/core/mlas/lib/power/SgemmKernelPower.cpp +++ b/onnxruntime/core/mlas/lib/power/SgemmKernelPower.cpp @@ -14,407 +14,7 @@ Abstract: multiply operation (SGEMM). --*/ - -#include "mlasi.h" - -// -// Templates to ensure that a loop is unrolled. -// - -template -struct MlasLoopUnrollStep -{ - template - MLAS_FORCEINLINE - static - void - Step( - IterationArgs&&... Arguments - ) - { - IterationType::template Iteration(Arguments...); - MlasLoopUnrollStep::template Step(Arguments...); - } -}; - -template -struct MlasLoopUnrollStep -{ - template - MLAS_FORCEINLINE - static - void - Step( - IterationArgs&&... - ) - { - // Terminate the loop. - } -}; - -template -struct MlasLoopUnroll -{ - template - MLAS_FORCEINLINE - void - operator()( - IterationArgs&&... Arguments - ) - { - MlasLoopUnrollStep::template Step(Arguments...); - } -}; - -// -// Templates used with loop unrolling to perform an action on one row of the -// output. -// - -struct MlasSgemmZeroAccumulators -{ - template - 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 - MLAS_FORCEINLINE - static - void - Iteration( - MLAS_FLOAT32X4 AElements[RowCount], - const float* A, - size_t lda - ) - { - AElements[Row] = MlasLoadFloat32x4(A + Row * lda); - } -}; - -struct MlasSgemmBroadcastAElements -{ - template - MLAS_FORCEINLINE - static - void - Iteration( - MLAS_FLOAT32X4 ABroadcast[RowCount], - const float* A, - size_t lda - ) - { - ABroadcast[Row] = MlasBroadcastFloat32x4(A + Row * lda); - } -}; - -template -struct MlasSgemmSplatAElements -{ - template - MLAS_FORCEINLINE - static - void - Iteration( - MLAS_FLOAT32X4 AElements[RowCount], - MLAS_FLOAT32X4 ABroadcast[RowCount] - ) - { - ABroadcast[Row] = vec_splat(AElements[Row], Lane); - } -}; - -struct MlasSgemmMultiplyAddRow -{ - template - 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 -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()(Accumulators, ABroadcast, BElements); -} - -struct MlasSgemmMultiplyAlphaRow -{ - template - MLAS_FORCEINLINE - static - void - Iteration( - MLAS_FLOAT32X4 Accumulators[4], - MLAS_FLOAT32X4 AlphaBroadcast - ) - { - Accumulators[Index] = MlasMultiplyFloat32x4(Accumulators[Index], AlphaBroadcast); - } -}; - -struct MlasSgemmMultiplyAlphaAddRow -{ - template - 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 - MLAS_FORCEINLINE - static - void - Iteration( - MLAS_FLOAT32X4 Accumulators[4], - float* C - ) - { - MlasStoreFloat32x4(C + Index * 4, Accumulators[Index]); - } -}; - -template -struct MlasSgemmStoreVector -{ - template - 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()(Accumulators[Row], AlphaBroadcast); - } else { - MlasLoopUnroll()(Accumulators[Row], AlphaBroadcast, c); - } - - MlasLoopUnroll()(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 - MLAS_FORCEINLINE - static - void - Iteration( - MLAS_FLOAT32X4 Accumulators[RowCount][4], - MLAS_FLOAT32X4 AlphaBroadcast - ) - { - Accumulators[Row][0] = MlasMultiplyFloat32x4(Accumulators[Row][0], AlphaBroadcast); - } -}; - -template -struct MlasSgemmStoreScalar -{ - template - 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(Accumulators[Row][0]); - - if (!ZeroMode) { - Value += *c; - } - - *c = Value; - } -}; - -template -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()(Accumulators); - - // - // Compute the output block. - // - - while (k >= 4) { - - MlasLoopUnroll()(AElements, a, lda); - - MlasLoopUnroll>()(AElements, ABroadcast); - MlasSgemmComputeBlock(Accumulators, ABroadcast, B); - - MlasLoopUnroll>()(AElements, ABroadcast); - MlasSgemmComputeBlock(Accumulators, ABroadcast, B + 16); - - MlasLoopUnroll>()(AElements, ABroadcast); - MlasSgemmComputeBlock(Accumulators, ABroadcast, B + 32); - - MlasLoopUnroll>()(AElements, ABroadcast); - MlasSgemmComputeBlock(Accumulators, ABroadcast, B + 48); - - a += 4; - B += 16 * 4; - k -= 4; - } - - while (k > 0) { - - MlasLoopUnroll()(ABroadcast, a, lda); - MlasSgemmComputeBlock(Accumulators, ABroadcast, B); - - a += 1; - B += 16; - k -= 1; - } - - if (CountN >= 16) { - - // - // Store the entire output block. - // - - MlasLoopUnroll>()(Accumulators, C, ldc, AlphaBroadcast, ZeroMode); - - } else { - - // - // Store the partial output block. - // - - if (CountN >= 12) { - MlasLoopUnroll>()(Accumulators, C, ldc, AlphaBroadcast, ZeroMode); - } else if (CountN >= 8) { - MlasLoopUnroll>()(Accumulators, C, ldc, AlphaBroadcast, ZeroMode); - } else if (CountN >= 4) { - MlasLoopUnroll>()(Accumulators, C, ldc, AlphaBroadcast, ZeroMode); - } - - // - // Store the remaining unaligned columns. - // - - C += (CountN & ~3); - CountN &= 3; - - if (CountN > 0) { - - MlasLoopUnroll()(Accumulators, AlphaBroadcast); - - MlasLoopUnroll>()(Accumulators, C, ldc, ZeroMode); - - if (CountN >= 2) { - MlasLoopUnroll>()(Accumulators, C, ldc, ZeroMode); - } - - if (CountN >= 3) { - MlasLoopUnroll>()(Accumulators, C, ldc, ZeroMode); - } - } - - break; - } - - C += 16; - CountN -= 16; - - } while (CountN > 0); - - return RowCount; -} +#include "SgemmKernelpower.h" size_t MLASCALL diff --git a/onnxruntime/core/mlas/lib/power/SgemmKernelpower.h b/onnxruntime/core/mlas/lib/power/SgemmKernelpower.h new file mode 100644 index 0000000000..1cd8d7dd16 --- /dev/null +++ b/onnxruntime/core/mlas/lib/power/SgemmKernelpower.h @@ -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 +struct MlasLoopUnrollStep +{ + template + MLAS_FORCEINLINE + static + void + Step( + IterationArgs&&... Arguments + ) + { + IterationType::template Iteration(Arguments...); + MlasLoopUnrollStep::template Step(Arguments...); + } +}; + +template +struct MlasLoopUnrollStep +{ + template + MLAS_FORCEINLINE + static + void + Step( + IterationArgs&&... + ) + { + // Terminate the loop. + } +}; + +template +struct MlasLoopUnroll +{ + template + MLAS_FORCEINLINE + void + operator()( + IterationArgs&&... Arguments + ) + { + MlasLoopUnrollStep::template Step(Arguments...); + } +}; + +// +// Templates used with loop unrolling to perform an action on one row of the +// output. +// + +struct MlasSgemmZeroAccumulators +{ + template + 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 + MLAS_FORCEINLINE + static + void + Iteration( + MLAS_FLOAT32X4 AElements[RowCount], + const float* A, + size_t lda + ) + { + AElements[Row] = MlasLoadFloat32x4(A + Row * lda); + } +}; + +struct MlasSgemmBroadcastAElements +{ + template + MLAS_FORCEINLINE + static + void + Iteration( + MLAS_FLOAT32X4 ABroadcast[RowCount], + const float* A, + size_t lda + ) + { + ABroadcast[Row] = MlasBroadcastFloat32x4(A + Row * lda); + } +}; + +template +struct MlasSgemmSplatAElements +{ + template + MLAS_FORCEINLINE + static + void + Iteration( + MLAS_FLOAT32X4 AElements[RowCount], + MLAS_FLOAT32X4 ABroadcast[RowCount] + ) + { + ABroadcast[Row] = vec_splat(AElements[Row], Lane); + } +}; + +struct MlasSgemmMultiplyAddRow +{ + template + 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 +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()(Accumulators, ABroadcast, BElements); +} + +struct MlasSgemmMultiplyAlphaRow +{ + template + MLAS_FORCEINLINE + static + void + Iteration( + MLAS_FLOAT32X4 Accumulators[4], + MLAS_FLOAT32X4 AlphaBroadcast + ) + { + Accumulators[Index] = MlasMultiplyFloat32x4(Accumulators[Index], AlphaBroadcast); + } +}; + +struct MlasSgemmMultiplyAlphaAddRow +{ + template + 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 + MLAS_FORCEINLINE + static + void + Iteration( + MLAS_FLOAT32X4 Accumulators[4], + float* C + ) + { + MlasStoreFloat32x4(C + Index * 4, Accumulators[Index]); + } +}; + +template +struct MlasSgemmStoreVector +{ + template + 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()(Accumulators[Row], AlphaBroadcast); + } else { + MlasLoopUnroll()(Accumulators[Row], AlphaBroadcast, c); + } + + MlasLoopUnroll()(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 + MLAS_FORCEINLINE + static + void + Iteration( + MLAS_FLOAT32X4 Accumulators[RowCount][4], + MLAS_FLOAT32X4 AlphaBroadcast + ) + { + Accumulators[Row][0] = MlasMultiplyFloat32x4(Accumulators[Row][0], AlphaBroadcast); + } +}; + +template +struct MlasSgemmStoreScalar +{ + template + 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(Accumulators[Row][0]); + + if (!ZeroMode) { + Value += *c; + } + + *c = Value; + } +}; + +template +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()(Accumulators); + + // + // Compute the output block. + // + + while (k >= 4) { + + MlasLoopUnroll()(AElements, a, lda); + + MlasLoopUnroll>()(AElements, ABroadcast); + MlasSgemmComputeBlock(Accumulators, ABroadcast, B); + + MlasLoopUnroll>()(AElements, ABroadcast); + MlasSgemmComputeBlock(Accumulators, ABroadcast, B + 16); + + MlasLoopUnroll>()(AElements, ABroadcast); + MlasSgemmComputeBlock(Accumulators, ABroadcast, B + 32); + + MlasLoopUnroll>()(AElements, ABroadcast); + MlasSgemmComputeBlock(Accumulators, ABroadcast, B + 48); + + a += 4; + B += 16 * 4; + k -= 4; + } + + while (k > 0) { + + MlasLoopUnroll()(ABroadcast, a, lda); + MlasSgemmComputeBlock(Accumulators, ABroadcast, B); + + a += 1; + B += 16; + k -= 1; + } + + if (CountN >= 16) { + + // + // Store the entire output block. + // + + MlasLoopUnroll>()(Accumulators, C, ldc, AlphaBroadcast, ZeroMode); + + } else { + + // + // Store the partial output block. + // + + if (CountN >= 12) { + MlasLoopUnroll>()(Accumulators, C, ldc, AlphaBroadcast, ZeroMode); + } else if (CountN >= 8) { + MlasLoopUnroll>()(Accumulators, C, ldc, AlphaBroadcast, ZeroMode); + } else if (CountN >= 4) { + MlasLoopUnroll>()(Accumulators, C, ldc, AlphaBroadcast, ZeroMode); + } + + // + // Store the remaining unaligned columns. + // + + C += (CountN & ~3); + CountN &= 3; + + if (CountN > 0) { + + MlasLoopUnroll()(Accumulators, AlphaBroadcast); + + MlasLoopUnroll>()(Accumulators, C, ldc, ZeroMode); + + if (CountN >= 2) { + MlasLoopUnroll>()(Accumulators, C, ldc, ZeroMode); + } + + if (CountN >= 3) { + MlasLoopUnroll>()(Accumulators, C, ldc, ZeroMode); + } + } + + break; + } + + C += 16; + CountN -= 16; + + } while (CountN > 0); + + return RowCount; +} + diff --git a/onnxruntime/core/mlas/lib/sgemm.cpp b/onnxruntime/core/mlas/lib/sgemm.cpp index 8934a8b520..e249486f5e 100644 --- a/onnxruntime/core/mlas/lib/sgemm.cpp +++ b/onnxruntime/core/mlas/lib/sgemm.cpp @@ -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);