[Flang][OpenMP] Access full list of entry block syms and vars (NFC)

[skatrak]

This patch adds methods to EntryBlockArgs to access the full list of entry block argument-related symbols and variables, in their standard order. This helps centralizing this logic in as few places as possible to avoid future inconsistencies.

[llvmbot]

@llvm/pr-subscribers-flang-fir-hlfir

Author: Sergio Afonso (skatrak)

Changes

This patch adds methods to EntryBlockArgs to access the full list of entry block argument-related symbols and variables, in their standard order. This helps centralizing this logic in as few places as possible to avoid future inconsistencies.

---

Full diff: https://github.com/llvm/llvm-project/pull/113681.diff

1 Files Affected:

• (modified) flang/lib/Lower/OpenMP/OpenMP.cpp (+16-5)

diff --git a/flang/lib/Lower/OpenMP/OpenMP.cpp b/flang/lib/Lower/OpenMP/OpenMP.cpp
index fc54da8babe63e..e2545a68241004 100644
--- a/flang/lib/Lower/OpenMP/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP/OpenMP.cpp
@@ -76,6 +76,18 @@ struct EntryBlockArgs {
            reduction.isValid() && taskReduction.isValid() &&
            useDeviceAddr.isValid() && useDevicePtr.isValid();
   }
+
+  auto getSyms() const {
+    return llvm::concat<const semantics::Symbol *const>(
+        inReduction.syms, map.syms, priv.syms, reduction.syms,
+        taskReduction.syms, useDeviceAddr.syms, useDevicePtr.syms);
+  }
+
+  auto getVars() const {
+    return llvm::concat<const mlir::Value>(
+        inReduction.vars, map.vars, priv.vars, reduction.vars,
+        taskReduction.vars, useDeviceAddr.vars, useDevicePtr.vars);
+  }
 };
 } // namespace
 
@@ -1506,8 +1518,7 @@ genParallelOp(lower::AbstractConverter &converter, lower::SymMap &symTable,
     genEntryBlock(converter, args, op->getRegion(0));
     bindEntryBlockArgs(
         converter, llvm::cast<mlir::omp::BlockArgOpenMPOpInterface>(op), args);
-    return llvm::to_vector(llvm::concat<const semantics::Symbol *const>(
-        args.priv.syms, args.reduction.syms));
+    return llvm::to_vector(args.getSyms());
   };
 
   assert((!enableDelayedPrivatization || dsp) &&
@@ -1581,11 +1592,11 @@ genSectionsOp(lower::AbstractConverter &converter, lower::SymMap &symTable,
   mlir::Operation *terminator =
       lower::genOpenMPTerminator(builder, sectionsOp, loc);
 
-  auto reductionCallback = [&](mlir::Operation *op) {
+  auto genRegionEntryCB = [&](mlir::Operation *op) {
     genEntryBlock(converter, args, op->getRegion(0));
     bindEntryBlockArgs(
         converter, llvm::cast<mlir::omp::BlockArgOpenMPOpInterface>(op), args);
-    return reductionSyms;
+    return llvm::to_vector(args.getSyms());
   };
 
   // Generate nested SECTION constructs.
@@ -1611,7 +1622,7 @@ genSectionsOp(lower::AbstractConverter &converter, lower::SymMap &symTable,
         OpWithBodyGenInfo(converter, symTable, semaCtx, loc, nestedEval,
                           llvm::omp::Directive::OMPD_section)
             .setClauses(&sectionQueue.begin()->clauses)
-            .setGenRegionEntryCb(reductionCallback),
+            .setGenRegionEntryCb(genRegionEntryCB),
         sectionQueue, sectionQueue.begin());
   }
 

[skatrak]

PR stack:

• [MLIR][OpenMP] Prevent composite omp.simd related crashes #113680
• [Flang][OpenMP] Access full list of entry block syms and vars (NFC) #113681
• [Flang][OpenMP] Properly bind arguments of composite operations #113682
• Revert "[Flang][OpenMP] Disable lowering of omp.simd reductions in co… #113683

[tblah]

LGTM, thanks!

Maybe sometime we could go a step further and provide a standard implementation of the genRegionEntryCB?

[skatrak]

Maybe sometime we could go a step further and provide a standard implementation of the genRegionEntryCB?

I agree, that would be nice to have. At this point, the only operation for which genOpWithBody is called with a callback that is different is omp.loop_nest. We may be able to provide a more elaborate default for the others, based on a pointer to the EntryBlockArgs structure added to the OpWithBodyGenInfo structure. I'll give that idea a try hopefully soon.