Repository navigation
[ARM CPU] Add MLFloat16 (FP16) support to CPU BiasGelu, MatMul and Gemm - #32370
abhijain1204fujitsu wants to merge 9 commits into
Conversation
|
@microsoft-github-policy-service agree [company="{fujitsu}"] |
There was a problem hiding this comment.
🟡 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.
| 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); |
| 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)>, |
| 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! |
| MLFloat16, | ||
| kCpuExecutionProvider, | ||
| KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<MLFloat16>()), | ||
| BiasGelu<MLFloat16, false>); |
| 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! |
| #ifdef MLAS_F16VEC_INTRINSICS_SUPPORTED | ||
| REGISTER_VERSIONED_UNARY_ELEMENTWISE_TYPED_KERNEL(Tanh, 6, 12, MLFloat16); | ||
| REGISTER_UNARY_ELEMENTWISE_TYPED_KERNEL(Tanh, 13, MLFloat16); |
| template <> | ||
| Status MatMul<MLFloat16>::Compute(OpKernelContext* ctx) const { |
a84da28 to
8d2da25
Compare
|
Hariharan Seshadri (@hariharans29) please support this PR. |
|
mirounga sir please support this PR. |
|
mirounga Sir please support this PR for f16 support. Thankyou. |
|
Hariharan Seshadri (@hariharans29) , mirounga Thanks ! |
|
abhijain1204fujitsu please read the following Contributor License Agreement(CLA). If you agree with the CLA, please reply with the following information.
Contributor License AgreementContribution License AgreementThis Contribution License Agreement (“Agreement”) is agreed to by the party signing below (“You”),
|
There was a problem hiding this comment.
🟡 Changes recommended
Several fusion, portability, test, and generated-documentation issues must be resolved before approval.
14 open findings
Prevent FP16 fusion when CPU kernel support is unavailable · New Preserve arithmetic type when removing normalization Casts · New The FP16 registration table stops at Gemm opset 12 even though this PR declares the opset-13 kernel… A schema-valid empty input such as[2, 0]with bias[0]makes bothelem_countandbias_len… Allocate only one scratch element for no-bias FastGelu · New Avoid strict-aliasing violation when reading MLFloat16 · New Remove or implement unreachable HGEMM prepack fallback · New Keep CPU Sigmoid FP16 fusion disabled · New Allowing every FP16 Gemm/activation pair is unsafe:FusedGemm<MLFloat16>calls the new… Regenerate operator kernel documentation for new FP16 kernels · New The existingMatMulFloat16test is compiled only for CUDA, CoreML, or XNNPACK builds, so an ARM… The existing FP16 Tanh test is guarded byUSE_CUDA || USE_COREML, so it is absent from the ARM… The new FP16 fusion path has no transformation or execution test; the existing Gemm-activation… The existing MLFloat16 BiasGelu test is compiled only with CUDA, DML, or WebGPU, so the newly added…
🧠 Review effort: Balanced
| #ifdef MLAS_F16VEC_INTRINSICS_SUPPORTED | ||
| if (data_type == ONNX_NAMESPACE::TensorProto_DataType_FLOAT16 && | ||
| !functors::IsFp16FusableActivation(next_node.OpType())) { | ||
| continue; | ||
| } | ||
| #endif |
| // 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())) { |
| BufferUniquePtr buffer = BufferUniquePtr( | ||
| alloc->Alloc(SafeInt<size_t>(sizeof(MLFloat16)) * elem_count * 2), | ||
| BufferDeleter(alloc)); |
| 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) |
| 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.
8deeb86 to
ff9a722
Compare
…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.
|
mirounga sir please support this PR further. |



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