diff --git a/libkineto/include/Config.h b/libkineto/include/Config.h index 9835ce0a9..493095204 100644 --- a/libkineto/include/Config.h +++ b/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& 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 selectedActivityTypes_; // The activity profiler settings are all on-demand diff --git a/libkineto/src/Config.cpp b/libkineto/src/Config.cpp index b1644458e..58c3f9d7e 100644 --- a/libkineto/src/Config.cpp +++ b/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)) { diff --git a/libkineto/src/RocmActivityProfiler.cpp b/libkineto/src/RocmActivityProfiler.cpp index 5e54f67ed..36fa01275 100644 --- a/libkineto/src/RocmActivityProfiler.cpp +++ b/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()); } diff --git a/libkineto/src/RocprofActivityApi.cpp b/libkineto/src/RocprofActivityApi.cpp index 165a2841e..2387dac4e 100644 --- a/libkineto/src/RocprofActivityApi.cpp +++ b/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; diff --git a/libkineto/src/RocprofActivityApi.h b/libkineto/src/RocprofActivityApi.h index 2ce39f3ca..c5fe441f0 100644 --- a/libkineto/src/RocprofActivityApi.h +++ b/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 handler, diff --git a/libkineto/src/RocprofLogger.cpp b/libkineto/src/RocprofLogger.cpp index 2988dec4e..0e484e826 100644 --- a/libkineto/src/RocprofLogger.cpp +++ b/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 lock(threadBuffersMutex_); + generation_.fetch_add(1, std::memory_order_release); + for (auto& buffer : threadBuffers_) { + std::lock_guard 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 lock(rowsMutex_); + rows_.clear(); + } + { + std::lock_guard 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 buffer; + + if (buffer == nullptr || + buffer->generation != dis->generation_.load(std::memory_order_acquire)) { + std::lock_guard lock(dis->threadBuffersMutex_); + const auto generation = dis->generation_.load(std::memory_order_relaxed); + if (buffer == nullptr || buffer->generation != generation) { + buffer = std::make_shared(generation); + dis->threadBuffers_.push_back(buffer); + } + } + return buffer.get(); +} + +void RocprofLogger::mergeThreadBuffers() { + std::lock_guard threadBuffersLock(threadBuffersMutex_); + if (threadBuffers_.empty()) { + return; } + generation_.fetch_add(1, std::memory_order_release); + std::lock_guard rowsLock(rowsMutex_); + std::lock_guard correlationsLock(externalCorrelationsMutex_); + uint64_t mergedRows = 0; + for (auto& buffer : threadBuffers_) { + std::lock_guard 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 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 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 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 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 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 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); diff --git a/libkineto/src/RocprofLogger.h b/libkineto/src/RocprofLogger.h index 6a6110aaf..9621097e0 100644 --- a/libkineto/src/RocprofLogger.h +++ b/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 rows; + std::vector> + 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 perThreadBuffers_{false}; + std::vector> threadBuffers_; + std::mutex threadBuffersMutex_; + std::atomic generation_{0}; + std::atomic totalRows_{0}; + bool externalCorrelationEnabled_{true}; bool logging_{false}; diff --git a/libkineto/test/ConfigTest.cpp b/libkineto/test/ConfigTest.cpp index ef4529cae..b0e5339d8 100644 --- a/libkineto/test/ConfigTest.cpp +++ b/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()); +}