[MLAS][KleidiAI] Integrate Depthwise 3x3 Conv Kernel

[JonathanC-ARM]

Description

This PR adds a Arm® KleidiAI™ backed f32 depthwise convolution path in MLAS and wires it into the CPU/KleidiAI convolution routing flow.
The new path targets eligible depthwise convolutions using the KleidiAI SME2 depthwise ukernel. It is integrated as a first-class KleidiAI convolution route so depthwise dispatch is handled consistently with the existing IGEMM / SGEMM fallback routing.

Changes

• Added a new KleidiAI f32 depthwise convolution implementation:
  • onnxruntime/core/mlas/lib/kleidiai/dw_conv_kleidiai.cpp
• Added depthwise ukernel plumbing through the KleidiAI ukernel interface.
• Added CMake integration for the new depthwise source.
• Added support checks for the new depthwise path.
• Integrated depthwise selection into the KleidiAI conv routing model via ConvRoute::Depthwise.
• Updated KleidiAI conv prepare/dispatch to route through:
  • Depthwise
  • IGemm
  • SGemmFallback
  • NoKleidiAi
• Updated NHWC transform gating so f32 Conv/FusedConv can be transformed when either dense or depthwise KleidiAI channels-last support is available.
• Added MLAS unit test coverage for the KleidiAI depthwise convolution path.

Supported depthwise shape envelope

The current KleidiAI depthwise route supports f32 depthwise convolution with:

• 2D convolution
• batch size 1
• depthwise multiplier 1
• 3x3 kernel
• stride 1
• dilation 1
• all-zero padding or all-one padding
• SME2 availability
• sufficient output dimensions for the KleidiAI tile shape

The implementation handles NHWC directly and can fall back to internal NCHW to NHWC conversion when needed.

Example Performance Results

Model	No KleidiAI	KleidiAI no DW	KleidiAI + DW	KleidiAI + DW vs No KleidiAI	KleidiAI + DW vs KleidiAI no DW	Latency reduction vs no DW
mobilenet_v1_f32.onnx	12.778 ms	4.567 ms	3.996 ms	3.20x faster	1.14x faster	12.5%
mobilenetv1_ssd_f32.onnx	29.348 ms	9.326 ms	8.538 ms	3.44x faster	1.09x faster	8.5%
retinaface_f32.onnx	70.761 ms	41.540 ms	37.511 ms	1.89x faster	1.11x faster	9.7%
deeplabv3_mobilenetv2_f32.onnx	119.709 ms	49.899 ms	49.398 ms	2.42x faster	1.01x faster	1.0%
de_efficientnetlitev3_f32.onnx	134.750 ms	40.252 ms	38.543 ms	3.50x faster	1.04x faster	4.2%

[Copilot]

Pull request overview

This PR integrates an Arm® KleidiAI™ SME2-backed FP32 depthwise 3x3 convolution path into MLAS, wires it into the existing KleidiAI CPU convolution routing (Depthwise → IGemm → SGemmFallback → NoKleidiAi), and enables NHWC graph transformation when the new depthwise channels-last capability is available. It also adds an MLAS unit test to validate both the direct KleidiAI depthwise entrypoint and end-to-end routing through MlasConv in channels-last mode.

Changes:

• Added a new KleidiAI FP32 depthwise convolution implementation and ukernel selection plumbing.
• Integrated a new ConvRoute::Depthwise into KleidiAI conv prepare/dispatch and enabled NHWC transform gating for supported depthwise shapes.
• Added MLAS unit tests covering direct depthwise invocation and MlasConv channels-last routing.

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
onnxruntime/test/mlas/unittest/test_dwconv_kleidiai.cpp	Adds unit tests for KleidiAI depthwise conv (direct call and MlasConv NHWC path).
onnxruntime/core/optimizer/nhwc_transformer.cc	Extends NHWC transform filtering to consider depthwise channels-last support constraints.
onnxruntime/core/mlas/lib/kleidiai/mlasi_kleidiai.h	Adds depthwise support/entrypoint declarations and introduces ConvRoute::Depthwise selection.
onnxruntime/core/mlas/lib/kleidiai/dw_conv_kleidiai.cpp	Implements the FP32 depthwise 3x3 kernel path, including layout conversions and packing.
onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp	Wires depthwise routing into KleidiAI conv prepare/dispatch and refactors IGEMM capability checks.
onnxruntime/core/mlas/lib/kai_ukernel_interface.h	Generalizes the ukernel wrapper type and adds a depthwise ukernel wrapper typedef and accessor.
onnxruntime/core/mlas/lib/kai_ukernel_interface.cpp	Registers the SME2 depthwise ukernel and exposes it via GetKleidiAIDepthwiseConvUKernel().
onnxruntime/core/mlas/lib/convolve.cpp	Enables MlasConvSupportsDepthwiseChannelsLast2DFloatKernel for KleidiAI-supported shapes to unlock NHWC transforms.
cmake/onnxruntime_mlas.cmake	Adds the new depthwise source file to the KleidiAI MLAS build.

[hariharans29]

Review — [MLAS][KleidiAI] Integrate Depthwise 3x3 Conv Kernel

What this PR does

Adds a KleidiAI SME2-backed FP32 depthwise 3×3 stride-1 conv path and wires it into the existing KleidiAI conv routing as a new first-class route (ConvRoute::Depthwise, ordered before IGemm / SGemmFallback / NoKleidiAi). Single commit 9a52413, +824 / −67 across 9 files.

Touched:

• New dw_conv_kleidiai.cpp (333 LOC) — the SME2 ukernel driver, weight repack, optional NCHW⇄NHWC staging, row-tile loop.
• convolve.cpp — MlasConvSupportsDepthwiseChannelsLast2DFloatKernel now actually returns true for the supported shape envelope (was hardcoded false with a TODO).
• convolve_kleidiai.cpp — refactors CheckCapabilitiesSme → CheckIgemmRouteCapabilities; moves SelectConvRoute up to MlasConvPrepare / MlasConv so the route decision is computed once.
• kai_ukernel_interface.cpp, kai_ukernel_interface.h — generalizes KaiMatmulKernel<T> to KaiKernel<T> (matmul becomes a thin alias), exposes KaiF32DepthwiseConvKernel + GetKleidiAIDepthwiseConvUKernel().
• mlasi_kleidiai.h — new ConvRoute::Depthwise enum + public DepthwiseConvKleidiAI / DepthwiseConvKleidiAISupported declarations.
• nhwc_transformer.cc — hoists input_channels_per_group extraction so depthwise gating sees it.
• New test_dwconv_kleidiai.cpp (305 LOC) — 5 tests covering NCHW/NHWC × zero/unit padding + end-to-end MlasConv channels-last.

Performance numbers (depthwise PR delta column): mobilenet_v1 +14%, mobilenetv1_ssd +9%, retinaface +11%, deeplab v3 +1%, efficientnet-lite +4%. Bigger on MobileNet-family (depthwise-heavy) which is exactly the expected shape.

What I like

• Routing structure is clean. SelectConvRoute returns one of four states; MlasConv dispatches with no surprises. The Depthwise branch returns true on success, false (→ generic MLAS fallback) on shape mismatch. Prior IGEMM/SGEMM paths preserved verbatim.
• CheckCapabilitiesSme → CheckIgemmRouteCapabilities rename is correct. The old function was overloaded — it computed the route AND validated IGEMM. The split makes the routing decision a single explicit call and lets each route validate its own constraints.
• Beta != 0.0f is rejected at the support check. The ukernel doesn't accumulate, so this is necessary.
• Padding gate is correct — the kernel supports zero-padding or unit-padding only, and the gate enforces uniform 4-corner padding (pads = [0,0,0,0] or [1,1,1,1]). Asymmetric or larger padding correctly falls through.
• The InputChannels == 1 && FilterCount == 1 route gate is right. That set covers true depthwise (per-group depthwise with multiplier 1) and the degenerate single-channel dense conv, which collapses to the same math. Wider channel-multiplier cases correctly skip the route.
• Unit tests parametrize NCHW vs NHWC and exercise both DepthwiseConvKleidiAI directly and end-to-end via MlasConv.
• MLAS backend kernel selector is honored end-to-end. Verified by walking both edges (see item 4 below).

Concerns (please address before merge)

1. Bounded TLS retention is missing. dw_conv_kleidiai.cpp:30 declares:

   thread_local DwconvTlsBuffers g_dwconv_tls;

   …with no shrink-on-scope-exit. PR #28786 explicitly established the MlasShrinkKleidiAIScratchIfTooLarge (8 MiB cap) + ScopedKaiHalfTlsCleanup pattern in mlasi_kleidiai.h precisely to avoid unbounded TLS growth from outlier-shape calls. The hazard here is real: feature_map_nhwc and nhwc_out for the NCHW fallback path can be batches × C × H × W × 4 B, easily tens to hundreds of MB per worker thread. Please use the same scoped-cleanup pattern.

2. dwconv global reference outside any namespace. At file scope, before namespace ArmKleidiAI:

   const KaiF32DepthwiseConvKernel& dwconv = GetKleidiAIDepthwiseConvUKernel();

   This is global-namespace pollution from a TU that already has namespace ArmKleidiAI. Move inside namespace ArmKleidiAI (or an anonymous detail namespace). It's only used internally.

3. DepthwiseConvKleidiAI is not self-contained — Copilot's existing comment is correct. The function is hard-wired to the SME2 3×3 s1 ukernel but accepts arbitrary filter_height / filter_width / pad_* / stride from the caller. If a future caller invokes it without first going through DepthwiseConvKleidiAISupported or SelectConvRoute, it will pass through bogus shapes and either silently produce wrong results or crash inside the ukernel. Add the same 3×3 / s1 / uniform-0-or-1-padding gate inside DepthwiseConvKleidiAI so the function fails closed. Trivial defense-in-depth, removes a real foot-gun.

4. Static-gate vs runtime-gate hardcoding divergence. Three places encode the tile dimensions:

   • convolve.cpp — hardcoded kKleidiAIDepthwiseRowsPerTile = 4, kKleidiAIDepthwiseColsPerTile = 4.
   • dw_conv_kleidiai.cpp runtime support check — dwconv.ukernel.get_m_step() for rows, hardcoded kDwconvColsPerTile = 4 for cols.
   • The kernel name …3x3_s1_4xc_sme2_mla implies m_step == 4 and 4×c tile, so they currently agree.

   If the ukernel selection ever swaps for a wider variant (e.g. 8xc), m_step changes underneath the runtime check but the static check still says 4. The NHWC transformer accepts the conv based on the static gate, but the runtime kernel rejects it, and the dispatcher falls back to generic MLAS — silent regression. Pull both tile constants from a single source (expose them via the wrapper, or define both in mlasi_kleidiai.h). Worth doing while the only caller is fresh.

On the MLAS backend kernel selector integration

I walked both edges of the selector path because this PR introduces a brand-new MLAS route that has to honor the use_kleidiai opt-out. Both edges are wired correctly — flagging it here as positive context, not as a concern:

• Runtime path: conv.cc:382 populates Parameters.BackendKernelSelectorConfig = &mlas_backend_kernel_selector_config_;. ArmKleidiAI::MlasConvPrepare calls SelectConvRoute, which routes through DepthwiseConvKleidiAISupported, which on dw_conv_kleidiai.cpp:138-140 checks Parameters->BackendKernelSelectorConfig->use_kleidiai and returns false on opt-out. If a user passes use_kleidiai=false, the depthwise route is declined at Prepare time and the dispatcher falls back to generic MLAS.
• Transformer path: nhwc_transformer.cc:404 gates the entire USE_KLEIDIAI F32 Conv→NHWC FusedConv registration block with if (mlas_backend_kernel_selector_config.use_kleidiai) { ... }. The FloatNhwcWrapperFilter that calls the new MlasConvSupportsDepthwiseChannelsLast2DFloatKernel only registers when KleidiAI is enabled, so the static-gate decision is also opt-out-aware.

Defense-in-depth nit (not a blocker): MlasConvSupportsDepthwiseChannelsLast2DFloatKernel doesn't take a selector parameter, so its safety depends on every future caller wrapping it in an if (use_kleidiai) guard like the NHWC transformer does today. The new depthwise helpers (DepthwiseConvKleidiAISupported / DepthwiseConvKleidiAI) read BackendKernelSelectorConfig directly from MLAS_CONV_PARAMETERS, which is the safer pattern. Consider giving MlasConvSupportsDepthwiseChannelsLast2DFloatKernel the same treatment — add a trailing const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* = nullptr parameter and short-circuit to false when cfg && !cfg->use_kleidiai. Follow-up, not a blocker for this PR.

Nits (do not block)

5. No threading in the row loop. DepthwiseConvKleidiAI doesn't even take a ThreadPool*. The SME2 ukernel parallelizes internally, and the benchmark numbers say single-threaded outer loop is fine for MobileNet-class workloads, but a // TODO: parallelize over (out_row tiles, channel groups) when worthwhile comment would document the deliberate choice for whoever profiles this later on a deeper network.

6. No prepacked-weights path. Every MlasConv call repacks (ConvertDepthwiseWeightsToHwcn + kai_run_rhs_dwconv_pack_*). For constant weights — the universal inference case — that's redundant work. The IGEMM path already plumbs through MLAS_CONV_PARAMETERS::PackedFilter. Worth wiring the depthwise path into the same prepack hook as a follow-up.

7. Test coverage is thin. 5 tests, all with 32 channels × 8×8 input. Worth adding:

   • Larger channel counts (≥ 128, plus one that's not a multiple of the SME2 vector length) to catch ukernel boundary bugs.
   • An output-too-small case (output H or W < 4) to confirm DepthwiseConvKleidiAISupported rejects it.
   • BatchCount > 1 to confirm the route is declined and the model still works.
   • One test that flips the selector (use_kleidiai = false in MLAS_CONV_PARAMETERS::BackendKernelSelectorConfig) and asserts MlasConvPrepare declines the depthwise route. Cheap to add and locks in the integration contract.

8. KaiMatmulKernel<T> → KaiKernel<T> alias. Clean refactor. Either add a one-line comment on the alias saying new code should use KaiKernel<T> directly, or drop the alias if no consumer depends on the old name. Minor.

9. convolve.cpp:1456-1459 override null-check is correct (if (Override == nullptr) return false;) but it's the one place that ties the static-gate decision to dispatch-table state. A one-line comment about that intent would help.

Bottom line

Approve once items 1–3 are addressed:

1. Use the existing MlasShrinkKleidiAIScratchIfTooLarge pattern for the dwconv TLS buffers — don't reinvent; the helper is already in mlasi_kleidiai.h from PR #28786.
2. Move the dwconv global into namespace ArmKleidiAI.
3. Add the 3×3 / s1 / dilation-1 / uniform-padding gate inside DepthwiseConvKleidiAI itself (Copilot's suggestion is exactly right).

Defer to follow-ups: tile-constant deduplication between static and runtime gates (item 4); threading of the row loop (item 5); prepacked weights via PackedFilter (item 6); wider test grid + selector-opt-out test (item 7); selector-as-parameter on MlasConvSupportsDepthwiseChannelsLast2DFloatKernel (item 4 in the selector section).

The architectural design is sound — the new route slots cleanly into the existing routing model, the unit tests exercise both direct entry and end-to-end dispatch, the performance numbers are believable and on the right networks, and the selector integration is well-formed. The three things to fix before merge are all small.