Opt-in per-thread activity buffers on the ROCm HIP path

libkineto/include/Config.h

@@ -78,6 +78,10 @@ class Config : public AbstractConfig {
     return perThreadBufferEnabled_;
   }
 
+  [[nodiscard]] bool rocprofPerThreadBufferEnabled() const {
+    return rocprofPerThreadBufferEnabled_;
+  }
+
   void setSelectedActivityTypes(const std::set<ActivityType>& types) {
     selectedActivityTypes_ = types;
   }
@@ -307,8 +311,12 @@ class Config : public AbstractConfig {
   // Activity profiler
   bool activityProfilerEnabled_;
 
-  // Enable per-thread buffer
-  bool perThreadBufferEnabled_;
+  // Enable per-thread buffer for CUDA
+  bool perThreadBufferEnabled_{false};
+
+  // Kept separate from perThreadBufferEnabled_ because existing configs already
+  // set that key true, and reusing it would flip every ROCm run onto the path.
+  bool rocprofPerThreadBufferEnabled_{false};
   std::set<ActivityType> selectedActivityTypes_;
 
   // The activity profiler settings are all on-demand

libkineto/src/Config.cpp

@@ -52,6 +52,8 @@ constexpr size_t kDefaultCuptiDeviceBufferPoolLimit(20);
 constexpr char kActivitiesEnabledKey[] = "ACTIVITIES_ENABLED";
 constexpr char kCuptiPerThreadBufferEnabledKey[] =
     "CUPTI_PER_THREAD_BUFFER_ENABLED";
+constexpr std::string_view kRocprofPerThreadBufferEnabledKey =
+    "ROCPROF_PER_THREAD_BUFFER_ENABLED";
 constexpr char kPerformanceMetricsKey[] = "PERFORMANCE_METRICS";
 constexpr char kPerformanceMetricsDeviceIdKey[] =
     "PERFORMANCE_METRICS_DEVICE_ID";
@@ -347,6 +349,8 @@ bool Config::handleOption(const std::string& name, std::string& val) {
     activityProfilerEnabled_ = toBool(val);
   } else if (!name.compare(kCuptiPerThreadBufferEnabledKey)) {
     perThreadBufferEnabled_ = toBool(val);
+  } else if (name == kRocprofPerThreadBufferEnabledKey) {
+    rocprofPerThreadBufferEnabled_ = toBool(val);
   } else if (!name.compare(kPerformanceMetricsKey)) {
     performanceMetricNames_ = splitAndTrim(val, ',');
   } else if (!name.compare(kPerformanceMetricsDeviceIdKey)) {

libkineto/src/RocmActivityProfiler.cpp

@@ -66,6 +66,7 @@ void RocmActivityProfiler::setMaxGpuBufferSize(int64_t size) {
 
 void RocmActivityProfiler::enableGpuTracing() {
   roc_.setMaxEvents(config().maxEvents());
+  roc_.setPerThreadBuffers(config().rocprofPerThreadBufferEnabled());
   roc_.enableActivities(derivedConfig_->profileActivityTypes());
 }
 

libkineto/src/RocprofActivityApi.cpp

@@ -71,6 +71,10 @@ void RocprofActivityApi::setMaxEvents(uint32_t maxEvents) {
   d->setMaxEvents(maxEvents);
 }
 
+void RocprofActivityApi::setPerThreadBuffers(bool enabled) {
+  d->setPerThreadBuffers(enabled);
+}
+
 void RocprofActivityApi::setMaxBufferSize([[maybe_unused]] int64_t size) {
   // FIXME: implement?
   // maxGpuBufferCount_ = 1 + size / kBufSize;
@@ -109,6 +113,8 @@ int RocprofActivityApi::processActivities(
 
   int count = 0;
 
+  d->mergeThreadBuffers();
+
   // Process all external correlations pairs
   for (int it = RocLogger::CorrelationDomain::begin;
        it < RocLogger::CorrelationDomain::end;

libkineto/src/RocprofActivityApi.h

@@ -46,6 +46,7 @@ class RocprofActivityApi {
   void teardownContext() {}
   void setTimeOffset(timestamp_t toffset);
   void setMaxEvents(uint32_t maxEvents);
+  void setPerThreadBuffers(bool enabled);
 
   virtual int processActivities(
       std::function<void(const rocprofBase*)> handler,

libkineto/src/RocprofLogger.cpp

@@ -507,25 +507,118 @@ void RocprofLogger::popCorrelationID(CorrelationDomain type) {
 }
 
 void RocprofLogger::clearLogs() {
-  // CuptiActivityProfiler clears this before the output Loggers use the data
-  // for (auto &row : rows_)
-  //  delete row;
-  rows_.clear();
-  for (int i = 0; i < CorrelationDomain::size; ++i) {
-    externalCorrelations_[i].clear();
+  {
+    std::lock_guard<std::mutex> lock(threadBuffersMutex_);
+    generation_.fetch_add(1, std::memory_order_release);
+    for (auto& buffer : threadBuffers_) {
+      std::lock_guard<std::mutex> bufferLock(buffer->mutex);
+      buffer->rows.clear();
+      for (auto& correlations : buffer->externalCorrelations) {
+        correlations.clear();
+      }
+    }
+    threadBuffers_.clear();
+    totalRows_.store(0, std::memory_order_relaxed);
+  }
+
+  {
+    std::lock_guard<std::mutex> lock(rowsMutex_);
+    rows_.clear();
+  }
+  {
+    std::lock_guard<std::mutex> lock(externalCorrelationsMutex_);
+    for (auto& correlations : externalCorrelations_) {
+      correlations.clear();
+    }
+  }
+}
+
+ThreadTraceBuffer* RocprofLogger::getThreadBuffer() {
+  static RocprofLogger* dis = &singleton();
+  // This owner keeps the returned pointer alive while clearLogs() invalidates
+  // and releases the registry's copy from another thread.
+  thread_local std::shared_ptr<ThreadTraceBuffer> buffer;
+
+  if (buffer == nullptr ||
+      buffer->generation != dis->generation_.load(std::memory_order_acquire)) {
+    std::lock_guard<std::mutex> lock(dis->threadBuffersMutex_);
+    const auto generation = dis->generation_.load(std::memory_order_relaxed);
+    if (buffer == nullptr || buffer->generation != generation) {
+      buffer = std::make_shared<ThreadTraceBuffer>(generation);
+      dis->threadBuffers_.push_back(buffer);
+    }
+  }
+  return buffer.get();
+}
+
+void RocprofLogger::mergeThreadBuffers() {
+  std::lock_guard<std::mutex> threadBuffersLock(threadBuffersMutex_);
+  if (threadBuffers_.empty()) {
+    return;
   }
+  generation_.fetch_add(1, std::memory_order_release);
+  std::lock_guard<std::mutex> rowsLock(rowsMutex_);
+  std::lock_guard<std::mutex> correlationsLock(externalCorrelationsMutex_);
+  uint64_t mergedRows = 0;
+  for (auto& buffer : threadBuffers_) {
+    std::lock_guard<std::mutex> bufferLock(buffer->mutex);
+    mergedRows += buffer->rows.size();
+    rows_.insert(rows_.end(), buffer->rows.begin(), buffer->rows.end());
+    buffer->rows.clear();
+    for (int i = 0; i < CorrelationDomain::size; ++i) {
+      auto& correlations = buffer->externalCorrelations[i];
+      externalCorrelations_[i].insert(
+          externalCorrelations_[i].end(),
+          correlations.begin(),
+          correlations.end());
+      correlations.clear();
+    }
+  }
+  totalRows_.fetch_sub(mergedRows, std::memory_order_relaxed);
+  threadBuffers_.clear();
 }
 
 void RocprofLogger::insert_row_to_buffer(rocprofBase* row) {
   RocprofLogger* dis = &singleton();
-  std::lock_guard<std::mutex> lock(dis->rowsMutex_);
-  if (dis->rows_.size() >= dis->maxBufferSize_) {
-    LOG_FIRST_N(WARNING, 10)
-        << "Exceeded max GPU buffer count (" << dis->rows_.size() << " > "
-        << dis->maxBufferSize_ << ") - terminating tracing";
+
+  // perThreadBuffers disabled
+  if (!dis->perThreadBuffers_.load(std::memory_order_relaxed)) {
+    std::lock_guard<std::mutex> lock(dis->rowsMutex_);
+    if (dis->rows_.size() >= dis->maxBufferSize_) {
+      LOG_FIRST_N(WARNING, 10)
+          << "Exceeded max GPU buffer count (" << dis->rows_.size() << " > "
+          << dis->maxBufferSize_ << ") - terminating tracing";
+      return;
+    }
+    dis->rows_.push_back(row);
+    return;
+  }
+
+  // perThreadBuffers enabled
+  while (true) {
+    auto* buffer = getThreadBuffer();
+    std::lock_guard<std::mutex> lock(buffer->mutex);
+    if (buffer->generation !=
+        dis->generation_.load(std::memory_order_acquire)) {
+      continue;
+    }
+
+    uint64_t rowCount = dis->totalRows_.load(std::memory_order_relaxed);
+    while (true) {
+      if (rowCount >= dis->maxBufferSize_) {
+        LOG_FIRST_N(WARNING, 10)
+            << "Exceeded max GPU buffer count (" << rowCount
+            << " >= " << dis->maxBufferSize_ << ") - terminating tracing";
+        return;
+      }
+      if (dis->totalRows_.compare_exchange_weak(
+              rowCount, rowCount + 1, std::memory_order_relaxed)) {
+        break;
+      }
+    }
+    buffer->rows.push_back(row);
     return;
   }
-  dis->rows_.push_back(row);
 }
 
 void RocprofLogger::code_object_callback(
@@ -663,15 +756,38 @@ void RocprofLogger::api_callback(
             endTime);
         insert_row_to_buffer(row);
       }
-      // External correlation
+      // Store external correlations
       static RocprofLogger* dis = &singleton();
-      for (int it = RocLogger::CorrelationDomain::begin;
-           it < RocLogger::CorrelationDomain::end;
-           ++it) {
-        if (t_externalIds[it].size() > 0) {
-          std::lock_guard<std::mutex> lock(dis->externalCorrelationsMutex_);
-          dis->externalCorrelations_[it].emplace_back(
-              record.correlation_id.internal, t_externalIds[it].back());
+
+      // perThreadBuffers enabled
+      if (dis->perThreadBuffers_.load(std::memory_order_relaxed)) {
+        while (true) {
+          auto* buffer = getThreadBuffer();
+          std::lock_guard<std::mutex> lock(buffer->mutex);
+          if (buffer->generation !=
+              dis->generation_.load(std::memory_order_acquire)) {
+            continue;
+          }
+          for (int it = RocLogger::CorrelationDomain::begin;
+               it < RocLogger::CorrelationDomain::end;
+               ++it) {
+            if (t_externalIds[it].size() > 0) {
+              buffer->externalCorrelations[it].emplace_back(
+                  record.correlation_id.internal, t_externalIds[it].back());
+            }
+          }
+          break;
+        }
+      } else {
+        // perThreadBuffers disabled
+        for (int it = RocLogger::CorrelationDomain::begin;
+             it < RocLogger::CorrelationDomain::end;
+             ++it) {
+          if (t_externalIds[it].size() > 0) {
+            std::lock_guard<std::mutex> lock(dis->externalCorrelationsMutex_);
+            dis->externalCorrelations_[it].emplace_back(
+                record.correlation_id.internal, t_externalIds[it].back());
+          }
         }
       }
     } // ROCPROFILER_CALLBACK_PHASE_EXIT
@@ -776,6 +892,10 @@ void RocprofLogger::setMaxEvents(uint32_t maxBufferSize) {
   maxBufferSize_ = maxBufferSize;
 }
 
+void RocprofLogger::setPerThreadBuffers(bool enabled) {
+  perThreadBuffers_.store(enabled, std::memory_order_relaxed);
+}
+
 void RocprofLogger::ensureRegistered() {
   int status = 0;
   rocprofiler_is_initialized(&status);

libkineto/src/RocprofLogger.h

@@ -24,6 +24,17 @@
 
 #include "RocLogger.h"
 
+struct ThreadTraceBuffer {
+  explicit ThreadTraceBuffer(uint64_t bufferGeneration)
+      : generation(bufferGeneration) {}
+
+  const uint64_t generation;
+  std::mutex mutex;
+  std::vector<rocprofBase*> rows;
+  std::vector<std::pair<uint64_t, uint64_t>>
+      externalCorrelations[RocLogger::CorrelationDomain::size];
+};
+
 class RocprofLogger {
  public:
   RocprofLogger();
@@ -42,6 +53,7 @@ class RocprofLogger {
   void stopLogging();
   void clearLogs();
   void setMaxEvents(uint32_t maxBufferSize);
+  void setPerThreadBuffers(bool enabled);
 
   static int toolInit(
       rocprofiler_client_finalize_t finalize_func,
@@ -62,6 +74,9 @@ class RocprofLogger {
 
   static void insert_row_to_buffer(rocprofBase* row);
 
+  static ThreadTraceBuffer* getThreadBuffer();
+  void mergeThreadBuffers();
+
   //
   static void api_callback(
       rocprofiler_callback_tracing_record_t record,
@@ -91,6 +106,12 @@ class RocprofLogger {
       externalCorrelations_[RocLogger::CorrelationDomain::size];
   std::mutex externalCorrelationsMutex_;
 
+  std::atomic<bool> perThreadBuffers_{false};
+  std::vector<std::shared_ptr<ThreadTraceBuffer>> threadBuffers_;
+  std::mutex threadBuffersMutex_;
+  std::atomic<uint64_t> generation_{0};
+  std::atomic<uint64_t> totalRows_{0};
+
   bool externalCorrelationEnabled_{true};
   bool logging_{false};
 

libkineto/test/ConfigTest.cpp

@@ -214,3 +214,19 @@ TEST(ParseTest, OnDemandLogFileRejectedOutsideAllowedDir) {
   EXPECT_TRUE(cfg.parse("ACTIVITIES_LOG_FILE=/tmp/../etc/cron.d/payload"));
   EXPECT_EQ(cfg.activitiesLogFile(), original);
 }
+
+TEST(ParseTest, RocprofPerThreadBuffer) {
+  Config cfg;
+  EXPECT_FALSE(cfg.rocprofPerThreadBufferEnabled());
+  EXPECT_TRUE(cfg.parse("ROCPROF_PER_THREAD_BUFFER_ENABLED=true"));
+  EXPECT_TRUE(cfg.rocprofPerThreadBufferEnabled());
+}
+
+// The two keys stay independent. Configs in the wild already set the CUPTI one,
+// so letting it reach the rocprof path would enable it on every ROCm run.
+TEST(ParseTest, CuptiPerThreadBufferDoesNotEnableRocprof) {
+  Config cfg;
+  EXPECT_TRUE(cfg.parse("CUPTI_PER_THREAD_BUFFER_ENABLED=true"));
+  EXPECT_TRUE(cfg.perThreadBufferEnabled());
+  EXPECT_FALSE(cfg.rocprofPerThreadBufferEnabled());
+}