From 06722c23bf1aaece4b219b1df24efc3f39c14f58 Mon Sep 17 00:00:00 2001 From: lucylq Date: Thu, 4 Sep 2025 21:58:17 -0700 Subject: [PATCH] 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] --- extension/pybindings/pybindings.cpp | 81 ++++++++++++++++++++--------- 1 file changed, 55 insertions(+), 26 deletions(-) diff --git a/extension/pybindings/pybindings.cpp b/extension/pybindings/pybindings.cpp index 2f414816aae..74fd13a67d5 100644 --- a/extension/pybindings/pybindings.cpp +++ b/extension/pybindings/pybindings.cpp @@ -160,8 +160,7 @@ void setup_output_storage( inline std::unique_ptr load_module_from_buffer( const void* ptr, size_t ptr_len, - bool enable_etdump, - size_t debug_buffer_size, + std::unique_ptr event_tracer, Program::Verification program_verification) { EXECUTORCH_SCOPE_PROF("load_module_from_buffer"); auto loader = std::make_unique(ptr, ptr_len); @@ -169,14 +168,13 @@ inline std::unique_ptr load_module_from_buffer( std::move(loader), nullptr, // memory_allocator nullptr, // temp_allocator - enable_etdump ? std::make_unique() : nullptr, + std::move(event_tracer), // event_tracer nullptr); // data_map_loader } inline std::unique_ptr load_module_from_file( const std::string& path, - bool enable_etdump, - size_t debug_buffer_size, + std::unique_ptr event_tracer, Program::Verification program_verification) { EXECUTORCH_SCOPE_PROF("load_module_from_file"); @@ -193,7 +191,7 @@ inline std::unique_ptr load_module_from_file( std::move(loader), nullptr, // memory_allocator nullptr, // temp_allocator - enable_etdump ? std::make_unique() : 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().data(), - py::len(buffer), - enable_etdump, - debug_buffer_size, - program_verification)) {} + : debug_buffer_size_(debug_buffer_size) { + std::unique_ptr event_tracer = enable_etdump + ? std::make_unique() + : nullptr; + if (enable_etdump && debug_buffer_size > 0) { + debug_buffer_ = std::make_unique(debug_buffer_size); + event_tracer->set_debug_buffer( + Span(debug_buffer_.get(), debug_buffer_size)); + event_tracer->set_event_tracer_debug_level( + EventTracerDebugLogLevel::kIntermediateOutputs); + } + module_ = load_module_from_buffer( + buffer.cast().data(), + py::len(buffer), + std::move(event_tracer), + program_verification); + } explicit PyModule( const void* ptr, @@ -503,12 +512,20 @@ 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 event_tracer = enable_etdump + ? std::make_unique() + : nullptr; + if (enable_etdump && debug_buffer_size > 0) { + debug_buffer_ = std::make_unique(debug_buffer_size); + event_tracer->set_debug_buffer( + Span(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, @@ -516,11 +533,20 @@ struct PyModule final { 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 event_tracer = enable_etdump + ? std::make_unique() + : nullptr; + if (enable_etdump && debug_buffer_size > 0) { + debug_buffer_ = std::make_unique(debug_buffer_size); + event_tracer->set_debug_buffer( + Span(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(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> output_tensors_; + // Hold onto the debug_buffer_ for the event_tracer. + std::unique_ptr 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(