Fix builds

Author: hariharans29

onnxruntime/test/mlas/unittest/test_fgemm.h

@@ -87,7 +87,7 @@ class FgemmPackedContext<double, false> {
       data[i].alpha = alpha;
       data[i].beta = beta;
     }
-    MlasGemmBatch(TransA, TransB, M, N, K, data.data(), BatchSize, threadpool, nullptr, nullptr);
+    MlasGemmBatch(TransA, TransB, M, N, K, data.data(), BatchSize, threadpool, nullptr);
   }
 };
 #endif

orttraining/orttraining/training_ops/cpu/nn/conv_grad.cc

@@ -169,7 +169,7 @@ Status ConvGrad<T>::Compute(OpKernelContext* context) const {
             skip_im2col ? Xdata + group_id * X_offset : col_buffer_data,
             1,
             dWdata + group_id * W_offset,
-            tp);
+            tp, &mlas_backend_kernel_selector_config_);
       }
     }
     if (dB) {
@@ -207,7 +207,7 @@ Status ConvGrad<T>::Compute(OpKernelContext* context) const {
             dYdata,
             0,
             col_buffer_data,
-            tp);
+            tp, &mlas_backend_kernel_selector_config_);
 
         if (kernel_rank == 2) {
           math::Col2im<T, CPUMathUtil, StorageOrder::NCHW>(

orttraining/orttraining/training_ops/cpu/nn/conv_grad.h

@@ -5,6 +5,8 @@
 
 #include "core/framework/op_kernel.h"
 #include "core/providers/cpu/nn/conv_attributes.h"
+#include "core/mlas/inc/mlas.h"
+#include "core/session/onnxruntime_session_options_config_keys.h"
 
 namespace onnxruntime {
 namespace contrib {
@@ -13,6 +15,8 @@ template <typename T>
 class ConvGrad final : public OpKernel {
  public:
   explicit ConvGrad(const OpKernelInfo& info) : OpKernel(info), conv_attrs_(info) {
+    mlas_backend_kernel_selector_config_.use_kleidiai =
+                              info.GetConfigOptions().GetConfigEntry(kOrtSessionOptionsMlasDisableKleidiai) != "1";
   }
 
   Status Compute(OpKernelContext* context) const override;
@@ -22,6 +26,7 @@ class ConvGrad final : public OpKernel {
 
  private:
   ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(ConvGrad);
+  MLAS_BACKEND_KERNEL_SELECTOR_CONFIG mlas_backend_kernel_selector_config_;
 };
 
 }  // namespace contrib

orttraining/orttraining/training_ops/cpu/op_gradients.cc

@@ -148,7 +148,7 @@ Status SoftmaxGrad<T>::Compute(OpKernelContext* context) const {
     concurrency::ThreadPool* tp = context->GetOperatorThreadPool();
     math::Gemm<float>(CblasNoTrans, CblasNoTrans, n, d, 1, -1,
                       scaledata, sum_multiplier_.data(), 1,
-                      dXdata, tp);
+                      dXdata, tp, &mlas_backend_kernel_selector_config_);
 
     math::Mul<float, CPUMathUtil>(gsl::narrow_cast<int>(Y.Shape().Size()), dXdata, Ydata, dXdata, nullptr);
   }

orttraining/orttraining/training_ops/cpu/op_gradients.h

@@ -4,6 +4,8 @@
 #pragma once
 
 #include "core/framework/op_kernel.h"
+#include "core/mlas/inc/mlas.h"
+#include "core/session/onnxruntime_session_options_config_keys.h"
 #include <cctype>
 
 namespace onnxruntime {
@@ -67,6 +69,8 @@ class SoftmaxGrad final : public OpKernel {
     opset_ = (node.OpType() == "SoftmaxGrad_13" || node.OpType() == "LogSoftmaxGrad_13") ? 13 : 1;
     axis_ = info.GetAttrOrDefault("axis", static_cast<int64_t>(opset_ < 13 ? 1 : -1));
     is_logsoftmaxgrad_ = node.OpType() == "LogSoftmaxGrad_13" || node.OpType() == "LogSoftmaxGrad";
+    mlas_backend_kernel_selector_config_.use_kleidiai =
+                              info.GetConfigOptions().GetConfigEntry(kOrtSessionOptionsMlasDisableKleidiai) != "1";
   }
 
   Status Compute(OpKernelContext* context) const override;
@@ -76,6 +80,7 @@ class SoftmaxGrad final : public OpKernel {
   int64_t axis_;
   int opset_;  // opset_ of the forward Softmax operator
   bool is_logsoftmaxgrad_;
+  MLAS_BACKEND_KERNEL_SELECTOR_CONFIG mlas_backend_kernel_selector_config_;
 };
 
 template <typename T>