From e4370ed6dcd76f55b3560c3d15061fac00aa2bed Mon Sep 17 00:00:00 2001 From: Adeem Jassani Date: Wed, 10 Jun 2026 17:28:08 -0400 Subject: [PATCH] Add native Perfetto (.pftrace) export to torch.profiler Expose a new `export_pftrace()` API on both `torch.profiler.profile` (`_KinetoProfile`) and the legacy `torch.autograd.profiler.profile`, letting users write Perfetto protobuf traces directly instead of going through Chrome JSON conversion. Python side mirrors `export_chrome_trace`, including `.gz` handling. C++ side adds `ProfilerResult::savePftrace` / `ActivityTraceWrapper::savePftrace`, which route through Kineto's `pftrace://` save protocol, plus the pybind binding `save_pftrace`. Build is gated behind the new `USE_KINETO_PERFETTO` option (default OFF), which is plumbed to Kineto's `KINETO_ENABLE_PERFETTO`. When Kineto is built without Perfetto support the existing JSON path is unaffected. Co-authored-by: Cursor --- CMakeLists.txt | 1 + cmake/Dependencies.cmake | 2 ++ torch/autograd/profiler.py | 9 +++++++++ torch/csrc/autograd/init.cpp | 1 + torch/csrc/autograd/profiler_kineto.cpp | 4 ++++ torch/csrc/autograd/profiler_kineto.h | 1 + torch/csrc/profiler/kineto_shim.cpp | 13 +++++++++++++ torch/csrc/profiler/kineto_shim.h | 1 + torch/profiler/profiler.py | 18 ++++++++++++++++++ 9 files changed, 50 insertions(+) diff --git a/CMakeLists.txt b/CMakeLists.txt index a8b655d9cf571..58e8a0811d91f 100644 --- a/CMakeLists.txt +++ b/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) diff --git a/cmake/Dependencies.cmake b/cmake/Dependencies.cmake index 2baa187e69acd..15c0f3cee0e79 100644 --- a/cmake/Dependencies.cmake +++ b/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 "") diff --git a/torch/autograd/profiler.py b/torch/autograd/profiler.py index 02eb59e632aba..1f1d57c20b4cf 100644 --- a/torch/autograd/profiler.py +++ b/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: diff --git a/torch/csrc/autograd/init.cpp b/torch/csrc/autograd/init.cpp index f9ae57d12270e..0978c7b2d7319 100644 --- a/torch/csrc/autograd/init.cpp +++ b/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) { diff --git a/torch/csrc/autograd/profiler_kineto.cpp b/torch/csrc/autograd/profiler_kineto.cpp index b2d01a213936f..80925930f6822 100644 --- a/torch/csrc/autograd/profiler_kineto.cpp +++ b/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* ProfilerResult:: traceActivities() { diff --git a/torch/csrc/autograd/profiler_kineto.h b/torch/csrc/autograd/profiler_kineto.h index 520bd2cd20b93..307852c1ecfa2 100644 --- a/torch/csrc/autograd/profiler_kineto.h +++ b/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* traceActivities(); #endif diff --git a/torch/csrc/profiler/kineto_shim.cpp b/torch/csrc/profiler/kineto_shim.cpp index 86ab9b70a2974..5a7f771bef561 100644 --- a/torch/csrc/profiler/kineto_shim.cpp +++ b/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 { diff --git a/torch/csrc/profiler/kineto_shim.h b/torch/csrc/profiler/kineto_shim.h index 831981478f11f..1b40071a03ea4 100644 --- a/torch/csrc/profiler/kineto_shim.h +++ b/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& get() { return trace_; diff --git a/torch/profiler/profiler.py b/torch/profiler/profiler.py index ea6fbf2e3bef5..d8bf616cad1d8 100644 --- a/torch/profiler/profiler.py +++ b/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