[ET-VK][EZ] Split Graph.h classes into multiple files

backends/vulkan/runtime/VulkanBackend.cpp

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

backends/vulkan/runtime/graph/Graph.cpp → backends/vulkan/runtime/graph/ComputeGraph.cpp

@@ -6,76 +6,14 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#include <executorch/backends/vulkan/runtime/graph/Graph.h>
+#include <executorch/backends/vulkan/runtime/graph/ComputeGraph.h>
 
 #include <executorch/backends/vulkan/runtime/graph/ops/Staging.h>
 
 namespace at {
 namespace native {
 namespace vulkan {
 
-//
-// SharedObject
-//
-
-void SharedObject::add_user(ComputeGraph* const graph, const ValueRef idx) {
-  vTensor& t = graph->get_val(idx).toTensor();
-
-  //
-  // Aggregate Memory Requirements
-  //
-
-  const VkMemoryRequirements mem_reqs = t.get_memory_requirements();
-  aggregate_memory_requirements.size =
-      std::max(mem_reqs.size, aggregate_memory_requirements.size);
-  aggregate_memory_requirements.alignment =
-      std::max(mem_reqs.alignment, aggregate_memory_requirements.alignment);
-  aggregate_memory_requirements.memoryTypeBits |= mem_reqs.memoryTypeBits;
-
-  //
-  // Aggregate Allocation Create Info
-  //
-
-  const VmaAllocationCreateInfo create_info = t.get_allocation_create_info();
-  // Clear out CREATE_STRATEGY bit flags in case of conflict
-  VmaAllocationCreateFlags clear_mask = ~VMA_ALLOCATION_CREATE_STRATEGY_MASK;
-  VmaAllocationCreateFlags create_flags = create_info.flags & clear_mask;
-  // Use the default allocation strategy
-  aggregate_create_info.flags = create_flags | api::DEFAULT_ALLOCATION_STRATEGY;
-
-  // Set the usage flag if it is currently not set
-  if (aggregate_create_info.usage == VMA_MEMORY_USAGE_UNKNOWN) {
-    aggregate_create_info.usage = create_info.usage;
-  }
-  // Otherwise check that there is no conflict regarding usage
-  VK_CHECK_COND(aggregate_create_info.usage == create_info.usage);
-  aggregate_create_info.requiredFlags |= create_info.requiredFlags;
-  aggregate_create_info.preferredFlags |= create_info.preferredFlags;
-
-  users.emplace_back(idx);
-}
-
-void SharedObject::allocate(ComputeGraph* const graph) {
-  if (aggregate_memory_requirements.size == 0) {
-    return;
-  }
-  allocation = graph->context()->adapter_ptr()->vma().create_allocation(
-      aggregate_memory_requirements, aggregate_create_info);
-}
-
-void SharedObject::bind_users(ComputeGraph* const graph) {
-  if (users.empty()) {
-    return;
-  }
-  for (const ValueRef idx : users) {
-    graph->get_val(idx).toTensor().bind_allocation(allocation);
-  }
-}
-
-//
-// ComputeGraph
-//
-
 ComputeGraph::ComputeGraph(GraphConfig config)
     : config_{config},
       context_{new api::Context(

backends/vulkan/runtime/graph/Graph.h → backends/vulkan/runtime/graph/ComputeGraph.h

@@ -17,85 +17,15 @@
 #include <ATen/native/vulkan/api/Types.h>
 
 #include <executorch/backends/vulkan/runtime/graph/Config.h>
+#include <executorch/backends/vulkan/runtime/graph/ExecuteNode.h>
+#include <executorch/backends/vulkan/runtime/graph/PrepackNode.h>
+#include <executorch/backends/vulkan/runtime/graph/SharedObject.h>
 #include <executorch/backends/vulkan/runtime/graph/Value.h>
 
 namespace at {
 namespace native {
 namespace vulkan {
 
-using ValueRef = int32_t;
-
-struct IOValueRef {
-  ValueRef value;
-  ValueRef staging;
-};
-
-class ComputeGraph;
-
-/*
- * Represents a single prepacking op in a ML model. In graph mode, ops will be
- * implemented in a derived class that implements encode, which will implement
- * encoding of shaders transferring necessary data (such as weights and biases)
- * to the GPU.
- */
-class PrepackNode {
-  friend class ComputeGraph;
-
- public:
-  PrepackNode(ValueRef tref, ValueRef packed) : tref_{tref}, packed_{packed} {}
-
-  virtual ~PrepackNode() = default;
-
- protected:
-  ValueRef tref_;
-  ValueRef packed_;
-
- public:
-  virtual void encode(ComputeGraph* graph) const = 0;
-};
-
-/*
- * 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 ExecuteNode {
-  friend class ComputeGraph;
-
- public:
-  ExecuteNode(ValueRef input, ValueRef output)
-      : inputs_{input}, outputs_{output} {}
-  ExecuteNode(
-      const std::vector<ValueRef>& inputs,
-      const std::vector<ValueRef>& outputs)
-      : inputs_(inputs), outputs_(outputs) {}
-
-  virtual ~ExecuteNode() = default;
-
- protected:
-  std::vector<ValueRef> inputs_;
-  std::vector<ValueRef> outputs_;
-
- public:
-  virtual void encode(ComputeGraph* graph) const = 0;
-};
-
-struct SharedObject {
-  friend class ComputeGraph;
-
-  explicit SharedObject() = default;
-
-  VkMemoryRequirements aggregate_memory_requirements;
-  VmaAllocationCreateInfo aggregate_create_info;
-  std::vector<ValueRef> users;
-  api::MemoryAllocation allocation;
-
-  void add_user(ComputeGraph* const graph, const ValueRef idx);
-  void allocate(ComputeGraph* const graph);
-  void bind_users(ComputeGraph* const graph);
-};
-
 /*
  * This is the core data structure used to execute Vulkan models in graph mode.
  * As opposed to ATen/eager mode where a command buffer is encoded every

backends/vulkan/runtime/graph/ExecuteNode.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
+
+#ifdef USE_VULKAN_API
+
+#include <ATen/native/vulkan/api/Context.h>
+#include <ATen/native/vulkan/api/Tensor.h>
+#include <ATen/native/vulkan/api/Types.h>
+
+#include <executorch/backends/vulkan/runtime/graph/Value.h>
+
+namespace at {
+namespace native {
+namespace vulkan {
+
+class ComputeGraph;
+
+/*
+ * 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 ExecuteNode {
+  friend class ComputeGraph;
+
+ public:
+  ExecuteNode(ValueRef input, ValueRef output)
+      : inputs_{input}, outputs_{output} {}
+  ExecuteNode(
+      const std::vector<ValueRef>& inputs,
+      const std::vector<ValueRef>& outputs)
+      : inputs_(inputs), outputs_(outputs) {}
+
+  virtual ~ExecuteNode() = default;
+
+ protected:
+  std::vector<ValueRef> inputs_;
+  std::vector<ValueRef> outputs_;
+
+ public:
+  virtual void encode(ComputeGraph* graph) const = 0;
+};
+
+} // namespace vulkan
+} // namespace native
+} // namespace at
+
+#endif /* USE_VULKAN_API */

backends/vulkan/runtime/graph/Functions.h

@@ -10,7 +10,7 @@
 
 #ifdef USE_VULKAN_API
 
-#include <executorch/backends/vulkan/runtime/graph/Graph.h>
+#include <executorch/backends/vulkan/runtime/graph/ComputeGraph.h>
 
 namespace at {
 namespace native {

backends/vulkan/runtime/graph/OperatorRegistry.h

@@ -10,7 +10,7 @@
 
 #ifdef USE_VULKAN_API
 
-#include <executorch/backends/vulkan/runtime/graph/Graph.h>
+#include <executorch/backends/vulkan/runtime/graph/ComputeGraph.h>
 
 #include <functional>
 #include <unordered_map>

backends/vulkan/runtime/graph/PrepackNode.h

@@ -0,0 +1,51 @@
+/*
+ * 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
+
+#ifdef USE_VULKAN_API
+
+#include <ATen/native/vulkan/api/Context.h>
+#include <ATen/native/vulkan/api/Tensor.h>
+#include <ATen/native/vulkan/api/Types.h>
+
+#include <executorch/backends/vulkan/runtime/graph/Value.h>
+
+namespace at {
+namespace native {
+namespace vulkan {
+
+class ComputeGraph;
+
+/*
+ * Represents a single prepacking op in a ML model. In graph mode, ops will be
+ * implemented in a derived class that implements encode, which will implement
+ * encoding of shaders transferring necessary data (such as weights and biases)
+ * to the GPU.
+ */
+class PrepackNode {
+  friend class ComputeGraph;
+
+ public:
+  PrepackNode(ValueRef tref, ValueRef packed) : tref_{tref}, packed_{packed} {}
+
+  virtual ~PrepackNode() = default;
+
+ protected:
+  ValueRef tref_;
+  ValueRef packed_;
+
+ public:
+  virtual void encode(ComputeGraph* graph) const = 0;
+};
+
+} // namespace vulkan
+} // namespace native
+} // namespace at
+
+#endif /* USE_VULKAN_API */

backends/vulkan/runtime/graph/SharedObject.cpp

@@ -0,0 +1,73 @@
+/*
+ * 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/SharedObject.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ComputeGraph.h>
+
+namespace at {
+namespace native {
+namespace vulkan {
+
+void SharedObject::add_user(ComputeGraph* const graph, const ValueRef idx) {
+  vTensor& t = graph->get_val(idx).toTensor();
+
+  //
+  // Aggregate Memory Requirements
+  //
+
+  const VkMemoryRequirements mem_reqs = t.get_memory_requirements();
+  aggregate_memory_requirements.size =
+      std::max(mem_reqs.size, aggregate_memory_requirements.size);
+  aggregate_memory_requirements.alignment =
+      std::max(mem_reqs.alignment, aggregate_memory_requirements.alignment);
+  aggregate_memory_requirements.memoryTypeBits |= mem_reqs.memoryTypeBits;
+
+  //
+  // Aggregate Allocation Create Info
+  //
+
+  const VmaAllocationCreateInfo create_info = t.get_allocation_create_info();
+  // Clear out CREATE_STRATEGY bit flags in case of conflict
+  VmaAllocationCreateFlags clear_mask = ~VMA_ALLOCATION_CREATE_STRATEGY_MASK;
+  VmaAllocationCreateFlags create_flags = create_info.flags & clear_mask;
+  // Use the default allocation strategy
+  aggregate_create_info.flags = create_flags | api::DEFAULT_ALLOCATION_STRATEGY;
+
+  // Set the usage flag if it is currently not set
+  if (aggregate_create_info.usage == VMA_MEMORY_USAGE_UNKNOWN) {
+    aggregate_create_info.usage = create_info.usage;
+  }
+  // Otherwise check that there is no conflict regarding usage
+  VK_CHECK_COND(aggregate_create_info.usage == create_info.usage);
+  aggregate_create_info.requiredFlags |= create_info.requiredFlags;
+  aggregate_create_info.preferredFlags |= create_info.preferredFlags;
+
+  users.emplace_back(idx);
+}
+
+void SharedObject::allocate(ComputeGraph* const graph) {
+  if (aggregate_memory_requirements.size == 0) {
+    return;
+  }
+  allocation = graph->context()->adapter_ptr()->vma().create_allocation(
+      aggregate_memory_requirements, aggregate_create_info);
+}
+
+void SharedObject::bind_users(ComputeGraph* const graph) {
+  if (users.empty()) {
+    return;
+  }
+  for (const ValueRef idx : users) {
+    graph->get_val(idx).toTensor().bind_allocation(allocation);
+  }
+}
+
+} // namespace vulkan
+} // namespace native
+} // namespace at