[ET-VK] Allow overwriting local workgroup size

Pull Request resolved: #4046

Introduce a `GraphConfig` toggle following the convention of `storage_type` and `memory_layout`.

Differential Revision: [D58957058](https://our.internmc.facebook.com/intern/diff/D58957058/)
ghstack-source-id: 231385519

Author: jorgep31415

backends/vulkan/runtime/graph/ComputeGraph.cpp

@@ -320,6 +320,10 @@ api::utils::uvec3 ComputeGraph::create_global_wg_size(const ValueRef idx) {
 }
 
 api::utils::uvec3 ComputeGraph::create_local_wg_size(const ValueRef idx) {
+  if (config_.enable_local_wg_size_override) {
+    return config_.local_wg_size_override;
+  }
+
   if (is_buffer_storage(idx)) {
     return {64u, 1u, 1u};
   }

backends/vulkan/runtime/graph/GraphConfig.cpp

@@ -60,6 +60,9 @@ GraphConfig::GraphConfig() {
   // QueryPool objects are used to measure execution times of individual shader
   // dispatches. By default, this functionality is disabled.
   enable_querypool = false;
+
+  enable_local_wg_size_override = false;
+  local_wg_size_override = {};
 }
 
 void GraphConfig::set_storage_type_override(api::StorageType storage_type) {
@@ -73,4 +76,10 @@ void GraphConfig::set_memory_layout_override(
   memory_layout_override = memory_layout;
 }
 
+void GraphConfig::set_local_wg_size_override(
+    const api::utils::uvec3& local_wg_size) {
+  enable_local_wg_size_override = true;
+  local_wg_size_override = local_wg_size;
+}
+
 } // namespace vkcompute

backends/vulkan/runtime/graph/GraphConfig.h

@@ -30,11 +30,15 @@ struct GraphConfig final {
 
   bool enable_querypool;
 
+  bool enable_local_wg_size_override;
+  api::utils::uvec3 local_wg_size_override;
+
   // Generate a default graph config with pre-configured settings
   explicit GraphConfig();
 
   void set_storage_type_override(api::StorageType storage_type);
   void set_memory_layout_override(api::GPUMemoryLayout memory_layout);
+  void set_local_wg_size_override(const api::utils::uvec3& local_wg_size);
 };
 
 } // namespace vkcompute