[Mlir-commits] [flang] [mlir] [mlir][OpenMP] Add copyprivate support to omp.single (PR #80477)
Leandro Lupori
llvmlistbot at llvm.org
Fri Feb 9 04:47:35 PST 2024
https://github.com/luporl updated https://github.com/llvm/llvm-project/pull/80477
>From 00a1487df65f9e55c4c6f96f75c931c550ed0b65 Mon Sep 17 00:00:00 2001
From: Leandro Lupori <leandro.lupori at linaro.org>
Date: Fri, 2 Feb 2024 15:15:23 -0300
Subject: [PATCH 1/3] [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: https://github.com/llvm/llvm-project/pull/73128
---
flang/lib/Lower/OpenMP.cpp | 4 +-
mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td | 10 ++
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp | 104 +++++++++++++++++-
mlir/test/Dialect/OpenMP/invalid.mlir | 48 +++++++-
mlir/test/Dialect/OpenMP/ops.mlir | 17 +++
5 files changed, 180 insertions(+), 3 deletions(-)
diff --git a/flang/lib/Lower/OpenMP.cpp b/flang/lib/Lower/OpenMP.cpp
index ad4cffc707535f..448eb1a14fa9fd 100644
--- a/flang/lib/Lower/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP.cpp
@@ -2607,6 +2607,7 @@ genSingleOp(Fortran::lower::AbstractConverter &converter,
const Fortran::parser::OmpClauseList &beginClauseList,
const Fortran::parser::OmpClauseList &endClauseList) {
llvm::SmallVector<mlir::Value> allocateOperands, allocatorOperands;
+ llvm::SmallVector<mlir::Value> copyPrivateVars;
mlir::UnitAttr nowaitAttr;
ClauseProcessor cp(converter, beginClauseList);
@@ -2620,7 +2621,8 @@ genSingleOp(Fortran::lower::AbstractConverter &converter,
OpWithBodyGenInfo(converter, currentLocation, eval)
.setGenNested(genNested)
.setClauses(&beginClauseList),
- allocateOperands, allocatorOperands, nowaitAttr);
+ allocateOperands, allocatorOperands, copyPrivateVars,
+ /*copyPrivateFuncs=*/nullptr, nowaitAttr);
}
static mlir::omp::TaskOp
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index ca363505485773..c3985ac4dfe396 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -378,10 +378,16 @@ def SingleOp : OpenMP_Op<"single", [AttrSizedOperandSegments]> {
master thread), in the context of its implicit task. The other threads
in the team, which do not execute the block, wait at an implicit barrier
at the end of the single construct unless a nowait clause is specified.
+
+ If copyprivate variables and functions are specified, then each thread
+ variable is updated with the variable value of the thread that executed
+ the single region, using the specified copy functions.
}];
let arguments = (ins Variadic<AnyType>:$allocate_vars,
Variadic<AnyType>:$allocators_vars,
+ Variadic<OpenMP_PointerLikeType>:$copyprivate_vars,
+ OptionalAttr<SymbolRefArrayAttr>:$copyprivate_funcs,
UnitAttr:$nowait);
let regions = (region AnyRegion:$region);
@@ -393,6 +399,10 @@ def SingleOp : OpenMP_Op<"single", [AttrSizedOperandSegments]> {
$allocators_vars, type($allocators_vars)
) `)`
|`nowait` $nowait
+ |`copyprivate` `(`
+ custom<CopyPrivateVarList>(
+ $copyprivate_vars, type($copyprivate_vars), $copyprivate_funcs
+ ) `)`
) $region attr-dict
}];
let hasVerifier = 1;
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 381f17d0804191..c874a9880ba253 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -505,6 +505,107 @@ static LogicalResult verifyReductionVarList(Operation *op,
return success();
}
+//===----------------------------------------------------------------------===//
+// Parser, printer and verifier for CopyPrivateVarList
+//===----------------------------------------------------------------------===//
+
+/// copyprivate-entry-list ::= copyprivate-entry
+/// | copyprivate-entry-list `,` copyprivate-entry
+/// copyprivate-entry ::= ssa-id `->` symbol-ref `:` type
+static ParseResult parseCopyPrivateVarList(
+ OpAsmParser &parser,
+ SmallVectorImpl<OpAsmParser::UnresolvedOperand> &operands,
+ SmallVectorImpl<Type> &types, ArrayAttr ©PrivateSymbols) {
+ SmallVector<SymbolRefAttr> copyPrivateFuncsVec;
+ if (failed(parser.parseCommaSeparatedList([&]() {
+ if (parser.parseOperand(operands.emplace_back()) ||
+ parser.parseArrow() ||
+ parser.parseAttribute(copyPrivateFuncsVec.emplace_back()) ||
+ parser.parseColonType(types.emplace_back()))
+ return failure();
+ return success();
+ })))
+ return failure();
+ SmallVector<Attribute> copyPrivateFuncs(copyPrivateFuncsVec.begin(),
+ copyPrivateFuncsVec.end());
+ copyPrivateSymbols = ArrayAttr::get(parser.getContext(), copyPrivateFuncs);
+ return success();
+}
+
+/// Print CopyPrivate clause
+static void printCopyPrivateVarList(OpAsmPrinter &p, Operation *op,
+ OperandRange copyPrivateVars,
+ TypeRange copyPrivateTypes,
+ std::optional<ArrayAttr> copyPrivateFuncs) {
+ assert(copyPrivateFuncs.has_value() || copyPrivateVars.empty());
+ for (unsigned i = 0, e = copyPrivateVars.size(); i < e; ++i) {
+ if (i != 0)
+ p << ", ";
+ p << copyPrivateVars[i] << " -> " << (*copyPrivateFuncs)[i] << " : "
+ << copyPrivateTypes[i];
+ }
+}
+
+/// Verifies CopyPrivate Clause
+static LogicalResult
+verifyCopyPrivateVarList(Operation *op, OperandRange copyPrivateVars,
+ std::optional<ArrayAttr> copyPrivateFuncs) {
+ if (!copyPrivateVars.empty()) {
+ if (!copyPrivateFuncs || copyPrivateFuncs->size() != copyPrivateVars.size())
+ return op->emitOpError() << "expected as many copyPrivate functions as "
+ "copyPrivate variables";
+ } else {
+ if (copyPrivateFuncs)
+ return op->emitOpError() << "unexpected copyPrivate functions";
+ return success();
+ }
+
+ for (auto args : llvm::zip(copyPrivateVars, *copyPrivateFuncs)) {
+ auto symbolRef = llvm::cast<SymbolRefAttr>(std::get<1>(args));
+ std::optional<std::variant<mlir::func::FuncOp, mlir::LLVM::LLVMFuncOp>>
+ funcOp;
+ if (mlir::func::FuncOp mlirFuncOp =
+ SymbolTable::lookupNearestSymbolFrom<mlir::func::FuncOp>(op,
+ symbolRef))
+ funcOp = mlirFuncOp;
+ else if (mlir::LLVM::LLVMFuncOp llvmFuncOp =
+ SymbolTable::lookupNearestSymbolFrom<mlir::LLVM::LLVMFuncOp>(
+ op, symbolRef))
+ funcOp = llvmFuncOp;
+
+ auto getNumArguments = [&] {
+ return std::visit([](auto &f) { return f.getNumArguments(); }, *funcOp);
+ };
+
+ auto getArgumentType = [&](unsigned i) {
+ return std::visit([i](auto &f) { return f.getArgumentTypes()[i]; },
+ *funcOp);
+ };
+
+ if (!funcOp)
+ return op->emitOpError() << "expected symbol reference " << symbolRef
+ << " to point to a copy function";
+
+ if (getNumArguments() != 2)
+ return op->emitOpError()
+ << "expected copy function " << symbolRef << " to have 2 operands";
+
+ Type argTy = getArgumentType(0);
+ if (argTy != getArgumentType(1))
+ return op->emitOpError() << "expected copy function " << symbolRef
+ << " arguments to have the same type";
+
+ Type varType = std::get<0>(args).getType();
+ if (argTy != varType)
+ return op->emitOpError()
+ << "expected copy function arguments' type (" << argTy
+ << ") to be the same as copyprivate variable's type (" << varType
+ << ")";
+ }
+
+ return success();
+}
+
//===----------------------------------------------------------------------===//
// Parser, printer and verifier for DependVarList
//===----------------------------------------------------------------------===//
@@ -1072,7 +1173,8 @@ LogicalResult SingleOp::verify() {
return emitError(
"expected equal sizes for allocate and allocator variables");
- return success();
+ return verifyCopyPrivateVarList(*this, getCopyprivateVars(),
+ getCopyprivateFuncs());
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/OpenMP/invalid.mlir b/mlir/test/Dialect/OpenMP/invalid.mlir
index 812b79e35595f0..cf84c8d9a55ec0 100644
--- a/mlir/test/Dialect/OpenMP/invalid.mlir
+++ b/mlir/test/Dialect/OpenMP/invalid.mlir
@@ -1284,7 +1284,53 @@ func.func @omp_single(%data_var : memref<i32>) -> () {
// expected-error @below {{expected equal sizes for allocate and allocator variables}}
"omp.single" (%data_var) ({
omp.barrier
- }) {operandSegmentSizes = array<i32: 1,0>} : (memref<i32>) -> ()
+ }) {operandSegmentSizes = array<i32: 1,0,0>} : (memref<i32>) -> ()
+ return
+}
+
+// -----
+
+func.func @omp_single_copyprivate(%data_var : memref<i32>) -> () {
+ // expected-error @below {{expected symbol reference @copy_func to point to a copy function}}
+ omp.single copyprivate(%data_var -> @copy_func : memref<i32>) {
+ omp.barrier
+ }
+ return
+}
+
+// -----
+
+func.func private @copy_func(memref<i32>)
+
+func.func @omp_single_copyprivate(%data_var : memref<i32>) -> () {
+ // expected-error @below {{expected copy function @copy_func to have 2 operands}}
+ omp.single copyprivate(%data_var -> @copy_func : memref<i32>) {
+ omp.barrier
+ }
+ return
+}
+
+// -----
+
+func.func private @copy_func(memref<i32>, memref<f32>)
+
+func.func @omp_single_copyprivate(%data_var : memref<i32>) -> () {
+ // expected-error @below {{expected copy function @copy_func arguments to have the same type}}
+ omp.single copyprivate(%data_var -> @copy_func : memref<i32>) {
+ omp.barrier
+ }
+ return
+}
+
+// -----
+
+func.func private @copy_func(memref<f32>, memref<f32>)
+
+func.func @omp_single_copyprivate(%data_var : memref<i32>) -> () {
+ // expected-error @below {{expected copy function arguments' type ('memref<f32>') to be the same as copyprivate variable's type ('memref<i32>')}}
+ omp.single copyprivate(%data_var -> @copy_func : memref<i32>) {
+ omp.barrier
+ }
return
}
diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index 65a704d18107b5..172e358a4237d4 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -1607,6 +1607,23 @@ func.func @omp_single_multiple_blocks() {
return
}
+func.func private @copy_i32(memref<i32>, memref<i32>)
+
+// CHECK-LABEL: func @omp_single_copyprivate
+func.func @omp_single_copyprivate(%data_var: memref<i32>) {
+ omp.parallel {
+ // CHECK: omp.single copyprivate(%{{.*}} -> @copy_i32 : memref<i32>) {
+ omp.single copyprivate(%data_var -> @copy_i32 : memref<i32>) {
+ "test.payload"() : () -> ()
+ // CHECK: omp.terminator
+ omp.terminator
+ }
+ // CHECK: omp.terminator
+ omp.terminator
+ }
+ return
+}
+
// CHECK-LABEL: @omp_task
// CHECK-SAME: (%[[bool_var:.*]]: i1, %[[i64_var:.*]]: i64, %[[i32_var:.*]]: i32, %[[data_var:.*]]: memref<i32>)
func.func @omp_task(%bool_var: i1, %i64_var: i64, %i32_var: i32, %data_var: memref<i32>) {
>From e098f3eb0b1817eabf4087eb196db6a6419e2182 Mon Sep 17 00:00:00 2001
From: Leandro Lupori <leandro.lupori at linaro.org>
Date: Tue, 6 Feb 2024 17:12:08 -0300
Subject: [PATCH 2/3] Use llvm::interleaveComma in printCopyPrivateVarList
---
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp | 15 ++++++++-------
1 file changed, 8 insertions(+), 7 deletions(-)
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index c874a9880ba253..525d9100462ca1 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -537,13 +537,14 @@ static void printCopyPrivateVarList(OpAsmPrinter &p, Operation *op,
OperandRange copyPrivateVars,
TypeRange copyPrivateTypes,
std::optional<ArrayAttr> copyPrivateFuncs) {
- assert(copyPrivateFuncs.has_value() || copyPrivateVars.empty());
- for (unsigned i = 0, e = copyPrivateVars.size(); i < e; ++i) {
- if (i != 0)
- p << ", ";
- p << copyPrivateVars[i] << " -> " << (*copyPrivateFuncs)[i] << " : "
- << copyPrivateTypes[i];
- }
+ if (!copyPrivateFuncs.has_value())
+ return;
+ llvm::interleaveComma(
+ llvm::zip(copyPrivateVars, *copyPrivateFuncs, copyPrivateTypes), p,
+ [&](const auto &args) {
+ p << std::get<0>(args) << " -> " << std::get<1>(args) << " : "
+ << std::get<2>(args);
+ });
}
/// Verifies CopyPrivate Clause
>From 365f7799aac88b62529d64a78ce6ebf605db66ce Mon Sep 17 00:00:00 2001
From: Leandro Lupori <leandro.lupori at linaro.org>
Date: Fri, 9 Feb 2024 09:42:38 -0300
Subject: [PATCH 3/3] Address reviewer's comments
---
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp | 24 +++++++++++---------
mlir/test/Dialect/OpenMP/invalid.mlir | 10 ++++++++
2 files changed, 23 insertions(+), 11 deletions(-)
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 525d9100462ca1..87d198a6ac3507 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -551,18 +551,20 @@ static void printCopyPrivateVarList(OpAsmPrinter &p, Operation *op,
static LogicalResult
verifyCopyPrivateVarList(Operation *op, OperandRange copyPrivateVars,
std::optional<ArrayAttr> copyPrivateFuncs) {
- if (!copyPrivateVars.empty()) {
- if (!copyPrivateFuncs || copyPrivateFuncs->size() != copyPrivateVars.size())
- return op->emitOpError() << "expected as many copyPrivate functions as "
- "copyPrivate variables";
- } else {
- if (copyPrivateFuncs)
- return op->emitOpError() << "unexpected copyPrivate functions";
+ size_t copyPrivateFuncsSize =
+ copyPrivateFuncs.has_value() ? copyPrivateFuncs->size() : 0;
+ if (copyPrivateFuncsSize != copyPrivateVars.size())
+ return op->emitOpError() << "inconsistent number of copyPrivate vars (= "
+ << copyPrivateVars.size()
+ << ") and functions (= " << copyPrivateFuncsSize
+ << "), both must be equal";
+ if (!copyPrivateFuncs.has_value())
return success();
- }
- for (auto args : llvm::zip(copyPrivateVars, *copyPrivateFuncs)) {
- auto symbolRef = llvm::cast<SymbolRefAttr>(std::get<1>(args));
+ for (auto copyPrivateVarAndFunc :
+ llvm::zip(copyPrivateVars, *copyPrivateFuncs)) {
+ auto symbolRef =
+ llvm::cast<SymbolRefAttr>(std::get<1>(copyPrivateVarAndFunc));
std::optional<std::variant<mlir::func::FuncOp, mlir::LLVM::LLVMFuncOp>>
funcOp;
if (mlir::func::FuncOp mlirFuncOp =
@@ -596,7 +598,7 @@ verifyCopyPrivateVarList(Operation *op, OperandRange copyPrivateVars,
return op->emitOpError() << "expected copy function " << symbolRef
<< " arguments to have the same type";
- Type varType = std::get<0>(args).getType();
+ Type varType = std::get<0>(copyPrivateVarAndFunc).getType();
if (argTy != varType)
return op->emitOpError()
<< "expected copy function arguments' type (" << argTy
diff --git a/mlir/test/Dialect/OpenMP/invalid.mlir b/mlir/test/Dialect/OpenMP/invalid.mlir
index cf84c8d9a55ec0..8f66af4b623fc8 100644
--- a/mlir/test/Dialect/OpenMP/invalid.mlir
+++ b/mlir/test/Dialect/OpenMP/invalid.mlir
@@ -1290,6 +1290,16 @@ func.func @omp_single(%data_var : memref<i32>) -> () {
// -----
+func.func @omp_single_copyprivate(%data_var : memref<i32>) -> () {
+ // expected-error @below {{inconsistent number of copyPrivate vars (= 1) and functions (= 0), both must be equal}}
+ "omp.single" (%data_var) ({
+ omp.barrier
+ }) {operandSegmentSizes = array<i32: 0,0,1>} : (memref<i32>) -> ()
+ return
+}
+
+// -----
+
func.func @omp_single_copyprivate(%data_var : memref<i32>) -> () {
// expected-error @below {{expected symbol reference @copy_func to point to a copy function}}
omp.single copyprivate(%data_var -> @copy_func : memref<i32>) {
More information about the Mlir-commits
mailing list