[MLIR][Arith] add and(a, or(a,b)) folder

[wsmoses]

No description provided.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-arith

Author: William Moses (wsmoses)

Changes

---

Full diff: https://github.com/llvm/llvm-project/pull/138998.diff

1 Files Affected:

• (modified) mlir/lib/Dialect/Arith/IR/ArithOps.cpp (+12)

diff --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
index 3b308716c84dc..7cf65cdd4f2da 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
@@ -896,6 +896,18 @@ OpFoldResult arith::AndIOp::fold(FoldAdaptor adaptor) {
   if (Value result = foldAndIofAndI(*this))
     return result;
 
+  /// and(a, or(a, b)) -> a
+   for (int i = 0; i < 2; i++) {
+     auto a = getOperand(1 - i);
+     if (auto orOp = getOperand(i).getDefiningOp<arith::OrIOp>()) {
+       for (int j = 0; j < 2; j++) {
+         if (orOp->getOperand(j) == a) {
+           return a;
+         }
+       }
+     }
+   }
+
   return constFoldBinaryOp<IntegerAttr>(
       adaptor.getOperands(),
       [](APInt a, const APInt &b) { return std::move(a) & b; });

[krzysz00]

Question - partly for the peanut gallery - should this be a canonicalization instead of a fold?

[kuhar]

nit: Spell out the full type here since it's not immediately obvious based on the RHS: https://llvm.org/docs/CodingStandards.html#use-auto-type-deduction-to-make-code-more-readable

[kuhar]

Use pre-increment: https://llvm.org/docs/CodingStandards.html#prefer-preincrement

[joker-eph]

I would really rather avoid raw loops entirely.

Can this just be written as:

for (Value operand : getOperands()) {
  if (auto orOp = operand..getDefiningOp<arith::OrIOp>()) {
  

(same for the second loop)

[wsmoses]

so the reason that might be hard is that we need both the original operand, and also the "other" operand. In the loop form that can be written as get operand[i] and operand[1-i], and/or operand[i] and i == 0 ? operand[1] : operand[0]. Not sure how to do that as a single foreach iterator

[joker-eph]

Oh I missed the 1-i in the code, this looks like all code obfuscation to me.

I would write this with a lambda instead, something like:

/// and(or(a, b), a) -> a
auto matchAndOr = [&] (Value lhs, Value rhs) {
  auto orOp = lhs.getDefiningOp<arith::OrIOp>();
   if (!orOp) return false;
   for (Value orOperand : orOp->getOperands())
     if (orOperand == rhs) return true;
   return false;
};

Value lhs = getOperand(0);
Value rhs = getOperand(1);
if (matchAndOr(lhs, rhs)) return rhs;

/// `and` is commutative, swap the operands: `and(a, or(a, b)) -> a`
if (matchAndOr(rhs, lhs)) return lhs;

[kuhar]

also here

[kuhar]

Question - partly for the peanut gallery - should this be a canonicalization instead of a fold?

I think we converged on a policy to prefer folds when possible: https://mlir.llvm.org/docs/Canonicalization/#when-to-use-the-fold-method-vs-rewriterpatterns-for-canonicalizations

[joker-eph]

Suggested change

	if (orOp->getOperand(j) == a) {
	return a;
	}
	if (orOp->getOperand(j) == a)
	return a;

[ThomasRaoux]

should it be arith.andi %a, %c?