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

[Sergio Afonso]

================
@@ -2313,72 +2396,112 @@ 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;
+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();
 
-  for (auto mapOp : mapVars) {
-    if (!mapOp.getDefiningOp())
-      return emitError(op->getLoc(), "missing map operation");
+  bool to = mapTypeToBool(mapTypeBits, ClauseMapFlags::to);
+  bool from = mapTypeToBool(mapTypeBits, ClauseMapFlags::from);
+  bool del = mapTypeToBool(mapTypeBits, ClauseMapFlags::del);
 
-    if (auto mapInfoOp = mapOp.getDefiningOp<mlir::omp::MapInfoOp>()) {
-      mlir::omp::ClauseMapFlags mapTypeBits = mapInfoOp.getMapType();
+  bool always = mapTypeToBool(mapTypeBits, ClauseMapFlags::always);
+  bool close = mapTypeToBool(mapTypeBits, ClauseMapFlags::close);
+  bool implicit = mapTypeToBool(mapTypeBits, ClauseMapFlags::implicit);
 
-      bool to = mapTypeToBool(mapTypeBits, ClauseMapFlags::to);
-      bool from = mapTypeToBool(mapTypeBits, ClauseMapFlags::from);
-      bool del = mapTypeToBool(mapTypeBits, ClauseMapFlags::del);
+  if ((isa<TargetDataOp>(op) || isa<TargetOp>(op)) && del)
+    return emitError(op->getLoc(),
+                     "to, from, tofrom and alloc map types are permitted");
 
-      bool always = mapTypeToBool(mapTypeBits, ClauseMapFlags::always);
-      bool close = mapTypeToBool(mapTypeBits, ClauseMapFlags::close);
-      bool implicit = mapTypeToBool(mapTypeBits, ClauseMapFlags::implicit);
+  if (isa<TargetEnterDataOp>(op) && (from || del))
+    return emitError(op->getLoc(), "to and alloc map types are permitted");
 
-      if ((isa<TargetDataOp>(op) || isa<TargetOp>(op)) && del)
-        return emitError(op->getLoc(),
-                         "to, from, tofrom 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<TargetEnterDataOp>(op) && (from || del))
-        return emitError(op->getLoc(), "to and alloc 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<TargetExitDataOp>(op) && to)
-        return emitError(op->getLoc(),
-                         "from, release and delete map types are permitted");
+    if (!to && !from) {
+      return emitError(op->getLoc(),
+                       "at least one of to or from map types must be "
+                       "specified, other map types are not 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");
-        }
+    auto updateVar = mapInfoOp.getVarPtr();
 
-        if (!to && !from) {
-          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) || (to && updateFromVars.contains(updateVar)) ||
+        (from && updateToVars.contains(updateVar))) {
+      return emitError(
+          op->getLoc(),
+          "either to or from map types can be specified, not both");
+    }
 
-        auto updateVar = mapInfoOp.getVarPtr();
+    if (always || close || implicit) {
+      return emitError(
+          op->getLoc(),
+          "present, mapper and iterator map type modifiers 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");
-        }
+    to ? updateToVars.insert(updateVar) : updateFromVars.insert(updateVar);
+  }
 
-        if (always || close || implicit) {
-          return emitError(
-              op->getLoc(),
-              "present, mapper and iterator map type modifiers are permitted");
-        }
+  return success();
+}
 
-        to ? updateToVars.insert(updateVar) : updateFromVars.insert(updateVar);
-      }
+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;
+
+  for (auto mapOp : mapVars) {
+    if (!mapOp.getDefiningOp())
+      return emitError(op->getLoc(), "missing map operation");
+
+    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 &region = iterOp.getRegion();
+    if (region.empty())
+      continue;
+
+    auto yieldOp = dyn_cast<mlir::omp::YieldOp>(region.front().getTerminator());
+    if (!yieldOp || yieldOp.getNumOperands() == 0)
+      continue;
----------------
skatrak wrote:

I see the `omp.iterator` op is marked as `SingleBlockImplicitTerminator`. What would be the use of an iterator (as a modifier for map or for any other clause) that returns nothing? It seems like it shouldn't have an implicit terminator, and that we should make sure it yields something. Let me know if I'm missing something.

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