[Openmp-commits] [llvm] [openmp] [OpenMP][offload] Fix single-thread SPMD cross-team reductions (PR #231843)
via Openmp-commits
openmp-commits at lists.llvm.org
Wed Oct 14 02:54:38 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-offload
Author: Robert Imschweiler (ro-i)
<details>
<summary>Changes</summary>
If a kernel gets SPMDized by OpenMPOpt but does not have a parallel region, it actually is only executed by a single thread. forceSingleThreadPerWorkgroupHelper() is responsible for that and had been introduced by https://reviews.llvm.org/D133120. The cross-team reduction trusted that SPMD mode means that we're dealing with all threads. Now, OpenMPOpt also limits such kernels to a single thread per team via the kernel environment, so that this assumption holds.
Claude assisted with this patch.
---
Full diff: https://github.com/llvm/llvm-project/pull/231843.diff
5 Files Affected:
- (modified) llvm/lib/Transforms/IPO/OpenMPOpt.cpp (+7-3)
- (modified) llvm/test/Transforms/OpenMP/always_inline_device.ll (+1-1)
- (modified) llvm/test/Transforms/OpenMP/get_hardware_num_threads_in_block_fold.ll (+2-2)
- (added) offload/test/offloading/xteam_reduction_no_parallel.c (+46)
- (modified) openmp/device/src/Reduction.cpp (+10-3)
``````````diff
diff --git a/llvm/lib/Transforms/IPO/OpenMPOpt.cpp b/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
index 019d649b8d5ff0..119b32aee2e8ef 100644
--- a/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
+++ b/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
@@ -4435,11 +4435,15 @@ struct AAKernelInfoFunction : AAKernelInfo {
Changed = ChangeStatus::CHANGED;
// Do not use instruction guards when no parallel is present inside
- // the target region.
- if (mayContainParallelRegion())
+ // the target region. Instead, only the main thread of each team executes
+ // the user code, so there is no need to launch more threads.
+ if (mayContainParallelRegion()) {
insertInstructionGuardsHelper(A);
- else
+ } else {
forceSingleThreadPerWorkgroupHelper(A);
+ setMaxThreadsOfKernelEnvironment(
+ ConstantInt::get(Type::getInt32Ty(Kernel->getContext()), 1));
+ }
// Adjust the global exec mode flag that tells the runtime what mode this
// kernel is executed in.
diff --git a/llvm/test/Transforms/OpenMP/always_inline_device.ll b/llvm/test/Transforms/OpenMP/always_inline_device.ll
index 9c5b19f7a6c88c..4d8edf7f1fdc14 100644
--- a/llvm/test/Transforms/OpenMP/always_inline_device.ll
+++ b/llvm/test/Transforms/OpenMP/always_inline_device.ll
@@ -15,7 +15,7 @@
; CHECK: @[[GLOB0:[0-9]+]] = private unnamed_addr constant [23 x i8] c"
; CHECK: @[[GLOB1:[0-9]+]] = private unnamed_addr constant %struct.ident_t { i32 0, i32 2, i32 0, i32 0, ptr @[[GLOB0]] }, align 8
; CHECK: @G = external global i8
-; CHECK: @kernel_environment = local_unnamed_addr constant %struct.KernelEnvironmentTy { %struct.ConfigurationEnvironmentTy { i8 0, i8 0, i8 3, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0 }, ptr @[[GLOB1]], ptr null }
+; CHECK: @kernel_environment = local_unnamed_addr constant %struct.KernelEnvironmentTy { %struct.ConfigurationEnvironmentTy { i8 0, i8 0, i8 3, i32 0, i32 1, i32 0, i32 0, i32 0, i32 0 }, ptr @[[GLOB1]], ptr null }
;.
define weak ptx_kernel void @__omp_offloading_fd02_c0934fc2_foo_l4(ptr %dyn) #0 {
; CHECK: Function Attrs: norecurse nounwind
diff --git a/llvm/test/Transforms/OpenMP/get_hardware_num_threads_in_block_fold.ll b/llvm/test/Transforms/OpenMP/get_hardware_num_threads_in_block_fold.ll
index 1e37431675ec75..71f548b8905c4a 100644
--- a/llvm/test/Transforms/OpenMP/get_hardware_num_threads_in_block_fold.ll
+++ b/llvm/test/Transforms/OpenMP/get_hardware_num_threads_in_block_fold.ll
@@ -13,8 +13,8 @@ target triple = "nvptx64"
;.
; CHECK: @G = external global i32
-; CHECK: @kernel0_kernel_environment = local_unnamed_addr constant %struct.KernelEnvironmentTy { %struct.ConfigurationEnvironmentTy { i8 0, i8 0, i8 3, i32 0, i32 666, i32 0, i32 777, i32 0, i32 0 }, ptr null, ptr null }
-; CHECK: @kernel1_kernel_environment = local_unnamed_addr constant %struct.KernelEnvironmentTy { %struct.ConfigurationEnvironmentTy { i8 0, i8 0, i8 3, i32 0, i32 666, i32 0, i32 777, i32 0, i32 0 }, ptr null, ptr null }
+; CHECK: @kernel0_kernel_environment = local_unnamed_addr constant %struct.KernelEnvironmentTy { %struct.ConfigurationEnvironmentTy { i8 0, i8 0, i8 3, i32 0, i32 1, i32 0, i32 777, i32 0, i32 0 }, ptr null, ptr null }
+; CHECK: @kernel1_kernel_environment = local_unnamed_addr constant %struct.KernelEnvironmentTy { %struct.ConfigurationEnvironmentTy { i8 0, i8 0, i8 3, i32 0, i32 1, i32 0, i32 777, i32 0, i32 0 }, ptr null, ptr null }
; CHECK: @kernel2_kernel_environment = local_unnamed_addr constant %struct.KernelEnvironmentTy { %struct.ConfigurationEnvironmentTy { i8 0, i8 0, i8 3, i32 0, i32 666, i32 0, i32 777, i32 0, i32 0 }, ptr null, ptr null }
; CHECK: @[[GLOB0:[0-9]+]] = private unnamed_addr constant [23 x i8] c"
; CHECK: @[[GLOB1:[0-9]+]] = private unnamed_addr constant %struct.ident_t { i32 0, i32 2, i32 0, i32 22, ptr @[[GLOB0]] }, align 8
diff --git a/offload/test/offloading/xteam_reduction_no_parallel.c b/offload/test/offloading/xteam_reduction_no_parallel.c
new file mode 100644
index 00000000000000..d0a2b1223c35da
--- /dev/null
+++ b/offload/test/offloading/xteam_reduction_no_parallel.c
@@ -0,0 +1,46 @@
+// RUN: %libomptarget-compile-run-and-check-generic
+// RUN: %libomptarget-compileopt-run-and-check-generic
+
+// Cross-team reductions in kernels without a parallel region. With
+// optimizations, such kernels are executed with only one thread per team
+// ("fake"-SPMD mode), which the cross-team reduction has to take into account.
+
+#include <stdio.h>
+
+#define N (1 << 20)
+
+int Y[N];
+
+#pragma omp declare target
+int getY(int *A, int I) { return A[I]; }
+#pragma omp end declare target
+
+int main(void) {
+ for (int I = 0; I < N; ++I)
+ Y[I] = 1;
+
+ long Sum = 0;
+#pragma omp target teams distribute reduction(+ : Sum) map(to : Y)
+ for (int I = 0; I < N; ++I)
+ Sum += Y[I];
+ // CHECK: teams distribute: 1048576
+ printf("teams distribute: %ld\n", Sum);
+
+ // The call to getY prevents emitting 'loop' as 'distribute parallel for'.
+ Sum = 0;
+#pragma omp target teams loop reduction(+ : Sum) map(to : Y)
+ for (int I = 0; I < N; ++I)
+ Sum += getY(Y, I);
+ // CHECK: teams loop: 1048576
+ printf("teams loop: %ld\n", Sum);
+
+ Sum = 0;
+#pragma omp target map(to : Y) map(tofrom : Sum)
+#pragma omp teams distribute reduction(+ : Sum)
+ for (int I = 0; I < N; ++I)
+ Sum += Y[I];
+ // CHECK: target + teams distribute: 1048576
+ printf("target + teams distribute: %ld\n", Sum);
+
+ return 0;
+}
diff --git a/openmp/device/src/Reduction.cpp b/openmp/device/src/Reduction.cpp
index 9d0a5121995af7..0b7f687d6a83e4 100644
--- a/openmp/device/src/Reduction.cpp
+++ b/openmp/device/src/Reduction.cpp
@@ -287,6 +287,16 @@ int32_t __kmpc_gpu_xteam_reduce_nowait(IdentTy *Loc, void *reduce_data,
// Make sure that global buffer is fresh.
fence::kernel(atomic::acquire);
+
+ // Use a loop with a constant stride for the single-thread case, so that it
+ // can be unrolled.
+ if (NumThreads == 1) {
+ glcpyFct(GlobalBuffer, 0, reduce_data);
+ for (uint32_t I = 1; I < NumTeams; ++I)
+ glredFct(GlobalBuffer, I, reduce_data);
+ return 1;
+ }
+
// Get the team values from the global buffer.
glcpyFct(GlobalBuffer, ThreadId, reduce_data);
// In case we have more teams than threads, we need to iterate over the
@@ -294,9 +304,6 @@ int32_t __kmpc_gpu_xteam_reduce_nowait(IdentTy *Loc, void *reduce_data,
for (uint32_t I = NumThreads + ThreadId; I < NumTeams; I += NumThreads)
glredFct(GlobalBuffer, I, reduce_data);
- if (NumThreads == 1)
- return 1;
-
return gpu_block_reduce<false>(reduce_data, shflFct, cpyFct, NumThreads,
ThreadId);
}
``````````
</details>
https://github.com/llvm/llvm-project/pull/231843
More information about the Openmp-commits
mailing list