From 0bfb660d02f877ea7619dd779cd941fe18b3af8d Mon Sep 17 00:00:00 2001
From: "Tung D. Le" <tung@jp.ibm.com>
Date: Sat, 8 Feb 2020 03:45:37 +0900
Subject: [PATCH] Import 2-argument Gemm as GemmNoBias (#68)

Co-authored-by: Gheorghe-Teodor Bercea <gt.bercea@gmail.com>
---
 doc/gen_doc.py                               |  1 +
 src/builder/frontend_dialect_transformer.cpp | 12 ++++++++++++
 src/builder/op_build_table.inc               |  2 +-
 3 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/doc/gen_doc.py b/doc/gen_doc.py
index bd126d5..2b150d6 100644
--- a/doc/gen_doc.py
+++ b/doc/gen_doc.py
@@ -36,6 +36,7 @@ special_attr_defaults = dict([
 special_op_handler = dict([
         ("Conv", "ImportNodeConv"),
         ("MaxPool", "ImportNodeMaxPool"),
+        ("Gemm", "ImportNodeGemm"),
         #("Transpose", "ImportNodeTranspose")
         ])
 
diff --git a/src/builder/frontend_dialect_transformer.cpp b/src/builder/frontend_dialect_transformer.cpp
index 7058cdc..10fc2c9 100644
--- a/src/builder/frontend_dialect_transformer.cpp
+++ b/src/builder/frontend_dialect_transformer.cpp
@@ -434,6 +434,18 @@ private:
     }
   }
 
+  /*!
+   * Special handle for Gemm operations.
+   */
+  void ImportNodeGemm(onnx::NodeProto node, int nIn, int nOut) {
+    int nOps = node.input().size();
+    if (nOps == 2) {
+      ImportNodeOneOut<mlir::ONNXGemmNoBiasOp>(node, 2, nOut);
+    } else {
+      ImportNodeOneOut<mlir::ONNXGemmOp>(node, nIn, nOut);
+    }
+  }
+
   void ImportNode(const onnx::NodeProto &node) {
     std::vector<mlir::Value> inputs;
     for (const auto &item : node.input()) {
diff --git a/src/builder/op_build_table.inc b/src/builder/op_build_table.inc
index 963b3b4..d7dea0f 100644
--- a/src/builder/op_build_table.inc
+++ b/src/builder/op_build_table.inc
@@ -98,7 +98,7 @@
     }else if (OpName == "GatherND") {
        ImportNodeOneOut<mlir::ONNXGatherNDOp>(node, 2, 1);
     }else if (OpName == "Gemm") {
-       ImportNodeOneOut<mlir::ONNXGemmOp>(node, 3, 1);
+       ImportNodeGemm(node, 3, 1);
     }else if (OpName == "GlobalAveragePool") {
        ImportNodeOneOut<mlir::ONNXGlobalAveragePoolOp>(node, 1, 1);
     }else if (OpName == "GlobalLpPool") {