diff --git a/examples/devtools/example_runner/example_runner.cpp b/examples/devtools/example_runner/example_runner.cpp index 43f7d7071c1..eafcd58c22e 100644 --- a/examples/devtools/example_runner/example_runner.cpp +++ b/examples/devtools/example_runner/example_runner.cpp @@ -24,6 +24,7 @@ #include #include +#include #include #include #include @@ -76,6 +77,7 @@ DEFINE_int32( 262144, // 256 KB "Size of the debug buffer in bytes to allocate for intermediate outputs and program outputs logging."); +using executorch::etdump::BufferDataSink; using executorch::etdump::ETDumpGen; using executorch::etdump::ETDumpResult; using executorch::extension::BufferDataLoader; @@ -216,27 +218,38 @@ int main(int argc, char** argv) { // be used by a single thread at at time, but it can be reused. // ETDumpGen etdump_gen; - Result method = - program->load_method(method_name, &memory_manager, &etdump_gen); - ET_CHECK_MSG( - method.ok(), - "Loading of method %s failed with status 0x%" PRIx32, - method_name, - static_cast(method.error())); - ET_LOG(Info, "Method loaded."); + // Malloc debug buffer and create if and only if we need to log intermediate + // tensor outputs void* debug_buffer = malloc(FLAGS_debug_buffer_size); + Result data_sink_ret = + BufferDataSink::create(debug_buffer, FLAGS_debug_buffer_size); + ET_CHECK_MSG( + data_sink_ret.ok(), + "Creating Datasink for etdump failed with status 0x%" PRIx32, + static_cast(data_sink_ret.error())); + + BufferDataSink* data_sink = &data_sink_ret.get(); + if (FLAGS_dump_intermediate_outputs) { - Span buffer((uint8_t*)debug_buffer, FLAGS_debug_buffer_size); - etdump_gen.set_debug_buffer(buffer); etdump_gen.set_event_tracer_debug_level( EventTracerDebugLogLevel::kIntermediateOutputs); + etdump_gen.set_data_sink(data_sink); } else if (FLAGS_dump_outputs) { - Span buffer((uint8_t*)debug_buffer, FLAGS_debug_buffer_size); - etdump_gen.set_debug_buffer(buffer); etdump_gen.set_event_tracer_debug_level( EventTracerDebugLogLevel::kProgramOutputs); + etdump_gen.set_data_sink(data_sink); } + + Result method = + program->load_method(method_name, &memory_manager, &etdump_gen); + ET_CHECK_MSG( + method.ok(), + "Loading of method %s failed with status 0x%" PRIx32, + method_name, + static_cast(method.error())); + ET_LOG(Info, "Method loaded."); + // Use the inputs embedded in the bundled program. status = executorch::bundled_program::load_bundled_input( *method, file_data.data(), FLAGS_testset_idx);