[flang-commits] [flang] 65b9b9a - Add Allocate Clause to MLIR Parallel Operation Definition

Irina Dobrescu via flang-commits flang-commits at lists.llvm.org
Wed Oct 14 09:14:08 PDT 2020


Author: Irina Dobrescu
Date: 2020-10-14T17:13:48+01:00
New Revision: 65b9b9aa5014ed40afec2beaefcbe73392d16e30

URL: https://github.com/llvm/llvm-project/commit/65b9b9aa5014ed40afec2beaefcbe73392d16e30
DIFF: https://github.com/llvm/llvm-project/commit/65b9b9aa5014ed40afec2beaefcbe73392d16e30.diff

LOG: Add Allocate Clause to MLIR Parallel Operation Definition

Differential Revision: https://reviews.llvm.org/D87684

Added: 
    

Modified: 
    flang/lib/Lower/OpenMP.cpp
    mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
    mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
    mlir/test/Dialect/OpenMP/ops.mlir

Removed: 
    


################################################################################
diff  --git a/flang/lib/Lower/OpenMP.cpp b/flang/lib/Lower/OpenMP.cpp
index 694ef7e33d76..169e225bdae9 100644
--- a/flang/lib/Lower/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP.cpp
@@ -109,7 +109,7 @@ genOMP(Fortran::lower::AbstractConverter &converter,
     auto parallelOp = firOpBuilder.create<mlir::omp::ParallelOp>(
         currentLocation, argTy, Value(), numThreads,
         defaultValue.dyn_cast_or_null<StringAttr>(), ValueRange(), ValueRange(),
-        ValueRange(), ValueRange(),
+        ValueRange(), ValueRange(), ValueRange(), ValueRange(),
         procBindValue.dyn_cast_or_null<StringAttr>());
     firOpBuilder.createBlock(&parallelOp.getRegion());
     auto &block = parallelOp.getRegion().back();

diff  --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 02bae9d43322..17845a598c31 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -65,6 +65,9 @@ def ParallelOp : OpenMP_Op<"parallel", [AttrSizedOperandSegments]> {
     are a variadic list of values that specify the data sharing attribute of
     those values.
 
+    The $allocators_vars and $allocate_vars parameters are a variadic list of values
+    that specify the memory allocator to be used to obtain storage for private values.
+
     The optional $proc_bind_val attribute controls the thread affinity for the execution
     of the parallel region.
   }];
@@ -76,12 +79,15 @@ def ParallelOp : OpenMP_Op<"parallel", [AttrSizedOperandSegments]> {
              Variadic<AnyType>:$firstprivate_vars,
              Variadic<AnyType>:$shared_vars,
              Variadic<AnyType>:$copyin_vars,
+             Variadic<AnyType>:$allocate_vars,
+             Variadic<AnyType>:$allocators_vars,
              OptionalAttr<ProcBindKind>:$proc_bind_val);
 
   let regions = (region AnyRegion:$region);
 
   let parser = [{ return parseParallelOp(parser, result); }];
   let printer = [{ return printParallelOp(p, *this); }];
+  let verifier = [{ return ::verifyParallelOp(*this); }];
 }
 
 def TerminatorOp : OpenMP_Op<"terminator", [Terminator]> {

diff  --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 84db8e798b2f..7ab4534f924a 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -41,7 +41,7 @@ void OpenMPDialect::initialize() {
 ///
 /// operand-and-type-list ::= `(` ssa-id-and-type-list `)`
 /// ssa-id-and-type-list ::= ssa-id-and-type |
-///                          ssa-id-and-type ',' ssa-id-and-type-list
+///                          ssa-id-and-type `,` ssa-id-and-type-list
 /// ssa-id-and-type ::= ssa-id `:` type
 static ParseResult
 parseOperandAndTypeList(OpAsmParser &parser,
@@ -65,6 +65,52 @@ parseOperandAndTypeList(OpAsmParser &parser,
   return success();
 }
 
+/// Parse an allocate clause with allocators and a list of operands with types.
+///
+/// operand-and-type-list ::= `(` allocate-operand-list `)`
+/// allocate-operand-list :: = allocate-operand |
+///                            allocator-operand `,` allocate-operand-list
+/// allocate-operand :: = ssa-id-and-type -> ssa-id-and-type
+/// ssa-id-and-type ::= ssa-id `:` type
+static ParseResult parseAllocateAndAllocator(
+    OpAsmParser &parser,
+    SmallVectorImpl<OpAsmParser::OperandType> &operandsAllocate,
+    SmallVectorImpl<Type> &typesAllocate,
+    SmallVectorImpl<OpAsmParser::OperandType> &operandsAllocator,
+    SmallVectorImpl<Type> &typesAllocator) {
+  if (parser.parseLParen())
+    return failure();
+
+  do {
+    OpAsmParser::OperandType operand;
+    Type type;
+
+    if (parser.parseOperand(operand) || parser.parseColonType(type))
+      return failure();
+    operandsAllocator.push_back(operand);
+    typesAllocator.push_back(type);
+    if (parser.parseArrow())
+      return failure();
+    if (parser.parseOperand(operand) || parser.parseColonType(type))
+      return failure();
+
+    operandsAllocate.push_back(operand);
+    typesAllocate.push_back(type);
+  } while (succeeded(parser.parseOptionalComma()));
+
+  if (parser.parseRParen())
+    return failure();
+
+  return success();
+}
+
+static LogicalResult verifyParallelOp(ParallelOp op) {
+  if (op.allocate_vars().size() != op.allocators_vars().size())
+    return op.emitError(
+        "expected equal sizes for allocate and allocator variables");
+  return success();
+}
+
 static void printParallelOp(OpAsmPrinter &p, ParallelOp op) {
   p << "omp.parallel";
 
@@ -84,10 +130,26 @@ static void printParallelOp(OpAsmPrinter &p, ParallelOp op) {
       }
     }
   };
+
+  // Print allocator and allocate parameters
+  auto printAllocateAndAllocator = [&p](OperandRange varsAllocate,
+                                        OperandRange varsAllocator) {
+    if (varsAllocate.empty())
+      return;
+
+    p << " allocate(";
+    for (unsigned i = 0; i < varsAllocate.size(); ++i) {
+      std::string separator = i == varsAllocate.size() - 1 ? ")" : ", ";
+      p << varsAllocator[i] << " : " << varsAllocator[i].getType() << " -> ";
+      p << varsAllocate[i] << " : " << varsAllocate[i].getType() << separator;
+    }
+  };
+
   printDataVars("private", op.private_vars());
   printDataVars("firstprivate", op.firstprivate_vars());
   printDataVars("shared", op.shared_vars());
   printDataVars("copyin", op.copyin_vars());
+  printAllocateAndAllocator(op.allocate_vars(), op.allocators_vars());
 
   if (auto def = op.default_val())
     p << " default(" << def->drop_front(3) << ")";
@@ -118,6 +180,7 @@ static ParseResult allowedOnce(OpAsmParser &parser, llvm::StringRef clause,
 /// firstprivate ::= `firstprivate` operand-and-type-list
 /// shared ::= `shared` operand-and-type-list
 /// copyin ::= `copyin` operand-and-type-list
+/// allocate ::= `allocate` operand-and-type `->` operand-and-type-list
 /// default ::= `default` `(` (`private` | `firstprivate` | `shared` | `none`)
 /// procBind ::= `proc_bind` `(` (`master` | `close` | `spread`) `)`
 ///
@@ -134,7 +197,11 @@ static ParseResult parseParallelOp(OpAsmParser &parser,
   SmallVector<Type, 4> sharedTypes;
   SmallVector<OpAsmParser::OperandType, 4> copyins;
   SmallVector<Type, 4> copyinTypes;
-  std::array<int, 6> segments{0, 0, 0, 0, 0, 0};
+  SmallVector<OpAsmParser::OperandType, 4> allocates;
+  SmallVector<Type, 4> allocateTypes;
+  SmallVector<OpAsmParser::OperandType, 4> allocators;
+  SmallVector<Type, 4> allocatorTypes;
+  std::array<int, 8> segments{0, 0, 0, 0, 0, 0, 0, 0};
   llvm::StringRef keyword;
   bool defaultVal = false;
   bool procBind = false;
@@ -145,6 +212,8 @@ static ParseResult parseParallelOp(OpAsmParser &parser,
   const int firstprivateClausePos = 3;
   const int sharedClausePos = 4;
   const int copyinClausePos = 5;
+  const int allocateClausePos = 6;
+  const int allocatorPos = 7;
   const llvm::StringRef opName = result.name.getStringRef();
 
   while (succeeded(parser.parseOptionalKeyword(&keyword))) {
@@ -192,6 +261,15 @@ static ParseResult parseParallelOp(OpAsmParser &parser,
       if (parseOperandAndTypeList(parser, copyins, copyinTypes))
         return failure();
       segments[copyinClausePos] = copyins.size();
+    } else if (keyword == "allocate") {
+      // fail if there was already another allocate clause
+      if (segments[allocateClausePos])
+        return allowedOnce(parser, "allocate", opName);
+      if (parseAllocateAndAllocator(parser, allocates, allocateTypes,
+                                    allocators, allocatorTypes))
+        return failure();
+      segments[allocateClausePos] = allocates.size();
+      segments[allocatorPos] = allocators.size();
     } else if (keyword == "default") {
       // fail if there was already another default clause
       if (defaultVal)
@@ -261,6 +339,18 @@ static ParseResult parseParallelOp(OpAsmParser &parser,
                              result.operands))
     return failure();
 
+  // Add allocate parameters
+  if (segments[allocateClausePos] &&
+      parser.resolveOperands(allocates, allocateTypes, allocates[0].location,
+                             result.operands))
+    return failure();
+
+  // Add allocator parameters
+  if (segments[allocatorPos] &&
+      parser.resolveOperands(allocators, allocatorTypes, allocators[0].location,
+                             result.operands))
+    return failure();
+
   result.addAttribute("operand_segment_sizes",
                       parser.getBuilder().getI32VectorAttr(segments));
 

diff  --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index c74f55b761c2..add50414d71c 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -52,37 +52,43 @@ func @omp_terminator() -> () {
 }
 
 func @omp_parallel(%data_var : memref<i32>, %if_cond : i1, %num_threads : si32) -> () {
-  // CHECK: omp.parallel if(%{{.*}}) num_threads(%{{.*}} : si32) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) shared(%{{.*}} : memref<i32>) copyin(%{{.*}} : memref<i32>)
-  "omp.parallel" (%if_cond, %num_threads, %data_var, %data_var, %data_var, %data_var) ({
+  // CHECK: omp.parallel if(%{{.*}}) num_threads(%{{.*}} : si32) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) shared(%{{.*}} : memref<i32>) copyin(%{{.*}} : memref<i32>) allocate(%{{.*}} : memref<i32> -> %{{.*}} : memref<i32>)
+  "omp.parallel" (%if_cond, %num_threads, %data_var, %data_var, %data_var, %data_var, %data_var, %data_var) ({
 
   // test without if condition
-  // CHECK: omp.parallel num_threads(%{{.*}} : si32) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) shared(%{{.*}} : memref<i32>) copyin(%{{.*}} : memref<i32>)
-    "omp.parallel"(%num_threads, %data_var, %data_var, %data_var, %data_var) ({
+  // CHECK: omp.parallel num_threads(%{{.*}} : si32) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) shared(%{{.*}} : memref<i32>) copyin(%{{.*}} : memref<i32>) allocate(%{{.*}} : memref<i32> -> %{{.*}} : memref<i32>)
+    "omp.parallel"(%num_threads, %data_var, %data_var, %data_var, %data_var, %data_var, %data_var) ({
       omp.terminator
-    }) {operand_segment_sizes = dense<[0,1,1,1,1,1]>: vector<6xi32>, default_val = "defshared"} : (si32, memref<i32>, memref<i32>, memref<i32>, memref<i32>) -> ()
+    }) {operand_segment_sizes = dense<[0,1,1,1,1,1,1,1]>: vector<8xi32>, default_val = "defshared"} : (si32, memref<i32>, memref<i32>, memref<i32>, memref<i32>, memref<i32>, memref<i32>) -> ()
 
   // CHECK: omp.barrier
     omp.barrier
 
   // test without num_threads
-  // CHECK: omp.parallel if(%{{.*}}) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) shared(%{{.*}} : memref<i32>) copyin(%{{.*}} : memref<i32>)
-    "omp.parallel"(%if_cond, %data_var, %data_var, %data_var, %data_var) ({
+  // CHECK: omp.parallel if(%{{.*}}) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) shared(%{{.*}} : memref<i32>) copyin(%{{.*}} : memref<i32>) allocate(%{{.*}} : memref<i32> -> %{{.*}} : memref<i32>)
+    "omp.parallel"(%if_cond, %data_var, %data_var, %data_var, %data_var, %data_var, %data_var) ({
+      omp.terminator
+    }) {operand_segment_sizes = dense<[1,0,1,1,1,1,1,1]> : vector<8xi32>} : (i1, memref<i32>, memref<i32>, memref<i32>, memref<i32>, memref<i32>, memref<i32>) -> ()
+
+  // test without allocate
+  // CHECK: omp.parallel if(%{{.*}}) num_threads(%{{.*}} : si32) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) shared(%{{.*}} : memref<i32>) copyin(%{{.*}} : memref<i32>)
+    "omp.parallel"(%if_cond, %num_threads, %data_var, %data_var, %data_var, %data_var) ({
       omp.terminator
-    }) {operand_segment_sizes = dense<[1,0,1,1,1,1]> : vector<6xi32>} : (i1, memref<i32>, memref<i32>, memref<i32>, memref<i32>) -> ()
+    }) {operand_segment_sizes = dense<[1,1,1,1,1,1,0,0]> : vector<8xi32>} : (i1, si32, memref<i32>, memref<i32>, memref<i32>, memref<i32>) -> ()
 
     omp.terminator
-  }) {operand_segment_sizes = dense<[1,1,1,1,1,1]> : vector<6xi32>, proc_bind_val = "spread"} : (i1, si32, memref<i32>, memref<i32>, memref<i32>, memref<i32>) -> ()
+  }) {operand_segment_sizes = dense<[1,1,1,1,1,1,1,1]> : vector<8xi32>, proc_bind_val = "spread"} : (i1, si32, memref<i32>, memref<i32>, memref<i32>, memref<i32>, memref<i32>, memref<i32>) -> ()
 
   // test with multiple parameters for single variadic argument
-  // CHECK: omp.parallel private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>, %{{.*}} : memref<i32>) shared(%{{.*}} : memref<i32>) copyin(%{{.*}} : memref<i32>)
-  "omp.parallel" (%data_var, %data_var, %data_var, %data_var, %data_var) ({
+  // CHECK: omp.parallel private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>, %{{.*}} : memref<i32>) shared(%{{.*}} : memref<i32>) copyin(%{{.*}} : memref<i32>) allocate(%{{.*}} : memref<i32> -> %{{.*}} : memref<i32>)
+  "omp.parallel" (%data_var, %data_var, %data_var, %data_var, %data_var, %data_var, %data_var) ({
     omp.terminator
-  }) {operand_segment_sizes = dense<[0,0,1,2,1,1]> : vector<6xi32>} : (memref<i32>, memref<i32>, memref<i32>, memref<i32>, memref<i32>) -> ()
+  }) {operand_segment_sizes = dense<[0,0,1,2,1,1,1,1]> : vector<8xi32>} : (memref<i32>, memref<i32>, memref<i32>, memref<i32>, memref<i32>, memref<i32>, memref<i32>) -> ()
 
   return
 }
 
-func @omp_parallel_pretty(%data_var : memref<i32>, %if_cond : i1, %num_threads : si32) -> () {
+func @omp_parallel_pretty(%data_var : memref<i32>, %if_cond : i1, %num_threads : si32, %allocator : si32) -> () {
   // CHECK: omp.parallel
   omp.parallel {
     omp.terminator
@@ -93,6 +99,11 @@ func @omp_parallel_pretty(%data_var : memref<i32>, %if_cond : i1, %num_threads :
     omp.terminator
   }
 
+  // CHECK: omp.parallel allocate(%{{.*}} : memref<i32> -> %{{.*}} : memref<i32>)
+  omp.parallel allocate(%data_var : memref<i32> -> %data_var : memref<i32>) {
+    omp.terminator
+  }
+
   // CHECK: omp.parallel private(%{{.*}} : memref<i32>, %{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>)
   omp.parallel private(%data_var : memref<i32>, %data_var : memref<i32>) firstprivate(%data_var : memref<i32>) {
     omp.terminator


        


More information about the flang-commits mailing list