[Mlir-commits] [mlir] [mlir] Add transformation to wrap scf::while in zero-trip-check (PR #81050)

Jerry Wu llvmlistbot at llvm.org
Thu Feb 8 10:47:11 PST 2024


https://github.com/pzread updated https://github.com/llvm/llvm-project/pull/81050

>From 9f12915ac163df5d7cc521df370d8be38f4d0dba Mon Sep 17 00:00:00 2001
From: Jerry Wu <cheyuw at google.com>
Date: Thu, 1 Feb 2024 19:57:26 +0000
Subject: [PATCH 1/4] Add scf::wrapWhileLoopInZeroTripCheck

---
 .../mlir/Dialect/SCF/Transforms/Transforms.h  |  36 +++++
 .../lib/Dialect/SCF/Transforms/CMakeLists.txt |   1 +
 .../SCF/Transforms/WrapInZeroTripCheck.cpp    | 123 ++++++++++++++++++
 .../wrap-while-loop-in-zero-trip-check.mlir   |  40 ++++++
 mlir/test/lib/Dialect/SCF/CMakeLists.txt      |   1 +
 .../SCF/TestSCFWrapInZeroTripCheck.cpp        |  58 +++++++++
 mlir/tools/mlir-opt/mlir-opt.cpp              |   2 +
 7 files changed, 261 insertions(+)
 create mode 100644 mlir/lib/Dialect/SCF/Transforms/WrapInZeroTripCheck.cpp
 create mode 100644 mlir/test/Dialect/SCF/wrap-while-loop-in-zero-trip-check.mlir
 create mode 100644 mlir/test/lib/Dialect/SCF/TestSCFWrapInZeroTripCheck.cpp

diff --git a/mlir/include/mlir/Dialect/SCF/Transforms/Transforms.h b/mlir/include/mlir/Dialect/SCF/Transforms/Transforms.h
index e91f9e4469ab72..9c3fff864c9d91 100644
--- a/mlir/include/mlir/Dialect/SCF/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/SCF/Transforms/Transforms.h
@@ -30,6 +30,7 @@ namespace scf {
 class IfOp;
 class ForOp;
 class ParallelOp;
+class WhileOp;
 
 /// Fuses all adjacent scf.parallel operations with identical bounds and step
 /// into one scf.parallel operations. Uses a naive aliasing and dependency
@@ -181,6 +182,41 @@ FailureOr<ForOp> pipelineForLoop(RewriterBase &rewriter, ForOp forOp,
                                  const PipeliningOption &options,
                                  bool *modifiedIR = nullptr);
 
+/// Create zero-trip-check for a `while` op and return the replaced loop op
+/// wrapped in the check. The loop is rotated to avoid evaluating the condition
+/// twice. It turns:
+///
+///   scf.while (%arg0 = %init) : (i32) -> i64 {
+///     %val = .., %arg0 : i64
+///     %cond = arith.cmpi .., %arg0 : i32
+///     scf.condition(%cond) %val : i64
+///   } do {
+///   ^bb0(%arg1: i64):
+///     %next = .., %arg1 : i32
+///     scf.yield %next : i32
+///   }
+///
+///  into:
+///
+///   %pre_val = .., %init : i64
+///   %pre_cond = arith.cmpi .., %init : i32
+///   scf.if %pre_cond -> i64 {
+///     %res = scf.while (%arg1 = %va0) : (i64) -> i64 {
+///       %next = .., %arg1 : i32
+///       %val = .., %next : i64
+///       %cond = arith.cmpi .., %next : i32
+///       scf.condition(%cond) %val : i64
+///     } do {
+///     ^bb0(%arg2: i64):
+///       %scf.yield %arg2 : i32
+///     }
+///     scf.yield %res : i64
+///   } else {
+///     scf.yield %pre_val : i64
+///   }
+FailureOr<WhileOp> wrapWhileLoopInZeroTripCheck(WhileOp whileOp,
+                                                RewriterBase &rewriter);
+
 } // namespace scf
 } // namespace mlir
 
diff --git a/mlir/lib/Dialect/SCF/Transforms/CMakeLists.txt b/mlir/lib/Dialect/SCF/Transforms/CMakeLists.txt
index fdaeb2fad9afa4..e5494205e086ac 100644
--- a/mlir/lib/Dialect/SCF/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/SCF/Transforms/CMakeLists.txt
@@ -13,6 +13,7 @@ add_mlir_dialect_library(MLIRSCFTransforms
   ParallelLoopTiling.cpp
   StructuralTypeConversions.cpp
   TileUsingInterface.cpp
+  WrapInZeroTripCheck.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/SCF
diff --git a/mlir/lib/Dialect/SCF/Transforms/WrapInZeroTripCheck.cpp b/mlir/lib/Dialect/SCF/Transforms/WrapInZeroTripCheck.cpp
new file mode 100644
index 00000000000000..162df34405cc67
--- /dev/null
+++ b/mlir/lib/Dialect/SCF/Transforms/WrapInZeroTripCheck.cpp
@@ -0,0 +1,123 @@
+//===- WrapInZeroTripCheck.cpp - Loop transforms to add zero-trip-check ---===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/SCF/Transforms/Transforms.h"
+#include "mlir/IR/IRMapping.h"
+#include "mlir/IR/PatternMatch.h"
+
+using namespace mlir;
+
+/// Create zero-trip-check for a `while` op and return the replaced loop op
+/// wrapped in the check. The loop is rotated to avoid evaluating the condition
+/// twice.
+///
+/// Given an example below:
+///
+///   scf.while (%arg0 = %init) : (i32) -> i64 {
+///     %val = .., %arg0 : i64
+///     %cond = arith.cmpi .., %arg0 : i32
+///     scf.condition(%cond) %val : i64
+///   } do {
+///   ^bb0(%arg1: i64):
+///     %next = .., %arg1 : i32
+///     scf.yield %next : i32
+///   }
+///
+/// First clone before block to the front of the loop:
+///
+///   %pre_val = .., %init : i64
+///   %pre_cond = arith.cmpi .., %init : i32
+///   scf.while (%arg0 = %init) : (i32) -> i64 {
+///     %val = .., %arg0 : i64
+///     %cond = arith.cmpi .., %arg0 : i32
+///     scf.condition(%cond) %val : i64
+///   } do {
+///   ^bb0(%arg1: i64):
+///     %next = .., %arg1 : i32
+///     scf.yield %next : i32
+///   }
+///
+/// Create `if` op with the condition, rotate and move the loop into the else
+/// branch:
+///
+///   %pre_val = .., %init : i64
+///   %pre_cond = arith.cmpi .., %init : i32
+///   scf.if %pre_cond -> i64 {
+///     %res = scf.while (%arg1 = %va0) : (i64) -> i64 {
+///       // Original after block
+///       %next = .., %arg1 : i32
+///       // Original before block
+///       %val = .., %next : i64
+///       %cond = arith.cmpi .., %next : i32
+///       scf.condition(%cond) %val : i64
+///     } do {
+///     ^bb0(%arg2: i64):
+///       %scf.yield %arg2 : i32
+///     }
+///     scf.yield %res : i64
+///   } else {
+///     scf.yield %pre_val : i64
+///   }
+FailureOr<scf::WhileOp>
+mlir::scf::wrapWhileLoopInZeroTripCheck(scf::WhileOp whileOp,
+                                        RewriterBase &rewriter) {
+  IRMapping mapper;
+  Block *beforeBlock = whileOp.getBeforeBody();
+  // Clone before block before the loop for zero-trip-check.
+  for (auto [arg, init] :
+       llvm::zip_equal(beforeBlock->getArguments(), whileOp.getInits())) {
+    mapper.map(arg, init);
+  }
+  rewriter.setInsertionPoint(whileOp);
+  for (auto &op : *beforeBlock) {
+    if (isa<scf::ConditionOp>(op)) {
+      break;
+    }
+    // Safe to clone everything as in a single block all defs have been cloned
+    // and added to mapper in order.
+    rewriter.insert(op.clone(mapper));
+  }
+
+  auto condOp = whileOp.getConditionOp();
+  auto clonedCondition = mapper.lookupOrDefault(condOp.getCondition());
+  auto clonedCondArgs = llvm::map_to_vector(
+      condOp.getArgs(), [&](Value arg) { return mapper.lookupOrDefault(arg); });
+
+  // Create zero-trip-check and move the while loop in.
+  scf::WhileOp newLoopOp = nullptr;
+  auto ifOp = rewriter.create<scf::IfOp>(
+      whileOp->getLoc(), clonedCondition,
+      [&](OpBuilder &builder, Location loc) {
+        // Then runs the while loop.
+        newLoopOp = builder.create<scf::WhileOp>(
+            loc, whileOp.getResultTypes(), clonedCondArgs,
+            [&](OpBuilder &builder, Location loc, ValueRange args) {
+              // Rotate and move the loop body into before block.
+              auto newBlock = builder.getBlock();
+              rewriter.mergeBlocks(whileOp.getAfterBody(), newBlock, args);
+              auto yieldOp = cast<scf::YieldOp>(newBlock->getTerminator());
+              rewriter.mergeBlocks(whileOp.getBeforeBody(), newBlock,
+                                   yieldOp.getResults());
+              rewriter.eraseOp(yieldOp);
+            },
+            [&](OpBuilder &builder, Location loc, ValueRange args) {
+              // Pass-through values.
+              builder.create<scf::YieldOp>(loc, args);
+            });
+        builder.create<scf::YieldOp>(loc, newLoopOp.getResults());
+      },
+      [&](OpBuilder &builder, Location loc) {
+        // Else returns the results from zero-trip-check.
+        builder.create<scf::YieldOp>(loc, clonedCondArgs);
+      });
+
+  rewriter.replaceOp(whileOp, ifOp);
+
+  return newLoopOp;
+}
diff --git a/mlir/test/Dialect/SCF/wrap-while-loop-in-zero-trip-check.mlir b/mlir/test/Dialect/SCF/wrap-while-loop-in-zero-trip-check.mlir
new file mode 100644
index 00000000000000..b87c6003ddd310
--- /dev/null
+++ b/mlir/test/Dialect/SCF/wrap-while-loop-in-zero-trip-check.mlir
@@ -0,0 +1,40 @@
+// RUN: mlir-opt %s -test-wrap-scf-while-loop-in-zero-trip-check -split-input-file  | FileCheck %s
+
+func.func @wrap_while_loop_in_zero_trip_check(%bound : i32) -> i32 {
+  %cst0 = arith.constant 0 : i32
+  %cst5 = arith.constant 5 : i32
+  %res:2 = scf.while (%iter = %cst0) : (i32) -> (i32, i32) {
+    %cond = arith.cmpi slt, %iter, %bound : i32
+    %inv = arith.addi %bound, %cst5 : i32
+    scf.condition(%cond) %iter, %inv : i32, i32
+  } do {
+  ^bb0(%arg1: i32, %arg2: i32):
+    %next = arith.addi %arg1, %arg2 : i32
+    scf.yield %next : i32
+  }
+  return %res#0 : i32
+}
+
+// CHECK-LABEL: func.func @wrap_while_loop_in_zero_trip_check(
+// CHECK-SAME:      %[[ARG0:.*]]: i32) -> i32 {
+// CHECK-DAG:     %[[C0:.*]] = arith.constant 0 : i32
+// CHECK-DAG:     %[[C5:.*]] = arith.constant 5 : i32
+// CHECK-DAG:     %[[PRE_COND:.*]] = arith.cmpi slt, %[[C0]], %[[ARG0]] : i32
+// CHECK-DAG:     %[[PRE_INV:.*]] = arith.addi %[[ARG0]], %[[C5]] : i32
+// CHECK:         %[[IF:.*]]:2 = scf.if %[[PRE_COND]] -> (i32, i32) {
+// CHECK:           %[[WHILE:.*]]:2 = scf.while (
+// CHECK-SAME:          %[[ARG1:.*]] = %[[C0]], %[[ARG2:.*]] = %[[PRE_INV]]
+// CHECK-SAME:      ) : (i32, i32) -> (i32, i32) {
+// CHECK:             %[[NEXT:.*]] = arith.addi %[[ARG1]], %[[ARG2]] : i32
+// CHECK:             %[[COND:.*]] = arith.cmpi slt, %[[NEXT]], %[[ARG0]] : i32
+// CHECK:             %[[INV:.*]] = arith.addi %[[ARG0]], %[[C5]] : i32
+// CHECK:             scf.condition(%[[COND]]) %[[NEXT]], %[[INV]] : i32, i32
+// CHECK:           } do {
+// CHECK:           ^bb0(%[[ARG3:.*]]: i32, %[[ARG4:.*]]: i32):
+// CHECK:             scf.yield %[[ARG3]], %[[ARG4]] : i32, i32
+// CHECK:           }
+// CHECK:           scf.yield %[[WHILE]]#0, %[[WHILE]]#1 : i32, i32
+// CHECK:         } else {
+// CHECK:           scf.yield %[[C0]], %[[PRE_INV]] : i32, i32
+// CHECK:         }
+// CHECK:         return %[[IF]]#0 : i32
diff --git a/mlir/test/lib/Dialect/SCF/CMakeLists.txt b/mlir/test/lib/Dialect/SCF/CMakeLists.txt
index 22c2f2388de69b..d93bd559151829 100644
--- a/mlir/test/lib/Dialect/SCF/CMakeLists.txt
+++ b/mlir/test/lib/Dialect/SCF/CMakeLists.txt
@@ -3,6 +3,7 @@ add_mlir_library(MLIRSCFTestPasses
   TestLoopParametricTiling.cpp
   TestLoopUnrolling.cpp
   TestSCFUtils.cpp
+  TestSCFWrapInZeroTripCheck.cpp
   TestWhileOpBuilder.cpp
 
   EXCLUDE_FROM_LIBMLIR
diff --git a/mlir/test/lib/Dialect/SCF/TestSCFWrapInZeroTripCheck.cpp b/mlir/test/lib/Dialect/SCF/TestSCFWrapInZeroTripCheck.cpp
new file mode 100644
index 00000000000000..b51ef03288436f
--- /dev/null
+++ b/mlir/test/lib/Dialect/SCF/TestSCFWrapInZeroTripCheck.cpp
@@ -0,0 +1,58 @@
+//===- TestWrapInZeroTripCheck.cpp -- Passes to test SCF zero-trip-check --===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements the passes to test wrap-in-zero-trip-check transforms on
+// SCF loop ops.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/SCF/Transforms/Transforms.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+
+namespace {
+
+struct TestWrapWhileLoopInZeroTripCheck
+    : public PassWrapper<TestWrapWhileLoopInZeroTripCheck,
+                         OperationPass<func::FuncOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestWrapWhileLoopInZeroTripCheck)
+
+  StringRef getArgument() const final {
+    return "test-wrap-scf-while-loop-in-zero-trip-check";
+  }
+  StringRef getDescription() const final {
+    return "test scf::wrapWhileLoopInZeroTripCheck";
+  }
+
+  void runOnOperation() override {
+    func::FuncOp func = getOperation();
+    MLIRContext *context = &getContext();
+    IRRewriter rewriter(context);
+    func.walk([&](scf::WhileOp op) {
+      auto result = scf::wrapWhileLoopInZeroTripCheck(op, rewriter);
+      if (failed(result)) {
+        // Ignore not implemented failure in tests. The expected output should
+        // catch problems (e.g. transformation doesn't happen).
+      }
+    });
+  }
+};
+
+} // namespace
+
+namespace mlir {
+namespace test {
+void registerTestSCFWrapInZeroTripCheckPasses() {
+  PassRegistration<TestWrapWhileLoopInZeroTripCheck>();
+}
+} // namespace test
+} // namespace mlir
diff --git a/mlir/tools/mlir-opt/mlir-opt.cpp b/mlir/tools/mlir-opt/mlir-opt.cpp
index 428bdd9691e095..8ca16f17f66e8e 100644
--- a/mlir/tools/mlir-opt/mlir-opt.cpp
+++ b/mlir/tools/mlir-opt/mlir-opt.cpp
@@ -129,6 +129,7 @@ void registerTestPreparationPassWithAllowedMemrefResults();
 void registerTestRecursiveTypesPass();
 void registerTestSCFUtilsPass();
 void registerTestSCFWhileOpBuilderPass();
+void registerTestSCFWrapInZeroTripCheckPasses();
 void registerTestShapeMappingPass();
 void registerTestSliceAnalysisPass();
 void registerTestTensorCopyInsertionPass();
@@ -251,6 +252,7 @@ void registerTestPasses() {
   mlir::test::registerTestRecursiveTypesPass();
   mlir::test::registerTestSCFUtilsPass();
   mlir::test::registerTestSCFWhileOpBuilderPass();
+  mlir::test::registerTestSCFWrapInZeroTripCheckPasses();
   mlir::test::registerTestShapeMappingPass();
   mlir::test::registerTestSliceAnalysisPass();
   mlir::test::registerTestTensorCopyInsertionPass();

>From 338dc37d23ea7274e8d06a54d92ed55f7cff646c Mon Sep 17 00:00:00 2001
From: Jerry Wu <cheyuw at google.com>
Date: Wed, 7 Feb 2024 23:00:23 +0000
Subject: [PATCH 2/4] Update comments

---
 mlir/include/mlir/Dialect/SCF/Transforms/Transforms.h   | 6 +++---
 mlir/lib/Dialect/SCF/Transforms/WrapInZeroTripCheck.cpp | 5 ++---
 2 files changed, 5 insertions(+), 6 deletions(-)

diff --git a/mlir/include/mlir/Dialect/SCF/Transforms/Transforms.h b/mlir/include/mlir/Dialect/SCF/Transforms/Transforms.h
index 9c3fff864c9d91..1c1803113c2320 100644
--- a/mlir/include/mlir/Dialect/SCF/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/SCF/Transforms/Transforms.h
@@ -182,9 +182,9 @@ FailureOr<ForOp> pipelineForLoop(RewriterBase &rewriter, ForOp forOp,
                                  const PipeliningOption &options,
                                  bool *modifiedIR = nullptr);
 
-/// Create zero-trip-check for a `while` op and return the replaced loop op
-/// wrapped in the check. The loop is rotated to avoid evaluating the condition
-/// twice. It turns:
+/// Create zero-trip-check around a `while` op and return the new loop op in the
+/// check. The while loop is rotated to avoid evaluating the condition twice. It
+/// turns:
 ///
 ///   scf.while (%arg0 = %init) : (i32) -> i64 {
 ///     %val = .., %arg0 : i64
diff --git a/mlir/lib/Dialect/SCF/Transforms/WrapInZeroTripCheck.cpp b/mlir/lib/Dialect/SCF/Transforms/WrapInZeroTripCheck.cpp
index 162df34405cc67..0e1a15c2bdbda9 100644
--- a/mlir/lib/Dialect/SCF/Transforms/WrapInZeroTripCheck.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/WrapInZeroTripCheck.cpp
@@ -13,9 +13,8 @@
 
 using namespace mlir;
 
-/// Create zero-trip-check for a `while` op and return the replaced loop op
-/// wrapped in the check. The loop is rotated to avoid evaluating the condition
-/// twice.
+/// Create zero-trip-check around a `while` op and return the new loop op in the
+/// check. The while loop is rotated to avoid evaluating the condition twice.
 ///
 /// Given an example below:
 ///

>From 25ba59be015a4d39f7c150efb1b480e9feb72548 Mon Sep 17 00:00:00 2001
From: Jerry Wu <cheyuw at google.com>
Date: Thu, 8 Feb 2024 18:37:11 +0000
Subject: [PATCH 3/4] Address comments

---
 .../SCF/Transforms/WrapInZeroTripCheck.cpp    |  8 +-
 .../wrap-while-loop-in-zero-trip-check.mlir   | 84 +++++++++++++++++--
 .../SCF/TestSCFWrapInZeroTripCheck.cpp        | 10 +--
 3 files changed, 89 insertions(+), 13 deletions(-)

diff --git a/mlir/lib/Dialect/SCF/Transforms/WrapInZeroTripCheck.cpp b/mlir/lib/Dialect/SCF/Transforms/WrapInZeroTripCheck.cpp
index 0e1a15c2bdbda9..c27f46aad23d9f 100644
--- a/mlir/lib/Dialect/SCF/Transforms/WrapInZeroTripCheck.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/WrapInZeroTripCheck.cpp
@@ -66,6 +66,8 @@ using namespace mlir;
 FailureOr<scf::WhileOp>
 mlir::scf::wrapWhileLoopInZeroTripCheck(scf::WhileOp whileOp,
                                         RewriterBase &rewriter) {
+  OpBuilder::InsertionGuard insertion_guard(rewriter);
+
   IRMapping mapper;
   Block *beforeBlock = whileOp.getBeforeBody();
   // Clone before block before the loop for zero-trip-check.
@@ -83,9 +85,9 @@ mlir::scf::wrapWhileLoopInZeroTripCheck(scf::WhileOp whileOp,
     rewriter.insert(op.clone(mapper));
   }
 
-  auto condOp = whileOp.getConditionOp();
-  auto clonedCondition = mapper.lookupOrDefault(condOp.getCondition());
-  auto clonedCondArgs = llvm::map_to_vector(
+  scf::ConditionOp condOp = whileOp.getConditionOp();
+  Value clonedCondition = mapper.lookupOrDefault(condOp.getCondition());
+  SmallVector<Value> clonedCondArgs = llvm::map_to_vector(
       condOp.getArgs(), [&](Value arg) { return mapper.lookupOrDefault(arg); });
 
   // Create zero-trip-check and move the while loop in.
diff --git a/mlir/test/Dialect/SCF/wrap-while-loop-in-zero-trip-check.mlir b/mlir/test/Dialect/SCF/wrap-while-loop-in-zero-trip-check.mlir
index b87c6003ddd310..7eb3abe71fcf01 100644
--- a/mlir/test/Dialect/SCF/wrap-while-loop-in-zero-trip-check.mlir
+++ b/mlir/test/Dialect/SCF/wrap-while-loop-in-zero-trip-check.mlir
@@ -16,18 +16,18 @@ func.func @wrap_while_loop_in_zero_trip_check(%bound : i32) -> i32 {
 }
 
 // CHECK-LABEL: func.func @wrap_while_loop_in_zero_trip_check(
-// CHECK-SAME:      %[[ARG0:.*]]: i32) -> i32 {
+// CHECK-SAME:      %[[BOUND:.*]]: i32) -> i32 {
 // CHECK-DAG:     %[[C0:.*]] = arith.constant 0 : i32
 // CHECK-DAG:     %[[C5:.*]] = arith.constant 5 : i32
-// CHECK-DAG:     %[[PRE_COND:.*]] = arith.cmpi slt, %[[C0]], %[[ARG0]] : i32
-// CHECK-DAG:     %[[PRE_INV:.*]] = arith.addi %[[ARG0]], %[[C5]] : i32
+// CHECK-DAG:     %[[PRE_COND:.*]] = arith.cmpi slt, %[[C0]], %[[BOUND]] : i32
+// CHECK-DAG:     %[[PRE_INV:.*]] = arith.addi %[[BOUND]], %[[C5]] : i32
 // CHECK:         %[[IF:.*]]:2 = scf.if %[[PRE_COND]] -> (i32, i32) {
 // CHECK:           %[[WHILE:.*]]:2 = scf.while (
 // CHECK-SAME:          %[[ARG1:.*]] = %[[C0]], %[[ARG2:.*]] = %[[PRE_INV]]
 // CHECK-SAME:      ) : (i32, i32) -> (i32, i32) {
 // CHECK:             %[[NEXT:.*]] = arith.addi %[[ARG1]], %[[ARG2]] : i32
-// CHECK:             %[[COND:.*]] = arith.cmpi slt, %[[NEXT]], %[[ARG0]] : i32
-// CHECK:             %[[INV:.*]] = arith.addi %[[ARG0]], %[[C5]] : i32
+// CHECK:             %[[COND:.*]] = arith.cmpi slt, %[[NEXT]], %[[BOUND]] : i32
+// CHECK:             %[[INV:.*]] = arith.addi %[[BOUND]], %[[C5]] : i32
 // CHECK:             scf.condition(%[[COND]]) %[[NEXT]], %[[INV]] : i32, i32
 // CHECK:           } do {
 // CHECK:           ^bb0(%[[ARG3:.*]]: i32, %[[ARG4:.*]]: i32):
@@ -38,3 +38,77 @@ func.func @wrap_while_loop_in_zero_trip_check(%bound : i32) -> i32 {
 // CHECK:           scf.yield %[[C0]], %[[PRE_INV]] : i32, i32
 // CHECK:         }
 // CHECK:         return %[[IF]]#0 : i32
+
+// -----
+
+func.func @wrap_do_while_loop_in_zero_trip_check(%bound : i32) -> i32 {
+  %cst0 = arith.constant 0 : i32
+  %true = arith.constant true
+  %cst5 = arith.constant 5 : i32
+  %res = scf.while (%iter = %cst0, %arg0 = %true) : (i32, i1) -> i32 {
+    scf.condition(%arg0) %iter : i32
+  } do {
+  ^bb0(%arg1: i32):
+    %next = arith.addi %arg1, %cst5 : i32
+    %cond = arith.cmpi slt, %next, %bound : i32
+    scf.yield %next, %cond : i32, i1
+  }
+  return %res : i32
+}
+
+// CHECK-LABEL: func.func @wrap_do_while_loop_in_zero_trip_check(
+// CHECK-SAME:      %[[BOUND:.*]]: i32) -> i32 {
+// CHECK-DAG:     %[[C0:.*]] = arith.constant 0 : i32
+// CHECK-DAG:     %[[TRUE:.*]] = arith.constant true
+// CHECK-DAG:     %[[C5:.*]] = arith.constant 5 : i32
+// CHECK:         %[[IF:.*]] = scf.if %[[TRUE]] -> (i32) {
+// CHECK:           %[[WHILE:.*]] = scf.while (%[[ARG1:.*]] = %[[C0]]) : (i32) -> i32 {
+// CHECK:             %[[NEXT:.*]] = arith.addi %[[ARG1]], %[[C5]] : i32
+// CHECK:             %[[COND:.*]] = arith.cmpi slt, %[[NEXT]], %[[BOUND]] : i32
+// CHECK:             scf.condition(%[[COND]]) %[[NEXT]] : i32
+// CHECK:           } do {
+// CHECK:           ^bb0(%[[ARG2:.*]]: i32):
+// CHECK:             scf.yield %[[ARG2]] : i32
+// CHECK:           }
+// CHECK:           scf.yield %[[WHILE]] : i32
+// CHECK:         } else {
+// CHECK:           scf.yield %[[C0]] : i32
+// CHECK:         }
+// CHECK:         return %[[IF]] : i32
+
+// -----
+
+func.func @wrap_while_loop_with_minimal_after_block(%bound : i32) -> i32 {
+  %cst0 = arith.constant 0 : i32
+  %cst5 = arith.constant 5 : i32
+  %res = scf.while (%iter = %cst0) : (i32) -> i32 {
+    %next = arith.addi %iter, %cst5 : i32
+    %cond = arith.cmpi slt, %next, %bound : i32
+    scf.condition(%cond) %next : i32
+  } do {
+  ^bb0(%arg1: i32):
+    scf.yield %arg1 : i32
+  }
+  return %res : i32
+}
+
+// CHECK-LABEL: func.func @wrap_while_loop_with_minimal_after_block(
+// CHECK-SAME:      %[[BOUND:.*]]: i32) -> i32 {
+// CHECK-DAG:     %[[C0:.*]] = arith.constant 0 : i32
+// CHECK-DAG:     %[[C5:.*]] = arith.constant 5 : i32
+// CHECK:         %[[PRE_NEXT:.*]] = arith.addi %[[C0]], %[[C5]] : i32
+// CHECK:         %[[PRE_COND:.*]] = arith.cmpi slt, %[[PRE_NEXT]], %[[BOUND]] : i32
+// CHECK:         %[[IF:.*]] = scf.if %[[PRE_COND]] -> (i32) {
+// CHECK:           %[[WHILE:.*]] = scf.while (%[[ARG1:.*]] = %[[PRE_NEXT]]) : (i32) -> i32 {
+// CHECK:             %[[NEXT:.*]] = arith.addi %[[ARG1]], %[[C5]] : i32
+// CHECK:             %[[COND:.*]] = arith.cmpi slt, %[[NEXT]], %[[BOUND]] : i32
+// CHECK:             scf.condition(%[[COND]]) %[[NEXT]] : i32
+// CHECK:           } do {
+// CHECK:           ^bb0(%[[ARG2:.*]]: i32):
+// CHECK:             scf.yield %[[ARG2]] : i32
+// CHECK:           }
+// CHECK:           scf.yield %[[WHILE]] : i32
+// CHECK:         } else {
+// CHECK:           scf.yield %[[PRE_NEXT]] : i32
+// CHECK:         }
+// CHECK:         return %[[IF]] : i32
diff --git a/mlir/test/lib/Dialect/SCF/TestSCFWrapInZeroTripCheck.cpp b/mlir/test/lib/Dialect/SCF/TestSCFWrapInZeroTripCheck.cpp
index b51ef03288436f..d6975d544a1fac 100644
--- a/mlir/test/lib/Dialect/SCF/TestSCFWrapInZeroTripCheck.cpp
+++ b/mlir/test/lib/Dialect/SCF/TestSCFWrapInZeroTripCheck.cpp
@@ -38,11 +38,11 @@ struct TestWrapWhileLoopInZeroTripCheck
     MLIRContext *context = &getContext();
     IRRewriter rewriter(context);
     func.walk([&](scf::WhileOp op) {
-      auto result = scf::wrapWhileLoopInZeroTripCheck(op, rewriter);
-      if (failed(result)) {
-        // Ignore not implemented failure in tests. The expected output should
-        // catch problems (e.g. transformation doesn't happen).
-      }
+      FailureOr<scf::WhileOp> result =
+          scf::wrapWhileLoopInZeroTripCheck(op, rewriter);
+      // Ignore not implemented failure in tests. The expected output should
+      // catch problems (e.g. transformation doesn't happen).
+      (void)result;
     });
   }
 };

>From 9e7ac7dd78c8b99f5e1513e415097cedfc787334 Mon Sep 17 00:00:00 2001
From: Jerry Wu <cheyuw at google.com>
Date: Thu, 8 Feb 2024 18:46:59 +0000
Subject: [PATCH 4/4] Refactor

---
 .../SCF/Transforms/WrapInZeroTripCheck.cpp    | 39 ++++++++++---------
 1 file changed, 21 insertions(+), 18 deletions(-)

diff --git a/mlir/lib/Dialect/SCF/Transforms/WrapInZeroTripCheck.cpp b/mlir/lib/Dialect/SCF/Transforms/WrapInZeroTripCheck.cpp
index c27f46aad23d9f..65793217a4c24c 100644
--- a/mlir/lib/Dialect/SCF/Transforms/WrapInZeroTripCheck.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/WrapInZeroTripCheck.cpp
@@ -90,31 +90,34 @@ mlir::scf::wrapWhileLoopInZeroTripCheck(scf::WhileOp whileOp,
   SmallVector<Value> clonedCondArgs = llvm::map_to_vector(
       condOp.getArgs(), [&](Value arg) { return mapper.lookupOrDefault(arg); });
 
+  // Create rotated while loop.
+  auto newLoopOp = rewriter.create<scf::WhileOp>(
+      whileOp.getLoc(), whileOp.getResultTypes(), clonedCondArgs,
+      [&](OpBuilder &builder, Location loc, ValueRange args) {
+        // Rotate and move the loop body into before block.
+        auto newBlock = builder.getBlock();
+        rewriter.mergeBlocks(whileOp.getAfterBody(), newBlock, args);
+        auto yieldOp = cast<scf::YieldOp>(newBlock->getTerminator());
+        rewriter.mergeBlocks(whileOp.getBeforeBody(), newBlock,
+                             yieldOp.getResults());
+        rewriter.eraseOp(yieldOp);
+      },
+      [&](OpBuilder &builder, Location loc, ValueRange args) {
+        // Pass-through values.
+        builder.create<scf::YieldOp>(loc, args);
+      });
+
   // Create zero-trip-check and move the while loop in.
-  scf::WhileOp newLoopOp = nullptr;
   auto ifOp = rewriter.create<scf::IfOp>(
-      whileOp->getLoc(), clonedCondition,
+      whileOp.getLoc(), clonedCondition,
       [&](OpBuilder &builder, Location loc) {
         // Then runs the while loop.
-        newLoopOp = builder.create<scf::WhileOp>(
-            loc, whileOp.getResultTypes(), clonedCondArgs,
-            [&](OpBuilder &builder, Location loc, ValueRange args) {
-              // Rotate and move the loop body into before block.
-              auto newBlock = builder.getBlock();
-              rewriter.mergeBlocks(whileOp.getAfterBody(), newBlock, args);
-              auto yieldOp = cast<scf::YieldOp>(newBlock->getTerminator());
-              rewriter.mergeBlocks(whileOp.getBeforeBody(), newBlock,
-                                   yieldOp.getResults());
-              rewriter.eraseOp(yieldOp);
-            },
-            [&](OpBuilder &builder, Location loc, ValueRange args) {
-              // Pass-through values.
-              builder.create<scf::YieldOp>(loc, args);
-            });
+        rewriter.moveOpBefore(newLoopOp, builder.getInsertionBlock(),
+                              builder.getInsertionPoint());
         builder.create<scf::YieldOp>(loc, newLoopOp.getResults());
       },
       [&](OpBuilder &builder, Location loc) {
-        // Else returns the results from zero-trip-check.
+        // Else returns the results from precondition.
         builder.create<scf::YieldOp>(loc, clonedCondArgs);
       });
 



More information about the Mlir-commits mailing list