[flang-commits] [flang] [flang][OpenMP] Fix wrong results for FORALL in a workshare construct (PR #211371)
Carlos Seo via flang-commits
flang-commits at lists.llvm.org
Fri Jul 24 08:57:34 PDT 2026
https://github.com/ceseo updated https://github.com/llvm/llvm-project/pull/211371
>From 438a6adfd7b7229424162bbc1269133cab52f9fa Mon Sep 17 00:00:00 2001
From: Carlos Seo <carlos.seo at linaro.org>
Date: Wed, 22 Jul 2026 13:46:44 -0300
Subject: [PATCH 1/3] [flang][OpenMP] Fix wrong results for FORALL in a
workshare construct
A FORALL in a workshare construct could produce wrong results
non-deterministically. This is caused by two issues in the workshare lowering:
1. A FORALL whose left-hand side may overlap its right-hand side is
lowered into two loop nests around a runtime value stack: the first nest
evaluates each right-hand side and pushes it, the second one fetches the
saved values back with a running counter. That counter lives in a
fir.alloca which, since omp.parallel is an alloca scope, is thread
private. The counter is read, incremented and written back from inside
the omp.single generated for the fetch, because the incremented value is
only available there. Only the thread which executed the omp.single
therefore bumped its own copy of the counter, and all the other threads
kept a stale one and refetched an already consumed element on the
following iterations.
Collect the thread local memory which is only updated by the thread
executing an omp.single and broadcast it with copyprivate, so that the
copies of the other threads stay in sync. As nowait and copyprivate are
mutually exclusive on a single construct, nowait is no longer set when
there is something to broadcast.
2. nowait was only suppressed when the immediately enclosing
operation was loop-like. A masked FORALL introduces a fir.if inside the
fir.do_loop, so the last omp.single or omp.wsloop of the fir.if body was
given nowait even though the loop may run it again, and even though
there was more work after the loop. Thread the information down the
recursion instead, so that only the work which is really last in the
whole omp.workshare region may rely on the barrier emitted at the end of
that region.
Fixes #209942
Fixes #209943
---
flang/lib/Optimizer/OpenMP/LowerWorkshare.cpp | 102 +++++++++++++++---
.../OpenMP/lower-workshare-nowait.mlir | 47 ++++++++
.../OpenMP/lower-workshare-thread-local.mlir | 69 +++++++++++-
3 files changed, 200 insertions(+), 18 deletions(-)
diff --git a/flang/lib/Optimizer/OpenMP/LowerWorkshare.cpp b/flang/lib/Optimizer/OpenMP/LowerWorkshare.cpp
index 2bc8a4b80589c..8f50b10a07917 100644
--- a/flang/lib/Optimizer/OpenMP/LowerWorkshare.cpp
+++ b/flang/lib/Optimizer/OpenMP/LowerWorkshare.cpp
@@ -259,6 +259,48 @@ static bool isSafeToParallelize(Operation *op) {
return false;
}
+// Collects the thread-local memory locations that op writes to and that
+// need to be broadcasted to other threads when op ends up being executed
+// by a single thread only.
+//
+// Some thread-local variables carry state which is logically shared by the
+// whole omp.workshare region even though each thread owns a copy of it.
+//
+// One example is the fetch counter of the temporary storage used to implement
+// FORALL: it is bumped from within an omp.single (because the value it is
+// bumped by is only available there), so the copies owned by the threads
+// which did not execute the omp.single would otherwise go stale and the
+// following iterations would fetch the wrong element. See issue #209942.
+//
+// Only the thread-local allocation itself is considered, so that a shallow
+// copy of it faithfully reproduces the update on the other threads.
+static void collectThreadLocalWrites(Operation *op,
+ llvm::SmallVectorImpl<Value> &vars) {
+ auto memEffects = dyn_cast<MemoryEffectOpInterface>(op);
+ if (!memEffects)
+ return;
+ SmallVector<MemoryEffects::EffectInstance> effects;
+ memEffects.getEffects(effects);
+ for (const MemoryEffects::EffectInstance &effect : effects) {
+ if (!isa<MemoryEffects::Write>(effect.getEffect()))
+ continue;
+ Value val = effect.getValue();
+ if (!val || !val.getDefiningOp<fir::AllocaOp>())
+ continue;
+ auto refTy = dyn_cast<fir::ReferenceType>(val.getType());
+ if (!refTy)
+ continue;
+ // createCopyFunc emits a load/store pair, so restrict this to types for
+ // which such a shallow copy is both legal and cheap.
+ mlir::Type eleTy = refTy.getEleTy();
+ if (!fir::isa_trivial(eleTy) && !fir::isa_box_type(eleTy))
+ continue;
+ if (!isOpenMPThreadLocalMemory(op, val))
+ continue;
+ vars.push_back(val);
+ }
+}
+
/// Simple shallow copies suffice for our purposes in this pass, so we implement
/// this simpler alternative to the full fledged `createCopyFunc` in the
/// frontend
@@ -339,9 +381,13 @@ static void cleanupBlock(Block *block) {
op.erase();
}
+// canUseNowait to check whether the work generated for sourceRegion is the
+// very last thing the omp.workshare region does, and thus whether the
+// synchronization of its last omp.single/omp.wsloop may be left to the
+// barrier emitted at the end of the omp.workshare region.
static void parallelizeRegion(Region &sourceRegion, Region &targetRegion,
IRMapping &rootMapping, Location loc,
- mlir::DominanceInfo &di) {
+ mlir::DominanceInfo &di, bool canUseNowait) {
OpBuilder rootBuilder(sourceRegion.getContext());
ModuleOp m = sourceRegion.getParentOfType<ModuleOp>();
OpBuilder copyFuncBuilder(m.getBodyRegion());
@@ -365,6 +411,9 @@ static void parallelizeRegion(Region &sourceRegion, Region &targetRegion,
OpBuilder parallelBuilder) -> std::pair<bool, SmallVector<Value>> {
IRMapping singleMapping = rootMapping;
SmallVector<Value> copyPrivate;
+ // Thread-local memory updated by the single thread only, which has to be
+ // broadcasted to the other threads to keep their copies in sync.
+ SmallVector<Value> threadLocalWrites;
bool allParallelized = true;
for (Operation &op : llvm::make_range(sr.begin, sr.end)) {
@@ -388,6 +437,9 @@ static void parallelizeRegion(Region &sourceRegion, Region &targetRegion,
assert(llvm::all_of(op.getResults(), [&](Value v) {
return !isTransitivelyUsedOutside(v, sr);
}));
+ // The operation only runs on the thread executing the omp.single,
+ // so the thread-local memory it updates has to be broadcasted.
+ collectThreadLocalWrites(&op, threadLocalWrites);
allParallelized = false;
}
} else if (auto alloca = dyn_cast<fir::AllocaOp>(&op)) {
@@ -399,6 +451,7 @@ static void parallelizeRegion(Region &sourceRegion, Region &targetRegion,
allParallelized = false;
} else {
singleBuilder.clone(op, singleMapping);
+ collectThreadLocalWrites(&op, threadLocalWrites);
// Prepare reloaded values for results of operations that cannot be
// safely parallelized and which are used after the region `sr`.
for (auto res : op.getResults()) {
@@ -413,6 +466,17 @@ static void parallelizeRegion(Region &sourceRegion, Region &targetRegion,
}
}
omp::TerminatorOp::create(singleBuilder, loc);
+
+ // Broadcast the thread-local state which only the thread executing the
+ // omp.single has updated. Values defined inside sr are remapped; values
+ // defined before it (e.g. hoisted allocas) are used as is.
+ llvm::SmallDenseSet<Value> seen(copyPrivate.begin(), copyPrivate.end());
+ for (Value v : threadLocalWrites) {
+ Value mapped = singleMapping.lookupOrDefault(v);
+ if (seen.insert(mapped).second)
+ copyPrivate.push_back(mapped);
+ }
+
return {allParallelized, copyPrivate};
};
@@ -425,7 +489,7 @@ static void parallelizeRegion(Region &sourceRegion, Region &targetRegion,
rootMapping.map(block.getArguments(), targetBlock->getArguments());
}
- auto handleOneBlock = [&](Block &block) {
+ auto handleOneBlock = [&](Block &block, bool blockCanUseNowait) {
Block &targetBlock = *rootMapping.lookup(&block);
rootBuilder.setInsertionPointToStart(&targetBlock);
Operation *terminator = block.getTerminator();
@@ -453,14 +517,10 @@ static void parallelizeRegion(Region &sourceRegion, Region &targetRegion,
;
for (auto [i, opOrSingle] : llvm::enumerate(regions)) {
- bool isLast = i + 1 == regions.size();
- // Make sure shared runtime calls are synchronized: disable `nowait`
- // insertion, and rely on the implicit barrier at the end of the
- // omp.workshare block. This applies to any loop-like operation
- // (fir.do_loop, fir.iterate_while, fir.do_concurrent.loop, etc.)
- // because iterations could overlap if nowait is used.
- if (isa<LoopLikeOpInterface>(block.getParentOp()))
- isLast = false;
+ // Only the very last piece of work of the whole omp.workshare region
+ // may use nowait and rely on the barrier emitted at the end of that
+ // region.
+ bool isLast = blockCanUseNowait && i + 1 == regions.size();
if (std::holds_alternative<SingleRegion>(opOrSingle)) {
OpBuilder singleBuilder(sourceRegion.getContext());
Block *singleBlock = new Block();
@@ -485,7 +545,10 @@ static void parallelizeRegion(Region &sourceRegion, Region &targetRegion,
delete singleBlock;
} else {
omp::SingleOperands singleOperands;
- if (isLast)
+ // nowait and copyprivate are mutually exclusive on a single
+ // construct: the broadcast relies on the barrier at the end of the
+ // region.
+ if (isLast && copyprivateVars.empty())
singleOperands.nowait = rootBuilder.getUnitAttr();
singleOperands.copyprivateVars = copyprivateVars;
cleanupBlock(singleBlock);
@@ -519,10 +582,15 @@ static void parallelizeRegion(Region &sourceRegion, Region &targetRegion,
clonedWslw->erase();
} else {
assert(mustParallelizeOp(op));
+ // A loop-like operation may run its region more than once, so the
+ // iterations of the work generated for it could overlap if nowait
+ // were used inside of it.
+ bool nestedCanUseNowait = isLast && !isa<LoopLikeOpInterface>(op);
Operation *cloned = rootBuilder.cloneWithoutRegions(*op, rootMapping);
for (auto [region, clonedRegion] :
llvm::zip(op->getRegions(), cloned->getRegions()))
- parallelizeRegion(region, clonedRegion, rootMapping, loc, di);
+ parallelizeRegion(region, clonedRegion, rootMapping, loc, di,
+ nestedCanUseNowait);
}
}
}
@@ -531,11 +599,13 @@ static void parallelizeRegion(Region &sourceRegion, Region &targetRegion,
};
if (sourceRegion.hasOneBlock()) {
- handleOneBlock(sourceRegion.front());
+ handleOneBlock(sourceRegion.front(), canUseNowait);
} else if (!sourceRegion.empty()) {
+ // With several blocks, no block is known to hold the last piece of work of
+ // the region, so none of them may use nowait.
auto &domTree = di.getDomTree(&sourceRegion);
for (auto node : llvm::breadth_first(domTree.getRootNode())) {
- handleOneBlock(*node->getBlock());
+ handleOneBlock(*node->getBlock(), /*blockCanUseNowait=*/false);
}
}
@@ -595,8 +665,8 @@ LogicalResult lowerWorkshare(mlir::omp::WorkshareOp wsOp, DominanceInfo &di) {
if (!wsOp.getNowait())
omp::BarrierOp::create(rootBuilder, loc);
- parallelizeRegion(wsOp.getRegion(), newOp.getRegion(), rootMapping, loc,
- di);
+ parallelizeRegion(wsOp.getRegion(), newOp.getRegion(), rootMapping, loc, di,
+ /*canUseNowait=*/true);
// Inline the contents of the placeholder workshare op into its parent
// block.
diff --git a/flang/test/Transforms/OpenMP/lower-workshare-nowait.mlir b/flang/test/Transforms/OpenMP/lower-workshare-nowait.mlir
index 940662e0bdccc..ced4f3f12296f 100644
--- a/flang/test/Transforms/OpenMP/lower-workshare-nowait.mlir
+++ b/flang/test/Transforms/OpenMP/lower-workshare-nowait.mlir
@@ -21,3 +21,50 @@ func.func @nowait(%arg0: !fir.ref<!fir.array<42xi32>>) {
}
return
}
+
+// -----
+
+// Check that nowait is not propagated into a region which is nested in
+// something that is not itself the last piece of work of the omp.workshare
+// region, or that may run more than once.
+
+// CHECK-LABEL: func.func @no_nowait_in_nested_conditional
+func.func @no_nowait_in_nested_conditional(%arg0: !fir.ref<i32>, %cond: i1) {
+ omp.parallel {
+ omp.workshare {
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ fir.do_loop %i = %c1 to %c10 step %c1 {
+ fir.if %cond {
+ omp.workshare.loop_wrapper {
+ omp.loop_nest (%j) : index = (%c1) to (%c10) inclusive step (%c1) {
+ "test.inner"(%j) : (index) -> ()
+ omp.yield
+ }
+ }
+ "test.side_effect"(%arg0) : (!fir.ref<i32>) -> ()
+ }
+ }
+ "test.after_loop"(%arg0) : (!fir.ref<i32>) -> ()
+ omp.terminator
+ }
+ omp.terminator
+ }
+ return
+}
+
+// CHECK: fir.do_loop
+// CHECK: fir.if
+// CHECK: omp.wsloop {
+// CHECK-NOT: nowait
+// CHECK: omp.single {
+// CHECK: "test.side_effect"
+// CHECK: omp.terminator
+// CHECK-NEXT: }
+// The work after the loop is the last one, so it may use nowait and rely on
+// the barrier at the end of the omp.workshare region.
+// CHECK: omp.single nowait {
+// CHECK: "test.after_loop"
+// CHECK: omp.terminator
+// CHECK-NEXT: }
+// CHECK-NEXT: omp.barrier
diff --git a/flang/test/Transforms/OpenMP/lower-workshare-thread-local.mlir b/flang/test/Transforms/OpenMP/lower-workshare-thread-local.mlir
index d6000c989515b..546779810ce90 100644
--- a/flang/test/Transforms/OpenMP/lower-workshare-thread-local.mlir
+++ b/flang/test/Transforms/OpenMP/lower-workshare-thread-local.mlir
@@ -317,11 +317,14 @@ func.func @thread_local_load_and_store() {
}
// The store to thread-local memory is parallelized (outside the single),
-// but the load remains inside the single to maintain synchronization.
+// but the load remains inside the single to maintain synchronization. The
+// store which depends on that load can only run on the thread executing the
+// single, so the thread-local memory it updates is broadcast with copyprivate
+// to keep the copies of the other threads in sync.
// CHECK: omp.parallel {
// CHECK-NEXT: %[[ALLOCA:.*]] = fir.alloca i32
-// CHECK: omp.single nowait {
+// CHECK: omp.single copyprivate(%[[ALLOCA]] -> @_workshare_copy_i32 : !fir.ref<i32>) {
// CHECK: fir.store {{.*}} to %[[ALLOCA]] : !fir.ref<i32>
// CHECK: fir.load %[[ALLOCA]] : !fir.ref<i32>
// CHECK: fir.store {{.*}} to %[[ALLOCA]] : !fir.ref<i32>
@@ -403,3 +406,65 @@ func.func @forall_pattern_in_workshare(%shared: !fir.ref<i32>) {
// CHECK: }
// CHECK: omp.barrier
// CHECK: }
+
+// Check the FORALL fetch-counter pattern: a thread-local counter which is
+// read, incremented and written back from inside an omp.single.
+//
+// !$omp workshare
+// forall (i=1:1)
+// forall (j=1:3)
+// a(:,i,j) = a(:,i,j) + 1
+// end forall
+// end forall
+// !$omp end workshare
+//
+// The increment can only be computed on the thread executing the omp.single,
+// so the counter must be broadcast with copyprivate. Otherwise the threads
+// which did not execute the omp.single keep a stale counter and fetch the
+// wrong element on the following iterations. See issue #209942.
+
+// CHECK-LABEL: func.func @forall_fetch_counter_in_workshare
+func.func @forall_fetch_counter_in_workshare(%stack: !fir.ref<i32>) {
+ omp.parallel {
+ %counter = fir.alloca i64 {pinned}
+ omp.workshare {
+ %c0_i64 = arith.constant 0 : i64
+ %c1_i64 = arith.constant 1 : i64
+ %c1 = arith.constant 1 : index
+ %c3 = arith.constant 3 : index
+ fir.store %c0_i64 to %counter : !fir.ref<i64>
+ fir.do_loop %iv = %c1 to %c3 step %c1 {
+ %idx = fir.load %counter : !fir.ref<i64>
+ %next = arith.addi %idx, %c1_i64 : i64
+ fir.store %next to %counter : !fir.ref<i64>
+ "test.fetch"(%stack, %idx) : (!fir.ref<i32>, i64) -> ()
+ omp.workshare.loop_wrapper {
+ omp.loop_nest (%j) : index = (%c1) to (%c3) inclusive step (%c1) {
+ "test.inner"(%j) : (index) -> ()
+ omp.yield
+ }
+ }
+ }
+ omp.terminator
+ }
+ omp.terminator
+ }
+ return
+}
+
+// CHECK: omp.parallel {
+// CHECK: %[[COUNTER:.*]] = fir.alloca i64 {pinned}
+// The reset of the counter is a write to thread-local memory whose operands
+// are all available, so it is parallelized and all threads run it.
+// CHECK: fir.store %{{.*}} to %[[COUNTER]] : !fir.ref<i64>
+// CHECK: fir.do_loop
+// CHECK: omp.single copyprivate(%[[COUNTER]] -> @_workshare_copy_i64 : !fir.ref<i64>) {
+// CHECK: %[[IDX:.*]] = fir.load %[[COUNTER]] : !fir.ref<i64>
+// CHECK: %[[NEXT:.*]] = arith.addi %[[IDX]], %{{.*}} : i64
+// CHECK: fir.store %[[NEXT]] to %[[COUNTER]] : !fir.ref<i64>
+// CHECK: "test.fetch"
+// CHECK: omp.terminator
+// CHECK-NEXT: }
+// The increment must not be repeated outside the single.
+// CHECK-NOT: fir.store {{.*}} to %[[COUNTER]]
+// CHECK: omp.wsloop {
>From 791ab3c1e98c1c3a11c7a85d96fa0ac55477e866 Mon Sep 17 00:00:00 2001
From: Carlos Seo <carlos.seo at linaro.org>
Date: Fri, 24 Jul 2026 11:45:57 -0300
Subject: [PATCH 2/3] [flang][OpenMP] Only broadcast thread-local state that is
read back
Refine the previous fix so that a thread-local location written from
within an omp.single is broadcasted with copyprivate only when some other
thread may actually read it back. A location that is written but never
read (for example a temporary updated in a terminal omp.single) does not
need to be broadcasted.
Collect, once per omp.workshare region, the thread-local allocations that
are read by the whole team and keep only those in the copyprivate list.
The enclosing omp.parallel is used as the scope of the read search so
that reads performed after the omp.workshare region are accounted for as
well; reads are matched by the identity of the alloca, mirroring the
write collection, so that only a direct load or store of the whole
thread-local location keeps it live for broadcast.
The FORALL fetch counter is still broadcasted, since it is read from
within the omp.single, and a new test checks that a write-only
thread-local location is left with a plain omp.single instead.
---
flang/lib/Optimizer/OpenMP/LowerWorkshare.cpp | 63 ++++++++++++++++---
.../OpenMP/lower-workshare-thread-local.mlir | 63 +++++++++++++++++++
2 files changed, 119 insertions(+), 7 deletions(-)
diff --git a/flang/lib/Optimizer/OpenMP/LowerWorkshare.cpp b/flang/lib/Optimizer/OpenMP/LowerWorkshare.cpp
index 8f50b10a07917..b545560c49134 100644
--- a/flang/lib/Optimizer/OpenMP/LowerWorkshare.cpp
+++ b/flang/lib/Optimizer/OpenMP/LowerWorkshare.cpp
@@ -301,6 +301,35 @@ static void collectThreadLocalWrites(Operation *op,
}
}
+// Collects into reads the thread-local allocations that are read anywhere in
+// scope. A thread-local location written from within an omp.single only needs
+// to be broadcasted if some other thread may later read it. The scope must be
+// a region executed by the whole team (i.e. the enclosing omp.parallel), so
+// that reads performed after the omp.workshare region are accounted for too.
+//
+// Reads are matched by the identity of the alloca, mirroring
+// collectThreadLocalWrites, so that only a direct load/store of the whole
+// thread-local location keeps it live for broadcasting.
+static void collectThreadLocalReads(Region &scope,
+ llvm::SmallDenseSet<Value> &reads) {
+ scope.walk([&](Operation *op) {
+ auto memEffects = dyn_cast<MemoryEffectOpInterface>(op);
+ if (!memEffects)
+ return;
+ SmallVector<MemoryEffects::EffectInstance> effects;
+ memEffects.getEffects(effects);
+ for (const MemoryEffects::EffectInstance &effect : effects) {
+ if (!isa<MemoryEffects::Read>(effect.getEffect()))
+ continue;
+ Value val = effect.getValue();
+ if (!val || !val.getDefiningOp<fir::AllocaOp>())
+ continue;
+ if (isOpenMPThreadLocalMemory(op, val))
+ reads.insert(val);
+ }
+ });
+}
+
/// Simple shallow copies suffice for our purposes in this pass, so we implement
/// this simpler alternative to the full fledged `createCopyFunc` in the
/// frontend
@@ -385,9 +414,11 @@ static void cleanupBlock(Block *block) {
// very last thing the omp.workshare region does, and thus whether the
// synchronization of its last omp.single/omp.wsloop may be left to the
// barrier emitted at the end of the omp.workshare region.
-static void parallelizeRegion(Region &sourceRegion, Region &targetRegion,
- IRMapping &rootMapping, Location loc,
- mlir::DominanceInfo &di, bool canUseNowait) {
+static void
+parallelizeRegion(Region &sourceRegion, Region &targetRegion,
+ IRMapping &rootMapping, Location loc, mlir::DominanceInfo &di,
+ bool canUseNowait,
+ const llvm::SmallDenseSet<Value> &threadLocalReads) {
OpBuilder rootBuilder(sourceRegion.getContext());
ModuleOp m = sourceRegion.getParentOfType<ModuleOp>();
OpBuilder copyFuncBuilder(m.getBodyRegion());
@@ -468,10 +499,15 @@ static void parallelizeRegion(Region &sourceRegion, Region &targetRegion,
omp::TerminatorOp::create(singleBuilder, loc);
// Broadcast the thread-local state which only the thread executing the
- // omp.single has updated. Values defined inside sr are remapped; values
- // defined before it (e.g. hoisted allocas) are used as is.
+ // omp.single has updated, but only when some other thread may actually read
+ // it back: a location that is never read (e.g. a write to a temporary in a
+ // terminal omp.single) does not need to be broadcasted. Values defined
+ // inside sr are remapped; values defined before it (e.g. hoisted allocas)
+ // are used as is.
llvm::SmallDenseSet<Value> seen(copyPrivate.begin(), copyPrivate.end());
for (Value v : threadLocalWrites) {
+ if (!threadLocalReads.contains(v))
+ continue;
Value mapped = singleMapping.lookupOrDefault(v);
if (seen.insert(mapped).second)
copyPrivate.push_back(mapped);
@@ -590,7 +626,7 @@ static void parallelizeRegion(Region &sourceRegion, Region &targetRegion,
for (auto [region, clonedRegion] :
llvm::zip(op->getRegions(), cloned->getRegions()))
parallelizeRegion(region, clonedRegion, rootMapping, loc, di,
- nestedCanUseNowait);
+ nestedCanUseNowait, threadLocalReads);
}
}
}
@@ -665,8 +701,21 @@ LogicalResult lowerWorkshare(mlir::omp::WorkshareOp wsOp, DominanceInfo &di) {
if (!wsOp.getNowait())
omp::BarrierOp::create(rootBuilder, loc);
+ // Compute the thread-local locations read by the whole team, so that only
+ // those get broadcasted out of the omp.single's below. The enclosing
+ // omp.parallel is used as the scope so that reads performed after the
+ // omp.workshare region are taken into account as well; if there is none,
+ // fall back to the innermost isolated-from-above ancestor.
+ llvm::SmallDenseSet<Value> threadLocalReads;
+ if (auto parallelOp = wsOp->getParentOfType<omp::ParallelOp>())
+ collectThreadLocalReads(parallelOp.getRegion(), threadLocalReads);
+ else if (Operation *top =
+ wsOp->getParentWithTrait<OpTrait::IsIsolatedFromAbove>())
+ for (Region &r : top->getRegions())
+ collectThreadLocalReads(r, threadLocalReads);
+
parallelizeRegion(wsOp.getRegion(), newOp.getRegion(), rootMapping, loc, di,
- /*canUseNowait=*/true);
+ /*canUseNowait=*/true, threadLocalReads);
// Inline the contents of the placeholder workshare op into its parent
// block.
diff --git a/flang/test/Transforms/OpenMP/lower-workshare-thread-local.mlir b/flang/test/Transforms/OpenMP/lower-workshare-thread-local.mlir
index 546779810ce90..c3ec65a76bc8d 100644
--- a/flang/test/Transforms/OpenMP/lower-workshare-thread-local.mlir
+++ b/flang/test/Transforms/OpenMP/lower-workshare-thread-local.mlir
@@ -468,3 +468,66 @@ func.func @forall_fetch_counter_in_workshare(%stack: !fir.ref<i32>) {
// The increment must not be repeated outside the single.
// CHECK-NOT: fir.store {{.*}} to %[[COUNTER]]
// CHECK: omp.wsloop {
+
+
+// -----
+
+// Check that a thread-local location written from within an omp.single but
+// never read back by the team is NOT broadcasted with copyprivate. The store
+// is not safe to parallelize on its own here because its value comes from a
+// shared load that must stay in the omp.single, so it ends up executed by a
+// single thread.
+
+// CHECK-LABEL: func.func @write_only_thread_local_not_broadcast
+func.func @write_only_thread_local_not_broadcast(%shared: !fir.ref<i32>) {
+ omp.parallel {
+ %tl = fir.alloca i32
+ omp.workshare {
+ %v = fir.load %shared : !fir.ref<i32>
+ fir.store %v to %tl : !fir.ref<i32>
+ omp.terminator
+ }
+ omp.terminator
+ }
+ return
+}
+
+// CHECK: omp.parallel {
+// CHECK-NEXT: %[[TL:.*]] = fir.alloca i32
+// The single carries no copyprivate: %[[TL]] is never read by the team.
+// CHECK: omp.single nowait {
+// CHECK-NOT: copyprivate
+// CHECK: fir.store %{{.*}} to %[[TL]] : !fir.ref<i32>
+// CHECK: omp.terminator
+// CHECK-NEXT: }
+// CHECK-NEXT: omp.barrier
+
+// -----
+
+// Same write from within an omp.single, but now the location is read back by
+// the whole team after the omp.workshare region: it must be broadcasted so the
+// threads which did not run the single do not observe a stale value.
+
+// CHECK-LABEL: func.func @written_then_read_thread_local_is_broadcast
+func.func @written_then_read_thread_local_is_broadcast(%shared: !fir.ref<i32>, %sink: !fir.ref<i32>) {
+ omp.parallel {
+ %tl = fir.alloca i32
+ omp.workshare {
+ %v = fir.load %shared : !fir.ref<i32>
+ fir.store %v to %tl : !fir.ref<i32>
+ omp.terminator
+ }
+ %r = fir.load %tl : !fir.ref<i32>
+ fir.store %r to %sink : !fir.ref<i32>
+ omp.terminator
+ }
+ return
+}
+
+// CHECK: omp.parallel {
+// CHECK-NEXT: %[[TL:.*]] = fir.alloca i32
+// CHECK: omp.single copyprivate(%[[TL]] -> @_workshare_copy_i32 : !fir.ref<i32>) {
+// CHECK: fir.store %{{.*}} to %[[TL]] : !fir.ref<i32>
+// CHECK: omp.terminator
+// CHECK-NEXT: }
+// CHECK: fir.load %[[TL]] : !fir.ref<i32>
>From f7ed6e942b9a598bd937e3b9a926f58b174a1be1 Mon Sep 17 00:00:00 2001
From: Carlos Seo <carlos.seo at linaro.org>
Date: Fri, 24 Jul 2026 12:13:50 -0300
Subject: [PATCH 3/3] [flang][OpenMP] Match thread-local reads and writes
through declares
The broadcast of thread-local state out of an omp.single matched the
memory locations by the identity of the value the load or store operates
on, requiring it to be the fir.alloca directly. Flang lowering routinely
inserts fir.declare/fir.convert between an allocation and its uses, so a
store to an alloca and a load from a fir.declare of that same alloca were
treated as unrelated. As a result a required broadcast could be dropped:
either the write was not collected, or the read that keeps it live was
not seen and the write was filtered out. A thread which did not execute
the omp.single would then observe a stale value.
Match reads and writes by their underlying thread-local allocation
instead. getOpenMPThreadLocalSource uses the alias analysis, which
already looks through declares, converts and reboxes, so that both sides
canonicalize to the same value. The broadcast still references the
underlying allocation, which is what the shallow copy function expects.
---
flang/lib/Optimizer/OpenMP/LowerWorkshare.cpp | 42 ++++++++++----
.../OpenMP/lower-workshare-thread-local.mlir | 55 +++++++++++++++++++
2 files changed, 85 insertions(+), 12 deletions(-)
diff --git a/flang/lib/Optimizer/OpenMP/LowerWorkshare.cpp b/flang/lib/Optimizer/OpenMP/LowerWorkshare.cpp
index b545560c49134..a1a7eb551aeb5 100644
--- a/flang/lib/Optimizer/OpenMP/LowerWorkshare.cpp
+++ b/flang/lib/Optimizer/OpenMP/LowerWorkshare.cpp
@@ -259,6 +259,23 @@ static bool isSafeToParallelize(Operation *op) {
return false;
}
+// Returns the underlying thread-local storage that mem refers to, or null if
+// mem is not thread-local. The alias analysis is used to look through
+// fir.declare/hlfir.declare, fir.convert, fir.rebox, etc., so that two
+// accesses of the same thread-local location yield the same value even if one
+// goes through such ops and the other does not. This is what makes it safe to
+// match the reads and writes of collect{Reads,Writes} against each other by
+// value identity: a store to an alloca and a load from a fir.declare of that
+// alloca map to the same key. Matching the raw effect value instead would
+// silently miss such accesses, dropping a required broadcast.
+static Value getOpenMPThreadLocalSource(Operation *op, Value mem) {
+ if (!isOpenMPThreadLocalMemory(op, mem))
+ return nullptr;
+ fir::AliasAnalysis aliasAnalysis;
+ return llvm::dyn_cast_if_present<mlir::Value>(
+ aliasAnalysis.getSource(mem).origin.u);
+}
+
// Collects the thread-local memory locations that op writes to and that
// need to be broadcasted to other threads when op ends up being executed
// by a single thread only.
@@ -272,7 +289,7 @@ static bool isSafeToParallelize(Operation *op) {
// which did not execute the omp.single would otherwise go stale and the
// following iterations would fetch the wrong element. See issue #209942.
//
-// Only the thread-local allocation itself is considered, so that a shallow
+// Only the underlying thread-local allocation is considered, so that a shallow
// copy of it faithfully reproduces the update on the other threads.
static void collectThreadLocalWrites(Operation *op,
llvm::SmallVectorImpl<Value> &vars) {
@@ -285,9 +302,12 @@ static void collectThreadLocalWrites(Operation *op,
if (!isa<MemoryEffects::Write>(effect.getEffect()))
continue;
Value val = effect.getValue();
- if (!val || !val.getDefiningOp<fir::AllocaOp>())
+ if (!val)
continue;
- auto refTy = dyn_cast<fir::ReferenceType>(val.getType());
+ Value source = getOpenMPThreadLocalSource(op, val);
+ if (!source)
+ continue;
+ auto refTy = dyn_cast<fir::ReferenceType>(source.getType());
if (!refTy)
continue;
// createCopyFunc emits a load/store pair, so restrict this to types for
@@ -295,9 +315,7 @@ static void collectThreadLocalWrites(Operation *op,
mlir::Type eleTy = refTy.getEleTy();
if (!fir::isa_trivial(eleTy) && !fir::isa_box_type(eleTy))
continue;
- if (!isOpenMPThreadLocalMemory(op, val))
- continue;
- vars.push_back(val);
+ vars.push_back(source);
}
}
@@ -307,9 +325,9 @@ static void collectThreadLocalWrites(Operation *op,
// a region executed by the whole team (i.e. the enclosing omp.parallel), so
// that reads performed after the omp.workshare region are accounted for too.
//
-// Reads are matched by the identity of the alloca, mirroring
-// collectThreadLocalWrites, so that only a direct load/store of the whole
-// thread-local location keeps it live for broadcasting.
+// Reads are matched by their underlying thread-local allocation, mirroring
+// collectThreadLocalWrites, so that a load through a fir.declare/fir.convert
+// still keeps the corresponding write live for broadcasting.
static void collectThreadLocalReads(Region &scope,
llvm::SmallDenseSet<Value> &reads) {
scope.walk([&](Operation *op) {
@@ -322,10 +340,10 @@ static void collectThreadLocalReads(Region &scope,
if (!isa<MemoryEffects::Read>(effect.getEffect()))
continue;
Value val = effect.getValue();
- if (!val || !val.getDefiningOp<fir::AllocaOp>())
+ if (!val)
continue;
- if (isOpenMPThreadLocalMemory(op, val))
- reads.insert(val);
+ if (Value source = getOpenMPThreadLocalSource(op, val))
+ reads.insert(source);
}
});
}
diff --git a/flang/test/Transforms/OpenMP/lower-workshare-thread-local.mlir b/flang/test/Transforms/OpenMP/lower-workshare-thread-local.mlir
index c3ec65a76bc8d..2319e9b81d580 100644
--- a/flang/test/Transforms/OpenMP/lower-workshare-thread-local.mlir
+++ b/flang/test/Transforms/OpenMP/lower-workshare-thread-local.mlir
@@ -531,3 +531,58 @@ func.func @written_then_read_thread_local_is_broadcast(%shared: !fir.ref<i32>, %
// CHECK: omp.terminator
// CHECK-NEXT: }
// CHECK: fir.load %[[TL]] : !fir.ref<i32>
+
+
+// -----
+
+// Check that the read which keeps a thread-local write live for broadcasting
+// is still recognized when it goes through a fir.declare of the allocation
+// (looking through declares/converts, as flang lowering routinely inserts
+// them). Matching the raw load/store value instead would miss this read and
+// drop the required broadcast.
+
+// CHECK-LABEL: func.func @broadcast_when_read_through_declare
+func.func @broadcast_when_read_through_declare(%shared: !fir.ref<i32>, %sink: !fir.ref<i32>) {
+ omp.parallel {
+ %tl = fir.alloca i32
+ %d = fir.declare %tl {uniq_name = "tl"} : (!fir.ref<i32>) -> !fir.ref<i32>
+ omp.workshare {
+ %v = fir.load %shared : !fir.ref<i32>
+ fir.store %v to %tl : !fir.ref<i32>
+ omp.terminator
+ }
+ %r = fir.load %d : !fir.ref<i32>
+ fir.store %r to %sink : !fir.ref<i32>
+ omp.terminator
+ }
+ return
+}
+
+// CHECK: %[[TL:.*]] = fir.alloca i32
+// The broadcast copies the underlying allocation, not the fir.declare.
+// CHECK: omp.single copyprivate(%[[TL]] -> @_workshare_copy_i32 : !fir.ref<i32>) {
+
+// -----
+
+// Same, but now the write from within the omp.single goes through a
+// fir.declare of the allocation while the read is direct.
+
+// CHECK-LABEL: func.func @broadcast_when_written_through_declare
+func.func @broadcast_when_written_through_declare(%shared: !fir.ref<i32>, %sink: !fir.ref<i32>) {
+ omp.parallel {
+ %tl = fir.alloca i32
+ %d = fir.declare %tl {uniq_name = "tl"} : (!fir.ref<i32>) -> !fir.ref<i32>
+ omp.workshare {
+ %v = fir.load %shared : !fir.ref<i32>
+ fir.store %v to %d : !fir.ref<i32>
+ omp.terminator
+ }
+ %r = fir.load %tl : !fir.ref<i32>
+ fir.store %r to %sink : !fir.ref<i32>
+ omp.terminator
+ }
+ return
+}
+
+// CHECK: %[[TL:.*]] = fir.alloca i32
+// CHECK: omp.single copyprivate(%[[TL]] -> @_workshare_copy_i32 : !fir.ref<i32>) {
More information about the flang-commits
mailing list