[Mlir-commits] [mlir] [mlir][WIP] Implement replaceWithZeroTripCheck for scf.while (PR #80349)

Jerry Wu llvmlistbot at llvm.org
Thu Feb 1 15:13:54 PST 2024


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

>From c07dce70b3337ff42e1ea28dd1c6294cdb8f32ef 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/3] Add replaceWithZeroTripCheck to LoopLikeOpInterface

---
 .../mlir/Interfaces/LoopLikeInterface.td      | 22 +++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/mlir/include/mlir/Interfaces/LoopLikeInterface.td b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
index e2ac85a3f7725..77409cb3a8274 100644
--- a/mlir/include/mlir/Interfaces/LoopLikeInterface.td
+++ b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
@@ -220,6 +220,28 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
       /*defaultImplementation=*/[{
         return ::mlir::failure();
       }]
+    >,
+    InterfaceMethod<[{
+        Add a zero-trip-check around the loop to check if the loop body is ever
+        run and return the new loop inside the check. The loop body is moved
+        over to the new loop. Returns "failure" if the loop doesn't support
+        this transformation.
+
+        After the transformation, the ops inserted to the parent region of the
+        loop are guaranteed to be run only if the loop body runs at least one
+        iteration.
+
+        Note: Ops in the loop body might be rearranged because of loop rotating
+        to maintain the semantic. Terminators might be removed/added during this
+        transformation.
+      }],
+      /*retTy=*/"::mlir::FailureOr<::mlir::LoopLikeOpInterface>",
+      /*methodName=*/"replaceWithZeroTripCheck",
+      /*args=*/(ins "::mlir::RewriterBase &":$rewriter),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return ::mlir::failure();
+      }]
     >
   ];
 

>From 18d70d975315e871ab38d575b69165a4a3612e7d Mon Sep 17 00:00:00 2001
From: Jerry Wu <cheyuw at google.com>
Date: Thu, 1 Feb 2024 21:51:09 +0000
Subject: [PATCH 2/3] Implement replaceWithZeroTripCheck for scf.while

---
 mlir/include/mlir/Dialect/SCF/IR/SCFOps.td |   4 +-
 mlir/lib/Dialect/SCF/IR/SCF.cpp            | 104 +++++++++++++++++++++
 2 files changed, 107 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
index b3d085bfff1af..7873020c5a181 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -939,7 +939,9 @@ def WhileOp : SCF_Op<"while",
     [DeclareOpInterfaceMethods<RegionBranchOpInterface,
         ["getEntrySuccessorOperands"]>,
      DeclareOpInterfaceMethods<LoopLikeOpInterface,
-        ["getRegionIterArgs", "getYieldedValuesMutable"]>,
+        ["getRegionIterArgs",
+         "getYieldedValuesMutable",
+         "replaceWithZeroTripCheck"]>,
      RecursiveMemoryEffects, SingleBlock]> {
   let summary = "a generic 'while' loop";
   let description = [{
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 9822ee522c6ed..58f13204f95cb 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -3254,6 +3254,110 @@ LogicalResult scf::WhileOp::verify() {
   return success(afterTerminator != nullptr);
 }
 
+/// Create zero-trip-check for a `while` op. 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:
+///
+///   %val0 = .., %init : i64
+///   %cond0 = 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:
+///
+///   %val0 = .., %init : i64
+///   %cond0 = arith.cmpi .., %init : i32
+///   scf.if %cond0 -> 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 %val0 : i64
+///   }
+FailureOr<LoopLikeOpInterface>
+scf::WhileOp::replaceWithZeroTripCheck(RewriterBase &rewriter) {
+  IRMapping mapper;
+  Block *beforeBlock = this->getBeforeBody();
+  // Clone before block before the loop for zero-trip-check.
+  for (auto [arg, init] :
+       llvm::zip_equal(beforeBlock->getArguments(), this->getInits())) {
+    mapper.map(arg, init);
+  }
+  rewriter.setInsertionPoint(*this);
+  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 = this->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 newLoop = nullptr;
+  auto ifOp = rewriter.create<scf::IfOp>(
+      this->getLoc(), clonedCondition,
+      [&](OpBuilder &builder, Location loc) {
+        // Then runs the while loop.
+        newLoop = builder.create<scf::WhileOp>(
+            loc, this->getResultTypes(), clonedCondArgs,
+            [&](OpBuilder &builder, Location loc, ValueRange args) {
+              // Rotate and move the loop body into before block.
+              auto newBlock = builder.getBlock();
+              rewriter.mergeBlocks(this->getAfterBody(), newBlock, args);
+              auto yieldOp = cast<scf::YieldOp>(newBlock->getTerminator());
+              rewriter.mergeBlocks(this->getBeforeBody(), newBlock,
+                                   yieldOp.getResults());
+              rewriter.eraseOp(yieldOp);
+            },
+            [&](OpBuilder &builder, Location loc, ValueRange args) {
+              // Pass-through values in after block.
+              builder.create<scf::YieldOp>(loc, args);
+            });
+        builder.create<scf::YieldOp>(loc, newLoop.getResults());
+      },
+      [&](OpBuilder &builder, Location loc) {
+        // Else returns the results from zero-trip-check.
+        builder.create<scf::YieldOp>(loc, clonedCondArgs);
+      });
+
+  rewriter.replaceOp(*this, ifOp);
+
+  return cast<LoopLikeOpInterface>(newLoop.getOperation());
+}
+
 namespace {
 /// Replace uses of the condition within the do block with true, since otherwise
 /// the block would not be evaluated.

>From 19084483b2f6f3dda1242d6973fccb2d75ce277b Mon Sep 17 00:00:00 2001
From: Jerry Wu <cheyuw at google.com>
Date: Thu, 1 Feb 2024 23:13:37 +0000
Subject: [PATCH 3/3] Add tests

---
 .../SCF/while-loop-zero-trip-check.mlir       | 40 +++++++++++++
 mlir/test/lib/Dialect/SCF/CMakeLists.txt      |  1 +
 .../lib/Dialect/SCF/TestLoopZeroTripCheck.cpp | 59 +++++++++++++++++++
 mlir/tools/mlir-opt/mlir-opt.cpp              |  2 +
 4 files changed, 102 insertions(+)
 create mode 100644 mlir/test/Dialect/SCF/while-loop-zero-trip-check.mlir
 create mode 100644 mlir/test/lib/Dialect/SCF/TestLoopZeroTripCheck.cpp

diff --git a/mlir/test/Dialect/SCF/while-loop-zero-trip-check.mlir b/mlir/test/Dialect/SCF/while-loop-zero-trip-check.mlir
new file mode 100644
index 0000000000000..f5f0a55ad4f16
--- /dev/null
+++ b/mlir/test/Dialect/SCF/while-loop-zero-trip-check.mlir
@@ -0,0 +1,40 @@
+// RUN: mlir-opt %s -test-scf-while-zero-trip-check -split-input-file  | FileCheck %s
+
+func.func @replace_scf_while_with_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 @replace_scf_while_with_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 22c2f2388de69..d704fe6fe81e3 100644
--- a/mlir/test/lib/Dialect/SCF/CMakeLists.txt
+++ b/mlir/test/lib/Dialect/SCF/CMakeLists.txt
@@ -2,6 +2,7 @@
 add_mlir_library(MLIRSCFTestPasses
   TestLoopParametricTiling.cpp
   TestLoopUnrolling.cpp
+  TestLoopZeroTripCheck.cpp
   TestSCFUtils.cpp
   TestWhileOpBuilder.cpp
 
diff --git a/mlir/test/lib/Dialect/SCF/TestLoopZeroTripCheck.cpp b/mlir/test/lib/Dialect/SCF/TestLoopZeroTripCheck.cpp
new file mode 100644
index 0000000000000..8e203b82a7fae
--- /dev/null
+++ b/mlir/test/lib/Dialect/SCF/TestLoopZeroTripCheck.cpp
@@ -0,0 +1,59 @@
+//===- TestLoopZeroTripCheck.cpp -- Pass to test replaceWithZeroTripCheck -===//
+//
+// 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 replaceWithZeroTripCheck for SCF
+// dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+
+namespace {
+
+struct TestSCFWhileZeroTripCheckPass
+    : public PassWrapper<TestSCFWhileZeroTripCheckPass,
+                         OperationPass<func::FuncOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestSCFWhileZeroTripCheckPass)
+
+  StringRef getArgument() const final {
+    return "test-scf-while-zero-trip-check";
+  }
+  StringRef getDescription() const final {
+    return "test replaceWithZeroTripCheck of scf.while";
+  }
+  explicit TestSCFWhileZeroTripCheckPass() = default;
+  TestSCFWhileZeroTripCheckPass(const TestSCFWhileZeroTripCheckPass &pass)
+      : PassWrapper(pass) {}
+
+  void runOnOperation() override {
+    func::FuncOp func = getOperation();
+    MLIRContext *context = &getContext();
+    IRRewriter rewriter(context);
+    func.walk([&](scf::WhileOp op) {
+      auto result = op.replaceWithZeroTripCheck(rewriter);
+      if (failed(result)) {
+        signalPassFailure();
+      }
+    });
+  }
+};
+
+} // namespace
+
+namespace mlir {
+namespace test {
+void registerTestLoopZeroTripCheckPass() {
+  PassRegistration<TestSCFWhileZeroTripCheckPass>();
+}
+} // namespace test
+} // namespace mlir
diff --git a/mlir/tools/mlir-opt/mlir-opt.cpp b/mlir/tools/mlir-opt/mlir-opt.cpp
index 428bdd9691e09..6ac3283bcb9d1 100644
--- a/mlir/tools/mlir-opt/mlir-opt.cpp
+++ b/mlir/tools/mlir-opt/mlir-opt.cpp
@@ -110,6 +110,7 @@ void registerTestLoopFusion();
 void registerTestCFGLoopInfoPass();
 void registerTestLoopMappingPass();
 void registerTestLoopUnrollingPass();
+void registerTestLoopZeroTripCheckPass();
 void registerTestLowerToLLVM();
 void registerTestLowerToNVVM();
 void registerTestMakeIsolatedFromAbovePass();
@@ -234,6 +235,7 @@ void registerTestPasses() {
   mlir::test::registerTestCFGLoopInfoPass();
   mlir::test::registerTestLoopMappingPass();
   mlir::test::registerTestLoopUnrollingPass();
+  mlir::test::registerTestLoopZeroTripCheckPass();
   mlir::test::registerTestLowerToLLVM();
   mlir::test::registerTestMakeIsolatedFromAbovePass();
   mlir::test::registerTestMatchReductionPass();



More information about the Mlir-commits mailing list