Update sync torch_musa 2.11.0.post2

libkineto/src/MuptiActivityApi.cpp

@@ -57,8 +57,8 @@ inline void reenableMuptiCallbacks_(std::shared_ptr<MuptiCallbackApi>& 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<MuptiActivityBufferMap> MuptiActivityApi::activityBuffers() {
   {
     std::lock_guard<std::mutex> 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<std::mutex> guard(mutex_);
-  auto it = allocatedGpuTraceBuffers_.find(buffer);
-  if (it == allocatedGpuTraceBuffers_.end()) {
-    LOG(ERROR) << "bufferCompleted called with unknown buffer: "
-               << (void*) buffer;
-    return;
-  }
+  {
+    std::lock_guard<std::mutex> guard(mutex_);
+    auto it = allocatedGpuTraceBuffers_.find(buffer);
+    if (it == allocatedGpuTraceBuffers_.end()) {
+      LOG(ERROR) << "bufferCompleted called with unknown buffer: "
+                 << static_cast<void*>(buffer);
+      return;
+    }
 
-  if (!readyGpuTraceBuffers_) {
-    readyGpuTraceBuffers_ = std::make_unique<MuptiActivityBufferMap>();
+    if (!readyGpuTraceBuffers_) {
+      readyGpuTraceBuffers_ = std::make_unique<MuptiActivityBufferMap>();
+    }
+    // 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();
   }

libkineto/src/MuptiActivityApi.h

@@ -90,6 +90,7 @@ class MuptiActivityApi {
   std::unique_ptr<MuptiActivityBufferMap> readyGpuTraceBuffers_;
   std::mutex mutex_;
   std::atomic<uint32_t> tracingEnabled_{0};
+  std::atomic<uint32_t> tearingDown_{0};
   bool externalCorrelationEnabled_{false};
 
 #ifdef HAS_MUPTI

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<const MUpti_ActivityExternalCorrelation*>(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<const MUpti_ActivityExternalCorrelation*>(record));
+      // Populated in buildProcessingState().
       break;
     case MUPTI_ACTIVITY_KIND_RUNTIME:
       handleRuntimeActivity(

libkineto/src/MuptiActivityProfiler.h

@@ -28,6 +28,7 @@
 #ifdef HAS_MUPTI
 #include <mupti.h>
 #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(

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> MuptiCallbackApi::singleton() {
-  static const std::shared_ptr<MuptiCallbackApi> instance = [] {
-    std::shared_ptr<MuptiCallbackApi> inst =
-        std::make_shared<MuptiCallbackApi>();
-    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<MuptiCallbackApi>(
+            std::make_shared<MuptiCallbackApi>());
+    return *instance;
 }
 
 void MuptiCallbackApi::initCallbackApi() {