Validate HalfGemm packed-B edge cases

- Add validation for null HalfGemm data parameters and reject invalid backend-native packed-B parameter combinations before the zero-K fast path.
- Make HalfGemm packed-B size calculation return a valid zero size for degenerate N/K shapes.
- Restore the KleidiAI header include needed by debug logging in mlasi_kleidiai.h.
- Update HalfGemm PackB tests to avoid exception-only gtest macros in no-exception builds, and add sentinel checks for early-return paths.

Signed-off-by: Cathal Lawlor <cathal.lawlor@arm.com>

Author: Laan33

onnxruntime/core/mlas/lib/halfgemm.cpp

@@ -92,7 +92,6 @@ TryGetHalfGemmBackendSelectorConfig(
 }
 #endif
 
-
 void
 MLASCALL
 MlasHalfGemmBatch(
@@ -108,7 +107,21 @@ MlasHalfGemmBatch(
         return;
     }
 
+    if (DataParams == nullptr) {
+        MLAS_THROW_EX(std::runtime_error, "HalfGemm requires non-null DataParams.");
+    }
+
     if (K == 0) {
+        for (size_t gemm_i = 0; gemm_i < BatchN; gemm_i++) {
+            const auto& data = DataParams[gemm_i];
+            if (data.BIsBackendNativePacked &&
+                (data.ldb != 0 || data.Bias != nullptr || data.OutputProcessor != nullptr)) {
+                MLAS_THROW_EX(
+                    std::runtime_error,
+                    "backend-native halfgemm packed B with K==0 requires ldb == 0, Bias == nullptr, and "
+                    "OutputProcessor == nullptr");
+            }
+        }
         MlasHalfGemmZeroKBatch(M, N, BatchN, DataParams);
         return;
     }

onnxruntime/core/mlas/lib/halfgemm.h

@@ -61,6 +61,11 @@ MlasHalfGemmTryGetPackedBSize(
     size_t* PackedBSize
     )
 {
+    if (N == 0 || K == 0) {
+        *PackedBSize = 0;
+        return true;
+    }
+
     if (PackedK == 0) {
         return false;
     }

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

@@ -8,6 +8,7 @@
 
 #include "../mlasi.h"
 #include <limits>
+#include <iostream>
 #include <vector>
 
 // Fix to ensure compatibility with MSVC build

onnxruntime/test/mlas/unittest/test_halfgemm.cpp

@@ -33,6 +33,13 @@ struct HalfGemmPackBPaddingKernel {
   static constexpr size_t PackedK = 4;
 };
 
+void ExpectBufferFilledWith(const std::vector<std::byte>& buffer, std::byte expected) {
+  const auto expected_value = std::to_integer<unsigned int>(expected);
+  for (size_t i = 0; i < buffer.size(); ++i) {
+    EXPECT_EQ(std::to_integer<unsigned int>(buffer[i]), expected_value) << "index=" << i;
+  }
+}
+
 }  // namespace
 
 #if defined(USE_KLEIDIAI)
@@ -487,25 +494,33 @@ TEST(HalfGemmPackB, ReturnsZeroOnOverflow) {
   EXPECT_EQ(MlasHalfGemmPackBSize(max, 2, true), size_t{0});
 }
 
-TEST(HalfGemmPackB, PackBReturnsOnOverflow) {
+TEST(HalfGemmPackB, PackBLeavesOutputUnchangedOnOverflow) {
   const size_t max = (std::numeric_limits<size_t>::max)();
   std::vector<MLAS_FP16> b(1);
-  std::vector<std::byte> packed_b(1);
+  constexpr std::byte sentinel{0x5A};
+  std::vector<std::byte> packed_b(16, sentinel);
 
-  EXPECT_NO_THROW(MlasHalfGemmPackB(max, 2, b.data(), max, packed_b.data()));
+  MlasHalfGemmPackB(max, 2, b.data(), max, packed_b.data());
+  ExpectBufferFilledWith(packed_b, sentinel);
 }
 
-TEST(HalfGemmPackB, PackBRejectsInvalidArguments) {
+TEST(HalfGemmPackB, PackBExitsEarlyForInvalidArguments) {
   constexpr size_t N = 5;
   constexpr size_t K = 3;
+  constexpr std::byte sentinel{0x5A};
 
   std::vector<MLAS_FP16> b(K * N);
-  std::vector<std::byte> packed_b(MlasHalfGemmPackBSize(N, K, false));
+  std::vector<std::byte> packed_b(MlasHalfGemmPackBSize(N, K, false), sentinel);
   ASSERT_FALSE(packed_b.empty());
 
-  EXPECT_NO_THROW(MlasHalfGemmPackB(N, K, b.data(), N - 1, packed_b.data()));
-  EXPECT_NO_THROW(MlasHalfGemmPackB(N, K, nullptr, N, packed_b.data()));
-  EXPECT_NO_THROW(MlasHalfGemmPackB(N, K, b.data(), N, nullptr));
+  MlasHalfGemmPackB(N, K, b.data(), N - 1, packed_b.data());
+  ExpectBufferFilledWith(packed_b, sentinel);
+
+  MlasHalfGemmPackB(N, K, nullptr, N, packed_b.data());
+  ExpectBufferFilledWith(packed_b, sentinel);
+
+  MlasHalfGemmPackB(N, K, b.data(), N, nullptr);
+  ExpectBufferFilledWith(packed_b, sentinel);
 }
 
 TEST(HalfGemmPackB, CopyPackBZeroPadsAlignedKTail) {
@@ -746,7 +761,7 @@ TEST(HalfGemmKleidiAIPath, PackedBFloatSingleThreadVariedShapesAndBiasWithoutOut
 // KleidiAI-specific packed-B uses a separate direct-consumption contract from
 // generic halfgemm PackB. Unsupported combinations fail at the public API
 // boundary because generic MLAS cannot consume this backend-native layout.
-TEST(HalfGemmKleidiAIPath, KleidiAIPackedBWithBiasThrows) {
+TEST(HalfGemmKleidiAIPath, KleidiAIPackedBWithBiasIsRejected) {
   if (!MlasFp16AccelerationSupported()) {
     GTEST_SKIP();
   }
@@ -790,7 +805,7 @@ TEST(HalfGemmKleidiAIPath, KleidiAIPackedBWithBiasThrows) {
 #endif
 }
 
-TEST(HalfGemmKleidiAIPath, KleidiAIPackedBWithOutputProcessorThrows) {
+TEST(HalfGemmKleidiAIPath, KleidiAIPackedBWithOutputProcessorIsRejected) {
   if (!MlasFp16AccelerationSupported()) {
     GTEST_SKIP();
   }
@@ -837,7 +852,7 @@ TEST(HalfGemmKleidiAIPath, KleidiAIPackedBWithOutputProcessorThrows) {
 #endif
 }
 
-TEST(HalfGemmKleidiAIPath, ZeroKFallsBack) {
+TEST(HalfGemmKleidiAIPath, ZeroKIsNotHandledByKleidiAIOverride) {
   if (GetMlasPlatform().MlasHalfGemmBatchOverride == nullptr) {
     GTEST_SKIP() << "KleidiAI halfgemm override unavailable";
   }