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

[via flang-commits]

================
@@ -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};
+      };
+
+      if (mapInfoOp.getMapperId()) {
+        auto childMapperOp =
+            SymbolTable::lookupNearestSymbolFrom<omp::DeclareMapperOp>(
+                mapInfoOp, mapInfoOp.getMapperIdAttr());
+        llvm::Expected<llvm::Function *> childMapperFn =
+            getOrCreateUserDefinedMapperFunc(
----------------
MattPD wrote:

When `getOrCreateUserDefinedMapperFunc` emits the child mapper (`mi` in the reproducer) for the first time, the child's `genMapInfoCB` calls `builder.restoreIP(codeGenIP)`, setting the insertion point of `builder` inside the child function. The insertion point of `builder` remains inside the child function after `getOrCreateUserDefinedMapperFunc` returns.

If another `map_iterated` entry follows, the `IteratorInfo` constructor for that entry passes `builder` to `calculateCanonicalLoopTripCount`. During that call, `updateToLocation` sets the OpenMPIRBuilder's own insertion point inside the child function too. As a result, translation emits the rest of the outer mapper into the child function.

Save the following as `repro.f90`:

```fortran
subroutine f(x)
  type :: inner
    integer :: q
  end type
  type :: t
    integer :: n
    type(inner) :: items(4)
    integer :: b(8)
  end type
  type(t) :: x
  !$omp declare mapper(mi: inner :: w) map(w)
  !$omp declare mapper(m: t :: v) map(mapper(mi), iterator(i=1:2): v%items(i)) map(iterator(j=1:v%n): v%b(j))
  !$omp target data map(mapper(m): x)
  !$omp end target data
end
```

Then run `flang -fc1 -emit-llvm -fopenmp -fopenmp-version=52 -fopenmp-targets=amdgcn-amd-amdhsa repro.f90 -o -`. Flang built with assertions prints `Basic Block in function '.omp_mapper._QQFfm' does not have terminator!` and then fails the assertion `cannot get DomTreeNode of block with different parent`. Before this PR, Flang rejected the same input with `not yet implemented: Support for iterator modifiers is not implemented yet`. If you swap the two `map` clauses, Flang compiles the reproducer.

Could an `llvm::IRBuilderBase::InsertPointGuard` around the `getOrCreateUserDefinedMapperFunc` call restore the insertion point of `builder`? A test with two `map_iterated` entries, the first carrying `mapper(...)`, would exercise child-mapper emission followed by a second iterated entry.

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