[ET-VK][ez][Refactor] Re-order DispatchNode arguments to match shader layout spec (#10700)

## Context

As title. Note that this PR was written mainly with a meta internal coding agent. This PR re-orders the arguments to `DispatchNode` to be more intuitive.

I asked the agent:

---

I want to change the constructor of the class

```
  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 = {},
      const std::vector<PushConstantDataInfo>& push_constants = {});
```

to instead be

```
  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 std::vector<PushConstantDataInfo>& push_constants = {},
      const vkapi::SpecVarList& spec_vars = {},
      const std::vector<ValueRef>& resize_args = {},
      const ResizeFunction& resize_fn = nullptr);
```

Can you make this change and update the callsites as well?

---

The motivation is to have the arguments match the order in which parameter UBOs, push constant blocks, and specialization variables are declared in a GLSL shader.

The order of `resize`_args` and `resize_fn` was also swapped in the interest of having the function pointer be the last argument. It will also make more sense in a following diff where a `DynamicDispatchNode` class is introduced, which will allow selecting a different compute shader depending on input sizes.

As a small additional change, I also asked the agent to

---

Go through all the files under `xplat/executorch/backends/vulkan/runtime/graph/ops/impl`

Change `vkapi::MemoryAccessType::WRITE` to `vkapi::kWrite` and `vkapi::MemoryAccessType::READ` to `vkapi::kRead` for each file.

---

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

Author: pytorchbot

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

@@ -21,10 +21,10 @@ DispatchNode::DispatchNode(
     const utils::uvec3& local_workgroup_size,
     const std::vector<ArgGroup>& args,
     const vkapi::ParamsBindList& params,
+    const std::vector<PushConstantDataInfo>& push_constants,
     const vkapi::SpecVarList& spec_vars,
-    const ResizeFunction& resize_fn,
     const std::vector<ValueRef>& resize_args,
-    const std::vector<PushConstantDataInfo>& push_constants)
+    const ResizeFunction& resize_fn)
     : ExecuteNode(resize_fn, resize_args, args, shader.kernel_name),
       shader_(shader),
       global_workgroup_size_(global_workgroup_size),

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

@@ -33,10 +33,10 @@ class DispatchNode final : public ExecuteNode {
       const utils::uvec3& local_workgroup_size,
       const std::vector<ArgGroup>& args,
       const vkapi::ParamsBindList& params,
+      const std::vector<PushConstantDataInfo>& push_constants = {},
       const vkapi::SpecVarList& spec_vars = {},
-      const ResizeFunction& resize_fn = nullptr,
       const std::vector<ValueRef>& resize_args = {},
-      const std::vector<PushConstantDataInfo>& push_constants = {});
+      const ResizeFunction& resize_fn = nullptr);
 
   ~DispatchNode() override = default;
 

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

@@ -94,16 +94,19 @@ void add_arange_node(
       graph.create_global_wg_size(out),
       graph.create_local_wg_size(out),
       // Inputs and Outputs
-      {{out, vkapi::MemoryAccessType::WRITE}},
+      {{out, vkapi::kWrite}},
       // Shader params buffers
       {t_out->sizes_ubo(),
        graph.create_params_buffer(start_val),
        graph.create_params_buffer(step_val)},
+      // Push Constants
+      {},
       // Specialization Constants
       {},
+      // Resize Args
+      {start, end, step},
       // Resizing Logic
-      resize_arange_node,
-      {start, end, step}));
+      resize_arange_node));
 }
 
 void arange(ComputeGraph& graph, const std::vector<ValueRef>& args) {

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

@@ -90,12 +90,19 @@ void add_native_batch_norm_node(
       VK_KERNEL_FROM_STR(kernel_name),
       graph.create_global_wg_size(out_ref),
       graph.create_local_wg_size(out_ref),
-      {{out_ref, vkapi::MemoryAccessType::WRITE},
-       {{in_ref, arg_weight, arg_bias, arg_mean, arg_var},
-        vkapi::MemoryAccessType::READ}},
+      {{out_ref, vkapi::kWrite},
+       {{in_ref, arg_weight, arg_bias, arg_mean, arg_var}, vkapi::kRead}},
       {t_out->logical_limits_ubo(),
        graph.create_params_buffer(epsilon),
-       graph.create_params_buffer(num_texel_per_batch)}));
+       graph.create_params_buffer(num_texel_per_batch)},
+      // Push Constants
+      {},
+      // Specialization Constants
+      {},
+      // Resize Args
+      {},
+      // Resizing Logic
+      nullptr));
 }
 
 void native_batch_norm(ComputeGraph& graph, const std::vector<ValueRef>& args) {

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

@@ -84,19 +84,20 @@ void add_binary_op_texture_node(
       graph.create_global_wg_size(out),
       graph.create_local_wg_size(out),
       // Inputs and Outputs
-      {{out, vkapi::MemoryAccessType::WRITE},
-       {{arg1, arg2}, vkapi::MemoryAccessType::READ}},
+      {{out, vkapi::kWrite}, {{arg1, arg2}, vkapi::kRead}},
       // Shader params buffers
       {},
-      // Specialization Constants
-      {t_out->hashed_layout(), t_in1->hashed_layout(), t_in2->hashed_layout()},
-      // Resizing Logic
-      resize_binary_op_node,
-      {},
+      // Push Constants
       {{graph.sizes_pc_of(out),
         graph.sizes_pc_of(arg1),
         graph.sizes_pc_of(arg2),
-        PushConstantDataInfo(&binary_ops_params, sizeof(binary_ops_params))}}));
+        PushConstantDataInfo(&binary_ops_params, sizeof(binary_ops_params))}},
+      // Specialization Constants
+      {t_out->hashed_layout(), t_in1->hashed_layout(), t_in2->hashed_layout()},
+      // Resize Args
+      {},
+      // Resizing Logic
+      resize_binary_op_node));
 }
 
 void add_binary_op_buffer_node(
@@ -127,17 +128,10 @@ void add_binary_op_buffer_node(
       graph.create_global_wg_size(out),
       graph.create_local_wg_size(out),
       // Inputs and Outputs
-      {{out, vkapi::MemoryAccessType::WRITE},
-       {{in1, in2}, vkapi::MemoryAccessType::READ}},
+      {{out, vkapi::kWrite}, {{in1, in2}, vkapi::kRead}},
       // Shader params buffers
       {},
-      // Specialization Constants
-      {graph.packed_dim_of(out),
-       graph.packed_dim_of(in1),
-       graph.packed_dim_of(in2)},
-      // Resizing Logic
-      resize_binary_op_node,
-      {},
+      // Push Constants
       {{
           graph.sizes_pc_of(in1),
           graph.sizes_pc_of(in2),
@@ -146,7 +140,15 @@ void add_binary_op_buffer_node(
           graph.strides_pc_of(in2),
           graph.numel_pc_of(out),
           PushConstantDataInfo(&alpha_val, sizeof(float)),
-      }}));
+      }},
+      // Specialization Constants
+      {graph.packed_dim_of(out),
+       graph.packed_dim_of(in1),
+       graph.packed_dim_of(in2)},
+      // Resize Args
+      {},
+      // Resizing Logic
+      resize_binary_op_node));
 }
 
 void add_binary_op_node(

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

@@ -50,8 +50,12 @@ void add_clone_node(
       {{out, vkapi::kWrite}, {in, vkapi::kRead}},
       // Parameter Buffers
       {t_out->logical_limits_ubo()},
+      // Push Constants
+      {},
       // Specialization Constants
       {},
+      // Resize Args
+      {},
       // Resizing Logic
       resize_clone_node));
 }
@@ -74,8 +78,12 @@ void add_image_to_buffer_node(
       {{buffer, vkapi::kWrite}, {image, vkapi::kRead}},
       // Parameter Buffers
       {graph.sizes_ubo(image), graph.strides_ubo(buffer)},
+      // Push Constants
+      {},
       // Specialization Constants
       {graph.hashed_layout_of(image)},
+      // Resize Args
+      {},
       // Resizing Logic
       resize_clone_node));
 }
@@ -98,8 +106,12 @@ void add_buffer_to_image_node(
       {{image, vkapi::kWrite}, {buffer, vkapi::kRead}},
       // Parameter Buffers
       {graph.sizes_ubo(image), graph.strides_ubo(buffer)},
+      // Push Constants
+      {},
       // Specialization Constants
       {graph.hashed_layout_of(image)},
+      // Resize Args
+      {},
       // Resizing Logic
       resize_clone_node));
 }

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

@@ -444,16 +444,17 @@ void add_conv2d_node(
       wg_size,
       graph.create_local_wg_size(wg_size),
       // Inputs and Outputs
-      {{out, vkapi::MemoryAccessType::WRITE},
-       {{in, arg_weight, arg_bias}, vkapi::MemoryAccessType::READ}},
+      {{out, vkapi::kWrite}, {{in, arg_weight, arg_bias}, vkapi::kRead}},
       // Shader params buffers
       param_buffers,
+      // Push Constants
+      push_constants,
       // Specialization Constants
       {},
-      // Resizing Logic
-      resize_conv2d_node,
+      // Resize Args
       {weight_data, stride, padding, dilation, transposed, output_padding},
-      push_constants));
+      // Resizing Logic
+      resize_conv2d_node));
 }
 
 void add_conv1d_node(
@@ -548,23 +549,25 @@ void add_conv1d_node(
       global_size,
       local_size,
       // Inputs and Outputs
-      {{out, vkapi::MemoryAccessType::WRITE},
-       {{in, arg_weight, arg_bias}, vkapi::MemoryAccessType::READ}},
+      {{out, vkapi::kWrite}, {{in, arg_weight, arg_bias}, vkapi::kRead}},
       // Shader params buffers
       {
           t_out->logical_limits_ubo(),
           t_in->sizes_ubo(),
           graph.create_params_buffer(kernel_params),
           graph.create_params_buffer(out_params),
       },
+      // Push Constants
+      {},
       // Specialization Constants
       {t_out->hashed_layout(),
        t_in->hashed_layout(),
        t_weight->hashed_layout(),
        t_bias->hashed_layout()},
+      // Resize Args
+      {weight, stride, padding, dilation},
       // Resizing Logic
-      resize_conv1d_node,
-      {weight, stride, padding, dilation}));
+      resize_conv1d_node));
 }
 
 void conv(ComputeGraph& graph, const std::vector<ValueRef>& args) {

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

@@ -50,19 +50,22 @@ void add_copy_offset_node(
       },
       // Parameter buffers
       {},
+      // Push Constants
+      {
+          PushConstantDataInfo(&range, sizeof(range), sizeof(ivec4)),
+          PushConstantDataInfo(&src_offset, sizeof(src_offset), sizeof(ivec4)),
+          PushConstantDataInfo(&dst_offset, sizeof(dst_offset), sizeof(ivec4)),
+      },
       // Specialization Constants
       {graph.hashed_layout_of(out),
        graph.hashed_layout_of(in),
        (calc_out_pos_using_src_chnl      ? 1
             : calc_in_pos_using_dst_chnl ? 2
                                          : 0)},
-      nullptr,
+      // Resize Args
       {},
-      {
-          PushConstantDataInfo(&range, sizeof(range), sizeof(ivec4)),
-          PushConstantDataInfo(&src_offset, sizeof(src_offset), sizeof(ivec4)),
-          PushConstantDataInfo(&dst_offset, sizeof(dst_offset), sizeof(ivec4)),
-      }));
+      // Resizing Logic
+      nullptr));
 }
 
 void add_copy_packed_dim_offset_node(
@@ -138,22 +141,25 @@ void add_copy_packed_dim_offset_node(
       graph.create_local_wg_size(global_wg_size),
       // Inputs and Outputs
       {
-          {out, vkapi::MemoryAccessType::WRITE},
-          {out, vkapi::MemoryAccessType::READ},
-          {in, vkapi::MemoryAccessType::READ},
+          {out, vkapi::kWrite},
+          {out, vkapi::kRead},
+          {in, vkapi::kRead},
       },
       // Parameter buffers
       {},
-      // Specialization Constants
-      {graph.hashed_layout_of(out), graph.hashed_layout_of(in)},
-      nullptr,
-      {},
+      // Push Constants
       {
           PushConstantDataInfo(
               &final_range, sizeof(final_range), sizeof(ivec4)),
           PushConstantDataInfo(&src_offset, sizeof(src_offset), sizeof(ivec4)),
           PushConstantDataInfo(&dst_offset, sizeof(dst_offset), sizeof(ivec4)),
-      }));
+      },
+      // Specialization Constants
+      {graph.hashed_layout_of(out), graph.hashed_layout_of(in)},
+      // Resize Args
+      {},
+      // Resizing Logic
+      nullptr));
 }
 
 void add_copy_channel_offset_node(
@@ -248,22 +254,24 @@ void add_copy_channel_offset_node(
         local_size,
         // Inputs and Outputs
         {
-            {out, vkapi::MemoryAccessType::WRITE},
-            {out, vkapi::MemoryAccessType::READ},
-            {in, vkapi::MemoryAccessType::READ},
+            {out, vkapi::kWrite},
+            {out, vkapi::kRead},
+            {in, vkapi::kRead},
         },
         // Parameter buffers
         {},
-        // Specialization Constants
-        {graph.hashed_layout_of(out), graph.hashed_layout_of(in)},
-        nullptr,
-        {},
+        // Push Constants
         {graph.sizes_pc_of(out),
          graph.sizes_pc_of(in),
          PushConstantDataInfo(&range_params, sizeof(range_params)),
          PushConstantDataInfo(&offset_params, sizeof(offset_params)),
-         PushConstantDataInfo(
-             &src_channel_offset, sizeof(src_channel_offset))}));
+         PushConstantDataInfo(&src_channel_offset, sizeof(src_channel_offset))},
+        // Specialization Constants
+        {graph.hashed_layout_of(out), graph.hashed_layout_of(in)},
+        // Resize Args
+        {},
+        // Resizing Logic
+        nullptr));
   }
 }
 

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

@@ -58,9 +58,16 @@ void add_embedding_node(
       {
           t_out->sizes_ubo(),
       },
+      // Push Constants
+      {},
+      // Specialization Constants
       {t_out->hashed_layout(),
        t_in->hashed_layout(),
-       t_weight->hashed_layout()}));
+       t_weight->hashed_layout()},
+      // Resize Args
+      {},
+      // Resizing Logic
+      nullptr));
 }
 
 void embedding(ComputeGraph& graph, const std::vector<ValueRef>& args) {

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

@@ -74,8 +74,12 @@ void add_flip_node(
           graph.sizes_ubo(out),
           graph.create_params_buffer(dim_bitmap),
       },
+      // Push Constants
+      {},
       // Specialization Constants
       {},
+      // Resize Args
+      {},
       // Resizing Logic
       resize_flip_node));
 }