[Mlir-commits] [mlir] [mlir][OpenMP] Set default address space for OpenMPIRBuilder (PR #158689)

[Nick Sarnie]

https://github.com/sarnex created https://github.com/llvm/llvm-project/pull/158689

None

>From d3dc73c5300ea7a118927b242cfa49ef0cb236bb Mon Sep 17 00:00:00 2001
From: "Sarnie, Nick" <nick.sarnie at intel.com>
Date: Mon, 15 Sep 2025 08:12:17 -0700
Subject: [PATCH] [mlir][OpenMP] Set default address space for OpenMPIRBuilder

Signed-off-by: Sarnie, Nick <nick.sarnie at intel.com>
---
 mlir/lib/Target/LLVMIR/ModuleTranslation.cpp  | 12 ++++++++---
 .../LLVMIR/openmp-target-default-as.mlir      | 21 +++++++++++++++++++
 2 files changed, 30 insertions(+), 3 deletions(-)
 create mode 100644 mlir/test/Target/LLVMIR/openmp-target-default-as.mlir

diff --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
index 9725359160a1a..50d67846ca42b 100644
--- a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
@@ -2246,18 +2246,24 @@ SmallVector<llvm::Value *> ModuleTranslation::lookupValues(ValueRange values) {
 llvm::OpenMPIRBuilder *ModuleTranslation::getOpenMPBuilder() {
   if (!ompBuilder) {
     ompBuilder = std::make_unique<llvm::OpenMPIRBuilder>(*llvmModule);
-    ompBuilder->initialize();
 
     // Flags represented as top-level OpenMP dialect attributes are set in
     // `OpenMPDialectLLVMIRTranslationInterface::amendOperation()`. Here we set
     // the default configuration.
-    ompBuilder->setConfig(llvm::OpenMPIRBuilderConfig(
+    llvm::OpenMPIRBuilderConfig config(
         /* IsTargetDevice = */ false, /* IsGPU = */ false,
         /* OpenMPOffloadMandatory = */ false,
         /* HasRequiresReverseOffload = */ false,
         /* HasRequiresUnifiedAddress = */ false,
         /* HasRequiresUnifiedSharedMemory = */ false,
-        /* HasRequiresDynamicAllocators = */ false));
+        /* HasRequiresDynamicAllocators = */ false);
+    unsigned int defaultAS = getLLVMModule()
+                               ->getDataLayout()
+                               .getProgramAddressSpace();
+    config.setDefaultTargetAS(defaultAS);
+    ompBuilder->setConfig(std::move(config));
+    ompBuilder->initialize();
+
   }
   return ompBuilder.get();
 }
diff --git a/mlir/test/Target/LLVMIR/openmp-target-default-as.mlir b/mlir/test/Target/LLVMIR/openmp-target-default-as.mlir
new file mode 100644
index 0000000000000..8344867d5fb7b
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/openmp-target-default-as.mlir
@@ -0,0 +1,21 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+// This tests that we correctly use the default program AS from the data layout.
+module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"dlti.program_memory_space", 4 : ui32>>, llvm.target_triple = "spirv64-intel", omp.is_target_device = true, omp.is_gpu = true} {
+
+// CHECK: @[[IDENT:.*]] = private unnamed_addr constant %s{{.*}} { i32 0, i32 2, i32 0, i32 22, ptr addrspace(4) addrspacecast (ptr @{{.*}} to ptr addrspace(4)) }, align 8
+
+ llvm.func @omp_target_region_() {
+    %0 = llvm.mlir.constant(20 : i32) : i32
+    %1 = llvm.mlir.constant(10 : i32) : i32
+    %2 = llvm.mlir.constant(1 : i64) : i64
+    %3 = llvm.alloca %2 x i32 {bindc_name = "a", in_type = i32, operandSegmentSizes = array<i32: 0, 0>, uniq_name = "_QFomp_target_regionEa"} : (i64) -> !llvm.ptr<5>
+    %4 = llvm.addrspacecast %3 : !llvm.ptr<5> to !llvm.ptr
+    llvm.store %1, %4 : i32, !llvm.ptr
+    %map = omp.map.info var_ptr(%4 : !llvm.ptr, i32)   map_clauses(tofrom) capture(ByRef) -> !llvm.ptr {name = ""}
+    omp.target map_entries(%map -> %arg : !llvm.ptr) {
+      omp.terminator
+    }
+    llvm.return
+  }
+}