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,8 +1195,23 @@ llvm::GlobalValue *CGNVCUDARuntime::getKernelHandle(llvm::Function *F, GlobalDecl GD) { auto Loc = KernelHandles.find(F->getName()); - if (Loc != KernelHandles.end()) - return Loc->second; + if (Loc != KernelHandles.end()) { + auto 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) { + // For HIP compilation the handle itself does not change, so we only need + // to update the Stub value. + KernelStubs[OldHandle] = F; + return OldHandle; + } + // For non-HIP compilation, erase the old Stub and fall-through to creating + // new entries. + KernelStubs.erase(OldHandle); + } if (!CGM.getLangOpts().HIP) { KernelHandles[F->getName()] = F; 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