Implement Graph node (#6037)

Summary:

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.

Differential Revision: D64080291

Author: Abhi-hpp

backends/vulkan/runtime/graph/ComputeGraph.cpp

@@ -573,7 +573,7 @@ void ComputeGraph::encode_execute() {
     shared_object.bind_users(this);
   }
 
-  for (std::unique_ptr<ExecuteNode>& node : execute_nodes_) {
+  for (std::unique_ptr<GraphNode>& node : execute_nodes_) {
     node->encode(this);
   }
 }
@@ -592,7 +592,7 @@ void ComputeGraph::resize_input(
 }
 
 void ComputeGraph::propagate_resize() {
-  for (std::unique_ptr<ExecuteNode>& node : execute_nodes_) {
+  for (std::unique_ptr<GraphNode>& node : execute_nodes_) {
     node->trigger_resize(this);
   }
 }

backends/vulkan/runtime/graph/ComputeGraph.h

@@ -20,7 +20,7 @@
 #include <executorch/backends/vulkan/runtime/graph/containers/SharedObject.h>
 #include <executorch/backends/vulkan/runtime/graph/containers/Value.h>
 
-#include <executorch/backends/vulkan/runtime/graph/ops/ExecuteNode.h>
+#include <executorch/backends/vulkan/runtime/graph/ops/GraphNode.h>
 #include <executorch/backends/vulkan/runtime/graph/ops/PrepackNode.h>
 
 namespace vkcompute {
@@ -178,7 +178,7 @@ class ComputeGraph final {
   std::vector<api::ParamsBuffer> param_ubos_;
 
   std::vector<std::unique_ptr<PrepackNode>> prepack_nodes_;
-  std::vector<std::unique_ptr<ExecuteNode>> execute_nodes_;
+  std::vector<std::unique_ptr<GraphNode>> execute_nodes_;
 
   std::vector<IOValueRef> inputs_;
   std::vector<IOValueRef> outputs_;
@@ -207,7 +207,7 @@ class ComputeGraph final {
     return prepack_nodes_;
   }
 
-  inline std::vector<std::unique_ptr<ExecuteNode>>& execute_nodes() {
+  inline std::vector<std::unique_ptr<GraphNode>>& execute_nodes() {
     return execute_nodes_;
   }
 

backends/vulkan/runtime/graph/Logging.cpp

@@ -154,9 +154,9 @@ void ComputeGraph::print_readable() {
             << std::endl;
 
   size_t node_idx = 0;
-  for (const std::unique_ptr<ExecuteNode>& node : execute_nodes()) {
+  for (const std::unique_ptr<GraphNode>& 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/ExecuteNode.cpp

@@ -24,29 +24,15 @@ ExecuteNode::ExecuteNode(
     const vkapi::SpecVarList& spec_vars,
     const ResizeFunction& resize_fn,
     const std::vector<ValueRef>& resize_args)
-    : shader_(shader),
+    : GraphNode(resize_fn, resize_args, args, shader.kernel_name),
+      shader_(shader),
       global_workgroup_size_(global_workgroup_size),
       local_workgroup_size_(local_workgroup_size),
-      args_(args),
       params_(params),
-      spec_vars_(spec_vars),
-      resize_fn_(resize_fn),
-      resize_args_(resize_args) {
+      spec_vars_(spec_vars) {
   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;

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

@@ -12,42 +12,19 @@
 
 #include <executorch/backends/vulkan/runtime/graph/containers/Value.h>
 
+#include <executorch/backends/vulkan/runtime/graph/ops/GraphNode.h>
+
 namespace vkcompute {
 
 class ComputeGraph;
 
 /*
- * Represents a group of shader arguments (images and/or buffers), with a common
- * access permission.
- */
-struct ArgGroup {
-  ArgGroup(const ValueRef ref, const vkapi::MemoryAccessFlags access)
-      : refs{ref}, access(access) {}
-
-  ArgGroup(
-      const std::vector<ValueRef>& refs,
-      const vkapi::MemoryAccessFlags access)
-      : refs(refs), access(access) {}
-
-  const std::vector<ValueRef> refs;
-  const vkapi::MemoryAccessFlags access;
-};
-
-/*
- * Represents a single execution op in a ML model. In graph mode, ops will be
- * implemented in a derived class that implements encode, which will implement
- * encoding of the shader corresponding to the op into the command buffer of a
- * ComputeGraph.
+ * Represents a single shader execution op in a ML model.
  */
-class ExecuteNode final {
+class ExecuteNode final : public GraphNode {
   friend class ComputeGraph;
 
  public:
-  using ResizeFunction = const std::function<void(
-      ComputeGraph*,
-      const std::vector<ArgGroup>&,
-      const std::vector<ValueRef>&)>;
-
   explicit ExecuteNode(
       ComputeGraph& graph,
       const vkapi::ShaderInfo& shader,
@@ -59,39 +36,16 @@ class ExecuteNode final {
       const ResizeFunction& resize_fn = nullptr,
       const std::vector<ValueRef>& resize_args = {});
 
-  /*
-   * This overload of the ExecuteNode 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 = {});
-
   ~ExecuteNode() = default;
 
-  void encode(ComputeGraph* graph);
-
-  inline void trigger_resize(ComputeGraph* graph) {
-    if (resize_fn_ != nullptr) {
-      resize_fn_(graph, args_, resize_args_);
-    }
-  }
-
-  inline void set_node_id(uint32_t node_id) {
-    node_id_ = node_id;
-  }
+  void encode(ComputeGraph* graph) override;
 
  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 {

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

@@ -0,0 +1,23 @@
+/*
+ * 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/GraphNode.h>
+
+namespace vkcompute {
+
+GraphNode::GraphNode(
+    const ResizeFunction& resize_fn,
+    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),
+      name_(name) {}
+
+} // namespace vkcompute

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

@@ -0,0 +1,88 @@
+/*
+ * 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>
+
+namespace vkcompute {
+
+class ComputeGraph;
+
+/*
+ * Represents a group of shader arguments (images and/or buffers), with a common
+ * access permission.
+ */
+struct ArgGroup {
+  ArgGroup(const ValueRef ref, const vkapi::MemoryAccessFlags access)
+      : refs{ref}, access(access) {}
+
+  ArgGroup(
+      const std::vector<ValueRef>& refs,
+      const vkapi::MemoryAccessFlags access)
+      : refs(refs), access(access) {}
+
+  const std::vector<ValueRef> refs;
+  const vkapi::MemoryAccessFlags access;
+};
+
+/*
+ * Represents a single execution op in a ML model. In graph mode, ops will be
+ * implemented in a derived class that implements encode, which will implement
+ * encoding of the shader corresponding to the op into the command buffer of a
+ * ComputeGraph.
+ */
+class GraphNode {
+  friend class ComputeGraph;
+
+ public:
+  using ResizeFunction = const std::function<void(
+      ComputeGraph*,
+      const std::vector<ArgGroup>&,
+      const std::vector<ValueRef>&)>;
+
+  /*
+   * This overload of the GraphNode 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 GraphNode(
+      const ResizeFunction& resize_fn = nullptr,
+      const std::vector<ValueRef>& resize_args = {},
+      const std::vector<ArgGroup>& args = {},
+      const std::string& name = "Graph Node");
+
+  virtual ~GraphNode() = default;
+
+  virtual void encode(ComputeGraph* graph) {}
+
+  inline void trigger_resize(ComputeGraph* graph) {
+    if (resize_fn_ != nullptr) {
+      resize_fn_(graph, args_, resize_args_);
+    }
+  }
+
+  inline void set_node_id(uint32_t node_id) {
+    node_id_ = node_id;
+  }
+
+  inline const std::string& name() const {
+    return name_;
+  }
+
+ protected:
+  uint32_t node_id_;
+  const ResizeFunction resize_fn_;
+  const std::vector<ValueRef> resize_args_;
+  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/ExecuteNode.h>
 
 #include <functional>
 #include <unordered_map>

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

@@ -159,7 +159,7 @@ void add_cache_slice_view_node(
 
   graph.get_tensor(cache_sliced)->virtual_resize(slice_sizes);
 
-  graph.execute_nodes().emplace_back(new ExecuteNode(
+  graph.execute_nodes().emplace_back(new GraphNode(
       resize_cache_slice_view_node,
       {cache, input_pos_symint, q_projected, cache_sliced}));
 }
@@ -328,7 +328,7 @@ void sdpa_with_kv_cache_impl(
       mat2_is_transposed);
 
   graph.execute_nodes().emplace_back(
-      new ExecuteNode(resize_sdpa_out, {q_projected, out}));
+      new GraphNode(resize_sdpa_out, {q_projected, out}));
 }
 
 REGISTER_OPERATORS {

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

@@ -264,7 +264,7 @@ void add_slice_view_node(
 
   graph.get_tensor(out_ref)->virtual_resize(new_out_sizes);
 
-  graph.execute_nodes().emplace_back(new ExecuteNode(
+  graph.execute_nodes().emplace_back(new GraphNode(
       resize_slice_view_node,
       {out_ref, in_ref, dim_ref, opt_start_ref, opt_end_ref, opt_step_ref}));
 }

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

@@ -5,7 +5,7 @@
  * 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/OperatorRegistry.h>
 #include <executorch/backends/vulkan/runtime/graph/ops/impl/Staging.h>
 
 #include <executorch/backends/vulkan/runtime/graph/ops/utils/StagingUtils.h>

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

@@ -66,7 +66,7 @@ void add_transpose_view_node(
   graph.get_tensor(out_ref)->virtual_clone(*in);
   graph.get_tensor(out_ref)->virtual_transpose(dim0, dim1);
 
-  graph.execute_nodes().emplace_back(new ExecuteNode(
+  graph.execute_nodes().emplace_back(new GraphNode(
       resize_transpose_view_node, {out_ref, input_ref, dim0_ref, dim1_ref}));
 }
 

backends/vulkan/test/vulkan_compute_api_test.cpp

@@ -25,6 +25,8 @@
 
 #include <executorch/backends/vulkan/test/utils/test_utils.h>
 
+#include <executorch/backends/vulkan/runtime/graph/ops/ExecuteNode.h>
+
 using namespace vkcompute;
 using namespace vkcompute::api;
 
@@ -1148,16 +1150,15 @@ TEST(VulkanComputeGraphTest, test_values_string) {
   EXPECT_TRUE(stored == "hello, world");
 }
 
-TEST(VulkanComputeGraphTest, empty_init_executenode_test) {
-  ExecuteNode node(nullptr, {});
-  EXPECT_FALSE(node);
+TEST(VulkanComputeGraphTest, empty_init_graphnode_test) {
+  GraphNode node(nullptr, {});
 
   GraphConfig config;
   ComputeGraph graph(config);
 
-  // Encode an empty ExecuteNode and check that command buffer encoding does not
+  // Encode an empty GraphNode and check that command buffer encoding does not
   // crash.
-  graph.execute_nodes().emplace_back(new ExecuteNode(nullptr, {}));
+  graph.execute_nodes().emplace_back(new GraphNode(nullptr, {}));
   EXPECT_NO_FATAL_FAILURE(graph.encode_execute());
 }