[clang] 2aa90da - [CUDA] Update cached kernel handle when the function instance changes.

[Artem Belevich via cfe-commits]

Author: Artem Belevich
Date: 2023-03-21T15:36:12-07:00
New Revision: 2aa90da012596712a4166e96d2a40fc90598c7fb

URL: https://github.com/llvm/llvm-project/commit/2aa90da012596712a4166e96d2a40fc90598c7fb
DIFF: https://github.com/llvm/llvm-project/commit/2aa90da012596712a4166e96d2a40fc90598c7fb.diff

LOG: [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

Added: 
    clang/test/CodeGenCUDA/bug-kerner-registration-reuse.cu

Modified: 
    clang/lib/CodeGen/CGCUDANV.cpp

Removed: 
    


################################################################################
diff  --git a/clang/lib/CodeGen/CGCUDANV.cpp b/clang/lib/CodeGen/CGCUDANV.cpp
index bb887df3e4e04..1f429e4305790 100644
--- a/clang/lib/CodeGen/CGCUDANV.cpp
+++ b/clang/lib/CodeGen/CGCUDANV.cpp
@@ -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;

diff  --git a/clang/test/CodeGenCUDA/bug-kerner-registration-reuse.cu b/clang/test/CodeGenCUDA/bug-kerner-registration-reuse.cu
new file mode 100644
index 0000000000000..e1fef41e25079
--- /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 <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