diff --git a/cmake/onnxruntime_mlas.cmake b/cmake/onnxruntime_mlas.cmake index dc62d1547b..eb6a586efa 100644 --- a/cmake/onnxruntime_mlas.cmake +++ b/cmake/onnxruntime_mlas.cmake @@ -283,6 +283,8 @@ else() if(POWER AND MLAS_SOURCE_IS_NOT_SET) set(mlas_platform_srcs ${MLAS_SRC_DIR}/power/SgemmKernelPower.cpp + ${MLAS_SRC_DIR}/dgemm.cpp + ${MLAS_SRC_DIR}/power/DgemmKernelPower.cpp ) check_cxx_compiler_flag("-mcpu=power10" HAS_POWER10) if(HAS_POWER10) diff --git a/onnxruntime/core/mlas/inc/mlas.h b/onnxruntime/core/mlas/inc/mlas.h index c2003dba55..01f22c25b4 100644 --- a/onnxruntime/core/mlas/inc/mlas.h +++ b/onnxruntime/core/mlas/inc/mlas.h @@ -69,7 +69,7 @@ Abstract: // Define the support levels for the target architecture. // -#if defined(MLAS_TARGET_AMD64) +#if defined(MLAS_TARGET_AMD64) || defined (MLAS_TARGET_POWER) #define MLAS_SUPPORTS_GEMM_DOUBLE #endif diff --git a/onnxruntime/core/mlas/lib/dgemm.cpp b/onnxruntime/core/mlas/lib/dgemm.cpp index d809e07d2a..c1d20dfba4 100644 --- a/onnxruntime/core/mlas/lib/dgemm.cpp +++ b/onnxruntime/core/mlas/lib/dgemm.cpp @@ -26,7 +26,7 @@ Abstract: #define MLAS_DGEMM_TRANSA_ROWS 12 -#ifdef MLAS_TARGET_AMD64 +#if defined (MLAS_TARGET_AMD64) || defined (MLAS_TARGET_POWER) void MlasDgemmMultiplyBeta( @@ -530,7 +530,7 @@ Return Value: size_t RowsHandled; -#if defined(MLAS_TARGET_AMD64_IX86) +#if defined(MLAS_TARGET_AMD64_IX86) || defined (MLAS_TARGET_POWER) RowsHandled = MlasPlatform.GemmDoubleKernel(A, B, C, CountK, CountM, CountN, lda, ldc, alpha, ZeroMode); #else if (ZeroMode) { diff --git a/onnxruntime/core/mlas/lib/mlasi.h b/onnxruntime/core/mlas/lib/mlasi.h index d7da513a75..bb77ccd889 100644 --- a/onnxruntime/core/mlas/lib/mlasi.h +++ b/onnxruntime/core/mlas/lib/mlasi.h @@ -498,6 +498,7 @@ extern "C" { #elif defined(MLAS_TARGET_POWER) MLAS_GEMM_FLOAT_KERNEL MlasSgemmKernel; MLAS_GEMM_FLOAT_KERNEL MlasSgemmKernelPOWER10; + MLAS_GEMM_DOUBLE_KERNEL MlasDgemmKernel; #else MLAS_GEMM_FLOAT_KERNEL MlasSgemmKernelZero; MLAS_GEMM_FLOAT_KERNEL MlasSgemmKernelAdd; @@ -728,7 +729,9 @@ struct MLAS_PLATFORM { #endif const MLAS_CONV_SYM_DISPATCH* ConvSymDispatch{nullptr}; - +#if defined(MLAS_TARGET_POWER) + MLAS_GEMM_DOUBLE_KERNEL* GemmDoubleKernel; +#endif #if defined(MLAS_TARGET_AMD64) MLAS_SGEMM_KERNEL_M1_ROUTINE* KernelM1Routine; MLAS_SGEMM_KERNEL_M1_ROUTINE* KernelM1TransposeBRoutine; @@ -1770,18 +1773,44 @@ MlasPowerOf2Float32x4(MLAS_FLOAT32X4 Vector) #if defined(MLAS_SSE2_INTRINSICS) typedef __m128d MLAS_FLOAT64X2; +#elif defined(MLAS_VSX_INTRINSICS) +typedef __vector double MLAS_FLOAT64X2; #else #define MLAS_FLOAT64X2_UNSUPPORTED #endif #ifndef MLAS_FLOAT64X2_UNSUPPORTED +#if defined(MLAS_VSX_INTRINSICS) +template +MLAS_FORCEINLINE +double +MlasExtractLaneFloat64x2(MLAS_FLOAT64X2 Vector) +{ + return Vector[Lane]; +} +MLAS_FORCEINLINE +MLAS_FLOAT64X2 +MlasMultiplyAddFloat64x2(MLAS_FLOAT64X2 Vector1, MLAS_FLOAT64X2 Vector2, MLAS_FLOAT64X2 Vector3) +{ + return vec_madd(Vector1, Vector2, Vector3); +} + +MLAS_FORCEINLINE +MLAS_FLOAT64X2 +MlasBroadcastFloat64x2(const double *Value) +{ + return MLAS_FLOAT64X2{*Value, *Value}; +} +#endif MLAS_FORCEINLINE MLAS_FLOAT64X2 MlasBroadcastFloat64x2(double Value) { #if defined(MLAS_SSE2_INTRINSICS) return _mm_set1_pd(Value); +#elif defined(MLAS_VSX_INTRINSICS) + return MLAS_FLOAT64X2{Value, Value}; #endif } @@ -1791,6 +1820,8 @@ MlasZeroFloat64x2(void) { #if defined(MLAS_SSE2_INTRINSICS) return _mm_setzero_pd(); +#elif defined(MLAS_VSX_INTRINSICS) + return MlasBroadcastFloat64x2(0.0f); #endif } @@ -1800,6 +1831,8 @@ MlasLoadFloat64x2(const double* Buffer) { #if defined(MLAS_SSE2_INTRINSICS) return _mm_loadu_pd(Buffer); +#elif defined(MLAS_VSX_INTRINSICS) + return vec_vsx_ld(0, Buffer); #endif } @@ -1809,6 +1842,8 @@ MlasStoreFloat64x2(double* Buffer, MLAS_FLOAT64X2 Vector) { #if defined(MLAS_SSE2_INTRINSICS) _mm_storeu_pd(Buffer, Vector); +#elif defined(MLAS_VSX_INTRINSICS) + vec_vsx_st(Vector, 0, Buffer); #endif } @@ -1818,6 +1853,8 @@ MlasStoreAlignedFloat64x2(double* Buffer, MLAS_FLOAT64X2 Vector) { #if defined(MLAS_SSE2_INTRINSICS) _mm_store_pd(Buffer, Vector); +#elif defined(MLAS_VSX_INTRINSICS) + vec_st(Vector, 0, Buffer); #endif } @@ -1827,6 +1864,8 @@ MlasMultiplyFloat64x2(MLAS_FLOAT64X2 Vector1, MLAS_FLOAT64X2 Vector2) { #if defined(MLAS_SSE2_INTRINSICS) return _mm_mul_pd(Vector1, Vector2); +#elif defined(MLAS_VSX_INTRINSICS) + return Vector1 * Vector2; #endif } diff --git a/onnxruntime/core/mlas/lib/platform.cpp b/onnxruntime/core/mlas/lib/platform.cpp index 561fab825e..7af4379aed 100644 --- a/onnxruntime/core/mlas/lib/platform.cpp +++ b/onnxruntime/core/mlas/lib/platform.cpp @@ -370,6 +370,7 @@ Return Value: #endif // MLAS_TARGET_ARM64 #if defined(MLAS_TARGET_POWER) this->GemmFloatKernel = MlasSgemmKernel; + this->GemmDoubleKernel = MlasDgemmKernel; #if defined(__linux__) && defined(POWER10) #if (defined(__GNUC__) && ((__GNUC__ > 10) || (__GNUC__== 10 && __GNUC_MINOR__ >= 2))) || \ (defined(__clang__) && (__clang_major__ >= 12)) diff --git a/onnxruntime/core/mlas/lib/power/DgemmKernelPower.cpp b/onnxruntime/core/mlas/lib/power/DgemmKernelPower.cpp new file mode 100644 index 0000000000..ba83175440 --- /dev/null +++ b/onnxruntime/core/mlas/lib/power/DgemmKernelPower.cpp @@ -0,0 +1,87 @@ +/*++ + +Copyright (c) Microsoft Corporation. All rights reserved. + +Licensed under the MIT License. + +Module Name: + + DgemmKernelPower.cpp + +Abstract: + + This module implements the kernels for the double precision matrix/matrix + multiply operation (DGEMM). + +--*/ +#include "DgemmKernelpower.h" + +size_t +MLASCALL +MlasDgemmKernel( + const double* A, + const double* B, + double* C, + size_t CountK, + size_t CountM, + size_t CountN, + size_t lda, + size_t ldc, + double 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 + MlasDgemmCopyPackB or MlasDgemmTransposePackB. + + 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 DGEMM 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_FLOAT64X2 AlphaBroadcast = MlasBroadcastFloat64x2(alpha); + + if (CountM >= 4) { + RowsHandled = MlasDgemmProcessCount<4>(A, B, C, CountK, CountN, lda, ldc, AlphaBroadcast, ZeroMode); + } else if (CountM >= 2) { + RowsHandled = MlasDgemmProcessCount<2>(A, B, C, CountK, CountN, lda, ldc, AlphaBroadcast, ZeroMode); + } else { + RowsHandled = MlasDgemmProcessCount<1>(A, B, C, CountK, CountN, lda, ldc, AlphaBroadcast, ZeroMode); + } + + return RowsHandled; +} diff --git a/onnxruntime/core/mlas/lib/power/DgemmKernelpower.h b/onnxruntime/core/mlas/lib/power/DgemmKernelpower.h new file mode 100644 index 0000000000..a7f780a22d --- /dev/null +++ b/onnxruntime/core/mlas/lib/power/DgemmKernelpower.h @@ -0,0 +1,399 @@ +/*++ + +Copyright (c) Microsoft Corporation. All rights reserved. + +Licensed under the MIT License. + +Module Name: + + DgemmKernelPower.cpp + +Abstract: + + This module implements the kernels for the single precision matrix/matrix + multiply operation (DGEMM). + +--*/ + +#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 MlasDgemmZeroAccumulators +{ + template + MLAS_FORCEINLINE + static + void + Iteration( + MLAS_FLOAT64X2 Accumulators[RowCount][4] + ) + { + Accumulators[Row][0] = MlasZeroFloat64x2(); + Accumulators[Row][1] = MlasZeroFloat64x2(); + Accumulators[Row][2] = MlasZeroFloat64x2(); + Accumulators[Row][3] = MlasZeroFloat64x2(); + } +}; + +struct MlasDgemmLoadAElements +{ + template + MLAS_FORCEINLINE + static + void + Iteration( + MLAS_FLOAT64X2 AElements[RowCount], + const double* A, + size_t lda + ) + { + AElements[Row] = MlasLoadFloat64x2(A + Row * lda); + } +}; + +struct MlasDgemmBroadcastAElements +{ + template + MLAS_FORCEINLINE + static + void + Iteration( + MLAS_FLOAT64X2 ABroadcast[RowCount], + const double* A, + size_t lda + ) + { + ABroadcast[Row] = MlasBroadcastFloat64x2(A + Row * lda); + } +}; + +template +struct MlasDgemmSplatAElements +{ + template + MLAS_FORCEINLINE + static + void + Iteration( + MLAS_FLOAT64X2 AElements[RowCount], + MLAS_FLOAT64X2 ABroadcast[RowCount] + ) + { + ABroadcast[Row] = vec_splat(AElements[Row], Lane); + } +}; + +struct MlasDgemmMultiplyAddRow +{ + template + MLAS_FORCEINLINE + static + void + Iteration( + MLAS_FLOAT64X2 Accumulators[RowCount][4], + MLAS_FLOAT64X2 ABroadcast[RowCount], + MLAS_FLOAT64X2 BElements[4] + ) + { + Accumulators[Row][0] = MlasMultiplyAddFloat64x2(ABroadcast[Row], BElements[0], Accumulators[Row][0]); + Accumulators[Row][1] = MlasMultiplyAddFloat64x2(ABroadcast[Row], BElements[1], Accumulators[Row][1]); + Accumulators[Row][2] = MlasMultiplyAddFloat64x2(ABroadcast[Row], BElements[2], Accumulators[Row][2]); + Accumulators[Row][3] = MlasMultiplyAddFloat64x2(ABroadcast[Row], BElements[3], Accumulators[Row][3]); + } +}; + +template +MLAS_FORCEINLINE +void +MlasDgemmComputeBlock( + MLAS_FLOAT64X2 Accumulators[RowCount][4], + MLAS_FLOAT64X2 ABroadcast[RowCount], + const double* B + ) +{ + MLAS_FLOAT64X2 BElements[4]; + + BElements[0] = MlasLoadFloat64x2(B); + BElements[1] = MlasLoadFloat64x2(B + 2); + BElements[2] = MlasLoadFloat64x2(B + 4); + BElements[3] = MlasLoadFloat64x2(B + 6); + + MlasLoopUnroll()(Accumulators, ABroadcast, BElements); +} + +struct MlasDgemmMultiplyAlphaRow +{ + template + MLAS_FORCEINLINE + static + void + Iteration( + MLAS_FLOAT64X2 Accumulators[4], + MLAS_FLOAT64X2 AlphaBroadcast + ) + { + Accumulators[Index] = MlasMultiplyFloat64x2(Accumulators[Index], AlphaBroadcast); + } +}; + +struct MlasDgemmMultiplyAlphaAddRow +{ + template + MLAS_FORCEINLINE + static + void + Iteration( + MLAS_FLOAT64X2 Accumulators[4], + MLAS_FLOAT64X2 AlphaBroadcast, + const double* C + ) + { + Accumulators[Index] = MlasMultiplyAddFloat64x2(Accumulators[Index], + AlphaBroadcast, MlasLoadFloat64x2(C + Index * 2)); + } +}; + +struct MlasDgemmStoreRow +{ + template + MLAS_FORCEINLINE + static + void + Iteration( + MLAS_FLOAT64X2 Accumulators[4], + double* C + ) + { + MlasStoreFloat64x2(C + Index * 2, Accumulators[Index]); + } +}; + +template +struct MlasDgemmStoreVector +{ + template + MLAS_FORCEINLINE + static + void + Iteration( + MLAS_FLOAT64X2 Accumulators[RowCount][4], + double* C, + size_t ldc, + MLAS_FLOAT64X2 AlphaBroadcast, + bool ZeroMode + ) + { + double* 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 MlasDgemmMultiplyAlphaTrailing +{ + template + MLAS_FORCEINLINE + static + void + Iteration( + MLAS_FLOAT64X2 Accumulators[RowCount][4], + MLAS_FLOAT64X2 AlphaBroadcast + ) + { + Accumulators[Row][0] = MlasMultiplyFloat64x2(Accumulators[Row][0], AlphaBroadcast); + } +}; + +template +struct MlasDgemmStoreScalar +{ + template + MLAS_FORCEINLINE + static + void + Iteration( + MLAS_FLOAT64X2 Accumulators[RowCount][4], + double* C, + size_t ldc, + bool ZeroMode + ) + { + double* c = C + Row * ldc + Lane; + double Value = MlasExtractLaneFloat64x2(Accumulators[Row][0]); + + if (!ZeroMode) { + Value += *c; + } + + *c = Value; + } +}; + +template +MLAS_FORCEINLINE +size_t +MlasDgemmProcessCount( + const double* A, + const double* B, + double* C, + size_t CountK, + size_t CountN, + size_t lda, + size_t ldc, + MLAS_FLOAT64X2 AlphaBroadcast, + bool ZeroMode + ) +{ + do { + + const double* a = A; + size_t k = CountK; + + MLAS_FLOAT64X2 Accumulators[RowCount][4]; + MLAS_FLOAT64X2 AElements[RowCount]; + MLAS_FLOAT64X2 ABroadcast[RowCount]; + + // + // Clear the block accumulators. + // + + MlasLoopUnroll()(Accumulators); + + // + // Compute the output block. + // + while (k >= 2) { + + MlasLoopUnroll()(AElements, a, lda); + + MlasLoopUnroll>()(AElements, ABroadcast); + MlasDgemmComputeBlock(Accumulators, ABroadcast, B); + + MlasLoopUnroll>()(AElements, ABroadcast); + MlasDgemmComputeBlock(Accumulators, ABroadcast, B + 8); + + a += 2; + B += 8 * 2; + k -= 2; + } + if (k > 0) { + + MlasLoopUnroll()(ABroadcast, a, lda); + MlasDgemmComputeBlock(Accumulators, ABroadcast, B); + + a += 1; + B += 8; + k -= 1; + } + + if (CountN >= 8) { + + // + // Store the entire output block. + // + + MlasLoopUnroll>()(Accumulators, C, ldc, AlphaBroadcast, ZeroMode); + + } else { + + // + // Store the partial output block. + // + + // + if (CountN >= 6) { + MlasLoopUnroll>()(Accumulators, C, ldc, AlphaBroadcast, ZeroMode); + } else if (CountN >= 4) { + MlasLoopUnroll>()(Accumulators, C, ldc, AlphaBroadcast, ZeroMode); + } else if (CountN >= 2) { + MlasLoopUnroll>()(Accumulators, C, ldc, AlphaBroadcast, ZeroMode); + } + // + // Store the remaining unaligned columns. + // + C += (CountN & ~1); + CountN &= 1; + + if (CountN > 0) { + + MlasLoopUnroll()(Accumulators, AlphaBroadcast); + + MlasLoopUnroll>()(Accumulators, C, ldc, ZeroMode); + } + + break; + } + + C += 8; + CountN -= 8; + + } while (CountN > 0); + + return RowCount; +} + diff --git a/onnxruntime/test/mlas/unittest/test_fgemm.h b/onnxruntime/test/mlas/unittest/test_fgemm.h index fab9a875bb..8fa7961842 100644 --- a/onnxruntime/test/mlas/unittest/test_fgemm.h +++ b/onnxruntime/test/mlas/unittest/test_fgemm.h @@ -56,7 +56,7 @@ class FgemmPackedContext { } }; -#ifdef MLAS_TARGET_AMD64 +#if defined(MLAS_TARGET_AMD64) || defined (MLAS_TARGET_POWER) template <> class FgemmPackedContext { public: