From 416ee3c4d26e9535576002b2c65a88c594f8e172 Mon Sep 17 00:00:00 2001 From: Tracy Sharpe <42477615+tracysh@users.noreply.github.com> Date: Mon, 22 Mar 2021 16:20:31 -0700 Subject: [PATCH] MLAS: add 32-bit transpose support (#7092) --- onnxruntime/core/mlas/inc/mlas.h | 9 + onnxruntime/core/mlas/lib/transpose.cpp | 253 ++++++++++++++---- .../core/providers/cpu/tensor/transpose.cc | 57 ++-- onnxruntime/test/mlas/unittest.cpp | 62 ++++- 4 files changed, 305 insertions(+), 76 deletions(-) diff --git a/onnxruntime/core/mlas/inc/mlas.h b/onnxruntime/core/mlas/inc/mlas.h index 9a23bfc561..92f0958be9 100644 --- a/onnxruntime/core/mlas/inc/mlas.h +++ b/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. // diff --git a/onnxruntime/core/mlas/lib/transpose.cpp b/onnxruntime/core/mlas/lib/transpose.cpp index 9a2c315c8c..3f12773050 100644 --- a/onnxruntime/core/mlas/lib/transpose.cpp +++ b/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 +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 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(Input), + reinterpret_cast(Output), + M, + N); } void diff --git a/onnxruntime/core/providers/cpu/tensor/transpose.cc b/onnxruntime/core/providers/cpu/tensor/transpose.cc index 646974f93c..7693b257ae 100644 --- a/onnxruntime/core/providers/cpu/tensor/transpose.cc +++ b/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 +struct has_mlas_transpose : std::false_type {}; + +template <> +struct has_mlas_transpose : std::true_type {}; + +template <> +struct has_mlas_transpose : 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 -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::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 std::enable_if::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& permutations, const Tensor& input, Tensor& output, - int64_t from, int64_t to, const TensorShape* input_shape_override = nullptr) { +void TransposeSingleAxisOutwards(const std::vector& 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& permutations, } template -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::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 std::enable_if::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& permutations, const Tensor& input, Tensor& output, - int64_t from, int64_t to, const TensorShape* input_shape_override = nullptr) { +void TransposeSingleAxisInwards(const std::vector& 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& permutations, } // `input_shape_override` overrides the shape of `input` for compute purposes. -static void SingleAxisTranspose(const std::vector& permutations, const Tensor& input, Tensor& output, - size_t from, size_t to, const TensorShape* input_shape_override = nullptr) { +void SingleAxisTranspose(const std::vector& 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& permutations, const T } } -static bool IsMovingSingleAxis(const std::vector& permutations, size_t& from, size_t& to) { +bool IsMovingSingleAxis(const std::vector& 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& permutations, size_t& return single_axis_moved; } +} // namespace + bool IsTransposeReshape(const std::vector& perm, const std::vector& 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). diff --git a/onnxruntime/test/mlas/unittest.cpp b/onnxruntime/test/mlas/unittest.cpp index 762175b228..41fe1eb6eb 100644 --- a/onnxruntime/test/mlas/unittest.cpp +++ b/onnxruntime/test/mlas/unittest.cpp @@ -2827,7 +2827,8 @@ public: } }; -class MlasQLinearGlobalAveragePoolU8Test : public MlasTestBase { +class MlasQLinearGlobalAveragePoolU8Test : public MlasTestBase +{ private: MatrixGuardBuffer BufferInput; MatrixGuardBuffer BufferOutput; @@ -3214,6 +3215,61 @@ public: } }; +template +class MlasTransposeTest : public MlasTestBase +{ +private: + MatrixGuardBuffer BufferInput; + MatrixGuardBuffer BufferOutput; + MatrixGuardBuffer 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>()->ExecuteShort(); onnxruntime::make_unique>()->ExecuteShort(); + printf("Transpose tests.\n"); + onnxruntime::make_unique>()->ExecuteShort(); + onnxruntime::make_unique>()->ExecuteShort(); + printf("Done.\n"); return 0;