[Mlir-commits] [mlir] [mlir][WIP] Implement replaceWithZeroTripCheck for scf.while (PR #80349)
Jerry Wu
llvmlistbot at llvm.org
Thu Feb 1 15:14:59 PST 2024
https://github.com/pzread updated https://github.com/llvm/llvm-project/pull/80349
>From 70f54b51bef87bde5e3f5ee067c0f2414d34e915 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 29f03571b10535ec84ee10e7bb22dbd0a7e6eb37 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 2bdb36416ddccc89b639d0849c57444199301e89 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