[Openmp-commits] [openmp] [OpenMP][offload] Fix wave32 handling of xteam reductions (PR #233096)

[Robert Imschweiler via Openmp-commits]

https://github.com/ro-i updated https://github.com/llvm/llvm-project/pull/233096

>From 5b782e91083068605922a90b60f2511333625951 Mon Sep 17 00:00:00 2001
From: Robert Imschweiler <robert.imschweiler at amd.com>
Date: Sat, 17 Oct 2026 09:54:05 -0500
Subject: [PATCH 1/2] [OpenMP][offload] Fix wave32 handling of xteam reductions

---
 openmp/device/src/Reduction.cpp | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/openmp/device/src/Reduction.cpp b/openmp/device/src/Reduction.cpp
index 0b7f687d6a83e..c0655501a9201 100644
--- a/openmp/device/src/Reduction.cpp
+++ b/openmp/device/src/Reduction.cpp
@@ -101,11 +101,14 @@ static uint32_t gpu_block_reduce(void *reduce_data, ShuffleReductFnTy shflFct,
   if constexpr (checkLiveness) {
     __kmpc_impl_lanemask_t Liveness = mapping::activemask();
     // Check for partial warp with non-contiguous lanes.
-    if (Liveness == lanes::All) {
+    // Consider that Liveness::ALL is a 64-bit value, so compare the actual
+    // number of set bits to handle 32bit waves correctly.
+    uint32_t NumLive = static_cast<uint32_t>(utils::popc(Liveness));
+    if (NumLive == mapping::getWarpSize()) {
       gpu_regular_warp_reduce(reduce_data, shflFct);
     } else if (!(Liveness & (Liveness + 1))) {
       // Partial warp but contiguous lanes.
-      gpu_irregular_warp_reduce(reduce_data, shflFct, utils::popc(Liveness),
+      gpu_irregular_warp_reduce(reduce_data, shflFct, NumLive,
                                 BlockThreadId % mapping::getWarpSize());
     } else {
       // Dispersed lanes. Only threads in L2 parallel region may enter here.

>From 2cd8489f1f16a5cc6617d3fb88bbdd22d43f68bd Mon Sep 17 00:00:00 2001
From: Robert Imschweiler <robert.imschweiler at amd.com>
Date: Tue, 20 Oct 2026 03:07:04 -0500
Subject: [PATCH 2/2] implement better solution

Co-authored-by: Joseph Huber <huberjn at outlook.com>
---
 openmp/device/src/Reduction.cpp | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/openmp/device/src/Reduction.cpp b/openmp/device/src/Reduction.cpp
index c0655501a9201..c20f1d10344d0 100644
--- a/openmp/device/src/Reduction.cpp
+++ b/openmp/device/src/Reduction.cpp
@@ -103,12 +103,12 @@ static uint32_t gpu_block_reduce(void *reduce_data, ShuffleReductFnTy shflFct,
     // Check for partial warp with non-contiguous lanes.
     // Consider that Liveness::ALL is a 64-bit value, so compare the actual
     // number of set bits to handle 32bit waves correctly.
-    uint32_t NumLive = static_cast<uint32_t>(utils::popc(Liveness));
-    if (NumLive == mapping::getWarpSize()) {
+    const uint64_t Full = ~uint64_t(0) >> (64 - mapping::getWarpSize());
+    if (Liveness == Full) {
       gpu_regular_warp_reduce(reduce_data, shflFct);
     } else if (!(Liveness & (Liveness + 1))) {
       // Partial warp but contiguous lanes.
-      gpu_irregular_warp_reduce(reduce_data, shflFct, NumLive,
+      gpu_irregular_warp_reduce(reduce_data, shflFct, utils::popc(Liveness),
                                 BlockThreadId % mapping::getWarpSize());
     } else {
       // Dispersed lanes. Only threads in L2 parallel region may enter here.