[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 507101.
tra added a comment.

Simplify invalidation logic. No need to replace HIP's handle variable.


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,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
Index: clang/lib/CodeGen/CGCUDANV.cpp
===================================================================
--- clang/lib/CodeGen/CGCUDANV.cpp
+++ 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;


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D146448.507101.patch
Type: text/x-patch
Size: 2812 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/cfe-commits/attachments/20230321/3c3cee18/attachment.bin>