Address review feedback

Signed-off-by: Jonathan Clohessy <Jonathan.Clohessy@arm.com>

Author: JonathanC-ARM

onnxruntime/contrib_ops/cpu/moe/moe_cpu.cc

@@ -550,6 +550,7 @@ Status MoE<MLFloat16>::ComputeGEMM(const MLFloat16* A, const MLFloat16* B, MLFlo
   } else {
     params.ldb = static_cast<size_t>(N);
   }
+  params.BackendKernelSelectorConfig = &mlas_backend_kernel_selector_config_;
 
   MlasHalfGemmBatch(static_cast<size_t>(M), static_cast<size_t>(N), static_cast<size_t>(K), 1, &params, nullptr);
   return Status::OK();

onnxruntime/core/mlas/lib/kleidiai/halfconv_kleidiai.cpp

@@ -32,7 +32,7 @@ struct KaiHalfConvTlsBuffers {
     std::vector<std::byte> packed_lhs;
 
     void ReleaseLargeBuffers() {
-        MlasShrinkKleidiAIScratchIfTooLarge(packed_lhs);
+        ArmKleidiAI::MlasShrinkKleidiAIScratchIfTooLarge(packed_lhs);
     }
 };
 

onnxruntime/core/mlas/lib/kleidiai/halfgemm_kleidiai.cpp

@@ -24,10 +24,10 @@ struct KaiHalfTlsBuffers {
     std::vector<std::byte> 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;

onnxruntime/core/mlas/lib/kleidiai/mlasi_kleidiai.h

@@ -8,7 +8,6 @@
 
 #include "../mlasi.h"
 #include <limits>
-#include <iostream>
 #include <vector>
 
 // 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 <iostream>
 #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 <typename T>
@@ -62,8 +63,6 @@ void MlasShrinkKleidiAIScratchIfTooLarge(std::vector<T>& 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();

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<size_t>(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<size_t>(output_count),

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