Depthwise channel pre-multiplication

Resolves: COMPMID-6337
Change-Id: Ie9097b3f56e8071426c621386a5988bd7f7e8ef2
Signed-off-by: Michael Tyler <michael.tyler@arm.com>
Reviewed-on: https://review.mlplatform.org/c/ml/ComputeLibrary/+/9852
Tested-by: Arm Jenkins <bsgcomp@arm.com>
Reviewed-by: Viet-Hoa Do <viet-hoa.do@arm.com>
Benchmark: Arm Jenkins <bsgcomp@arm.com>

3 files changed, 14 insertions, 10 deletions

diff --git a/src/core/NEON/kernels/arm_conv/depthwise/interleaves/a64_s8q_3x3_dot.cpp b/src/core/NEON/kernels/arm_conv/depthwise/interleaves/a64_s8q_3x3_dot.cpp
index 5e4bf99120..3de4bdc1fb 100644
--- a/src/core/NEON/kernels/arm_conv/depthwise/interleaves/a64_s8q_3x3_dot.cpp
+++ b/src/core/NEON/kernels/arm_conv/depthwise/interleaves/a64_s8q_3x3_dot.cpp
@@ -42,7 +42,7 @@ size_t interleave_a64_s8q_3x3_dot::get_packed_size(const DepthwiseArgs &args)
 {
   // We store 7 vectors for every <vector_of_ints> of channels.
   const unsigned int n = arm_gemm::roundup(
-    arm_gemm::iceildiv((long unsigned int) args.input_channels,
+    arm_gemm::iceildiv((long unsigned int) args.input_channels * args.channel_multiplier,
                        get_vector_length<int32_t>(arm_gemm::VLType::None)), 4lu
   );
   return n * 7 * get_vector_length<int8_t>(arm_gemm::VLType::None);
diff --git a/src/core/NEON/kernels/arm_conv/depthwise/interleaves/generic.cpp b/src/core/NEON/kernels/arm_conv/depthwise/interleaves/generic.cpp
index 056f08d037..dc505a013d 100644
--- a/src/core/NEON/kernels/arm_conv/depthwise/interleaves/generic.cpp
+++ b/src/core/NEON/kernels/arm_conv/depthwise/interleaves/generic.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022 Arm Limited.
+ * Copyright (c) 2022-2023 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -32,11 +32,11 @@ namespace interleaves {
 
 PackingArguments::PackingArguments(
   unsigned int kernel_rows, unsigned int kernel_cols, size_t weight_element_size,
-  bool include_bias, size_t bias_element_size,
+  bool include_bias, size_t bias_element_size, bool premultiply,
   arm_gemm::VLType vl_type, size_t accumulator_element_size, unsigned int accumulator_depth_vl,
   std::function<bool(unsigned int, unsigned int &, unsigned int &)> get_weight_pos
 ) : kernel_rows(kernel_rows), kernel_cols(kernel_cols), weight_element_size(weight_element_size),
-    include_bias(include_bias), bias_element_size(bias_element_size),
+    include_bias(include_bias), bias_element_size(bias_element_size), premultiply(premultiply),
     vl_type(vl_type), accumulator_element_size(accumulator_element_size), accumulator_depth_vl(accumulator_depth_vl),
     get_weight_pos(get_weight_pos)
 {
@@ -46,7 +46,7 @@ size_t get_storage_size_generic(const PackingArguments &packing_args, const Dept
 {
   // If the channel multiplier is greater than one, then we treat this as a
   // repeated packing of `channel_multiplier`-sized problems.
-  if (args.channel_multiplier > 1)
+  if (args.channel_multiplier > 1 && !packing_args.premultiply)
   {
     DepthwiseArgs args_per_input_channel(args);
     args_per_input_channel.input_channels = args.channel_multiplier;
@@ -58,7 +58,7 @@ size_t get_storage_size_generic(const PackingArguments &packing_args, const Dept
   const unsigned int vl =
     packing_args.accumulator_depth_vl *
     arm_gemm::utils::get_vector_length<uint8_t>(packing_args.vl_type) / packing_args.accumulator_element_size;
-  const unsigned int n_packs = arm_gemm::iceildiv(args.input_channels, vl);
+  const unsigned int n_packs = arm_gemm::iceildiv(args.input_channels * args.channel_multiplier, vl);
   const auto pack_size = (packing_args.include_bias ? packing_args.bias_element_size : 0) +
                          packing_args.kernel_points() * packing_args.weight_element_size;
   return n_packs * pack_size * vl;
@@ -81,7 +81,7 @@ void pack_parameters_generic(
 
   // If the channel multiplier is greater than one, then we treat this as a
   // repeated packing of `channel_multiplier`-sized problems.
-  if (args.channel_multiplier > 1)
+  if (args.channel_multiplier > 1 && !packing_args.premultiply)
   {
     // Get a modified copy of the depthwise arguments
     DepthwiseArgs args_per_input_channel(args);
@@ -107,17 +107,19 @@ void pack_parameters_generic(
     return;
   }
 
+  auto input_channels = args.input_channels * args.channel_multiplier;
+
   // Finalise the weight strides
-  ld_weight_col = (ld_weight_col == 0) ? args.input_channels : ld_weight_col;
+  ld_weight_col = (ld_weight_col == 0) ? input_channels : ld_weight_col;
   ld_weight_row = (ld_weight_row == 0) ? packing_args.kernel_cols * ld_weight_col : ld_weight_row;
 
   const unsigned int vl =
     packing_args.accumulator_depth_vl *
     arm_gemm::utils::get_vector_length<uint8_t>(packing_args.vl_type) / packing_args.accumulator_element_size;
 
-  for (unsigned int n = 0; n < args.input_channels; n += vl)
+  for (unsigned int n = 0; n < input_channels; n += vl)
   {
-    const unsigned int todo = std::min(vl, args.input_channels - n);
+    const unsigned int todo = std::min(vl, input_channels - n);
 
     if (packing_args.include_bias)
     {
diff --git a/src/core/NEON/kernels/arm_conv/depthwise/interleaves/generic.hpp b/src/core/NEON/kernels/arm_conv/depthwise/interleaves/generic.hpp
index 756c50b98c..1842f10150 100644
--- a/src/core/NEON/kernels/arm_conv/depthwise/interleaves/generic.hpp
+++ b/src/core/NEON/kernels/arm_conv/depthwise/interleaves/generic.hpp
@@ -40,6 +40,7 @@ struct PackingArguments
   const size_t weight_element_size;
   const bool include_bias;
   const size_t bias_element_size;
+  const bool premultiply;
   arm_gemm::VLType vl_type;
   const size_t accumulator_element_size;
   const unsigned int accumulator_depth_vl;
@@ -53,6 +54,7 @@ struct PackingArguments
     size_t weight_element_size,
     bool include_bias,
     size_t bias_element_size,
+    bool premultiply,
     arm_gemm::VLType vl_type,
     size_t accumulator_element_size,
     unsigned int accumulator_depth_vl,