From 0e19bd4dd411c3a40ad9ba0ecabe570cd14acae4 Mon Sep 17 00:00:00 2001
From: Brice Videau <bvideau@anl.gov>
Date: Thu, 31 Aug 2023 14:40:46 -0500
Subject: [PATCH] Add stub implementation of CUDA symbols to avoid segfaulting.
 (#140)

* Add stub implementation of CUDA symbols to avoid segfaulting.

* Try and be smart.

* Change at_exit location.
---
 cuda/cuda_model.rb                 |  8 +++--
 cuda/gen_cuda.rb                   | 46 ++++++++++++++++++++++++-----
 cuda/tracer_cuda_helpers.include.c | 47 ++++++++++++++++++++----------
 3 files changed, 76 insertions(+), 25 deletions(-)

diff --git a/cuda/cuda_model.rb b/cuda/cuda_model.rb
index 92d1af38..0d49bf1d 100644
--- a/cuda/cuda_model.rb
+++ b/cuda/cuda_model.rb
@@ -34,8 +34,6 @@
 gen_struct_map(typedefs, structs)
 gen_ffi_type_map(typedefs)
 
-INIT_FUNCTIONS = /cuInit|cuDriverGetVersion|cuGetExportTable|cuDeviceGetCount|cuGetProcAddress/
-
 HEX_INT_TYPES.push("CUdeviceptr")
 
 class TracepointParameter
@@ -286,6 +284,12 @@ def upper_snake_case(str)
   }
 EOF
 
+register_epilogue "cuInit", <<EOF
+  if (_retval == CUDA_SUCCESS) {
+    _init_cuda();
+  }
+EOF
+
 # cuGetProcAddress*
 command_names = $cuda_commands.collect(&:name).to_set
 pt_condition = "((flags & CU_GET_PROC_ADDRESS_PER_THREAD_DEFAULT_STREAM) && !(flags & CU_GET_PROC_ADDRESS_LEGACY_STREAM))"
diff --git a/cuda/gen_cuda.rb b/cuda/gen_cuda.rb
index 06a7f1de..a556f61a 100644
--- a/cuda/gen_cuda.rb
+++ b/cuda/gen_cuda.rb
@@ -11,6 +11,8 @@
 #include "cuda_properties.h"
 #include "utlist.h"
 #include "uthash.h"
+
+static void _init_tracer(void);
 EOF
 #puts <<EOF
 ##include <ffi.h>
@@ -20,7 +22,37 @@
   puts "#define #{CUDA_POINTER_NAMES[c]} #{c.pointer_name}"
 }
 
-($cuda_commands + $cuda_exports_commands).each { |c|
+($cuda_commands).each { |c|
+  puts <<EOF
+
+static #{YAMLCAst::Declaration::new(name: c.name + "_unsupp", type: c.function.type)} {
+  #{c.parameters.map(&:name).map { |n| "(void)#{n};" }.join("\n  ")}
+  fprintf(stderr, "THAPI: #{c.name} was called, but it is unsupported by the driver\\n");
+  return CUDA_ERROR_NOT_SUPPORTED;
+}
+static #{YAMLCAst::Declaration::new(name: c.name + "_uninit", type: c.function.type)};
+#{c.decl_pointer(c.pointer_type_name)};
+static #{c.pointer_type_name} #{CUDA_POINTER_NAMES[c]} = (void *)&#{c.name}_uninit;
+static #{YAMLCAst::Declaration::new(name: c.name + "_uninit", type: c.function.type)} {
+  #{c.parameters.map(&:name).map { |n| "(void)#{n};" }.join("\n  ")}
+  _init_tracer();
+EOF
+  params = c.parameters.collect(&:name)
+  if c.has_return_type?
+    puts <<EOF
+  return #{CUDA_POINTER_NAMES[c]}(#{params.join(", ")});
+EOF
+  else
+    puts <<EOF
+  #{CUDA_POINTER_NAMES[c]}(#{params.join(", ")});
+EOF
+  end
+  puts <<EOF
+}
+EOF
+}
+
+($cuda_exports_commands).each { |c|
   puts <<EOF
 
 #{c.decl_pointer(c.pointer_type_name)};
@@ -47,8 +79,11 @@
   puts <<EOF
 
   #{CUDA_POINTER_NAMES[c]} = (#{c.pointer_type_name})(intptr_t)dlsym(handle, "#{c.name}");
-  if (!#{CUDA_POINTER_NAMES[c]} && verbose)
-    fprintf(stderr, "Missing symbol #{c.name}!\\n");
+  if (!#{CUDA_POINTER_NAMES[c]}) {
+    #{CUDA_POINTER_NAMES[c]} = &#{c.name}_unsupp;
+    if (verbose)
+      fprintf(stderr, "THAPI: Missing symbol #{c.name}!\\n");
+  }
 EOF
 }
 
@@ -147,11 +182,6 @@
   puts <<EOF
 #{c.decl} {
 EOF
-  if c.init?
-    puts <<EOF
-  _init_tracer();
-EOF
-  end
   common_block.call(c, provider)
   if c.has_return_type?
     puts <<EOF
diff --git a/cuda/tracer_cuda_helpers.include.c b/cuda/tracer_cuda_helpers.include.c
index 53ede7ec..1e4bf9db 100644
--- a/cuda/tracer_cuda_helpers.include.c
+++ b/cuda/tracer_cuda_helpers.include.c
@@ -480,10 +480,6 @@ static void _lib_cleanup() {
   }
 }
 
-static pthread_once_t _init = PTHREAD_ONCE_INIT;
-static __thread volatile int in_init = 0;
-static volatile int _initialized = 0;
-
 static void _load_tracer(void) {
   char *s = NULL;
   void *handle = NULL;
@@ -503,7 +499,7 @@ static void _load_tracer(void) {
   }
 
   if( !handle ) {
-    fprintf(stderr, "Failure: could not load cuda library!\n");
+    fprintf(stderr, "THAPI: Failure: could not load cuda library!\n");
     exit(1);
   }
 
@@ -512,32 +508,53 @@ static void _load_tracer(void) {
     verbose = 1;
 
   find_cuda_symbols(handle, verbose);
-  CU_INIT_PTR(0);
   find_cuda_extensions();
 
-  _dump_properties();
-
   s = getenv("LTTNG_UST_CUDA_PROFILE");
   if (s)
     _do_profile = 1;
   if(tracepoint_enabled(lttng_ust_cuda_exports, export_called))
     _do_trace_export_tables = 1;
 
-  if (_do_profile)
-    atexit(&_lib_cleanup);
 }
 
-static inline void _init_tracer(void) {
+static pthread_once_t _init_tracer_once = PTHREAD_ONCE_INIT;
+static __thread volatile int _in_init = 0;
+static volatile int _initialized = 0;
+
+static void _init_tracer(void) {
   if( __builtin_expect (_initialized, 1) )
     return;
   /* Avoid reentrancy */
-  if (!in_init) {
-    in_init=1;
+  if (!_in_init) {
+    _in_init = 1;
     __sync_synchronize();
-    pthread_once(&_init, _load_tracer);
+    pthread_once(&_init_tracer_once, _load_tracer);
     __sync_synchronize();
-    in_init=0;
+    _in_init = 0;
   }
   _initialized = 1;
 }
 
+static void _load_cuda(void) {
+  _dump_properties();
+  if (_do_profile)
+    atexit(&_lib_cleanup);
+}
+
+static pthread_once_t _init_cuda_once = PTHREAD_ONCE_INIT;
+static __thread volatile int _in_init_cuda = 0;
+static volatile int _initialized_cuda = 0;
+
+static void _init_cuda(void) {
+  if( __builtin_expect (_initialized_cuda, 1) )
+    return;
+  if (!_in_init_cuda) {
+    _in_init_cuda = 1;
+    __sync_synchronize();
+    pthread_once(&_init_cuda_once, _load_cuda);
+    __sync_synchronize();
+    _in_init_cuda = 0;
+  }
+  _initialized_cuda = 1;
+}