Add debug buffer support for pybindings

Add debug buffer to PyModule. Pass this to the event tracer and collect it back later.

Differential Revision: [D81751067](https://our.internmc.facebook.com/intern/diff/D81751067/)

[ghstack-poisoned]

Author: lucylq

extension/pybindings/pybindings.cpp

@@ -160,23 +160,21 @@ void setup_output_storage(
 inline std::unique_ptr<Module> load_module_from_buffer(
     const void* ptr,
     size_t ptr_len,
-    bool enable_etdump,
-    size_t debug_buffer_size,
+    std::unique_ptr<runtime::EventTracer> event_tracer,
     Program::Verification program_verification) {
   EXECUTORCH_SCOPE_PROF("load_module_from_buffer");
   auto loader = std::make_unique<BufferDataLoader>(ptr, ptr_len);
   return std::make_unique<Module>(
       std::move(loader),
       nullptr, // memory_allocator
       nullptr, // temp_allocator
-      enable_etdump ? std::make_unique<torch::executor::ETDumpGen>() : nullptr,
+      std::move(event_tracer), // event_tracer
       nullptr); // data_map_loader
 }
 
 inline std::unique_ptr<Module> load_module_from_file(
     const std::string& path,
-    bool enable_etdump,
-    size_t debug_buffer_size,
+    std::unique_ptr<runtime::EventTracer> event_tracer,
     Program::Verification program_verification) {
   EXECUTORCH_SCOPE_PROF("load_module_from_file");
 
@@ -193,7 +191,7 @@ inline std::unique_ptr<Module> load_module_from_file(
       std::move(loader),
       nullptr, // memory_allocator
       nullptr, // temp_allocator
-      enable_etdump ? std::make_unique<torch::executor::ETDumpGen>() : nullptr,
+      std::move(event_tracer), // event_tracer
       nullptr); // data_map_loader
 }
 
@@ -489,12 +487,23 @@ struct PyModule final {
       size_t debug_buffer_size = 0,
       Program::Verification program_verification =
           Program::Verification::InternalConsistency)
-      : module_(load_module_from_buffer(
-            buffer.cast<std::string_view>().data(),
-            py::len(buffer),
-            enable_etdump,
-            debug_buffer_size,
-            program_verification)) {}
+      : debug_buffer_size_(debug_buffer_size) {
+    std::unique_ptr<torch::executor::ETDumpGen> event_tracer = enable_etdump
+        ? std::make_unique<torch::executor::ETDumpGen>()
+        : nullptr;
+    if (enable_etdump && debug_buffer_size > 0) {
+      debug_buffer_ = std::make_unique<uint8_t[]>(debug_buffer_size);
+      event_tracer->set_debug_buffer(
+          Span<uint8_t>(debug_buffer_.get(), debug_buffer_size));
+      event_tracer->set_event_tracer_debug_level(
+          EventTracerDebugLogLevel::kIntermediateOutputs);
+    }
+    module_ = load_module_from_buffer(
+        buffer.cast<std::string_view>().data(),
+        py::len(buffer),
+        std::move(event_tracer),
+        program_verification);
+  }
 
   explicit PyModule(
       const void* ptr,
@@ -503,24 +512,41 @@ struct PyModule final {
       size_t debug_buffer_size = 0,
       Program::Verification program_verification =
           Program::Verification::InternalConsistency)
-      : module_(load_module_from_buffer(
-            ptr,
-            ptr_len,
-            enable_etdump,
-            debug_buffer_size,
-            program_verification)) {}
+      : debug_buffer_size_(debug_buffer_size) {
+    std::unique_ptr<torch::executor::ETDumpGen> event_tracer = enable_etdump
+        ? std::make_unique<torch::executor::ETDumpGen>()
+        : nullptr;
+    if (enable_etdump && debug_buffer_size > 0) {
+      debug_buffer_ = std::make_unique<uint8_t[]>(debug_buffer_size);
+      event_tracer->set_debug_buffer(
+          Span<uint8_t>(debug_buffer_.get(), debug_buffer_size));
+      event_tracer->set_event_tracer_debug_level(
+          EventTracerDebugLogLevel::kIntermediateOutputs);
+    }
+    module_ = load_module_from_buffer(
+        ptr, ptr_len, std::move(event_tracer), program_verification);
+  }
 
   explicit PyModule(
       const std::string& path,
       bool enable_etdump,
       size_t debug_buffer_size = 0,
       Program::Verification program_verification =
           Program::Verification::InternalConsistency)
-      : module_(load_module_from_file(
-            path,
-            enable_etdump,
-            debug_buffer_size,
-            program_verification)) {}
+      : debug_buffer_size_(debug_buffer_size) {
+    std::unique_ptr<torch::executor::ETDumpGen> event_tracer = enable_etdump
+        ? std::make_unique<torch::executor::ETDumpGen>()
+        : nullptr;
+    if (enable_etdump && debug_buffer_size > 0) {
+      debug_buffer_ = std::make_unique<uint8_t[]>(debug_buffer_size);
+      event_tracer->set_debug_buffer(
+          Span<uint8_t>(debug_buffer_.get(), debug_buffer_size));
+      event_tracer->set_event_tracer_debug_level(
+          EventTracerDebugLogLevel::kIntermediateOutputs);
+    }
+    module_ = load_module_from_file(
+        path, std::move(event_tracer), program_verification);
+  }
 
   PyModule(const PyModule&) = delete;
   PyModule& operator=(const PyModule&) = delete;
@@ -693,10 +719,9 @@ struct PyModule final {
         // 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 = module_->debug_buffer();
-        if (debug_buffer.size() > 0) {
+        if (debug_buffer_ && debug_buffer_size_ > 0) {
           write_data_to_file(
-              debug_buffer_path_str, debug_buffer.data(), debug_buffer.size());
+              debug_buffer_path_str, debug_buffer_.get(), debug_buffer_size_);
         }
       }
     } else {
@@ -746,6 +771,10 @@ struct PyModule final {
   // bundled programs.
   std::vector<std::optional<TensorPtr>> output_tensors_;
 
+  // Hold onto the debug_buffer_ for the event_tracer.
+  std::unique_ptr<uint8_t[]> debug_buffer_;
+  size_t debug_buffer_size_;
+
   void allocate_output_tensors(const std::string& method_name) {
     auto method_meta_result = module_->method_meta(method_name);
     THROW_IF_ERROR(