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 SME2 or else they are no-ops.
       // We check that here too before attempting to use them.
-      if (!CPUIDInfo::GetCPUIDInfo().HasArm_SME2()) {
+      if (!ArmKleidiAI::SMEInfo::CanUseSME2) {
         can_use_dynamic_quant_mlas_ = false;
       }
 

onnxruntime/core/mlas/lib/convolve.cpp

@@ -938,9 +938,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;
@@ -1201,7 +1201,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

@@ -407,7 +407,7 @@ Return Value:
         ~(BufferAlignment - 1);
     // If this gemm B argument is used in a dynamically quantization gemm operation we can optimize for
     // this use case. Concat both packed representations for later decision. This allows for cases later
-    // where we still have the prepack at the cost of some memory otherwise we can use the qgemm quantization 
+    // where we still have the prepack at the cost of some memory otherwise we can use the qgemm quantization
     // for better performance
     return AlignedBytesRequired + MlasDynamicQgemmPackBSize(N, K);
 }
@@ -425,7 +425,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: