[Openmp-commits] [clang] [llvm] [mlir] [openmp] [OpenMP][offload] Cross-team reductions with variable number of teams (PR #195102)

[Johannes Doerfert via Openmp-commits]

================
@@ -143,11 +159,16 @@ GenericKernelTy::getKernelLaunchEnvironment(
   LocalKLE.DynCGroupMemFb = DynBlockMemConf.Fallback;
   LocalKLE.ReductionBuffer = nullptr;
 
-  if (KernelEnvironment.Configuration.ReductionDataSize &&
-      KernelEnvironment.Configuration.ReductionBufferLength) {
+  if (NeedsReductionBuffer) {
+    // Size the teams-reduction buffer. ReductionBufferLength is the
+    // statically known upper bound on the number of teams. When it is 0
+    // (i.e. num_teams was not a compile-time constant), we size the buffer
+    // to exactly match the number of teams for this launch (NumBlocks0).
+    uint32_t BufferElements = RedCfg.ReductionBufferLength
+                                  ? RedCfg.ReductionBufferLength
+                                  : NumBlocks0;
----------------
jdoerfert wrote:

You can grab KernelArgs.NumTeams[0] here, no? Later we need to look into making this work for >1 dim too.

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