[CIR][HIP] Use GlobalView to access stub (#1957)

Author: koparasy

clang/lib/CIR/CodeGen/CIRGenCUDANV.cpp

@@ -265,12 +265,20 @@ void CIRGenNVCUDARuntime::emitDeviceStubBodyNew(CIRGenFunction &cgf,
 
 void CIRGenNVCUDARuntime::emitDeviceStub(CIRGenFunction &cgf, cir::FuncOp fn,
                                          FunctionArgList &args) {
+
   if (auto globalOp =
           llvm::dyn_cast<cir::GlobalOp>(KernelHandles[fn.getSymName()])) {
-    auto symbol = mlir::FlatSymbolRefAttr::get(fn.getSymNameAttr());
-    // Set the initializer for the global
-    cgm.setInitializer(globalOp, symbol);
+    auto &builder = cgm.getBuilder();
+    auto fnPtrTy = globalOp.getSymType();
+    auto sym = mlir::FlatSymbolRefAttr::get(fn.getSymNameAttr());
+    auto gv = cir::GlobalViewAttr::get(fnPtrTy, sym);
+
+    globalOp->setAttr("initial_value", gv);
+    globalOp->removeAttr("sym_visibility");
+    globalOp->setAttr("alignment", builder.getI64IntegerAttr(
+                                       cgm.getPointerAlign().getQuantity()));
   }
+
   // CUDA 9.0 changed the way to launch kernels.
   if (CudaFeatureEnabled(cgm.getTarget().getSDKVersion(),
                          CudaFeature::CUDA_USES_NEW_LAUNCH) ||
@@ -322,12 +330,11 @@ mlir::Operation *CIRGenNVCUDARuntime::getKernelHandle(cir::FuncOp fn,
             cgm, fn->getLoc(), globalName,
             builder.getPointerTo(fn.getFunctionType()), true,
             cir::AddressSpace::Default,
-            /*insertPoint=*/nullptr, fn.getLinkage());
+            /*insertPoint=*/nullptr);
       });
 
   globalOp->setAttr("alignment", builder.getI64IntegerAttr(
                                      cgm.getPointerAlign().getQuantity()));
-  globalOp->setAttr("visibility", fn->getAttr("sym_visibility"));
 
   // Store references
   KernelHandles[fn.getSymName()] = globalOp;

clang/test/CIR/CodeGen/HIP/simple.cpp

@@ -5,10 +5,15 @@
 // RUN:            -emit-cir %s -o %t.cir
 // RUN: FileCheck --check-prefix=CIR-HOST --input-file=%t.cir %s
 
-// RUN: %clang_cc1 -triple=amdgcn-amd-amdhsa -x hip \
+// RUN: %clang_cc1 -triple=amdgcn-amd-amdhsa -x hip -fclangir \
 // RUN:            -fcuda-is-device -fhip-new-launch-api \
 // RUN:              -emit-cir %s -o %t.cir
 // RUN: FileCheck --check-prefix=CIR-DEVICE --input-file=%t.cir %s
+//
+// RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -fclangir \
+// RUN:            -x hip -emit-llvm -fhip-new-launch-api \
+// RUN:            %s -o %t.ll
+// RUN: FileCheck --check-prefix=LLVM-HOST --input-file=%t.ll %s
 
 // Attribute for global_fn
 // CIR-HOST: [[Kernel:#[a-zA-Z_0-9]+]] = {{.*}}#cir.cu.kernel_name<_Z9global_fni>{{.*}}
@@ -25,14 +30,21 @@ __device__ void device_fn(int* a, double b, float c) {}
 __global__ void global_fn(int a) {}
 // CIR-DEVICE: @_Z9global_fni
 
-// CIR-HOST: cir.alloca {{.*}}"kernel_args"
+// CIR-HOST: @_Z24__device_stub__global_fni{{.*}}extra([[Kernel]])
+// CIR-HOST: %[[#CIRKernelArgs:]] = cir.alloca {{.*}}"kernel_args"
+// CIR-HOST: %[[#Decayed:]] = cir.cast array_to_ptrdecay %[[#CIRKernelArgs]]
 // CIR-HOST: cir.call @__hipPopCallConfiguration
-
-// Host access the global stub instead of the functiond evice stub.
-// The stub has the mangled name of the function
-// CIR-HOST: cir.get_global @_Z9global_fni
+// CIR-HOST: cir.get_global @_Z9global_fni : !cir.ptr<!cir.ptr<!cir.func<(!s32i)>>>
 // CIR-HOST: cir.call @hipLaunchKernel
 
+// LLVM-HOST: void @_Z24__device_stub__global_fni
+// LLVM-HOST: %[[#KernelArgs:]] = alloca [1 x ptr], i64 1, align 16
+// LLVM-HOST: %[[#GEP1:]] = getelementptr ptr, ptr %[[#KernelArgs]], i32 0
+// LLVM-HOST: %[[#GEP2:]] = getelementptr [1 x ptr], ptr %[[#KernelArgs]], i32 0, i64 0
+// LLVM-HOST: call i32 @__hipPopCallConfiguration
+// LLVM-HOST: call i32 @hipLaunchKernel(ptr @_Z9global_fni 
+
+
 int main() {
   global_fn<<<1, 1>>>(1);
 }
@@ -49,4 +61,20 @@ int main() {
 // CIR-HOST:   cir.call @_Z24__device_stub__global_fni([[Arg]])
 // CIR-HOST: }
 
+// LLVM-HOST: define dso_local i32 @main
+// LLVM-HOST: alloca %struct.dim3
+// LLVM-HOST: alloca %struct.dim3
+// LLVM-HOST: call void @_ZN4dim3C1Ejjj
+// LLVM-HOST: call void @_ZN4dim3C1Ejjj
+// LLVM-HOST: %[[#ConfigOK:]] = call i32 @__hipPushCallConfiguration
+// LLVM-HOST: %[[#ConfigCond:]] = icmp ne i32 %[[#ConfigOK]], 0
+// LLVM-HOST: br i1 %[[#ConfigCond]], label %[[#Good:]], label %[[#Bad:]]
+// LLVM-HOST: [[#Good]]:
+// LLVM-HOST:   br label %[[#End:]]
+// LLVM-HOST: [[#Bad]]:
+// LLVM-HOST:   call void @_Z24__device_stub__global_fni(i32 1)
+// LLVM-HOST:   br label %[[#End:]]
+// LLVM-HOST: [[#End]]:
+// LLVM-HOST:   %[[#]] = load i32
+// LLVM-HOST:   ret i32