From 7e2a6af82c887915500aac64d77f93f1e79e8d28 Mon Sep 17 00:00:00 2001
From: Jakub Kuderski <kubak@google.com>
Date: Tue, 13 Dec 2022 13:04:37 -0500
Subject: [PATCH] [mlir][arith] Clean up arith canonicalization td rules

- Use `ConstantLikeMatcher` instead of matching `Arith_ConstantOp`
- Reflow some defs to fit the line limit

Suggested by @jpienaar and @Mogball in https://reviews.llvm.org/D139778.

Reviewed By: Mogball

Differential Revision: https://reviews.llvm.org/D139941
---
 mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td | 66 ++++++++++++----------
 1 file changed, 37 insertions(+), 29 deletions(-)

diff --git a/mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td b/mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td
index cf2a767..a2f45c2 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td
+++ b/mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td
@@ -31,22 +31,22 @@ def SubIntAttrs : NativeCodeCall<"subIntegerAttrs($_builder, $0, $1, $2)">;
 // addi(addi(x, c0), c1) -> addi(x, c0 + c1)
 def AddIAddConstant :
     Pat<(Arith_AddIOp:$res
-          (Arith_AddIOp $x, (Arith_ConstantOp APIntAttr:$c0)),
-          (Arith_ConstantOp APIntAttr:$c1)),
+          (Arith_AddIOp $x, (ConstantLikeMatcher APIntAttr:$c0)),
+          (ConstantLikeMatcher APIntAttr:$c1)),
         (Arith_AddIOp $x, (Arith_ConstantOp (AddIntAttrs $res, $c0, $c1)))>;
 
 // addi(subi(x, c0), c1) -> addi(x, c1 - c0)
 def AddISubConstantRHS :
     Pat<(Arith_AddIOp:$res
-          (Arith_SubIOp $x, (Arith_ConstantOp APIntAttr:$c0)),
-          (Arith_ConstantOp APIntAttr:$c1)),
+          (Arith_SubIOp $x, (ConstantLikeMatcher APIntAttr:$c0)),
+          (ConstantLikeMatcher APIntAttr:$c1)),
         (Arith_AddIOp $x, (Arith_ConstantOp (SubIntAttrs $res, $c1, $c0)))>;
 
 // addi(subi(c0, x), c1) -> subi(c0 + c1, x)
 def AddISubConstantLHS :
     Pat<(Arith_AddIOp:$res
-          (Arith_SubIOp (Arith_ConstantOp APIntAttr:$c0), $x),
-          (Arith_ConstantOp APIntAttr:$c1)),
+          (Arith_SubIOp (ConstantLikeMatcher APIntAttr:$c0), $x),
+          (ConstantLikeMatcher APIntAttr:$c1)),
         (Arith_SubIOp (Arith_ConstantOp (AddIntAttrs $res, $c0, $c1)), $x)>;
 
 //===----------------------------------------------------------------------===//
@@ -67,43 +67,43 @@ def AddUIExtendedToAddI:
 // subi(addi(x, c0), c1) -> addi(x, c0 - c1)
 def SubIRHSAddConstant :
     Pat<(Arith_SubIOp:$res
-          (Arith_AddIOp $x, (Arith_ConstantOp APIntAttr:$c0)),
-          (Arith_ConstantOp APIntAttr:$c1)),
+          (Arith_AddIOp $x, (ConstantLikeMatcher APIntAttr:$c0)),
+          (ConstantLikeMatcher APIntAttr:$c1)),
         (Arith_AddIOp $x, (Arith_ConstantOp (SubIntAttrs $res, $c0, $c1)))>;
 
 // subi(c1, addi(x, c0)) -> subi(c1 - c0, x)
 def SubILHSAddConstant :
     Pat<(Arith_SubIOp:$res
-          (Arith_ConstantOp APIntAttr:$c1),
-          (Arith_AddIOp $x, (Arith_ConstantOp APIntAttr:$c0))),
+          (ConstantLikeMatcher APIntAttr:$c1),
+          (Arith_AddIOp $x, (ConstantLikeMatcher APIntAttr:$c0))),
         (Arith_SubIOp (Arith_ConstantOp (SubIntAttrs $res, $c1, $c0)), $x)>;
 
 // subi(subi(x, c0), c1) -> subi(x, c0 + c1)
 def SubIRHSSubConstantRHS :
     Pat<(Arith_SubIOp:$res
-          (Arith_SubIOp $x, (Arith_ConstantOp APIntAttr:$c0)),
-          (Arith_ConstantOp APIntAttr:$c1)),
+          (Arith_SubIOp $x, (ConstantLikeMatcher APIntAttr:$c0)),
+          (ConstantLikeMatcher APIntAttr:$c1)),
         (Arith_SubIOp $x, (Arith_ConstantOp (AddIntAttrs $res, $c0, $c1)))>;
 
 // subi(subi(c0, x), c1) -> subi(c0 - c1, x)
 def SubIRHSSubConstantLHS :
     Pat<(Arith_SubIOp:$res
-          (Arith_SubIOp (Arith_ConstantOp APIntAttr:$c0), $x),
-          (Arith_ConstantOp APIntAttr:$c1)),
+          (Arith_SubIOp (ConstantLikeMatcher APIntAttr:$c0), $x),
+          (ConstantLikeMatcher APIntAttr:$c1)),
         (Arith_SubIOp (Arith_ConstantOp (SubIntAttrs $res, $c0, $c1)), $x)>;
 
 // subi(c1, subi(x, c0)) -> subi(c0 + c1, x)
 def SubILHSSubConstantRHS :
     Pat<(Arith_SubIOp:$res
-          (Arith_ConstantOp APIntAttr:$c1),
-          (Arith_SubIOp $x, (Arith_ConstantOp APIntAttr:$c0))),
+          (ConstantLikeMatcher APIntAttr:$c1),
+          (Arith_SubIOp $x, (ConstantLikeMatcher APIntAttr:$c0))),
         (Arith_SubIOp (Arith_ConstantOp (AddIntAttrs $res, $c0, $c1)), $x)>;
 
 // subi(c1, subi(c0, x)) -> addi(x, c1 - c0)
 def SubILHSSubConstantLHS :
     Pat<(Arith_SubIOp:$res
-          (Arith_ConstantOp APIntAttr:$c1),
-          (Arith_SubIOp (Arith_ConstantOp APIntAttr:$c0), $x)),
+          (ConstantLikeMatcher APIntAttr:$c1),
+          (Arith_SubIOp (ConstantLikeMatcher APIntAttr:$c0), $x)),
         (Arith_AddIOp $x, (Arith_ConstantOp (SubIntAttrs $res, $c1, $c0)))>;
 
 // subi(subi(a, b), a) -> subi(0, b)
@@ -145,7 +145,7 @@ def InvertPredicate : NativeCodeCall<"invertPredicate($0)">;
 def XOrINotCmpI :
     Pat<(Arith_XOrIOp
           (Arith_CmpIOp $pred, $a, $b),
-          (Arith_ConstantOp ConstantAttr<I1Attr, "1">)),
+          (ConstantLikeMatcher ConstantAttr<I1Attr, "1">)),
         (Arith_CmpIOp (InvertPredicate $pred), $a, $b)>;
 
 //===----------------------------------------------------------------------===//
@@ -159,7 +159,9 @@ def CmpIExtSI :
           (Arith_ExtSIOp $b)),
         (Arith_CmpIOp $pred, $a, $b),
         [(Constraint<CPred<"$0.getType() == $1.getType()">> $a, $b),
-         (Constraint<CPred<"$0.getValue() == arith::CmpIPredicate::eq || $0.getValue() == arith::CmpIPredicate::ne">> $pred)]>;
+         (Constraint<
+            CPred<"$0.getValue() == arith::CmpIPredicate::eq || "
+                  "$0.getValue() == arith::CmpIPredicate::ne">> $pred)]>;
 
 // cmpi(== or !=, a ext iNN, b ext iNN) == cmpi(== or !=, a, b)
 def CmpIExtUI :
@@ -168,7 +170,9 @@ def CmpIExtUI :
           (Arith_ExtUIOp $b)),
         (Arith_CmpIOp $pred, $a, $b),
         [(Constraint<CPred<"$0.getType() == $1.getType()">> $a, $b),
-         (Constraint<CPred<"$0.getValue() == arith::CmpIPredicate::eq || $0.getValue() == arith::CmpIPredicate::ne">> $pred)]>;
+         (Constraint<
+            CPred<"$0.getValue() == arith::CmpIPredicate::eq || "
+                  "$0.getValue() == arith::CmpIPredicate::ne">> $pred)]>;
 
 //===----------------------------------------------------------------------===//
 // IndexCastOp
@@ -221,13 +225,15 @@ def ExtSIOfExtUI :
 
 // and extui(x), extui(y) -> extui(and(x,y))
 def AndOfExtUI :
-    Pat<(Arith_AndIOp (Arith_ExtUIOp $x), (Arith_ExtUIOp $y)), (Arith_ExtUIOp (Arith_AndIOp $x, $y)),
-      [(Constraint<CPred<"$0.getType() == $1.getType()">> $x, $y)]>;
+    Pat<(Arith_AndIOp (Arith_ExtUIOp $x), (Arith_ExtUIOp $y)),
+        (Arith_ExtUIOp (Arith_AndIOp $x, $y)),
+        [(Constraint<CPred<"$0.getType() == $1.getType()">> $x, $y)]>;
 
 // and extsi(x), extsi(y) -> extsi(and(x,y))
 def AndOfExtSI :
-    Pat<(Arith_AndIOp (Arith_ExtSIOp $x), (Arith_ExtSIOp $y)), (Arith_ExtSIOp (Arith_AndIOp $x, $y)),
-      [(Constraint<CPred<"$0.getType() == $1.getType()">> $x, $y)]>;
+    Pat<(Arith_AndIOp (Arith_ExtSIOp $x), (Arith_ExtSIOp $y)),
+        (Arith_ExtSIOp (Arith_AndIOp $x, $y)),
+        [(Constraint<CPred<"$0.getType() == $1.getType()">> $x, $y)]>;
 
 //===----------------------------------------------------------------------===//
 // OrIOp
@@ -235,13 +241,15 @@ def AndOfExtSI :
 
 // or extui(x), extui(y) -> extui(or(x,y))
 def OrOfExtUI :
-    Pat<(Arith_OrIOp (Arith_ExtUIOp $x), (Arith_ExtUIOp $y)), (Arith_ExtUIOp (Arith_OrIOp $x, $y)),
-      [(Constraint<CPred<"$0.getType() == $1.getType()">> $x, $y)]>;
+    Pat<(Arith_OrIOp (Arith_ExtUIOp $x), (Arith_ExtUIOp $y)),
+        (Arith_ExtUIOp (Arith_OrIOp $x, $y)),
+        [(Constraint<CPred<"$0.getType() == $1.getType()">> $x, $y)]>;
 
 // or extsi(x), extsi(y) -> extsi(or(x,y))
 def OrOfExtSI :
-    Pat<(Arith_OrIOp (Arith_ExtSIOp $x), (Arith_ExtSIOp $y)), (Arith_ExtSIOp (Arith_OrIOp $x, $y)),
-      [(Constraint<CPred<"$0.getType() == $1.getType()">> $x, $y)]>;
+    Pat<(Arith_OrIOp (Arith_ExtSIOp $x), (Arith_ExtSIOp $y)),
+        (Arith_ExtSIOp (Arith_OrIOp $x, $y)),
+        [(Constraint<CPred<"$0.getType() == $1.getType()">> $x, $y)]>;
 
 //===----------------------------------------------------------------------===//
 // MulFOp
-- 
2.7.4