[ET-VK] Add workgroup sizes to ShaderData

Pull Request resolved: pytorch#7875

Information is critical for shader optimization.
ghstack-source-id: 262751884
@exported-using-ghexport

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

Co-authored-by: jorgep31415 <[email protected]>

Author: 2 people

backends/vulkan/runtime/VulkanBackend.cpp

@@ -602,7 +602,9 @@ class VulkanBackend final : public ::executorch::runtime::BackendInterface {
           event_name.c_str(),
           /* delegate_debug_id = */ -1,
           r.start_time_ns,
-          r.end_time_ns);
+          r.end_time_ns,
+          (void*)(&r.metadata),
+          sizeof(r.metadata));
     }
 #endif // ET_EVENT_TRACER_ENABLED
 

backends/vulkan/runtime/vk_api/QueryPool.cpp

@@ -176,6 +176,7 @@ std::string stringize(const VkExtent3D& extents) {
      << "}";
   return ss.str();
 }
+
 std::vector<ShaderResult> QueryPool::get_shader_timestamp_data() {
   if (querypool_ == VK_NULL_HANDLE) {
     return {};
@@ -188,7 +189,16 @@ std::vector<ShaderResult> QueryPool::get_shader_timestamp_data() {
         .dispatch_id = entry.dispatch_id,
         .start_time_ns = entry.start_time_ns,
         .end_time_ns = entry.end_time_ns,
-    });
+        .metadata = ShaderMetadata{
+            .global_workgroup_size =
+                {entry.global_workgroup_size.width,
+                 entry.global_workgroup_size.height,
+                 entry.global_workgroup_size.depth},
+            .local_workgroup_size =
+                {entry.local_workgroup_size.width,
+                 entry.local_workgroup_size.height,
+                 entry.local_workgroup_size.depth},
+        }});
   }
   return shader_result;
 }

backends/vulkan/runtime/vk_api/QueryPool.h

@@ -26,11 +26,17 @@
 namespace vkcompute {
 namespace vkapi {
 
+struct ShaderMetadata final {
+  const uint32_t global_workgroup_size[3];
+  const uint32_t local_workgroup_size[3];
+};
+
 struct ShaderResult final {
   const std::string kernel_name;
   const uint32_t dispatch_id;
   const uint64_t start_time_ns;
   const uint64_t end_time_ns;
+  const ShaderMetadata metadata;
 };
 
 struct QueryPoolConfig final {