[Mlir-commits] [flang] [llvm] [mlir] [flang][mlir][OpenMP] Support iterator modifier in declare mapper map clause (PR #215485)

[llvmlistbot at llvm.org]

================
@@ -8597,9 +8629,73 @@ emitUserDefinedMapper(Operation *op, llvm::IRBuilderBase &builder,
     genMapInfos(builder, moduleTranslation, dl, combinedInfo, mapData,
                 targetDirective);
 
-    // Drop the mapping that is no longer necessary so that the same region
-    // can be processed multiple times.
-    moduleTranslation.forgetMapping(declMapperOp.getRegion());
+    // Components produced by an `iterator` modifier on the mapper's own map
+    // clause: their count is only known at runtime, so they are recorded as
+    // dynamic segments and pushed one at a time (via
+    // __tgt_push_mapper_component) from a generated loop, mirroring how the
+    // depend/affinity clauses already handle iterator-produced entries.
+    for (mlir::Value iterVal : declMapperInfoOp.getMapIterated()) {
+      auto itersOp = iterVal.getDefiningOp<mlir::omp::IteratorOp>();
+      assert(itersOp && "map_iterated value must be defined by omp.iterator");
+      // Lowering a scalar subscript (e.g. involving LEN_TRIM) can introduce
+      // control flow inside the iterator, so its `omp.yield` is not
+      // necessarily in the region's first block.
+      mlir::omp::YieldOp yieldOp = itersOp.getYieldOp();
+      assert(yieldOp && "expected exactly one omp.yield in the iterator body");
+      auto mapInfoOp =
+          yieldOp.getResults()[0].getDefiningOp<mlir::omp::MapInfoOp>();
+      assert(mapInfoOp && "expected an omp.map.info in the map_iterated body");
+
+      llvm::Value *sizeVal =
+          builder.getInt64(dl.getTypeSize(mapInfoOp.getVarPtrType()));
+      llvm::omp::OpenMPOffloadMappingFlags mapType =
+          convertClauseMapFlags(mapInfoOp.getMapType());
+      bool hasAttachPtr =
+          (mapType & llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_ATTACH) !=
+          llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_NONE;
+
+      IteratorInfo iterInfo(itersOp, moduleTranslation, builder);
+      llvm::Value *count = iterInfo.getTotalTrips();
+
+      combinedInfo.DynamicSegments.push_back(
+          MapInfosTy::DynamicSegment{count, {}, mapType, hasAttachPtr});
+      MapInfosTy::DynamicSegment &seg = combinedInfo.DynamicSegments.back();
+      // Capture `iterInfo` by value: it already holds the lb/ub/step/trip
+      // llvm::Value*s looked up while `declMapperOp`'s region mapping was
+      // still valid. GenEntry runs later (from within
+      // OpenMPIRBuilder::createIteratorLoop), after the mapping for that
+      // region has been dropped below via forgetMapping, so re-deriving a
+      // fresh IteratorInfo from `itersOp` at that point would fail to find
+      // the lb/ub/step values in the (already cleared) ModuleTranslation
+      // value map.
+      seg.GenEntry = [&moduleTranslation, itersOp, mapInfoOp, sizeVal,
+                      iterInfo](llvm::IRBuilderBase &b,
+                                llvm::Value *linearIV) mutable
+          -> llvm::Expected<llvm::SmallVector<llvm::Value *, 3>> {
+        if (failed(convertIteratorRegion(
+                linearIV, iterInfo, itersOp.getRegion(), b, moduleTranslation)))
+          return llvm::make_error<PreviouslyReportedError>();
+        // `mapInfoOp` (the omp.yield operand) has no translated value of its
+        // own; its runtime address is its var_ptr operand.
+        llvm::Value *addr =
+            moduleTranslation.lookupValue(mapInfoOp.getVarPtr());
+        moduleTranslation.forgetMapping(itersOp.getRegion());
+        return llvm::SmallVector<llvm::Value *, 3>{addr, addr, sizeVal};
----------------
MattPD wrote:

The `{addr, addr, sizeVal}` result makes the mapper push each iterated element as a separate component, with base and begin both set to the element address and without `MEMBER_OF` bits. For `map(iterator(i=1:2): v%a(i))`, the mapper pushes two independent 4-byte components. For the non-iterator form `map(v%a(1), v%a(2))`, the mapper instead pushes a component spanning both members, followed by two member components carrying `MEMBER_OF`.

Without the spanning component, I would expect libomptarget to allocate each element separately. For a `target` construct mapping `x` through this mapper, `processDataBefore` calls `getTgtPtrBegin(&x, sizeof(x))` to look up the kernel argument. The lookup succeeds even when the mapping found for `&x` ends before `&x + sizeof(x)`. Device code accessing `x%a(2)` through `x` would then use an address outside the first element's 4-byte allocation. I have not run this on a device.

The grouping question seems separate from the difference between element pointers and base-plus-bounds entries mentioned in https://github.com/llvm/llvm-project/pull/215485#issuecomment-6350645151, because it concerns how the runtime groups components.

Should the mapper group the iterated components under a spanning component with `MEMBER_OF` bits, as it does for the non-iterator members?

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