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

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)<br/> **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)<br/> **indices** = tensor(int64)|
 |Gelu|*in* X:**T**<br> *out* Y:**T**|20+|**T** = tensor(float)|
-|Gemm|*in* A:**T**<br> *in* B:**T**<br> *in* C:**T**<br> *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**<br> *in* B:**T**<br> *in* C:**T**<br> *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**<br> *out* Y:**T**|22+|**T** = tensor(float)|
 |||[1, 21]|**T** = tensor(float)|
 |GlobalLpPool|*in* X:**T**<br> *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**<br> *in* B:**T**<br> *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**<br> *in* B:**T**<br> *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**<br> *in* B:**T2**<br> *in* a_zero_point:**T1**<br> *in* b_zero_point:**T2**<br> *out* Y:**T3**|10+|**T1** = tensor(int8), tensor(uint8)<br/> **T2** = tensor(int8), tensor(uint8)<br/> **T3** = tensor(int32)|
 |Max|*in* data_0:**T**<br> *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)|

onnxruntime/contrib_ops/cpu/moe/moe_cpu.cc

@@ -536,7 +536,7 @@ Status MoE<float>::ComputeGEMM(const float* A, const float* B, float* C,
 template <>
 Status MoE<MLFloat16>::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<size_t>(K);
   params.C = C;

onnxruntime/core/framework/prepacked_weights.h

@@ -20,6 +20,10 @@ struct PrePackedWeights final {
   InlinedVector<IAllocatorUniquePtr<void>> buffers_;  // cache pre-packed buffers associated with the kernel
   InlinedVector<size_t> 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;
 

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;
 }
 

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

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

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;

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;
     }

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

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

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<ONNX_NAMESPACE::TensorProto_DataType>(left_type))) {
+    if (!IsAllowedFusedMatMulDataType(static_cast<ONNX_NAMESPACE::TensorProto_DataType>(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<ONNX_NAMESPACE::TensorProto_DataType>(right_type))) {
+    if (!IsAllowedFusedMatMulDataType(static_cast<ONNX_NAMESPACE::TensorProto_DataType>(right_type),
+                                      node.GetExecutionProviderType())) {
       continue;
     }