Implement Graph node (#6037)

Summary:
Pull Request resolved: #6037

Introduce Graph node which will be the parent class for nodes. This allows us to have different nodes that correspond to different vk command buffer functions including shader dispatch.

Reviewed By: SS-JIA

Differential Revision: D64080291

fbshipit-source-id: b49540b99c8901f18645c86d3f5a9d274e5191b5

Author: Abhi-hpp

backends/vulkan/runtime/graph/Logging.cpp

@@ -156,7 +156,7 @@ void ComputeGraph::print_readable() {
   size_t node_idx = 0;
   for (const std::unique_ptr<ExecuteNode>& node : execute_nodes()) {
     std::cout << std::setw(6) << node_idx;
-    std::cout << std::setw(32) << node->shader_.kernel_name;
+    std::cout << std::setw(32) << node->name();
 
     std::stringstream read_s;
     for (const ArgGroup& arg_group : node->args_) {

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

@@ -0,0 +1,66 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/backends/vulkan/runtime/graph/ops/DispatchNode.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ComputeGraph.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ops/utils/BindingUtils.h>
+
+namespace vkcompute {
+
+DispatchNode::DispatchNode(
+    ComputeGraph& graph,
+    const vkapi::ShaderInfo& shader,
+    const utils::uvec3& global_workgroup_size,
+    const utils::uvec3& local_workgroup_size,
+    const std::vector<ArgGroup>& args,
+    const vkapi::ParamsBindList& params,
+    const vkapi::SpecVarList& spec_vars,
+    const ResizeFunction& resize_fn,
+    const std::vector<ValueRef>& resize_args)
+    : ExecuteNode(resize_fn, resize_args, args, shader.kernel_name),
+      shader_(shader),
+      global_workgroup_size_(global_workgroup_size),
+      local_workgroup_size_(local_workgroup_size),
+      params_(params),
+      spec_vars_(spec_vars) {
+  graph.update_descriptor_counts(shader, /*execute = */ true);
+}
+
+void DispatchNode::encode(ComputeGraph* graph) {
+  if (!shader_) {
+    return;
+  }
+  api::Context* const context = graph->context();
+  vkapi::PipelineBarrier pipeline_barrier{};
+
+  std::unique_lock<std::mutex> cmd_lock = context->dispatch_lock();
+
+  context->report_shader_dispatch_start(
+      shader_.kernel_name,
+      global_workgroup_size_,
+      local_workgroup_size_,
+      node_id_);
+
+  vkapi::DescriptorSet descriptor_set =
+      context->get_descriptor_set(shader_, local_workgroup_size_, spec_vars_);
+
+  uint32_t idx = 0;
+  idx = bind_values_to_descriptor_set(
+      graph, args_, pipeline_barrier, descriptor_set, idx);
+
+  bind_params_to_descriptor_set(params_, descriptor_set, idx);
+
+  context->register_shader_dispatch(
+      descriptor_set, pipeline_barrier, shader_, global_workgroup_size_);
+
+  context->report_shader_dispatch_end();
+}
+
+} // namespace vkcompute

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

@@ -0,0 +1,56 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <executorch/backends/vulkan/runtime/api/api.h>
+
+#include <executorch/backends/vulkan/runtime/graph/containers/Value.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ops/ExecuteNode.h>
+
+namespace vkcompute {
+
+class ComputeGraph;
+
+/*
+ * Represents a single shader execution op in a ML model.
+ */
+class DispatchNode final : public ExecuteNode {
+  friend class ComputeGraph;
+
+ public:
+  explicit DispatchNode(
+      ComputeGraph& graph,
+      const vkapi::ShaderInfo& shader,
+      const utils::uvec3& global_workgroup_size,
+      const utils::uvec3& local_workgroup_size,
+      const std::vector<ArgGroup>& args,
+      const vkapi::ParamsBindList& params,
+      const vkapi::SpecVarList& spec_vars = {},
+      const ResizeFunction& resize_fn = nullptr,
+      const std::vector<ValueRef>& resize_args = {});
+
+  ~DispatchNode() override = default;
+
+  void encode(ComputeGraph* graph) override;
+
+ protected:
+  const vkapi::ShaderInfo shader_;
+  const utils::uvec3 global_workgroup_size_;
+  const utils::uvec3 local_workgroup_size_;
+  const vkapi::ParamsBindList params_;
+  const vkapi::SpecVarList spec_vars_;
+
+ public:
+  operator bool() const {
+    return shader_;
+  }
+};
+
+} // namespace vkcompute

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

@@ -8,73 +8,14 @@
 
 #include <executorch/backends/vulkan/runtime/graph/ops/ExecuteNode.h>
 
-#include <executorch/backends/vulkan/runtime/graph/ComputeGraph.h>
-
-#include <executorch/backends/vulkan/runtime/graph/ops/utils/BindingUtils.h>
-
 namespace vkcompute {
-
 ExecuteNode::ExecuteNode(
-    ComputeGraph& graph,
-    const vkapi::ShaderInfo& shader,
-    const utils::uvec3& global_workgroup_size,
-    const utils::uvec3& local_workgroup_size,
-    const std::vector<ArgGroup>& args,
-    const vkapi::ParamsBindList& params,
-    const vkapi::SpecVarList& spec_vars,
     const ResizeFunction& resize_fn,
-    const std::vector<ValueRef>& resize_args)
-    : shader_(shader),
-      global_workgroup_size_(global_workgroup_size),
-      local_workgroup_size_(local_workgroup_size),
+    const std::vector<ValueRef>& resize_args,
+    const std::vector<ArgGroup>& args,
+    const std::string& name)
+    : resize_fn_(resize_fn),
+      resize_args_(resize_args),
       args_(args),
-      params_(params),
-      spec_vars_(spec_vars),
-      resize_fn_(resize_fn),
-      resize_args_(resize_args) {
-  graph.update_descriptor_counts(shader, /*execute = */ true);
-}
-
-ExecuteNode::ExecuteNode(
-    const ResizeFunction& resize_fn,
-    const std::vector<ValueRef>& resize_args)
-    : shader_(),
-      global_workgroup_size_({0u, 0u, 0u}),
-      local_workgroup_size_({0u, 0u, 0u}),
-      args_(),
-      params_(),
-      spec_vars_(),
-      resize_fn_(resize_fn),
-      resize_args_(resize_args) {}
-
-void ExecuteNode::encode(ComputeGraph* graph) {
-  if (!shader_) {
-    return;
-  }
-  api::Context* const context = graph->context();
-  vkapi::PipelineBarrier pipeline_barrier{};
-
-  std::unique_lock<std::mutex> cmd_lock = context->dispatch_lock();
-
-  context->report_shader_dispatch_start(
-      shader_.kernel_name,
-      global_workgroup_size_,
-      local_workgroup_size_,
-      node_id_);
-
-  vkapi::DescriptorSet descriptor_set =
-      context->get_descriptor_set(shader_, local_workgroup_size_, spec_vars_);
-
-  uint32_t idx = 0;
-  idx = bind_values_to_descriptor_set(
-      graph, args_, pipeline_barrier, descriptor_set, idx);
-
-  bind_params_to_descriptor_set(params_, descriptor_set, idx);
-
-  context->register_shader_dispatch(
-      descriptor_set, pipeline_barrier, shader_, global_workgroup_size_);
-
-  context->report_shader_dispatch_end();
-}
-
+      name_(name) {}
 } // namespace vkcompute

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

@@ -39,7 +39,7 @@ struct ArgGroup {
  * encoding of the shader corresponding to the op into the command buffer of a
  * ComputeGraph.
  */
-class ExecuteNode final {
+class ExecuteNode {
   friend class ComputeGraph;
 
  public:
@@ -48,29 +48,22 @@ class ExecuteNode final {
       const std::vector<ArgGroup>&,
       const std::vector<ValueRef>&)>;
 
-  explicit ExecuteNode(
-      ComputeGraph& graph,
-      const vkapi::ShaderInfo& shader,
-      const utils::uvec3& global_workgroup_size,
-      const utils::uvec3& local_workgroup_size,
-      const std::vector<ArgGroup>& args,
-      const vkapi::ParamsBindList& params,
-      const vkapi::SpecVarList& spec_vars = {},
-      const ResizeFunction& resize_fn = nullptr,
-      const std::vector<ValueRef>& resize_args = {});
-
   /*
-   * This overload of the ExecuteNode constructor is used to register ops which
+   * This overload of the DispatchNode constructor is used to register ops which
    * update a tensor view. No shader is dispatched, but the node still needs to
    * update the view's sizes and strides after a resize.
    */
   explicit ExecuteNode(
       const ResizeFunction& resize_fn = nullptr,
-      const std::vector<ValueRef>& resize_args = {});
+      const std::vector<ValueRef>& resize_args = {},
+      const std::vector<ArgGroup>& args = {},
+      const std::string& name = "Graph Node");
 
-  ~ExecuteNode() = default;
+  virtual ~ExecuteNode() = default;
 
-  void encode(ComputeGraph* graph);
+  virtual void encode(ComputeGraph* graph) {
+    (void)graph;
+  }
 
   inline void trigger_resize(ComputeGraph* graph) {
     if (resize_fn_ != nullptr) {
@@ -82,21 +75,16 @@ class ExecuteNode final {
     node_id_ = node_id;
   }
 
+  inline const std::string& name() const {
+    return name_;
+  }
+
  protected:
   uint32_t node_id_;
-  const vkapi::ShaderInfo shader_;
-  const utils::uvec3 global_workgroup_size_;
-  const utils::uvec3 local_workgroup_size_;
-  const std::vector<ArgGroup> args_;
-  const vkapi::ParamsBindList params_;
-  const vkapi::SpecVarList spec_vars_;
   const ResizeFunction resize_fn_;
   const std::vector<ValueRef> resize_args_;
-
- public:
-  operator bool() const {
-    return shader_;
-  }
+  const std::vector<ArgGroup> args_;
+  const std::string name_;
 };
 
 } // namespace vkcompute

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

@@ -9,6 +9,7 @@
 #pragma once
 
 #include <executorch/backends/vulkan/runtime/graph/ComputeGraph.h>
+#include <executorch/backends/vulkan/runtime/graph/ops/DispatchNode.h>
 
 #include <functional>
 #include <unordered_map>

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

@@ -6,7 +6,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#include <executorch/backends/vulkan/runtime/graph/ops/ExecuteNode.h>
+#include <executorch/backends/vulkan/runtime/graph/ops/DispatchNode.h>
 
 #include <executorch/backends/vulkan/runtime/graph/ComputeGraph.h>
 

backends/vulkan/runtime/graph/ops/impl/Arange.cpp

@@ -88,7 +88,7 @@ void add_arange_node(
   kernel_name.reserve(kShaderNameReserve);
   add_dtype_suffix(kernel_name, *t_out);
 
-  graph.execute_nodes().emplace_back(new ExecuteNode(
+  graph.execute_nodes().emplace_back(new DispatchNode(
       graph,
       VK_KERNEL_FROM_STR(kernel_name),
       graph.create_global_wg_size(out),

backends/vulkan/runtime/graph/ops/impl/BatchNorm.cpp

@@ -80,7 +80,7 @@ void add_native_batch_norm_node(
   int32_t num_texel_per_batch =
       utils::div_up_4((dim_at<kChannel4D>(t_in->sizes())));
 
-  graph.execute_nodes().emplace_back(new ExecuteNode(
+  graph.execute_nodes().emplace_back(new DispatchNode(
       graph,
       VK_KERNEL_FROM_STR(kernel_name),
       graph.create_global_wg_size(out_ref),

backends/vulkan/runtime/graph/ops/impl/BinaryOp.cpp

@@ -75,7 +75,7 @@ void add_binary_op_node(
   kernel_name += op_name;
   add_dtype_suffix(kernel_name, *t_out);
 
-  graph.execute_nodes().emplace_back(new ExecuteNode(
+  graph.execute_nodes().emplace_back(new DispatchNode(
       graph,
       VK_KERNEL_FROM_STR(kernel_name),
       graph.create_global_wg_size(out),