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

[orlmon01]

Description

A path for MLAS to support NHWC Convolutions without the need for transposes was added in PR: #26834
This PR expands those changes to also support Depthwise Convolutions via the same pathway

What changed:

• The shared NHWC capability gate in onnxruntime/core/mlas/lib/convolve.cpp:1348 stopped requiring GroupCount == 1. It now allows GroupCount > 1 only when the op is true depthwise, meaning filters_per_group ==
  1.
• The NHWC transformer in onnxruntime/core/optimizer/nhwc_transformer.cc:162 was updated to pass the real group value and compute filter_count per group instead of hard-coding group 1. That is what lets grouped
  depthwise Conv/FusedConv nodes get rewritten to com.microsoft.NhwcFusedConv.
• The KleidiAI execution path in onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp:553 learned how to handle grouped NHWC tensors by:
  • gathering one group’s channels out of interleaved NHWC input into a temporary contiguous buffer,
  • running the existing per-group kernel,
  • scattering that group’s output channels back into interleaved NHWC output.
• Tests were added for a working NHWC depthwise case in onnxruntime/test/contrib_ops/fused_conv_test.cc:466, and transformer tests were updated to verify both the new positive case and the expected skip cases in
  onnxruntime/test/optimizer/nhwc_transformer_test.cc:416.

Added performance benchmark tests to allow for comparison between the new NHWC path and the old NCHW default.
Sample output:

-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
Benchmark                                                                                                                                     Time             CPU   Iterations
-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
SCONV_NCHW/KleidiAiNhwcComparison_NchwBaseline/Rank:2/N:1/G:1/Cpg:64/Fpg:64/I:56/56/K:3/3/P:1/1/1/1/S:1/1/D:1/1/real_time                508509 ns       508507 ns         1374
SCONV_NCHW/KleidiAiNhwcComparison_NchwBaseline/Rank:2/N:1/G:1/Cpg:128/Fpg:128/I:28/28/K:3/3/P:1/1/1/1/S:1/1/D:1/1/real_time              700573 ns       700386 ns          997
SCONV_NCHW/KleidiAiNhwcComparison_NchwBaseline/Rank:2/N:1/G:64/Cpg:1/Fpg:1/I:56/56/K:3/3/P:1/1/1/1/S:1/1/D:1/1/real_time                6471094 ns      6471114 ns          132
SCONV_NCHW/KleidiAiNhwcComparison_NchwBaseline/Rank:2/N:1/G:72/Cpg:1/Fpg:1/I:48/80/K:3/3/P:1/1/1/1/S:2/2/D:1/1/real_time                3768969 ns      3767797 ns          217
SCONV_NHWC_KLEIDIAI/KleidiAiNhwcComparison_NhwcFastPath/Rank:2/N:1/G:1/Cpg:64/Fpg:64/I:56/56/K:3/3/P:1/1/1/1/S:1/1/D:1/1/real_time       414198 ns       414197 ns         1688
SCONV_NHWC_KLEIDIAI/KleidiAiNhwcComparison_NhwcFastPath/Rank:2/N:1/G:1/Cpg:128/Fpg:128/I:28/28/K:3/3/P:1/1/1/1/S:1/1/D:1/1/real_time     652454 ns       652454 ns         1074
SCONV_NHWC_KLEIDIAI/KleidiAiNhwcComparison_NhwcFastPath/Rank:2/N:1/G:64/Cpg:1/Fpg:1/I:56/56/K:3/3/P:1/1/1/1/S:1/1/D:1/1/real_time       6032947 ns      6032940 ns          117
SCONV_NHWC_KLEIDIAI/KleidiAiNhwcComparison_NhwcFastPath/Rank:2/N:1/G:72/Cpg:1/Fpg:1/I:48/80/K:3/3/P:1/1/1/1/S:2/2/D:1/1/real_time       3022041 ns      3018352 ns          227

[orlmon01]

@microsoft-github-policy-service agree company="Arm"

[Copilot]

Pull request overview

Expands the existing MLAS/KleidiAI NHWC “no-transpose” convolution fast path to support true depthwise convolutions (grouped conv where filters-per-group == 1), and wires that capability through the NHWC transformer plus adds test/benchmark coverage.

Changes:

• Relax MLAS NHWC capability gating to allow GroupCount > 1 only for true depthwise (FilterCount-per-group == 1).
• Update NHWC transformer filtering to pass the real group count and compute per-group filter count.
• Extend KleidiAI NHWC execution to handle grouped NHWC tensors via per-group gather/compute/scatter, plus add unit tests and a benchmark comparison.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
onnxruntime/core/mlas/lib/convolve.cpp	Updates NHWC capability gate to allow depthwise grouped convs.
onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp	Implements grouped-NHWC handling by gathering/scattering channels per group.
onnxruntime/core/optimizer/nhwc_transformer.cc	Passes group count + per-group filter count into the NHWC fast-path capability check.
onnxruntime/core/providers/cpu/nn/conv.h	Broadens KleidiAI fast-path compilation guard to MLAS_TARGET_ARM64.
onnxruntime/core/providers/cpu/nn/conv.cc	Same guard update for KleidiAI fast-path code.
onnxruntime/test/optimizer/nhwc_transformer_test.cc	Adds/updates tests validating depthwise enablement and expected skip cases.
onnxruntime/test/contrib_ops/fused_conv_test.cc	Adds an NHWC depthwise FusedConv correctness test (conditionally enabled).
onnxruntime/test/mlas/bench/bench_sconv.cpp	Adds benchmark cases comparing NCHW baseline vs NHWC KleidiAI fast path, including depthwise shapes.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated 3 comments.

[hariharans29]

Review of PR #28565 — [MLAS] Update the NHWC sans transposes path to also support Depthwise convolutions

Overall direction is good — extending the existing NHWC no-transpose KleidiAI/SME path to depthwise is worthwhile, and the gate relaxation in MlasConvSupportsSymmetricChannelsLast2DFloatKernel plus the per-group gather/compute/scatter in ConvolveSme is the right shape of change. However, there is at least one real correctness bug for non-identity activations, and several smaller issues worth tightening before merge.

Likely bugs / correctness

1. (HIGH — Copilot was right) Activation is only applied to a single group's worth of output.
   ArmKleidiAI::MlasConv ends with:

   MlasActivation(Parameters->Activation, Output, nullptr,
                  Parameters->FilterCount, Parameters->OutputSize,
                  Parameters->OutputSize);

   With GroupCount > 1, the real output is GroupCount * FilterCount * OutputSize floats laid out as NHWC interleaved channels. Parameters->FilterCount is per-group, so this call only touches FilterCount * OutputSize floats — and because the layout is interleaved ([H*W, G*F]), even those floats are not a coherent contiguous channel slab. For non-identity activations (Relu, Clip, Sigmoid, HardSigmoid, LeakyRelu, …) the result will be silently wrong for any value that should have been clipped/saturated but lives outside the first FilterCount * OutputSize elements.

   This is not caught by Cpu_NhwcDepthwiseConv2D_SymmetricPadding because the test uses Relu on a tensor whose true output is already entirely non-negative (X and W are non-negative, so all sums are ≥ 0). Please add at least one depthwise test that produces negative pre-activation values (e.g. signed weights or a Clip/LeakyRelu activation) so this regresses if the fix is removed.

   Fix sketch: scale the M dim of the activation call by GroupCount, or (cleaner) apply activation per-group inside ConvolveSme on the scattered NHWC output using the appropriate M = OutputSize, N = GroupCount * FilterCount view.

2. NCHW + grouped is now silently accepted by the capability gate but the execution path was not adapted/tested.
   MlasConvSupportsSymmetricChannelsLast2DFloatKernel no longer requires GroupCount == 1, and it does not key off ChannelsLast. In practice today, Conv<float>::Compute only routes to the KleidiAI fast path when channels_last_ is true, so this is probably benign — but the gate is named "ChannelsLast" yet now accepts grouped configurations regardless of layout. Either:

   • tighten the gate to require channels-last for grouped (GroupCount > 1 ⇒ ChannelsLast), since the per-group gather/scatter in ConvolveSme is conditioned on input_is_channels_last && groups > 1, or
   • explicitly document why the NCHW grouped case is still safe (the unmodified in += ci*ih*iw; out += m*co per-group stepping, with result = tmp_mlas_aligned only when need_transpose, has not been exercised by any new test in this PR).

3. Depthwise + kernel < 3x3 is now allowed by the relaxed branch in convolve.cpp:

   if (!is_depthwise && (FilterCount <= 1 || KernelShape[0] < 3 || KernelShape[1] < 3)) {
       return false;
   }

   The prior gate explicitly excluded < 3 kernels because the SME imatmul path was not validated there. For depthwise the gate is now bypassed entirely. Is depthwise 1x1 (or 1x3 / 3x1) actually correct and beneficial on this kernel? If not validated, please keep the kernel-size floor for depthwise too.

4. MlasConvSupportsSymmetricChannelsLast2DFloatKernel does not validate FilterCount for the depthwise branch. With is_depthwise == true the FilterCount <= 1 check is skipped — that is the intended semantics (FilterCount == 1 per group), but combined with the unsigned arithmetic in callers this means a bogus FilterCount == 0 could now pass the gate. Worth an explicit FilterCount == 1 assert/return for the depthwise branch instead of relying on the caller.

Performance

1. Dead-ish conditional in the in/out advancement in ConvolveSme:

   const bool use_temp_output = grouped_channels_last || need_transpose;
   ...
   if (!grouped_channels_last) {
       in  += ci * ih * iw;
       out += use_temp_output ? 0 : m * co;
       if (need_transpose) {
           out += m * co;
       }
   }

   Inside the !grouped_channels_last branch, use_temp_output == need_transpose, so the net effect is always out += m * co — same as the pre-PR code. The ternary is misleading. Either drop it or rewrite as out += m * co; plain.

2. Two consecutive identical writes to result:

   if (need_transpose) {
       result = tmp_mlas_aligned;
   }
   if (grouped_channels_last) {
       result = tmp_mlas_aligned;
   }

   Collapse to if (need_transpose || grouped_channels_last) result = tmp_mlas_aligned;.

3. Per-pixel std::copy_n(src, ci, ...) gather/scatter is single-threaded and degenerates badly for depthwise (ci == 1) — you get one function-call worth of work per pixel per group, run sequentially across all groups. For Mobile-style depthwise with groups == 64..128 and H*W in the thousands this is non-trivial. Worth:

   • using a manual strided loop or memcpy(dst, src, ci*sizeof(float)) (which the compiler will inline for ci == 1),
   • and/or wrapping the per-group loop in MlasTrySimpleParallel so the gather/compute/scatter pipeline is overlapped across groups.

   Your own benchmark already shows the depthwise wins are modest (6.47ms → 6.03ms and 3.77ms → 3.02ms) compared to the dense conv wins; the gather/scatter overhead is plausibly the reason.

4. input_group_buffer is sized to ih * iw * ci even when groups == 1 and the gather path is unused. It's now hoisted out of the loop (good), but the resize happens unconditionally inside the if (grouped_channels_last) so that's already fine — just confirm the buffer is not touched on the non-grouped path. (It isn't, looks fine.)

5. bench_sconv.cpp SCONV_NHWC_KLEIDIAI is ~140 lines of near-duplicate of SCONV_NCHW. Consider factoring shared shape/setup into a helper rather than copy-pasting; otherwise the next gating change has to be made in three places.

Style / hygiene

1. Two near-identical HasFloatNhwcNoTransposeSupport helpers now exist — one in test/contrib_ops/fused_conv_test.cc and one in test/optimizer/nhwc_transformer_test.cc. They differ only in signature and a couple of asserts. Please dedupe into a single header under test/util/include/ or test/mlas/. Both copies also have the same missing-upper-bound issue Copilot flagged on narrow<size_t>(group) — fix once, share.

2. input_channels_total / output_channels_total / input_base / output_base are computed unconditionally but only used under grouped_channels_last. Move them inside the if for clarity (and to avoid spurious ci * groups overflow concerns on contrived inputs).

3. #include <limits> added to nhwc_transformer.cc is correct; while you're there, the new kSizeTMax constant is fine but static_cast<uint64_t>(group) > std::numeric_limits<size_t>::max() would inline equally well without the named constant. Minor.

4. 8 commits with several merge-from-main commits — please squash on merge so git log shows one logical change.

5. Test name ConvDepthwiseFloat_UsesHelperCapability is vague — ConvDepthwiseFloat_RewrittenToNhwcFusedConv (or similar) reads better.

6. KleidiAI NHWC benchmark requires 2D convolution with batch size 1. — minor, but SkipWithError will mark the bench as failed in CI dashboards. state.SkipWithMessage(...) (benchmark ≥ 1.8) or just early-return without recording an error is friendlier for matrix benchmarks.

Things that look good

• Capability-gate relaxation is appropriately narrow (GroupCount > 1 && FilterCount == 1 only) — doesn't accidentally open the door for arbitrary grouped convs.
• NHWC transformer correctly recomputes filter_count = total_filter_count / group_count and passes through the real group_count. The size_t-overflow guard added in the latest commit is the right call.
• #if defined(USE_KLEIDIAI) && defined(MLAS_TARGET_ARM64) simplification (dropping the redundant __aarch64__ && __linux__) is a nice cleanup and subsumes Remove ifdef for linux targets - NHWC Transformer in Conv op #28564.
• Negative transformer tests (ConvDepthwiseFloat_AsymmetricPaddingSkipsNhwc, ConvGroupedFloat_SkipNhwc, ConvDepthwiseMultiplier2Float_SkipNhwc) cover the most important "don't rewrite" cases.
• Benchmark coverage for both the dense and depthwise envelope is a nice touch.

Summary

Required before merge:

• Fix the MlasActivation extent for GroupCount > 1 (item 1 under Bugs). Add a test that would catch a non-identity activation regression (e.g. Relu with signed weights, or Clip).

Strongly recommended:

• Decide whether to tighten the capability gate so grouped is only accepted with ChannelsLast (item 2), or add a test that actually exercises the NCHW+grouped path.
• Keep the KernelShape >= 3 floor for depthwise unless you've validated 1x1/1x3/3x1 (item 3).
• Collapse the dead use_temp_output / duplicate result = logic (Perf 1–2).
• Dedupe the two HasFloatNhwcNoTransposeSupport test helpers (Style 1).
• Squash the 8 commits on merge.

Nice-to-have / nits:

• Parallelize / memcpy-ify the per-group gather/scatter (Perf 3).
• Factor SCONV_NHWC_KLEIDIAI shared setup with SCONV_NCHW (Perf 5).
• Rename ConvDepthwiseFloat_UsesHelperCapability (Style 5).

[orlmon01]

Sorry, for the delay. Some of the refactors caused test errors on our CI. Should all be good now. :)

[hariharans29]

Re-review of PR #28565 — second pass

Verdict: Approve (pending green CI on the most-recent merge from main). Every required item from the prior review is addressed, the new commits read cleanly, and the previously-broken activation path is now exercised by a regression test.

---

Items from the prior review — status

Required-before-merge

1. Activation extent for GroupCount > 1 — FIXED.
   onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp (around line 789):

   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);

   Total element count activation_rows * activation_cols is now OutputSize * G * F for the grouped-NHWC case (vs. F * OutputSize before), which is the correct slab size. All supported MLAS_ACTIVATION kinds (Relu / LeakyRelu / Tanh / Logistic / Clip / HardSigmoid) are stateless elementwise, so the M × N interpretation is fine regardless of NHWC layout — only the touched element count matters.

   The new regression test Cpu_NhwcDepthwiseConv2D_Relu_NegativePreActivation uses signed weights to force negative pre-activations, so a future regression that re-shrinks the touched slab would now actually fail (the old Cpu_NhwcDepthwiseConv2D_SymmetricPadding couldn't catch it because all pre-activations were non-negative). Good test.

Strongly recommended

2. Tighten capability gate for grouped → ChannelsLast only — DONE.
   convolve_kleidiai.cpp::CheckCapabilitiesSme (around line 121-126):

   if (Parameters->GroupCount > 1 && !Parameters->ChannelsLast) {
     return false;
   }

   That's the cleanest of the two options I suggested — the gate now matches what ConvolveSme actually handles. Good.

3. Keep KernelShape >= 3 floor for depthwise — DONE.
   convolve.cpp (around line 1398-1408):

   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;

   The kernel-size floor now applies to both branches, and the explicit FilterCount == 1 check for the depthwise branch addresses my item Incremental updates. #4 (the depthwise branch no longer relies on the caller to validate filter count).

4. Collapse dead use_temp_output / duplicate result = — DONE.
   Single combined check:

   if (need_transpose || grouped_channels_last) {
     result = tmp_mlas_aligned;
   }

   And the pointer-advancement is gated on !grouped_channels_last instead of the misleading ternary.

Nice-to-have

5. Gather/scatter — partial. CopyChannelBlock replaces std::copy_n with memcpy and a channels == 1 fast path — that's the depthwise hot case and the inlined single-float copy avoids the call. Not parallelized across groups (no MlasTrySimpleParallel around the for (size_t g = 0; g < groups; ++g)), so the per-group serial work pattern remains. Fine for this PR; the perf wins on depthwise benchmarks (6.47ms → 6.03ms, 3.77ms → 3.02ms) are real either way.

6. Hoist input_group_buffer allocation — DONE (commit 2c43663).

7. narrow<size_t>(group) overflow guard — DONE in nhwc_transformer.cc:

   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);

   Note: the same guard is not added to the test helper HasFloatNhwcNoTransposeSupport or its sibling in fused_conv_test.cc — both only check group <= 0. In practice the tests pass literal small ints (≤ 8), so this is benign. Worth a one-line addition for consistency, but not blocking.

Not addressed (cosmetic only — no objection)

• The two HasFloatNhwcNoTransposeSupport helpers in fused_conv_test.cc and nhwc_transformer_test.cc are still duplicated. Author kept them separate (signatures differ slightly). Acceptable.
• Test name ConvDepthwiseFloat_UsesHelperCapability retained. Not a blocker.
• bench_sconv.cpp SCONV_NHWC_KLEIDIAI is still a near-duplicate of SCONV_NCHW. Benchmark scaffolding; fine.

---

New thing worth a sanity-check

The new CopyChannelBlock uses static inline at file scope — that's the expected idiom here. No concerns.

The input_channels_total = ci * groups and output_channels_total = co * groups arithmetic in ConvolveSme is unguarded size_t multiplication. Both factors come from a graph whose shapes have already passed narrow<size_t> and the ChannelsLast capability gate, so realistic shapes can't overflow size_t on 64-bit ARM64 (the only target this code compiles into). Fine.

---

Bottom line

Solid follow-up. The activation correctness bug is gone, the gate is tightened to match what the kernel actually handles, the kernel-floor and filter-count contracts are explicit, and the cleanups landed. Once the most-recent merge-from-main commit (44e844a) finishes its CI matrix green, this is ready to merge. Please squash the 15 commits on merge so git log shows one logical change.

[orlmon01]

Cheers Hari,
How do I squash the commits on merge on github. Is it something I have to do locally and push it or is it something that is done by whoever does the merging on github?

[hariharans29]

Cheers Hari, How do I squash the commits on merge on github. Is it something I have to do locally and push it or is it something that is done by whoever does the merging on github?

Oh I will just squash and merge. Don't worry - that comment was from my automated review tool :)

Thanks for this PR !