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

[Tom Eccles via flang-commits]

================
@@ -259,6 +259,95 @@ 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.
+//
+// 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 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) {
+  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)
+      continue;
+    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
+    // 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;
+    vars.push_back(source);
+  }
+}
+
+// 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 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) {
+    auto memEffects = dyn_cast<MemoryEffectOpInterface>(op);
+    if (!memEffects)
+      return;
----------------
tblah wrote:

Codex pointed out that it actually does matter if we find operations with unknown memory effects when collecting reads. This could happen for example for opaque function calls.

In that case there might be a read in the function all so we need a way to express that *all* values might have been read in this case.

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