diff --git a/libkineto/src/MuptiActivityApi.cpp b/libkineto/src/MuptiActivityApi.cpp index 629af9b..9a76f2f 100644 --- a/libkineto/src/MuptiActivityApi.cpp +++ b/libkineto/src/MuptiActivityApi.cpp @@ -57,8 +57,8 @@ inline void reenableMuptiCallbacks_(std::shared_ptr& cbapi_) { #endif MuptiActivityApi& MuptiActivityApi::singleton() { - static MuptiActivityApi instance; - return instance; + static auto* instance = new MuptiActivityApi(); + return *instance; } void MuptiActivityApi::pushCorrelationID(int id, CorrelationFlowType type) { @@ -179,6 +179,9 @@ std::unique_ptr MuptiActivityApi::activityBuffers() { { std::lock_guard guard(mutex_); if (allocatedGpuTraceBuffers_.empty()) { + if (readyGpuTraceBuffers_) { + return std::move(readyGpuTraceBuffers_); + } return nullptr; } } @@ -272,25 +275,29 @@ void MuptiActivityApi::bufferCompleted( size_t /* unused */, size_t validSize) { - std::lock_guard guard(mutex_); - auto it = allocatedGpuTraceBuffers_.find(buffer); - if (it == allocatedGpuTraceBuffers_.end()) { - LOG(ERROR) << "bufferCompleted called with unknown buffer: " - << (void*) buffer; - return; - } + { + std::lock_guard guard(mutex_); + auto it = allocatedGpuTraceBuffers_.find(buffer); + if (it == allocatedGpuTraceBuffers_.end()) { + LOG(ERROR) << "bufferCompleted called with unknown buffer: " + << static_cast(buffer); + return; + } - if (!readyGpuTraceBuffers_) { - readyGpuTraceBuffers_ = std::make_unique(); + if (!readyGpuTraceBuffers_) { + readyGpuTraceBuffers_ = std::make_unique(); + } + // Set valid size of buffer before moving to ready map + it->second->setSize(validSize); + (*readyGpuTraceBuffers_)[it->first] = std::move(it->second); + allocatedGpuTraceBuffers_.erase(it); } - // Set valid size of buffer before moving to ready map - it->second->setSize(validSize); - (*readyGpuTraceBuffers_)[it->first] = std::move(it->second); - allocatedGpuTraceBuffers_.erase(it); // report any records dropped from the queue; to avoid unnecessary mupti // API calls, we make it report only in verbose mode (it doesn't happen // often in our testing anyways) + // Can't hold mutex_ during this call, since muptiActivityGetNumDroppedRecords + // can acquire MUPTI's internal lock while MUPTI callbacks acquire mutex_. if (VLOG_IS_ON(1)) { size_t dropped = 0; MUPTI_CALL(muptiActivityGetNumDroppedRecords(ctx, streamId, &dropped)); @@ -399,7 +406,11 @@ void MuptiActivityApi::teardownContext() { if (!tracingEnabled_) { return; } + if (tearingDown_) { + return; + } if (muptiTearDown_()) { + tearingDown_ = 1; LOG(INFO) << "teardownMupti starting"; // PyTorch Profiler is synchronous, so teardown needs to be run async in this thread. @@ -409,6 +420,7 @@ void MuptiActivityApi::teardownContext() { cbapi_->initCallbackApi(); if (!cbapi_->initSuccess()) { LOG(WARNING) << "MUPTI Callback failed to init, skipping teardown"; + tearingDown_ = 0; return; } } @@ -417,6 +429,7 @@ void MuptiActivityApi::teardownContext() { status = status && cbapi_->enableCallbackDomain(MUPTI_CB_DOMAIN_DRIVER_API); if (!status) { LOG(WARNING) << "MUPTI Callback failed to enable for domain, skipping teardown"; + tearingDown_ = 0; return; } @@ -442,6 +455,7 @@ void MuptiActivityApi::teardownContext() { reenableMuptiCallbacks_(cbapi_); } cbapi_.reset(); + tearingDown_ = 0; }); teardownThread.detach(); } diff --git a/libkineto/src/MuptiActivityApi.h b/libkineto/src/MuptiActivityApi.h index 8005bf6..75bf637 100644 --- a/libkineto/src/MuptiActivityApi.h +++ b/libkineto/src/MuptiActivityApi.h @@ -90,6 +90,7 @@ class MuptiActivityApi { std::unique_ptr readyGpuTraceBuffers_; std::mutex mutex_; std::atomic tracingEnabled_{0}; + std::atomic tearingDown_{0}; bool externalCorrelationEnabled_{false}; #ifdef HAS_MUPTI diff --git a/libkineto/src/MuptiActivityProfiler.cpp b/libkineto/src/MuptiActivityProfiler.cpp index c2624d6..28189ba 100644 --- a/libkineto/src/MuptiActivityProfiler.cpp +++ b/libkineto/src/MuptiActivityProfiler.cpp @@ -294,6 +294,13 @@ void MuptiActivityProfiler::processTraceInternal(ActivityLogger& logger) { addOverheadSample(flushOverhead_, mupti_.flushOverhead); } if (traceBuffers_->gpu) { + // Pass 1: Preprocess all raw records to populate correlation and + // context lookup state. + buildProcessingState(*traceBuffers_->gpu); + + // Pass 2: Materialize activities. buildProcessingState() has already + // populated correlation and context lookup state; + // EXTERNAL_CORRELATION is a no-op in handleMuptiActivity. const auto count_and_size = mupti_.processActivities( *traceBuffers_->gpu, std::bind( @@ -390,6 +397,20 @@ void MuptiActivityProfiler::processCpuTrace( } #ifdef HAS_MUPTI +void MuptiActivityProfiler::buildProcessingState( + MuptiActivityBufferMap& buffers) { + mupti_.processActivities(buffers, [this](const MUpti_Activity* record) { + switch (record->kind) { + case MUPTI_ACTIVITY_KIND_EXTERNAL_CORRELATION: + handleCorrelationActivity( + reinterpret_cast(record)); + break; + default: + break; + } + }); +} + inline void MuptiActivityProfiler::handleCorrelationActivity( const MUpti_ActivityExternalCorrelation* correlation) { if (correlation->externalKind == MUPTI_EXTERNAL_CORRELATION_KIND_CUSTOM0) { @@ -788,8 +809,7 @@ void MuptiActivityProfiler::handleMuptiActivity( ActivityLogger* logger) { switch (record->kind) { case MUPTI_ACTIVITY_KIND_EXTERNAL_CORRELATION: - handleCorrelationActivity( - reinterpret_cast(record)); + // Populated in buildProcessingState(). break; case MUPTI_ACTIVITY_KIND_RUNTIME: handleRuntimeActivity( diff --git a/libkineto/src/MuptiActivityProfiler.h b/libkineto/src/MuptiActivityProfiler.h index af62a35..d24f595 100644 --- a/libkineto/src/MuptiActivityProfiler.h +++ b/libkineto/src/MuptiActivityProfiler.h @@ -28,6 +28,7 @@ #ifdef HAS_MUPTI #include #include "MuptiActivity.h" +#include "MuptiActivityBuffer.h" #endif // HAS_MUPTI #include "ThreadUtil.h" @@ -382,6 +383,7 @@ class MuptiActivityProfiler { void handleMuptiActivity(const MUpti_Activity* record, ActivityLogger* logger); // Process specific GPU activity types + void buildProcessingState(MuptiActivityBufferMap& buffers); void handleCorrelationActivity( const MUpti_ActivityExternalCorrelation* correlation); void handleRuntimeActivity( diff --git a/libkineto/src/MuptiCallbackApi.cpp b/libkineto/src/MuptiCallbackApi.cpp index 9d9d9f5..1785906 100644 --- a/libkineto/src/MuptiCallbackApi.cpp +++ b/libkineto/src/MuptiCallbackApi.cpp @@ -65,10 +65,10 @@ static void callback_switchboard( // below statement is likey going to call a mutex // on the singleton access -//#if defined(REAL_MUSA_VERSION) && (REAL_MUSA_VERSION >= 40303) -// MuptiCallbackApi::singleton()->__callback_switchboard( -// domain, cbid, cbInfo); -//#endif +#if defined(REAL_MUSA_VERSION) && (REAL_MUSA_VERSION >= 40303) + MuptiCallbackApi::singleton()->__callback_switchboard( + domain, cbid, cbInfo); +#endif } @@ -155,12 +155,13 @@ void MuptiCallbackApi::__callback_switchboard( } std::shared_ptr MuptiCallbackApi::singleton() { - static const std::shared_ptr instance = [] { - std::shared_ptr inst = - std::make_shared(); - return inst; - }(); - return instance; + // MUPTI may invoke resource callbacks during process shutdown, after static + // destruction has started. Intentionally leak the singleton so callback + // state remains valid until process termination. + static const auto* instance = + new std::shared_ptr( + std::make_shared()); + return *instance; } void MuptiCallbackApi::initCallbackApi() {