[MLIR][SCF] Canonicalize while statement whose cmp condition is recomputed in the after region

[wsmoses]

Summary

Given a while loop whose condition is given by a cmp, don't recomputed the comparison (or its inverse) in the after region, instead use a constant since the original condition must be true if we branched to the after region.

[mehdi_amini]

mlir/lib/Dialect/SCF/SCF.cpp

2468

Shouldn't this be handled in CSE: the expression in the do block should be hoisted in the condition block here?

[wsmoses]

mlir/lib/Dialect/SCF/SCF.cpp

2468

Will CSE actually hoist from the do block into the while? Running this test on CSE, implies no (though that would certainly be desirable).

Regardless, this also handles the case where the condition is the negation.

[wsmoses]

mlir/lib/Dialect/SCF/SCF.cpp

2468

I think part of the reason CSE can't perform this hoisting is because the argument's aren't completely the same (e.g. the later comparison uses the block arguments, which are only defined at the start of the after region).

[rriddle]

mlir/lib/Dialect/SCF/SCF.cpp

2496–2497

Do you have coverage for when this gets hit twice?

[mehdi_amini]

mlir/lib/Dialect/SCF/SCF.cpp

2468

Will CSE actually hoist from the do block into the while? Running this test on CSE, implies no (though that would certainly be desirable).

Right, I didn't mean that CSE already implements it, just that this is where it belongs :)

2483

2496–2497

I suspect we'll be missing an make_early_inc_range for the for (auto &u : std::get<1>(tup).getUses()) {

[wsmoses]

Address comments

[rriddle]

mlir/lib/Dialect/SCF/SCF.cpp

2482

Can you negate this and use early continue to reduce nesting?

2494–2496

Is it possible for u.getOwner() to have multiple uses of the value?

[wsmoses]

Address comments

[wsmoses]

mlir/lib/Dialect/SCF/SCF.cpp

2494–2496

No, as the owner is the comparison cmp2, which has its rhs equal to the comparison in the before block's rhs (and thus cannot be a blockArg to the after region).

[mehdi_amini]

mlir/lib/Dialect/SCF/SCF.cpp

2494–2496

I don't think so because of the condition line 2489? (if (cmp2.getRhs() != cmp.getRhs()))

By the way this make the pattern only catch one side isn't it? If I just reverse the operands we'll bail out:

  %condition = arith.cmpi ne, %arg0, %val : i32
  scf.condition(%condition) %val : i32
} do {
^bb0(%val2: i32):
  %condition2 = arith.cmpi ne,  %arg0, %val2 : i32
  %negcondition2 = arith.cmpi eq, %arg0, %val2 : i32
  "test.use"(%condition2, %negcondition2, %val2) : (i1, i1, i32) -> ()

Also it seems like it won't handle cases where both operands come from the condition block.

[wsmoses]

Handle both sides of comparison

mlir/include/mlir/Dialect/Arithmetic/IR/Arithmetic.h

 /// Returns the identity value associated with an AtomicRMWKind op.
 Value getIdentityValue(AtomicRMWKind op, Type resultType, OpBuilder &builder,
                        Location loc);
 
 /// Returns the value obtained by applying the reduction operation kind
 /// associated with a binary AtomicRMWKind op to `lhs` and `rhs`.
 Value getReductionOp(AtomicRMWKind op, OpBuilder &builder, Location loc,
                      Value lhs, Value rhs);
+
+arith::CmpIPredicate invertPredicate(arith::CmpIPredicate pred);
 } // namespace arith
 } // namespace mlir
 
 #endif // MLIR_DIALECT_ARITHMETIC_IR_ARITHMETIC_H_

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

 static IntegerAttr subIntegerAttrs(PatternRewriter &builder, Value res,
                                    Attribute lhs, Attribute rhs) {
   return builder.getIntegerAttr(res.getType(),
                                 lhs.cast<IntegerAttr>().getInt() -
                                     rhs.cast<IntegerAttr>().getInt());
 }
 
 /// Invert an integer comparison predicate.
-static arith::CmpIPredicate invertPredicate(arith::CmpIPredicate pred) {
+arith::CmpIPredicate arith::invertPredicate(arith::CmpIPredicate pred) {
   switch (pred) {
   case arith::CmpIPredicate::eq:
     return arith::CmpIPredicate::ne;
   case arith::CmpIPredicate::ne:
     return arith::CmpIPredicate::eq;
   case arith::CmpIPredicate::slt:
     return arith::CmpIPredicate::sge;
   case arith::CmpIPredicate::sle:

mlir/lib/Dialect/SCF/SCF.cpp

LogicalResult matchAndRewrite(WhileOp op,
 
     Block &afterBlock = op.getAfter().front();
     rewriter.mergeBlocks(&afterBlock, &newAfterBlock, newAfterBlockArgs);
 
     rewriter.replaceOp(op, newResults);
     return success();
   }
 };
+
+/// Replace operations equivalent to the condition in the do block with true,
+/// since otherwise the block would not be evaluated.
+///
+/// scf.while (..) : (i32, ...) -> ... {
+///  %z = ... : i32
+///  %condition = cmpi pred %z, %a
+///  scf.condition(%condition) %z : i32, ...
+/// } do {
+/// ^bb0(%arg0: i32, ...):
+///    %condition2 = cmpi pred %arg0, %a
+///    use(%condition2)
+///    ...
+///
+/// becomes
+/// scf.while (..) : (i32, ...) -> ... {
+///  %z = ... : i32
+///  %condition = cmpi pred %z, %a
+///  scf.condition(%condition) %z : i32, ...
+/// } do {
+/// ^bb0(%arg0: i32, ...):
+///    use(%true)
+///    ...
+struct WhileCmpCond : public OpRewritePattern<scf::WhileOp> {
+  using OpRewritePattern<scf::WhileOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(scf::WhileOp op,
+                                PatternRewriter &rewriter) const override {
+    using namespace scf;
+    auto cond = op.getConditionOp();
+    auto cmp = cond.getCondition().getDefiningOp<arith::CmpIOp>();
+    if (!cmp)
+      return failure();
+    bool changed = false;
+    for (auto tup :
+         llvm::zip(cond.getArgs(), op.getAfter().front().getArguments())) {
+      for (size_t opIdx = 0; opIdx < 2; opIdx++) {
+        if (std::get<0>(tup) != cmp.getOperand(opIdx))
+          continue;
+        for (OpOperand &u :
+             llvm::make_early_inc_range(std::get<1>(tup).getUses())) {
+          auto cmp2 = dyn_cast<arith::CmpIOp>(u.getOwner());
+          if (!cmp2)
+            continue;
+          // For a binary operator 1-opIdx gets the other side.
+          if (cmp2.getOperand(1 - opIdx) != cmp.getOperand(1 - opIdx))
+            continue;
+          bool samePredicate;
+          if (cmp2.getPredicate() == cmp.getPredicate())
+            samePredicate = true;
+          else if (cmp2.getPredicate() ==
+                   arith::invertPredicate(cmp.getPredicate()))
+            samePredicate = false;
+          else
+            continue;
+
+          rewriter.replaceOpWithNewOp<arith::ConstantIntOp>(cmp2, samePredicate,
+                                                            1);
+          changed = true;
+        }
+      }
+    }
+    return success(changed);
+  }
+};
 } // namespace
 
 void WhileOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
                                           MLIRContext *context) {
-  results.insert<WhileConditionTruth, WhileUnusedResult>(context);
+  results.insert<WhileConditionTruth, WhileUnusedResult, WhileCmpCond>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/SCF/SCFOps.cpp.inc"

mlir/test/Dialect/SCF/canonicalize.mlir

 // CHECK-NEXT:           scf.condition(%[[cmp]]) %[[val]] : i32
 // CHECK-NEXT:         } do {
 // CHECK-NEXT:         ^bb0(%[[arg:.*]]: i32):  // no predecessors
 // CHECK-NEXT:           "test.use"(%[[arg]]) : (i32) -> ()
 // CHECK-NEXT:           scf.yield
 // CHECK-NEXT:         }
 // CHECK-NEXT:         return %[[res]] : i32
 
+// CHECK-LABEL: @while_cmp_lhs
+func @while_cmp_lhs(%arg0 : i32) {
+  %0 = scf.while () : () -> i32 {
+    %val = "test.val"() : () -> i32
+    %condition = arith.cmpi ne, %val, %arg0 : i32
+    scf.condition(%condition) %val : i32
+  } do {
+  ^bb0(%val2: i32):
+    %condition2 = arith.cmpi ne, %val2, %arg0 : i32
+    %negcondition2 = arith.cmpi eq, %val2, %arg0 : i32
+    "test.use"(%condition2, %negcondition2, %val2) : (i1, i1, i32) -> ()
+    scf.yield
+  }
+  return
+}
+// CHECK-DAG:         %[[true:.+]] = arith.constant true
+// CHECK-DAG:         %[[false:.+]] = arith.constant false
+// CHECK-DAG:         %{{.+}} = scf.while : () -> i32 {
+// CHECK-NEXT:         %[[val:.+]] = "test.val"
+// CHECK-NEXT:         %[[cmp:.+]] = arith.cmpi ne, %[[val]], %arg0 : i32
+// CHECK-NEXT:           scf.condition(%[[cmp]]) %[[val]] : i32
+// CHECK-NEXT:         } do {
+// CHECK-NEXT:         ^bb0(%arg1: i32):  // no predecessors
+// CHECK-NEXT:           "test.use"(%[[true]], %[[false]], %arg1) : (i1, i1, i32) -> ()
+// CHECK-NEXT:           scf.yield
+// CHECK-NEXT:         }
+
+// CHECK-LABEL: @while_cmp_rhs
+func @while_cmp_rhs(%arg0 : i32) {
+  %0 = scf.while () : () -> i32 {
+    %val = "test.val"() : () -> i32
+    %condition = arith.cmpi ne, %arg0, %val : i32
+    scf.condition(%condition) %val : i32
+  } do {
+  ^bb0(%val2: i32):
+    %condition2 = arith.cmpi ne, %arg0, %val2 : i32
+    %negcondition2 = arith.cmpi eq, %arg0, %val2 : i32
+    "test.use"(%condition2, %negcondition2, %val2) : (i1, i1, i32) -> ()
+    scf.yield
+  }
+  return
+}
+// CHECK-DAG:         %[[true:.+]] = arith.constant true
+// CHECK-DAG:         %[[false:.+]] = arith.constant false
+// CHECK-DAG:         %{{.+}} = scf.while : () -> i32 {
+// CHECK-NEXT:         %[[val:.+]] = "test.val"
+// CHECK-NEXT:         %[[cmp:.+]] = arith.cmpi ne, %arg0, %[[val]] : i32
+// CHECK-NEXT:           scf.condition(%[[cmp]]) %[[val]] : i32
+// CHECK-NEXT:         } do {
+// CHECK-NEXT:         ^bb0(%arg1: i32):  // no predecessors
+// CHECK-NEXT:           "test.use"(%[[true]], %[[false]], %arg1) : (i1, i1, i32) -> ()
+// CHECK-NEXT:           scf.yield
+// CHECK-NEXT:         }
+
 // -----
 
 // CHECK-LABEL: @combineIfs
 func @combineIfs(%arg0 : i1, %arg2: i64) -> (i32, i32) {
   %res = scf.if %arg0 -> i32 {
     %v = "test.firstCodeTrue"() : () -> i32
     scf.yield %v : i32
   } else {