[flang-commits] [flang] 1ee8805 - [flang][OpenMP] Lower array-section workdistribute assign to element loop (#225595)
via flang-commits
flang-commits at lists.llvm.org
Wed Oct 14 22:18:43 PDT 2026
Author: Chaitanya
Date: 2026-10-15T10:48:37+05:30
New Revision: 1ee8805f01a663b7f2b09215a8eab5ab2d15ce35
URL: https://github.com/llvm/llvm-project/commit/1ee8805f01a663b7f2b09215a8eab5ab2d15ce35
DIFF: https://github.com/llvm/llvm-project/commit/1ee8805f01a663b7f2b09215a8eab5ab2d15ce35.diff
LOG: [flang][OpenMP] Lower array-section workdistribute assign to element loop (#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+.
- Identify and lower `array-to-array` assigns to two unordered loops
through a heap temporary (src -> tmp -> dest), using `fir.array_coor` so
strides are honored. This handles overlapping sections.
- Runtime-call lowering now runs before fission and emits at the call
site, so each loop gets its own kernel.
- Rename the scalar broadcast helper to
`replaceScalarToArrayAssignWithUnorderedDoLoop`.
- Tests: Strided array assign and overlapping same-array assign tests
are added.
Assisted by: cladue opus 4.8
Added:
flang/test/Transforms/OpenMP/lower-workdistribute-runtime-assign-array-alloc.mlir
flang/test/Transforms/OpenMP/lower-workdistribute-runtime-assign-array-overlap.mlir
flang/test/Transforms/OpenMP/lower-workdistribute-runtime-assign-array.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 70d4b0efe58f6d..1b20ae4ee35132 100644
--- a/flang/lib/Optimizer/OpenMP/LowerWorkdistribute.cpp
+++ b/flang/lib/Optimizer/OpenMP/LowerWorkdistribute.cpp
@@ -416,33 +416,28 @@ workdistributeDoLower(omp::WorkdistributeOp workdistribute,
return false;
}
-/// Check if the enclosed type in fir.ref is fir.box and fir.box encloses array
-static bool isEnclosedTypeRefToBoxArray(Type type) {
- // Check if it's a reference type
- if (auto refType = dyn_cast<fir::ReferenceType>(type)) {
- // Get the referenced type (should be fir.box)
- auto referencedType = refType.getEleTy();
- // Check if referenced type is a box
- if (auto boxType = dyn_cast<fir::BoxType>(referencedType)) {
- // Get the boxed type and check if it's an array
- auto boxedType = boxType.getEleTy();
- // Check if boxed type is a sequence (array)
- return isa<fir::SequenceType>(boxedType);
- }
- }
- return false;
+/// Returns the box type of a FortranAAssign operand if it describes an array of
+/// known rank, peeling a fir.ref when \p isRef is set. Returns null otherwise.
+/// Allocatable and pointer descriptors (box of heap/ptr array) are accepted.
+static fir::BaseBoxType getAssignArrayBoxType(Type type, bool isRef) {
+ if (isRef) {
+ if (!fir::isBoxAddress(type))
+ return {};
+ type = fir::unwrapRefType(type);
+ }
+ auto boxTy = dyn_cast<fir::BaseBoxType>(type);
+ if (!boxTy)
+ return {};
+ auto seqTy =
+ dyn_cast_or_null<fir::SequenceType>(fir::extractSequenceType(boxTy));
+ if (!seqTy || seqTy.hasUnknownShape())
+ return {};
+ return boxTy;
}
-/// Check if the enclosed type in fir.box is scalar (not array)
-static bool isEnclosedTypeBoxScalar(Type type) {
- // Check if it's a box type
- if (auto boxType = dyn_cast<fir::BoxType>(type)) {
- // Get the boxed type
- auto boxedType = boxType.getEleTy();
- // Check if boxed type is NOT a sequence (array)
- return !isa<fir::SequenceType>(boxedType);
- }
- return false;
+/// Returns the sequence type described by an array box value.
+static fir::SequenceType getBoxSequenceType(Value arrayBox) {
+ return cast<fir::SequenceType>(fir::extractSequenceType(arrayBox.getType()));
}
/// Check if the FortranAAssign call has src as scalar and dest as array
@@ -464,9 +459,32 @@ static bool isFortranAssignSrcScalarAndDestArray(fir::CallOp callOp) {
auto destOrigType = destConvert.getValue().getType();
// Check if src is scalar and dest is array
- bool srcIsScalar = isEnclosedTypeBoxScalar(srcOrigType);
- bool destIsArray = isEnclosedTypeRefToBoxArray(destOrigType);
- return srcIsScalar && destIsArray;
+ auto srcBoxTy = dyn_cast<fir::BoxType>(srcOrigType);
+ return srcBoxTy && !srcBoxTy.isArray() &&
+ getAssignArrayBoxType(destOrigType, /*isRef=*/true);
+}
+
+/// 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;
+ fir::BaseBoxType srcBoxTy =
+ getAssignArrayBoxType(srcConvert.getValue().getType(), /*isRef=*/false);
+ fir::BaseBoxType destBoxTy =
+ getAssignArrayBoxType(destConvert.getValue().getType(), /*isRef=*/true);
+ if (!srcBoxTy || !destBoxTy)
+ return false;
+ // The raw element load/store only handles trivial, matching element types.
+ // Derived types need finalization and characters need length parameters.
+ Type srcEleTy = srcBoxTy.unwrapInnerType();
+ return fir::isa_trivial(srcEleTy) && srcEleTy == destBoxTy.unwrapInnerType();
}
/// Convert a flat index to multi-dimensional indices for an array box
@@ -483,9 +501,7 @@ static SmallVector<Value> convertFlatToMultiDim(OpBuilder &builder,
Location loc, Value flatIdx,
Value arrayBox) {
// Get array type and rank
- auto boxType = cast<fir::BoxType>(arrayBox.getType());
- auto seqType = cast<fir::SequenceType>(boxType.getEleTy());
- int rank = seqType.getDimension();
+ int rank = getBoxSequenceType(arrayBox).getDimension();
// Get all extents
SmallVector<Value> extents;
@@ -518,9 +534,7 @@ static SmallVector<Value> convertFlatToMultiDim(OpBuilder &builder,
/// (totalElems = extent(1) * extent(2) * ... * extent(n))
static Value CalculateTotalElements(OpBuilder &builder, Location loc,
Value arrayBox) {
- auto boxType = cast<fir::BoxType>(arrayBox.getType());
- auto seqType = cast<fir::SequenceType>(boxType.getEleTy());
- int rank = seqType.getDimension();
+ int rank = getBoxSequenceType(arrayBox).getDimension();
Value totalElems = nullptr;
for (int i = 0; i < rank; ++i) {
@@ -536,11 +550,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,
+ fir::CallOp callOp) {
auto destConvert = callOp.getOperand(0).getDefiningOp<fir::ConvertOp>();
auto srcConvert = callOp.getOperand(1).getDefiningOp<fir::ConvertOp>();
@@ -563,7 +577,7 @@ static void replaceWithUnorderedDoLoop(OpBuilder &builder, Location loc,
}
}
- builder.setInsertionPoint(teamsOp);
+ builder.setInsertionPoint(callOp);
// Load destination array box (if it's a reference)
Value arrayBox = destBox;
if (isa<fir::ReferenceType>(destBox.getType()))
@@ -576,11 +590,10 @@ static void replaceWithUnorderedDoLoop(OpBuilder &builder, Location loc,
auto c0 = arith::ConstantIndexOp::create(builder, loc, 0);
auto c1 = arith::ConstantIndexOp::create(builder, loc, 1);
Value totalElems = CalculateTotalElements(builder, loc, arrayBox);
+ // fir.do_loop upper bound is inclusive, so iterate [0, totalElems - 1].
+ Value ub = arith::SubIOp::create(builder, loc, totalElems, c1);
- auto *workdistributeBlock = &workdistribute.getRegion().front();
- builder.setInsertionPointToStart(workdistributeBlock);
- // Create single unordered loop for flattened array
- auto doLoop = fir::DoLoopOp::create(builder, loc, c0, totalElems, c1, true);
+ auto doLoop = fir::DoLoopOp::create(builder, loc, c0, ub, c1, true);
Block *loopBlock = &doLoop.getRegion().front();
builder.setInsertionPointToStart(doLoop.getBody());
@@ -595,10 +608,84 @@ 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 two unordered do
+/// loops that copy src into a heap temporary and then the temporary into dest.
+/// 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,
+ fir::CallOp callOp) {
+ auto destConvert = callOp.getOperand(0).getDefiningOp<fir::ConvertOp>();
+ auto srcConvert = callOp.getOperand(1).getDefiningOp<fir::ConvertOp>();
+
+ builder.setInsertionPoint(callOp);
+ Value destBox = getAssignArrayBox(builder, loc, destConvert.getValue());
+ Value srcBox = getAssignArrayBox(builder, loc, srcConvert.getValue());
+
+ // Element type comes from the destination sequence.
+ Type eleTy = getBoxSequenceType(destBox).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);
+ // fir.do_loop upper bound is inclusive, so iterate [0, totalElems - 1].
+ Value ub = arith::SubIOp::create(builder, loc, totalElems, c1);
+
+ // Fortran evaluates the whole RHS before storing, and src and dest may
+ // overlap. Fission later puts each loop in its own kernel, so every read
+ // finishes before any write.
+ auto tmpTy =
+ fir::SequenceType::get({fir::SequenceType::getUnknownExtent()}, eleTy);
+ Value tmp = fir::AllocMemOp::create(builder, loc, tmpTy, ValueRange{},
+ ValueRange{totalElems});
+
+ // Copies between box[indices] and tmp[flatIdx] over the flattened dest index
+ // space. Dest and src conform, so one index set fits both.
+ auto genCopyLoop = [&](Value box, bool toTmp) {
+ builder.setInsertionPoint(callOp);
+ auto doLoop = fir::DoLoopOp::create(builder, loc, c0, ub, c1, true);
+ builder.setInsertionPointToStart(doLoop.getBody());
+ Value flatIdx = doLoop.getInductionVar();
+ SmallVector<Value> indices =
+ convertFlatToMultiDim(builder, loc, flatIdx, destBox);
+ Value boxPtr = fir::ArrayCoorOp::create(
+ builder, loc, eleRefTy, box, nullptr, nullptr, indices, ValueRange{});
+ Value tmpPtr = fir::CoordinateOp::create(builder, loc, eleRefTy, tmp,
+ ValueRange{flatIdx});
+ Value from = toTmp ? boxPtr : tmpPtr;
+ Value to = toTmp ? tmpPtr : boxPtr;
+ Value value = fir::LoadOp::create(builder, loc, from);
+ fir::StoreOp::create(builder, loc, value, to);
+ };
+ genCopyLoop(srcBox, /*toTmp=*/true);
+ genCopyLoop(destBox, /*toTmp=*/false);
+
+ builder.setInsertionPoint(callOp);
+ fir::FreeMemOp::create(builder, loc, tmp);
+}
+
/// 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 +707,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 +717,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,
+ 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,
+ 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();
}
}
}
@@ -1852,27 +1957,29 @@ class LowerWorkdistributePass
if (verify.wasInterrupted())
return signalPassFailure();
- auto fission =
+ // Runs before fission so the loops it emits get split into their own
+ // teams regions.
+ auto rtCallLower =
moduleOp->walk([&](mlir::omp::WorkdistributeOp workdistribute) {
- auto res = fissionWorkdistribute(workdistribute);
+ auto res = workdistributeRuntimeCallLower(workdistribute,
+ targetOpsToProcess);
if (failed(res))
return WalkResult::interrupt();
changed |= *res;
return WalkResult::advance();
});
- if (fission.wasInterrupted())
+ if (rtCallLower.wasInterrupted())
return signalPassFailure();
- auto rtCallLower =
+ auto fission =
moduleOp->walk([&](mlir::omp::WorkdistributeOp workdistribute) {
- auto res = workdistributeRuntimeCallLower(workdistribute,
- targetOpsToProcess);
+ auto res = fissionWorkdistribute(workdistribute);
if (failed(res))
return WalkResult::interrupt();
changed |= *res;
return WalkResult::advance();
});
- if (rtCallLower.wasInterrupted())
+ if (fission.wasInterrupted())
return signalPassFailure();
moduleOp->walk([&](mlir::omp::WorkdistributeOp workdistribute) {
diff --git a/flang/test/Transforms/OpenMP/lower-workdistribute-runtime-assign-array-alloc.mlir b/flang/test/Transforms/OpenMP/lower-workdistribute-runtime-assign-array-alloc.mlir
new file mode 100644
index 00000000000000..69547b782108c9
--- /dev/null
+++ b/flang/test/Transforms/OpenMP/lower-workdistribute-runtime-assign-array-alloc.mlir
@@ -0,0 +1,59 @@
+// RUN: fir-opt --lower-workdistribute %s | FileCheck %s
+
+// An array-to-array _FortranAAssignSimple between allocatable descriptors
+// (box of a heap array) must take the element-wise lowering too.
+
+// Example Fortran code:
+// real, allocatable :: x(:), y(:)
+// !$omp target teams workdistribute
+// y = x
+// !$omp end target teams workdistribute
+
+// CHECK-LABEL: func.func @alloc_assign(
+// CHECK: omp.target_allocmem
+// CHECK: omp.target kernel_type
+// CHECK: omp.loop_nest
+// CHECK: %[[SRC:.*]] = fir.array_coor {{.*}} : (!fir.box<!fir.heap<!fir.array<?xf32>>>, index) -> !fir.ref<f32>
+// CHECK: %[[TMP:.*]] = fir.coordinate_of {{.*}} -> !fir.ref<f32>
+// CHECK: %[[VAL:.*]] = fir.load %[[SRC]] : !fir.ref<f32>
+// CHECK: fir.store %[[VAL]] to %[[TMP]] : !fir.ref<f32>
+// CHECK: omp.target kernel_type
+// CHECK: omp.loop_nest
+// CHECK: %[[DST:.*]] = fir.array_coor {{.*}} : (!fir.box<!fir.heap<!fir.array<?xf32>>>, index) -> !fir.ref<f32>
+// CHECK: %[[TMP2:.*]] = fir.coordinate_of {{.*}} -> !fir.ref<f32>
+// CHECK: %[[VAL2:.*]] = fir.load %[[TMP2]] : !fir.ref<f32>
+// CHECK: fir.store %[[VAL2]] to %[[DST]] : !fir.ref<f32>
+// CHECK: omp.target_freemem
+// CHECK-NOT: fir.call @_FortranAAssignSimple
+
+module attributes {llvm.target_triple = "amdgcn-amd-amdhsa", omp.is_gpu = true, omp.is_target_device = true} {
+func.func @alloc_assign(%x : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, %y : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) {
+ %mapx = 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>>>>
+ %mapy = 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>>>>
+ omp.target kernel_type(generic) map_entries(%mapx -> %argx, %mapy -> %argy : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) {
+ omp.teams {
+ %dtmp = fir.alloca !fir.box<!fir.heap<!fir.array<?xf32>>> {pinned}
+ omp.workdistribute {
+ %xbox = fir.load %argx : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ %ybox = fir.load %argy : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ fir.store %ybox to %dtmp : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ %str = fir.address_of(@_QQcl) : !fir.ref<!fir.char<1,2>>
+ %line = arith.constant 9 : i32
+ %destc = fir.convert %dtmp : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> !fir.ref<!fir.box<none>>
+ %srcc = fir.convert %xbox : (!fir.box<!fir.heap<!fir.array<?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>
+}
+}
diff --git a/flang/test/Transforms/OpenMP/lower-workdistribute-runtime-assign-array-overlap.mlir b/flang/test/Transforms/OpenMP/lower-workdistribute-runtime-assign-array-overlap.mlir
new file mode 100644
index 00000000000000..92c5260f44acec
--- /dev/null
+++ b/flang/test/Transforms/OpenMP/lower-workdistribute-runtime-assign-array-overlap.mlir
@@ -0,0 +1,72 @@
+// RUN: fir-opt --lower-workdistribute %s | FileCheck %s
+
+// An array-to-array _FortranAAssign whose sections overlap must copy through a
+// temporary. The copy is split into two kernels so every read of the source
+// finishes before any write to the destination.
+
+// Example Fortran code:
+// !$omp target teams workdistribute
+// a(2:10) = a(1:9)
+// !$omp end target teams workdistribute
+
+// CHECK-LABEL: func.func @overlap_assign(
+// CHECK: omp.target_data
+// CHECK: omp.target_allocmem
+// CHECK: omp.target kernel_type
+// CHECK: omp.loop_nest
+// CHECK: %[[SRC:.*]] = fir.array_coor
+// CHECK: %[[TMP:.*]] = fir.coordinate_of
+// CHECK: %[[VAL:.*]] = fir.load %[[SRC]]
+// CHECK: fir.store %[[VAL]] to %[[TMP]]
+// CHECK: omp.target kernel_type
+// CHECK: omp.loop_nest
+// CHECK: %[[DST:.*]] = fir.array_coor
+// CHECK: %[[TMP2:.*]] = fir.coordinate_of
+// CHECK: %[[VAL2:.*]] = fir.load %[[TMP2]]
+// CHECK: fir.store %[[VAL2]] to %[[DST]]
+// CHECK: omp.target_freemem
+// CHECK-NOT: fir.call @_FortranAAssign
+
+module attributes {llvm.target_triple = "amdgcn-amd-amdhsa", omp.is_gpu = true, omp.is_target_device = true} {
+func.func @overlap_assign(%a : !fir.ref<!fir.array<?xf32>>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : 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)
+ %mapa = omp.map.info var_ptr(%a : !fir.ref<!fir.array<?xf32>>, f32) map_clauses(implicit, tofrom) capture(ByRef) bounds(%bnd0) name("a") -> !fir.ref<!fir.array<?xf32>>
+ omp.target kernel_type(generic) map_entries(%mapa -> %arga : !fir.ref<!fir.array<?xf32>>) {
+ %e0 = arith.constant 10 : index
+ %shape = fir.shape %e0 : (index) -> !fir.shape<1>
+ %da = fir.declare %arga(%shape) uniq_name("a") : (!fir.ref<!fir.array<?xf32>>, !fir.shape<1>) -> !fir.ref<!fir.array<?xf32>>
+ omp.teams {
+ %dtmp = fir.alloca !fir.box<!fir.array<?xf32>> {pinned}
+ omp.workdistribute {
+ %one = arith.constant 1 : index
+ %two = arith.constant 2 : index
+ %nine = arith.constant 9 : index
+ %srcslice = fir.slice %one, %nine, %one : (index, index, index) -> !fir.slice<1>
+ %dstslice = fir.slice %two, %e0, %one : (index, index, index) -> !fir.slice<1>
+ %srcbox = fir.embox %da(%shape) [%srcslice] : (!fir.ref<!fir.array<?xf32>>, !fir.shape<1>, !fir.slice<1>) -> !fir.box<!fir.array<?xf32>>
+ %dstbox = fir.embox %da(%shape) [%dstslice] : (!fir.ref<!fir.array<?xf32>>, !fir.shape<1>, !fir.slice<1>) -> !fir.box<!fir.array<?xf32>>
+ fir.store %dstbox to %dtmp : !fir.ref<!fir.box<!fir.array<?xf32>>>
+ %str = fir.address_of(@_QQcl) : !fir.ref<!fir.char<1,2>>
+ %line = arith.constant 9 : i32
+ %destc = fir.convert %dtmp : (!fir.ref<!fir.box<!fir.array<?xf32>>>) -> !fir.ref<!fir.box<none>>
+ %srcc = fir.convert %srcbox : (!fir.box<!fir.array<?xf32>>) -> !fir.box<none>
+ %strc = fir.convert %str : (!fir.ref<!fir.char<1,2>>) -> !fir.ref<i8>
+ fir.call @_FortranAAssign(%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 @_FortranAAssign(!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>
+}
+}
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..a2c2d5426289e0
--- /dev/null
+++ b/flang/test/Transforms/OpenMP/lower-workdistribute-runtime-assign-array.mlir
@@ -0,0 +1,91 @@
+// 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. The
+// copy must address through the (strided) section descriptors. It goes
+// through a heap temporary in two kernels, src to tmp and then tmp to dest.
+
+// Example Fortran code:
+// !$omp target teams workdistribute
+// a(1:10:2,:) = b(1:10:2,:)
+// !$omp end target teams workdistribute
+
+// CHECK-LABEL: func.func @array_assign(
+// CHECK: omp.target_data
+// CHECK: omp.target_allocmem
+
+// First kernel reads the stride-2 section of b into the temporary.
+// CHECK: omp.target kernel_type
+// CHECK: %[[DB:.*]] = fir.declare {{.*}} uniq_name("b")
+// CHECK: %[[STEP:.*]] = arith.constant 2 : index
+// CHECK: %[[SLICE:.*]] = fir.slice {{[^,]*}}, {{[^,]*}}, %[[STEP]], {{.*}} -> !fir.slice<2>
+// CHECK: %[[SRCBOX:.*]] = fir.embox %[[DB]]({{.*}}) [%[[SLICE]]]
+// CHECK: omp.loop_nest
+// CHECK: %[[SRC:.*]] = fir.array_coor %[[SRCBOX]] {{.*}} : (!fir.box<!fir.array<?x?xf32>>, index, index) -> !fir.ref<f32>
+// CHECK: %[[TMP:.*]] = fir.coordinate_of {{.*}} -> !fir.ref<f32>
+// CHECK: %[[VAL:.*]] = fir.load %[[SRC]] : !fir.ref<f32>
+// CHECK: fir.store %[[VAL]] to %[[TMP]] : !fir.ref<f32>
+
+// Second kernel writes the temporary into the stride-2 section of a.
+// CHECK: omp.target kernel_type
+// CHECK: %[[DA:.*]] = fir.declare {{.*}} uniq_name("a")
+// CHECK: %[[STEP2:.*]] = arith.constant 2 : index
+// CHECK: %[[SLICE2:.*]] = fir.slice {{[^,]*}}, {{[^,]*}}, %[[STEP2]], {{.*}} -> !fir.slice<2>
+// CHECK: %[[DSTBOX:.*]] = fir.embox %[[DA]]({{.*}}) [%[[SLICE2]]]
+// CHECK: omp.loop_nest
+// CHECK: %[[DST:.*]] = fir.array_coor %[[DSTBOX]] {{.*}} : (!fir.box<!fir.array<?x?xf32>>, index, index) -> !fir.ref<f32>
+// CHECK: %[[TMP2:.*]] = fir.coordinate_of {{.*}} -> !fir.ref<f32>
+// CHECK: %[[VAL2:.*]] = fir.load %[[TMP2]] : !fir.ref<f32>
+// CHECK: fir.store %[[VAL2]] to %[[DST]] : !fir.ref<f32>
+// CHECK: omp.target_freemem
+// 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 {
+ // Stride-2 sections in dim 0, so the copy must address through the descriptors.
+ %lb = arith.constant 1 : index
+ %st = arith.constant 2 : index
+ %slice = fir.slice %lb, %e0, %st, %lb, %e1, %lb : (index, index, index, index, index, index) -> !fir.slice<2>
+ %srcbox = fir.embox %db(%shape) [%slice] : (!fir.ref<!fir.array<?x?xf32>>, !fir.shape<2>, !fir.slice<2>) -> !fir.box<!fir.array<?x?xf32>>
+ %dstbox = fir.embox %da(%shape) [%slice] : (!fir.ref<!fir.array<?x?xf32>>, !fir.shape<2>, !fir.slice<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