1 files changed, 16 insertions, 6 deletions
diff --git a/src/core/NEON/kernels/arm_conv/depthwise/depthwise_depthfirst_multiplier.hpp b/src/core/NEON/kernels/arm_conv/depthwise/depthwise_depthfirst_multiplier.hpp
index 3d305b6d18..b93caa2aaa 100644
--- a/src/core/NEON/kernels/arm_conv/depthwise/depthwise_depthfirst_multiplier.hpp
+++ b/src/core/NEON/kernels/arm_conv/depthwise/depthwise_depthfirst_multiplier.hpp
@@ -42,7 +42,7 @@ class DepthfirstMultiplierStrategy : public DepthwiseDepthfirstStrategyCommon<TI
   {
     return interleaves::PackingArguments(
       args.kernel_rows, args.kernel_cols, sizeof(TWeight),
-      true, sizeof(TAccum),
+      true, sizeof(TAccum), this->uses_premultiply(),
       this->get_vl_type(),
       sizeof(TAccum), 1,
       [args] (unsigned int pos, unsigned int &x, unsigned int &y) -> bool
@@ -57,6 +57,10 @@ class DepthfirstMultiplierStrategy : public DepthwiseDepthfirstStrategyCommon<TI
       }
     );
   }
+  
+  bool uses_premultiply() const override {
+    return false;
+  }
 
   public:
   using Parent::Parent;
@@ -192,7 +196,7 @@ class GenericDepthfirstMultiplierStrategy : public DepthwiseDepthfirstStrategyCo
   {
     return interleaves::PackingArguments(
       args.kernel_rows, args.kernel_cols, sizeof(TWeight),
-      false, sizeof(TAccum),
+      false, sizeof(TAccum), this->uses_premultiply(),
       this->get_vl_type(),
       sizeof(TAccum), 1,
       [args] (unsigned int pos, unsigned int &x, unsigned int &y) -> bool
@@ -207,6 +211,10 @@ class GenericDepthfirstMultiplierStrategy : public DepthwiseDepthfirstStrategyCo
       }
     );
   }
+  
+  bool uses_premultiply() const override {
+    return false;
+  }
 
   public:
   GenericDepthfirstMultiplierStrategy(KernelStrategyType *kern, const DepthwiseArgs &args)
@@ -483,6 +491,10 @@ class DepthwiseDepthfirstMultiplier : public DepthfirstDriver<TInput, TWeight, T
   OutputStage m_os;  // Copy of the output parameters
   const void *m_bias = nullptr;  // Copy of the bias (should we need it)
 
+  bool uses_premultiply() const override {
+    return false;
+  }
+
   public:
   DepthwiseDepthfirstMultiplier(StratType *const strat, const DepthwiseArgs &args, const OutputStage &os = {})
   : DepthfirstDriver<TInput, TWeight, TOutput>(strat, args), m_os(os)
@@ -506,17 +518,15 @@ class DepthwiseDepthfirstMultiplier : public DepthfirstDriver<TInput, TWeight, T
     depthwise_depthfirst::stash_bias(m_os, biases);
   }
 
-  size_t get_working_size_per_thread(const unsigned int n_input_channels) const override
+  size_t get_working_size_per_thread() const override
   {
     DepthwiseArgs args(this->m_args);
-    args.input_channels = n_input_channels;
     return WorkspaceManager::get_sizeof_workspace(WorkspaceArgs<IDepthfirstStrategy, OutputStage>(this->m_strat.get(), args, m_os));
   }
 
-  void initialise_working_space(void *buffer, unsigned int n_input_channels) const override
+  void initialise_working_space(void *buffer) const override
   {
     DepthwiseArgs args(this->m_args);
-    args.input_channels = n_input_channels;
     return WorkspaceManager::initialise(buffer, WorkspaceArgs<IDepthfirstStrategy, OutputStage>(this->m_strat.get(), args, m_os));
   }