From db054d711518b6527048aa06196e443ff2e1dada Mon Sep 17 00:00:00 2001 From: Uday Bondhugula Date: Thu, 9 Apr 2020 15:16:24 +0530 Subject: [PATCH] [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 --- mlir/docs/Traits.md | 10 ++++++++++ mlir/include/mlir/Dialect/GPU/GPUOps.td | 3 ++- mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td | 3 ++- mlir/include/mlir/Dialect/SPIRV/SPIRVStructureOps.td | 2 +- mlir/include/mlir/Dialect/StandardOps/IR/Ops.td | 6 ++++-- mlir/include/mlir/IR/Function.h | 8 +++++--- mlir/include/mlir/IR/OpBase.td | 2 ++ mlir/include/mlir/IR/OpDefinition.h | 16 ++++++++++++++++ mlir/lib/Dialect/StandardOps/IR/Ops.cpp | 17 ++++++++++++++--- mlir/test/IR/invalid-ops.mlir | 8 ++++++++ 10 files changed, 64 insertions(+), 11 deletions(-) diff --git a/mlir/docs/Traits.md b/mlir/docs/Traits.md index 9877d90..1035b6c 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 981f397..5d91ff6 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 54b2e94..eac6637 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, OptionalAttr:$personality, diff --git a/mlir/include/mlir/Dialect/SPIRV/SPIRVStructureOps.td b/mlir/include/mlir/Dialect/SPIRV/SPIRVStructureOps.td index 3bcbdc8..c9f27b9 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, + AutomaticAllocationScope, DeclareOpInterfaceMethods, 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 227c006..55478a4 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 49c0c7c..1b2dd19 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 { +class FuncOp + : public Op { public: using Op::Op; using Op::print; diff --git a/mlir/include/mlir/IR/OpBase.td b/mlir/include/mlir/IR/OpBase.td index 09cea1b..58c9643 100644 --- a/mlir/include/mlir/IR/OpBase.td +++ b/mlir/include/mlir/IR/OpBase.td @@ -1587,6 +1587,8 @@ class PredOpTrait : 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 4901f5b..38403e8 100644 --- a/mlir/include/mlir/IR/OpDefinition.h +++ b/mlir/include/mlir/IR/OpDefinition.h @@ -1026,6 +1026,22 @@ public: } }; +/// 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 +class AutomaticAllocationScope + : public TraitBase { +public: + static LogicalResult verifyTrait(Operation *op) { + if (op->hasTrait()) + 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 struct SingleBlockImplicitTerminator { diff --git a/mlir/lib/Dialect/StandardOps/IR/Ops.cpp b/mlir/lib/Dialect/StandardOps/IR/Ops.cpp index 2f0f4b1..bbe04cf 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 static LogicalResult verify(AllocLikeOp op) { - static_assert(std::is_same::value || - std::is_same::value, + static_assert(llvm::is_one_of::value, "applies to only alloc or alloca"); auto memRefType = op.getResult().getType().template dyn_cast(); 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::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()) + 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 3f0804a..c555776 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@-1 {{requires an ancestor op with AutomaticAllocationScope trait}} + return +}) : () -> () -- 2.7.4