Update on "[ET-VK][Op Redesign][4/n] Merge ArithmeticNode into ExecuteNode"

This diff moves the logic of `ArithmeticNode` into its corresponding OpFunction `add_arithmetic_node()` and the `ExecuteNode` class.

Our aim is to remove all derived classes of `ExecuteNode`, i.e., to make `ExecuteNode` a final class. All operator-specific logic will be handled in the OpFunction.

Note the next change will move `StagingNode` into its OpFunction + this new ExecuteNode implementation. Until then, we can't tidy up the `ExecuteNode` class fully. Finally, we leave a few task TODOs.

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

[ghstack-poisoned]

Author: jorgep31415

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

@@ -41,13 +41,13 @@ class ExecuteNode {
       const api::utils::uvec3& local_workgroup_size,
       const std::vector<ValueRef>& outputs,
       const std::vector<ValueRef>& inputs,
-      const api::UniformParamsBuffer& params)
+      api::UniformParamsBuffer&& params)
       : shader_(shader),
         global_workgroup_size_(global_workgroup_size),
         local_workgroup_size_(local_workgroup_size),
         outputs_(outputs),
         inputs_(inputs),
-        params_(params) {}
+        params_(std::move(params)) {}
 
   virtual ~ExecuteNode() = default;
 

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

@@ -82,7 +82,7 @@ void add_arithmetic_node(
   api::UniformParamsBuffer params(graph.context(), block);
 
   graph.execute_nodes().emplace_back(new ExecuteNode(
-      shader, global_size, local_size, {out}, {arg1, arg2}, params));
+      shader, global_size, local_size, {out}, {arg1, arg2}, std::move(params)));
 }
 
 ArithmeticPrepack::ArithmeticPrepack(const ValueRef tref, const ValueRef packed)

backends/vulkan/test/vulkan_compute_api_test.cpp

@@ -526,7 +526,8 @@ TEST(VulkanComputeGraphTest, test_simple_shared_objects) {
       api::kFloat,
       /*shared_object_idx = */ 4);
 
-  // Allocation count will be 2 (1 staging buffer for each input tensor)
+  // Allocation count will be 2:
+  // 1 staging buffer for each input tensor
   EXPECT_TRUE(get_vma_allocation_count() == 2);
 
   ValueRef c = add_arithmetic_node(
@@ -542,8 +543,10 @@ TEST(VulkanComputeGraphTest, test_simple_shared_objects) {
       api::kFloat,
       /*shared_object_idx = */ 2);
 
-  // Allocation count will be 3 (1 staging buffer for each input tensor)
-  EXPECT_TRUE(get_vma_allocation_count() == 3);
+  // Allocation count will be 4, two are new:
+  // 1 uniform buffer for arithmetic shader params
+  // 1 staging buffer for the input tensor
+  EXPECT_TRUE(get_vma_allocation_count() == 4);
 
   ValueRef e = add_arithmetic_node(
       graph,
@@ -557,14 +560,16 @@ TEST(VulkanComputeGraphTest, test_simple_shared_objects) {
   out.value = e;
   out.staging = graph.set_output_tensor(out.value);
 
-  // Allocation count will be 4 (1 staging buffer for each I/O tensor)
-  EXPECT_TRUE(get_vma_allocation_count() == 4);
+  // Allocation count will be 6, three are new:
+  // 1 uniform buffer for arithmetic shader params
+  // 1 staging buffer for the input tensor
+  EXPECT_TRUE(get_vma_allocation_count() == 6);
 
   graph.encode_execute();
 
   // Allocation count will be 13:
   // 4 staging buffers for each I/O tensor
-  // 6 uniform buffers to store args for each shader dispatch
+  // 6 uniform buffers to store params for each shader dispatch
   // 3 shared objects to back tensor memory
   EXPECT_TRUE(get_vma_allocation_count() == 13);