[Mlir-commits] [llvm] [mlir] [flang][openmp] Fix GPU byref reduction descriptor initialization (PR #178934)
Sunil Shrestha
llvmlistbot at llvm.org
Wed Feb 4 07:55:25 PST 2026
https://github.com/sshrestha-aa updated https://github.com/llvm/llvm-project/pull/178934
>From d6bf3483d764181985550c70933480aa79036c8e Mon Sep 17 00:00:00 2001
From: Sunil Shrestha <sshrestha at pe28vega.hpc.amslabs.hpecorp.net>
Date: Wed, 14 Jan 2026 18:07:44 -0600
Subject: [PATCH] [flang][openmp] Fix GPU byref reduction descriptor
initialization
When generating GPU reduction code for arrays passed by reference,
only the base_ptr field was initialized in the shuffled descriptor,
leaving extent, stride, and rank fields uninitialized. This caused
garbage metadata to be passed to user reduction combiners, resulting
in incorrect iteration bounds and crashes on GPU targets.
Fix by copying the entire source descriptor and then updating the
base_ptr to point to thread-private storage. This preserves all
metadata (extents, strides, rank) while correctly pointing to the
shuffled data location.
The fix applies to three reduction helper functions:
- _omp_reduction_shuffle_and_reduce_func (warp-level shuffle)
- _omp_reduction_list_to_global_reduce_func (block-to-global)
- _omp_reduction_global_to_list_copy_func (global-to-block)
Fixes multi-dimensional array reductions on GPU target regions with
teams distribute parallel for directives.
---
.../llvm/Frontend/OpenMP/OMPIRBuilder.h | 15 ++
llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp | 79 ++++++++---
...distribute-reduction-array-descriptor.mlir | 129 ++++++++++++++++++
3 files changed, 207 insertions(+), 16 deletions(-)
create mode 100644 mlir/test/Target/LLVMIR/omptarget-teams-distribute-reduction-array-descriptor.mlir
diff --git a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
index 037fcaa863fe7..778c3cc72281b 100644
--- a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
+++ b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
@@ -1935,6 +1935,21 @@ class OpenMPIRBuilder {
/// Get the function name of a reduction function.
std::string getReductionFuncName(StringRef Name) const;
+ /// Generate a Fortran descriptor for array reductions
+ ///
+ /// \param DescriptorAddr Address of the descriptor to initialize
+ /// \param DataPtr Pointer to the actual data the descriptor should reference
+ /// \param ElemType Type of elements in the array (may be array type)
+ /// \param DescriptorType Type of the descriptor structure
+ /// \param DataPtrPtrGen Callback to get the base_ptr field in the descriptor
+ ///
+ /// \return Error if DataPtrPtrGen fails, otherwise success.
+ InsertPointOrErrorTy generateReductionDescriptor(
+ Value *DescriptorAddr, Value *DataPtr, Value *SrcDescriptorAddr,
+ Type *DescriptorType,
+ function_ref<InsertPointOrErrorTy(InsertPointTy, Value *, Value *&)>
+ DataPtrPtrGen);
+
/// Emits reduction function.
/// \param ReducerName Name of the function calling the reduction.
/// \param ReductionInfos Array type containing the ReductionOps.
diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index 99f98d078a482..79ce0909228bb 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -3075,19 +3075,16 @@ Error OpenMPIRBuilder::emitReductionListCopy(
RemoteLaneOffset, ReductionArrayTy, IsByRefElem);
if (IsByRefElem) {
- Value *GEP;
- InsertPointOrErrorTy GenResult =
- RI.DataPtrPtrGen(Builder.saveIP(),
- Builder.CreatePointerBitCastOrAddrSpaceCast(
- DestAlloca, Builder.getPtrTy(), ".ascast"),
- GEP);
+ // Copy descriptor from source and update base_ptr to shuffled data
+ Value *DestDescriptorAddr = Builder.CreatePointerBitCastOrAddrSpaceCast(
+ DestAlloca, Builder.getPtrTy(), ".ascast");
+
+ InsertPointOrErrorTy GenResult = generateReductionDescriptor(
+ DestDescriptorAddr, LocalStorage, SrcElementAddr,
+ RI.ByRefAllocatedType, RI.DataPtrPtrGen);
if (!GenResult)
return GenResult.takeError();
-
- Builder.CreateStore(Builder.CreatePointerBitCastOrAddrSpaceCast(
- LocalStorage, Builder.getPtrTy(), ".ascast"),
- GEP);
}
} else {
switch (RI.EvaluationKind) {
@@ -3577,6 +3574,37 @@ Expected<Function *> OpenMPIRBuilder::emitShuffleAndReduceFunction(
return SarFunc;
}
+OpenMPIRBuilder::InsertPointOrErrorTy
+OpenMPIRBuilder::generateReductionDescriptor(
+ Value *DescriptorAddr, Value *DataPtr, Value *SrcDescriptorAddr,
+ Type *DescriptorType,
+ function_ref<InsertPointOrErrorTy(InsertPointTy, Value *, Value *&)>
+ DataPtrPtrGen) {
+
+ // Copy the source descriptor to preserve all metadata (rank, extents,
+ // strides, etc.)
+ Value *DescriptorSize =
+ Builder.getInt64(M.getDataLayout().getTypeStoreSize(DescriptorType));
+ Builder.CreateMemCpy(
+ DescriptorAddr, M.getDataLayout().getPrefTypeAlign(DescriptorType),
+ SrcDescriptorAddr, M.getDataLayout().getPrefTypeAlign(DescriptorType),
+ DescriptorSize);
+
+ // Update the base pointer field to point to the local shuffled data
+ Value *DataPtrField;
+ InsertPointOrErrorTy GenResult =
+ DataPtrPtrGen(Builder.saveIP(), DescriptorAddr, DataPtrField);
+
+ if (!GenResult)
+ return GenResult.takeError();
+
+ Builder.CreateStore(Builder.CreatePointerBitCastOrAddrSpaceCast(
+ DataPtr, Builder.getPtrTy(), ".ascast"),
+ DataPtrField);
+
+ return Builder.saveIP();
+}
+
Expected<Function *> OpenMPIRBuilder::emitListToGlobalCopyFunction(
ArrayRef<ReductionInfo> ReductionInfos, Type *ReductionsBufferTy,
AttributeList FuncAttrs, ArrayRef<bool> IsByRef) {
@@ -3789,15 +3817,24 @@ Expected<Function *> OpenMPIRBuilder::emitListToGlobalReduceFunction(
ReductionsBufferTy, BufferVD, 0, En.index());
if (!IsByRef.empty() && IsByRef[En.index()]) {
- Value *ByRefDataPtr;
+ // Get source descriptor from the reduce list argument
+ Value *ReduceList =
+ Builder.CreateLoad(Builder.getPtrTy(), ReduceListArgAddrCast);
+ Value *SrcElementPtrPtr =
+ Builder.CreateInBoundsGEP(RedListArrayTy, ReduceList,
+ {ConstantInt::get(IndexTy, 0),
+ ConstantInt::get(IndexTy, En.index())});
+ Value *SrcDescriptorAddr =
+ Builder.CreateLoad(Builder.getPtrTy(), SrcElementPtrPtr);
+ // Copy descriptor from source and update base_ptr to global buffer data
InsertPointOrErrorTy GenResult =
- RI.DataPtrPtrGen(Builder.saveIP(), ByRefAlloc, ByRefDataPtr);
+ generateReductionDescriptor(ByRefAlloc, GlobValPtr, SrcDescriptorAddr,
+ RI.ByRefAllocatedType, RI.DataPtrPtrGen);
if (!GenResult)
return GenResult.takeError();
- Builder.CreateStore(GlobValPtr, ByRefDataPtr);
Builder.CreateStore(ByRefAlloc, TargetElementPtrPtr);
} else {
Builder.CreateStore(GlobValPtr, TargetElementPtrPtr);
@@ -4023,13 +4060,23 @@ Expected<Function *> OpenMPIRBuilder::emitGlobalToListReduceFunction(
ReductionsBufferTy, BufferVD, 0, En.index());
if (!IsByRef.empty() && IsByRef[En.index()]) {
- Value *ByRefDataPtr;
+ // Get source descriptor from the reduce list
+ Value *ReduceListVal =
+ Builder.CreateLoad(Builder.getPtrTy(), ReduceListArgAddrCast);
+ Value *SrcElementPtrPtr =
+ Builder.CreateInBoundsGEP(RedListArrayTy, ReduceListVal,
+ {ConstantInt::get(IndexTy, 0),
+ ConstantInt::get(IndexTy, En.index())});
+ Value *SrcDescriptorAddr =
+ Builder.CreateLoad(Builder.getPtrTy(), SrcElementPtrPtr);
+
+ // Copy descriptor from source and update base_ptr to global buffer data
InsertPointOrErrorTy GenResult =
- RI.DataPtrPtrGen(Builder.saveIP(), ByRefAlloc, ByRefDataPtr);
+ generateReductionDescriptor(ByRefAlloc, GlobValPtr, SrcDescriptorAddr,
+ RI.ByRefAllocatedType, RI.DataPtrPtrGen);
if (!GenResult)
return GenResult.takeError();
- Builder.CreateStore(GlobValPtr, ByRefDataPtr);
Builder.CreateStore(ByRefAlloc, TargetElementPtrPtr);
} else {
Builder.CreateStore(GlobValPtr, TargetElementPtrPtr);
diff --git a/mlir/test/Target/LLVMIR/omptarget-teams-distribute-reduction-array-descriptor.mlir b/mlir/test/Target/LLVMIR/omptarget-teams-distribute-reduction-array-descriptor.mlir
new file mode 100644
index 0000000000000..84b4a0e71c36f
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/omptarget-teams-distribute-reduction-array-descriptor.mlir
@@ -0,0 +1,129 @@
+// RUN: mlir-translate -mlir-to-llvmir -split-input-file %s | FileCheck --check-prefixes=AMDGCN,NVPTX %s
+
+// Minimal MLIR to exercise array byref reduction descriptor handling in
+// target teams distribute parallel do.
+
+module attributes {dlti.dl_spec = #dlti.dl_spec<"dlti.alloca_memory_space" = 5 : ui64, "dlti.global_memory_space" = 1 : ui64>, llvm.target_triple = "amdgcn-amd-amdhsa", omp.is_gpu = true, omp.is_target_device = true} {
+ omp.declare_reduction @add_reduction_byref_box_4xi32 : !llvm.ptr attributes {byref_element_type = !llvm.array<4 x i32>} alloc {
+ %0 = llvm.mlir.constant(1 : i64) : i64
+ %1 = llvm.alloca %0 x !llvm.struct<(ptr, i64, i32, i8, i8, i8, i8, array<1 x array<3 x i64>>)> : (i64) -> !llvm.ptr<5>
+ %2 = llvm.addrspacecast %1 : !llvm.ptr<5> to !llvm.ptr
+ omp.yield(%2 : !llvm.ptr)
+ } init {
+ ^bb0(%arg0: !llvm.ptr, %arg1: !llvm.ptr):
+ omp.yield(%arg1 : !llvm.ptr)
+ } combiner {
+ ^bb0(%arg0: !llvm.ptr, %arg1: !llvm.ptr):
+ omp.yield(%arg0 : !llvm.ptr)
+ } data_ptr_ptr {
+ ^bb0(%arg0: !llvm.ptr):
+ %0 = llvm.getelementptr %arg0[0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<(ptr, i64, i32, i8, i8, i8, i8, array<1 x array<3 x i64>>)>
+ omp.yield(%0 : !llvm.ptr)
+ }
+
+ llvm.func @test_array_reduction_() attributes {omp.declare_target = #omp.declaretarget<device_type = (host), capture_clause = (to)>} {
+ %0 = llvm.mlir.constant(1 : i64) : i64
+ %1 = llvm.alloca %0 x !llvm.array<4 x i32> : (i64) -> !llvm.ptr<5>
+ %2 = llvm.addrspacecast %1 : !llvm.ptr<5> to !llvm.ptr
+ %3 = omp.map.info var_ptr(%2 : !llvm.ptr, !llvm.array<4 x i32>) map_clauses(tofrom) capture(ByRef) -> !llvm.ptr {name = "red_array"}
+ omp.target map_entries(%3 -> %arg0 : !llvm.ptr) {
+ %4 = llvm.mlir.constant(1 : i32) : i32
+ %5 = llvm.mlir.constant(1000 : i32) : i32
+ omp.teams reduction(byref @add_reduction_byref_box_4xi32 %arg0 -> %arg1 : !llvm.ptr) {
+ omp.parallel {
+ omp.distribute {
+ omp.wsloop {
+ omp.loop_nest (%iv) : i32 = (%4) to (%5) inclusive step (%4) {
+ omp.yield
+ }
+ } {omp.composite}
+ } {omp.composite}
+ omp.terminator
+ } {omp.composite}
+ omp.terminator
+ }
+ omp.terminator
+ }
+ llvm.return
+ }
+}
+
+// Verify descriptor is copied via memcpy and base_ptr is updated in all helpers
+// AMDGCN-LABEL: define internal void @_omp_reduction_shuffle_and_reduce_func
+// AMDGCN: call void @llvm.memcpy{{.*}}(ptr {{.*}}, ptr {{.*}}, i64 {{[0-9]+}}, i1 false)
+// AMDGCN: getelementptr {{.*}} ptr {{%.*}}, i32 0, i32 0
+// AMDGCN: store ptr {{%.*}}, ptr
+
+// AMDGCN-LABEL: define internal void @_omp_reduction_list_to_global_reduce_func
+// AMDGCN: call void @llvm.memcpy{{.*}}(ptr {{.*}}, ptr {{.*}}, i64 {{[0-9]+}}, i1 false)
+// AMDGCN: getelementptr {{.*}} ptr {{%.*}}, i32 0, i32 0
+// AMDGCN: store ptr {{%.*}}, ptr
+
+// AMDGCN-LABEL: define internal void @_omp_reduction_global_to_list_copy_func
+// AMDGCN: call void @llvm.memcpy{{.*}}(ptr {{.*}}, ptr {{.*}}, i64 {{[0-9]+}}, i1 false)
+// AMDGCN: getelementptr {{.*}} ptr {{%.*}}, i32 0, i32 0
+// AMDGCN: store ptr {{%.*}}, ptr
+
+// -----
+
+module attributes {llvm.target_triple = "nvptx64-nvidia-cuda", omp.is_gpu = true, omp.is_target_device = true} {
+ omp.declare_reduction @add_reduction_byref_box_4xi32 : !llvm.ptr attributes {byref_element_type = !llvm.array<4 x i32>} alloc {
+ %0 = llvm.mlir.constant(1 : i64) : i64
+ %1 = llvm.alloca %0 x !llvm.struct<(ptr, i64, i32, i8, i8, i8, i8, array<1 x array<3 x i64>>)> : (i64) -> !llvm.ptr<5>
+ %2 = llvm.addrspacecast %1 : !llvm.ptr<5> to !llvm.ptr
+ omp.yield(%2 : !llvm.ptr)
+ } init {
+ ^bb0(%arg0: !llvm.ptr, %arg1: !llvm.ptr):
+ omp.yield(%arg1 : !llvm.ptr)
+ } combiner {
+ ^bb0(%arg0: !llvm.ptr, %arg1: !llvm.ptr):
+ omp.yield(%arg0 : !llvm.ptr)
+ } data_ptr_ptr {
+ ^bb0(%arg0: !llvm.ptr):
+ %0 = llvm.getelementptr %arg0[0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<(ptr, i64, i32, i8, i8, i8, i8, array<1 x array<3 x i64>>)>
+ omp.yield(%0 : !llvm.ptr)
+ }
+
+ llvm.func @test_array_reduction_() attributes {omp.declare_target = #omp.declaretarget<device_type = (host), capture_clause = (to)>} {
+ %0 = llvm.mlir.constant(1 : i64) : i64
+ %1 = llvm.alloca %0 x !llvm.array<4 x i32> : (i64) -> !llvm.ptr<5>
+ %2 = llvm.addrspacecast %1 : !llvm.ptr<5> to !llvm.ptr
+ %3 = omp.map.info var_ptr(%2 : !llvm.ptr, !llvm.array<4 x i32>) map_clauses(tofrom) capture(ByRef) -> !llvm.ptr {name = "red_array"}
+ omp.target map_entries(%3 -> %arg0 : !llvm.ptr) {
+ %4 = llvm.mlir.constant(1 : i32) : i32
+ %5 = llvm.mlir.constant(1000 : i32) : i32
+ omp.teams reduction(byref @add_reduction_byref_box_4xi32 %arg0 -> %arg1 : !llvm.ptr) {
+ omp.parallel {
+ omp.distribute {
+ omp.wsloop {
+ omp.loop_nest (%iv) : i32 = (%4) to (%5) inclusive step (%4) {
+ omp.yield
+ }
+ } {omp.composite}
+ } {omp.composite}
+ omp.terminator
+ } {omp.composite}
+ omp.terminator
+ }
+ omp.terminator
+ }
+ llvm.return
+ }
+}
+
+// Verify descriptor is copied via memcpy and base_ptr is updated in all helpers
+// NVPTX-LABEL: define internal void @_omp_reduction_shuffle_and_reduce_func
+// NVPTX: call void @llvm.memcpy{{.*}}(ptr {{.*}}, ptr {{.*}}, i64 {{[0-9]+}}, i1 false)
+// NVPTX: getelementptr {{.*}} ptr {{%.*}}, i32 0, i32 0
+// NVPTX: store ptr {{%.*}}, ptr
+
+// NVPTX-LABEL: define internal void @_omp_reduction_list_to_global_reduce_func
+// NVPTX: call void @llvm.memcpy{{.*}}(ptr {{.*}}, ptr {{.*}}, i64 {{[0-9]+}}, i1 false)
+// NVPTX: getelementptr {{.*}} ptr {{%.*}}, i32 0, i32 0
+// NVPTX: store ptr {{%.*}}, ptr
+
+// NVPTX-LABEL: define internal void @_omp_reduction_global_to_list_copy_func
+// NVPTX: call void @llvm.memcpy{{.*}}(ptr {{.*}}, ptr {{.*}}, i64 {{[0-9]+}}, i1 false)
+// NVPTX: getelementptr {{.*}} ptr {{%.*}}, i32 0, i32 0
+// NVPTX: store ptr {{%.*}}, ptr
+
More information about the Mlir-commits
mailing list