POWER: Add Dgemm kernel for POWER processor (#9459)

* POWER: Add Dgemm kernel for POWER processor

This patch adds new dgemm kernel specific to POWER processor.

* POWER: Restrict new functions to VSX in header

* Remove warning check in header

* POWER: Dgemm Adjust indentation

Fixing indentation based on review comments.

Co-authored-by: Rajalakshmi Srinivasaraghavan <rajis@linux.ibm.com>

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)

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
 

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) {

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<unsigned Lane>
+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
 }
 

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))

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;
+}

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<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 MlasDgemmZeroAccumulators
+{
+    template<size_t RowCount, size_t Row>
+    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<size_t RowCount, size_t Row>
+    MLAS_FORCEINLINE
+    static
+    void
+    Iteration(
+        MLAS_FLOAT64X2 AElements[RowCount],
+        const double* A,
+        size_t lda
+        )
+    {
+        AElements[Row] = MlasLoadFloat64x2(A + Row * lda);
+    }
+};
+
+struct MlasDgemmBroadcastAElements
+{
+    template<size_t RowCount, size_t Row>
+    MLAS_FORCEINLINE
+    static
+    void
+    Iteration(
+        MLAS_FLOAT64X2 ABroadcast[RowCount],
+        const double* A,
+        size_t lda
+        )
+    {
+        ABroadcast[Row] = MlasBroadcastFloat64x2(A + Row * lda);
+    }
+};
+
+template<unsigned Lane>
+struct MlasDgemmSplatAElements
+{
+    template<size_t RowCount, size_t Row>
+    MLAS_FORCEINLINE
+    static
+    void
+    Iteration(
+        MLAS_FLOAT64X2 AElements[RowCount],
+        MLAS_FLOAT64X2 ABroadcast[RowCount]
+        )
+    {
+        ABroadcast[Row] = vec_splat(AElements[Row], Lane);
+    }
+};
+
+struct MlasDgemmMultiplyAddRow
+{
+    template<size_t RowCount, size_t Row>
+    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<size_t RowCount>
+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<RowCount, MlasDgemmMultiplyAddRow>()(Accumulators, ABroadcast, BElements);
+}
+
+struct MlasDgemmMultiplyAlphaRow
+{
+    template<size_t Count, size_t Index>
+    MLAS_FORCEINLINE
+    static
+    void
+    Iteration(
+        MLAS_FLOAT64X2 Accumulators[4],
+        MLAS_FLOAT64X2 AlphaBroadcast
+        )
+    {
+        Accumulators[Index] = MlasMultiplyFloat64x2(Accumulators[Index], AlphaBroadcast);
+    }
+};
+
+struct MlasDgemmMultiplyAlphaAddRow
+{
+    template<size_t Count, size_t Index>
+    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<size_t Count, size_t Index>
+    MLAS_FORCEINLINE
+    static
+    void
+    Iteration(
+        MLAS_FLOAT64X2 Accumulators[4],
+        double* C
+        )
+    {
+        MlasStoreFloat64x2(C + Index * 2, Accumulators[Index]);
+    }
+};
+
+template<size_t VectorCount>
+struct MlasDgemmStoreVector
+{
+    template<size_t RowCount, size_t Row>
+    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<VectorCount, MlasDgemmMultiplyAlphaRow>()(Accumulators[Row], AlphaBroadcast);
+        } else {
+            MlasLoopUnroll<VectorCount, MlasDgemmMultiplyAlphaAddRow>()(Accumulators[Row], AlphaBroadcast, c);
+        }
+        MlasLoopUnroll<VectorCount, MlasDgemmStoreRow>()(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<size_t RowCount, size_t Row>
+    MLAS_FORCEINLINE
+    static
+    void
+    Iteration(
+        MLAS_FLOAT64X2 Accumulators[RowCount][4],
+        MLAS_FLOAT64X2 AlphaBroadcast
+        )
+    {
+        Accumulators[Row][0] = MlasMultiplyFloat64x2(Accumulators[Row][0], AlphaBroadcast);
+    }
+};
+
+template<unsigned Lane>
+struct MlasDgemmStoreScalar
+{
+    template<size_t RowCount, size_t Row>
+    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<Lane>(Accumulators[Row][0]);
+
+        if (!ZeroMode) {
+            Value += *c;
+        }
+
+        *c = Value;
+    }
+};
+
+template<size_t RowCount>
+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<RowCount, MlasDgemmZeroAccumulators>()(Accumulators);
+
+        //
+        // Compute the output block.
+        //
+        while (k >= 2) {
+
+            MlasLoopUnroll<RowCount, MlasDgemmLoadAElements>()(AElements, a, lda);
+
+            MlasLoopUnroll<RowCount, MlasDgemmSplatAElements<0>>()(AElements, ABroadcast);
+            MlasDgemmComputeBlock<RowCount>(Accumulators, ABroadcast, B);
+
+            MlasLoopUnroll<RowCount, MlasDgemmSplatAElements<1>>()(AElements, ABroadcast);
+            MlasDgemmComputeBlock<RowCount>(Accumulators, ABroadcast, B + 8);
+
+            a += 2;
+            B += 8 * 2;
+            k -= 2;
+        }
+        if (k > 0) {
+
+            MlasLoopUnroll<RowCount, MlasDgemmBroadcastAElements>()(ABroadcast, a, lda);
+            MlasDgemmComputeBlock<RowCount>(Accumulators, ABroadcast, B);
+
+            a += 1;
+            B += 8;
+            k -= 1;
+        }
+
+        if (CountN >= 8) {
+
+            //
+            // Store the entire output block.
+            //
+
+            MlasLoopUnroll<RowCount, MlasDgemmStoreVector<4>>()(Accumulators, C, ldc, AlphaBroadcast, ZeroMode);
+
+        } else {
+
+            //
+            // Store the partial output block.
+            //
+
+            //
+            if (CountN >= 6) {
+                MlasLoopUnroll<RowCount, MlasDgemmStoreVector<3>>()(Accumulators, C, ldc, AlphaBroadcast, ZeroMode);
+            } else if (CountN >= 4) {
+                MlasLoopUnroll<RowCount, MlasDgemmStoreVector<2>>()(Accumulators, C, ldc, AlphaBroadcast, ZeroMode);
+            } else if (CountN >= 2) {
+                MlasLoopUnroll<RowCount, MlasDgemmStoreVector<1>>()(Accumulators, C, ldc, AlphaBroadcast, ZeroMode);
+            }
+            //
+            // Store the remaining unaligned columns.
+            //
+            C += (CountN & ~1);
+            CountN &= 1;
+
+            if (CountN > 0) {
+
+                MlasLoopUnroll<RowCount, MlasDgemmMultiplyAlphaTrailing>()(Accumulators, AlphaBroadcast);
+
+                MlasLoopUnroll<RowCount, MlasDgemmStoreScalar<0>>()(Accumulators, C, ldc, ZeroMode);
+            }
+
+            break;
+        }
+
+        C += 8;
+        CountN -= 8;
+
+    } while (CountN > 0);
+
+    return RowCount;
+}
+

onnxruntime/test/mlas/unittest/test_fgemm.h

@@ -56,7 +56,7 @@ class FgemmPackedContext<float, false> {
   }
 };
 
-#ifdef MLAS_TARGET_AMD64
+#if defined(MLAS_TARGET_AMD64) || defined (MLAS_TARGET_POWER)
 template <>
 class FgemmPackedContext<double, false> {
  public: