[flang][MLIR][OpenMP] Add support for `target update` directive. #75047

ergawy · 2023-12-11T13:33:20Z

Add an op in the OMP dialect to model the target update direcive. This change reuses the MapInfoOp used by other device directive to model map clauses but verifies that the restrictions imposed by the target update directive are respected.

llvmbot · 2023-12-11T13:33:51Z

@llvm/pr-subscribers-mlir-llvm
@llvm/pr-subscribers-mlir-openmp

@llvm/pr-subscribers-mlir

Author: Kareem Ergawy (ergawy)

Changes

Summary:

Add an op in the OMP dialect to model the target update direcive. This change reuses the MapInfoOp used by other device directive to model map clauses but verifies that the restrictions imposed by the target update directive are respected.

Full diff: https://github.com/llvm/llvm-project/pull/75047.diff

4 Files Affected:

(modified) mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td (+43)
(modified) mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp (+43)
(modified) mlir/test/Dialect/OpenMP/invalid.mlir (+73)
(modified) mlir/test/Dialect/OpenMP/ops.mlir (+12)

diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 8ff5380f71ad45..68e7ef9be8bfb3 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -1370,6 +1370,49 @@ def Target_ExitDataOp: OpenMP_Op<"target_exit_data",
   let hasVerifier = 1;
 }
 
+//===---------------------------------------------------------------------===//
+// 2.14.6 target update data Construct
+//===---------------------------------------------------------------------===//
+
+def Target_UpdateDataOp: OpenMP_Op<"target_update_data",
+                                                 [AttrSizedOperandSegments]>{
+  let  summary = "target update data construct";
+  let description = [{
+    The target update directive makes the corresponding list items in the device
+    data environment consistent with their original list items, according to the
+    specified motion clauses. The target update construct is a stand-alone
+    directive.
+
+    The optional $if_expr parameter specifies a boolean result of a
+    conditional check. If this value is 1 or is not provided then the target
+    region runs on a device, if it is 0 then the target region is executed
+    on the host device.
+
+    The optional $device parameter specifies the device number for the
+    target region.
+
+    The optional $nowait eliminates the implicit barrier so the parent
+    task can make progress even if the target task is not yet completed.
+
+    TODO: depend clause
+  }];
+
+  let arguments = (ins Optional<I1>:$if_expr,
+                       Optional<AnyInteger>:$device,
+                       UnitAttr:$nowait,
+                       Variadic<AnyType>:$motion_operands);
+
+  let assemblyFormat = [{
+    oilist(`if` `(` $if_expr `:` type($if_expr) `)`
+    | `device` `(` $device `:` type($device) `)`
+    | `nowait` $nowait
+    | `motion_entries` `(` $motion_operands `:` type($motion_operands) `)`
+    ) attr-dict
+   }];
+
+  let hasVerifier = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // 2.14.5 target construct
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 20df0099cbd24d..c952cec1b273c7 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -691,6 +691,7 @@ static ParseResult parseMapClause(OpAsmParser &parser, IntegerAttr &mapType) {
     if (parser.parseKeyword(&mapTypeMod))
       return failure();
 
+    // Map-type-modifiers
     if (mapTypeMod == "always")
       mapTypeBits |= llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_ALWAYS;
 
@@ -703,6 +704,7 @@ static ParseResult parseMapClause(OpAsmParser &parser, IntegerAttr &mapType) {
     if (mapTypeMod == "present")
       mapTypeBits |= llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_PRESENT;
 
+    // Map-types
     if (mapTypeMod == "to")
       mapTypeBits |= llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_TO;
 
@@ -882,6 +884,7 @@ static ParseResult parseCaptureType(OpAsmParser &parser,
 }
 
 static LogicalResult verifyMapClause(Operation *op, OperandRange mapOperands) {
+  bool foundRequiredMapTypes = false;
 
   for (auto mapOp : mapOperands) {
     if (!mapOp.getDefiningOp())
@@ -898,6 +901,7 @@ static LogicalResult verifyMapClause(Operation *op, OperandRange mapOperands) {
 
       uint64_t mapTypeBits = MapInfoOp.getMapType().value();
 
+      // Map-types
       bool to = mapTypeToBitFlag(
           mapTypeBits, llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_TO);
       bool from = mapTypeToBitFlag(
@@ -905,6 +909,14 @@ static LogicalResult verifyMapClause(Operation *op, OperandRange mapOperands) {
       bool del = mapTypeToBitFlag(
           mapTypeBits, llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_DELETE);
 
+      // Map-type-modifiers
+      bool always = mapTypeToBitFlag(
+          mapTypeBits, llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_ALWAYS);
+      bool close = mapTypeToBitFlag(
+          mapTypeBits, llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_CLOSE);
+      bool implicit = mapTypeToBitFlag(
+          mapTypeBits, llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_IMPLICIT);
+
       if ((isa<DataOp>(op) || isa<TargetOp>(op)) && del)
         return emitError(op->getLoc(),
                          "to, from, tofrom and alloc map types are permitted");
@@ -915,11 +927,38 @@ static LogicalResult verifyMapClause(Operation *op, OperandRange mapOperands) {
       if (isa<ExitDataOp>(op) && to)
         return emitError(op->getLoc(),
                          "from, release and delete map types are permitted");
+
+      if (isa<UpdateDataOp>(op)) {
+        if (del) {
+          return emitError(op->getLoc(),
+                           "at least one of to or from map types must be "
+                           "specified, other map types are not permitted.");
+        }
+
+        if (to | from) {
+          foundRequiredMapTypes = true;
+        }
+      }
+
+      // Check UpdateDataOp's valid map-type-modifiers.
+      if (isa<UpdateDataOp>(op) && (always | close | implicit)) {
+        return emitError(
+            op->getLoc(),
+            "present, mapper and iterator map type modifiers are permitted");
+      }
     } else {
       emitError(op->getLoc(), "map argument is not a map entry operation");
     }
   }
 
+  if (isa<UpdateDataOp>(op)) {
+    if (!foundRequiredMapTypes) {
+      return emitError(op->getLoc(),
+                       "At least one of to or from map types must be "
+                       "specified. Other map types are not permitted.");
+    }
+  }
+
   return success();
 }
 
@@ -940,6 +979,10 @@ LogicalResult ExitDataOp::verify() {
   return verifyMapClause(*this, getMapOperands());
 }
 
+LogicalResult UpdateDataOp::verify() {
+  return verifyMapClause(*this, getMotionOperands());
+}
+
 LogicalResult TargetOp::verify() {
   return verifyMapClause(*this, getMapOperands());
 }
diff --git a/mlir/test/Dialect/OpenMP/invalid.mlir b/mlir/test/Dialect/OpenMP/invalid.mlir
index e54808f6cfdee5..af19215010947c 100644
--- a/mlir/test/Dialect/OpenMP/invalid.mlir
+++ b/mlir/test/Dialect/OpenMP/invalid.mlir
@@ -1658,4 +1658,77 @@ func.func @omp_target_exit_data(%map1: memref<?xi32>) {
   return
 }
 
+// -----
+
+func.func @omp_target_update_data_if(%if_cond : i1) {
+  // expected-error @below {{`if` clause can appear at most once in the expansion of the oilist directive}}
+  omp.target_update_data if(%if_cond : i1) if(%if_cond : i1)
+  return
+}
+
+// -----
+
+func.func @omp_target_update_data_device(%device : si32) {
+  // expected-error @below {{`device` clause can appear at most once in the expansion of the oilist directive}}
+  omp.target_update_data device(%device : si32) device(%device : si32)
+  return
+}
+
+// -----
+
+func.func @omp_target_update_data_nowait() {
+  // expected-error @below {{`nowait` clause can appear at most once in the expansion of the oilist directive}}
+  omp.target_update_data nowait nowait
+  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 = ""}
+
+  // expected-error @below {{at least one of to or from map types must be specified, other map types are not permitted.}}
+  omp.target_update_data motion_entries(%mapv : memref<?xi32>)
+  return
+}
+
+// -----
+
+func.func @omp_target_update_invalid_motion_type_2(%map1 : memref<?xi32>) {
+  %mapv = omp.map_info var_ptr(%map1 : memref<?xi32>, tensor<?xi32>) map_clauses(delete) capture(ByRef) -> memref<?xi32> {name = ""}
+
+  // expected-error @below {{at least one of to or from map types must be specified, other map types are not permitted.}}
+  omp.target_update_data motion_entries(%mapv : memref<?xi32>)
+  return
+}
+
+// -----
+
+func.func @omp_target_update_invalid_motion_modifier(%map1 : memref<?xi32>) {
+  %mapv = omp.map_info var_ptr(%map1 : memref<?xi32>, tensor<?xi32>) map_clauses(always, to) capture(ByRef) -> memref<?xi32> {name = ""}
+
+  // expected-error @below {{present, mapper and iterator map type modifiers are permitted}}
+  omp.target_update_data motion_entries(%mapv : memref<?xi32>)
+  return
+}
+
+// -----
+
+func.func @omp_target_update_invalid_motion_modifier_2(%map1 : memref<?xi32>) {
+  %mapv = omp.map_info var_ptr(%map1 : memref<?xi32>, tensor<?xi32>) map_clauses(close, to) capture(ByRef) -> memref<?xi32> {name = ""}
+
+  // expected-error @below {{present, mapper and iterator map type modifiers are permitted}}
+  omp.target_update_data motion_entries(%mapv : memref<?xi32>)
+  return
+}
+
+// -----
+
+func.func @omp_target_update_invalid_motion_modifier_3(%map1 : memref<?xi32>) {
+  %mapv = omp.map_info var_ptr(%map1 : memref<?xi32>, tensor<?xi32>) map_clauses(implicit, to) capture(ByRef) -> memref<?xi32> {name = ""}
+
+  // expected-error @below {{present, mapper and iterator map type modifiers are permitted}}
+  omp.target_update_data motion_entries(%mapv : memref<?xi32>)
+  return
+}
 llvm.mlir.global internal @_QFsubEx() : i32
diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index 4d88d9ac86fe16..b0a6a5ac0a3fb9 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -2082,3 +2082,15 @@ func.func @omp_targets_with_map_bounds(%arg0: !llvm.ptr, %arg1: !llvm.ptr) -> ()
 
     return
 }
+
+// CHECK-LABEL: omp_target_update_data
+func.func @omp_target_update_data (%if_cond : i1, %device : si32, %map1: memref<?xi32>) -> () {
+    %mapv_from = omp.map_info var_ptr(%map1 : memref<?xi32>, tensor<?xi32>) map_clauses(from) capture(ByRef) -> memref<?xi32> {name = ""}
+
+    %mapv_to = omp.map_info var_ptr(%map1 : memref<?xi32>, tensor<?xi32>) map_clauses(present, to) capture(ByRef) -> memref<?xi32> {name = ""}
+
+    // CHECK: omp.target_update_data if(%[[VAL_0:.*]] : i1) device(%[[VAL_1:.*]] : si32) nowait motion_entries(%{{.*}}, %{{.*}} : memref<?xi32>, memref<?xi32>)
+    omp.target_update_data if(%if_cond : i1) device(%device : si32) nowait motion_entries(%mapv_from , %mapv_to : memref<?xi32>, memref<?xi32>)
+    return
+}
+

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

agozillon

LGTM aside from the one Nit that can be ignored if you like.

I'd wait for one more review before landing however!

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

mlir/test/Dialect/OpenMP/invalid.mlir

ergawy · 2023-12-13T11:25:18Z

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

+        if (to & from) {
+          return emitError(
+              op->getLoc(),
+              "either to or from map types can be specified, not both");
+        }


I think we also need verify whether 2 separate map_info ops refer to the same value. I will double check that later today.

When creating the updateOp, it's more likely that you'd create a new mapInfoOp along with it. However, I'm not sure if that should be an enforced rule. Can you explain why you think they should always be separate?

I think @TIFitis is right. You should check that your motion operands comes from MapInfoOp but I think you can have a mix of to and from. Otherwise, how to you lower smth like !$omp target update to(a,b,c) from(x,y,z)

However, I'm not sure if that should be an enforced rule.

I am not trying to treat this as in enforced rule. What I meant is that if you have input where the same value is used in both to and from map_info ops, we should error out. For example, see the following IR:

func.func @omp_target_update_invalid_motion_modifier_4(%map1 : memref<?xi32>) { %mapv = omp.map_info var_ptr(%map1 : memref<?xi32>, tensor<?xi32>) map_clauses(to) capture(ByRef) -> memref<?xi32> {name = ""} %mapv2 = omp.map_info var_ptr(%map1 : memref<?xi32>, tensor<?xi32>) map_clauses(from) capture(ByRef) -> memref<?xi32> {name = ""} // expected-error @below {{either to or from map types can be specified, not both}} omp.target_update_data motion_entries(%mapv, %mapv2 : memref<?xi32>, memref<?xi32>) return }

According to the spec: A list item can only appear in a to or from clause, but not in both. (see).

For the equivalent Fortran code:

subroutine foo integer :: a(1024) integer :: b(1024) logical :: i !$omp target update from(a) to(a, b) nowait end subroutine foo

Here is what flang-new currently produces:

error: Semantic errors in /tmp/test.f90 /tmp/test.f90:17:29: because: 'a' appears in the FROM clause. !$omp target update from(a) to(a, b) nowait ^ /tmp/test.f90:17:35: error: A list item ('a') can only appear in a TO or FROM clause, but not in both. !$omp target update from(a) to(a, b) nowait ^ /tmp/test.f90:17:35: because: 'a' appears in the TO clause. !$omp target update from(a) to(a, b) nowait ^

which is the correct behavior that we still need to reflect in the IR.

Otherwise, how to you lower smth like !$omp target update to(a,b,c) from(x,y,z)

That should be totally fine. Since each map_info op models motion/mapping information for a single value, we will have one map_info op for each one of these variables and we are currently able to handle it fine.

Let me know if I misunderstood either of you :).

I updated the PR to reflect what I wrote above.

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

TIFitis · 2023-12-13T16:46:27Z

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

+                           "specified, other map types are not permitted");
+        }
+
+        if (to & from) {


Did you mean to use && instead of & here?

Regardless, I think we want to check for negation of exclusive or here so it should be if(to == from) then emit error.

This way we can get also get rid of the other if condition below.

Did you mean to use && instead of & here?

My eyes are playing tricks on me, I used & to be consistent with the other checks in the same function; I thought other checks in the function use the same as well. For bools both are equivalent. Anyway, changed it to &&.

Regardless, I think we want to check for negation of exclusive or here so it should be if(to == from) then emit error.

Wouldn't it be better to have separate error messages for each situation? This is the main reason I separated them.

These error messages would usually only be seen by developers working on MLIR and not end users of OpenMP, so usually we don't care about being too descriptive. However, it doesn't hurt, so you can leave them both in there if you want 😁

Summary: Add an op in the OMP dialect to model the `target update` direcive. This change reuses the `MapInfoOp` used by other device directive to model `map` clauses but verifies that the restrictions imposed by the `target update` directive are respected.

TIFitis

LGTM, Thanks for the patch!

…75345) Emits MLIR op corresponding to `!$omp target update` directive. So far, only motion types: `to` and `from` are supported. Motion modifiers: `present`, `mapper`, and `iterator` are not supported yet. This is a follow up to #75047 & #75159, only the last commit is relevant to this PR.

llvmbot added mlir mlir:openmp flang:openmp labels Dec 11, 2023

ergawy requested a review from raghavendhra December 11, 2023 13:33

ergawy requested review from kiranchandramohan, agozillon, TIFitis and clementval December 11, 2023 13:33

agozillon reviewed Dec 11, 2023

View reviewed changes

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td Show resolved Hide resolved

agozillon approved these changes Dec 11, 2023

View reviewed changes

TIFitis reviewed Dec 11, 2023

View reviewed changes

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td Outdated Show resolved Hide resolved

ergawy force-pushed the omp_target_upate branch from 51f4aab to 37dc256 Compare December 12, 2023 04:22

llvmbot added the mlir:llvm label Dec 12, 2023

ergawy force-pushed the omp_target_upate branch from 71ee3df to 37dc256 Compare December 12, 2023 07:13

ergawy mentioned this pull request Dec 12, 2023

[MLIR][OpenMP][Offload] Lower target update op to DeviceRT #75159

Merged

ergawy force-pushed the omp_target_upate branch from 37dc256 to a026752 Compare December 12, 2023 10:19

TIFitis reviewed Dec 12, 2023

View reviewed changes

mlir/test/Dialect/OpenMP/invalid.mlir Outdated Show resolved Hide resolved

ergawy force-pushed the omp_target_upate branch 2 times, most recently from 2a3f2e2 to bb1c8d1 Compare December 13, 2023 04:45

ergawy commented Dec 13, 2023

View reviewed changes

ergawy mentioned this pull request Dec 13, 2023

[flang][MLIR][OpenMP] Emit UpdateDataOp from !$omp target update #75345

Merged

ergawy force-pushed the omp_target_upate branch from bb1c8d1 to 7b20e42 Compare December 13, 2023 14:51

TIFitis reviewed Dec 13, 2023

View reviewed changes

ergawy force-pushed the omp_target_upate branch from 7b20e42 to 19ec0f4 Compare December 14, 2023 04:51

ergawy force-pushed the omp_target_upate branch from 19ec0f4 to 302eaaf Compare December 14, 2023 05:49

TIFitis approved these changes Dec 14, 2023

View reviewed changes

ergawy merged commit 2ab926d into llvm:main Dec 14, 2023
4 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[flang][MLIR][OpenMP] Add support for `target update` directive. #75047

[flang][MLIR][OpenMP] Add support for `target update` directive. #75047

ergawy commented Dec 11, 2023 •

edited

Loading

llvmbot commented Dec 11, 2023 •

edited

Loading

agozillon left a comment

ergawy Dec 13, 2023

TIFitis Dec 13, 2023

clementval Dec 13, 2023

ergawy Dec 14, 2023

ergawy Dec 14, 2023

TIFitis Dec 13, 2023

ergawy Dec 14, 2023

TIFitis Dec 14, 2023 •

edited

Loading

TIFitis left a comment

[flang][MLIR][OpenMP] Add support for target update directive. #75047

[flang][MLIR][OpenMP] Add support for target update directive. #75047

Conversation

ergawy commented Dec 11, 2023 • edited Loading

llvmbot commented Dec 11, 2023 • edited Loading

agozillon left a comment

Choose a reason for hiding this comment

ergawy Dec 13, 2023

Choose a reason for hiding this comment

TIFitis Dec 13, 2023

Choose a reason for hiding this comment

clementval Dec 13, 2023

Choose a reason for hiding this comment

ergawy Dec 14, 2023

Choose a reason for hiding this comment

ergawy Dec 14, 2023

Choose a reason for hiding this comment

TIFitis Dec 13, 2023

Choose a reason for hiding this comment

ergawy Dec 14, 2023

Choose a reason for hiding this comment

TIFitis Dec 14, 2023 • edited Loading

Choose a reason for hiding this comment

TIFitis left a comment

Choose a reason for hiding this comment

[flang][MLIR][OpenMP] Add support for `target update` directive. #75047

[flang][MLIR][OpenMP] Add support for `target update` directive. #75047

ergawy commented Dec 11, 2023 •

edited

Loading

llvmbot commented Dec 11, 2023 •

edited

Loading

TIFitis Dec 14, 2023 •

edited

Loading