Update const for kernel interface and sme checks

Signed-off-by: Jonathan Clohessy <[email protected]>

Author: JonathanC-ARM

onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_matmul.cc

@@ -15,6 +15,10 @@
 #include <algorithm>
 #include <vector>
 
+#if defined(USE_KLEIDIAI) && !defined(_MSC_VER)
+#include "core/mlas/lib/kleidiai/mlasi_kleidiai.h"
+#endif
+
 namespace onnxruntime {
 namespace contrib {
 
@@ -215,7 +219,7 @@ class DynamicQuantizeMatMul final : public MatMulIntegerToFloatBase {
 
       // Currently, MlasDynamicQGemmBatch() and associated functions require SME or else they are no-ops.
       // We check that here too before attempting to use them.
-      if (!CPUIDInfo::GetCPUIDInfo().HasArm_SME()) {
+      if (!ArmKleidiAI::SMEInfo::CanUseSME2) {
         can_use_dynamic_quant_mlas_ = false;
       }
 

onnxruntime/core/mlas/lib/convolve.cpp

@@ -862,9 +862,9 @@ Return Value:
 --*/
 {
     // Override
-    if(GetMlasPlatform().MlasConvOverride != nullptr &&
+    if(ArmKleidiAI::SMEInfo::CanUseSME2 && GetMlasPlatform().MlasConvOverride != nullptr &&
         GetMlasPlatform().MlasConvOverride(Parameters,Input,Filter,Bias,WorkingBuffer,Output,ThreadPool)){
-    return;
+        return;
     }
 
     const size_t FilterCount = Parameters->FilterCount;
@@ -1101,7 +1101,7 @@ Return Value:
 --*/
 {
     // Override
-    if (GetMlasPlatform().MlasConvPrepareOverride != nullptr &&
+    if (ArmKleidiAI::SMEInfo::CanUseSME2 && GetMlasPlatform().MlasConvPrepareOverride != nullptr &&
         GetMlasPlatform().MlasConvPrepareOverride(Parameters, Dimensions, BatchCount, GroupCount, InputChannels,
         InputShape,KernelShape,DilationShape, Padding, StrideShape, OutputShape, FilterCount,
         Activation, WorkingBufferSize, Beta, ThreadPool)){

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

@@ -27,8 +27,8 @@ struct KaiTlsBuffers {
 };
 static thread_local KaiTlsBuffers g_kai_tls;
 
-kai_matmul_clamp_f32_f32p_f32p_ukernel sgemm_gemm = GetKleidiAISGemmUKernel();
-kai_matmul_clamp_f32_f32_f32p_ukernel sgemm_gemv = GetKleidiAISGemvUKernel();
+const kai_matmul_clamp_f32_f32p_f32p_ukernel sgemm_gemm = GetKleidiAISGemmUKernel();
+const kai_matmul_clamp_f32_f32_f32p_ukernel sgemm_gemv = GetKleidiAISGemvUKernel();
 
 
 // Helpers for GEMV

onnxruntime/core/mlas/lib/qgemm.cpp

@@ -423,7 +423,7 @@ MlasDynamicQgemmPackB(
 {
 #if defined(USE_KLEIDIAI) && !defined(_MSC_VER)
     //No fallback
-    if(MLAS_CPUIDINFO::GetCPUIDInfo().HasArm_SME()){//Still require this since no override
+    if (ArmKleidiAI::SMEInfo::CanUseSME2) {//Still require this since no override
         ArmKleidiAI::MlasDynamicQgemmPackB(N, K, B, Scales, Bias, PackedB);
     }
 #endif

onnxruntime/test/mlas/unittest/test_dynamic_qgemm.cpp

@@ -9,6 +9,7 @@
 
 #include "test_util.h"
 #include "core/mlas/lib/mlasi.h"  // for MLAS_CPUIDINFO
+#include "core/mlas/lib/kleidiai/mlasi_kleidiai.h"
 
 class MlasDynamicQgemmTest {
  private:
@@ -21,7 +22,7 @@ class MlasDynamicQgemmTest {
  public:
   void Test(size_t M, size_t N, size_t K, size_t BatchSize) {
     // Currently, MlasDynamicQGemmBatch() and associated functions require SME or else they are no-ops.
-    if (!MLAS_CPUIDINFO::GetCPUIDInfo().HasArm_SME()) {
+    if (!ArmKleidiAI::SMEInfo::CanUseSME2) {
       GTEST_SKIP() << "MlasDynamicQGemmBatch() requires ARM64 SME but it was not detected. Skipping test.";
     }