Try and detect when the loader is initializing through us and forward…

… dispatch tables in this case.
argonne-lcf · Feb 5, 2025 · 0a5442e · 0a5442e
1 parent 98af8e3
commit 0a5442e
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 9 deletions.
diff --git a/ze/tracer_ze_helpers.include.c b/ze/tracer_ze_helpers.include.c
@@ -625,11 +625,12 @@ static void _on_destroy_command_list(ze_command_list_handle_t command_list) {
 }
 
 static pthread_once_t _init = PTHREAD_ONCE_INIT;
-static __thread volatile int in_init = 0;
+static __thread volatile int _in_init = 0;
+static volatile unsigned int _in_loader_init = 0;
 static volatile unsigned int _initialized = 0;
 
 static pthread_once_t _init_dump = PTHREAD_ONCE_INIT;
-static __thread volatile int in_init_dump = 0;
+static __thread volatile int _in_init_dump = 0;
 static volatile unsigned int _initialized_dump = 0;
 
 static inline int _do_state() {
@@ -859,12 +860,12 @@ static inline 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);
     __sync_synchronize();
-    in_init=0;
+    _in_init=0;
   }
   _initialized = 1;
 }
@@ -873,12 +874,12 @@ static inline void _init_tracer_dump(void) {
   if( __builtin_expect (_initialized_dump, 1) )
     return;
   /* Avoid reentrancy */
-  if (!in_init_dump) {
-    in_init_dump=1;
+  if (!_in_init_dump) {
+    _in_init_dump=1;
     __sync_synchronize();
     pthread_once(&_init_dump, _load_tracer_dump);
     __sync_synchronize();
-    in_init_dump=0;
+    _in_init_dump=0;
   }
   _initialized_dump = 1;
 }
diff --git a/ze/ze_model.rb b/ze/ze_model.rb
@@ -387,7 +387,7 @@ def upper_snake_case(str)
   child_types = STRUCT_MAP.select { |k, _| k.match(parent_type) }
   str = <<EOF
   #{c.type} _retval;
-  if (!_do_ddi_table_forward && pDdiTable && version <= ZE_API_VERSION_CURRENT) {
+  if (!_do_ddi_table_forward && !_in_loader_init && pDdiTable && version <= ZE_API_VERSION_CURRENT) {
 EOF
   str << "   "
   str << child_types.reverse_each.collect { |k, v|
@@ -485,3 +485,11 @@ def get_ffi_type(a)
     _dump_memory_info_ctx(hContext, *pptr);
   }
 EOF
+
+register_prologue "zeLoaderInit", <<EOF
+  _in_loader_init = 1;
+EOF
+
+register_epilogue "zeInit", <<EOF
+  _in_loader_init = 0;
+EOF