[ET-VK][runtime][3/3] Enforce device-safe dispatch geometry

backends/vulkan/runtime/api/Context.cpp

@@ -73,16 +73,11 @@ void Context::cmd_reset_querypool() {
 
 void Context::report_shader_dispatch_start(
     const std::string& shader_name,
-    const utils::uvec3& global_wg_size,
-    const LocalWorkGroup& local_wg_size,
+    const GlobalWorkGrid& gwg,
+    const LocalWorkGroup& lwg,
     const uint32_t dispatch_id) {
   if (querypool_) {
-    querypool_.shader_profile_begin(
-        cmd_,
-        dispatch_id,
-        shader_name,
-        vkapi::create_extent3d(global_wg_size),
-        vkapi::create_extent3d((utils::uvec3)local_wg_size));
+    querypool_.shader_profile_begin(cmd_, dispatch_id, shader_name, gwg, lwg);
   }
 }
 
@@ -133,7 +128,7 @@ void Context::check_device_capabilities(const vkapi::ShaderInfo& shader) {
 
 vkapi::DescriptorSet Context::get_descriptor_set(
     const vkapi::ShaderInfo& shader_descriptor,
-    const LocalWorkGroup& local_workgroup_size,
+    const LocalWorkGroup& lwg,
     const vkapi::SpecVarList& additional_constants,
     const uint32_t push_constants_size) {
   VkDescriptorSetLayout shader_layout =
@@ -142,10 +137,7 @@ vkapi::DescriptorSet Context::get_descriptor_set(
   VkPipelineLayout pipeline_layout =
       pipeline_layout_cache().retrieve(shader_layout, push_constants_size);
 
-  vkapi::SpecVarList spec_constants = {
-      SV(local_workgroup_size[0u]),
-      SV(local_workgroup_size[1u]),
-      SV(local_workgroup_size[2u])};
+  vkapi::SpecVarList spec_constants = {SV(lwg[0u]), SV(lwg[1u]), SV(lwg[2u])};
 
   spec_constants.append(additional_constants);
 
@@ -158,7 +150,7 @@ vkapi::DescriptorSet Context::get_descriptor_set(
        spec_constants,
        resolved_required_subgroup_size});
 
-  cmd_.bind_pipeline(pipeline, pipeline_layout, local_workgroup_size);
+  cmd_.bind_pipeline(pipeline, pipeline_layout, lwg);
 
   return descriptor_pool().get_descriptor_set(
       shader_layout, shader_descriptor.kernel_layout);
@@ -168,30 +160,35 @@ void Context::register_shader_dispatch(
     const vkapi::DescriptorSet& descriptors,
     vkapi::PipelineBarrier& pipeline_barrier,
     const vkapi::ShaderInfo& shader_descriptor,
-    const utils::uvec3& global_workgroup_size,
+    const GlobalWorkGrid& gwg,
+    const LocalWorkGroup& lwg,
     const void* push_constants_data,
     const uint32_t push_constants_size) {
+  vkapi::Adapter* const adapter = adapter_ptr();
+  const LocalWorkGroup& dispatch_lwg =
+      gwg.required_lwg_size().is_valid() ? gwg.required_lwg_size() : lwg;
+  dispatch_lwg.validate(
+      adapter->max_compute_workgroup_size(),
+      adapter->max_compute_workgroup_invocations());
+  gwg.validate(
+      dispatch_lwg,
+      adapter->max_compute_workgroup_count(),
+      shader_descriptor.out_tile_size);
+
   // Adjust the global workgroup size based on the output tile size
-  uint32_t global_wg_w = utils::div_up(
-      global_workgroup_size[0u], shader_descriptor.out_tile_size[0u]);
-  uint32_t global_wg_h = utils::div_up(
-      global_workgroup_size[1u], shader_descriptor.out_tile_size[1u]);
-  uint32_t global_wg_d = utils::div_up(
-      global_workgroup_size[2u], shader_descriptor.out_tile_size[2u]);
+  uint32_t gwg_w = utils::div_up(gwg[0u], shader_descriptor.out_tile_size[0u]);
+  uint32_t gwg_h = utils::div_up(gwg[1u], shader_descriptor.out_tile_size[1u]);
+  uint32_t gwg_d = utils::div_up(gwg[2u], shader_descriptor.out_tile_size[2u]);
 
   // Submitting a global work group size of 0 is undefined behaviour. If this is
   // detected then submit a single workgroup instead.
-  if (global_wg_w == 0u || global_wg_h == 0u || global_wg_d == 0u) {
-    global_wg_w = 1u;
-    global_wg_h = 1u;
-    global_wg_d = 1u;
+  if (gwg_w == 0u || gwg_h == 0u || gwg_d == 0u) {
+    gwg_w = 1u;
+    gwg_h = 1u;
+    gwg_d = 1u;
   }
 
-  const utils::uvec3 effective_global_wg = {
-      global_wg_w,
-      global_wg_h,
-      global_wg_d,
-  };
+  const GlobalWorkGrid effective_gwg({gwg_w, gwg_h, gwg_d}, kExplicitWorkGrid);
 
   cmd_.bind_descriptors(descriptors.get_bind_handle());
   cmd_.insert_barrier(pipeline_barrier);
@@ -205,7 +202,7 @@ void Context::register_shader_dispatch(
         pipeline_layout, push_constants_data, push_constants_size);
   }
 
-  cmd_.dispatch(effective_global_wg);
+  cmd_.dispatch(effective_gwg, dispatch_lwg);
 }
 
 void Context::register_barrier(vkapi::PipelineBarrier& pipeline_barrier) {
@@ -324,11 +321,8 @@ VkPipeline Context::get_shader_pipeline(
   VkPipelineLayout pipeline_layout =
       pipeline_layout_cache().retrieve(shader_layout, push_constants_size);
 
-  const LocalWorkGroup local_workgroup_size(4u, 4u, 1u);
-  vkapi::SpecVarList spec_constants = {
-      SV(local_workgroup_size[0u]),
-      SV(local_workgroup_size[1u]),
-      SV(local_workgroup_size[2u])};
+  const LocalWorkGroup lwg(4u, 4u, 1u);
+  vkapi::SpecVarList spec_constants = {SV(lwg[0u]), SV(lwg[1u]), SV(lwg[2u])};
 
   spec_constants.append(additional_constants);
 

backends/vulkan/runtime/api/Context.h

@@ -151,8 +151,8 @@ class Context final {
    */
   void report_shader_dispatch_start(
       const std::string& shader_name,
-      const utils::uvec3& global_wg_size,
-      const LocalWorkGroup& local_wg_size,
+      const GlobalWorkGrid& gwg,
+      const LocalWorkGroup& lwg,
       const uint32_t dispatch_id = UINT32_MAX);
 
   /*
@@ -197,15 +197,16 @@ class Context final {
 
   inline vkapi::DescriptorSet get_descriptor_set(
       const vkapi::ShaderInfo& shader_descriptor,
-      const LocalWorkGroup& local_work_group_size) {
-    return get_descriptor_set(shader_descriptor, local_work_group_size, {}, 0u);
+      const LocalWorkGroup& lwg) {
+    return get_descriptor_set(shader_descriptor, lwg, {}, 0u);
   }
 
   void register_shader_dispatch(
       const vkapi::DescriptorSet&,
       vkapi::PipelineBarrier&,
       const vkapi::ShaderInfo&,
-      const utils::uvec3&,
+      const GlobalWorkGrid&,
+      const LocalWorkGroup&,
       const void* = nullptr,
       const uint32_t = 0);
 
@@ -220,8 +221,8 @@ class Context final {
   bool submit_compute_job(
       const vkapi::ShaderInfo&,
       vkapi::PipelineBarrier&,
-      const utils::uvec3&,
-      const utils::uvec3&,
+      const GlobalWorkGrid&,
+      const LocalWorkGroup&,
       const vkapi::SpecVarList&,
       VkFence fence_handle,
       const uint32_t dispatch_id,
@@ -334,8 +335,8 @@ template <typename... Arguments>
 inline bool Context::submit_compute_job(
     const vkapi::ShaderInfo& shader,
     vkapi::PipelineBarrier& pipeline_barrier,
-    const utils::uvec3& global_work_group,
-    const utils::uvec3& local_work_group_size,
+    const GlobalWorkGrid& gwg,
+    const LocalWorkGroup& lwg,
     const vkapi::SpecVarList& specialization_constants,
     VkFence fence_handle,
     const uint32_t dispatch_id,
@@ -368,20 +369,17 @@ inline bool Context::submit_compute_job(
 
   set_cmd();
 
+  const LocalWorkGroup& dispatch_lwg =
+      gwg.required_lwg_size().is_valid() ? gwg.required_lwg_size() : lwg;
+
   report_shader_dispatch_start(
-      shader.kernel_name,
-      global_work_group,
-      LocalWorkGroup(local_work_group_size),
-      dispatch_id);
+      shader.kernel_name, gwg, dispatch_lwg, 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,
-      LocalWorkGroup(local_work_group_size),
-      specialization_constants,
-      0u);
+  vkapi::DescriptorSet descriptor_set =
+      get_descriptor_set(shader, dispatch_lwg, specialization_constants, 0u);
 
   detail::bind(
       descriptor_set,
@@ -390,7 +388,7 @@ inline bool Context::submit_compute_job(
 
   // Factor out template parameter independent code to minimize code bloat.
   register_shader_dispatch(
-      descriptor_set, pipeline_barrier, shader, global_work_group);
+      descriptor_set, pipeline_barrier, shader, gwg, dispatch_lwg);
 
   report_shader_dispatch_end();
 

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 LocalWorkGroup& local_workgroup_size,
+    const LocalWorkGroup& lwg,
     const vkapi::SpecVarList& spec_vars,
     const std::vector<PushConstantDataInfo>& push_constants) {
   VkDescriptorSetLayout shader_layout =
@@ -857,10 +857,7 @@ void ComputeGraph::register_pipeline_to_create(
     pc_offset += pc.write(pc_data.data(), pc_offset, kMaxPushConstantSize);
   }
 
-  vkapi::SpecVarList spec_constants = {
-      SV(local_workgroup_size[0u]),
-      SV(local_workgroup_size[1u]),
-      SV(local_workgroup_size[2u])};
+  vkapi::SpecVarList spec_constants = {SV(lwg[0u]), SV(lwg[1u]), SV(lwg[2u])};
 
   spec_constants.append(spec_vars);
 
@@ -889,64 +886,46 @@ void ComputeGraph::register_pipeline_to_create(
   pipeline_descriptors_.insert(desc);
 }
 
-utils::uvec3 ComputeGraph::create_global_wg_size(const ValueRef idx) {
+GlobalWorkGrid ComputeGraph::create_gwg(const ValueRef idx) {
   if (is_buffer_storage(idx)) {
-    return {uint32_t(numel_of(idx)), 1u, 1u};
+    return create_linear_gwg(utils::safe_downcast<uint64_t>(numel_of(idx)));
   }
-  return logical_limits_of(idx);
-}
 
-utils::uvec3 ComputeGraph::create_local_wg_size(
-    const utils::uvec3 global_wg_size) {
-  if (config_.enable_local_wg_size_override) {
-    return config_.local_wg_size_override;
-  }
+  return GlobalWorkGrid(
+      utils::make_uvec3(logical_limits_of(idx)), kTextureExtentsWorkGrid);
+}
 
-  // array containing axis index and global workgroup size
-  std::pair<uint32_t, uint32_t> global_wg_size_desc[] = {
-      {0u, global_wg_size[0]},
-      {1u, global_wg_size[1]},
-      {2u, global_wg_size[2]}};
+GlobalWorkGrid ComputeGraph::create_linear_gwg(const uint64_t numel) {
+  vkapi::Adapter* const adapter = context()->adapter_ptr();
+  GlobalWorkGrid gwg(
+      {utils::safe_downcast<uint32_t>(numel), 1u, 1u}, kLinearWorkGrid);
+  gwg.wrap_linear_dispatch(
+      adapter->max_compute_workgroup_count(),
+      adapter->recommended_lwg_nthreads());
+  return gwg;
+}
 
-  // sort the global workgroup size in descending order
-  if (global_wg_size_desc[0].second < global_wg_size_desc[1].second) {
-    std::swap(global_wg_size_desc[0], global_wg_size_desc[1]);
-  }
-  if (global_wg_size_desc[1].second < global_wg_size_desc[2].second) {
-    std::swap(global_wg_size_desc[1], global_wg_size_desc[2]);
-  }
-  if (global_wg_size_desc[0].second < global_wg_size_desc[1].second) {
-    std::swap(global_wg_size_desc[0], global_wg_size_desc[1]);
+LocalWorkGroup ComputeGraph::create_lwg(const GlobalWorkGrid& gwg) {
+  if (gwg.required_lwg_size().is_valid()) {
+    return gwg.required_lwg_size();
   }
 
-  utils::uvec3 local_group_size = {
-      8,
-      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};
+  vkapi::Adapter* const adapter = context()->adapter_ptr();
 
-  if (global_wg_size_desc[2u].second == 1) {
-    if (global_wg_size_desc[1u].second == 1) {
-      local_group_size[0u] = 64;
-      local_group_size[1u] = 1;
-    } else if (global_wg_size_desc[1u].second % 4 == 0) {
-      local_group_size[0u] = 16;
-      local_group_size[1u] = 4;
-    } else {
-      local_group_size[0u] = 32;
-      local_group_size[1u] = 2;
-    }
+  utils::uvec3 shape_weights{
+      gwg[0] > 1u ? 1u : 0u, gwg[1] > 1u ? 1u : 0u, gwg[2] > 1u ? 1u : 0u};
+  if (shape_weights == utils::uvec3{0u, 0u, 0u}) {
+    shape_weights[0] = 1u;
   }
 
-  return {
-      local_group_size[global_wg_size_desc[0].first],
-      local_group_size[global_wg_size_desc[1].first],
-      local_group_size[global_wg_size_desc[2].first]};
+  LocalWorkGroup lwg(
+      LwgShape(shape_weights), adapter->recommended_lwg_nthreads());
+  lwg.fit_to_global(gwg);
+  return lwg;
 }
 
-utils::uvec3 ComputeGraph::create_local_wg_size(const ValueRef idx) {
-  return create_local_wg_size(create_global_wg_size(idx));
+LocalWorkGroup ComputeGraph::create_lwg(const ValueRef idx) {
+  return create_lwg(create_gwg(idx));
 }
 
 void ComputeGraph::bind_tensor_to_descriptor_set(

backends/vulkan/runtime/graph/ComputeGraph.h

@@ -1012,7 +1012,7 @@ class ComputeGraph final {
 
   void register_pipeline_to_create(
       const vkapi::ShaderInfo& shader_info,
-      const LocalWorkGroup& local_workgroup_size,
+      const LocalWorkGroup& lwg,
       const vkapi::SpecVarList& spec_vars,
       const std::vector<PushConstantDataInfo>& push_constants);
 
@@ -1025,37 +1025,36 @@ class ComputeGraph final {
   //
 
   /*
-   * Create a global workgroup size for a given `api::vTensor` value assuming
-   * that every shader invocation calculates one texel element of the output
-   * tensor.
+   * Create a global invocation size for a given `api::vTensor` value assuming
+   * that every shader invocation calculates one element of the output tensor.
    *
    * For tensors that use texture storage, the image extents of the
    * `api::vTensor` will be used to set the global workgroup size.
    *
-   * For tensor that use buffer storage, the number of texels in the texel
-   * buffer will be used to set the x component of the global workgroup size.
-   * All other components will be set to 1 (i.e. {ntexels, 1, 1} will be
-   * returned).
+   * Buffer tensors use linear dispatch intent. Oversized X dimensions are
+   * wrapped across X and Y according to device workgroup-count limits.
    */
-  utils::uvec3 create_global_wg_size(const ValueRef idx);
+  GlobalWorkGrid create_gwg(const ValueRef idx);
+
+  GlobalWorkGrid create_linear_gwg(const uint64_t numel);
 
   /*
-   * Suggest a local workgroup size for a given global workgroup size.
+   * Suggest a local workgroup size for a given global invocation size.
    *
    * The local workgroup size will be formed to try and minimize the number of
    * inactive invocations.
    *
-   * Currently, the local workgroup size is hard-coded to contain a total of 64
-   * shader invocations. In the future, this value can be configured.
+   * Linear dispatches return their binding hint. Other dispatches use a shape
+   * heuristic targeting the adapter's recommended thread count.
    */
-  utils::uvec3 create_local_wg_size(const utils::uvec3 global_wg_size);
+  LocalWorkGroup create_lwg(const GlobalWorkGrid& gwg);
 
   /*
    * Convenience function to suggest a local workgroup size for a given
    * `api::vTensor` value, assuming that every shader invocation calculates one
    * texel element of the output tensor.
    */
-  utils::uvec3 create_local_wg_size(const ValueRef idx);
+  LocalWorkGroup create_lwg(const ValueRef idx);
 
   void bind_tensor_to_descriptor_set(
       const ValueRef ref,