Skip to content

[MLAS][KleidiAI] Add MatMul, Gemm, and Conv FP16 execution paths#29709

Open
JonathanC-ARM wants to merge 12 commits into
microsoft:mainfrom
JonathanC-ARM:fp16-split/03-cpu-fp16-kernels
Open

[MLAS][KleidiAI] Add MatMul, Gemm, and Conv FP16 execution paths#29709
JonathanC-ARM wants to merge 12 commits into
microsoft:mainfrom
JonathanC-ARM:fp16-split/03-cpu-fp16-kernels

Conversation

@JonathanC-ARM

@JonathanC-ARM JonathanC-ARM commented Jul 14, 2026

Copy link
Copy Markdown
Contributor

Description

This PR builds on fp16-split/02-mlas-half-api-squashed by connecting its FP16 MLAS APIs and KleidiAI backends to the CPU Execution Provider.

It:

- Registers CPU `MLFloat16` kernels for:
  - `MatMul` opsets 1–8, 9–12, and 13+
  - `Gemm` opsets 7–8, 9–10, 11–12, and 13+
- Adds a dedicated FP16 `MatMul` implementation which:
  - Preserves ONNX broadcasting and batched execution through `MatMulComputeHelper`.
  - Handles empty outputs and zero-K inputs.
  - Uses the existing HGEMM path for eligible small, unpacked matrices.
  - Otherwise dispatches through `MlasHalfGemmBatch`.
  - Carries the configured MLAS backend selector through every batch entry.
  - Pre-packs constant two-dimensional B matrices into the selected backend’s native layout when supported.
- Keeps backend-native FP16 MatMul packs owned by the kernel. These layouts depend on the active backend, so they are not placed in the cross-session shared-prepack cache without suitable layout metadata.
- Adds an explicit kernel-owned packed-weights marker to `PrePackedWeights`. Session initialization continues to reject broken kernels which report successful prepacking without either providing cacheable buffers or explicitly identifying a valid kernel-owned pack.
- Extends FP16 `Gemm` to use `MlasHalfGemmBatch` for supported no-transpose, `alpha == 1` cases:
  - No effective bias, including a missing C input or `beta == 0`.
  - `beta == 1` with an MLAS-compatible N-element bias.
  - Unsupported transpose, scaling, or bias-broadcast combinations retain the existing Eigen fallback.
- Adds an MLAS HalfConv fast path to the existing FP16 CPU Conv implementation:
  - Uses `MlasHalfConvPrepare` to determine runtime and backend eligibility for two-dimensional convolutions without a fused Sum input.
  - Runs eligible convolutions through `MlasHalfConv`.
  - Supports NCHW and NHWC layouts and forwards padding, stride, dilation, activation, and workspace information.
  - Pre-packs eligible constant filters using `MlasHalfConvPackWeightsAndBias`.
  - Includes a compatible constant bias in the native pack when safe.
  - Falls back to the existing reordered-weight/im2col plus HalfGemm implementation when HalfConv is unavailable or rejects the configuration.
- Keeps combined HalfConv weight-and-bias packs session-local because their contents depend on the bias while the existing shared-prepack key identifies only W. Generic Conv weight packs remain shareable across sessions.
- Centralizes HalfGemm backend-selector handling in the public MLAS dispatcher and updates the affected MatMul, Gemm, Conv, and MoE call sites.
- Updates the CPU operator documentation to advertise FP16 support for `MatMul` and `Gemm`.

Test coverage includes:

- FP16 Gemm with no C input and with `beta == 0`.
- FP16 MatMul with constant and dynamic B inputs.
- MatMul backend-native B prepacking.
- Kernel-owned MatMul packs remaining intentionally unshared across sessions.
- Session-state validation for broken prepacking implementations.
- HalfConv-eligible NCHW and NHWC convolutions.
- HalfConv execution with and without bias.
- Correct fallback behavior when KleidiAI is explicitly disabled.
- Safe sharing of generic Conv prepacked weights while keeping bias-dependent HalfConv packs local.
- Updated HalfGemm selector and backend-native packed-B contract coverage.

Motivation and Context

This is pr 3 of 5 as per #28786 splitting suggestion
The preceding PR introduces the MLAS primitives for FP16 HalfGemm, backend-native B packing, and HalfConv, but standard ONNX CPU operators do not yet consume them.

This PR connects those primitives to CPU MatMul, Gemm, and Conv, making native FP16 execution and KleidiAI acceleration available to the operator layer while retaining portable fallback behavior for unsupported hardware, shapes, attributes, or backend configurations.

The prepacking changes avoid repeatedly packing constant weights while preventing backend-specific or bias-dependent layouts from being incorrectly reused through the shared prepack cache.

This PR is limited to CPU kernel registration and execution. The follow-up InsertCastTransformer PR introduces the opt-in policy and shape heuristics that decide when graph optimization should preserve FP16 CPU nodes. Apple ARM64 build enablement is also handled separately.

Laan33 and others added 12 commits July 13, 2026 16:42
Avoid holding a reference into lhs_ptrs_cache_by_pad across the lookup/update sequence so the cache remains keyed by the current pad buffer identity.

Source-commit: 923422f

Signed-off-by: Cathal Lawlor <cathal.lawlor@arm.com>
Native fp16 accumulation paths need a slightly wider tolerance, and WebGPU needs a separate bound for this coverage.

Source-commit: 2483990

Source-commit: 08f4c6a8c51dc018b4440c63d279bf5995d3386a

Signed-off-by: Cathal Lawlor <cathal.lawlor@arm.com>
The Resize(13) MLFloat16 test should not run against EPs that do not provide this kernel.

Source-commit: 2483990

Signed-off-by: Cathal Lawlor <cathal.lawlor@arm.com>
Allow the fp16 FusedConvWithSum transformer test to tolerate native fp16 numerical drift.

Source-commit: 2483990

Signed-off-by: Cathal Lawlor <cathal.lawlor@arm.com>
These internal-domain NHWC fp16 pool tests target ORT CPU/MLAS coverage and should not be offered to unrelated registered EPs.

Source-commit: a504a0c

Signed-off-by: Cathal Lawlor <cathal.lawlor@arm.com>
Signed-off-by: Cathal Lawlor <cathal.lawlor@arm.com>
- Add validation for null HalfGemm data parameters and reject invalid backend-native packed-B parameter combinations before the zero-K fast path.
- Make HalfGemm packed-B size calculation return a valid zero size for degenerate N/K shapes.
- Restore the KleidiAI header include needed by debug logging in mlasi_kleidiai.h.
- Update HalfGemm PackB tests to avoid exception-only gtest macros in no-exception builds, and add sentinel checks for early-return paths.

Signed-off-by: Cathal Lawlor <cathal.lawlor@arm.com>
Set MLAS_CONV_PARAMETERS::ChannelsLast in the HalfConv prepare path so
Prepare fully initializes the convolution parameter layout state.

Add a focused MLAS test that verifies Prepare overwrites stale layout flag
values for both NCHW and NHWC HalfConv configurations.

Signed-off-by: Cathal Lawlor <cathal.lawlor@arm.com>
Signed-off-by: Jonathan Clohessy <Jonathan.Clohessy@arm.com>
Signed-off-by: Jonathan Clohessy <Jonathan.Clohessy@arm.com>
Register CPU fp16 MatMul/Gemm kernels and route fp16 Conv through the native half-conv/prepack paths using the MLAS APIs from the previous branch.

This carries the current branch session_state.cc relaxation as WIP; the maintainer-requested sentinel modeling for backend-owned prepack buffers still needs to replace it before review.

Source-commit: 923422f

Source-commit: a504a0c

Source-commit: 152ef31

Source-commit: b81f7ab9f2e35b85d9e4f825c31aa4ddaad7e875

Signed-off-by: Cathal Lawlor <cathal.lawlor@arm.com>
Introduce an explicit kernel-owned packed-weights marker so fp16 MatMul can
report successful backend-owned prepacking without relaxing the existing shared
prepack invariant for broken kernels.

Update fp16 MatMul, Gemm, Conv, Moe, and MLAS test call sites for the branch-2
half-GEMM API shape where backend selector config is carried through
MLAS_HALF_GEMM_DATA_PARAMS.

Clarify HalfConv working-buffer sizing as bytes, zero-initialize
MLAS_CONV_PARAMETERS, and add/adjust tests for the fp16 MatMul kernel-owned
prepack path, broken prepack invariant, half-GEMM selector handling, and
HalfConv byte-size reporting.

Signed-off-by: Cathal Lawlor <cathal.lawlor@arm.com>
@azure-pipelines

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants