[Mlir-commits] [mlir] 205ee34 - [flang][openmp] Fix GPU byref reduction descriptor initialization (#178934)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Feb 5 02:26:51 PST 2026
Author: Sunil Shrestha
Date: 2026-02-05T11:26:44+01:00
New Revision: 205ee34317c0fbc97dac39608e2cca414c2d8491
URL: https://github.com/llvm/llvm-project/commit/205ee34317c0fbc97dac39608e2cca414c2d8491
DIFF: https://github.com/llvm/llvm-project/commit/205ee34317c0fbc97dac39608e2cca414c2d8491.diff
LOG: [flang][openmp] Fix GPU byref reduction descriptor initialization (#178934)
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.
Co-authored-by: Sunil Shrestha <sshrestha at pe28vega.hpc.amslabs.hpecorp.net>
Added:
mlir/test/Target/LLVMIR/omptarget-teams-distribute-reduction-array-descriptor.mlir
Modified:
llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
Removed:
################################################################################
diff --git a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
index 8ed4e0ba14502..f73e10c97e642 100644
--- a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
+++ b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
@@ -1939,6 +1939,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 cefd066d185bb..464ec5b5a2ece 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -3100,19 +3100,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) {
@@ -3602,6 +3599,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) {
@@ -3814,15 +3842,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);
@@ -4048,13 +4085,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