[CUDA] Update cached kernel handle when the function instance changes.

Fixes clang crash caused by a stale function pointer.

The bug has been present for a pretty long time, but we were lucky not to
trigger it until  D140663.

Differential Revision: https://reviews.llvm.org/D146448
This commit is contained in:
Artem Belevich 2023-03-20 11:18:53 -07:00
parent c5f6339c12
commit 2aa90da012
2 changed files with 55 additions and 2 deletions

View File

@ -1195,8 +1195,23 @@ llvm::Function *CGNVCUDARuntime::finalizeModule() {
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;

View File

@ -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 <typename T>
struct S { T t; };
template <typename T>
static __global__ void Kernel(S<T>) {}
// For some reason it takes three or more instantiations of Kernel to trigger a
// crash during CUDA compilation.
auto x = &Kernel<double>;
auto y = &Kernel<float>;
auto z = &Kernel<int>;
// This triggers HIP-specific code path.
void func (){
Kernel<short><<<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