[flang-commits] [flang] [flang][OpenMP] Lower array-section workdistribute assign to element loop (PR #225595)
via flang-commits
flang-commits at lists.llvm.org
Tue Sep 22 22:09:49 PDT 2026
https://github.com/skc7 created https://github.com/llvm/llvm-project/pull/225595
Summary:
- At -O0 an array-section assignment inside `omp target teams workdistribute` (e.g. residuum(0,:,:) = residuum(n,:,:)) stays as a `_FortranAAssign`/`_FortranAAssignSimple` runtime call, since HLFIR only inlines it to explicit loops at -O1+.
- Detect the `array-to-array` case and lower it, like the existing `scalar-to-array` path, to an `unordered do loop` copying element by element via `fir.array_coor` on both descriptors, so each section's strides and bounds are honored.
- Rename the scalar broadcast helper to `replaceScalarToArrayAssignWithUnorderedDoLoop`.
- Emit an error for `FortranAAssign` argument shapes that have no lowering instead of silently falling back to the flat memcpy.
- Added a regression test checking array-to-array assign lowers to a `fir.array_coor` element copy and not `omp_target_memcpy`.
Assisted by: cladue opus 4.8
>From a9191d34f0b2abc4205fc7d4a2636c028ee44205 Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Wed, 23 Sep 2026 09:59:44 +0530
Subject: [PATCH] [flang][OpenMP] Lower array-section workdistribute assign to
element loop
---
.../Optimizer/OpenMP/LowerWorkdistribute.cpp | 122 ++++++++++++++++--
...r-workdistribute-runtime-assign-array.mlir | 69 ++++++++++
2 files changed, 180 insertions(+), 11 deletions(-)
create mode 100644 flang/test/Transforms/OpenMP/lower-workdistribute-runtime-assign-array.mlir
diff --git a/flang/lib/Optimizer/OpenMP/LowerWorkdistribute.cpp b/flang/lib/Optimizer/OpenMP/LowerWorkdistribute.cpp
index 5af2d1ddb5f506..0be7125ea7bf74 100644
--- a/flang/lib/Optimizer/OpenMP/LowerWorkdistribute.cpp
+++ b/flang/lib/Optimizer/OpenMP/LowerWorkdistribute.cpp
@@ -445,6 +445,13 @@ static bool isEnclosedTypeBoxScalar(Type type) {
return false;
}
+/// Check if the type is a fir.box that encloses an array (by value, no ref).
+static bool isEnclosedTypeBoxArray(Type type) {
+ if (auto boxType = dyn_cast<fir::BoxType>(type))
+ return isa<fir::SequenceType>(boxType.getEleTy());
+ return false;
+}
+
/// Check if the FortranAAssign call has src as scalar and dest as array
static bool isFortranAssignSrcScalarAndDestArray(fir::CallOp callOp) {
if (callOp.getNumOperands() < 2)
@@ -469,6 +476,21 @@ static bool isFortranAssignSrcScalarAndDestArray(fir::CallOp callOp) {
return srcIsScalar && destIsArray;
}
+/// Check if the FortranAAssign call has both src and dest as array descriptors.
+/// This is the array-section copy case (e.g. a(0,:,:) = b(n,:,:)) that a flat
+/// omp_target_memcpy would get wrong for strided sections. Matching the runtime
+/// signature, dest is a reference to a box and src is a box passed by value.
+static bool isFortranAssignSrcArrayAndDestArray(fir::CallOp callOp) {
+ if (callOp.getNumOperands() < 2)
+ return false;
+ auto srcConvert = callOp.getOperand(1).getDefiningOp<fir::ConvertOp>();
+ auto destConvert = callOp.getOperand(0).getDefiningOp<fir::ConvertOp>();
+ if (!srcConvert || !destConvert)
+ return false;
+ return isEnclosedTypeBoxArray(srcConvert.getValue().getType()) &&
+ isEnclosedTypeRefToBoxArray(destConvert.getValue().getType());
+}
+
/// Convert a flat index to multi-dimensional indices for an array box
/// Example: 2D array with shape (2,4)
/// Col 1 Col 2 Col 3 Col 4
@@ -536,11 +558,11 @@ static Value CalculateTotalElements(OpBuilder &builder, Location loc,
return totalElems;
}
-/// Replace the FortranAAssign runtime call with an unordered do loop
-static void replaceWithUnorderedDoLoop(OpBuilder &builder, Location loc,
- omp::TeamsOp teamsOp,
- omp::WorkdistributeOp workdistribute,
- fir::CallOp callOp) {
+/// Replace a scalar-to-array FortranAAssign (broadcast) runtime call with an
+/// unordered do loop that stores the scalar into every element.
+static void replaceScalarToArrayAssignWithUnorderedDoLoop(
+ OpBuilder &builder, Location loc, omp::TeamsOp teamsOp,
+ omp::WorkdistributeOp workdistribute, fir::CallOp callOp) {
auto destConvert = callOp.getOperand(0).getDefiningOp<fir::ConvertOp>();
auto srcConvert = callOp.getOperand(1).getDefiningOp<fir::ConvertOp>();
@@ -595,10 +617,70 @@ static void replaceWithUnorderedDoLoop(OpBuilder &builder, Location loc,
fir::StoreOp::create(builder, loc, scalar, elemPtr);
}
+/// Return the array descriptor (fir.box) value behind a FortranAAssign arg.
+/// The arg is the address of a descriptor temp: prefer the box value stored
+/// into it, otherwise load the reference.
+static Value getAssignArrayBox(OpBuilder &builder, Location loc, Value box) {
+ if (auto alloca = box.getDefiningOp<fir::AllocaOp>()) {
+ for (auto *user : alloca->getUsers())
+ if (auto storeOp = dyn_cast<fir::StoreOp>(user)) {
+ box = storeOp.getValue();
+ break;
+ }
+ }
+ if (isa<fir::ReferenceType>(box.getType()))
+ box = fir::LoadOp::create(builder, loc, box);
+ return box;
+}
+
+/// Replace an array-to-array FortranAAssign runtime call with an unordered do
+/// loop that copies element by element. Addressing goes through fir.array_coor
+/// on both descriptors, so each section's own strides and bounds are honored -
+/// unlike a flat memcpy, this is correct for strided sections.
+static void replaceArrayToArrayAssignWithUnorderedDoLoop(
+ OpBuilder &builder, Location loc, omp::TeamsOp teamsOp,
+ omp::WorkdistributeOp workdistribute, fir::CallOp callOp) {
+ auto destConvert = callOp.getOperand(0).getDefiningOp<fir::ConvertOp>();
+ auto srcConvert = callOp.getOperand(1).getDefiningOp<fir::ConvertOp>();
+
+ builder.setInsertionPoint(teamsOp);
+ Value destBox = getAssignArrayBox(builder, loc, destConvert.getValue());
+ Value srcBox = getAssignArrayBox(builder, loc, srcConvert.getValue());
+
+ // Element type comes from the destination sequence.
+ auto destBoxType = cast<fir::BoxType>(destBox.getType());
+ auto destSeqType = cast<fir::SequenceType>(destBoxType.getEleTy());
+ Type eleTy = destSeqType.getEleTy();
+ auto eleRefTy = fir::ReferenceType::get(eleTy);
+
+ auto c0 = arith::ConstantIndexOp::create(builder, loc, 0);
+ auto c1 = arith::ConstantIndexOp::create(builder, loc, 1);
+ Value totalElems = CalculateTotalElements(builder, loc, destBox);
+
+ auto *workdistributeBlock = &workdistribute.getRegion().front();
+ builder.setInsertionPointToStart(workdistributeBlock);
+ // Single flattened loop: dest and src conform, so one index set fits both.
+ auto doLoop = fir::DoLoopOp::create(builder, loc, c0, totalElems, c1, true);
+ builder.setInsertionPointToStart(doLoop.getBody());
+
+ auto flatIdx = doLoop.getRegion().front().getArgument(0);
+ SmallVector<Value> indices =
+ convertFlatToMultiDim(builder, loc, flatIdx, destBox);
+
+ auto srcPtr =
+ fir::ArrayCoorOp::create(builder, loc, eleRefTy, srcBox, nullptr, nullptr,
+ ValueRange{indices}, ValueRange{});
+ Value value = fir::LoadOp::create(builder, loc, srcPtr);
+ auto destPtr =
+ fir::ArrayCoorOp::create(builder, loc, eleRefTy, destBox, nullptr,
+ nullptr, ValueRange{indices}, ValueRange{});
+ fir::StoreOp::create(builder, loc, value, destPtr);
+}
+
/// workdistributeRuntimeCallLower method finds the runtime calls
-/// nested in teams {workdistribute{}} and
-/// lowers FortranAAssign to unordered do loop if src is scalar and dest is
-/// array. Other runtime calls are not handled currently.
+/// nested in teams {workdistribute{}} and lowers FortranAAssign to an
+/// unordered do loop for scalar-to-array and array-to-array assigns.
+/// Unsupported assign shapes error out. Other runtime calls are left as is.
static FailureOr<bool>
workdistributeRuntimeCallLower(omp::WorkdistributeOp workdistribute,
SetVector<omp::TargetOp> &targetOpsToProcess) {
@@ -620,6 +702,9 @@ workdistributeRuntimeCallLower(omp::WorkdistributeOp workdistribute,
bool changed = false;
// Get the target op parent of teams
omp::TargetOp targetOp = dyn_cast<omp::TargetOp>(teams->getParentOp());
+ // Runtime-call lowering only applies inside omp.target.
+ if (!targetOp)
+ return false;
SmallVector<Operation *> opsToErase;
for (auto &op : workdistribute.getOps()) {
if (isRuntimeCall(&op)) {
@@ -627,13 +712,28 @@ workdistributeRuntimeCallLower(omp::WorkdistributeOp workdistribute,
fir::CallOp runtimeCall = cast<fir::CallOp>(op);
auto funcName = runtimeCall.getCallee()->getRootReference().getValue();
if (isFortranAssignCall(funcName)) {
- if (isFortranAssignSrcScalarAndDestArray(runtimeCall) && targetOp) {
+ if (isFortranAssignSrcScalarAndDestArray(runtimeCall)) {
// Record the target ops to process later
targetOpsToProcess.insert(targetOp);
- replaceWithUnorderedDoLoop(rewriter, loc, teams, workdistribute,
- runtimeCall);
+ replaceScalarToArrayAssignWithUnorderedDoLoop(
+ rewriter, loc, teams, workdistribute, runtimeCall);
+ opsToErase.push_back(&op);
+ changed = true;
+ } else if (isFortranAssignSrcArrayAndDestArray(runtimeCall)) {
+ // Array-section copy: element-wise loop honors strides, unlike the
+ // flat omp_target_memcpy fallback used otherwise.
+ targetOpsToProcess.insert(targetOp);
+ replaceArrayToArrayAssignWithUnorderedDoLoop(
+ rewriter, loc, teams, workdistribute, runtimeCall);
opsToErase.push_back(&op);
changed = true;
+ } else {
+ // Recognized runtime call, but its argument shape has no lowering.
+ emitError(runtimeCall->getLoc(),
+ "Runtime call " + funcName +
+ " with this argument shape is not supported in "
+ "workdistribute yet.\n");
+ return failure();
}
}
}
diff --git a/flang/test/Transforms/OpenMP/lower-workdistribute-runtime-assign-array.mlir b/flang/test/Transforms/OpenMP/lower-workdistribute-runtime-assign-array.mlir
new file mode 100644
index 00000000000000..f3f7741e669494
--- /dev/null
+++ b/flang/test/Transforms/OpenMP/lower-workdistribute-runtime-assign-array.mlir
@@ -0,0 +1,69 @@
+// RUN: fir-opt --lower-workdistribute %s | FileCheck %s
+
+// An array-to-array _FortranAAssign in target teams workdistribute must lower
+// to an element-wise fir.array_coor copy, not a flat omp_target_memcpy.
+
+// Example Fortran code:
+// !$omp target teams workdistribute
+// a(:,:) = b(:,:)
+// !$omp end target teams workdistribute
+
+// CHECK-LABEL: func.func @array_assign(
+// CHECK: omp.target_data
+// CHECK: omp.target
+// CHECK: omp.teams
+// CHECK: omp.parallel
+// CHECK: omp.distribute
+// CHECK: omp.wsloop
+// CHECK: omp.loop_nest
+// CHECK: %[[SRC:.*]] = fir.array_coor {{.*}} : (!fir.box<!fir.array<?x?xf32>>, index, index) -> !fir.ref<f32>
+// CHECK: %[[VAL:.*]] = fir.load %[[SRC]] : !fir.ref<f32>
+// CHECK: %[[DST:.*]] = fir.array_coor {{.*}} : (!fir.box<!fir.array<?x?xf32>>, index, index) -> !fir.ref<f32>
+// CHECK: fir.store %[[VAL]] to %[[DST]] : !fir.ref<f32>
+// CHECK-NOT: omp_target_memcpy
+
+module attributes {llvm.target_triple = "amdgcn-amd-amdhsa", omp.is_gpu = true, omp.is_target_device = true} {
+func.func @array_assign(%a : !fir.ref<!fir.array<?x?xf32>>, %b : !fir.ref<!fir.array<?x?xf32>>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ %c20 = arith.constant 20 : index
+ %ub0 = arith.subi %c10, %c1 : index
+ %bnd0 = omp.map.bounds lower_bound(%c0 : index) upper_bound(%ub0 : index) extent(%c10 : index) stride(%c1 : index) start_idx(%c1 : index)
+ %ub1 = arith.subi %c20, %c1 : index
+ %bnd1 = omp.map.bounds lower_bound(%c0 : index) upper_bound(%ub1 : index) extent(%c20 : index) stride(%c1 : index) start_idx(%c1 : index)
+ %mapa = omp.map.info var_ptr(%a : !fir.ref<!fir.array<?x?xf32>>, f32) map_clauses(implicit, tofrom) capture(ByRef) bounds(%bnd0, %bnd1) name("a") -> !fir.ref<!fir.array<?x?xf32>>
+ %mapb = omp.map.info var_ptr(%b : !fir.ref<!fir.array<?x?xf32>>, f32) map_clauses(implicit, tofrom) capture(ByRef) bounds(%bnd0, %bnd1) name("b") -> !fir.ref<!fir.array<?x?xf32>>
+ omp.target kernel_type(generic) map_entries(%mapa -> %arga, %mapb -> %argb : !fir.ref<!fir.array<?x?xf32>>, !fir.ref<!fir.array<?x?xf32>>) {
+ // omp.target is isolated from above, so re-declare the extents here.
+ %e0 = arith.constant 10 : index
+ %e1 = arith.constant 20 : index
+ %shape = fir.shape %e0, %e1 : (index, index) -> !fir.shape<2>
+ %da = fir.declare %arga(%shape) {uniq_name = "a"} : (!fir.ref<!fir.array<?x?xf32>>, !fir.shape<2>) -> !fir.ref<!fir.array<?x?xf32>>
+ %db = fir.declare %argb(%shape) {uniq_name = "b"} : (!fir.ref<!fir.array<?x?xf32>>, !fir.shape<2>) -> !fir.ref<!fir.array<?x?xf32>>
+ omp.teams {
+ %dtmp = fir.alloca !fir.box<!fir.array<?x?xf32>> {pinned}
+ omp.workdistribute {
+ %srcbox = fir.embox %db(%shape) : (!fir.ref<!fir.array<?x?xf32>>, !fir.shape<2>) -> !fir.box<!fir.array<?x?xf32>>
+ %dstbox = fir.embox %da(%shape) : (!fir.ref<!fir.array<?x?xf32>>, !fir.shape<2>) -> !fir.box<!fir.array<?x?xf32>>
+ fir.store %dstbox to %dtmp : !fir.ref<!fir.box<!fir.array<?x?xf32>>>
+ %str = fir.address_of(@_QQcl) : !fir.ref<!fir.char<1,2>>
+ %line = arith.constant 13 : i32
+ %destc = fir.convert %dtmp : (!fir.ref<!fir.box<!fir.array<?x?xf32>>>) -> !fir.ref<!fir.box<none>>
+ %srcc = fir.convert %srcbox : (!fir.box<!fir.array<?x?xf32>>) -> !fir.box<none>
+ %strc = fir.convert %str : (!fir.ref<!fir.char<1,2>>) -> !fir.ref<i8>
+ fir.call @_FortranAAssignSimple(%destc, %srcc, %strc, %line) : (!fir.ref<!fir.box<none>>, !fir.box<none>, !fir.ref<i8>, i32) -> ()
+ omp.terminator
+ }
+ omp.terminator
+ }
+ omp.terminator
+ }
+ return
+}
+func.func private @_FortranAAssignSimple(!fir.ref<!fir.box<none>>, !fir.box<none>, !fir.ref<i8>, i32) attributes {fir.runtime}
+fir.global linkonce @_QQcl constant : !fir.char<1,2> {
+ %0 = fir.string_lit "f\00"(2) : !fir.char<1,2>
+ fir.has_value %0 : !fir.char<1,2>
+}
+}
More information about the flang-commits
mailing list