support etdump generation in executorch.runtime

devtools/etdump/TARGETS

@@ -29,6 +29,7 @@ runtime.python_library(
     },
     visibility = [
         "//executorch/devtools/...",
+        "//executorch/runtime/test/...",
     ],
     deps = [
         "fbsource//third-party/pypi/setuptools:setuptools",

extension/pybindings/pybindings.cpp

@@ -1449,7 +1449,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,
@@ -1474,6 +1474,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 set to ON or with "
+          "buck run --config executorch.event_tracer_enabled=true");
+    }
+  }
+
  private:
   std::shared_ptr<ProgramMemory> memory_;
   std::shared_ptr<ProgramState> state_;
@@ -1706,6 +1739,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

@@ -31,12 +31,55 @@
 
 .. code-block:: text
 
-    Program methods: ('forward', 'forward2')
+    Program methods: {'forward'}
     Ran forward((tensor([[1., 1.],
             [1., 1.]]), tensor([[1., 1.],
             [1., 1.]])))
-      outputs: [tensor([[1., 1.],
-            [1., 1.]])]
+      outputs: [tensor([[2., 2.],
+            [2., 2.]])]
+
+Example usage with ETDump generation:
+
+.. code-block:: python
+
+    from pathlib import Path
+    import os
+
+    import torch
+    from executorch.runtime import Verification, Runtime, Program, Method
+
+    # Create program with etdump generation enabled
+    et_runtime: Runtime = Runtime.get()
+    program: Program = et_runtime.load_program(
+        Path("/tmp/program.pte"),
+        verification=Verification.Minimal,
+        enable_etdump=True,
+        debug_buffer_size=1e7, # A large buffer size to ensure that all debug info is captured
+    )
+
+    # Load method and execute
+    forward: Method = program.load_method("forward")
+    inputs = (torch.ones(2, 2), torch.ones(2, 2))
+    outputs = forward.execute(inputs)
+
+    # Write etdump result to file
+    etdump_file = "/tmp/etdump_output.etdp"
+    debug_file = "/tmp/debug_output.bin"
+    program.write_etdump_result_to_file(etdump_file, debug_file)
+
+    # Check that files were created
+    print(f"ETDump file created: {os.path.exists(etdump_file)}")
+    print(f"Debug file created: {os.path.exists(debug_file)}")
+    print("Directory contents:", os.listdir("/tmp"))
+
+Example output:
+
+.. code-block:: text
+
+    Program methods: {'forward'}
+    ETDump file created: True
+    Debug file created: True
+    Directory contents: ['program.pte', 'etdump_output.etdp', 'debug_output.bin']
 """
 
 import functools
@@ -137,6 +180,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 +253,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 +268,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)

runtime/test/TARGETS

@@ -8,5 +8,9 @@ runtime.python_test(
     deps = [
         "//executorch/extension/pybindings/test:make_test",
         "//executorch/runtime:runtime",
+        "//executorch/devtools/etdump:serialize",
     ],
+    labels = ci.labels(
+        ci.buckconfig("executorch.event_tracer_enabled", "true"),
+    ),
 )

runtime/test/test_runtime.py

@@ -16,7 +16,8 @@
     ModuleMulti,
 )
 from executorch.runtime import Runtime, Verification
-
+import os
+from executorch.devtools.etdump.serialize import deserialize_from_etdump_flatcc
 
 class RuntimeTest(unittest.TestCase):
     def test_smoke(self):
@@ -76,3 +77,63 @@ def test_add(program):
             with open(f.name, "rb") as f:
                 program = runtime.load_program(f.read())
                 test_add(program)
+
+    def test_etdump_generation(self):
+        """Test etdump generation by creating a program with etdump enabled and verifying the output."""
+
+        ep, inputs = create_program(ModuleAdd())
+        runtime = Runtime.get()
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            # Save the program to a file
+            program_path = os.path.join(temp_dir, "test_program.pte")
+            with open(program_path, "wb") as f:
+                f.write(ep.buffer)
+
+            # Load program with etdump generation enabled
+            program = runtime.load_program(
+                program_path,
+                verification=Verification.Minimal,
+                enable_etdump=True,
+                debug_buffer_size=int(1e7),  # Large buffer size to ensure all debug info is captured
+            )
+
+            # Execute the method
+            method = program.load_method("forward")
+            outputs = method.execute(inputs)
+
+            # Verify the computation is correct
+            self.assertTrue(torch.allclose(outputs[0], inputs[0] + inputs[1]))
+
+            # Write etdump result to files
+            etdump_path = os.path.join(temp_dir, "etdump_output.etdp")
+            debug_path = os.path.join(temp_dir, "debug_output.bin")
+            program.write_etdump_result_to_file(etdump_path, debug_path)
+
+            # Check that files were created
+            self.assertTrue(os.path.exists(etdump_path), f"ETDump file not created at {etdump_path}")
+            self.assertTrue(os.path.exists(debug_path), f"Debug file not created at {debug_path}")
+
+            # Verify the etdump file is not empty
+            etdump_size = os.path.getsize(etdump_path)
+            self.assertGreater(etdump_size, 0, "ETDump file is empty")
+
+            # Read and deserialize the etdump file to verify its structure
+            with open(etdump_path, "rb") as f:
+                etdump_data = f.read()
+
+            # Deserialize the etdump and check its header/structure
+            etdump = deserialize_from_etdump_flatcc(etdump_data)
+
+            # Verify ETDump header properties
+            self.assertIsInstance(etdump.version, int, "ETDump version should be an integer")
+            self.assertGreaterEqual(etdump.version, 0, "ETDump version should be non-negative")
+
+            # Verify run_data structure
+            self.assertIsInstance(etdump.run_data, list, "ETDump run_data should be a list")
+            self.assertGreater(len(etdump.run_data), 0, "ETDump should contain at least one run data entry")
+
+            # Check the first run_data entry
+            run_data = etdump.run_data[0]
+            self.assertIsInstance(run_data.events, list, "Run data should contain events list")
+            self.assertGreater(len(run_data.events), 0, "Run data should contain at least one events")