Refactor

lib/JitEngineDevice.hpp

@@ -583,10 +583,17 @@ void JitEngineDevice<ImplT>::registerFunction(void *Handle, void *Kernel,
                                               int32_t NumRCs) {
   DBG(Logger::logs("proteus")
       << "Register function " << Kernel << " To Handle " << Handle << "\n");
-  assert(!KernelToHandleMap.contains(Kernel) &&
-         "Expected kernel inserted only once in the map");
-  //if (KernelToHandleMap.contains(Kernel))
-  //  return;
+  // NOTE: HIP RDC might call multiple times the registerFunction for the same
+  // kernel, which has weak linkage, when it comes from different translation
+  // units. Either the first or the second call can prevail and should be
+  // equivalent. We let the first one prevail.
+  if (KernelToHandleMap.contains(Kernel)) {
+    DBG(Logger::logs("proteus")
+        << "Warning: duplicate register function for kernel " +
+               std::string(KernelName)
+        << "\n");
+    return;
+  }
   KernelToHandleMap[Kernel] = Handle;
 
   JITKernelInfoMap[Kernel] =

pass/ProteusPass.cpp

@@ -250,8 +250,10 @@ class ProteusJitPassImpl {
       }
 
       if (JitFunctionInfoMap.contains(Fn)) {
-        // continue;
-        FATAL_ERROR("Duplicate jit annotation for Fn " + Fn->getName());
+        DEBUG(Logger::logs("proteus-pass")
+              << "Warning: Duplicate jit annotation for Fn " + Fn->getName() +
+                     "\n");
+        continue;
       }
 
       DEBUG(Logger::logs("proteus-pass")

tests/gpu/file1_kernel_launcher.cpp

@@ -1,24 +1,26 @@
+// clang-format off
 // RUN: ./multi_file_launcher.%ext | FileCheck %s --check-prefixes=CHECK,CHECK-FIRST
 // Second run uses the object cache.
 // RUN: ./multi_file_launcher.%ext | FileCheck %s --check-prefixes=CHECK,CHECK-SECOND
+// clang-format on
 #include <climits>
 #include <cstdio>
 
 #include "gpu_common.h"
 #include "launcher.hpp"
 
-
 void foo();
 
 int main() {
-  launcher(my_kernel_body);
+  gpuErrCheck(launcher(kernel_body));
   gpuErrCheck(gpuDeviceSynchronize());
   foo();
+  gpuErrCheck(gpuDeviceSynchronize());
   return 0;
 }
 
-// CHECK: File1 Kernel
-// CHECK: File2 Kernel
+// CHECK: Kernel body
+// CHECK: Kernel body
 // CHECK: JitCache hits 0 total 2
 // CHECK: HashValue {{[0-9]+}} NumExecs 1 NumHits 0
 // CHECK: HashValue {{[0-9]+}} NumExecs 1 NumHits 0

tests/gpu/file2_kernel_launcher.cpp

@@ -3,6 +3,4 @@
 #include "gpu_common.h"
 #include "launcher.hpp"
 
-void foo() {
-  launcher(my_kernel_body);
-}
+void foo() { gpuErrCheck(launcher(kernel_body)); }

tests/gpu/gpu_common.h

@@ -29,9 +29,3 @@
       abort();                                                                 \
     }                                                                          \
   }
-
-struct kernel_body {
-  __device__ void operator()() { printf("Kernel"); }
-};
-
-const kernel_body my_kernel_body{};

tests/gpu/launcher.hpp

@@ -1,13 +1,16 @@
-template<typename LB>
-__global__ 
-__attribute__((annotate("jit")))
-void kernel(LB lb) {
+struct kernel_body_t {
+  __device__ void operator()() { printf("Kernel body"); }
+};
+
+const kernel_body_t kernel_body{};
+
+template <typename LB>
+__global__ __attribute__((annotate("jit"))) void kernel(LB lb) {
   lb();
 }
 
-template <typename T> 
-gpuError_t launcher(T lb) {
-   auto func = reinterpret_cast<const void*>( & kernel<T> );
-   void *args[] = {(void*)&lb };
+template <typename T> gpuError_t launcher(T lb) {
+  auto func = reinterpret_cast<const void *>(&kernel<T>);
+  void *args[] = {(void *)&lb};
   return gpuLaunchKernel(func, 1, 1, args, 0, 0);
 }