optimize QLinearConv depthwise convolutions (#5605)

onnxruntime/core/mlas/inc/mlas.h

@@ -393,6 +393,20 @@ MlasConv(
     MLAS_THREADPOOL* ThreadPool
     );
 
+template<typename FilterType>
+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.
 //

onnxruntime/core/mlas/lib/amd64/QgemmU8S8KernelAvxVnni.asm

@@ -99,7 +99,7 @@ ComputeBlock MACRO ColumnCount, RowCount, VectorOffset, BroadcastOffset
         EmitIfCountGE RowCount, 3, <vpbroadcastd ymm2,DWORD PTR [rcx+r9*2+BroadcastOffset]>
         EmitIfCountGE RowCount, 3, <MultiplyAccumulateRow ColumnCount, ymm8, ymm9>
         EmitIfCountGE RowCount, 4, <vpbroadcastd ymm2,DWORD PTR [rbx+BroadcastOffset]>
-        EmitIfCountGE RowCount, 4, <MultiplyAccumulateRow ColumnCount, ymm10, ymm11
+        EmitIfCountGE RowCount, 4, <MultiplyAccumulateRow ColumnCount, ymm10, ymm11>
         EmitIfCountGE RowCount, 5, <vpbroadcastd ymm2,DWORD PTR [rbx+r9+BroadcastOffset]>
         EmitIfCountGE RowCount, 5, <MultiplyAccumulateRow ColumnCount, ymm12, ymm13>
         EmitIfCountGE RowCount, 6, <vpbroadcastd ymm2,DWORD PTR [rbx+r9*2+BroadcastOffset]>
@@ -295,4 +295,4 @@ ProcessCountM5:
 
         NESTED_END MlasGemmU8S8KernelAvxVnni, _TEXT
 
-        END
+        END

onnxruntime/core/mlas/lib/convolve.cpp

@@ -1257,3 +1257,127 @@ Return Value:
         *WorkingBufferSize = TargetThreadCount * MLAS_CONV_WORKING_BUFFER_SIZE_PER_THREAD;
     }
 }
+
+template<typename FilterType>
+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<FilterType>::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<int8_t>(
+    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<uint8_t>(
+    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
+    );

onnxruntime/core/providers/cpu/nn/qlinearconv.cc

@@ -252,7 +252,11 @@ Status QLinearConv<uint8_t>::Compute(OpKernelContext* context) const {
 template <>
 class QLinearConv<int8_t> : public OpKernel {
  public:
-  explicit QLinearConv<int8_t>(const OpKernelInfo& info) : OpKernel(info), conv_attrs_(info), is_W_packed_(false) {}
+  explicit QLinearConv<int8_t>(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<int8_t> : 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<int8_t>::PrePack(const Tensor& tensor, int input_idx, bool& i
 
   const auto* Wdata = static_cast<const uint8_t*>(tensor.DataRaw());
   W_shape_ = shape;
+  is_W_signed_ = tensor.IsDataType<int8_t>();
 
   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<uint8_t*>(alloc->Alloc(SafeInt<size_t>(group_count) * packed_W_size_));
-    packed_W_buffer_ = BufferUniquePtr(packed_W, BufferDeleter(alloc));
+    if (packed_W_size_ != 0) {
+      auto* packed_W = static_cast<uint8_t*>(alloc->Alloc(SafeInt<size_t>(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<uint8_t*>(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<uint8_t*>(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<int8_t>::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<uint8_t*>(reordered_W_buffer_.get());
+    } else {
+      // Weight tensor was not constant or prepacking is disabled.
+      reordered_W = static_cast<uint8_t*>(alloc->Alloc(SafeInt<size_t>(sizeof(uint8_t)) * W_shape.Size()));
+      reordered_W_buffer = BufferUniquePtr(reordered_W, BufferDeleter(alloc));
+      ReorderFilter(static_cast<const uint8_t*>(W->DataRaw()),
+                    reordered_W,
+                    static_cast<size_t>(M),
+                    static_cast<size_t>(W_shape[1]),
+                    static_cast<size_t>(kernel_size));
+      is_W_signed = W->IsDataType<int8_t>();
+    }
+  }
+
+  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<size_t>(sizeof(int32_t)) * Y_offset);
@@ -499,30 +544,6 @@ Status QLinearConv<int8_t>::Compute(OpKernelContext* context) const {
   auto* transpose_output = static_cast<uint8_t*>(alloc->Alloc(SafeInt<size_t>(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<uint8_t*>(reordered_W_buffer_.get());
-    } else {
-      // Weight tensor was not constant or prepacking is disabled.
-      reordered_W = static_cast<uint8_t*>(alloc->Alloc(SafeInt<size_t>(sizeof(uint8_t)) * W_shape.Size()));
-      reordered_W_buffer = BufferUniquePtr(reordered_W, BufferDeleter(alloc));
-      ReorderFilter(static_cast<const uint8_t*>(W->DataRaw()),
-                    reordered_W,
-                    static_cast<size_t>(M),
-                    static_cast<size_t>(group_input_channels),
-                    static_cast<size_t>(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<int8_t>::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<int8_t*>(reordered_W),
+                              static_cast<int8_t>(0),
+                              worker_gemm_output,
+                              static_cast<size_t>(group_output_channels),
+                              static_cast<size_t>(output_count),
+                              static_cast<size_t>(kernel_size));
+          } else {
+            MlasConvDepthwise(worker_gemm_input,
+                              X_zero_point_value,
+                              reordered_W,
+                              static_cast<uint8_t>(0),
+                              worker_gemm_output,
+                              static_cast<size_t>(group_output_channels),
+                              static_cast<size_t>(output_count),
+                              static_cast<size_t>(kernel_size));
+          }
+        } else {
 #ifdef MLAS_SUPPORTS_PACKED_GEMM_U8X8
-        if (packed_W_buffer_) {
-          MlasGemm(static_cast<size_t>(output_count),
-                   static_cast<size_t>(group_output_channels),
-                   static_cast<size_t>(kernel_dim),
-                   worker_gemm_input,
-                   static_cast<size_t>(kernel_dim),
-                   X_zero_point_value,
-                   static_cast<const int8_t*>(packed_W_buffer_.get()) + group_id * packed_W_size_,
-                   0,
-                   true,
-                   worker_gemm_output,
-                   static_cast<size_t>(group_output_channels),
-                   nullptr);
-        } else
+          if (packed_W_buffer_) {
+            MlasGemm(static_cast<size_t>(output_count),
+                     static_cast<size_t>(group_output_channels),
+                     static_cast<size_t>(kernel_dim),
+                     worker_gemm_input,
+                     static_cast<size_t>(kernel_dim),
+                     X_zero_point_value,
+                     static_cast<const int8_t*>(packed_W_buffer_.get()) + group_id * packed_W_size_,
+                     0,
+                     is_W_signed,
+                     worker_gemm_output,
+                     static_cast<size_t>(group_output_channels),
+                     nullptr);
+          } else
 #endif
-        {
-          MlasGemm(static_cast<size_t>(output_count),
-                   static_cast<size_t>(group_output_channels),
-                   static_cast<size_t>(kernel_dim),
-                   worker_gemm_input,
-                   static_cast<size_t>(kernel_dim),
-                   X_zero_point_value,
-                   reordered_W + group_id * group_output_channels,
-                   static_cast<size_t>(M),
-                   0,
-                   true,
-                   worker_gemm_output,
-                   static_cast<size_t>(group_output_channels),
-                   nullptr);
+          {
+            MlasGemm(static_cast<size_t>(output_count),
+                     static_cast<size_t>(group_output_channels),
+                     static_cast<size_t>(kernel_dim),
+                     worker_gemm_input,
+                     static_cast<size_t>(kernel_dim),
+                     X_zero_point_value,
+                     reordered_W + group_id * group_output_channels,
+                     static_cast<size_t>(M),
+                     0,
+                     is_W_signed,
+                     worker_gemm_output,
+                     static_cast<size_t>(group_output_channels),
+                     nullptr);
+          }
         }
 
         if (output_scales.size() == 1) {

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<uint8_t, int8_t> 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<uint8_t, int8_t> 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