Enable in xplat codebase

libkineto/include/Config.h

@@ -541,4 +541,9 @@ constexpr char kUseDaemonEnvVar[] = "KINETO_USE_DAEMON";
 
 bool isDaemonEnvVarSet();
 
+// Returns a reference to the protobuf trace enabled flag.
+// This allows the flag to be set externally (e.g., from JustKnobs in FBConfig)
+// and read in other components (e.g., ChromeTraceLogger).
+bool& get_protobuf_trace_enabled();
+
 } // namespace libkineto

libkineto/src/Config.cpp

@@ -631,4 +631,11 @@ void Config::setActivityDependentConfig() {
   AbstractConfig::setActivityDependentConfig();
 }
 
+// Returns a reference to the protobuf trace enabled flag.
+// Default is false. FBConfig will set this based on JustKnobs.
+bool& get_protobuf_trace_enabled() {
+  static bool _protobuf_trace_enabled = false;
+  return _protobuf_trace_enabled;
+}
+
 } // namespace KINETO_NAMESPACE

libkineto/src/PerfettoTraceBuilder.cpp

@@ -14,7 +14,6 @@
 #include "Logger.h"
 #include "TraceSpan.h"
 #include "output_base.h"
-#include "parfait/protos/perfetto/trace/perfetto_trace.pb.h"
 
 namespace KINETO_NAMESPACE {
 

libkineto/src/PerfettoTraceBuilder.h

@@ -8,17 +8,17 @@
 
 #pragma once
 
-#include <parfait/protos/perfetto/trace/perfetto_trace.pb.h>
+#include <kineto/libkineto/src/perfetto_trace.pb.h>
 #include <string>
 #include <unordered_map>
 #include "output_base.h"
 
 namespace KINETO_NAMESPACE {
 
 // PerfettoTraceBuilder is a helper class that constructs Perfetto protobuf
-// traces It provides similar interface to ActivityLogger but builds a Perfetto
-// trace object This class is designed to be used alongside ChromeTraceLogger to
-// generate both JSON and Perfetto protobuf traces
+// traces It provides similar interface to ActivityLogger but builds a
+// Perfetto trace object This class is designed to be used alongside
+// ChromeTraceLogger to generate both JSON and Perfetto protobuf traces
 class PerfettoTraceBuilder {
  public:
   PerfettoTraceBuilder();
@@ -76,7 +76,8 @@ class PerfettoTraceBuilder {
   // Map from (deviceId, resourceId, threadId) to track UUID
   std::unordered_map<std::string, TrackDescriptor*> trackUuids_;
 
-  // Map from deviceId to process TrackDescriptor pointer for later name updates
+  // Map from deviceId to process TrackDescriptor pointer for later name
+  // updates
   std::unordered_map<int64_t, TrackDescriptor*> deviceProcessDescriptors_;
 
   // Map from Main Thread to child thread

libkineto/src/output_json.cpp

@@ -15,9 +15,10 @@
 #include <fstream>
 #include <iterator>
 #include "Config.h"
-#include "TraceSpan.h"
-
 #include "Logger.h"
+// PerfettoTraceBuilder is only available in fbcode builds (in fb/ directory)
+#include "PerfettoTraceBuilder.h"
+#include "TraceSpan.h"
 
 namespace KINETO_NAMESPACE {
 
@@ -114,7 +115,7 @@ void ChromeTraceLogger::metadataToJSON(
     const std::unordered_map<std::string, std::string>& metadata) {
   for (const auto& [k, v] : metadata) {
     std::string sanitizedValue = v;
-    // There is a seperate mechanism for recording distributedInfo in on-demand
+    // There is a separate mechanism for recording distributedInfo in on-demand
     // so add a guard to prevent "double counting" in auto-trace.
     if (k == "distributedInfo") {
       distInfo_.distInfo_present_ = true;
@@ -135,6 +136,9 @@ void ChromeTraceLogger::handleTraceStart(
   if (!traceOf_) {
     return;
   }
+  if (perfettoBuilder_) {
+    perfettoBuilder_->handleTraceStart(metadata, device_properties);
+  }
   std::string display_unit = "ms";
 #ifdef DISPLAY_TRACE_IN_NS
   display_unit = "ns";
@@ -181,10 +185,26 @@ void ChromeTraceLogger::finalizeMemoryTrace(const std::string&, const Config&) {
   LOG(INFO) << "finalizeMemoryTrace not implemented for ChromeTraceLogger";
 }
 
+std::string ChromeTraceLogger::getPerfettoFileName() const {
+  // Replace .json extension with .pb
+  std::string pbFileName = fileName_;
+  size_t pos = pbFileName.rfind(".json");
+  if (pos != std::string::npos) {
+    pbFileName.replace(pos, 5, ".pb");
+  } else {
+    pbFileName += ".pb";
+  }
+  return pbFileName;
+}
+
 ChromeTraceLogger::ChromeTraceLogger(const std::string& traceFileName) {
   fileName_ = traceFileName.empty() ? defaultFileName() : traceFileName;
   traceOf_.clear(std::ios_base::badbit);
   openTraceFile();
+  if (get_protobuf_trace_enabled()) {
+    perfettoBuilder_ = std::make_unique<PerfettoTraceBuilder>();
+    LOG(INFO) << "initing perfetto trace builder";
+  }
 }
 
 void ChromeTraceLogger::handleDeviceInfo(
@@ -194,6 +214,10 @@ void ChromeTraceLogger::handleDeviceInfo(
     return;
   }
 
+  if (perfettoBuilder_) {
+    perfettoBuilder_->handleDeviceInfo(info);
+  }
+
   // M is for metadata
   // process_name needs a pid and a name arg
   // clang-format off
@@ -233,6 +257,10 @@ void ChromeTraceLogger::handleResourceInfo(
     return;
   }
 
+  if (perfettoBuilder_) {
+    perfettoBuilder_->handleResourceInfo(info);
+  }
+
   // M is for metadata
   // thread_name needs a pid and a name arg
   // clang-format off
@@ -260,6 +288,11 @@ void ChromeTraceLogger::handleResourceInfo(
 void ChromeTraceLogger::handleOverheadInfo(
     const OverheadInfo& info,
     int64_t time) {
+  // Also send to Perfetto trace builder
+  if (perfettoBuilder_) {
+    perfettoBuilder_->handleOverheadInfo(info, time);
+  }
+
   if (!traceOf_) {
     return;
   }
@@ -293,6 +326,10 @@ void ChromeTraceLogger::handleTraceSpan(const TraceSpan& span) {
     return;
   }
 
+  if (perfettoBuilder_) {
+    perfettoBuilder_->handleTraceSpan(span);
+  }
+
   uint64_t start = transToRelativeTime(span.startTime);
 
   // If endTime is 0 and start time is non-zero, dur can overflow. Add
@@ -377,6 +414,10 @@ void ChromeTraceLogger::handleActivity(const libkineto::ITraceActivity& op) {
     return;
   }
 
+  if (perfettoBuilder_) {
+    perfettoBuilder_->handleActivity(op);
+  }
+
   if (op.type() == ActivityType::CPU_INSTANT_EVENT) {
     handleGenericInstantEvent(op);
     return;
@@ -731,6 +772,18 @@ void ChromeTraceLogger::finalizeTrace(
 }})JSON", fileName_);
 
   traceOf_.close();
+
+  // Write Perfetto trace to separate file alongside JSON
+  if (perfettoBuilder_) {
+    std::string perfettoFileName = getPerfettoFileName();
+    if (perfettoBuilder_->writeToFile(perfettoFileName)) {
+      LOG(INFO) << "Perfetto protobuf trace written alongside JSON trace";
+    } else {
+      LOG(WARNING) << "Failed to write Perfetto protobuf trace to " << perfettoFileName;
+    }
+  }
+
+
   // On some systems, rename() fails if the destination file exists.
   // So, remove the destination file first.
   remove(fileName_.c_str());

libkineto/src/output_json.h

@@ -21,6 +21,7 @@
 // @lint-ignore-every CLANGTIDY facebook-hte-RelativeInclude
 #include "ActivityBuffers.h"
 #include "GenericTraceActivity.h"
+#include "PerfettoTraceBuilder.h"
 #include "output_base.h"
 #include "time_since_epoch.h"
 
@@ -115,6 +116,7 @@ class ChromeTraceLogger : public libkineto::ActivityLogger {
 
   void addOnDemandDistMetadata();
 
+  std::string getPerfettoFileName() const;
   std::string fileName_;
   std::string tempFileName_;
   std::ofstream traceOf_;
@@ -123,6 +125,7 @@ class ChromeTraceLogger : public libkineto::ActivityLogger {
   // pg_name, value is pgConfig that will be used to populate pg_config in
   // distributedInfo of trace
   std::unordered_map<std::string, pgConfig> pgMap = {};
+  std::unique_ptr<PerfettoTraceBuilder> perfettoBuilder_ = nullptr;
 };
 
 // std::chrono header start