Intermediate output logging enablement - Runtime APIs (#3822)

Summary:
Pull Request resolved: #3822

This diff creates and implements the APIs that log intermediate outputs from delegates into etdump file.
Design doc: https://docs.google.com/document/d/1qGHsgd-roqtxPz4CrUlqGrKaAtbaf9bArziMuwHD0So/edit

bypass-github-export-checks
bypass-github-pytorch-ci-checks
bypass-github-executorch-ci-checks

Reviewed By: tarun292

Differential Revision: D56941407

fbshipit-source-id: df985bba0e7ada6917afed5ba305d2a43e49d06b

Author: Olivia-liu

runtime/core/event_tracer.h

@@ -6,6 +6,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#include <executorch/runtime/core/array_ref.h>
 #include <executorch/runtime/core/evalue.h>
 #include <executorch/runtime/platform/platform.h>
 #include <stdlib.h>
@@ -256,6 +257,106 @@ class EventTracer {
       const EValue& evalue,
       LoggedEValueType evalue_type) = 0;
 
+  /**
+   * Log an intermediate tensor output from a delegate.
+   *
+   * @param[in] name Human readable name for the delegate event. This name has
+   * to be the same name that was passed in during the Debug delegate mapping
+   * generation in the export/ahead-of-time process. If indices and not names
+   * are used by this delegate to identify ops executed in the backend then
+   * nullptr can be passed in. Users calling this interface do not need to keep
+   * the memory pointed to by this pointer around. The string must be copied
+   * over into internal memory during this call.
+   * @param[in] delegate_debug_index The id of the delegate event. If string
+   * based names are used by this delegate to identify ops executed in the
+   * backend then kUnsetDebugHandle should be passed in here.
+   * @param[in] output The tensor type output to be logged.
+   */
+  virtual void log_intermediate_output_delegate(
+      const char* name,
+      DebugHandle delegate_debug_index,
+      const Tensor& output) = 0;
+
+  /**
+   * Log an intermediate tensor array output from a delegate.
+   *
+   * @param[in] name Human readable name for the delegate event. This name has
+   * to be the same name that was passed in during the Debug delegate mapping
+   * generation in the export/ahead-of-time process. If indices and not names
+   * are used by this delegate to identify ops executed in the backend then
+   * nullptr can be passed in. Users calling this interface do not need to keep
+   * the memory pointed to by this pointer around. The string must be copied
+   * over into internal memory during this call.
+   * @param[in] delegate_debug_index The id of the delegate event. If string
+   * based names are used by this delegate to identify ops executed in the
+   * backend then kUnsetDebugHandle should be passed in here.
+   * @param[in] output The tensor array type output to be logged.
+   */
+  virtual void log_intermediate_output_delegate(
+      const char* name,
+      DebugHandle delegate_debug_index,
+      const ArrayRef<Tensor> output) = 0;
+
+  /**
+   * Log an intermediate int output from a delegate.
+   *
+   * @param[in] name Human readable name for the delegate event. This name has
+   * to be the same name that was passed in during the Debug delegate mapping
+   * generation in the export/ahead-of-time process. If indices and not names
+   * are used by this delegate to identify ops executed in the backend then
+   * nullptr can be passed in. Users calling this interface do not need to keep
+   * the memory pointed to by this pointer around. The string must be copied
+   * over into internal memory during this call.
+   * @param[in] delegate_debug_index The id of the delegate event. If string
+   * based names are used by this delegate to identify ops executed in the
+   * backend then kUnsetDebugHandle should be passed in here.
+   * @param[in] output The int type output to be logged.
+   */
+  virtual void log_intermediate_output_delegate(
+      const char* name,
+      DebugHandle delegate_debug_index,
+      const int& output) = 0;
+
+  /**
+   * Log an intermediate bool output from a delegate.
+   *
+   * @param[in] name Human readable name for the delegate event. This name has
+   * to be the same name that was passed in during the Debug delegate mapping
+   * generation in the export/ahead-of-time process. If indices and not names
+   * are used by this delegate to identify ops executed in the backend then
+   * nullptr can be passed in. Users calling this interface do not need to keep
+   * the memory pointed to by this pointer around. The string must be copied
+   * over into internal memory during this call.
+   * @param[in] delegate_debug_index The id of the delegate event. If string
+   * based names are used by this delegate to identify ops executed in the
+   * backend then kUnsetDebugHandle should be passed in here.
+   * @param[in] output The bool type output to be logged.
+   */
+  virtual void log_intermediate_output_delegate(
+      const char* name,
+      DebugHandle delegate_debug_index,
+      const bool& output) = 0;
+
+  /**
+   * Log an intermediate double output from a delegate.
+   *
+   * @param[in] name Human readable name for the delegate event. This name has
+   * to be the same name that was passed in during the Debug delegate mapping
+   * generation in the export/ahead-of-time process. If indices and not names
+   * are used by this delegate to identify ops executed in the backend then
+   * nullptr can be passed in. Users calling this interface do not need to keep
+   * the memory pointed to by this pointer around. The string must be copied
+   * over into internal memory during this call.
+   * @param[in] delegate_debug_index The id of the delegate event. If string
+   * based names are used by this delegate to identify ops executed in the
+   * backend then kUnsetDebugHandle should be passed in here.
+   * @param[in] output The double type output to be logged.
+   */
+  virtual void log_intermediate_output_delegate(
+      const char* name,
+      DebugHandle delegate_debug_index,
+      const double& output) = 0;
+
   /**
    * Helper function to set the chain id ands debug handle. Users have two
    * options, the first is that they can directly pass in the chain id and debug

runtime/core/event_tracer_hooks_delegate.h

@@ -145,5 +145,47 @@ inline void event_tracer_log_profiling_delegate(
 #endif
 }
 
+/**
+ * This templated interfaces can be called in a loop etc. to log any number of
+ * debug events that are part of this delegate. Supported values types are int,
+ * bool, double, tensor and array of tensors. Can be left in production code as
+ * these hooks compile conditionally.
+ *
+ * @param[in] event_tracer The event tracer instance that is doing the logging.
+ * @param[in] name Human readable name for the delegate event. This name has
+ * to be the same name that was passed in during the Debug delegate mapping
+ * generation in the export/ahead-of-time process. If indices and not names
+ * are used by this delegate to identify ops executed in the backend then
+ * nullptr can be passed in. Users calling this interface do not need to keep
+ * the memory pointed to by this pointer around. The string must
+ * be copied over into internal memory during this call.
+ * @param[in] delegate_debug_id The id of the delegate event. If string
+ * based names are used by this delegate to identify ops executed in the
+ * backend then -1 should be passed in here.
+ * @param[in] output The output to be logged.
+ */
+template <typename T>
+inline void event_tracer_log_output_delegate(
+    EventTracer* event_tracer,
+    const char* name,
+    DebugHandle delegate_debug_id,
+    const T& output) {
+#ifdef ET_EVENT_TRACER_ENABLED
+  if (event_tracer) {
+    static_assert(
+        std::is_same<T, int>::value || std::is_same<T, bool>::value ||
+            std::is_same<T, double>::value || std::is_same<T, Tensor>::value ||
+            std::is_same<T, ArrayRef<Tensor>>::value,
+        "Unsupported type for intermediate output");
+    event_tracer->log_intermediate_output_delegate(
+        name, delegate_debug_id, output);
+  }
+#else //! ET_EVENT_TRACER_ENABLED
+  (void)name;
+  (void)delegate_debug_id;
+  (void)output;
+#endif
+}
+
 } // namespace executor
 } // namespace torch

runtime/core/test/event_tracer_test.cpp

@@ -88,6 +88,51 @@ class DummyEventTracer : public EventTracer {
     (void)metadata_len;
   }
 
+  void log_intermediate_output_delegate(
+      const char* name,
+      DebugHandle delegate_debug_index,
+      const Tensor& output) override {
+    (void)name;
+    (void)delegate_debug_index;
+    (void)output;
+  }
+
+  void log_intermediate_output_delegate(
+      const char* name,
+      DebugHandle delegate_debug_index,
+      const ArrayRef<Tensor> output) override {
+    (void)name;
+    (void)delegate_debug_index;
+    (void)output;
+  }
+
+  void log_intermediate_output_delegate(
+      const char* name,
+      DebugHandle delegate_debug_index,
+      const int& output) override {
+    (void)name;
+    (void)delegate_debug_index;
+    (void)output;
+  }
+
+  virtual void log_intermediate_output_delegate(
+      const char* name,
+      DebugHandle delegate_debug_index,
+      const bool& output) override {
+    (void)name;
+    (void)delegate_debug_index;
+    (void)output;
+  }
+
+  virtual void log_intermediate_output_delegate(
+      const char* name,
+      DebugHandle delegate_debug_index,
+      const double& output) override {
+    (void)name;
+    (void)delegate_debug_index;
+    (void)output;
+  }
+
   void log_evalue(const EValue& evalue, LoggedEValueType evalue_type) override {
     logged_evalue_ = evalue;
     logged_evalue_type_ = evalue_type;

sdk/etdump/etdump_flatcc.cpp

@@ -288,6 +288,124 @@ void ETDumpGen::log_profiling_delegate(
   etdump_RunData_events_push_end(builder);
 }
 
+void ETDumpGen::log_intermediate_output_delegate(
+    const char* name,
+    DebugHandle delegate_debug_index,
+    const Tensor& output) {
+  log_intermediate_output_delegate_helper(name, delegate_debug_index, output);
+}
+
+void ETDumpGen::log_intermediate_output_delegate(
+    const char* name,
+    DebugHandle delegate_debug_index,
+    const ArrayRef<Tensor> output) {
+  log_intermediate_output_delegate_helper(name, delegate_debug_index, output);
+}
+
+void ETDumpGen::log_intermediate_output_delegate(
+    const char* name,
+    DebugHandle delegate_debug_index,
+    const int& output) {
+  log_intermediate_output_delegate_helper(name, delegate_debug_index, output);
+}
+
+void ETDumpGen::log_intermediate_output_delegate(
+    const char* name,
+    DebugHandle delegate_debug_index,
+    const bool& output) {
+  log_intermediate_output_delegate_helper(name, delegate_debug_index, output);
+}
+
+void ETDumpGen::log_intermediate_output_delegate(
+    const char* name,
+    DebugHandle delegate_debug_index,
+    const double& output) {
+  log_intermediate_output_delegate_helper(name, delegate_debug_index, output);
+}
+
+template <typename T>
+void ETDumpGen::log_intermediate_output_delegate_helper(
+    const char* name,
+    DebugHandle delegate_debug_index,
+    const T& output) {
+  ET_CHECK_MSG(
+      (name == nullptr) ^ (delegate_debug_index == -1),
+      "Only name or delegate_debug_index can be valid. Check DelegateMappingBuilder documentation for more details.");
+  if (debug_buffer.empty()) {
+    ET_CHECK_MSG(0, "Must pre-set debug buffer with set_debug_buffer()\n");
+    return;
+  }
+
+  check_ready_to_add_events();
+  int64_t string_id = name != nullptr ? create_string_entry(name) : -1;
+
+  etdump_DebugEvent_start(builder);
+
+  etdump_DebugEvent_chain_index_add(builder, chain_id_);
+  etdump_DebugEvent_instruction_id_add(builder, debug_handle_);
+  if (string_id == -1) {
+    etdump_DebugEvent_delegate_debug_id_int_add(builder, delegate_debug_index);
+  } else {
+    etdump_DebugEvent_delegate_debug_id_str_add(builder, string_id);
+  }
+
+  // Check the type of `output` then call the corresponding logging functions
+  if constexpr (std::is_same<T, Tensor>::value) {
+    long offset = copy_tensor_to_debug_buffer(output);
+    etdump_Tensor_ref_t tensor_ref = add_tensor_entry(builder, output, offset);
+
+    etdump_Value_start(builder);
+    etdump_Value_val_add(builder, etdump_ValueType_Tensor);
+    etdump_Value_tensor_add(builder, tensor_ref);
+
+  } else if constexpr (std::is_same<T, ArrayRef<Tensor>>::value) {
+    etdump_Tensor_vec_start(builder);
+    for (size_t i = 0; i < output.size(); ++i) {
+      long offset = copy_tensor_to_debug_buffer(output[i]);
+      etdump_Tensor_vec_push(
+          builder, add_tensor_entry(builder, output[i], offset));
+    }
+    etdump_Tensor_vec_ref_t tensor_vec_ref = etdump_Tensor_vec_end(builder);
+    etdump_TensorList_ref_t tensor_list_ref =
+        etdump_TensorList_create(builder, tensor_vec_ref);
+
+    etdump_Value_start(builder);
+    etdump_Value_val_add(builder, etdump_ValueType_TensorList);
+    etdump_Value_tensor_list_add(builder, tensor_list_ref);
+  } else if constexpr (std::is_same<T, int>::value) {
+    auto int_ref = etdump_Int_create(builder, output);
+
+    etdump_Value_start(builder);
+    etdump_Value_val_add(builder, etdump_ValueType_Int);
+    etdump_Value_int_value_add(builder, int_ref);
+  } else if constexpr (std::is_same<T, double>::value) {
+    auto double_ref = etdump_Double_create(builder, output);
+
+    etdump_Value_start(builder);
+    etdump_Value_double_value_add(builder, double_ref);
+    etdump_Value_val_add(builder, etdump_ValueType_Double);
+  } else if constexpr (std::is_same<T, bool>::value) {
+    flatbuffers_bool_t flatbuffer_bool_val =
+        output ? FLATBUFFERS_TRUE : FLATBUFFERS_FALSE;
+    auto bool_ref = etdump_Bool_create(builder, flatbuffer_bool_val);
+
+    etdump_Value_start(builder);
+    etdump_Value_bool_value_add(builder, bool_ref);
+    etdump_Value_val_add(builder, etdump_ValueType_Bool);
+  } else {
+    ET_CHECK_MSG(0, "Unsupported output type for intermediate logging\n");
+  }
+
+  auto value_ref = etdump_Value_end(builder);
+  etdump_DebugEvent_debug_entry_add(builder, value_ref);
+
+  etdump_DebugEvent_ref_t debug_event = etdump_DebugEvent_end(builder);
+
+  etdump_RunData_events_push_start(builder);
+  etdump_Event_debug_event_add(builder, debug_event);
+  etdump_RunData_events_push_end(builder);
+}
+
 void ETDumpGen::end_profiling(EventTracerEntry prof_entry) {
   et_timestamp_t end_time = et_pal_current_ticks();
   ET_CHECK_MSG(