[ET-VK] Pass detailed op information to event tracer (pytorch#16266)

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at
bottom):
* pytorch#16267
* __->__ pytorch#16266

## Context

Currently, when using the event tracer API to log shader execution
times, only the shader name is recorded as the event name. However, this
provides very minimal context to use when interpreting profiling data.
For example, if we see that a convolution shader is running slow, it's
is impossible to know from the profiling data alone what the
input/output sizes were, what convolution parameters (i.e.
stride/padding/dilation) were used, etc.

## Changes

This diff makes it so that for each shader dispatch, a JSON is recorded
as the event name instead which contains the complete details of all the
arguments the operator was invoked with.

The JSON will look something like

```json
    {
      "name": "aten.where.self",
      "args": [
        {
          "type": "TENSOR",
          "value_ref": 25,
          "dtype": "Bool",
          "sizes": [
            1,
            1,
            1,
            8
          ],
          "storage": "TEXTURE_3D",
          "packed_dim": 2
        },
        {
          "type": "TENSOR",
          "value_ref": 30,
          "dtype": "Float",
          "sizes": [
            1,
            6,
            43,
            8
          ],
          "storage": "TEXTURE_3D",
          "packed_dim": 2
        },
        {
          "type": "TENSOR",
          "value_ref": 32,
          "dtype": "Float",
          "sizes": [
            1,
            6,
            43,
            8
          ],
          "storage": "TEXTURE_3D",
          "packed_dim": 2
        },
        {
          "type": "TENSOR",
          "value_ref": 33,
          "dtype": "Float",
          "sizes": [
            1,
            6,
            43,
            8
          ],
          "storage": "TEXTURE_3D",
          "packed_dim": 2
        }
      ]
    }
```

Then, when processing the profiling data, the JSON can be post-processed
to provide useful information to contextualize the shader execution
times, for example:

* Memory Throughput
* GFLOPS

Differential Revision:
[D84646748](https://our.internmc.facebook.com/intern/diff/D84646748/)

Co-authored-by: ssjia <ssjia@devvm1479.ncg0.facebook.com>

Author: SS-JIA

backends/vulkan/runtime/VulkanBackend.cpp

@@ -19,6 +19,7 @@
 #include <executorch/runtime/core/error.h>
 #include <executorch/runtime/core/evalue.h>
 #ifdef ET_EVENT_TRACER_ENABLED
+#include <executorch/backends/vulkan/runtime/graph/Logging.h>
 #include <executorch/runtime/core/event_tracer_hooks_delegate.h>
 #endif // ET_EVENT_TRACER_ENABLED
 #include <executorch/runtime/core/exec_aten/util/tensor_util.h>
@@ -422,6 +423,13 @@ class GraphBuilder {
         args.push_back(get_fb_id_valueref(static_cast<int>(arg_fb_id)));
       }
 
+#ifdef ET_EVENT_TRACER_ENABLED
+      std::string operator_json =
+          make_operator_json(compute_graph_, op_name, args);
+      set_and_get_current_operator_json(operator_json);
+      get_current_operator_count(true);
+#endif // ET_EVENT_TRACER_ENABLED
+
       auto vkFn = VK_GET_OP_FN(op_name);
       vkFn(*compute_graph_, args);
     }
@@ -431,6 +439,9 @@ class GraphBuilder {
     for (const uint32_t fb_id : *flatbuffer_->output_ids()) {
       const ValueRef ref = get_fb_id_valueref(fb_id);
       if (compute_graph_->val_is_tensor(ref)) {
+#ifdef ET_EVENT_TRACER_ENABLED
+        get_current_operator_count(true);
+#endif // ET_EVENT_TRACER_ENABLED
         compute_graph_->set_output_tensor(
             ref, get_staging_scalar_type_of(fb_id));
       } else {
@@ -694,16 +705,14 @@ class VulkanBackend final : public ::executorch::runtime::BackendInterface {
     compute_graph->context()->querypool().extract_results();
     for (const auto& r :
          compute_graph->context()->querypool().get_shader_timestamp_data()) {
-      std::string event_name =
-          r.kernel_name + "_" + std::to_string(r.dispatch_id);
+      std::string event_name = "{" + r.kernel_name +
+          ", \"dispatch_id\": " + std::to_string(r.dispatch_id) + "}";
       event_tracer_log_profiling_delegate(
           event_tracer,
           event_name.c_str(),
           /* delegate_debug_id = */ -1,
           r.start_time_ns,
-          r.end_time_ns,
-          (void*)(&r.metadata),
-          sizeof(r.metadata));
+          r.end_time_ns);
     }
 #endif // ET_EVENT_TRACER_ENABLED
 

backends/vulkan/runtime/graph/ComputeGraph.cpp

@@ -15,6 +15,24 @@
 
 #include <executorch/backends/vulkan/runtime/graph/ops/utils/StagingUtils.h>
 
+#ifdef ET_EVENT_TRACER_ENABLED
+std::string& set_and_get_current_operator_json(const std::string& json) {
+  static std::string current_operator_json;
+  if (json.size() > 0) {
+    current_operator_json = json;
+  }
+  return current_operator_json;
+}
+
+size_t get_current_operator_count(const bool increment) {
+  static int count = 0;
+  if (increment) {
+    count++;
+  }
+  return count;
+}
+#endif /* ET_EVENT_TRACER_ENABLED */
+
 namespace vkcompute {
 
 //

backends/vulkan/runtime/graph/ComputeGraph.h

@@ -25,6 +25,11 @@
 #include <executorch/backends/vulkan/runtime/graph/ops/ExecuteNode.h>
 #include <executorch/backends/vulkan/runtime/graph/ops/PrepackNode.h>
 
+#ifdef ET_EVENT_TRACER_ENABLED
+std::string& set_and_get_current_operator_json(const std::string& json);
+size_t get_current_operator_count(const bool increment = false);
+#endif
+
 namespace vkcompute {
 
 // Define valid scalar types that the Value class can

backends/vulkan/runtime/graph/Logging.cpp

@@ -17,6 +17,82 @@
 
 namespace vkcompute {
 
+std::ostream& operator<<(std::ostream& os, const std::vector<int64_t>& sizes) {
+  if (sizes.size() == 0) {
+    os << "[]";
+    return os;
+  }
+  os << "[";
+  for (int i = 0; i < sizes.size() - 1; ++i) {
+    os << sizes.at(i) << ", ";
+  }
+  os << sizes.at(sizes.size() - 1);
+  os << "]";
+  return os;
+}
+
+std::string make_arg_json(ComputeGraph* const compute_graph, ValueRef arg) {
+  std::stringstream ss;
+  ss << "{\"type\": \"" << compute_graph->get_val_type(arg) << "\", ";
+  ss << "\"value_ref\": " << arg;
+  if (compute_graph->val_is_tensor(arg)) {
+    ss << ", \"dtype\": \"";
+    ss << compute_graph->dtype_of(arg) << "\"";
+    ss << ", \"sizes\": ";
+    ss << compute_graph->sizes_of(arg);
+    ss << ", \"storage\": \"";
+    ss << compute_graph->storage_type_of(arg) << "\"";
+    ss << ", \"packed_dim\": ";
+    ss << compute_graph->packed_dim_of(arg);
+  } else if (compute_graph->val_is_tref(arg)) {
+    ss << ", \"sizes\": ";
+    ss << compute_graph->sizes_of(arg);
+    ss << ", \"dtype\": \"";
+    ss << compute_graph->dtype_of(arg) << "\"";
+  } else if (compute_graph->val_is_value_list(arg)) {
+    ValueListPtr val_list = compute_graph->get_value_list(arg);
+    ss << ", \"values\": [";
+    for (const ValueRef& value : *val_list) {
+      ss << value << ", ";
+    }
+    ss << "]";
+  } else if (compute_graph->val_is_int_list(arg)) {
+    ss << ", \"values\": ";
+    ss << *compute_graph->get_int_list(arg);
+  } else if (compute_graph->val_is_int(arg)) {
+    ss << ", \"value\": ";
+    ss << compute_graph->get_int(arg);
+  } else if (compute_graph->val_is_double(arg)) {
+    ss << ", \"value\": ";
+    ss << compute_graph->get_double(arg);
+  } else if (compute_graph->val_is_bool(arg)) {
+    ss << ", \"value\": ";
+    ss << compute_graph->get_bool(arg);
+  } else if (compute_graph->val_is_symint(arg)) {
+    ss << ", \"value\": ";
+    ss << compute_graph->read_symint(arg);
+  }
+  ss << "}";
+
+  return ss.str();
+}
+
+std::string make_operator_json(
+    ComputeGraph* const compute_graph,
+    std::string& op_name,
+    std::vector<ValueRef>& args) {
+  std::stringstream ss;
+  ss << "\"name\": \"" << op_name << "\", \"args\": [";
+  for (size_t i = 0; i < args.size(); ++i) {
+    ss << make_arg_json(compute_graph, args[i]);
+    if (i + 1 < args.size()) {
+      ss << ", ";
+    }
+  }
+  ss << "]";
+  return ss.str();
+}
+
 void ComputeGraph::print_readable() {
   std::set<ValueRef> input_set;
   for (const IOValueRef& io_val : inputs()) {

backends/vulkan/runtime/graph/Logging.h

@@ -10,6 +10,8 @@
 
 #include <executorch/backends/vulkan/runtime/api/api.h>
 
+#include <executorch/backends/vulkan/runtime/graph/ComputeGraph.h>
+
 #include <optional>
 #include <ostream>
 #include <vector>
@@ -42,6 +44,8 @@ inline std::ostream& operator<<(std::ostream& os, const utils::ivec4& v) {
   return utils::operator<<(os, v);
 }
 
+std::ostream& operator<<(std::ostream& os, const std::vector<int64_t>& sizes);
+
 template <typename T>
 inline std::ostream& operator<<(std::ostream& os, const std::optional<T>& opt) {
   os << "[";
@@ -52,4 +56,11 @@ inline std::ostream& operator<<(std::ostream& os, const std::optional<T>& opt) {
   return os;
 }
 
+std::string make_arg_json(ComputeGraph* const compute_graph, ValueRef arg);
+
+std::string make_operator_json(
+    ComputeGraph* const compute_graph,
+    std::string& op_name,
+    std::vector<ValueRef>& args);
+
 } // namespace vkcompute

backends/vulkan/runtime/graph/ops/DispatchNode.cpp

@@ -60,8 +60,21 @@ void DispatchNode::encode(ComputeGraph* graph) {
 
   write_push_constant_data();
 
+#ifdef ET_EVENT_TRACER_ENABLED
+  std::string event_name;
+  if (!operator_json.empty()) {
+    event_name += "\"operator\": {" + operator_json + "}, ";
+  }
+  event_name += "\"kernel_name\": \"" + shader_.kernel_name + "\", ";
+  event_name += "\"operator_id\": " + std::to_string(operator_count);
+#endif
+
   context->report_shader_dispatch_start(
+#ifdef ET_EVENT_TRACER_ENABLED
+      event_name,
+#else
       shader_.kernel_name,
+#endif
       global_workgroup_size_,
       local_workgroup_size_,
       node_id_);

backends/vulkan/runtime/graph/ops/ExecuteNode.cpp

@@ -20,7 +20,12 @@ ExecuteNode::ExecuteNode(
       resize_args_(resize_args),
       args_(args),
       name_(name),
-      has_data_dependent_shape_(has_data_dependent_shape) {}
+      has_data_dependent_shape_(has_data_dependent_shape) {
+#ifdef ET_EVENT_TRACER_ENABLED
+  operator_json = set_and_get_current_operator_json("");
+  operator_count = get_current_operator_count();
+#endif
+}
 
 bool ExecuteNode::trigger_resize(ComputeGraph* graph) {
   bool any_arg_updated = was_any_arg_updated(graph);

backends/vulkan/runtime/graph/ops/ExecuteNode.h

@@ -89,6 +89,11 @@ class ExecuteNode {
   const std::vector<ArgGroup> args_;
   const std::string name_;
   bool has_data_dependent_shape_ = false;
+
+#ifdef ET_EVENT_TRACER_ENABLED
+  std::string operator_json;
+  size_t operator_count = 0;
+#endif
 };
 
 } // namespace vkcompute