[MLAS][KleidiAI] Integrate Depthwise 3x3 Conv Kernel by JonathanC-ARM · Pull Request #29441 · microsoft/onnxruntime

Jonathan Clohessy (JonathanC-ARM) · 2026-06-30T15:32:37Z

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.

Hariharan Seshadri (hariharans29) · 2026-06-30T16:37:05Z

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)

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

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.
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.
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.
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.
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:

Use the existing MlasShrinkKleidiAIScratchIfTooLarge pattern for the dwconv TLS buffers — don't reinvent; the helper is already in mlasi_kleidiai.h from PR #28786.
Move the dwconv global into namespace ArmKleidiAI.
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.

Copilot

🟡 Changes recommended

The new route has disabled-backend correctness issues, regresses symbolic dense transformation, and leaves an SME2 optimizer test stale.

Get a fresh assessment by requesting another Copilot review.

Review details

Files reviewed: 11/11 changed files
Comments generated: 3
Review effort level: Balanced

Melike Kaptan (melkap01-Arm) · 2026-09-21T13:05:40Z

+  const bool depthwise_nhwc_fastpath =
+      wants_channels_last && !sum_present && MlasConvSupportsDepthwiseChannelsLast2DFloatKernel(kernel_rank, narrow<size_t>(N), group_count, input_channels_per_group, input_shape_size_t.data(), kernel_shape_size_t.data(), dilations_size_t.data(), pads_size_t.data(), strides_size_t.data(), filter_count_per_group,
+                                                                                                /*Beta*/ 0.0f);


Melike Kaptan (melkap01-Arm) · 2026-09-21T13:05:58Z

+        parameters.StrideShape[dim] = StrideShape[dim];
+    }
+
+    return ArmKleidiAI::DepthwiseConvKleidiAISupported(&parameters);


Melike Kaptan (melkap01-Arm) · 2026-09-21T13:06:16Z

      !TryGetDimValueAsSizeT(*input_shape, 2, input_spatial_shape[0]) ||
      !TryGetDimValueAsSizeT(*input_shape, 3, input_spatial_shape[1]) ||
      !TryGetDimValueAsSizeT(*weight_shape, 0, total_filter_count) ||
+      !TryGetDimValueAsSizeT(*weight_shape, 1, input_channels_per_group) ||


Melike Kaptan (melkap01-Arm) · 2026-09-18T17:14:33Z

After threading added , threaded vs non threaded version is benchmarked:

Apple M4 Pro Machine:

Model	Threads	Without Threading	With Threading	Change
MobileNet V1	1	3.960 [3.867–4.003]	3.978 [3.902–4.024]	−0.46%
MobileNet V1	2	3.086 [3.049–3.140]	3.026 [2.970–3.124]	+1.97%
MobileNet V1	3	2.748 [2.726–2.752]	2.685 [2.667–2.765]	+2.30%
MobileNet V1	4	2.512 [2.491–2.615]	2.459 [2.418–2.559]	+2.13%
MobileNet V1 SSD	1	8.313 [8.238–8.414]	8.397 [8.314–8.424]	−1.01%
MobileNet V1 SSD	2	6.613 [6.411–6.858]	6.368 [6.265–6.670]	+3.70%
MobileNet V1 SSD	3	5.878 [5.678–6.009]	5.583 [5.549–5.657]	+5.01%
MobileNet V1 SSD	4	5.214 [5.110–5.295]	5.013 [4.979–5.104]	+3.86%
RetinaFace	1	36.471 [35.542–36.848]	36.247 [35.592–36.400]	+0.62%
RetinaFace	2	25.003 [24.596–25.254]	24.931 [24.565–24.961]	+0.29%
RetinaFace	3	19.990 [19.881–20.589]	19.937 [19.646–20.484]	+0.26%
RetinaFace	4	17.896 [17.782–18.184]	17.673 [17.613–19.647]	+1.25%
DeepLabV3 MobileNetV2	1	48.116 [46.896–48.419]	48.580 [47.723–48.990]	−0.96%
DeepLabV3 MobileNetV2	2	33.948 [33.020–34.150]	34.066 [33.645–34.387]	−0.35%
DeepLabV3 MobileNetV2	3	28.773 [28.576–29.102]	28.966 [28.407–29.147]	−0.67%
DeepLabV3 MobileNetV2	4	25.947 [25.668–26.050]	25.922 [25.881–26.404]	+0.10%
EfficientNet Lite V3	1	38.013 [37.482–38.344]	38.026 [37.543–38.180]	−0.03%
EfficientNet Lite V3	2	23.461 [23.256–23.518]	23.474 [23.336–23.688]	−0.06%
EfficientNet Lite V3	3	24.854 [23.985–25.715]	24.267 [24.145–24.826]	+2.36%
EfficientNet Lite V3	4	21.034 [20.571–21.782]	21.583 [19.686–22.163]	−2.61%

Android X300 Device

Model	Threads	Before (ms)	Threaded (ms)	Average Δ	P50 Δ
EfficientNet-LiteV3	1	108.796	108.818	−0.02%	−0.41%
EfficientNet-LiteV3	2	81.794	83.102	−1.60%	−0.17%
EfficientNet-LiteV3	3	71.949	73.473	−2.12%	−1.42%
EfficientNet-LiteV3	4	69.659	69.824	−0.24%	−0.13%
DeepLabV3-MobileNetV2	1	133.556	132.837	+0.54%	+0.40%
DeepLabV3-MobileNetV2	2	104.554	101.806	+2.63%	+0.25%
DeepLabV3-MobileNetV2	3	84.752	84.999	−0.29%	−0.46%
DeepLabV3-MobileNetV2	4	74.067	73.730	+0.45%	−0.13%
MobileNetV1	1	13.380	13.459	−0.59%	−0.05%
MobileNetV1	2	10.043	10.683	−6.38%*	+1.93%
MobileNetV1	3	9.037	8.832	+2.27%	+2.13%
MobileNetV1	4	7.835	7.866	−0.39%	+0.52%
MobileNetV1-SSD	1	27.641	27.634	+0.03%	+0.77%
MobileNetV1-SSD	2	22.354	21.601	+3.37%*	+0.99%
MobileNetV1-SSD	3	19.654	19.677	−0.12%	+2.43%
MobileNetV1-SSD	4	16.843	16.751	+0.54%	−1.21%
RetinaFace	1	110.256	110.709	−0.41%	−0.79%
RetinaFace	2	81.428	80.714	+0.88%*	+0.70%
RetinaFace	3	67.029	68.508	−2.21%	−1.70%
RetinaFace	4	56.462	59.701	−5.74%	−6.21%

Melike Kaptan (melkap01-Arm) · 2026-10-02T10:32:17Z

All points in the review are addressed.

Copilot

Copilot review overview

🔵 Needs a closer look

The hardware-specific SME2 kernel, packed-data format, and vector-length-aware threading require final validation on supported Arm hardware.

Review effort: Balanced
Findings: 1 · 2

Open (3)

This marks the input as a native NHWC fast path even when session.disable_kleidiai=1. For a…
Requiring weight dimension 1 here regresses the existing dense NHWC route: a legal dense Conv with…
This newly makes the SME2-supported {1,8,7,7} / {8,1,3,3}, group-8 case return true, but…

Melike Kaptan (melkap01-Arm) · 2026-10-19T08:26:25Z

Copilot review overview

🔵 Needs a closer look

The hardware-specific SME2 kernel, packed-data format, and vector-length-aware threading require final validation on supported Arm hardware.

Review effort: Balanced Findings: 1 · 2

Open (3)

This marks the input as a native NHWC fast path even when session.disable_kleidiai=1. For a…

Requiring weight dimension 1 here regresses the existing dense NHWC route: a legal dense Conv with…

This newly makes the SME2-supported {1,8,7,7} / {8,1,3,3}, group-8 case return true, but…

Those points are already covered and fixed in the previous commits. I checked there is no fix required at the moment. I am assuming , answering the co pilot for the review but not closing the review item caused them to be listed again here. Appreciate if you see the above answers to the individual review items and run another CI.

Copilot

🟡 Changes recommended

Finite clamping changes convolution results, and the updated transformer test fails when NCHWc takes precedence.

5 open issues

This marks the input as a native NHWC fast path even when session.disable_kleidiai=1. For a…
Preserve infinite and overflow results in depthwise Conv · New
Disable NCHWc transformation in NHWC capability test · New
Requiring weight dimension 1 here regresses the existing dense NHWC route: a legal dense Conv with…
This newly makes the SME2-supported {1,8,7,7} / {8,1,3,3}, group-8 case return true, but…

_{Review effort: Balanced}

Melike Kaptan (melkap01-Arm) · 2026-10-27T15:09:48Z

+                                  -std::numeric_limits<float>::max(),
+                                  std::numeric_limits<float>::max(),


Melike Kaptan (melkap01-Arm) · 2026-10-27T15:10:52Z

+    const bool expect_nhwc = HasFloatNhwcNoTransposeSupport({1, 8, 7, 7}, {8, 1, 3, 3}, {}, {}, {}, 8);
+    EXPECT_EQ(op_to_count["Conv"] + op_to_count["com.microsoft.nchwc.Conv"], expect_nhwc ? 0 : 1);
+    EXPECT_EQ(op_to_count["com.microsoft.NhwcFusedConv"], expect_nhwc ? 1 : 0);
+    EXPECT_EQ(op_to_count["Transpose"], expect_nhwc ? 2 : 0);


Copilot

🟡 Changes recommended

The selected KleidiAI kernel converts a quiet-NaN bias to negative infinity in eligible depthwise convolutions, requiring a NaN-preserving path and regression coverage.

6 open issues

_{Review effort: Balanced}

Melike Kaptan (melkap01-Arm) · 2026-10-29T14:10:01Z

+                                  -std::numeric_limits<float>::infinity(),
+                                  std::numeric_limits<float>::infinity(),


fallback provided in the last commit. Done

Copilot

🟡 Changes recommended

The fallback still permits incorrect NaN handling for biasless inputs, weights, and invalid arithmetic, while the threaded test unconditionally fails in standalone MLAS builds without a thread pool.

8 open issues

_{Review effort: Balanced}

Melike Kaptan (melkap01-Arm) · 2026-11-05T16:05:01Z

+  auto backend_kernel_selector_config = mlas_backend_kernel_selector_config_;
+#if defined(USE_KLEIDIAI) && defined(MLAS_TARGET_ARM64)
+  const bool is_depthwise_multiplier_one = input_channels_per_group == 1 && filter_count_per_group == 1;
+  if (backend_kernel_selector_config.use_kleidiai && is_depthwise_multiplier_one && B != nullptr) {


NaN preserving clamping provided, fallback removed. Done

Melike Kaptan (melkap01-Arm) · 2026-11-05T16:04:34Z

+}
+
+TEST(MlasKleidiDepthwiseTest, ThreadedRowTilesWithTailNhwc) {
+  ASSERT_NE(GetMlasThreadPool(), nullptr);


Signed-off-by: Jonathan Clohessy <Jonathan.Clohessy@arm.com>

Signed-off-by: melkap01 <melike.kaptan@arm.com>

…ts, refactor cleaned Signed-off-by: melkap01 <melike.kaptan@arm.com>

Signed-off-by: melkap01 <melike.kaptan@arm.com>

Copilot

🔵 Needs a closer look

The SME2 integration requires hardware-backed validation of row/channel tails, packed-weight vector-length handling, and NaN-preserving clamp behavior that this source-only review could not verify.

8 open issues

_{Review effort: Balanced}

Signed-off-by: melkap01 <melike.kaptan@arm.com>

Hariharan Seshadri (hariharans29) requested a review from Copilot June 30, 2026 16:03

Copilot started reviewing on behalf of Hariharan Seshadri (hariharans29) June 30, 2026 16:04

View session

Copilot AI reviewed Jun 30, 2026

View reviewed changes

Comment thread onnxruntime/core/mlas/lib/kleidiai/dw_conv_kleidiai.cpp Resolved

Melike Kaptan (melkap01-Arm) force-pushed the jonclo01/dw_conv_integration branch from 9a52413 to 80524ec Compare September 9, 2026 11:36

mirounga self-requested a review September 15, 2026 16:45

mirounga previously approved these changes Sep 15, 2026

View reviewed changes

Melike Kaptan (melkap01-Arm) dismissed mirounga’s stale review via 715cc76 September 16, 2026 18:27

mirounga enabled auto-merge (squash) September 17, 2026 14:32

auto-merge was automatically disabled September 18, 2026 09:41
Head branch was pushed to by a user without write access

mirounga requested a balanced review from Copilot September 18, 2026 15:30

Copilot started reviewing on behalf of mirounga September 18, 2026 15:30

View session

Copilot AI reviewed Sep 18, 2026

View reviewed changes

mirounga requested a balanced review from Copilot October 6, 2026 16:05

Copilot started reviewing on behalf of mirounga October 6, 2026 16:13

View session

Copilot AI reviewed Oct 6, 2026

View reviewed changes

mirounga requested a balanced review from Copilot October 22, 2026 16:38

Copilot started reviewing on behalf of mirounga October 22, 2026 16:41

View session

Copilot AI reviewed Oct 22, 2026

View reviewed changes

mirounga requested a balanced review from Copilot October 28, 2026 15:22

Copilot started reviewing on behalf of mirounga October 28, 2026 15:23

View session

Copilot AI reviewed Oct 28, 2026

View reviewed changes

mirounga requested a balanced review from Copilot November 2, 2026 17:32

Copilot started reviewing on behalf of mirounga November 2, 2026 17:32

View session

Copilot AI reviewed Nov 2, 2026

View reviewed changes

Jonathan Clohessy (JonathanC-ARM) and others added 2 commits November 5, 2026 16:07

Integrate DW Conv Kernel

2dcfc42

Signed-off-by: Jonathan Clohessy <Jonathan.Clohessy@arm.com>

review points implemented

2150dd9

Signed-off-by: melkap01 <melike.kaptan@arm.com>

Melike Kaptan (melkap01-Arm) added 6 commits November 5, 2026 16:13

pre packed weights added, tests improved to cover larger channel coun…

fa579e7

…ts, refactor cleaned Signed-off-by: melkap01 <melike.kaptan@arm.com>

threading added into the DepthwiseConvKleidiAI

eb628da

Signed-off-by: melkap01 <melike.kaptan@arm.com>

co pilot review comments adressed

5cae7e4

Signed-off-by: melkap01 <melike.kaptan@arm.com>

codex review points on clamping addressed

a0b21b2

Signed-off-by: melkap01 <melike.kaptan@arm.com>

NaN bias case fallback provided

28cd39f

Signed-off-by: melkap01 <melike.kaptan@arm.com>

Preserve NaNs in KleidiAI depthwise convolution

a43c69d

Signed-off-by: melkap01 <melike.kaptan@arm.com>

Melike Kaptan (melkap01-Arm) force-pushed the jonclo01/dw_conv_integration branch from 3e8b1ea to a43c69d Compare November 5, 2026 16:28

mirounga requested a balanced review from Copilot November 5, 2026 16:35

Copilot started reviewing on behalf of mirounga November 5, 2026 16:36

View session

Copilot AI reviewed Nov 5, 2026

View reviewed changes

Melike Kaptan (melkap01-Arm) added 2 commits November 5, 2026 17:18

lintrunner run

32411d4

Signed-off-by: melkap01 <melike.kaptan@arm.com>

Merge branch 'microsoft:main' into jonclo01/dw_conv_integration

b46f439

		-std::numeric_limits<float>::max(),
		std::numeric_limits<float>::max(),

		-std::numeric_limits<float>::infinity(),
		std::numeric_limits<float>::infinity(),

Conversation

Jonathan Clohessy (JonathanC-ARM) commented Jun 30, 2026

Description

Changes

Supported depthwise shape envelope

Example Performance Results

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Hariharan Seshadri (hariharans29) commented Jun 30, 2026

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

What this PR does

What I like

Concerns (please address before merge)

On the MLAS backend kernel selector integration

Nits (do not block)

Bottom line

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Melike Kaptan (melkap01-Arm) commented Sep 18, 2026

Uh oh!

Melike Kaptan (melkap01-Arm) commented Oct 2, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🔵 Needs a closer look

Uh oh!

Melike Kaptan (melkap01-Arm) commented Oct 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Copilot review overview

🔵 Needs a closer look

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🔵 Needs a closer look

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

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

Melike Kaptan (melkap01-Arm) commented Oct 19, 2026 •

edited

Loading