[mlir][arith] Fold `and(a, and(a, b))` to `and(a, b)`

[Hardcode84]

ping

[Hardcode84]

rebase

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

OpFoldResult arith::RemSIOp::fold(ArrayRef<Attribute> operands) {
 
   return div0 ? Attribute() : result;
 }
 
 //===----------------------------------------------------------------------===//
 // AndIOp
 //===----------------------------------------------------------------------===//
 
+/// Fold `and(a, and(a, b))` to `and(a, b)`
+static Value foldAndIofAndI(arith::AndIOp op) {
+  for (bool reversePrev : {false, true}) {
+    auto prev = (reversePrev ? op.getRhs() : op.getLhs())
+                    .getDefiningOp<arith::AndIOp>();
+    if (!prev)
+      continue;
+
+    Value other = (reversePrev ? op.getLhs() : op.getRhs());
+    if (other != prev.getLhs() && other != prev.getRhs())
+      continue;
+
+    return prev.getResult();
+  }
+  return {};
+}
+
 OpFoldResult arith::AndIOp::fold(ArrayRef<Attribute> operands) {
   /// and(x, 0) -> 0
   if (matchPattern(getRhs(), m_Zero()))
     return getRhs();
   /// and(x, allOnes) -> x
   APInt intValue;
   if (matchPattern(getRhs(), m_ConstantInt(&intValue)) && intValue.isAllOnes())
     return getLhs();
   /// and(x, not(x)) -> 0
   if (matchPattern(getRhs(), m_Op<XOrIOp>(matchers::m_Val(getLhs()),
                                           m_ConstantInt(&intValue))) &&
       intValue.isAllOnes())
     return IntegerAttr::get(getType(), 0);
   /// and(not(x), x) -> 0
   if (matchPattern(getLhs(), m_Op<XOrIOp>(matchers::m_Val(getRhs()),
                                           m_ConstantInt(&intValue))) &&
       intValue.isAllOnes())
     return IntegerAttr::get(getType(), 0);
 
+  /// and(a, and(a, b)) -> and(a, b)
+  if (Value result = foldAndIofAndI(*this))
+    return result;
+
   return constFoldBinaryOp<IntegerAttr>(
       operands, [](APInt a, const APInt &b) { return std::move(a) & b; });
 }
 
 //===----------------------------------------------------------------------===//
 // OrIOp
 //===----------------------------------------------------------------------===//
 

mlir/test/Dialect/Arith/canonicalize.mlir

 //       CHECK:   return %arg0
 func.func @xorxor3(%a : i32, %b : i32) -> i32 {
   %c = arith.xori %b, %a : i32
   %res = arith.xori %b, %c : i32
   return %res : i32
 }
 
 // -----
+
+/// and(a, and(a, b)) -> and(a, b)
+
+// CHECK-LABEL: @andand0
+//  CHECK-SAME:   (%[[A:.*]]: i32, %[[B:.*]]: i32)
+//       CHECK:   %[[RES:.*]] = arith.andi %[[A]], %[[B]] : i32
+//       CHECK:   return %[[RES]]
+func.func @andand0(%a : i32, %b : i32) -> i32 {
+  %c = arith.andi %a, %b : i32
+  %res = arith.andi %a, %c : i32
+  return %res : i32
+}
+
+// CHECK-LABEL: @andand1
+//  CHECK-SAME:   (%[[A:.*]]: i32, %[[B:.*]]: i32)
+//       CHECK:   %[[RES:.*]] = arith.andi %[[A]], %[[B]] : i32
+//       CHECK:   return %[[RES]]
+func.func @andand1(%a : i32, %b : i32) -> i32 {
+  %c = arith.andi %a, %b : i32
+  %res = arith.andi %c, %a : i32
+  return %res : i32
+}
+
+// CHECK-LABEL: @andand2
+//  CHECK-SAME:   (%[[A:.*]]: i32, %[[B:.*]]: i32)
+//       CHECK:   %[[RES:.*]] = arith.andi %[[A]], %[[B]] : i32
+//       CHECK:   return %[[RES]]
+func.func @andand2(%a : i32, %b : i32) -> i32 {
+  %c = arith.andi %a, %b : i32
+  %res = arith.andi %b, %c : i32
+  return %res : i32
+}
+
+// CHECK-LABEL: @andand3
+//  CHECK-SAME:   (%[[A:.*]]: i32, %[[B:.*]]: i32)
+//       CHECK:   %[[RES:.*]] = arith.andi %[[A]], %[[B]] : i32
+//       CHECK:   return %[[RES]]
+func.func @andand3(%a : i32, %b : i32) -> i32 {
+  %c = arith.andi %a, %b : i32
+  %res = arith.andi %c, %b : i32
+  return %res : i32
+}