[mlir] Fix memory leaks after #81759 #82762
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This commit fixes memory leaks that were introduced by llvm#81759. The way ops and blocks are erased changed slightly. The leaks were caused by an incorrect implementation of op builders: blocks must be created with the supplied builder object. Otherwise, they are not properly tracked by the dialect conversion and can leak during rollback.
|
@llvm/pr-subscribers-mlir-scf @llvm/pr-subscribers-mlir-gpu Author: Matthias Springer (matthias-springer) ChangesThis commit fixes memory leaks that were introduced by #81759. The way ops and blocks are erased changed slightly. The leaks were caused by an incorrect implementation of op builders: blocks must be created with the supplied builder object. Otherwise, they are not properly tracked by the dialect conversion and can leak during rollback. Full diff: https://github.com/llvm/llvm-project/pull/82762.diff 2 Files Affected:
diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
index 30b6cd74147e6f..33ce5c159db4f9 100644
--- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
+++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
@@ -648,6 +648,8 @@ void LaunchOp::build(OpBuilder &builder, OperationState &result,
TypeRange workgroupAttributions,
TypeRange privateAttributions, Value clusterSizeX,
Value clusterSizeY, Value clusterSizeZ) {
+ OpBuilder::InsertionGuard g(builder);
+
// Add a WorkGroup attribution attribute. This attribute is required to
// identify private attributions in the list of block argguments.
result.addAttribute(getNumWorkgroupAttributionsAttrName(),
@@ -674,7 +676,7 @@ void LaunchOp::build(OpBuilder &builder, OperationState &result,
// attributions, where the first kNumConfigRegionAttributes arguments have
// `index` type and the rest have the same types as the data operands.
Region *kernelRegion = result.addRegion();
- Block *body = new Block();
+ Block *body = builder.createBlock(kernelRegion);
// TODO: Allow passing in proper locations here.
for (unsigned i = 0; i < kNumConfigRegionAttributes; ++i)
body->addArgument(builder.getIndexType(), result.location);
@@ -683,7 +685,6 @@ void LaunchOp::build(OpBuilder &builder, OperationState &result,
body->addArgument(argTy, result.location);
for (Type argTy : privateAttributions)
body->addArgument(argTy, result.location);
- kernelRegion->push_back(body);
// Fill OperandSegmentSize Attribute.
SmallVector<int32_t, 11> segmentSizes(11, 1);
segmentSizes.front() = asyncDependencies.size();
@@ -1325,6 +1326,8 @@ void GPUFuncOp::build(OpBuilder &builder, OperationState &result,
TypeRange workgroupAttributions,
TypeRange privateAttributions,
ArrayRef<NamedAttribute> attrs) {
+ OpBuilder::InsertionGuard g(builder);
+
result.addAttribute(SymbolTable::getSymbolAttrName(),
builder.getStringAttr(name));
result.addAttribute(getFunctionTypeAttrName(result.name),
@@ -1333,7 +1336,7 @@ void GPUFuncOp::build(OpBuilder &builder, OperationState &result,
builder.getI64IntegerAttr(workgroupAttributions.size()));
result.addAttributes(attrs);
Region *body = result.addRegion();
- Block *entryBlock = new Block;
+ Block *entryBlock = builder.createBlock(body);
// TODO: Allow passing in proper locations here.
for (Type argTy : type.getInputs())
@@ -1342,8 +1345,6 @@ void GPUFuncOp::build(OpBuilder &builder, OperationState &result,
entryBlock->addArgument(argTy, result.location);
for (Type argTy : privateAttributions)
entryBlock->addArgument(argTy, result.location);
-
- body->getBlocks().push_back(entryBlock);
}
/// Parses a GPU function memory attribution.
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 119df9acd9e9e3..233e702dbb2292 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -306,17 +306,18 @@ void ConditionOp::getSuccessorRegions(
void ForOp::build(OpBuilder &builder, OperationState &result, Value lb,
Value ub, Value step, ValueRange iterArgs,
BodyBuilderFn bodyBuilder) {
+ OpBuilder::InsertionGuard guard(builder);
+
result.addOperands({lb, ub, step});
result.addOperands(iterArgs);
for (Value v : iterArgs)
result.addTypes(v.getType());
Type t = lb.getType();
Region *bodyRegion = result.addRegion();
- bodyRegion->push_back(new Block);
- Block &bodyBlock = bodyRegion->front();
- bodyBlock.addArgument(t, result.location);
+ Block *bodyBlock = builder.createBlock(bodyRegion);
+ bodyBlock->addArgument(t, result.location);
for (Value v : iterArgs)
- bodyBlock.addArgument(v.getType(), v.getLoc());
+ bodyBlock->addArgument(v.getType(), v.getLoc());
// Create the default terminator if the builder is not provided and if the
// iteration arguments are not provided. Otherwise, leave this to the caller
@@ -325,9 +326,9 @@ void ForOp::build(OpBuilder &builder, OperationState &result, Value lb,
ForOp::ensureTerminator(*bodyRegion, builder, result.location);
} else if (bodyBuilder) {
OpBuilder::InsertionGuard guard(builder);
- builder.setInsertionPointToStart(&bodyBlock);
- bodyBuilder(builder, result.location, bodyBlock.getArgument(0),
- bodyBlock.getArguments().drop_front());
+ builder.setInsertionPointToStart(bodyBlock);
+ bodyBuilder(builder, result.location, bodyBlock->getArgument(0),
+ bodyBlock->getArguments().drop_front());
}
}
|
|
@llvm/pr-subscribers-mlir Author: Matthias Springer (matthias-springer) ChangesThis commit fixes memory leaks that were introduced by #81759. The way ops and blocks are erased changed slightly. The leaks were caused by an incorrect implementation of op builders: blocks must be created with the supplied builder object. Otherwise, they are not properly tracked by the dialect conversion and can leak during rollback. Full diff: https://github.com/llvm/llvm-project/pull/82762.diff 2 Files Affected:
diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
index 30b6cd74147e6f..33ce5c159db4f9 100644
--- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
+++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
@@ -648,6 +648,8 @@ void LaunchOp::build(OpBuilder &builder, OperationState &result,
TypeRange workgroupAttributions,
TypeRange privateAttributions, Value clusterSizeX,
Value clusterSizeY, Value clusterSizeZ) {
+ OpBuilder::InsertionGuard g(builder);
+
// Add a WorkGroup attribution attribute. This attribute is required to
// identify private attributions in the list of block argguments.
result.addAttribute(getNumWorkgroupAttributionsAttrName(),
@@ -674,7 +676,7 @@ void LaunchOp::build(OpBuilder &builder, OperationState &result,
// attributions, where the first kNumConfigRegionAttributes arguments have
// `index` type and the rest have the same types as the data operands.
Region *kernelRegion = result.addRegion();
- Block *body = new Block();
+ Block *body = builder.createBlock(kernelRegion);
// TODO: Allow passing in proper locations here.
for (unsigned i = 0; i < kNumConfigRegionAttributes; ++i)
body->addArgument(builder.getIndexType(), result.location);
@@ -683,7 +685,6 @@ void LaunchOp::build(OpBuilder &builder, OperationState &result,
body->addArgument(argTy, result.location);
for (Type argTy : privateAttributions)
body->addArgument(argTy, result.location);
- kernelRegion->push_back(body);
// Fill OperandSegmentSize Attribute.
SmallVector<int32_t, 11> segmentSizes(11, 1);
segmentSizes.front() = asyncDependencies.size();
@@ -1325,6 +1326,8 @@ void GPUFuncOp::build(OpBuilder &builder, OperationState &result,
TypeRange workgroupAttributions,
TypeRange privateAttributions,
ArrayRef<NamedAttribute> attrs) {
+ OpBuilder::InsertionGuard g(builder);
+
result.addAttribute(SymbolTable::getSymbolAttrName(),
builder.getStringAttr(name));
result.addAttribute(getFunctionTypeAttrName(result.name),
@@ -1333,7 +1336,7 @@ void GPUFuncOp::build(OpBuilder &builder, OperationState &result,
builder.getI64IntegerAttr(workgroupAttributions.size()));
result.addAttributes(attrs);
Region *body = result.addRegion();
- Block *entryBlock = new Block;
+ Block *entryBlock = builder.createBlock(body);
// TODO: Allow passing in proper locations here.
for (Type argTy : type.getInputs())
@@ -1342,8 +1345,6 @@ void GPUFuncOp::build(OpBuilder &builder, OperationState &result,
entryBlock->addArgument(argTy, result.location);
for (Type argTy : privateAttributions)
entryBlock->addArgument(argTy, result.location);
-
- body->getBlocks().push_back(entryBlock);
}
/// Parses a GPU function memory attribution.
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 119df9acd9e9e3..233e702dbb2292 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -306,17 +306,18 @@ void ConditionOp::getSuccessorRegions(
void ForOp::build(OpBuilder &builder, OperationState &result, Value lb,
Value ub, Value step, ValueRange iterArgs,
BodyBuilderFn bodyBuilder) {
+ OpBuilder::InsertionGuard guard(builder);
+
result.addOperands({lb, ub, step});
result.addOperands(iterArgs);
for (Value v : iterArgs)
result.addTypes(v.getType());
Type t = lb.getType();
Region *bodyRegion = result.addRegion();
- bodyRegion->push_back(new Block);
- Block &bodyBlock = bodyRegion->front();
- bodyBlock.addArgument(t, result.location);
+ Block *bodyBlock = builder.createBlock(bodyRegion);
+ bodyBlock->addArgument(t, result.location);
for (Value v : iterArgs)
- bodyBlock.addArgument(v.getType(), v.getLoc());
+ bodyBlock->addArgument(v.getType(), v.getLoc());
// Create the default terminator if the builder is not provided and if the
// iteration arguments are not provided. Otherwise, leave this to the caller
@@ -325,9 +326,9 @@ void ForOp::build(OpBuilder &builder, OperationState &result, Value lb,
ForOp::ensureTerminator(*bodyRegion, builder, result.location);
} else if (bodyBuilder) {
OpBuilder::InsertionGuard guard(builder);
- builder.setInsertionPointToStart(&bodyBlock);
- bodyBuilder(builder, result.location, bodyBlock.getArgument(0),
- bodyBlock.getArguments().drop_front());
+ builder.setInsertionPointToStart(bodyBlock);
+ bodyBuilder(builder, result.location, bodyBlock->getArgument(0),
+ bodyBlock->getArguments().drop_front());
}
}
|
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This commit fixes memory leaks that were introduced by #81759. The way ops and blocks are erased changed slightly.
The leaks were caused by an incorrect implementation of op builders: blocks must be created with the supplied builder object. Otherwise, they are not properly tracked by the dialect conversion and can leak during rollback.