[mlir][scf] More WhileOp canonicalizations

[Hardcode84]

Summary

Remove duplicated ConditonOp args, remove unused init/yield args

[Hardcode84]

rebase

[Hardcode84]

rebase

[Hardcode84]

ping

[Hardcode84]

rebase

[Hardcode84]

ping

[ftynse]

mlir/lib/Dialect/SCF/IR/SCF.cpp

3811

Nit: prefer notifyMatchFailure with an explanatory message.

3815

Nit: feel free to add builder overloads that don't take callbacks.

3826

auto &&, enumerate will eventually return references.

3863

Ditto for auto &&.

[Hardcode84]

rebase, review comments

[Hardcode84]

mlir/lib/Dialect/SCF/IR/SCF.cpp

3815

Will do in separate review

[mehdi_amini]

mlir/lib/Dialect/SCF/IR/SCF.cpp

3796

This isn't used before the early exit, can you move this closer to its use?

3798

There is a getConditionOp() method on WhileOp.

3801

Shall we reserve() these before the loop?

I'm also not sure if we could do better to filter out the "do nothing case" (which is the most important one to optimize for).
For example we could do a first pass with a SmallPtrSet just to check if there are duplicated argument and bail out immediately?

3841

mergeBlocks is pretty expensive: it is possible to steal the block as the new body of the region instead? Isn't the destination block empty?
Actually mergeBlocks() should be optimized for the case where the destination block is empty! (maybe it is already?)

mlir/test/Dialect/SCF/canonicalize.mlir

1249

Do we need that many checks? Seems like a CHECK-NOT: test.val1 would do the job here, potentially after a CHECK: scf.while : () -> i32. The rest seems spurious.

[Hardcode84]

There is a another issue, there is an existing WhileUnusedArg pattern, but it is not added anywhere, and test which was supposed to check it passes because of unrelated pattern. I'll do another PR to clean thing up

[Hardcode84]

mlir/test/Dialect/SCF/canonicalize.mlir

1249

I still want to check general structure of the loop is preserved.

[Hardcode84]

mlir/lib/Dialect/SCF/IR/SCF.cpp

3841

mergeBlocks just doing replaceAllUsesWith + getOperations().splice, is it really that expensive?

mlir/lib/Dialect/SCF/IR/SCF.cpp

LogicalResult matchAndRewrite(WhileOp op,
     replacement.getAfter().takeBody(op.getAfter());
     rewriter.replaceOp(op, replacement.getResults());
 
     rewriter.setInsertionPoint(yield);
     rewriter.replaceOpWithNewOp<YieldOp>(yield, newYields);
     return success();
   }
 };
+
+/// Remove duplicated ConditionOp args.
+struct WhileRemoveDuplicatedResults : public OpRewritePattern<WhileOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(WhileOp op,
+                                PatternRewriter &rewriter) const override {
+    Block &beforeBlock = op.getBefore().front();
+    Block &afterBlock = op.getAfter().front();
+
+    auto condOp = cast<ConditionOp>(beforeBlock.getTerminator());
+    ValueRange condOpArgs = condOp.getArgs();
+    llvm::SmallDenseMap<Value, unsigned> argsMap;
+    SmallVector<Value> newArgs;
+    for (auto arg : condOpArgs) {
+      if (!argsMap.count(arg)) {
+        auto pos = static_cast<unsigned>(argsMap.size());
+        argsMap.insert({arg, pos});
+        newArgs.emplace_back(arg);
+      }
+    }
+
+    if (argsMap.size() == condOpArgs.size())
+      return rewriter.notifyMatchFailure(op, "No results to remove");
+
+    ValueRange argsRange(newArgs);
+    auto emptyBuilder = [](OpBuilder &, Location, ValueRange) {
+      // Nothing
+    };
+
+    Location loc = op.getLoc();
+    auto newWhileOp = rewriter.create<scf::WhileOp>(
+        loc, argsRange.getTypes(), op.getInits(), emptyBuilder, emptyBuilder);
+    Block &newBeforeBlock = newWhileOp.getBefore().front();
+    Block &newAfterBlock = newWhileOp.getAfter().front();
+
+    SmallVector<Value> afterArgsMapping;
+    SmallVector<Value> resultsMapping;
+    for (auto &&[i, arg] : llvm::enumerate(condOpArgs)) {
+      auto it = argsMap.find(arg);
+      assert(it != argsMap.end());
+      auto pos = it->second;
+      afterArgsMapping.emplace_back(newAfterBlock.getArgument(pos));
+      resultsMapping.emplace_back(newWhileOp->getResult(pos));
+    }
+
+    OpBuilder::InsertionGuard g(rewriter);
+    rewriter.setInsertionPoint(condOp);
+    rewriter.replaceOpWithNewOp<ConditionOp>(condOp, condOp.getCondition(),
+                                             argsRange);
+
+    rewriter.mergeBlocks(&beforeBlock, &newBeforeBlock,
+                         newBeforeBlock.getArguments());
+    rewriter.mergeBlocks(&afterBlock, &newAfterBlock, afterArgsMapping);
+    rewriter.replaceOp(op, resultsMapping);
+    return success();
+  }
+};
+
+/// Remove unused init/yield args.
+struct WhileRemoveUnusedArgs : public mlir::OpRewritePattern<WhileOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(WhileOp op,
+                                PatternRewriter &rewriter) const override {
+    Block &beforeBlock = op.getBefore().front();
+    Block &afterBlock = op.getAfter().front();
+
+    auto yield = cast<YieldOp>(afterBlock.getTerminator());
+
+    llvm::BitVector argsToRemove;
+    SmallVector<Value> newInits;
+    SmallVector<Value> newYieldArgs;
+
+    bool changed = false;
+    for (auto &&[arg, init, yieldArg] : llvm::zip(
+             beforeBlock.getArguments(), op.getInits(), yield.getResults())) {
+      bool empty = arg.use_empty();
+      argsToRemove.push_back(empty);
+      if (empty) {
+        changed = true;
+        continue;
+      }
+
+      newInits.emplace_back(init);
+      newYieldArgs.emplace_back(yieldArg);
+    }
+
+    if (!changed)
+      return rewriter.notifyMatchFailure(op, "No args to remove");
+
+    beforeBlock.eraseArguments(argsToRemove);
+
+    auto emptyBuilder = [](OpBuilder &, Location, ValueRange) {
+      // Nothing
+    };
+
+    Location loc = op.getLoc();
+    auto newWhileOp = rewriter.create<WhileOp>(
+        loc, op->getResultTypes(), newInits, emptyBuilder, emptyBuilder);
+    Block &newBeforeBlock = newWhileOp.getBefore().front();
+    Block &newAfterBlock = newWhileOp.getAfter().front();
+
+    OpBuilder::InsertionGuard g(rewriter);
+    rewriter.setInsertionPoint(yield);
+    rewriter.replaceOpWithNewOp<YieldOp>(yield, newYieldArgs);
+
+    rewriter.mergeBlocks(&beforeBlock, &newBeforeBlock,
+                         newBeforeBlock.getArguments());
+    rewriter.mergeBlocks(&afterBlock, &newAfterBlock,
+                         newAfterBlock.getArguments());
+    rewriter.replaceOp(op, newWhileOp.getResults());
+    return success();
+  }
+};
 } // namespace
 
 void WhileOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                           MLIRContext *context) {
   results.add<RemoveLoopInvariantArgsFromBeforeBlock,
               RemoveLoopInvariantValueYielded, WhileConditionTruth,
-              WhileCmpCond, WhileUnusedResult>(context);
+              WhileCmpCond, WhileUnusedResult, WhileRemoveDuplicatedResults,
+              WhileRemoveUnusedArgs>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // IndexSwitchOp
 //===----------------------------------------------------------------------===//
 
 /// Parse the case regions and values.
 static ParseResult

mlir/test/Dialect/SCF/canonicalize.mlir

 // CHECK-NEXT:         } do {
 // CHECK-NEXT:         ^bb0(%arg1: i32):
 // CHECK-NEXT:           "test.use"(%[[true]], %[[false]], %arg1) : (i1, i1, i32) -> ()
 // CHECK-NEXT:           scf.yield
 // CHECK-NEXT:         }
 
 // -----
 
+// CHECK-LABEL: @while_duplicated_res
+func.func @while_duplicated_res() -> (i32, i32) {
+  %0:2 = scf.while () : () -> (i32, i32) {
+    %val = "test.val"() : () -> i32
+    %condition = "test.condition"() : () -> i1
+    scf.condition(%condition) %val, %val : i32, i32
+  } do {
+  ^bb0(%val2: i32, %val3: i32):
+    "test.use"(%val2, %val3) : (i32, i32) -> ()
+    scf.yield
+  }
+  return %0#0, %0#1: i32, i32
+}
+// CHECK:         %[[RES:.*]] = scf.while : () -> i32 {
+// CHECK:         %[[VAL:.*]] = "test.val"() : () -> i32
+// CHECK:         %[[COND:.*]] = "test.condition"() : () -> i1
+// CHECK:           scf.condition(%[[COND]]) %[[VAL]] : i32
+// CHECK:         } do {
+// CHECK:         ^bb0(%[[ARG:.*]]: i32):
+// CHECK:           "test.use"(%[[ARG]], %[[ARG]]) : (i32, i32) -> ()
+// CHECK:           scf.yield
+// CHECK:         }
+// CHECK:         return %[[RES]], %[[RES]] : i32, i32
+
+// -----
+
+// CHECK-LABEL: @while_unused_arg
+func.func @while_unused_arg(%val0: i32) -> i32 {
+  %0 = scf.while (%val1 = %val0) : (i32) -> i32 {
+    %val = "test.val"() : () -> i32
+    %condition = "test.condition"() : () -> i1
+    scf.condition(%condition) %val: i32
+  } do {
+  ^bb0(%val2: i32):
+    "test.use"(%val2) : (i32) -> ()
+    %val1 = "test.val1"() : () -> i32
+    scf.yield %val1 : i32
+  }
+  return %0 : i32
+}
+// CHECK:         %[[RES:.*]] = scf.while : () -> i32 {
+// CHECK:         %[[VAL:.*]] = "test.val"() : () -> i32
+// CHECK:         %[[COND:.*]] = "test.condition"() : () -> i1
+// CHECK:           scf.condition(%[[COND]]) %[[VAL]] : i32
+// CHECK:         } do {
+// CHECK:         ^bb0(%[[ARG:.*]]: i32):
+// CHECK:           "test.use"(%[[ARG]]) : (i32) -> ()
+// CHECK:           scf.yield
+// CHECK:         }
+// CHECK:         return %[[RES]] : i32
+
+
+// -----
+
 // CHECK-LABEL: @combineIfs
 func.func @combineIfs(%arg0 : i1, %arg2: i64) -> (i32, i32) {
   %res = scf.if %arg0 -> i32 {
     %v = "test.firstCodeTrue"() : () -> i32
     scf.yield %v : i32
   } else {
     %v2 = "test.firstCodeFalse"() : () -> i32
     scf.yield %v2 : i32

mlir/test/Dialect/SCF/one-shot-bufferize.mlir

 //       CHECK:     memref.dealloc %[[alloc2]]
 //       CHECK:     scf.yield %[[cloned]]
 //       CHECK:   }
 func.func @scf_while_iter_arg_result_mismatch(%arg0: tensor<5xi1>,
                                               %arg1: tensor<5xi1>,
                                               %arg2: index) {
   scf.while (%arg3 = %arg1) : (tensor<5xi1>) -> () {
     %0 = tensor.extract %arg0[%arg2] : tensor<5xi1>
+    %1 = tensor.extract %arg3[%arg2] : tensor<5xi1>
+    "dummy.use"(%1) : (i1) -> ()
     scf.condition(%0)
   } do {
     %0 = "dummy.some_op"() : () -> index
     %1 = "dummy.another_op"() : () -> i1
     %2 = tensor.insert %1 into %arg0[%0] : tensor<5xi1>
     scf.yield %2 : tensor<5xi1>
   }
   return