Skip to content

Commit 90b2556

Browse files
zero9178zero9178
committed
[mlir][NFC] Expose the CFGToSCFInterface for the ControlFlowToSCF pass
This is useful for any downstream users who may just want to use slightly different ops than the pass or need to take into account other ops in the input while still dealing with ControlFlow ops for the most part. This also helps writing test implementations for `transformCFGToSCF`. Since the implementation is now public, comments noting the post conditions of the methods have been added to ensure downstream users can stay compatible with the existing implementation. Differential Revision: https://reviews.llvm.org/D157971
1 parent ca68a7f commit 90b2556

File tree

2 files changed

+172
-121
lines changed

2 files changed

+172
-121
lines changed

‎mlir/include/mlir/Conversion/ControlFlowToSCF/ControlFlowToSCF.h

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,60 @@
1515

1616
#include <memory>
1717

18+
#include "mlir/Transforms/CFGToSCF.h"
19+
1820
namespace mlir {
1921
class Pass;
2022

23+
/// Implementation of `CFGToSCFInterface` used to lift Control Flow Dialect
24+
/// operations to SCF Dialect operations.
25+
class ControlFlowToSCFTransformation : public CFGToSCFInterface {
26+
public:
27+
/// Creates an `scf.if` op if `controlFlowCondOp` is a `cf.cond_br` op or
28+
/// an `scf.index_switch` if `controlFlowCondOp` is a `cf.switch`.
29+
/// Returns failure otherwise.
30+
FailureOr<Operation *> createStructuredBranchRegionOp(
31+
OpBuilder &builder, Operation *controlFlowCondOp, TypeRange resultTypes,
32+
MutableArrayRef<Region> regions) override;
33+
34+
/// Creates an `scf.yield` op returning the given results.
35+
LogicalResult
36+
createStructuredBranchRegionTerminatorOp(Location loc, OpBuilder &builder,
37+
Operation *branchRegionOp,
38+
ValueRange results) override;
39+
40+
/// Creates an `scf.while` op. The loop body is made the before-region of the
41+
/// while op and terminated with an `scf.condition` op. The after-region does
42+
/// nothing but forward the iteration variables.
43+
FailureOr<Operation *>
44+
createStructuredDoWhileLoopOp(OpBuilder &builder, Operation *replacedOp,
45+
ValueRange loopVariablesInit, Value condition,
46+
ValueRange loopVariablesNextIter,
47+
Region &&loopBody) override;
48+
49+
/// Creates an `arith.constant` with an i32 attribute of the given value.
50+
Value getCFGSwitchValue(Location loc, OpBuilder &builder,
51+
unsigned value) override;
52+
53+
/// Creates a `cf.switch` op with the given cases and flag.
54+
void createCFGSwitchOp(Location loc, OpBuilder &builder, Value flag,
55+
ArrayRef<unsigned> caseValues,
56+
BlockRange caseDestinations,
57+
ArrayRef<ValueRange> caseArguments, Block *defaultDest,
58+
ValueRange defaultArgs) override;
59+
60+
/// Creates a `ub.poison` op of the given type.
61+
Value getUndefValue(Location loc, OpBuilder &builder, Type type) override;
62+
63+
/// Creates a `func.return` op with poison for each of the return values of
64+
/// the function. It is guaranteed to be directly within the function body.
65+
/// TODO: This can be made independent of the `func` dialect once the UB
66+
/// dialect has a `ub.unreachable` op.
67+
FailureOr<Operation *> createUnreachableTerminator(Location loc,
68+
OpBuilder &builder,
69+
Region &region) override;
70+
};
71+
2172
#define GEN_PASS_DECL_LIFTCONTROLFLOWTOSCFPASS
2273
#include "mlir/Conversion/Passes.h.inc"
2374

‎mlir/lib/Conversion/ControlFlowToSCF/ControlFlowToSCF.cpp

Lines changed: 121 additions & 121 deletions
Original file line numberDiff line numberDiff line change
@@ -29,132 +29,132 @@ namespace mlir {
2929

3030
using namespace mlir;
3131

32-
namespace {
33-
34-
class ControlFlowToSCFTransformation : public CFGToSCFInterface {
35-
public:
36-
FailureOr<Operation *> createStructuredBranchRegionOp(
37-
OpBuilder &builder, Operation *controlFlowCondOp, TypeRange resultTypes,
38-
MutableArrayRef<Region> regions) override {
39-
if (auto condBrOp = dyn_cast<cf::CondBranchOp>(controlFlowCondOp)) {
40-
assert(regions.size() == 2);
41-
auto ifOp = builder.create<scf::IfOp>(
42-
controlFlowCondOp->getLoc(), resultTypes, condBrOp.getCondition());
43-
ifOp.getThenRegion().takeBody(regions[0]);
44-
ifOp.getElseRegion().takeBody(regions[1]);
45-
return ifOp.getOperation();
46-
}
47-
48-
if (auto switchOp = dyn_cast<cf::SwitchOp>(controlFlowCondOp)) {
49-
// `getCFGSwitchValue` returns an i32 that we need to convert to index
50-
// fist.
51-
auto cast = builder.create<arith::IndexCastUIOp>(
52-
controlFlowCondOp->getLoc(), builder.getIndexType(),
53-
switchOp.getFlag());
54-
SmallVector<int64_t> cases;
55-
if (auto caseValues = switchOp.getCaseValues())
56-
llvm::append_range(
57-
cases, llvm::map_range(*caseValues, [](const llvm::APInt &apInt) {
58-
return apInt.getZExtValue();
59-
}));
60-
61-
assert(regions.size() == cases.size() + 1);
62-
63-
auto indexSwitchOp = builder.create<scf::IndexSwitchOp>(
64-
controlFlowCondOp->getLoc(), resultTypes, cast, cases, cases.size());
65-
66-
indexSwitchOp.getDefaultRegion().takeBody(regions[0]);
67-
for (auto &&[targetRegion, sourceRegion] :
68-
llvm::zip(indexSwitchOp.getCaseRegions(), llvm::drop_begin(regions)))
69-
targetRegion.takeBody(sourceRegion);
70-
71-
return indexSwitchOp.getOperation();
72-
}
73-
74-
controlFlowCondOp->emitOpError(
75-
"Cannot convert unknown control flow op to structured control flow");
76-
return failure();
77-
}
78-
79-
LogicalResult
80-
createStructuredBranchRegionTerminatorOp(Location loc, OpBuilder &builder,
81-
Operation *branchRegionOp,
82-
ValueRange results) override {
83-
builder.create<scf::YieldOp>(loc, results);
84-
return success();
85-
}
86-
87-
FailureOr<Operation *>
88-
createStructuredDoWhileLoopOp(OpBuilder &builder, Operation *replacedOp,
89-
ValueRange loopVariablesInit, Value condition,
90-
ValueRange loopVariablesNextIter,
91-
Region &&loopBody) override {
92-
Location loc = replacedOp->getLoc();
93-
auto whileOp = builder.create<scf::WhileOp>(
94-
loc, loopVariablesInit.getTypes(), loopVariablesInit);
95-
96-
whileOp.getBefore().takeBody(loopBody);
97-
98-
builder.setInsertionPointToEnd(&whileOp.getBefore().back());
99-
// `getCFGSwitchValue` returns a i32. We therefore need to truncate the
100-
// condition to i1 first. It is guaranteed to be either 0 or 1 already.
101-
builder.create<scf::ConditionOp>(
102-
loc,
103-
builder.create<arith::TruncIOp>(loc, builder.getI1Type(), condition),
104-
loopVariablesNextIter);
105-
106-
auto *afterBlock = new Block;
107-
whileOp.getAfter().push_back(afterBlock);
108-
afterBlock->addArguments(
109-
loopVariablesInit.getTypes(),
110-
SmallVector<Location>(loopVariablesInit.size(), loc));
111-
builder.setInsertionPointToEnd(afterBlock);
112-
builder.create<scf::YieldOp>(loc, afterBlock->getArguments());
113-
114-
return whileOp.getOperation();
32+
FailureOr<Operation *>
33+
ControlFlowToSCFTransformation::createStructuredBranchRegionOp(
34+
OpBuilder &builder, Operation *controlFlowCondOp, TypeRange resultTypes,
35+
MutableArrayRef<Region> regions) {
36+
if (auto condBrOp = dyn_cast<cf::CondBranchOp>(controlFlowCondOp)) {
37+
assert(regions.size() == 2);
38+
auto ifOp = builder.create<scf::IfOp>(controlFlowCondOp->getLoc(),
39+
resultTypes, condBrOp.getCondition());
40+
ifOp.getThenRegion().takeBody(regions[0]);
41+
ifOp.getElseRegion().takeBody(regions[1]);
42+
return ifOp.getOperation();
11543
}
11644

117-
Value getCFGSwitchValue(Location loc, OpBuilder &builder,
118-
unsigned int value) override {
119-
return builder.create<arith::ConstantOp>(loc,
120-
builder.getI32IntegerAttr(value));
45+
if (auto switchOp = dyn_cast<cf::SwitchOp>(controlFlowCondOp)) {
46+
// `getCFGSwitchValue` returns an i32 that we need to convert to index
47+
// fist.
48+
auto cast = builder.create<arith::IndexCastUIOp>(
49+
controlFlowCondOp->getLoc(), builder.getIndexType(),
50+
switchOp.getFlag());
51+
SmallVector<int64_t> cases;
52+
if (auto caseValues = switchOp.getCaseValues())
53+
llvm::append_range(
54+
cases, llvm::map_range(*caseValues, [](const llvm::APInt &apInt) {
55+
return apInt.getZExtValue();
56+
}));
57+
58+
assert(regions.size() == cases.size() + 1);
59+
60+
auto indexSwitchOp = builder.create<scf::IndexSwitchOp>(
61+
controlFlowCondOp->getLoc(), resultTypes, cast, cases, cases.size());
62+
63+
indexSwitchOp.getDefaultRegion().takeBody(regions[0]);
64+
for (auto &&[targetRegion, sourceRegion] :
65+
llvm::zip(indexSwitchOp.getCaseRegions(), llvm::drop_begin(regions)))
66+
targetRegion.takeBody(sourceRegion);
67+
68+
return indexSwitchOp.getOperation();
12169
}
12270

123-
void createCFGSwitchOp(Location loc, OpBuilder &builder, Value flag,
124-
ArrayRef<unsigned int> caseValues,
125-
BlockRange caseDestinations,
126-
ArrayRef<ValueRange> caseArguments, Block *defaultDest,
127-
ValueRange defaultArgs) override {
128-
builder.create<cf::SwitchOp>(loc, flag, defaultDest, defaultArgs,
129-
llvm::to_vector_of<int32_t>(caseValues),
130-
caseDestinations, caseArguments);
131-
}
132-
133-
Value getUndefValue(Location loc, OpBuilder &builder, Type type) override {
134-
return builder.create<ub::PoisonOp>(loc, type, nullptr);
135-
}
71+
controlFlowCondOp->emitOpError(
72+
"Cannot convert unknown control flow op to structured control flow");
73+
return failure();
74+
}
75+
76+
LogicalResult
77+
ControlFlowToSCFTransformation::createStructuredBranchRegionTerminatorOp(
78+
Location loc, OpBuilder &builder, Operation *branchRegionOp,
79+
ValueRange results) {
80+
builder.create<scf::YieldOp>(loc, results);
81+
return success();
82+
}
83+
84+
FailureOr<Operation *>
85+
ControlFlowToSCFTransformation::createStructuredDoWhileLoopOp(
86+
OpBuilder &builder, Operation *replacedOp, ValueRange loopVariablesInit,
87+
Value condition, ValueRange loopVariablesNextIter, Region &&loopBody) {
88+
Location loc = replacedOp->getLoc();
89+
auto whileOp = builder.create<scf::WhileOp>(loc, loopVariablesInit.getTypes(),
90+
loopVariablesInit);
91+
92+
whileOp.getBefore().takeBody(loopBody);
93+
94+
builder.setInsertionPointToEnd(&whileOp.getBefore().back());
95+
// `getCFGSwitchValue` returns a i32. We therefore need to truncate the
96+
// condition to i1 first. It is guaranteed to be either 0 or 1 already.
97+
builder.create<scf::ConditionOp>(
98+
loc, builder.create<arith::TruncIOp>(loc, builder.getI1Type(), condition),
99+
loopVariablesNextIter);
100+
101+
auto *afterBlock = new Block;
102+
whileOp.getAfter().push_back(afterBlock);
103+
afterBlock->addArguments(
104+
loopVariablesInit.getTypes(),
105+
SmallVector<Location>(loopVariablesInit.size(), loc));
106+
builder.setInsertionPointToEnd(afterBlock);
107+
builder.create<scf::YieldOp>(loc, afterBlock->getArguments());
108+
109+
return whileOp.getOperation();
110+
}
111+
112+
Value ControlFlowToSCFTransformation::getCFGSwitchValue(Location loc,
113+
OpBuilder &builder,
114+
unsigned int value) {
115+
return builder.create<arith::ConstantOp>(loc,
116+
builder.getI32IntegerAttr(value));
117+
}
118+
119+
void ControlFlowToSCFTransformation::createCFGSwitchOp(
120+
Location loc, OpBuilder &builder, Value flag,
121+
ArrayRef<unsigned int> caseValues, BlockRange caseDestinations,
122+
ArrayRef<ValueRange> caseArguments, Block *defaultDest,
123+
ValueRange defaultArgs) {
124+
builder.create<cf::SwitchOp>(loc, flag, defaultDest, defaultArgs,
125+
llvm::to_vector_of<int32_t>(caseValues),
126+
caseDestinations, caseArguments);
127+
}
128+
129+
Value ControlFlowToSCFTransformation::getUndefValue(Location loc,
130+
OpBuilder &builder,
131+
Type type) {
132+
return builder.create<ub::PoisonOp>(loc, type, nullptr);
133+
}
134+
135+
FailureOr<Operation *>
136+
ControlFlowToSCFTransformation::createUnreachableTerminator(Location loc,
137+
OpBuilder &builder,
138+
Region &region) {
139+
140+
// TODO: This should create a `ub.unreachable` op. Once such an operation
141+
// exists to make the pass independent of the func dialect. For now just
142+
// return poison values.
143+
auto funcOp = dyn_cast<func::FuncOp>(region.getParentOp());
144+
if (!funcOp)
145+
return emitError(loc, "Expected '")
146+
<< func::FuncOp::getOperationName() << "' as top level operation";
147+
148+
return builder
149+
.create<func::ReturnOp>(
150+
loc, llvm::map_to_vector(funcOp.getResultTypes(),
151+
[&](Type type) {
152+
return getUndefValue(loc, builder, type);
153+
}))
154+
.getOperation();
155+
}
136156

137-
FailureOr<Operation *> createUnreachableTerminator(Location loc,
138-
OpBuilder &builder,
139-
Region &region) override {
140-
141-
// TODO: This should create a `ub.unreachable` op. Once such an operation
142-
// exists to make the pass can be made independent of the func
143-
// dialect. For now just return poison values.
144-
auto funcOp = dyn_cast<func::FuncOp>(region.getParentOp());
145-
if (!funcOp)
146-
return emitError(loc, "Expected '")
147-
<< func::FuncOp::getOperationName() << "' as top level operation";
148-
149-
return builder
150-
.create<func::ReturnOp>(
151-
loc, llvm::map_to_vector(funcOp.getResultTypes(),
152-
[&](Type type) {
153-
return getUndefValue(loc, builder, type);
154-
}))
155-
.getOperation();
156-
}
157-
};
157+
namespace {
158158

159159
struct LiftControlFlowToSCF
160160
: public impl::LiftControlFlowToSCFPassBase<LiftControlFlowToSCF> {

0 commit comments

Comments
 (0)