diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md
index 5f2da22aabe65..916ead7621810 100644
--- a/docs/OperatorKernels.md
+++ b/docs/OperatorKernels.md
@@ -174,10 +174,10 @@ The **OpSet Version** column uses the following notation:
|||12|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**indices** = tensor(int64)|
|||11|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**indices** = tensor(int64)|
|Gelu|*in* X:**T**
*out* Y:**T**|20+|**T** = tensor(float)|
-|Gemm|*in* A:**T**
*in* B:**T**
*in* C:**T**
*out* Y:**T**|13+|**T** = tensor(double), tensor(float)|
-|||[11, 12]|**T** = tensor(double), tensor(float)|
-|||[9, 10]|**T** = tensor(double), tensor(float)|
-|||[7, 8]|**T** = tensor(double), tensor(float)|
+|Gemm|*in* A:**T**
*in* B:**T**
*in* C:**T**
*out* Y:**T**|13+|**T** = tensor(double), tensor(float), tensor(float16)|
+|||[11, 12]|**T** = tensor(double), tensor(float), tensor(float16)|
+|||[9, 10]|**T** = tensor(double), tensor(float), tensor(float16)|
+|||[7, 8]|**T** = tensor(double), tensor(float), tensor(float16)|
|GlobalAveragePool|*in* X:**T**
*out* Y:**T**|22+|**T** = tensor(float)|
|||[1, 21]|**T** = tensor(float)|
|GlobalLpPool|*in* X:**T**
*out* Y:**T**|2+|**T** = tensor(float)|
@@ -258,9 +258,9 @@ The **OpSet Version** column uses the following notation:
|||[18, 21]|**T** = tensor(float)|
|||[11, 17]|**T** = tensor(float)|
|||[2, 10]|**T** = tensor(float)|
-|MatMul|*in* A:**T**
*in* B:**T**
*out* Y:**T**|13+|**T** = tensor(double), tensor(float), tensor(int32), tensor(int64), tensor(uint32), tensor(uint64)|
-|||[9, 12]|**T** = tensor(double), tensor(float), tensor(int32), tensor(int64), tensor(uint32), tensor(uint64)|
-|||[1, 8]|**T** = tensor(double), tensor(float)|
+|MatMul|*in* A:**T**
*in* B:**T**
*out* Y:**T**|13+|**T** = tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(uint32), tensor(uint64)|
+|||[9, 12]|**T** = tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(uint32), tensor(uint64)|
+|||[1, 8]|**T** = tensor(double), tensor(float), tensor(float16)|
|MatMulInteger|*in* A:**T1**
*in* B:**T2**
*in* a_zero_point:**T1**
*in* b_zero_point:**T2**
*out* Y:**T3**|10+|**T1** = tensor(int8), tensor(uint8)
**T2** = tensor(int8), tensor(uint8)
**T3** = tensor(int32)|
|Max|*in* data_0:**T**
*out* max:**T**|13+|**T** = tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(int8), tensor(uint32), tensor(uint64), tensor(uint8)|
|||12|**T** = tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(int8), tensor(uint32), tensor(uint64), tensor(uint8)|
diff --git a/onnxruntime/contrib_ops/cpu/moe/moe_cpu.cc b/onnxruntime/contrib_ops/cpu/moe/moe_cpu.cc
index 5db2d2e6256d2..9528995f0d7da 100644
--- a/onnxruntime/contrib_ops/cpu/moe/moe_cpu.cc
+++ b/onnxruntime/contrib_ops/cpu/moe/moe_cpu.cc
@@ -536,7 +536,7 @@ Status MoE::ComputeGEMM(const float* A, const float* B, float* C,
template <>
Status MoE::ComputeGEMM(const MLFloat16* A, const MLFloat16* B, MLFloat16* C,
int64_t M, int64_t K, int64_t N, bool transpose_B) const {
- MLAS_HALF_GEMM_DATA_PARAMS params;
+ MLAS_HALF_GEMM_DATA_PARAMS params{};
params.A = A;
params.lda = static_cast(K);
params.C = C;
diff --git a/onnxruntime/core/framework/prepacked_weights.h b/onnxruntime/core/framework/prepacked_weights.h
index 9695be1e0554c..1a0dada115a72 100644
--- a/onnxruntime/core/framework/prepacked_weights.h
+++ b/onnxruntime/core/framework/prepacked_weights.h
@@ -20,6 +20,10 @@ struct PrePackedWeights final {
InlinedVector> buffers_; // cache pre-packed buffers associated with the kernel
InlinedVector buffer_sizes_; // cache sizes of pre-packed buffers (in bytes)
+ // Set when PrePack successfully produced kernel-owned packed weights that
+ // intentionally cannot be stored in a shared pre-packed weights container.
+ bool has_kernel_owned_packed_weights_{false};
+
// Produces a hash of the buffers stored in the given instance of this class
HashValue GetHash() const;
diff --git a/onnxruntime/core/framework/prepacked_weights_container.cc b/onnxruntime/core/framework/prepacked_weights_container.cc
index 7c832a0ac2691..9fcf46e61837b 100644
--- a/onnxruntime/core/framework/prepacked_weights_container.cc
+++ b/onnxruntime/core/framework/prepacked_weights_container.cc
@@ -15,6 +15,7 @@ PrePackedWeights PrePackedWeights::CreateReferringCopy() const {
}
copy.buffer_sizes_ = buffer_sizes_;
+ copy.has_kernel_owned_packed_weights_ = has_kernel_owned_packed_weights_;
return copy;
}
diff --git a/onnxruntime/core/framework/session_state.cc b/onnxruntime/core/framework/session_state.cc
index 241eb8362ddfa..5d71011197378 100644
--- a/onnxruntime/core/framework/session_state.cc
+++ b/onnxruntime/core/framework/session_state.cc
@@ -525,12 +525,15 @@ Status SessionState::PrepackConstantInitializedTensors(
&weights_to_be_filled_in));
if (is_packed) {
- // BUG CHECK: Ensure that the kernel has filled in the pre-packed weight
- // to be cached if the weight was pre-packed
- ORT_ENFORCE(weights_to_be_filled_in.buffers_.size() > 0,
- "The kernel corresponding to the node ", node.Name(),
- " doesn't have an implementation that can cache computed pre-packed weights");
+ // BUG CHECK: Ensure that a kernel either filled in the pre-packed weights
+ // to be cached, or explicitly marked the packed weights as kernel-owned.
+ ORT_RETURN_IF_NOT(!weights_to_be_filled_in.buffers_.empty() ||
+ weights_to_be_filled_in.has_kernel_owned_packed_weights_,
+ "The kernel corresponding to the node ", node.Name(),
+ " doesn't have an implementation that can cache computed pre-packed weights");
+ }
+ if (is_packed && !weights_to_be_filled_in.has_kernel_owned_packed_weights_) {
const auto& op_type = node.OpType();
// Sanity check
diff --git a/onnxruntime/core/mlas/lib/halfgemm.cpp b/onnxruntime/core/mlas/lib/halfgemm.cpp
index f45b77faac00e..8f140d33a363a 100644
--- a/onnxruntime/core/mlas/lib/halfgemm.cpp
+++ b/onnxruntime/core/mlas/lib/halfgemm.cpp
@@ -125,7 +125,7 @@ MlasHalfGemmBatch(
if (TryGetHalfGemmBackendSelectorConfig(BatchN, DataParams, BackendKernelSelectorConfig) &&
GetMlasPlatform().MlasHalfGemmBatchOverride != nullptr &&
GetMlasPlatform().MlasHalfGemmBatchOverride(
- M, N, K, BatchN, DataParams, ThreadPool, BackendKernelSelectorConfig)) {
+ M, N, K, BatchN, DataParams, ThreadPool)) {
return;
}
@@ -472,6 +472,10 @@ MlasHGemmSupported(
CBLAS_TRANSPOSE TransA,
CBLAS_TRANSPOSE TransB
) {
+ if (!MlasFp16AccelerationSupported()) {
+ return false;
+ }
+
auto* dispatch = GetMlasPlatform().HGemmDispatch;
if (TransA == CblasNoTrans && TransB == CblasTrans) {
return dispatch &&
diff --git a/onnxruntime/core/mlas/lib/halfgemm.h b/onnxruntime/core/mlas/lib/halfgemm.h
index fbe208f5b0e64..74ee22bf12e59 100644
--- a/onnxruntime/core/mlas/lib/halfgemm.h
+++ b/onnxruntime/core/mlas/lib/halfgemm.h
@@ -595,7 +595,10 @@ const MLAS_HALFGEMM_DISPATCH*
MlasHalfGemmGetDispatch()
{
#if defined(MLAS_F16VEC_INTRINSICS_SUPPORTED) && defined(MLAS_TARGET_ARM64)
- return &MlasHalfGemmDispatchNeon;
+ if (MLAS_CPUIDINFO::GetCPUIDInfo().HasFp16VectorAcceleration()) {
+ return &MlasHalfGemmDispatchNeon;
+ }
+ return &MlasHalfGemmDispatchDefault;
#elif defined(MLAS_TARGET_RISCV64) && defined(MLAS_USE_RVV_ZVFH)
if (MLAS_CPUIDINFO::GetCPUIDInfo().HasFp16VectorAcceleration()) {
return &MlasHalfGemmDispatchRvv;
diff --git a/onnxruntime/core/mlas/lib/kleidiai/halfgemm_kleidiai.cpp b/onnxruntime/core/mlas/lib/kleidiai/halfgemm_kleidiai.cpp
index 3baed91b7ad9a..7d57bf8bd7aec 100644
--- a/onnxruntime/core/mlas/lib/kleidiai/halfgemm_kleidiai.cpp
+++ b/onnxruntime/core/mlas/lib/kleidiai/halfgemm_kleidiai.cpp
@@ -129,12 +129,8 @@ ArmKleidiAI::MlasHalfGemmBatch(
size_t K,
size_t BatchN,
const MLAS_HALF_GEMM_DATA_PARAMS* DataParams,
- MLAS_THREADPOOL* ThreadPool,
- const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig
+ MLAS_THREADPOOL* ThreadPool
) {
- if (BackendKernelSelectorConfig != nullptr && !BackendKernelSelectorConfig->use_kleidiai) {
- return false;
- }
if (BatchN == 0 || M == 0 || N == 0) {
return true;
}
diff --git a/onnxruntime/core/mlas/lib/kleidiai/mlasi_kleidiai.h b/onnxruntime/core/mlas/lib/kleidiai/mlasi_kleidiai.h
index a03e840e15713..9ce5751026997 100644
--- a/onnxruntime/core/mlas/lib/kleidiai/mlasi_kleidiai.h
+++ b/onnxruntime/core/mlas/lib/kleidiai/mlasi_kleidiai.h
@@ -395,8 +395,7 @@ MlasHalfGemmBatch(
size_t K,
size_t BatchN,
const MLAS_HALF_GEMM_DATA_PARAMS* DataParams,
- MLAS_THREADPOOL* ThreadPool,
- const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig
+ MLAS_THREADPOOL* ThreadPool
);
size_t
diff --git a/onnxruntime/core/mlas/lib/mlasi.h b/onnxruntime/core/mlas/lib/mlasi.h
index ec0ed1c0a465c..55f8b8e82f6de 100644
--- a/onnxruntime/core/mlas/lib/mlasi.h
+++ b/onnxruntime/core/mlas/lib/mlasi.h
@@ -1055,8 +1055,7 @@ bool
size_t K,
size_t BatchN,
const MLAS_HALF_GEMM_DATA_PARAMS* DataParams,
- MLAS_THREADPOOL* ThreadPool,
- const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig);
+ MLAS_THREADPOOL* ThreadPool);
typedef
size_t
diff --git a/onnxruntime/core/optimizer/matmul_transpose_fusion.cc b/onnxruntime/core/optimizer/matmul_transpose_fusion.cc
index ad678d5384c66..7c0b24af96897 100644
--- a/onnxruntime/core/optimizer/matmul_transpose_fusion.cc
+++ b/onnxruntime/core/optimizer/matmul_transpose_fusion.cc
@@ -188,8 +188,14 @@ static Node* ReorderCastAndTranspose(Graph& graph, Node* cast,
return &new_transpose;
}
-// Check whether the element_type is an allowed FusedMatMul data type or not.
-constexpr static bool IsAllowedFusedMatMulDataType(ONNX_NAMESPACE::TensorProto_DataType element_type) {
+// Check whether the element_type is supported by FusedMatMul for the assigned EP.
+static bool IsAllowedFusedMatMulDataType(ONNX_NAMESPACE::TensorProto_DataType element_type,
+ std::string_view execution_provider_type) {
+ if (execution_provider_type == kCpuExecutionProvider) {
+ return element_type == ONNX_NAMESPACE::TensorProto_DataType_FLOAT ||
+ element_type == ONNX_NAMESPACE::TensorProto_DataType_DOUBLE;
+ }
+
return element_type == ONNX_NAMESPACE::TensorProto_DataType_FLOAT ||
element_type == ONNX_NAMESPACE::TensorProto_DataType_FLOAT16 ||
element_type == ONNX_NAMESPACE::TensorProto_DataType_DOUBLE ||
@@ -306,13 +312,15 @@ Status MatmulTransposeFusion::ApplyImpl(Graph& graph, bool& modified, int graph_
NodeArg* left_input = node.MutableInputDefs()[0];
auto left_type = left_input->TypeAsProto()->tensor_type().elem_type();
- if (!IsAllowedFusedMatMulDataType(static_cast(left_type))) {
+ if (!IsAllowedFusedMatMulDataType(static_cast(left_type),
+ node.GetExecutionProviderType())) {
continue;
}
NodeArg* right_input = node.MutableInputDefs()[1];
auto right_type = right_input->TypeAsProto()->tensor_type().elem_type();
- if (!IsAllowedFusedMatMulDataType(static_cast(right_type))) {
+ if (!IsAllowedFusedMatMulDataType(static_cast(right_type),
+ node.GetExecutionProviderType())) {
continue;
}
diff --git a/onnxruntime/core/providers/cpu/cpu_execution_provider.cc b/onnxruntime/core/providers/cpu/cpu_execution_provider.cc
index 86200a92f4ff4..41108424fc3da 100644
--- a/onnxruntime/core/providers/cpu/cpu_execution_provider.cc
+++ b/onnxruntime/core/providers/cpu/cpu_execution_provider.cc
@@ -177,11 +177,13 @@ class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDoma
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 7, 21, Atan);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 7, 8, float, Gemm);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 7, 8, double, Gemm);
+class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 7, 8, MLFloat16, Gemm);
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 1, 10, Hardmax);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 1, 10, float, LogSoftmax);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 1, 10, double, LogSoftmax);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 1, 8, float, MatMul);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 1, 8, double, MatMul);
+class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 1, 8, MLFloat16, MatMul);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 1, 10, float, Softmax);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 1, 10, double, Softmax);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 1, 9, float, TopK);
@@ -397,8 +399,10 @@ class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOn
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, 10, Flatten);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, 10, float, Gemm);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, 10, double, Gemm);
+class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, 10, MLFloat16, Gemm);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, 12, float, MatMul);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, 12, double, MatMul);
+class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, 12, MLFloat16, MatMul);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, 12, int32_t, MatMul);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, 12, int64_t, MatMul);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, 13, float,
@@ -575,6 +579,7 @@ class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDoma
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 12, ScatterND);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 12, float, Gemm);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 12, double, Gemm);
+class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 12, MLFloat16, Gemm);
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 12, GatherElements);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, uint8_t, BitShift);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, uint32_t, BitShift);
@@ -709,8 +714,10 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain,
#endif
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, float, Gemm);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, double, Gemm);
+class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, MLFloat16, Gemm);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, float, MatMul);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, double, MatMul);
+class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, MLFloat16, MatMul);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, int32_t, MatMul);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, int64_t, MatMul);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, Min);
@@ -1748,6 +1755,8 @@ Status RegisterOnnxOperatorKernels(KernelRegistry& kernel_registry) {
float, Gemm)>,
BuildKernelCreateInfo,
+ BuildKernelCreateInfo,
BuildKernelCreateInfo,
BuildKernelCreateInfo,
BuildKernelCreateInfo,
+ BuildKernelCreateInfo,
BuildKernelCreateInfo,
BuildKernelCreateInfo,
BuildKernelCreateInfo,
+ BuildKernelCreateInfo,
BuildKernelCreateInfo,
BuildKernelCreateInfo,
+ BuildKernelCreateInfo,
BuildKernelCreateInfo,
BuildKernelCreateInfo,
BuildKernelCreateInfo,
+ BuildKernelCreateInfo,
BuildKernelCreateInfo,
BuildKernelCreateInfo,
BuildKernelCreateInfo,
+ BuildKernelCreateInfo,
BuildKernelCreateInfo,
BuildKernelCreateInfo,
BuildKernelCreateInfo,
BuildKernelCreateInfo,
+ BuildKernelCreateInfo,
BuildKernelCreateInfo,
BuildKernelCreateInfo,
BuildKernelCreateInfo,
diff --git a/onnxruntime/core/providers/cpu/fp16/fp16_conv.cc b/onnxruntime/core/providers/cpu/fp16/fp16_conv.cc
index 8aee133e39c45..ac3e691d59835 100644
--- a/onnxruntime/core/providers/cpu/fp16/fp16_conv.cc
+++ b/onnxruntime/core/providers/cpu/fp16/fp16_conv.cc
@@ -9,6 +9,9 @@
#ifdef MLAS_F16VEC_INTRINSICS_SUPPORTED
+#include
+
+#include "core/common/narrow.h"
#include "core/common/safeint.h"
#include "core/common/float16.h"
#include "core/framework/op_kernel.h"
@@ -48,6 +51,12 @@ class FusedConvFp16 final : public OpKernel {
ORT_ENFORCE(GetFusedActivationAttr(info, activation_).IsOK());
channels_last_ = (info.GetKernelDef().OpName() == "NhwcFusedConv");
SetupMlasBackendKernelSelectorFromConfigOptions(mlas_backend_kernel_selector_config_, info.GetConfigOptions());
+ const auto& input_defs = info.node().InputDefs();
+ has_bias_input_ = input_defs.size() >= 3 && input_defs[2]->Exists();
+ const Tensor* bias = nullptr;
+ if (has_bias_input_ && info.TryGetConstantInput(2, &bias)) {
+ constant_B_ = bias;
+ }
}
Status Compute(OpKernelContext* context) const override;
@@ -97,7 +106,11 @@ class FusedConvFp16 final : public OpKernel {
MLAS_ACTIVATION activation_;
MLAS_BACKEND_KERNEL_SELECTOR_CONFIG mlas_backend_kernel_selector_config_;
ConvAttributes conv_attrs_;
+ bool has_bias_input_{false};
+ const Tensor* constant_B_{nullptr};
TensorShape W_shape_;
+ BufferUniquePtr packed_halfconv_weights_and_bias_buffer_;
+ size_t packed_halfconv_weights_and_bias_size_{0};
BufferUniquePtr packed_W_buffer_;
size_t packed_W_size_{0};
bool is_W_packed_{false};
@@ -142,8 +155,53 @@ Status FusedConvFp16::PrePack(const Tensor& tensor, int input_idx, AllocatorPtr
const size_t kernel_dim = group_input_channels * kernel_size;
bool share_prepacked_weights = (prepacked_weights != nullptr);
+ const bool has_valid_constant_halfconv_bias =
+ constant_B_ != nullptr &&
+ !share_prepacked_weights &&
+ constant_B_->Shape().NumDimensions() == 1 &&
+ constant_B_->Shape()[0] == static_cast(output_channels);
const bool is_depthwise_conv = (group_input_channels == 1 && group_output_channels == 1);
+
+ if (group_count == 1 &&
+ rank == 4 &&
+ output_channels > 1 &&
+ activation_.ActivationKind == MlasIdentityActivation &&
+ (!has_bias_input_ || has_valid_constant_halfconv_bias)) {
+ std::array halfconv_kernel_shape{shape[2], shape[3]};
+ std::array halfconv_dilations{1, 1};
+ if (!conv_attrs_.dilations.empty()) {
+ halfconv_dilations = {conv_attrs_.dilations[0], conv_attrs_.dilations[1]};
+ }
+
+ packed_halfconv_weights_and_bias_size_ = MlasHalfConvPackWeightsAndBiasSize(
+ output_channels,
+ group_input_channels,
+ halfconv_kernel_shape.data(),
+ halfconv_dilations.data(),
+ &mlas_backend_kernel_selector_config_);
+ if (packed_halfconv_weights_and_bias_size_ != 0) {
+ auto* packed_halfconv_weights_and_bias = alloc->Alloc(packed_halfconv_weights_and_bias_size_);
+ BufferUniquePtr packed_halfconv_weights_and_bias_buffer(
+ packed_halfconv_weights_and_bias, BufferDeleter(alloc));
+ const auto* bias_data = constant_B_ != nullptr ? constant_B_->Data() : nullptr;
+ if (MlasHalfConvPackWeightsAndBias(
+ output_channels,
+ group_input_channels,
+ halfconv_kernel_shape.data(),
+ halfconv_dilations.data(),
+ Wdata,
+ bias_data,
+ packed_halfconv_weights_and_bias,
+ nullptr,
+ &mlas_backend_kernel_selector_config_)) {
+ packed_halfconv_weights_and_bias_buffer_ = std::move(packed_halfconv_weights_and_bias_buffer);
+ } else {
+ packed_halfconv_weights_and_bias_size_ = 0;
+ }
+ }
+ }
+
// Don't pack the filter buffer if the MlasConvDepthwise path is used.
if (!is_depthwise_conv) {
packed_W_size_ = MlasHalfGemmPackBSize(group_output_channels, kernel_dim, false);
@@ -188,6 +246,9 @@ Status FusedConvFp16::PrePack(const Tensor& tensor, int input_idx, AllocatorPtr
}
if (share_prepacked_weights) {
+ // `packed_halfconv_weights_and_bias_buffer_` is session-local only. It encodes a bias
+ // choice even when the bias is null, so it must not be shared under the
+ // W-only prepack cache key.
prepacked_weights->buffers_.push_back(nullptr); // packed_W_buffer_ is nullptr
prepacked_weights->buffer_sizes_.push_back(0);
}
@@ -225,12 +286,13 @@ Status FusedConvFp16::UseSharedPrePackedBuffers(std::vector& pr
used_shared_buffers = true;
- if (prepacked_buffers.size() == 1) { // This means that only packed_W_ exists
+ if (prepacked_buffers.size() == 1) { // only packed_W_ exists
packed_W_buffer_ = std::move(prepacked_buffers[0]);
- } else if (prepacked_buffers.size() == 2) { // This means that only reordered_W_ exists
- // Enforce that the first "placeholder" buffer is nullptr
+ } else if (prepacked_buffers.size() == 2) { // placeholder + reordered_W_
ORT_ENFORCE(prepacked_buffers[0].get() == nullptr);
reordered_W_buffer_ = std::move(prepacked_buffers[1]);
+ } else {
+ ORT_ENFORCE(false, "Unexpected number of shared prepacked fp16 conv buffers.");
}
return Status::OK();
@@ -299,6 +361,63 @@ Status FusedConvFp16::Compute(OpKernelContext* context) const {
AllocatorPtr alloc;
ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&alloc));
+ concurrency::ThreadPool* thread_pool = context->GetOperatorThreadPool();
+
+ const auto* Bdata = B != nullptr ? B->Data() : nullptr;
+
+ if (Sum == nullptr && kernel_rank == 2) {
+ MLAS_CONV_PARAMETERS parameters{};
+
+ size_t working_buffer_size_in_bytes = 0;
+ if (MlasHalfConvPrepare(¶meters,
+ kernel_rank,
+ narrow(N),
+ narrow(conv_attrs_.group),
+ narrow(C / conv_attrs_.group),
+ input_shape.GetDims().data(),
+ kernel_shape.data(),
+ dilations.data(),
+ pads.data(),
+ strides.data(),
+ output_shape.GetDims().data(),
+ narrow(M / conv_attrs_.group),
+ &activation_,
+ &working_buffer_size_in_bytes,
+ 0.0f,
+ channels_last_,
+ thread_pool,
+ &mlas_backend_kernel_selector_config_)) {
+ const MLFloat16* halfconv_filter = nullptr;
+ const MLFloat16* halfconv_bias = Bdata;
+ bool halfconv_filter_and_bias_are_packed = false;
+
+ if (packed_halfconv_weights_and_bias_buffer_ != nullptr) {
+ halfconv_filter = static_cast(packed_halfconv_weights_and_bias_buffer_.get());
+ halfconv_filter_and_bias_are_packed = true;
+ halfconv_bias = nullptr;
+ } else if (W != nullptr) {
+ halfconv_filter = W->Data();
+ }
+
+ if (halfconv_filter != nullptr) {
+ auto* working_data = working_buffer_size_in_bytes > 0
+ ? alloc->Alloc(working_buffer_size_in_bytes)
+ : nullptr;
+ BufferUniquePtr working_buffer(working_data, BufferDeleter(alloc));
+
+ if (MlasHalfConv(¶meters,
+ X->Data(),
+ halfconv_filter,
+ halfconv_filter_and_bias_are_packed,
+ halfconv_bias,
+ static_cast(working_buffer.get()),
+ Y->MutableData(),
+ thread_pool)) {
+ return Status::OK();
+ }
+ }
+ }
+ }
// Handle the case of a dynamic weight filter.
BufferUniquePtr reordered_W_buffer;
@@ -340,7 +459,6 @@ Status FusedConvFp16::Compute(OpKernelContext* context) const {
const int64_t col_buffer_size = kernel_dim * output_image_size;
const auto* Xdata = X->Data();
- const auto* Bdata = B != nullptr ? B->Data() : nullptr;
auto* Ydata = Y->MutableData();
const auto* sum_data = Sum != nullptr ? Sum->Data() : nullptr;
@@ -390,8 +508,6 @@ Status FusedConvFp16::Compute(OpKernelContext* context) const {
padding_data.resize(static_cast(C), MLFloat16());
}
- concurrency::ThreadPool* thread_pool = context->GetOperatorThreadPool();
-
/*************************************
* Thread partition idea: we are essentially partition a GEMM A[M,K] x B[K,N].
* Here B contains the conv filters, which are usually not big, so we assume
@@ -547,7 +663,7 @@ Status FusedConvFp16::Compute(OpKernelContext* context) const {
const auto* gemm_add = add_src == nullptr ? nullptr : worker_addsrc + group_id * group_output_channels;
MLAS_HALF_GEMM_ACTIVATION_PROCESSOR act(activation_, gemm_add);
- MLAS_HALF_GEMM_DATA_PARAMS gemm_params;
+ MLAS_HALF_GEMM_DATA_PARAMS gemm_params{};
gemm_params.A = AData;
gemm_params.lda = lda;
if (packed_W_buffer_) {
diff --git a/onnxruntime/core/providers/cpu/math/gemm.cc b/onnxruntime/core/providers/cpu/math/gemm.cc
index e12761ac61631..b3092a6b54988 100644
--- a/onnxruntime/core/providers/cpu/math/gemm.cc
+++ b/onnxruntime/core/providers/cpu/math/gemm.cc
@@ -205,31 +205,42 @@ void Gemm_MLFloat16(CBLAS_TRANSPOSE trans_a, CBLAS_TRANSPOSE trans_b,
if (c_data == nullptr)
beta = onnxruntime::MLFloat16::Zero;
-#ifdef MLAS_F16VEC_INTRINSICS_SUPPORTED
- bool support_mlas = false;
+
+ // Guard against using generic half GEMM when no accelerated implementation is
+ // available. Native packing support currently also signals an accelerated
+ // backend path.
+ const bool has_accelerated_half_gemm =
+ MlasFp16AccelerationSupported() ||
+ MlasHalfGemmNativePackBSize(CblasNoTrans, CblasNoTrans,
+ static_cast(N), static_cast(K),
+ mlas_backend_kernel_selector_config) != 0;
+ bool support_mlas_bias = false;
if (c_shape == nullptr) {
- support_mlas = true;
+ support_mlas_bias = true;
} else if (c_shape->NumDimensions() == 1 && (*c_shape)[0] == N) {
- support_mlas = true;
- } else if (c_shape->NumDimensions() == 2 && (((*c_shape)[0] == 1 && (*c_shape)[1] == N) || ((*c_shape)[0] == N && (*c_shape)[1] == 1))) {
- support_mlas = true;
+ support_mlas_bias = true;
+ } else if (c_shape->NumDimensions() == 2 &&
+ (((*c_shape)[0] == 1 && (*c_shape)[1] == N) || ((*c_shape)[0] == N && (*c_shape)[1] == 1))) {
+ support_mlas_bias = true;
}
- if (trans_a == CblasNoTrans && trans_b == CblasNoTrans && support_mlas && alpha.ToFloat() == 1.0 && beta.ToFloat() == 1.0) {
- MLAS_HALF_GEMM_DATA_PARAMS data;
+ const bool use_mlas_no_bias = beta == onnxruntime::MLFloat16::Zero;
+ const bool use_mlas_bias = beta == onnxruntime::MLFloat16::One && support_mlas_bias;
+ if (has_accelerated_half_gemm && trans_a == CblasNoTrans && trans_b == CblasNoTrans &&
+ alpha == onnxruntime::MLFloat16::One && (use_mlas_no_bias || use_mlas_bias)) {
+ MLAS_HALF_GEMM_DATA_PARAMS data{};
data.A = a_data;
data.lda = K;
data.B = b_data;
data.ldb = N;
data.C = y_data;
data.ldc = N;
- if (c_shape != nullptr) {
+ if (use_mlas_bias && c_shape != nullptr) {
data.Bias = c_data;
}
data.BackendKernelSelectorConfig = mlas_backend_kernel_selector_config;
MlasHalfGemmBatch(M, N, K, 1, &data, thread_pool);
return;
}
-#endif
// Fallback to Eigen
// Broadcast the bias as needed if bias is given
GemmBroadcastBias(M, N, beta, c_data, c_shape, y_data);
diff --git a/onnxruntime/core/providers/cpu/math/matmul.cc b/onnxruntime/core/providers/cpu/math/matmul.cc
index 3117001eae5b1..250117e784ff2 100644
--- a/onnxruntime/core/providers/cpu/math/matmul.cc
+++ b/onnxruntime/core/providers/cpu/math/matmul.cc
@@ -7,6 +7,9 @@
#include "core/providers/cpu/math/matmul_helper.h"
#include "core/util/math.h"
#include "core/util/math_cpuonly.h"
+#include
+#include
+#include
namespace onnxruntime {
@@ -24,6 +27,13 @@ ONNX_CPU_OPERATOR_VERSIONED_TYPED_KERNEL(
KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()),
MatMul);
+ONNX_CPU_OPERATOR_VERSIONED_TYPED_KERNEL(
+ MatMul,
+ 1, 8,
+ MLFloat16,
+ KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()),
+ MatMul);
+
// opset 9 supports more types
ONNX_CPU_OPERATOR_VERSIONED_TYPED_KERNEL(
MatMul,
@@ -41,6 +51,14 @@ ONNX_CPU_OPERATOR_VERSIONED_TYPED_KERNEL(
KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()),
MatMul);
+ONNX_CPU_OPERATOR_VERSIONED_TYPED_KERNEL(
+ MatMul,
+ 9,
+ 12,
+ MLFloat16,
+ KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()),
+ MatMul);
+
ONNX_CPU_OPERATOR_VERSIONED_TYPED_KERNEL(
MatMul,
9,
@@ -73,6 +91,13 @@ ONNX_CPU_OPERATOR_TYPED_KERNEL(
KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()),
MatMul);
+ONNX_CPU_OPERATOR_TYPED_KERNEL(
+ MatMul,
+ 13,
+ MLFloat16,
+ KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()),
+ MatMul);
+
ONNX_CPU_OPERATOR_TYPED_KERNEL(
MatMul,
13,
@@ -107,9 +132,8 @@ Status MatMul::Compute(OpKernelContext* ctx) const {
if (helper.K() == 0) {
// When we have (M, 0, N) then the inputs are empty, but the output should
// be filled out with zeros.
- EigenMatrixMapRowMajor dest(y->MutableData(),
- narrow(helper.M()), narrow(helper.N()));
- dest.setZero();
+ auto output_span = gsl::make_span(y->MutableData(), narrow(y->Shape().Size()));
+ std::fill(output_span.begin(), output_span.end(), T{});
return Status::OK();
}
@@ -236,6 +260,45 @@ bool GemmPackBBfloat16(AllocatorPtr& alloc,
}
#endif
+bool GemmPackBHalfNative(AllocatorPtr& alloc,
+ const Tensor& tensor_b,
+ IAllocatorUniquePtr& packed_b,
+ size_t& packed_b_size,
+ TensorShape& b_shape,
+ const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* mlas_backend_kernel_selector_config) {
+ // Only handle the common case of a 2D weight matrix. Additional matrices
+ // could be handled by stacking the packed buffers.
+ if (tensor_b.Shape().NumDimensions() != 2) {
+ return false;
+ }
+
+ b_shape = tensor_b.Shape();
+
+ const size_t K = static_cast(b_shape[0]);
+ const size_t N = static_cast(b_shape[1]);
+
+ packed_b_size = MlasHalfGemmNativePackBSize(CblasNoTrans, CblasNoTrans, N, K,
+ mlas_backend_kernel_selector_config);
+ if (packed_b_size == 0) {
+ return false;
+ }
+
+ packed_b = IAllocator::MakeUniquePtr(alloc, packed_b_size, true);
+ auto* packed_b_data = packed_b.get();
+ memset(packed_b_data, 0, packed_b_size);
+
+ if (!MlasHalfGemmNativePackB(CblasNoTrans, CblasNoTrans, N, K,
+ reinterpret_cast(tensor_b.Data()), N, packed_b_data,
+ mlas_backend_kernel_selector_config)) {
+ packed_b.reset();
+ packed_b_size = 0;
+ b_shape = TensorShape();
+ return false;
+ }
+
+ return true;
+}
+
Status MatMul::PrePack(const Tensor& tensor, int input_idx, /*out*/ AllocatorPtr alloc,
/*out*/ bool& is_packed,
/*out*/ PrePackedWeights* prepacked_weights) {
@@ -278,6 +341,124 @@ Status MatMul::UseSharedPrePackedBuffers(std::vector& pr
return Status::OK();
}
+Status MatMul::PrePack(const Tensor& tensor, int input_idx, /*out*/ AllocatorPtr alloc,
+ /*out*/ bool& is_packed,
+ /*out*/ PrePackedWeights* prepacked_weights) {
+ is_packed = false;
+ if (input_idx == 1) {
+ size_t packed_b_size = 0;
+ is_packed = GemmPackBHalfNative(alloc, tensor, packed_b_, packed_b_size, b_shape_,
+ &mlas_backend_kernel_selector_config_);
+ // The native fp16 packed-B layout depends on the active MLAS backend selector.
+ // Keep it owned by this kernel until shared prepacked weights carry layout metadata.
+ if (is_packed && prepacked_weights != nullptr) {
+ prepacked_weights->has_kernel_owned_packed_weights_ = true;
+ }
+ }
+
+ return Status::OK();
+}
+
+Status MatMul::UseSharedPrePackedBuffers(std::vector& prepacked_buffers,
+ gsl::span prepacked_buffer_sizes,
+ int input_idx,
+ /*out*/ bool& used_shared_buffers) {
+ ORT_UNUSED_PARAMETER(prepacked_buffers);
+ ORT_UNUSED_PARAMETER(prepacked_buffer_sizes);
+ ORT_UNUSED_PARAMETER(input_idx);
+
+ // Native fp16 packed-B buffers are backend-layout-specific. Decline shared
+ // buffers until the shared prepack cache can validate that layout.
+ used_shared_buffers = false;
+ return Status::OK();
+}
+
+Status MatMul::Compute(OpKernelContext* ctx) const {
+ concurrency::ThreadPool* thread_pool = ctx->GetOperatorThreadPool();
+
+ const Tensor* a = ctx->Input(0);
+ const Tensor* b = packed_b_ ? nullptr : ctx->Input(1);
+ const auto& b_shape = b ? b->Shape() : b_shape_;
+
+ MatMulComputeHelper helper;
+ ORT_RETURN_IF_ERROR(helper.Compute(a->Shape(), b_shape));
+ Tensor* y = ctx->Output(0, helper.OutputShape());
+
+ if (y->Shape().Size() == 0) {
+ return Status::OK();
+ }
+
+ if (helper.K() == 0) {
+ auto output_span = gsl::make_span(y->MutableData(), narrow(y->Shape().Size()));
+ std::fill(output_span.begin(), output_span.end(), MLFloat16{});
+ return Status::OK();
+ }
+
+ const auto* a_data = a->Data();
+ const auto* b_data = b ? b->Data() : nullptr;
+ auto* y_data = y->MutableData();
+
+ const size_t max_len = helper.OutputOffsets().size();
+ const size_t M = static_cast(helper.M());
+ const size_t N = static_cast(helper.N());
+ const size_t K = static_cast(helper.K());
+ const size_t lda = helper.Lda(false);
+ const size_t ldb = helper.Ldb(false);
+
+ // Guard against using generic half GEMM when no accelerated implementation is
+ // available. Native packing support currently also signals an accelerated
+ // backend path.
+ const bool has_accelerated_half_gemm =
+ MlasFp16AccelerationSupported() ||
+ MlasHalfGemmNativePackBSize(CblasNoTrans, CblasNoTrans, N, K,
+ &mlas_backend_kernel_selector_config_) != 0;
+ if (!has_accelerated_half_gemm && packed_b_ == nullptr) {
+ for (size_t i = 0; i < max_len; ++i) {
+ math::MatMul(narrow(M), narrow(N), narrow(K),
+ a_data + helper.LeftOffsets()[i],
+ b_data + helper.RightOffsets()[i],
+ y_data + helper.OutputOffsets()[i],
+ thread_pool,
+ &mlas_backend_kernel_selector_config_);
+ }
+ return Status::OK();
+ }
+
+ if (M <= 2 && packed_b_ == nullptr && MlasHGemmSupported(CblasNoTrans, CblasNoTrans)) {
+ const auto alpha = MLFloat16(1.0f);
+ const auto beta = MLFloat16(0.0f);
+ InlinedVector data(max_len);
+ for (size_t i = 0; i < max_len; i++) {
+ data[i].A = a_data + helper.LeftOffsets()[i];
+ data[i].lda = lda;
+ data[i].B = b_data + helper.RightOffsets()[i];
+ data[i].ldb = ldb;
+ data[i].C = y_data + helper.OutputOffsets()[i];
+ data[i].ldc = N;
+ data[i].alpha = alpha.val;
+ data[i].beta = beta.val;
+ }
+
+ MlasGemmBatch(CblasNoTrans, CblasNoTrans, M, N, K, data.data(), max_len, thread_pool);
+ return Status::OK();
+ }
+
+ InlinedVector data(max_len);
+ for (size_t i = 0; i < max_len; i++) {
+ data[i].A = a_data + helper.LeftOffsets()[i];
+ data[i].lda = lda;
+ data[i].B = packed_b_ ? packed_b_.get() : static_cast(b_data + helper.RightOffsets()[i]);
+ data[i].ldb = packed_b_ ? 0 : ldb;
+ data[i].C = y_data + helper.OutputOffsets()[i];
+ data[i].ldc = N;
+ data[i].BIsBackendNativePacked = static_cast(packed_b_);
+ data[i].BackendKernelSelectorConfig = &mlas_backend_kernel_selector_config_;
+ }
+
+ MlasHalfGemmBatch(M, N, K, max_len, data.data(), thread_pool);
+ return Status::OK();
+}
+
Status MatMul::Compute(OpKernelContext* ctx) const {
concurrency::ThreadPool* thread_pool = ctx->GetOperatorThreadPool();
diff --git a/onnxruntime/core/providers/cpu/math/matmul.h b/onnxruntime/core/providers/cpu/math/matmul.h
index 7f7c06e8eedb0..a14c7719d57d0 100644
--- a/onnxruntime/core/providers/cpu/math/matmul.h
+++ b/onnxruntime/core/providers/cpu/math/matmul.h
@@ -118,4 +118,28 @@ class MatMul final : public OpKernel {
#endif
};
+template <>
+class MatMul final : public OpKernel {
+ public:
+ MatMul(const OpKernelInfo& info) : OpKernel(info) {
+ SetupMlasBackendKernelSelectorFromConfigOptions(mlas_backend_kernel_selector_config_, info.GetConfigOptions());
+ }
+
+ Status PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc,
+ /*out*/ bool& is_packed,
+ /*out*/ PrePackedWeights* prepacked_weights) override;
+
+ Status UseSharedPrePackedBuffers(std::vector& prepacked_buffers,
+ gsl::span /*prepacked_buffer_sizes*/,
+ int input_idx,
+ /*out*/ bool& used_shared_buffers) override;
+
+ Status Compute(OpKernelContext* context) const override;
+
+ private:
+ TensorShape b_shape_;
+ IAllocatorUniquePtr packed_b_;
+ MLAS_BACKEND_KERNEL_SELECTOR_CONFIG mlas_backend_kernel_selector_config_;
+};
+
} // namespace onnxruntime
diff --git a/onnxruntime/core/util/math_cpu.cc b/onnxruntime/core/util/math_cpu.cc
index 608d30c12b587..e1e994cd5b983 100644
--- a/onnxruntime/core/util/math_cpu.cc
+++ b/onnxruntime/core/util/math_cpu.cc
@@ -190,6 +190,45 @@ void MatMul(ptrdiff_t M, ptrdiff_t N, ptrdiff_t K, const float* A, const
MlasGemm(CblasNoTrans, CblasNoTrans, M, N, K, 1.f, A, K, B, N, 0.f, C, N, threadpool, mlas_backend_kernel_selector_config);
}
+template <>
+void MatMul(ptrdiff_t M, ptrdiff_t N, ptrdiff_t K, const MLFloat16* A, const MLFloat16* B, MLFloat16* C, ThreadPool* threadpool,
+ const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* mlas_backend_kernel_selector_config) {
+ // Guard against using generic half GEMM when no accelerated implementation is
+ // available. Native packing support currently also signals an accelerated
+ // backend path.
+ const bool has_accelerated_half_gemm =
+ MlasFp16AccelerationSupported() ||
+ MlasHalfGemmNativePackBSize(CblasNoTrans, CblasNoTrans,
+ static_cast(N), static_cast(K),
+ mlas_backend_kernel_selector_config) != 0;
+ if (has_accelerated_half_gemm) {
+ MLAS_HALF_GEMM_DATA_PARAMS data{};
+ data.A = A;
+ data.lda = static_cast(K);
+ data.B = B;
+ data.ldb = static_cast(N);
+ data.C = C;
+ data.ldc = static_cast(N);
+ data.BackendKernelSelectorConfig = mlas_backend_kernel_selector_config;
+ MlasHalfGemmBatch(static_cast(M), static_cast(N), static_cast(K), 1, &data, threadpool);
+ return;
+ }
+
+#if defined(__GNUC__)
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wstrict-aliasing"
+#endif
+ auto C_mat = EigenMatrixMap(reinterpret_cast(C), N, M);
+ // Accumulate the fallback in fp32 and round only the result to fp16.
+ C_mat.noalias() =
+ (ConstEigenMatrixMap(reinterpret_cast(B), N, K).cast() *
+ ConstEigenMatrixMap(reinterpret_cast(A), K, M).cast())
+ .cast();
+#if defined(__GNUC__)
+#pragma GCC diagnostic pop
+#endif
+}
+
#ifdef MLAS_SUPPORTS_GEMM_DOUBLE
template <>
void MatMul(ptrdiff_t M, ptrdiff_t N, ptrdiff_t K, const double* A, const double* B, double* C, ThreadPool* threadpool,
diff --git a/onnxruntime/test/framework/session_state_test.cc b/onnxruntime/test/framework/session_state_test.cc
index 35d698cfef726..19c237a1a7a2b 100644
--- a/onnxruntime/test/framework/session_state_test.cc
+++ b/onnxruntime/test/framework/session_state_test.cc
@@ -700,7 +700,28 @@ class PrePackingTestOpKernel : public OpKernel {
IAllocatorUniquePtr weight_packed_;
};
-static void CreateSimpleGraph(Graph& graph) {
+class BrokenPrePackingTestOpKernel : public OpKernel {
+ public:
+ BrokenPrePackingTestOpKernel(const OpKernelInfo& info) : OpKernel(info) {}
+
+ Status Compute(OpKernelContext* context) const override {
+ ORT_UNUSED_PARAMETER(context);
+ return Status::OK();
+ }
+
+ Status PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc,
+ /*out*/ bool& is_packed, /*out*/ PrePackedWeights* prepacked_weights) override {
+ ORT_UNUSED_PARAMETER(tensor);
+ ORT_UNUSED_PARAMETER(input_idx);
+ ORT_UNUSED_PARAMETER(alloc);
+ ORT_UNUSED_PARAMETER(prepacked_weights);
+
+ is_packed = true;
+ return Status::OK();
+ }
+};
+
+static void CreateSimpleGraph(Graph& graph, const std::string& op_type = "PrePackingTest") {
// node creation and placement
TypeProto type;
type.mutable_tensor_type()->set_elem_type(TensorProto_DataType_FLOAT);
@@ -716,7 +737,7 @@ static void CreateSimpleGraph(Graph& graph) {
onnxruntime::NodeArg output_arg("node_0_output_0", &type);
outputs.push_back(&output_arg);
- graph.AddNode("node_0", "PrePackingTest", "node 0", inputs, outputs);
+ graph.AddNode("node_0", op_type, "node 0", inputs, outputs);
// add an initializer
ONNX_NAMESPACE::TensorProto tensor;
@@ -845,6 +866,11 @@ void RegisterPrePackingTestSchemaOnce() {
.Input(0, "Input_0", "input 0", "tensor(float)")
.Input(1, "Input_1", "input 1", "tensor(float)")
.Output(0, "output_0", "docstr for output_0.", "tensor(float)");
+ ONNX_OPERATOR_SCHEMA(BrokenPrePackingTest)
+ .SetDoc("Faking broken Node for PrePacking")
+ .Input(0, "Input_0", "input 0", "tensor(float)")
+ .Input(1, "Input_1", "input 1", "tensor(float)")
+ .Output(0, "output_0", "docstr for output_0.", "tensor(float)");
});
}
} // namespace
@@ -956,6 +982,19 @@ class SessionStateTestSharedInitalizersWithPrePacking : public ::testing::Test {
KernelCreateInfo(std::move(kernel_def),
[](FuncManager&, const OpKernelInfo& info, std::unique_ptr& out) -> Status { out = std::make_unique(info); return Status::OK(); })));
+ auto broken_kernel_def = KernelDefBuilder()
+ .SetName("BrokenPrePackingTest")
+ .Provider(kCpuExecutionProvider)
+ .SinceVersion(1)
+ .Build();
+
+ ASSERT_STATUS_OK(kernel_registry->Register(
+ KernelCreateInfo(std::move(broken_kernel_def),
+ [](FuncManager&, const OpKernelInfo& info, std::unique_ptr& out) -> Status {
+ out = std::make_unique(info);
+ return Status::OK();
+ })));
+
kernel_registry_manager.RegisterKernelRegistry(kernel_registry);
}
};
@@ -1188,6 +1227,46 @@ TEST_F(SessionStateTestSharedInitalizersWithPrePacking, test3) {
ASSERT_EQ(session_state_2.GetUsedSharedPrePackedWeightCounter(), static_cast(1));
}
+TEST_F(SessionStateTestSharedInitalizersWithPrePacking, BrokenKernelWithoutCacheableBuffersFails) {
+ SessionOptions sess_options;
+ sess_options.enable_mem_pattern = true;
+ sess_options.execution_mode = ExecutionMode::ORT_SEQUENTIAL;
+ sess_options.use_deterministic_compute = false;
+ sess_options.enable_mem_reuse = true;
+ sess_options.config_options.configurations[kOrtSessionOptionsConfigDisablePrepacking] = "0";
+
+ OrtMemoryInfo mem_info(CPU, OrtDeviceAllocator);
+ std::vector float_data(1, 1);
+ auto value = std::make_unique();
+ Tensor::InitOrtValue(DataTypeImpl::GetType(), TensorShape(std::vector{1}),
+ reinterpret_cast(float_data.data()), mem_info, *value);
+
+ ASSERT_STATUS_OK(sess_options.AddInitializer("node_0_input_1", value.get()));
+
+ PrepackedWeightsContainer prepacked_weights_container;
+
+ Model model("graph_main", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(),
+ domain_to_version, std::vector(),
+ DefaultLoggingManager().DefaultLogger());
+
+ CreateSimpleGraph(model.MainGraph(), "BrokenPrePackingTest");
+ PlaceAllNodesToCPUEP(model.MainGraph());
+ SessionState session_state(model.MainGraph(),
+ execution_providers,
+ tp.get(),
+ nullptr, /*inter_op_thread_pool*/
+ dtm,
+ edlm,
+ DefaultLoggingManager().DefaultLogger(),
+ profiler,
+ sess_options,
+ &prepacked_weights_container);
+
+ ASSERT_STATUS_NOT_OK_AND_HAS_SUBSTR(
+ session_state.FinalizeSessionState(std::basic_string(), kernel_registry_manager),
+ "doesn't have an implementation that can cache computed pre-packed weights");
+}
+
// Pre-packing enabled + shared initializers +
// pre-packed weights container + subgraphs =
// caching enabled in pre-packed weights used in subgraphs
diff --git a/onnxruntime/test/mlas/unittest/test_halfgemm.cpp b/onnxruntime/test/mlas/unittest/test_halfgemm.cpp
index 149a1ddda0848..20f7defc7048d 100644
--- a/onnxruntime/test/mlas/unittest/test_halfgemm.cpp
+++ b/onnxruntime/test/mlas/unittest/test_halfgemm.cpp
@@ -54,8 +54,7 @@ TestHalfGemmBatchOverride(
size_t,
size_t BatchN,
const MLAS_HALF_GEMM_DATA_PARAMS* DataParams,
- MLAS_THREADPOOL*,
- const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG*) {
+ MLAS_THREADPOOL*) {
g_test_halfgemm_override_called = true;
for (size_t batch = 0; batch < BatchN; ++batch) {
@@ -799,7 +798,7 @@ TEST(HalfGemmKleidiAIPath, KleidiAIPackedBWithBiasIsRejected) {
data.ldc = N;
data.BIsBackendNativePacked = true;
- ASSERT_FALSE(ArmKleidiAI::MlasHalfGemmBatch(M, N, K, 1, &data, nullptr, nullptr));
+ ASSERT_FALSE(ArmKleidiAI::MlasHalfGemmBatch(M, N, K, 1, &data, nullptr));
#if !defined(ORT_NO_EXCEPTIONS)
EXPECT_THROW(MlasHalfGemmBatch(M, N, K, 1, &data, nullptr), std::runtime_error);
#endif
@@ -846,7 +845,7 @@ TEST(HalfGemmKleidiAIPath, KleidiAIPackedBWithOutputProcessorIsRejected) {
data.BIsBackendNativePacked = true;
data.OutputProcessor = &output_processor;
- ASSERT_FALSE(ArmKleidiAI::MlasHalfGemmBatch(M, N, K, 1, &data, nullptr, nullptr));
+ ASSERT_FALSE(ArmKleidiAI::MlasHalfGemmBatch(M, N, K, 1, &data, nullptr));
#if !defined(ORT_NO_EXCEPTIONS)
EXPECT_THROW(MlasHalfGemmBatch(M, N, K, 1, &data, nullptr), std::runtime_error);
#endif
@@ -871,7 +870,7 @@ TEST(HalfGemmKleidiAIPath, ZeroKIsNotHandledByKleidiAIOverride) {
data.C = reinterpret_cast(C.data());
data.ldc = N;
- const bool handled = ArmKleidiAI::MlasHalfGemmBatch(M, N, K, 1, &data, nullptr, nullptr);
+ const bool handled = ArmKleidiAI::MlasHalfGemmBatch(M, N, K, 1, &data, nullptr);
ASSERT_FALSE(handled);
}
diff --git a/onnxruntime/test/mlas/unittest/test_halfgemm.h b/onnxruntime/test/mlas/unittest/test_halfgemm.h
index 0ce058dc3c1eb..dbc85be30d2bd 100644
--- a/onnxruntime/test/mlas/unittest/test_halfgemm.h
+++ b/onnxruntime/test/mlas/unittest/test_halfgemm.h
@@ -122,7 +122,7 @@ class MlasHalfGemmTest : public MlasTestBase {
if (enforce_kleidiai_override) {
ASSERT_NE(GetMlasPlatform().MlasHalfGemmBatchOverride, nullptr);
const bool handled = GetMlasPlatform().MlasHalfGemmBatchOverride(
- M, N, K, BatchSize, GemmParameters.data(), threadpool_, nullptr);
+ M, N, K, BatchSize, GemmParameters.data(), threadpool_);
ASSERT_TRUE(handled);
} else {
MlasHalfGemmBatch(M, N, K, BatchSize, GemmParameters.data(), threadpool_);
diff --git a/onnxruntime/test/optimizer/compute_optimizer_test.cc b/onnxruntime/test/optimizer/compute_optimizer_test.cc
index 08c7a0700030f..4e2e9c5e74b6b 100644
--- a/onnxruntime/test/optimizer/compute_optimizer_test.cc
+++ b/onnxruntime/test/optimizer/compute_optimizer_test.cc
@@ -1773,9 +1773,8 @@ TEST(ComputeOptimizerTests, GatherRobertaE2E) {
ASSERT_TRUE(expected_ort_values.size() == actual_ort_values.size());
- // "expected 0.793675 (3f4b2e44), got 0.79232 (3f4ad584), diff: 0.00135422, tol=0.000179367 idx=4276.
- // 1713 of 8192 differ"
- // Loose the atol a bit because we see the MatMuls results differ once we move Gather before it.
+ // Moving Gather before MatMul changes the MatMul shapes and may alter the
+ // floating-point accumulation order, so allow a small numerical difference.
constexpr double per_sample_tolerance = 2e-3;
constexpr double relative_per_sample_tolerance = 2e-3;
for (size_t i = 0; i < expected_ort_values.size(); i++) {
diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc
index 32b7f0334c4e8..cb7f5ce1f88ff 100644
--- a/onnxruntime/test/optimizer/graph_transform_test.cc
+++ b/onnxruntime/test/optimizer/graph_transform_test.cc
@@ -3878,6 +3878,32 @@ TEST_F(GraphTransformationTests, TransposeMatmulFusion) {
ASSERT_TRUE(op_to_count["com.microsoft.FusedMatMul"] == 1);
}
+TEST_F(GraphTransformationTests, TransposeMatmulNoFusionForCpuFp16) {
+ auto build_test_case = [](ModelTestBuilder& builder) {
+ auto* q = builder.MakeInput({{2, 4, 8, 16}});
+ auto* k = builder.MakeInput({{2, 4, 6, 16}});
+ auto* k_transposed = builder.MakeIntermediate({{2, 4, 16, 6}});
+ auto* output = builder.MakeOutput({{2, 4, 8, 6}});
+
+ builder.AddNode("Transpose", {k}, {k_transposed})
+ .AddAttribute("perm", std::vector{0, 1, 3, 2});
+ builder.AddNode("MatMul", {q, k_transposed}, {output})
+ .SetExecutionProviderType(kCpuExecutionProvider);
+ };
+
+ auto check_unfused = [](Graph& graph) {
+ auto op_to_count = CountOpsInGraph(graph);
+ TEST_RETURN_IF_NOT(op_to_count["Transpose"] == 1);
+ TEST_RETURN_IF_NOT(op_to_count["MatMul"] == 1);
+ TEST_RETURN_IF_NOT(op_to_count["com.microsoft.FusedMatMul"] == 0);
+ return Status::OK();
+ };
+
+ ASSERT_STATUS_OK(TestGraphTransformer(
+ build_test_case, 13, *logger_, std::make_unique(),
+ TransformerLevel::Level2, 1, check_unfused, check_unfused));
+}
+
TEST_F(GraphTransformationTests, TransposeCastMatmulFusion) {
const std::vector model_uris = {
MODEL_FOLDER "fusion/transpose_cast_matmul_4d_fusion0.onnx", // Test fusion from the right input
diff --git a/onnxruntime/test/providers/cpu/math/gemm_test.cc b/onnxruntime/test/providers/cpu/math/gemm_test.cc
index 9effdd7e5fb6e..05f9ac2b18224 100644
--- a/onnxruntime/test/providers/cpu/math/gemm_test.cc
+++ b/onnxruntime/test/providers/cpu/math/gemm_test.cc
@@ -118,6 +118,52 @@ TEST(GemmOpTest, GemmNoTrans_f16) {
ConvertFloatToMLFloat16(A.data(), f_A.data(), 8);
ConvertFloatToMLFloat16(B.data(), f_B.data(), 12);
+ {
+ // Missing C uses effective beta == 0.
+ std::vector f_Y(6);
+ std::vector Y{19.3f, -1.4f, -26.9f,
+ -19.3f, 1.4f, 26.9f};
+ ConvertFloatToMLFloat16(Y.data(), f_Y.data(), 6);
+
+ OpTester test("Gemm", 13);
+
+ test.AddAttribute("transA", (int64_t)0);
+ test.AddAttribute("transB", (int64_t)0);
+ test.AddAttribute("alpha", 1.0f);
+ test.AddAttribute("beta", 0.0f);
+ test.AddInput("A", {2, 4}, f_A);
+ test.AddInput("B", {4, 3}, f_B, true);
+ test.AddOutput("Y", {2, 3}, f_Y);
+ test.SetOutputTolerance(0.005f);
+ test.ConfigExcludeEps({kTensorrtExecutionProvider}) // TensorRT: fp16 is not supported
+ .Config(run_with_tunable_op)
+ .RunWithConfig();
+ }
+ {
+ // beta == 0 ignores C, even when C has the full output shape.
+ std::vector f_Y(6);
+ std::vector Y{19.3f, -1.4f, -26.9f,
+ -19.3f, 1.4f, 26.9f};
+ ConvertFloatToMLFloat16(Y.data(), f_Y.data(), 6);
+
+ std::vector f_C(6);
+ ConvertFloatToMLFloat16(C.data(), f_C.data(), 6);
+
+ OpTester test("Gemm", 13);
+
+ test.AddAttribute("transA", (int64_t)0);
+ test.AddAttribute("transB", (int64_t)0);
+ test.AddAttribute("alpha", 1.0f);
+ test.AddAttribute("beta", 0.0f);
+ test.AddInput("A", {2, 4}, f_A);
+ test.AddInput("B", {4, 3}, f_B);
+ test.AddInput("C", {2, 3}, f_C);
+ test.AddOutput("Y", {2, 3}, f_Y);
+ test.SetOutputTolerance(0.005f);
+ test.ConfigExcludeEps({kTensorrtExecutionProvider}) // TensorRT: fp16 is not supported
+ .Config(run_with_tunable_op)
+ .RunWithConfig();
+ }
{
// bias has same shape as output
std::vector f_Y(6);
diff --git a/onnxruntime/test/providers/cpu/math/matmul_test.cc b/onnxruntime/test/providers/cpu/math/matmul_test.cc
index 734cc62459c28..adef2a7cb17dc 100644
--- a/onnxruntime/test/providers/cpu/math/matmul_test.cc
+++ b/onnxruntime/test/providers/cpu/math/matmul_test.cc
@@ -397,6 +397,23 @@ TEST(MathOpTest, MatMulFloatType) {
RunMatMulTest(7, false, true);
}
+TEST(MathOpTest, MatMulFloat16Cpu) {
+ // M > 2 and a non-constant B exercise the regular HalfGemm dispatch path.
+ OpTester test("MatMul", 14);
+ test.AddInput("A", {3, 2},
+ {MLFloat16(1.0f), MLFloat16(2.0f),
+ MLFloat16(3.0f), MLFloat16(4.0f),
+ MLFloat16(5.0f), MLFloat16(6.0f)});
+ test.AddInput("B", {2, 2},
+ {MLFloat16(7.0f), MLFloat16(8.0f),
+ MLFloat16(9.0f), MLFloat16(10.0f)});
+ test.AddOutput("Y", {3, 2},
+ {MLFloat16(25.0f), MLFloat16(28.0f),
+ MLFloat16(57.0f), MLFloat16(64.0f),
+ MLFloat16(89.0f), MLFloat16(100.0f)});
+ test.ConfigEp(DefaultCpuExecutionProvider()).RunWithConfig();
+}
+
#if defined(USE_CUDA) || defined(USE_COREML) || defined(USE_XNNPACK)
TEST(MathOpTest, MatMulFloat16) {
#ifdef USE_CUDA
@@ -520,6 +537,66 @@ TEST(MathOpTest, MatMul_Float16) {
}
#endif
+#ifdef MLAS_F16VEC_INTRINSICS_SUPPORTED
+TEST(MathOpTest, MatMulFloat16SharedInitializerWithKernelOwnedPrepack) {
+ if (!MlasFp16AccelerationSupported()) {
+ GTEST_SKIP() << "Native CPU fp16 MatMul runtime support is unavailable.";
+ }
+
+ std::vector a = {
+ MLFloat16(1.0f), MLFloat16(2.0f), MLFloat16(3.0f), MLFloat16(4.0f),
+ MLFloat16(-1.0f), MLFloat16(-2.0f), MLFloat16(-3.0f), MLFloat16(-4.0f)};
+ std::vector b(12, MLFloat16(1.0f));
+ std::vector y = {
+ MLFloat16(10.0f), MLFloat16(10.0f), MLFloat16(10.0f),
+ MLFloat16(-10.0f), MLFloat16(-10.0f), MLFloat16(-10.0f)};
+
+ OpTester test("MatMul", 14);
+ test.AddInput("A", {2, 4}, a);
+ test.AddInput("B", {4, 3}, b, true);
+ test.AddOutput("Y", {2, 3}, y);
+
+ OrtValue b_initializer;
+ Tensor::InitOrtValue(DataTypeImpl::GetType(), TensorShape({4, 3}),
+ b.data(), OrtMemoryInfo(CPU, OrtAllocatorType::OrtDeviceAllocator), b_initializer);
+
+ SessionOptions so;
+ ASSERT_EQ(so.AddInitializer("B", &b_initializer), Status::OK());
+
+ test.EnableSharingOfPrePackedWeightsAcrossSessions();
+
+ auto cpu_ep = []() -> std::vector> {
+ std::vector> execution_providers;
+ execution_providers.push_back(DefaultCpuExecutionProvider());
+ return execution_providers;
+ };
+
+ size_t number_of_pre_packed_weights_counter_session_1 = 0;
+ size_t number_of_shared_pre_packed_weights_counter = 0;
+
+ test.Config(so)
+ .Config(run_with_tunable_op)
+ .ConfigEps(cpu_ep())
+ .RunWithConfig(&number_of_pre_packed_weights_counter_session_1, &number_of_shared_pre_packed_weights_counter);
+
+ if (number_of_pre_packed_weights_counter_session_1 == 0) {
+ GTEST_SKIP() << "Native CPU fp16 MatMul packed-B path is unavailable.";
+ }
+
+ ASSERT_EQ(test.GetNumPrePackedWeightsShared(), static_cast(0));
+ ASSERT_EQ(number_of_shared_pre_packed_weights_counter, static_cast(0));
+
+ size_t number_of_pre_packed_weights_counter_session_2 = 0;
+ test.Config(so)
+ .Config(run_with_tunable_op)
+ .ConfigEps(cpu_ep())
+ .RunWithConfig(&number_of_pre_packed_weights_counter_session_2, &number_of_shared_pre_packed_weights_counter);
+
+ ASSERT_EQ(number_of_pre_packed_weights_counter_session_2, number_of_pre_packed_weights_counter_session_1);
+ ASSERT_EQ(number_of_shared_pre_packed_weights_counter, static_cast(0));
+}
+#endif
+
#if defined(USE_CUDA) || defined(USE_DNNL)
TEST(MathOpTest, MatMul_bfloat16) {
#ifdef USE_CUDA
diff --git a/onnxruntime/test/providers/cpu/nn/conv_fp16_test.cc b/onnxruntime/test/providers/cpu/nn/conv_fp16_test.cc
index 843d925ed6638..03b1eac1e91bf 100644
--- a/onnxruntime/test/providers/cpu/nn/conv_fp16_test.cc
+++ b/onnxruntime/test/providers/cpu/nn/conv_fp16_test.cc
@@ -10,6 +10,7 @@
#include "test/common/random_generator.h"
#include "test/providers/provider_test_utils.h"
#include "default_providers.h"
+#include "core/session/onnxruntime_session_options_config_keys.h"
using namespace std;
namespace onnxruntime {
@@ -47,7 +48,8 @@ void TestConvFp16Op(const ConvOpAndTestAttributes& attributes,
OpTester::ExpectResult expect_result = OpTester::ExpectResult::kExpectSuccess,
const std::string& err_str = "",
int opset = 11,
- float rel_error = 0.002f) {
+ float rel_error = 0.002f,
+ bool disable_kleidiai = false) {
std::unique_ptr tester;
if (!attributes.activation.empty()) {
tester = std::make_unique("NhwcFusedConv", 1, onnxruntime::kMSDomain);
@@ -89,6 +91,12 @@ void TestConvFp16Op(const ConvOpAndTestAttributes& attributes,
tester->AddOutput("Y", expected_output_shape, expected_output, /*no sort*/ false, rel_error, 0.0f);
+ if (disable_kleidiai) {
+ SessionOptions session_options;
+ ASSERT_STATUS_OK(session_options.config_options.AddConfigEntry(kOrtSessionOptionsMlasDisableKleidiAi, "1"));
+ tester->Config(session_options);
+ }
+
std::unordered_set excluded_providers(attributes.excluded_providers);
// Disable TensorRT because weight as input is not supported
excluded_providers.insert(kTensorrtExecutionProvider);
@@ -369,6 +377,158 @@ TEST(ConvFp16Test, Conv2D_1) {
TestConvFp16Op(attrs, {X, W}, {X_shape, W_shape}, expected_vals, Y_shape, true);
}
+TEST(ConvFp16Test, Conv2D_KleidiAiImatmulEligibleNoBias) {
+ ConvOpAndTestAttributes attrs = {
+ "", // auto_pad
+ vector{1, 1}, // dilations
+ 1, // group
+ vector{3, 3}, // kernel_shape
+ vector{1, 1, 1, 1}, // pads
+ vector{1, 1}, // strides
+ {} // excluded EPs
+ };
+
+ vector X = {
+ MLFloat16(1.0f), MLFloat16(2.0f), MLFloat16(3.0f), MLFloat16(4.0f),
+ MLFloat16(5.0f), MLFloat16(6.0f), MLFloat16(7.0f), MLFloat16(8.0f),
+ MLFloat16(9.0f), MLFloat16(10.0f), MLFloat16(11.0f), MLFloat16(12.0f),
+ MLFloat16(13.0f), MLFloat16(14.0f), MLFloat16(15.0f), MLFloat16(16.0f)};
+ vector X_shape = {1, 1, 4, 4};
+ vector W = {
+ MLFloat16(1.0f), MLFloat16(1.0f), MLFloat16(1.0f),
+ MLFloat16(1.0f), MLFloat16(1.0f), MLFloat16(1.0f),
+ MLFloat16(1.0f), MLFloat16(1.0f), MLFloat16(1.0f),
+ MLFloat16(0.5f), MLFloat16(0.5f), MLFloat16(0.5f),
+ MLFloat16(0.5f), MLFloat16(0.5f), MLFloat16(0.5f),
+ MLFloat16(0.5f), MLFloat16(0.5f), MLFloat16(0.5f)};
+ vector W_shape = {2, 1, 3, 3};
+ vector Y_shape = {1, 2, 4, 4};
+ auto expected_vals = {
+ MLFloat16(14.0f), MLFloat16(24.0f), MLFloat16(30.0f), MLFloat16(22.0f),
+ MLFloat16(33.0f), MLFloat16(54.0f), MLFloat16(63.0f), MLFloat16(45.0f),
+ MLFloat16(57.0f), MLFloat16(90.0f), MLFloat16(99.0f), MLFloat16(69.0f),
+ MLFloat16(46.0f), MLFloat16(72.0f), MLFloat16(78.0f), MLFloat16(54.0f),
+ MLFloat16(7.0f), MLFloat16(12.0f), MLFloat16(15.0f), MLFloat16(11.0f),
+ MLFloat16(16.5f), MLFloat16(27.0f), MLFloat16(31.5f), MLFloat16(22.5f),
+ MLFloat16(28.5f), MLFloat16(45.0f), MLFloat16(49.5f), MLFloat16(34.5f),
+ MLFloat16(23.0f), MLFloat16(36.0f), MLFloat16(39.0f), MLFloat16(27.0f)};
+
+ TestConvFp16Op(attrs, {X, W}, {X_shape, W_shape}, expected_vals, Y_shape);
+ TestConvFp16Op(attrs, {X, W}, {X_shape, W_shape}, expected_vals, Y_shape, true);
+}
+
+TEST(ConvFp16Test, Conv2D_KleidiAiImatmulEligibleBiasAndDisabledFallback) {
+ ConvOpAndTestAttributes attrs = {
+ "", // auto_pad
+ vector{1, 1}, // dilations
+ 1, // group
+ vector{3, 3}, // kernel_shape
+ vector{1, 1, 1, 1}, // pads
+ vector{1, 1}, // strides
+ {} // excluded EPs
+ };
+
+ vector X = {
+ MLFloat16(1.0f), MLFloat16(2.0f), MLFloat16(3.0f), MLFloat16(4.0f),
+ MLFloat16(5.0f), MLFloat16(6.0f), MLFloat16(7.0f), MLFloat16(8.0f),
+ MLFloat16(9.0f), MLFloat16(10.0f), MLFloat16(11.0f), MLFloat16(12.0f),
+ MLFloat16(13.0f), MLFloat16(14.0f), MLFloat16(15.0f), MLFloat16(16.0f)};
+ vector X_shape = {1, 1, 4, 4};
+ vector W = {
+ MLFloat16(1.0f), MLFloat16(1.0f), MLFloat16(1.0f),
+ MLFloat16(1.0f), MLFloat16(1.0f), MLFloat16(1.0f),
+ MLFloat16(1.0f), MLFloat16(1.0f), MLFloat16(1.0f),
+ MLFloat16(0.5f), MLFloat16(0.5f), MLFloat16(0.5f),
+ MLFloat16(0.5f), MLFloat16(0.5f), MLFloat16(0.5f),
+ MLFloat16(0.5f), MLFloat16(0.5f), MLFloat16(0.5f)};
+ vector W_shape = {2, 1, 3, 3};
+ vector B = {MLFloat16(1.0f), MLFloat16(-2.0f)};
+ vector B_shape = {2};
+ vector Y_shape = {1, 2, 4, 4};
+ auto expected_vals = {
+ MLFloat16(15.0f), MLFloat16(25.0f), MLFloat16(31.0f), MLFloat16(23.0f),
+ MLFloat16(34.0f), MLFloat16(55.0f), MLFloat16(64.0f), MLFloat16(46.0f),
+ MLFloat16(58.0f), MLFloat16(91.0f), MLFloat16(100.0f), MLFloat16(70.0f),
+ MLFloat16(47.0f), MLFloat16(73.0f), MLFloat16(79.0f), MLFloat16(55.0f),
+ MLFloat16(5.0f), MLFloat16(10.0f), MLFloat16(13.0f), MLFloat16(9.0f),
+ MLFloat16(14.5f), MLFloat16(25.0f), MLFloat16(29.5f), MLFloat16(20.5f),
+ MLFloat16(26.5f), MLFloat16(43.0f), MLFloat16(47.5f), MLFloat16(32.5f),
+ MLFloat16(21.0f), MLFloat16(34.0f), MLFloat16(37.0f), MLFloat16(25.0f)};
+
+ TestConvFp16Op(attrs, {X, W, B}, {X_shape, W_shape, B_shape}, expected_vals, Y_shape);
+
+ constexpr bool weight_is_initializer = true;
+ TestConvFp16Op(attrs, {X, W, B}, {X_shape, W_shape, B_shape}, expected_vals, Y_shape, weight_is_initializer);
+
+ constexpr bool disable_kleidiai = true;
+ TestConvFp16Op(attrs, {X, W, B}, {X_shape, W_shape, B_shape}, expected_vals, Y_shape, weight_is_initializer,
+ OpTester::ExpectResult::kExpectSuccess, "", 11, 0.002f, disable_kleidiai);
+}
+
+TEST(ConvFp16Test, NhwcFusedConv2D_KleidiAiImatmulEligibleBiasAndDisabledFallback) {
+#if !defined(__aarch64__) && !defined(_M_ARM64)
+ GTEST_SKIP() << "Native CPU fp16 Conv runtime support is only tested on Arm64.";
+#else
+ if (!MlasFp16AccelerationSupported()) {
+ GTEST_SKIP() << "Native CPU fp16 Conv runtime support is unavailable.";
+ }
+#endif
+
+ auto run_test = [](bool disable_kleidiai) {
+ OpTester test("NhwcFusedConv", 1, onnxruntime::kMSDomain);
+ test.AddAttribute("group", static_cast(1));
+ test.AddAttribute("kernel_shape", vector{3, 3});
+ test.AddAttribute("pads", vector{1, 1, 1, 1});
+ test.AddAttribute("strides", vector{1, 1});
+ test.AddAttribute("dilations", vector{1, 1});
+
+ vector X = {
+ MLFloat16(1.0f), MLFloat16(2.0f), MLFloat16(3.0f), MLFloat16(4.0f),
+ MLFloat16(5.0f), MLFloat16(6.0f), MLFloat16(7.0f), MLFloat16(8.0f),
+ MLFloat16(9.0f), MLFloat16(10.0f), MLFloat16(11.0f), MLFloat16(12.0f),
+ MLFloat16(13.0f), MLFloat16(14.0f), MLFloat16(15.0f), MLFloat16(16.0f)};
+ vector X_shape = {1, 4, 4, 1};
+ vector W = {
+ MLFloat16(1.0f), MLFloat16(1.0f), MLFloat16(1.0f),
+ MLFloat16(1.0f), MLFloat16(1.0f), MLFloat16(1.0f),
+ MLFloat16(1.0f), MLFloat16(1.0f), MLFloat16(1.0f),
+ MLFloat16(0.5f), MLFloat16(0.5f), MLFloat16(0.5f),
+ MLFloat16(0.5f), MLFloat16(0.5f), MLFloat16(0.5f),
+ MLFloat16(0.5f), MLFloat16(0.5f), MLFloat16(0.5f)};
+ vector W_shape = {2, 1, 3, 3};
+ vector B = {MLFloat16(1.0f), MLFloat16(-2.0f)};
+ vector B_shape = {2};
+ vector Y_shape = {1, 4, 4, 2};
+ auto expected_vals = {
+ MLFloat16(15.0f), MLFloat16(5.0f), MLFloat16(25.0f), MLFloat16(10.0f),
+ MLFloat16(31.0f), MLFloat16(13.0f), MLFloat16(23.0f), MLFloat16(9.0f),
+ MLFloat16(34.0f), MLFloat16(14.5f), MLFloat16(55.0f), MLFloat16(25.0f),
+ MLFloat16(64.0f), MLFloat16(29.5f), MLFloat16(46.0f), MLFloat16(20.5f),
+ MLFloat16(58.0f), MLFloat16(26.5f), MLFloat16(91.0f), MLFloat16(43.0f),
+ MLFloat16(100.0f), MLFloat16(47.5f), MLFloat16(70.0f), MLFloat16(32.5f),
+ MLFloat16(47.0f), MLFloat16(21.0f), MLFloat16(73.0f), MLFloat16(34.0f),
+ MLFloat16(79.0f), MLFloat16(37.0f), MLFloat16(55.0f), MLFloat16(25.0f)};
+
+ test.AddInput("X", X_shape, X);
+ test.AddInput("W", W_shape, W, true);
+ test.AddInput("B", B_shape, B, true);
+ test.AddOutput("Y", Y_shape, expected_vals, /*no sort*/ false, 0.002f, 0.0f);
+
+ if (disable_kleidiai) {
+ SessionOptions session_options;
+ ASSERT_STATUS_OK(session_options.config_options.AddConfigEntry(kOrtSessionOptionsMlasDisableKleidiAi, "1"));
+ test.Config(session_options);
+ }
+
+ std::vector> execution_providers;
+ execution_providers.push_back(DefaultCpuExecutionProvider());
+ test.ConfigEps(std::move(execution_providers)).RunWithConfig();
+ };
+
+ run_test(false);
+ run_test(true);
+}
+
TEST(ConvFp16Test, Conv2D_2) {
ConvOpAndTestAttributes attrs = {
"", // auto_pad
@@ -1496,6 +1656,165 @@ TEST(ConvFp16Test, SharedPrepackedWeights) {
}
}
+TEST(ConvFp16Test, SharedPrepackedWeights_HalfConvEligible_NoBias) {
+ OpTester test("Conv", 11);
+ test.AddAttribute("group", static_cast(1));
+ test.AddAttribute("kernel_shape", vector{3, 3});
+ test.AddAttribute("pads", vector{1, 1, 1, 1});
+ test.AddAttribute("strides", vector{1, 1});
+ test.AddAttribute("dilations", vector{1, 1});
+
+ vector X = {MLFloat16(1.0f), MLFloat16(2.0f), MLFloat16(3.0f), MLFloat16(4.0f),
+ MLFloat16(5.0f), MLFloat16(6.0f), MLFloat16(7.0f), MLFloat16(8.0f),
+ MLFloat16(9.0f), MLFloat16(10.0f), MLFloat16(11.0f), MLFloat16(12.0f),
+ MLFloat16(13.0f), MLFloat16(14.0f), MLFloat16(15.0f), MLFloat16(16.0f)};
+ vector X_shape = {1, 1, 4, 4};
+ vector W = {MLFloat16(1.0f), MLFloat16(1.0f), MLFloat16(1.0f), MLFloat16(1.0f), MLFloat16(1.0f),
+ MLFloat16(1.0f), MLFloat16(1.0f), MLFloat16(1.0f), MLFloat16(1.0f), MLFloat16(0.5f),
+ MLFloat16(0.5f), MLFloat16(0.5f), MLFloat16(0.5f), MLFloat16(0.5f), MLFloat16(0.5f),
+ MLFloat16(0.5f), MLFloat16(0.5f), MLFloat16(0.5f)};
+ vector W_shape = {2, 1, 3, 3};
+ vector Y_shape = {1, 2, 4, 4};
+ auto expected_vals = {
+ MLFloat16(14.0f), MLFloat16(24.0f), MLFloat16(30.0f), MLFloat16(22.0f), MLFloat16(33.0f), MLFloat16(54.0f),
+ MLFloat16(63.0f), MLFloat16(45.0f), MLFloat16(57.0f), MLFloat16(90.0f), MLFloat16(99.0f), MLFloat16(69.0f),
+ MLFloat16(46.0f), MLFloat16(72.0f), MLFloat16(78.0f), MLFloat16(54.0f), MLFloat16(7.0f), MLFloat16(12.0f),
+ MLFloat16(15.0f), MLFloat16(11.0f), MLFloat16(16.5f), MLFloat16(27.0f), MLFloat16(31.5f), MLFloat16(22.5f),
+ MLFloat16(28.5f), MLFloat16(45.0f), MLFloat16(49.5f), MLFloat16(34.5f), MLFloat16(23.0f), MLFloat16(36.0f),
+ MLFloat16(39.0f), MLFloat16(27.0f)};
+
+ test.AddInput("X", X_shape, X);
+ test.AddInput("W", W_shape, W, true);
+ test.AddOutput("Y", Y_shape, expected_vals, /*no sort*/ false, 0.002f, 0.0f);
+
+ OrtValue w;
+ Tensor::InitOrtValue(DataTypeImpl::GetType(), TensorShape(W_shape),
+ W.data(), OrtMemoryInfo(CPU, OrtAllocatorType::OrtDeviceAllocator), w);
+
+ SessionOptions so;
+ ASSERT_EQ(so.AddInitializer("W", &w), Status::OK());
+
+ test.EnableSharingOfPrePackedWeightsAcrossSessions();
+
+ auto cpu_ep = []() -> std::vector> {
+ std::vector> execution_providers;
+ execution_providers.push_back(DefaultCpuExecutionProvider());
+ return execution_providers;
+ };
+
+ size_t number_of_pre_packed_weights_counter_session_1 = 0;
+ size_t number_of_shared_pre_packed_weights_counter = 0;
+
+ {
+ test.Config(so)
+ .Config(run_with_tunable_op)
+ .ConfigEps(cpu_ep())
+ .RunWithConfig(&number_of_pre_packed_weights_counter_session_1, &number_of_shared_pre_packed_weights_counter);
+ ASSERT_EQ(number_of_shared_pre_packed_weights_counter, static_cast(0));
+ }
+
+ const auto number_of_elements_in_shared_prepacked_buffers_container = test.GetNumPrePackedWeightsShared();
+
+ if (number_of_pre_packed_weights_counter_session_1 == 0) {
+ GTEST_SKIP() << "No pre-packed weights were produced.";
+ }
+
+ ASSERT_EQ(number_of_elements_in_shared_prepacked_buffers_container, static_cast(1));
+
+ {
+ size_t number_of_pre_packed_weights_counter_session_2 = 0;
+ test.Config(so)
+ .Config(run_with_tunable_op)
+ .ConfigEps(cpu_ep())
+ .RunWithConfig(&number_of_pre_packed_weights_counter_session_2, &number_of_shared_pre_packed_weights_counter);
+
+ ASSERT_GE(number_of_pre_packed_weights_counter_session_1, number_of_shared_pre_packed_weights_counter);
+ ASSERT_GE(number_of_pre_packed_weights_counter_session_2, number_of_shared_pre_packed_weights_counter);
+ ASSERT_EQ(number_of_shared_pre_packed_weights_counter, static_cast(1));
+ }
+}
+
+TEST(ConvFp16Test, SharedPrepackedWeights_HalfConvEligible_BiasNotShared) {
+ OpTester test("Conv", 11);
+ test.AddAttribute("group", static_cast(1));
+ test.AddAttribute("kernel_shape", vector{3, 3});
+ test.AddAttribute("pads", vector