[Mlir-commits] [mlir] 670aeec - [MLIR][OpenMP][SCF] Mark parallel regions as allocation scopes
William S. Moses
llvmlistbot at llvm.org
Fri Feb 18 08:06:37 PST 2022
Author: William S. Moses
Date: 2022-02-18T11:06:32-05:00
New Revision: 670aeece51ae0e4029d078fe3a5f06b50d42f15c
URL: https://github.com/llvm/llvm-project/commit/670aeece51ae0e4029d078fe3a5f06b50d42f15c
DIFF: https://github.com/llvm/llvm-project/commit/670aeece51ae0e4029d078fe3a5f06b50d42f15c.diff
LOG: [MLIR][OpenMP][SCF] Mark parallel regions as allocation scopes
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.
Reviewed By: ftynse
Differential Revision: https://reviews.llvm.org/D119743
Added:
Modified:
mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
mlir/include/mlir/Dialect/GPU/GPUOps.td
mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
mlir/include/mlir/Dialect/SCF/SCFOps.td
mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index 5e1b910c8b68..914973df5594 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
@@ -107,7 +107,7 @@ def AffineApplyOp : Affine_Op<"apply", [NoSideEffect]> {
}
def AffineForOp : Affine_Op<"for",
- [ImplicitAffineTerminator, RecursiveSideEffects,
+ [AutomaticAllocationScope, ImplicitAffineTerminator, RecursiveSideEffects,
DeclareOpInterfaceMethods<LoopLikeOpInterface>]> {
let summary = "for operation";
let description = [{
@@ -608,7 +608,7 @@ def AffineMaxOp : AffineMinMaxOpBase<"max", [NoSideEffect]> {
}
def AffineParallelOp : Affine_Op<"parallel",
- [ImplicitAffineTerminator, RecursiveSideEffects,
+ [AutomaticAllocationScope, ImplicitAffineTerminator, RecursiveSideEffects,
DeclareOpInterfaceMethods<LoopLikeOpInterface>, MemRefsNormalizable]> {
let summary = "multi-index parallel band operation";
let description = [{
diff --git a/mlir/include/mlir/Dialect/GPU/GPUOps.td b/mlir/include/mlir/Dialect/GPU/GPUOps.td
index 5d25892175b9..ecddae6b10d0 100644
--- a/mlir/include/mlir/Dialect/GPU/GPUOps.td
+++ b/mlir/include/mlir/Dialect/GPU/GPUOps.td
@@ -439,7 +439,7 @@ def GPU_LaunchFuncOp : GPU_Op<"launch_func",
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)>,
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index ddeb698fb2a2..51ae0dc5bea2 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -80,7 +80,8 @@ def ClauseDefault : I32EnumAttr<
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 = [{
diff --git a/mlir/include/mlir/Dialect/SCF/SCFOps.td b/mlir/include/mlir/Dialect/SCF/SCFOps.td
index a2218d13ab0f..122b13f44900 100644
--- a/mlir/include/mlir/Dialect/SCF/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/SCFOps.td
@@ -110,7 +110,7 @@ def ExecuteRegionOp : SCF_Op<"execute_region"> {
}
def ForOp : SCF_Op<"for",
- [DeclareOpInterfaceMethods<LoopLikeOpInterface>,
+ [AutomaticAllocationScope, DeclareOpInterfaceMethods<LoopLikeOpInterface>,
DeclareOpInterfaceMethods<RegionBranchOpInterface>,
SingleBlockImplicitTerminator<"scf::YieldOp">,
RecursiveSideEffects]> {
@@ -404,7 +404,8 @@ def IfOp : SCF_Op<"if",
}
def ParallelOp : SCF_Op<"parallel",
- [AttrSizedOperandSegments,
+ [AutomaticAllocationScope,
+ AttrSizedOperandSegments,
DeclareOpInterfaceMethods<LoopLikeOpInterface>,
RecursiveSideEffects,
SingleBlockImplicitTerminator<"scf::YieldOp">]> {
diff --git a/mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir b/mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir
index 15b70caa930f..ab6e7a0fa049 100644
--- a/mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir
+++ b/mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir
@@ -83,7 +83,6 @@ func @materialize_read_1d_partially_specialized(%dyn1 : index, %dyn2 : index, %d
// 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
@@ -94,6 +93,7 @@ func @materialize_read(%M: index, %N: index, %O: index, %P: index) {
// 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]])
@@ -149,7 +149,6 @@ func @materialize_read(%M: index, %N: index, %O: index, %P: index) {
// 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
@@ -161,6 +160,7 @@ func @materialize_write(%M: index, %N: index, %O: index, %P: index) {
// 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]] {
More information about the Mlir-commits
mailing list