[Mlir-commits] [mlir] db054d7 - [MLIR] Introduce an op trait that defines a new scope for auto allocation
Uday Bondhugula
llvmlistbot at llvm.org
Thu Apr 9 23:36:51 PDT 2020
Author: Uday Bondhugula
Date: 2020-04-10T12:05:52+05:30
New Revision: db054d711518b6527048aa06196e443ff2e1dada
URL: https://github.com/llvm/llvm-project/commit/db054d711518b6527048aa06196e443ff2e1dada
DIFF: https://github.com/llvm/llvm-project/commit/db054d711518b6527048aa06196e443ff2e1dada.diff
LOG: [MLIR] Introduce an op trait that defines a new scope for auto allocation
Introduce a new operation property / trait (AutomaticAllocationScope)
for operations with regions that define a new scope for automatic allocations;
such allocations (typically realized on stack) are automatically freed when
control leaves such ops' regions. std.alloca's are freed at the closest
surrounding op that has this trait. All FunctionLike operations should normally
have this trait.
Differential Revision: https://reviews.llvm.org/D77787
Added:
Modified:
mlir/docs/Traits.md
mlir/include/mlir/Dialect/GPU/GPUOps.td
mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
mlir/include/mlir/Dialect/SPIRV/SPIRVStructureOps.td
mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
mlir/include/mlir/IR/Function.h
mlir/include/mlir/IR/OpBase.td
mlir/include/mlir/IR/OpDefinition.h
mlir/lib/Dialect/StandardOps/IR/Ops.cpp
mlir/test/IR/invalid-ops.mlir
Removed:
################################################################################
diff --git a/mlir/docs/Traits.md b/mlir/docs/Traits.md
index 9877d902982d..1035b6cd6e78 100644
--- a/mlir/docs/Traits.md
+++ b/mlir/docs/Traits.md
@@ -135,6 +135,16 @@ section goes as follows:
* `Header`
- (`C++ class` -- `ODS class`(if applicable))
+### AutomaticAllocationScope
+
+* `OpTrait::AutomaticAllocationScope` -- `AutomaticAllocationScope`
+
+This trait is carried by region holding operations that define a new scope for
+automatic allocation. Such allocations are automatically freed when control is
+transferred back from the regions of such operations. As an example, allocations
+performed by std.alloca are automatically freed when control leaves the region
+of its closest surrounding op that has the trait AutomaticAllocationScope.
+
### Broadcastable
* `OpTrait::ResultsBroadcastableShape` -- `ResultsBroadcastableShape`
diff --git a/mlir/include/mlir/Dialect/GPU/GPUOps.td b/mlir/include/mlir/Dialect/GPU/GPUOps.td
index 981f39756455..5d91ff6ac545 100644
--- a/mlir/include/mlir/Dialect/GPU/GPUOps.td
+++ b/mlir/include/mlir/Dialect/GPU/GPUOps.td
@@ -85,7 +85,8 @@ def GPU_ThreadIdOp : GPU_IndexOp<"thread_id"> {
}];
}
-def GPU_GPUFuncOp : GPU_Op<"func", [FunctionLike, IsolatedFromAbove, Symbol]> {
+def GPU_GPUFuncOp : GPU_Op<"func", [AutomaticAllocationScope, FunctionLike,
+ IsolatedFromAbove, Symbol]> {
let summary = "Function executable on a GPU";
let description = [{
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
index 54b2e9429bc3..eac663795453 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
@@ -660,7 +660,8 @@ def LLVM_GlobalOp
}
def LLVM_LLVMFuncOp
- : LLVM_ZeroResultOp<"func", [IsolatedFromAbove, FunctionLike, Symbol]>,
+ : LLVM_ZeroResultOp<"func", [AutomaticAllocationScope, IsolatedFromAbove,
+ FunctionLike, Symbol]>,
Arguments<(ins DefaultValuedAttr<Linkage,
"Linkage::External">:$linkage,
OptionalAttr<FlatSymbolRefAttr>:$personality,
diff --git a/mlir/include/mlir/Dialect/SPIRV/SPIRVStructureOps.td b/mlir/include/mlir/Dialect/SPIRV/SPIRVStructureOps.td
index 3bcbdc8069a1..c9f27b971bb7 100644
--- a/mlir/include/mlir/Dialect/SPIRV/SPIRVStructureOps.td
+++ b/mlir/include/mlir/Dialect/SPIRV/SPIRVStructureOps.td
@@ -197,7 +197,7 @@ def SPV_EntryPointOp : SPV_Op<"EntryPoint", [InModuleScope]> {
// -----
def SPV_FuncOp : SPV_Op<"func", [
- DeclareOpInterfaceMethods<CallableOpInterface>,
+ AutomaticAllocationScope, DeclareOpInterfaceMethods<CallableOpInterface>,
FunctionLike, InModuleScope, IsolatedFromAbove, Symbol
]> {
let summary = "Declare or define a function";
diff --git a/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td b/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
index 227c006fe4fc..55478a4ae9c2 100644
--- a/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
+++ b/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
@@ -326,8 +326,10 @@ def AllocaOp : AllocLikeOp<"alloca"> {
let summary = "stack memory allocation operation";
let description = [{
The `alloca` operation allocates memory on the stack, to be automatically
- released when the stack frame is discarded. The amount of memory allocated
- is specified by its memref and additional operands. For example:
+ released when control transfers back from the region of its closest
+ surrounding operation with a AutomaticAllocationScope trait. The amount of
+ memory allocated is specified by its memref and additional operands. For
+ example:
```mlir
%0 = alloca() : memref<8x64xf32>
diff --git a/mlir/include/mlir/IR/Function.h b/mlir/include/mlir/IR/Function.h
index 49c0c7c61284..1b2dd19eedab 100644
--- a/mlir/include/mlir/IR/Function.h
+++ b/mlir/include/mlir/IR/Function.h
@@ -30,9 +30,11 @@ namespace mlir {
/// implicitly capture global values, and all external references must use
/// Function arguments or attributes that establish a symbolic connection(e.g.
/// symbols referenced by name via a string attribute).
-class FuncOp : public Op<FuncOp, OpTrait::ZeroOperands, OpTrait::ZeroResult,
- OpTrait::IsIsolatedFromAbove, OpTrait::Symbol,
- OpTrait::FunctionLike, CallableOpInterface::Trait> {
+class FuncOp
+ : public Op<FuncOp, OpTrait::ZeroOperands, OpTrait::ZeroResult,
+ OpTrait::IsIsolatedFromAbove, OpTrait::Symbol,
+ OpTrait::FunctionLike, OpTrait::AutomaticAllocationScope,
+ CallableOpInterface::Trait> {
public:
using Op::Op;
using Op::print;
diff --git a/mlir/include/mlir/IR/OpBase.td b/mlir/include/mlir/IR/OpBase.td
index 09cea1b1f5ea..58c9643f3efb 100644
--- a/mlir/include/mlir/IR/OpBase.td
+++ b/mlir/include/mlir/IR/OpBase.td
@@ -1587,6 +1587,8 @@ class PredOpTrait<string descr, Pred pred> : OpTrait {
Pred predicate = pred;
}
+// Op defines an automatic allocation scope.
+def AutomaticAllocationScope : NativeOpTrait<"AutomaticAllocationScope">;
// Op supports operand broadcast behavior.
def ResultsBroadcastableShape :
NativeOpTrait<"ResultsBroadcastableShape">;
diff --git a/mlir/include/mlir/IR/OpDefinition.h b/mlir/include/mlir/IR/OpDefinition.h
index 4901f5b232b7..38403e86f952 100644
--- a/mlir/include/mlir/IR/OpDefinition.h
+++ b/mlir/include/mlir/IR/OpDefinition.h
@@ -1026,6 +1026,22 @@ class IsIsolatedFromAbove
}
};
+/// A trait of region holding operations that define a new scope for automatic
+/// allocations, i.e., allocations that are freed when control is transferred
+/// back from the operation's region. Any operations performing such allocations
+/// (for eg. std.alloca) will have their allocations automatically freed at
+/// their closest enclosing operation with this trait.
+template <typename ConcreteType>
+class AutomaticAllocationScope
+ : public TraitBase<ConcreteType, AutomaticAllocationScope> {
+public:
+ static LogicalResult verifyTrait(Operation *op) {
+ if (op->hasTrait<ZeroRegion>())
+ return op->emitOpError("is expected to have regions");
+ return success();
+ }
+};
+
/// This class provides APIs and verifiers for ops with regions having a single
/// block that must terminate with `TerminatorOpType`.
template <typename TerminatorOpType> struct SingleBlockImplicitTerminator {
diff --git a/mlir/lib/Dialect/StandardOps/IR/Ops.cpp b/mlir/lib/Dialect/StandardOps/IR/Ops.cpp
index 2f0f4b1447c0..bbe04cf6e1f8 100644
--- a/mlir/lib/Dialect/StandardOps/IR/Ops.cpp
+++ b/mlir/lib/Dialect/StandardOps/IR/Ops.cpp
@@ -295,8 +295,7 @@ static ParseResult parseAllocLikeOp(OpAsmParser &parser,
template <typename AllocLikeOp>
static LogicalResult verify(AllocLikeOp op) {
- static_assert(std::is_same<AllocLikeOp, AllocOp>::value ||
- std::is_same<AllocLikeOp, AllocaOp>::value,
+ static_assert(llvm::is_one_of<AllocLikeOp, AllocOp, AllocaOp>::value,
"applies to only alloc or alloca");
auto memRefType = op.getResult().getType().template dyn_cast<MemRefType>();
if (!memRefType)
@@ -321,7 +320,19 @@ static LogicalResult verify(AllocLikeOp op) {
for (auto operandType : op.getOperandTypes())
if (!operandType.isIndex())
return op.emitOpError("requires operands to be of type Index");
- return success();
+
+ if (std::is_same<AllocLikeOp, AllocOp>::value)
+ return success();
+
+ // An alloca op needs to have an ancestor with an allocation scope trait.
+ auto *parentOp = op.getParentOp();
+ while (parentOp) {
+ if (parentOp->template hasTrait<OpTrait::AutomaticAllocationScope>())
+ return success();
+ parentOp = parentOp->getParentOp();
+ }
+ return op.emitOpError(
+ "requires an ancestor op with AutomaticAllocationScope trait");
}
namespace {
diff --git a/mlir/test/IR/invalid-ops.mlir b/mlir/test/IR/invalid-ops.mlir
index 3f0804aff96a..c555776d9d35 100644
--- a/mlir/test/IR/invalid-ops.mlir
+++ b/mlir/test/IR/invalid-ops.mlir
@@ -1158,3 +1158,11 @@ func @assume_alignment(%0: memref<4x4xf16>) {
std.assume_alignment %0, 0 : memref<4x4xf16>
return
}
+
+// -----
+
+"alloca_without_scoped_alloc_parent"() ( {
+ std.alloca() : memref<1xf32>
+ // expected-error at -1 {{requires an ancestor op with AutomaticAllocationScope trait}}
+ return
+}) : () -> ()
More information about the Mlir-commits
mailing list