[Mlir-commits] [mlir] 1a87171 - [OpenMPIRBuilder] Use generic pointers for outlined teams arguments (#222176)

[llvmlistbot at llvm.org]

Author: Gheorghe-Teodor Bercea
Date: 2026-09-16T23:04:20-04:00
New Revision: 1a87171c50553ccf767c3b51a516b81748c4e0d1

URL: https://github.com/llvm/llvm-project/commit/1a87171c50553ccf767c3b51a516b81748c4e0d1
DIFF: https://github.com/llvm/llvm-project/commit/1a87171c50553ccf767c3b51a516b81748c4e0d1.diff

LOG: [OpenMPIRBuilder] Use generic pointers for outlined teams arguments (#222176)

createFakeIntVal models the thread and bound id arguments of the
function
outlined for a teams region as allocas. When the target's alloca address
space is not zero these were passed on as-is, so the outlined function
took
addrspace(5) pointers while __kmpc_fork_teams calls it with generic
ones.

Cast the alloca to address space zero when it is passed as a pointer,
which
only affects targets with a non-zero alloca address space such as
AMDGPU.

Added: 
    

Modified: 
    llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
    mlir/test/Target/LLVMIR/openmp-teams.mlir

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index 0747b943d1ba4..8b196c461f195 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -475,6 +475,13 @@ Value *createFakeIntVal(IRBuilderBase &Builder,
 
   if (AsPtr) {
     FakeVal = FakeValAddr;
+    // The runtime passes these extra arguments to the outlined function as
+    // generic pointers, so cast away a non-zero alloca address space.
+    if (FakeValAddr->getAddressSpace() != 0) {
+      FakeVal = cast<Instruction>(Builder.CreateAddrSpaceCast(
+          FakeValAddr, Builder.getPtrTy(), Name + ".ascast"));
+      ToBeDeleted.push_back(FakeVal);
+    }
   } else {
     FakeVal = Builder.CreateLoad(IntTy, FakeValAddr, Name + ".val");
     ToBeDeleted.push_back(FakeVal);

diff  --git a/mlir/test/Target/LLVMIR/openmp-teams.mlir b/mlir/test/Target/LLVMIR/openmp-teams.mlir
index 4690b51122beb..a7bc92f8b99b9 100644
--- a/mlir/test/Target/LLVMIR/openmp-teams.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-teams.mlir
@@ -311,3 +311,29 @@ llvm.func @teams_if_with_num_teams(%condition: i1, %numTeamsLower: i32, %numTeam
     llvm.call @afterTeams() : () -> ()
     llvm.return
 }
+
+// -----
+
+// Check that the thread and bound id arguments of the outlined function are
+// generic pointers even when allocas are created in a non-zero address space.
+// The cast that produces them is deleted along with the fake allocas once the
+// region has been outlined, so the signature below is the only place it can be
+// observed.
+module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"dlti.alloca_memory_space", 5 : ui32>>} {
+    llvm.func @foo()
+
+    // CHECK-LABEL: @omp_teams_alloca_addrspace
+    // CHECK: call void {{.*}} @__kmpc_fork_teams(ptr @{{.+}}, i32 0, ptr @[[OUTLINED_FN:.+]])
+    llvm.func @omp_teams_alloca_addrspace() {
+        omp.teams {
+            llvm.call @foo() : () -> ()
+            omp.terminator
+        }
+        llvm.return
+    }
+
+    // CHECK:      define internal void @[[OUTLINED_FN]]
+    // CHECK-SAME: (ptr %global.tid.ptr, ptr %bound.tid.ptr)
+    // CHECK:   call void @foo()
+    // CHECK:   ret void
+}