diff --git a/clang/lib/CodeGen/CGCUDANV.cpp b/clang/lib/CodeGen/CGCUDANV.cpp --- a/clang/lib/CodeGen/CGCUDANV.cpp +++ b/clang/lib/CodeGen/CGCUDANV.cpp @@ -1195,15 +1195,26 @@ llvm::GlobalValue *CGNVCUDARuntime::getKernelHandle(llvm::Function *F, GlobalDecl GD) { auto Loc = KernelHandles.find(F->getName()); - if (Loc != KernelHandles.end()) - return Loc->second; + llvm::GlobalValue *OldHandle = nullptr; + if (Loc != KernelHandles.end()) { + OldHandle = Loc->second; + if (KernelStubs[OldHandle] == F) + return OldHandle; + // We've found the function name, but F itself has changed, so we need to + // update the references. + } if (!CGM.getLangOpts().HIP) { KernelHandles[F->getName()] = F; + if (OldHandle) + KernelStubs.erase(OldHandle); KernelStubs[F] = F; return F; } + // Make sure that the old variable name is no longer used. + if (OldHandle) + OldHandle->setName(""); auto *Var = new llvm::GlobalVariable( TheModule, F->getType(), /*isConstant=*/true, F->getLinkage(), /*Initializer=*/nullptr, @@ -1213,6 +1224,11 @@ Var->setDSOLocal(F->isDSOLocal()); Var->setVisibility(F->getVisibility()); CGM.maybeSetTrivialComdat(*GD.getDecl(), *Var); + if (OldHandle) { + KernelStubs.erase(OldHandle); + OldHandle->replaceAllUsesWith(Var); + OldHandle->eraseFromParent(); + } KernelHandles[F->getName()] = Var; KernelStubs[Var] = F; return Var; diff --git a/clang/test/CodeGenCUDA/bug-kerner-registration-reuse.cu b/clang/test/CodeGenCUDA/bug-kerner-registration-reuse.cu new file mode 100644 --- /dev/null +++ b/clang/test/CodeGenCUDA/bug-kerner-registration-reuse.cu @@ -0,0 +1,38 @@ +// RUN: echo -n "GPU binary would be here." > %t +// RUN: %clang_cc1 -triple x86_64-linux-gnu -emit-llvm %s \ +// RUN: -target-sdk-version=11.0 -fcuda-include-gpubinary %t -o - \ +// RUN: | FileCheck %s --check-prefixes CUDA +// RUN: %clang_cc1 -triple x86_64-linux-gnu -emit-llvm %s -x hip \ +// RUN: -fcuda-include-gpubinary %t -o - \ +// RUN: | FileCheck %s --check-prefixes HIP + +#include "Inputs/cuda.h" + +template +struct S { T t; }; + +template + static __global__ void Kernel(S) {} + +// For some reason it takes three or more instantiations of Kernel to trigger a +// crash during CUDA compilation. +auto x = &Kernel; +auto y = &Kernel; +auto z = &Kernel; + +// This triggers HIP-specific code path. +void func (){ + Kernel<<<1,1>>>({1}); +} + +// CUDA-LABEL: @__cuda_register_globals( +// CUDA: call i32 @__cudaRegisterFunction(ptr %0, ptr @_ZL21__device_stub__KernelIdEv1SIT_E +// CUDA: call i32 @__cudaRegisterFunction(ptr %0, ptr @_ZL21__device_stub__KernelIfEv1SIT_E +// CUDA: call i32 @__cudaRegisterFunction(ptr %0, ptr @_ZL21__device_stub__KernelIiEv1SIT_E +// CUDA: ret void + +// HIP-LABEL: @__hip_register_globals( +// HIP: call i32 @__hipRegisterFunction(ptr %0, ptr @_ZL6KernelIdEv1SIT_E +// HIP: call i32 @__hipRegisterFunction(ptr %0, ptr @_ZL6KernelIfEv1SIT_E +// HIP: call i32 @__hipRegisterFunction(ptr %0, ptr @_ZL6KernelIiEv1SIT_E +// HIP: ret void