From 0c36057404f7cc091eeac92020010a01f023c753 Mon Sep 17 00:00:00 2001
From: TUNG LEDUC <TUNG@jp.ibm.com>
Date: Tue, 19 Nov 2019 11:23:46 +0900
Subject: [PATCH] [MLIR] Add a constraint for the rewriting rule of add_matmult
 to fullgemm (#374)

* Add a constraint for the rewriting rule of matmul_add to fullgemm

* Edit comments
---
 src/compiler/pass/onnx_combine.td | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/src/compiler/pass/onnx_combine.td b/src/compiler/pass/onnx_combine.td
index bfd7905..ceb3c7e 100644
--- a/src/compiler/pass/onnx_combine.td
+++ b/src/compiler/pass/onnx_combine.td
@@ -24,12 +24,15 @@ include "dialect/onnx/onnx.td"
 ///    dag benefitsAdded = (addBenefit 0)
 /// >;
 
+def HasOneUse : Constraint<CPred<"$0->hasOneUse()">>;
+
 //===----------------------------------------------------------------------===//
 // Pattern-Match and Rewrite
 //===----------------------------------------------------------------------===//
 
-// onnx.add(onnx.matmult(%X, %Y), %Z) = onnx.gemm(%X, %Y, %Z)
-def MulAddToGemmOptPattern : Pat<(ONNXAddOp(ONNXMatMulOp $m1, $m2), $m3),
-                                 (ONNXFullGemmOp $m1, $m2, $m3)>;
+// onnx.add(onnx.matmul(%X, %Y), %Z) = onnx.full_gemm(%X, %Y, %Z)
+def MulAddToGemmOptPattern : Pat<(ONNXAddOp (ONNXMatMulOp:$res $m1, $m2), $m3),
+                                 (ONNXFullGemmOp $m1, $m2, $m3),
+				 [(HasOneUse $res)]>;
 
 #endif // ONNX_COMBINE