Merge branch 'main' into fix_windows

pytorch · Aug 26, 2024 · 405d100 · 405d100
2 parents c162fa0 + 464bccf
commit 405d100
Show file tree

Hide file tree

Showing 6 changed files with 13 additions and 64 deletions.
diff --git a/.github/workflows/libkineto_ci.yml b/.github/workflows/libkineto_ci.yml
diff --git a/libkineto/src/CuptiActivityProfiler.cpp b/libkineto/src/CuptiActivityProfiler.cpp
@@ -1001,7 +1001,7 @@ void CuptiActivityProfiler::configure(
     // presumably because structures are allocated and initialized, callbacks
     // are activated etc. After a while the overhead decreases and stabilizes.
     // It's therefore useful to perform some warmup before starting recording.
-    LOG(INFO) << "Enabling GPU tracing";
+    LOG(INFO) << "Enabling GPU tracing with max CUPTI buffer size " << config_->activitiesMaxGpuBufferSize() / 1024 / 1024 << "MB)";
     cupti_.setMaxBufferSize(config_->activitiesMaxGpuBufferSize());
     time_point<system_clock> timestamp;
     if (VLOG_IS_ON(1)) {
@@ -1174,6 +1174,8 @@ const time_point<system_clock> CuptiActivityProfiler::performRunLoopStep(
         std::lock_guard<std::mutex> guard(mutex_);
         stopTraceInternal(now);
         resetInternal();
+        LOG(ERROR) << "State: Warmup stopped by CUPTI. (Buffer size configured is " << config_->activitiesMaxGpuBufferSize() / 1024 / 1024 << "MB)";
+        UST_LOGGER_MARK_COMPLETED(kWarmUpStage);
         VLOG(0) << "Warmup -> WaitForRequest";
         break;
       }
@@ -1222,7 +1224,10 @@ const time_point<system_clock> CuptiActivityProfiler::performRunLoopStep(
         }
 
 #if defined(HAS_CUPTI) || defined(HAS_ROCTRACER)
-        ecs_.cupti_stopped_early = cupti_.stopCollection;
+        if (cupti_.stopCollection) {
+          ecs_.cupti_stopped_early = cupti_.stopCollection;
+          LOG(ERROR) << "State: CollectTrace stopped by CUPTI. (Buffer size configured is " << config_->activitiesMaxGpuBufferSize() / 1024 / 1024 << "MB)";
+        }
 #endif // HAS_CUPTI || HAS_ROCTRACER
 
         std::lock_guard<std::mutex> guard(mutex_);

diff --git a/libkineto/src/RoctracerLogger.cpp b/libkineto/src/RoctracerLogger.cpp
@@ -18,6 +18,7 @@
 #include "ThreadUtil.h"
 #include "Demangle.h"
 
+using namespace libkineto;
 using namespace std::chrono;
 
 class Flush

diff --git a/libkineto/src/RoctracerLogger.h b/libkineto/src/RoctracerLogger.h
@@ -63,8 +63,6 @@ static timestamp_t timespec_to_ns(const timespec& time) {
   return ((timestamp_t)time.tv_sec * 1000000000) + time.tv_nsec;
 }
 
-using namespace libkineto;
-
 class ApiIdList {
  public:
   ApiIdList();

diff --git a/libkineto/src/init.cpp b/libkineto/src/init.cpp
@@ -225,7 +225,7 @@ void libkineto_init(bool cpuOnly, bool logOnError) {
 }
 
 // The cuda driver calls this function if the CUDA_INJECTION64_PATH environment
-// variable is set
+// variable is set. Should be skipped if unset or CUDA_INJECTION64_PATH=none.
 int InitializeInjection(void) {
   LOG(INFO) << "Injection mode: Initializing libkineto";
   libkineto_init(false /*cpuOnly*/, true /*logOnError*/);

diff --git a/libkineto/src/output_json.cpp b/libkineto/src/output_json.cpp
@@ -335,9 +335,10 @@ void ChromeTraceLogger::handleActivity(
   }
 
   std::string arg_values = "";
-  if (op.correlationId() != 0) {
-    arg_values.append(fmt::format("\"External id\": {}",
-      op.linkedActivity() ? op.linkedActivity()->correlationId() : op.correlationId()));
+  if (op.linkedActivity() && op.linkedActivity()->correlationId() != 0) {
+    arg_values.append(fmt::format("\"External id\": {}", op.linkedActivity()->correlationId()));
+  } else if (op.correlationId() != 0) {
+    arg_values.append(fmt::format("\"External id\": {}", op.correlationId()));
   }
   std::string op_metadata = op.metadataJson();
   sanitizeStrForJSON(op_metadata);