[ET-VK][runtime][2/3] Add typed dispatch size abstractions

backends/vulkan/runtime/api/Context.cpp

@@ -74,7 +74,7 @@ void Context::cmd_reset_querypool() {
 void Context::report_shader_dispatch_start(
     const std::string& shader_name,
     const utils::uvec3& global_wg_size,
-    const utils::WorkgroupSize& local_wg_size,
+    const LocalWorkGroup& local_wg_size,
     const uint32_t dispatch_id) {
   if (querypool_) {
     querypool_.shader_profile_begin(
@@ -133,7 +133,7 @@ void Context::check_device_capabilities(const vkapi::ShaderInfo& shader) {
 
 vkapi::DescriptorSet Context::get_descriptor_set(
     const vkapi::ShaderInfo& shader_descriptor,
-    const utils::WorkgroupSize& local_workgroup_size,
+    const LocalWorkGroup& local_workgroup_size,
     const vkapi::SpecVarList& additional_constants,
     const uint32_t push_constants_size) {
   VkDescriptorSetLayout shader_layout =
@@ -324,7 +324,7 @@ VkPipeline Context::get_shader_pipeline(
   VkPipelineLayout pipeline_layout =
       pipeline_layout_cache().retrieve(shader_layout, push_constants_size);
 
-  const utils::WorkgroupSize local_workgroup_size(4u, 4u, 1u);
+  const LocalWorkGroup local_workgroup_size(4u, 4u, 1u);
   vkapi::SpecVarList spec_constants = {
       SV(local_workgroup_size[0u]),
       SV(local_workgroup_size[1u]),

backends/vulkan/runtime/api/Context.h

@@ -16,6 +16,7 @@
 #include <executorch/backends/vulkan/runtime/vk_api/Adapter.h>
 #include <executorch/backends/vulkan/runtime/vk_api/Command.h>
 #include <executorch/backends/vulkan/runtime/vk_api/Descriptor.h>
+#include <executorch/backends/vulkan/runtime/vk_api/DispatchGrid.h>
 #include <executorch/backends/vulkan/runtime/vk_api/Fence.h>
 #include <executorch/backends/vulkan/runtime/vk_api/QueryPool.h>
 #include <executorch/backends/vulkan/runtime/vk_api/Runtime.h>
@@ -151,7 +152,7 @@ class Context final {
   void report_shader_dispatch_start(
       const std::string& shader_name,
       const utils::uvec3& global_wg_size,
-      const utils::WorkgroupSize& local_wg_size,
+      const LocalWorkGroup& local_wg_size,
       const uint32_t dispatch_id = UINT32_MAX);
 
   /*
@@ -190,13 +191,13 @@ class Context final {
 
   vkapi::DescriptorSet get_descriptor_set(
       const vkapi::ShaderInfo&,
-      const utils::WorkgroupSize&,
+      const LocalWorkGroup&,
       const vkapi::SpecVarList&,
       const uint32_t push_constants_size);
 
   inline vkapi::DescriptorSet get_descriptor_set(
       const vkapi::ShaderInfo& shader_descriptor,
-      const utils::WorkgroupSize& local_work_group_size) {
+      const LocalWorkGroup& local_work_group_size) {
     return get_descriptor_set(shader_descriptor, local_work_group_size, {}, 0u);
   }
 
@@ -370,15 +371,15 @@ inline bool Context::submit_compute_job(
   report_shader_dispatch_start(
       shader.kernel_name,
       global_work_group,
-      utils::WorkgroupSize(local_work_group_size),
+      LocalWorkGroup(local_work_group_size),
       dispatch_id);
 
   // Factor out template parameter independent code to minimize code bloat.
   // Note that push constants are not exposed yet via this API, therefore the
   // push constants size is assumed to be 0.
   vkapi::DescriptorSet descriptor_set = get_descriptor_set(
       shader,
-      utils::WorkgroupSize(local_work_group_size),
+      LocalWorkGroup(local_work_group_size),
       specialization_constants,
       0u);
 

backends/vulkan/runtime/graph/ComputeGraph.cpp

@@ -845,7 +845,7 @@ void ComputeGraph::update_descriptor_counts(
 
 void ComputeGraph::register_pipeline_to_create(
     const vkapi::ShaderInfo& shader_info,
-    const utils::WorkgroupSize& local_workgroup_size,
+    const LocalWorkGroup& local_workgroup_size,
     const vkapi::SpecVarList& spec_vars,
     const std::vector<PushConstantDataInfo>& push_constants) {
   VkDescriptorSetLayout shader_layout =
@@ -921,8 +921,10 @@ utils::uvec3 ComputeGraph::create_local_wg_size(
 
   utils::uvec3 local_group_size = {
       8,
-      std::max(1u, std::min(4u, global_wg_size_desc[1].second)),
-      std::max(1u, std::min(2u, global_wg_size_desc[2].second))};
+      global_wg_size_desc[1].second >= 4u       ? 4u
+          : global_wg_size_desc[1].second >= 2u ? 2u
+                                                : 1u,
+      global_wg_size_desc[2].second >= 2u ? 2u : 1u};
 
   if (global_wg_size_desc[2u].second == 1) {
     if (global_wg_size_desc[1u].second == 1) {

backends/vulkan/runtime/graph/ComputeGraph.h

@@ -26,6 +26,8 @@
 #include <executorch/backends/vulkan/runtime/graph/ops/ExecuteNode.h>
 #include <executorch/backends/vulkan/runtime/graph/ops/PrepackNode.h>
 
+#include <executorch/backends/vulkan/runtime/vk_api/DispatchGrid.h>
+
 #ifdef ET_EVENT_TRACER_ENABLED
 std::string& set_and_get_current_operator_json(const std::string& json);
 size_t get_current_operator_count(const bool increment = false);
@@ -1010,7 +1012,7 @@ class ComputeGraph final {
 
   void register_pipeline_to_create(
       const vkapi::ShaderInfo& shader_info,
-      const utils::WorkgroupSize& local_workgroup_size,
+      const LocalWorkGroup& local_workgroup_size,
       const vkapi::SpecVarList& spec_vars,
       const std::vector<PushConstantDataInfo>& push_constants);
 

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

@@ -44,7 +44,7 @@ void BlitNode::encode(ComputeGraph* graph) {
   kernel_name += vkapi::to_string(graph->dtype_of(dst_));
 
   context->report_shader_dispatch_start(
-      kernel_name, utils::uvec3(), utils::WorkgroupSize(), node_id_);
+      kernel_name, utils::uvec3(), LocalWorkGroup(), node_id_);
 
   context->register_blit(
       pipeline_barrier,

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

@@ -15,6 +15,8 @@
 
 #include <executorch/backends/vulkan/runtime/graph/ops/ExecuteNode.h>
 
+#include <executorch/backends/vulkan/runtime/vk_api/DispatchGrid.h>
+
 namespace vkcompute {
 
 class ComputeGraph;
@@ -49,7 +51,7 @@ class DispatchNode : public ExecuteNode {
  protected:
   vkapi::ShaderInfo shader_;
   utils::uvec3 global_workgroup_size_;
-  utils::WorkgroupSize local_workgroup_size_;
+  LocalWorkGroup local_workgroup_size_;
   const vkapi::ParamsBindList params_;
   const vkapi::SpecVarList spec_vars_;
   const std::vector<PushConstantDataInfo> push_constants_;

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

@@ -39,7 +39,7 @@ DynamicDispatchNode::DynamicDispatchNode(
       pick_local_wg_fn_(pick_local_wg_fn) {
   global_workgroup_size_ =
       pick_global_wg_fn(&graph, shader_, args, resize_args);
-  local_workgroup_size_ = utils::WorkgroupSize(pick_local_wg_fn(
+  local_workgroup_size_ = LocalWorkGroup(pick_local_wg_fn(
       &graph, shader_, global_workgroup_size_, args, resize_args));
 
   // Calculate dispatch grid similar to Context.cpp register_shader_dispatch
@@ -76,7 +76,7 @@ DynamicDispatchNode::DynamicDispatchNode(
       pick_local_wg_fn_(pick_local_wg_fn) {
   global_workgroup_size_ =
       pick_global_wg_fn(&graph, shader_, args, resize_args);
-  local_workgroup_size_ = utils::WorkgroupSize(pick_local_wg_fn(
+  local_workgroup_size_ = LocalWorkGroup(pick_local_wg_fn(
       &graph, shader_, global_workgroup_size_, args, resize_args));
   // Calculate the work group grid that will be dispatched
   wg_dispatch_grid_ = {
@@ -119,8 +119,7 @@ bool DynamicDispatchNode::trigger_resize(ComputeGraph* graph) {
   if (pick_local_wg_fn_) {
     utils::uvec3 new_local_wg_uvec3 = pick_local_wg_fn_(
         graph, shader_, global_workgroup_size_, args_, resize_args_);
-    utils::WorkgroupSize new_local_wg =
-        utils::WorkgroupSize(new_local_wg_uvec3);
+    LocalWorkGroup new_local_wg = LocalWorkGroup(new_local_wg_uvec3);
     if (local_workgroup_size_ != new_local_wg) {
       local_workgroup_size_ = new_local_wg;
       dispatch_params_changed = true;

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

@@ -13,6 +13,8 @@
 #include <executorch/backends/vulkan/runtime/graph/containers/PushConstantData.h>
 #include <executorch/backends/vulkan/runtime/graph/containers/Value.h>
 
+#include <executorch/backends/vulkan/runtime/vk_api/DispatchGrid.h>
+
 namespace vkcompute {
 
 class ComputeGraph;
@@ -52,7 +54,7 @@ class PrepackNode final {
   uint32_t node_id_;
   const vkapi::ShaderInfo shader_;
   const utils::uvec3 global_workgroup_size_;
-  const utils::WorkgroupSize local_workgroup_size_;
+  const LocalWorkGroup local_workgroup_size_;
   const ValueRef tref_;
   const ValueRef packed_;
   const vkapi::ParamsBindList params_;

backends/vulkan/runtime/utils/VecUtils.h

@@ -501,59 +501,5 @@ inline int64_t multiply_integers(const C& container) {
   return multiply_integers(container.begin(), container.end());
 }
 
-class WorkgroupSize final {
-  uint32_t val;
-
- public:
-  explicit WorkgroupSize() : val(0) {}
-  explicit WorkgroupSize(const uint32_t x, const uint32_t y, const uint32_t z) {
-    // shift numbers by multiple of 11 bits, since each local workgroup axis can
-    // be 1024 at most and which is 0x400. only z axis can't store 1024, because
-    // it would overflow uint32_t storage.
-    if (z == 1024) {
-      throw std::runtime_error(
-          "Workgroup size in z axis cannot be 1024 because it would overflow uint32_t storage");
-    }
-    val = x | (y << 11) | (z << 22);
-  }
-
-  explicit WorkgroupSize(const uvec3& vec) {
-    // shift numbers by multiple of 11 bits, since each local workgroup axis can
-    // be 1024 at most and which is 0x400. only z axis can't store 1024, because
-    // it would overflow uint32_t storage.
-    if (vec[2u] == 1024) {
-      throw std::runtime_error(
-          "Workgroup size in z axis cannot be 1024 because it would overflow uint32_t storage");
-    }
-    val = vec[0u] | (vec[1u] << 11) | (vec[2u] << 22);
-  }
-
-  explicit inline operator uvec3() const {
-    return {
-        val & 0x7ffu,
-        (val >> 11) & 0x7ffu,
-        (val >> 22),
-    };
-  }
-
-  explicit inline operator uint32_t() const {
-    return val;
-  }
-
-  inline constexpr uint32_t operator[](const int idx) const {
-    return (val >> (11 * idx)) & 0x7ffu;
-  }
-
-  // Equality operator
-  bool operator==(const WorkgroupSize& other) const {
-    return val == other.val;
-  }
-
-  // Inequality operator (optional, for completeness)
-  bool operator!=(const WorkgroupSize& other) const {
-    return !(*this == other);
-  }
-};
-
 } // namespace utils
 } // namespace vkcompute

backends/vulkan/runtime/vk_api/Command.cpp

@@ -81,7 +81,7 @@ void CommandBuffer::end() {
 void CommandBuffer::bind_pipeline(
     VkPipeline pipeline,
     VkPipelineLayout pipeline_layout,
-    const utils::WorkgroupSize local_workgroup_size) {
+    const LocalWorkGroup& local_workgroup_size) {
   VK_CHECK_COND(
       state_ == CommandBuffer::State::RECORDING,
       "Vulkan CommandBuffer: called bind_pipeline() on a command buffer whose state "

backends/vulkan/runtime/vk_api/Command.h

@@ -12,7 +12,7 @@
 
 #include <executorch/backends/vulkan/runtime/vk_api/vk_api.h>
 
-#include <executorch/backends/vulkan/runtime/utils/VecUtils.h>
+#include <executorch/backends/vulkan/runtime/vk_api/DispatchGrid.h>
 
 #include <executorch/backends/vulkan/runtime/vk_api/Descriptor.h>
 #include <executorch/backends/vulkan/runtime/vk_api/Pipeline.h>
@@ -51,7 +51,7 @@ class CommandBuffer final {
   struct Bound {
     VkPipeline pipeline;
     VkPipelineLayout pipeline_layout;
-    utils::WorkgroupSize local_workgroup_size;
+    LocalWorkGroup local_workgroup_size;
     VkDescriptorSet descriptors;
 
     explicit Bound()
@@ -63,7 +63,7 @@ class CommandBuffer final {
     inline void reset() {
       pipeline = VK_NULL_HANDLE;
       pipeline_layout = VK_NULL_HANDLE;
-      local_workgroup_size = utils::WorkgroupSize{0u, 0u, 0u};
+      local_workgroup_size = LocalWorkGroup{0u, 0u, 0u};
       descriptors = VK_NULL_HANDLE;
     }
   };
@@ -89,7 +89,7 @@ class CommandBuffer final {
   void begin();
   void end();
 
-  void bind_pipeline(VkPipeline, VkPipelineLayout, const utils::WorkgroupSize);
+  void bind_pipeline(VkPipeline, VkPipelineLayout, const LocalWorkGroup&);
   void bind_descriptors(VkDescriptorSet);
   void set_push_constants(VkPipelineLayout, const void*, uint32_t);