Merge pull request #1070 from pasaulais/pa/dlopen-cupti

[CUDA] Dynamically load the CUPTI library when tracing

source/adapters/cuda/CMakeLists.txt

@@ -76,6 +76,10 @@ else()
  message(WARNING "CUDA adapter USM pools are disabled, set UMF_ENABLE_POOL_TRACKING to enable them")
 endif()
 
+if (CUDA_cupti_LIBRARY)
+ target_compile_definitions("ur_adapter_cuda" PRIVATE CUPTI_LIB_PATH="${CUDA_cupti_LIBRARY}")
+endif()
+
 target_link_libraries(${TARGET_NAME} PRIVATE
  ${PROJECT_NAME}::headers
  ${PROJECT_NAME}::common

source/adapters/cuda/adapter.cpp

@@ -11,13 +11,12 @@
 #include <ur_api.h>
 
 #include "common.hpp"
-
-void enableCUDATracing();
-void disableCUDATracing();
+#include "tracing.hpp"
 
 struct ur_adapter_handle_t_ {
  std::atomic<uint32_t> RefCount = 0;
  std::mutex Mutex;
+ struct cuda_tracing_context_t_ *TracingCtx = nullptr;
 };
 
 ur_adapter_handle_t_ adapter{};
@@ -28,7 +27,8 @@ urAdapterGet(uint32_t NumEntries, ur_adapter_handle_t *phAdapters,
  if (NumEntries > 0 && phAdapters) {
  std::lock_guard<std::mutex> Lock{adapter.Mutex};
  if (adapter.RefCount++ == 0) {
- enableCUDATracing();
+ adapter.TracingCtx = createCUDATracingContext();
+ enableCUDATracing(adapter.TracingCtx);
  }
 
  *phAdapters = &adapter;
@@ -50,7 +50,9 @@ UR_APIEXPORT ur_result_t UR_APICALL urAdapterRetain(ur_adapter_handle_t) {
 UR_APIEXPORT ur_result_t UR_APICALL urAdapterRelease(ur_adapter_handle_t) {
  std::lock_guard<std::mutex> Lock{adapter.Mutex};
  if (--adapter.RefCount == 0) {
- disableCUDATracing();
+ disableCUDATracing(adapter.TracingCtx);
+ freeCUDATracingContext(adapter.TracingCtx);
+ adapter.TracingCtx = nullptr;
  }
  return UR_RESULT_SUCCESS;
 }

source/adapters/cuda/tracing.cpp

@@ -18,27 +18,77 @@
 #include <cupti.h>
 #endif // XPTI_ENABLE_INSTRUMENTATION
 
+#include "tracing.hpp"
+#include "ur_lib_loader.hpp"
 #include <exception>
 #include <iostream>
 
+#ifdef XPTI_ENABLE_INSTRUMENTATION
+using tracing_event_t = xpti_td *;
+using subscriber_handle_t = CUpti_SubscriberHandle;
+
+using cuptiSubscribe_fn = CUPTIAPI
+CUptiResult (*)(CUpti_SubscriberHandle *subscriber, CUpti_CallbackFunc callback,
+ void *userdata);
+
+using cuptiUnsubscribe_fn = CUPTIAPI
+CUptiResult (*)(CUpti_SubscriberHandle subscriber);
+
+using cuptiEnableDomain_fn = CUPTIAPI
+CUptiResult (*)(uint32_t enable, CUpti_SubscriberHandle subscriber,
+ CUpti_CallbackDomain domain);
+
+using cuptiEnableCallback_fn = CUPTIAPI
+CUptiResult (*)(uint32_t enable, CUpti_SubscriberHandle subscriber,
+ CUpti_CallbackDomain domain, CUpti_CallbackId cbid);
+
+#define LOAD_CUPTI_SYM(p, lib, x) \
+ p.x = (cupti##x##_fn)ur_loader::LibLoader::getFunctionPtr(lib.get(), \
+ "cupti" #x);
+
+#else
+using tracing_event_t = void *;
+using subscriber_handle_t = void *;
+using cuptiSubscribe_fn = void *;
+using cuptiUnsubscribe_fn = void *;
+using cuptiEnableDomain_fn = void *;
+using cuptiEnableCallback_fn = void *;
+#endif // XPTI_ENABLE_INSTRUMENTATION
+
+struct cupti_table_t_ {
+ cuptiSubscribe_fn Subscribe = nullptr;
+ cuptiUnsubscribe_fn Unsubscribe = nullptr;
+ cuptiEnableDomain_fn EnableDomain = nullptr;
+ cuptiEnableCallback_fn EnableCallback = nullptr;
+
+ bool isInitialized() const;
+};
+
+struct cuda_tracing_context_t_ {
+ tracing_event_t CallEvent = nullptr;
+ tracing_event_t DebugEvent = nullptr;
+ subscriber_handle_t Subscriber = nullptr;
+ ur_loader::LibLoader::Lib Library;
+ cupti_table_t_ Cupti;
+};
+
 #ifdef XPTI_ENABLE_INSTRUMENTATION
 constexpr auto CUDA_CALL_STREAM_NAME = "sycl.experimental.cuda.call";
 constexpr auto CUDA_DEBUG_STREAM_NAME = "sycl.experimental.cuda.debug";
 
 thread_local uint64_t CallCorrelationID = 0;
 thread_local uint64_t DebugCorrelationID = 0;
 
-static xpti_td *GCallEvent = nullptr;
-static xpti_td *GDebugEvent = nullptr;
-
 constexpr auto GVerStr = "0.1";
 constexpr int GMajVer = 0;
 constexpr int GMinVer = 1;
 
-static void cuptiCallback(void *, CUpti_CallbackDomain, CUpti_CallbackId CBID,
- const void *CBData) {
+static void cuptiCallback(void *UserData, CUpti_CallbackDomain,
+ CUpti_CallbackId CBID, const void *CBData) {
  if (xptiTraceEnabled()) {
  const auto *CBInfo = static_cast<const CUpti_CallbackData *>(CBData);
+ cuda_tracing_context_t_ *Ctx =
+ static_cast<cuda_tracing_context_t_ *>(UserData);
 
  if (CBInfo->callbackSite == CUPTI_API_ENTER) {
  CallCorrelationID = xptiGetUniqueId();
@@ -57,22 +107,95 @@ static void cuptiCallback(void *, CUpti_CallbackDomain, CUpti_CallbackId CBID,
  uint8_t CallStreamID = xptiRegisterStream(CUDA_CALL_STREAM_NAME);
  uint8_t DebugStreamID = xptiRegisterStream(CUDA_DEBUG_STREAM_NAME);
 
- xptiNotifySubscribers(CallStreamID, TraceType, GCallEvent, nullptr,
+ xptiNotifySubscribers(CallStreamID, TraceType, Ctx->CallEvent, nullptr,
  CallCorrelationID, FuncName);
 
  xpti::function_with_args_t Payload{
  FuncID, FuncName, const_cast<void *>(CBInfo->functionParams),
  CBInfo->functionReturnValue, CBInfo->context};
- xptiNotifySubscribers(DebugStreamID, TraceTypeArgs, GDebugEvent, nullptr,
- DebugCorrelationID, &Payload);
+ xptiNotifySubscribers(DebugStreamID, TraceTypeArgs, Ctx->DebugEvent,
+ nullptr, DebugCorrelationID, &Payload);
  }
 }
 #endif
 
+cuda_tracing_context_t_ *createCUDATracingContext() {
+#ifdef XPTI_ENABLE_INSTRUMENTATION
+ if (!xptiTraceEnabled())
+ return nullptr;
+ return new cuda_tracing_context_t_;
+#else
+ return nullptr;
+#endif // XPTI_ENABLE_INSTRUMENTATION
+}
+
+void freeCUDATracingContext(cuda_tracing_context_t_ *Ctx) {
+#ifdef XPTI_ENABLE_INSTRUMENTATION
+ unloadCUDATracingLibrary(Ctx);
+ delete Ctx;
+#else
+ (void)Ctx;
+#endif // XPTI_ENABLE_INSTRUMENTATION
+}
+
+bool cupti_table_t_::isInitialized() const {
+ return Subscribe && Unsubscribe && EnableDomain && EnableCallback;
+}
+
+bool loadCUDATracingLibrary(cuda_tracing_context_t_ *Ctx) {
+#if defined(XPTI_ENABLE_INSTRUMENTATION) && defined(CUPTI_LIB_PATH)
+ if (!Ctx)
+ return false;
+ if (Ctx->Library)
+ return true;
+ auto Lib{ur_loader::LibLoader::loadAdapterLibrary(CUPTI_LIB_PATH)};
+ if (!Lib)
+ return false;
+ cupti_table_t_ Table;
+ LOAD_CUPTI_SYM(Table, Lib, Subscribe)
+ LOAD_CUPTI_SYM(Table, Lib, Unsubscribe)
+ LOAD_CUPTI_SYM(Table, Lib, EnableDomain)
+ LOAD_CUPTI_SYM(Table, Lib, EnableCallback)
+ if (!Table.isInitialized()) {
+ return false;
+ }
+ Ctx->Library = std::move(Lib);
+ Ctx->Cupti = Table;
+ return true;
+#else
+ (void)Ctx;
+ return false;
+#endif // XPTI_ENABLE_INSTRUMENTATION && CUPTI_LIB_PATH
+}
+
+void unloadCUDATracingLibrary(cuda_tracing_context_t_ *Ctx) {
+#ifdef XPTI_ENABLE_INSTRUMENTATION
+ if (!Ctx)
+ return;
+ Ctx->Library.reset();
+ Ctx->Cupti = cupti_table_t_();
+#else
+ (void)Ctx;
+#endif // XPTI_ENABLE_INSTRUMENTATION
+}
+
 void enableCUDATracing() {
 #ifdef XPTI_ENABLE_INSTRUMENTATION
  if (!xptiTraceEnabled())
  return;
+ static cuda_tracing_context_t_ *Ctx = nullptr;
+ if (!Ctx)
+ Ctx = createCUDATracingContext();
+ enableCUDATracing(Ctx);
+#endif
+}
+
+void enableCUDATracing(cuda_tracing_context_t_ *Ctx) {
+#ifdef XPTI_ENABLE_INSTRUMENTATION
+ if (!Ctx || !xptiTraceEnabled())
+ return;
+ else if (!loadCUDATracingLibrary(Ctx))
+ return;
 
  xptiRegisterStream(CUDA_CALL_STREAM_NAME);
  xptiInitialize(CUDA_CALL_STREAM_NAME, GMajVer, GMinVer, GVerStr);
@@ -81,31 +204,39 @@ void enableCUDATracing() {
 
  uint64_t Dummy;
  xpti::payload_t CUDAPayload("CUDA Plugin Layer");
- GCallEvent =
+ Ctx->CallEvent =
  xptiMakeEvent("CUDA Plugin Layer", &CUDAPayload,
  xpti::trace_algorithm_event, xpti_at::active, &Dummy);
 
  xpti::payload_t CUDADebugPayload("CUDA Plugin Debug Layer");
- GDebugEvent =
+ Ctx->DebugEvent =
  xptiMakeEvent("CUDA Plugin Debug Layer", &CUDADebugPayload,
  xpti::trace_algorithm_event, xpti_at::active, &Dummy);
 
- CUpti_SubscriberHandle Subscriber;
- cuptiSubscribe(&Subscriber, cuptiCallback, nullptr);
- cuptiEnableDomain(1, Subscriber, CUPTI_CB_DOMAIN_DRIVER_API);
- cuptiEnableCallback(0, Subscriber, CUPTI_CB_DOMAIN_DRIVER_API,
- CUPTI_DRIVER_TRACE_CBID_cuGetErrorString);
- cuptiEnableCallback(0, Subscriber, CUPTI_CB_DOMAIN_DRIVER_API,
- CUPTI_DRIVER_TRACE_CBID_cuGetErrorName);
+ Ctx->Cupti.Subscribe(&Ctx->Subscriber, cuptiCallback, Ctx);
+ Ctx->Cupti.EnableDomain(1, Ctx->Subscriber, CUPTI_CB_DOMAIN_DRIVER_API);
+ Ctx->Cupti.EnableCallback(0, Ctx->Subscriber, CUPTI_CB_DOMAIN_DRIVER_API,
+ CUPTI_DRIVER_TRACE_CBID_cuGetErrorString);
+ Ctx->Cupti.EnableCallback(0, Ctx->Subscriber, CUPTI_CB_DOMAIN_DRIVER_API,
+ CUPTI_DRIVER_TRACE_CBID_cuGetErrorName);
+#else
+ (void)Ctx;
 #endif
 }
 
-void disableCUDATracing() {
+void disableCUDATracing(cuda_tracing_context_t_ *Ctx) {
 #ifdef XPTI_ENABLE_INSTRUMENTATION
- if (!xptiTraceEnabled())
+ if (!Ctx || !xptiTraceEnabled())
  return;
 
+ if (Ctx->Subscriber && Ctx->Cupti.isInitialized()) {
+ Ctx->Cupti.Unsubscribe(Ctx->Subscriber);
+ Ctx->Subscriber = nullptr;
+ }
+
  xptiFinalize(CUDA_CALL_STREAM_NAME);
  xptiFinalize(CUDA_DEBUG_STREAM_NAME);
+#else
+ (void)Ctx;
 #endif // XPTI_ENABLE_INSTRUMENTATION
 }

source/adapters/cuda/tracing.hpp

@@ -0,0 +1,24 @@
+//===--------- tracing.hpp - CUDA Host API Tracing -------------------------==//
+//
+// Copyright (C) 2023 Intel Corporation
+//
+// Part of the Unified-Runtime Project, under the Apache License v2.0 with LLVM
+// Exceptions. See LICENSE.TXT
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+struct cuda_tracing_context_t_;
+
+cuda_tracing_context_t_ *createCUDATracingContext();
+void freeCUDATracingContext(cuda_tracing_context_t_ *Ctx);
+
+bool loadCUDATracingLibrary(cuda_tracing_context_t_ *Ctx);
+void unloadCUDATracingLibrary(cuda_tracing_context_t_ *Ctx);
+
+void enableCUDATracing(cuda_tracing_context_t_ *Ctx);
+void disableCUDATracing(cuda_tracing_context_t_ *Ctx);
+
+// Deprecated. Will be removed once pi_cuda has been updated to use the variant
+// that takes a context pointer.
+void enableCUDATracing();

source/common/ur_lib_loader.hpp

@@ -27,8 +27,9 @@ class LibLoader {
  void operator()(HMODULE handle) { freeAdapterLibrary(handle); }
  };
 
- static std::unique_ptr<HMODULE, lib_dtor>
- loadAdapterLibrary(const char *name);
+ using Lib = std::unique_ptr<HMODULE, lib_dtor>;
+
+ static Lib loadAdapterLibrary(const char *name);
 
  static void freeAdapterLibrary(HMODULE handle);