Merge pull request #2000 from igchor/tracing_fix

[Loader] Fix xpti finalization order

Author: omarahmed1111

source/loader/layers/tracing/ur_tracing_layer.cpp

@@ -29,36 +29,19 @@ constexpr auto STREAM_VER_MINOR = UR_MINOR_VERSION(UR_API_VERSION_CURRENT);
 // Unfortunately this doesn't match the semantics of XPTI, which can be initialized
 // and finalized exactly once. To workaround this, XPTI is globally initialized on
 // first use and finalized in the destructor.
-class XptiContext {
-    XptiContext() {
-        xptiFrameworkInitialize();
-        inited = true;
-    }
-
-    ~XptiContext() {
-        xptiFrameworkFinalize();
-        inited = false;
-    }
-
-    // Accessing this after destruction is technically UB, but if we get there,
-    // it means something is calling UR after it has been destroyed at program
-    // exit.
-    std::atomic_bool inited;
-
-  public:
-    static bool running() {
-        static XptiContext context;
-        return context.inited;
-    }
+struct XptiContextManager {
+    XptiContextManager() { xptiFrameworkInitialize(); }
+    ~XptiContextManager() { xptiFrameworkFinalize(); }
 };
 
+static std::shared_ptr<XptiContextManager> xptiContextManagerGlobal = [] {
+    return std::make_shared<XptiContextManager>();
+}();
 static thread_local xpti_td *activeEvent;
 
 ///////////////////////////////////////////////////////////////////////////////
 context_t::context_t() : logger(logger::create_logger("tracing", true, true)) {
-    if (!XptiContext::running()) {
-        return;
-    }
+    this->xptiContextManager = xptiContextManagerGlobal;
 
     call_stream_id = xptiRegisterStream(CALL_STREAM_NAME);
     std::ostringstream streamv;
@@ -69,20 +52,12 @@ context_t::context_t() : logger(logger::create_logger("tracing", true, true)) {
 
 void context_t::notify(uint16_t trace_type, uint32_t id, const char *name,
                        void *args, ur_result_t *resultp, uint64_t instance) {
-    if (!XptiContext::running()) {
-        return;
-    }
-
     xpti::function_with_args_t payload{id, name, args, resultp, nullptr};
     xptiNotifySubscribers(call_stream_id, trace_type, nullptr, activeEvent,
                           instance, &payload);
 }
 
 uint64_t context_t::notify_begin(uint32_t id, const char *name, void *args) {
-    if (!XptiContext::running()) {
-        return 0;
-    }
-
     if (auto loc = codelocData.get_codeloc()) {
         xpti::payload_t payload =
             xpti::payload_t(loc->functionName, loc->sourceFile, loc->lineNumber,
@@ -101,20 +76,10 @@ uint64_t context_t::notify_begin(uint32_t id, const char *name, void *args) {
 
 void context_t::notify_end(uint32_t id, const char *name, void *args,
                            ur_result_t *resultp, uint64_t instance) {
-    if (!XptiContext::running()) {
-        return;
-    }
-
     notify((uint16_t)xpti::trace_point_type_t::function_with_args_end, id, name,
            args, resultp, instance);
 }
 
 ///////////////////////////////////////////////////////////////////////////////
-context_t::~context_t() {
-    if (!XptiContext::running()) {
-        return;
-    }
-
-    xptiFinalize(CALL_STREAM_NAME);
-}
+context_t::~context_t() { xptiFinalize(CALL_STREAM_NAME); }
 } // namespace ur_tracing_layer

source/loader/layers/tracing/ur_tracing_layer.hpp

@@ -21,6 +21,8 @@
 #define TRACING_COMP_NAME "tracing layer"
 
 namespace ur_tracing_layer {
+struct XptiContextManager;
+
 ///////////////////////////////////////////////////////////////////////////////
 class __urdlllocal context_t : public proxy_layer_context_t,
                                public AtomicSingleton<context_t> {
@@ -47,6 +49,8 @@ class __urdlllocal context_t : public proxy_layer_context_t,
     uint8_t call_stream_id;
 
     inline static const std::string name = "UR_LAYER_TRACING";
+
+    std::shared_ptr<XptiContextManager> xptiContextManager;
 };
 
 context_t *getContext();