[clang] [llvm] [clang][OpenMP][SPIR-V] Fix addrspace of globals and global constants (PR #134399)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-flang-openmp

Author: Nick Sarnie (sarnex)

<details>
<summary>Changes</summary>

SPIR-V has strict address space rules, globals cannot be in the default address space. Normal globals should be in addrspace 1 (which is what we get from `opencl_global` in the SPIR-V address space map) and 2 for global constants (`opencl_constant` in the SPIR-V address space map)

This is similar to what was done for HIPSPV.

The OMPIRBuilder change was required for lit tests to pass, we were missing an addrspacecast.

---
Full diff: https://github.com/llvm/llvm-project/pull/134399.diff


3 Files Affected:

- (modified) clang/lib/CodeGen/CodeGenModule.cpp (+9) 
- (added) clang/test/OpenMP/spirv_target_addrspace.c (+20) 
- (modified) llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp (+6) 


``````````diff
diff --git a/clang/lib/CodeGen/CodeGenModule.cpp b/clang/lib/CodeGen/CodeGenModule.cpp
index 8f9cf965af2b9..cc6d726445cbb 100644
--- a/clang/lib/CodeGen/CodeGenModule.cpp
+++ b/clang/lib/CodeGen/CodeGenModule.cpp
@@ -5384,6 +5384,11 @@ LangAS CodeGenModule::GetGlobalVarAddressSpace(const VarDecl *D) {
     LangAS AS;
     if (OpenMPRuntime->hasAllocateAttributeForGlobalVar(D, AS))
       return AS;
+    if (LangOpts.OpenMPIsTargetDevice && getTriple().isSPIRV())
+      // SPIR-V globals should map to CrossWorkGroup instead of default
+      // AS, as generic/no address space is invalid. This is similar
+      // to what is done for HIPSPV.
+      return LangAS::opencl_global;
   }
   return getTargetCodeGenInfo().getGlobalVarAddressSpace(*this, D);
 }
@@ -5402,6 +5407,10 @@ LangAS CodeGenModule::GetGlobalConstantAddressSpace() const {
     // UniformConstant storage class is not viable as pointers to it may not be
     // casted to Generic pointers which are used to model HIP's "flat" pointers.
     return LangAS::cuda_device;
+  if (LangOpts.OpenMPIsTargetDevice && getTriple().isSPIRV())
+    // OpenMP SPIR-V global constants should map to UniformConstant, different
+    // from the HIPSPV case above.
+    return LangAS::opencl_constant;
   if (auto AS = getTarget().getConstantAddressSpace())
     return *AS;
   return LangAS::Default;
diff --git a/clang/test/OpenMP/spirv_target_addrspace.c b/clang/test/OpenMP/spirv_target_addrspace.c
new file mode 100644
index 0000000000000..8430a30efe0c8
--- /dev/null
+++ b/clang/test/OpenMP/spirv_target_addrspace.c
@@ -0,0 +1,20 @@
+// RUN: %clang_cc1 -fopenmp -x c++ -triple x86_64-unknown-unknown -fopenmp-targets=spirv64 -emit-llvm-bc %s -o %t-host.bc
+// RUN: %clang_cc1 -O0 -fopenmp -fopenmp-targets=spirv64 -fopenmp-is-target-device -triple spirv64 -fopenmp-host-ir-file-path %t-host.bc -emit-llvm %s -o - | FileCheck %s
+
+extern int printf(char[]);
+
+#pragma omp declare target
+// CHECK: @global = addrspace(1) global i32 0, align 4
+// CHECK: @.str = private unnamed_addr addrspace(2) constant [4 x i8] c"foo\00", align 1
+int global = 0;
+#pragma omp end declare target
+int main() {
+  // CHECK: = call i32 @__kmpc_target_init(ptr addrspacecast (ptr addrspace(1) @__omp_offloading_{{.*}}_kernel_environment to ptr), ptr %{{.*}})
+#pragma omp target
+  {
+    for(int i = 0; i < 5; i++)
+      global++;
+    printf("foo");
+  }
+  return global;
+}
diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index 68b1fa42934ad..998702c1af3cd 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -6295,6 +6295,12 @@ OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createTargetInit(
           : ConstantExpr::getAddrSpaceCast(KernelEnvironmentGV,
                                            KernelEnvironmentPtr);
   Value *KernelLaunchEnvironment = DebugKernelWrapper->getArg(0);
+  Type *KernelLaunchEnvParamTy = Fn->getFunctionType()->getParamType(1);
+  KernelLaunchEnvironment =
+      KernelLaunchEnvironment->getType() == KernelLaunchEnvParamTy
+          ? KernelLaunchEnvironment
+          : Builder.CreateAddrSpaceCast(KernelLaunchEnvironment,
+                                        KernelLaunchEnvParamTy);
   CallInst *ThreadKind =
       Builder.CreateCall(Fn, {KernelEnvironment, KernelLaunchEnvironment});
 

``````````

</details>


https://github.com/llvm/llvm-project/pull/134399