diff --git a/libkineto/CMakeLists.txt b/libkineto/CMakeLists.txt index f17f8b072..5a13dbb71 100644 --- a/libkineto/CMakeLists.txt +++ b/libkineto/CMakeLists.txt @@ -27,6 +27,9 @@ set(KINETO_LIBRARY_TYPE "default" CACHE STRING set_property(CACHE KINETO_LIBRARY_TYPE PROPERTY STRINGS default shared static) option(KINETO_BUILD_TESTS "Build kineto unit tests" ON) option(KINETO_BUILD_BENCHMARKS "Build kineto benchmarks" OFF) +option(KINETO_ENABLE_PERFETTO "Enable native Perfetto/pftrace trace export" OFF) +set(PERFETTO_SOURCE_DIR "" CACHE PATH + "Path to a Perfetto checkout containing sdk/perfetto.h and sdk/perfetto.cc") set(LIBKINETO_DIR "${CMAKE_CURRENT_SOURCE_DIR}") set(LIBKINETO_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/src") @@ -213,6 +216,33 @@ elseif(KINETO_BACKEND STREQUAL "xpu") endif() endif() +if(KINETO_ENABLE_PERFETTO) + if(NOT PERFETTO_SOURCE_DIR) + if(DEFINED ENV{PERFETTO_SOURCE_DIR}) + set(PERFETTO_SOURCE_DIR "$ENV{PERFETTO_SOURCE_DIR}") + elseif(EXISTS "${LIBKINETO_THIRDPARTY_DIR}/perfetto/sdk/perfetto.h") + set(PERFETTO_SOURCE_DIR "${LIBKINETO_THIRDPARTY_DIR}/perfetto") + endif() + endif() + if(NOT EXISTS "${PERFETTO_SOURCE_DIR}/sdk/perfetto.h" OR + NOT EXISTS "${PERFETTO_SOURCE_DIR}/sdk/perfetto.cc") + message(FATAL_ERROR + "KINETO_ENABLE_PERFETTO=ON requires PERFETTO_SOURCE_DIR to point to " + "a Perfetto checkout containing sdk/perfetto.h and sdk/perfetto.cc.") + endif() + add_library(kineto_perfetto STATIC + "${PERFETTO_SOURCE_DIR}/sdk/perfetto.cc") + target_include_directories(kineto_perfetto SYSTEM PUBLIC + "${PERFETTO_SOURCE_DIR}/sdk") + set_target_properties(kineto_perfetto PROPERTIES + POSITION_INDEPENDENT_CODE ON) + if(MSVC) + target_compile_options(kineto_perfetto PRIVATE /w) + else() + target_compile_options(kineto_perfetto PRIVATE -w -O0 -g0) + endif() + list(APPEND KINETO_DEFINITIONS "KINETO_ENABLE_PERFETTO") +endif() target_compile_definitions(kineto_base PUBLIC "${KINETO_DEFINITIONS}") target_compile_options(kineto_base PRIVATE "${KINETO_COMPILE_OPTIONS}") target_compile_definitions(kineto_api PUBLIC "${KINETO_DEFINITIONS}") @@ -249,6 +279,9 @@ target_include_directories(kineto_base SYSTEM PUBLIC $ $) target_link_libraries(kineto_base PRIVATE $) +if(KINETO_ENABLE_PERFETTO) + target_link_libraries(kineto_base PUBLIC kineto_perfetto) +endif() target_include_directories(kineto_api PUBLIC $ @@ -265,6 +298,9 @@ elseif(NOT KINETO_LIBRARY_TYPE STREQUAL "default") endif() add_library(kineto $ $) +if(KINETO_ENABLE_PERFETTO) + target_link_libraries(kineto $) +endif() if(BUILD_SHARED_LIBS) set_property(TARGET kineto_base PROPERTY POSITION_INDEPENDENT_CODE ON) set_property(TARGET kineto_api PROPERTY POSITION_INDEPENDENT_CODE ON) diff --git a/libkineto/libkineto_defs.bzl b/libkineto/libkineto_defs.bzl index 8750d4092..ea1bb8ee3 100644 --- a/libkineto/libkineto_defs.bzl +++ b/libkineto/libkineto_defs.bzl @@ -82,6 +82,7 @@ def get_libkineto_cpu_only_srcs(with_api = True): "src/init.cpp", "src/output_csv.cpp", "src/output_json.cpp", + "src/output_perfetto.cpp", ] + (get_libkineto_api_srcs() if with_api else []) def get_libkineto_public_headers(): diff --git a/libkineto/src/ActivityProfilerController.cpp b/libkineto/src/ActivityProfilerController.cpp index 6dc0f9ca9..ffb231dd4 100644 --- a/libkineto/src/ActivityProfilerController.cpp +++ b/libkineto/src/ActivityProfilerController.cpp @@ -32,6 +32,9 @@ #include "output_json.h" #include "output_membuf.h" +#ifdef KINETO_ENABLE_PERFETTO +#include "output_perfetto.h" +#endif #include "Logger.h" @@ -111,6 +114,11 @@ static ActivityLoggerFactory initLoggerFactory() { factory.addProtocol("file", [](const std::string& url) { return std::unique_ptr(new ChromeTraceLogger(url)); }); +#ifdef KINETO_ENABLE_PERFETTO + factory.addProtocol("pftrace", [](const std::string& url) { + return std::unique_ptr(new PerfettoTraceLogger(url)); + }); +#endif return factory; } diff --git a/libkineto/src/ActivityTrace.h b/libkineto/src/ActivityTrace.h index d7849bed9..e02f02f3c 100644 --- a/libkineto/src/ActivityTrace.h +++ b/libkineto/src/ActivityTrace.h @@ -15,6 +15,9 @@ #include "ActivityTraceInterface.h" #include "output_json.h" #include "output_membuf.h" +#ifdef KINETO_ENABLE_PERFETTO +#include "output_perfetto.h" +#endif namespace libkineto { @@ -31,7 +34,7 @@ class ActivityTrace : public ActivityTraceInterface { std::string prefix; // if no protocol is specified, default to file if (url.find("://") == url.npos) { - prefix = "file://"; + prefix = url.ends_with(".pftrace") ? "pftrace://" : "file://"; } memLogger_->setChromeLogger(loggerFactory_.makeLogger(prefix + url)); memLogger_->log(*memLogger_->getChromeLogger()); diff --git a/libkineto/src/output_perfetto.cpp b/libkineto/src/output_perfetto.cpp new file mode 100644 index 000000000..844824bae --- /dev/null +++ b/libkineto/src/output_perfetto.cpp @@ -0,0 +1,474 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#ifndef KINETO_ENABLE_PERFETTO +namespace libkineto { +namespace { +[[maybe_unused]] constexpr int kPerfettoTraceLoggerDisabled = 0; +} // namespace +} // namespace libkineto +#else + +#include "output_perfetto.h" + +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "ActivityType.h" +#include "Config.h" +#include "Logger.h" +#include "ThreadUtil.h" +#include "TraceSpan.h" + +PERFETTO_DEFINE_CATEGORIES( + perfetto::Category("cpu_op").SetDescription("CPU operator"), + perfetto::Category("user_annotation").SetDescription("User annotation"), + perfetto::Category("gpu_user_annotation").SetDescription("GPU user annotation"), + perfetto::Category("gpu_memcpy").SetDescription("GPU memory copy"), + perfetto::Category("gpu_memset").SetDescription("GPU memory set"), + perfetto::Category("kernel").SetDescription("GPU kernel"), + perfetto::Category("external_correlation").SetDescription("External correlation"), + perfetto::Category("cuda_runtime").SetDescription("CUDA/HIP runtime event"), + perfetto::Category("cuda_driver").SetDescription("CUDA/HIP driver event"), + perfetto::Category("cpu_instant_event").SetDescription("CPU instant event"), + perfetto::Category("python_function").SetDescription("Python function"), + perfetto::Category("overhead").SetDescription("Profiler overhead"), + perfetto::Category("mtia_runtime").SetDescription("MTIA runtime event"), + perfetto::Category("mtia_ccp_events").SetDescription("MTIA CCP event"), + perfetto::Category("mtia_insight").SetDescription("MTIA insight event"), + perfetto::Category("cuda_sync").SetDescription("CUDA/HIP synchronization"), + perfetto::Category("cuda_event").SetDescription("CUDA/HIP event"), + perfetto::Category("mtia_counters").SetDescription("MTIA counter"), + perfetto::Category("glow_runtime").SetDescription("Glow runtime event"), + perfetto::Category("cuda_profiler_range").SetDescription("CUDA profiler range"), + perfetto::Category("hpu_op").SetDescription("HPU operation"), + perfetto::Category("xpu_runtime").SetDescription("XPU runtime event"), + perfetto::Category("xpu_driver").SetDescription("XPU driver event"), + perfetto::Category("collective_comm").SetDescription("Collective communication"), + perfetto::Category("privateuse1_runtime").SetDescription("PrivateUse1 runtime event"), + perfetto::Category("privateuse1_driver").SetDescription("PrivateUse1 driver event"), + perfetto::Category("xpu_scope_profiler").SetDescription("XPU scope profiler")); + +PERFETTO_TRACK_EVENT_STATIC_STORAGE(); + +namespace KINETO_NAMESPACE { +namespace { + +constexpr size_t kBufferSizeKb = 102400; +constexpr size_t kSharedMemorySizeKb = 64; + +#ifdef __linux__ +constexpr const char* kDefaultLogFileFmt = "/tmp/libkineto_activities_{}.pftrace"; +#else +constexpr const char* kDefaultLogFileFmt = "libkineto_activities_{}.pftrace"; +#endif + +uint64_t stableTrackId(const std::string& key) { + return std::hash{}(key); +} + +std::string makeTrackKey(std::string_view prefix, int64_t device, int64_t resource) { + std::ostringstream os; + os << prefix << ':' << device << ':' << resource; + return os.str(); +} + +bool isDeviceActivity(ActivityType type) { + switch (type) { + case ActivityType::GPU_USER_ANNOTATION: + case ActivityType::GPU_MEMCPY: + case ActivityType::GPU_MEMSET: + case ActivityType::CONCURRENT_KERNEL: + case ActivityType::MTIA_CCP_EVENTS: + case ActivityType::MTIA_INSIGHT: + case ActivityType::MTIA_COUNTERS: + case ActivityType::XPU_SCOPE_PROFILER: + return true; + default: + return false; + } +} + +void addDebugAnnotation( + perfetto::EventContext& ctx, + const std::string& name, + const std::string& value) { + if (value.empty()) { + return; + } + auto* annotation = ctx.event()->add_debug_annotations(); + annotation->set_name(name); + annotation->set_string_value(value); +} + +void addDebugAnnotation( + perfetto::EventContext& ctx, + const std::string& name, + int64_t value) { + auto* annotation = ctx.event()->add_debug_annotations(); + annotation->set_name(name); + annotation->set_int_value(value); +} + +void addActivityDebugAnnotations( + perfetto::EventContext& ctx, + const ITraceActivity& activity, + int64_t beginNs, + int64_t endNs) { + addDebugAnnotation(ctx, "activity_type", toString(activity.type())); + addDebugAnnotation(ctx, "begin_ns", beginNs); + addDebugAnnotation(ctx, "end_ns", endNs); + addDebugAnnotation(ctx, "device", activity.deviceId()); + addDebugAnnotation(ctx, "resource", activity.resourceId()); + addDebugAnnotation(ctx, "thread_id", activity.getThreadId()); + addDebugAnnotation(ctx, "correlation_id", activity.correlationId()); + addDebugAnnotation(ctx, "metadata_json", activity.metadataJson()); +} + +void initializePerfettoOnce() { + static std::once_flag once; + std::call_once(once, []() { + perfetto::TracingInitArgs args; + args.backends |= perfetto::kInProcessBackend; + args.shmem_size_hint_kb = kSharedMemorySizeKb; + perfetto::Tracing::Initialize(args); + perfetto::TrackEvent::Register(); + }); +} + +std::string defaultFileName() { + return fmt::format(kDefaultLogFileFmt, processId()); +} + +} // namespace + +PerfettoTraceLogger::PerfettoTraceLogger(const std::string& traceFileName) + : fileName_(traceFileName.empty() ? defaultFileName() : traceFileName) {} + +PerfettoTraceLogger::~PerfettoTraceLogger() { + if (started_ && !finalized_) { + try { + writeTraceFile(); + } catch (const std::exception& e) { + LOG(ERROR) << "Failed to finalize Perfetto trace: " << e.what(); + } + } +} + +void PerfettoTraceLogger::handleTraceStart( + const std::unordered_map& metadata, + const std::string& device_properties) { + if (started_) { + return; + } + + initializePerfettoOnce(); + + perfetto::TraceConfig cfg; + auto* buffer = cfg.add_buffers(); + buffer->set_size_kb(kBufferSizeKb); + buffer->set_fill_policy( + perfetto::protos::gen::TraceConfig_BufferConfig_FillPolicy_DISCARD); + + perfetto::protos::gen::TrackEventConfig trackEventConfig; + auto* dataSource = cfg.add_data_sources()->mutable_config(); + dataSource->set_name("track_event"); + dataSource->set_track_event_config_raw(trackEventConfig.SerializeAsString()); + + tracingSession_ = perfetto::Tracing::NewTrace(); + tracingSession_->Setup(cfg); + tracingSession_->StartBlocking(); + started_ = true; + + auto metadataTrack = perfetto::Track{stableTrackId("kineto:metadata")}; + auto descriptor = metadataTrack.Serialize(); + descriptor.set_name("Kineto metadata"); + perfetto::TrackEvent::SetTrackDescriptor(metadataTrack, descriptor); + tracks_.emplace( + "kineto:metadata", std::make_unique(metadataTrack)); + + TRACE_EVENT_BEGIN( + "user_annotation", + "Kineto trace metadata", + metadataTrack, + static_cast(0), + [&](perfetto::EventContext ctx) { + addDebugAnnotation(ctx, "device_properties", device_properties); + for (const auto& [key, value] : metadata) { + addDebugAnnotation(ctx, key, value); + } + }); + TRACE_EVENT_END("user_annotation", metadataTrack, static_cast(1)); +} + +void PerfettoTraceLogger::handleDeviceInfo(const DeviceInfo& info, int64_t /*time*/) { + trackForDevice(info.id, info.name); +} + +void PerfettoTraceLogger::handleResourceInfo( + const ResourceInfo& info, + int64_t /*time*/) { + std::ostringstream name; + name << info.name << " (device " << info.deviceId << ", resource " << info.id << ")"; + trackForDevice(info.id, name.str()); +} + +void PerfettoTraceLogger::handleOverheadInfo( + const OverheadInfo& info, + int64_t time) { + auto track = perfetto::Track{stableTrackId("kineto:overhead")}; + TRACE_EVENT_BEGIN( + "overhead", + perfetto::DynamicString(info.name.c_str()), + track, + static_cast(time)); + TRACE_EVENT_END("overhead", track, static_cast(time + 1)); +} + +void PerfettoTraceLogger::handleTraceSpan(const TraceSpan& span) { + auto track = perfetto::Track{stableTrackId("kineto:trace_span")}; + auto descriptor = track.Serialize(); + descriptor.set_name("Kineto trace spans"); + perfetto::TrackEvent::SetTrackDescriptor(track, descriptor); + TRACE_EVENT_BEGIN( + "user_annotation", + perfetto::DynamicString(span.name.c_str()), + track, + static_cast(span.startTime), + "iteration", + span.iteration); + TRACE_EVENT_END("user_annotation", track, static_cast(span.endTime)); +} + +void PerfettoTraceLogger::handleGenericActivity( + const GenericTraceActivity& activity) { + handleActivity(activity); +} + +void PerfettoTraceLogger::handleActivity(const ITraceActivity& activity) { + if (!started_) { + handleTraceStart({}, ""); + } + + auto name = activity.name(); + auto beginNs = activity.timestamp(); + auto endNs = beginNs + std::max(activity.duration(), 1); + const auto& track = trackForActivity(activity); + + const auto flowId = activity.flowId() != 0 ? activity.flowId() + : activity.linkedActivity() ? activity.linkedActivity()->correlationId() + : 0; + +#define KINETO_TRACE_ACTIVITY_FOR_CATEGORY(category_literal) \ + do { \ + if (flowId != 0) { \ + TRACE_EVENT_BEGIN( \ + category_literal, \ + perfetto::DynamicString(name.c_str()), \ + track, \ + static_cast(beginNs), \ + perfetto::Flow::ProcessScoped(flowId), \ + [&](perfetto::EventContext ctx) { \ + addActivityDebugAnnotations(ctx, activity, beginNs, endNs); \ + }); \ + } else { \ + TRACE_EVENT_BEGIN( \ + category_literal, \ + perfetto::DynamicString(name.c_str()), \ + track, \ + static_cast(beginNs), \ + [&](perfetto::EventContext ctx) { \ + addActivityDebugAnnotations(ctx, activity, beginNs, endNs); \ + }); \ + } \ + TRACE_EVENT_END(category_literal, track, static_cast(endNs)); \ + } while (false) + + switch (activity.type()) { + case ActivityType::CPU_OP: + KINETO_TRACE_ACTIVITY_FOR_CATEGORY("cpu_op"); + break; + case ActivityType::USER_ANNOTATION: + KINETO_TRACE_ACTIVITY_FOR_CATEGORY("user_annotation"); + break; + case ActivityType::GPU_USER_ANNOTATION: + KINETO_TRACE_ACTIVITY_FOR_CATEGORY("gpu_user_annotation"); + break; + case ActivityType::GPU_MEMCPY: + KINETO_TRACE_ACTIVITY_FOR_CATEGORY("gpu_memcpy"); + break; + case ActivityType::GPU_MEMSET: + KINETO_TRACE_ACTIVITY_FOR_CATEGORY("gpu_memset"); + break; + case ActivityType::CONCURRENT_KERNEL: + KINETO_TRACE_ACTIVITY_FOR_CATEGORY("kernel"); + break; + case ActivityType::EXTERNAL_CORRELATION: + KINETO_TRACE_ACTIVITY_FOR_CATEGORY("external_correlation"); + break; + case ActivityType::CUDA_RUNTIME: + KINETO_TRACE_ACTIVITY_FOR_CATEGORY("cuda_runtime"); + break; + case ActivityType::CUDA_DRIVER: + KINETO_TRACE_ACTIVITY_FOR_CATEGORY("cuda_driver"); + break; + case ActivityType::CPU_INSTANT_EVENT: + KINETO_TRACE_ACTIVITY_FOR_CATEGORY("cpu_instant_event"); + break; + case ActivityType::PYTHON_FUNCTION: + KINETO_TRACE_ACTIVITY_FOR_CATEGORY("python_function"); + break; + case ActivityType::OVERHEAD: + KINETO_TRACE_ACTIVITY_FOR_CATEGORY("overhead"); + break; + case ActivityType::MTIA_RUNTIME: + KINETO_TRACE_ACTIVITY_FOR_CATEGORY("mtia_runtime"); + break; + case ActivityType::MTIA_CCP_EVENTS: + KINETO_TRACE_ACTIVITY_FOR_CATEGORY("mtia_ccp_events"); + break; + case ActivityType::MTIA_INSIGHT: + KINETO_TRACE_ACTIVITY_FOR_CATEGORY("mtia_insight"); + break; + case ActivityType::CUDA_SYNC: + KINETO_TRACE_ACTIVITY_FOR_CATEGORY("cuda_sync"); + break; + case ActivityType::CUDA_EVENT: + KINETO_TRACE_ACTIVITY_FOR_CATEGORY("cuda_event"); + break; + case ActivityType::MTIA_COUNTERS: + KINETO_TRACE_ACTIVITY_FOR_CATEGORY("mtia_counters"); + break; + case ActivityType::GLOW_RUNTIME: + KINETO_TRACE_ACTIVITY_FOR_CATEGORY("glow_runtime"); + break; + case ActivityType::CUDA_PROFILER_RANGE: + KINETO_TRACE_ACTIVITY_FOR_CATEGORY("cuda_profiler_range"); + break; + case ActivityType::HPU_OP: + KINETO_TRACE_ACTIVITY_FOR_CATEGORY("hpu_op"); + break; + case ActivityType::XPU_RUNTIME: + KINETO_TRACE_ACTIVITY_FOR_CATEGORY("xpu_runtime"); + break; + case ActivityType::XPU_DRIVER: + KINETO_TRACE_ACTIVITY_FOR_CATEGORY("xpu_driver"); + break; + case ActivityType::COLLECTIVE_COMM: + KINETO_TRACE_ACTIVITY_FOR_CATEGORY("collective_comm"); + break; + case ActivityType::PRIVATEUSE1_RUNTIME: + KINETO_TRACE_ACTIVITY_FOR_CATEGORY("privateuse1_runtime"); + break; + case ActivityType::PRIVATEUSE1_DRIVER: + KINETO_TRACE_ACTIVITY_FOR_CATEGORY("privateuse1_driver"); + break; + case ActivityType::XPU_SCOPE_PROFILER: + KINETO_TRACE_ACTIVITY_FOR_CATEGORY("xpu_scope_profiler"); + break; + case ActivityType::ENUM_COUNT: + KINETO_TRACE_ACTIVITY_FOR_CATEGORY("user_annotation"); + break; + } + +#undef KINETO_TRACE_ACTIVITY_FOR_CATEGORY +} + +const perfetto::Track& PerfettoTraceLogger::trackForActivity( + const ITraceActivity& activity) { + const auto device = activity.deviceId(); + const auto resource = activity.resourceId(); + const auto deviceActivity = isDeviceActivity(activity.type()); + auto key = makeTrackKey(deviceActivity ? "device" : "thread", device, resource); + auto it = tracks_.find(key); + if (it != tracks_.end()) { + return *it->second; + } + + std::ostringstream name; + if (deviceActivity) { + name << "Device " << device << " resource " << resource; + } else { + name << "Thread " << activity.getThreadId() << " (device " << device + << ", resource " << resource << ")"; + } + + auto track = std::make_unique(stableTrackId(key)); + auto descriptor = track->Serialize(); + descriptor.set_name(name.str()); + perfetto::TrackEvent::SetTrackDescriptor(*track, descriptor); + auto emplaceResult = tracks_.emplace(std::move(key), std::move(track)); + return *emplaceResult.first->second; +} + +const perfetto::Track& PerfettoTraceLogger::trackForDevice( + int64_t deviceId, + std::string name) { + auto key = fmt::format("device-info:{}", deviceId); + auto it = tracks_.find(key); + if (it != tracks_.end()) { + return *it->second; + } + + auto track = std::make_unique(stableTrackId(key)); + auto descriptor = track->Serialize(); + descriptor.set_name(std::move(name)); + perfetto::TrackEvent::SetTrackDescriptor(*track, descriptor); + auto emplaceResult = tracks_.emplace(std::move(key), std::move(track)); + return *emplaceResult.first->second; +} + +void PerfettoTraceLogger::finalizeTrace( + const Config& /*config*/, + std::unique_ptr /*buffers*/, + int64_t /*endTime*/, + std::unordered_map>& /*metadata*/) { + if (!started_ || finalized_) { + return; + } + writeTraceFile(); +} + +void PerfettoTraceLogger::finalizeMemoryTrace( + const std::string&, + const Config&) { + LOG(INFO) << "finalizeMemoryTrace not implemented for PerfettoTraceLogger"; +} + +void PerfettoTraceLogger::emitFlowIfNeeded(const ITraceActivity& /*activity*/) {} + +void PerfettoTraceLogger::writeTraceFile() { + perfetto::TrackEvent::Flush(); + tracingSession_->FlushBlocking(); + tracingSession_->StopBlocking(); + auto traceData = std::vector{tracingSession_->ReadTraceBlocking()}; + if (traceData.empty()) { + throw std::runtime_error("Perfetto trace contains no data"); + } + + auto out = std::ofstream{fileName_, std::ios::out | std::ios::binary | std::ios::trunc}; + if (!out) { + throw std::runtime_error(fmt::format("Failed to open {}", fileName_)); + } + out.write(traceData.data(), static_cast(traceData.size())); + finalized_ = true; +} + +} // namespace KINETO_NAMESPACE + +#endif // KINETO_ENABLE_PERFETTO diff --git a/libkineto/src/output_perfetto.h b/libkineto/src/output_perfetto.h new file mode 100644 index 000000000..8c866f2e4 --- /dev/null +++ b/libkineto/src/output_perfetto.h @@ -0,0 +1,71 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#ifdef KINETO_ENABLE_PERFETTO + +#include +#include +#include +#include +#include + +#include "GenericTraceActivity.h" +#include "output_base.h" + +namespace perfetto { +class TracingSession; +class Track; +} // namespace perfetto + +namespace KINETO_NAMESPACE { + +class Config; +struct TraceSpan; + +class PerfettoTraceLogger : public libkineto::ActivityLogger { + public: + explicit PerfettoTraceLogger(const std::string& traceFileName); + ~PerfettoTraceLogger() override; + + void handleDeviceInfo(const DeviceInfo& info, int64_t time) override; + void handleOverheadInfo(const OverheadInfo& info, int64_t time) override; + void handleResourceInfo(const ResourceInfo& info, int64_t time) override; + void handleTraceSpan(const TraceSpan& span) override; + void handleActivity(const ITraceActivity& activity) override; + void handleGenericActivity(const GenericTraceActivity& activity) override; + + void handleTraceStart( + const std::unordered_map& metadata, + const std::string& device_properties) override; + + void finalizeTrace( + const Config& config, + std::unique_ptr buffers, + int64_t endTime, + std::unordered_map>& metadata) override; + + void finalizeMemoryTrace(const std::string&, const Config&) override; + + private: + const perfetto::Track& trackForActivity(const ITraceActivity& activity); + const perfetto::Track& trackForDevice(int64_t deviceId, std::string name); + void emitFlowIfNeeded(const ITraceActivity& activity); + void writeTraceFile(); + + std::string fileName_; + bool started_{false}; + bool finalized_{false}; + std::unique_ptr tracingSession_; + std::unordered_map> tracks_; +}; + +} // namespace KINETO_NAMESPACE + +#endif // KINETO_ENABLE_PERFETTO diff --git a/libkineto/test/PerfettoTraceLoggerTest.cpp b/libkineto/test/PerfettoTraceLoggerTest.cpp new file mode 100644 index 000000000..1333261b7 --- /dev/null +++ b/libkineto/test/PerfettoTraceLoggerTest.cpp @@ -0,0 +1,59 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include +#include +#include +#include + +#include "ActivityBuffers.h" +#include "ActivityType.h" +#include "Config.h" +#include "GenericTraceActivity.h" +#include "ThreadUtil.h" +#include "TraceSpan.h" +#include "output_perfetto.h" + +namespace libkineto { +namespace { + +TEST(PerfettoTraceLoggerTest, WritesNonEmptyPftraceFile) { + auto tracePath = std::filesystem::temp_directory_path() / + ("kineto_perfetto_logger_test_" + std::to_string(processId()) + ".pftrace"); + std::filesystem::remove(tracePath); + + { + PerfettoTraceLogger logger(tracePath.string()); + logger.handleTraceStart({{"unit_test", "true"}}, ""); + + TraceSpan span(1000, 5000, "Test trace"); + GenericTraceActivity op(span, ActivityType::CPU_OP, "aten::add"); + op.startTime = 1000; + op.endTime = 3000; + op.device = processId(); + op.resource = systemThreadId(); + op.threadId = systemThreadId(); + op.id = 7; + op.addMetadataQuoted("Input type", "float"); + + logger.handleActivity(op); + + Config cfg; + std::unordered_map> metadata; + logger.finalizeTrace(cfg, nullptr, 5000, metadata); + } + + ASSERT_TRUE(std::filesystem::exists(tracePath)); + EXPECT_GT(std::filesystem::file_size(tracePath), 0); + std::filesystem::remove(tracePath); +} + +} // namespace +} // namespace libkineto