From fb9ce18e88b9fe3c7bb6f98e80587aad522f58da Mon Sep 17 00:00:00 2001
From: Edward Chen <18449977+edgchen1@users.noreply.github.com>
Date: Wed, 21 Aug 2024 09:15:58 -0700
Subject: [PATCH] Add K=0 check to MatMul<float>::Compute() specialization.
 (#21803)

Add K=0 check to `MatMul<float>::Compute()` specialization.
Add unit test to cover both primary template and float specialization.
---
 onnxruntime/core/providers/cpu/math/matmul.cc |  8 ++++++
 .../test/providers/cpu/math/matmul_test.cc    | 26 ++++++++++++-------
 2 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/onnxruntime/core/providers/cpu/math/matmul.cc b/onnxruntime/core/providers/cpu/math/matmul.cc
index e95ac44af1..6a71283f9d 100644
--- a/onnxruntime/core/providers/cpu/math/matmul.cc
+++ b/onnxruntime/core/providers/cpu/math/matmul.cc
@@ -238,6 +238,14 @@ Status MatMul<float>::Compute(OpKernelContext* ctx) const {
   if (y->Shape().Size() == 0)
     return Status::OK();
 
+  if (helper.K() == 0) {
+    // When we have (M, 0, N) then the inputs are empty, but the output should
+    // be filled out with zeros.
+    auto output_span = y->MutableDataAsSpan<float>();
+    std::fill(output_span.begin(), output_span.end(), float{});
+    return Status::OK();
+  }
+
   const auto* a_data = a->Data<float>();
   const auto* b_data = b ? b->Data<float>() : nullptr;
   auto* y_data = y->MutableData<float>();
diff --git a/onnxruntime/test/providers/cpu/math/matmul_test.cc b/onnxruntime/test/providers/cpu/math/matmul_test.cc
index b7ae0a9f0d..9037056085 100644
--- a/onnxruntime/test/providers/cpu/math/matmul_test.cc
+++ b/onnxruntime/test/providers/cpu/math/matmul_test.cc
@@ -219,18 +219,16 @@ TEST(MathOpTest, MatMulUint64Type) {
   RunMatMulTest<uint64_t>(9);
 }
 
-TEST(MathOpTest, MatMul_ZeroK) {
+template <typename T>
+void RunMatMulZeroKTest() {
   // test with empty inputs and zero filled output
-  constexpr const std::array<float, 0> empty_input{};
-  const std::vector<float> expected_output{0, 0, 0, 0,
-                                           0, 0, 0, 0,
-                                           0, 0, 0, 0,
-                                           0, 0, 0, 0};
-  OpTester test("MatMul", 14);
+  constexpr const std::array<T, 0> empty_input{};
+  const std::vector<T> expected_output(4 * 4, T{});
+  OpTester test("MatMul", 13);
 
-  test.AddInput<float>("A", {4, 0}, empty_input);
-  test.AddInput<float>("B", {0, 4}, empty_input);
-  test.AddOutput<float>("Y", {4, 4}, expected_output);
+  test.AddInput<T>("A", {4, 0}, empty_input);
+  test.AddInput<T>("B", {0, 4}, empty_input);
+  test.AddOutput<T>("Y", {4, 4}, expected_output);
 
   // No special case is implemented.
   test.ConfigExcludeEps({kCoreMLExecutionProvider, kNnapiExecutionProvider,
@@ -240,6 +238,14 @@ TEST(MathOpTest, MatMul_ZeroK) {
       .RunWithConfig();
 }
 
+TEST(MathOpTest, MatMulZeroKFloatType) {
+  RunMatMulZeroKTest<float>();
+}
+
+TEST(MathOpTest, MatMulZeroKInt32Type) {
+  RunMatMulZeroKTest<int32_t>();
+}
+
 #if defined(USE_CUDA) || defined(USE_ROCM)
 TEST(MathOpTest, MatMul_Float16) {
 #ifdef USE_CUDA