[MLIR][Arith] Fold repeated xor and trunc

[wsmoses]

Summary

This patch adds two folds. One for a repeated xor (e.g. xor(xor(x, a), a)) and one for a repeated trunc (e.g. trunc(trunc(x))).

[mehdi_amini]

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

117 ↗

(On Diff #396558)

We should match any constant instead of directly Arith_ConstantOp, but more importantly I think that a reassociation like A(B(x, c0), c1) = A(x, B(c1, c2)) deserves to be generic and implemented with a "reassociative" trait.

147 ↗

(On Diff #396558)

Seems like this could be a fold?
Fold is allowed to mutate the current op in-place instead of returning a constant.

[wsmoses]

Change to folders

[wsmoses]

Fix rebase error

[ftynse]

This more focused version LGTM. The "Reassociative" trait suggested by Mehdi is still interesting as a separate patch.

[mehdi_amini]

Updated version LGTM as well :)

[bondhugula]

Can you please address these? Looks like these were missed by other reviewers.

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

866–867

We shouldn't be checking things that are guaranteed by the op's verifier. It's unnecessary and would lead to a bloat if done consistently.

875–876

This is leading to a redundant call to getDefiningOp -- there isn't a benefit to using matchPattern here. Just use getDefiningOp<arith::TruncIOp>() -- eliminates an extra check.

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

 
 OpFoldResult arith::XOrIOp::fold(ArrayRef<Attribute> operands) {
   /// xor(x, 0) -> x
   if (matchPattern(getRhs(), m_Zero()))
     return getLhs();
   /// xor(x, x) -> 0
   if (getLhs() == getRhs())
     return Builder(getContext()).getZeroAttr(getType());
+  /// xor(xor(x, a), a) -> x
+  if (arith::XOrIOp prev = getLhs().getDefiningOp<arith::XOrIOp>())
+    if (prev.getRhs() == getRhs())
+      return prev.getLhs();
 
   return constFoldBinaryOp<IntegerAttr>(
       operands, [](APInt a, const APInt &b) { return std::move(a) ^ b; });
 }
 
 void arith::XOrIOp::getCanonicalizationPatterns(
     OwningRewritePatternList &patterns, MLIRContext *context) {
   patterns.insert<XOrINotCmpI>(context);
bool arith::ExtFOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
   return checkWidthChangeCast<std::greater, FloatType>(inputs, outputs);
 }
 
 //===----------------------------------------------------------------------===//
 // TruncIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::TruncIOp::fold(ArrayRef<Attribute> operands) {
+  assert(operands.size() == 1 && "unary operation takes one operand");
+
   // trunci(zexti(a)) -> a
   // trunci(sexti(a)) -> a
   if (matchPattern(getOperand(), m_Op<arith::ExtUIOp>()) ||
       matchPattern(getOperand(), m_Op<arith::ExtSIOp>()))
     return getOperand().getDefiningOp()->getOperand(0);
 
-  assert(operands.size() == 1 && "unary operation takes one operand");
+  // trunci(trunci(a)) -> trunci(a))
+  if (matchPattern(getOperand(), m_Op<arith::TruncIOp>())) {
+    setOperand(getOperand().getDefiningOp()->getOperand(0));
+    return getResult();
+  }
 
   if (!operands[0])
     return {};
 
   if (auto lhs = operands[0].dyn_cast<IntegerAttr>()) {
     return IntegerAttr::get(
         getType(), lhs.getValue().trunc(getType().getIntOrFloatBitWidth()));
   }

mlir/test/Dialect/Arithmetic/canonicalize.mlir

 //       CHECK:   %[[cres:.+]] = arith.constant -2 : i16
 //       CHECK:   return %[[cres]]
 func @truncConstant(%arg0: i8) -> i16 {
   %c-2 = arith.constant -2 : i32
   %tr = arith.trunci %c-2 : i32 to i16
   return %tr : i16
 }
 
+// CHECK-LABEL: @truncTrunc
+//       CHECK:   %[[cres:.+]] = arith.trunci %arg0 : i64 to i8
+//       CHECK:   return %[[cres]]
+func @truncTrunc(%arg0: i64) -> i8 {
+  %tr1 = arith.trunci %arg0 : i64 to i32
+  %tr2 = arith.trunci %tr1 : i32 to i8
+  return %tr2 : i8
+}
+
 // CHECK-LABEL: @truncFPConstant
 //       CHECK:   %[[cres:.+]] = arith.constant 1.000000e+00 : bf16
 //       CHECK:   return %[[cres]]
 func @truncFPConstant() -> bf16 {
   %cst = arith.constant 1.000000e+00 : f32
   %0 = arith.truncf %cst : f32 to bf16
   return %0 : bf16
 }
func @notCmpUGE(%arg0: i8, %arg1: i8) -> i1 {
   %true = arith.constant true
   %cmp = arith.cmpi "uge", %arg0, %arg1 : i8
   %ncmp = arith.xori %cmp, %true : i1
   return %ncmp : i1
 }
 
 // -----
 
+// CHECK-LABEL: @xorxor(
+//       CHECK-NOT: xori
+//       CHECK:   return %arg0
+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: @bitcastSameType(
 // CHECK-SAME: %[[ARG:[a-zA-Z0-9_]*]]
 func @bitcastSameType(%arg : f32) -> f32 {
   // CHECK: return %[[ARG]]
   %res = arith.bitcast %arg : f32 to f32
   return %res : f32
 }