[Mlir-commits] [mlir] 81a79ee - [mlir] Add function for checking if a block is inside a loop
Tom Eccles
llvmlistbot at llvm.org
Fri Feb 10 08:25:58 PST 2023
Author: Tom Eccles
Date: 2023-02-10T16:14:17Z
New Revision: 81a79ee446fe499e217144ae1cd505ede9021712
URL: https://github.com/llvm/llvm-project/commit/81a79ee446fe499e217144ae1cd505ede9021712
DIFF: https://github.com/llvm/llvm-project/commit/81a79ee446fe499e217144ae1cd505ede9021712.diff
LOG: [mlir] Add function for checking if a block is inside a loop
This function returns whether a block is nested inside of a loop. There
can be three kinds of loop:
1) The block is nested inside of a LoopLikeOpInterface
2) The block is nested inside another block which is in a loop
3) There is a cycle in the control flow graph
This will be useful for Flang's stack arrays pass, which moves array
allocations from the heap to the stack. Special handling is needed when
allocations occur inside of loops to ensure additional stack space is
not allocated on each loop iteration.
Differential Revision: https://reviews.llvm.org/D141401
Added:
mlir/test/Interfaces/LoopLikeInterface/test-block-loop.mlir
mlir/test/lib/Interfaces/LoopLikeInterface/CMakeLists.txt
mlir/test/lib/Interfaces/LoopLikeInterface/TestBlockInLoop.cpp
Modified:
mlir/include/mlir/Interfaces/LoopLikeInterface.td
mlir/lib/Interfaces/LoopLikeInterface.cpp
mlir/test/lib/Interfaces/CMakeLists.txt
mlir/tools/mlir-opt/CMakeLists.txt
mlir/tools/mlir-opt/mlir-opt.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Interfaces/LoopLikeInterface.td b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
index 95ba4511c2cd2..df7315c4ecf04 100644
--- a/mlir/include/mlir/Interfaces/LoopLikeInterface.td
+++ b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
@@ -97,6 +97,13 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
}]
>,
];
+
+ let extraClassDeclaration = [{
+ /// Returns if a block is inside a loop (within the current function). This
+ /// can either be because the block is nested inside a LoopLikeInterface, or
+ /// because the control flow graph is cyclic
+ static bool blockIsInLoop(Block *block);
+ }];
}
#endif // MLIR_INTERFACES_LOOPLIKEINTERFACE
diff --git a/mlir/lib/Interfaces/LoopLikeInterface.cpp b/mlir/lib/Interfaces/LoopLikeInterface.cpp
index 40d505228cb0d..49f9648cea806 100644
--- a/mlir/lib/Interfaces/LoopLikeInterface.cpp
+++ b/mlir/lib/Interfaces/LoopLikeInterface.cpp
@@ -7,8 +7,47 @@
//===----------------------------------------------------------------------===//
#include "mlir/Interfaces/LoopLikeInterface.h"
+#include "mlir/IR/FunctionInterfaces.h"
+#include "llvm/ADT/DenseSet.h"
using namespace mlir;
/// Include the definitions of the loop-like interfaces.
#include "mlir/Interfaces/LoopLikeInterface.cpp.inc"
+
+bool LoopLikeOpInterface::blockIsInLoop(Block *block) {
+ Operation *parent = block->getParentOp();
+
+ // The block could be inside a loop-like operation
+ if (isa<LoopLikeOpInterface>(parent) ||
+ parent->getParentOfType<LoopLikeOpInterface>())
+ return true;
+
+ // This block might be nested inside another block, which is in a loop
+ if (!isa<FunctionOpInterface>(parent))
+ if (mlir::Block *parentBlock = parent->getBlock())
+ if (blockIsInLoop(parentBlock))
+ return true;
+
+ // Or the block could be inside a control flow graph loop:
+ // A block is in a control flow graph loop if it can reach itself in a graph
+ // traversal
+ DenseSet<Block *> visited;
+ SmallVector<Block *> stack;
+ stack.push_back(block);
+ while (!stack.empty()) {
+ Block *current = stack.pop_back_val();
+ auto [it, inserted] = visited.insert(current);
+ if (!inserted) {
+ // loop detected
+ if (current == block)
+ return true;
+ continue;
+ }
+
+ stack.reserve(stack.size() + current->getNumSuccessors());
+ for (Block *successor : current->getSuccessors())
+ stack.push_back(successor);
+ }
+ return false;
+}
diff --git a/mlir/test/Interfaces/LoopLikeInterface/test-block-loop.mlir b/mlir/test/Interfaces/LoopLikeInterface/test-block-loop.mlir
new file mode 100644
index 0000000000000..2d4a775f0e931
--- /dev/null
+++ b/mlir/test/Interfaces/LoopLikeInterface/test-block-loop.mlir
@@ -0,0 +1,143 @@
+// RUN: mlir-opt %s --mlir-disable-threading -test-block-is-in-loop 2>&1 | FileCheck %s
+
+module {
+ // Test function with only one bb
+ func.func @simple() {
+ func.return
+ }
+// CHECK: Block is not in a loop
+// CHECK-NEXT: ^bb0:
+
+ // Test simple loop bb0 -> bb0
+ func.func @loopForever() {
+ ^bb0:
+ cf.br ^bb1
+ ^bb1:
+ cf.br ^bb1
+ }
+// CHECK: Block is not in a loop
+// CHECK-NEXT: ^bb0:
+// CHECK: Block is in a loop
+// CHECK-NEXT: ^bb1:
+
+ // Test bb0 -> bb1 -> bb2 -> bb1
+ func.func @loopForever2() {
+ ^bb0:
+ cf.br ^bb1
+ ^bb1:
+ cf.br ^bb2
+ ^bb2:
+ cf.br ^bb1
+ }
+// CHECK: Block is not in a loop
+// CHECK-NEXT: ^bb0:
+// CHECK: Block is in a loop
+// CHECK-NEXT: ^bb1:
+// CHECK: Block is in a loop
+// CHECK-NEXT: ^bb2:
+
+ // Test conditional branch without loop
+ // bb0 -> bb1 -> {bb2, bb3}
+ func.func @noLoop(%arg0: i1) {
+ cf.br ^bb1
+ ^bb1:
+ cf.cond_br %arg0, ^bb2, ^bb3
+ ^bb2:
+ func.return
+ ^bb3:
+ func.return
+ }
+// CHECK: Block is not in a loop
+// CHECK-NEXT: ^bb0(%arg0: i1)
+// CHECK: Block is not in a loop
+// CHECK-NEXT: ^bb1:
+// CHECK: Block is not in a loop
+// CHECK-NEXT: ^bb2:
+// CHECK: Block is not in a loop
+// CHECK-NEXT: ^bb3:
+
+ // test multiple loops
+ // bb0 -> bb1 -> bb2 -> bb3 { -> bb2} -> bb4 { -> bb1 } -> bb5
+ func.func @multipleLoops(%arg0: i1, %arg1: i1) {
+ cf.br ^bb1
+ ^bb1:
+ cf.br ^bb2
+ ^bb2:
+ cf.br ^bb3
+ ^bb3:
+ cf.cond_br %arg0, ^bb4, ^bb2
+ ^bb4:
+ cf.cond_br %arg1, ^bb1, ^bb5
+ ^bb5:
+ return
+ }
+// CHECK: Block is not in a loop
+// CHECK-NEXT: ^bb0(%arg0: i1, %arg1: i1)
+// CHECK: Block is in a loop
+// CHECK-NEXT: ^bb1:
+// CHECK: Block is in a loop
+// CHECK-NEXT: ^bb2:
+// CHECK: Block is in a loop
+// CHECK-NEXT: ^bb3:
+// CHECK: Block is in a loop
+// CHECK-NEXT: ^bb4:
+// CHECK: Block is not in a loop
+// CHECK-NEXT: ^bb5:
+
+ // test derived from real Flang output
+ func.func @_QPblockTest0(%arg0: i1, %arg1: i1) {
+ cf.br ^bb1
+ ^bb1: // 2 preds: ^bb0, ^bb4
+ cf.cond_br %arg0, ^bb2, ^bb5
+ ^bb2: // pred: ^bb1
+ cf.cond_br %arg1, ^bb3, ^bb4
+ ^bb3: // pred: ^bb2
+ return
+ ^bb4: // pred: ^bb2
+ cf.br ^bb1
+ ^bb5: // pred: ^bb1
+ return
+ }
+// CHECK: Block is not in a loop
+// CHECK-NEXT: ^bb0(%arg0: i1, %arg1: i1)
+// CHECK: Block is in a loop
+// CHECK-NEXT: ^bb1:
+// CHECK: Block is in a loop
+// CHECK-NEXT: ^bb2:
+// CHECK: Block is not in a loop
+// CHECK-NEXT: ^bb3:
+// CHECK: Block is in a loop
+// CHECK-NEXT: ^bb4:
+// CHECK: Block is not in a loop
+// CHECK-NEXT: ^bb5:
+
+// check nested blocks
+ func.func @check_alloc_in_loop(%counter : i64) {
+ cf.br ^bb1(%counter: i64)
+ ^bb1(%lv : i64):
+ %cm1 = arith.constant -1 : i64
+ %rem = arith.addi %lv, %cm1 : i64
+ %zero = arith.constant 0 : i64
+ %p = arith.cmpi eq, %rem, %zero : i64
+ cf.cond_br %p, ^bb3, ^bb2
+ ^bb2:
+ scf.execute_region -> () {
+ %c1 = arith.constant 1 : i64
+ scf.yield
+ }
+ cf.br ^bb1(%rem: i64)
+ ^bb3:
+ return
+ }
+// CHECK: Block is not in a loop
+// CHECK-NEXT: ^bb0(%arg0: i64):
+// CHECK: Block is in a loop
+// CHECK-NEXT: ^bb1(%0: i64)
+// CHECK: Block is in a loop
+// CHECK-NEXT: ^bb0:
+// CHECK-NEXT: %c1_i64
+// CHECK: Block is in a loop
+// CHECK-NEXT: ^bb2:
+// CHECK: Block is not in a loop
+// CHECK-NEXT: ^bb3:
+}
diff --git a/mlir/test/lib/Interfaces/CMakeLists.txt b/mlir/test/lib/Interfaces/CMakeLists.txt
index 4a0567ab46423..6a21ed10eec6f 100644
--- a/mlir/test/lib/Interfaces/CMakeLists.txt
+++ b/mlir/test/lib/Interfaces/CMakeLists.txt
@@ -1 +1,2 @@
+add_subdirectory(LoopLikeInterface)
add_subdirectory(TilingInterface)
diff --git a/mlir/test/lib/Interfaces/LoopLikeInterface/CMakeLists.txt b/mlir/test/lib/Interfaces/LoopLikeInterface/CMakeLists.txt
new file mode 100644
index 0000000000000..f20219e00cb86
--- /dev/null
+++ b/mlir/test/lib/Interfaces/LoopLikeInterface/CMakeLists.txt
@@ -0,0 +1,10 @@
+add_mlir_library(MLIRLoopLikeInterfaceTestPasses
+ TestBlockInLoop.cpp
+
+ EXCLUDE_FROM_LIBMLIR
+
+ LINK_LIBS PUBLIC
+ MLIRPass
+ MLIRLoopLikeInterface
+ MLIRFuncDialect
+ )
diff --git a/mlir/test/lib/Interfaces/LoopLikeInterface/TestBlockInLoop.cpp b/mlir/test/lib/Interfaces/LoopLikeInterface/TestBlockInLoop.cpp
new file mode 100644
index 0000000000000..d695195064a6a
--- /dev/null
+++ b/mlir/test/lib/Interfaces/LoopLikeInterface/TestBlockInLoop.cpp
@@ -0,0 +1,47 @@
+//===- TestBlockInLoop.cpp - Pass to test mlir::blockIsInLoop -------------===//
+//
+// 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/Func/IR/FuncOps.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/Interfaces/LoopLikeInterface.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/Support/raw_ostream.h"
+
+using namespace mlir;
+
+namespace {
+/// This is a test pass that tests Blocks's isInLoop method by checking if each
+/// block in a function is in a loop and outputing if it is
+struct IsInLoopPass
+ : public PassWrapper<IsInLoopPass, OperationPass<func::FuncOp>> {
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(IsInLoopPass)
+
+ StringRef getArgument() const final { return "test-block-is-in-loop"; }
+ StringRef getDescription() const final {
+ return "Test mlir::blockIsInLoop()";
+ }
+
+ void runOnOperation() override {
+ mlir::func::FuncOp func = getOperation();
+ func.walk([](mlir::Block *block) {
+ llvm::outs() << "Block is ";
+ if (LoopLikeOpInterface::blockIsInLoop(block))
+ llvm::outs() << "in a loop\n";
+ else
+ llvm::outs() << "not in a loop\n";
+ block->print(llvm::outs());
+ llvm::outs() << "\n";
+ });
+ }
+};
+
+} // namespace
+
+namespace mlir {
+void registerLoopLikeInterfaceTestPasses() { PassRegistration<IsInLoopPass>(); }
+} // namespace mlir
diff --git a/mlir/tools/mlir-opt/CMakeLists.txt b/mlir/tools/mlir-opt/CMakeLists.txt
index 0c81891c80113..fb7b4e77c5ec8 100644
--- a/mlir/tools/mlir-opt/CMakeLists.txt
+++ b/mlir/tools/mlir-opt/CMakeLists.txt
@@ -21,6 +21,7 @@ if(MLIR_INCLUDE_TESTS)
MLIRFuncTestPasses
MLIRGPUTestPasses
MLIRLinalgTestPasses
+ MLIRLoopLikeInterfaceTestPasses
MLIRMathTestPasses
MLIRMemRefTestPasses
MLIRNVGPUTestPasses
diff --git a/mlir/tools/mlir-opt/mlir-opt.cpp b/mlir/tools/mlir-opt/mlir-opt.cpp
index 33f16241b2e49..5a07004baf38f 100644
--- a/mlir/tools/mlir-opt/mlir-opt.cpp
+++ b/mlir/tools/mlir-opt/mlir-opt.cpp
@@ -33,6 +33,7 @@ void registerConvertToTargetEnvPass();
void registerCloneTestPasses();
void registerPassManagerTestPass();
void registerPrintSpirvAvailabilityPass();
+void registerLoopLikeInterfaceTestPasses();
void registerShapeFunctionTestPasses();
void registerSideEffectTestPasses();
void registerSliceAnalysisTestPass();
@@ -140,6 +141,7 @@ void registerTestPasses() {
registerConvertToTargetEnvPass();
registerPassManagerTestPass();
registerPrintSpirvAvailabilityPass();
+ registerLoopLikeInterfaceTestPasses();
registerShapeFunctionTestPasses();
registerSideEffectTestPasses();
registerSliceAnalysisTestPass();
More information about the Mlir-commits
mailing list