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

Deallocate replaced Var.


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,15 +1195,28 @@
 llvm::GlobalValue *CGNVCUDARuntime::getKernelHandle(llvm::Function *F,
                                                     GlobalDecl GD) {
   auto Loc = KernelHandles.find(F->getName());
-  if (Loc != KernelHandles.end())
-    return Loc->second;
+  llvm::GlobalValue *OldHandle = nullptr;
+  if (Loc != KernelHandles.end()) {
+    OldHandle = Loc->second;
+    auto OrigF = KernelStubs[OldHandle];
+    if (OrigF == 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) {
     KernelHandles[F->getName()] = F;
+    if (OldHandle)
+      KernelStubs.erase(OldHandle);
     KernelStubs[F] = F;
     return F;
   }
 
+  // Make sure that the old variable name is no longer used.
+  if (OldHandle)
+    OldHandle->setName("");
   auto *Var = new llvm::GlobalVariable(
       TheModule, F->getType(), /*isConstant=*/true, F->getLinkage(),
       /*Initializer=*/nullptr,
@@ -1213,6 +1226,11 @@
   Var->setDSOLocal(F->isDSOLocal());
   Var->setVisibility(F->getVisibility());
   CGM.maybeSetTrivialComdat(*GD.getDecl(), *Var);
+  if (OldHandle) {
+    KernelStubs.erase(OldHandle);
+    OldHandle->replaceAllUsesWith(Var);
+    OldHandle->eraseFromParent();
+  }
   KernelHandles[F->getName()] = Var;
   KernelStubs[Var] = F;
   return Var;


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