Skip to content

Commit 65d0767

Browse files
matthias-springermgehre-amd
matthias-springer
authored and
mgehre-amd
committed
[mlir] Use OpBuilder::createBlock in op builders and patterns (llvm#82770)
When creating a new block in (conversion) rewrite patterns, `OpBuilder::createBlock` must be used. Otherwise, no `notifyBlockInserted` notification is sent to the listener. Note: The dialect conversion relies on listener notifications to keep track of IR modifications. Creating blocks without the builder API can lead to memory leaks during rollback.
1 parent a96e62b commit 65d0767

File tree

20 files changed

+71
-88
lines changed

20 files changed

+71
-88
lines changed

‎mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1439,7 +1439,7 @@ def LLVM_LLVMFuncOp : LLVM_Op<"func", [
14391439
let extraClassDeclaration = [{
14401440
// Add an entry block to an empty function, and set up the block arguments
14411441
// to match the signature of the function.
1442-
Block *addEntryBlock();
1442+
Block *addEntryBlock(OpBuilder &builder);
14431443

14441444
bool isVarArg() { return getFunctionType().isVarArg(); }
14451445

‎mlir/include/mlir/Dialect/SPIRV/IR/SPIRVControlFlowOps.td

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -285,7 +285,7 @@ def SPIRV_LoopOp : SPIRV_Op<"mlir.loop", [InFunctionScope]> {
285285

286286
// Adds an empty entry block and loop merge block containing one
287287
// spirv.mlir.merge op.
288-
void addEntryAndMergeBlock();
288+
void addEntryAndMergeBlock(OpBuilder &builder);
289289
}];
290290

291291
let hasOpcode = 0;
@@ -427,7 +427,7 @@ def SPIRV_SelectionOp : SPIRV_Op<"mlir.selection", [InFunctionScope]> {
427427
Block *getMergeBlock();
428428

429429
/// Adds a selection merge block containing one spirv.mlir.merge op.
430-
void addMergeBlock();
430+
void addMergeBlock(OpBuilder &builder);
431431

432432
/// Creates a spirv.mlir.selection op for `if (<condition>) then { <thenBody> }`
433433
/// with `builder`. `builder`'s insertion point will remain at after the

‎mlir/include/mlir/Interfaces/FunctionInterfaces.td

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -131,6 +131,7 @@ def FunctionOpInterface : OpInterface<"FunctionOpInterface", [
131131
static void buildWithEntryBlock(
132132
OpBuilder &builder, OperationState &state, StringRef name, Type type,
133133
ArrayRef<NamedAttribute> attrs, TypeRange inputTypes) {
134+
OpBuilder::InsertionGuard g(builder);
134135
state.addAttribute(SymbolTable::getSymbolAttrName(),
135136
builder.getStringAttr(name));
136137
state.addAttribute(ConcreteOp::getFunctionTypeAttrName(state.name),
@@ -139,8 +140,7 @@ def FunctionOpInterface : OpInterface<"FunctionOpInterface", [
139140

140141
// Add the function body.
141142
Region *bodyRegion = state.addRegion();
142-
Block *body = new Block();
143-
bodyRegion->push_back(body);
143+
Block *body = builder.createBlock(bodyRegion);
144144
for (Type input : inputTypes)
145145
body->addArgument(input, state.location);
146146
}

‎mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -259,7 +259,7 @@ static void addResumeFunction(ModuleOp module) {
259259
kResume, LLVM::LLVMFunctionType::get(voidTy, {ptrType}));
260260
resumeOp.setPrivate();
261261

262-
auto *block = resumeOp.addEntryBlock();
262+
auto *block = resumeOp.addEntryBlock(moduleBuilder);
263263
auto blockBuilder = ImplicitLocOpBuilder::atBlockEnd(loc, block);
264264

265265
blockBuilder.create<LLVM::CoroResumeOp>(resumeOp.getArgument(0));

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

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -98,12 +98,10 @@ ControlFlowToSCFTransformation::createStructuredDoWhileLoopOp(
9898
loc, builder.create<arith::TruncIOp>(loc, builder.getI1Type(), condition),
9999
loopVariablesNextIter);
100100

101-
auto *afterBlock = new Block;
102-
whileOp.getAfter().push_back(afterBlock);
101+
Block *afterBlock = builder.createBlock(&whileOp.getAfter());
103102
afterBlock->addArguments(
104103
loopVariablesInit.getTypes(),
105104
SmallVector<Location>(loopVariablesInit.size(), loc));
106-
builder.setInsertionPointToEnd(afterBlock);
107105
builder.create<scf::YieldOp>(loc, afterBlock->getArguments());
108106

109107
return whileOp.getOperation();

‎mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -135,7 +135,7 @@ static void wrapForExternalCallers(OpBuilder &rewriter, Location loc,
135135
propagateArgResAttrs(rewriter, !!resultStructType, funcOp, wrapperFuncOp);
136136

137137
OpBuilder::InsertionGuard guard(rewriter);
138-
rewriter.setInsertionPointToStart(wrapperFuncOp.addEntryBlock());
138+
rewriter.setInsertionPointToStart(wrapperFuncOp.addEntryBlock(rewriter));
139139

140140
SmallVector<Value, 8> args;
141141
size_t argOffset = resultStructType ? 1 : 0;
@@ -203,7 +203,7 @@ static void wrapExternalFunction(OpBuilder &builder, Location loc,
203203

204204
// The wrapper that we synthetize here should only be visible in this module.
205205
newFuncOp.setLinkage(LLVM::Linkage::Private);
206-
builder.setInsertionPointToStart(newFuncOp.addEntryBlock());
206+
builder.setInsertionPointToStart(newFuncOp.addEntryBlock(builder));
207207

208208
// Get a ValueRange containing arguments.
209209
FunctionType type = cast<FunctionType>(funcOp.getFunctionType());

‎mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -520,9 +520,7 @@ struct GlobalMemrefOpLowering
520520
global, arrayTy, global.getConstant(), linkage, global.getSymName(),
521521
initialValue, alignment, *addressSpace);
522522
if (!global.isExternal() && global.isUninitialized()) {
523-
Block *blk = new Block();
524-
newGlobal.getInitializerRegion().push_back(blk);
525-
rewriter.setInsertionPointToStart(blk);
523+
rewriter.createBlock(&newGlobal.getInitializerRegion());
526524
Value undef[] = {
527525
rewriter.create<LLVM::UndefOp>(global.getLoc(), arrayTy)};
528526
rewriter.create<LLVM::ReturnOp>(global.getLoc(), undef);

‎mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -138,14 +138,13 @@ struct ForOpConversion final : SCFToSPIRVPattern<scf::ForOp> {
138138
// from header to merge.
139139
auto loc = forOp.getLoc();
140140
auto loopOp = rewriter.create<spirv::LoopOp>(loc, spirv::LoopControl::None);
141-
loopOp.addEntryAndMergeBlock();
141+
loopOp.addEntryAndMergeBlock(rewriter);
142142

143143
OpBuilder::InsertionGuard guard(rewriter);
144144
// Create the block for the header.
145-
auto *header = new Block();
146-
// Insert the header.
147-
loopOp.getBody().getBlocks().insert(getBlockIt(loopOp.getBody(), 1),
148-
header);
145+
Block *header = rewriter.createBlock(&loopOp.getBody(),
146+
getBlockIt(loopOp.getBody(), 1));
147+
rewriter.setInsertionPointAfter(loopOp);
149148

150149
// Create the new induction variable to use.
151150
Value adapLowerBound = adaptor.getLowerBound();
@@ -342,7 +341,7 @@ struct WhileOpConversion final : SCFToSPIRVPattern<scf::WhileOp> {
342341
ConversionPatternRewriter &rewriter) const override {
343342
auto loc = whileOp.getLoc();
344343
auto loopOp = rewriter.create<spirv::LoopOp>(loc, spirv::LoopControl::None);
345-
loopOp.addEntryAndMergeBlock();
344+
loopOp.addEntryAndMergeBlock(rewriter);
346345

347346
Region &beforeRegion = whileOp.getBefore();
348347
Region &afterRegion = whileOp.getAfter();

‎mlir/lib/Dialect/Affine/IR/AffineOps.cpp

Lines changed: 14 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1809,6 +1809,8 @@ void AffineForOp::build(OpBuilder &builder, OperationState &result,
18091809
"upper bound operand count does not match the affine map");
18101810
assert(step > 0 && "step has to be a positive integer constant");
18111811

1812+
OpBuilder::InsertionGuard guard(builder);
1813+
18121814
// Set variadic segment sizes.
18131815
result.addAttribute(
18141816
getOperandSegmentSizeAttr(),
@@ -1837,12 +1839,11 @@ void AffineForOp::build(OpBuilder &builder, OperationState &result,
18371839
// Create a region and a block for the body. The argument of the region is
18381840
// the loop induction variable.
18391841
Region *bodyRegion = result.addRegion();
1840-
bodyRegion->push_back(new Block);
1841-
Block &bodyBlock = bodyRegion->front();
1842+
Block *bodyBlock = builder.createBlock(bodyRegion);
18421843
Value inductionVar =
1843-
bodyBlock.addArgument(builder.getIndexType(), result.location);
1844+
bodyBlock->addArgument(builder.getIndexType(), result.location);
18441845
for (Value val : iterArgs)
1845-
bodyBlock.addArgument(val.getType(), val.getLoc());
1846+
bodyBlock->addArgument(val.getType(), val.getLoc());
18461847

18471848
// Create the default terminator if the builder is not provided and if the
18481849
// iteration arguments are not provided. Otherwise, leave this to the caller
@@ -1851,9 +1852,9 @@ void AffineForOp::build(OpBuilder &builder, OperationState &result,
18511852
ensureTerminator(*bodyRegion, builder, result.location);
18521853
} else if (bodyBuilder) {
18531854
OpBuilder::InsertionGuard guard(builder);
1854-
builder.setInsertionPointToStart(&bodyBlock);
1855+
builder.setInsertionPointToStart(bodyBlock);
18551856
bodyBuilder(builder, result.location, inductionVar,
1856-
bodyBlock.getArguments().drop_front());
1857+
bodyBlock->getArguments().drop_front());
18571858
}
18581859
}
18591860

@@ -2890,18 +2891,20 @@ void AffineIfOp::build(OpBuilder &builder, OperationState &result,
28902891
TypeRange resultTypes, IntegerSet set, ValueRange args,
28912892
bool withElseRegion) {
28922893
assert(resultTypes.empty() || withElseRegion);
2894+
OpBuilder::InsertionGuard guard(builder);
2895+
28932896
result.addTypes(resultTypes);
28942897
result.addOperands(args);
28952898
result.addAttribute(getConditionAttrStrName(), IntegerSetAttr::get(set));
28962899

28972900
Region *thenRegion = result.addRegion();
2898-
thenRegion->push_back(new Block());
2901+
builder.createBlock(thenRegion);
28992902
if (resultTypes.empty())
29002903
AffineIfOp::ensureTerminator(*thenRegion, builder, result.location);
29012904

29022905
Region *elseRegion = result.addRegion();
29032906
if (withElseRegion) {
2904-
elseRegion->push_back(new Block());
2907+
builder.createBlock(elseRegion);
29052908
if (resultTypes.empty())
29062909
AffineIfOp::ensureTerminator(*elseRegion, builder, result.location);
29072910
}
@@ -3688,6 +3691,7 @@ void AffineParallelOp::build(OpBuilder &builder, OperationState &result,
36883691
"expected upper bound maps to have as many inputs as upper bound "
36893692
"operands");
36903693

3694+
OpBuilder::InsertionGuard guard(builder);
36913695
result.addTypes(resultTypes);
36923696

36933697
// Convert the reductions to integer attributes.
@@ -3733,11 +3737,11 @@ void AffineParallelOp::build(OpBuilder &builder, OperationState &result,
37333737

37343738
// Create a region and a block for the body.
37353739
auto *bodyRegion = result.addRegion();
3736-
auto *body = new Block();
3740+
Block *body = builder.createBlock(bodyRegion);
3741+
37373742
// Add all the block arguments.
37383743
for (unsigned i = 0, e = steps.size(); i < e; ++i)
37393744
body->addArgument(IndexType::get(builder.getContext()), result.location);
3740-
bodyRegion->push_back(body);
37413745
if (resultTypes.empty())
37423746
ensureTerminator(*bodyRegion, builder, result.location);
37433747
}

‎mlir/lib/Dialect/Async/IR/Async.cpp

Lines changed: 6 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@ void ExecuteOp::getSuccessorRegions(RegionBranchPoint point,
6868
void ExecuteOp::build(OpBuilder &builder, OperationState &result,
6969
TypeRange resultTypes, ValueRange dependencies,
7070
ValueRange operands, BodyBuilderFn bodyBuilder) {
71-
71+
OpBuilder::InsertionGuard guard(builder);
7272
result.addOperands(dependencies);
7373
result.addOperands(operands);
7474

@@ -87,26 +87,21 @@ void ExecuteOp::build(OpBuilder &builder, OperationState &result,
8787

8888
// Add a body region with block arguments as unwrapped async value operands.
8989
Region *bodyRegion = result.addRegion();
90-
bodyRegion->push_back(new Block);
91-
Block &bodyBlock = bodyRegion->front();
90+
Block *bodyBlock = builder.createBlock(bodyRegion);
9291
for (Value operand : operands) {
9392
auto valueType = llvm::dyn_cast<ValueType>(operand.getType());
94-
bodyBlock.addArgument(valueType ? valueType.getValueType()
95-
: operand.getType(),
96-
operand.getLoc());
93+
bodyBlock->addArgument(valueType ? valueType.getValueType()
94+
: operand.getType(),
95+
operand.getLoc());
9796
}
9897

9998
// Create the default terminator if the builder is not provided and if the
10099
// expected result is empty. Otherwise, leave this to the caller
101100
// because we don't know which values to return from the execute op.
102101
if (resultTypes.empty() && !bodyBuilder) {
103-
OpBuilder::InsertionGuard guard(builder);
104-
builder.setInsertionPointToStart(&bodyBlock);
105102
builder.create<async::YieldOp>(result.location, ValueRange());
106103
} else if (bodyBuilder) {
107-
OpBuilder::InsertionGuard guard(builder);
108-
builder.setInsertionPointToStart(&bodyBlock);
109-
bodyBuilder(builder, result.location, bodyBlock.getArguments());
104+
bodyBuilder(builder, result.location, bodyBlock->getArguments());
110105
}
111106
}
112107

‎mlir/lib/Dialect/EmitC/IR/EmitC.cpp

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -262,20 +262,20 @@ LogicalResult ExpressionOp::verify() {
262262

263263
void ForOp::build(OpBuilder &builder, OperationState &result, Value lb,
264264
Value ub, Value step, BodyBuilderFn bodyBuilder) {
265+
OpBuilder::InsertionGuard g(builder);
265266
result.addOperands({lb, ub, step});
266267
Type t = lb.getType();
267268
Region *bodyRegion = result.addRegion();
268-
bodyRegion->push_back(new Block);
269-
Block &bodyBlock = bodyRegion->front();
270-
bodyBlock.addArgument(t, result.location);
269+
Block *bodyBlock = builder.createBlock(bodyRegion);
270+
bodyBlock->addArgument(t, result.location);
271271

272272
// Create the default terminator if the builder is not provided.
273273
if (!bodyBuilder) {
274274
ForOp::ensureTerminator(*bodyRegion, builder, result.location);
275275
} else {
276276
OpBuilder::InsertionGuard guard(builder);
277-
builder.setInsertionPointToStart(&bodyBlock);
278-
bodyBuilder(builder, result.location, bodyBlock.getArgument(0));
277+
builder.setInsertionPointToStart(bodyBlock);
278+
bodyBuilder(builder, result.location, bodyBlock->getArgument(0));
279279
}
280280
}
281281

‎mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2166,11 +2166,10 @@ LogicalResult ShuffleVectorOp::verify() {
21662166
//===----------------------------------------------------------------------===//
21672167

21682168
// Add the entry block to the function.
2169-
Block *LLVMFuncOp::addEntryBlock() {
2169+
Block *LLVMFuncOp::addEntryBlock(OpBuilder &builder) {
21702170
assert(empty() && "function already has an entry block");
2171-
2172-
auto *entry = new Block;
2173-
push_back(entry);
2171+
OpBuilder::InsertionGuard g(builder);
2172+
Block *entry = builder.createBlock(&getBody());
21742173

21752174
// FIXME: Allow passing in proper locations for the entry arguments.
21762175
LLVMFunctionType type = getFunctionType();

‎mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -132,12 +132,10 @@ struct MoveInitOperandsToInput : public OpRewritePattern<GenericOp> {
132132
newIndexingMaps, genericOp.getIteratorTypesArray(),
133133
/*bodyBuild=*/nullptr, linalg::getPrunedAttributeList(genericOp));
134134

135+
OpBuilder::InsertionGuard guard(rewriter);
135136
Region &region = newOp.getRegion();
136-
Block *block = new Block();
137-
region.push_back(block);
137+
Block *block = rewriter.createBlock(&region);
138138
IRMapping mapper;
139-
OpBuilder::InsertionGuard guard(rewriter);
140-
rewriter.setInsertionPointToStart(block);
141139
for (auto bbarg : genericOp.getRegionInputArgs())
142140
mapper.map(bbarg, block->addArgument(bbarg.getType(), loc));
143141

‎mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -178,11 +178,9 @@ static void generateFusedElementwiseOpRegion(
178178
// Build the region of the fused op.
179179
Block &producerBlock = producer->getRegion(0).front();
180180
Block &consumerBlock = consumer->getRegion(0).front();
181-
Block *fusedBlock = new Block();
182-
fusedOp.getRegion().push_back(fusedBlock);
183-
IRMapping mapper;
184181
OpBuilder::InsertionGuard guard(rewriter);
185-
rewriter.setInsertionPointToStart(fusedBlock);
182+
Block *fusedBlock = rewriter.createBlock(&fusedOp.getRegion());
183+
IRMapping mapper;
186184

187185
// 2. Add an index operation for every fused loop dimension and use the
188186
// `consumerToProducerLoopsMap` to map the producer indices.

‎mlir/lib/Dialect/Linalg/Utils/Utils.cpp

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -275,14 +275,13 @@ GenericOp makeTransposeOp(OpBuilder &b, Location loc, Value inputTensor,
275275
auto transposeOp =
276276
b.create<GenericOp>(loc, resultTensorType, inputTensor, outputTensor,
277277
indexingMaps, iteratorTypes);
278-
Region &body = transposeOp.getRegion();
279-
body.push_back(new Block());
280-
body.front().addArguments({elementType, elementType}, {loc, loc});
281278

282279
// Create the body of the transpose operation.
283280
OpBuilder::InsertionGuard g(b);
284-
b.setInsertionPointToEnd(&body.front());
285-
b.create<YieldOp>(loc, transposeOp.getRegion().front().getArgument(0));
281+
Region &body = transposeOp.getRegion();
282+
Block *bodyBlock = b.createBlock(&body, /*insertPt=*/{},
283+
{elementType, elementType}, {loc, loc});
284+
b.create<YieldOp>(loc, bodyBlock->getArgument(0));
286285
return transposeOp;
287286
}
288287

‎mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1420,6 +1420,7 @@ OpFoldResult ExtractStridedMetadataOp::getConstifiedMixedOffset() {
14201420

14211421
void GenericAtomicRMWOp::build(OpBuilder &builder, OperationState &result,
14221422
Value memref, ValueRange ivs) {
1423+
OpBuilder::InsertionGuard g(builder);
14231424
result.addOperands(memref);
14241425
result.addOperands(ivs);
14251426

@@ -1428,7 +1429,7 @@ void GenericAtomicRMWOp::build(OpBuilder &builder, OperationState &result,
14281429
result.addTypes(elementType);
14291430

14301431
Region *bodyRegion = result.addRegion();
1431-
bodyRegion->push_back(new Block());
1432+
builder.createBlock(bodyRegion);
14321433
bodyRegion->addArgument(elementType, memref.getLoc());
14331434
}
14341435
}

‎mlir/lib/Dialect/SPIRV/IR/ControlFlowOps.cpp

Lines changed: 8 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -365,12 +365,11 @@ Block *LoopOp::getMergeBlock() {
365365
return &getBody().back();
366366
}
367367

368-
void LoopOp::addEntryAndMergeBlock() {
368+
void LoopOp::addEntryAndMergeBlock(OpBuilder &builder) {
369369
assert(getBody().empty() && "entry and merge block already exist");
370-
getBody().push_back(new Block());
371-
auto *mergeBlock = new Block();
372-
getBody().push_back(mergeBlock);
373-
OpBuilder builder = OpBuilder::atBlockEnd(mergeBlock);
370+
OpBuilder::InsertionGuard g(builder);
371+
builder.createBlock(&getBody());
372+
builder.createBlock(&getBody());
374373

375374
// Add a spirv.mlir.merge op into the merge block.
376375
builder.create<spirv::MergeOp>(getLoc());
@@ -525,11 +524,10 @@ Block *SelectionOp::getMergeBlock() {
525524
return &getBody().back();
526525
}
527526

528-
void SelectionOp::addMergeBlock() {
527+
void SelectionOp::addMergeBlock(OpBuilder &builder) {
529528
assert(getBody().empty() && "entry and merge block already exist");
530-
auto *mergeBlock = new Block();
531-
getBody().push_back(mergeBlock);
532-
OpBuilder builder = OpBuilder::atBlockEnd(mergeBlock);
529+
OpBuilder::InsertionGuard guard(builder);
530+
builder.createBlock(&getBody());
533531

534532
// Add a spirv.mlir.merge op into the merge block.
535533
builder.create<spirv::MergeOp>(getLoc());
@@ -542,7 +540,7 @@ SelectionOp::createIfThen(Location loc, Value condition,
542540
auto selectionOp =
543541
builder.create<spirv::SelectionOp>(loc, spirv::SelectionControl::None);
544542

545-
selectionOp.addMergeBlock();
543+
selectionOp.addMergeBlock(builder);
546544
Block *mergeBlock = selectionOp.getMergeBlock();
547545
Block *thenBlock = nullptr;
548546

0 commit comments

Comments
 (0)