[flang-commits] [flang] 096ebfa - [flang][OpenMP] Fix dangling map.info members in workdistribute target split (#225440)
via flang-commits
flang-commits at lists.llvm.org
Wed Oct 14 22:15:26 PDT 2026
Author: Chaitanya
Date: 2026-10-15T10:45:20+05:30
New Revision: 096ebfa432b92831f2d3324629181d23d1776d1a
URL: https://github.com/llvm/llvm-project/commit/096ebfa432b92831f2d3324629181d23d1776d1a
DIFF: https://github.com/llvm/llvm-project/commit/096ebfa432b92831f2d3324629181d23d1776d1a.diff
LOG: [flang][OpenMP] Fix dangling map.info members in workdistribute target split (#225440)
**Summary**:
- Flang `lower-workdistribute` pass splits an `omp.target` that contains
a `workdistribute` into an `omp.target_data { omp.target }` nest,
cloning every `omp.map.info` for the inner target.
- A Fortran allocatable/pointer is mapped as two linked maps: a
descriptor map whose `members` operand points at its data map.
- The old clone copied `members` verbatim, so the inner descriptor clone
still referenced the outer/original data map instead of its inner
sibling.
- **Fix**: Clone with an `IRMapping` so `members` operands are
retargeted to the inner clones, and clone in dependency order.
- Add a regression test exercising an allocatable descriptor/data member
pair.
Assisted by: claude opus 4.8
Added:
flang/test/Transforms/OpenMP/lower-workdistribute-target-map-members.mlir
Modified:
flang/lib/Optimizer/OpenMP/LowerWorkdistribute.cpp
Removed:
################################################################################
diff --git a/flang/lib/Optimizer/OpenMP/LowerWorkdistribute.cpp b/flang/lib/Optimizer/OpenMP/LowerWorkdistribute.cpp
index 25cf4661d7933f..70d4b0efe58f6d 100644
--- a/flang/lib/Optimizer/OpenMP/LowerWorkdistribute.cpp
+++ b/flang/lib/Optimizer/OpenMP/LowerWorkdistribute.cpp
@@ -31,6 +31,7 @@
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "mlir/Transforms/RegionUtils.h"
+#include "llvm/ADT/SmallPtrSet.h"
#include <mlir/Dialect/Arith/IR/Arith.h>
#include <mlir/Dialect/LLVMIR/LLVMTypes.h>
#include <mlir/Dialect/Utils/IndexingUtils.h>
@@ -719,33 +720,60 @@ FailureOr<omp::TargetOp> splitTargetData(omp::TargetOp targetOp,
}
rewriter.setInsertionPoint(targetOp);
- SmallVector<Value> innerMapInfos;
+ // Pre-sized so inner clones stay aligned with the inner target's block args.
+ SmallVector<Value> innerMapInfos(mapInfos.size());
SmallVector<Value> outerMapInfos;
- // Create new mapinfo ops for the inner target region
- for (auto mapInfo : mapInfos) {
- mlir::omp::ClauseMapFlags originalMapType = mapInfo.getMapType();
- auto originalCaptureType = mapInfo.getMapCaptureType();
- mlir::omp::ClauseMapFlags newMapType;
- mlir::omp::VariableCaptureKind newCaptureType;
- // For bycopy, we keep the same map type and capture type
- // For byref, we change the map type to none and keep the capture type
- if (originalCaptureType == mlir::omp::VariableCaptureKind::ByCopy) {
- newMapType = originalMapType;
- newCaptureType = originalCaptureType;
- } else if (originalCaptureType == mlir::omp::VariableCaptureKind::ByRef) {
- newMapType = mlir::omp::ClauseMapFlags::storage;
- newCaptureType = originalCaptureType;
- outerMapInfos.push_back(mapInfo);
- } else {
- emitError(targetOp->getLoc(), "Unhandled case");
- return failure();
+
+ // A mapinfo's "members" may point at another mapinfo in this set (an
+ // allocatable's descriptor map references its data map). Cloning with
+ // outerToInner retargets those operands to the inner clones, so a member must
+ // be cloned before the parent that references it. The pending set plus
+ // fixed-point loop is a topological sort without an explicit graph.
+ mlir::IRMapping outerToInner;
+ llvm::SmallPtrSet<Operation *, 8> pending;
+ for (auto mapInfo : mapInfos)
+ pending.insert(mapInfo);
+ for (bool progress = true; progress;) {
+ progress = false;
+ for (auto mapInfo : mapInfos) {
+ if (!pending.contains(mapInfo))
+ continue;
+ // Defer until every member also being split has been cloned.
+ if (llvm::any_of(mapInfo.getMembers(), [&](Value member) {
+ return pending.contains(member.getDefiningOp());
+ }))
+ continue;
+
+ mlir::omp::ClauseMapFlags originalMapType = mapInfo.getMapType();
+ auto originalCaptureType = mapInfo.getMapCaptureType();
+ mlir::omp::ClauseMapFlags newMapType = originalMapType;
+ // ByRef is split: inner target keeps only storage, outer data region
+ // keeps the original map. ByCopy is unchanged.
+ if (originalCaptureType == mlir::omp::VariableCaptureKind::ByRef) {
+ newMapType = mlir::omp::ClauseMapFlags::storage;
+ outerMapInfos.push_back(mapInfo);
+ } else if (originalCaptureType !=
+ mlir::omp::VariableCaptureKind::ByCopy) {
+ emitError(targetOp->getLoc(), "Unhandled case");
+ return failure();
+ }
+
+ // clone(op, outerToInner) remaps "members" to sibling inner clones.
+ auto innerMapInfo =
+ cast<omp::MapInfoOp>(rewriter.clone(*mapInfo, outerToInner));
+ innerMapInfo.setMapTypeAttr(
+ rewriter.getAttr<omp::ClauseMapFlagsAttr>(newMapType));
+ pending.erase(mapInfo);
+ progress = true;
}
- auto innerMapInfo = cast<omp::MapInfoOp>(rewriter.clone(*mapInfo));
- innerMapInfo.setMapTypeAttr(
- rewriter.getAttr<omp::ClauseMapFlagsAttr>(newMapType));
- innerMapInfo.setMapCaptureType(newCaptureType);
- innerMapInfos.push_back(innerMapInfo.getResult());
}
+ // Loop stalls with maps still pending only if their members form a cycle.
+ assert(pending.empty() &&
+ "cyclic mapinfo members: cannot topologically order clones");
+
+ // map_entries may repeat a map result, resolve every index to its clone.
+ for (auto [idx, mapInfo] : llvm::enumerate(mapInfos))
+ innerMapInfos[idx] = outerToInner.lookup(mapInfo.getResult());
rewriter.setInsertionPoint(targetOp);
auto device = targetOp.getDevice();
diff --git a/flang/test/Transforms/OpenMP/lower-workdistribute-target-map-members.mlir b/flang/test/Transforms/OpenMP/lower-workdistribute-target-map-members.mlir
new file mode 100644
index 00000000000000..21eebd741ccbee
--- /dev/null
+++ b/flang/test/Transforms/OpenMP/lower-workdistribute-target-map-members.mlir
@@ -0,0 +1,47 @@
+// RUN: fir-opt --lower-workdistribute %s | FileCheck %s
+
+// splitTargetData must retarget a cloned map's "members" to the inner clone.
+// A map result repeated in map_entries must resolve both slots to one clone.
+
+module attributes {llvm.target_triple = "amdgcn-amd-amdhsa", omp.is_gpu = true, omp.is_target_device = true} {
+ func.func @map_members_split(%box : !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>, %x : !fir.ref<i32>) {
+ %base_off = fir.box_offset %box base_addr : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>) -> !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>
+ %data_map = omp.map.info var_ptr(%box : !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>, !fir.box<!fir.heap<!fir.array<?xi32>>>) map_clauses(tofrom) capture(ByRef) var_ptr_ptr(%base_off : !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>, !fir.array<?xi32>) name("") -> !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>
+ %desc_map = omp.map.info var_ptr(%box : !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>, !fir.box<!fir.heap<!fir.array<?xi32>>>) map_clauses(to) capture(ByRef) members(%data_map : [0] : !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>) name("arr") -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>
+ %x_map = omp.map.info var_ptr(%x : !fir.ref<i32>, i32) map_clauses(tofrom) capture(ByRef) name("x") -> !fir.ref<i32>
+ omp.target kernel_type(generic) map_entries(%data_map -> %arg_data, %desc_map -> %arg_desc, %x_map -> %arg_x, %x_map -> %arg_x2 : !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>, !fir.ref<i32>, !fir.ref<i32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c9 = arith.constant 9 : index
+ %val = arith.constant 42 : i32
+ omp.teams {
+ omp.workdistribute {
+ fir.do_loop %iv = %c0 to %c9 step %c1 unordered {
+ fir.store %val to %arg_x : !fir.ref<i32>
+ }
+ omp.terminator
+ }
+ omp.terminator
+ }
+ omp.terminator
+ }
+ return
+ }
+}
+
+// CHECK-LABEL: func.func @map_members_split(
+
+// Original maps drive the outer target_data host<->device movement.
+// CHECK: %[[DATA_MAP:.*]] = omp.map.info {{.*}}map_clauses(tofrom) capture(ByRef) var_ptr_ptr({{.*}}) -> !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>
+// CHECK: %[[DESC_MAP:.*]] = omp.map.info {{.*}}map_clauses(to) capture(ByRef) members(%[[DATA_MAP]] : [0] : {{.*}}) -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>
+
+// Inner clones are downgraded to "storage". The descriptor clone must point at
+// the inner DATA clone, not the outer DATA_MAP above.
+// CHECK: %[[DATA_MAP_INNER:.*]] = omp.map.info {{.*}}map_clauses(storage) capture(ByRef) var_ptr_ptr({{.*}}) -> !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>
+// CHECK: %[[DESC_MAP_INNER:.*]] = omp.map.info {{.*}}map_clauses(storage) capture(ByRef) members(%[[DATA_MAP_INNER]] : [0] : {{.*}}) -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>
+
+// A repeated map result is cloned once and both inner slots reuse that clone.
+// CHECK: %[[X_MAP_INNER:.*]] = omp.map.info {{.*}}map_clauses(storage) capture(ByRef) name("x") -> !fir.ref<i32>
+
+// CHECK: omp.target_data map_entries({{.*}}%[[DATA_MAP]]{{.*}}%[[DESC_MAP]]{{.*}}) {
+// CHECK: omp.target {{.*}}map_entries({{.*}}%[[DATA_MAP_INNER]]{{.*}}%[[DESC_MAP_INNER]]{{.*}}%[[X_MAP_INNER]]{{.*}}%[[X_MAP_INNER]]{{.*}}) {
More information about the flang-commits
mailing list