Skip to content

[ARM CPU] Add MLFloat16 (FP16) support to CPU BiasGelu, MatMul and Gemm - #32370

Open
abhijain1204fujitsu wants to merge 9 commits into
microsoft:mainfrom
MonakaResearch:nikhil/fusion_support_f16
Open

abhijain1204fujitsu wants to merge 9 commits into
microsoft:mainfrom
MonakaResearch:nikhil/fusion_support_f16

Conversation

@abhijain1204fujitsu

@abhijain1204fujitsu abhijain1204fujitsu commented Sep 2, 2026 •

Copy link
Copy Markdown

Description

Add FP16 CPU support for BiasGelu, MatMul, Gemm, FusedGemm, and Tanh

This PR enables FP16 transformer execution on CPU without unnecessary float fallbacks by:

  • Adding FP16 support for BiasGelu, MatMul, Gemm, FusedGemm, and Tanh.
  • Registering missing MLFloat16 kernels across supported opsets.
  • Leveraging MLAS HGEMM/HalfGemm implementations where available, with fallback paths preserved.
  • Enabling FP16 Gemm+Activation fusion

Tests: extend GemmOpTypedTests to MLFloat16. The expected-value
computation now accumulates in float and converts once, matching the
kernel's rounding behaviour.

This PR is a joint contribution by:
Nikhil Sharma (@nikhilfujitsu), Aruna K (@akote123) & abhijain1204fujitsu

Copilot AI balanced review requested due to automatic review settings September 2, 2026 02:37
@abhijain1204fujitsu

Copy link
Copy Markdown
Author

@microsoft-github-policy-service agree [company="{fujitsu}"]

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Changes recommended

Opset-13 Gemm is not registered, unsupported FP16 activations can be fused, and several new CPU paths lack targeted tests.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Adds ARM CPU FP16 support for key transformer operators using MLAS acceleration and fallback paths.

Changes:

  • Adds FP16 MatMul, Gemm, BiasGelu, Tanh, and FusedGemm kernels.
  • Enables FP16 Gemm-activation fusion.
  • Extends Gemm typed tests to MLFloat16.
File summaries
File Description
onnxruntime/test/providers/cpu/math/gemm_test.cc Adds MLFloat16 Gemm coverage.
onnxruntime/core/providers/cpu/math/matmul.cc Implements FP16 MatMul.
onnxruntime/core/providers/cpu/math/gemm.cc Expands FP16 Gemm dispatch.
onnxruntime/core/providers/cpu/cpu_execution_provider.cc Registers FP16 CPU kernels.
onnxruntime/core/providers/cpu/activation/activations.h Declares FP16 Tanh specialization.
onnxruntime/core/providers/cpu/activation/activations.cc Implements and registers FP16 Tanh.
onnxruntime/core/optimizer/gemm_activation_fusion.cc Allows FP16 Gemm fusion.
onnxruntime/contrib_ops/cpu/fused_gemm.cc Registers FP16 FusedGemm.
onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc Registers FP16 contrib kernels.
onnxruntime/contrib_ops/cpu/bert/bias_gelu.cc Implements FP16 BiasGelu.
Review details
  • Files reviewed: 10/10 changed files
  • Comments generated: 7
  • Review effort level: Balanced

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Comment on lines +171 to +188
const MLFloat16* bias_data = bias->Data<MLFloat16>();
int64_t bias_len = bias->Shape().Size();

AllocatorPtr alloc;
ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&alloc));
BufferUniquePtr buffer = BufferUniquePtr(
alloc->Alloc(SafeInt<size_t>(sizeof(MLFloat16)) * elem_count),
BufferDeleter(alloc));
MLFloat16* tmp_data = static_cast<MLFloat16*>(buffer.get());

int64_t task_count = elem_count / bias_len;
concurrency::ThreadPool::TryBatchParallelFor(
context->GetOperatorThreadPool(), static_cast<int32_t>(task_count),
[&](ptrdiff_t task_idx) {
const MLFloat16* p_input = input_data + task_idx * bias_len;
MLFloat16* p_output = output_data + task_idx * bias_len;
MLFloat16* p_tmp = tmp_data + task_idx * bias_len;
AddBiasGelu(p_input, bias_data, p_tmp, p_output, bias_len);
Comment on lines +3759 to +3762
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, 10,
MLFloat16, Gemm)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 12,
MLFloat16, Gemm)>,
Comment on lines +61 to +63
if (data_type != ONNX_NAMESPACE::TensorProto_DataType_FLOAT &&
data_type != ONNX_NAMESPACE::TensorProto_DataType_FLOAT16) {
// FusedGemm is only registered for float and MLFloat16 data types in fused_gemm.cc!
Comment on lines +33 to +36
MLFloat16,
kCpuExecutionProvider,
KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<MLFloat16>()),
BiasGelu<MLFloat16, false>);
Comment on lines +61 to +63
if (data_type != ONNX_NAMESPACE::TensorProto_DataType_FLOAT &&
data_type != ONNX_NAMESPACE::TensorProto_DataType_FLOAT16) {
// FusedGemm is only registered for float and MLFloat16 data types in fused_gemm.cc!
Comment on lines +69 to +71
#ifdef MLAS_F16VEC_INTRINSICS_SUPPORTED
REGISTER_VERSIONED_UNARY_ELEMENTWISE_TYPED_KERNEL(Tanh, 6, 12, MLFloat16);
REGISTER_UNARY_ELEMENTWISE_TYPED_KERNEL(Tanh, 13, MLFloat16);
Comment on lines +212 to +213
template <>
Status MatMul<MLFloat16>::Compute(OpKernelContext* ctx) const {
@nikhilfujitsu

Copy link
Copy Markdown
Contributor

Hariharan Seshadri (@hariharans29) please support this PR.

@nikhilfujitsu

Copy link
Copy Markdown
Contributor

mirounga sir please support this PR.

@mirounga
mirounga self-requested a review September 16, 2026 17:20
@nikhilfujitsu

Copy link
Copy Markdown
Contributor

mirounga Sir please support this PR for f16 support. Thankyou.

@abhijain1204fujitsu

Copy link
Copy Markdown
Author

Hariharan Seshadri (@hariharans29) , mirounga
Kindly support to review the PR

Thanks !

@mirounga
mirounga requested a balanced review from Copilot October 13, 2026 20:10
@microsoft-github-policy-service

Copy link
Copy Markdown
Contributor

abhijain1204fujitsu please read the following Contributor License Agreement(CLA). If you agree with the CLA, please reply with the following information.

@microsoft-github-policy-service agree [company="{your company}"]

Options:

  • (default - no company specified) I have sole ownership of intellectual property rights to my Submissions and I am not making Submissions in the course of work for my employer.
@microsoft-github-policy-service agree
  • (when company given) I am making Submissions in the course of work for my employer (or my employer has intellectual property rights in my Submissions by contract or applicable law). I have permission from my employer to make Submissions and enter into this Agreement on behalf of my employer. By signing below, the defined term “You” includes me and my employer.
@microsoft-github-policy-service agree company="Microsoft"
Contributor License Agreement

Contribution License Agreement

This Contribution License Agreement (“Agreement”) is agreed to by the party signing below (“You”),
and conveys certain license rights to Microsoft Corporation and its affiliates (“Microsoft”) for Your
contributions to Microsoft open source projects. This Agreement is effective as of the latest signature
date below.

  1. Definitions.
    “Code” means the computer software code, whether in human-readable or machine-executable form,
    that is delivered by You to Microsoft under this Agreement.
    “Project” means any of the projects owned or managed by Microsoft and offered under a license
    approved by the Open Source Initiative (www.opensource.org).
    “Submit” is the act of uploading, submitting, transmitting, or distributing code or other content to any
    Project, including but not limited to communication on electronic mailing lists, source code control
    systems, and issue tracking systems that are managed by, or on behalf of, the Project for the purpose of
    discussing and improving that Project, but excluding communication that is conspicuously marked or
    otherwise designated in writing by You as “Not a Submission.”
    “Submission” means the Code and any other copyrightable material Submitted by You, including any
    associated comments and documentation.
  2. Your Submission. You must agree to the terms of this Agreement before making a Submission to any
    Project. This Agreement covers any and all Submissions that You, now or in the future (except as
    described in Section 4 below), Submit to any Project.
  3. Originality of Work. You represent that each of Your Submissions is entirely Your original work.
    Should You wish to Submit materials that are not Your original work, You may Submit them separately
    to the Project if You (a) retain all copyright and license information that was in the materials as You
    received them, (b) in the description accompanying Your Submission, include the phrase “Submission
    containing materials of a third party:” followed by the names of the third party and any licenses or other
    restrictions of which You are aware, and (c) follow any other instructions in the Project’s written
    guidelines concerning Submissions.
  4. Your Employer. References to “employer” in this Agreement include Your employer or anyone else
    for whom You are acting in making Your Submission, e.g. as a contractor, vendor, or agent. If Your
    Submission is made in the course of Your work for an employer or Your employer has intellectual
    property rights in Your Submission by contract or applicable law, You must secure permission from Your
    employer to make the Submission before signing this Agreement. In that case, the term “You” in this
    Agreement will refer to You and the employer collectively. If You change employers in the future and
    desire to Submit additional Submissions for the new employer, then You agree to sign a new Agreement
    and secure permission from the new employer before Submitting those Submissions.
  5. Licenses.
  • Copyright License. You grant Microsoft, and those who receive the Submission directly or
    indirectly from Microsoft, a perpetual, worldwide, non-exclusive, royalty-free, irrevocable license in the
    Submission to reproduce, prepare derivative works of, publicly display, publicly perform, and distribute
    the Submission and such derivative works, and to sublicense any or all of the foregoing rights to third
    parties.
  • Patent License. You grant Microsoft, and those who receive the Submission directly or
    indirectly from Microsoft, a perpetual, worldwide, non-exclusive, royalty-free, irrevocable license under
    Your patent claims that are necessarily infringed by the Submission or the combination of the
    Submission with the Project to which it was Submitted to make, have made, use, offer to sell, sell and
    import or otherwise dispose of the Submission alone or with the Project.
  • Other Rights Reserved. Each party reserves all rights not expressly granted in this Agreement.
    No additional licenses or rights whatsoever (including, without limitation, any implied licenses) are
    granted by implication, exhaustion, estoppel or otherwise.
  1. Representations and Warranties. You represent that You are legally entitled to grant the above
    licenses. You represent that each of Your Submissions is entirely Your original work (except as You may
    have disclosed under Section 3). You represent that You have secured permission from Your employer to
    make the Submission in cases where Your Submission is made in the course of Your work for Your
    employer or Your employer has intellectual property rights in Your Submission by contract or applicable
    law. If You are signing this Agreement on behalf of Your employer, You represent and warrant that You
    have the necessary authority to bind the listed employer to the obligations contained in this Agreement.
    You are not expected to provide support for Your Submission, unless You choose to do so. UNLESS
    REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING, AND EXCEPT FOR THE WARRANTIES
    EXPRESSLY STATED IN SECTIONS 3, 4, AND 6, THE SUBMISSION PROVIDED UNDER THIS AGREEMENT IS
    PROVIDED WITHOUT WARRANTY OF ANY KIND, INCLUDING, BUT NOT LIMITED TO, ANY WARRANTY OF
    NONINFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
  2. Notice to Microsoft. You agree to notify Microsoft in writing of any facts or circumstances of which
    You later become aware that would make Your representations in this Agreement inaccurate in any
    respect.
  3. Information about Submissions. You agree that contributions to Projects and information about
    contributions may be maintained indefinitely and disclosed publicly, including Your name and other
    information that You submit with Your Submission.
  4. Governing Law/Jurisdiction. This Agreement is governed by the laws of the State of Washington, and
    the parties consent to exclusive jurisdiction and venue in the federal courts sitting in King County,
    Washington, unless no federal subject matter jurisdiction exists, in which case the parties consent to
    exclusive jurisdiction and venue in the Superior Court of King County, Washington. The parties waive all
    defenses of lack of personal jurisdiction and forum non-conveniens.
  5. Entire Agreement/Assignment. This Agreement is the entire agreement between the parties, and
    supersedes any and all prior agreements, understandings or communications, written or oral, between
    the parties relating to the subject matter hereof. This Agreement may be assigned by Microsoft.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Comment on lines +73 to +78
#ifdef MLAS_F16VEC_INTRINSICS_SUPPORTED
if (data_type == ONNX_NAMESPACE::TensorProto_DataType_FLOAT16 &&
!functors::IsFp16FusableActivation(next_node.OpType())) {
continue;
}
#endif
Comment on lines +559 to +562
// CPU kernels need X, scale and bias to be the same type.
if (reduce_mean_node.GetExecutionProviderType() == kCpuExecutionProvider &&
x_input->TypeAsProto()->tensor_type().elem_type() == ONNX_NAMESPACE::TensorProto_DataType_FLOAT16) {
(x_input->TypeAsProto()->tensor_type().elem_type() != scale->TypeAsProto()->tensor_type().elem_type() ||
x_input->TypeAsProto()->tensor_type().elem_type() != bias->TypeAsProto()->tensor_type().elem_type())) {
Comment on lines +231 to +233
BufferUniquePtr buffer = BufferUniquePtr(
alloc->Alloc(SafeInt<size_t>(sizeof(MLFloat16)) * elem_count * 2),
BufferDeleter(alloc));
Comment on lines +259 to +260
data.alpha = *reinterpret_cast<const _mlas_fp16*>(&alpha);
data.beta = *reinterpret_cast<const _mlas_fp16*>(&beta);
size_t packed_b_size = 0;
is_packed = GemmPackBHalfNative(alloc, tensor, packed_b_, packed_b_size, b_shape_,
&mlas_backend_kernel_selector_config_);
#if defined(MLAS_HGEMM_PACKB_SUPPORTED)
}

#if defined(USE_CUDA) || defined(USE_COREML)
#if defined(USE_CUDA) || defined(USE_COREML) || defined(MLAS_F16VEC_INTRINSICS_SUPPORTED)
Comment on lines +398 to +402
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MLFloat16, BiasGelu)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Gelu)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MLFloat16, Gelu)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, FastGelu)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MLFloat16, FastGelu)>,
…ath for Gemm

Extends the CPU EP's FP16 coverage to the operators that an FP16 transformer
graph still falls back to float for, and routes Gemm through the HGEMM
kernels. MatMul and the FP16 Gemm kernel registrations are already on main,
so this builds on those rather than duplicating them.

Operators
  * BiasGelu (kMSDomain): split the registration into typed float / MLFloat16
    variants and add MLFloat16 specializations of AddBiasGelu and Compute.
    They accumulate in float and round once per element; Compute drops the
    FastGelu no-bias branch, which the non-approximation variant never takes.
  * FusedGemm: register the MLFloat16 variant so GemmActivationFusion can
    fuse FP16 Gemm+activation subgraphs.
  * Tanh: register the MLFloat16 CPU kernel (opsets 6-12 and 13) and route it
    to the FP16 MLAS tanh, so FP16 activations no longer force a cast to
    float.
  * Gemm: Gemm_MLFloat16 now tries MlasHGemmSupported/MlasGemmBatch first,
    which handles TransA/TransB and arbitrary alpha/beta with the bias
    pre-broadcast into C, before falling back to the existing
    MlasHalfGemmBatch fast path and then to Eigen.

Gemm+activation fusion
  IsFusableActivation() accepts more activations than
  ElementWiseRangedTransform<MLFloat16>::Create() can build, so an FP16
  Gemm->Relu or Gemm->LeakyRelu graph was fused into a FusedGemm<MLFloat16>
  that threw NOT_IMPLEMENTED at session initialization. Relu and LeakyRelu
  were the reachable cases: they are the only other fusable activations with
  an MLFloat16 CPU kernel, so the rest are cast to FP32 and never fuse.

  The MLFloat16 factory is extended to the activations that do have MLFloat16
  kernels (Relu, LeakyRelu, Tanh), and the FP16 fusion is gated on the
  factory's own list so the two cannot drift. IsFp16FusableActivation() lives
  in a new activation_fp16_support.h, deliberately free of Eigen so
  core/optimizer can include it without pulling in the CPU kernel headers.

MLAS FP16 NEON fixes
  * MlasLoadPartialFloat16x4()/MlasStorePartialFloat16x4() decompose their
    length into its 1- and 2-bits, so len == 4 matched nothing: the load
    returned zeros and the store wrote nothing. HPackB_B_Kernel() hits this
    in its CountN <= 4 tail, so B packed as all zeros whenever the N tail was
    exactly 4, and any FP16 GEMM with M >= 3 (the packed-B path) and
    N % 8 == 4 silently produced zeros for that 8-column block. A sweep over
    M 1-4 x N 1-20 x K 1-6 found 36 such shapes. The existing NEON HGEMM
    tests never covered it: their N values are 1, 15, 17, 31, 32, 33, 79 and
    263, none congruent to 4 mod 8. Both helpers now treat len >= 4 as a full
    4-lane access, and the pack tail zero-fills lanes 4-7 of each 8-wide row
    instead of leaving them uninitialized.
  * Gemm_MLFloat16 zeroes y_data when beta == 0 rather than relying on the
    kernel to overwrite it. HGemm_B_Kernel_Complicated folds beta in as
    fma(mul(C, beta), accu, alpha) at every store, so it reads C whatever
    beta is, and a stale NaN bit pattern survived 0 * NaN. Only alpha != 1
    reaches that path, which is why the existing FP16 MLAS entry points
    (alpha == 1 only) never saw it. This showed up as
    GemmOpTypedTests/2.TestGemmZeroBeta, TestGemmZeroAlphaBeta and
    TestGemmWithAlphaOpset11 failing about 1 run in 20.

Tests
  The existing FP16 tests are guarded on CUDA/DML/CoreML/WebGPU/XNNPACK, so
  none of this had CPU coverage.
  * graph_transform_test.cc: fusion is applied for Relu/LeakyRelu/Tanh and
    refused for Sigmoid/Elu/Softplus/Softsign/HardSigmoid/Selu.
  * element_wise_ops_test.cc: FusedGemmTest.MLFloat16Activations constructs
    and runs FusedGemm<MLFloat16> for each supported activation, so a future
    drift between the fusion and the factory fails here rather than in a
    model. The BiasGelu FP16 tests now build for
    MLAS_F16VEC_INTRINSICS_SUPPORTED, plus an empty-tensor case for the
    MLFloat16 Compute specialization.
  * activation_op_test.cc, matmul_test.cc: build the existing FP16 tests for
    MLAS_F16VEC_INTRINSICS_SUPPORTED, and add a zero-K MatMul FP16 case.
  * gemm_test.cc: extend GemmOpTypedTests to MLFloat16; the expected values
    now accumulate in float and convert once, matching the kernel.
Found by building and testing the branch on x86-64, where
MLAS_F16VEC_INTRINSICS_SUPPORTED is not defined.

  * Pow: only add MLFloat16 to the CPU default types when
    MLAS_F16VEC_INTRINSICS_SUPPORTED is defined, like the other new fp16
    kernels (Sqrt, ReduceMean, Softmax, Tanh, FusedGemm). Otherwise Pow
    stayed fp16 while the rest of an fp16 subgraph was cast to fp32, which
    broke SimplifiedLayerNormFusion on that fp32 subgraph
    (SimplifiedLayerNormFusionOptimizationLoopSharedPowExponent).
  * SimplifiedLayerNormFusion: on CPU, if absorbing the leading Cast would
    give the fused node mixed input/scale types, fuse the post-Cast subgraph
    and keep the Cast, as before the CPU leading-Cast handling was added.
  * insert_cast_transformer_test: Fp16FusedNodeWithNoCpuKernelForcedToFp32
    relied on BiasGelu having no fp16 CPU kernel, which this branch adds.
    Use QuickGelu, which still only has a float CPU kernel.
  * docs/OperatorKernels.md: add tensor(float16) to the CPU BiasGelu,
    FastGelu, FusedMatMul and Gelu rows.
  * clang-format fixes in gemm.cc, reduction_ops.cc and gelu.cc.
  * GemmActivationFusion: skip FP16 Gemm when MLAS_F16VEC_INTRINSICS_SUPPORTED
    is not defined, since FusedGemm<MLFloat16> is only registered under it.
  * FastGelu<MLFloat16>: the no-bias path only needs one scratch element per
    input; allocate two only when there is a bias.
  * Gemm_MLFloat16: pass alpha/beta through MLFloat16::val instead of reading
    the object through a uint16_t pointer.
  * activation_op_test: keep Sigmoid_fp16 under the CUDA/CoreML guard. There
    is no CPU MLFloat16 Sigmoid kernel, so on CPU builds it ran nothing.
…PORTED

The field is only read and written inside MLAS_HGEMM_PACKB_SUPPORTED
blocks in matmul.cc. Without that macro, Clang rejects the unused private
field (-Wunused-private-field), which breaks every Clang-based CI build.
@nikhilfujitsu

Copy link
Copy Markdown
Contributor

mirounga sir please support this PR further.

This branch has not been deployed

No deployments
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.

5 participants