[clang][OpenMP][SPIR-V] Fix addrspace of globals

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

Author: sarnex

clang/lib/Basic/Targets/SPIR.h

@@ -57,10 +57,11 @@ static const unsigned SPIRDefIsPrivMap[] = {
 // Used by both the SPIR and SPIR-V targets.
 static const unsigned SPIRDefIsGenMap[] = {
     4, // Default
-    // OpenCL address space values for this map are dummy and they can't be used
-    0, // opencl_global
+    // Some OpenCL address space values for this map are dummy and they can't be
+    // used
+    1, // opencl_global
     0, // opencl_local
-    0, // opencl_constant
+    2, // opencl_constant
     0, // opencl_private
     0, // opencl_generic
     0, // opencl_global_device
@@ -216,7 +217,7 @@ class LLVM_LIBRARY_VISIBILITY BaseSPIRTargetInfo : public TargetInfo {
         /*DefaultIsGeneric=*/Opts.SYCLIsDevice ||
         // The address mapping from HIP/CUDA language for device code is only
         // defined for SPIR-V.
-        (getTriple().isSPIRV() && Opts.CUDAIsDevice));
+        (getTriple().isSPIRV() && (Opts.CUDAIsDevice || Opts.OpenMPIsTargetDevice)));
   }
 
   void setSupportedOpenCLOpts() override {

clang/lib/CodeGen/CGOpenMPRuntime.cpp

@@ -2524,6 +2524,16 @@ void CGOpenMPRuntime::emitForDispatchInit(
                       Args);
 }
 
+llvm::Value *CGOpenMPRuntime::createRuntimeFunctionArgAddrSpaceCast(
+    CodeGenFunction &CGF, llvm::FunctionCallee RuntimeFcn, size_t ArgIdx,
+    llvm::Value *Arg) {
+  llvm::Type *ParamTy = RuntimeFcn.getFunctionType()->getParamType(ArgIdx);
+  llvm::Type *ArgTy = Arg->getType();
+  if (!ParamTy->isPointerTy())
+    return Arg;
+  return CGF.Builder.CreateAddrSpaceCast(Arg, ParamTy);
+}
+
 void CGOpenMPRuntime::emitForDispatchDeinit(CodeGenFunction &CGF,
                                             SourceLocation Loc) {
   if (!CGF.HaveInsertPoint())
@@ -2572,12 +2582,18 @@ static void emitForStaticInitCall(
       ThreadId,
       CGF.Builder.getInt32(addMonoNonMonoModifier(CGF.CGM, Schedule, M1,
                                                   M2)), // Schedule type
-      Values.IL.emitRawPointer(CGF),                    // &isLastIter
-      Values.LB.emitRawPointer(CGF),                    // &LB
-      Values.UB.emitRawPointer(CGF),                    // &UB
-      Values.ST.emitRawPointer(CGF),                    // &Stride
-      CGF.Builder.getIntN(Values.IVSize, 1),            // Incr
-      Chunk                                             // Chunk
+      CGOpenMPRuntime::createRuntimeFunctionArgAddrSpaceCast(
+          CGF, ForStaticInitFunction, 3,
+          Values.IL.emitRawPointer(CGF)), // &isLastIter
+      CGOpenMPRuntime::createRuntimeFunctionArgAddrSpaceCast(
+          CGF, ForStaticInitFunction, 4, Values.LB.emitRawPointer(CGF)), // &LB
+      CGOpenMPRuntime::createRuntimeFunctionArgAddrSpaceCast(
+          CGF, ForStaticInitFunction, 5, Values.UB.emitRawPointer(CGF)), // &UB
+      CGOpenMPRuntime::createRuntimeFunctionArgAddrSpaceCast(
+          CGF, ForStaticInitFunction, 6,
+          Values.ST.emitRawPointer(CGF)),    // &Stride
+      CGF.Builder.getIntN(Values.IVSize, 1), // Incr
+      Chunk                                  // Chunk
   };
   CGF.EmitRuntimeCall(ForStaticInitFunction, Args);
 }

clang/lib/CodeGen/CGOpenMPRuntime.h

@@ -1657,6 +1657,11 @@ class CGOpenMPRuntime {
 
   /// Returns true if the variable is a local variable in untied task.
   bool isLocalVarInUntiedTask(CodeGenFunction &CGF, const VarDecl *VD) const;
+
+  static llvm::Value *
+  createRuntimeFunctionArgAddrSpaceCast(CodeGenFunction &CGF,
+                                        llvm::FunctionCallee RuntimeFcn,
+                                        size_t ArgIdx, llvm::Value *Arg);
 };
 
 /// Class supports emissionof SIMD-only code.

clang/lib/CodeGen/CGOpenMPRuntimeGPU.cpp

@@ -1217,11 +1217,13 @@ void CGOpenMPRuntimeGPU::emitParallelCall(CodeGenFunction &CGF,
     CGBuilderTy &Bld = CGF.Builder;
     llvm::Value *NumThreadsVal = NumThreads;
     llvm::Function *WFn = WrapperFunctionsMap[OutlinedFn];
+    llvm::FunctionCallee RuntimeFn = OMPBuilder.getOrCreateRuntimeFunction(
+        CGM.getModule(), OMPRTL___kmpc_parallel_51);
     llvm::Value *ID = llvm::ConstantPointerNull::get(CGM.Int8PtrTy);
     if (WFn)
       ID = Bld.CreateBitOrPointerCast(WFn, CGM.Int8PtrTy);
-    llvm::Value *FnPtr = Bld.CreateBitOrPointerCast(OutlinedFn, CGM.Int8PtrTy);
-
+    llvm::Value *FnPtr = Bld.CreateAddrSpaceCast(OutlinedFn, CGM.Int8PtrTy);
+    FnPtr = Bld.CreateBitOrPointerCast(FnPtr, CGM.Int8PtrTy);
     // Create a private scope that will globalize the arguments
     // passed from the outside of the target region.
     // TODO: Is that needed?
@@ -1268,14 +1270,15 @@ void CGOpenMPRuntimeGPU::emitParallelCall(CodeGenFunction &CGF,
         IfCondVal,
         NumThreadsVal,
         llvm::ConstantInt::get(CGF.Int32Ty, -1),
-        FnPtr,
-        ID,
-        Bld.CreateBitOrPointerCast(CapturedVarsAddrs.emitRawPointer(CGF),
-                                   CGF.VoidPtrPtrTy),
+        createRuntimeFunctionArgAddrSpaceCast(CGF, RuntimeFn, 5, FnPtr),
+        createRuntimeFunctionArgAddrSpaceCast(CGF, RuntimeFn, 6, ID),
+        createRuntimeFunctionArgAddrSpaceCast(
+            CGF, RuntimeFn, 7,
+            Bld.CreateBitOrPointerCast(CapturedVarsAddrs.emitRawPointer(CGF),
+                                       CGF.VoidPtrPtrTy)),
         llvm::ConstantInt::get(CGM.SizeTy, CapturedVars.size())};
-    CGF.EmitRuntimeCall(OMPBuilder.getOrCreateRuntimeFunction(
-                            CGM.getModule(), OMPRTL___kmpc_parallel_51),
-                        Args);
+
+    CGF.EmitRuntimeCall(RuntimeFn, Args);
   };
 
   RegionCodeGenTy RCG(ParallelGen);

clang/lib/CodeGen/CGStmtOpenMP.cpp

@@ -585,6 +585,8 @@ static llvm::Function *emitOutlinedFunctionPrologue(
     F->removeFnAttr(llvm::Attribute::NoInline);
     F->addFnAttr(llvm::Attribute::AlwaysInline);
   }
+  if (CGM.getTriple().isSPIRV())
+    F->setCallingConv(llvm::CallingConv::SPIR_FUNC);
 
   // Generate the function.
   CGF.StartFunction(CD, Ctx.VoidTy, F, FuncInfo, TargetArgs,

clang/test/OpenMP/spirv_target_addrspace.c

@@ -0,0 +1,21 @@
+// 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 -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 fcn(const char[]);
+
+#pragma omp declare target
+// CHECK: @global = addrspace(1) global i32 0, align 4
+// CHECK: @.str = private unnamed_addr addrspace(1) 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 < 1024; i++)
+      global++;
+    fcn("foo");
+  }
+  return global;
+}
+

clang/test/OpenMP/spirv_target_addrspace_simd.c

@@ -0,0 +1,23 @@
+// 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 -fopenmp -fopenmp-targets=spirv64 -fopenmp-is-target-device -triple spirv64 -fopenmp-host-ir-file-path %t-host.bc -emit-llvm %s -o - | FileCheck %s
+
+int main() {
+  int x = 0;
+
+#pragma omp target teams distribute parallel for simd
+  for(int i = 0; i < 1024; i++)
+    x+=i;
+  return x;
+}
+
+// CHECK: @[[#STRLOC:]] = private unnamed_addr addrspace(1) constant [{{.*}} x i8] c{{.*}}, align 1
+// CHECK: @[[#IDENT:]] = private unnamed_addr addrspace(1) constant %struct.ident_t { i32 {{.*}}, i32 2050, i32 {{.*}}, i32 {{.*}}, ptr addrspacecast (ptr addrspace(1) @[[#STRLOC]] to ptr) }, align 8
+// CHECK: define internal spir_func void @__omp_offloading_{{.*}}_omp_outlined(ptr addrspace(4) noalias noundef {{.*}}., ptr addrspace(4) noalias noundef {{.*}}, i64 noundef {{.*}}) #{{.*}} {
+// CHECK: = load ptr addrspace(4), ptr addrspace(4) %{{.*}}, align 8
+// CHECK: = load i32, ptr addrspace(4) %{{.*}}, align 4
+// CHECK: = addrspacecast ptr addrspace(4) %{{.*}} to ptr
+// CHECK: = addrspacecast ptr addrspace(4) %{{.*}} to ptr
+// CHECK: = addrspacecast ptr addrspace(4) %{{.*}} to ptr
+// CHECK: = addrspacecast ptr addrspace(4) %{{.*}} to ptr
+// CHECK: call spir_func void @__kmpc_distribute_static_init{{.*}}(ptr addrspacecast (ptr addrspace(1) @[[#IDENT]] to ptr), i32 %{{.*}}, i32 {{.*}}, ptr %{{.*}}, ptr %{{.*}}, ptr %{{.*}}, ptr %{{.*}}, i32 {{.*}}, i32 %{{.*}})
+// CHECK: call spir_func void @__kmpc_distribute_static_fini{{.*}}(ptr addrspacecast (ptr addrspace(1) @[[#IDENT]] to ptr), i32 %{{.*}})

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

@@ -910,6 +910,14 @@ Constant *OpenMPIRBuilder::getOrCreateIdent(Constant *SrcLocStr,
                              ConstantInt::get(Int32, uint32_t(LocFlags)),
                              ConstantInt::get(Int32, Reserve2Flags),
                              ConstantInt::get(Int32, SrcLocStrSize), SrcLocStr};
+
+    size_t SrcLocStrArgIdx = 4;
+    if (OpenMPIRBuilder::Ident->getElementType(SrcLocStrArgIdx)
+            ->getPointerAddressSpace() !=
+        IdentData[SrcLocStrArgIdx]->getType()->getPointerAddressSpace())
+      IdentData[SrcLocStrArgIdx] = ConstantExpr::getAddrSpaceCast(
+          SrcLocStr, OpenMPIRBuilder::Ident->getElementType(SrcLocStrArgIdx));
+
     Constant *Initializer =
         ConstantStruct::get(OpenMPIRBuilder::Ident, IdentData);
 
@@ -950,8 +958,9 @@ Constant *OpenMPIRBuilder::getOrCreateSrcLocStr(StringRef LocStr,
           GV.getInitializer() == Initializer)
         return SrcLocStr = ConstantExpr::getPointerCast(&GV, Int8Ptr);
 
-    SrcLocStr = Builder.CreateGlobalString(LocStr, /* Name */ "",
-                                           /* AddressSpace */ 0, &M);
+    SrcLocStr = Builder.CreateGlobalString(
+        LocStr, /* Name */ "",
+        M.getDataLayout().getDefaultGlobalsAddressSpace(), &M);
   }
   return SrcLocStr;
 }