MLAS: add 32-bit transpose support (#7092)

onnxruntime/core/mlas/inc/mlas.h

@@ -539,6 +539,15 @@ MlasTranspose(
     size_t N
     );
 
+void
+MLASCALL
+MlasTranspose(
+    const float* Input,
+    float* Output,
+    size_t M,
+    size_t N
+    );
+
 //
 // Buffer reordering routines.
 //

onnxruntime/core/mlas/lib/transpose.cpp

@@ -16,7 +16,37 @@ Abstract:
 
 #include "mlasi.h"
 
-#if defined(MLAS_SSE2_INTRINSICS) || defined(MLAS_NEON_INTRINSICS)
+#if defined(MLAS_SSE2_INTRINSICS)
+
+MLAS_FORCEINLINE
+void
+MlasTranspose4x4Block(
+    const uint32_t* Input,
+    size_t InputStride,
+    uint32_t* Output,
+    size_t OutputStride
+    )
+{
+    __m128i a0 = _mm_loadu_si128((const __m128i*)&Input[InputStride * 0]);
+    __m128i a1 = _mm_loadu_si128((const __m128i*)&Input[InputStride * 1]);
+    __m128i a2 = _mm_loadu_si128((const __m128i*)&Input[InputStride * 2]);
+    __m128i a3 = _mm_loadu_si128((const __m128i*)&Input[InputStride * 3]);
+
+    __m128i b0 = _mm_unpacklo_epi32(a0, a2);
+    __m128i b1 = _mm_unpackhi_epi32(a0, a2);
+    __m128i b2 = _mm_unpacklo_epi32(a1, a3);
+    __m128i b3 = _mm_unpackhi_epi32(a1, a3);
+
+    __m128i c0 = _mm_unpacklo_epi32(b0, b2);
+    __m128i c1 = _mm_unpackhi_epi32(b0, b2);
+    __m128i c2 = _mm_unpacklo_epi32(b1, b3);
+    __m128i c3 = _mm_unpackhi_epi32(b1, b3);
+
+    _mm_storeu_si128((__m128i*)&Output[OutputStride * 0], c0);
+    _mm_storeu_si128((__m128i*)&Output[OutputStride * 1], c1);
+    _mm_storeu_si128((__m128i*)&Output[OutputStride * 2], c2);
+    _mm_storeu_si128((__m128i*)&Output[OutputStride * 3], c3);
+}
 
 MLAS_FORCEINLINE
 void
@@ -26,34 +56,7 @@ MlasTranspose8x8Block(
     uint8_t* Output,
     size_t OutputStride
     )
-/*++
-
-Routine Description:
-
-    This routine transposes an eight by eight element block from the input
-    matrix to the output matrix.
-
-Arguments:
-
-    Input - Supplies the input buffer.
-
-    InputStride - Supplies the number of elements between rows of the input
-        matrix.
-
-    Output - Supplies the output buffer.
-
-    OutputStride - Supplies the number of elements between rows of the output
-        matrix.
-
-Return Value:
-
-    None.
-
---*/
 {
-
-#if defined(MLAS_SSE2_INTRINSICS)
-
     __m128i a0 = _mm_loadl_epi64((const __m128i*)&Input[InputStride * 0]);
     __m128i a1 = _mm_loadl_epi64((const __m128i*)&Input[InputStride * 1]);
     __m128i b0 = _mm_unpacklo_epi8(a0, a1);
@@ -90,9 +93,45 @@ Return Value:
     __m128 d3 = _mm_castsi128_ps(_mm_unpackhi_epi32(c1, c3));
     _mm_storel_pi((__m64*)&Output[OutputStride * 6], d3);
     _mm_storeh_pi((__m64*)&Output[OutputStride * 7], d3);
+}
 
 #elif defined(MLAS_NEON_INTRINSICS)
 
+MLAS_FORCEINLINE
+void
+MlasTranspose4x4Block(
+    const uint32_t* Input,
+    size_t InputStride,
+    uint32_t* Output,
+    size_t OutputStride
+    )
+{
+    uint32x4_t a0 = vld1q_u32(&Input[InputStride * 0]);
+    uint32x4_t a1 = vld1q_u32(&Input[InputStride * 1]);
+    uint32x4_t a2 = vld1q_u32(&Input[InputStride * 2]);
+    uint32x4_t a3 = vld1q_u32(&Input[InputStride * 3]);
+
+    uint32x4x2_t b0 = vzipq_u32(a0, a2);
+    uint32x4x2_t b1 = vzipq_u32(a1, a3);
+
+    uint32x4x2_t c0 = vzipq_u32(b0.val[0], b1.val[0]);
+    uint32x4x2_t c1 = vzipq_u32(b0.val[1], b1.val[1]);
+
+    vst1q_u32(&Output[OutputStride * 0], c0.val[0]);
+    vst1q_u32(&Output[OutputStride * 1], c0.val[1]);
+    vst1q_u32(&Output[OutputStride * 2], c1.val[0]);
+    vst1q_u32(&Output[OutputStride * 3], c1.val[1]);
+}
+
+MLAS_FORCEINLINE
+void
+MlasTranspose8x8Block(
+    const uint8_t* Input,
+    size_t InputStride,
+    uint8_t* Output,
+    size_t OutputStride
+    )
+{
     uint8x8_t a0 = vld1_u8(&Input[InputStride * 0]);
     uint8x8_t a1 = vld1_u8(&Input[InputStride * 1]);
     uint8x8x2_t b0 = vzip_u8(a0, a1);
@@ -127,39 +166,71 @@ Return Value:
     vst1_u8(&Output[OutputStride * 5], vreinterpret_u8_u32(d2.val[1]));
     vst1_u8(&Output[OutputStride * 6], vreinterpret_u8_u32(d3.val[0]));
     vst1_u8(&Output[OutputStride * 7], vreinterpret_u8_u32(d3.val[1]));
-
-#endif
-
 }
 
 #endif
 
+template<typename ElementType>
+MLAS_FORCEINLINE
+void
+MlasTranspose4xNVector(
+    const ElementType* Input,
+    size_t InputStride,
+    ElementType* Output,
+    size_t OutputStride
+    )
+{
+    ElementType a0 = Input[InputStride * 0];
+    ElementType a1 = Input[InputStride * 1];
+    ElementType a2 = Input[InputStride * 2];
+    ElementType a3 = Input[InputStride * 3];
+
+    Output[OutputStride * 0] = a0;
+    Output[OutputStride * 1] = a1;
+    Output[OutputStride * 2] = a2;
+    Output[OutputStride * 3] = a3;
+}
+
+template<typename ElementType>
 MLAS_FORCEINLINE
 void
 MlasTranspose8xNVector(
-    const uint8_t* Input,
+    const ElementType* Input,
     size_t InputStride,
-    uint8_t* Output,
+    ElementType* Output,
     size_t OutputStride
     )
+{
+    MlasTranspose4xNVector(&Input[InputStride * 0], InputStride, &Output[OutputStride * 0], OutputStride);
+    MlasTranspose4xNVector(&Input[InputStride * 4], InputStride, &Output[OutputStride * 4], OutputStride);
+}
+
+void
+MLASCALL
+MlasTranspose(
+    const uint32_t* Input,
+    uint32_t* Output,
+    size_t M,
+    size_t N
+    )
 /*++
 
 Routine Description:
 
-    This routine transposes an eight element vector from the input matrix
-    to the output matrix.
+    This routine transposes the input matrix (M rows by N columns) to the
+    output matrix (N rows by M columns).
 
 Arguments:
 
     Input - Supplies the input buffer.
 
-    InputStride - Supplies the number of elements between rows of the input
-        matrix.
-
     Output - Supplies the output buffer.
 
-    OutputStride - Supplies the number of elements between rows of the output
-        matrix.
+    M - Supplies the number of rows for the input matrix and the number of
+        columns for the output matrix.
+
+    N - Supplies the number of columns for the input matrix and the number of
+        rows for the output matrix.
 
 Return Value:
 
@@ -167,25 +238,95 @@ Return Value:
 
 --*/
 {
-    uint8_t a0 = Input[InputStride * 0];
-    uint8_t a1 = Input[InputStride * 1];
-    uint8_t a2 = Input[InputStride * 2];
-    uint8_t a3 = Input[InputStride * 3];
+    size_t n = N;
 
-    Output[OutputStride * 0] = a0;
-    Output[OutputStride * 1] = a1;
-    Output[OutputStride * 2] = a2;
-    Output[OutputStride * 3] = a3;
+    //
+    // Transpose elements from the input matrix to the output matrix 4 columns
+    // at a time.
+    //
 
-    uint8_t a4 = Input[InputStride * 4];
-    uint8_t a5 = Input[InputStride * 5];
-    uint8_t a6 = Input[InputStride * 6];
-    uint8_t a7 = Input[InputStride * 7];
+    while (n >= 4) {
 
-    Output[OutputStride * 4] = a4;
-    Output[OutputStride * 5] = a5;
-    Output[OutputStride * 6] = a6;
-    Output[OutputStride * 7] = a7;
+        const uint32_t* s = Input;
+        uint32_t* d = Output;
+        size_t m = M;
+
+#if defined(MLAS_SSE2_INTRINSICS) || defined(MLAS_NEON_INTRINSICS)
+
+        while (m >= 4) {
+
+            MlasTranspose4x4Block(s, N, d, M);
+
+            s += N * 4;
+            d += 4;
+            m -= 4;
+        }
+
+#endif
+
+        while (m > 0) {
+
+            MlasTranspose4xNVector(s, 1, d, M);
+
+            s += N;
+            d += 1;
+            m -= 1;
+        }
+
+        Input += 4;
+        Output += M * 4;
+        n -= 4;
+    }
+
+    //
+    // Transpose elements from the input matrix to the output matrix for the
+    // remaining columns.
+    //
+
+    while (n > 0) {
+
+        const uint32_t* s = Input;
+        uint32_t* d = Output;
+        size_t m = M;
+
+        while (m >= 4) {
+
+            MlasTranspose4xNVector(s, N, d, 1);
+
+            s += N * 4;
+            d += 4;
+            m -= 4;
+        }
+
+        while (m > 0) {
+
+            d[0] = s[0];
+
+            s += N;
+            d += 1;
+            m -= 1;
+        }
+
+        Input += 1;
+        Output += M;
+        n -= 1;
+    }
+}
+
+void
+MLASCALL
+MlasTranspose(
+    const float* Input,
+    float* Output,
+    size_t M,
+    size_t N
+    )
+{
+    MlasTranspose(
+        reinterpret_cast<const uint32_t*>(Input),
+        reinterpret_cast<uint32_t*>(Output),
+        M,
+        N);
 }
 
 void

onnxruntime/core/providers/cpu/tensor/transpose.cc

@@ -340,11 +340,23 @@ We use memcpy if the block size is larger.
 We fall back to the default implementation in all other cases, and if the input is std::string.
 */
 
+namespace {
+
+template <typename T>
+struct has_mlas_transpose : std::false_type {};
+
+template <>
+struct has_mlas_transpose<uint8_t> : std::true_type {};
+
+template <>
+struct has_mlas_transpose<uint32_t> : std::true_type {};
+
 // moving a single axis outwards where the read/write size is a power of 2 and between 8 and 64 bits.
 template <typename T>
-static void SimpleTransposeSingleAxisOutwards(const T* input_data, T* output_data,
-                                              int64_t num_loops, int64_t num_writers,
-                                              int64_t writes_per_loop, int64_t writes_per_writer_per_loop) {
+typename std::enable_if<!has_mlas_transpose<T>::value, void>::type
+SimpleTransposeSingleAxisOutwards(const T* input_data, T* output_data,
+                                  int64_t num_loops, int64_t num_writers,
+                                  int64_t writes_per_loop, int64_t writes_per_writer_per_loop) {
   const T* end;
   for (int64_t l = 0; l < num_loops; ++l) {
     T* output_for_first_writer = output_data;
@@ -367,9 +379,11 @@ static void SimpleTransposeSingleAxisOutwards(const T* input_data, T* output_dat
   }
 }
 
-static void SimpleTransposeSingleAxisOutwards(const uint8_t* input_data, uint8_t* output_data,
-                                              int64_t num_loops, int64_t num_writers,
-                                              int64_t writes_per_loop, int64_t writes_per_writer_per_loop) {
+template <typename T>
+typename std::enable_if<has_mlas_transpose<T>::value, void>::type
+SimpleTransposeSingleAxisOutwards(const T* input_data, T* output_data,
+                                  int64_t num_loops, int64_t num_writers,
+                                  int64_t writes_per_loop, int64_t writes_per_writer_per_loop) {
   for (int64_t l = 0; l < num_loops; ++l) {
     MlasTranspose(input_data,
                   output_data,
@@ -381,8 +395,8 @@ static void SimpleTransposeSingleAxisOutwards(const uint8_t* input_data, uint8_t
 }
 
 //  `input_shape_override` overrides the shape of `input` for compute purposes.
-static void TransposeSingleAxisOutwards(const std::vector<size_t>& permutations, const Tensor& input, Tensor& output,
-                                        int64_t from, int64_t to, const TensorShape* input_shape_override = nullptr) {
+void TransposeSingleAxisOutwards(const std::vector<size_t>& permutations, const Tensor& input, Tensor& output,
+                                 int64_t from, int64_t to, const TensorShape* input_shape_override = nullptr) {
   ORT_UNUSED_PARAMETER(permutations);
 
   const auto& input_shape = input_shape_override ? *input_shape_override : input.Shape();
@@ -449,9 +463,10 @@ static void TransposeSingleAxisOutwards(const std::vector<size_t>& permutations,
 }
 
 template <typename T>
-static void SimpleTransposeSingleAxisInwards(const T* input_data, T* output_data,
-                                             int64_t num_loops, int64_t num_readers,
-                                             int64_t reads_per_loop, int64_t reads_per_reader_per_loop) {
+typename std::enable_if<!has_mlas_transpose<T>::value, void>::type
+SimpleTransposeSingleAxisInwards(const T* input_data, T* output_data,
+                                 int64_t num_loops, int64_t num_readers,
+                                 int64_t reads_per_loop, int64_t reads_per_reader_per_loop) {
   T* end;
   for (int64_t l = 0; l < num_loops; ++l) {
     const T* input_for_first_reader = input_data;
@@ -473,9 +488,11 @@ static void SimpleTransposeSingleAxisInwards(const T* input_data, T* output_data
   }
 }
 
-static void SimpleTransposeSingleAxisInwards(const uint8_t* input_data, uint8_t* output_data,
-                                             int64_t num_loops, int64_t num_readers,
-                                             int64_t reads_per_loop, int64_t reads_per_reader_per_loop) {
+template <typename T>
+typename std::enable_if<has_mlas_transpose<T>::value, void>::type
+SimpleTransposeSingleAxisInwards(const T* input_data, T* output_data,
+                                 int64_t num_loops, int64_t num_readers,
+                                 int64_t reads_per_loop, int64_t reads_per_reader_per_loop) {
   for (int64_t l = 0; l < num_loops; ++l) {
     MlasTranspose(input_data,
                   output_data,
@@ -488,8 +505,8 @@ static void SimpleTransposeSingleAxisInwards(const uint8_t* input_data, uint8_t*
 
 // moving a single axis inwards where the read/write size is a power of 2 and between 8 and 64 bits.
 //  `input_shape_override` overrides the shape of `input` for compute purposes.
-static void TransposeSingleAxisInwards(const std::vector<size_t>& permutations, const Tensor& input, Tensor& output,
-                                       int64_t from, int64_t to, const TensorShape* input_shape_override = nullptr) {
+void TransposeSingleAxisInwards(const std::vector<size_t>& permutations, const Tensor& input, Tensor& output,
+                                int64_t from, int64_t to, const TensorShape* input_shape_override = nullptr) {
   ORT_UNUSED_PARAMETER(permutations);
 
   const auto& input_shape = input_shape_override ? *input_shape_override : input.Shape();
@@ -557,8 +574,8 @@ static void TransposeSingleAxisInwards(const std::vector<size_t>& permutations,
 }
 
 //  `input_shape_override` overrides the shape of `input` for compute purposes.
-static void SingleAxisTranspose(const std::vector<size_t>& permutations, const Tensor& input, Tensor& output,
-                                size_t from, size_t to, const TensorShape* input_shape_override = nullptr) {
+void SingleAxisTranspose(const std::vector<size_t>& permutations, const Tensor& input, Tensor& output,
+                         size_t from, size_t to, const TensorShape* input_shape_override = nullptr) {
   if (from > to) {
     TransposeSingleAxisOutwards(permutations, input, output, from, to, input_shape_override);
   } else {
@@ -566,7 +583,7 @@ static void SingleAxisTranspose(const std::vector<size_t>& permutations, const T
   }
 }
 
-static bool IsMovingSingleAxis(const std::vector<size_t>& permutations, size_t& from, size_t& to) {
+bool IsMovingSingleAxis(const std::vector<size_t>& permutations, size_t& from, size_t& to) {
   // if a single axis moved to an outer dimension, the values should be one lower than the index until the slot the
   // axis was moved from, and equal to the index after that.
   // e.g. axis 3 moves out to 1 would be: 0, 3, 1, 2, 4
@@ -635,6 +652,8 @@ static bool IsMovingSingleAxis(const std::vector<size_t>& permutations, size_t&
   return single_axis_moved;
 }
 
+}  // namespace
+
 bool IsTransposeReshape(const std::vector<size_t>& perm, const std::vector<int64_t>& input_dims) {
   // As long as the dims with values > 1 stay in the same order, it's a reshape.
   // Example: Shape=(1,1,1024,4096) -> perm=(2,0,3,1).

onnxruntime/test/mlas/unittest.cpp

@@ -2827,7 +2827,8 @@ public:
     }
 };
 
-class MlasQLinearGlobalAveragePoolU8Test : public MlasTestBase {
+class MlasQLinearGlobalAveragePoolU8Test : public MlasTestBase
+{
 private:
     MatrixGuardBuffer<uint8_t> BufferInput;
     MatrixGuardBuffer<uint8_t> BufferOutput;
@@ -3214,6 +3215,61 @@ public:
     }
 };
 
+template<typename ElementType>
+class MlasTransposeTest : public MlasTestBase
+{
+private:
+    MatrixGuardBuffer<ElementType> BufferInput;
+    MatrixGuardBuffer<ElementType> BufferOutput;
+    MatrixGuardBuffer<ElementType> BufferOutputReference;
+
+    void
+    Test(
+        size_t M,
+        size_t N
+        )
+    {
+        ElementType* Input = BufferInput.GetBuffer(M * N);
+        ElementType* Output = BufferOutput.GetBuffer(M * N);
+        ElementType* OutputReference = BufferOutputReference.GetBuffer(M * N);
+
+        MlasTranspose(Input, Output, M, N);
+        ReferenceTranspose(Input, OutputReference, M, N);
+
+        if (memcmp(Output, OutputReference, M * N * sizeof(ElementType)) != 0) {
+            printf("mismatch: %zd,%zd (element size %zd)\n", M, N, sizeof(ElementType));
+        }
+    }
+
+    void
+    ReferenceTranspose(
+        const ElementType* Input,
+        ElementType* Output,
+        size_t M,
+        size_t N
+        )
+    {
+        for (size_t m = 0; m < M; m++) {
+            for (size_t n = 0; n < N; n++) {
+                Output[n * M + m] = Input[m * N + n];
+            }
+        }
+    }
+
+public:
+    void
+    ExecuteShort(
+        void
+        ) override
+    {
+        for (size_t m = 1; m <= 32; m++) {
+            for (size_t n = 1; n <= 32; n++) {
+                Test(m, n);
+            }
+        }
+    }
+};
+
 void
 RunThreadedTests(
     void
@@ -3334,6 +3390,10 @@ main(
     onnxruntime::make_unique<MlasQuantizeLinearTest<int8_t>>()->ExecuteShort();
     onnxruntime::make_unique<MlasQuantizeLinearTest<uint8_t>>()->ExecuteShort();
 
+    printf("Transpose tests.\n");
+    onnxruntime::make_unique<MlasTransposeTest<uint8_t>>()->ExecuteShort();
+    onnxruntime::make_unique<MlasTransposeTest<uint32_t>>()->ExecuteShort();
+
     printf("Done.\n");
 
     return 0;