[MLIR][SCF] Assume uses of condition in the body of scf.while is true

[mehdi_amini]

mlir/lib/Dialect/SCF/SCF.cpp

1865

I have a representation question here: why are we passing the condition to the body in the first place??

[wsmoses]

In theory, it would be nice for this not to be generated (hence why this makes a good canonicalization). In practice, this can end up being generated by an MLIR C/C++ frontend (https://github.com/wsmoses/Polygeist) as an intermediate state or simply because the source code was written like it. This consequently hinders subsequent optimization.

This isn't the particular case that triggered this to be created, but for example:

c

bool b;
while(b = condition()) {
  // canonicalized out;
  b |= knownFalse();
  if (b) {
    ...
  }
}

[mehdi_amini]

In D101801#2735333, @wsmoses wrote:

In theory, it would be nice for this not to be generated (hence why this makes a good canonicalization).

To me it does not makes sense that the canonical form of the operation included a dead block argument.

In practice, this can end up being generated by an MLIR C/C++ frontend (https://github.com/wsmoses/Polygeist) as an intermediate state or simply because the source code was written like it. This consequently hinders subsequent optimization.

This isn't the particular case that triggered this to be created, but for example:

c

bool b;
while(b = condition()) {
  // canonicalized out;
  b |= knownFalse();
  if (b) {
    ...
  }
}

This example does not seems like a good justification to me: it isn't in SSA form.

[wsmoses]

Ah I misunderstood your comment. The canonicalized version of the code has an unused argument because the types of the do region must match the types of the result. I suppose in this instance we could also replace uses of the result of the while with false and remove the argument entirely. Would it be better to do that here or a dead-arg elimination canonicalization (standalone pass for is intended to be submitted after this).

[mehdi_amini]

OK I think I misread the pattern the first time, forget my initial remark!

mlir/lib/Dialect/SCF/SCF.cpp

1871

Please add an accessor on WhileOp for this instead.

(in general if we have a guarantee of SingleBlock region, I rather not expose "raw" Region API and write .front() or back() everywhere, and same below).

1872

Can you move this in the rewriter.create call? That way we don't pay this when the pattern does not match.

[mehdi_amini]

mlir/lib/Dialect/SCF/SCF.cpp

1886

The constantTrue could be initialized inside this condition I think.
Another way to handle these kind of lazy initializations is with a lambda:

Value constantTrue = nullptr;
auto getConstantTrue() = [&] {
  if (!constantTrue) 
    constantTrue = rewriter.create<mlir::ConstantOp>(
              op.getLoc(), ty, rewriter.getBoolAttr(true));
  return constantTrue;
};


for (...)
  if (...)
    op.replaceAllUsesWith(getConstantTrue())

[wsmoses]

Refactor while accessors

mlir/include/mlir/Dialect/SCF/SCFOps.td

def WhileOp : SCF_Op<"while",
   }];
 
   let arguments = (ins Variadic<AnyType>:$inits);
   let results = (outs Variadic<AnyType>:$results);
   let regions = (region SizedRegion<1>:$before, SizedRegion<1>:$after);
 
   let extraClassDeclaration = [{
     OperandRange getSuccessorEntryOperands(unsigned index);
+    ConditionOp getConditionOp();
+    Block::BlockArgListType getAfterArguments();
   }];
+
+  let hasCanonicalizer = 1;
 }
 
 def YieldOp : SCF_Op<"yield", [NoSideEffect, ReturnLike, Terminator,
                                ParentOneOf<["IfOp, ForOp", "ParallelOp",
                                             "WhileOp"]>]> {
   let summary = "loop yield and termination operation";
   let description = [{
     "scf.yield" yields an SSA value from the SCF dialect op region and

mlir/lib/Dialect/SCF/SCF.cpp

 
 OperandRange WhileOp::getSuccessorEntryOperands(unsigned index) {
   assert(index == 0 &&
          "WhileOp is expected to branch only to the first region");
 
   return inits();
 }
 
+ConditionOp WhileOp::getConditionOp() {
+  return cast<ConditionOp>(before().front().getTerminator());
+}
+
+Block::BlockArgListType WhileOp::getAfterArguments() {
+  return after().front().getArguments();
+}
+
 void WhileOp::getSuccessorRegions(Optional<unsigned> index,
                                   ArrayRef<Attribute> operands,
                                   SmallVectorImpl<RegionSuccessor> &regions) {
   (void)operands;
 
   if (!index.hasValue()) {
     regions.emplace_back(&before(), before().getArguments());
     return;
if (failed(verifyTypeRangesMatch(
     return failure();
 
   auto afterTerminator = verifyAndGetTerminator<scf::YieldOp>(
       op, op.after(),
       "expects the 'after' region to terminate with 'scf.yield'");
   return success(afterTerminator != nullptr);
 }
 
+namespace {
+/// Replace uses of the condition within the do block with true, since otherwise
+/// the block would not be evaluated.
+///
+/// scf.while (..) : (i1, ...) -> ... {
+///  %condition = call @evaluate_condition() : () -> i1
+///  scf.condition(%condition) %condition : i1, ...
+/// } do {
+/// ^bb0(%arg0: i1, ...):
+///    use(%arg0)
+///    ...
+///
+/// becomes
+/// scf.while (..) : (i1, ...) -> ... {
+///  %condition = call @evaluate_condition() : () -> i1
+///  scf.condition(%condition) %condition : i1, ...
+/// } do {
+/// ^bb0(%arg0: i1, ...):
+///    use(%true)
+///    ...
+struct WhileConditionTruth : public OpRewritePattern<WhileOp> {
+  using OpRewritePattern<WhileOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(WhileOp op,
+                                PatternRewriter &rewriter) const override {
+    auto term = op.getConditionOp();
+
+    // These variables serve to prevent creating duplicate constants
+    // and hold constant true or false values.
+    Value constantTrue = nullptr;
+
+    bool replaced = false;
+    for (auto yieldedAndBlockArgs :
+         llvm::zip(term.args(), op.getAfterArguments())) {
+      if (std::get<0>(yieldedAndBlockArgs) == term.condition()) {
+        if (!std::get<1>(yieldedAndBlockArgs).use_empty()) {
+          if (!constantTrue)
+            constantTrue = rewriter.create<mlir::ConstantOp>(
+                op.getLoc(), term.condition().getType(),
+                rewriter.getBoolAttr(true));
+
+          std::get<1>(yieldedAndBlockArgs).replaceAllUsesWith(constantTrue);
+          replaced = true;
+        }
+      }
+    }
+    return success(replaced);
+  }
+};
+} // namespace
+
+void WhileOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
+                                          MLIRContext *context) {
+  results.insert<WhileConditionTruth>(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:     %[[if:.+]] = scf.if %arg0 -> (i32) {
 // CHECK-NEXT:       %[[sv1:.+]] = "test.get_some_value"() : () -> i32
 // CHECK-NEXT:       scf.yield %[[sv1]] : i32
 // CHECK-NEXT:     } else {
 // CHECK-NEXT:       %[[sv2:.+]] = "test.get_some_value"() : () -> i32
 // CHECK-NEXT:       scf.yield %[[sv2]] : i32
 // CHECK-NEXT:     }
 // CHECK-NEXT:     return %[[if]], %arg1 : i32, i64
+
+
+// CHECK-LABEL: @while_cond_true
+func @while_cond_true() {
+  %0 = scf.while () : () -> i1 {
+    %condition = "test.condition"() : () -> i1
+    scf.condition(%condition) %condition : i1
+  } do {
+  ^bb0(%arg0: i1):
+    "test.use"(%arg0) : (i1) -> ()
+    scf.yield
+  }
+  return
+}
+// CHECK-NEXT:         %[[true:.+]] = constant true
+// CHECK-NEXT:         %{{.+}} = scf.while : () -> i1 {
+// CHECK-NEXT:           %[[cmp:.+]] = "test.condition"() : () -> i1
+// CHECK-NEXT:           scf.condition(%[[cmp]]) %[[cmp]] : i1
+// CHECK-NEXT:         } do {
+// CHECK-NEXT:         ^bb0(%arg0: i1):  // no predecessors
+// CHECK-NEXT:           "test.use"(%[[true]]) : (i1) -> ()
+// CHECK-NEXT:           scf.yield
+// CHECK-NEXT:         }