[flang-commits] [flang] [mlir] [MLIR][OpenMP] Add `private` clause to `omp.parallel` (PR #81452)
Kareem Ergawy via flang-commits
flang-commits at lists.llvm.org
Tue Feb 13 02:10:14 PST 2024
https://github.com/ergawy updated https://github.com/llvm/llvm-project/pull/81452
>From b3b3eba27063632504b4edcdc1d7fa2c1a7e4455 Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy at amd.com>
Date: Sun, 11 Feb 2024 09:37:59 -0600
Subject: [PATCH] [MLIR][OpenMP] Add `private` clause to `omp.parallel`
Extends the `omp.parallel` op by adding a `private` clause to model
[first]private variables. This uses the `omp.private` op to map
privatized variables to their corresponding privatizers.
Example `omp.private` op with `private` variable:
```
omp.parallel private(@x.privatizer %arg0 -> %arg1 : !llvm.ptr) {
// ... use %arg1 ...
omp.terminator
}
```
Whether the variable is private or firstprivate is determined by the
attributes of the corresponding `omp.private` op.
---
flang/lib/Lower/OpenMP.cpp | 3 +-
mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td | 6 +-
.../Conversion/SCFToOpenMP/SCFToOpenMP.cpp | 4 +-
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp | 149 ++++++++++++++----
mlir/test/Dialect/OpenMP/invalid.mlir | 56 +++++++
mlir/test/Dialect/OpenMP/ops.mlir | 10 +-
mlir/test/Dialect/OpenMP/roundtrip.mlir | 53 +++++++
7 files changed, 237 insertions(+), 44 deletions(-)
diff --git a/flang/lib/Lower/OpenMP.cpp b/flang/lib/Lower/OpenMP.cpp
index e5887620d503b9..791f2c29610205 100644
--- a/flang/lib/Lower/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP.cpp
@@ -2640,7 +2640,8 @@ genParallelOp(Fortran::lower::AbstractConverter &converter,
? nullptr
: mlir::ArrayAttr::get(converter.getFirOpBuilder().getContext(),
reductionDeclSymbols),
- procBindKindAttr);
+ procBindKindAttr, /*private_vars=*/llvm::SmallVector<mlir::Value>{},
+ /*privatizers=*/nullptr);
}
static mlir::omp::SectionOp
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 44f3e5b8dbc361..61ea0b2385b404 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -270,7 +270,9 @@ def ParallelOp : OpenMP_Op<"parallel", [
Variadic<AnyType>:$allocators_vars,
Variadic<OpenMP_PointerLikeType>:$reduction_vars,
OptionalAttr<SymbolRefArrayAttr>:$reductions,
- OptionalAttr<ProcBindKindAttr>:$proc_bind_val);
+ OptionalAttr<ProcBindKindAttr>:$proc_bind_val,
+ Variadic<AnyType>:$private_vars,
+ OptionalAttr<SymbolRefArrayAttr>:$privatizers);
let regions = (region AnyRegion:$region);
@@ -291,7 +293,7 @@ def ParallelOp : OpenMP_Op<"parallel", [
$allocators_vars, type($allocators_vars)
) `)`
| `proc_bind` `(` custom<ClauseAttr>($proc_bind_val) `)`
- ) custom<ParallelRegion>($region, $reduction_vars, type($reduction_vars), $reductions) attr-dict
+ ) custom<ParallelRegion>($region, $reduction_vars, type($reduction_vars), $reductions, $private_vars, type($private_vars), $privatizers) attr-dict
}];
let hasVerifier = 1;
}
diff --git a/mlir/lib/Conversion/SCFToOpenMP/SCFToOpenMP.cpp b/mlir/lib/Conversion/SCFToOpenMP/SCFToOpenMP.cpp
index 2f8b3f7e11de15..4f22272d9aa3d1 100644
--- a/mlir/lib/Conversion/SCFToOpenMP/SCFToOpenMP.cpp
+++ b/mlir/lib/Conversion/SCFToOpenMP/SCFToOpenMP.cpp
@@ -420,7 +420,9 @@ struct ParallelOpLowering : public OpRewritePattern<scf::ParallelOp> {
/* allocators_vars = */ llvm::SmallVector<Value>{},
/* reduction_vars = */ llvm::SmallVector<Value>{},
/* reductions = */ ArrayAttr{},
- /* proc_bind_val = */ omp::ClauseProcBindKindAttr{});
+ /* proc_bind_val = */ omp::ClauseProcBindKindAttr{},
+ /* private_vars = */ ValueRange(),
+ /* privatizers = */ nullptr);
{
OpBuilder::InsertionGuard guard(rewriter);
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index ef08bd87efc93a..252be467273169 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -429,68 +429,98 @@ static void printScheduleClause(OpAsmPrinter &p, Operation *op,
// Parser, printer and verifier for ReductionVarList
//===----------------------------------------------------------------------===//
-ParseResult
-parseReductionClause(OpAsmParser &parser, Region ®ion,
- SmallVectorImpl<OpAsmParser::UnresolvedOperand> &operands,
- SmallVectorImpl<Type> &types, ArrayAttr &reductionSymbols,
- SmallVectorImpl<OpAsmParser::Argument> &privates) {
- if (failed(parser.parseOptionalKeyword("reduction")))
- return failure();
-
+ParseResult parseClauseWithRegionArgs(
+ OpAsmParser &parser, Region ®ion,
+ SmallVectorImpl<OpAsmParser::UnresolvedOperand> &operands,
+ SmallVectorImpl<Type> &types, ArrayAttr &symbols,
+ SmallVectorImpl<OpAsmParser::Argument> ®ionPrivateArgs) {
SmallVector<SymbolRefAttr> reductionVec;
+ unsigned regionArgOffset = regionPrivateArgs.size();
if (failed(
parser.parseCommaSeparatedList(OpAsmParser::Delimiter::Paren, [&]() {
if (parser.parseAttribute(reductionVec.emplace_back()) ||
parser.parseOperand(operands.emplace_back()) ||
parser.parseArrow() ||
- parser.parseArgument(privates.emplace_back()) ||
+ parser.parseArgument(regionPrivateArgs.emplace_back()) ||
parser.parseColonType(types.emplace_back()))
return failure();
return success();
})))
return failure();
- for (auto [prv, type] : llvm::zip_equal(privates, types)) {
+ auto *argsBegin = regionPrivateArgs.begin();
+ MutableArrayRef argsSubrange(argsBegin + regionArgOffset,
+ argsBegin + regionArgOffset + types.size());
+ for (auto [prv, type] : llvm::zip_equal(argsSubrange, types)) {
prv.type = type;
}
SmallVector<Attribute> reductions(reductionVec.begin(), reductionVec.end());
- reductionSymbols = ArrayAttr::get(parser.getContext(), reductions);
+ symbols = ArrayAttr::get(parser.getContext(), reductions);
return success();
}
-static void printReductionClause(OpAsmPrinter &p, Operation *op, Region ®ion,
- ValueRange operands, TypeRange types,
- ArrayAttr reductionSymbols) {
- p << "reduction(";
- llvm::interleaveComma(llvm::zip_equal(reductionSymbols, operands,
- region.front().getArguments(), types),
- p, [&p](auto t) {
- auto [sym, op, arg, type] = t;
- p << sym << " " << op << " -> " << arg << " : "
- << type;
- });
+static void printClauseWithRegionArgs(OpAsmPrinter &p, Operation *op,
+ Region ®ion, StringRef clauseName,
+ ValueRange operands, TypeRange types,
+ ArrayAttr symbols,
+ unsigned regionArgOffset) {
+ p << clauseName << "(";
+
+ auto *argsBegin = region.front().getArguments().begin();
+ MutableArrayRef argsSubrange(argsBegin + regionArgOffset,
+ argsBegin + regionArgOffset + types.size());
+ llvm::interleaveComma(
+ llvm::zip_equal(symbols, operands, argsSubrange, types), p, [&p](auto t) {
+ auto [sym, op, arg, type] = t;
+ p << sym << " " << op << " -> " << arg << " : " << type;
+ });
p << ") ";
}
-static ParseResult
-parseParallelRegion(OpAsmParser &parser, Region ®ion,
- SmallVectorImpl<OpAsmParser::UnresolvedOperand> &operands,
- SmallVectorImpl<Type> &types, ArrayAttr &reductionSymbols) {
+static ParseResult parseParallelRegion(
+ OpAsmParser &parser, Region ®ion,
+ SmallVectorImpl<OpAsmParser::UnresolvedOperand> &reductionVarOperands,
+ SmallVectorImpl<Type> &reductionVarTypes, ArrayAttr &reductionSymbols,
+ llvm::SmallVectorImpl<OpAsmParser::UnresolvedOperand> &privateVarOperands,
+ llvm::SmallVectorImpl<Type> &privateVarsTypes,
+ ArrayAttr &privatizerSymbols) {
+ llvm::SmallVector<OpAsmParser::Argument> regionPrivateArgs;
+
+ if (succeeded(parser.parseOptionalKeyword("reduction"))) {
+ if (failed(parseClauseWithRegionArgs(parser, region, reductionVarOperands,
+ reductionVarTypes, reductionSymbols,
+ regionPrivateArgs)))
+ return failure();
+ }
- llvm::SmallVector<OpAsmParser::Argument> privates;
- if (succeeded(parseReductionClause(parser, region, operands, types,
- reductionSymbols, privates)))
- return parser.parseRegion(region, privates);
+ if (succeeded(parser.parseOptionalKeyword("private"))) {
+ if (failed(parseClauseWithRegionArgs(parser, region, privateVarOperands,
+ privateVarsTypes, privatizerSymbols,
+ regionPrivateArgs)))
+ return failure();
+ }
- return parser.parseRegion(region);
+ return parser.parseRegion(region, regionPrivateArgs);
}
static void printParallelRegion(OpAsmPrinter &p, Operation *op, Region ®ion,
- ValueRange operands, TypeRange types,
- ArrayAttr reductionSymbols) {
+ ValueRange reductionVarOperands,
+ TypeRange reductionVarTypes,
+ ArrayAttr reductionSymbols,
+ ValueRange privateVarOperands,
+ TypeRange privateVarTypes,
+ ArrayAttr privatizerSymbols) {
if (reductionSymbols)
- printReductionClause(p, op, region, operands, types, reductionSymbols);
+ printClauseWithRegionArgs(p, op, region, "reduction", reductionVarOperands,
+ reductionVarTypes, reductionSymbols,
+ /*regionArgOffset=*/0);
+
+ if (privatizerSymbols)
+ printClauseWithRegionArgs(p, op, region, "private", privateVarOperands,
+ privateVarTypes, privatizerSymbols,
+ reductionVarOperands.size());
+
p.printRegion(region, /*printEntryBlockArgs=*/false);
}
@@ -1057,14 +1087,63 @@ void ParallelOp::build(OpBuilder &builder, OperationState &state,
builder, state, /*if_expr_var=*/nullptr, /*num_threads_var=*/nullptr,
/*allocate_vars=*/ValueRange(), /*allocators_vars=*/ValueRange(),
/*reduction_vars=*/ValueRange(), /*reductions=*/nullptr,
- /*proc_bind_val=*/nullptr);
+ /*proc_bind_val=*/nullptr, /*private_vars=*/ValueRange(),
+ /*privatizers=*/nullptr);
state.addAttributes(attributes);
}
+static LogicalResult verifyPrivateVarList(ParallelOp &op) {
+ auto privateVars = op.getPrivateVars();
+ auto privatizers = op.getPrivatizersAttr();
+
+ if (privateVars.empty() && (privatizers == nullptr || privatizers.empty()))
+ return success();
+
+ auto numPrivateVars = privateVars.size();
+ auto numPrivatizers = (privatizers == nullptr) ? 0 : privatizers.size();
+
+ if (numPrivateVars != numPrivatizers)
+ return op.emitError() << "inconsistent number of private variables and "
+ "privatizer op symbols, private vars: "
+ << numPrivateVars
+ << " vs. privatizer op symbols: " << numPrivatizers;
+
+ for (auto privateVarInfo : llvm::zip(privateVars, privatizers)) {
+ Type varType = std::get<0>(privateVarInfo).getType();
+ SymbolRefAttr privatizerSym =
+ std::get<1>(privateVarInfo).cast<SymbolRefAttr>();
+ PrivateClauseOp privatizerOp =
+ SymbolTable::lookupNearestSymbolFrom<PrivateClauseOp>(op,
+ privatizerSym);
+
+ if (privatizerOp == nullptr)
+ return op.emitError() << "failed to lookup privatizer op with symbol: '"
+ << privatizerSym << "'";
+
+ Type privatizerType = privatizerOp.getType();
+
+ if (varType != privatizerType)
+ return op.emitError()
+ << "type mismatch between a "
+ << (privatizerOp.getDataSharingType() ==
+ DataSharingClauseType::Private
+ ? "private"
+ : "firstprivate")
+ << " variable and its privatizer op, var type: " << varType
+ << " vs. privatizer op type: " << privatizerType;
+ }
+
+ return success();
+}
+
LogicalResult ParallelOp::verify() {
if (getAllocateVars().size() != getAllocatorsVars().size())
return emitError(
"expected equal sizes for allocate and allocator variables");
+
+ if (failed(verifyPrivateVarList(*this)))
+ return failure();
+
return verifyReductionVarList(*this, getReductions(), getReductionVars());
}
diff --git a/mlir/test/Dialect/OpenMP/invalid.mlir b/mlir/test/Dialect/OpenMP/invalid.mlir
index 59b42390b206f1..87e5aa6e92a887 100644
--- a/mlir/test/Dialect/OpenMP/invalid.mlir
+++ b/mlir/test/Dialect/OpenMP/invalid.mlir
@@ -1801,3 +1801,59 @@ omp.private {type = firstprivate} @x.privatizer : f32 alloc {
^bb0(%arg0: f32):
omp.yield(%arg0 : f32)
}
+
+// -----
+
+func.func @private_type_mismatch(%arg0: index) {
+// expected-error @below {{type mismatch between a private variable and its privatizer op, var type: 'index' vs. privatizer op type: '!llvm.ptr'}}
+ omp.parallel private(@var1.privatizer %arg0 -> %arg2 : index) {
+ omp.terminator
+ }
+
+ return
+}
+
+omp.private {type = private} @var1.privatizer : !llvm.ptr alloc {
+^bb0(%arg0: !llvm.ptr):
+ omp.yield(%arg0 : !llvm.ptr)
+}
+
+// -----
+
+func.func @firstprivate_type_mismatch(%arg0: index) {
+ // expected-error @below {{type mismatch between a firstprivate variable and its privatizer op, var type: 'index' vs. privatizer op type: '!llvm.ptr'}}
+ omp.parallel private(@var1.privatizer %arg0 -> %arg2 : index) {
+ omp.terminator
+ }
+
+ return
+}
+
+omp.private {type = firstprivate} @var1.privatizer : !llvm.ptr alloc {
+^bb0(%arg0: !llvm.ptr):
+ omp.yield(%arg0 : !llvm.ptr)
+} copy {
+^bb0(%arg0: !llvm.ptr, %arg1: !llvm.ptr):
+ omp.yield(%arg0 : !llvm.ptr)
+}
+
+// -----
+
+func.func @undefined_privatizer(%arg0: index) {
+ // expected-error @below {{failed to lookup privatizer op with symbol: '@var1.privatizer'}}
+ omp.parallel private(@var1.privatizer %arg0 -> %arg2 : index) {
+ omp.terminator
+ }
+
+ return
+}
+
+// -----
+func.func @undefined_privatizer(%arg0: !llvm.ptr) {
+ // expected-error @below {{inconsistent number of private variables and privatizer op symbols, private vars: 1 vs. privatizer op symbols: 2}}
+ "omp.parallel"(%arg0) <{operandSegmentSizes = array<i32: 0, 0, 0, 0, 0, 1>, privatizers = [@x.privatizer, @y.privatizer]}> ({
+ ^bb0(%arg2: !llvm.ptr):
+ omp.terminator
+ }) : (!llvm.ptr) -> ()
+ return
+}
diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index 651405964c0675..0bfe0e4f51059b 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -59,7 +59,7 @@ func.func @omp_parallel(%data_var : memref<i32>, %if_cond : i1, %num_threads : i
// CHECK: omp.parallel num_threads(%{{.*}} : i32) allocate(%{{.*}} : memref<i32> -> %{{.*}} : memref<i32>)
"omp.parallel"(%num_threads, %data_var, %data_var) ({
omp.terminator
- }) {operandSegmentSizes = array<i32: 0,1,1,1,0>} : (i32, memref<i32>, memref<i32>) -> ()
+ }) {operandSegmentSizes = array<i32: 0,1,1,1,0,0>} : (i32, memref<i32>, memref<i32>) -> ()
// CHECK: omp.barrier
omp.barrier
@@ -68,22 +68,22 @@ func.func @omp_parallel(%data_var : memref<i32>, %if_cond : i1, %num_threads : i
// CHECK: omp.parallel if(%{{.*}}) allocate(%{{.*}} : memref<i32> -> %{{.*}} : memref<i32>)
"omp.parallel"(%if_cond, %data_var, %data_var) ({
omp.terminator
- }) {operandSegmentSizes = array<i32: 1,0,1,1,0>} : (i1, memref<i32>, memref<i32>) -> ()
+ }) {operandSegmentSizes = array<i32: 1,0,1,1,0,0>} : (i1, memref<i32>, memref<i32>) -> ()
// test without allocate
// CHECK: omp.parallel if(%{{.*}}) num_threads(%{{.*}} : i32)
"omp.parallel"(%if_cond, %num_threads) ({
omp.terminator
- }) {operandSegmentSizes = array<i32: 1,1,0,0,0>} : (i1, i32) -> ()
+ }) {operandSegmentSizes = array<i32: 1,1,0,0,0,0>} : (i1, i32) -> ()
omp.terminator
- }) {operandSegmentSizes = array<i32: 1,1,1,1,0>, proc_bind_val = #omp<procbindkind spread>} : (i1, i32, memref<i32>, memref<i32>) -> ()
+ }) {operandSegmentSizes = array<i32: 1,1,1,1,0,0>, proc_bind_val = #omp<procbindkind spread>} : (i1, i32, memref<i32>, memref<i32>) -> ()
// test with multiple parameters for single variadic argument
// CHECK: omp.parallel allocate(%{{.*}} : memref<i32> -> %{{.*}} : memref<i32>)
"omp.parallel" (%data_var, %data_var) ({
omp.terminator
- }) {operandSegmentSizes = array<i32: 0,0,1,1,0>} : (memref<i32>, memref<i32>) -> ()
+ }) {operandSegmentSizes = array<i32: 0,0,1,1,0,0>} : (memref<i32>, memref<i32>) -> ()
return
}
diff --git a/mlir/test/Dialect/OpenMP/roundtrip.mlir b/mlir/test/Dialect/OpenMP/roundtrip.mlir
index 2553442638ee84..9340e0b5fdf971 100644
--- a/mlir/test/Dialect/OpenMP/roundtrip.mlir
+++ b/mlir/test/Dialect/OpenMP/roundtrip.mlir
@@ -1,5 +1,16 @@
// RUN: mlir-opt -verify-diagnostics %s | mlir-opt | FileCheck %s
+// CHECK-LABEL: parallel_op_privatizers
+func.func @parallel_op_privatizers(%arg0: !llvm.ptr, %arg1: !llvm.ptr) {
+ // CHECK: omp.parallel private(@x.privatizer %arg0 -> %arg2 : !llvm.ptr, @y.privatizer %arg1 -> %arg3 : !llvm.ptr)
+ omp.parallel private(@x.privatizer %arg0 -> %arg2 : !llvm.ptr, @y.privatizer %arg1 -> %arg3 : !llvm.ptr) {
+ %0 = llvm.load %arg2 : !llvm.ptr -> i32
+ %1 = llvm.load %arg3 : !llvm.ptr -> i32
+ omp.terminator
+ }
+ return
+}
+
// CHECK: omp.private {type = private} @x.privatizer : !llvm.ptr alloc {
omp.private {type = private} @x.privatizer : !llvm.ptr alloc {
// CHECK: ^bb0(%arg0: {{.*}}):
@@ -19,3 +30,45 @@ omp.private {type = firstprivate} @y.privatizer : !llvm.ptr alloc {
omp.yield(%arg0 : !llvm.ptr)
}
+// CHECK-LABEL: parallel_op_reduction_and_private
+func.func @parallel_op_reduction_and_private(%priv_var: !llvm.ptr, %priv_var2: !llvm.ptr, %reduc_var: !llvm.ptr, %reduc_var2: !llvm.ptr) {
+ // CHECK: omp.parallel
+ // CHECK-SAME: reduction(
+ // CHECK-SAME: @add_f32 %[[reduc_var:[0-9a-z]+]] -> %[[reduc_arg:[0-9a-z]+]] : !llvm.ptr,
+ // CHECK-SAME: @add_f32 %[[reduc_var2:[0-9a-z]+]] -> %[[reduc_arg2:[0-9a-z]+]] : !llvm.ptr)
+ //
+ // CHECK-SAME: private(
+ // CHECK-SAME: @x.privatizer %[[priv_var:[0-9a-z]+]] -> %[[priv_arg:[0-9a-z]+]] : !llvm.ptr,
+ // CHECK-SAME: @y.privatizer %[[priv_var2:[0-9a-z]+]] -> %[[priv_arg2:[0-9a-z]+]] : !llvm.ptr)
+ omp.parallel reduction(@add_f32 %reduc_var -> %reduc_arg : !llvm.ptr, @add_f32 %reduc_var2 -> %reduc_arg2 : !llvm.ptr)
+ private(@x.privatizer %priv_var -> %priv_arg : !llvm.ptr, @y.privatizer %priv_var2 -> %priv_arg2 : !llvm.ptr) {
+ // CHECK: llvm.load %[[priv_arg]]
+ %0 = llvm.load %priv_arg : !llvm.ptr -> f32
+ // CHECK: llvm.load %[[priv_arg2]]
+ %1 = llvm.load %priv_arg2 : !llvm.ptr -> f32
+ // CHECK: llvm.load %[[reduc_arg]]
+ %2 = llvm.load %reduc_arg : !llvm.ptr -> f32
+ // CHECK: llvm.load %[[reduc_arg2]]
+ %3 = llvm.load %reduc_arg2 : !llvm.ptr -> f32
+ omp.terminator
+ }
+ return
+}
+
+omp.reduction.declare @add_f32 : f32
+init {
+^bb0(%arg: f32):
+ %0 = arith.constant 0.0 : f32
+ omp.yield (%0 : f32)
+}
+combiner {
+^bb1(%arg0: f32, %arg1: f32):
+ %1 = arith.addf %arg0, %arg1 : f32
+ omp.yield (%1 : f32)
+}
+atomic {
+^bb2(%arg2: !llvm.ptr, %arg3: !llvm.ptr):
+ %2 = llvm.load %arg3 : !llvm.ptr -> f32
+ llvm.atomicrmw fadd %arg2, %2 monotonic : !llvm.ptr, f32
+ omp.yield
+}
More information about the flang-commits
mailing list