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

[tra]

Summary

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.

[dcastagna]

Thank you Artem!

clang/lib/CodeGen/CGCUDANV.cpp

1201

The different F is because sometimes at first we use an F with incomplete types as a key and later we use a different F (even if it's coming from the same function) with the instantiated types (usually after we encounter a statement where the size of the template parameter is necessary, like a function call).
Another option could be to always instantiate the complete type for F when we call GetOrCreateLLVMFunction (IIRC) if the function has a global attribute.

If you prefer this solution, the patch LGTM.

[tra]

clang/lib/CodeGen/CGCUDANV.cpp

1201

Instantiating a complete function when we normally would not sounds a bit too invasive and may buy us unintended side effects. E.g. I can see it potentially causing code explosion in some cases, where we previously could get by with just a function decl.

Fixing the cache locally maintains the global status quo which seems to be about the right balance between improving the situation, the effort involved, and the potential risk of the fix.

The way we track kernels could potentially be improved (I think it would be better to track them explicitly vs relying on observing the functions flying by), but that should be a separate patch. Right now I just want to fix the regression w/o affecting anything else.

[tra]

Fixed a typo in the test.

[yaxunl]

Do you know whether the issue happens to HIP too? Thanks.

[tra]

In D146448#4207591, @yaxunl wrote:

Do you know whether the issue happens to HIP too? Thanks.

Possibly, but I'm unable to reproduce the crash: https://godbolt.org/z/Gsx9sWG4v

The value cached by KernelHandles in HIP is the the manually created 'Var' and not the function itself. I don't think it ever changes outside of this function.

However, KernelStubs would still point to the original F, and that may potentially be left dangling. Perhaps I should always update both handles and stubs if F changes.

[tra]

In D146448#4207800, @tra wrote:

However, KernelStubs would still point to the original F, and that may potentially be left dangling. Perhaps I should always update both handles and stubs if F changes.

Yes, HIP is susceptibel to having the dangling pointer to F. What saves us is that KernelStubs happens to be used only *before* the cached functions are reinstantiated, so we never use those dangling pointers in case of HIP.
OK, looks like I'll need to fix HIP, too.

[tra]

Improved kernel handle cache invalidation.

[tra]

Deallocate replaced Var.

[tra]

Cosmetic cleanup

[tra]

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

[tra]

I've updated the patch with a more robust version of the handle/stub invalidation.

[yaxunl]

LGTM. Thanks!

clang/lib/CodeGen/CGCUDANV.cpp

else
     return makeModuleCtorFunction();
 
   return nullptr;
 }
 
 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;
     KernelStubs[F] = F;
     return F;
   }
 
   auto *Var = new llvm::GlobalVariable(

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

+// 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