[mlir]Add Allocate Clause to OMP Parallel Operation Definition

[ftynse]

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

68–70

Nit: "variable" is unclear in MLIR context, we only have values.

Also please terminate sentences with a dot, this is required for documentation, comments, etc. except error messages.

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

73

Nit: please use backticks consistently (here and above)

137

Nit: prefer early-returns, i.e. if (varAllocate.empty()) return; <current-code-with-less-indentation>

138–140

Nit: this does not look properly formatted

141–145

mlir::interleaveComma will make it easier to read

200–205

SmallVector is re-exported into the mlir namespace, please drop llvm:: here and above.

342

Nit: drop trivial braces

[DavidTruby]

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

88

Is it possible for us to "fix" the type of allocate, since we know it'll always be an intptr_t? Genuine question, I'm not sure if this is possible/desirable in mlir.

138–140

I also checked this because it looked odd, but this is how clang-format leaves this for me as well.

200–205

Shall we remove the llvm:: prefix on the existing variables as a separate patch, so as not to expand the scope of this one? I can submit that patch since those SmallVectors were added by me in the first place.

mlir/test/Dialect/OpenMP/ops.mlir

102

Question in general for these clauses, is it possible to elide any of these types or do they need to be specified each time?

This applies to the other parallel clauses as well as just this one: when I was writing those I couldn't find a way to elide the types.

Generally I think it doesn't make those ones that much more verbose, but here it does make them quite verbose due to the necessary "->" syntax etc. We might find this verbosity gets even worse for future clauses on other omp directives.

[ftynse]

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

88

MLIR's equivalent of intptr_t is index, but I suppose this also needs to be compatible with LLVM dialect integers.

138–140

Would it re-split lines if this had been p << " " << "allocate" << "("; ? It also looks like p << " allocate("; is just fine here.

200–205

Yes please!

mlir/test/Dialect/OpenMP/ops.mlir

102

It's possible to omit types in the pretty syntax if the parser can infer them within the scope of one operation. E.g., load/store expect subscript types to be index so they can just call parser.getBuilder().getIndexType() instead of requiring it to be present in the assembly form.

[DavidTruby]

mlir/test/Dialect/OpenMP/ops.mlir

102

So, it is possible to omit them in e.g. the if clause that I added before (since we know that the condition is always an i1) but not possible for example in copyin where the type is AnyType? Even though that type is set when that variable is passed into the function itself?

[ftynse]

mlir/test/Dialect/OpenMP/ops.mlir

102

Ultimately you need to be able to construct an operation given only the information present in the assembly syntax. At parsing time, the operand values may not have been created yet, so you cannot query their type.

[Rin]

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

138–140

I'll try this and see if it formats it weird again.

[Rin]

Address review comments

[Rin]

Fix Failing Test

[ftynse]

There are review comments that remain unaddressed on this. Please address or comment as to why you chose not to address them.

[Rin]

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

141–145

I've tried to use that and it seemed to make the code more complex and so l decided not to use it.

342

When adding the allocate and allocator parameters I followed the structure of the already implemented clauses, which use braces, so I thought it would be a good idea to keep those the same.

[Rin]

Address review comments.

[kiranchandramohan]

Comments inline.

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

135

There is a mismatch in the order of arguments between the caller and the callee here.
In the caller the first argument is allocate_vars and the second is allocator_vars. Here in the callee the order is reversed.

342

I believe this suggestion is as per llvm coding standards.
https://llvm.org/docs/CodingStandards.html#don-t-use-braces-on-simple-single-statement-bodies-of-if-else-loop-statements

Agree that there is other code which also needs to be changed. Perhaps in a different patch.

[Rin]

Address review comments

[DavidTruby]

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

200–205

This is done for the existing ones in D88025

[Rin]

Add verifier

[kiranchandramohan]

LGTM subject to addressing one comment inline. Would it be possible to add an error check for allocate parsing?
Please wait for approval from other reviewers.

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

90

Nit question: Should this be renamed to verifyParallelOp?

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

70–73

Unless I am missing something, this needs correction. Suggestion below, please consult other examples and suitably modify before use.

/// operand-and-type-list ::= `(` allocate-operand-list `)`
/// allocate-operand-list :: = allocate-operand | 
                                               allocate-operand `,` allocate-operand-list
/// allocate-operand :: = ssa-id-and-type -> ssa-id-and-type
/// ssa-id-and-type ::= ssa-id `:` type

[Rin]

Address review comments

[SouraVX]

Since it's already approved, feel free to remove test case modification and commit no need for another revision :)
EDIT: if other reviewers comments are already addressed.

flang/unittests/Lower/OpenMPLoweringTest.cpp

86

This unit-tests was removed in https://reviews.llvm.org/rGbe1197c403b22291e35cbc5e96788860ceabd40c

[Rin]

@ftynse Does the patch look alright?

[Rin]

Add aditional check and remove unnecessary test

[Rin]

flang/lib/Lower/OpenMP.cpp

112 ↗

(On Diff #297263)

Here these are just a placeholder for the allocate and allocator. They will be replaced with the actual lowering implementation for the allocate clause in my next patch.

[SouraVX]

flang/lib/Lower/OpenMP.cpp

112 ↗

(On Diff #297263)

This may cause conflict in FIR-dev. Could you please take a look ?

flang/unittests/Lower/OpenMPLoweringTest.cpp

struct Fortran::parser::OmpSimpleStandaloneDirective parallelDirective(
       llvm::omp::Directive::OMPD_parallel);
 
   // Check and lower the `!OMP parallel` node to `ParallelOp` operation of
   // OpenMPDialect.
   EXPECT_EQ(parallelDirective.v, llvm::omp::Directive::OMPD_parallel);
   auto insertPt = mlirOpBuilder->saveInsertionPoint();
   llvm::ArrayRef<mlir::Type> argTy;
   mlir::ValueRange range;
-  llvm::SmallVector<int32_t, 6> operandSegmentSizes(6 /*Size=*/, 0 /*Value=*/);
+  llvm::SmallVector<int32_t, 8> operandSegmentSizes(8 /*Size=*/, 0 /*Value=*/);
   // create and insert the operation.
   auto parallelOp = mlirOpBuilder->create<mlir::omp::ParallelOp>(
       mlirOpBuilder->getUnknownLoc(), argTy, range);
   parallelOp.setAttr(mlir::omp::ParallelOp::getOperandSegmentSizeAttr(),
       mlirOpBuilder->getI32VectorAttr(operandSegmentSizes));
   parallelOp.getRegion().push_back(new mlir::Block{});
   auto &block = parallelOp.getRegion().back();
   mlirOpBuilder->setInsertionPointToStart(&block);
   // ensure the block is well-formed.
   mlirOpBuilder->create<mlir::omp::TerminatorOp>(
       mlirOpBuilder->getUnknownLoc());
   mlirOpBuilder->restoreInsertionPoint(insertPt);
   EXPECT_EQ(succeeded(parallelOp.verify()), true);
 }
 
 // main() from gtest_main

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

let description = [{
     The optional $default_val attribute specifies the default data sharing attribute
     of variables used in the parallel region that are not passed explicitly as parameters
     to the operation.
 
     The $private_vars, $firstprivate_vars, $shared_vars and $copyin_vars parameters
     are a variadic list of variables that specify the data sharing attribute of
     those variables.
 
+    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.
   }];
 
   let arguments = (ins Optional<AnyType>:$if_expr_var,
              Optional<AnyType>:$num_threads_var,
              OptionalAttr<ClauseDefault>:$default_val,
              Variadic<AnyType>:$private_vars,
              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 ::verifyEqualSize(*this); }];
 }
 
 def TerminatorOp : OpenMP_Op<"terminator", [Terminator]> {
   let summary = "terminator for OpenMP regions";
   let description = [{
     A terminator operation for regions that appear in the body of OpenMP
     operation.  These regions are not expected to return any value so the
     terminator takes no operands. The terminator op returns control to the

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

 //===----------------------------------------------------------------------===//
 // ParallelOp
 //===----------------------------------------------------------------------===//
 
 /// Parse a list of operands with types.
 ///
 /// 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,
                         SmallVectorImpl<OpAsmParser::OperandType> &operands,
                         SmallVectorImpl<Type> &types) {
   if (parser.parseLParen())
     return failure();
 
   do {
     OpAsmParser::OperandType operand;
     Type type;
     if (parser.parseOperand(operand) || parser.parseColonType(type))
       return failure();
     operands.push_back(operand);
     types.push_back(type);
   } while (succeeded(parser.parseOptionalComma()));
 
   if (parser.parseRParen())
     return failure();
 
   return success();
 }
 
+/// Parse an allocate clause with allocators and a list of operands with types.
+///
+/// operand-and-type ::= `(` ssa-id-and-type `)` `->`
+/// 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 `:` 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 verifyEqualSize(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";
 
   if (auto ifCond = op.if_expr_var())
     p << " if(" << ifCond << " : " << ifCond.getType() << ")";
 
   if (auto threads = op.num_threads_var())
     p << " num_threads(" << threads << " : " << threads.getType() << ")";
 
   // Print private, firstprivate, shared and copyin parameters
   auto printDataVars = [&p](StringRef name, OperandRange vars) {
     if (vars.size()) {
       p << " " << name << "(";
       for (unsigned i = 0; i < vars.size(); ++i) {
         std::string separator = i == vars.size() - 1 ? ")" : ", ";
         p << vars[i] << " : " << vars[i].getType() << separator;
       }
     }
   };
+
+  // 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) << ")";
 
   if (auto bind = op.proc_bind_val())
     p << " proc_bind(" << bind << ")";
 
   p.printRegion(op.getRegion());

 /// clause ::= if | numThreads | private | firstprivate | shared | copyin |
 ///            default | procBind
 /// if ::= `if` `(` ssa-id `)`
 /// numThreads ::= `num_threads` `(` ssa-id-and-type `)`
 /// private ::= `private` operand-and-type-list
 /// 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`) `)`
 ///
 /// Note that each clause can only appear once in the clase-list.
 static ParseResult parseParallelOp(OpAsmParser &parser,
                                    OperationState &result) {
   std::pair<OpAsmParser::OperandType, Type> ifCond;
   std::pair<OpAsmParser::OperandType, Type> numThreads;
   llvm::SmallVector<OpAsmParser::OperandType, 4> privates;
   llvm::SmallVector<Type, 4> privateTypes;
   llvm::SmallVector<OpAsmParser::OperandType, 4> firstprivates;
   llvm::SmallVector<Type, 4> firstprivateTypes;
   llvm::SmallVector<OpAsmParser::OperandType, 4> shareds;
   llvm::SmallVector<Type, 4> sharedTypes;
   llvm::SmallVector<OpAsmParser::OperandType, 4> copyins;
   llvm::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;
 
   const int ifClausePos = 0;
   const int numThreadsClausePos = 1;
   const int privateClausePos = 2;
   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))) {
     if (keyword == "if") {
       // Fail if there was already another if condition
       if (segments[ifClausePos])
         return allowedOnce(parser, "if", opName);
       if (parser.parseLParen() || parser.parseOperand(ifCond.first) ||
if (keyword == "if") {
       segments[sharedClausePos] = shareds.size();
     } else if (keyword == "copyin") {
       // fail if there was already another copyin clause
       if (segments[copyinClausePos])
         return allowedOnce(parser, "copyin", opName);
       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)
         return allowedOnce(parser, "default", opName);
       defaultVal = true;
       llvm::StringRef defval;
       if (parser.parseLParen() || parser.parseKeyword(&defval) ||
           parser.parseRParen())
static ParseResult parseParallelOp(OpAsmParser &parser,
   }
 
   // Add copyin parameters
   if (segments[copyinClausePos]) {
     parser.resolveOperands(copyins, copyinTypes, copyins[0].location,
                            result.operands);
   }
 
+  // Add allocate parameters
+  if (segments[allocateClausePos])
+    parser.resolveOperands(allocates, allocateTypes, allocates[0].location,
+                           result.operands);
+
+  // Add allocator parameters
+  if (segments[allocatorPos])
+    parser.resolveOperands(allocators, allocatorTypes, allocators[0].location,
+                           result.operands);
+
   result.addAttribute("operand_segment_sizes",
                       parser.getBuilder().getI32VectorAttr(segments));
 
   Region *body = result.addRegion();
   llvm::SmallVector<OpAsmParser::OperandType, 4> regionArgs;
   llvm::SmallVector<Type, 4> regionArgTypes;
   if (parser.parseRegion(*body, regionArgs, regionArgTypes))
     return failure();

mlir/test/Dialect/OpenMP/ops.mlir

 }
 
 func @omp_terminator() -> () {
   // CHECK: omp.terminator
   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]> : vector<6xi32>} : (i1, memref<i32>, memref<i32>, memref<i32>, memref<i32>) -> ()
+    }) {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,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,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,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
   }
 
   // CHECK: omp.parallel num_threads(%{{.*}} : si32)
   omp.parallel num_threads(%num_threads : si32) {
     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
   }
 
   // CHECK omp.parallel shared(%{{.*}} : memref<i32>) copyin(%{{.*}} : memref<i32>, %{{.*}} : memref<i32>)
   omp.parallel shared(%data_var : memref<i32>) copyin(%data_var : memref<i32>, %data_var : memref<i32>) {
     omp.parallel if(%if_cond: i1) {