diff --git a/mlir/docs/Traits.md b/mlir/docs/Traits.md --- a/mlir/docs/Traits.md +++ b/mlir/docs/Traits.md @@ -135,6 +135,16 @@ * `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 --- a/mlir/include/mlir/Dialect/GPU/GPUOps.td +++ b/mlir/include/mlir/Dialect/GPU/GPUOps.td @@ -85,7 +85,8 @@ }]; } -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 --- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td +++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td @@ -660,7 +660,8 @@ } 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 --- a/mlir/include/mlir/Dialect/SPIRV/SPIRVStructureOps.td +++ b/mlir/include/mlir/Dialect/SPIRV/SPIRVStructureOps.td @@ -197,7 +197,7 @@ // ----- 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 --- a/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td +++ b/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td @@ -326,8 +326,10 @@ 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 --- a/mlir/include/mlir/IR/Function.h +++ b/mlir/include/mlir/IR/Function.h @@ -30,9 +30,11 @@ /// 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 --- a/mlir/include/mlir/IR/OpBase.td +++ b/mlir/include/mlir/IR/OpBase.td @@ -1587,6 +1587,8 @@ 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 --- a/mlir/include/mlir/IR/OpDefinition.h +++ b/mlir/include/mlir/IR/OpDefinition.h @@ -1026,6 +1026,22 @@ } }; +/// 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 --- a/mlir/lib/Dialect/StandardOps/IR/Ops.cpp +++ b/mlir/lib/Dialect/StandardOps/IR/Ops.cpp @@ -295,8 +295,7 @@ 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 @@ 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 --- a/mlir/test/IR/invalid-ops.mlir +++ b/mlir/test/IR/invalid-ops.mlir @@ -1158,3 +1158,11 @@ 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 +}) : () -> ()