[MLIR] Introduce a trait that defines a new scope for auto allocation

[bondhugula]

Summary

Introduce a new operation 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.

[ftynse]

Would it be possible to implement this without modifying Operation, i.e. as a simple trait rather than a trait based on a core operation property? Unlike "IsTerminator" and "IsolatedFromAbove", we don't have core mechanisms (which are the block verifier and the symbol resolver in the mentioned cases) that would mandate exposing this trait in Operation APIs.

mlir/include/mlir/Dialect/GPU/GPUOps.td

89

Nit: can we keep this sorted alphabetically?

[bondhugula]

In D77787#1971695, @ftynse wrote:

Would it be possible to implement this without modifying Operation, i.e. as a simple trait rather than a trait based on a core operation property? Unlike "IsTerminator" and "IsolatedFromAbove", we don't have core mechanisms (which are the block verifier and the symbol resolver in the mentioned cases) that would mandate exposing this trait in Operation APIs.

The 'OperationProperty' enum class is missing its doc comment entirely - on how it's different from a trait that is not an OperationProperty. We should fix that first - the necessary information can't be determined locally without documentation.

[ftynse]

Most traits are not an OperationProperty, it suffices to have a quick look in OpBase.td and see ~30 different traits being used. While we should add documentation to this enum (ping @rriddle), this is not precluding you from just defining a trait and landing it.

[bondhugula]

In D77787#1971994, @ftynse wrote:

Most traits are not an OperationProperty, it suffices to have a quick look in OpBase.td and see ~30 different traits being used. While we should add documentation to this enum

This justification doesn't make much sense, and can't be the basis to make a choice for a new trait. The traits are vastly different in what they mean.

(ping @rriddle), this is not precluding you from just defining a trait and landing it.

Just using a trait that doesn't use an OperationProperty bit is pretty straightforward, but simply landing the patch quickly isn't the point. What's more important here is to have the understanding and information documented on when something is a trait and an 'OperationProperty' as well - so that future contributors adding traits do the right thing. This should go into the doc comment on 'OperationProperty'. I'm assuming you aren't blocked by this revision.

[bondhugula]

Address review comments.

[bondhugula]

Sort trait list.

[ftynse]

In D77787#1972235, @bondhugula wrote:

Just using a trait that doesn't use an OperationProperty bit is pretty straightforward, but simply landing the patch quickly isn't the point. What's more important here is to have the understanding and information documented on when something is a trait and an 'OperationProperty' as well - for the benefit of those adding traits in the future. This should go into the doc comment on 'OperationProperty'. I'm assuming you aren't blocked by this revision.

We are in a situation where we have upstream code for IR that has unclear/unspecified behavior, which was pointed out by reviewers. I consider this to be an issue similar to a new bug being introduced by a change, i.e. we may want to revert the offending change. The present diff specifies this behavior, so I consider it a fix to a problem upstream. As such, it should land once accepted instead being blocked on a missing comment, that has been missing for more than a year.

[bondhugula]

In D77787#1972255, @ftynse wrote:

In D77787#1972235, @bondhugula wrote:

Just using a trait that doesn't use an OperationProperty bit is pretty straightforward, but simply landing the patch quickly isn't the point. What's more important here is to have the understanding and information documented on when something is a trait and an 'OperationProperty' as well - for the benefit of those adding traits in the future. This should go into the doc comment on 'OperationProperty'. I'm assuming you aren't blocked by this revision.

We are in a situation where we have upstream code for IR that has unclear/unspecified behavior, which was pointed out by reviewers. I consider this to be an issue similar to a new bug being introduced by a change, i.e. we may want to revert the offending change. The present diff specifies this behavior, so I consider it a fix to a problem upstream. As such, it should land once accepted instead being blocked on a missing comment, that has been missing for more than a year.

There is still a blocking review on this - I don't think I can commit it yet.

[rriddle]

In D77787#1972235, @bondhugula wrote:

In D77787#1971994, @ftynse wrote:

Most traits are not an OperationProperty, it suffices to have a quick look in OpBase.td and see ~30 different traits being used. While we should add documentation to this enum

This justification doesn't make much sense, and can't be the basis to make a choice for a new trait. The traits are vastly different in what they mean.

(ping @rriddle), this is not precluding you from just defining a trait and landing it.

Just using a trait that doesn't use an OperationProperty bit is pretty straightforward, but simply landing the patch quickly isn't the point. What's more important here is to have the understanding and information documented on when something is a trait and an 'OperationProperty' as well - so that future contributors adding traits do the right thing. This should go into the doc comment on 'OperationProperty'. I'm assuming you aren't blocked by this revision.

In reality I've considered just removing OperationProperty. No new fields have been added(some have been removed) after I added the opaque hasTrait, because there is no need to try and bake all of these things into operation properties. The only benefit to having OperationProperty is that it is O(1) to check if an operation has a specific trait, but there are traits querying more frequently that aren't in that list.

mlir/docs/Traits.md

140

Add the -- ODS class, see above.

mlir/include/mlir/IR/OpDefinition.h

1039

What happens if the operation has a variadic number of regions?

mlir/include/mlir/IR/Operation.h

435 ↗

(On Diff #256324)

This seems unrelated.

[rriddle]

In D77787#1972365, @rriddle wrote:

In D77787#1972235, @bondhugula wrote:

In D77787#1971994, @ftynse wrote:

Most traits are not an OperationProperty, it suffices to have a quick look in OpBase.td and see ~30 different traits being used. While we should add documentation to this enum

This justification doesn't make much sense, and can't be the basis to make a choice for a new trait. The traits are vastly different in what they mean.

(ping @rriddle), this is not precluding you from just defining a trait and landing it.

Just using a trait that doesn't use an OperationProperty bit is pretty straightforward, but simply landing the patch quickly isn't the point. What's more important here is to have the understanding and information documented on when something is a trait and an 'OperationProperty' as well - so that future contributors adding traits do the right thing. This should go into the doc comment on 'OperationProperty'. I'm assuming you aren't blocked by this revision.

In reality I've considered just removing OperationProperty. No new fields have been added(some have been removed) after I added the opaque hasTrait, because there is no need to try and bake all of these things into operation properties. The only benefit to having OperationProperty is that it is O(1) to check if an operation has a specific trait, but there are traits querying more frequently that aren't in that list.

Another thing to mention, is that I have been considering splitting the huge OpDefinitions file into a better sliced traits list. Interfaces have already been split into Interfaces/ so it makes sense to move some of the different categories of traits there. In that future, only a few traits would remain within IR/ which would give a better sense of what is "core" to operations.

[bondhugula]

mlir/include/mlir/IR/OpDefinition.h

1039

That should be fine. The allocations of each are still to be freed when control transfers from that respective region to the op.

[rriddle]

mlir/include/mlir/IR/OpDefinition.h

1039

Can we remove this check then? Otherwise, specific forms of operations would be legal and others not.

[bondhugula]

In D77787#1972365, @rriddle wrote:

In D77787#1972235, @bondhugula wrote:

In D77787#1971994, @ftynse wrote:

Most traits are not an OperationProperty, it suffices to have a quick look in OpBase.td and see ~30 different traits being used. While we should add documentation to this enum

This justification doesn't make much sense, and can't be the basis to make a choice for a new trait. The traits are vastly different in what they mean.

(ping @rriddle), this is not precluding you from just defining a trait and landing it.

Just using a trait that doesn't use an OperationProperty bit is pretty straightforward, but simply landing the patch quickly isn't the point. What's more important here is to have the understanding and information documented on when something is a trait and an 'OperationProperty' as well - so that future contributors adding traits do the right thing. This should go into the doc comment on 'OperationProperty'. I'm assuming you aren't blocked by this revision.

In reality I've considered just removing OperationProperty. No new fields have been added(some have been removed) after I added the opaque hasTrait, because there is no need to try and bake all of these things into operation properties. The only benefit to having OperationProperty is that it is O(1) to check if an operation has a specific trait, but there are traits querying more frequently that aren't in that list.

This is exactly what I was trying to find out - I think I now recall this from a discussion/question I had on the op interfaces proposal last year: there are sort of three tiers: op interfaces, traits, and operation properties. Isn't the O(1) check still useful (for eg. have a O large constant wouldn't be good at the innermost loop of say CSE like opts where you'd check whether the operands commute)? You should just have the top three over there? Isn't the number of bits preferably restricted to three?

[bondhugula]

mlir/include/mlir/IR/OpDefinition.h

1039

Ops that always have zero regions can't meaningfully have scoped allocations. We could just allow it for VariadicRegions instead by basing this check on the ZeroRegion trait?

[rriddle]

mlir/include/mlir/IR/OpDefinition.h

1039

I was wondering about the situations where we may end up in such a case, e.g. a switch-like dispatch operation that ends up having no cases(either by optimization, or generated by the front end). I would just rather avoid the situation that an optimization pass ends up generating a verifier failure for otherwise valid IR. Debugging that seems like it would be a pain.

ZeroRegion trait SGTM.

[bondhugula]

Address review comments; use ZeroRegion trait + other minor fixes.

mlir/docs/Traits.md

 MLIR provides a suite of traits that provide various functionalities that are
 common across many different operations. Below is a list of some key traits that
 may be used directly by any dialect. The format of the header for each trait
 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`
 
 This trait adds the property that the operation is known to have
 [broadcast-compatible](https://docs.scipy.org/doc/numpy/user/basics.broadcasting.html)
 operands and its result types' shape is the broadcast compatible with the shape
 of the broadcasted operands. Specifically, starting from the most varying

mlir/include/mlir/Dialect/GPU/GPUOps.td

let description = [{
     Example:
 
     ```mlir
     %tIdX = "gpu.thread_id"() {dimension = "x"} : () -> (index)
     ```
   }];
 }
 
-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 = [{
     Defines a function that can be executed on a GPU. This supports memory
     attribution and its body has a particular execution model.
 
     GPU functions are either kernels (as indicated by the `kernel` attribute) or
     regular functions. The former can be launched from the host side, while the

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

def LLVM_GlobalOp
   }];
 
   let printer = "printGlobalOp(p, *this);";
   let parser = "return parseGlobalOp(parser, result);";
   let verifier = "return ::verify(*this);";
 }
 
 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,
                      OptionalAttr<ArrayAttr>:$passthrough)> {
   let summary = "LLVM dialect function, has wrapped LLVM IR function type";
 
   let regions = (region AnyRegion:$body);
 

mlir/include/mlir/Dialect/SPIRV/SPIRVStructureOps.td

let builders = [OpBuilder<[{Builder *builder, OperationState &state,
                               spirv::ExecutionModel executionModel,
                               spirv::FuncOp function,
                               ArrayRef<Attribute> interfaceVars}]>];
 }
 
 // -----
 
 def SPV_FuncOp : SPV_Op<"func", [
-    DeclareOpInterfaceMethods<CallableOpInterface>,
+    AutomaticAllocationScope, DeclareOpInterfaceMethods<CallableOpInterface>,
     FunctionLike, InModuleScope, IsolatedFromAbove, Symbol
   ]> {
   let summary = "Declare or define a function";
 
   let description = [{
     This op declares or defines a SPIR-V function using one region, which
     contains one or more blocks.
 

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

 //===----------------------------------------------------------------------===//
 // AllocaOp
 //===----------------------------------------------------------------------===//
 
 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>
     ```
 
     The optional list of dimension operands are bound to the dynamic dimensions
     specified in its memref type. In the example below, the SSA value '%d' is
     bound to the second dimension of the memref (which is dynamic).

mlir/include/mlir/IR/Function.h

 // Function Operation.
 //===--------------------------------------------------------------------===//
 
 /// FuncOp represents a function, or an operation containing one region that
 /// forms a CFG(Control Flow Graph). The region of a function is not allowed to
 /// 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::FunctionLike, OpTrait::AutomaticAllocationScope,
+                CallableOpInterface::Trait> {
 public:
   using Op::Op;
   using Op::print;
 
   static StringRef getOperationName() { return "func"; }
 
   static FuncOp create(Location location, StringRef name, FunctionType type,
                        ArrayRef<NamedAttribute> attrs = {});

mlir/include/mlir/IR/OpBase.td

 }
 
 // PredOpTrait is an op trait implemented by way of a predicate on the op.
 class PredOpTrait<string descr, Pred pred> : OpTrait {
   string description = descr;
   Pred predicate = pred;
 }
 
+// Op defines an automatic allocation scope.
+def AutomaticAllocationScope : NativeOpTrait<"AutomaticAllocationScope">;
 // Op supports operand broadcast behavior.
 def ResultsBroadcastableShape :
   NativeOpTrait<"ResultsBroadcastableShape">;
 // TODO: Alias of the above, remove post integrate.
 def Broadcastable : NativeOpTrait<"ResultsBroadcastableShape">;
 // X op Y == Y op X
 def Commutative  : NativeOpTrait<"IsCommutative">;
 // Op behaves like a constant.

mlir/include/mlir/IR/OpDefinition.h

public:
   static LogicalResult verifyTrait(Operation *op) {
     for (auto &region : op->getRegions())
       if (!region.isIsolatedFromAbove(op->getLoc()))
         return failure();
     return success();
   }
 };
 
+/// 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 {
   template <typename ConcreteType>
   class Impl : public TraitBase<ConcreteType, Impl> {
   public:
     static LogicalResult verifyTrait(Operation *op) {
       for (unsigned i = 0, e = op->getNumRegions(); i < e; ++i) {

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

return parser.emitError(parser.getNameLoc())
            << "dimension operand count does not equal memref dynamic dimension "
               "count";
   result.types.push_back(type);
   return success();
 }
 
 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)
     return op.emitOpError("result must be a memref");
 
   unsigned numSymbols = 0;
   if (!memRefType.getAffineMaps().empty()) {
     // Store number of symbols used in affine map (used in subsequent check).
     AffineMap affineMap = memRefType.getAffineMaps()[0];
     numSymbols = affineMap.getNumSymbols();
   }
 
   // Check that the total number of operands matches the number of symbols in
   // the affine map, plus the number of dynamic dimensions specified in the
   // memref type.
   unsigned numDynamicDims = memRefType.getNumDynamicDims();
   if (op.getNumOperands() != numDynamicDims + numSymbols)
     return op.emitOpError(
         "operand count does not equal dimension plus symbol operand count");
 
   // Verify that all operands are of type Index.
   for (auto operandType : op.getOperandTypes())
     if (!operandType.isIndex())
       return op.emitOpError("requires operands to be of type Index");
+
+  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 {
 /// Fold constant dimensions into an alloc like operation.
 template <typename AllocLikeOp>
 struct SimplifyAllocConst : public OpRewritePattern<AllocLikeOp> {
   using OpRewritePattern<AllocLikeOp>::OpRewritePattern;
 

mlir/test/IR/invalid-ops.mlir

 // -----
 
 // 0 alignment value.
 func @assume_alignment(%0: memref<4x4xf16>) {
   // expected-error@+1 {{'std.assume_alignment' op attribute 'alignment' failed to satisfy constraint: 32-bit signless integer attribute whose value is positive}}
   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
+}) : () -> ()