[mlir][scf] Add simple LICM pattern for scf.while

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

@@ -3810,6 +3810,36 @@ struct WhileRemoveUnusedArgs : public OpRewritePattern<WhileOp> {
   }
 };
 
+/// Simple Loop Invariant Code Motion pattern for `scf.while` op.
+/// `scf.while` to `scf.for` uplifting expects `before` block consisting of
+/// single `cmp` op.
+/// Pattern moves ops from `before` block, doesn't visit nested regions.
+struct SCFWhileLICM : public OpRewritePattern<WhileOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(WhileOp loop,
+                                PatternRewriter &rewriter) const override {
+    bool changed = false;
+
+    DominanceInfo dom;
+    Block *body = loop.getBeforeBody();
+    for (Operation &op :
+         llvm::make_early_inc_range(body->without_terminator())) {
+      if (llvm::any_of(op.getOperands(), [&](Value arg) {
+            return !dom.properlyDominates(arg, loop);
+          }))
+        continue;
+
+      if (!isMemoryEffectFree(&op))
+        continue;
+
+      rewriter.updateRootInPlace(&op, [&]() { op.moveBefore(loop); });
+      changed = true;
+    }
+    return success(changed);
+  }
+};
+
 /// Remove duplicated ConditionOp args.
 struct WhileRemoveDuplicatedResults : public OpRewritePattern<WhileOp> {
   using OpRewritePattern::OpRewritePattern;
@@ -3879,7 +3909,7 @@ void WhileOp::getCanonicalizationPatterns(RewritePatternSet &results,
   results.add<RemoveLoopInvariantArgsFromBeforeBlock,
               RemoveLoopInvariantValueYielded, WhileConditionTruth,
               WhileCmpCond, WhileUnusedResult, WhileRemoveDuplicatedResults,
-              WhileRemoveUnusedArgs>(context);
+              SCFWhileLICM, WhileRemoveUnusedArgs>(context);
 }
 
 //===----------------------------------------------------------------------===//

mlir/test/Dialect/SCF/canonicalize.mlir

@@ -1022,10 +1022,10 @@ func.func @while_loop_invariant_argument_different_order(%arg : tensor<i32>) ->
 // CHECK-SAME: (%[[ARG:.+]]: tensor<i32>)
 // CHECK:    %[[ZERO:.*]] = arith.constant dense<0>
 // CHECK:    %[[ONE:.*]] = arith.constant dense<1>
+// CHECK:    %[[COND:.*]] = arith.cmpi sgt, %[[ARG]], %[[ZERO]]
+// CHECK:    %[[COND1:.*]] = tensor.extract %[[COND]][]
 // CHECK:    %[[WHILE:.*]]:2 = scf.while (%[[ARG1:.*]] = %[[ONE]], %[[ARG4:.*]] = %[[ZERO]])
-// CHECK:       arith.cmpi sgt, %[[ARG]], %[[ZERO]]
-// CHECK:       tensor.extract %{{.*}}[]
-// CHECK:       scf.condition(%{{.*}}) %[[ARG1]], %[[ARG4]]
+// CHECK:       scf.condition(%[[COND1]]) %[[ARG1]], %[[ARG4]]
 // CHECK:    } do {
 // CHECK:     ^{{.*}}(%{{.*}}: tensor<i32>, %{{.*}}: tensor<i32>):
 // CHECK:       scf.yield %[[ZERO]], %[[ONE]]
@@ -1144,6 +1144,29 @@ func.func @while_duplicated_res() -> (i32, i32) {
 // CHECK:         }
 // CHECK:         return %[[RES]], %[[RES]] : i32, i32
 
+// -----
+
+func.func @while_licm(%arg1: i32, %arg2: i32, %arg3: i32) {
+  scf.while () : () -> () {
+    %val0 = arith.addi %arg1, %arg2 : i32
+    %val = arith.addi %val0, %arg3 : i32
+    %condition = "test.condition"(%val) : (i32) -> i1
+    scf.condition(%condition)
+  } do {
+  ^bb0():
+    "test.test"() : () -> ()
+    scf.yield
+  }
+  return
+}
+// CHECK-LABEL: @while_licm
+//  CHECK-SAME:  (%[[ARG1:.*]]: i32, %[[ARG2:.*]]: i32, %[[ARG3:.*]]: i32)
+//       CHECK:  %[[VAL0:.*]] = arith.addi %[[ARG1]], %[[ARG2]] : i32
+//       CHECK:  %[[VAL1:.*]] = arith.addi %[[VAL0]], %[[ARG3]] : i32
+//       CHECK:  scf.while
+//  CHECK-NEXT:  %[[COND:.*]] = "test.condition"(%[[VAL1]]) : (i32) -> i1
+//  CHECK-NEXT:  scf.condition(%[[COND]])
+
 
 // -----