[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

@@ -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`

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 = [{

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,

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";

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>

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;

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">;

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 <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 {

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 {

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
+}) : () -> ()