[flang-commits] [flang] 1ac96b0 - [flang][OpenMP] Recompute descriptor across workdistribute target fission (#225607)
via flang-commits
flang-commits at lists.llvm.org
Thu Oct 15 02:21:42 PDT 2026
Author: Chaitanya
Date: 2026-10-15T14:51:36+05:30
New Revision: 1ac96b020a4ee54f36b4fcfc3ee3482e3ee72728
URL: https://github.com/llvm/llvm-project/commit/1ac96b020a4ee54f36b4fcfc3ee3482e3ee72728
DIFF: https://github.com/llvm/llvm-project/commit/1ac96b020a4ee54f36b4fcfc3ee3482e3ee72728.diff
LOG: [flang][OpenMP] Recompute descriptor across workdistribute target fission (#225607)
Issue:
- When lowering `omp target teams workdistribute`, the target-fission
cached a Fortran allocatable's `fir.box` descriptor by value across the
host/device split, freezing a host base_addr into the kernel (the flat
to/from copy never re-attaches the pointer), so the device indexed a
host address and hit a memory access fault.
Fix :
Mark descriptor (`fir.box`) loads as recomputable in
`isRecomputableAfterFission`, so the descriptor chain is rebuilt inside
each split target from its mapped-and-attached descriptor block arg
(device base_addr) instead of being cached.
This applies only when both of these hold:
- The load reads one of the target's map block arguments directly,
possibly through `fir.declare` or `fir.convert`.
- Nothing writes that descriptor between the load and the split.
Assisted by claude opus 4.8
Added:
flang/test/Lower/OpenMP/workdistribute-saxpy-allocatable.f90
flang/test/Transforms/OpenMP/lower-workdistribute-fission-recompute-box.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 1b20ae4ee3513..97d4f5b278f28 100644
--- a/flang/lib/Optimizer/OpenMP/LowerWorkdistribute.cpp
+++ b/flang/lib/Optimizer/OpenMP/LowerWorkdistribute.cpp
@@ -1012,8 +1012,58 @@ static bool usedOutsideSplit(Value v, Operation *split) {
return false;
}
+/// Returns the map block argument of \p targetOp that \p addr is derived from
+/// through fir.declare and fir.convert, or nullptr otherwise.
+static BlockArgument getMappedBlockArg(Value addr, omp::TargetOp targetOp) {
+ while (Operation *def = addr.getDefiningOp()) {
+ if (auto declare = dyn_cast<fir::DeclareOp>(def))
+ addr = declare.getMemref();
+ else if (auto convert = dyn_cast<fir::ConvertOp>(def))
+ addr = convert.getValue();
+ else
+ return nullptr;
+ }
+ auto argIface = cast<omp::BlockArgOpenMPOpInterface>(*targetOp);
+ auto arg = cast<BlockArgument>(addr);
+ return llvm::is_contained(argIface.getMapBlockArgs(), arg) ? arg : nullptr;
+}
+
+/// Returns true if \p addr, or a fir.declare or fir.convert of it, has a user
+/// other than fir.load that runs after \p load and no later than
+/// \p splitBefore. The split op counts since the post target recomputes after
+/// it runs.
+static bool isWrittenAfterLoad(Value addr, fir::LoadOp load,
+ Operation *splitBefore) {
+ Block *targetBlock = splitBefore->getBlock();
+ for (Operation *user : addr.getUsers()) {
+ if (isa<fir::LoadOp>(user))
+ continue;
+ if (isa<fir::DeclareOp, fir::ConvertOp>(user)) {
+ if (isWrittenAfterLoad(user->getResult(0), load, splitBefore))
+ return true;
+ continue;
+ }
+ Operation *ancestor = targetBlock->findAncestorOpInBlock(*user);
+ if (ancestor && load->isBeforeInBlock(ancestor) &&
+ !splitBefore->isBeforeInBlock(ancestor))
+ return true;
+ }
+ return false;
+}
+
/// isRecomputableAfterFission checks if an operation can be recomputed
static bool isRecomputableAfterFission(Operation *op, Operation *splitBefore) {
+ // A descriptor load must be recomputed from the mapped descriptor in each
+ // split target. Caching the box by value captures a host base_addr that the
+ // flat to/from copy cannot re-attach to the device data. Only safe when the
+ // descriptor is not written between the load and the split op.
+ if (auto load = dyn_cast<fir::LoadOp>(op)) {
+ if (isa<fir::BaseBoxType>(load.getType())) {
+ auto targetOp = cast<omp::TargetOp>(splitBefore->getParentOp());
+ BlockArgument arg = getMappedBlockArg(load.getMemref(), targetOp);
+ return arg && !isWrittenAfterLoad(arg, load, splitBefore);
+ }
+ }
// If the op has side effects, it cannot be recomputed.
// We consider fir.declare as having no side effects.
return isa<fir::DeclareOp>(op) || isMemoryEffectFree(op);
diff --git a/flang/test/Lower/OpenMP/workdistribute-saxpy-allocatable.f90 b/flang/test/Lower/OpenMP/workdistribute-saxpy-allocatable.f90
new file mode 100644
index 0000000000000..a946f31c4b0bb
--- /dev/null
+++ b/flang/test/Lower/OpenMP/workdistribute-saxpy-allocatable.f90
@@ -0,0 +1,43 @@
+! RUN: %flang_fc1 -emit-fir -O0 -fopenmp -fopenmp-version=60 %s -o - | FileCheck %s
+! RUN: %flang_fc1 -emit-fir -O3 -fopenmp -fopenmp-version=60 %s -o - | FileCheck %s
+
+! Allocatable descriptors crossing the workdistribute target fission are
+! reloaded from their maps in the kernel, not cached by value.
+
+! CHECK-LABEL: func @_QPsaxpy_allocatable
+! CHECK: %[[Y_MAP:.*]] = omp.map.info {{.*}}map_clauses(storage) capture(ByRef) members({{.*}}) name("y")
+
+! Scalars and temp array pointers may be cached, but no descriptor is.
+! CHECK: omp.target_data
+! CHECK-NOT: !fir.box<{{.*}}name("__flang_workdistribute
+
+! The kernel computing a * x + y reloads y's box from its mapped descriptor.
+! CHECK: omp.target kernel_type(spmd) {{.*}}map_entries(%[[Y_MAP]] -> %[[Y_ARG:[^,]+]],
+! CHECK: %[[Y_DECL:.*]] = fir.declare %[[Y_ARG]] uniq_name("_QFsaxpy_allocatableEy")
+! CHECK: %[[Y_BOX:.*]] = fir.load %[[Y_DECL]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+! CHECK: omp.teams {
+! CHECK-NEXT: omp.parallel {
+! CHECK-NEXT: omp.distribute {
+! CHECK-NEXT: omp.wsloop {
+! CHECK-NEXT: omp.loop_nest
+! CHECK: fir.box_addr %[[Y_BOX]]
+
+! Kernels after the host realloc store still reload y from its map.
+! CHECK-NOT: !fir.box<{{.*}}name("__flang_workdistribute
+! CHECK: fir.store %{{.*}} to %{{.*}} : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+! CHECK: omp.target kernel_type(spmd) {{.*}}map_entries(%[[Y_MAP]] -> %[[Y_COPY_ARG:[^,]+]],
+! CHECK: %[[Y_COPY_DECL:.*]] = fir.declare %[[Y_COPY_ARG]] uniq_name("_QFsaxpy_allocatableEy")
+! CHECK: fir.load %[[Y_COPY_DECL]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+! CHECK: omp.teams {
+! CHECK-NEXT: omp.parallel {
+! CHECK-NEXT: omp.distribute {
+! CHECK-NEXT: omp.wsloop {
+! CHECK-NEXT: omp.loop_nest
+subroutine saxpy_allocatable(a, x, y)
+ real :: a
+ real, allocatable :: x(:), y(:)
+
+ !$omp target teams workdistribute
+ y = a * x + y
+ !$omp end target teams workdistribute
+end subroutine saxpy_allocatable
diff --git a/flang/test/Transforms/OpenMP/lower-workdistribute-fission-recompute-box.mlir b/flang/test/Transforms/OpenMP/lower-workdistribute-fission-recompute-box.mlir
new file mode 100644
index 0000000000000..5aa80ccd6b61f
--- /dev/null
+++ b/flang/test/Transforms/OpenMP/lower-workdistribute-fission-recompute-box.mlir
@@ -0,0 +1,298 @@
+// RUN: fir-opt --lower-workdistribute %s | FileCheck %s
+
+// A Fortran allocatable descriptor (fir.box) crossing the workdistribute target
+// fission must be recomputed inside the isolated target from its mapped
+// descriptor, not cached by value via __flang_workdistribute_to/from. Caching
+// the box would freeze a host base_addr into the device kernel.
+
+// The kernel reloads the box from its own mapped descriptor argument.
+// CHECK-LABEL: func.func @recompute_descriptor(
+// CHECK-SAME: %[[X:.*]]: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) {
+// CHECK: %[[X_MAP:.*]] = omp.map.info var_ptr(%[[X]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.box<!fir.heap<!fir.array<?xf32>>>) map_clauses(tofrom) capture(ByRef) name("x") -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+// CHECK: %[[X_STORAGE:.*]] = omp.map.info var_ptr(%[[X]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.box<!fir.heap<!fir.array<?xf32>>>) map_clauses(storage) capture(ByRef) name("x") -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+// CHECK: omp.target_data map_entries(%[[X_MAP]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) {
+// CHECK-NOT: __flang_workdistribute
+// CHECK: %[[MLIR_0:.*]] = llvm.mlir.constant(0 : i32) : i32
+// CHECK: %[[CONSTANT_0:.*]] = arith.constant 0 : index
+// CHECK: %[[CONSTANT_1:.*]] = arith.constant 1 : index
+// CHECK: %[[CONSTANT_2:.*]] = arith.constant 9 : index
+// CHECK: %[[CONSTANT_3:.*]] = arith.constant 5.000000e+00 : f32
+// CHECK: %[[LOAD_0:.*]] = fir.load %[[X]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+// CHECK: omp.target kernel_type(spmd) map_entries(%[[X_STORAGE]] -> %[[X_ARG:.*]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) {
+// CHECK: %[[CONSTANT_4:.*]] = arith.constant 0 : index
+// CHECK: %[[CONSTANT_5:.*]] = arith.constant 1 : index
+// CHECK: %[[CONSTANT_6:.*]] = arith.constant 9 : index
+// CHECK: %[[CONSTANT_7:.*]] = arith.constant 5.000000e+00 : f32
+// CHECK: %[[BOX:.*]] = fir.load %[[X_ARG]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+// CHECK: omp.teams {
+// CHECK: omp.parallel {
+// CHECK: omp.distribute {
+// CHECK: omp.wsloop {
+// CHECK: omp.loop_nest (%[[IV:.*]]) : index = (%[[CONSTANT_4]]) to (%[[CONSTANT_6]]) inclusive step (%[[CONSTANT_5]]) {
+// CHECK: %[[BOX_ADDR_0:.*]] = fir.box_addr %[[BOX]] : (!fir.box<!fir.heap<!fir.array<?xf32>>>) -> !fir.heap<!fir.array<?xf32>>
+// CHECK: %[[COORDINATE_OF_0:.*]] = fir.coordinate_of %[[BOX_ADDR_0]], %[[IV]] : (!fir.heap<!fir.array<?xf32>>, index) -> !fir.ref<f32>
+// CHECK: fir.store %[[CONSTANT_7]] to %[[COORDINATE_OF_0]] : !fir.ref<f32>
+// CHECK: omp.yield
+// CHECK: }
+// CHECK: } {omp.composite}
+// CHECK: } {omp.composite}
+// CHECK: omp.terminator
+// CHECK: } {omp.composite}
+// CHECK: omp.terminator
+// CHECK: } {omp.combined}
+// CHECK: omp.terminator
+// CHECK: } {omp.combined}
+// CHECK: omp.terminator
+// CHECK: }
+// CHECK: return
+// CHECK: }
+
+module attributes {llvm.target_triple = "amdgcn-amd-amdhsa", omp.is_gpu = true, omp.is_target_device = true} {
+func.func @recompute_descriptor(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) {
+ %map = omp.map.info var_ptr(%arg0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.box<!fir.heap<!fir.array<?xf32>>>) map_clauses(tofrom) capture(ByRef) name("x") -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ omp.target kernel_type(generic) map_entries(%map -> %barg : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c9 = arith.constant 9 : index
+ %cst = arith.constant 5.000000e+00 : f32
+ // The descriptor load crosses the split - it must be recomputed, not cached.
+ %box = fir.load %barg : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ omp.teams {
+ omp.workdistribute {
+ fir.do_loop %iv = %c0 to %c9 step %c1 unordered {
+ %addr = fir.box_addr %box : (!fir.box<!fir.heap<!fir.array<?xf32>>>) -> !fir.heap<!fir.array<?xf32>>
+ %coor = fir.coordinate_of %addr, %iv : (!fir.heap<!fir.array<?xf32>>, index) -> !fir.ref<f32>
+ fir.store %cst to %coor : !fir.ref<f32>
+ }
+ omp.terminator
+ }
+ omp.terminator
+ }
+ omp.terminator
+ }
+ return
+}
+
+// A store before the load, like the realloc of an allocatable LHS, does not change what a re-read sees.
+// CHECK-LABEL: func.func @recompute_after_store(
+// CHECK-SAME: %[[X:.*]]: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>,
+// CHECK-SAME: %[[Y:.*]]: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) {
+// CHECK: %[[X_MAP:.*]] = omp.map.info var_ptr(%[[X]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.box<!fir.heap<!fir.array<?xf32>>>) map_clauses(tofrom) capture(ByRef) name("x") -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+// CHECK: %[[Y_MAP:.*]] = omp.map.info var_ptr(%[[Y]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.box<!fir.heap<!fir.array<?xf32>>>) map_clauses(tofrom) capture(ByRef) name("y") -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+// CHECK: %[[X_STORAGE:.*]] = omp.map.info var_ptr(%[[X]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.box<!fir.heap<!fir.array<?xf32>>>) map_clauses(storage) capture(ByRef) name("x") -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+// CHECK: %[[Y_STORAGE:.*]] = omp.map.info var_ptr(%[[Y]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.box<!fir.heap<!fir.array<?xf32>>>) map_clauses(storage) capture(ByRef) name("y") -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+// CHECK: omp.target_data map_entries(%[[X_MAP]], %[[Y_MAP]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) {
+// CHECK-NOT: __flang_workdistribute
+// CHECK: %[[MLIR_0:.*]] = llvm.mlir.constant(0 : i32) : i32
+// CHECK: %[[CONSTANT_0:.*]] = arith.constant 0 : index
+// CHECK: %[[CONSTANT_1:.*]] = arith.constant 1 : index
+// CHECK: %[[CONSTANT_2:.*]] = arith.constant 9 : index
+// CHECK: %[[CONSTANT_3:.*]] = arith.constant 5.000000e+00 : f32
+// CHECK: %[[OTHER:.*]] = fir.load %[[Y]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+// CHECK: fir.store %[[OTHER]] to %[[X]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+// CHECK: %[[LOAD_1:.*]] = fir.load %[[X]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+// CHECK: omp.target kernel_type(spmd) map_entries(%[[X_STORAGE]] -> %[[X_ARG:.*]], %[[Y_STORAGE]] -> %[[Y_ARG:.*]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) {
+// CHECK: %[[CONSTANT_4:.*]] = arith.constant 0 : index
+// CHECK: %[[CONSTANT_5:.*]] = arith.constant 1 : index
+// CHECK: %[[CONSTANT_6:.*]] = arith.constant 9 : index
+// CHECK: %[[CONSTANT_7:.*]] = arith.constant 5.000000e+00 : f32
+// CHECK: %[[BOX:.*]] = fir.load %[[X_ARG]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+// CHECK: omp.teams {
+// CHECK: omp.parallel {
+// CHECK: omp.distribute {
+// CHECK: omp.wsloop {
+// CHECK: omp.loop_nest (%[[IV:.*]]) : index = (%[[CONSTANT_4]]) to (%[[CONSTANT_6]]) inclusive step (%[[CONSTANT_5]]) {
+// CHECK: %[[BOX_ADDR_0:.*]] = fir.box_addr %[[BOX]] : (!fir.box<!fir.heap<!fir.array<?xf32>>>) -> !fir.heap<!fir.array<?xf32>>
+// CHECK: %[[COORDINATE_OF_0:.*]] = fir.coordinate_of %[[BOX_ADDR_0]], %[[IV]] : (!fir.heap<!fir.array<?xf32>>, index) -> !fir.ref<f32>
+// CHECK: fir.store %[[CONSTANT_7]] to %[[COORDINATE_OF_0]] : !fir.ref<f32>
+// CHECK: omp.yield
+// CHECK: }
+// CHECK: } {omp.composite}
+// CHECK: } {omp.composite}
+// CHECK: omp.terminator
+// CHECK: } {omp.composite}
+// CHECK: omp.terminator
+// CHECK: } {omp.combined}
+// CHECK: omp.terminator
+// CHECK: } {omp.combined}
+// CHECK: omp.terminator
+// CHECK: }
+// CHECK: return
+// CHECK: }
+func.func @recompute_after_store(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, %arg1: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) {
+ %map0 = omp.map.info var_ptr(%arg0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.box<!fir.heap<!fir.array<?xf32>>>) map_clauses(tofrom) capture(ByRef) name("x") -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ %map1 = omp.map.info var_ptr(%arg1 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.box<!fir.heap<!fir.array<?xf32>>>) map_clauses(tofrom) capture(ByRef) name("y") -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ omp.target kernel_type(generic) map_entries(%map0 -> %bx, %map1 -> %by : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c9 = arith.constant 9 : index
+ %cst = arith.constant 5.000000e+00 : f32
+ %other = fir.load %by : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ fir.store %other to %bx : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ %box = fir.load %bx : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ omp.teams {
+ omp.workdistribute {
+ fir.do_loop %iv = %c0 to %c9 step %c1 unordered {
+ %addr = fir.box_addr %box : (!fir.box<!fir.heap<!fir.array<?xf32>>>) -> !fir.heap<!fir.array<?xf32>>
+ %coor = fir.coordinate_of %addr, %iv : (!fir.heap<!fir.array<?xf32>>, index) -> !fir.ref<f32>
+ fir.store %cst to %coor : !fir.ref<f32>
+ }
+ omp.terminator
+ }
+ omp.terminator
+ }
+ omp.terminator
+ }
+ return
+}
+
+// A box loaded from a local copy must be cached, since the initializing store is not cloned.
+// CHECK-LABEL: func.func @cache_local_descriptor(
+// CHECK-SAME: %[[X:.*]]: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) {
+// CHECK: %[[X_MAP:.*]] = omp.map.info var_ptr(%[[X]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.box<!fir.heap<!fir.array<?xf32>>>) map_clauses(tofrom) capture(ByRef) name("x") -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+// CHECK: %[[X_STORAGE:.*]] = omp.map.info var_ptr(%[[X]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.box<!fir.heap<!fir.array<?xf32>>>) map_clauses(storage) capture(ByRef) name("x") -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+// CHECK: omp.target_data map_entries(%[[X_MAP]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) {
+// CHECK: %[[CACHE:.*]] = fir.alloca !fir.box<!fir.heap<!fir.array<?xf32>>>
+// CHECK: %[[CACHE_FROM:.*]] = omp.map.info var_ptr(%[[CACHE]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.box<!fir.heap<!fir.array<?xf32>>>) map_clauses(from) capture(ByRef) name("__flang_workdistribute_from") -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+// CHECK: %[[CACHE_TO:.*]] = omp.map.info var_ptr(%[[CACHE]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.box<!fir.heap<!fir.array<?xf32>>>) map_clauses(to) capture(ByRef) name("__flang_workdistribute_to") -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+// CHECK: %[[MLIR_0:.*]] = llvm.mlir.constant(0 : i32) : i32
+// CHECK: %[[CONSTANT_0:.*]] = arith.constant 0 : index
+// CHECK: %[[CONSTANT_1:.*]] = arith.constant 1 : index
+// CHECK: %[[CONSTANT_2:.*]] = arith.constant 9 : index
+// CHECK: %[[CONSTANT_3:.*]] = arith.constant 5.000000e+00 : f32
+// CHECK: %[[LOCAL:.*]] = fir.alloca !fir.box<!fir.heap<!fir.array<?xf32>>>
+// CHECK: %[[INIT:.*]] = fir.load %[[X]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+// CHECK: fir.store %[[INIT]] to %[[LOCAL]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+// CHECK: %[[LOCAL_BOX:.*]] = fir.load %[[LOCAL]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+// CHECK: fir.store %[[LOCAL_BOX]] to %[[CACHE]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+// CHECK: omp.target kernel_type(spmd) map_entries(%[[X_STORAGE]] -> %[[X_ARG:.*]], %[[CACHE_TO]] -> %[[CACHE_ARG:.*]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) {
+// CHECK: %[[BOX:.*]] = fir.load %[[CACHE_ARG]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+// CHECK: %[[CONSTANT_4:.*]] = arith.constant 0 : index
+// CHECK: %[[CONSTANT_5:.*]] = arith.constant 1 : index
+// CHECK: %[[CONSTANT_6:.*]] = arith.constant 9 : index
+// CHECK: %[[CONSTANT_7:.*]] = arith.constant 5.000000e+00 : f32
+// CHECK: omp.teams {
+// CHECK: omp.parallel {
+// CHECK: omp.distribute {
+// CHECK: omp.wsloop {
+// CHECK: omp.loop_nest (%[[IV:.*]]) : index = (%[[CONSTANT_4]]) to (%[[CONSTANT_6]]) inclusive step (%[[CONSTANT_5]]) {
+// CHECK: %[[BOX_ADDR_0:.*]] = fir.box_addr %[[BOX]] : (!fir.box<!fir.heap<!fir.array<?xf32>>>) -> !fir.heap<!fir.array<?xf32>>
+// CHECK: %[[COORDINATE_OF_0:.*]] = fir.coordinate_of %[[BOX_ADDR_0]], %[[IV]] : (!fir.heap<!fir.array<?xf32>>, index) -> !fir.ref<f32>
+// CHECK: fir.store %[[CONSTANT_7]] to %[[COORDINATE_OF_0]] : !fir.ref<f32>
+// CHECK: omp.yield
+// CHECK: }
+// CHECK: } {omp.composite}
+// CHECK: } {omp.composite}
+// CHECK: omp.terminator
+// CHECK: } {omp.composite}
+// CHECK: omp.terminator
+// CHECK: } {omp.combined}
+// CHECK: omp.terminator
+// CHECK: } {omp.combined}
+// CHECK: omp.terminator
+// CHECK: }
+// CHECK: return
+// CHECK: }
+func.func @cache_local_descriptor(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) {
+ %map = omp.map.info var_ptr(%arg0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.box<!fir.heap<!fir.array<?xf32>>>) map_clauses(tofrom) capture(ByRef) name("x") -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ omp.target kernel_type(generic) map_entries(%map -> %barg : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c9 = arith.constant 9 : index
+ %cst = arith.constant 5.000000e+00 : f32
+ %local = fir.alloca !fir.box<!fir.heap<!fir.array<?xf32>>>
+ %init = fir.load %barg : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ fir.store %init to %local : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ %box = fir.load %local : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ omp.teams {
+ omp.workdistribute {
+ fir.do_loop %iv = %c0 to %c9 step %c1 unordered {
+ %addr = fir.box_addr %box : (!fir.box<!fir.heap<!fir.array<?xf32>>>) -> !fir.heap<!fir.array<?xf32>>
+ %coor = fir.coordinate_of %addr, %iv : (!fir.heap<!fir.array<?xf32>>, index) -> !fir.ref<f32>
+ fir.store %cst to %coor : !fir.ref<f32>
+ }
+ omp.terminator
+ }
+ omp.terminator
+ }
+ omp.terminator
+ }
+ return
+}
+
+// A mapped box that is stored to after the load must be cached to keep the loaded value.
+// CHECK-LABEL: func.func @cache_written_descriptor(
+// CHECK-SAME: %[[X:.*]]: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>,
+// CHECK-SAME: %[[Y:.*]]: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) {
+// CHECK: %[[X_MAP:.*]] = omp.map.info var_ptr(%[[X]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.box<!fir.heap<!fir.array<?xf32>>>) map_clauses(tofrom) capture(ByRef) name("x") -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+// CHECK: %[[Y_MAP:.*]] = omp.map.info var_ptr(%[[Y]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.box<!fir.heap<!fir.array<?xf32>>>) map_clauses(tofrom) capture(ByRef) name("y") -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+// CHECK: %[[X_STORAGE:.*]] = omp.map.info var_ptr(%[[X]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.box<!fir.heap<!fir.array<?xf32>>>) map_clauses(storage) capture(ByRef) name("x") -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+// CHECK: %[[Y_STORAGE:.*]] = omp.map.info var_ptr(%[[Y]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.box<!fir.heap<!fir.array<?xf32>>>) map_clauses(storage) capture(ByRef) name("y") -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+// CHECK: omp.target_data map_entries(%[[X_MAP]], %[[Y_MAP]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) {
+// CHECK: %[[CACHE:.*]] = fir.alloca !fir.box<!fir.heap<!fir.array<?xf32>>>
+// CHECK: %[[CACHE_FROM:.*]] = omp.map.info var_ptr(%[[CACHE]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.box<!fir.heap<!fir.array<?xf32>>>) map_clauses(from) capture(ByRef) name("__flang_workdistribute_from") -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+// CHECK: %[[CACHE_TO:.*]] = omp.map.info var_ptr(%[[CACHE]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.box<!fir.heap<!fir.array<?xf32>>>) map_clauses(to) capture(ByRef) name("__flang_workdistribute_to") -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+// CHECK: %[[MLIR_0:.*]] = llvm.mlir.constant(0 : i32) : i32
+// CHECK: %[[CONSTANT_0:.*]] = arith.constant 0 : index
+// CHECK: %[[CONSTANT_1:.*]] = arith.constant 1 : index
+// CHECK: %[[CONSTANT_2:.*]] = arith.constant 9 : index
+// CHECK: %[[CONSTANT_3:.*]] = arith.constant 5.000000e+00 : f32
+// CHECK: %[[OLD_BOX:.*]] = fir.load %[[X]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+// CHECK: %[[OTHER:.*]] = fir.load %[[Y]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+// CHECK: fir.store %[[OTHER]] to %[[X]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+// CHECK: fir.store %[[OLD_BOX]] to %[[CACHE]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+// CHECK: omp.target kernel_type(spmd) map_entries(%[[X_STORAGE]] -> %[[X_ARG:.*]], %[[Y_STORAGE]] -> %[[Y_ARG:.*]], %[[CACHE_TO]] -> %[[CACHE_ARG:.*]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) {
+// CHECK: %[[BOX:.*]] = fir.load %[[CACHE_ARG]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+// CHECK: %[[CONSTANT_4:.*]] = arith.constant 0 : index
+// CHECK: %[[CONSTANT_5:.*]] = arith.constant 1 : index
+// CHECK: %[[CONSTANT_6:.*]] = arith.constant 9 : index
+// CHECK: %[[CONSTANT_7:.*]] = arith.constant 5.000000e+00 : f32
+// CHECK: omp.teams {
+// CHECK: omp.parallel {
+// CHECK: omp.distribute {
+// CHECK: omp.wsloop {
+// CHECK: omp.loop_nest (%[[IV:.*]]) : index = (%[[CONSTANT_4]]) to (%[[CONSTANT_6]]) inclusive step (%[[CONSTANT_5]]) {
+// CHECK: %[[BOX_ADDR_0:.*]] = fir.box_addr %[[BOX]] : (!fir.box<!fir.heap<!fir.array<?xf32>>>) -> !fir.heap<!fir.array<?xf32>>
+// CHECK: %[[COORDINATE_OF_0:.*]] = fir.coordinate_of %[[BOX_ADDR_0]], %[[IV]] : (!fir.heap<!fir.array<?xf32>>, index) -> !fir.ref<f32>
+// CHECK: fir.store %[[CONSTANT_7]] to %[[COORDINATE_OF_0]] : !fir.ref<f32>
+// CHECK: omp.yield
+// CHECK: }
+// CHECK: } {omp.composite}
+// CHECK: } {omp.composite}
+// CHECK: omp.terminator
+// CHECK: } {omp.composite}
+// CHECK: omp.terminator
+// CHECK: } {omp.combined}
+// CHECK: omp.terminator
+// CHECK: } {omp.combined}
+// CHECK: omp.terminator
+// CHECK: }
+// CHECK: return
+// CHECK: }
+func.func @cache_written_descriptor(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, %arg1: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) {
+ %map0 = omp.map.info var_ptr(%arg0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.box<!fir.heap<!fir.array<?xf32>>>) map_clauses(tofrom) capture(ByRef) name("x") -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ %map1 = omp.map.info var_ptr(%arg1 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.box<!fir.heap<!fir.array<?xf32>>>) map_clauses(tofrom) capture(ByRef) name("y") -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ omp.target kernel_type(generic) map_entries(%map0 -> %bx, %map1 -> %by : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c9 = arith.constant 9 : index
+ %cst = arith.constant 5.000000e+00 : f32
+ %box = fir.load %bx : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ %other = fir.load %by : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ fir.store %other to %bx : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ omp.teams {
+ omp.workdistribute {
+ fir.do_loop %iv = %c0 to %c9 step %c1 unordered {
+ %addr = fir.box_addr %box : (!fir.box<!fir.heap<!fir.array<?xf32>>>) -> !fir.heap<!fir.array<?xf32>>
+ %coor = fir.coordinate_of %addr, %iv : (!fir.heap<!fir.array<?xf32>>, index) -> !fir.ref<f32>
+ fir.store %cst to %coor : !fir.ref<f32>
+ }
+ omp.terminator
+ }
+ omp.terminator
+ }
+ omp.terminator
+ }
+ return
+}
+}
More information about the flang-commits
mailing list