[MLIR][Arith] Canonicalize xor with ext

[Miss_Grape]

ping

[Miss_Grape]

ping

mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td

 // not(cmpi(pred, a, b)) -> cmpi(~pred, a, b), where not(x) is xori(x, 1)
 def InvertPredicate : NativeCodeCall<"invertPredicate($0)">;
 def XOrINotCmpI :
     Pat<(Arith_XOrIOp
           (Arith_CmpIOp $pred, $a, $b),
           (ConstantLikeMatcher ConstantAttr<I1Attr, "1">)),
         (Arith_CmpIOp (InvertPredicate $pred), $a, $b)>;
 
+// xor extui(x), extui(y) -> extui(xor(x,y))
+def XOrIOfExtUI :
+    Pat<(Arith_XOrIOp (Arith_ExtUIOp $x), (Arith_ExtUIOp $y)), (Arith_ExtUIOp (Arith_XOrIOp $x, $y)),
+      [(Constraint<CPred<"$0.getType() == $1.getType()">> $x, $y)]>;
+
+// xor extsi(x), extsi(y) -> extsi(xor(x,y))
+def XOrIOfExtSI :
+    Pat<(Arith_XOrIOp (Arith_ExtSIOp $x), (Arith_ExtSIOp $y)), (Arith_ExtSIOp (Arith_XOrIOp $x, $y)),
+      [(Constraint<CPred<"$0.getType() == $1.getType()">> $x, $y)]>;
+
 //===----------------------------------------------------------------------===//
 // CmpIOp
 //===----------------------------------------------------------------------===//
 
 // cmpi(== or !=, a ext iNN, b ext iNN) == cmpi(== or !=, a, b)
 def CmpIExtSI :
     Pat<(Arith_CmpIOp $pred,
           (Arith_ExtSIOp $a),

mlir/lib/Dialect/Arith/IR/ArithOps.cpp

OpFoldResult arith::XOrIOp::fold(ArrayRef<Attribute> operands) {
   }
 
   return constFoldBinaryOp<IntegerAttr>(
       operands, [](APInt a, const APInt &b) { return std::move(a) ^ b; });
 }
 
 void arith::XOrIOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                                 MLIRContext *context) {
-  patterns.add<XOrINotCmpI>(context);
+  patterns.add<XOrINotCmpI, XOrIOfExtUI, XOrIOfExtSI>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // NegFOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::NegFOp::fold(ArrayRef<Attribute> operands) {
   /// negf(negf(x)) -> x

mlir/test/Dialect/Arith/canonicalize.mlir

 //       CHECK:   return %arg0
 func.func @xorxor(%cmp : i1) -> i1 {
   %true = arith.constant true
   %ncmp = arith.xori %cmp, %true : i1
   %nncmp = arith.xori %ncmp, %true : i1
   return %nncmp : i1
 }
 
+// CHECK-LABEL: @xorOfExtSI
+//       CHECK:  %[[comb:.+]] = arith.xori %arg0, %arg1 : i8
+//       CHECK:  %[[ext:.+]] = arith.extsi %[[comb]] : i8 to i64
+//       CHECK:   return %[[ext]]
+func.func @xorOfExtSI(%arg0: i8, %arg1: i8) -> i64 {
+  %ext0 = arith.extsi %arg0 : i8 to i64
+  %ext1 = arith.extsi %arg1 : i8 to i64
+  %res = arith.xori %ext0, %ext1 : i64
+  return %res : i64
+}
+
+// CHECK-LABEL: @xorOfExtUI
+//       CHECK:  %[[comb:.+]] = arith.xori %arg0, %arg1 : i8
+//       CHECK:  %[[ext:.+]] = arith.extui %[[comb]] : i8 to i64
+//       CHECK:   return %[[ext]]
+func.func @xorOfExtUI(%arg0: i8, %arg1: i8) -> i64 {
+  %ext0 = arith.extui %arg0 : i8 to i64
+  %ext1 = arith.extui %arg1 : i8 to i64
+  %res = arith.xori %ext0, %ext1 : i64
+  return %res : i64
+}
+
 // -----
 
 // CHECK-LABEL: @bitcastSameType(
 // CHECK-SAME: %[[ARG:[a-zA-Z0-9_]*]]
 func.func @bitcastSameType(%arg : f32) -> f32 {
   // CHECK: return %[[ARG]]
   %res = arith.bitcast %arg : f32 to f32
   return %res : f32