support etdump generation in executorch.runtime

Summary: make et.runtime support etdump generation.

Differential Revision: D80373976

Author: Gasoonjia

extension/pybindings/pybindings.cpp

@@ -1434,7 +1434,7 @@ struct PyProgram final {
 
   std::unique_ptr<PyMethod> load_method(const std::string& method_name) {
     Result<Method> res = state_->program_->load_method(
-        method_name.c_str(), memory_->mem_manager());
+        method_name.c_str(), memory_->mem_manager(), event_tracer_.get());
     THROW_IF_ERROR(
         res.error(),
         "Failed to load method %s, error: 0x:%" PRIx32,
@@ -1459,6 +1459,39 @@ struct PyProgram final {
     return std::make_unique<PyMethodMeta>(state_, std::move(res.get()));
   }
 
+  bool has_etdump() {
+    return static_cast<bool>(event_tracer_);
+  }
+
+  void write_etdump_result_to_file(
+      const std::string& path,
+      const py::object& debug_buffer_path) {
+    if (!has_etdump()) {
+      throw std::runtime_error("No etdump found");
+    }
+    auto& etdump = *event_tracer_;
+    etdump_result result = etdump.get_etdump_data();
+    if (result.buf != nullptr && result.size > 0) {
+      write_data_to_file(path, result.buf, result.size);
+      free(result.buf);
+      if (debug_buffer_size_ > 0 &&
+          py::isinstance<py::str>(debug_buffer_path)) {
+        // Also write out the debug buffer to a separate file if requested.
+        std::string debug_buffer_path_str =
+            py::cast<std::string>(debug_buffer_path);
+        const auto debug_buffer = get_etdump_debug_buffer();
+        write_data_to_file(
+            debug_buffer_path_str, debug_buffer.data(), debug_buffer.size());
+      }
+    } else {
+      ET_LOG(
+          Info,
+          "No etdump data found, try rebuilding with "
+          "the CMake option EXECUTORCH_ENABLE_EVENT_TRACER or with "
+          "buck run --config executorch.event_tracer_enabled=true");
+    }
+  }
+
  private:
   std::shared_ptr<ProgramMemory> memory_;
   std::shared_ptr<ProgramState> state_;
@@ -1685,6 +1718,13 @@ PYBIND11_MODULE(EXECUTORCH_PYTHON_MODULE_NAME, m) {
           "method_meta",
           &PyProgram::method_meta,
           py::arg("method_name"),
+          call_guard)
+      .def("has_etdump", &PyProgram::has_etdump, call_guard)
+      .def(
+          "write_etdump_result_to_file",
+          &PyProgram::write_etdump_result_to_file,
+          py::arg("path"),
+          py::arg("debug_buffer_path") = py::none(),
           call_guard);
   py::class_<PyMethod>(m, "ExecuTorchMethod")
       .def("set_inputs", &PyMethod::set_inputs, py::arg("inputs"), call_guard)

runtime/__init__.py

@@ -137,6 +137,15 @@ def metadata(self, method_name: str) -> MethodMeta:
         """
         return self._program.method_meta(method_name)
 
+    def write_etdump_result_to_file(self, etdump_path: str, debug_buffer_path: str) -> None:
+        """Writes the etdump and debug result to a file.
+
+        Args:
+            etdump_path: The path to the etdump file.
+            debug_buffer_path: The path to the debug buffer file.
+        """
+        self._program.write_etdump_result_to_file(etdump_path, debug_buffer_path)
+
 
 class BackendRegistry:
     """The registry of backends that are available to the runtime."""
@@ -201,6 +210,8 @@ def load_program(
         data: Union[bytes, bytearray, BinaryIO, Path, str],
         *,
         verification: Verification = Verification.InternalConsistency,
+        enable_etdump: bool = False,
+        debug_buffer_size: int = 0,
     ) -> Program:
         """Loads an ExecuTorch program from a PTE binary.
 
@@ -214,8 +225,8 @@ def load_program(
         if isinstance(data, (Path, str)):
             p = self._legacy_module._load_program(
                 str(data),
-                enable_etdump=False,
-                debug_buffer_size=0,
+                enable_etdump=enable_etdump,
+                debug_buffer_size=debug_buffer_size,
                 program_verification=verification,
             )
             return Program(p, data=None)