[MLIR][OpenMP] Set default address space for OpenMPIRBuilder (#158689)

Extension of llvm/llvm-project#158152 for MLIR.

---------

Signed-off-by: Sarnie, Nick <[email protected]>

Author: sarnex

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

@@ -2246,18 +2246,22 @@ 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();
 }

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
+  }
+}