From 373177ddd3f438698ebc9e84bd9f2932d4ae5a9f Mon Sep 17 00:00:00 2001
From: "Tang, Cheng" <souptc@gmail.com>
Date: Thu, 24 Jan 2019 18:24:32 -0800
Subject: [PATCH] gemm with empty input (#350)

* gemm with empty input

* check negative

* align the comments
---
 onnxruntime/core/providers/cpu/math/gemm.h      |  3 +++
 .../core/providers/cpu/math/gemm_helper.h       |  3 ++-
 .../test/providers/cpu/math/gemm_test.cc        | 17 +++++++++++++++++
 3 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/onnxruntime/core/providers/cpu/math/gemm.h b/onnxruntime/core/providers/cpu/math/gemm.h
index c060685b4f..c72c5bc1e0 100644
--- a/onnxruntime/core/providers/cpu/math/gemm.h
+++ b/onnxruntime/core/providers/cpu/math/gemm.h
@@ -42,6 +42,9 @@ class Gemm : public OpKernel {
     int64_t N = helper.N();
     int64_t K = helper.K();
     auto Y = context->Output(0, TensorShape({M, N}));
+    // if input is emtpy tensor, return directly as nothing need to be calculated.
+    if (M == 0 || N == 0)
+      return Status::OK();
     T_Y* y_data = Y->template MutableData<T_Y>();
 
     //bias
diff --git a/onnxruntime/core/providers/cpu/math/gemm_helper.h b/onnxruntime/core/providers/cpu/math/gemm_helper.h
index 6c0e7c22d6..46bd5e734f 100644
--- a/onnxruntime/core/providers/cpu/math/gemm_helper.h
+++ b/onnxruntime/core/providers/cpu/math/gemm_helper.h
@@ -40,7 +40,8 @@ class GemmHelper {
     if (!IsValidBroadcast(bias, M_, N_))
       status_ = common::Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, "Gemm: Invalid bias shape for broadcast");
 
-    ORT_ENFORCE(M_ > 0 && N_ > 0 && K_ > 0);
+    // it is possible the input is empty tensor, for example the output of roipool in fast rcnn.
+    ORT_ENFORCE(M_ >= 0 && K_ > 0 && N_ >= 0);
   }
 
   int64_t M() const { return M_; }
diff --git a/onnxruntime/test/providers/cpu/math/gemm_test.cc b/onnxruntime/test/providers/cpu/math/gemm_test.cc
index c915f2215f..8c63564a86 100644
--- a/onnxruntime/test/providers/cpu/math/gemm_test.cc
+++ b/onnxruntime/test/providers/cpu/math/gemm_test.cc
@@ -161,5 +161,22 @@ TEST(MathOpTest, GemmFalseBroadcast) {
   test.Run();
 }
 
+TEST(MathOpTest, GemmEmptyTensor) {
+  OpTester test("Gemm");
+
+  test.AddAttribute("transA", static_cast<int64_t>(0));
+  test.AddAttribute("transB", static_cast<int64_t>(0));
+  test.AddAttribute("alpha", 1.0f);
+  test.AddAttribute("beta", 1.0f);
+
+  test.AddInput<float>("A", {0, 4},
+                       {});
+  test.AddInput<float>("B", {4, 3}, std::vector<float>(12, 1.0f));
+  test.AddInput<float>("C", {3}, std::vector<float>(3, 1.0f));
+  test.AddOutput<float>("Y", {0, 3},
+                        {});
+  test.Run();
+}
+
 }  // namespace test
 }  // namespace onnxruntime