[MLIR][OpenMP][SCF] Mark parallel regions as allocation scopes

[wsmoses]

Summary

MLIR has the notion of allocation scopes which specify that stack allocations (e.g. memref.alloca, llvm.alloca) should be freed or equivalently aren't available at the end of the corresponding region.
Currently neither OpenMP parallel nor SCF parallel regions have the notion of such a scope.

This clearly makes sense for an OpenMP parallel as this is implemented in with a new function which outlines the region, and clearly any allocations in that newly outlined function have a lifetime that ends at the return of the function, by definition.

While SCF.parallel doesn't have a guaranteed runtime which it is implemented with, this similarly makes sense for SCF.parallel since otherwise an allocation within an SCF.parallel will needlessly continue to allocate stack memory that isn't cleaned up until the function (or other allocation scope op) which contains the SCF.parallel returns. This means that it is impossible to represent thread or iteration-local memory without causing a stack blow-up. In the case that this stack-blow-up behavior is intended, this can be equivalently represented with an allocation outside of the SCF.parallel with a size equal to the number of iterations.

[mehdi_amini]

This seems like something to handle uniformly across the SCF constructs I think.
The question in general is: should a stack allocation be able to escape this construct? I would think that none of the SCF construct would want this.

[wsmoses]

I definitely think there's justification for the other SCF ops to have this property, but I think this can be separate from the parallel ops question, where without this it is impossible to represent a truly thread-local allocation.

[wsmoses]

Add Affine parallel and GPU launch parallel regions.

[mehdi_amini]

I don't quite get why it is different for scf.parallel than it is for scf.for? The argument you had was " SCF.parallel will needlessly continue to allocate stack memory that isn't cleaned up until the function (or other allocation scope op) which contains the SCF.parallel returns", but is it the same for all iterations of scf.for (or scf.while) isn't it?

Why is there any specificity to the threading aspect of it?

[wsmoses]

Oh the logic as to why this is desirable as the semantic default still applies to an SCF.for, and so on.

The only use cases of behavior other than this is where the stack allocation escapes the broader op, which there could be a user of for generic SCF ops.

For most parallel operations, these semantics are not even possible (e.g. stack allocations in an omp.parallel cannot persist until the containing mlir function because the omp parallel is implemented with an outlined function, similar with a gpu launch).

Since there's the possibility of some use in other areas, my thought was to begin with a change to places which make little sense, though I'm completely happy adding this elsewhere.

[herhut]

My concern with only doing the scf.parallel operation is that we do lower it to scf.for loops in some cases and we would have to model this behavior there somehow. How would we do this? Do we have a suitable high-level abstraction that allows to insert an allocation scope? Alternatively, we could mark the scf.for loop the same. That seems a reasonable choice but scf.for supports returning values, so we might actually have escaping stack allocations currently. Is this a model we want to support?

So, to move this forward, we either have to extend scf.for in the same way or come up with a lowering strategy from scf.parallel to scf.for that captures this behavior.

[ftynse]

There is a scoping op https://mlir.llvm.org/docs/Dialects/MemRef/#memrefalloca_scope-mlirmemrefallocascopeop if we need to use it in lowerings.

This being said, I think we can also construct an example that escapes an allocaed value from scf.parallel (either store it in memory, or have a sort of "assignment reduction").

[wsmoses]

This being said, I think we can also construct an example that escapes an allocaed value from scf.parallel (either store it in memory, or have a sort of "assignment reduction").

The example of this would be

%escaping = alloca memref<N x memref<i64>>
scf.parallel %i = 0 .. N {
   %internal = alloca memref<i64>
   store %escaping[%i] = %internal
}
use(%escaping)

This, however, cannot be implemented with an omp lowering and I would argue is already semantically invalid.

Again the reason being that the OpenMP lowering (below) cannot preserve a stack allocation outside of an omp.parallel:

%escaping = alloca memref<N x memref<i64>>
omp.parallel {
  omp.for %i = 0 .. N {
     %internal = alloca memref<i64>
     store %escaping[%i] = %internal
   }
}
use(%escaping)

becomes

__kmpc_fork(outlined, escaping)
use(escaping);

void outlined(i64** escaping_cap) {
   kmpc_for(&start, &end);
   for(int i=start; i<end; i++) {
       // This will have its lifetime end at the end of the outlined call (e.g. omp.parallel). 
       %internal = alloca
       escaping_cap[i] = %internal;
   }
}

However @herhut as you say it could be implemented by an scf.for lowering. I concur with Alex that we could use the memref alloca scope when lowering to an scf.for. [1]

[1] As a minor aside, we could (but shouldn't) still use the existing scf.parallel to scf.for lowering and overallocate (e.g. continuing to extend the lifetime of the allocations), even if we could have better optimized the lowering by restricting the lifetime. That is to say, extending the lifetime of a variable should always be legal (although bad for stack size).

[herhut]

Thanks for your reference to the scoping operation @ftynse. I had the feeling we had something like it but could not find it.

Regarding extension of lifetimes. That is a good question. For example, would it be a valid canonicalization pattern to remove an alloca scope operation that has no alloca operations inside? Or would we keep it around in case something gets lowered to an alloca later? Is an allocation scope a property in itself that we should keep up?

We could either lower scf.parallel to an scf.for with allocation scope if the loop contains an alloca. Or we make the lowering illegal if the scf.parallel has an alloca operation. The latter at least preserves the status quo and avoids surprises.

And, @wsmoses, I agree that extending would be correct but it changes the stack requirement quite dramatically, especially in the context of loops. So if we widen the scope during lowering from one iteration to all iterations, that would be an issue (as you also mention). In the land of large tensors, any extension of lifetime can become problematic, even without loops.

[mehdi_amini]

This, however, cannot be implemented with an omp lowering and I would argue is already semantically invalid.
Again the reason being that the OpenMP lowering (below) cannot preserve a stack allocation outside of an omp.parallel:

I'm not comfortable using the lowering to OpenMP to justify the semantics of a higher level construct though.
In particular that wouldn't be a good reason to differentiate scf.for and scf.parallel here, I still see these as highly similar and I believe the argument for having them being "auto allocation scope" applies well to both as far as I can tell.

[wsmoses]

I mean I'm fine having it apply to scf.for as well. Presumably we would also want this to apply to affine.for, etc. Anything else?

[wsmoses]

Add to Affine.for and scf.for

[mehdi_amini]

I'm fine with this as-is, but @bondhugula should probably confirm that AffineFor can also get this right now.

[wsmoses]

Fix test

[ftynse]

I think we always intended affine loops to be allocation scopes.

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

def AffineApplyOp : Affine_Op<"apply", [NoSideEffect]> {
 
   let hasCanonicalizer = 1;
   let hasCustomAssemblyFormat = 1;
   let hasFolder = 1;
   let hasVerifier = 1;
 }
 
 def AffineForOp : Affine_Op<"for",
-    [ImplicitAffineTerminator, RecursiveSideEffects,
+    [AutomaticAllocationScope, ImplicitAffineTerminator, RecursiveSideEffects,
      DeclareOpInterfaceMethods<LoopLikeOpInterface>]> {
   let summary = "for operation";
   let description = [{
     Syntax:
 
     ```
     operation   ::= `affine.for` ssa-id `=` lower-bound `to` upper-bound
                     (`step` integer-literal)? `{` op* `}`
let description = [{
 
     ```mlir
     %0 = affine.max (d0) -> (1000, d0 + 512) (%i0) : index
     ```
   }];
 }
 
 def AffineParallelOp : Affine_Op<"parallel",
-    [ImplicitAffineTerminator, RecursiveSideEffects,
+    [AutomaticAllocationScope, ImplicitAffineTerminator, RecursiveSideEffects,
      DeclareOpInterfaceMethods<LoopLikeOpInterface>, MemRefsNormalizable]> {
   let summary = "multi-index parallel band operation";
   let description = [{
     The "affine.parallel" operation represents a hyper-rectangular affine
     parallel band, defining zero or more SSA values for its induction variables.
     It has one region capturing the parallel band body. The induction variables
     are represented as arguments of this region. These SSA values always have
     type index, which is the size of the machine word. The strides, represented

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

let assemblyFormat = [{
       `blocks` `in` ` ` `(`$gridSizeX`,` $gridSizeY`,` $gridSizeZ`)`
       `threads` `in` ` ` `(`$blockSizeX`,` $blockSizeY`,` $blockSizeZ`)`
       (`dynamic_shared_memory_size` $dynamicSharedMemorySize^)?
       custom<LaunchFuncOperands>($operands, type($operands)) attr-dict
   }];
   let hasVerifier = 1;
 }
 
-def GPU_LaunchOp : GPU_Op<"launch">,
+def GPU_LaunchOp : GPU_Op<"launch", [AutomaticAllocationScope]>,
     Arguments<(ins Index:$gridSizeX, Index:$gridSizeY, Index:$gridSizeZ,
                Index:$blockSizeX, Index:$blockSizeY, Index:$blockSizeZ,
                Optional<I32>:$dynamicSharedMemorySize)>,
     Results<(outs)> {
   let summary = "GPU kernel launch operation";
 
   let description = [{
     Launch a kernel on the specified grid of thread blocks. The body of the

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

def ClauseDefault : I32EnumAttr<
     [ClauseDefaultPrivate, ClauseDefaultFirstPrivate, ClauseDefaultShared,
      ClauseDefaultNone]> {
   let genSpecializedAttr = 0;
   let cppNamespace = "::mlir::omp";
 }
 def ClauseDefaultAttr : EnumAttr<OpenMP_Dialect, ClauseDefault,
                                  "clause_default">;
 
-def ParallelOp : OpenMP_Op<"parallel", [AttrSizedOperandSegments,
+def ParallelOp : OpenMP_Op<"parallel", [AutomaticAllocationScope, 
+                                        AttrSizedOperandSegments,
                  DeclareOpInterfaceMethods<OutlineableOpenMPOpInterface>]> {
   let summary = "parallel construct";
   let description = [{
     The parallel construct includes a region of code which is to be executed
     by a team of threads.
 
     The optional $if_expr_var parameter specifies a boolean result of a
     conditional check. If this value is 1 or is not provided then the parallel

mlir/include/mlir/Dialect/SCF/SCFOps.td

def ExecuteRegionOp : SCF_Op<"execute_region"> {
   let hasCanonicalizer = 1;
   let hasCustomAssemblyFormat = 1;
 
   let hasFolder = 0;
   let hasVerifier = 1;
 }
 
 def ForOp : SCF_Op<"for",
-      [DeclareOpInterfaceMethods<LoopLikeOpInterface>,
+      [AutomaticAllocationScope, DeclareOpInterfaceMethods<LoopLikeOpInterface>,
        DeclareOpInterfaceMethods<RegionBranchOpInterface>,
        SingleBlockImplicitTerminator<"scf::YieldOp">,
        RecursiveSideEffects]> {
   let summary = "for operation";
   let description = [{
     The "scf.for" operation represents a loop taking 3 SSA value as operands
     that represent the lower bound, upper bound and step respectively. The
     operation defines an SSA value for its induction variable. It has one
def IfOp : SCF_Op<"if",
   }];
   let hasFolder = 1;
   let hasCanonicalizer = 1;
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
 }
 
 def ParallelOp : SCF_Op<"parallel",
-    [AttrSizedOperandSegments,
+    [AutomaticAllocationScope,
+     AttrSizedOperandSegments,
      DeclareOpInterfaceMethods<LoopLikeOpInterface>,
      RecursiveSideEffects,
      SingleBlockImplicitTerminator<"scf::YieldOp">]> {
   let summary = "parallel for operation";
   let description = [{
     The "scf.parallel" operation represents a loop nest taking 4 groups of SSA
     values as operands that represent the lower bounds, upper bounds, steps and
     initial values, respectively. The operation defines a variadic number of

mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir

 
 // -----
 
 // CHECK: #[[$ADD:map.*]] = affine_map<(d0, d1) -> (d0 + d1)>
 
 // CHECK-LABEL: func @materialize_read(%{{.*}}: index, %{{.*}}: index, %{{.*}}: index, %{{.*}}: index) {
 func @materialize_read(%M: index, %N: index, %O: index, %P: index) {
   %f0 = arith.constant 0.0: f32
-  // CHECK-DAG:  %[[ALLOC:.*]] = memref.alloca() : memref<vector<5x4x3xf32>>
   // CHECK-DAG:  %[[C0:.*]] = arith.constant 0 : index
   // CHECK-DAG:  %[[C1:.*]] = arith.constant 1 : index
   // CHECK-DAG:  %[[C3:.*]] = arith.constant 3 : index
   // CHECK-DAG:  %[[C4:.*]] = arith.constant 4 : index
   // CHECK-DAG:  %[[C5:.*]] = arith.constant 5 : index
   // CHECK:      %{{.*}} = memref.alloc(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) : memref<?x?x?x?xf32>
   // CHECK-NEXT:  affine.for %[[I0:.*]] = 0 to %{{.*}} step 3 {
   // CHECK-NEXT:    affine.for %[[I1:.*]] = 0 to %{{.*}} {
   // CHECK-NEXT:      affine.for %[[I2:.*]] = 0 to %{{.*}} {
   // CHECK-NEXT:        affine.for %[[I3:.*]] = 0 to %{{.*}} step 5 {
+  // CHECK:               %[[ALLOC:.*]] = memref.alloca() : memref<vector<5x4x3xf32>>
   // CHECK:               scf.for %[[I4:.*]] = %[[C0]] to %[[C5]] step %[[C1]] {
   // CHECK:                 scf.if
   // CHECK:                   %[[L3:.*]] = affine.apply #[[$ADD]](%[[I3]], %[[I4]])
   // CHECK:                   scf.for %[[I5:.*]] = %[[C0]] to %[[C4]] step %[[C1]] {
   // CHECK:                     %[[VEC:.*]] = scf.for %[[I6:.*]] = %[[C0]] to %[[C3]] step %[[C1]] {{.*}} -> (vector<3xf32>) {
   // CHECK:                       %[[L0:.*]] = affine.apply #[[$ADD]](%[[I0]], %[[I6]])
   // CHECK:                       scf.if {{.*}} -> (vector<3xf32>) {
   // CHECK-NEXT:                    %[[SCAL:.*]] = memref.load %{{.*}}[%[[L0]], %[[I1]], %[[I2]], %[[L3]]] : memref<?x?x?x?xf32>

 }
 
 // -----
 
 // CHECK: #[[$ADD:map.*]] = affine_map<(d0, d1) -> (d0 + d1)>
 
 // CHECK-LABEL:func @materialize_write(%{{.*}}: index, %{{.*}}: index, %{{.*}}: index, %{{.*}}: index) {
 func @materialize_write(%M: index, %N: index, %O: index, %P: index) {
-  // CHECK-DAG:  %[[ALLOC:.*]] = memref.alloca() : memref<vector<5x4x3xf32>>
   // CHECK-DAG:  %{{.*}} = arith.constant dense<1.000000e+00> : vector<5x4x3xf32>
   // CHECK-DAG:  %[[C0:.*]] = arith.constant 0 : index
   // CHECK-DAG:  %[[C1:.*]] = arith.constant 1 : index
   // CHECK-DAG:  %[[C3:.*]] = arith.constant 3 : index
   // CHECK-DAG:  %[[C4:.*]] = arith.constant 4 : index
   // CHECK-DAG:  %[[C5:.*]] = arith.constant 5 : index
   // CHECK:      %{{.*}} = memref.alloc(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) : memref<?x?x?x?xf32>
   // CHECK-NEXT: affine.for %[[I0:.*]] = 0 to %{{.*}} step 3 {
   // CHECK-NEXT:   affine.for %[[I1:.*]] = 0 to %{{.*}} step 4 {
   // CHECK-NEXT:     affine.for %[[I2:.*]] = 0 to %{{.*}} {
   // CHECK-NEXT:       affine.for %[[I3:.*]] = 0 to %{{.*}} step 5 {
+  // CHECK:              %[[ALLOC:.*]] = memref.alloca() : memref<vector<5x4x3xf32>>
   // CHECK:              memref.store %{{.*}}, %[[ALLOC]][] : memref<vector<5x4x3xf32>>
   // CHECK:              %[[VECTOR_VIEW1:.*]] = vector.type_cast %[[ALLOC]] : memref<vector<5x4x3xf32>> to memref<5xvector<4x3xf32>>
   // CHECK:              scf.for %[[I4:.*]] = %[[C0]] to %[[C5]] step %[[C1]] {
   // CHECK:                scf.if
   // CHECK:                  %[[S3:.*]] = affine.apply #[[$ADD]](%[[I3]], %[[I4]])
   // CHECK:                  %[[VECTOR_VIEW2:.*]] = vector.type_cast %[[VECTOR_VIEW1]] : memref<5xvector<4x3xf32>> to memref<5x4xvector<3xf32>>
   // CHECK:                  scf.for %[[I5:.*]] = %[[C0]] to %[[C4]] step %[[C1]] {
   // CHECK:                    scf.if