[MLIR][Arith] Canonicalize cmpi of extui/extsi

[wsmoses]

Summary

Canonicalize cmpi(eq, ext a, ext b) and cmpi(ne, ext a, ext b)

[wsmoses]

Fix string literal split

[ftynse]

mlir/lib/Dialect/Arithmetic/IR/ArithmeticCanonicalization.td

119

Isn't type equality here guaranteed by cmpi semantics?

[wsmoses]

mlir/lib/Dialect/Arithmetic/IR/ArithmeticCanonicalization.td

119

Not necessarily, since you could exti(i8 -> i32) == exti(i16 -> i32), the two inner types may not be equal.

mlir/include/mlir/Dialect/Arithmetic/IR/ArithmeticOps.td

let builders = [
     }]>
   ];
 
   let extraClassDeclaration = [{
     static arith::CmpIPredicate getPredicateByName(StringRef name);
   }];
 
   let hasFolder = 1;
+  let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // CmpFOp
 //===----------------------------------------------------------------------===//
 
 def Arith_CmpFOp : Arith_CompareOp<"cmpf"> {
   let summary = "floating-point comparison operation";

mlir/lib/Dialect/Arithmetic/IR/ArithmeticCanonicalization.td

 def InvertPredicate : NativeCodeCall<"invertPredicate($0)">;
 def XOrINotCmpI :
     Pat<(Arith_XOrIOp
           (Arith_CmpIOp $pred, $a, $b),
           (Arith_ConstantOp ConstantAttr<I1Attr, "1">)),
         (Arith_CmpIOp (InvertPredicate $pred), $a, $b)>;
 
 //===----------------------------------------------------------------------===//
+// CmpIOp
+//===----------------------------------------------------------------------===//
+
+// cmpi(== or !=, a ext iNN, b ext iNN) == cmpi(== or !=, a, b)
+def CmpIExtSI :
+    Pat<(Arith_CmpIOp $pred,
+          (Arith_ExtSIOp $a),
+          (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)]>;
+
+// cmpi(== or !=, a ext iNN, b ext iNN) == cmpi(== or !=, a, b)
+def CmpIExtUI :
+    Pat<(Arith_CmpIOp $pred,
+          (Arith_ExtUIOp $a),
+          (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)]>;
+
+//===----------------------------------------------------------------------===//
 // IndexCastOp
 //===----------------------------------------------------------------------===//
 
 // index_cast(index_cast(x)) -> x, if dstType == srcType.
 def IndexCastOfIndexCast :
     Pat<(Arith_IndexCastOp:$res (Arith_IndexCastOp $x)),
         (replaceWithValue $x),
         [(Constraint<CPred<"$0.getType() == $1.getType()">> $res, $x)]>;

mlir/lib/Dialect/Arithmetic/IR/ArithmeticOps.cpp

OpFoldResult arith::CmpIOp::fold(ArrayRef<Attribute> operands) {
   auto rhs = operands.back().dyn_cast_or_null<IntegerAttr>();
   if (!lhs || !rhs)
     return {};
 
   auto val = applyCmpPredicate(getPredicate(), lhs.getValue(), rhs.getValue());
   return BoolAttr::get(getContext(), val);
 }
 
+void arith::CmpIOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
+                                                MLIRContext *context) {
+  patterns.insert<CmpIExtSI, CmpIExtUI>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // CmpFOp
 //===----------------------------------------------------------------------===//
 
 /// Compute `lhs` `pred` `rhs`, where `pred` is one of the known floating point
 /// comparison predicates.
 bool mlir::arith::applyCmpPredicate(arith::CmpFPredicate predicate,
                                     const APFloat &lhs, const APFloat &rhs) {

mlir/test/Dialect/Arithmetic/canonicalize.mlir

 func @extSIOfExtSI(%arg0: i1) -> i64 {
   %ext1 = arith.extsi %arg0 : i1 to i8
   %ext2 = arith.extsi %ext1 : i8 to i64
   return %ext2 : i64
 }
 
 // -----
 
+// CHECK-LABEL: @cmpIExtSINE
+//       CHECK:  %[[comb:.+]] = arith.cmpi ne, %arg0, %arg1 : i8
+//       CHECK:   return %[[comb]]
+func @cmpIExtSINE(%arg0: i8, %arg1: i8) -> i1 {
+  %ext0 = arith.extsi %arg0 : i8 to i64
+  %ext1 = arith.extsi %arg1 : i8 to i64
+  %res = arith.cmpi ne, %ext0, %ext1 : i64
+  return %res : i1
+}
+
+// CHECK-LABEL: @cmpIExtSIEQ
+//       CHECK:  %[[comb:.+]] = arith.cmpi eq, %arg0, %arg1 : i8
+//       CHECK:   return %[[comb]]
+func @cmpIExtSIEQ(%arg0: i8, %arg1: i8) -> i1 {
+  %ext0 = arith.extsi %arg0 : i8 to i64
+  %ext1 = arith.extsi %arg1 : i8 to i64
+  %res = arith.cmpi eq, %ext0, %ext1 : i64
+  return %res : i1
+}
+
+// CHECK-LABEL: @cmpIExtUINE
+//       CHECK:  %[[comb:.+]] = arith.cmpi ne, %arg0, %arg1 : i8
+//       CHECK:   return %[[comb]]
+func @cmpIExtUINE(%arg0: i8, %arg1: i8) -> i1 {
+  %ext0 = arith.extui %arg0 : i8 to i64
+  %ext1 = arith.extui %arg1 : i8 to i64
+  %res = arith.cmpi ne, %ext0, %ext1 : i64
+  return %res : i1
+}
+
+// CHECK-LABEL: @cmpIExtUIEQ
+//       CHECK:  %[[comb:.+]] = arith.cmpi eq, %arg0, %arg1 : i8
+//       CHECK:   return %[[comb]]
+func @cmpIExtUIEQ(%arg0: i8, %arg1: i8) -> i1 {
+  %ext0 = arith.extui %arg0 : i8 to i64
+  %ext1 = arith.extui %arg1 : i8 to i64
+  %res = arith.cmpi eq, %ext0, %ext1 : i64
+  return %res : i1
+}
+
+// -----
+
 // CHECK-LABEL: @andOfExtSI
 //       CHECK:  %[[comb:.+]] = arith.andi %arg0, %arg1 : i8
 //       CHECK:  %[[ext:.+]] = arith.extsi %[[comb]] : i8 to i64
 //       CHECK:   return %[[ext]]
 func @andOfExtSI(%arg0: i8, %arg1: i8) -> i64 {
   %ext0 = arith.extsi %arg0 : i8 to i64
   %ext1 = arith.extsi %arg1 : i8 to i64
   %res = arith.andi %ext0, %ext1 : i64