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

[Artem Belevich via Phabricator via cfe-commits]

tra updated this revision to Diff 506719.
tra added a comment.

Fixed a typo in the test.


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D146448/new/

https://reviews.llvm.org/D146448

Files:
  clang/lib/CodeGen/CGCUDANV.cpp
  clang/test/CodeGenCUDA/bug-kerner-registration-reuse.cu


Index: clang/test/CodeGenCUDA/bug-kerner-registration-reuse.cu
===================================================================
--- /dev/null
+++ clang/test/CodeGenCUDA/bug-kerner-registration-reuse.cu
@@ -0,0 +1,24 @@
+// 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
+
+#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.
+auto x = &Kernel<double>;
+auto y = &Kernel<float>;
+auto z = &Kernel<int>;
+
+// CHECK-LABEL: @__cuda_register_globals(
+// CHECK:  call i32 @__cudaRegisterFunction(ptr %0, ptr @_ZL21__device_stub__KernelIdEv1SIT_E
+// CHECK:  call i32 @__cudaRegisterFunction(ptr %0, ptr @_ZL21__device_stub__KernelIfEv1SIT_E
+// CHECK:  call i32 @__cudaRegisterFunction(ptr %0, ptr @_ZL21__device_stub__KernelIiEv1SIT_E
+// CHECK: ret void
Index: clang/lib/CodeGen/CGCUDANV.cpp
===================================================================
--- clang/lib/CodeGen/CGCUDANV.cpp
+++ clang/lib/CodeGen/CGCUDANV.cpp
@@ -1195,8 +1195,16 @@
 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()) {
+    if (Loc->second == F || CGM.getLangOpts().HIP)
+      return Loc->second;
+    // non-HIP compilation may end up with a different F and need to have
+    // handles and stubs updated.
+    KernelStubs.erase(Loc->second);
+    KernelStubs[F] = F;
+    KernelHandles[F->getName()] = F;
+    return F;
+  }
 
   if (!CGM.getLangOpts().HIP) {
     KernelHandles[F->getName()] = F;


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D146448.506719.patch
Type: text/x-patch
Size: 1957 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/cfe-commits/attachments/20230320/b46d24d0/attachment.bin>