Skip to content

Commit 65b9b9a

Browse files
author
Irina Dobrescu
committed
Add Allocate Clause to MLIR Parallel Operation Definition
Differential Revision: https://reviews.llvm.org/D87684
1 parent 4212533 commit 65b9b9a

File tree

4 files changed

+123
-16
lines changed

4 files changed

+123
-16
lines changed

‎flang/lib/Lower/OpenMP.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -109,7 +109,7 @@ genOMP(Fortran::lower::AbstractConverter &converter,
109109
auto parallelOp = firOpBuilder.create<mlir::omp::ParallelOp>(
110110
currentLocation, argTy, Value(), numThreads,
111111
defaultValue.dyn_cast_or_null<StringAttr>(), ValueRange(), ValueRange(),
112-
ValueRange(), ValueRange(),
112+
ValueRange(), ValueRange(), ValueRange(), ValueRange(),
113113
procBindValue.dyn_cast_or_null<StringAttr>());
114114
firOpBuilder.createBlock(&parallelOp.getRegion());
115115
auto &block = parallelOp.getRegion().back();

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,9 @@ def ParallelOp : OpenMP_Op<"parallel", [AttrSizedOperandSegments]> {
6565
are a variadic list of values that specify the data sharing attribute of
6666
those values.
6767

68+
The $allocators_vars and $allocate_vars parameters are a variadic list of values
69+
that specify the memory allocator to be used to obtain storage for private values.
70+
6871
The optional $proc_bind_val attribute controls the thread affinity for the execution
6972
of the parallel region.
7073
}];
@@ -76,12 +79,15 @@ def ParallelOp : OpenMP_Op<"parallel", [AttrSizedOperandSegments]> {
7679
Variadic<AnyType>:$firstprivate_vars,
7780
Variadic<AnyType>:$shared_vars,
7881
Variadic<AnyType>:$copyin_vars,
82+
Variadic<AnyType>:$allocate_vars,
83+
Variadic<AnyType>:$allocators_vars,
7984
OptionalAttr<ProcBindKind>:$proc_bind_val);
8085

8186
let regions = (region AnyRegion:$region);
8287

8388
let parser = [{ return parseParallelOp(parser, result); }];
8489
let printer = [{ return printParallelOp(p, *this); }];
90+
let verifier = [{ return ::verifyParallelOp(*this); }];
8591
}
8692

8793
def TerminatorOp : OpenMP_Op<"terminator", [Terminator]> {

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

Lines changed: 92 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ void OpenMPDialect::initialize() {
4141
///
4242
/// operand-and-type-list ::= `(` ssa-id-and-type-list `)`
4343
/// ssa-id-and-type-list ::= ssa-id-and-type |
44-
/// ssa-id-and-type ',' ssa-id-and-type-list
44+
/// ssa-id-and-type `,` ssa-id-and-type-list
4545
/// ssa-id-and-type ::= ssa-id `:` type
4646
static ParseResult
4747
parseOperandAndTypeList(OpAsmParser &parser,
@@ -65,6 +65,52 @@ parseOperandAndTypeList(OpAsmParser &parser,
6565
return success();
6666
}
6767

68+
/// Parse an allocate clause with allocators and a list of operands with types.
69+
///
70+
/// operand-and-type-list ::= `(` allocate-operand-list `)`
71+
/// allocate-operand-list :: = allocate-operand |
72+
/// allocator-operand `,` allocate-operand-list
73+
/// allocate-operand :: = ssa-id-and-type -> ssa-id-and-type
74+
/// ssa-id-and-type ::= ssa-id `:` type
75+
static ParseResult parseAllocateAndAllocator(
76+
OpAsmParser &parser,
77+
SmallVectorImpl<OpAsmParser::OperandType> &operandsAllocate,
78+
SmallVectorImpl<Type> &typesAllocate,
79+
SmallVectorImpl<OpAsmParser::OperandType> &operandsAllocator,
80+
SmallVectorImpl<Type> &typesAllocator) {
81+
if (parser.parseLParen())
82+
return failure();
83+
84+
do {
85+
OpAsmParser::OperandType operand;
86+
Type type;
87+
88+
if (parser.parseOperand(operand) || parser.parseColonType(type))
89+
return failure();
90+
operandsAllocator.push_back(operand);
91+
typesAllocator.push_back(type);
92+
if (parser.parseArrow())
93+
return failure();
94+
if (parser.parseOperand(operand) || parser.parseColonType(type))
95+
return failure();
96+
97+
operandsAllocate.push_back(operand);
98+
typesAllocate.push_back(type);
99+
} while (succeeded(parser.parseOptionalComma()));
100+
101+
if (parser.parseRParen())
102+
return failure();
103+
104+
return success();
105+
}
106+
107+
static LogicalResult verifyParallelOp(ParallelOp op) {
108+
if (op.allocate_vars().size() != op.allocators_vars().size())
109+
return op.emitError(
110+
"expected equal sizes for allocate and allocator variables");
111+
return success();
112+
}
113+
68114
static void printParallelOp(OpAsmPrinter &p, ParallelOp op) {
69115
p << "omp.parallel";
70116

@@ -84,10 +130,26 @@ static void printParallelOp(OpAsmPrinter &p, ParallelOp op) {
84130
}
85131
}
86132
};
133+
134+
// Print allocator and allocate parameters
135+
auto printAllocateAndAllocator = [&p](OperandRange varsAllocate,
136+
OperandRange varsAllocator) {
137+
if (varsAllocate.empty())
138+
return;
139+
140+
p << " allocate(";
141+
for (unsigned i = 0; i < varsAllocate.size(); ++i) {
142+
std::string separator = i == varsAllocate.size() - 1 ? ")" : ", ";
143+
p << varsAllocator[i] << " : " << varsAllocator[i].getType() << " -> ";
144+
p << varsAllocate[i] << " : " << varsAllocate[i].getType() << separator;
145+
}
146+
};
147+
87148
printDataVars("private", op.private_vars());
88149
printDataVars("firstprivate", op.firstprivate_vars());
89150
printDataVars("shared", op.shared_vars());
90151
printDataVars("copyin", op.copyin_vars());
152+
printAllocateAndAllocator(op.allocate_vars(), op.allocators_vars());
91153

92154
if (auto def = op.default_val())
93155
p << " default(" << def->drop_front(3) << ")";
@@ -118,6 +180,7 @@ static ParseResult allowedOnce(OpAsmParser &parser, llvm::StringRef clause,
118180
/// firstprivate ::= `firstprivate` operand-and-type-list
119181
/// shared ::= `shared` operand-and-type-list
120182
/// copyin ::= `copyin` operand-and-type-list
183+
/// allocate ::= `allocate` operand-and-type `->` operand-and-type-list
121184
/// default ::= `default` `(` (`private` | `firstprivate` | `shared` | `none`)
122185
/// procBind ::= `proc_bind` `(` (`master` | `close` | `spread`) `)`
123186
///
@@ -134,7 +197,11 @@ static ParseResult parseParallelOp(OpAsmParser &parser,
134197
SmallVector<Type, 4> sharedTypes;
135198
SmallVector<OpAsmParser::OperandType, 4> copyins;
136199
SmallVector<Type, 4> copyinTypes;
137-
std::array<int, 6> segments{0, 0, 0, 0, 0, 0};
200+
SmallVector<OpAsmParser::OperandType, 4> allocates;
201+
SmallVector<Type, 4> allocateTypes;
202+
SmallVector<OpAsmParser::OperandType, 4> allocators;
203+
SmallVector<Type, 4> allocatorTypes;
204+
std::array<int, 8> segments{0, 0, 0, 0, 0, 0, 0, 0};
138205
llvm::StringRef keyword;
139206
bool defaultVal = false;
140207
bool procBind = false;
@@ -145,6 +212,8 @@ static ParseResult parseParallelOp(OpAsmParser &parser,
145212
const int firstprivateClausePos = 3;
146213
const int sharedClausePos = 4;
147214
const int copyinClausePos = 5;
215+
const int allocateClausePos = 6;
216+
const int allocatorPos = 7;
148217
const llvm::StringRef opName = result.name.getStringRef();
149218

150219
while (succeeded(parser.parseOptionalKeyword(&keyword))) {
@@ -192,6 +261,15 @@ static ParseResult parseParallelOp(OpAsmParser &parser,
192261
if (parseOperandAndTypeList(parser, copyins, copyinTypes))
193262
return failure();
194263
segments[copyinClausePos] = copyins.size();
264+
} else if (keyword == "allocate") {
265+
// fail if there was already another allocate clause
266+
if (segments[allocateClausePos])
267+
return allowedOnce(parser, "allocate", opName);
268+
if (parseAllocateAndAllocator(parser, allocates, allocateTypes,
269+
allocators, allocatorTypes))
270+
return failure();
271+
segments[allocateClausePos] = allocates.size();
272+
segments[allocatorPos] = allocators.size();
195273
} else if (keyword == "default") {
196274
// fail if there was already another default clause
197275
if (defaultVal)
@@ -261,6 +339,18 @@ static ParseResult parseParallelOp(OpAsmParser &parser,
261339
result.operands))
262340
return failure();
263341

342+
// Add allocate parameters
343+
if (segments[allocateClausePos] &&
344+
parser.resolveOperands(allocates, allocateTypes, allocates[0].location,
345+
result.operands))
346+
return failure();
347+
348+
// Add allocator parameters
349+
if (segments[allocatorPos] &&
350+
parser.resolveOperands(allocators, allocatorTypes, allocators[0].location,
351+
result.operands))
352+
return failure();
353+
264354
result.addAttribute("operand_segment_sizes",
265355
parser.getBuilder().getI32VectorAttr(segments));
266356

‎mlir/test/Dialect/OpenMP/ops.mlir

Lines changed: 24 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -52,37 +52,43 @@ func @omp_terminator() -> () {
5252
}
5353

5454
func @omp_parallel(%data_var : memref<i32>, %if_cond : i1, %num_threads : si32) -> () {
55-
// CHECK: omp.parallel if(%{{.*}}) num_threads(%{{.*}} : si32) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) shared(%{{.*}} : memref<i32>) copyin(%{{.*}} : memref<i32>)
56-
"omp.parallel" (%if_cond, %num_threads, %data_var, %data_var, %data_var, %data_var) ({
55+
// CHECK: omp.parallel if(%{{.*}}) num_threads(%{{.*}} : si32) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) shared(%{{.*}} : memref<i32>) copyin(%{{.*}} : memref<i32>) allocate(%{{.*}} : memref<i32> -> %{{.*}} : memref<i32>)
56+
"omp.parallel" (%if_cond, %num_threads, %data_var, %data_var, %data_var, %data_var, %data_var, %data_var) ({
5757

5858
// test without if condition
59-
// CHECK: omp.parallel num_threads(%{{.*}} : si32) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) shared(%{{.*}} : memref<i32>) copyin(%{{.*}} : memref<i32>)
60-
"omp.parallel"(%num_threads, %data_var, %data_var, %data_var, %data_var) ({
59+
// CHECK: omp.parallel num_threads(%{{.*}} : si32) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) shared(%{{.*}} : memref<i32>) copyin(%{{.*}} : memref<i32>) allocate(%{{.*}} : memref<i32> -> %{{.*}} : memref<i32>)
60+
"omp.parallel"(%num_threads, %data_var, %data_var, %data_var, %data_var, %data_var, %data_var) ({
6161
omp.terminator
62-
}) {operand_segment_sizes = dense<[0,1,1,1,1,1]>: vector<6xi32>, default_val = "defshared"} : (si32, memref<i32>, memref<i32>, memref<i32>, memref<i32>) -> ()
62+
}) {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>) -> ()
6363

6464
// CHECK: omp.barrier
6565
omp.barrier
6666

6767
// test without num_threads
68-
// CHECK: omp.parallel if(%{{.*}}) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) shared(%{{.*}} : memref<i32>) copyin(%{{.*}} : memref<i32>)
69-
"omp.parallel"(%if_cond, %data_var, %data_var, %data_var, %data_var) ({
68+
// CHECK: omp.parallel if(%{{.*}}) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) shared(%{{.*}} : memref<i32>) copyin(%{{.*}} : memref<i32>) allocate(%{{.*}} : memref<i32> -> %{{.*}} : memref<i32>)
69+
"omp.parallel"(%if_cond, %data_var, %data_var, %data_var, %data_var, %data_var, %data_var) ({
70+
omp.terminator
71+
}) {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>) -> ()
72+
73+
// test without allocate
74+
// CHECK: omp.parallel if(%{{.*}}) num_threads(%{{.*}} : si32) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) shared(%{{.*}} : memref<i32>) copyin(%{{.*}} : memref<i32>)
75+
"omp.parallel"(%if_cond, %num_threads, %data_var, %data_var, %data_var, %data_var) ({
7076
omp.terminator
71-
}) {operand_segment_sizes = dense<[1,0,1,1,1,1]> : vector<6xi32>} : (i1, memref<i32>, memref<i32>, memref<i32>, memref<i32>) -> ()
77+
}) {operand_segment_sizes = dense<[1,1,1,1,1,1,0,0]> : vector<8xi32>} : (i1, si32, memref<i32>, memref<i32>, memref<i32>, memref<i32>) -> ()
7278

7379
omp.terminator
74-
}) {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>) -> ()
80+
}) {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>) -> ()
7581

7682
// test with multiple parameters for single variadic argument
77-
// CHECK: omp.parallel private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>, %{{.*}} : memref<i32>) shared(%{{.*}} : memref<i32>) copyin(%{{.*}} : memref<i32>)
78-
"omp.parallel" (%data_var, %data_var, %data_var, %data_var, %data_var) ({
83+
// CHECK: omp.parallel private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>, %{{.*}} : memref<i32>) shared(%{{.*}} : memref<i32>) copyin(%{{.*}} : memref<i32>) allocate(%{{.*}} : memref<i32> -> %{{.*}} : memref<i32>)
84+
"omp.parallel" (%data_var, %data_var, %data_var, %data_var, %data_var, %data_var, %data_var) ({
7985
omp.terminator
80-
}) {operand_segment_sizes = dense<[0,0,1,2,1,1]> : vector<6xi32>} : (memref<i32>, memref<i32>, memref<i32>, memref<i32>, memref<i32>) -> ()
86+
}) {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>) -> ()
8187

8288
return
8389
}
8490

85-
func @omp_parallel_pretty(%data_var : memref<i32>, %if_cond : i1, %num_threads : si32) -> () {
91+
func @omp_parallel_pretty(%data_var : memref<i32>, %if_cond : i1, %num_threads : si32, %allocator : si32) -> () {
8692
// CHECK: omp.parallel
8793
omp.parallel {
8894
omp.terminator
@@ -93,6 +99,11 @@ func @omp_parallel_pretty(%data_var : memref<i32>, %if_cond : i1, %num_threads :
9399
omp.terminator
94100
}
95101

102+
// CHECK: omp.parallel allocate(%{{.*}} : memref<i32> -> %{{.*}} : memref<i32>)
103+
omp.parallel allocate(%data_var : memref<i32> -> %data_var : memref<i32>) {
104+
omp.terminator
105+
}
106+
96107
// CHECK: omp.parallel private(%{{.*}} : memref<i32>, %{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>)
97108
omp.parallel private(%data_var : memref<i32>, %data_var : memref<i32>) firstprivate(%data_var : memref<i32>) {
98109
omp.terminator

0 commit comments

Comments
 (0)