[MLAS][KleidiAI] Add MLAS Arm64 half GEMM and convolution support

[Laan33]

Description

Adds MLAS half-precision GEMM and convolution support for Arm64 KleidiAI paths.

This change:

• Adds public MLAS HalfGemm backend-native packed-B APIs.
• Adds public MLAS HalfConv prepare, execute, and packed weights+bias APIs.
• Wires KleidiAI HalfGemm and HalfConv overrides into the MLAS platform dispatch.
• Adds SME/SME2 FP16 HGEMM and FP16 IMATMUL ukernel selection.
• Adds MLAS unit coverage for packed-B behavior, selector opt-out, zero-K handling, HalfConv prepare behaviour, and varied HalfGemm shapes/bias cases.

Motivation and Context

This is the second MR in the FP16 split and introduces the MLAS API and Arm64 backend plumbing needed for accelerated FP16 CPU kernels.

The later CPU operator changes can use these MLAS HalfGemm and HalfConv entry points without carrying KleidiAI-specific details at the operator layer. The backend selector support also preserves an explicit fallback path when KleidiAI should not be used for a given call.

Note: This PR is the second offshoot from 28487

[Copilot]

Pull request overview

This PR expands ONNX Runtime’s MLAS Arm64 KleidiAI integration to support FP16 (half) GEMM and convolution via new public MLAS APIs, KleidiAI override wiring, and additional kernel selection logic, with accompanying unit tests and some test-harness exclusions/tolerance tweaks.

Changes:

• Adds public MLAS FP16 HalfGemm packed-B APIs (generic + backend-native) and FP16 HalfConv prepare/execute/pack APIs.
• Wires KleidiAI FP16 HalfGemm/HalfConv overrides into MLAS platform dispatch and adds SME/SME2 FP16 ukernel selection.
• Extends MLAS/unit and ORT operator tests to cover new behaviors (packed-B contracts, selector opt-out, zero-K handling, HalfConv prepare behavior), plus minor EP exclusions/tolerance adjustments.

Reviewed changes

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

Show a summary per file

File	Description
onnxruntime/test/providers/cpu/tensor/resize_op_test.cc	Excludes additional EPs for FP16 Resize(13) coverage gaps.
onnxruntime/test/optimizer/nhwc_transformer_test.cc	Adjusts transformer tester invocation to use explicit tolerances/opset.
onnxruntime/test/mlas/unittest/test_util.h	Ensures guard buffer state is fully reset (guard pointer).
onnxruntime/test/mlas/unittest/test_halfgemm.h	Adds FP16-native comparison tolerance + new execution modes for HalfGemm tests.
onnxruntime/test/mlas/unittest/test_halfgemm.cpp	Adds HalfGemm tests for selector routing, packed-B behavior, overflow handling, and more shape/bias cases.
onnxruntime/test/mlas/unittest/test_conv2d.cpp	Adds HalfConv prepare tests for selector config handling and working-buffer sizing.
onnxruntime/test/contrib_ops/nhwc_pool_in_op_test.cc	Excludes non-owning EPs from internal NHWC fp16 pool tests.
onnxruntime/test/contrib_ops/matmul_4bits_test.cc	Adjusts abs error tolerance for FP16, conditional on WebGPU.
onnxruntime/test/contrib_ops/attention_op_test.cc	Disables WebGPU for a FP16 Attention test that uses unsupported mask_index.
onnxruntime/core/mlas/lib/platform.cpp	Enables KleidiAI overrides under SME or SME2 and wires HalfGemm/HalfConv overrides.
onnxruntime/core/mlas/lib/mlasi.h	Adds size_t overflow helpers used by KleidiAI/packing paths.
onnxruntime/core/mlas/lib/kleidiai/sgemm_kleidiai.cpp	Switches overflow checks to shared helpers.
onnxruntime/core/mlas/lib/kleidiai/sbgemm_kleidiai.cpp	Switches overflow checks to shared helpers.
onnxruntime/core/mlas/lib/kleidiai/qgemm_kleidiai.cpp	Renames BatchSize→BatchN for clarity/consistency and updates uses.
onnxruntime/core/mlas/lib/kleidiai/mlasi_kleidiai.h	Refactors includes and adds scratch-buffer shrinking helper declarations.
onnxruntime/core/mlas/lib/kleidiai/halfgemm_kleidiai.cpp	Adds KleidiAI FP16 HalfGemm implementation and native RHS pack.
onnxruntime/core/mlas/lib/kleidiai/halfconv_kleidiai.cpp	Adds KleidiAI FP16 convolution implementation (IMATMUL-based) and packing.
onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp	Minor cache-lookup refactor for indirection table cache.
onnxruntime/core/mlas/lib/kai_ukernel_interface.h	Adds FP16 IMATMUL/HGEMM kernel wrapper types and selection APIs.
onnxruntime/core/mlas/lib/kai_ukernel_interface.cpp	Implements FP16 kernel selection (SME vs SME2) and new ukernel wrappers.
onnxruntime/core/mlas/lib/halfgemm.h	Adds shared packed-B sizing helper + implements generic CopyPackB with padding.
onnxruntime/core/mlas/lib/halfgemm.cpp	Adds zero-K behavior, selector-config routing, native-pack APIs, and enables CopyPackB dispatch.
onnxruntime/core/mlas/lib/halfgemm_kernel_neon.cpp	Enables CopyPackB dispatch for NEON halfgemm.
onnxruntime/core/mlas/lib/halfconv.cpp	Adds public dispatch wrappers for HalfConv prepare/execute/pack APIs.
onnxruntime/core/mlas/inc/mlas.h	Exposes new HalfGemm/HalfConv APIs and extends parameter structs for new flags/config.
cmake/onnxruntime_mlas.cmake	Adds new MLAS sources (halfconv + KleidiAI halfgemm/halfconv).

---

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

[Copilot]

Pull request overview

Copilot reviewed 26 out of 26 changed files in this pull request and generated 2 comments.

[Copilot]

Pull request overview

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

[hariharans29]

Review of PR #28786 — [MLAS][KleidiAI] Add MLAS Arm64 half GEMM and convolution support

Verdict: structurally sound, recommend approve after addressing a few points below. This is a substantial PR (+3,274/−65, 24 files, 8 commits, 85/85 checks OK) that lays the MLAS API surface for FP16 GEMM/Conv with KleidiAI as the first backend. The contracts are well thought-out, the validation discipline is appropriate, and the test additions are non-trivial. The size makes line-by-line audit of halfconv_kleidiai.cpp (1,041 LOC) impractical from the diff view alone — flagging that explicitly.

What I like

1. Native-packed vs generic-packed B contract is clean.

The MLAS_HALF_GEMM_DATA_PARAMS split into BIsPacked (generic MLAS pre-pack) and BIsBackendNativePacked (backend-direct-consumption layout) with explicit invariants documented at include/onnxruntime/core/mlas/inc/mlas.h:1923-1930:

When true, B must be produced by MlasHalfGemmNativePackB, ldb must be 0, Bias must be nullptr, and OutputProcessor must be nullptr.

Validated both at entry (zero-K path) and in the KleidiAI override. The generic kernel path explicitly throws when fed a native-packed B — preventing silent miscompute from contract violation. This is the right call for a "trust-the-pointer" packed-layout API.

2. Defensive batch-entry pre-validation in the KleidiAI override.

onnxruntime/core/mlas/lib/kleidiai/halfgemm_kleidiai.cpp:152-166 validates all batch entries before executing any of them. Without this, a malformed entry at index N+1 after writing entries 0..N would leave partially-correct outputs that the fallback then can't safely re-execute (subsequent fallback would either re-execute correct entries — racy — or skip them — wrong). The all-or-nothing precheck is exactly right.

3. Overflow-checked size arithmetic is consolidated and reused.

The new MlasMultiplyOverflowsSizeT / MlasAddOverflowsSizeT in onnxruntime/core/mlas/lib/mlasi.h:104-142 replace the prior mul_overflow_size_t_builtin (which was tucked inside mlasi_kleidiai.h). Migrated cleanly across sbgemm_kleidiai.cpp and sgemm_kleidiai.cpp. Used pervasively in the new MlasHalfGemmTryGetPackedBSize (which is itself a real improvement — previously (K + PackedK - 1) & ~(PackedK - 1) and the subsequent N * AlignedK * FP16_SIZE could silently overflow). Good hygiene.

4. Zero-K and zero-batch handling are explicit and tested.

MlasHalfGemmBatch now has an explicit early-out for BatchN == 0 || M == 0 || N == 0 and a K == 0 path that fills C with Bias (or zero) and runs OutputProcessor. Previously these were undefined.

5. Default and Neon dispatches now actually support generic packed-B.

The dispatches at halfgemm.cpp:741-742 and halfgemm_kernel_neon.cpp:182 wire MlasHalfGemmCopyPackB<...> (was nullptr previously). Before this PR, MlasHalfGemmPackBSize would return 0 on these dispatches (no CopyPackBRoutine) and any caller-side packing of B for default/Neon would silently no-op. Quietly significant fix.

6. TLS scratch with bounded retention.

The thread_local KaiHalfTlsBuffers + ScopedKaiHalfTlsCleanup + MlasShrinkKleidiAIScratchIfTooLarge (8 MiB cap) pattern in halfgemm_kleidiai.cpp avoids per-call heap churn without unbounded thread-local growth on outlier-shape calls. The needs_rhs_packing flag avoids allocating the RHS scratch when the entire batch is BIsBackendNativePacked — nice optimization.

7. Backend selector config plumbing.

MLAS_BACKEND_KERNEL_SELECTOR_CONFIG::use_kleidiai = false is honored both in the public-dispatch wrapper (halfconv.cpp returns false before consulting platform) and inside the override entry points. Conv2d_HalfConv.PrepareRespectsBackendSelectorConfig test gates this. The "consensus" function TryGetHalfGemmBackendSelectorConfig correctly returns false if any entry has use_kleidiai == false while picking the first non-null config to forward.

8. KleidiAI dispatch table rollout was correctly extended to SME2.

platform.cpp:707 widens the conditional from HasArm_SME() to HasArm_SME() || HasArm_SME2(). Necessary because SME2-only CPUs would have been skipping all KleidiAI overrides previously. Good catch.

Concerns (please address before merge)

1. MLAS_THROW_EX introductions in MlasHalfGemmBatch change established MLAS error semantics.

The new throws at halfgemm.cpp:111, halfgemm.cpp:119-122, and halfgemm.cpp:144-146 break with the MLAS convention of returning silently on invalid inputs (most other public MLAS APIs return / set output to 0 rather than throw). Three separate exception sites for what are programmer-error contract violations (null DataParams, malformed K==0 native-packed entry, native-packed B fed to generic kernels) is heavy. Three options ranked by my preference:

• Best: convert these to assert(...) (debug-only) plus a silent fallback (skip the bad entry, or return). Cheaper and matches the rest of the MLAS surface.
• Acceptable: keep them but document the throw contract in the MlasHalfGemmBatch header comment. ORT code calling MLAS has been historically permitted to ignore these failure modes.
• Risky: leave undocumented. ORT inference paths that don't wrap MLAS calls in try/catch can now propagate exceptions to user code in ways they previously could not.

The "native-packed B fed to generic kernels" case has the strongest justification (it's a real memory-safety hazard, not just a contract violation) — I'd keep that throw.

2. MlasHalfGemmPackB silently no-ops on malformed inputs.

Newly added at halfgemm.cpp:254-257:

if (B == nullptr || PackedB == nullptr || ldb < N ||
    dispatch->CopyPackBRoutine == nullptr || MlasHalfGemmPackBSize(N, K, false) == 0) {
  return;
}

If a caller passes ldb < N (an actual bug) the function returns without writing PackedB, the caller then runs MlasHalfGemmBatch against uninitialized memory and silently produces garbage. At minimum, this should assert(B != nullptr && PackedB != nullptr && ldb >= N) in debug builds. Compare with the throw-style elsewhere in this same PR — the asymmetry is suspicious.

3. matmul_4bits_test.cc tolerance bump 0.055 → 0.065 (and 0.1 for WebGPU) lacks justification.

} else if constexpr (std::is_same<AType, MLFloat16>::value) {
#if defined(USE_WEBGPU)
  base_opts.output_abs_error = 0.1f;
#else
  base_opts.output_abs_error = 0.065f;  // was 0.055f
#endif
}

A ~20% tolerance relaxation in an op's golden test is a code-smell unless the rationale is captured. Two questions:

• Why does fp16 path drift more after this PR? Is it because the KleidiAI native HGEMM path uses fp16 (not fp32) accumulation and is now wired in by default on SME hardware? If so, please add an inline comment: // Relaxed for KleidiAI native fp16 GEMM path which accumulates in fp16 (see PR #28786).
• Why is WebGPU's tolerance now even looser (0.1 vs 0.065)? Is this a pre-existing WebGPU drift that this PR is just now codifying as it carves the path out? If pre-existing, mention as such in the comment.

Tolerance bumps without rationale are how regressions hide.

4. halfconv_kleidiai.cpp (1,041 LOC) was not reviewed in detail.

GitHub's diff view collapsed this file as "large diff." A maintainer review of this kind of PR really wants eyes on the depthwise/non-depthwise dispatch, the packed weights+bias layout invariants, and the FP16 IMATMUL kernel driver loop. Either expand the file in the PR review UI, or split the file into a separate follow-up PR with its own line-level review. As is, I'm trusting it on the strength of:

• 85/85 CI checks green
• The dispatch contracts and selector config are clean at the boundary I did read
• New Conv2d_HalfConv.PrepareRespectsBackendSelectorConfig, PrepareReportsWorkingBufferSizeInBytes, and PrepareInitializesLayoutFlags tests cover the wrapper's behavioral contract

If you want me to deep-dive that file, I'd do it as a separate review pass — just say the word.

5. Inconsistency: halfgemm uses #if defined(USE_KLEIDIAI); halfconv does not.

The halfgemm public APIs (MlasHalfGemmNativePackBSize, MlasHalfGemmNativePackB) gate their body with #if defined(USE_KLEIDIAI) and MLAS_UNREFERENCED_PARAMETER the args otherwise. The halfconv public APIs in halfconv.cpp do not — they just check GetMlasPlatform().MlasHalfConvPrepareOverride == nullptr and return false. The latter is cleaner (matches convolve.cpp's pattern and doesn't bake USE_KLEIDIAI knowledge into the dispatch layer). Suggest harmonizing halfgemm to the same style; the override-null check already serves the same purpose.

Nits (do not block)

6. convolve_kleidiai.cpp LHS-cache fix lacks a bug description.

Commit 4e72126 inlines the cache lookup:

-auto& lhs_ptrs_cache = lhs_ptrs_cache_by_pad[cur_pad_ptr];
-if (auto found = lhs_ptrs_cache.find(key); found != lhs_ptrs_cache.end()) {
+if (auto found = lhs_ptrs_cache_by_pad[cur_pad_ptr].find(key); found != lhs_ptrs_cache_by_pad[cur_pad_ptr].end()) {

If the bug was that the auto& reference was being invalidated by intervening unordered_map mutation (rehash), then the inlined-double-lookup form just hides the problem — the reference style is more efficient and not inherently wrong. A one-line code comment naming the actual hazard (or a commit message explaining it) would help future readers avoid re-introducing the bug. If the bug was something else (e.g., the lhs_ptrs_cache_by_pad[cur_pad_ptr] was an unintended insertion due to operator[] on a not-yet-present key), then the inlined form has the same bug — both lines call operator[].

7. Hardcoded clamp_min/clamp_max = ±infinity will fight fused-activation.

In halfgemm_kleidiai.cpp:180-181 — fine for vanilla GEMM. If MLAS_HALF_GEMM_DATA_PARAMS ever gains a fused-activation field, the wiring point is here. A // TODO: plumb through MLAS_ACTIVATION when MLAS_HALF_GEMM_DATA_PARAMS supports it comment is worth its weight.

8. MLAS_CONV_PARAMETERS::InputOutputChannelsLast vs ChannelsLast.

Two channels-last booleans on the same struct (ChannelsLast was pre-existing). The relationship and disjoint meaning isn't documented in mlas.h. Add a one-line comment: e.g., // InputOutputChannelsLast = true means both input and output tensors are NHWC; ChannelsLast = true means the compute kernel produces NHWC output (which may differ when input is NCHW and a transpose is fused). Whatever the actual semantics are.

9. The resize_op_test.cc "example_ep" exclusion uses a string literal.

excluded_providers.insert("example_ep");

example_ep is the autoep test plugin name used throughout test/autoep/test_model_package.cc — it's referred to by string literal there too, so this is at least consistent. But ideally there'd be a constexpr char kExampleAutoEpName[] = "example_ep" somewhere in the autoep test utilities to centralize it. Out of scope for this PR; just noting.

10. nhwc_pool_in_op_test.cc correctly explains the EP exclusion.

// These tests target ORT's CPU implementation of the MS-internal NHWC fp16 pool
// ops. Do not offer the models to EPs that may be registered by the test
// harness but do not own this internal-domain CPU/MLAS coverage.
const std::unordered_set<std::string> kNhwcFp16PoolExcludedProviders{
    kCoreMLExecutionProvider,
    kTensorrtExecutionProvider,
};

Exactly the style of explanation I wanted in matmul_4bits_test.cc for the tolerance bump. Carry this style there.

Bottom line

Approve once:

• The tolerance bump in matmul_4bits_test.cc gets an explanatory comment.
• The throw vs. silent-no-op API style is reconciled (or at least documented in MlasHalfGemmBatch's doc-comment).
• The MlasHalfGemmPackB malformed-input case gets a debug assert (or comment).

Suggest deferring/splitting:

• Optional: split halfconv_kleidiai.cpp into a separate PR with its own line-level review, given its size.

The architectural design — split native vs generic packed-B contract, overflow-checked size arithmetic in mlasi.h, defensive batch pre-validation, TLS scratch with bounded retention, selector-config plumbing — is solid. Quietly significant fixes alongside the main work (default + Neon dispatch CopyPackB wiring, SME2-CPU coverage for all KleidiAI overrides, KleidiAI conv LHS-cache fix). This is good infrastructure for the FP16 follow-on PRs.

[Copilot]

Pull request overview

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

[hariharans29]

Review — PR #28786 [MLAS][KleidiAI] Add MLAS Arm64 half GEMM and convolution support

Verdict: structurally sound, recommend approve after addressing the FP16 consumer-wiring gap and a small handful of source-cleanliness items below. This is a substantial PR (~+3,200 / −65 across 24 files) that lays the MLAS API surface for FP16 GEMM/Conv with KleidiAI as the first backend. The contracts are well thought-out, the validation discipline is appropriate, and the test additions are non-trivial. The size makes line-by-line audit of halfconv_kleidiai.cpp (1,041 LOC) impractical from the diff view alone — flagging that explicitly.

What I like

1. Native-packed vs generic-packed B contract is clean. The MLAS_HALF_GEMM_DATA_PARAMS split into BIsPacked (generic MLAS pre-pack) and BIsBackendNativePacked (backend-direct-consumption layout) with explicit invariants documented at mlas.h:1923-1930:

   When true, B must be produced by MlasHalfGemmNativePackB, ldb must be 0, Bias must be nullptr, and OutputProcessor must be nullptr.

   Validated both at entry (zero-K path) and in the KleidiAI override. The generic kernel path explicitly throws when fed a native-packed B — preventing silent miscompute from contract violation. This is the right call for a "trust-the-pointer" packed-layout API.

2. Defensive batch-entry pre-validation in the KleidiAI override. halfgemm_kleidiai.cpp:152-166 validates all batch entries before executing any of them. Without this, a malformed entry at index N+1 after writing entries 0..N would leave partially-correct outputs that the fallback then can't safely re-execute (subsequent fallback would either re-execute correct entries — racy — or skip them — wrong). The all-or-nothing precheck is exactly right.

3. Overflow-checked size arithmetic is consolidated and reused. The new MlasMultiplyOverflowsSizeT / MlasAddOverflowsSizeT in mlasi.h:104-142 replace the prior mul_overflow_size_t_builtin (which was tucked inside mlasi_kleidiai.h). Migrated cleanly across sbgemm_kleidiai.cpp and sgemm_kleidiai.cpp. Used pervasively in the new MlasHalfGemmTryGetPackedBSize — a real improvement, since previously (K + PackedK - 1) & ~(PackedK - 1) and the subsequent N * AlignedK * FP16_SIZE could silently overflow. Good hygiene.

4. Zero-K and zero-batch handling are explicit and tested. MlasHalfGemmBatch now has an explicit early-out for BatchN == 0 || M == 0 || N == 0 and a K == 0 path that fills C with Bias (or zero) and runs OutputProcessor. Previously these were undefined.

5. Default and Neon dispatches now actually support generic packed-B. The dispatches at halfgemm.cpp:741-742 and halfgemm_kernel_neon.cpp:182 wire MlasHalfGemmCopyPackB<...> (was nullptr previously). Before this PR, MlasHalfGemmPackBSize would return 0 on these dispatches and any caller-side packing of B for default/Neon would silently no-op. Quietly significant fix.

6. TLS scratch with bounded retention. The thread_local KaiHalfTlsBuffers + ScopedKaiHalfTlsCleanup + MlasShrinkKleidiAIScratchIfTooLarge (8 MiB cap) pattern in halfgemm_kleidiai.cpp avoids per-call heap churn without unbounded thread-local growth on outlier-shape calls. The needs_rhs_packing flag avoids allocating the RHS scratch when the entire batch is BIsBackendNativePacked — nice optimization. This becomes a reusable pattern for follow-on KleidiAI work; please keep it general rather than fp16-specific.

7. Backend selector config plumbing inside MLAS is good. MLAS_BACKEND_KERNEL_SELECTOR_CONFIG::use_kleidiai = false is honored in the public-dispatch wrapper (halfconv.cpp returns false before consulting platform), and the TryGetHalfGemmBackendSelectorConfig consensus function in halfgemm.cpp correctly returns false if any entry has use_kleidiai == false while picking the first non-null config to forward. Conv2d_HalfConv.PrepareRespectsBackendSelectorConfig test gates the conv side. (Consumer-side wiring is the issue — see concern Set up CI with Azure Pipelines #1 below.)

8. KleidiAI dispatch table rollout was correctly extended to SME2. platform.cpp:707 widens the conditional from HasArm_SME() to HasArm_SME() || HasArm_SME2(). Necessary because SME2-only CPUs would have been skipping all KleidiAI overrides previously.

9. API hardening details are good. MlasHalfGemmBatch's doc-comment formally documents the MLAS_THROW_EX contract for contract violations that cannot be safely ignored. MlasHalfGemmPackB has a debug assert (assert(N == 0 || K == 0 || (B != nullptr && PackedB != nullptr && ldb >= N));) before the silent-return guard — debug catches bugs, production keeps the silent return that matches the rest of MLAS. Test tolerances that needed loosening to accommodate native fp16 accumulation (matmul_4bits_test.cc) carry an inline comment naming the why ("native HGEMM on SME; see PR [MLAS][KleidiAI] Add MLAS Arm64 half GEMM and convolution support #28786"), and WebGPU's separate looser bound has its own provider-specific justification. This is the discipline I want to see.

Concerns (please address before merge)

1. Blocker — MLAS_BACKEND_KERNEL_SELECTOR_CONFIG plumbing does not reach the existing FP16 op consumers

The PR threads MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* through every new/changed MLAS API in mlas.h (the Native packing APIs, all 4 of the HalfConv APIs, and the new per-entry MLAS_HALF_GEMM_DATA_PARAMS::BackendKernelSelectorConfig). The dispatch is correctly wired inside MLAS (item 7 above). But the three existing in-tree consumers of MlasHalfGemmBatch were not updated to populate the new field, so after this PR lands the user-facing use_kleidiai=false session option is silently ignored for every FP16 GEMM/Conv/MoE call:

1. onnxruntime/core/providers/cpu/math/gemm.cc:228 — standard Gemm op (FP16 path). The surrounding class already holds mlas_backend_kernel_selector_config_ (set up in gemm.h:30) — the constructor side is done. Only the call site is missing data.BackendKernelSelectorConfig = &mlas_backend_kernel_selector_config_;.
2. onnxruntime/core/providers/cpu/fp16/fp16_conv.cc:562 — FP16 NHWC conv op family. This class does not currently hold an MLAS_BACKEND_KERNEL_SELECTOR_CONFIG member. Needs all three: declare the member, call SetupMlasBackendKernelSelectorFromConfigOptions(...) in the constructor, set gemm_params.BackendKernelSelectorConfig. The include is mlas_backend_kernel_selector_config_utils.h.
3. onnxruntime/contrib_ops/cpu/moe/moe_cpu.cc:554 — contrib MoE op FP16 GEMM call. MoEBase (the parent class) already holds the member via moe_base_cpu.h:57. Same one-line fix as gemm.cc.

The reference pattern lives at matmul_nbits.cc:115 (member + ctor wiring) — every MLAS call site in that file threads &mlas_backend_kernel_selector_config_. The same pattern is what's missing on the FP16 side.

This isn't an architectural issue with the PR — the API surface is correct. It's a consumer-wiring omission, and without it the new BackendKernelSelectorConfig field on MLAS_HALF_GEMM_DATA_PARAMS is dead weight for every in-tree caller, and a user who sets the session option to disable KleidiAI gets KleidiAI anyway. The contract that "passing this option disables KleidiAI" is what makes the opt-out trustable.

Two acceptable resolutions:

• (Preferred) Land the three call-site fixes in this PR — each is one to a few lines.
• (Acceptable) Land them in an explicit follow-up PR named in the description of this one, so reviewers know the FP16 op consumers are explicitly tracked rather than forgotten.

A small ORT-level test on SME/SME2 ARM64 that runs an FP16 model twice (with and without the KleidiAI opt-out option) and asserts both paths produce matching output would also catch any future regression of this kind. Currently no test in tree would catch the opt-out being silently ignored on these three consumers.

2. Global-namespace pollution: MaximumRetainedKleidiAIScratchBytes and MlasShrinkKleidiAIScratchIfTooLarge

These are KleidiAI-specific helpers defined in the global namespace in mlasi_kleidiai.h, a widely-included header. They should live in namespace ArmKleidiAI (or a detail sub-namespace) along with everything else KleidiAI-related in this file. Worth fixing now — once these names ship in a public-ish header, moving them later becomes a breaking change for any internal consumer that took a dependency on the global names.

3. Stale MLAS_UNREFERENCED_PARAMETER(ThreadPool) in ArmKleidiAI::MlasHalfGemmBatch

The MLAS_UNREFERENCED_PARAMETER(ThreadPool) is dead/wrong — ThreadPool is actively used later in the function via MlasTrySimpleParallel. Should be removed; the macro is a real comprehension hazard for the next reader trying to trace threading behavior.

A separate Copilot suggestion in the same area was to also drop MLAS_UNREFERENCED_PARAMETER(BackendKernelSelectorConfig) and add an explicit short-circuit on use_kleidiai == false. Functionally that's defense-in-depth only — the public dispatch already filters via TryGetHalfGemmBackendSelectorConfig, which returns false when any batch entry has use_kleidiai == false, so the override is never invoked in that case. Cheap to add; nice-to-have, not required.

4. halfconv_kleidiai.cpp (1,041 LOC) was not reviewed in detail

GitHub's diff view collapsed this file as "large diff." A maintainer review of this kind of PR really wants eyes on the depthwise/non-depthwise dispatch, the packed weights+bias layout invariants, and the FP16 IMATMUL kernel driver loop. Either expand the file in the PR review UI, or split the file into a separate follow-up PR with its own line-level review. As is, I'm trusting it on the strength of CI, the dispatch contracts and selector config being clean at the boundary I did read, and the new Conv2d_HalfConv.PrepareRespectsBackendSelectorConfig, PrepareReportsWorkingBufferSizeInBytes, and PrepareInitializesLayoutFlags tests covering the wrapper's behavioral contract. If you want a dedicated pass on this file, say the word and I'll book it as a separate review.

Nits (do not block)

5. Inconsistency: halfgemm uses #if defined(USE_KLEIDIAI); halfconv does not. The halfgemm public APIs (MlasHalfGemmNativePackBSize, MlasHalfGemmNativePackB) gate their body with #if defined(USE_KLEIDIAI) and MLAS_UNREFERENCED_PARAMETER the args otherwise. The halfconv public APIs in halfconv.cpp just check GetMlasPlatform().MlasHalfConvPrepareOverride == nullptr and return false. The latter is cleaner (matches convolve.cpp's pattern and doesn't bake USE_KLEIDIAI knowledge into the dispatch layer). Suggest harmonizing halfgemm to the same style; the override-null check already serves the same purpose.

6. <iostream> unconditional include in mlasi_kleidiai.h. The file's existing logging path conditionally includes <iostream> under #if KLEIDIAI_DEBUG_LOGGING || KLEIDIAI_KERNEL_LOGGING. Adding it again unconditionally at top-of-file expands compile time across every TU that pulls this header. Drop it or scope it.

7. convolve_kleidiai.cpp LHS-cache fix lacks a bug description. The cache-lookup inlining:

   -auto& lhs_ptrs_cache = lhs_ptrs_cache_by_pad[cur_pad_ptr];
   -if (auto found = lhs_ptrs_cache.find(key); found != lhs_ptrs_cache.end()) {
   +if (auto found = lhs_ptrs_cache_by_pad[cur_pad_ptr].find(key); found != lhs_ptrs_cache_by_pad[cur_pad_ptr].end()) {

   If the bug was that the auto& reference was being invalidated by intervening unordered_map mutation (rehash), the inlined-double-lookup form just hides the problem — the reference style is more efficient and not inherently wrong. A one-line code comment naming the actual hazard would help future readers avoid re-introducing the bug. If the bug was something else (e.g., lhs_ptrs_cache_by_pad[cur_pad_ptr] was an unintended insertion via operator[] on a not-yet-present key), then the inlined form has the same bug — both lines call operator[].

8. Hardcoded clamp_min/clamp_max = ±infinity will fight fused-activation. In halfgemm_kleidiai.cpp:180-181 — fine for vanilla GEMM. If MLAS_HALF_GEMM_DATA_PARAMS ever gains a fused-activation field, this is the wiring point. A // TODO: plumb through MLAS_ACTIVATION when MLAS_HALF_GEMM_DATA_PARAMS supports it comment is worth its weight.

9. MLAS_CONV_PARAMETERS::InputOutputChannelsLast vs ChannelsLast. Two channels-last booleans on the same struct (ChannelsLast was pre-existing). The relationship and disjoint meaning isn't documented in mlas.h. Add a one-line comment, e.g.: // InputOutputChannelsLast = true means both input and output tensors are NHWC; ChannelsLast = true means the compute kernel produces NHWC output (which may differ when input is NCHW and a transpose is fused). Whatever the actual semantics are.

10. resize_op_test.cc's "example_ep" exclusion uses a string literal. example_ep is the autoep test plugin name used throughout test/autoep/test_model_package.cc — it's referred to by string literal there too, so this is at least consistent. Ideally there'd be a constexpr char kExampleAutoEpName[] = "example_ep" somewhere in the autoep test utilities to centralize it. Out of scope for this PR; just noting.

11. nhwc_pool_in_op_test.cc correctly explains the EP exclusion. Carry the same explanatory style to any future EP-exclusion sites:

    // These tests target ORT's CPU implementation of the MS-internal NHWC fp16 pool
    // ops. Do not offer the models to EPs that may be registered by the test
    // harness but do not own this internal-domain CPU/MLAS coverage.

Bottom line

Approve once:

1. The BackendKernelSelectorConfig plumbing reaches the three existing FP16 consumers (gemm.cc, fp16_conv.cc, moe_cpu.cc) — either in this PR or in a clearly named follow-up that's referenced from this PR's description.
2. MaximumRetainedKleidiAIScratchBytes / MlasShrinkKleidiAIScratchIfTooLarge are moved into namespace ArmKleidiAI (or a detail namespace).
3. The stale MLAS_UNREFERENCED_PARAMETER(ThreadPool) is dropped from ArmKleidiAI::MlasHalfGemmBatch.
4. CI is fully green on the merge commit.

Deferrable to follow-ups:

• Deep review of halfconv_kleidiai.cpp.
• USE_KLEIDIAI halfgemm/halfconv gating-style consistency.
• <iostream> include cleanup.
• The defense-in-depth use_kleidiai short-circuit inside the halfgemm override.
• An ORT-level FP16 model test that flips the KleidiAI opt-out and asserts identical output, so the consumer-wiring contract is locked in by CI.
• The remaining stylistic nits (LHS-cache comment, fused-activation TODO, ChannelsLast documentation, example_ep constant).

The architectural design — split native vs generic packed-B contract, overflow-checked size arithmetic in mlasi.h, defensive batch pre-validation, TLS scratch with bounded retention, selector-config plumbing through MLAS — is solid. Quietly significant fixes alongside the main work (default + Neon dispatch CopyPackB wiring, SME2-CPU coverage for all KleidiAI overrides, KleidiAI conv LHS-cache fix). This is good infrastructure for the FP16 follow-on PRs; the one blocker is just making sure the new opt-out actually reaches the three existing FP16 callers.

[hariharans29]

Re-review: PR #28786 — [MLAS][KleidiAI] Add MLAS Arm64 half GEMM and convolution support

Verdict — approve. The FP16 consumer-wiring blocker from my prior pass is fixed, the two other "please address before merge" items are resolved, and the design continues to hold up. Head is a8861ad ("Address review feedback"), +3,263/−32 across 25 files; CI is at 71/86 (down from 85/85 previously — the regressed jobs look infra/flaky at a glance rather than KleidiAI-related, worth confirming before merge but not blocking). Corrections to my prior review are noted.

What was resolved since my prior review

1. Blocker Set up CI with Azure Pipelines #1 — BackendKernelSelectorConfig plumbing now reaches the FP16 consumers. All three call sites are wired:

   • onnxruntime/core/providers/cpu/math/gemm.cc:228 — data.BackendKernelSelectorConfig = mlas_backend_kernel_selector_config; on the FP16 Gemm path.
   • onnxruntime/core/providers/cpu/fp16/fp16_conv.cc:50, :98, :564 — new MLAS_BACKEND_KERNEL_SELECTOR_CONFIG member, SetupMlasBackendKernelSelectorFromConfigOptions(...) in the constructor, and per-call assignment on gemm_params.
   • onnxruntime/contrib_ops/cpu/moe/moe_cpu.cc:553 — one-line addition using the pre-existing MoEBase::mlas_backend_kernel_selector_config_.

   The user-facing use_kleidiai=false session option is now trustable across FP16 GEMM/Conv/MoE.

2. Concern update HighLevelDesign.md #3 — stale MLAS_UNREFERENCED_PARAMETER in ArmKleidiAI::MlasHalfGemmBatch removed. The MLAS_UNREFERENCED_PARAMETER(ThreadPool) and MLAS_UNREFERENCED_PARAMETER(BackendKernelSelectorConfig) lines are gone, replaced with an explicit early-return:

   if (BackendKernelSelectorConfig != nullptr && !BackendKernelSelectorConfig->use_kleidiai) {
       return false;
   }

   This is defense-in-depth on top of TryGetHalfGemmBackendSelectorConfig's consensus filter in the public MlasHalfGemmBatch dispatch. Good.

3. Nit Set up CI with Azure Pipelines #7 — convolve_kleidiai.cpp LHS-cache fix now has an explanatory comment at convolve_kleidiai.cpp:462-464. The reader now knows the pad-buffer-realloc invariant the cache is guarding.

4. Nit Fix build #8 — hardcoded clamp_min/clamp_max now has a TODO at halfgemm_kleidiai.cpp:181-182:

   // TODO: Plumb MLAS_ACTIVATION through this call site if MLAS_HALF_GEMM_DATA_PARAMS
   // grows fused activation support.
   

5. Nit Enable Mac pipeline #9 — ChannelsLast vs InputOutputChannelsLast disambiguated in mlas.h. Both fields now carry doc-comments explaining their relationship for the HalfConv case. mlas.h:896-899, mlas.h:918-923.

Correction to my prior review

Concern #2 was wrong on my part — MaximumRetainedKleidiAIScratchBytes and MlasShrinkKleidiAIScratchIfTooLarge are already declared inside namespace ArmKleidiAI at mlasi_kleidiai.h:53-64 (the namespace ArmKleidiAI { opens at line 53, closes at end of file, and the helpers sit between). Call sites in the KleidiAI backend use ArmKleidiAI::MlasShrinkKleidiAIScratchIfTooLarge(...). I misread the file structure; Copilot's AI review made the same mistake. No action needed.

What still stands (non-blocking)

6. halfconv_kleidiai.cpp (1,041 LOC) not deep-audited from the diff view. Structurally I still trust it on the strength of CI, clean dispatch/selector boundaries, and the Conv2d_HalfConv.* prepare tests. If a dedicated pass is desired I'm happy to book one against the merged file post-merge.

7. USE_KLEIDIAI gating asymmetry. halfgemm's new MlasHalfGemmNativePackBSize/MlasHalfGemmNativePackB wrap their bodies in #if defined(USE_KLEIDIAI); the halfconv equivalents in halfconv.cpp just check the platform-override function pointer for null. The latter matches convolve.cpp and is the cleaner idiom — worth harmonizing when this area is next touched.

8. <iostream> unconditional include in mlasi_kleidiai.h. Copilot flagged this and it's still there. The file's existing logging path conditionally includes <iostream> under #if KLEIDIAI_DEBUG_LOGGING || KLEIDIAI_KERNEL_LOGGING; the unconditional include at top-of-file is duplicative and expands compile time in every TU that pulls this header. Trivial follow-up.

9. resize_op_test.cc "example_ep" is a bare string literal. Consistent with test/autoep/test_model_package.cc which does the same, so no change in this PR. A constexpr char kExampleAutoEpName[] in the autoep test utilities would be nice one day.

10. OpCount behavior for a fused vs unfused pattern remains covered elsewhere — no regression concern here for this PR.

CI

Head a8861ad is at 71/86 checks OK vs the 85/85 the branch previously reached. Please confirm the regressed jobs are infra/flakes (the check-annotations exposed in the files-changed view are on unrelated files — QgemmU8X8KernelAvx2.asm epilog-length warnings and a model_package/src/manifest_parser.cc strerror warning — neither is introduced by this PR) rather than KleidiAI/FP16-related before merge.

Bottom line

The three "please address before merge" items from my prior review are resolved; five of six nits are also addressed. No new blockers. Ready to merge once CI is confirmed green (or the reds are triaged as unrelated). The remaining nits — USE_KLEIDIAI gating consistency, unconditional <iostream>, and a deep pass on halfconv_kleidiai.cpp — are all fair follow-up candidates.

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.