[flang-commits] [flang] [flang][OpenMP] Fix wrong results for FORALL in a workshare construct (PR #211371)

[Tom Eccles via flang-commits]

================
@@ -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();
----------------
tblah wrote:

fir::unwrapRefType

https://github.com/llvm/llvm-project/pull/211371