[MLAS] Update the NHWC sans transposes path to also support Depthwise convolutions

onnxruntime/core/mlas/lib/convolve.cpp

@@ -1377,7 +1377,7 @@ MlasConvSupportsSymmetricChannelsLast2DFloatKernel(
         return false;
     }
 
-    if (Dimensions != 2 || BatchCount != 1 || GroupCount != 1 || Beta != 0.0f) {
+    if (Dimensions != 2 || BatchCount != 1 || Beta != 0.0f) {
         return false;
     }
 
@@ -1395,7 +1395,16 @@ MlasConvSupportsSymmetricChannelsLast2DFloatKernel(
         return false;
     }
 
-    if (FilterCount <= 1 || KernelShape[0] < 3 || KernelShape[1] < 3) {
+    const bool is_depthwise = GroupCount > 1;
+    if (is_depthwise) {
+        if (FilterCount != 1) {
+            return false;
+        }
+    } else if (FilterCount <= 1) {
+        return false;
+    }
+
+    if (KernelShape[0] < 3 || KernelShape[1] < 3) {
         return false;
     }
 

onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp

@@ -7,6 +7,7 @@
 #include <cassert>
 #include <algorithm>
 #include <cstddef>
+#include <cstring>
 #include <functional>
 #include <array>
 #include <vector>
@@ -92,6 +93,15 @@ static constexpr size_t ComputeConvOutSize(const size_t L, const size_t K, const
     return 0;
 }
 
+static inline void CopyChannelBlock(float* dst, const float* src, size_t channels) {
+    if (channels == 1) {
+        *dst = *src;
+        return;
+    }
+
+    std::memcpy(dst, src, channels * sizeof(float));
+}
+
 static size_t ComputeMlasWorkingBufferSize(const size_t co,
                                            const size_t ih, const size_t iw,
                                            const size_t kh, const size_t kw,
@@ -109,6 +119,12 @@ static size_t ComputeMlasWorkingBufferSize(const size_t co,
 }
 
 static bool CheckCapabilitiesSme(const MLAS_CONV_PARAMETERS* Parameters) {
+    // Grouped support in this override is only implemented for channels-last
+    // layout. The generic grouped path still assumes contiguous per-group CHW.
+    if (Parameters->GroupCount > 1 && !Parameters->ChannelsLast) {
+        return false;
+    }
+
     if (!MlasConvSupportsSymmetricChannelsLast2DFloatKernel(
             Parameters->Dimensions,
             Parameters->BatchCount,
@@ -550,15 +566,36 @@ static void ConvolveSme(const size_t co, //channels out
     dim[1] = MlasDivRoundup(m, m_step);
     dim[2] = MlasDivRoundup(co, n_step);
 
+    const bool grouped_channels_last = input_is_channels_last && groups > 1;
+    const size_t input_channels_total = ci * groups;
+    const size_t output_channels_total = co * groups;
+    const float* input_base = in;
+    float* output_base = out;
+    std::vector<float> input_group_buffer;
+    if (grouped_channels_last) {
+        input_group_buffer.resize(ih * iw * ci);
+    }
+
     for (size_t g = 0; g < groups; ++g) {
+        const float* input_group = in;
+        if (grouped_channels_last) {
+            for (size_t pixel = 0; pixel < ih * iw; ++pixel) {
+                const float* src = input_base + pixel * input_channels_total + g * ci;
+                CopyChannelBlock(input_group_buffer.data() + pixel * ci, src, ci);
+            }
+            input_group = input_group_buffer.data();
+        }
 
         auto result = out;
         const bool need_transpose = (!input_is_channels_last) && (co > 1);
-        if (need_transpose) {
+        if (need_transpose || grouped_channels_last) {
             result = tmp_mlas_aligned;
         }
 
-        auto lhs = LhsPackImageDataSme(ci, ih, iw, d_kh, d_kw, sh, sw, padding, in, input_is_channels_last, ThreadPool);
+        auto lhs = LhsPackImageDataSme(ci, ih, iw, d_kh, d_kw, sh, sw, padding,
+                                      input_group,
+                                      input_is_channels_last,
+                                      ThreadPool);
         const std::byte* rhs_data = packed_rhs ? packed_rhs + g * packed_rhs_group_stride : nullptr;
         std::unique_ptr<std::byte[]> rhs_storage;
         if (rhs_data == nullptr) {
@@ -613,13 +650,23 @@ static void ConvolveSme(const size_t co, //channels out
             );
         });
 
+        if (grouped_channels_last) {
+            for (size_t pixel = 0; pixel < m; ++pixel) {
+                float* dst = output_base + pixel * output_channels_total + g * co;
+                const float* src = result + pixel * co;
+                CopyChannelBlock(dst, src, co);
+            }
+        }
+
         if (need_transpose) {
             //Note: this could be absorbed into post conv activation
             MlasTranspose(tmp_mlas_aligned, out, m, co, ThreadPool);
         }
 
-        in += ci * ih * iw;
-        out += m * co;
+        if (!grouped_channels_last) {
+            in += ci * ih * iw;
+            out += m * co;
+        }
         weights += co * ci * kh * kw;
         if(bias){
             bias += co;
@@ -739,7 +786,10 @@ ArmKleidiAI::MlasConv(
                 Parameters->PackedFilterGroupStride,
                 Input, Output, WorkingBuffer, Parameters->ChannelsLast, ThreadPool);
 
-    MlasActivation(Parameters->Activation, Output, nullptr, Parameters->FilterCount, Parameters->OutputSize,
-                   Parameters->OutputSize);
+    const bool grouped_channels_last = Parameters->ChannelsLast && Parameters->GroupCount > 1;
+    const size_t activation_rows = grouped_channels_last ? Parameters->OutputSize : Parameters->FilterCount;
+    const size_t activation_cols =
+        grouped_channels_last ? Parameters->GroupCount * Parameters->FilterCount : Parameters->OutputSize;
+    MlasActivation(Parameters->Activation, Output, nullptr, activation_rows, activation_cols, activation_cols);
     return true;
 }

onnxruntime/core/optimizer/nhwc_transformer.cc

@@ -5,6 +5,7 @@
 #include <array>
 #include <cstdint>
 #include <deque>
+#include <limits>
 #include <vector>
 #include "core/common/cpuid_info.h"
 #include "core/graph/constants.h"
@@ -191,27 +192,37 @@ bool FloatNhwcWrapperFilter(const onnx_transpose_optimization::api::GraphRef& gr
   }
 
   const auto group = node.GetAttributeInt("group").value_or(1);
-  if (group != 1) {
+  if (group <= 0) {
     return false;
   }
+  constexpr uint64_t kSizeTMax = static_cast<uint64_t>(std::numeric_limits<size_t>::max());
+  if (static_cast<uint64_t>(group) > kSizeTMax) {
+    return false;
+  }
+  const auto group_count = narrow<size_t>(group);
 
   std::array<size_t, 2> input_spatial_shape{};
   std::array<size_t, 2> kernel_spatial_shape{};
   std::array<size_t, 2> dilations{1, 1};
   std::array<size_t, 2> strides{1, 1};
   std::array<size_t, 4> pads{};
   size_t batch_count = 0;
-  size_t filter_count = 0;
+  size_t total_filter_count = 0;
 
   if (!TryGetDimValueAsSizeT(*input_shape, 0, batch_count) ||
       !TryGetDimValueAsSizeT(*input_shape, 2, input_spatial_shape[0]) ||
       !TryGetDimValueAsSizeT(*input_shape, 3, input_spatial_shape[1]) ||
-      !TryGetDimValueAsSizeT(*weight_shape, 0, filter_count) ||
+      !TryGetDimValueAsSizeT(*weight_shape, 0, total_filter_count) ||
       !TryGetDimValueAsSizeT(*weight_shape, 2, kernel_spatial_shape[0]) ||
       !TryGetDimValueAsSizeT(*weight_shape, 3, kernel_spatial_shape[1])) {
     return false;
   }
 
+  if (total_filter_count == 0 || total_filter_count % group_count != 0) {
+    return false;
+  }
+  const size_t filter_count = total_filter_count / group_count;
+
   const auto dilations_opt = node.GetAttributeInts("dilations");
   if (dilations_opt.has_value() && !TryReadPositiveInts(*dilations_opt, dilations)) {
     return false;
@@ -229,7 +240,7 @@ bool FloatNhwcWrapperFilter(const onnx_transpose_optimization::api::GraphRef& gr
   return MlasConvSupportsSymmetricChannelsLast2DFloatKernel(
       /*Dimensions*/ 2,
       batch_count,
-      /*GroupCount*/ 1,
+      group_count,
       input_spatial_shape.data(),
       kernel_spatial_shape.data(),
       dilations.data(),

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

@@ -23,7 +23,7 @@
 #include "core/common/safeint.h"
 #include "core/util/math_cpuonly.h"
 
-#if defined(USE_KLEIDIAI) && defined(__aarch64__) && defined(__linux__)
+#if defined(USE_KLEIDIAI) && defined(MLAS_TARGET_ARM64)
 #include "core/mlas/lib/kleidiai/mlasi_kleidiai.h"
 #endif
 
@@ -191,7 +191,7 @@ Status Conv<T>::Compute(OpKernelContext* context) const {
   return Status::OK();
 }
 
-#if defined(USE_KLEIDIAI) && defined(__aarch64__) && defined(__linux__)
+#if defined(USE_KLEIDIAI) && defined(MLAS_TARGET_ARM64)
 Status Conv<float>::EnsurePackedChannelsLastFilter(concurrency::ThreadPool* thread_pool,
                                                    size_t filter_count_per_group,
                                                    size_t input_channels_per_group,
@@ -329,7 +329,7 @@ Status Conv<float>::Compute(OpKernelContext* context) const {
           narrow<size_t>(M / conv_attrs_.group),
           /*Beta*/ 0.0f);
 
-#if defined(USE_KLEIDIAI) && defined(__aarch64__) && defined(__linux__)
+#if defined(USE_KLEIDIAI) && defined(MLAS_TARGET_ARM64)
   if (nhwc_fastpath && can_cache_packed_filter_) {
     ORT_RETURN_IF_ERROR(EnsurePackedChannelsLastFilter(thread_pool,
                                                        narrow<size_t>(M / conv_attrs_.group),
@@ -385,7 +385,7 @@ Status Conv<float>::Compute(OpKernelContext* context) const {
                     nhwc_fastpath ? 0.0f : Beta,
                     thread_pool);
 
-#if defined(USE_KLEIDIAI) && defined(__aarch64__) && defined(__linux__)
+#if defined(USE_KLEIDIAI) && defined(MLAS_TARGET_ARM64)
     if (nhwc_fastpath && packed_filter_ != nullptr) {
       Parameters.FilterIsPacked = true;
       Parameters.PackedFilter = packed_filter_.get();

onnxruntime/core/providers/cpu/nn/conv.h

@@ -31,7 +31,7 @@ class Conv<float> : public OpKernel {
     activation_.ActivationKind = MlasIdentityActivation;
     SetupMlasBackendKernelSelectorFromConfigOptions(mlas_backend_kernel_selector_config_, info.GetConfigOptions());
 
-#if defined(USE_KLEIDIAI) && defined(__aarch64__) && defined(__linux__)
+#if defined(USE_KLEIDIAI) && defined(MLAS_TARGET_ARM64)
     if (channels_last_) {
       const auto& input_defs = info.node().InputDefs();
       const bool has_bias_input = input_defs.size() >= 3 && input_defs[2] != nullptr;
@@ -56,7 +56,7 @@ class Conv<float> : public OpKernel {
   ConvAttributes conv_attrs_;
   bool channels_last_{false};
 
-#if defined(USE_KLEIDIAI) && defined(__aarch64__) && defined(__linux__)
+#if defined(USE_KLEIDIAI) && defined(MLAS_TARGET_ARM64)
  private:
   Status EnsurePackedChannelsLastFilter(concurrency::ThreadPool* thread_pool,
                                         size_t filter_count_per_group,