From b68e98e0b0c13bde4954885cc1c998b104686a03 Mon Sep 17 00:00:00 2001 From: Tracy Sharpe <42477615+tracysh@users.noreply.github.com> Date: Wed, 28 Oct 2020 16:42:53 -0700 Subject: [PATCH] optimize QLinearConv depthwise convolutions (#5605) --- onnxruntime/core/mlas/inc/mlas.h | 14 ++ .../mlas/lib/amd64/QgemmU8S8KernelAvxVnni.asm | 4 +- onnxruntime/core/mlas/lib/convolve.cpp | 124 +++++++++++ .../core/providers/cpu/nn/qlinearconv.cc | 203 +++++++++++------- .../providers/cpu/nn/qlinearconv_op_test.cc | 24 +++ 5 files changed, 287 insertions(+), 82 deletions(-) diff --git a/onnxruntime/core/mlas/inc/mlas.h b/onnxruntime/core/mlas/inc/mlas.h index 5e44308ded..eb5724ae39 100644 --- a/onnxruntime/core/mlas/inc/mlas.h +++ b/onnxruntime/core/mlas/inc/mlas.h @@ -393,6 +393,20 @@ MlasConv( MLAS_THREADPOOL* ThreadPool ); +template +void +MLASCALL +MlasConvDepthwise( + const uint8_t* Input, + uint8_t InputZeroPoint, + const FilterType* Filter, + FilterType FilterZeroPoint, + int32_t* Output, + size_t Channels, + size_t OutputCount, + size_t KernelSize + ); + // // Pooling routines. // diff --git a/onnxruntime/core/mlas/lib/amd64/QgemmU8S8KernelAvxVnni.asm b/onnxruntime/core/mlas/lib/amd64/QgemmU8S8KernelAvxVnni.asm index 613fd7b410..bd847336eb 100644 --- a/onnxruntime/core/mlas/lib/amd64/QgemmU8S8KernelAvxVnni.asm +++ b/onnxruntime/core/mlas/lib/amd64/QgemmU8S8KernelAvxVnni.asm @@ -99,7 +99,7 @@ ComputeBlock MACRO ColumnCount, RowCount, VectorOffset, BroadcastOffset EmitIfCountGE RowCount, 3, EmitIfCountGE RowCount, 3, EmitIfCountGE RowCount, 4, - EmitIfCountGE RowCount, 4, EmitIfCountGE RowCount, 5, EmitIfCountGE RowCount, 5, EmitIfCountGE RowCount, 6, @@ -295,4 +295,4 @@ ProcessCountM5: NESTED_END MlasGemmU8S8KernelAvxVnni, _TEXT - END \ No newline at end of file + END diff --git a/onnxruntime/core/mlas/lib/convolve.cpp b/onnxruntime/core/mlas/lib/convolve.cpp index 72141607cf..df78949e06 100644 --- a/onnxruntime/core/mlas/lib/convolve.cpp +++ b/onnxruntime/core/mlas/lib/convolve.cpp @@ -1257,3 +1257,127 @@ Return Value: *WorkingBufferSize = TargetThreadCount * MLAS_CONV_WORKING_BUFFER_SIZE_PER_THREAD; } } + +template +void +MLASCALL +MlasConvDepthwise( + const uint8_t* Input, + uint8_t InputZeroPoint, + const FilterType* Filter, + FilterType FilterZeroPoint, + int32_t* Output, + size_t Channels, + size_t OutputCount, + size_t KernelSize + ) +{ +#if defined(MLAS_SSE2_INTRINSICS) + const __m128i ZeroVector = _mm_setzero_si128(); + const __m128i InputZeroPointVector = _mm_set1_epi16(InputZeroPoint); + const __m128i FilterZeroPointVector = _mm_set1_epi16(FilterZeroPoint); +#endif + + while (OutputCount > 0) { + + size_t ChannelOffset = 0; + size_t c = Channels; + +#if defined(MLAS_SSE2_INTRINSICS) + + while (c >= 8) { + + __m128i Accumulator0 = _mm_setzero_si128(); + __m128i Accumulator1 = _mm_setzero_si128(); + size_t ChannelKernelOffset = ChannelOffset; + + for (size_t k = 0; k < KernelSize; k++) { + + __m128i InputVector = _mm_loadl_epi64((const __m128i*)&Input[ChannelKernelOffset]); + __m128i FilterVector = _mm_loadl_epi64((const __m128i*)&Filter[ChannelKernelOffset]); + + InputVector = _mm_unpacklo_epi8(InputVector, ZeroVector); + + if (std::is_signed::value) { + FilterVector = _mm_srai_epi16(_mm_unpacklo_epi8(ZeroVector, FilterVector), 8); + } else { + FilterVector = _mm_unpacklo_epi8(FilterVector, ZeroVector); + } + + InputVector = _mm_sub_epi16(InputVector, InputZeroPointVector); + FilterVector = _mm_sub_epi16(FilterVector, FilterZeroPointVector); + + // N.B. Emulate PMULLD functionality on SSE2 by computing the low + // and high parts of the result and interleaving the results. + __m128i MultiplyLowWords = _mm_mullo_epi16(InputVector, FilterVector); + __m128i MultiplyHighWords = _mm_mulhi_epi16(InputVector, FilterVector); + __m128i Multiply0 = _mm_unpacklo_epi16(MultiplyLowWords, MultiplyHighWords); + __m128i Multiply1 = _mm_unpackhi_epi16(MultiplyLowWords, MultiplyHighWords); + + Accumulator0 = _mm_add_epi32(Accumulator0, Multiply0); + Accumulator1 = _mm_add_epi32(Accumulator1, Multiply1); + ChannelKernelOffset += Channels; + } + + _mm_storeu_si128((__m128i*)&Output[0], Accumulator0); + _mm_storeu_si128((__m128i*)&Output[4], Accumulator1); + Output += 8; + + ChannelOffset += 8; + c -= 8; + } + +#endif + + while (c > 0) { + + int32_t Accumulator = 0; + size_t ChannelKernelOffset = ChannelOffset; + + for (size_t k = 0; k < KernelSize; k++) { + + int32_t InputValue = int32_t(Input[ChannelKernelOffset]) - InputZeroPoint; + int32_t FilterValue = int32_t(Filter[ChannelKernelOffset]) - FilterZeroPoint; + + Accumulator += InputValue * FilterValue; + ChannelKernelOffset += Channels; + } + + *Output++ = Accumulator; + + ChannelOffset += 1; + c -= 1; + } + + Input += Channels * KernelSize; + OutputCount -= 1; + } +} + +template +void +MLASCALL +MlasConvDepthwise( + const uint8_t* Input, + uint8_t InputZeroPoint, + const int8_t* Filter, + int8_t FilterZeroPoint, + int32_t* Output, + size_t Channels, + size_t OutputCount, + size_t KernelSize + ); + +template +void +MLASCALL +MlasConvDepthwise( + const uint8_t* Input, + uint8_t InputZeroPoint, + const uint8_t* Filter, + uint8_t FilterZeroPoint, + int32_t* Output, + size_t Channels, + size_t OutputCount, + size_t KernelSize + ); diff --git a/onnxruntime/core/providers/cpu/nn/qlinearconv.cc b/onnxruntime/core/providers/cpu/nn/qlinearconv.cc index 1d28da229c..328233ad6b 100644 --- a/onnxruntime/core/providers/cpu/nn/qlinearconv.cc +++ b/onnxruntime/core/providers/cpu/nn/qlinearconv.cc @@ -252,7 +252,11 @@ Status QLinearConv::Compute(OpKernelContext* context) const { template <> class QLinearConv : public OpKernel { public: - explicit QLinearConv(const OpKernelInfo& info) : OpKernel(info), conv_attrs_(info), is_W_packed_(false) {} + explicit QLinearConv(const OpKernelInfo& info) : OpKernel(info), + conv_attrs_(info), + is_W_signed_(false), + is_W_packed_(false) { + } Status Compute(OpKernelContext* context) const override; Status PrePack(const Tensor& tensor, int input_idx, bool& is_packed) override; @@ -271,6 +275,7 @@ class QLinearConv : public OpKernel { size_t packed_W_size_; #endif BufferUniquePtr reordered_W_buffer_; + bool is_W_signed_; bool is_W_packed_; }; @@ -330,36 +335,40 @@ Status QLinearConv::PrePack(const Tensor& tensor, int input_idx, bool& i const auto* Wdata = static_cast(tensor.DataRaw()); W_shape_ = shape; + is_W_signed_ = tensor.IsDataType(); auto alloc = Info().GetAllocator(0, OrtMemTypeDefault); #ifdef MLAS_SUPPORTS_PACKED_GEMM_U8X8 - packed_W_size_ = MlasGemmPackBSize(group_output_channels, kernel_dim, true); + // Don't pack the filter buffer if the MlasConvDepthwise path is used. + if (group_input_channels != 1 && group_output_channels != 1) { + packed_W_size_ = MlasGemmPackBSize(group_output_channels, kernel_dim, true); - if (packed_W_size_ != 0) { - auto* packed_W = static_cast(alloc->Alloc(SafeInt(group_count) * packed_W_size_)); - packed_W_buffer_ = BufferUniquePtr(packed_W, BufferDeleter(alloc)); + if (packed_W_size_ != 0) { + auto* packed_W = static_cast(alloc->Alloc(SafeInt(group_count) * packed_W_size_)); + packed_W_buffer_ = BufferUniquePtr(packed_W, BufferDeleter(alloc)); - // Allocate a temporary buffer to hold the reordered oihw->ohwi filter for - // a single group. - // - // Note: The size of this buffer is less than or equal to the size of the original - // weight tensor, so the allocation size is guaranteed to fit inside size_t. - auto* group_reordered_W = static_cast(alloc->Alloc(group_output_channels * group_input_channels * kernel_size)); - BufferUniquePtr group_reordered_W_buffer(group_reordered_W, BufferDeleter(alloc)); + // Allocate a temporary buffer to hold the reordered oihw->ohwi filter for + // a single group. + // + // Note: The size of this buffer is less than or equal to the size of the original + // weight tensor, so the allocation size is guaranteed to fit inside size_t. + auto* group_reordered_W = static_cast(alloc->Alloc(group_output_channels * group_input_channels * kernel_size)); + BufferUniquePtr group_reordered_W_buffer(group_reordered_W, BufferDeleter(alloc)); - const size_t W_offset = group_output_channels * kernel_dim; + const size_t W_offset = group_output_channels * kernel_dim; - for (int64_t group_id = 0; group_id < conv_attrs_.group; ++group_id) { - ReorderFilter(Wdata, group_reordered_W, group_output_channels, group_input_channels, kernel_size); - MlasGemmPackB(group_output_channels, kernel_dim, group_reordered_W, group_output_channels, true, packed_W); - packed_W += packed_W_size_; - Wdata += W_offset; + for (int64_t group_id = 0; group_id < conv_attrs_.group; ++group_id) { + ReorderFilter(Wdata, group_reordered_W, group_output_channels, group_input_channels, kernel_size); + MlasGemmPackB(group_output_channels, kernel_dim, group_reordered_W, group_output_channels, is_W_signed_, packed_W); + packed_W += packed_W_size_; + Wdata += W_offset; + } + + is_W_packed_ = true; + is_packed = true; + return Status::OK(); } - - is_W_packed_ = true; - is_packed = true; - return Status::OK(); } #endif @@ -471,18 +480,54 @@ Status QLinearConv::Compute(OpKernelContext* context) const { const int64_t output_image_size = output_shape.Size(); const int64_t kernel_size = TensorShape(kernel_shape).Size(); - const int64_t group_count = conv_attrs_.group; - const int64_t group_input_channels = W_shape[1]; - const int64_t group_output_channels = M / group_count; + AllocatorPtr alloc; + ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&alloc)); + + // Handle the case of a dynamic weight filter. + BufferUniquePtr reordered_W_buffer; + uint8_t* reordered_W = nullptr; + bool use_reordered_W = true; + bool is_W_signed = is_W_signed_; +#ifdef MLAS_SUPPORTS_PACKED_GEMM_U8X8 + if (packed_W_buffer_) { + use_reordered_W = false; + } +#endif + if (use_reordered_W) { + if (reordered_W_buffer_) { + reordered_W = static_cast(reordered_W_buffer_.get()); + } else { + // Weight tensor was not constant or prepacking is disabled. + reordered_W = static_cast(alloc->Alloc(SafeInt(sizeof(uint8_t)) * W_shape.Size())); + reordered_W_buffer = BufferUniquePtr(reordered_W, BufferDeleter(alloc)); + ReorderFilter(static_cast(W->DataRaw()), + reordered_W, + static_cast(M), + static_cast(W_shape[1]), + static_cast(kernel_size)); + is_W_signed = W->IsDataType(); + } + } + + int64_t group_count = conv_attrs_.group; + int64_t group_input_channels = W_shape[1]; + int64_t group_output_channels = M / group_count; + + // Test for depthwise convolution. + const bool is_depthwise_conv = (use_reordered_W && group_input_channels == 1 && group_output_channels == 1); + if (is_depthwise_conv) { + // Update the input and output channels to the number of groups in order to + // reuse as much of the below standard convolution path. + group_input_channels = group_count; + group_output_channels = group_count; + group_count = 1; + } const int64_t X_offset = group_input_channels * input_image_size; const int64_t Y_offset = group_output_channels * output_image_size; const int64_t kernel_dim = group_input_channels * kernel_size; const int64_t col_buffer_size = kernel_dim * output_image_size; - AllocatorPtr alloc; - ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&alloc)); - // Use an intermediate int32_t buffer for the GEMM computation before // requantizing to the output type. auto gemm_output_data = alloc->Alloc(SafeInt(sizeof(int32_t)) * Y_offset); @@ -499,30 +544,6 @@ Status QLinearConv::Compute(OpKernelContext* context) const { auto* transpose_output = static_cast(alloc->Alloc(SafeInt(sizeof(uint8_t)) * Y_offset)); BufferUniquePtr transpose_output_buffer(transpose_output, BufferDeleter(alloc)); - // Handle the case of a dynamic weight filter. - BufferUniquePtr reordered_W_buffer; - uint8_t* reordered_W = nullptr; - bool use_reordered_W = true; -#ifdef MLAS_SUPPORTS_PACKED_GEMM_U8X8 - if (packed_W_buffer_) { - use_reordered_W = false; - } -#endif - if (use_reordered_W) { - if (reordered_W_buffer_) { - reordered_W = static_cast(reordered_W_buffer_.get()); - } else { - // Weight tensor was not constant or prepacking is disabled. - reordered_W = static_cast(alloc->Alloc(SafeInt(sizeof(uint8_t)) * W_shape.Size())); - reordered_W_buffer = BufferUniquePtr(reordered_W, BufferDeleter(alloc)); - ReorderFilter(static_cast(W->DataRaw()), - reordered_W, - static_cast(M), - static_cast(group_input_channels), - static_cast(kernel_size)); - } - } - // Pointwise convolutions can use the original input tensor in place, // otherwise a temporary buffer is required for the im2col transform. BufferUniquePtr col_buffer; @@ -596,36 +617,58 @@ Status QLinearConv::Compute(OpKernelContext* context) const { auto* worker_gemm_output = gemm_output + output_start * group_output_channels; auto* worker_transpose_output = transpose_output + output_start * group_output_channels; + if (is_depthwise_conv) { + if (is_W_signed) { + MlasConvDepthwise(worker_gemm_input, + X_zero_point_value, + reinterpret_cast(reordered_W), + static_cast(0), + worker_gemm_output, + static_cast(group_output_channels), + static_cast(output_count), + static_cast(kernel_size)); + } else { + MlasConvDepthwise(worker_gemm_input, + X_zero_point_value, + reordered_W, + static_cast(0), + worker_gemm_output, + static_cast(group_output_channels), + static_cast(output_count), + static_cast(kernel_size)); + } + } else { #ifdef MLAS_SUPPORTS_PACKED_GEMM_U8X8 - if (packed_W_buffer_) { - MlasGemm(static_cast(output_count), - static_cast(group_output_channels), - static_cast(kernel_dim), - worker_gemm_input, - static_cast(kernel_dim), - X_zero_point_value, - static_cast(packed_W_buffer_.get()) + group_id * packed_W_size_, - 0, - true, - worker_gemm_output, - static_cast(group_output_channels), - nullptr); - } else + if (packed_W_buffer_) { + MlasGemm(static_cast(output_count), + static_cast(group_output_channels), + static_cast(kernel_dim), + worker_gemm_input, + static_cast(kernel_dim), + X_zero_point_value, + static_cast(packed_W_buffer_.get()) + group_id * packed_W_size_, + 0, + is_W_signed, + worker_gemm_output, + static_cast(group_output_channels), + nullptr); + } else #endif - { - MlasGemm(static_cast(output_count), - static_cast(group_output_channels), - static_cast(kernel_dim), - worker_gemm_input, - static_cast(kernel_dim), - X_zero_point_value, - reordered_W + group_id * group_output_channels, - static_cast(M), - 0, - true, - worker_gemm_output, - static_cast(group_output_channels), - nullptr); + { + MlasGemm(static_cast(output_count), + static_cast(group_output_channels), + static_cast(kernel_dim), + worker_gemm_input, + static_cast(kernel_dim), + X_zero_point_value, + reordered_W + group_id * group_output_channels, + static_cast(M), + 0, + is_W_signed, + worker_gemm_output, + static_cast(group_output_channels), + nullptr); + } } if (output_scales.size() == 1) { diff --git a/onnxruntime/test/providers/cpu/nn/qlinearconv_op_test.cc b/onnxruntime/test/providers/cpu/nn/qlinearconv_op_test.cc index a3341b0978..a48e411a1b 100644 --- a/onnxruntime/test/providers/cpu/nn/qlinearconv_op_test.cc +++ b/onnxruntime/test/providers/cpu/nn/qlinearconv_op_test.cc @@ -589,6 +589,30 @@ TEST(QLinearConvTest, Conv2D_U8S8_Groups_PerChannel) { test.Run(); } +TEST(QLinearConvTest, Conv2D_U8S8_Depthwise5x5) { + QLinearConvOpTester test; + test.GenerateRandomInput({1, 24, 25, 25}, .03f, 12); + test.GenerateRandomWeights({24, 1, 5, 5}, .10f, 0); + test.GenerateRandomBias(); + test.SetPads({2, 2, 2, 2}); + test.SetGroups(24); + test.SetOutputScaleAndZeroPoint(.76f, 88); + test.Run(); +} + +TEST(QLinearConvTest, Conv2D_U8S8_Depthwise1x1) { + // Tests the combination of using the depthwise convolution path along with the + // pointed convolution optimization that avoids im2col. + QLinearConvOpTester test; + test.GenerateRandomInput({1, 27, 18, 18}, .03f, 12); + test.GenerateRandomInput({1, 27, 4, 4}, .03f, 12); + test.GenerateRandomWeights({27, 1, 1, 1}, .05f, 0); + test.GenerateRandomBias(); + test.SetGroups(27); + test.SetOutputScaleAndZeroPoint(.24f, 88); + test.Run(); +} + #endif } // namespace