Address review comments

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

Author: JonathanC-ARM

onnxruntime/core/mlas/inc/mlas.h

@@ -893,6 +893,10 @@ struct MLAS_CONV_PARAMETERS {
     size_t BatchCount;
     size_t GroupCount;
     size_t InputChannels;
+    /**
+     * Existing Conv layout flag from MlasConvPrepare; for HalfConv this matches
+     * InputOutputChannelsLast.
+     */
     bool ChannelsLast;
     size_t InputShape[3];
     size_t KernelShape[3];
@@ -911,6 +915,11 @@ struct MLAS_CONV_PARAMETERS {
     const void* PackedFilter = nullptr;
     size_t PackedFilterGroupStride = 0;
     bool FilterIsPacked = false;
+    /**
+     * HalfConv-specific: true means caller-provided input and output tensors
+     * are channels-last (NHWC); false means NCHW public tensors with any NHWC
+     * staging handled internally.
+     */
     bool InputOutputChannelsLast = false;
     union {
         struct {
@@ -1952,6 +1961,11 @@ struct MLAS_HALF_GEMM_DATA_PARAMS {
  *        Backend-native packed B is a constrained direct-consumption layout
  *        and does not support runtime Bias or OutputProcessor.
  *
+ * Uses MLAS_THROW_EX(std::runtime_error, ...) for contract violations that
+ * cannot be safely ignored: non-empty work with null DataParams, malformed
+ * backend-native packed-B parameters on the K == 0 path, or backend-native
+ * packed B reaching the generic MLAS kernels when no backend override consumes it.
+ *
  * Note:  We only support uniform batching, so shapes and types of the
  *        input must be same across all parameter blocks.
  *
@@ -1996,6 +2010,9 @@ MlasHalfGemmPackBSize(
  * @brief For half precision GEMM, pack the right hand
  *        side matrix B
  *
+ * @pre For non-zero N and K, B and PackedB must be non-null and ldb must be at
+ *      least N.
+ *
  * @param[in]  N        Number of columns
  * @param[in]  K        Number of rows
  * @param[in]  B        Address of matrix B

onnxruntime/core/mlas/lib/halfgemm.cpp

@@ -19,9 +19,6 @@ Module Name:
 #include "mlas_float16.h"
 
 #include "halfgemm.h"
-#if defined(USE_KLEIDIAI)
-#include "kleidiai/mlasi_kleidiai.h"
-#endif
 
 #include <exception>
 
@@ -63,8 +60,7 @@ MlasFp16AccelerationSupported()
 #endif
 }
 
-#if defined(USE_KLEIDIAI)
-bool
+static bool
 TryGetHalfGemmBackendSelectorConfig(
     size_t BatchN,
     const MLAS_HALF_GEMM_DATA_PARAMS* DataParams,
@@ -90,7 +86,6 @@ TryGetHalfGemmBackendSelectorConfig(
 
     return true;
 }
-#endif
 
 void
 MLASCALL
@@ -126,15 +121,13 @@ MlasHalfGemmBatch(
         return;
     }
 
-#if defined(USE_KLEIDIAI)
     const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig = nullptr;
     if (TryGetHalfGemmBackendSelectorConfig(BatchN, DataParams, BackendKernelSelectorConfig) &&
         GetMlasPlatform().MlasHalfGemmBatchOverride != nullptr &&
         GetMlasPlatform().MlasHalfGemmBatchOverride(
             M, N, K, BatchN, DataParams, ThreadPool, BackendKernelSelectorConfig)) {
         return;
     }
-#endif
 
     // BIsPacked denotes the generic MLAS halfgemm packed-B layout and can be
     // consumed here. Backend-native packed layouts are separate and must not
@@ -251,6 +244,7 @@ MlasHalfGemmPackB(
     )
 {
     const auto* dispatch = MlasHalfGemmGetDispatch();
+    assert(N == 0 || K == 0 || (B != nullptr && PackedB != nullptr && ldb >= N));
     if (B == nullptr || PackedB == nullptr || ldb < N ||
         dispatch->CopyPackBRoutine == nullptr || MlasHalfGemmPackBSize(N, K, false) == 0) {
         return;
@@ -269,18 +263,11 @@ MlasHalfGemmNativePackBSize(
     const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig
     )
 {
-#if defined(USE_KLEIDIAI)
     if ((!BackendKernelSelectorConfig || BackendKernelSelectorConfig->use_kleidiai) &&
         GetMlasPlatform().MlasHalfGemmPackBSizeOverride != nullptr) {
         return GetMlasPlatform().MlasHalfGemmPackBSizeOverride(TransA, TransB, N, K);
     }
-#else
-    MLAS_UNREFERENCED_PARAMETER(TransA);
-    MLAS_UNREFERENCED_PARAMETER(TransB);
-    MLAS_UNREFERENCED_PARAMETER(N);
-    MLAS_UNREFERENCED_PARAMETER(K);
-    MLAS_UNREFERENCED_PARAMETER(BackendKernelSelectorConfig);
-#endif
+
     return 0;
 }
 
@@ -297,21 +284,11 @@ MlasHalfGemmNativePackB(
     const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig
     )
 {
-#if defined(USE_KLEIDIAI)
     if ((!BackendKernelSelectorConfig || BackendKernelSelectorConfig->use_kleidiai) &&
         GetMlasPlatform().MlasHalfGemmPackBOverride != nullptr) {
         return GetMlasPlatform().MlasHalfGemmPackBOverride(TransA, TransB, N, K, B, ldb, PackedB);
     }
-#else
-    MLAS_UNREFERENCED_PARAMETER(TransA);
-    MLAS_UNREFERENCED_PARAMETER(TransB);
-    MLAS_UNREFERENCED_PARAMETER(N);
-    MLAS_UNREFERENCED_PARAMETER(K);
-    MLAS_UNREFERENCED_PARAMETER(B);
-    MLAS_UNREFERENCED_PARAMETER(ldb);
-    MLAS_UNREFERENCED_PARAMETER(PackedB);
-    MLAS_UNREFERENCED_PARAMETER(BackendKernelSelectorConfig);
-#endif
+
     return false;
 }
 

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

@@ -460,8 +460,8 @@ static std::shared_ptr<const void*[]> GetOrCreateLhsPtrTableSme(const size_t ci,
     };
 
     // Cache of computed LHS pointer offsets. thread_local to prevent interference from parallel sessions.
-    // Entries include pointers to the pad buffer for out-of-bounds pixels, so group by pad buffer identity and
-    // invalidate the old group if the grow-only pad buffer is reallocated.
+    // Entries include pointers to the pad buffer for out-of-bounds pixels; if the grow-only pad buffer
+    // reallocates, erase the old group so stale pointer tables cannot reference the previous allocation.
     using LhsPtrsCache = std::unordered_map<LhsCacheKey, std::shared_ptr<const void*[]>>;
     thread_local std::unordered_map<const float*, LhsPtrsCache> lhs_ptrs_cache_by_pad;
     thread_local const float* last_pad_ptr = nullptr;

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

@@ -177,6 +177,8 @@ ArmKleidiAI::MlasHalfGemmBatch(
         return false;
     }
 
+    // TODO: Plumb MLAS_ACTIVATION through this call site if MLAS_HALF_GEMM_DATA_PARAMS
+    // grows fused activation support.
     const float clamp_min = -std::numeric_limits<float>::infinity();
     const float clamp_max = std::numeric_limits<float>::infinity();
 

onnxruntime/core/mlas/lib/mlasi.h

@@ -101,26 +101,6 @@ Module Name:
 #define MLAS_FORCEINLINE __attribute__ ((always_inline)) inline
 #endif
 
-MLAS_FORCEINLINE
-bool
-MlasMultiplyOverflowsSizeT(
-    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
-    if (b != 0 && a > (std::numeric_limits<size_t>::max)() / b) {
-        return true;
-    }
-    *out = a * b;
-    return false;
-}
-
 MLAS_FORCEINLINE
 bool
 MlasAddOverflowsSizeT(

onnxruntime/test/contrib_ops/matmul_4bits_test.cc

@@ -307,7 +307,10 @@ void TestMatMulNBitsTyped(std::optional<float> abs_error = std::nullopt,
   } else if (base_opts.accuracy_level == 4) {
     base_opts.output_abs_error = 0.1f;
   } else if constexpr (std::is_same<AType, MLFloat16>::value) {
+    // The fp16 provider paths compare against a float reference while native kernels may accumulate
+    // in fp16 (for example native HGEMM on SME; see PR #28786), so allow slightly wider drift.
 #if defined(USE_WEBGPU)
+    // WebGPU's fp16 path has additional provider-specific rounding drift for these quantized matmul cases.
     base_opts.output_abs_error = 0.1f;
 #else
     base_opts.output_abs_error = 0.065f;