SDK example runner changes to suppport intermediate tensor logging (#1307)

Summary:
Pull Request resolved: #1307

This diff adds support to the `sdk_example_runner` to take in command line parameters to enable intermediate/program output logging to etdump. Along with that we take in additional parameters for debug_buffer path and debug_buffer size. The debug_buffer is where the tensor and scalar data is stored (just the raw data not the metadata).

Reviewed By: Jack-Khuu

Differential Revision: D51599690

fbshipit-source-id: 3c4a00edd4db62000cc14990427f0144a83c3278

Author: tarun292

examples/sdk/sdk_example_runner/sdk_example_runner.cpp

@@ -61,6 +61,23 @@ DEFINE_bool(
     false,
     "Print the output of the ET model to stdout, if needs.");
 
+DEFINE_bool(dump_outputs, false, "Dump outputs to etdump file");
+
+DEFINE_bool(
+    dump_intermediate_outputs,
+    false,
+    "Dump intermediate outputs to etdump file.");
+
+DEFINE_string(
+    debug_output_path,
+    "debug_output.bin",
+    "Path to dump debug outputs to.");
+
+DEFINE_int32(
+    debug_buffer_size,
+    262144, // 256 KB
+    "Size of the debug buffer in bytes to allocate for intermediate outputs and program outputs logging.");
+
 using namespace torch::executor;
 using torch::executor::util::FileDataLoader;
 
@@ -199,6 +216,18 @@ int main(int argc, char** argv) {
       method.error());
   ET_LOG(Info, "Method loaded.");
 
+  void* debug_buffer = malloc(FLAGS_debug_buffer_size);
+  if (FLAGS_dump_intermediate_outputs) {
+    Span<uint8_t> buffer((uint8_t*)debug_buffer, FLAGS_debug_buffer_size);
+    etdump_gen.set_debug_buffer(buffer);
+    etdump_gen.set_event_tracer_debug_level(
+        EventTracerDebugLogLevel::kIntermediateOutputs);
+  } else if (FLAGS_dump_outputs) {
+    Span<uint8_t> buffer((uint8_t*)debug_buffer, FLAGS_debug_buffer_size);
+    etdump_gen.set_debug_buffer(buffer);
+    etdump_gen.set_event_tracer_debug_level(
+        EventTracerDebugLogLevel::kProgramOutputs);
+  }
   // Prepare the inputs.
   exec_aten::ArrayRef<void*> inputs;
   // Use the inputs embedded in the bundled program.
@@ -264,5 +293,12 @@ int main(int argc, char** argv) {
     ET_LOG(Info, "Model verified successfully.");
   }
 
+  if (FLAGS_dump_outputs || FLAGS_dump_intermediate_outputs) {
+    FILE* f = fopen(FLAGS_debug_output_path.c_str(), "w+");
+    fwrite((uint8_t*)debug_buffer, 1, FLAGS_debug_buffer_size, f);
+    fclose(f);
+  }
+  free(debug_buffer);
+
   return 0;
 }

runtime/executor/method.cpp

@@ -1035,7 +1035,7 @@ Error Method::experimental_reset_execution() {
 void Method::log_outputs() {
 #ifdef ET_EVENT_TRACER_ENABLED
   if (event_tracer_ != nullptr) {
-    if (event_tracer->event_tracer_debug_level() >=
+    if (event_tracer_->event_tracer_debug_level() >=
         EventTracerDebugLogLevel::kProgramOutputs) {
       for (size_t i = 0; i < outputs_size(); i++) {
         internal::event_tracer_log_evalue_output(event_tracer_, get_output(i));

sdk/bundled_program/bundled_program.cpp

@@ -16,6 +16,7 @@
 #include <ATen/ATen.h>
 #endif // USE_ATEN_LIB
 
+#include <executorch/runtime/core/event_tracer_hooks.h>
 #include <executorch/runtime/core/exec_aten/util/dim_order_util.h>
 #include <executorch/runtime/core/memory_allocator.h>
 #include <executorch/runtime/executor/method.h>
@@ -293,6 +294,9 @@ __ET_NODISCARD Error LoadBundledInput(
         status);
   }
 
+  internal::event_tracer_set_bundled_input_index(
+      method.get_event_tracer(), testset_idx);
+
   return Error::Ok;
 }