[MLAS] KleidiAI fix igemm regression (#28571)

### Description
<!-- Describe your changes. -->

This PR fixes a convolution performance regression affecting some OCR
models with large-kernel convolutions when the KleidiAI SME IGEMM
convolution path is selected.

The change has 2 parts:
1. updates to the KleidiAI IGEMM LHS packing to pack rows in bounded
chunks instead of packing the full LHS buffer up front, which reduces
memory usage and improves cache locality for large convolutions,
2. a new route selection function `ArmKleidiAI::SelectConvRoute` that
decides between `Igemm`, `GemmFallback` and `None` based on convolution
parameters and a workload size-based heuristic.

The function `CheckCapabilitiesSme` runs `SelectConvRoute` and only
returns true if the selected route is `Igemm`. The patch also adds a
standard GEMM fallback to the `ConvRoute` possibilities, and runs
`MlasGemm` if said fallback is selected. If the function selects `None`,
then the convolution falls back to `MlasSgemmOperation`.

### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->

Fixes #27633.

---------

Signed-off-by: Qxiang Xu <Qixiang.Xu@arm.com>
Signed-off-by: Jonathan Clohessy <Jonathan.Clohessy@arm.com>
Signed-off-by: Martin Klacer <martin.klacer@arm.com>
Co-authored-by: Jonathan Clohessy <Jonathan.Clohessy@arm.com>
Co-authored-by: Damien Dooley <damien.dooley@arm.com>

Author: 3 people

include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h

@@ -498,6 +498,14 @@ static const char* const kOrtSessionOptionsMlasLutGemm = "mlas.use_lut_gemm";
 // - "1": Disable KleidiAI kernels even if available.
 static const char* const kOrtSessionOptionsMlasDisableKleidiAi = "mlas.disable_kleidiai";
 
+// Power-user tuning option for the Arm® KleidiAI™ SME IGEMM convolution route on Arm64.
+// For 2D convolutions where both SME IGEMM and the MlasGemm SGEMM fallback are valid routes, work is estimated as:
+//  output_h * output_w * input_channels * dilated_kernel_h * dilated_kernel_w * filter_count.
+// Work above this threshold routes through the SGEMM fallback; work at or below it stays on IGEMM.
+// "0" or unset uses the MLAS default heuristic, intended for typical workloads.
+// This option exists for perf experimentation; the default may be retuned in future ORT releases.
+static const char* const kOrtSessionOptionsMlasKleidiAiConvIgemmMaxWork = "mlas.kleidiai.conv_igemm_max_work";
+
 // When converting DQ + MatMul -> MatMulNBits, the accuracy level of the MatMulNBits is controlled by this option.
 // Refer to MatMulNBits op schema for more details.
 // If not provided, default is 4.

onnxruntime/core/mlas/inc/mlas.h

@@ -207,6 +207,7 @@ MlasActivation(
 
 struct MLAS_BACKEND_KERNEL_SELECTOR_CONFIG {
     bool use_kleidiai = true; /**< Flag to use KleidiAI backend kernels if available */
+    size_t kleidiai_conv_igemm_max_work = 0; /**< Optional SME IGEMM route threshold override; 0 uses default */
 };
 
 //

onnxruntime/core/mlas/lib/convolve.cpp

@@ -574,6 +574,9 @@ Return Value:
     const size_t FilterCount = Parameters->FilterCount;
     const size_t OutputSize = Parameters->OutputSize;
     const size_t K = Parameters->K;
+    const bool use_gemm_batch_override =
+        GetMlasPlatform().MlasConvSGemmRouteOverride != nullptr &&
+        GetMlasPlatform().MlasConvSGemmRouteOverride(Parameters) == MlasConvSGemmRouteDispatch;
 
     //
     // Compute the strides to step through slices of the local segment.
@@ -637,9 +640,15 @@ Return Value:
                     SegmentStartN + n, CountN);
             }
 
-            MlasSgemmOperation(CblasNoTrans, CblasNoTrans, FilterCount, CountN,
-                CountK, 1.0f, Filter + k, K, ColumnBuffer, CountN, beta,
-                SegmentOutput, OutputSize);
+            if (use_gemm_batch_override) {
+                MlasGemm(CblasNoTrans, CblasNoTrans, FilterCount, CountN,
+                    CountK, 1.0f, Filter + k, K, ColumnBuffer, CountN, beta,
+                    SegmentOutput, OutputSize, nullptr, Parameters->BackendKernelSelectorConfig);
+            } else {
+                MlasSgemmOperation(CblasNoTrans, CblasNoTrans, FilterCount, CountN,
+                    CountK, 1.0f, Filter + k, K, ColumnBuffer, CountN, beta,
+                    SegmentOutput, OutputSize);
+            }
 
             beta = 1.0f;
         }
@@ -1385,7 +1394,12 @@ MlasConvSupportsDenseChannelsLast2DFloatKernel(
         return false;
     }
 
-    if (Padding[0] != Padding[2] || Padding[1] != Padding[3]) {
+    // The channels-last float convolution path is only implemented by the Arm® KleidiAI™
+    // SME conv override, which currently uses a single padding value in its LHS indirection table.
+    // Require uniform padding until separate H/W padding is supported there.
+    if (Padding[0] != Padding[2] ||
+        Padding[1] != Padding[3] ||
+        Padding[0] != Padding[1]) {
         return false;
     }
 

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

@@ -7,7 +7,6 @@
 #pragma once
 
 #include "../mlasi.h"
-#include <iostream>
 
 // Fix to ensure compatibility with MSVC build
 #if defined(_MSC_VER)
@@ -25,6 +24,7 @@
 #endif
 
 #if KLEIDIAI_DEBUG_LOGGING ||KLEIDIAI_KERNEL_LOGGING
+#include <iostream>
 #define KLEIDIAI_LOG(tag, msg) \
     do { \
         std::cout << "[KLEIDIAI " << tag << "]: " << __FILE__ << " : " << __LINE__ << " : " << msg << std::endl; \
@@ -56,6 +56,143 @@ inline const bool UseSME2 = MLAS_CPUIDINFO::GetCPUIDInfo().HasArm_SME2();
 inline const bool UseSME = MLAS_CPUIDINFO::GetCPUIDInfo().HasArm_SME();
 inline const std::string_view vendor_name = MLAS_CPUIDINFO::GetCPUIDInfo().GetCPUVendor();
 
+// Selects the convolution route for Arm® KleidiAI™
+enum class ConvRoute {
+    NoKleidiAi,    // decline the conv, caller runs unchanged
+    IGemm,         // handle the whole conv via SME IGEMM kernel
+    SGemmFallback, // decline IGEMM, but still route the per-segment SGEMM slices through MlasGemm
+                   // so that the Arm® KleidiAI™ SGEMM backend override can pick them up
+};
+
+struct ConvRouteSelection {
+    ConvRoute route = ConvRoute::NoKleidiAi;
+    size_t effective_kernel_h = 0;
+    size_t effective_kernel_w = 0;
+};
+
+// Heuristic default for SME IGEMM selection. Work is estimated as
+// output_m * effective_k * filter_count; larger workloads use the SGEMM-backed
+// fallback to avoid IGEMM packing overhead on large effective-k convolutions.
+inline constexpr size_t ConvIgemmMaxWorkDefault = 1'000'000ULL;
+
+inline bool TryComputeDilatedKernelSize(size_t dilation, size_t kernel, size_t* result) {
+    if (dilation == 0 || kernel == 0) return false;
+
+    // using formula: dilated_kernel_size = dilation * (kernel - 1) + 1
+    size_t scaled_kernel;
+    if (MlasMultiplyOverflowsSizeT(kernel - 1, dilation, &scaled_kernel)) return false;
+
+    if (scaled_kernel == SIZE_MAX) return false;
+
+    *result = scaled_kernel + 1;
+    return true;
+}
+
+inline bool TryComputeConvOutputSize(size_t input, size_t kernel, size_t padding, size_t stride, size_t* result) {
+
+    if (stride == 0) return false;
+
+    size_t double_padding;
+    if (MlasMultiplyOverflowsSizeT(padding, 2, &double_padding)) return false;
+
+    if (double_padding > (SIZE_MAX - input)) return false;
+    const size_t padded_input = double_padding + input;
+
+    if (padded_input < kernel) return false;
+
+    // using formula: output_size = ((2*padding + input - kernel) / stride) + 1
+    const size_t output_minus_one = (padded_input - kernel) / stride;
+    if (output_minus_one == SIZE_MAX) return false;
+    *result = output_minus_one + 1;
+    return true;
+}
+
+inline ConvRouteSelection SelectConvRoute(const MLAS_CONV_PARAMETERS* Parameters) {
+    if ((Parameters->Dimensions != 2) ||
+        (Parameters->BatchCount != 1) ||
+        (Parameters->Beta != 0.f) ||
+        (Parameters->Padding[0] != Parameters->Padding[1]) ||
+        (Parameters->Padding[0] != Parameters->Padding[2]) ||
+        (Parameters->Padding[0] != Parameters->Padding[3])) {
+        return ConvRouteSelection{};
+    }
+
+    size_t effective_kernel_h;
+    size_t effective_kernel_w;
+    if (!TryComputeDilatedKernelSize(Parameters->DilationShape[0], Parameters->KernelShape[0], &effective_kernel_h) ||
+        !TryComputeDilatedKernelSize(Parameters->DilationShape[1], Parameters->KernelShape[1], &effective_kernel_w)) {
+        return ConvRouteSelection{};
+    }
+
+    size_t output_m;
+    size_t output_h_size;
+    size_t output_w_size;
+    if (!TryComputeConvOutputSize(Parameters->InputShape[0],
+                                  effective_kernel_h,
+                                  Parameters->Padding[0],
+                                  Parameters->StrideShape[0],
+                                  &output_h_size) ||
+        !TryComputeConvOutputSize(Parameters->InputShape[1],
+                                  effective_kernel_w,
+                                  Parameters->Padding[1],
+                                  Parameters->StrideShape[1],
+                                  &output_w_size) ||
+        MlasMultiplyOverflowsSizeT(output_h_size, output_w_size, &output_m)) {
+        return ConvRouteSelection{};
+    }
+
+    if (output_m == 0) {
+        return ConvRouteSelection{};
+    }
+
+    if (Parameters->KernelShape[0] < 3 || Parameters->KernelShape[1] < 3) {
+        return ConvRouteSelection{};
+    }
+
+    const auto filter_count = Parameters->FilterCount;
+    if (filter_count == 0) {
+        return ConvRouteSelection{};
+    }
+
+    size_t effective_k;
+    if (MlasMultiplyOverflowsSizeT(Parameters->InputChannels, effective_kernel_h, &effective_k) ||
+        MlasMultiplyOverflowsSizeT(effective_k, effective_kernel_w, &effective_k)) {
+        return ConvRouteSelection{};
+    }
+    if (effective_k == 0) {
+        return ConvRouteSelection{};
+    }
+
+    // Currently, the fallback routes assume NCHW layout, so keep valid NHWC convolutions on IGEMM
+    if (Parameters->ChannelsLast) {
+        return ConvRouteSelection{ConvRoute::IGemm, effective_kernel_h, effective_kernel_w};
+    }
+
+    if (filter_count == 1) {
+        // Single-output-channel convolutions do not fill the SME IGEMM N tile efficiently,
+        // so defer to the default convolution route.
+        return ConvRouteSelection{};
+    }
+
+    size_t total_work;
+    if (MlasMultiplyOverflowsSizeT(output_m, effective_k, &total_work) ||
+        MlasMultiplyOverflowsSizeT(total_work, filter_count, &total_work)) {
+        // Total work calculation overflow means total work is too large, fall back
+        return ConvRouteSelection{ConvRoute::SGemmFallback, effective_kernel_h, effective_kernel_w};
+    }
+
+    const size_t conv_igemm_max_work =
+        Parameters->BackendKernelSelectorConfig != nullptr &&
+        Parameters->BackendKernelSelectorConfig->kleidiai_conv_igemm_max_work != 0
+            ? Parameters->BackendKernelSelectorConfig->kleidiai_conv_igemm_max_work
+            : ConvIgemmMaxWorkDefault;
+    if (total_work > conv_igemm_max_work) {
+        return ConvRouteSelection{ConvRoute::SGemmFallback, effective_kernel_h, effective_kernel_w};
+    }
+
+    return ConvRouteSelection{ConvRoute::IGemm, effective_kernel_h, effective_kernel_w};
+}
+
 // Buffer packing routines.
 //
 size_t
@@ -224,38 +361,18 @@ MlasConvSymmetricChannelsLast2DFloatPackW(
     size_t PackedFilterGroupStride,
     MLAS_THREADPOOL* ThreadPool
     );
-}
-
-/*++
-
-Routine Description:
-
-    This routine determines if a wraparound will occur when multiplying two size_t variables
-    Uses __builtin_mul_overflow if available on the current system and if not falls back
-    to a default implementation to check this wraparound.
-
-Arguments:
-
-    a - Supplies the first number to be muliplied.
 
-    b - Supplies the second number to be muliplied.
-
-    out - pointer to a size_t which acts as the return value in success cases.
-
-Return Value:
-
-    Returns false if the operation was successful
-    Returns true if wraparound of size_t was detected
-
---*/
-inline bool mul_overflow_size_t_builtin(size_t a, size_t b, size_t* out) {
-#if defined(__has_builtin)
-#  if __has_builtin(__builtin_mul_overflow)
-    return __builtin_mul_overflow(a, b, out);
-#  endif
-#endif
-    // Fallback to manual check if builtin not available
-    if (b != 0 && a > SIZE_MAX / b) return true;
-    if (out) *out = a * b;
-    return false;
+inline MLAS_CONV_SGEMM_ROUTE
+MLASCALL
+MlasConvSGemmRoute(const MLAS_CONV_PARAMETERS* Parameters) {
+    if (Parameters->BackendKernelSelectorConfig &&
+        !Parameters->BackendKernelSelectorConfig->use_kleidiai) {
+        return MlasConvSGemmRouteDirect;
+    }
+
+    return ArmKleidiAI::SelectConvRoute(Parameters).route == ArmKleidiAI::ConvRoute::SGemmFallback
+        ? MlasConvSGemmRouteDispatch
+        : MlasConvSGemmRouteDirect;
 }
+}
+

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

@@ -279,7 +279,7 @@ Return Value:
     LhsPackedStride = kai_get_lhs_packed_size_lhs_pack_bf16p2vlx2_f32_sme(M, K, mr, kr, sr);
 
     size_t lhs_resize = 0;
-    if (mul_overflow_size_t_builtin(LhsPackedStride, BatchSize, &lhs_resize))
+    if (MlasMultiplyOverflowsSizeT(LhsPackedStride, BatchSize, &lhs_resize))
     {
         // size_t wraparound detected for LhsPackedStride, fallback to MLAS
         return false;
@@ -304,7 +304,7 @@ Return Value:
         // Multithread pack lhs and rhs
         RhsPackedStride = ArmKleidiAI::MlasSBGemmPackBSize(TransA, TransB, N, K);
         size_t rhs_resize = 0;
-        if (mul_overflow_size_t_builtin(RhsPackedStride, BatchSize, &rhs_resize))
+        if (MlasMultiplyOverflowsSizeT(RhsPackedStride, BatchSize, &rhs_resize))
         {
             // size_t wraparound detected for RhsPackedStride, fallback to MLAS
             return false;
@@ -354,7 +354,7 @@ Return Value:
     // Pre-check maximum tile size to avoid per-iteration overflow inside the parallel loop.
     // Any TileSizeM/TileSizeN used below will be <= m_step/n_step respectively.
     size_t max_tile_elems = 0;
-    if (mul_overflow_size_t_builtin(m_step, n_step, &max_tile_elems)) {
+    if (MlasMultiplyOverflowsSizeT(m_step, n_step, &max_tile_elems)) {
         // size_t wraparound detected for tile size, fallback to MLAS
         return false;
     }

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

@@ -542,7 +542,7 @@ Return Value:
     LhsPackedStride = kai_get_lhs_packed_size_lhs_pack_f32p2vlx1_f32_sme(M, K, mr, kr, sr);
 
     size_t lhs_resize = 0;
-    if(mul_overflow_size_t_builtin(LhsPackedStride, BatchSize, &lhs_resize))
+    if(MlasMultiplyOverflowsSizeT(LhsPackedStride, BatchSize, &lhs_resize))
     {
         // size_t wraparound detected for LhsPackedStride, fallback to MLAS
         return false;
@@ -568,7 +568,7 @@ Return Value:
         // Multithread pack lhs and rhs
         RhsPackedStride = ArmKleidiAI::MlasGemmPackBSize(TransA, TransB, N, K);
         size_t rhs_resize = 0;
-        if (mul_overflow_size_t_builtin(RhsPackedStride, BatchSize, &rhs_resize))
+        if (MlasMultiplyOverflowsSizeT(RhsPackedStride, BatchSize, &rhs_resize))
         {
             // size_t wraparound detected for RhsPackedStride, fallback to MLAS
             return false;
@@ -616,7 +616,7 @@ Return Value:
     // Pre-check maximum tile size to avoid per-iteration overflow inside the parallel loop.
     // Any TileSizeM/TileSizeN used below will be <= m_step/n_step respectively.
     size_t max_tile_elems = 0;
-    if (mul_overflow_size_t_builtin(m_step, n_step, &max_tile_elems)) {
+    if (MlasMultiplyOverflowsSizeT(m_step, n_step, &max_tile_elems)) {
         // size_t wraparound detected for tile size, fallback to MLAS
         return false;
     }

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

@@ -119,6 +119,33 @@ Module Name:
 
 #define MLAS_UNREFERENCED_PARAMETER(parameter) ((void)(parameter))
 
+//
+// Reports whether multiplying two size_t values would overflow.
+//
+
+MLAS_FORCEINLINE
+bool
+MlasMultiplyOverflowsSizeT(
+    size_t a,
+    size_t b,
+    size_t* out
+    )
+{
+#if defined(__has_builtin)
+#if __has_builtin(__builtin_mul_overflow)
+    size_t result;
+    return __builtin_mul_overflow(a, b, out != nullptr ? out : &result);
+#endif
+#endif
+    if (b != 0 && a > std::numeric_limits<size_t>::max() / b) {
+        return true;
+    }
+    if (out != nullptr) {
+        *out = a * b;
+    }
+    return false;
+}
+
 #ifdef MLAS_NO_EXCEPTION
 
 MLAS_FORCEINLINE void
@@ -891,6 +918,18 @@ bool
     MLAS_THREADPOOL* ThreadPool
     );
 
+enum MLAS_CONV_SGEMM_ROUTE {
+    MlasConvSGemmRouteDirect,   // call MlasSgemmOperation directly
+    MlasConvSGemmRouteDispatch, // call MlasGemm so backend SGEMM overrides may be selected
+};
+
+typedef
+MLAS_CONV_SGEMM_ROUTE
+(MLASCALL MLAS_CONV_SGEMM_ROUTE_OVERRIDE)(
+    const MLAS_CONV_PARAMETERS* Parameters
+    );
+
+
 typedef
 bool
 (MLASCALL MLAS_SGEMM_BATCH_OVERRIDE)(
@@ -1520,6 +1559,7 @@ struct MLAS_PLATFORM {
     // MLAS Conv overrides
     MLAS_CONV_PREPARE_FLOAT_OVERRIDE* MlasConvPrepareOverride = nullptr;
     MLAS_CONV_FLOAT_OVERRIDE* MlasConvOverride = nullptr;
+    MLAS_CONV_SGEMM_ROUTE_OVERRIDE* MlasConvSGemmRouteOverride = nullptr;
 #if defined(__aarch64__) && defined(__linux__)
     // SBGemm overrides
     MLAS_SBGEMM_BATCH_OVERRIDE* MlasSBGemmBatchOverride = nullptr;