[Mlir-commits] [mlir] e71369f - [mlir][OpenMP] Add copyprivate support to omp.single (#80477)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Feb 15 03:29:30 PST 2024
Author: Leandro Lupori
Date: 2024-02-15T08:29:24-03:00
New Revision: e71369fcb039b7a5f555862f49867ee45ec1ccf5
URL: https://github.com/llvm/llvm-project/commit/e71369fcb039b7a5f555862f49867ee45ec1ccf5
DIFF: https://github.com/llvm/llvm-project/commit/e71369fcb039b7a5f555862f49867ee45ec1ccf5.diff
LOG: [mlir][OpenMP] Add copyprivate support to omp.single (#80477)
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
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/invalid.mlir
mlir/test/Dialect/OpenMP/ops.mlir
Removed:
################################################################################
diff --git a/flang/lib/Lower/OpenMP.cpp b/flang/lib/Lower/OpenMP.cpp
index 24f91765cb439b..7519da844eebba 100644
--- a/flang/lib/Lower/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP.cpp
@@ -2665,6 +2665,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, semaCtx, beginClauseList);
@@ -2678,7 +2679,8 @@ genSingleOp(Fortran::lower::AbstractConverter &converter,
OpWithBodyGenInfo(converter, semaCtx, 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 0adf186ae0c7e9..6a799228d74894 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -466,10 +466,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);
@@ -481,6 +487,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 13fc01d58eced5..d1bff4ee70152f 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -573,6 +573,110 @@ 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) {
+ 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
+static LogicalResult
+verifyCopyPrivateVarList(Operation *op, OperandRange copyPrivateVars,
+ std::optional<ArrayAttr> copyPrivateFuncs) {
+ 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 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 =
+ 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>(copyPrivateVarAndFunc).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
//===----------------------------------------------------------------------===//
@@ -1152,7 +1256,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 523a4038b7c32a..7965d6dc284203 100644
--- a/mlir/test/Dialect/OpenMP/invalid.mlir
+++ b/mlir/test/Dialect/OpenMP/invalid.mlir
@@ -1255,7 +1255,63 @@ 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 {{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>) {
+ 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 99ca802089b8d3..d92a554caf77ca 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -1622,6 +1622,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>) {
More information about the Mlir-commits
mailing list