[flang-commits] [flang] [mlir] [mlir][OpenMP] Add iterator support to map/motion clause (PR #197047)

[Sergio Afonso via flang-commits]

================
@@ -2308,87 +2308,121 @@ static ParseResult parseCaptureType(OpAsmParser &parser,
   return success();
 }
 
-static LogicalResult verifyMapClause(Operation *op, OperandRange mapVars) {
-  llvm::DenseSet<mlir::TypedValue<mlir::omp::PointerLikeType>> updateToVars;
-  llvm::DenseSet<mlir::TypedValue<mlir::omp::PointerLikeType>> updateFromVars;
-
-  for (auto mapOp : mapVars) {
-    if (!mapOp.getDefiningOp())
-      return emitError(op->getLoc(), "missing map operation");
-
-    if (auto mapInfoOp = mapOp.getDefiningOp<mlir::omp::MapInfoOp>()) {
-      mlir::omp::ClauseMapFlags mapTypeBits = mapInfoOp.getMapType();
-
-      bool to = mapTypeToBool(mapTypeBits, ClauseMapFlags::to);
-      bool from = mapTypeToBool(mapTypeBits, ClauseMapFlags::from);
-      bool del = mapTypeToBool(mapTypeBits, ClauseMapFlags::del);
-
-      bool always = mapTypeToBool(mapTypeBits, ClauseMapFlags::always);
-      bool close = mapTypeToBool(mapTypeBits, ClauseMapFlags::close);
-      bool implicit = mapTypeToBool(mapTypeBits, ClauseMapFlags::implicit);
-      bool attach = mapTypeToBool(mapTypeBits, ClauseMapFlags::attach);
+static LogicalResult verifyMapInfoForMapClause(
+    Operation *op, mlir::omp::MapInfoOp mapInfoOp,
+    llvm::DenseSet<mlir::TypedValue<mlir::omp::PointerLikeType>> &updateToVars,
+    llvm::DenseSet<mlir::TypedValue<mlir::omp::PointerLikeType>>
+        &updateFromVars) {
+  mlir::omp::ClauseMapFlags mapTypeBits = mapInfoOp.getMapType();
+
+  bool to = mapTypeToBool(mapTypeBits, ClauseMapFlags::to);
+  bool from = mapTypeToBool(mapTypeBits, ClauseMapFlags::from);
+  bool del = mapTypeToBool(mapTypeBits, ClauseMapFlags::del);
+
+  bool always = mapTypeToBool(mapTypeBits, ClauseMapFlags::always);
+  bool close = mapTypeToBool(mapTypeBits, ClauseMapFlags::close);
+  bool implicit = mapTypeToBool(mapTypeBits, ClauseMapFlags::implicit);
+  bool attach = mapTypeToBool(mapTypeBits, ClauseMapFlags::attach);
+
+  if ((isa<TargetDataOp>(op) || isa<TargetOp>(op)) && del)
+    return emitError(op->getLoc(),
+                     "to, from, tofrom and alloc map types are permitted");
+
+  if (isa<TargetEnterDataOp>(op) && (from || del))
+    return emitError(op->getLoc(), "to and alloc map types are permitted");
+
+  if (isa<TargetExitDataOp>(op) && to)
+    return emitError(op->getLoc(),
+                     "from, release and delete map types are permitted");
+
+  if (isa<TargetUpdateOp>(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 ((isa<TargetDataOp>(op) || isa<TargetOp>(op)) && del)
-        return emitError(op->getLoc(),
-                         "to, from, tofrom and alloc map types are permitted");
+    if (!to && !from && !attach) {
+      return emitError(op->getLoc(),
+                       "at least one of to or from or attach map types must be "
+                       "specified, other map types are not permitted");
+    }
 
-      if (isa<TargetEnterDataOp>(op) && (from || del))
-        return emitError(op->getLoc(), "to and alloc map types are permitted");
+    auto updateVar = mapInfoOp.getVarPtr();
 
-      if (isa<TargetExitDataOp>(op) && to)
-        return emitError(op->getLoc(),
-                         "from, release and delete map types are permitted");
+    if ((to && from) || (to && updateFromVars.contains(updateVar)) ||
+        (from && updateToVars.contains(updateVar))) {
+      return emitError(
+          op->getLoc(),
+          "either to or from map types can be specified, not both");
+    }
 
-      if (isa<TargetUpdateOp>(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 (always || close || implicit) {
+      return emitError(
+          op->getLoc(),
+          "present, mapper and iterator map type modifiers are permitted");
+    }
 
-        if (!to && !from && !attach) {
-          return emitError(
-              op->getLoc(),
-              "at least one of to or from or attach map types must be "
-              "specified, other map types are not permitted");
-        }
+    // It's possible we have an attach map, in which case if there is no to
+    // or from tied to it, we skip insertion.
+    if (to || from) {
+      to ? updateToVars.insert(updateVar) : updateFromVars.insert(updateVar);
+    }
+  }
 
-        auto updateVar = mapInfoOp.getVarPtr();
+  if ((mapInfoOp.getVarPtrPtr() && !mapInfoOp.getVarPtrPtrType()) ||
+      (!mapInfoOp.getVarPtrPtr() && mapInfoOp.getVarPtrPtrType())) {
+    return emitError(op->getLoc(),
+                     "if varPtrPtr or varPtrPtrType is specified, then both "
+                     "must be present");
+  }
 
-        if ((to && from) || (to && updateFromVars.contains(updateVar)) ||
-            (from && updateToVars.contains(updateVar))) {
-          return emitError(
-              op->getLoc(),
-              "either to or from map types can be specified, not both");
-        }
+  return success();
+}
 
-        if (always || close || implicit) {
-          return emitError(
-              op->getLoc(),
-              "present, mapper and iterator map type modifiers are permitted");
-        }
+static LogicalResult verifyMapClause(Operation *op, OperandRange mapVars,
+                                     OperandRange mapIterated) {
+  llvm::DenseSet<mlir::TypedValue<mlir::omp::PointerLikeType>> updateToVars;
+  llvm::DenseSet<mlir::TypedValue<mlir::omp::PointerLikeType>> updateFromVars;
 
-        // It's possible we have an attach map, in which case if there is no to
-        // or from tied to it, we skip insertion.
-        if (to || from) {
-          to ? updateToVars.insert(updateVar)
-             : updateFromVars.insert(updateVar);
-        }
-      }
+  for (auto mapOp : mapVars) {
+    if (!mapOp.getDefiningOp())
+      return emitError(op->getLoc(), "missing map operation");
 
-      if ((mapInfoOp.getVarPtrPtr() && !mapInfoOp.getVarPtrPtrType()) ||
-          (!mapInfoOp.getVarPtrPtr() && mapInfoOp.getVarPtrPtrType())) {
-        return emitError(
-            op->getLoc(),
-            "if varPtrPtr or varPtrPtrType is specified, then both "
-            "must be present");
-      }
+    if (auto mapInfoOp = mapOp.getDefiningOp<mlir::omp::MapInfoOp>()) {
+      if (failed(verifyMapInfoForMapClause(op, mapInfoOp, updateToVars,
+                                           updateFromVars)))
+        return failure();
     } else if (!isa<DeclareMapperInfoOp>(op)) {
       return emitError(op->getLoc(),
                        "map argument is not a map entry operation");
     }
   }
 
+  // Verify iterated map entries.
+  for (auto iterVal : mapIterated) {
+    auto iterOp = iterVal.getDefiningOp<mlir::omp::IteratorOp>();
+    if (!iterOp)
+      return op->emitOpError() << "'map_iterated' arguments must be defined by "
+                                  "'omp.iterator' ops";
+
+    // Check that the iterator body yields a value defined by omp.map.info.
+    auto yieldOp = dyn_cast<mlir::omp::YieldOp>(
+        iterOp.getRegion().front().getTerminator());
+    if (!yieldOp || yieldOp.getResults().empty())
+      continue;
----------------
skatrak wrote:

This is already verified by the `IteratorOp` verifier and traits.

https://github.com/llvm/llvm-project/pull/197047