Add native Perfetto (.pftrace) export to torch.profiler

CMakeLists.txt

@@ -280,6 +280,7 @@ cmake_dependent_option(USE_CUDSS "Use cuDSS" ON "USE_CUDA" OFF)
 cmake_dependent_option(USE_CUFILE "Use cuFile" ON "USE_CUDA AND NOT WIN32" OFF)
 option(USE_FBGEMM "Use FBGEMM (quantized 8-bit server operators)" ON)
 option(USE_KINETO "Use Kineto profiling library" ON)
+option(USE_KINETO_PERFETTO "Build Kineto with native Perfetto/pftrace export support" OFF)
 option(USE_CUPTI_SO "Use CUPTI as a shared library" ON)
 option(USE_GFLAGS "Use GFLAGS" OFF)
 option(USE_GLOG "Use GLOG" OFF)

cmake/Dependencies.cmake

@@ -1692,12 +1692,14 @@ if(USE_KINETO)
   set(KINETO_SOURCE_DIR "${CAFFE2_THIRD_PARTY_ROOT}/kineto/libkineto" CACHE STRING "")
   set(KINETO_BUILD_TESTS OFF CACHE BOOL "")
   set(KINETO_LIBRARY_TYPE "static" CACHE STRING "")
+  set(KINETO_ENABLE_PERFETTO ${USE_KINETO_PERFETTO} CACHE BOOL "" FORCE)
 
   message(STATUS "Configuring Kineto dependency:")
   message(STATUS "  KINETO_SOURCE_DIR = ${KINETO_SOURCE_DIR}")
   message(STATUS "  KINETO_BUILD_TESTS = ${KINETO_BUILD_TESTS}")
   message(STATUS "  KINETO_LIBRARY_TYPE = ${KINETO_LIBRARY_TYPE}")
   message(STATUS "  KINETO_BACKEND = ${KINETO_BACKEND}")
+  message(STATUS "  KINETO_ENABLE_PERFETTO = ${KINETO_ENABLE_PERFETTO}")
 
   if(KINETO_BACKEND STREQUAL "rocm")
     if("$ENV{ROCM_SOURCE_DIR}" STREQUAL "")

torch/autograd/profiler.py

@@ -549,6 +549,15 @@ def export_chrome_trace(self, path, metadata=None, use_python_export=False):
 
     export_chrome_trace.__doc__ = EventList.export_chrome_trace.__doc__
 
+    def export_pftrace(self, path):
+        """
+        Exports the collected Kineto trace in Perfetto pftrace format. Only
+        supported when Kineto was built with Perfetto export support.
+        """
+        if not kineto_available():
+            raise AssertionError("export_pftrace() requires Kineto support")
+        self.kineto_results.save_pftrace(path)  # type: ignore[union-attr]
+
     def export_stacks(self, path: str, metric: str = "self_cpu_time_total"):
         self._ensure_function_events()
         if self._function_events is None:

torch/csrc/autograd/init.cpp

@@ -410,6 +410,7 @@ PyObject* THPAutograd_initExtension(PyObject* _unused, PyObject* unused) {
       .def("experimental_event_tree", &ProfilerResult::event_tree)
 #ifdef USE_KINETO
       .def("save", &ProfilerResult::save)
+      .def("save_pftrace", &ProfilerResult::savePftrace)
       .def(
           "trace_activities",
           [](py::object self) {

torch/csrc/autograd/profiler_kineto.cpp

@@ -1368,6 +1368,10 @@ void ProfilerResult::save(const std::string& path) {
   trace_->save(path);
 }
 
+void ProfilerResult::savePftrace(const std::string& path) {
+  trace_->savePftrace(path);
+}
+
 #ifdef USE_KINETO
 const std::vector<const libkineto::ITraceActivity*>* ProfilerResult::
     traceActivities() {

torch/csrc/autograd/profiler_kineto.h

@@ -125,6 +125,7 @@ struct TORCH_API ProfilerResult {
   }
 
   void save(const std::string& path);
+  void savePftrace(const std::string& path);
 #ifdef USE_KINETO
   const std::vector<const libkineto::ITraceActivity*>* traceActivities();
 #endif

torch/csrc/profiler/kineto_shim.cpp

@@ -202,6 +202,19 @@ void ActivityTraceWrapper::save(const std::string& path) {
 #endif // USE_KINETO
 }
 
+void ActivityTraceWrapper::savePftrace(const std::string& path) {
+#ifdef USE_KINETO
+  TORCH_CHECK(!saved_, "Trace is already saved.");
+  TORCH_CHECK(trace_ != nullptr, "Missing trace.")
+  trace_->save("pftrace://" + path);
+  saved_ = true;
+#else
+  TORCH_CHECK(
+      false,
+      "Saving a pftrace requires using torch.profiler with Kineto support (USE_KINETO=1)");
+#endif // USE_KINETO
+}
+
 namespace {
 // Handles processing of Experimental Config options for Kineto
 class ExperimentalConfigWrapper {

torch/csrc/profiler/kineto_shim.h

@@ -95,6 +95,7 @@ struct ActivityTraceWrapper {
   ActivityTraceWrapper() = default;
   explicit operator bool() const;
   void save(const std::string& path);
+  void savePftrace(const std::string& path);
 
   const std::unique_ptr<interface_trace_t>& get() {
     return trace_;

torch/profiler/profiler.py

@@ -445,6 +445,24 @@ def export_chrome_trace(self, path: str, use_python_export: bool = False):
         else:
             self.profiler.export_chrome_trace(path, self._trace_metadata)
 
+    def export_pftrace(self, path: str):
+        """
+        Exports the collected trace in Perfetto pftrace format. If a schedule is
+        used, only the last cycle is exported.
+        """
+        if self.profiler is None:
+            raise AssertionError(
+                "Profiler must be initialized before exporting pftrace"
+            )
+        if path.endswith(".gz"):
+            with tempfile.NamedTemporaryFile("w+b", suffix=".pftrace") as fp:
+                retvalue = self.profiler.export_pftrace(fp.name)
+                with open(fp.name, "rb") as fin, gzip.open(path, "wb") as fout:
+                    fout.writelines(fin)
+            return retvalue
+        else:
+            return self.profiler.export_pftrace(path)
+
     def export_stacks(self, path: str, metric: str = "self_cpu_time_total"):
         """Save stack traces to a file