Intermediate output logging enablement

Summary: Design doc: https://docs.google.com/document/d/1qGHsgd-roqtxPz4CrUlqGrKaAtbaf9bArziMuwHD0So/edit

Differential Revision: D56941407

Author: Olivia-liu

runtime/core/event_tracer.h

@@ -256,6 +256,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

sdk/etdump/etdump_flatcc.cpp

@@ -288,6 +288,126 @@ 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()) {
+    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,
+        "This EValue type = %s is not yet supported for logging\n",
+        typeid(T).name());
+  }
+
+  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(

sdk/etdump/etdump_flatcc.h

@@ -97,6 +97,45 @@ class ETDumpGen : public EventTracer {
       const EValue& evalue,
       LoggedEValueType evalue_type =
           LoggedEValueType::kIntermediateOutput) override;
+  /**
+   * Log an intermediate tensor output from a delegate.
+   */
+  virtual void log_intermediate_output_delegate(
+      const char* name,
+      DebugHandle delegate_debug_index,
+      const Tensor& output) override;
+
+  /**
+   * Log an intermediate tensor array output from a delegate.
+   */
+  virtual void log_intermediate_output_delegate(
+      const char* name,
+      DebugHandle delegate_debug_index,
+      const ArrayRef<Tensor> output) override;
+
+  /**
+   * Log an intermediate int output from a delegate.
+   */
+  virtual void log_intermediate_output_delegate(
+      const char* name,
+      DebugHandle delegate_debug_index,
+      const int& output) override;
+
+  /**
+   * Log an intermediate bool output from a delegate.
+   */
+  virtual void log_intermediate_output_delegate(
+      const char* name,
+      DebugHandle delegate_debug_index,
+      const bool& output) override;
+
+  /**
+   * Log an intermediate double output from a delegate.
+   */
+  virtual void log_intermediate_output_delegate(
+      const char* name,
+      DebugHandle delegate_debug_index,
+      const double& output) override;
   void set_debug_buffer(Span<uint8_t> buffer);
   etdump_result get_etdump_data();
   size_t get_num_blocks();
@@ -115,6 +154,16 @@ class ETDumpGen : public EventTracer {
   void check_ready_to_add_events();
   int64_t create_string_entry(const char* name);
   size_t copy_tensor_to_debug_buffer(exec_aten::Tensor tensor);
+
+  /**
+   * Templated helper function used to log various types of intermediate output.
+   * Supported types include tensor, tensor array, int, bool and double.
+   */
+  template <typename T>
+  void log_intermediate_output_delegate_helper(
+      const char* name,
+      DebugHandle delegate_debug_index,
+      const T& output);
 };
 
 } // namespace executor

sdk/etdump/etdump_schema_flatcc.fbs

@@ -70,6 +70,12 @@ table DebugEvent {
   // The order of these entries is expected to be the same as the parameters
   // that are passed into a operator.
   debug_entry:Value;
+
+  // Integer based delegate debug identifier.
+  delegate_debug_id_int:int = -1;
+
+  // String based delegate debug identifier.
+  delegate_debug_id_str:string;
 }
 
 // All the details pertaining to an allocation done in the runtime. The main

sdk/etdump/schema_flatcc.py

@@ -95,6 +95,8 @@ class Value:
 class DebugEvent:
     chain_index: int
     instruction_id: int
+    delegate_debug_id_int: Optional[int]
+    delegate_debug_id_str: Optional[str]
     debug_entry: Value
 
 

sdk/etdump/tests/serialize_test.py

@@ -85,6 +85,8 @@ def get_sample_etdump_flatcc() -> flatcc.ETDumpFlatCC:
                         debug_event=flatcc.DebugEvent(
                             chain_index=1,
                             instruction_id=0,
+                            delegate_debug_id_str="56",
+                            delegate_debug_id_int=-1,
                             debug_entry=flatcc.Value(
                                 val=flatcc.ValueType.TENSOR.value,
                                 tensor=flatcc.Tensor(