support etdump generation in executorch.runtime (#14205)

This PR was created by the merge bot to help merge the original PR into
the main branch.
ghstack PR number: #14172 by
@Gasoonjia
^ Please use this as the source of truth for the PR details, comments,
and reviews
ghstack PR base:
https://github.com/pytorch/executorch/tree/gh/gasoonjia/41/base
ghstack PR head:
https://github.com/pytorch/executorch/tree/gh/gasoonjia/41/head
Merge bot PR base: https://github.com/pytorch/executorch/tree/main
Merge bot PR head:
https://github.com/pytorch/executorch/tree/gh/gasoonjia/41/orig
@diff-train-skip-merge

Co-authored-by: gasoonjia <[email protected]>

Author: pytorchbot

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

@@ -1296,7 +1296,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,
@@ -1321,6 +1321,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_;
@@ -1554,6 +1587,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)

pytest.ini

@@ -8,12 +8,12 @@ addopts =
     --capture=sys
     # don't suppress warnings, but don't shove them all to the end either
     -p no:warnings
-    
+
     # === TEST DIRECTORIES TO RUN ===
-    
+
     # ci/scripts
     .ci/scripts/tests
-    
+
     # backends
     backends/apple/coreml/test
     backends/test/harness/tests
@@ -28,15 +28,15 @@ addopts =
     --ignore=backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py
     # Ignore backends/test root - WIP testing infra, see https://github.com/pytorch/executorch/discussions/11140
     --ignore=backends/test
-    
+
     # codegen
     codegen/test
-    
+
     # devtools
     devtools/
     # Ignore test with missing dependencies
     --ignore=devtools/visualization/visualization_utils_test.py
-    
+
     # examples
     examples/models/test
     examples/models/llama/tests
@@ -53,7 +53,7 @@ addopts =
     --ignore=examples/models/llava/test/test_pte.py
     # Ignore failing llava tests (missing accelerate dependency)
     --ignore=examples/models/llava/test/test_llava.py
-    
+
     # exir
     exir/
     # Ignore tests with missing custom_ops_generated_lib dependencies
@@ -73,11 +73,11 @@ addopts =
     --ignore=exir/operator/test/test_operator.py
     --ignore=exir/tests/test_common.py
     --ignore=exir/tests/test_op_convert.py
-    
+
     # export
     export/tests
     --ignore=export/tests/test_export_stages.py
-    
+
     # extension
     extension/
     # Ignore tests with missing dependencies or build issues
@@ -87,28 +87,29 @@ addopts =
     --ignore=extension/llm/tokenizers/third-party/sentencepiece/python/test/sentencepiece_test.py
     # Ignore failing tokenizer tests
     --ignore=extension/llm/tokenizers/test/test_tekken_python.py
-    
+
     # kernels
     kernels/prim_ops/test
     kernels/quantized
     kernels/test/test_case_gen.py
     # Ignore test depending on test-only cpp ops lib
     --ignore=kernels/quantized/test/test_quant_dequant_per_token.py
-    
+
     # profiler
     profiler/
     # Ignore test with missing dependencies
     --ignore=profiler/test/test_profiler_e2e.py
-    
+
     # runtime
     runtime
-    
+    # Ignore tests with missing compiler dependencies
+    --ignore=runtime/test/test_runtime_etdump_gen.py
     # test
     test/
     # Ignore tests with missing dependencies
     --ignore=test/end2end/test_end2end.py
     --ignore=test/end2end/test_temp_allocator_fix.py
-    
+
     # tools
     tools/cmake
 

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,17 @@ 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 +255,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 +270,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,16 @@ runtime.python_test(
     deps = [
         "//executorch/extension/pybindings/test:make_test",
         "//executorch/runtime:runtime",
+        "//executorch/devtools/etdump:serialize",
+    ],
+)
+
+runtime.python_test(
+    name = "test_runtime_etdump_gen",
+    srcs = ["test_runtime_etdump_gen.py"],
+    deps = [
+        "//executorch/extension/pybindings/test:make_test",
+        "//executorch/runtime:runtime",
+        "//executorch/devtools/etdump:serialize",
     ],
 )

runtime/test/test_runtime_etdump_gen.py

@@ -0,0 +1,97 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+import tempfile
+import unittest
+
+import torch
+from executorch.devtools.etdump.serialize import deserialize_from_etdump_flatcc
+
+from executorch.extension.pybindings.test.make_test import create_program, ModuleAdd
+from executorch.runtime import Runtime, Verification
+
+
+class RuntimeETDumpGenTest(unittest.TestCase):
+    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"
+            )