From 8f7e6010205ef470ce072dc4de6b54922dee1f63 Mon Sep 17 00:00:00 2001 From: Cathal Lawlor Date: Wed, 27 May 2026 10:24:39 +0100 Subject: [PATCH 01/10] Fix KleidiAI conv LHS pointer cache lookup 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: 923422fea2556b7c19f3b4ef7166173a612aa0d7 Signed-off-by: Cathal Lawlor --- onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp b/onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp index 4f62ab24648ef..1367ba1d3d7f4 100644 --- a/onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp +++ b/onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp @@ -459,12 +459,13 @@ static std::shared_ptr GetOrCreateLhsPtrTableSme(const size_t ci, 1, 1 }; - // Setting up LHS ptr cache and tracking value of last passed pad_ptr + // Cache of computed LHS pointer offsets. thread_local to prevent interference from parallel sessions. + // Entries include pointers to the pad buffer for out-of-bounds pixels, so group by pad buffer identity and + // invalidate the old group if the grow-only pad buffer is reallocated. using LhsPtrsCache = std::unordered_map>; thread_local std::unordered_map lhs_ptrs_cache_by_pad; thread_local const float* last_pad_ptr = nullptr; - // If pad_ptr moved (vector reallocation), drop only the old group to avoid accumulating unreachable entries. const float* cur_pad_ptr = pad_ptr; if (last_pad_ptr != nullptr && last_pad_ptr != cur_pad_ptr) { lhs_ptrs_cache_by_pad.erase(last_pad_ptr); From f1f83ecd3e4b6a3006f6fba3915d69cc073a70e3 Mon Sep 17 00:00:00 2001 From: Cathal Lawlor Date: Wed, 27 May 2026 10:24:55 +0100 Subject: [PATCH 02/10] Adjust fp16 MatMulNBits tolerance Native fp16 accumulation paths need a slightly wider tolerance, and WebGPU needs a separate bound for this coverage. Source-commit: 248399052eda5d3798b48cc0ea4d13a3a2aa9a75 Source-commit: 08f4c6a8c51dc018b4440c63d279bf5995d3386a Signed-off-by: Cathal Lawlor --- onnxruntime/test/contrib_ops/matmul_4bits_test.cc | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc index 51895938ce3aa..de7cab33febf8 100644 --- a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc @@ -328,7 +328,11 @@ void TestMatMulNBitsTyped(std::optional abs_error = std::nullopt, } else if (base_opts.accuracy_level == 4) { base_opts.output_abs_error = 0.1f; } else if constexpr (std::is_same::value) { - base_opts.output_abs_error = 0.055f; +#if defined(USE_WEBGPU) + base_opts.output_abs_error = 0.1f; +#else + base_opts.output_abs_error = 0.065f; +#endif } else { base_opts.output_abs_error = 0.05f; } From 89e41bc42ac1f84879fdae0fe0f88bf31f6535db Mon Sep 17 00:00:00 2001 From: Cathal Lawlor Date: Wed, 27 May 2026 10:24:55 +0100 Subject: [PATCH 03/10] Exclude EPs without fp16 Resize coverage The Resize(13) MLFloat16 test should not run against EPs that do not provide this kernel. Source-commit: 248399052eda5d3798b48cc0ea4d13a3a2aa9a75 Signed-off-by: Cathal Lawlor --- onnxruntime/test/providers/cpu/tensor/resize_op_test.cc | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/onnxruntime/test/providers/cpu/tensor/resize_op_test.cc b/onnxruntime/test/providers/cpu/tensor/resize_op_test.cc index 662ec5693e6c4..2620c486a4e38 100644 --- a/onnxruntime/test/providers/cpu/tensor/resize_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/resize_op_test.cc @@ -3,6 +3,7 @@ #include #include +#include #include "gtest/gtest.h" #include "test/providers/provider_test_utils.h" #include "test/util/include/default_providers.h" @@ -254,6 +255,14 @@ TYPED_TEST(ResizeOpTest, ResizeOpLinearDownSampleTest_4DBilinear) { // QNN: result diff // TRT: Segmentation fault in A100 std::unordered_set excluded_providers({kQnnExecutionProvider}); + if constexpr (std::is_same_v) { + // These EPs do not provide this Resize(13) MLFloat16 kernel. + excluded_providers.insert(kNnapiExecutionProvider); + excluded_providers.insert(kXnnpackExecutionProvider); + excluded_providers.insert(kCoreMLExecutionProvider); + excluded_providers.insert(kTensorrtExecutionProvider); + excluded_providers.insert("example_ep"); + } test.Run(OpTester::ExpectResult::kExpectSuccess, "", ExcludeTrtOnA100(excluded_providers)); }; From 23ed6ac5806d1e3cea0036166cd135e862ab5c2a Mon Sep 17 00:00:00 2001 From: Cathal Lawlor Date: Wed, 27 May 2026 10:24:55 +0100 Subject: [PATCH 04/10] Loosen NHWC fp16 fused-conv tolerance Allow the fp16 FusedConvWithSum transformer test to tolerate native fp16 numerical drift. Source-commit: 248399052eda5d3798b48cc0ea4d13a3a2aa9a75 Signed-off-by: Cathal Lawlor --- onnxruntime/test/optimizer/nhwc_transformer_test.cc | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/onnxruntime/test/optimizer/nhwc_transformer_test.cc b/onnxruntime/test/optimizer/nhwc_transformer_test.cc index eb55f1d7d46ab..b9d8d478fdff2 100644 --- a/onnxruntime/test/optimizer/nhwc_transformer_test.cc +++ b/onnxruntime/test/optimizer/nhwc_transformer_test.cc @@ -1131,7 +1131,10 @@ TEST_F(NhwcTransformerTestsFp16, FusedConvWithSumFp16) { TransformerTester(build_test_case, check_nhwc_graph, TransformerLevel::Level2, - TransformerLevel::Level3); + TransformerLevel::Level3, + /*opset_version*/ 12, + /*per_sample_tolerance*/ 0.02, + /*relative_per_sample_tolerance*/ 0.02); } TEST_F(NhwcTransformerTestsFp16, ConvMaxPoolFp16) { From 576ce4a90b72a3f4c0b7b6baf8979f77c8c77d78 Mon Sep 17 00:00:00 2001 From: Cathal Lawlor Date: Wed, 27 May 2026 10:24:55 +0100 Subject: [PATCH 05/10] Constrain NHWC fp16 pool tests to owning EPs These internal-domain NHWC fp16 pool tests target ORT CPU/MLAS coverage and should not be offered to unrelated registered EPs. Source-commit: a504a0c3c2e11f45764a002e8d60785472b92b92 Signed-off-by: Cathal Lawlor --- .../test/contrib_ops/nhwc_pool_in_op_test.cc | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/onnxruntime/test/contrib_ops/nhwc_pool_in_op_test.cc b/onnxruntime/test/contrib_ops/nhwc_pool_in_op_test.cc index e40e635c79a26..c58d813e1b061 100644 --- a/onnxruntime/test/contrib_ops/nhwc_pool_in_op_test.cc +++ b/onnxruntime/test/contrib_ops/nhwc_pool_in_op_test.cc @@ -7,6 +7,7 @@ #include #include +#include #include "core/util/math.h" #include "core/mlas/inc/mlas.h" @@ -17,6 +18,16 @@ namespace onnxruntime { namespace test { #ifdef MLAS_F16VEC_INTRINSICS_SUPPORTED +namespace { + +// 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 kNhwcFp16PoolExcludedProviders{ + kCoreMLExecutionProvider, + kTensorrtExecutionProvider, +}; + class NhwcFp16PoolOpTester { private: bool is_max_pool_; // max or average pool @@ -181,10 +192,12 @@ class NhwcFp16PoolOpTester { if (!dilations_.empty()) { test.AddAttribute("dilations", dilations_); } - test.Run(OpTester::ExpectResult::kExpectSuccess, ""); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", kNhwcFp16PoolExcludedProviders); } }; +} // namespace + TEST(NhwcFp16PoolOpTest, MaxPool1D) { for (int64_t channels = 1; channels < 94; channels++) { NhwcFp16PoolOpTester test(true); @@ -303,7 +316,7 @@ TEST(NhwcFp16PoolOpTest, AvgPoolIncludePadPixel) { test.AddInput("X", x_dims, x_vals); test.AddOutput("Y", expected_dims, expected_vals); - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", kNhwcFp16PoolExcludedProviders); } TEST(NhwcFp16PoolOpTest, GlobalAveragePool) { @@ -508,7 +521,7 @@ TEST(NhwcFp16PoolOpTest, GlobalAveragePool) { test.AddInput("X", x_dims, x_vals); test.AddOutput("Y", expected_dims, expected_vals); - test.Run(); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", kNhwcFp16PoolExcludedProviders); } #endif From 6dde25e9596761bb59149c0b20e9424e8fbd79b4 Mon Sep 17 00:00:00 2001 From: Cathal Lawlor Date: Tue, 2 Jun 2026 15:22:10 +0100 Subject: [PATCH 06/10] Add MLAS Arm64 half API support Signed-off-by: Cathal Lawlor --- cmake/onnxruntime_mlas.cmake | 3 + onnxruntime/core/mlas/inc/mlas.h | 113 +- onnxruntime/core/mlas/lib/halfconv.cpp | 158 +++ onnxruntime/core/mlas/lib/halfgemm.cpp | 164 ++- onnxruntime/core/mlas/lib/halfgemm.h | 94 +- .../core/mlas/lib/halfgemm_kernel_neon.cpp | 2 +- .../core/mlas/lib/kai_ukernel_interface.cpp | 34 + .../core/mlas/lib/kai_ukernel_interface.h | 18 +- .../mlas/lib/kleidiai/halfconv_kleidiai.cpp | 1040 +++++++++++++++++ .../mlas/lib/kleidiai/halfgemm_kleidiai.cpp | 300 +++++ .../core/mlas/lib/kleidiai/mlasi_kleidiai.h | 109 +- onnxruntime/core/mlas/lib/mlasi.h | 133 +++ onnxruntime/core/mlas/lib/platform.cpp | 9 +- .../test/mlas/unittest/test_conv2d.cpp | 95 ++ .../test/mlas/unittest/test_halfgemm.cpp | 746 ++++++++++++ .../test/mlas/unittest/test_halfgemm.h | 144 ++- onnxruntime/test/mlas/unittest/test_util.h | 2 + 17 files changed, 3140 insertions(+), 24 deletions(-) create mode 100644 onnxruntime/core/mlas/lib/halfconv.cpp create mode 100644 onnxruntime/core/mlas/lib/kleidiai/halfconv_kleidiai.cpp create mode 100644 onnxruntime/core/mlas/lib/kleidiai/halfgemm_kleidiai.cpp diff --git a/cmake/onnxruntime_mlas.cmake b/cmake/onnxruntime_mlas.cmake index 319c897b8349c..7876ec8d1cff2 100644 --- a/cmake/onnxruntime_mlas.cmake +++ b/cmake/onnxruntime_mlas.cmake @@ -23,6 +23,7 @@ onnxruntime_add_static_library(onnxruntime_mlas ${MLAS_SRC_DIR}/threading.cpp ${MLAS_SRC_DIR}/sgemm.cpp ${MLAS_SRC_DIR}/halfgemm.cpp + ${MLAS_SRC_DIR}/halfconv.cpp ${MLAS_SRC_DIR}/qgemm.cpp ${MLAS_SRC_DIR}/qdwconv.cpp ${MLAS_SRC_DIR}/convolve.cpp @@ -326,6 +327,8 @@ function(setup_kleidiai) target_sources(onnxruntime_mlas PRIVATE ${MLAS_SRC_DIR}/kai_ukernel_interface.cpp ${MLAS_SRC_DIR}/kleidiai/sgemm_kleidiai.cpp + ${MLAS_SRC_DIR}/kleidiai/halfgemm_kleidiai.cpp + ${MLAS_SRC_DIR}/kleidiai/halfconv_kleidiai.cpp ${MLAS_SRC_DIR}/kleidiai/sbgemm_kleidiai.cpp ${MLAS_SRC_DIR}/kleidiai/convolve_kleidiai.cpp ${MLAS_SRC_DIR}/kleidiai/qgemm_kleidiai.cpp diff --git a/onnxruntime/core/mlas/inc/mlas.h b/onnxruntime/core/mlas/inc/mlas.h index cbca2d85a97a4..f9f5cc44e8d52 100644 --- a/onnxruntime/core/mlas/inc/mlas.h +++ b/onnxruntime/core/mlas/inc/mlas.h @@ -911,6 +911,7 @@ struct MLAS_CONV_PARAMETERS { const void* PackedFilter = nullptr; size_t PackedFilterGroupStride = 0; bool FilterIsPacked = false; + bool InputOutputChannelsLast = false; union { struct { CBLAS_TRANSPOSE TransB; @@ -1935,11 +1936,21 @@ struct MLAS_HALF_GEMM_DATA_PARAMS { const MLAS_HALF_GEMM_POSTPROCESSOR* OutputProcessor = nullptr; bool AIsfp32 = false; /**< matrix A is fp32, needs to be casted into fp16*/ bool BIsfp32 = false; /**< matrix B is fp32, needs to be casted into fp16*/ + bool BIsPacked = false; /**< matrix B is pre-packed by MlasHalfGemmPackB/MlasHalfGemmConvertPackB */ + const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig = nullptr; + /** + * Matrix B uses a backend-specific direct-consumption packed layout. + * When true, B must be produced by MlasHalfGemmNativePackB, ldb must be 0, + * Bias must be nullptr, and OutputProcessor must be nullptr. + */ + bool BIsBackendNativePacked = false; }; /** * @brief Half precision Batched GEMM: C = A * B + Bias - * Either A or B can be fp32 or fp16 + * Either A or B can be fp32 or fp16. + * Backend-native packed B is a constrained direct-consumption layout + * and does not support runtime Bias or OutputProcessor. * * Note: We only support uniform batching, so shapes and types of the * input must be same across all parameter blocks. @@ -2001,6 +2012,106 @@ MlasHalfGemmPackB( void* PackedB ); +/** + * @brief For half precision GEMM, returns the size of a backend-native + * direct-consumption packing buffer for right hand side B. + * + * The returned layout is only valid when MlasHalfGemmBatch consumes + * MLAS_HALF_GEMM_DATA_PARAMS with BIsBackendNativePacked set to true, ldb set + * to 0, Bias set to nullptr, and OutputProcessor set to nullptr. Returns 0 + * when no backend-native format is available for the given transpose/shape/config. + */ +size_t +MLASCALL +MlasHalfGemmNativePackBSize( + CBLAS_TRANSPOSE TransA, + CBLAS_TRANSPOSE TransB, + size_t N, + size_t K, + const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig + ); + +/** + * @brief For half precision GEMM, pack right hand side B into the + * backend-native direct-consumption format. + * + * Returns false when the backend-native format is unavailable for the given + * transpose/shape/config. The resulting buffer must be passed to + * MlasHalfGemmBatch with ldb set to 0, BIsBackendNativePacked set to true, + * Bias set to nullptr, and OutputProcessor set to nullptr. + */ +bool +MLASCALL +MlasHalfGemmNativePackB( + CBLAS_TRANSPOSE TransA, + CBLAS_TRANSPOSE TransB, + size_t N, + size_t K, + const MLAS_FP16* B, + size_t ldb, + void* PackedB, + const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig + ); + +bool +MLASCALL +MlasHalfConvPrepare( + MLAS_CONV_PARAMETERS* Parameters, + size_t Dimensions, + size_t BatchCount, + size_t GroupCount, + size_t InputChannels, + const int64_t* InputShape, + const int64_t* KernelShape, + const int64_t* DilationShape, + const int64_t* Padding, + const int64_t* StrideShape, + const int64_t* OutputShape, + size_t FilterCount, + const MLAS_ACTIVATION* Activation, + size_t* WorkingBufferSize, + float Beta, + bool InputOutputChannelsLast, + MLAS_THREADPOOL* ThreadPool, + const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig); + +bool +MLASCALL +MlasHalfConv( + const MLAS_CONV_PARAMETERS* Parameters, + const MLAS_FP16* Input, + const MLAS_FP16* Filter, + // When true, Filter must point to a packed weights+bias buffer produced by + // MlasHalfConvPackWeightsAndBias and Bias must be nullptr. + bool FilterAndBiasArePacked, + const MLAS_FP16* Bias, + MLAS_FP16* WorkingBuffer, + MLAS_FP16* Output, + MLAS_THREADPOOL* ThreadPool); + +size_t +MLASCALL +MlasHalfConvPackWeightsAndBiasSize( + size_t FilterCount, + size_t InputChannels, + const int64_t* KernelShape, + const int64_t* DilationShape, + const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig); + +bool +MLASCALL +MlasHalfConvPackWeightsAndBias( + size_t FilterCount, + size_t InputChannels, + const int64_t* KernelShape, + const int64_t* DilationShape, + const MLAS_FP16* Filter, + // Optional bias to bake into the packed direct-consumption buffer. + const MLAS_FP16* Bias, + void* PackedWeightsAndBias, + MLAS_THREADPOOL* ThreadPool, + const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig); + /** * @brief For half precision GEMM, convert the float matrix B * to half precision and pack it into a packing buffer diff --git a/onnxruntime/core/mlas/lib/halfconv.cpp b/onnxruntime/core/mlas/lib/halfconv.cpp new file mode 100644 index 0000000000000..22fa48dbc2874 --- /dev/null +++ b/onnxruntime/core/mlas/lib/halfconv.cpp @@ -0,0 +1,158 @@ +// +// SPDX-FileCopyrightText: Copyright 2026 Arm Limited and/or its affiliates +// +// SPDX-License-Identifier: MIT +// + +/*++ + +Module Name: + + halfconv.cpp + +Abstract: + + This module implements public dispatch wrappers for optional half precision + convolution backends. + +--*/ + +#include "mlasi.h" + +bool +MLASCALL +MlasHalfConvPrepare( + MLAS_CONV_PARAMETERS* Parameters, + size_t Dimensions, + size_t BatchCount, + size_t GroupCount, + size_t InputChannels, + const int64_t* InputShape, + const int64_t* KernelShape, + const int64_t* DilationShape, + const int64_t* Padding, + const int64_t* StrideShape, + const int64_t* OutputShape, + size_t FilterCount, + const MLAS_ACTIVATION* Activation, + size_t* WorkingBufferSize, + float Beta, + bool InputOutputChannelsLast, + MLAS_THREADPOOL* ThreadPool, + const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig + ) +{ + if (GetMlasPlatform().MlasHalfConvPrepareOverride == nullptr) { + return false; + } + + return GetMlasPlatform().MlasHalfConvPrepareOverride( + Parameters, + Dimensions, + BatchCount, + GroupCount, + InputChannels, + InputShape, + KernelShape, + DilationShape, + Padding, + StrideShape, + OutputShape, + FilterCount, + Activation, + WorkingBufferSize, + Beta, + InputOutputChannelsLast, + ThreadPool, + BackendKernelSelectorConfig); +} + +bool +MLASCALL +MlasHalfConv( + const MLAS_CONV_PARAMETERS* Parameters, + const MLAS_FP16* Input, + const MLAS_FP16* Filter, + bool FilterAndBiasArePacked, + const MLAS_FP16* Bias, + MLAS_FP16* WorkingBuffer, + MLAS_FP16* Output, + MLAS_THREADPOOL* ThreadPool + ) +{ + if (GetMlasPlatform().MlasHalfConvOverride == nullptr) { + return false; + } + + if (FilterAndBiasArePacked && Bias != nullptr) { + return false; + } + + return GetMlasPlatform().MlasHalfConvOverride( + Parameters, + Input, + Filter, + FilterAndBiasArePacked, + Bias, + WorkingBuffer, + Output, + ThreadPool); +} + +size_t +MLASCALL +MlasHalfConvPackWeightsAndBiasSize( + size_t FilterCount, + size_t InputChannels, + const int64_t* KernelShape, + const int64_t* DilationShape, + const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig + ) +{ + if (BackendKernelSelectorConfig != nullptr && !BackendKernelSelectorConfig->use_kleidiai) { + return 0; + } + + if (GetMlasPlatform().MlasHalfConvPackWeightsAndBiasSizeOverride == nullptr) { + return 0; + } + + return GetMlasPlatform().MlasHalfConvPackWeightsAndBiasSizeOverride( + FilterCount, + InputChannels, + KernelShape, + DilationShape); +} + +bool +MLASCALL +MlasHalfConvPackWeightsAndBias( + size_t FilterCount, + size_t InputChannels, + const int64_t* KernelShape, + const int64_t* DilationShape, + const MLAS_FP16* Filter, + const MLAS_FP16* Bias, + void* PackedWeightsAndBias, + MLAS_THREADPOOL* ThreadPool, + const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig + ) +{ + if (BackendKernelSelectorConfig != nullptr && !BackendKernelSelectorConfig->use_kleidiai) { + return false; + } + + if (GetMlasPlatform().MlasHalfConvPackWeightsAndBiasOverride == nullptr) { + return false; + } + + return GetMlasPlatform().MlasHalfConvPackWeightsAndBiasOverride( + FilterCount, + InputChannels, + KernelShape, + DilationShape, + Filter, + Bias, + PackedWeightsAndBias, + ThreadPool); +} diff --git a/onnxruntime/core/mlas/lib/halfgemm.cpp b/onnxruntime/core/mlas/lib/halfgemm.cpp index 05cde92d9f9d7..f0771c97b2148 100644 --- a/onnxruntime/core/mlas/lib/halfgemm.cpp +++ b/onnxruntime/core/mlas/lib/halfgemm.cpp @@ -19,9 +19,38 @@ Module Name: #include "mlas_float16.h" #include "halfgemm.h" +#if defined(USE_KLEIDIAI) +#include "kleidiai/mlasi_kleidiai.h" +#endif #include +static void +MlasHalfGemmZeroKBatch( + const size_t M, + const size_t N, + const size_t BatchN, + const MLAS_HALF_GEMM_DATA_PARAMS* DataParams + ) +{ + for (size_t gemm_i = 0; gemm_i < BatchN; gemm_i++) { + const auto* Data = &DataParams[gemm_i]; + auto* C = reinterpret_cast(Data->C); + const auto* Bias = reinterpret_cast(Data->Bias); + const size_t ldc = Data->ldc; + + for (size_t m = 0; m < M; m++) { + for (size_t n = 0; n < N; n++) { + C[m * ldc + n] = Bias == nullptr ? MLAS_FP16::FromBits(0) : Bias[n]; + } + } + + if (Data->OutputProcessor != nullptr) { + Data->OutputProcessor->Process(Data->C, 0, 0, M, N, ldc); + } + } +} + bool MLASCALL MlasFp16AccelerationSupported() { @@ -34,6 +63,35 @@ MlasFp16AccelerationSupported() #endif } +#if defined(USE_KLEIDIAI) +bool +TryGetHalfGemmBackendSelectorConfig( + size_t BatchN, + const MLAS_HALF_GEMM_DATA_PARAMS* DataParams, + const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG*& BackendKernelSelectorConfig + ) +{ + BackendKernelSelectorConfig = nullptr; + + for (size_t gemm_i = 0; gemm_i < BatchN; ++gemm_i) { + const auto* config = DataParams[gemm_i].BackendKernelSelectorConfig; + if (config == nullptr) { + continue; + } + + if (!config->use_kleidiai) { + return false; + } + + if (BackendKernelSelectorConfig == nullptr) { + BackendKernelSelectorConfig = config; + } + } + + return true; +} +#endif + void MLASCALL @@ -46,6 +104,36 @@ MlasHalfGemmBatch( MLAS_THREADPOOL* ThreadPool ) { + if (BatchN == 0 || M == 0 || N == 0) { + return; + } + + if (K == 0) { + MlasHalfGemmZeroKBatch(M, N, BatchN, DataParams); + return; + } + +#if defined(USE_KLEIDIAI) + const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig = nullptr; + if (TryGetHalfGemmBackendSelectorConfig(BatchN, DataParams, BackendKernelSelectorConfig) && + GetMlasPlatform().MlasHalfGemmBatchOverride != nullptr && + GetMlasPlatform().MlasHalfGemmBatchOverride( + M, N, K, BatchN, DataParams, ThreadPool, BackendKernelSelectorConfig)) { + return; + } +#endif + + // BIsPacked denotes the generic MLAS halfgemm packed-B layout and can be + // consumed here. Backend-native packed layouts are separate and must not + // silently fall through to the generic kernels. + for (size_t gemm_i = 0; gemm_i < BatchN; gemm_i++) { + if (DataParams[gemm_i].BIsBackendNativePacked) { + MLAS_THROW_EX( + std::runtime_error, + "backend-native halfgemm packed B is not supported by generic MLAS halfgemm"); + } + } + const MLAS_HALFGEMM_DISPATCH* dispatch = MlasHalfGemmGetDispatch(); MLAS_HALFGEMM_OPERATION* operation = dispatch->Operation; @@ -130,12 +218,13 @@ MlasHalfGemmPackBSize( // No packing routine provided return 0; } - const size_t AlignedK = (K + PackedK - 1) & ~(PackedK - 1); - const size_t BytesRequired = N * AlignedK * FP16_SIZE + padding; - const size_t BufferAlignment = MlasGetPreferredBufferAlignment(); - const size_t AlignedBytesRequired = - (BytesRequired + BufferAlignment - 1) & ~(BufferAlignment - 1); - return AlignedBytesRequired; + + size_t packed_b_size = 0; + if (!MlasHalfGemmTryGetPackedBSize(N, K, PackedK, padding, &packed_b_size)) { + return 0; + } + + return packed_b_size; } void @@ -149,9 +238,70 @@ MlasHalfGemmPackB( ) { const auto* dispatch = MlasHalfGemmGetDispatch(); + if (B == nullptr || PackedB == nullptr || ldb < N || + dispatch->CopyPackBRoutine == nullptr || MlasHalfGemmPackBSize(N, K, false) == 0) { + return; + } + dispatch->CopyPackBRoutine((_mlas_fp16_*)PackedB, (const _mlas_fp16_*)B, ldb, N, K); } +size_t +MLASCALL +MlasHalfGemmNativePackBSize( + CBLAS_TRANSPOSE TransA, + CBLAS_TRANSPOSE TransB, + size_t N, + size_t K, + const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig + ) +{ +#if defined(USE_KLEIDIAI) + if ((!BackendKernelSelectorConfig || BackendKernelSelectorConfig->use_kleidiai) && + GetMlasPlatform().MlasHalfGemmPackBSizeOverride != nullptr) { + return GetMlasPlatform().MlasHalfGemmPackBSizeOverride(TransA, TransB, N, K); + } +#else + MLAS_UNREFERENCED_PARAMETER(TransA); + MLAS_UNREFERENCED_PARAMETER(TransB); + MLAS_UNREFERENCED_PARAMETER(N); + MLAS_UNREFERENCED_PARAMETER(K); + MLAS_UNREFERENCED_PARAMETER(BackendKernelSelectorConfig); +#endif + return 0; +} + +bool +MLASCALL +MlasHalfGemmNativePackB( + CBLAS_TRANSPOSE TransA, + CBLAS_TRANSPOSE TransB, + size_t N, + size_t K, + const MLAS_FP16* B, + size_t ldb, + void* PackedB, + const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig + ) +{ +#if defined(USE_KLEIDIAI) + if ((!BackendKernelSelectorConfig || BackendKernelSelectorConfig->use_kleidiai) && + GetMlasPlatform().MlasHalfGemmPackBOverride != nullptr) { + return GetMlasPlatform().MlasHalfGemmPackBOverride(TransA, TransB, N, K, B, ldb, PackedB); + } +#else + MLAS_UNREFERENCED_PARAMETER(TransA); + MLAS_UNREFERENCED_PARAMETER(TransB); + MLAS_UNREFERENCED_PARAMETER(N); + MLAS_UNREFERENCED_PARAMETER(K); + MLAS_UNREFERENCED_PARAMETER(B); + MLAS_UNREFERENCED_PARAMETER(ldb); + MLAS_UNREFERENCED_PARAMETER(PackedB); + MLAS_UNREFERENCED_PARAMETER(BackendKernelSelectorConfig); +#endif + return false; +} + void MLASCALL MlasHalfGemmConvertPackB( @@ -576,7 +726,7 @@ MlasGemmBatch( const MLAS_HALFGEMM_DISPATCH MlasHalfGemmDispatchDefault = { MlasHalfGemmOperation, - nullptr, + MlasHalfGemmCopyPackB, MlasHalfGemmConvertPackB, MLAS_HALF_GEMM_KERNEL_DEFAULT::PackedK, MLAS_HALF_GEMM_KERNEL_DEFAULT::KernelMaxM, diff --git a/onnxruntime/core/mlas/lib/halfgemm.h b/onnxruntime/core/mlas/lib/halfgemm.h index 3f63e00f05f12..ceaca2b17f85c 100644 --- a/onnxruntime/core/mlas/lib/halfgemm.h +++ b/onnxruntime/core/mlas/lib/halfgemm.h @@ -32,9 +32,10 @@ Module Name: #pragma once -#include #include -#include +#include +#include +#include #include "mlasi.h" #include "mlas_float16.h" @@ -50,6 +51,44 @@ struct MLAS_HALF_GEMM_STRIDES { size_t K; }; +MLAS_FORCEINLINE +bool +MlasHalfGemmTryGetPackedBSize( + size_t N, + size_t K, + size_t PackedK, + size_t Padding, + size_t* PackedBSize + ) +{ + if (PackedK == 0) { + return false; + } + + size_t aligned_k_input = 0; + if (MlasAddOverflowsSizeT(K, PackedK - 1, &aligned_k_input)) { + return false; + } + + const size_t AlignedK = aligned_k_input & ~(PackedK - 1); + + size_t BytesRequired = 0; + if (MlasMultiplyOverflowsSizeT(N, AlignedK, &BytesRequired) || + MlasMultiplyOverflowsSizeT(BytesRequired, sizeof(_mlas_fp16_), &BytesRequired) || + MlasAddOverflowsSizeT(BytesRequired, Padding, &BytesRequired)) { + return false; + } + + const size_t BufferAlignment = MlasGetPreferredBufferAlignment(); + size_t aligned_bytes_input = 0; + if (MlasAddOverflowsSizeT(BytesRequired, BufferAlignment - 1, &aligned_bytes_input)) { + return false; + } + + *PackedBSize = aligned_bytes_input & ~(BufferAlignment - 1); + return true; +} + /** * @brief Packing function for fp16 B matrix * @@ -71,12 +110,51 @@ MlasHalfGemmCopyPackB( size_t CountK ) { - MLAS_UNREFERENCED_PARAMETER(D); - MLAS_UNREFERENCED_PARAMETER(B); - MLAS_UNREFERENCED_PARAMETER(ldb); - MLAS_UNREFERENCED_PARAMETER(CountN); - MLAS_UNREFERENCED_PARAMETER(CountK); - // No packing needed by default + size_t aligned_count_k_input = 0; + if (MlasAddOverflowsSizeT(CountK, KernelType::PackedK - 1, &aligned_count_k_input)) { + assert(!"MlasHalfGemmCopyPackB aligned K overflow"); + return; + } + const size_t AlignedCountK = aligned_count_k_input & ~(KernelType::PackedK - 1); + size_t PaddingCountK = AlignedCountK - CountK; + + if (ldb == CountN) { + size_t elements_to_copy = 0; + size_t bytes_to_copy = 0; + if (MlasMultiplyOverflowsSizeT(CountK, CountN, &elements_to_copy) || + MlasMultiplyOverflowsSizeT(elements_to_copy, sizeof(_mlas_fp16_), &bytes_to_copy)) { + assert(!"MlasHalfGemmCopyPackB size overflow"); + return; + } + std::memcpy(D, B, bytes_to_copy); + if (PaddingCountK > 0) { + size_t padding_bytes = 0; + if (MlasMultiplyOverflowsSizeT(PaddingCountK, CountN, &padding_bytes) || + MlasMultiplyOverflowsSizeT(padding_bytes, sizeof(_mlas_fp16_), &padding_bytes)) { + assert(!"MlasHalfGemmCopyPackB padding size overflow"); + return; + } + std::memset(D + elements_to_copy, 0, padding_bytes); + } + return; + } + + size_t row_bytes = 0; + if (MlasMultiplyOverflowsSizeT(CountN, sizeof(_mlas_fp16_), &row_bytes)) { + assert(!"MlasHalfGemmCopyPackB row size overflow"); + return; + } + while (CountK > 0) { + std::memcpy(D, B, row_bytes); + B += ldb; + D += CountN; + CountK--; + } + while (PaddingCountK > 0) { + std::memset(D, 0, row_bytes); + D += CountN; + PaddingCountK--; + } } /** diff --git a/onnxruntime/core/mlas/lib/halfgemm_kernel_neon.cpp b/onnxruntime/core/mlas/lib/halfgemm_kernel_neon.cpp index d7f5a90b00589..f9bb3da1c5b57 100644 --- a/onnxruntime/core/mlas/lib/halfgemm_kernel_neon.cpp +++ b/onnxruntime/core/mlas/lib/halfgemm_kernel_neon.cpp @@ -179,7 +179,7 @@ MlasHalfGemmKernel( const MLAS_HALFGEMM_DISPATCH MlasHalfGemmDispatchNeon = { MlasHalfGemmOperation, - nullptr, + MlasHalfGemmCopyPackB, MlasHalfGemmConvertPackB, MLAS_HALF_GEMM_KERNEL_NEON::PackedK, MLAS_HALF_GEMM_KERNEL_NEON::KernelMaxM, diff --git a/onnxruntime/core/mlas/lib/kai_ukernel_interface.cpp b/onnxruntime/core/mlas/lib/kai_ukernel_interface.cpp index 6ee80594c6b49..a352fc85abbd1 100644 --- a/onnxruntime/core/mlas/lib/kai_ukernel_interface.cpp +++ b/onnxruntime/core/mlas/lib/kai_ukernel_interface.cpp @@ -28,6 +28,7 @@ #include "kai/ukernels/matmul/matmul_clamp_f32_f32_f32p/kai_matmul_clamp_f32_f32_f32p2vlx1b_1x8vl_sme_mla.h" // IMATMUL #include "kai/ukernels/matmul/imatmul_clamp_f32_f32p_f32p/kai_imatmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa.h" +#include "kai/ukernels/matmul/imatmul_clamp_f16_f16p_f16p/kai_imatmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa.h" // SME2 kernels // GEMM/QGEMM/SBGEMM @@ -40,6 +41,11 @@ #include "kai/ukernels/matmul/matmul_clamp_f32_f32_f32p/kai_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla.h" // IMATMUL #include "kai/ukernels/matmul/imatmul_clamp_f32_f32p_f32p/kai_imatmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme2_mopa.h" +#include "kai/ukernels/matmul/imatmul_clamp_f16_f16p_f16p/kai_imatmul_clamp_f16_f16p2vlx2_f16p2vlx2_2vlx2vl_sme2_mopa.h" + +// FP16 HGEMM kernels +#include "kai/ukernels/matmul/matmul_clamp_f16_f16_f16p/kai_matmul_clamp_f16_f16_f16p2vlx2b_1x8vl_sme_mla.h" +#include "kai/ukernels/matmul/matmul_clamp_f16_f16_f16p/kai_matmul_clamp_f16_f16_f16p2vlx2b_1x16vl_sme2_dot.h" #if defined(ENABLE_QMX_KERNELS) // QMX kernels (optional) @@ -229,6 +235,12 @@ const KaiF32IMatmulKernel imatmul_conv_sme = const KaiF32IMatmulKernel imatmul_conv_sme2 = KAI_WRAP_UKERNEL_RUN_IMATMUL_PACKED_7(imatmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme2_mopa); +const KaiF16IMatmulKernel imatmul_f16_conv_sme = + KAI_WRAP_UKERNEL_RUN_IMATMUL_PACKED_7(imatmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa); + +const KaiF16IMatmulKernel imatmul_f16_conv_sme2 = + KAI_WRAP_UKERNEL_RUN_IMATMUL_PACKED_7(imatmul_clamp_f16_f16p2vlx2_f16p2vlx2_2vlx2vl_sme2_mopa); + const KaiBF16SBgemmKernel sbgemm_gemm_sme2 = KAI_WRAP_UKERNEL_RUN_MATMUL_11(matmul_clamp_f32_bf16p2vlx2_bf16p2vlx2_2vlx2vl_sme2_mopa); @@ -246,6 +258,12 @@ const KaiDynamicQGemmKernel qgemm_gemm_sme = const KaiDynamicQGemmKernel qgemm_gemm_sme2 = KAI_WRAP_UKERNEL_RUN_MATMUL_11(matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme2_mopa); +const KaiF16HgemmKernel hgemm_sme = + KAI_WRAP_UKERNEL_RUN_MATMUL_10_LHS_OFFSET(matmul_clamp_f16_f16_f16p2vlx2b_1x8vl_sme_mla); + +const KaiF16HgemmKernel hgemm_sme2 = + KAI_WRAP_UKERNEL_RUN_MATMUL_10_LHS_PACKED_OFFSET(matmul_clamp_f16_f16_f16p2vlx2b_1x16vl_sme2_dot); + #if defined(ENABLE_QMX_KERNELS) @@ -350,6 +368,14 @@ const KaiF32IMatmulKernel& GetKleidiAIF32IMatmulUKernel() { } } +const KaiF16IMatmulKernel& GetKleidiAIF16IMatmulUKernel() { + if (MLAS_CPUIDINFO::GetCPUIDInfo().HasArm_SME2()) { + return imatmul_f16_conv_sme2; + } else { + return imatmul_f16_conv_sme; + } +} + const KaiDynamicQGemmKernel& GetKleidiAIQGemmUKernel() { if (MLAS_CPUIDINFO::GetCPUIDInfo().HasArm_SME2()) { return qgemm_gemm_sme2; @@ -372,3 +398,11 @@ const KaiBF16SBgemmKernel& GetKleidiAISBGemmUKernel() { // Currently only SME2 variant exists for bfloat16/SBGEMM kernel return sbgemm_gemm_sme2; } + +const KaiF16HgemmKernel& GetKleidiAIHgemmUKernel() { + if (MLAS_CPUIDINFO::GetCPUIDInfo().HasArm_SME2()) { + return hgemm_sme2; + } else { + return hgemm_sme; + } +} diff --git a/onnxruntime/core/mlas/lib/kai_ukernel_interface.h b/onnxruntime/core/mlas/lib/kai_ukernel_interface.h index 155ecf1762b3b..9b42e30a969d0 100644 --- a/onnxruntime/core/mlas/lib/kai_ukernel_interface.h +++ b/onnxruntime/core/mlas/lib/kai_ukernel_interface.h @@ -18,6 +18,10 @@ #include "kai/ukernels/matmul/imatmul_clamp_f32_f32p_f32p/kai_imatmul_clamp_f32_f32p_f32p_interface.h" +#include "kai/ukernels/matmul/imatmul_clamp_f16_f16p_f16p/kai_imatmul_clamp_f16_f16p_f16p_interface.h" + +#include "kai/ukernels/matmul/matmul_clamp_f16_f16_f16p/kai_matmul_clamp_f16_f16_f16p_interface.h" + // Wrapper type that carries a stable "name" alongside the KAI ukernel interface. // This avoids needing to infer which underlying microkernel was selected from a function pointer. template @@ -41,8 +45,14 @@ using KaiDynamicQGemmKernel = KaiMatmulKernel; +// Wrapper for FP16 IMATMUL kernels used by the KleidiAI convolution implementation. +using KaiF16IMatmulKernel = KaiMatmulKernel; + using KaiBF16SBgemmKernel = KaiMatmulKernel; +// Wrapper for FP16 HGEMM kernels producing FP16 output. +using KaiF16HgemmKernel = KaiMatmulKernel; + // Returns the selected Qnbit GEMM ukernel based on runtime CPU capabilities. const KaiQnbitGemmKernel& GetKleidiAIGemmUKernel(); @@ -61,5 +71,11 @@ const KaiF32SgemvKernel& GetKleidiAISGemvUKernel(); // Returns the selected FP32 IMATMUL ukernel used by the KleidiAI convolution implementation. const KaiF32IMatmulKernel& GetKleidiAIF32IMatmulUKernel(); +// Returns the selected FP16 IMATMUL ukernel used by the KleidiAI convolution implementation. +const KaiF16IMatmulKernel& GetKleidiAIF16IMatmulUKernel(); + // Returns the selected BF16 SBGEMM ukernel used by the KleidiAI based on runtime CPU capabilities. -const KaiBF16SBgemmKernel& GetKleidiAISBGemmUKernel(); \ No newline at end of file +const KaiBF16SBgemmKernel& GetKleidiAISBGemmUKernel(); + +// Returns the selected FP16 HGEMM ukernel based on runtime CPU capabilities. +const KaiF16HgemmKernel& GetKleidiAIHgemmUKernel(); diff --git a/onnxruntime/core/mlas/lib/kleidiai/halfconv_kleidiai.cpp b/onnxruntime/core/mlas/lib/kleidiai/halfconv_kleidiai.cpp new file mode 100644 index 0000000000000..ee487437f336d --- /dev/null +++ b/onnxruntime/core/mlas/lib/kleidiai/halfconv_kleidiai.cpp @@ -0,0 +1,1040 @@ +// +// SPDX-FileCopyrightText: Copyright 2026 Arm Limited and/or its affiliates +// +// SPDX-License-Identifier: MIT +// + +#include +#include +#include +#include +#include +#include + +#include "kai/ukernels/matmul/pack/kai_lhs_imatmul_pack_x16p2vlx2_x16p_sme.h" +#include "kai/ukernels/matmul/pack/kai_rhs_imatmul_pack_kxn_x16p2vlx2b_x16_x16_sme.h" +#include "kai_ukernel_interface.h" +#include "mlasi_kleidiai.h" + +namespace +{ + +// Cache-oriented heuristics for the chunked IMATMUL path. These bound the +// packed LHS working set so we only split when the single-pass pack/compute +// footprint is large enough that chunking can improve locality. +constexpr size_t AutomaticMaximumLhsChunkBytes = 2 * 1024 * 1024; +constexpr size_t MinimumLhsPackedBytesForAutomaticChunking = 8 * 1024 * 1024; +constexpr size_t MinimumEffectiveKForAutomaticChunking = 1024; +// Small filter counts typically do not amortize the extra chunking overhead. +constexpr size_t MaximumFilterCountForAutomaticChunking = 32; + +struct KaiHalfConvTlsBuffers { + std::vector packed_lhs; + + void ReleaseLargeBuffers() { + MlasShrinkKleidiAIScratchIfTooLarge(packed_lhs); + } +}; + +struct ScopedKaiHalfConvTlsCleanup { + KaiHalfConvTlsBuffers& buffers; + + ~ScopedKaiHalfConvTlsCleanup() { + buffers.ReleaseLargeBuffers(); + } +}; + +thread_local KaiHalfConvTlsBuffers g_kai_half_conv_tls; + +template +bool TryResizeVector(std::vector& buffer, size_t size) { + if (size > buffer.max_size()) { + return false; + } + buffer.resize(size); + return true; +} + +bool +TryComputeKernelSize( + size_t dilation, + size_t kernel, + size_t& dilated_kernel +) +{ + if (dilation == 0 || kernel == 0) { + return false; + } + + size_t scaled_kernel = 0; + if (MlasMultiplyOverflowsSizeT(dilation, kernel, &scaled_kernel) || scaled_kernel < (dilation - 1)) { + return false; + } + + dilated_kernel = scaled_kernel - (dilation - 1); + return true; +} + +bool +TryComputeConvOutSize( + size_t input, + size_t kernel, + size_t padding, + size_t stride, + size_t& output +) +{ + output = 0; + if (stride == 0) { + return false; + } + + size_t total_padding = 0; + size_t padded_input = 0; + if (MlasMultiplyOverflowsSizeT(padding, 2, &total_padding) || + MlasAddOverflowsSizeT(input, total_padding, &padded_input)) { + return false; + } + + if (padded_input < kernel) { + return true; + } + + output = ((padded_input - kernel) / stride) + 1; + return true; +} + +bool +TryComputeOutputSize( + size_t input_height, + size_t input_width, + size_t kernel_height, + size_t kernel_width, + size_t padding_height, + size_t padding_width, + size_t stride_height, + size_t stride_width, + size_t& output_height, + size_t& output_width, + size_t& output_size +) +{ + if (!TryComputeConvOutSize(input_height, kernel_height, padding_height, stride_height, output_height) || + !TryComputeConvOutSize(input_width, kernel_width, padding_width, stride_width, output_width) || + MlasMultiplyOverflowsSizeT(output_height, output_width, &output_size)) { + return false; + } + + return true; +} + +bool +TryComputeOutputSize( + size_t input_height, + size_t input_width, + size_t kernel_height, + size_t kernel_width, + size_t padding_height, + size_t padding_width, + size_t stride_height, + size_t stride_width, + size_t& output_size +) +{ + size_t output_height = 0; + size_t output_width = 0; + return TryComputeOutputSize( + input_height, + input_width, + kernel_height, + kernel_width, + padding_height, + padding_width, + stride_height, + stride_width, + output_height, + output_width, + output_size); +} + +size_t +SelectMaximumLhsChunkBytes( + size_t full_lhs_size, + size_t filter_count, + size_t effective_k +) +{ + // The bounded-LHS path is most useful when LHS packing is cache-hostile and + // there are few output channels, so full-LHS reuse across N tiles is limited. + if (full_lhs_size >= MinimumLhsPackedBytesForAutomaticChunking && + filter_count <= MaximumFilterCountForAutomaticChunking && + effective_k >= MinimumEffectiveKForAutomaticChunking) { + return AutomaticMaximumLhsChunkBytes; + } + + return 0; +} + +bool +IsPaddingSymmetric2D(const MLAS_CONV_PARAMETERS* parameters) +{ + return parameters->Padding[0] == parameters->Padding[1] && + parameters->Padding[0] == parameters->Padding[2] && + parameters->Padding[0] == parameters->Padding[3]; +} + +bool +CheckCapabilitiesSme(const MLAS_CONV_PARAMETERS* parameters) +{ + if (parameters == nullptr) { + return false; + } + + if (parameters->BackendKernelSelectorConfig != nullptr && + !parameters->BackendKernelSelectorConfig->use_kleidiai) { + KLEIDIAI_DEBUG_LOG("User explicitly disabled KleidiAI, returning false from MlasHalfConv."); + return false; + } + + if ((parameters->Dimensions != 2) || + (parameters->BatchCount != 1) || + (parameters->GroupCount != 1) || + (parameters->Beta != 0.0f) || + !(parameters->Activation == nullptr || + parameters->Activation->ActivationKind == MlasIdentityActivation) || + !IsPaddingSymmetric2D(parameters)) { + KLEIDIAI_DEBUG_LOG("MlasHalfConv capability check failed: unsupported configuration."); + return false; + } + + size_t d_kh = 0; + size_t d_kw = 0; + size_t output_size = 0; + if (!TryComputeKernelSize(parameters->DilationShape[0], parameters->KernelShape[0], d_kh) || + !TryComputeKernelSize(parameters->DilationShape[1], parameters->KernelShape[1], d_kw) || + !TryComputeOutputSize( + parameters->InputShape[0], + parameters->InputShape[1], + d_kh, + d_kw, + parameters->Padding[0], + parameters->Padding[1], + parameters->StrideShape[0], + parameters->StrideShape[1], + output_size)) { + return false; + } + + if (output_size == 0 || + output_size != parameters->OutputSize || + parameters->InputChannels == 0 || + parameters->FilterCount == 0 || + parameters->FilterCount == 1 || + parameters->KernelShape[0] < 3 || + parameters->KernelShape[1] < 3) { + KLEIDIAI_DEBUG_LOG("MlasHalfConv capability check failed: shape/heuristic gating."); + return false; + } + + return true; +} + +bool +GetPackedFilterSize( + size_t filter_count, + size_t input_channels, + const int64_t* kernel_shape, + const int64_t* dilation_shape, + size_t* packed_size +) +{ + if (packed_size == nullptr || + kernel_shape == nullptr || + dilation_shape == nullptr || + filter_count <= 1 || + input_channels == 0 || + kernel_shape[0] < 3 || + kernel_shape[1] < 3 || + dilation_shape[0] <= 0 || + dilation_shape[1] <= 0) { + return false; + } + + size_t d_kh = 0; + size_t d_kw = 0; + size_t k_chunk_count = 0; + if (!TryComputeKernelSize(static_cast(dilation_shape[0]), static_cast(kernel_shape[0]), d_kh) || + !TryComputeKernelSize(static_cast(dilation_shape[1]), static_cast(kernel_shape[1]), d_kw) || + MlasMultiplyOverflowsSizeT(d_kh, d_kw, &k_chunk_count)) { + return false; + } + + *packed_size = kai_get_rhs_packed_size_rhs_imatmul_pack_kxn_x16p2vlx2b_x16_x16_sme( + filter_count, k_chunk_count, input_channels + ); + return *packed_size != 0; +} + +bool +NchwToNhwc( + const MLAS_FP16* input, + size_t channels, + size_t height, + size_t width, + std::vector& output +) +{ + size_t element_count = 0; + if (MlasMultiplyOverflowsSizeT(channels, height, &element_count) || + MlasMultiplyOverflowsSizeT(element_count, width, &element_count)) { + return false; + } + output.resize(element_count); + + if (input == nullptr) { + return false; + } + + for (size_t c = 0; c < channels; ++c) { + for (size_t h = 0; h < height; ++h) { + for (size_t w = 0; w < width; ++w) { + output[(h * width + w) * channels + c] = input[(c * height + h) * width + w]; + } + } + } + + return true; +} + +bool +FillIndirectionTable( + const MLAS_FP16* input_nhwc, + const MLAS_FP16* pad, + size_t input_channels, + size_t input_height, + size_t input_width, + size_t kernel_height, + size_t kernel_width, + size_t stride_height, + size_t stride_width, + size_t padding, + size_t output_size, + std::vector& indirection +) +{ + const auto& imatmul = GetKleidiAIF16IMatmulUKernel(); + const size_t m_step = imatmul.ukernel.get_m_step(); + size_t lhs_ptrs_k = 0; + if (MlasMultiplyOverflowsSizeT(kernel_height, kernel_width, &lhs_ptrs_k)) { + return false; + } + + size_t lhs_ptrs_m = 0; + if (MlasMultiplyOverflowsSizeT(m_step, MlasDivRoundup(output_size, m_step), &lhs_ptrs_m)) { + return false; + } + + size_t table_size = 0; + if (MlasMultiplyOverflowsSizeT(lhs_ptrs_k, lhs_ptrs_m, &table_size)) { + return false; + } + indirection.resize(table_size); + + std::fill(indirection.begin(), indirection.end(), pad); + + auto ptr_offset = [lhs_ptrs_k, m_step](size_t k, size_t m) { + return ((m / m_step) * lhs_ptrs_k * m_step) + (k * m_step) + (m % m_step); + }; + + auto pixel_ptr = [=](size_t h, size_t w) -> const void* { + if (h < padding || w < padding) { + return pad; + } + + h -= padding; + w -= padding; + if (h >= input_height || w >= input_width) { + return pad; + } + + return input_nhwc + (h * input_width + w) * input_channels; + }; + + size_t output_height = 0; + size_t output_width = 0; + size_t computed_output_size = 0; + if (!TryComputeOutputSize( + input_height, + input_width, + kernel_height, + kernel_width, + padding, + padding, + stride_height, + stride_width, + output_height, + output_width, + computed_output_size) || + computed_output_size != output_size) { + return false; + } + + size_t m = 0; + for (size_t oh = 0; oh < output_height; ++oh) { + for (size_t ow = 0; ow < output_width; ++ow, ++m) { + size_t k = 0; + const size_t input_base_h = oh * stride_height; + const size_t input_base_w = ow * stride_width; + for (size_t kh = 0; kh < kernel_height; ++kh) { + for (size_t kw = 0; kw < kernel_width; ++kw, ++k) { + indirection[ptr_offset(k, m)] = pixel_ptr(input_base_h + kh, input_base_w + kw); + } + } + } + } + + return m == output_size; +} + +bool +PrepareLhsInput( + const MLAS_CONV_PARAMETERS* parameters, + const MLAS_FP16* input, + size_t output_size, + std::vector& input_nhwc, + std::vector& pad, + std::vector& indirection +) +{ + size_t d_kh = 0; + size_t d_kw = 0; + if (!TryComputeKernelSize(parameters->DilationShape[0], parameters->KernelShape[0], d_kh) || + !TryComputeKernelSize(parameters->DilationShape[1], parameters->KernelShape[1], d_kw)) { + return false; + } + + const size_t input_channels = parameters->InputChannels; + + const MLAS_FP16* input_nhwc_data = input; + if (!parameters->InputOutputChannelsLast) { + if (!NchwToNhwc( + input, + input_channels, + parameters->InputShape[0], + parameters->InputShape[1], + input_nhwc + )) { + return false; + } + input_nhwc_data = input_nhwc.data(); + } + + pad.resize(input_channels); + std::fill(pad.begin(), pad.end(), MLAS_FP16::FromBits(0)); + + return FillIndirectionTable( + input_nhwc_data, + pad.data(), + input_channels, + parameters->InputShape[0], + parameters->InputShape[1], + d_kh, + d_kw, + parameters->StrideShape[0], + parameters->StrideShape[1], + parameters->Padding[0], + output_size, + indirection + ); +} + +bool +PackFilter( + size_t filter_count, + size_t input_channels, + const int64_t* kernel_shape, + const int64_t* dilation_shape, + const MLAS_FP16* filter, + const MLAS_FP16* bias, + void* packed_filter +) +{ + if (filter == nullptr || packed_filter == nullptr || + kernel_shape == nullptr || dilation_shape == nullptr || + filter_count <= 1 || input_channels == 0 || + kernel_shape[0] < 3 || kernel_shape[1] < 3 || + dilation_shape[0] <= 0 || dilation_shape[1] <= 0) { + return false; + } + + const size_t kernel_height = static_cast(kernel_shape[0]); + const size_t kernel_width = static_cast(kernel_shape[1]); + const size_t dilation_height = static_cast(dilation_shape[0]); + const size_t dilation_width = static_cast(dilation_shape[1]); + size_t d_kh = 0; + size_t d_kw = 0; + size_t k_chunk_count = 0; + if (!TryComputeKernelSize(dilation_height, kernel_height, d_kh) || + !TryComputeKernelSize(dilation_width, kernel_width, d_kw) || + MlasMultiplyOverflowsSizeT(d_kh, d_kw, &k_chunk_count)) { + return false; + } + + size_t reordered_size = 0; + if (MlasMultiplyOverflowsSizeT(k_chunk_count, input_channels, &reordered_size) || + MlasMultiplyOverflowsSizeT(reordered_size, filter_count, &reordered_size)) { + return false; + } + + std::vector reordered_filter; + reordered_filter.resize(reordered_size); + std::fill(reordered_filter.begin(), reordered_filter.end(), MLAS_FP16::FromBits(0)); + + for (size_t oc = 0; oc < filter_count; ++oc) { + for (size_t ic = 0; ic < input_channels; ++ic) { + for (size_t kh = 0; kh < kernel_height; ++kh) { + for (size_t kw = 0; kw < kernel_width; ++kw) { + const size_t src = ((oc * input_channels + ic) * kernel_height + kh) * kernel_width + kw; + const size_t dk = ((kh * dilation_height) * d_kw + (kw * dilation_width)) * input_channels + ic; + reordered_filter[dk * filter_count + oc] = filter[src]; + } + } + } + } + + std::vector zero_bias; + const MLAS_FP16* bias_data = bias; + if (bias_data == nullptr) { + zero_bias.resize(filter_count); + std::fill(zero_bias.begin(), zero_bias.end(), MLAS_FP16::FromBits(0)); + bias_data = zero_bias.data(); + } + + KLEIDIAI_KERNEL_LOG("kai_run_rhs_imatmul_pack_kxn_x16p2vlx2b_x16_x16_sme" << " N=" << filter_count << " k_chunk_count=" << k_chunk_count << " k_chunk_length=" << input_channels << " rhs_stride_row=" << (filter_count * sizeof(MLAS_FP16))); + kai_run_rhs_imatmul_pack_kxn_x16p2vlx2b_x16_x16_sme( + filter_count, + k_chunk_count, + input_channels, + filter_count * sizeof(MLAS_FP16), + reordered_filter.data(), + bias_data, + packed_filter + ); + + return true; +} + +bool +ConvolveSme( + const MLAS_CONV_PARAMETERS* parameters, + const MLAS_FP16* input, + const MLAS_FP16* filter, + bool filter_and_bias_are_packed, + const MLAS_FP16* bias, + MLAS_FP16* working_buffer, + MLAS_FP16* output, + MLAS_THREADPOOL* thread_pool +) +{ + if (input == nullptr || filter == nullptr || output == nullptr || + (!parameters->InputOutputChannelsLast && working_buffer == nullptr)) { + return false; + } + + size_t d_kh = 0; + size_t d_kw = 0; + size_t output_size = 0; + if (!TryComputeKernelSize(parameters->DilationShape[0], parameters->KernelShape[0], d_kh) || + !TryComputeKernelSize(parameters->DilationShape[1], parameters->KernelShape[1], d_kw) || + !TryComputeOutputSize( + parameters->InputShape[0], + parameters->InputShape[1], + d_kh, + d_kw, + parameters->Padding[0], + parameters->Padding[1], + parameters->StrideShape[0], + parameters->StrideShape[1], + output_size)) { + return false; + } + + std::vector input_nhwc; + std::vector pad; + std::vector indirection; + if (!PrepareLhsInput(parameters, input, output_size, input_nhwc, pad, indirection)) { + return false; + } + + std::vector packed_filter_buffer; + const std::byte* packed_filter = reinterpret_cast(filter); + if (!filter_and_bias_are_packed) { + const std::array kernel_shape{ + static_cast(parameters->KernelShape[0]), + static_cast(parameters->KernelShape[1]) + }; + const std::array dilation_shape{ + static_cast(parameters->DilationShape[0]), + static_cast(parameters->DilationShape[1]) + }; + size_t packed_filter_size = 0; + if (!GetPackedFilterSize( + parameters->FilterCount, + parameters->InputChannels, + kernel_shape.data(), + dilation_shape.data(), + &packed_filter_size + )) { + return false; + } + packed_filter_buffer.resize(packed_filter_size); + if (!PackFilter( + parameters->FilterCount, + parameters->InputChannels, + kernel_shape.data(), + dilation_shape.data(), + filter, + bias, + packed_filter_buffer.data() + )) { + return false; + } + packed_filter = packed_filter_buffer.data(); + } + + const auto& imatmul = GetKleidiAIF16IMatmulUKernel(); + const size_t base_n_step = imatmul.ukernel.get_n_step(); + const size_t base_m_step = imatmul.ukernel.get_m_step(); + size_t n_step = base_n_step; + size_t m_step = base_m_step; + const size_t filter_count = parameters->FilterCount; + const size_t input_channels = parameters->InputChannels; + + std::array dim{ + MlasDivRoundup(output_size, m_step), + MlasDivRoundup(filter_count, n_step) + }; + + size_t tile_count = 0; + if (MlasMultiplyOverflowsSizeT(dim[0], dim[1], &tile_count)) { + return false; + } + + const size_t required_tiles = std::min( + static_cast(MlasGetMaximumThreadCount(thread_pool)), + tile_count + ); + + if (required_tiles == 0) { + return false; + } + + const size_t original_dim0 = dim[0]; + const size_t original_dim1 = dim[1]; + size_t scaled_dim0 = 0; + size_t scaled_dim1 = 0; + if (MlasMultiplyOverflowsSizeT(required_tiles, original_dim0, &scaled_dim0) || + MlasMultiplyOverflowsSizeT(required_tiles, original_dim1, &scaled_dim1)) { + return false; + } + + dim[0] = MlasDivRoundup(scaled_dim0, tile_count); + dim[1] = MlasDivRoundup(scaled_dim1, tile_count); + + size_t new_m_step = 0; + size_t new_n_step = 0; + if (MlasMultiplyOverflowsSizeT(m_step, MlasDivRoundup(MlasDivRoundup(output_size, dim[0]), m_step), &new_m_step) || + MlasMultiplyOverflowsSizeT(n_step, MlasDivRoundup(MlasDivRoundup(filter_count, dim[1]), n_step), &new_n_step)) { + return false; + } + m_step = new_m_step; + n_step = new_n_step; + + dim[0] = MlasDivRoundup(output_size, m_step); + dim[1] = MlasDivRoundup(filter_count, n_step); + size_t finalized_tile_count = 0; + if (MlasMultiplyOverflowsSizeT(dim[0], dim[1], &finalized_tile_count)) { + return false; + } + + const float clamp_min = -std::numeric_limits::infinity(); + const float clamp_max = std::numeric_limits::infinity(); + size_t dst_stride = 0; + if (MlasMultiplyOverflowsSizeT(filter_count, sizeof(MLAS_FP16), &dst_stride)) { + return false; + } + + MLAS_FP16* destination = parameters->InputOutputChannelsLast ? output : working_buffer; + + size_t kernel_chunk_count = 0; + size_t effective_k = 0; + if (MlasMultiplyOverflowsSizeT(d_kh, d_kw, &kernel_chunk_count) || + MlasMultiplyOverflowsSizeT(kernel_chunk_count, input_channels, &effective_k)) { + return false; + } + + const size_t full_lhs_size = kai_get_lhs_packed_size_lhs_imatmul_pack_x16p2vlx2_x16p_sme( + output_size, kernel_chunk_count, input_channels + ); + if (full_lhs_size == 0) { + return false; + } + + const size_t maximum_lhs_chunk_bytes = SelectMaximumLhsChunkBytes( + full_lhs_size, + filter_count, + effective_k + ); + + if (maximum_lhs_chunk_bytes == 0 || full_lhs_size <= maximum_lhs_chunk_bytes) { + std::vector packed_lhs; + packed_lhs.resize(full_lhs_size); + + KLEIDIAI_KERNEL_LOG("kai_run_lhs_imatmul_pack_x16p2vlx2_x16p_sme" + << " M=" << output_size + << " k_chunk_count=" << kernel_chunk_count + << " k_chunk_length=" << input_channels); + kai_run_lhs_imatmul_pack_x16p2vlx2_x16p_sme( + output_size, + kernel_chunk_count, + input_channels, + indirection.data(), + 0, + pad.data(), + packed_lhs.data() + ); + + std::atomic ok{true}; + MlasTrySimpleParallel(thread_pool, static_cast(finalized_tile_count), [&](ptrdiff_t tid) { + if (!ok.load(std::memory_order_relaxed)) { + return; + } + + const size_t m_idx = (static_cast(tid) / dim[1]) * m_step; + const size_t n_idx = (static_cast(tid) % dim[1]) * n_step; + const size_t tile_m = std::min(m_step, output_size - m_idx); + const size_t tile_n = std::min(n_step, filter_count - n_idx); + + const std::byte* lhs_tile = + packed_lhs.data() + imatmul.ukernel.get_lhs_packed_offset(m_idx, kernel_chunk_count, input_channels); + const std::byte* rhs_tile = + packed_filter + imatmul.ukernel.get_rhs_packed_offset(n_idx, kernel_chunk_count, input_channels); + + size_t dst_elements = 0; + size_t dst_bytes = 0; + if (MlasMultiplyOverflowsSizeT(m_idx, filter_count, &dst_elements) || + MlasAddOverflowsSizeT(dst_elements, n_idx, &dst_elements) || + MlasMultiplyOverflowsSizeT(dst_elements, sizeof(MLAS_FP16), &dst_bytes)) { + ok.store(false, std::memory_order_relaxed); + return; + } + + std::byte* dst_tile = reinterpret_cast(destination) + dst_bytes; + + KLEIDIAI_KERNEL_LOG(imatmul.name << " M=" << tile_m + << " N=" << tile_n + << " k_chunk_count=" << kernel_chunk_count + << " k_chunk_length=" << input_channels); + imatmul.ukernel.run_imatmul( + tile_m, + tile_n, + kernel_chunk_count, + input_channels, + lhs_tile, + rhs_tile, + dst_tile, + dst_stride, + clamp_min, + clamp_max + ); + }); + + if (!ok.load(std::memory_order_relaxed)) { + return false; + } + } else { + const size_t bytes_per_m_step = kai_get_lhs_packed_size_lhs_imatmul_pack_x16p2vlx2_x16p_sme( + base_m_step, kernel_chunk_count, input_channels + ); + if (bytes_per_m_step == 0) { + return false; + } + + const size_t max_m_steps_per_chunk = std::min( + MlasDivRoundup(output_size, base_m_step), + std::max(1, maximum_lhs_chunk_bytes / bytes_per_m_step) + ); + size_t lhs_chunk_m = 0; + if (MlasMultiplyOverflowsSizeT(base_m_step, max_m_steps_per_chunk, &lhs_chunk_m)) { + return false; + } + + const size_t lhs_chunk_size = kai_get_lhs_packed_size_lhs_imatmul_pack_x16p2vlx2_x16p_sme( + lhs_chunk_m, kernel_chunk_count, input_channels + ); + if (lhs_chunk_size == 0 || lhs_chunk_size > std::vector().max_size()) { + return false; + } + + const size_t m_chunk_count = MlasDivRoundup(output_size, lhs_chunk_m); + std::atomic ok{true}; + MlasTrySimpleParallel(thread_pool, static_cast(m_chunk_count), [&](ptrdiff_t tid) { + if (!ok.load(std::memory_order_relaxed)) { + return; + } + + const size_t global_m_idx = static_cast(tid) * lhs_chunk_m; + const size_t chunk_m = std::min(lhs_chunk_m, output_size - global_m_idx); + size_t indirection_offset = 0; + size_t indirection_tiles = 0; + if (MlasMultiplyOverflowsSizeT(global_m_idx / base_m_step, kernel_chunk_count, &indirection_tiles) || + MlasMultiplyOverflowsSizeT(indirection_tiles, base_m_step, &indirection_offset)) { + ok.store(false, std::memory_order_relaxed); + return; + } + + auto& packed_lhs = g_kai_half_conv_tls.packed_lhs; + ScopedKaiHalfConvTlsCleanup cleanup{g_kai_half_conv_tls}; + if (!TryResizeVector(packed_lhs, lhs_chunk_size)) { + ok.store(false, std::memory_order_relaxed); + return; + } + + KLEIDIAI_KERNEL_LOG("kai_run_lhs_imatmul_pack_x16p2vlx2_x16p_sme" + << " M=" << chunk_m + << " k_chunk_count=" << kernel_chunk_count + << " k_chunk_length=" << input_channels); + kai_run_lhs_imatmul_pack_x16p2vlx2_x16p_sme( + chunk_m, + kernel_chunk_count, + input_channels, + indirection.data() + indirection_offset, + 0, + pad.data(), + packed_lhs.data() + ); + + for (size_t n_tile_idx = 0; n_tile_idx < dim[1]; ++n_tile_idx) { + const size_t n_idx = n_tile_idx * n_step; + const size_t tile_n = std::min(n_step, filter_count - n_idx); + const std::byte* rhs_tile = + packed_filter + imatmul.ukernel.get_rhs_packed_offset(n_idx, kernel_chunk_count, input_channels); + + size_t dst_elements = 0; + size_t dst_bytes = 0; + if (MlasMultiplyOverflowsSizeT(global_m_idx, filter_count, &dst_elements) || + MlasAddOverflowsSizeT(dst_elements, n_idx, &dst_elements) || + MlasMultiplyOverflowsSizeT(dst_elements, sizeof(MLAS_FP16), &dst_bytes)) { + ok.store(false, std::memory_order_relaxed); + return; + } + std::byte* dst_tile = reinterpret_cast(destination) + dst_bytes; + + KLEIDIAI_KERNEL_LOG(imatmul.name << " M=" << chunk_m + << " N=" << tile_n + << " k_chunk_count=" << kernel_chunk_count + << " k_chunk_length=" << input_channels); + imatmul.ukernel.run_imatmul( + chunk_m, + tile_n, + kernel_chunk_count, + input_channels, + packed_lhs.data(), + rhs_tile, + dst_tile, + dst_stride, + clamp_min, + clamp_max + ); + } + }); + + if (!ok.load(std::memory_order_relaxed)) { + return false; + } + } + + if (parameters->InputOutputChannelsLast) { + return true; + } + + MlasTranspose(working_buffer, output, output_size, filter_count, thread_pool); + return true; +} + +} // namespace + +bool + MLASCALL + ArmKleidiAI::MlasHalfConvPrepare( + MLAS_CONV_PARAMETERS* Parameters, + size_t Dimensions, + size_t BatchCount, + size_t GroupCount, + size_t InputChannels, + const int64_t* InputShape, + const int64_t* KernelShape, + const int64_t* DilationShape, + const int64_t* Padding, + const int64_t* StrideShape, + const int64_t* OutputShape, + size_t FilterCount, + const MLAS_ACTIVATION* Activation, + size_t* WorkingBufferSize, + float Beta, + bool InputOutputChannelsLast, + MLAS_THREADPOOL* ThreadPool, + const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig + ) +{ + MLAS_UNREFERENCED_PARAMETER(ThreadPool); + + if (Parameters == nullptr || + InputShape == nullptr || + KernelShape == nullptr || + DilationShape == nullptr || + Padding == nullptr || + StrideShape == nullptr || + OutputShape == nullptr || + WorkingBufferSize == nullptr || + Dimensions != 2) { + return false; + } + + if (BackendKernelSelectorConfig != nullptr && + !BackendKernelSelectorConfig->use_kleidiai) { + KLEIDIAI_DEBUG_LOG("User explicitly disabled KleidiAI, returning false from MlasHalfConvPrepare."); + return false; + } + + Parameters->BackendKernelSelectorConfig = BackendKernelSelectorConfig; + Parameters->Activation = Activation; + Parameters->Dimensions = Dimensions; + Parameters->BatchCount = BatchCount; + Parameters->GroupCount = GroupCount; + Parameters->InputChannels = InputChannels; + Parameters->FilterCount = FilterCount; + Parameters->Beta = Beta; + Parameters->InputOutputChannelsLast = InputOutputChannelsLast; + + size_t input_size = 1; + size_t output_size = 1; + size_t k = InputChannels; + size_t dilated_kernel_size = InputChannels; + for (size_t dim = 0; dim < Dimensions; ++dim) { + if (InputShape[dim] <= 0 || + OutputShape[dim] <= 0 || + KernelShape[dim] <= 0 || + DilationShape[dim] <= 0 || + StrideShape[dim] <= 0 || + Padding[dim] < 0 || + Padding[dim + Dimensions] < 0) { + return false; + } + + Parameters->InputShape[dim] = static_cast(InputShape[dim]); + Parameters->OutputShape[dim] = static_cast(OutputShape[dim]); + Parameters->KernelShape[dim] = static_cast(KernelShape[dim]); + Parameters->DilationShape[dim] = static_cast(DilationShape[dim]); + Parameters->Padding[dim] = static_cast(Padding[dim]); + Parameters->Padding[dim + Dimensions] = static_cast(Padding[dim + Dimensions]); + Parameters->StrideShape[dim] = static_cast(StrideShape[dim]); + + size_t dilated_kernel_dim = 0; + if (!TryComputeKernelSize(Parameters->DilationShape[dim], Parameters->KernelShape[dim], dilated_kernel_dim)) { + return false; + } + + if (MlasMultiplyOverflowsSizeT(input_size, Parameters->InputShape[dim], &input_size) || + MlasMultiplyOverflowsSizeT(output_size, Parameters->OutputShape[dim], &output_size) || + MlasMultiplyOverflowsSizeT(k, Parameters->KernelShape[dim], &k) || + MlasMultiplyOverflowsSizeT(dilated_kernel_size, dilated_kernel_dim, &dilated_kernel_size)) { + return false; + } + } + + Parameters->InputSize = input_size; + Parameters->OutputSize = output_size; + Parameters->K = k; + Parameters->ThreadCount = MlasGetMaximumThreadCount(ThreadPool); + + if (!CheckCapabilitiesSme(Parameters)) { + return false; + } + + size_t working_elements = 0; + if (MlasMultiplyOverflowsSizeT(Parameters->OutputSize, Parameters->FilterCount, &working_elements)) { + return false; + } + if (Parameters->InputOutputChannelsLast) { + *WorkingBufferSize = 0; + return true; + } + + if (MlasMultiplyOverflowsSizeT(working_elements, sizeof(MLAS_FP16), WorkingBufferSize)) { + return false; + } + return true; +} + +bool + MLASCALL + ArmKleidiAI::MlasHalfConv( + const MLAS_CONV_PARAMETERS* Parameters, + const MLAS_FP16* Input, + const MLAS_FP16* Filter, + bool FilterAndBiasArePacked, + const MLAS_FP16* Bias, + MLAS_FP16* WorkingBuffer, + MLAS_FP16* Output, + MLAS_THREADPOOL* ThreadPool + ) +{ + if (!CheckCapabilitiesSme(Parameters)) { + return false; + } + + return ConvolveSme(Parameters, Input, Filter, FilterAndBiasArePacked, Bias, WorkingBuffer, Output, ThreadPool); +} + +size_t + MLASCALL + ArmKleidiAI::MlasHalfConvPackWeightsAndBiasSize( + size_t FilterCount, + size_t InputChannels, + const int64_t* KernelShape, + const int64_t* DilationShape + ) +{ + size_t packed_size = 0; + if (!GetPackedFilterSize(FilterCount, InputChannels, KernelShape, DilationShape, &packed_size)) { + return 0; + } + return packed_size; +} + +bool + MLASCALL + ArmKleidiAI::MlasHalfConvPackWeightsAndBias( + size_t FilterCount, + size_t InputChannels, + const int64_t* KernelShape, + const int64_t* DilationShape, + const MLAS_FP16* Filter, + const MLAS_FP16* Bias, + void* PackedWeightsAndBias, + MLAS_THREADPOOL* ThreadPool + ) +{ + MLAS_UNREFERENCED_PARAMETER(ThreadPool); + + return PackFilter( + FilterCount, + InputChannels, + KernelShape, + DilationShape, + Filter, + Bias, + PackedWeightsAndBias + ); +} diff --git a/onnxruntime/core/mlas/lib/kleidiai/halfgemm_kleidiai.cpp b/onnxruntime/core/mlas/lib/kleidiai/halfgemm_kleidiai.cpp new file mode 100644 index 0000000000000..b443164ce69a1 --- /dev/null +++ b/onnxruntime/core/mlas/lib/kleidiai/halfgemm_kleidiai.cpp @@ -0,0 +1,300 @@ +// +// SPDX-FileCopyrightText: Copyright 2026 Arm Limited and/or its affiliates +// +// SPDX-License-Identifier: MIT +// + +#include +#include +#include +#include +#include "mlas.h" + +#include "mlasi_kleidiai.h" + +#include "kai_ukernel_interface.h" + +#include "kai/ukernels/matmul/pack/kai_rhs_pack_kxn_x16p2vlx2b_x16_x16_sme.h" + +namespace { +struct KaiHalfTlsBuffers { + std::vector lhs_converted; + std::vector rhs_converted; + std::vector bias_zero; + std::vector rhs_packed; + + void ReleaseLargeBuffers() { + MlasShrinkKleidiAIScratchIfTooLarge(lhs_converted); + MlasShrinkKleidiAIScratchIfTooLarge(rhs_converted); + MlasShrinkKleidiAIScratchIfTooLarge(bias_zero); + MlasShrinkKleidiAIScratchIfTooLarge(rhs_packed); + } +}; + +struct ScopedKaiHalfTlsCleanup { + KaiHalfTlsBuffers& buffers; + + ~ScopedKaiHalfTlsCleanup() { + buffers.ReleaseLargeBuffers(); + } +}; + +thread_local KaiHalfTlsBuffers g_kai_half_tls; + +template +bool TryResizeVector(std::vector& buffer, size_t size) { + if (size > buffer.max_size()) { + return false; + } + buffer.resize(size); + return true; +} + +static inline void ConvertFloatMatrixToHalf( + const float* src, + MLAS_FP16* dst, + size_t rows, + size_t cols, + size_t src_ld) { + for (size_t r = 0; r < rows; ++r) { + MlasConvertFloatToHalfBuffer(src + r * src_ld, dst + r * cols, cols); + } +} +} // namespace + +size_t +MLASCALL +ArmKleidiAI::MlasHalfGemmKleidiAIPackBSize( + CBLAS_TRANSPOSE TransA, + CBLAS_TRANSPOSE TransB, + size_t N, + size_t K +) { + if (TransA != CblasNoTrans || TransB != CblasNoTrans || N == 0 || K == 0) { + return 0; + } + + return kai_get_rhs_packed_size_rhs_pack_kxn_x16p2vlx2b_x16_x16_sme(N, K); +} + +bool +MLASCALL +ArmKleidiAI::MlasHalfGemmKleidiAIPackB( + CBLAS_TRANSPOSE TransA, + CBLAS_TRANSPOSE TransB, + size_t N, + size_t K, + const MLAS_FP16* B, + size_t ldb, + void* PackedB +) { + if (TransA != CblasNoTrans || TransB != CblasNoTrans) { + return false; + } + + if (PackedB == nullptr || B == nullptr || N == 0 || K == 0 || ldb < N) { + return false; + } + + const size_t packed_rhs_size = ArmKleidiAI::MlasHalfGemmKleidiAIPackBSize(TransA, TransB, N, K); + if (packed_rhs_size == 0) { + return false; + } + + std::vector zero_bias(N, MLAS_FP16::FromBits(0)); + + size_t ldb_bytes = 0; + if (MlasMultiplyOverflowsSizeT(ldb, sizeof(MLAS_FP16), &ldb_bytes)) { + return false; + } + + const auto& hgemm = GetKleidiAIHgemmUKernel(); + kai_run_rhs_pack_kxn_x16p2vlx2b_x16_x16_sme( + 1, N, K, hgemm.ukernel.get_nr(), hgemm.ukernel.get_kr(), hgemm.ukernel.get_sr(), ldb_bytes, + B, + zero_bias.data(), + nullptr, + PackedB, + 0, + nullptr); + + return true; +} + +bool +MLASCALL +ArmKleidiAI::MlasHalfGemmBatch( + size_t M, + size_t N, + size_t K, + size_t BatchN, + const MLAS_HALF_GEMM_DATA_PARAMS* DataParams, + MLAS_THREADPOOL* ThreadPool, + const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig +) { + if (BatchN == 0 || M == 0 || N == 0) { + return true; + } + if (K == 0) { + return false; + } + if (DataParams == nullptr) { + return false; + } + + ScopedKaiHalfTlsCleanup cleanup{g_kai_half_tls}; + + MLAS_UNREFERENCED_PARAMETER(ThreadPool); + MLAS_UNREFERENCED_PARAMETER(BackendKernelSelectorConfig); + // Validate all batch entries up front so we never partially execute and then + // fall back (which would corrupt results for the already-written outputs). + bool needs_rhs_packing = false; + for (size_t b = 0; b < BatchN; ++b) { + const auto& data = DataParams[b]; + if (data.OutputProcessor != nullptr) { + return false; + } + if (data.BIsBackendNativePacked && data.Bias != nullptr) { + return false; + } + if (data.BIsBackendNativePacked && data.ldb != 0) { + return false; + } + // Native-packed RHS is consumed directly below. Only allocate the + // runtime RHS packing scratch when at least one batch entry needs it. + needs_rhs_packing = needs_rhs_packing || !data.BIsBackendNativePacked; + } + + const auto& hgemm = GetKleidiAIHgemmUKernel(); + const size_t n_step = hgemm.ukernel.get_n_step(); + const size_t nr = hgemm.ukernel.get_nr(); + const size_t kr = hgemm.ukernel.get_kr(); + const size_t sr = hgemm.ukernel.get_sr(); + KLEIDIAI_KERNEL_LOG(hgemm.name); + + const size_t packed_rhs_size = kai_get_rhs_packed_size_rhs_pack_kxn_x16p2vlx2b_x16_x16_sme(N, K); + if (packed_rhs_size == 0) { + return false; + } + + const float clamp_min = -std::numeric_limits::infinity(); + const float clamp_max = std::numeric_limits::infinity(); + + if (needs_rhs_packing && !TryResizeVector(g_kai_half_tls.rhs_packed, packed_rhs_size)) { + return false; + } + + for (size_t b = 0; b < BatchN; ++b) { + const auto& data = DataParams[b]; + + const MLAS_FP16* lhs_base = reinterpret_cast(data.A); + const MLAS_FP16* rhs_base = reinterpret_cast(data.B); + const std::byte* rhs_packed = nullptr; + size_t lhs_ld = data.lda; + size_t rhs_ld = data.ldb; + + if (data.AIsfp32) { + size_t lhs_elements = 0; + if (MlasMultiplyOverflowsSizeT(M, K, &lhs_elements) || + !TryResizeVector(g_kai_half_tls.lhs_converted, lhs_elements)) { + return false; + } + ConvertFloatMatrixToHalf( + reinterpret_cast(data.A), + g_kai_half_tls.lhs_converted.data(), + M, K, data.lda); + lhs_base = g_kai_half_tls.lhs_converted.data(); + lhs_ld = K; + } + + if (data.BIsBackendNativePacked) { + rhs_packed = reinterpret_cast(data.B); + } else if (data.ldb == 0) { + // Prepacked B from MlasHalfGemmPackB/MlasHalfGemmConvertPackB. + // For the current default halfgemm dispatch this is a row-major + // fp16 KxN buffer with leading dimension N. It is not the native + // KleidiAI RHS-packed layout, so this path falls back to packing + // it into KleidiAI format before execution. + rhs_ld = N; + } else if (data.BIsfp32) { + size_t rhs_elements = 0; + if (MlasMultiplyOverflowsSizeT(K, N, &rhs_elements) || + !TryResizeVector(g_kai_half_tls.rhs_converted, rhs_elements)) { + return false; + } + ConvertFloatMatrixToHalf( + reinterpret_cast(data.B), + g_kai_half_tls.rhs_converted.data(), + K, N, data.ldb); + rhs_base = g_kai_half_tls.rhs_converted.data(); + rhs_ld = N; + } + + if (rhs_packed == nullptr) { + auto* rhs_packed_buffer = g_kai_half_tls.rhs_packed.data(); + + size_t ldb_bytes = 0; + if (MlasMultiplyOverflowsSizeT(rhs_ld, sizeof(MLAS_FP16), &ldb_bytes)) { + return false; + } + if (data.Bias == nullptr) { + if (!TryResizeVector(g_kai_half_tls.bias_zero, N)) { + return false; + } + std::fill(g_kai_half_tls.bias_zero.begin(), g_kai_half_tls.bias_zero.end(), MLAS_FP16::FromBits(0)); + } + + kai_run_rhs_pack_kxn_x16p2vlx2b_x16_x16_sme( + 1, N, K, nr, kr, sr, ldb_bytes, + rhs_base, + data.Bias != nullptr ? data.Bias : g_kai_half_tls.bias_zero.data(), + nullptr, + rhs_packed_buffer, + 0, + nullptr); + rhs_packed = rhs_packed_buffer; + } + + size_t lda_bytes = 0; + if (MlasMultiplyOverflowsSizeT(lhs_ld, sizeof(MLAS_FP16), &lda_bytes)) { + return false; + } + size_t dst_stride_bytes = 0; + if (MlasMultiplyOverflowsSizeT(data.ldc, sizeof(MLAS_FP16), &dst_stride_bytes)) { + return false; + } + + MlasTrySimpleParallel(ThreadPool, static_cast(M), [&](ptrdiff_t m_idx) { + const size_t m = static_cast(m_idx); + const auto* lhs = lhs_base + m * lhs_ld; + auto* dst_row = data.C + m * data.ldc; + const auto* rhs_packed_base = rhs_packed; + // The selected KleidiAI HGEMM micro-kernel is 1xN by design. + // We execute one output row per call and parallelize over rows. + constexpr size_t kernel_m = 1; + for (size_t n_idx = 0; n_idx < N; n_idx += n_step) { + const size_t tile_n = std::min(n_step, N - n_idx); + const auto* rhs_tile = rhs_packed_base + hgemm.ukernel.get_rhs_packed_offset(n_idx, K); + auto* dst_tile = reinterpret_cast( + reinterpret_cast(dst_row) + + hgemm.ukernel.get_dst_offset(0, n_idx, dst_stride_bytes)); + + hgemm.ukernel.run_matmul( + kernel_m, + tile_n, + K, + lhs, + lda_bytes, + rhs_tile, + dst_tile, + dst_stride_bytes, + sizeof(MLAS_FP16), + clamp_min, + clamp_max); + } + }); + + } + + return true; +} diff --git a/onnxruntime/core/mlas/lib/kleidiai/mlasi_kleidiai.h b/onnxruntime/core/mlas/lib/kleidiai/mlasi_kleidiai.h index ae000168f7004..f692d4e3d1ca6 100644 --- a/onnxruntime/core/mlas/lib/kleidiai/mlasi_kleidiai.h +++ b/onnxruntime/core/mlas/lib/kleidiai/mlasi_kleidiai.h @@ -7,6 +7,8 @@ #pragma once #include "../mlasi.h" +#include +#include // Fix to ensure compatibility with MSVC build #if defined(_MSC_VER) @@ -49,6 +51,16 @@ #define KLEIDIAI_KERNEL_LOG(msg) #endif +constexpr size_t MaximumRetainedKleidiAIScratchBytes = 8 * 1024 * 1024; + +template +void MlasShrinkKleidiAIScratchIfTooLarge(std::vector& buffer) +{ + if (buffer.capacity() > MaximumRetainedKleidiAIScratchBytes / sizeof(T)) { + std::vector().swap(buffer); + } +} + namespace ArmKleidiAI { // By default we should try for SME2 first before falling back to SME. @@ -374,5 +386,100 @@ MlasConvSGemmRoute(const MLAS_CONV_PARAMETERS* Parameters) { ? MlasConvSGemmRouteDispatch : MlasConvSGemmRouteDirect; } -} +bool +MLASCALL +MlasHalfGemmBatch( + size_t M, + size_t N, + size_t K, + size_t BatchN, + const MLAS_HALF_GEMM_DATA_PARAMS* DataParams, + MLAS_THREADPOOL* ThreadPool, + const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig + ); + +size_t +MLASCALL +MlasHalfGemmKleidiAIPackBSize( + CBLAS_TRANSPOSE TransA, + CBLAS_TRANSPOSE TransB, + size_t N, + size_t K + ); + +// Packs B into the native KleidiAI RHS-packed layout for the supported +// halfgemm configuration only. This differs from the generic MLAS halfgemm +// prepacked-B format produced by MlasHalfGemmPackB and +// MlasHalfGemmConvertPackB, so generic MLAS prepacked weights may need to be +// repacked into this layout before running the KleidiAI halfgemm path. +// Unsupported transpose combinations return false/0 so the caller can fall +// back to the generic MLAS path. +bool +MLASCALL +MlasHalfGemmKleidiAIPackB( + CBLAS_TRANSPOSE TransA, + CBLAS_TRANSPOSE TransB, + size_t N, + size_t K, + const MLAS_FP16* B, + size_t ldb, + void* PackedB + ); + +bool +MLASCALL +MlasHalfConvPrepare(MLAS_CONV_PARAMETERS* Parameters, + size_t Dimensions, + size_t BatchCount, + size_t GroupCount, + size_t InputChannels, + const int64_t* InputShape, + const int64_t* KernelShape, + const int64_t* DilationShape, + const int64_t* Padding, + const int64_t* StrideShape, + const int64_t* OutputShape, + size_t FilterCount, + const MLAS_ACTIVATION* Activation, + size_t* WorkingBufferSize, + float Beta, + bool InputOutputChannelsLast, + MLAS_THREADPOOL* ThreadPool, + const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig); + +bool +MLASCALL +MlasHalfConv( + const MLAS_CONV_PARAMETERS* Parameters, + const MLAS_FP16* Input, + const MLAS_FP16* Filter, + bool FilterAndBiasArePacked, + const MLAS_FP16* Bias, + MLAS_FP16* WorkingBuffer, + MLAS_FP16* Output, + MLAS_THREADPOOL* ThreadPool + ); + +size_t +MLASCALL +MlasHalfConvPackWeightsAndBiasSize( + size_t FilterCount, + size_t InputChannels, + const int64_t* KernelShape, + const int64_t* DilationShape + ); + +bool +MLASCALL +MlasHalfConvPackWeightsAndBias( + size_t FilterCount, + size_t InputChannels, + const int64_t* KernelShape, + const int64_t* DilationShape, + const MLAS_FP16* Filter, + const MLAS_FP16* Bias, + void* PackedWeightsAndBias, + MLAS_THREADPOOL* ThreadPool + ); +} diff --git a/onnxruntime/core/mlas/lib/mlasi.h b/onnxruntime/core/mlas/lib/mlasi.h index 7e8c5e931b0ef..85a30373611c6 100644 --- a/onnxruntime/core/mlas/lib/mlasi.h +++ b/onnxruntime/core/mlas/lib/mlasi.h @@ -101,6 +101,46 @@ Module Name: #define MLAS_FORCEINLINE __attribute__ ((always_inline)) inline #endif +MLAS_FORCEINLINE +bool +MlasMultiplyOverflowsSizeT( + size_t a, + size_t b, + size_t* out + ) +{ +#if defined(__has_builtin) +#if __has_builtin(__builtin_mul_overflow) + return __builtin_mul_overflow(a, b, out); +#endif +#endif + if (b != 0 && a > (std::numeric_limits::max)() / b) { + return true; + } + *out = a * b; + return false; +} + +MLAS_FORCEINLINE +bool +MlasAddOverflowsSizeT( + size_t a, + size_t b, + size_t* out + ) +{ +#if defined(__has_builtin) +#if __has_builtin(__builtin_add_overflow) + return __builtin_add_overflow(a, b, out); +#endif +#endif + if (a > (std::numeric_limits::max)() - b) { + return true; + } + *out = a + b; + return false; +} + // // Macro to tag globals as internal data shared with kernels written in // assembly. These globals are marked with having hidden visibility to avoid @@ -1018,6 +1058,90 @@ bool void* PackedB); #endif +typedef +bool +(MLASCALL MLAS_HALF_GEMM_BATCH_OVERRIDE)( + size_t M, + size_t N, + size_t K, + size_t BatchN, + const MLAS_HALF_GEMM_DATA_PARAMS* DataParams, + MLAS_THREADPOOL* ThreadPool, + const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig); + +typedef +size_t +(MLASCALL MLAS_HALF_GEMM_PACK_B_SIZE_OVERRIDE)( + CBLAS_TRANSPOSE TransA, + CBLAS_TRANSPOSE TransB, + size_t N, + size_t K); + +typedef +bool +(MLASCALL MLAS_HALF_GEMM_PACK_B_OVERRIDE)( + CBLAS_TRANSPOSE TransA, + CBLAS_TRANSPOSE TransB, + size_t N, + size_t K, + const MLAS_FP16* B, + size_t ldb, + void* PackedB); + +typedef +bool +(MLASCALL MLAS_HALF_CONV_PREPARE_OVERRIDE)( + MLAS_CONV_PARAMETERS* Parameters, + size_t Dimensions, + size_t BatchCount, + size_t GroupCount, + size_t InputChannels, + const int64_t* InputShape, + const int64_t* KernelShape, + const int64_t* DilationShape, + const int64_t* Padding, + const int64_t* StrideShape, + const int64_t* OutputShape, + size_t FilterCount, + const MLAS_ACTIVATION* Activation, + size_t* WorkingBufferSize, + float Beta, + bool InputOutputChannelsLast, + MLAS_THREADPOOL* ThreadPool, + const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig); + +typedef +bool +(MLASCALL MLAS_HALF_CONV_OVERRIDE)( + const MLAS_CONV_PARAMETERS* Parameters, + const MLAS_FP16* Input, + const MLAS_FP16* Filter, + bool FilterAndBiasArePacked, + const MLAS_FP16* Bias, + MLAS_FP16* WorkingBuffer, + MLAS_FP16* Output, + MLAS_THREADPOOL* ThreadPool); + +typedef +size_t +(MLASCALL MLAS_HALF_CONV_PACK_WEIGHTS_AND_BIAS_SIZE_OVERRIDE)( + size_t FilterCount, + size_t InputChannels, + const int64_t* KernelShape, + const int64_t* DilationShape); + +typedef +bool +(MLASCALL MLAS_HALF_CONV_PACK_WEIGHTS_AND_BIAS_OVERRIDE)( + size_t FilterCount, + size_t InputChannels, + const int64_t* KernelShape, + const int64_t* DilationShape, + const MLAS_FP16* Filter, + const MLAS_FP16* Bias, + void* PackedWeightsAndBias, + MLAS_THREADPOOL* ThreadPool); + extern "C" { #if defined(MLAS_TARGET_AMD64_IX86) @@ -1556,6 +1680,15 @@ struct MLAS_PLATFORM { MLAS_DYNAMIC_QGEMM_BATCH_OVERRIDE* MlasDynamicQGemmBatchOverride = nullptr; MLAS_DYNAMIC_QGEMM_PACK_B_SIZE_OVERRIDE* MlasDynamicQGemmPackBSizeOverride = nullptr; MLAS_DYNAMIC_QGEMM_PACK_B_OVERRIDE* MlasDynamicQGemmPackBOverride = nullptr; + // MLAS HalfGemm overrides + MLAS_HALF_GEMM_BATCH_OVERRIDE* MlasHalfGemmBatchOverride = nullptr; + MLAS_HALF_GEMM_PACK_B_SIZE_OVERRIDE* MlasHalfGemmPackBSizeOverride = nullptr; + MLAS_HALF_GEMM_PACK_B_OVERRIDE* MlasHalfGemmPackBOverride = nullptr; + // MLAS HalfConv overrides + MLAS_HALF_CONV_PREPARE_OVERRIDE* MlasHalfConvPrepareOverride = nullptr; + MLAS_HALF_CONV_OVERRIDE* MlasHalfConvOverride = nullptr; + MLAS_HALF_CONV_PACK_WEIGHTS_AND_BIAS_SIZE_OVERRIDE* MlasHalfConvPackWeightsAndBiasSizeOverride = nullptr; + MLAS_HALF_CONV_PACK_WEIGHTS_AND_BIAS_OVERRIDE* MlasHalfConvPackWeightsAndBiasOverride = nullptr; // MLAS Conv overrides MLAS_CONV_PREPARE_FLOAT_OVERRIDE* MlasConvPrepareOverride = nullptr; MLAS_CONV_FLOAT_OVERRIDE* MlasConvOverride = nullptr; diff --git a/onnxruntime/core/mlas/lib/platform.cpp b/onnxruntime/core/mlas/lib/platform.cpp index 54d47449a23a6..796929450810c 100644 --- a/onnxruntime/core/mlas/lib/platform.cpp +++ b/onnxruntime/core/mlas/lib/platform.cpp @@ -704,13 +704,20 @@ Return Value: } #if defined(USE_KLEIDIAI) - if(MLAS_CPUIDINFO::GetCPUIDInfo().HasArm_SME()){ + if(MLAS_CPUIDINFO::GetCPUIDInfo().HasArm_SME() || MLAS_CPUIDINFO::GetCPUIDInfo().HasArm_SME2()){ this->MlasSGemmBatchOverride = ArmKleidiAI::MlasGemmBatch; this->MlasSGemmPackBSizeOverride = ArmKleidiAI::MlasGemmPackBSize; this->MlasSGemmPackBOverride = ArmKleidiAI::MlasGemmPackB; this->MlasDynamicQGemmBatchOverride = ArmKleidiAI::MlasDynamicQGemmBatch; this->MlasDynamicQGemmPackBSizeOverride = ArmKleidiAI::MlasDynamicQGemmPackBSize; this->MlasDynamicQGemmPackBOverride = ArmKleidiAI::MlasDynamicQGemmPackB; + this->MlasHalfGemmBatchOverride = ArmKleidiAI::MlasHalfGemmBatch; + this->MlasHalfGemmPackBSizeOverride = ArmKleidiAI::MlasHalfGemmKleidiAIPackBSize; + this->MlasHalfGemmPackBOverride = ArmKleidiAI::MlasHalfGemmKleidiAIPackB; + this->MlasHalfConvPrepareOverride = ArmKleidiAI::MlasHalfConvPrepare; + this->MlasHalfConvOverride = ArmKleidiAI::MlasHalfConv; + this->MlasHalfConvPackWeightsAndBiasSizeOverride = ArmKleidiAI::MlasHalfConvPackWeightsAndBiasSize; + this->MlasHalfConvPackWeightsAndBiasOverride = ArmKleidiAI::MlasHalfConvPackWeightsAndBias; this->MlasConvPrepareOverride = ArmKleidiAI::MlasConvPrepare; this->MlasConvOverride = ArmKleidiAI::MlasConv; this->MlasConvSGemmRouteOverride = ArmKleidiAI::MlasConvSGemmRoute; diff --git a/onnxruntime/test/mlas/unittest/test_conv2d.cpp b/onnxruntime/test/mlas/unittest/test_conv2d.cpp index 091d4ee833f8f..0ce78c816f68a 100644 --- a/onnxruntime/test/mlas/unittest/test_conv2d.cpp +++ b/onnxruntime/test/mlas/unittest/test_conv2d.cpp @@ -4,6 +4,101 @@ #include "test_conv2d.h" #include "test_conv2d_fixture.h" +TEST(Conv2d_HalfConv, PrepareRespectsBackendSelectorConfig) { + const int64_t input_shape[] = {5, 5}; + const int64_t kernel_shape[] = {3, 3}; + const int64_t dilation_shape[] = {1, 1}; + const int64_t padding[] = {1, 1, 1, 1}; + const int64_t stride_shape[] = {1, 1}; + const int64_t output_shape[] = {5, 5}; + + MLAS_ACTIVATION activation; + activation.ActivationKind = MlasIdentityActivation; + + MLAS_CONV_PARAMETERS parameters{}; + size_t working_buffer_size = 0; + if (!MlasHalfConvPrepare(¶meters, + 2, + 1, + 1, + 4, + input_shape, + kernel_shape, + dilation_shape, + padding, + stride_shape, + output_shape, + 8, + &activation, + &working_buffer_size, + 0.0f, + false, + nullptr, + nullptr)) { + GTEST_SKIP() << "HalfConv prepare path unavailable"; + } + + MLAS_BACKEND_KERNEL_SELECTOR_CONFIG selector_config; + selector_config.use_kleidiai = false; + + MLAS_CONV_PARAMETERS disabled_parameters{}; + EXPECT_FALSE(MlasHalfConvPrepare(&disabled_parameters, + 2, + 1, + 1, + 4, + input_shape, + kernel_shape, + dilation_shape, + padding, + stride_shape, + output_shape, + 8, + &activation, + &working_buffer_size, + 0.0f, + false, + nullptr, + &selector_config)); +} + +TEST(Conv2d_HalfConv, PrepareReportsWorkingBufferSizeInBytes) { + const int64_t input_shape[] = {5, 5}; + const int64_t kernel_shape[] = {3, 3}; + const int64_t dilation_shape[] = {1, 1}; + const int64_t padding[] = {1, 1, 1, 1}; + const int64_t stride_shape[] = {1, 1}; + const int64_t output_shape[] = {5, 5}; + + MLAS_ACTIVATION activation; + activation.ActivationKind = MlasIdentityActivation; + + MLAS_CONV_PARAMETERS parameters{}; + size_t working_buffer_size = 0; + if (!MlasHalfConvPrepare(¶meters, + 2, + 1, + 1, + 4, + input_shape, + kernel_shape, + dilation_shape, + padding, + stride_shape, + output_shape, + 8, + &activation, + &working_buffer_size, + 0.0f, + false, + nullptr, + nullptr)) { + GTEST_SKIP() << "HalfConv prepare path unavailable"; + } + + EXPECT_EQ(working_buffer_size, parameters.OutputSize * parameters.FilterCount * sizeof(uint16_t)); +} + static size_t Conv2dRegistLongExecute() { size_t count = MlasLongExecuteTests>::RegisterLongExecute(); if (GetMlasThreadPool() != nullptr) { diff --git a/onnxruntime/test/mlas/unittest/test_halfgemm.cpp b/onnxruntime/test/mlas/unittest/test_halfgemm.cpp index aafdcc14c0028..be0c291d482fc 100644 --- a/onnxruntime/test/mlas/unittest/test_halfgemm.cpp +++ b/onnxruntime/test/mlas/unittest/test_halfgemm.cpp @@ -15,6 +15,320 @@ Module Name: --*/ #include "test_halfgemm.h" +#include "core/mlas/lib/halfgemm.h" +#if defined(USE_KLEIDIAI) +#include "core/mlas/lib/mlasi.h" +#include "core/mlas/lib/kleidiai/mlasi_kleidiai.h" +#endif + +#include +#include +#include +#include +#include + +namespace { + +struct HalfGemmPackBPaddingKernel { + static constexpr size_t PackedK = 4; +}; + +} // namespace + +#if defined(USE_KLEIDIAI) +namespace { + +bool g_test_halfgemm_override_called = false; + +bool MLASCALL +TestHalfGemmBatchOverride( + size_t M, + size_t N, + size_t, + size_t BatchN, + const MLAS_HALF_GEMM_DATA_PARAMS* DataParams, + MLAS_THREADPOOL*, + const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG*) { + g_test_halfgemm_override_called = true; + + for (size_t batch = 0; batch < BatchN; ++batch) { + auto* c = DataParams[batch].C; + for (size_t m = 0; m < M; ++m) { + for (size_t n = 0; n < N; ++n) { + c[m * DataParams[batch].ldc + n] = onnxruntime::MLFloat16(1.0f); + } + } + } + + return true; +} + +void ReferenceHalfGemm( + size_t M, + size_t N, + size_t K, + const MLFp16* A, + const MLFp16* B, + MLFp16* C) { + for (size_t m = 0; m < M; ++m) { + for (size_t n = 0; n < N; ++n) { + float sum = 0.0f; + for (size_t k = 0; k < K; ++k) { + sum += float(A[m * K + k]) * float(B[k * N + n]); + } + C[m * N + n] = MLFp16(sum); + } + } +} + +struct HalfGemmOverrideGuard { + explicit HalfGemmOverrideGuard(MLAS_HALF_GEMM_BATCH_OVERRIDE* replacement) + : original_(GetMlasPlatform().MlasHalfGemmBatchOverride) { + GetMlasPlatform().MlasHalfGemmBatchOverride = replacement; + } + + ~HalfGemmOverrideGuard() { + GetMlasPlatform().MlasHalfGemmBatchOverride = original_; + } + + private: + MLAS_HALF_GEMM_BATCH_OVERRIDE* original_; +}; + +} // namespace + +TEST(HalfGemmKleidiAISelector, DisableKleidiAIBypassesOverride) { + if (!MlasFp16AccelerationSupported()) { + GTEST_SKIP() << "HalfGemm FP16 acceleration not available."; + } + + constexpr size_t M = 5; + constexpr size_t N = 7; + constexpr size_t K = 9; + + std::vector A(M * K); + std::vector B(K * N); + std::vector C(M * N, MLFp16(0.0f)); + std::vector CReference(M * N, MLFp16(0.0f)); + + SmallFloatFill(A.data(), A.size()); + SmallFloatFill(B.data(), B.size()); + + MLAS_HALF_GEMM_DATA_PARAMS data{}; + data.A = A.data(); + data.B = B.data(); + data.C = reinterpret_cast(C.data()); + data.lda = K; + data.ldb = N; + data.ldc = N; + + HalfGemmOverrideGuard guard(TestHalfGemmBatchOverride); + + g_test_halfgemm_override_called = false; + MlasHalfGemmBatch(M, N, K, 1, &data, nullptr); + ASSERT_TRUE(g_test_halfgemm_override_called); + for (const auto& value : C) { + ASSERT_EQ(float(value), 1.0f); + } + + std::fill(C.begin(), C.end(), MLFp16(0.0f)); + ReferenceHalfGemm(M, N, K, A.data(), B.data(), CReference.data()); + + MLAS_BACKEND_KERNEL_SELECTOR_CONFIG selector_config; + selector_config.use_kleidiai = false; + data.BackendKernelSelectorConfig = &selector_config; + + g_test_halfgemm_override_called = false; + MlasHalfGemmBatch(M, N, K, 1, &data, nullptr); + ASSERT_FALSE(g_test_halfgemm_override_called); + + for (size_t i = 0; i < C.size(); ++i) { + ASSERT_TRUE(CloseEnough(float(C[i]), float(CReference[i]))) << "index=" << i; + } +} + +TEST(HalfGemmKleidiAISelector, DisableKleidiAIInAnyBatchBypassesOverride) { + if (!MlasFp16AccelerationSupported()) { + GTEST_SKIP() << "HalfGemm FP16 acceleration not available."; + } + + constexpr size_t M = 5; + constexpr size_t N = 7; + constexpr size_t K = 9; + constexpr size_t BatchN = 2; + + std::vector A(BatchN * M * K); + std::vector B(BatchN * K * N); + std::vector C(BatchN * M * N, MLFp16(0.0f)); + std::vector CReference(BatchN * M * N, MLFp16(0.0f)); + SmallFloatFill(A.data(), A.size()); + SmallFloatFill(B.data(), B.size()); + + std::array data{}; + for (size_t batch = 0; batch < BatchN; ++batch) { + data[batch].A = A.data() + batch * M * K; + data[batch].B = B.data() + batch * K * N; + data[batch].C = reinterpret_cast(C.data() + batch * M * N); + data[batch].lda = K; + data[batch].ldb = N; + data[batch].ldc = N; + ReferenceHalfGemm(M, N, K, + A.data() + batch * M * K, + B.data() + batch * K * N, + CReference.data() + batch * M * N); + } + + MLAS_BACKEND_KERNEL_SELECTOR_CONFIG selector_config; + selector_config.use_kleidiai = false; + data[1].BackendKernelSelectorConfig = &selector_config; + + HalfGemmOverrideGuard guard(TestHalfGemmBatchOverride); + g_test_halfgemm_override_called = false; + MlasHalfGemmBatch(M, N, K, BatchN, data.data(), nullptr); + ASSERT_FALSE(g_test_halfgemm_override_called); + + for (size_t i = 0; i < C.size(); ++i) { + ASSERT_TRUE(CloseEnough(float(C[i]), float(CReference[i]))) << "index=" << i; + } +} +#endif + +namespace { + +struct HalfGemmCase { + const char* name; + size_t M; + size_t N; + size_t K; + size_t Batch; + bool has_bias; +}; + +template +void RunHalfGemmCases(const HalfGemmCase* test_cases, size_t num_cases, Runner run_case) { + for (size_t i = 0; i < num_cases; ++i) { + const auto& test_case = test_cases[i]; + SCOPED_TRACE(testing::Message() + << test_case.name + << " Batch=" << test_case.Batch + << " M=" << test_case.M + << " N=" << test_case.N + << " K=" << test_case.K + << " hasBias=" << test_case.has_bias); + run_case(test_case); + } +} + +void ReferenceHalfGemmPackedCompatibility( + size_t M, + size_t N, + size_t K, + const MLFp16* A, + const MLFp16* B, + MLFp16* C) { + for (size_t m = 0; m < M; ++m) { + for (size_t n = 0; n < N; ++n) { + float sum = 0.0f; + for (size_t k = 0; k < K; ++k) { + MLFp16 down(float(A[m * K + k]) * float(B[k * N + n]) + sum); + sum = float(down); + } + C[m * N + n] = MLFp16(sum); + } + } +} + +constexpr HalfGemmCase kNativeFp16Cases[] = { + {"WideNBatch1WithBias", 43, 500, 401, 1, true}, + {"WideNBatch3NoBias", 43, 500, 401, 3, false}, + {"VectorLikeBatch3WithBias", 1, 32, 79, 3, true}, + {"RectangularBatch1WithBias", 64, 48, 80, 1, true}, + {"RectangularBatch1NoBias", 64, 48, 80, 1, false}, +}; + +template +void RunNativeFp16WithoutOutputProcessorCases(const HalfGemmCase* test_cases, size_t num_cases) { + RunHalfGemmCases(test_cases, num_cases, [](const HalfGemmCase& test_case) { + MlasHalfGemmTest test; + test.TestNativeFp16WithoutOutputProcessor( + test_case.M, test_case.N, test_case.K, test_case.Batch, test_case.has_bias); + }); +} + +constexpr HalfGemmCase kKleidiAIPathNonPackedCases[] = { + {"WideNBatch3WithBias", 43, 500, 401, 3, true}, + {"WideNBatch3NoBias", 43, 500, 401, 3, false}, + {"RectangularBatch1WithBias", 64, 48, 80, 1, true}, + {"RectangularBatch1NoBias", 64, 48, 80, 1, false}, +}; + +constexpr HalfGemmCase kKleidiAIPathPackedCases[] = { + {"WideNBatch1WithBias", 43, 500, 401, 1, true}, + {"WideNBatch1NoBias", 43, 500, 401, 1, false}, + {"RectangularBatch1WithBias", 64, 48, 80, 1, true}, + {"RectangularBatch1NoBias", 64, 48, 80, 1, false}, +}; + +template +void RunKleidiAIWithoutOutputProcessorCases(const HalfGemmCase* test_cases, size_t num_cases) { + RunHalfGemmCases(test_cases, num_cases, [](const HalfGemmCase& test_case) { + MlasHalfGemmTest test; + test.TestKleidiAIWithoutOutputProcessor( + test_case.M, test_case.N, test_case.K, test_case.Batch, test_case.has_bias); + }); +} + +} // namespace + +TEST(HalfGemm, ZeroKInitializesBiasAndRunsOutputProcessor) { + constexpr size_t M = 3; + constexpr size_t N = 4; + constexpr size_t K = 0; + + std::vector Bias{MLFp16(1.0f), MLFp16(-2.0f), MLFp16(3.5f), MLFp16(0.25f)}; + std::vector C(M * N, MLFp16(-9.0f)); + std::vector CFloat(M * N, -9.0f); + + MLAS_ACTIVATION act; + act.ActivationKind = MlasIdentityActivation; + MLAS_HALF_GEMM_2FLOAT_PROCESSOR output_processor(act, CFloat.data(), N); + + MLAS_HALF_GEMM_DATA_PARAMS data{}; + data.Bias = reinterpret_cast(Bias.data()); + data.C = reinterpret_cast(C.data()); + data.ldc = N; + data.AIsfp32 = true; + data.OutputProcessor = &output_processor; + + MlasHalfGemmBatch(M, N, K, 1, &data, nullptr); + + for (size_t m = 0; m < M; ++m) { + for (size_t n = 0; n < N; ++n) { + const size_t index = m * N + n; + ASSERT_EQ(float(C[index]), float(Bias[n])) << "index=" << index; + ASSERT_EQ(CFloat[index], float(Bias[n])) << "index=" << index; + } + } +} + +TEST(HalfGemm, ZeroKInitializesZeroWithoutBias) { + constexpr size_t M = 3; + constexpr size_t N = 4; + constexpr size_t K = 0; + + std::vector C(M * N, MLFp16(-9.0f)); + + MLAS_HALF_GEMM_DATA_PARAMS data{}; + data.C = reinterpret_cast(C.data()); + data.ldc = N; + data.BIsfp32 = true; + + MlasHalfGemmBatch(M, N, K, 1, &data, nullptr); + + for (const auto& value : C) { + ASSERT_EQ(float(value), 0.0f); + } +} // // Short Execute() test helper to register each test separately by all parameters. @@ -166,3 +480,435 @@ static UNUSED_VARIABLE bool added_to_main = AddTestRegister([](bool is_short_exe } return HalfGemmRegistLongExecute() > 0; }); + +TEST(HalfGemmPackB, ReturnsZeroOnOverflow) { + const size_t max = (std::numeric_limits::max)(); + EXPECT_EQ(MlasHalfGemmPackBSize(1, max, true), size_t{0}); + EXPECT_EQ(MlasHalfGemmPackBSize(max, 2, true), size_t{0}); +} + +TEST(HalfGemmPackB, PackBReturnsOnOverflow) { + const size_t max = (std::numeric_limits::max)(); + std::vector b(1); + std::vector packed_b(1); + + EXPECT_NO_THROW(MlasHalfGemmPackB(max, 2, b.data(), max, packed_b.data())); +} + +TEST(HalfGemmPackB, PackBRejectsInvalidArguments) { + constexpr size_t N = 5; + constexpr size_t K = 3; + + std::vector b(K * N); + std::vector packed_b(MlasHalfGemmPackBSize(N, K, false)); + ASSERT_FALSE(packed_b.empty()); + + EXPECT_NO_THROW(MlasHalfGemmPackB(N, K, b.data(), N - 1, packed_b.data())); + EXPECT_NO_THROW(MlasHalfGemmPackB(N, K, nullptr, N, packed_b.data())); + EXPECT_NO_THROW(MlasHalfGemmPackB(N, K, b.data(), N, nullptr)); +} + +TEST(HalfGemmPackB, CopyPackBZeroPadsAlignedKTail) { + constexpr size_t N = 5; + constexpr size_t K = 3; + constexpr size_t AlignedK = 4; + + std::vector<_mlas_fp16_> b(K * N); + for (size_t i = 0; i < b.size(); ++i) { + b[i] = static_cast<_mlas_fp16_>(i + 1); + } + + constexpr _mlas_fp16_ stale_tail_value = 0xFFFF; + std::vector<_mlas_fp16_> packed(AlignedK * N, stale_tail_value); + + MlasHalfGemmCopyPackB( + packed.data(), + b.data(), + N, + N, + K); + + for (size_t n = 0; n < N; ++n) { + EXPECT_EQ(packed[K * N + n], 0) << "n=" << n; + } +} + +TEST(HalfGemmPackB, GenericPackedBFlagRunsOnFallback) { + if (!MlasFp16AccelerationSupported()) { + GTEST_SKIP(); + } + + constexpr size_t M = 5; + constexpr size_t N = 7; + constexpr size_t K = 9; + + std::vector A(M * K); + std::vector B(K * N); + std::vector C(M * N, MLFp16(0.0f)); + std::vector CReference(M * N, MLFp16(0.0f)); + + SmallFloatFill(A.data(), A.size()); + SmallFloatFill(B.data(), B.size()); + + const size_t packed_b_size = MlasHalfGemmPackBSize(N, K, false); + if (packed_b_size == 0) { + GTEST_SKIP(); + } + + std::vector packed_b(packed_b_size); + MlasHalfGemmPackB(N, K, reinterpret_cast(B.data()), N, packed_b.data()); + + MLAS_HALF_GEMM_DATA_PARAMS data{}; + data.A = A.data(); + data.B = packed_b.data(); + data.C = reinterpret_cast(C.data()); + data.lda = K; + data.ldb = 0; + data.ldc = N; + data.BIsPacked = true; + + MLAS_BACKEND_KERNEL_SELECTOR_CONFIG selector_config{}; + selector_config.use_kleidiai = false; + data.BackendKernelSelectorConfig = &selector_config; + + MlasHalfGemmBatch(M, N, K, 1, &data, nullptr); + ReferenceHalfGemmPackedCompatibility(M, N, K, A.data(), B.data(), CReference.data()); + + for (size_t i = 0; i < C.size(); ++i) { + ASSERT_TRUE(CloseEnough(float(C[i]), float(CReference[i]))) << "index=" << i; + } +} + +TEST(HalfGemmKleidiAINativeFp16, NoPackSingleThreadWithoutOutputProcessor) { + if (!MlasFp16AccelerationSupported()) { + GTEST_SKIP(); + } + + MlasHalfGemmTest test; + test.TestNativeFp16WithoutOutputProcessor(43, 500, 401, 1, true); +} + +TEST(HalfGemmKleidiAINativeFp16, NoPackSingleThreadWithoutOutputProcessorBatch3) { + if (!MlasFp16AccelerationSupported()) { + GTEST_SKIP(); + } + + MlasHalfGemmTest test; + test.TestNativeFp16WithoutOutputProcessor(43, 500, 401, 3, true); +} + +TEST(HalfGemmKleidiAINativeFp16, NoPackThreadedWithoutOutputProcessor) { + if (!MlasFp16AccelerationSupported()) { + GTEST_SKIP(); + } + if (GetMlasThreadPool() == nullptr) { + GTEST_SKIP(); + } + + MlasHalfGemmTest test; + test.TestNativeFp16WithoutOutputProcessor(43, 500, 401, 1, true); +} + +TEST(HalfGemmKleidiAINativeFp16, NoPackThreadedWithoutOutputProcessorBatch3) { + if (!MlasFp16AccelerationSupported()) { + GTEST_SKIP(); + } + if (GetMlasThreadPool() == nullptr) { + GTEST_SKIP(); + } + + MlasHalfGemmTest test; + test.TestNativeFp16WithoutOutputProcessor(43, 500, 401, 3, true); +} + +TEST(HalfGemmKleidiAINativeFp16, NoPackSingleThreadWithoutOutputProcessorVariedShapesAndBias) { + if (!MlasFp16AccelerationSupported()) { + GTEST_SKIP(); + } + + RunNativeFp16WithoutOutputProcessorCases(kNativeFp16Cases, std::size(kNativeFp16Cases)); +} + +TEST(HalfGemmKleidiAINativeFp16, NoPackThreadedWithoutOutputProcessorVariedShapesAndBias) { + if (!MlasFp16AccelerationSupported()) { + GTEST_SKIP(); + } + if (GetMlasThreadPool() == nullptr) { + GTEST_SKIP(); + } + + RunNativeFp16WithoutOutputProcessorCases(kNativeFp16Cases, std::size(kNativeFp16Cases)); +} + +TEST(HalfGemmKleidiAIPath, Fp32AConversionSingleThreadBatch3WithoutOutputProcessor) { + if (!MlasFp16AccelerationSupported()) { + GTEST_SKIP(); + } + + MlasHalfGemmTest test; + test.TestKleidiAIWithoutOutputProcessor(43, 500, 401, 3, true); +} + +TEST(HalfGemmKleidiAIPath, Fp32AConversionSingleThreadVariedShapesAndBiasWithoutOutputProcessor) { + if (!MlasFp16AccelerationSupported()) { + GTEST_SKIP(); + } + + RunKleidiAIWithoutOutputProcessorCases( + kKleidiAIPathNonPackedCases, std::size(kKleidiAIPathNonPackedCases)); +} + +TEST(HalfGemmKleidiAIPath, Fp32BConversionSingleThreadBatch3WithoutOutputProcessor) { + if (!MlasFp16AccelerationSupported()) { + GTEST_SKIP(); + } + + MlasHalfGemmTest test; + test.TestKleidiAIWithoutOutputProcessor(43, 500, 401, 3, true); +} + +TEST(HalfGemmKleidiAIPath, Fp32BConversionSingleThreadVariedShapesAndBiasWithoutOutputProcessor) { + if (!MlasFp16AccelerationSupported()) { + GTEST_SKIP(); + } + + RunKleidiAIWithoutOutputProcessorCases( + kKleidiAIPathNonPackedCases, std::size(kKleidiAIPathNonPackedCases)); +} + +TEST(HalfGemmKleidiAIPath, Fp32ABConversionSingleThreadBatch3WithoutOutputProcessor) { + if (!MlasFp16AccelerationSupported()) { + GTEST_SKIP(); + } + + MlasHalfGemmTest test; + test.TestKleidiAIWithoutOutputProcessor(43, 500, 401, 3, true); +} + +TEST(HalfGemmKleidiAIPath, Fp32ABConversionSingleThreadVariedShapesAndBiasWithoutOutputProcessor) { + if (!MlasFp16AccelerationSupported()) { + GTEST_SKIP(); + } + + RunKleidiAIWithoutOutputProcessorCases( + kKleidiAIPathNonPackedCases, std::size(kKleidiAIPathNonPackedCases)); +} + +TEST(HalfGemmKleidiAIPath, PackedBFp16SingleThreadWithoutOutputProcessor) { + if (!MlasFp16AccelerationSupported()) { + GTEST_SKIP(); + } + if (MlasHalfGemmPackBSize(128, 128, false) == 0) { + GTEST_SKIP(); + } + + MlasHalfGemmTest test; + test.TestKleidiAIWithoutOutputProcessor(43, 500, 401, 1, true); +} + +TEST(HalfGemmKleidiAIPath, PackedBFp16SingleThreadVariedShapesAndBiasWithoutOutputProcessor) { + if (!MlasFp16AccelerationSupported()) { + GTEST_SKIP(); + } + if (MlasHalfGemmPackBSize(128, 128, false) == 0) { + GTEST_SKIP(); + } + + RunKleidiAIWithoutOutputProcessorCases( + kKleidiAIPathPackedCases, std::size(kKleidiAIPathPackedCases)); +} + +TEST(HalfGemmKleidiAIPath, PackedBFloatSingleThreadWithoutOutputProcessor) { + if (!MlasFp16AccelerationSupported()) { + GTEST_SKIP(); + } + if (MlasHalfGemmPackBSize(128, 128, true) == 0) { + GTEST_SKIP(); + } + + MlasHalfGemmTest test; + test.TestKleidiAIWithoutOutputProcessor(43, 500, 401, 1, true); +} + +TEST(HalfGemmKleidiAIPath, PackedBFloatSingleThreadVariedShapesAndBiasWithoutOutputProcessor) { + if (!MlasFp16AccelerationSupported()) { + GTEST_SKIP(); + } + if (MlasHalfGemmPackBSize(128, 128, true) == 0) { + GTEST_SKIP(); + } + + RunKleidiAIWithoutOutputProcessorCases( + kKleidiAIPathPackedCases, std::size(kKleidiAIPathPackedCases)); +} + +#if defined(USE_KLEIDIAI) +// KleidiAI-specific packed-B uses a separate direct-consumption contract from +// generic halfgemm PackB. Unsupported combinations fail at the public API +// boundary because generic MLAS cannot consume this backend-native layout. +TEST(HalfGemmKleidiAIPath, KleidiAIPackedBWithBiasThrows) { + if (!MlasFp16AccelerationSupported()) { + GTEST_SKIP(); + } + if (GetMlasPlatform().MlasHalfGemmBatchOverride == nullptr) { + GTEST_SKIP() << "KleidiAI halfgemm override unavailable"; + } + + constexpr size_t M = 5; + constexpr size_t N = 7; + constexpr size_t K = 9; + + std::vector A(M * K); + std::vector B(K * N); + std::vector Bias(N); + std::vector C(M * N, MLFp16(0.0f)); + + SmallFloatFill(A.data(), A.size()); + SmallFloatFill(B.data(), B.size()); + SmallFloatFill(Bias.data(), Bias.size()); + + const size_t packed_b_size = ArmKleidiAI::MlasHalfGemmKleidiAIPackBSize(CblasNoTrans, CblasNoTrans, N, K); + ASSERT_NE(packed_b_size, size_t{0}); + + std::vector packed_b(packed_b_size); + ASSERT_TRUE(ArmKleidiAI::MlasHalfGemmKleidiAIPackB( + CblasNoTrans, CblasNoTrans, N, K, reinterpret_cast(B.data()), N, packed_b.data())); + + MLAS_HALF_GEMM_DATA_PARAMS data{}; + data.A = A.data(); + data.B = packed_b.data(); + data.Bias = reinterpret_cast(Bias.data()); + data.C = reinterpret_cast(C.data()); + data.lda = K; + data.ldb = 0; + data.ldc = N; + data.BIsBackendNativePacked = true; + + ASSERT_FALSE(ArmKleidiAI::MlasHalfGemmBatch(M, N, K, 1, &data, nullptr, nullptr)); +#if !defined(ORT_NO_EXCEPTIONS) + EXPECT_THROW(MlasHalfGemmBatch(M, N, K, 1, &data, nullptr), std::runtime_error); +#endif +} + +TEST(HalfGemmKleidiAIPath, KleidiAIPackedBWithOutputProcessorThrows) { + if (!MlasFp16AccelerationSupported()) { + GTEST_SKIP(); + } + if (GetMlasPlatform().MlasHalfGemmBatchOverride == nullptr) { + GTEST_SKIP() << "KleidiAI halfgemm override unavailable"; + } + + constexpr size_t M = 5; + constexpr size_t N = 7; + constexpr size_t K = 9; + + std::vector A(M * K); + std::vector B(K * N); + std::vector C(M * N, MLFp16(0.0f)); + std::vector CFloat(M * N, 0.0f); + + SmallFloatFill(A.data(), A.size()); + SmallFloatFill(B.data(), B.size()); + + const size_t packed_b_size = ArmKleidiAI::MlasHalfGemmKleidiAIPackBSize(CblasNoTrans, CblasNoTrans, N, K); + ASSERT_NE(packed_b_size, size_t{0}); + + std::vector packed_b(packed_b_size); + ASSERT_TRUE(ArmKleidiAI::MlasHalfGemmKleidiAIPackB( + CblasNoTrans, CblasNoTrans, N, K, reinterpret_cast(B.data()), N, packed_b.data())); + + MLAS_ACTIVATION act; + act.ActivationKind = MlasIdentityActivation; + MLAS_HALF_GEMM_2FLOAT_PROCESSOR output_processor(act, CFloat.data(), N); + + MLAS_HALF_GEMM_DATA_PARAMS data{}; + data.A = A.data(); + data.B = packed_b.data(); + data.C = reinterpret_cast(C.data()); + data.lda = K; + data.ldb = 0; + data.ldc = N; + data.BIsBackendNativePacked = true; + data.OutputProcessor = &output_processor; + + ASSERT_FALSE(ArmKleidiAI::MlasHalfGemmBatch(M, N, K, 1, &data, nullptr, nullptr)); +#if !defined(ORT_NO_EXCEPTIONS) + EXPECT_THROW(MlasHalfGemmBatch(M, N, K, 1, &data, nullptr), std::runtime_error); +#endif +} + +TEST(HalfGemmKleidiAIPath, ZeroKFallsBack) { + if (GetMlasPlatform().MlasHalfGemmBatchOverride == nullptr) { + GTEST_SKIP() << "KleidiAI halfgemm override unavailable"; + } + + constexpr size_t M = 5; + constexpr size_t N = 7; + constexpr size_t K = 0; + + std::vector Bias(N); + std::vector C(M * N, MLFp16(1.0f)); + + SmallFloatFill(Bias.data(), Bias.size()); + + MLAS_HALF_GEMM_DATA_PARAMS data{}; + data.Bias = reinterpret_cast(Bias.data()); + data.C = reinterpret_cast(C.data()); + data.ldc = N; + + const bool handled = ArmKleidiAI::MlasHalfGemmBatch(M, N, K, 1, &data, nullptr, nullptr); + ASSERT_FALSE(handled); +} + +TEST(HalfGemmKleidiAIPath, KleidiAIPackedBSizeRejectsUnsupportedTranspose) { + if (GetMlasPlatform().MlasHalfGemmBatchOverride == nullptr) { + GTEST_SKIP() << "KleidiAI halfgemm override unavailable"; + } + + constexpr size_t N = 7; + constexpr size_t K = 9; + + EXPECT_EQ(ArmKleidiAI::MlasHalfGemmKleidiAIPackBSize(CblasTrans, CblasNoTrans, N, K), size_t{0}); + EXPECT_EQ(ArmKleidiAI::MlasHalfGemmKleidiAIPackBSize(CblasNoTrans, CblasTrans, N, K), size_t{0}); +} + +TEST(HalfGemmKleidiAIPath, KleidiAIPackedBRejectsUnsupportedTranspose) { + if (GetMlasPlatform().MlasHalfGemmBatchOverride == nullptr) { + GTEST_SKIP() << "KleidiAI halfgemm override unavailable"; + } + + constexpr size_t N = 7; + constexpr size_t K = 9; + + std::vector B(K * N); + SmallFloatFill(B.data(), B.size()); + + const size_t packed_b_size = ArmKleidiAI::MlasHalfGemmKleidiAIPackBSize(CblasNoTrans, CblasNoTrans, N, K); + ASSERT_NE(packed_b_size, size_t{0}); + + std::vector packed_b(packed_b_size); + EXPECT_FALSE(ArmKleidiAI::MlasHalfGemmKleidiAIPackB( + CblasTrans, CblasNoTrans, N, K, reinterpret_cast(B.data()), N, packed_b.data())); + EXPECT_FALSE(ArmKleidiAI::MlasHalfGemmKleidiAIPackB( + CblasNoTrans, CblasTrans, N, K, reinterpret_cast(B.data()), N, packed_b.data())); +} + +TEST(HalfGemmKleidiAIPath, KleidiAIPackedBRejectsInvalidLeadingDimension) { + if (GetMlasPlatform().MlasHalfGemmBatchOverride == nullptr) { + GTEST_SKIP() << "KleidiAI halfgemm override unavailable"; + } + + constexpr size_t N = 7; + constexpr size_t K = 9; + + std::vector B(K * N); + SmallFloatFill(B.data(), B.size()); + + const size_t packed_b_size = ArmKleidiAI::MlasHalfGemmKleidiAIPackBSize(CblasNoTrans, CblasNoTrans, N, K); + ASSERT_NE(packed_b_size, size_t{0}); + + std::vector packed_b(packed_b_size); + EXPECT_FALSE(ArmKleidiAI::MlasHalfGemmKleidiAIPackB( + CblasNoTrans, CblasNoTrans, N, K, reinterpret_cast(B.data()), N - 1, packed_b.data())); +} +#endif diff --git a/onnxruntime/test/mlas/unittest/test_halfgemm.h b/onnxruntime/test/mlas/unittest/test_halfgemm.h index 4db5c2bebca40..0ce058dc3c1eb 100644 --- a/onnxruntime/test/mlas/unittest/test_halfgemm.h +++ b/onnxruntime/test/mlas/unittest/test_halfgemm.h @@ -17,6 +17,9 @@ Module Name: #pragma once #include "test_fp16.h" +#include "core/mlas/lib/mlasi.h" +#include +#include /** * @brief Test class for half precision GEMM @@ -26,6 +29,21 @@ Module Name: template class MlasHalfGemmTest : public MlasTestBase { private: + // Native FP16 is validated against the FP32 reference path + // rather than the existing stepwise-FP16 oracle, so these backend-specific + // tests use a separate tolerance. + static bool CloseEnoughNativeFp16(float got, float ref) { + constexpr float abs_tol = 0.03125f; + constexpr float rel_tol = 0.005f; + + const float diff = std::fabs(got - ref); + if (diff <= abs_tol) { + return true; + } + + return diff <= rel_tol * std::max(std::fabs(got), std::fabs(ref)); + } + MatrixGuardBuffer BufferBPacked; MatrixGuardBuffer BufferA; MatrixGuardBuffer BufferB; @@ -60,7 +78,9 @@ class MlasHalfGemmTest : public MlasTestBase { const MLFp16* Bias, MLFp16* C, size_t ldc, - float* Cfloat) { + float* Cfloat, + bool use_output_processor = true, + bool enforce_kleidiai_override = false) { MLAS_ACTIVATION act; act.ActivationKind = MlasIdentityActivation; std::vector Converters; @@ -84,17 +104,29 @@ class MlasHalfGemmTest : public MlasTestBase { ASSERT_EQ(BatchSize, size_t(1)) << "Packing B not supported in batching yet!"; params.B = PackB(N, K, B, ldb); params.ldb = 0; + params.BIsPacked = true; } else { params.B = B + (K * N * i); params.ldb = ldb; } params.AIsfp32 = std::is_same::value; params.BIsfp32 = std::is_same::value; - Converters.emplace_back(act, Cfloat + (M * N * i), N); - params.OutputProcessor = &(Converters[i]); + if (use_output_processor) { + Converters.emplace_back(act, Cfloat + (M * N * i), N); + params.OutputProcessor = &(Converters[i]); + } else { + params.OutputProcessor = nullptr; + } } - MlasHalfGemmBatch(M, N, K, BatchSize, GemmParameters.data(), threadpool_); + if (enforce_kleidiai_override) { + ASSERT_NE(GetMlasPlatform().MlasHalfGemmBatchOverride, nullptr); + const bool handled = GetMlasPlatform().MlasHalfGemmBatchOverride( + M, N, K, BatchSize, GemmParameters.data(), threadpool_, nullptr); + ASSERT_TRUE(handled); + } else { + MlasHalfGemmBatch(M, N, K, BatchSize, GemmParameters.data(), threadpool_); + } } void ReferenceQgemm(size_t M, @@ -153,6 +185,60 @@ class MlasHalfGemmTest : public MlasTestBase { } } + void ReferenceMlasGemmFp32(size_t M, + size_t N, + size_t K, + size_t BatchSize, + const AType* A, + const BType* B, + const MLFp16* Bias, + float* C) { + MatrixGuardBuffer buffer_a_fp32{}; + MatrixGuardBuffer buffer_b_fp32{}; + MatrixGuardBuffer buffer_c_fp32{}; + + float* AFloat = buffer_a_fp32.GetBuffer(M * K * BatchSize); + float* BFloat = buffer_b_fp32.GetBuffer(K * N * BatchSize); + float* CFloat = buffer_c_fp32.GetBuffer(M * N * BatchSize, true); + + for (size_t i = 0; i < M * K * BatchSize; ++i) { + AFloat[i] = float(A[i]); + } + + for (size_t i = 0; i < K * N * BatchSize; ++i) { + BFloat[i] = float(B[i]); + } + + for (size_t batch = 0; batch < BatchSize; ++batch) { + MlasGemm( + CblasNoTrans, CblasNoTrans, M, N, K, + 1.0f, + AFloat + batch * (M * K), K, + BFloat + batch * (K * N), N, + 0.0f, + CFloat + batch * (M * N), N, + threadpool_, + nullptr); + } + + for (size_t batch = 0; batch < BatchSize; batch++) { + for (size_t m = 0; m < M; m++) { + for (size_t n = 0; n < N; n++) { + const size_t idx = (M * N * batch) + (m * N) + n; + float sum = CFloat[idx]; + if (Bias != nullptr) { + sum += float(Bias[n]); + } + C[idx] = float(MLFp16(sum)); + } + } + + if (Bias) { + Bias += N; + } + } + } + public: MlasHalfGemmTest() : threadpool_(Threaded ? GetMlasThreadPool() : nullptr) {} @@ -200,6 +286,56 @@ class MlasHalfGemmTest : public MlasTestBase { } } + void TestKleidiAIWithoutOutputProcessor(size_t M, size_t N, size_t K, size_t BatchSize, bool withBias) { + if (GetMlasPlatform().MlasHalfGemmBatchOverride == nullptr) { + GTEST_SKIP() << "KleidiAI halfgemm override unavailable"; + } + + const AType* A = BufferA.GetFilledBuffer(K * M * BatchSize + 16, SmallFloatFill); + const BType* B = BufferB.GetFilledBuffer(N * K * BatchSize + 16, SmallFloatFill); + + const MLFp16* Bias = nullptr; + if (withBias) { + Bias = BufferBias.GetFilledBuffer(N * BatchSize + 16, SmallFloatFill); + } + + MLFp16* C = BufferC.GetFilledBuffer(N * M * BatchSize, SmallFloatFill); + float* Cfloat = BufferFloatC.GetBuffer(N * M * BatchSize, true); + float* CReference = BufferCReference.GetFilledBuffer( + N * M * BatchSize, + [](float* start, size_t size) { + std::fill_n(start, size, -1.0f); + }); + MatrixGuardBuffer buffer_fp32_reference{}; + float* CFp32Reference = buffer_fp32_reference.GetFilledBuffer( + N * M * BatchSize, + [](float* start, size_t size) { + std::fill_n(start, size, -1.0f); + }); + + this->CallGemm(M, N, K, BatchSize, A, K, B, N, Bias, C, N, Cfloat, false, true); + ReferenceQgemm(M, N, K, BatchSize, A, B, Bias, CReference); + ReferenceMlasGemmFp32(M, N, K, BatchSize, A, B, Bias, CFp32Reference); + + for (size_t batch = 0, f = 0; batch < BatchSize; batch++) { + for (size_t m = 0; m < M; m++) { + for (size_t n = 0; n < N; n++, f++) { + ASSERT_TRUE(CloseEnoughNativeFp16(float(C[f]), CFp32Reference[f])) << "@[" << batch << "x" << m << "x" << n << "], " + << "Batch=" << BatchSize << "M=" << M << ", N=" << N << ", K=" << K + << " got=" << float(C[f]) + << " fp32=" << CFp32Reference[f] + << " stepwise=" << CReference[f]; + } + } + } + } + + void TestNativeFp16WithoutOutputProcessor(size_t M, size_t N, size_t K, size_t BatchSize, bool withBias) { + static_assert(std::is_same_v); + static_assert(std::is_same_v); + TestKleidiAIWithoutOutputProcessor(M, N, K, BatchSize, withBias); + } + private: public: static const char* GetTestSuiteName() { diff --git a/onnxruntime/test/mlas/unittest/test_util.h b/onnxruntime/test/mlas/unittest/test_util.h index a000e353f370d..8c31ad943e319 100644 --- a/onnxruntime/test/mlas/unittest/test_util.h +++ b/onnxruntime/test/mlas/unittest/test_util.h @@ -44,6 +44,7 @@ class MatrixGuardBuffer { _BaseBuffer = nullptr; _BaseBufferSize = 0; _ElementsAllocated = 0; + _GuardAddress = nullptr; } ~MatrixGuardBuffer(void) { @@ -150,6 +151,7 @@ class MatrixGuardBuffer { } _ElementsAllocated = 0; + _GuardAddress = nullptr; } private: From faccbbf53e763624fb55f5a453c8eae80e28aabb Mon Sep 17 00:00:00 2001 From: Cathal Lawlor Date: Thu, 11 Jun 2026 10:50:58 +0100 Subject: [PATCH 07/10] Validate HalfGemm packed-B edge cases - 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 --- onnxruntime/core/mlas/lib/halfgemm.cpp | 15 +++++++- onnxruntime/core/mlas/lib/halfgemm.h | 5 +++ .../core/mlas/lib/kleidiai/mlasi_kleidiai.h | 1 + .../test/mlas/unittest/test_halfgemm.cpp | 37 +++++++++++++------ 4 files changed, 46 insertions(+), 12 deletions(-) diff --git a/onnxruntime/core/mlas/lib/halfgemm.cpp b/onnxruntime/core/mlas/lib/halfgemm.cpp index f0771c97b2148..00f34d28c0e29 100644 --- a/onnxruntime/core/mlas/lib/halfgemm.cpp +++ b/onnxruntime/core/mlas/lib/halfgemm.cpp @@ -92,7 +92,6 @@ TryGetHalfGemmBackendSelectorConfig( } #endif - void MLASCALL MlasHalfGemmBatch( @@ -108,7 +107,21 @@ MlasHalfGemmBatch( return; } + if (DataParams == nullptr) { + MLAS_THROW_EX(std::runtime_error, "HalfGemm requires non-null DataParams."); + } + if (K == 0) { + for (size_t gemm_i = 0; gemm_i < BatchN; gemm_i++) { + const auto& data = DataParams[gemm_i]; + if (data.BIsBackendNativePacked && + (data.ldb != 0 || data.Bias != nullptr || data.OutputProcessor != nullptr)) { + MLAS_THROW_EX( + std::runtime_error, + "backend-native halfgemm packed B with K==0 requires ldb == 0, Bias == nullptr, and " + "OutputProcessor == nullptr"); + } + } MlasHalfGemmZeroKBatch(M, N, BatchN, DataParams); return; } diff --git a/onnxruntime/core/mlas/lib/halfgemm.h b/onnxruntime/core/mlas/lib/halfgemm.h index ceaca2b17f85c..fbe208f5b0e64 100644 --- a/onnxruntime/core/mlas/lib/halfgemm.h +++ b/onnxruntime/core/mlas/lib/halfgemm.h @@ -61,6 +61,11 @@ MlasHalfGemmTryGetPackedBSize( size_t* PackedBSize ) { + if (N == 0 || K == 0) { + *PackedBSize = 0; + return true; + } + if (PackedK == 0) { return false; } diff --git a/onnxruntime/core/mlas/lib/kleidiai/mlasi_kleidiai.h b/onnxruntime/core/mlas/lib/kleidiai/mlasi_kleidiai.h index f692d4e3d1ca6..0c429d04ac519 100644 --- a/onnxruntime/core/mlas/lib/kleidiai/mlasi_kleidiai.h +++ b/onnxruntime/core/mlas/lib/kleidiai/mlasi_kleidiai.h @@ -8,6 +8,7 @@ #include "../mlasi.h" #include +#include #include // Fix to ensure compatibility with MSVC build diff --git a/onnxruntime/test/mlas/unittest/test_halfgemm.cpp b/onnxruntime/test/mlas/unittest/test_halfgemm.cpp index be0c291d482fc..149a1ddda0848 100644 --- a/onnxruntime/test/mlas/unittest/test_halfgemm.cpp +++ b/onnxruntime/test/mlas/unittest/test_halfgemm.cpp @@ -33,6 +33,13 @@ struct HalfGemmPackBPaddingKernel { static constexpr size_t PackedK = 4; }; +void ExpectBufferFilledWith(const std::vector& buffer, std::byte expected) { + const auto expected_value = std::to_integer(expected); + for (size_t i = 0; i < buffer.size(); ++i) { + EXPECT_EQ(std::to_integer(buffer[i]), expected_value) << "index=" << i; + } +} + } // namespace #if defined(USE_KLEIDIAI) @@ -487,25 +494,33 @@ TEST(HalfGemmPackB, ReturnsZeroOnOverflow) { EXPECT_EQ(MlasHalfGemmPackBSize(max, 2, true), size_t{0}); } -TEST(HalfGemmPackB, PackBReturnsOnOverflow) { +TEST(HalfGemmPackB, PackBLeavesOutputUnchangedOnOverflow) { const size_t max = (std::numeric_limits::max)(); std::vector b(1); - std::vector packed_b(1); + constexpr std::byte sentinel{0x5A}; + std::vector packed_b(16, sentinel); - EXPECT_NO_THROW(MlasHalfGemmPackB(max, 2, b.data(), max, packed_b.data())); + MlasHalfGemmPackB(max, 2, b.data(), max, packed_b.data()); + ExpectBufferFilledWith(packed_b, sentinel); } -TEST(HalfGemmPackB, PackBRejectsInvalidArguments) { +TEST(HalfGemmPackB, PackBExitsEarlyForInvalidArguments) { constexpr size_t N = 5; constexpr size_t K = 3; + constexpr std::byte sentinel{0x5A}; std::vector b(K * N); - std::vector packed_b(MlasHalfGemmPackBSize(N, K, false)); + std::vector packed_b(MlasHalfGemmPackBSize(N, K, false), sentinel); ASSERT_FALSE(packed_b.empty()); - EXPECT_NO_THROW(MlasHalfGemmPackB(N, K, b.data(), N - 1, packed_b.data())); - EXPECT_NO_THROW(MlasHalfGemmPackB(N, K, nullptr, N, packed_b.data())); - EXPECT_NO_THROW(MlasHalfGemmPackB(N, K, b.data(), N, nullptr)); + MlasHalfGemmPackB(N, K, b.data(), N - 1, packed_b.data()); + ExpectBufferFilledWith(packed_b, sentinel); + + MlasHalfGemmPackB(N, K, nullptr, N, packed_b.data()); + ExpectBufferFilledWith(packed_b, sentinel); + + MlasHalfGemmPackB(N, K, b.data(), N, nullptr); + ExpectBufferFilledWith(packed_b, sentinel); } TEST(HalfGemmPackB, CopyPackBZeroPadsAlignedKTail) { @@ -746,7 +761,7 @@ TEST(HalfGemmKleidiAIPath, PackedBFloatSingleThreadVariedShapesAndBiasWithoutOut // KleidiAI-specific packed-B uses a separate direct-consumption contract from // generic halfgemm PackB. Unsupported combinations fail at the public API // boundary because generic MLAS cannot consume this backend-native layout. -TEST(HalfGemmKleidiAIPath, KleidiAIPackedBWithBiasThrows) { +TEST(HalfGemmKleidiAIPath, KleidiAIPackedBWithBiasIsRejected) { if (!MlasFp16AccelerationSupported()) { GTEST_SKIP(); } @@ -790,7 +805,7 @@ TEST(HalfGemmKleidiAIPath, KleidiAIPackedBWithBiasThrows) { #endif } -TEST(HalfGemmKleidiAIPath, KleidiAIPackedBWithOutputProcessorThrows) { +TEST(HalfGemmKleidiAIPath, KleidiAIPackedBWithOutputProcessorIsRejected) { if (!MlasFp16AccelerationSupported()) { GTEST_SKIP(); } @@ -837,7 +852,7 @@ TEST(HalfGemmKleidiAIPath, KleidiAIPackedBWithOutputProcessorThrows) { #endif } -TEST(HalfGemmKleidiAIPath, ZeroKFallsBack) { +TEST(HalfGemmKleidiAIPath, ZeroKIsNotHandledByKleidiAIOverride) { if (GetMlasPlatform().MlasHalfGemmBatchOverride == nullptr) { GTEST_SKIP() << "KleidiAI halfgemm override unavailable"; } From 87f7bf753e123578f81d2cac9938f87b606f2a48 Mon Sep 17 00:00:00 2001 From: Cathal Lawlor Date: Fri, 12 Jun 2026 18:06:54 +0100 Subject: [PATCH 08/10] Initialize HalfConv layout flags 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 --- .../mlas/lib/kleidiai/halfconv_kleidiai.cpp | 1 + .../test/mlas/unittest/test_conv2d.cpp | 65 +++++++++++++++++++ 2 files changed, 66 insertions(+) diff --git a/onnxruntime/core/mlas/lib/kleidiai/halfconv_kleidiai.cpp b/onnxruntime/core/mlas/lib/kleidiai/halfconv_kleidiai.cpp index ee487437f336d..fb4d7b3063eb5 100644 --- a/onnxruntime/core/mlas/lib/kleidiai/halfconv_kleidiai.cpp +++ b/onnxruntime/core/mlas/lib/kleidiai/halfconv_kleidiai.cpp @@ -915,6 +915,7 @@ bool Parameters->InputChannels = InputChannels; Parameters->FilterCount = FilterCount; Parameters->Beta = Beta; + Parameters->ChannelsLast = InputOutputChannelsLast; Parameters->InputOutputChannelsLast = InputOutputChannelsLast; size_t input_size = 1; diff --git a/onnxruntime/test/mlas/unittest/test_conv2d.cpp b/onnxruntime/test/mlas/unittest/test_conv2d.cpp index 0ce78c816f68a..55d0370d9076a 100644 --- a/onnxruntime/test/mlas/unittest/test_conv2d.cpp +++ b/onnxruntime/test/mlas/unittest/test_conv2d.cpp @@ -99,6 +99,71 @@ TEST(Conv2d_HalfConv, PrepareReportsWorkingBufferSizeInBytes) { EXPECT_EQ(working_buffer_size, parameters.OutputSize * parameters.FilterCount * sizeof(uint16_t)); } +TEST(Conv2d_HalfConv, PrepareInitializesLayoutFlags) { + const int64_t input_shape[] = {5, 5}; + const int64_t kernel_shape[] = {3, 3}; + const int64_t dilation_shape[] = {1, 1}; + const int64_t padding[] = {1, 1, 1, 1}; + const int64_t stride_shape[] = {1, 1}; + const int64_t output_shape[] = {5, 5}; + + MLAS_ACTIVATION activation; + activation.ActivationKind = MlasIdentityActivation; + + MLAS_CONV_PARAMETERS nchw_parameters{}; + nchw_parameters.ChannelsLast = true; + nchw_parameters.InputOutputChannelsLast = true; + size_t nchw_working_buffer_size = 0; + if (!MlasHalfConvPrepare(&nchw_parameters, + 2, + 1, + 1, + 4, + input_shape, + kernel_shape, + dilation_shape, + padding, + stride_shape, + output_shape, + 8, + &activation, + &nchw_working_buffer_size, + 0.0f, + false, + nullptr, + nullptr)) { + GTEST_SKIP() << "HalfConv prepare path unavailable"; + } + + EXPECT_FALSE(nchw_parameters.ChannelsLast); + EXPECT_FALSE(nchw_parameters.InputOutputChannelsLast); + + MLAS_CONV_PARAMETERS nhwc_parameters{}; + size_t nhwc_working_buffer_size = 0; + ASSERT_TRUE(MlasHalfConvPrepare(&nhwc_parameters, + 2, + 1, + 1, + 4, + input_shape, + kernel_shape, + dilation_shape, + padding, + stride_shape, + output_shape, + 8, + &activation, + &nhwc_working_buffer_size, + 0.0f, + true, + nullptr, + nullptr)); + + EXPECT_TRUE(nhwc_parameters.ChannelsLast); + EXPECT_TRUE(nhwc_parameters.InputOutputChannelsLast); + EXPECT_EQ(nhwc_working_buffer_size, size_t{0}); +} + static size_t Conv2dRegistLongExecute() { size_t count = MlasLongExecuteTests>::RegisterLongExecute(); if (GetMlasThreadPool() != nullptr) { From 3015fc5cc67d793d3632643f13b042860c88bb24 Mon Sep 17 00:00:00 2001 From: Jonathan Clohessy Date: Tue, 30 Jun 2026 11:48:14 +0100 Subject: [PATCH 09/10] Address review comments Signed-off-by: Jonathan Clohessy --- onnxruntime/core/mlas/inc/mlas.h | 17 ++++++++++ onnxruntime/core/mlas/lib/halfgemm.cpp | 31 +++---------------- .../mlas/lib/kleidiai/convolve_kleidiai.cpp | 4 +-- .../mlas/lib/kleidiai/halfgemm_kleidiai.cpp | 2 ++ onnxruntime/core/mlas/lib/mlasi.h | 20 ------------ .../test/contrib_ops/matmul_4bits_test.cc | 3 ++ 6 files changed, 28 insertions(+), 49 deletions(-) diff --git a/onnxruntime/core/mlas/inc/mlas.h b/onnxruntime/core/mlas/inc/mlas.h index f9f5cc44e8d52..0f9cbd10aa3cd 100644 --- a/onnxruntime/core/mlas/inc/mlas.h +++ b/onnxruntime/core/mlas/inc/mlas.h @@ -893,6 +893,10 @@ struct MLAS_CONV_PARAMETERS { size_t BatchCount; size_t GroupCount; size_t InputChannels; + /** + * Existing Conv layout flag from MlasConvPrepare; for HalfConv this matches + * InputOutputChannelsLast. + */ bool ChannelsLast; size_t InputShape[3]; size_t KernelShape[3]; @@ -911,6 +915,11 @@ struct MLAS_CONV_PARAMETERS { const void* PackedFilter = nullptr; size_t PackedFilterGroupStride = 0; bool FilterIsPacked = false; + /** + * HalfConv-specific: true means caller-provided input and output tensors + * are channels-last (NHWC); false means NCHW public tensors with any NHWC + * staging handled internally. + */ bool InputOutputChannelsLast = false; union { struct { @@ -1952,6 +1961,11 @@ struct MLAS_HALF_GEMM_DATA_PARAMS { * Backend-native packed B is a constrained direct-consumption layout * and does not support runtime Bias or OutputProcessor. * + * Uses MLAS_THROW_EX(std::runtime_error, ...) for contract violations that + * cannot be safely ignored: non-empty work with null DataParams, malformed + * backend-native packed-B parameters on the K == 0 path, or backend-native + * packed B reaching the generic MLAS kernels when no backend override consumes it. + * * Note: We only support uniform batching, so shapes and types of the * input must be same across all parameter blocks. * @@ -1996,6 +2010,9 @@ MlasHalfGemmPackBSize( * @brief For half precision GEMM, pack the right hand * side matrix B * + * @pre For non-zero N and K, B and PackedB must be non-null and ldb must be at + * least N. + * * @param[in] N Number of columns * @param[in] K Number of rows * @param[in] B Address of matrix B diff --git a/onnxruntime/core/mlas/lib/halfgemm.cpp b/onnxruntime/core/mlas/lib/halfgemm.cpp index 00f34d28c0e29..f45b77faac00e 100644 --- a/onnxruntime/core/mlas/lib/halfgemm.cpp +++ b/onnxruntime/core/mlas/lib/halfgemm.cpp @@ -19,9 +19,6 @@ Module Name: #include "mlas_float16.h" #include "halfgemm.h" -#if defined(USE_KLEIDIAI) -#include "kleidiai/mlasi_kleidiai.h" -#endif #include @@ -63,8 +60,7 @@ MlasFp16AccelerationSupported() #endif } -#if defined(USE_KLEIDIAI) -bool +static bool TryGetHalfGemmBackendSelectorConfig( size_t BatchN, const MLAS_HALF_GEMM_DATA_PARAMS* DataParams, @@ -90,7 +86,6 @@ TryGetHalfGemmBackendSelectorConfig( return true; } -#endif void MLASCALL @@ -126,7 +121,6 @@ MlasHalfGemmBatch( return; } -#if defined(USE_KLEIDIAI) const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig = nullptr; if (TryGetHalfGemmBackendSelectorConfig(BatchN, DataParams, BackendKernelSelectorConfig) && GetMlasPlatform().MlasHalfGemmBatchOverride != nullptr && @@ -134,7 +128,6 @@ MlasHalfGemmBatch( M, N, K, BatchN, DataParams, ThreadPool, BackendKernelSelectorConfig)) { return; } -#endif // BIsPacked denotes the generic MLAS halfgemm packed-B layout and can be // consumed here. Backend-native packed layouts are separate and must not @@ -251,6 +244,7 @@ MlasHalfGemmPackB( ) { const auto* dispatch = MlasHalfGemmGetDispatch(); + assert(N == 0 || K == 0 || (B != nullptr && PackedB != nullptr && ldb >= N)); if (B == nullptr || PackedB == nullptr || ldb < N || dispatch->CopyPackBRoutine == nullptr || MlasHalfGemmPackBSize(N, K, false) == 0) { return; @@ -269,18 +263,11 @@ MlasHalfGemmNativePackBSize( const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig ) { -#if defined(USE_KLEIDIAI) if ((!BackendKernelSelectorConfig || BackendKernelSelectorConfig->use_kleidiai) && GetMlasPlatform().MlasHalfGemmPackBSizeOverride != nullptr) { return GetMlasPlatform().MlasHalfGemmPackBSizeOverride(TransA, TransB, N, K); } -#else - MLAS_UNREFERENCED_PARAMETER(TransA); - MLAS_UNREFERENCED_PARAMETER(TransB); - MLAS_UNREFERENCED_PARAMETER(N); - MLAS_UNREFERENCED_PARAMETER(K); - MLAS_UNREFERENCED_PARAMETER(BackendKernelSelectorConfig); -#endif + return 0; } @@ -297,21 +284,11 @@ MlasHalfGemmNativePackB( const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig ) { -#if defined(USE_KLEIDIAI) if ((!BackendKernelSelectorConfig || BackendKernelSelectorConfig->use_kleidiai) && GetMlasPlatform().MlasHalfGemmPackBOverride != nullptr) { return GetMlasPlatform().MlasHalfGemmPackBOverride(TransA, TransB, N, K, B, ldb, PackedB); } -#else - MLAS_UNREFERENCED_PARAMETER(TransA); - MLAS_UNREFERENCED_PARAMETER(TransB); - MLAS_UNREFERENCED_PARAMETER(N); - MLAS_UNREFERENCED_PARAMETER(K); - MLAS_UNREFERENCED_PARAMETER(B); - MLAS_UNREFERENCED_PARAMETER(ldb); - MLAS_UNREFERENCED_PARAMETER(PackedB); - MLAS_UNREFERENCED_PARAMETER(BackendKernelSelectorConfig); -#endif + return false; } diff --git a/onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp b/onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp index 1367ba1d3d7f4..0fe529d9193a5 100644 --- a/onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp +++ b/onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp @@ -460,8 +460,8 @@ static std::shared_ptr GetOrCreateLhsPtrTableSme(const size_t ci, }; // Cache of computed LHS pointer offsets. thread_local to prevent interference from parallel sessions. - // Entries include pointers to the pad buffer for out-of-bounds pixels, so group by pad buffer identity and - // invalidate the old group if the grow-only pad buffer is reallocated. + // Entries include pointers to the pad buffer for out-of-bounds pixels; if the grow-only pad buffer + // reallocates, erase the old group so stale pointer tables cannot reference the previous allocation. using LhsPtrsCache = std::unordered_map>; thread_local std::unordered_map lhs_ptrs_cache_by_pad; thread_local const float* last_pad_ptr = nullptr; diff --git a/onnxruntime/core/mlas/lib/kleidiai/halfgemm_kleidiai.cpp b/onnxruntime/core/mlas/lib/kleidiai/halfgemm_kleidiai.cpp index b443164ce69a1..e654c843081f3 100644 --- a/onnxruntime/core/mlas/lib/kleidiai/halfgemm_kleidiai.cpp +++ b/onnxruntime/core/mlas/lib/kleidiai/halfgemm_kleidiai.cpp @@ -177,6 +177,8 @@ ArmKleidiAI::MlasHalfGemmBatch( return false; } + // TODO: Plumb MLAS_ACTIVATION through this call site if MLAS_HALF_GEMM_DATA_PARAMS + // grows fused activation support. const float clamp_min = -std::numeric_limits::infinity(); const float clamp_max = std::numeric_limits::infinity(); diff --git a/onnxruntime/core/mlas/lib/mlasi.h b/onnxruntime/core/mlas/lib/mlasi.h index 85a30373611c6..699c123e71930 100644 --- a/onnxruntime/core/mlas/lib/mlasi.h +++ b/onnxruntime/core/mlas/lib/mlasi.h @@ -101,26 +101,6 @@ Module Name: #define MLAS_FORCEINLINE __attribute__ ((always_inline)) inline #endif -MLAS_FORCEINLINE -bool -MlasMultiplyOverflowsSizeT( - size_t a, - size_t b, - size_t* out - ) -{ -#if defined(__has_builtin) -#if __has_builtin(__builtin_mul_overflow) - return __builtin_mul_overflow(a, b, out); -#endif -#endif - if (b != 0 && a > (std::numeric_limits::max)() / b) { - return true; - } - *out = a * b; - return false; -} - MLAS_FORCEINLINE bool MlasAddOverflowsSizeT( diff --git a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc index de7cab33febf8..e01d902235fd6 100644 --- a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc @@ -328,7 +328,10 @@ void TestMatMulNBitsTyped(std::optional abs_error = std::nullopt, } else if (base_opts.accuracy_level == 4) { base_opts.output_abs_error = 0.1f; } else if constexpr (std::is_same::value) { + // The fp16 provider paths compare against a float reference while native kernels may accumulate + // in fp16 (for example native HGEMM on SME; see PR #28786), so allow slightly wider drift. #if defined(USE_WEBGPU) + // WebGPU's fp16 path has additional provider-specific rounding drift for these quantized matmul cases. base_opts.output_abs_error = 0.1f; #else base_opts.output_abs_error = 0.065f; From a8861ad7ef33f88238b5f4db833cafb74e1dbfc5 Mon Sep 17 00:00:00 2001 From: Jonathan Clohessy Date: Mon, 13 Jul 2026 16:54:51 +0100 Subject: [PATCH 10/10] Address review feedback Signed-off-by: Jonathan Clohessy --- onnxruntime/contrib_ops/cpu/moe/moe_cpu.cc | 1 + .../core/mlas/lib/kleidiai/halfconv_kleidiai.cpp | 2 +- .../core/mlas/lib/kleidiai/halfgemm_kleidiai.cpp | 13 +++++++------ onnxruntime/core/mlas/lib/kleidiai/mlasi_kleidiai.h | 7 +++---- onnxruntime/core/providers/cpu/fp16/fp16_conv.cc | 4 ++++ onnxruntime/core/providers/cpu/math/gemm.cc | 1 + 6 files changed, 17 insertions(+), 11 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/moe/moe_cpu.cc b/onnxruntime/contrib_ops/cpu/moe/moe_cpu.cc index c3194f1c57ba7..5db2d2e6256d2 100644 --- a/onnxruntime/contrib_ops/cpu/moe/moe_cpu.cc +++ b/onnxruntime/contrib_ops/cpu/moe/moe_cpu.cc @@ -550,6 +550,7 @@ Status MoE::ComputeGEMM(const MLFloat16* A, const MLFloat16* B, MLFlo } else { params.ldb = static_cast(N); } + params.BackendKernelSelectorConfig = &mlas_backend_kernel_selector_config_; MlasHalfGemmBatch(static_cast(M), static_cast(N), static_cast(K), 1, ¶ms, nullptr); return Status::OK(); diff --git a/onnxruntime/core/mlas/lib/kleidiai/halfconv_kleidiai.cpp b/onnxruntime/core/mlas/lib/kleidiai/halfconv_kleidiai.cpp index fb4d7b3063eb5..b4b1c2ff19547 100644 --- a/onnxruntime/core/mlas/lib/kleidiai/halfconv_kleidiai.cpp +++ b/onnxruntime/core/mlas/lib/kleidiai/halfconv_kleidiai.cpp @@ -32,7 +32,7 @@ struct KaiHalfConvTlsBuffers { std::vector packed_lhs; void ReleaseLargeBuffers() { - MlasShrinkKleidiAIScratchIfTooLarge(packed_lhs); + ArmKleidiAI::MlasShrinkKleidiAIScratchIfTooLarge(packed_lhs); } }; diff --git a/onnxruntime/core/mlas/lib/kleidiai/halfgemm_kleidiai.cpp b/onnxruntime/core/mlas/lib/kleidiai/halfgemm_kleidiai.cpp index e654c843081f3..3baed91b7ad9a 100644 --- a/onnxruntime/core/mlas/lib/kleidiai/halfgemm_kleidiai.cpp +++ b/onnxruntime/core/mlas/lib/kleidiai/halfgemm_kleidiai.cpp @@ -24,10 +24,10 @@ struct KaiHalfTlsBuffers { std::vector rhs_packed; void ReleaseLargeBuffers() { - MlasShrinkKleidiAIScratchIfTooLarge(lhs_converted); - MlasShrinkKleidiAIScratchIfTooLarge(rhs_converted); - MlasShrinkKleidiAIScratchIfTooLarge(bias_zero); - MlasShrinkKleidiAIScratchIfTooLarge(rhs_packed); + ArmKleidiAI::MlasShrinkKleidiAIScratchIfTooLarge(lhs_converted); + ArmKleidiAI::MlasShrinkKleidiAIScratchIfTooLarge(rhs_converted); + ArmKleidiAI::MlasShrinkKleidiAIScratchIfTooLarge(bias_zero); + ArmKleidiAI::MlasShrinkKleidiAIScratchIfTooLarge(rhs_packed); } }; @@ -132,6 +132,9 @@ ArmKleidiAI::MlasHalfGemmBatch( MLAS_THREADPOOL* ThreadPool, const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig ) { + if (BackendKernelSelectorConfig != nullptr && !BackendKernelSelectorConfig->use_kleidiai) { + return false; + } if (BatchN == 0 || M == 0 || N == 0) { return true; } @@ -144,8 +147,6 @@ ArmKleidiAI::MlasHalfGemmBatch( ScopedKaiHalfTlsCleanup cleanup{g_kai_half_tls}; - MLAS_UNREFERENCED_PARAMETER(ThreadPool); - MLAS_UNREFERENCED_PARAMETER(BackendKernelSelectorConfig); // Validate all batch entries up front so we never partially execute and then // fall back (which would corrupt results for the already-written outputs). bool needs_rhs_packing = false; diff --git a/onnxruntime/core/mlas/lib/kleidiai/mlasi_kleidiai.h b/onnxruntime/core/mlas/lib/kleidiai/mlasi_kleidiai.h index 0c429d04ac519..a03e840e15713 100644 --- a/onnxruntime/core/mlas/lib/kleidiai/mlasi_kleidiai.h +++ b/onnxruntime/core/mlas/lib/kleidiai/mlasi_kleidiai.h @@ -8,7 +8,6 @@ #include "../mlasi.h" #include -#include #include // Fix to ensure compatibility with MSVC build @@ -26,7 +25,7 @@ #define KLEIDIAI_KERNEL_LOGGING 0 #endif -#if KLEIDIAI_DEBUG_LOGGING ||KLEIDIAI_KERNEL_LOGGING +#if KLEIDIAI_DEBUG_LOGGING || KLEIDIAI_KERNEL_LOGGING #include #define KLEIDIAI_LOG(tag, msg) \ do { \ @@ -52,6 +51,8 @@ #define KLEIDIAI_KERNEL_LOG(msg) #endif +namespace ArmKleidiAI { + constexpr size_t MaximumRetainedKleidiAIScratchBytes = 8 * 1024 * 1024; template @@ -62,8 +63,6 @@ void MlasShrinkKleidiAIScratchIfTooLarge(std::vector& buffer) } } -namespace ArmKleidiAI { - // By default we should try for SME2 first before falling back to SME. inline const bool UseSME2 = MLAS_CPUIDINFO::GetCPUIDInfo().HasArm_SME2(); inline const bool UseSME = MLAS_CPUIDINFO::GetCPUIDInfo().HasArm_SME(); diff --git a/onnxruntime/core/providers/cpu/fp16/fp16_conv.cc b/onnxruntime/core/providers/cpu/fp16/fp16_conv.cc index 08dbc46213f65..8aee133e39c45 100644 --- a/onnxruntime/core/providers/cpu/fp16/fp16_conv.cc +++ b/onnxruntime/core/providers/cpu/fp16/fp16_conv.cc @@ -12,6 +12,7 @@ #include "core/common/safeint.h" #include "core/common/float16.h" #include "core/framework/op_kernel.h" +#include "core/providers/cpu/mlas_backend_kernel_selector_config_utils.h" #include "core/providers/cpu/nn/conv_attributes.h" #include "contrib_ops/cpu/fused_activation.h" @@ -46,6 +47,7 @@ class FusedConvFp16 final : public OpKernel { FusedConvFp16(const OpKernelInfo& info) : OpKernel(info), conv_attrs_(info) { ORT_ENFORCE(GetFusedActivationAttr(info, activation_).IsOK()); channels_last_ = (info.GetKernelDef().OpName() == "NhwcFusedConv"); + SetupMlasBackendKernelSelectorFromConfigOptions(mlas_backend_kernel_selector_config_, info.GetConfigOptions()); } Status Compute(OpKernelContext* context) const override; @@ -93,6 +95,7 @@ class FusedConvFp16 final : public OpKernel { } MLAS_ACTIVATION activation_; + MLAS_BACKEND_KERNEL_SELECTOR_CONFIG mlas_backend_kernel_selector_config_; ConvAttributes conv_attrs_; TensorShape W_shape_; BufferUniquePtr packed_W_buffer_; @@ -558,6 +561,7 @@ Status FusedConvFp16::Compute(OpKernelContext* context) const { gemm_params.ldc = static_cast(M); gemm_params.Bias = Bdata; gemm_params.OutputProcessor = (!channels_last_ && sum_data) ? nullptr : &act; // process fused activation and add + gemm_params.BackendKernelSelectorConfig = &mlas_backend_kernel_selector_config_; MlasHalfGemmBatch( static_cast(output_count), diff --git a/onnxruntime/core/providers/cpu/math/gemm.cc b/onnxruntime/core/providers/cpu/math/gemm.cc index c0da9aec1e1b1..e12761ac61631 100644 --- a/onnxruntime/core/providers/cpu/math/gemm.cc +++ b/onnxruntime/core/providers/cpu/math/gemm.cc @@ -225,6 +225,7 @@ void Gemm_MLFloat16(CBLAS_TRANSPOSE trans_a, CBLAS_TRANSPOSE trans_b, if (c_shape != nullptr) { data.Bias = c_data; } + data.BackendKernelSelectorConfig = mlas_backend_kernel_selector_config; MlasHalfGemmBatch(M, N, K, 1, &data, thread_pool); return; }