Skip to content

Commit dcee187

Browse files
tblahtblah
committed
[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
1 parent de4321c commit dcee187

File tree

8 files changed

+251
-0
lines changed

8 files changed

+251
-0
lines changed

‎mlir/include/mlir/Interfaces/LoopLikeInterface.td

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,14 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
9797
}]
9898
>,
9999
];
100+
101+
let extraClassDeclaration = [{
102+
/// Returns if a block is inside a loop (within the current function). This
103+
/// can be either because the block is nested inside a LoopLikeInterface or
104+
/// because the block is nested inside a LoopLikeInterface or because
105+
/// the control flow graph is cyclic
106+
static bool blockIsInLoop(Block *block);
107+
}];
100108
}
101109

102110
#endif // MLIR_INTERFACES_LOOPLIKEINTERFACE

‎mlir/lib/Interfaces/LoopLikeInterface.cpp

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,48 @@
77
//===----------------------------------------------------------------------===//
88

99
#include "mlir/Interfaces/LoopLikeInterface.h"
10+
#include "mlir/Dialect/Func/IR/FuncOps.h"
11+
#include "llvm/ADT/DenseSet.h"
1012

1113
using namespace mlir;
1214

1315
/// Include the definitions of the loop-like interfaces.
1416
#include "mlir/Interfaces/LoopLikeInterface.cpp.inc"
17+
18+
bool LoopLikeOpInterface::blockIsInLoop(Block *block) {
19+
Operation *parent = block->getParentOp();
20+
21+
// The block could be inside a loop-like operation
22+
if (isa<LoopLikeOpInterface>(parent) ||
23+
parent->getParentOfType<LoopLikeOpInterface>())
24+
return true;
25+
26+
// This block might be nested inside another block, which is in a loop
27+
if (!isa<FunctionOpInterface>(parent)) {
28+
if (blockIsInLoop(parent->getBlock())) {
29+
return true;
30+
}
31+
}
32+
33+
// Or the block could be inside a control flow graph loop:
34+
// A block is in a control flow graph loop if it can reach itself in a graph
35+
// traversal
36+
DenseSet<Block *> visited;
37+
SmallVector<Block *> stack;
38+
stack.push_back(block);
39+
while (!stack.empty()) {
40+
Block *current = stack.pop_back_val();
41+
auto [it, inserted] = visited.insert(current);
42+
if (!inserted) {
43+
// loop detected
44+
if (current == block)
45+
return true;
46+
continue;
47+
}
48+
49+
stack.reserve(stack.size() + current->getNumSuccessors());
50+
for (Block *successor : current->getSuccessors())
51+
stack.push_back(successor);
52+
}
53+
return false;
54+
}

‎mlir/test/Interfaces/LoopLikeInterface/test-block-loop.mlir

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,143 @@
1+
// RUN: mlir-opt %s --mlir-disable-threading -test-block-is-in-loop 2>&1 | FileCheck %s
2+
3+
module {
4+
// Test function with only one bb
5+
func.func @simple() {
6+
func.return
7+
}
8+
// CHECK: Block is not in a loop
9+
// CHECK-NEXT: ^bb0:
10+
11+
// Test simple loop bb0 -> bb0
12+
func.func @loopForever() {
13+
^bb0:
14+
cf.br ^bb1
15+
^bb1:
16+
cf.br ^bb1
17+
}
18+
// CHECK: Block is not in a loop
19+
// CHECK-NEXT: ^bb0:
20+
// CHECK: Block is in a loop
21+
// CHECK-NEXT: ^bb1:
22+
23+
// Test bb0 -> bb1 -> bb2 -> bb1
24+
func.func @loopForever2() {
25+
^bb0:
26+
cf.br ^bb1
27+
^bb1:
28+
cf.br ^bb2
29+
^bb2:
30+
cf.br ^bb1
31+
}
32+
// CHECK: Block is not in a loop
33+
// CHECK-NEXT: ^bb0:
34+
// CHECK: Block is in a loop
35+
// CHECK-NEXT: ^bb1:
36+
// CHECK: Block is in a loop
37+
// CHECK-NEXT: ^bb2:
38+
39+
// Test conditional branch without loop
40+
// bb0 -> bb1 -> {bb2, bb3}
41+
func.func @noLoop(%arg0: i1) {
42+
cf.br ^bb1
43+
^bb1:
44+
cf.cond_br %arg0, ^bb2, ^bb3
45+
^bb2:
46+
func.return
47+
^bb3:
48+
func.return
49+
}
50+
// CHECK: Block is not in a loop
51+
// CHECK-NEXT: ^bb0(%arg0: i1)
52+
// CHECK: Block is not in a loop
53+
// CHECK-NEXT: ^bb1:
54+
// CHECK: Block is not in a loop
55+
// CHECK-NEXT: ^bb2:
56+
// CHECK: Block is not in a loop
57+
// CHECK-NEXT: ^bb3:
58+
59+
// test multiple loops
60+
// bb0 -> bb1 -> bb2 -> bb3 { -> bb2} -> bb4 { -> bb1 } -> bb5
61+
func.func @multipleLoops(%arg0: i1, %arg1: i1) {
62+
cf.br ^bb1
63+
^bb1:
64+
cf.br ^bb2
65+
^bb2:
66+
cf.br ^bb3
67+
^bb3:
68+
cf.cond_br %arg0, ^bb4, ^bb2
69+
^bb4:
70+
cf.cond_br %arg1, ^bb1, ^bb5
71+
^bb5:
72+
return
73+
}
74+
// CHECK: Block is not in a loop
75+
// CHECK-NEXT: ^bb0(%arg0: i1, %arg1: i1)
76+
// CHECK: Block is in a loop
77+
// CHECK-NEXT: ^bb1:
78+
// CHECK: Block is in a loop
79+
// CHECK-NEXT: ^bb2:
80+
// CHECK: Block is in a loop
81+
// CHECK-NEXT: ^bb3:
82+
// CHECK: Block is in a loop
83+
// CHECK-NEXT: ^bb4:
84+
// CHECK: Block is not in a loop
85+
// CHECK-NEXT: ^bb5:
86+
87+
// test derived from real Flang output
88+
func.func @_QPblockTest0(%arg0: i1, %arg1: i1) {
89+
cf.br ^bb1
90+
^bb1: // 2 preds: ^bb0, ^bb4
91+
cf.cond_br %arg0, ^bb2, ^bb5
92+
^bb2: // pred: ^bb1
93+
cf.cond_br %arg1, ^bb3, ^bb4
94+
^bb3: // pred: ^bb2
95+
return
96+
^bb4: // pred: ^bb2
97+
cf.br ^bb1
98+
^bb5: // pred: ^bb1
99+
return
100+
}
101+
// CHECK: Block is not in a loop
102+
// CHECK-NEXT: ^bb0(%arg0: i1, %arg1: i1)
103+
// CHECK: Block is in a loop
104+
// CHECK-NEXT: ^bb1:
105+
// CHECK: Block is in a loop
106+
// CHECK-NEXT: ^bb2:
107+
// CHECK: Block is not in a loop
108+
// CHECK-NEXT: ^bb3:
109+
// CHECK: Block is in a loop
110+
// CHECK-NEXT: ^bb4:
111+
// CHECK: Block is not in a loop
112+
// CHECK-NEXT: ^bb5:
113+
114+
// check nested blocks
115+
func.func @check_alloc_in_loop(%counter : i64) {
116+
cf.br ^bb1(%counter: i64)
117+
^bb1(%lv : i64):
118+
%cm1 = arith.constant -1 : i64
119+
%rem = arith.addi %lv, %cm1 : i64
120+
%zero = arith.constant 0 : i64
121+
%p = arith.cmpi eq, %rem, %zero : i64
122+
cf.cond_br %p, ^bb3, ^bb2
123+
^bb2:
124+
scf.execute_region -> () {
125+
%c1 = arith.constant 1 : i64
126+
scf.yield
127+
}
128+
cf.br ^bb1(%rem: i64)
129+
^bb3:
130+
return
131+
}
132+
// CHECK: Block is not in a loop
133+
// CHECK-NEXT: ^bb0(%arg0: i64):
134+
// CHECK: Block is in a loop
135+
// CHECK-NEXT: ^bb1(%0: i64)
136+
// CHECK: Block is in a loop
137+
// CHECK-NEXT: ^bb0:
138+
// CHECK-NEXT: %c1_i64
139+
// CHECK: Block is in a loop
140+
// CHECK-NEXT: ^bb2:
141+
// CHECK: Block is not in a loop
142+
// CHECK-NEXT: ^bb3:
143+
}

‎mlir/test/lib/Interfaces/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1 +1,2 @@
1+
add_subdirectory(LoopLikeInterface)
12
add_subdirectory(TilingInterface)

‎mlir/test/lib/Interfaces/LoopLikeInterface/CMakeLists.txt

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
add_mlir_library(MLIRLoopLikeInterfaceTestPasses
2+
TestBlockInLoop.cpp
3+
4+
EXCLUDE_FROM_LIBMLIR
5+
6+
LINK_LIBS PUBLIC
7+
MLIRPass
8+
MLIRLoopLikeInterface
9+
)

‎mlir/test/lib/Interfaces/LoopLikeInterface/TestBlockInLoop.cpp

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
//===- TestBlockInLoop.cpp - Pass to test mlir::blockIsInLoop -------------===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
9+
#include "mlir/Dialect/Func/IR/FuncOps.h"
10+
#include "mlir/IR/BuiltinOps.h"
11+
#include "mlir/Interfaces/LoopLikeInterface.h"
12+
#include "mlir/Pass/Pass.h"
13+
#include "llvm/Support/raw_ostream.h"
14+
15+
using namespace mlir;
16+
17+
namespace {
18+
/// This is a test pass that tests Blocks's isInLoop method by checking if each
19+
/// block in a function is in a loop and outputing if it is
20+
struct IsInLoopPass
21+
: public PassWrapper<IsInLoopPass, OperationPass<func::FuncOp>> {
22+
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(IsInLoopPass)
23+
24+
StringRef getArgument() const final { return "test-block-is-in-loop"; }
25+
StringRef getDescription() const final {
26+
return "Test mlir::blockIsInLoop()";
27+
}
28+
29+
void runOnOperation() override {
30+
mlir::func::FuncOp func = getOperation();
31+
func.walk([](mlir::Block *block) {
32+
llvm::outs() << "Block is ";
33+
if (LoopLikeOpInterface::blockIsInLoop(block))
34+
llvm::outs() << "in a loop\n";
35+
else
36+
llvm::outs() << "not in a loop\n";
37+
block->print(llvm::outs());
38+
llvm::outs() << "\n";
39+
});
40+
}
41+
};
42+
43+
} // namespace
44+
45+
namespace mlir {
46+
void registerLoopLikeInterfaceTestPasses() { PassRegistration<IsInLoopPass>(); }
47+
} // namespace mlir

‎mlir/tools/mlir-opt/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ if(MLIR_INCLUDE_TESTS)
2121
MLIRFuncTestPasses
2222
MLIRGPUTestPasses
2323
MLIRLinalgTestPasses
24+
MLIRLoopLikeInterfaceTestPasses
2425
MLIRMathTestPasses
2526
MLIRMemRefTestPasses
2627
MLIRNVGPUTestPasses

‎mlir/tools/mlir-opt/mlir-opt.cpp

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ void registerConvertToTargetEnvPass();
3333
void registerCloneTestPasses();
3434
void registerPassManagerTestPass();
3535
void registerPrintSpirvAvailabilityPass();
36+
void registerLoopLikeInterfaceTestPasses();
3637
void registerShapeFunctionTestPasses();
3738
void registerSideEffectTestPasses();
3839
void registerSliceAnalysisTestPass();
@@ -140,6 +141,7 @@ void registerTestPasses() {
140141
registerConvertToTargetEnvPass();
141142
registerPassManagerTestPass();
142143
registerPrintSpirvAvailabilityPass();
144+
registerLoopLikeInterfaceTestPasses();
143145
registerShapeFunctionTestPasses();
144146
registerSideEffectTestPasses();
145147
registerSliceAnalysisTestPass();

0 commit comments

Comments
 (0)