Skip to content

Commit 41f60a3

Browse files
committed
[mlir][OpenMP] Add copyprivate support to omp.single
This adds a new custom CopyPrivateVarList to the single operation. Each list item is formed by a reference to the variable to be updated, its type and the function to be used to perform the copy. It will be translated to LLVM IR using OpenMP builder, that will use the information in the copyprivate list to call __kmpc_copyprivate. This is patch 2 of 4, to add support for COPYPRIVATE in Flang. Original PR: llvm#73128
1 parent 5ac2320 commit 41f60a3

File tree

5 files changed

+180
-3
lines changed

5 files changed

+180
-3
lines changed

flang/lib/Lower/OpenMP.cpp

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2590,6 +2590,7 @@ genSingleOp(Fortran::lower::AbstractConverter &converter,
25902590
const Fortran::parser::OmpClauseList &beginClauseList,
25912591
const Fortran::parser::OmpClauseList &endClauseList) {
25922592
llvm::SmallVector<mlir::Value> allocateOperands, allocatorOperands;
2593+
llvm::SmallVector<mlir::Value> copyPrivateVars;
25932594
mlir::UnitAttr nowaitAttr;
25942595

25952596
ClauseProcessor cp(converter, beginClauseList);
@@ -2602,7 +2603,8 @@ genSingleOp(Fortran::lower::AbstractConverter &converter,
26022603
return genOpWithBody<mlir::omp::SingleOp>(
26032604
{converter, eval, genNested, currentLocation,
26042605
/*outerCombined=*/false, &beginClauseList},
2605-
allocateOperands, allocatorOperands, nowaitAttr);
2606+
allocateOperands, allocatorOperands, copyPrivateVars,
2607+
/*copyPrivateFuncs=*/nullptr, nowaitAttr);
26062608
}
26072609

26082610
static mlir::omp::TaskOp

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -378,10 +378,16 @@ def SingleOp : OpenMP_Op<"single", [AttrSizedOperandSegments]> {
378378
master thread), in the context of its implicit task. The other threads
379379
in the team, which do not execute the block, wait at an implicit barrier
380380
at the end of the single construct unless a nowait clause is specified.
381+
382+
If copyprivate variables and functions are specified, then each thread
383+
variable is updated with the variable value of the thread that executed
384+
the single region, using the specified copy functions.
381385
}];
382386

383387
let arguments = (ins Variadic<AnyType>:$allocate_vars,
384388
Variadic<AnyType>:$allocators_vars,
389+
Variadic<OpenMP_PointerLikeType>:$copyprivate_vars,
390+
OptionalAttr<SymbolRefArrayAttr>:$copyprivate_funcs,
385391
UnitAttr:$nowait);
386392

387393
let regions = (region AnyRegion:$region);
@@ -393,6 +399,10 @@ def SingleOp : OpenMP_Op<"single", [AttrSizedOperandSegments]> {
393399
$allocators_vars, type($allocators_vars)
394400
) `)`
395401
|`nowait` $nowait
402+
|`copyprivate` `(`
403+
custom<CopyPrivateVarList>(
404+
$copyprivate_vars, type($copyprivate_vars), $copyprivate_funcs
405+
) `)`
396406
) $region attr-dict
397407
}];
398408
let hasVerifier = 1;

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

Lines changed: 103 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -505,6 +505,107 @@ static LogicalResult verifyReductionVarList(Operation *op,
505505
return success();
506506
}
507507

508+
//===----------------------------------------------------------------------===//
509+
// Parser, printer and verifier for CopyPrivateVarList
510+
//===----------------------------------------------------------------------===//
511+
512+
/// copyprivate-entry-list ::= copyprivate-entry
513+
/// | copyprivate-entry-list `,` copyprivate-entry
514+
/// copyprivate-entry ::= ssa-id `->` symbol-ref `:` type
515+
static ParseResult parseCopyPrivateVarList(
516+
OpAsmParser &parser,
517+
SmallVectorImpl<OpAsmParser::UnresolvedOperand> &operands,
518+
SmallVectorImpl<Type> &types, ArrayAttr &copyPrivateSymbols) {
519+
SmallVector<SymbolRefAttr> copyPrivateFuncsVec;
520+
if (failed(parser.parseCommaSeparatedList([&]() {
521+
if (parser.parseOperand(operands.emplace_back()) ||
522+
parser.parseArrow() ||
523+
parser.parseAttribute(copyPrivateFuncsVec.emplace_back()) ||
524+
parser.parseColonType(types.emplace_back()))
525+
return failure();
526+
return success();
527+
})))
528+
return failure();
529+
SmallVector<Attribute> copyPrivateFuncs(copyPrivateFuncsVec.begin(),
530+
copyPrivateFuncsVec.end());
531+
copyPrivateSymbols = ArrayAttr::get(parser.getContext(), copyPrivateFuncs);
532+
return success();
533+
}
534+
535+
/// Print CopyPrivate clause
536+
static void printCopyPrivateVarList(OpAsmPrinter &p, Operation *op,
537+
OperandRange copyPrivateVars,
538+
TypeRange copyPrivateTypes,
539+
std::optional<ArrayAttr> copyPrivateFuncs) {
540+
assert(copyPrivateFuncs.has_value() || copyPrivateVars.empty());
541+
for (unsigned i = 0, e = copyPrivateVars.size(); i < e; ++i) {
542+
if (i != 0)
543+
p << ", ";
544+
p << copyPrivateVars[i] << " -> " << (*copyPrivateFuncs)[i] << " : "
545+
<< copyPrivateTypes[i];
546+
}
547+
}
548+
549+
/// Verifies CopyPrivate Clause
550+
static LogicalResult
551+
verifyCopyPrivateVarList(Operation *op, OperandRange copyPrivateVars,
552+
std::optional<ArrayAttr> copyPrivateFuncs) {
553+
if (!copyPrivateVars.empty()) {
554+
if (!copyPrivateFuncs || copyPrivateFuncs->size() != copyPrivateVars.size())
555+
return op->emitOpError() << "expected as many copyPrivate functions as "
556+
"copyPrivate variables";
557+
} else {
558+
if (copyPrivateFuncs)
559+
return op->emitOpError() << "unexpected copyPrivate functions";
560+
return success();
561+
}
562+
563+
for (auto args : llvm::zip(copyPrivateVars, *copyPrivateFuncs)) {
564+
auto symbolRef = llvm::cast<SymbolRefAttr>(std::get<1>(args));
565+
std::optional<std::variant<mlir::func::FuncOp, mlir::LLVM::LLVMFuncOp>>
566+
funcOp;
567+
if (mlir::func::FuncOp mlirFuncOp =
568+
SymbolTable::lookupNearestSymbolFrom<mlir::func::FuncOp>(op,
569+
symbolRef))
570+
funcOp = mlirFuncOp;
571+
else if (mlir::LLVM::LLVMFuncOp llvmFuncOp =
572+
SymbolTable::lookupNearestSymbolFrom<mlir::LLVM::LLVMFuncOp>(
573+
op, symbolRef))
574+
funcOp = llvmFuncOp;
575+
576+
auto getNumArguments = [&] {
577+
return std::visit([](auto &f) { return f.getNumArguments(); }, *funcOp);
578+
};
579+
580+
auto getArgumentType = [&](unsigned i) {
581+
return std::visit([i](auto &f) { return f.getArgumentTypes()[i]; },
582+
*funcOp);
583+
};
584+
585+
if (!funcOp)
586+
return op->emitOpError() << "expected symbol reference " << symbolRef
587+
<< " to point to a copy function";
588+
589+
if (getNumArguments() != 2)
590+
return op->emitOpError()
591+
<< "expected copy function " << symbolRef << " to have 2 operands";
592+
593+
Type argTy = getArgumentType(0);
594+
if (argTy != getArgumentType(1))
595+
return op->emitOpError() << "expected copy function " << symbolRef
596+
<< " arguments to have the same type";
597+
598+
Type varType = std::get<0>(args).getType();
599+
if (argTy != varType)
600+
return op->emitOpError()
601+
<< "expected copy function arguments' type (" << argTy
602+
<< ") to be the same as copyprivate variable's type (" << varType
603+
<< ")";
604+
}
605+
606+
return success();
607+
}
608+
508609
//===----------------------------------------------------------------------===//
509610
// Parser, printer and verifier for DependVarList
510611
//===----------------------------------------------------------------------===//
@@ -1072,7 +1173,8 @@ LogicalResult SingleOp::verify() {
10721173
return emitError(
10731174
"expected equal sizes for allocate and allocator variables");
10741175

1075-
return success();
1176+
return verifyCopyPrivateVarList(*this, getCopyprivateVars(),
1177+
getCopyprivateFuncs());
10761178
}
10771179

10781180
//===----------------------------------------------------------------------===//

mlir/test/Dialect/OpenMP/invalid.mlir

Lines changed: 47 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1284,7 +1284,53 @@ func.func @omp_single(%data_var : memref<i32>) -> () {
12841284
// expected-error @below {{expected equal sizes for allocate and allocator variables}}
12851285
"omp.single" (%data_var) ({
12861286
omp.barrier
1287-
}) {operandSegmentSizes = array<i32: 1,0>} : (memref<i32>) -> ()
1287+
}) {operandSegmentSizes = array<i32: 1,0,0>} : (memref<i32>) -> ()
1288+
return
1289+
}
1290+
1291+
// -----
1292+
1293+
func.func @omp_single_copyprivate(%data_var : memref<i32>) -> () {
1294+
// expected-error @below {{expected symbol reference @copy_func to point to a copy function}}
1295+
omp.single copyprivate(%data_var -> @copy_func : memref<i32>) {
1296+
omp.barrier
1297+
}
1298+
return
1299+
}
1300+
1301+
// -----
1302+
1303+
func.func private @copy_func(memref<i32>)
1304+
1305+
func.func @omp_single_copyprivate(%data_var : memref<i32>) -> () {
1306+
// expected-error @below {{expected copy function @copy_func to have 2 operands}}
1307+
omp.single copyprivate(%data_var -> @copy_func : memref<i32>) {
1308+
omp.barrier
1309+
}
1310+
return
1311+
}
1312+
1313+
// -----
1314+
1315+
func.func private @copy_func(memref<i32>, memref<f32>)
1316+
1317+
func.func @omp_single_copyprivate(%data_var : memref<i32>) -> () {
1318+
// expected-error @below {{expected copy function @copy_func arguments to have the same type}}
1319+
omp.single copyprivate(%data_var -> @copy_func : memref<i32>) {
1320+
omp.barrier
1321+
}
1322+
return
1323+
}
1324+
1325+
// -----
1326+
1327+
func.func private @copy_func(memref<f32>, memref<f32>)
1328+
1329+
func.func @omp_single_copyprivate(%data_var : memref<i32>) -> () {
1330+
// expected-error @below {{expected copy function arguments' type ('memref<f32>') to be the same as copyprivate variable's type ('memref<i32>')}}
1331+
omp.single copyprivate(%data_var -> @copy_func : memref<i32>) {
1332+
omp.barrier
1333+
}
12881334
return
12891335
}
12901336

mlir/test/Dialect/OpenMP/ops.mlir

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1607,6 +1607,23 @@ func.func @omp_single_multiple_blocks() {
16071607
return
16081608
}
16091609

1610+
func.func private @copy_i32(memref<i32>, memref<i32>)
1611+
1612+
// CHECK-LABEL: func @omp_single_copyprivate
1613+
func.func @omp_single_copyprivate(%data_var: memref<i32>) {
1614+
omp.parallel {
1615+
// CHECK: omp.single copyprivate(%{{.*}} -> @copy_i32 : memref<i32>) {
1616+
omp.single copyprivate(%data_var -> @copy_i32 : memref<i32>) {
1617+
"test.payload"() : () -> ()
1618+
// CHECK: omp.terminator
1619+
omp.terminator
1620+
}
1621+
// CHECK: omp.terminator
1622+
omp.terminator
1623+
}
1624+
return
1625+
}
1626+
16101627
// CHECK-LABEL: @omp_task
16111628
// CHECK-SAME: (%[[bool_var:.*]]: i1, %[[i64_var:.*]]: i64, %[[i32_var:.*]]: i32, %[[data_var:.*]]: memref<i32>)
16121629
func.func @omp_task(%bool_var: i1, %i64_var: i64, %i32_var: i32, %data_var: memref<i32>) {

0 commit comments

Comments
 (0)