[flang-commits] [flang] [mlir] [mlir][openmp][flang] - MLIR support for the depend clause (omp dialect) in offloading directives (PR #80626)
Pranav Bhandarkar via flang-commits
flang-commits at lists.llvm.org
Sun Feb 4 17:33:52 PST 2024
https://github.com/bhandarkar-pranav created https://github.com/llvm/llvm-project/pull/80626
This patch adds support for the depend clause in a number of OpenMP directives/constructs related to offloading. Specifically, it adds the handling of the depend clause when it is used with the following constructs
- target
- target enter data
- target update data
- target exit data
Prior to this, one would get an error like the one shown below if one were to use the 'depend' clause with the target directive
"not yet implemented: Unhandled clause DEPEND in TARGET construct"
Note, this patch adds support only for lowering from PFT to MLIR. A subsequent PR will handle conversion from MLIR to LLVMIR
>From 66b500bbf137a6d1d7bde0087d77444725844e73 Mon Sep 17 00:00:00 2001
From: Pranav Bhandarkar <pranav.bhandarkar at amd.com>
Date: Mon, 8 Jan 2024 15:46:01 -0600
Subject: [PATCH] [mlir][openmp][flang] - MLIR support for the depend clause
(omp dialect) in offloading directives
This patch adds support for the depend clause in a number of OpenMP directives/constructs
related to offloading. Specifically, it adds the handling of the depend clause when it
is used with the following constructs
- target
- target enter data
- target update data
- target exit data
Prior to this, one would get an error like the one shown below if one were to use
the 'depend' clause with the target directive
"not yet implemented: Unhandled clause DEPEND in TARGET construct"
Note, this patch adds support only for lowering from PFT to MLIR. A subsequent
PR will handle conversion from MLIR to LLVMIR
---
flang/lib/Lower/OpenMP.cpp | 43 ++++++++++---
mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td | 37 +++++++++--
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp | 22 +++++--
mlir/test/Dialect/OpenMP/invalid.mlir | 37 +++++++++++
mlir/test/Dialect/OpenMP/ops.mlir | 61 ++++++++++++++++++-
5 files changed, 179 insertions(+), 21 deletions(-)
diff --git a/flang/lib/Lower/OpenMP.cpp b/flang/lib/Lower/OpenMP.cpp
index be2117efbabc0..83165a4abeef7 100644
--- a/flang/lib/Lower/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP.cpp
@@ -1753,6 +1753,21 @@ bool ClauseProcessor::processDepend(
return findRepeatableClause<ClauseTy::Depend>(
[&](const ClauseTy::Depend *dependClause,
const Fortran::parser::CharBlock &) {
+ Fortran::parser::DumpTree(llvm::errs(), *dependClause);
+ if (std::holds_alternative<Fortran::parser::OmpDependClause::Source>(
+ dependClause->v.u) ||
+ std::holds_alternative<Fortran::parser::OmpDependClause::Sink>(
+ dependClause->v.u)) {
+ fir::emitFatalError(
+ converter.getCurrentLocation(),
+ "Attempt to use Source and Sink in the " +
+ llvm::StringRef(Fortran::parser::ParseTreeDumper::GetNodeName(
+ *dependClause))
+ .upper() +
+ " clause of a target or task directive",
+ false);
+ }
+
const std::list<Fortran::parser::Designator> &depVal =
std::get<std::list<Fortran::parser::Designator>>(
std::get<Fortran::parser::OmpDependClause::InOut>(
@@ -2662,7 +2677,8 @@ genEnterExitUpdateDataOp(Fortran::lower::AbstractConverter &converter,
Fortran::lower::StatementContext stmtCtx;
mlir::Value ifClauseOperand, deviceOperand;
mlir::UnitAttr nowaitAttr;
- llvm::SmallVector<mlir::Value> mapOperands;
+ llvm::SmallVector<mlir::Value> mapOperands, dependOperands;
+ llvm::SmallVector<mlir::Attribute> dependTypeOperands;
Fortran::parser::OmpIfClause::DirectiveNameModifier directiveName;
llvm::omp::Directive directive;
@@ -2697,12 +2713,15 @@ genEnterExitUpdateDataOp(Fortran::lower::AbstractConverter &converter,
cp.processMap(currentLocation, directive, semanticsContext, stmtCtx,
mapOperands);
}
+ cp.processDepend(dependTypeOperands, dependOperands);
- cp.processTODO<Fortran::parser::OmpClause::Depend>(currentLocation,
- directive);
-
- return firOpBuilder.create<OpTy>(currentLocation, ifClauseOperand,
- deviceOperand, nowaitAttr, mapOperands);
+ return firOpBuilder.create<OpTy>(
+ currentLocation, ifClauseOperand, deviceOperand,
+ dependTypeOperands.empty()
+ ? nullptr
+ : mlir::ArrayAttr::get(converter.getFirOpBuilder().getContext(),
+ dependTypeOperands),
+ dependOperands, nowaitAttr, mapOperands);
}
// This functions creates a block for the body of the targetOp's region. It adds
@@ -2867,7 +2886,8 @@ genTargetOp(Fortran::lower::AbstractConverter &converter,
Fortran::lower::StatementContext stmtCtx;
mlir::Value ifClauseOperand, deviceOperand, threadLimitOperand;
mlir::UnitAttr nowaitAttr;
- llvm::SmallVector<mlir::Value> mapOperands;
+ llvm::SmallVector<mlir::Attribute> dependTypeOperands;
+ llvm::SmallVector<mlir::Value> mapOperands, dependOperands;
llvm::SmallVector<mlir::Type> mapSymTypes;
llvm::SmallVector<mlir::Location> mapSymLocs;
llvm::SmallVector<const Fortran::semantics::Symbol *> mapSymbols;
@@ -2880,8 +2900,8 @@ genTargetOp(Fortran::lower::AbstractConverter &converter,
cp.processNowait(nowaitAttr);
cp.processMap(currentLocation, directive, semanticsContext, stmtCtx,
mapOperands, &mapSymTypes, &mapSymLocs, &mapSymbols);
+ cp.processDepend(dependTypeOperands, dependOperands);
cp.processTODO<Fortran::parser::OmpClause::Private,
- Fortran::parser::OmpClause::Depend,
Fortran::parser::OmpClause::Firstprivate,
Fortran::parser::OmpClause::IsDevicePtr,
Fortran::parser::OmpClause::HasDeviceAddr,
@@ -2891,7 +2911,6 @@ genTargetOp(Fortran::lower::AbstractConverter &converter,
Fortran::parser::OmpClause::UsesAllocators,
Fortran::parser::OmpClause::Defaultmap>(
currentLocation, llvm::omp::Directive::OMPD_target);
-
// 5.8.1 Implicit Data-Mapping Attribute Rules
// The following code follows the implicit data-mapping rules to map all the
// symbols used inside the region that have not been explicitly mapped using
@@ -2962,7 +2981,11 @@ genTargetOp(Fortran::lower::AbstractConverter &converter,
auto targetOp = converter.getFirOpBuilder().create<mlir::omp::TargetOp>(
currentLocation, ifClauseOperand, deviceOperand, threadLimitOperand,
- nowaitAttr, mapOperands);
+ dependTypeOperands.empty()
+ ? nullptr
+ : mlir::ArrayAttr::get(converter.getFirOpBuilder().getContext(),
+ dependTypeOperands),
+ dependOperands, nowaitAttr, mapOperands);
genBodyOfTargetOp(converter, eval, genNested, targetOp, mapSymTypes,
mapSymLocs, mapSymbols, currentLocation);
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 451828ec4ba77..9626a41db0c5b 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -693,7 +693,7 @@ def ClauseTaskDependInOut : I32EnumAttrCase<"taskdependinout", 2>;
def ClauseTaskDepend : I32EnumAttr<
"ClauseTaskDepend",
- "task depend clause",
+ "depend clause in a target or task construct",
[ClauseTaskDependIn, ClauseTaskDependOut, ClauseTaskDependInOut]> {
let genSpecializedAttr = 0;
let cppNamespace = "::mlir::omp";
@@ -1351,11 +1351,17 @@ def Target_EnterDataOp: OpenMP_Op<"target_enter_data",
The $map_types specifies the types and modifiers for the map clause.
- TODO: depend clause and map_type_modifier values iterator and mapper.
+ The `depends` and `depend_vars` arguments are variadic lists of values
+ that specify the dependencies of this particular target task in relation to
+ other tasks.
+
+ TODO: map_type_modifier values iterator and mapper.
}];
let arguments = (ins Optional<I1>:$if_expr,
Optional<AnyInteger>:$device,
+ OptionalAttr<TaskDependArrayAttr>:$depends,
+ Variadic<OpenMP_PointerLikeType>:$depend_vars,
UnitAttr:$nowait,
Variadic<AnyType>:$map_operands);
@@ -1364,6 +1370,7 @@ def Target_EnterDataOp: OpenMP_Op<"target_enter_data",
| `device` `(` $device `:` type($device) `)`
| `nowait` $nowait
| `map_entries` `(` $map_operands `:` type($map_operands) `)`
+ | `depend` `(` custom<DependVarList>($depend_vars, type($depend_vars), $depends) `)`
) attr-dict
}];
@@ -1397,11 +1404,17 @@ def Target_ExitDataOp: OpenMP_Op<"target_exit_data",
The $map_types specifies the types and modifiers for the map clause.
- TODO: depend clause and map_type_modifier values iterator and mapper.
+ The `depends` and `depend_vars` arguments are variadic lists of values
+ that specify the dependencies of this particular target task in relation to
+ other tasks.
+
+ TODO: map_type_modifier values iterator and mapper.
}];
let arguments = (ins Optional<I1>:$if_expr,
Optional<AnyInteger>:$device,
+ OptionalAttr<TaskDependArrayAttr>:$depends,
+ Variadic<OpenMP_PointerLikeType>:$depend_vars,
UnitAttr:$nowait,
Variadic<AnyType>:$map_operands);
@@ -1410,6 +1423,7 @@ def Target_ExitDataOp: OpenMP_Op<"target_exit_data",
| `device` `(` $device `:` type($device) `)`
| `nowait` $nowait
| `map_entries` `(` $map_operands `:` type($map_operands) `)`
+ | `depend` `(` custom<DependVarList>($depend_vars, type($depend_vars), $depends) `)`
) attr-dict
}];
@@ -1447,11 +1461,16 @@ def Target_UpdateDataOp: OpenMP_Op<"target_update_data",
during verification to make sure the restrictions for target update are
respected.
- TODO: depend clause
+ The `depends` and `depend_vars` arguments are variadic lists of values
+ that specify the dependencies of this particular target task in relation to
+ other tasks.
+
}];
let arguments = (ins Optional<I1>:$if_expr,
Optional<AnyInteger>:$device,
+ OptionalAttr<TaskDependArrayAttr>:$depends,
+ Variadic<OpenMP_PointerLikeType>:$depend_vars,
UnitAttr:$nowait,
Variadic<OpenMP_PointerLikeType>:$motion_operands);
@@ -1460,6 +1479,7 @@ def Target_UpdateDataOp: OpenMP_Op<"target_update_data",
| `device` `(` $device `:` type($device) `)`
| `nowait` $nowait
| `motion_entries` `(` $motion_operands `:` type($motion_operands) `)`
+ | `depend` `(` custom<DependVarList>($depend_vars, type($depend_vars), $depends) `)`
) attr-dict
}];
@@ -1488,13 +1508,19 @@ def TargetOp : OpenMP_Op<"target",[IsolatedFromAbove, OutlineableOpenMPOpInterfa
The optional $nowait elliminates the implicit barrier so the parent task can make progress
even if the target task is not yet completed.
- TODO: is_device_ptr, depend, defaultmap, in_reduction
+ The `depends` and `depend_vars` arguments are variadic lists of values
+ that specify the dependencies of this particular target task in relation to
+ other tasks.
+
+ TODO: is_device_ptr, defaultmap, in_reduction
}];
let arguments = (ins Optional<I1>:$if_expr,
Optional<AnyInteger>:$device,
Optional<AnyInteger>:$thread_limit,
+ OptionalAttr<TaskDependArrayAttr>:$depends,
+ Variadic<OpenMP_PointerLikeType>:$depend_vars,
UnitAttr:$nowait,
Variadic<AnyType>:$map_operands);
@@ -1506,6 +1532,7 @@ def TargetOp : OpenMP_Op<"target",[IsolatedFromAbove, OutlineableOpenMPOpInterfa
| `thread_limit` `(` $thread_limit `:` type($thread_limit) `)`
| `nowait` $nowait
| `map_entries` `(` custom<MapEntries>($map_operands, type($map_operands)) `)`
+ | `depend` `(` custom<DependVarList>($depend_vars, type($depend_vars), $depends) `)`
) $region attr-dict
}];
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 13cc16125a273..2ee9f11251291 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -561,7 +561,7 @@ static LogicalResult verifyDependVarList(Operation *op,
return op->emitOpError() << "expected as many depend values"
" as depend variables";
} else {
- if (depends)
+ if (depends && !depends->empty())
return op->emitOpError() << "unexpected depend values";
return success();
}
@@ -965,19 +965,31 @@ LogicalResult DataOp::verify() {
}
LogicalResult EnterDataOp::verify() {
- return verifyMapClause(*this, getMapOperands());
+ LogicalResult verifyDependVars =
+ verifyDependVarList(*this, getDepends(), getDependVars());
+ return failed(verifyDependVars) ? verifyDependVars
+ : verifyMapClause(*this, getMapOperands());
}
LogicalResult ExitDataOp::verify() {
- return verifyMapClause(*this, getMapOperands());
+ LogicalResult verifyDependVars =
+ verifyDependVarList(*this, getDepends(), getDependVars());
+ return failed(verifyDependVars) ? verifyDependVars
+ : verifyMapClause(*this, getMapOperands());
}
LogicalResult UpdateDataOp::verify() {
- return verifyMapClause(*this, getMotionOperands());
+ LogicalResult verifyDependVars =
+ verifyDependVarList(*this, getDepends(), getDependVars());
+ return failed(verifyDependVars) ? verifyDependVars
+ : verifyMapClause(*this, getMotionOperands());
}
LogicalResult TargetOp::verify() {
- return verifyMapClause(*this, getMapOperands());
+ LogicalResult verifyDependVars =
+ verifyDependVarList(*this, getDepends(), getDependVars());
+ return failed(verifyDependVars) ? verifyDependVars
+ : verifyMapClause(*this, getMapOperands());
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/OpenMP/invalid.mlir b/mlir/test/Dialect/OpenMP/invalid.mlir
index 812b79e35595f..481aa950ceeef 100644
--- a/mlir/test/Dialect/OpenMP/invalid.mlir
+++ b/mlir/test/Dialect/OpenMP/invalid.mlir
@@ -1651,6 +1651,15 @@ func.func @omp_target_enter_data(%map1: memref<?xi32>) {
// -----
+func.func @omp_target_enter_data_depend(%a: memref<?xi32>) {
+ %0 = omp.map_info var_ptr(%a: memref<?xi32>, tensor<?xi32>) map_clauses(to) capture(ByRef) -> memref<?xi32>
+ // expected-error @below {{op expected as many depend values as depend variables}}
+ omp.target_enter_data map_entries(%0: memref<?xi32> ) {operandSegmentSizes = array<i32: 0, 0, 1, 0>}
+ return
+}
+
+// -----
+
func.func @omp_target_exit_data(%map1: memref<?xi32>) {
%mapv = omp.map_info var_ptr(%map1 : memref<?xi32>, tensor<?xi32>) map_clauses(to) capture(ByRef) -> memref<?xi32> {name = ""}
// expected-error @below {{from, release and delete map types are permitted}}
@@ -1660,6 +1669,15 @@ func.func @omp_target_exit_data(%map1: memref<?xi32>) {
// -----
+func.func @omp_target_exit_data_depend(%a: memref<?xi32>) {
+ %0 = omp.map_info var_ptr(%a: memref<?xi32>, tensor<?xi32>) map_clauses(from) capture(ByRef) -> memref<?xi32>
+ // expected-error @below {{op expected as many depend values as depend variables}}
+ omp.target_exit_data map_entries(%0: memref<?xi32> ) {operandSegmentSizes = array<i32: 0, 0, 1, 0>}
+ return
+}
+
+// -----
+
func.func @omp_target_update_invalid_motion_type(%map1 : memref<?xi32>) {
%mapv = omp.map_info var_ptr(%map1 : memref<?xi32>, tensor<?xi32>) map_clauses(exit_release_or_enter_alloc) capture(ByRef) -> memref<?xi32> {name = ""}
@@ -1732,6 +1750,25 @@ llvm.mlir.global internal @_QFsubEx() : i32
// -----
+func.func @omp_target_update_data_depend(%a: memref<?xi32>) {
+ %0 = omp.map_info var_ptr(%a: memref<?xi32>, tensor<?xi32>) map_clauses(to) capture(ByRef) -> memref<?xi32>
+ // expected-error @below {{op expected as many depend values as depend variables}}
+ omp.target_update_data motion_entries(%0: memref<?xi32> ) {operandSegmentSizes = array<i32: 0, 0, 1, 0>}
+ return
+}
+
+// -----
+
+func.func @omp_target_depend(%data_var: memref<i32>) {
+ // expected-error @below {{op expected as many depend values as depend variables}}
+ "omp.target"(%data_var) ({
+ "omp.terminator"() : () -> ()
+ }) {depends = [], operandSegmentSizes = array<i32: 0, 0, 0, 1, 0>} : (memref<i32>) -> ()
+ "func.return"() : () -> ()
+}
+
+// -----
+
func.func @omp_distribute(%data_var : memref<i32>) -> () {
// expected-error @below {{expected equal sizes for allocate and allocator variables}}
"omp.distribute"(%data_var) <{operandSegmentSizes = array<i32: 0, 1, 0>}> ({
diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index ccf72ae31d439..89331f2dde7f2 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -517,7 +517,7 @@ func.func @omp_target(%if_cond : i1, %device : si32, %num_threads : i32, %map1:
"omp.target"(%if_cond, %device, %num_threads) ({
// CHECK: omp.terminator
omp.terminator
- }) {nowait, operandSegmentSizes = array<i32: 1,1,1,0>} : ( i1, si32, i32 ) -> ()
+ }) {nowait, operandSegmentSizes = array<i32: 1,1,1,0,0>} : ( i1, si32, i32 ) -> ()
// Test with optional map clause.
// CHECK: %[[MAP_A:.*]] = omp.map_info var_ptr(%[[VAL_1:.*]] : memref<?xi32>, tensor<?xi32>) map_clauses(tofrom) capture(ByRef) -> memref<?xi32> {name = ""}
@@ -1710,6 +1710,18 @@ func.func @omp_task_depend(%arg0: memref<i32>, %arg1: memref<i32>) {
return
}
+
+// CHECK-LABEL: @omp_target_depend
+// CHECK-SAME: (%arg0: memref<i32>, %arg1: memref<i32>) {
+func.func @omp_target_depend(%arg0: memref<i32>, %arg1: memref<i32>) {
+ // CHECK: omp.target depend(taskdependin -> %arg0 : memref<i32>, taskdependin -> %arg1 : memref<i32>, taskdependinout -> %arg0 : memref<i32>) {
+ omp.target depend(taskdependin -> %arg0 : memref<i32>, taskdependin -> %arg1 : memref<i32>, taskdependinout -> %arg0 : memref<i32>) {
+ // CHECK: omp.terminator
+ omp.terminator
+ } {operandSegmentSizes = array<i32: 0,0,0,3,0>}
+ return
+}
+
func.func @omp_threadprivate() {
%0 = arith.constant 1 : i32
%1 = arith.constant 2 : i32
@@ -2124,3 +2136,50 @@ func.func @omp_target_update_data (%if_cond : i1, %device : si32, %map1: memref<
return
}
+// CHECK-LABEL: func @omp_target_enter_update_exit_data_depend
+// CHECK-SAME:([[ARG0:%.*]]: memref<?xi32>, [[ARG1:%.*]]: memref<?xi32>, [[ARG2:%.*]]: memref<?xi32>) {
+func.func @omp_target_enter_update_exit_data_depend(%a: memref<?xi32>, %b: memref<?xi32>, %c: memref<?xi32>) {
+// CHECK-NEXT: [[MAP0:%.*]] = omp.map_info
+// CHECK-NEXT: [[MAP1:%.*]] = omp.map_info
+// CHECK-NEXT: [[MAP2:%.*]] = omp.map_info
+ %map_a = omp.map_info var_ptr(%a: memref<?xi32>, tensor<?xi32>) map_clauses(to) capture(ByRef) -> memref<?xi32>
+ %map_b = omp.map_info var_ptr(%b: memref<?xi32>, tensor<?xi32>) map_clauses(from) capture(ByRef) -> memref<?xi32>
+ %map_c = omp.map_info var_ptr(%c: memref<?xi32>, tensor<?xi32>) map_clauses(exit_release_or_enter_alloc) capture(ByRef) -> memref<?xi32>
+
+ // Do some work on the host that writes to 'a'
+ omp.task depend(taskdependout -> %a : memref<?xi32>) {
+ "test.foo"(%a) : (memref<?xi32>) -> ()
+ omp.terminator
+ }
+
+ // Then map that over to the target
+ // CHECK: omp.target_enter_data nowait map_entries([[MAP0]], [[MAP2]] : memref<?xi32>, memref<?xi32>) depend(taskdependin -> [[ARG0]] : memref<?xi32>)
+ omp.target_enter_data nowait map_entries(%map_a, %map_c: memref<?xi32>, memref<?xi32>) depend(taskdependin -> %a: memref<?xi32>)
+
+ // Compute 'b' on the target and copy it back
+ omp.target map_entries(%map_b -> %arg0 : memref<?xi32>) {
+ ^bb0(%arg0: memref<?xi32>) :
+ "test.foo"(%arg0) : (memref<?xi32>) -> ()
+ omp.terminator
+ }
+
+ // Update 'a' on the host using 'b'
+ omp.task depend(taskdependout -> %a: memref<?xi32>){
+ "test.bar"(%a, %b) : (memref<?xi32>, memref<?xi32>) -> ()
+ }
+
+ // Copy the updated 'a' onto the target
+ // CHECK: omp.target_update_data nowait motion_entries([[MAP0]] : memref<?xi32>) depend(taskdependin -> [[ARG0]] : memref<?xi32>)
+ omp.target_update_data motion_entries(%map_a : memref<?xi32>) depend(taskdependin -> %a : memref<?xi32>) nowait
+
+ // Compute 'c' on the target and copy it back
+ %map_c_from = omp.map_info var_ptr(%c: memref<?xi32>, tensor<?xi32>) map_clauses(from) capture(ByRef) -> memref<?xi32>
+ omp.target map_entries(%map_a -> %arg0, %map_c_from -> %arg1 : memref<?xi32>, memref<?xi32>) depend(taskdependout -> %c : memref<?xi32>) {
+ ^bb0(%arg0 : memref<?xi32>, %arg1 : memref<?xi32>) :
+ "test.foobar"() : ()->()
+ omp.terminator
+ }
+ // CHECK: omp.target_exit_data map_entries([[MAP2]] : memref<?xi32>) depend(taskdependin -> [[ARG2]] : memref<?xi32>)
+ omp.target_exit_data map_entries(%map_c : memref<?xi32>) depend(taskdependin -> %c : memref<?xi32>)
+ return
+}
More information about the flang-commits
mailing list