[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
This commit is contained in:
TUNG LEDUC 2019-11-19 11:23:46 +09:00 committed by Tian Jin
parent 701c7d7508
commit 0c36057404
1 changed files with 6 additions and 3 deletions

View File

@ -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