[ET-VK][ez] Fix using a temporary variable when creating ComputePipeline

## Context

Fix changes introduced in #8634 (D70021032).

This change decoupled local work group size from additional specialization constants. As part of this change, when creating a `VkComputePipeline` a temporary `SpecVarList` is created to merge the `WorkgroupSize` with additional specialization constants.

However, this can be an issue with some Vulkan drivers because the `SpecVarList` is a temporary, and thus will be destroyed at the end of the function call. The pointer stored in the `VkSpecializationInfo` will be invalidated, leading to undefined behaviour.


This diff fixes this by restoring the behaviour of `ComputePipeline::Descriptor` storing the `specialization_constants` directly.

Also fix the fact that the `VkSpecializationMapEntry` vector was also a temporary when creating a `VkComputePipeline` by storing it in `ComputePipeline`.

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

[ghstack-poisoned]

Author: SS-JIA

backends/vulkan/runtime/api/Context.cpp

@@ -124,11 +124,17 @@ 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])};
+
+  spec_constants.append(additional_constants);
+
   VkPipeline pipeline = pipeline_cache().retrieve(
       {pipeline_layout_cache().retrieve(shader_layout, push_constants_size),
        shader_cache().retrieve(shader_descriptor),
-       additional_constants,
-       local_workgroup_size});
+       spec_constants});
 
   cmd_.bind_pipeline(pipeline, pipeline_layout, local_workgroup_size);
 

backends/vulkan/runtime/vk_api/Pipeline.cpp

@@ -275,23 +275,13 @@ ComputePipeline::ComputePipeline(
     const ComputePipeline::Descriptor& descriptor,
     VkPipelineCache pipeline_cache)
     : device_(device), handle_{VK_NULL_HANDLE} {
-  SpecVarList specialization_constants;
-
-  specialization_constants.reserve(
-      3 + descriptor.specialization_constants.size());
-  specialization_constants.append(descriptor.local_wg_size[0]);
-  specialization_constants.append(descriptor.local_wg_size[1]);
-  specialization_constants.append(descriptor.local_wg_size[2]);
-
-  specialization_constants.append(descriptor.specialization_constants);
-  const std::vector<VkSpecializationMapEntry> map_entries =
-      specialization_constants.generate_map_entries();
+  map_entries_ = descriptor.specialization_constants.generate_map_entries();
 
   const VkSpecializationInfo specialization_info{
-      specialization_constants.size(), // mapEntryCount
-      map_entries.data(), // pMapEntries
-      specialization_constants.data_nbytes(), // dataSize
-      specialization_constants.data(), // pData
+      descriptor.specialization_constants.size(), // mapEntryCount
+      map_entries_.data(), // pMapEntries
+      descriptor.specialization_constants.data_nbytes(), // dataSize
+      descriptor.specialization_constants.data(), // pData
   };
 
   const VkPipelineShaderStageCreateInfo shader_stage_create_info{
@@ -330,7 +320,9 @@ ComputePipeline::ComputePipeline(
 }
 
 ComputePipeline::ComputePipeline(ComputePipeline&& other) noexcept
-    : device_(other.device_), handle_(other.handle_) {
+    : device_(other.device_),
+      handle_(other.handle_),
+      map_entries_(std::move(other.map_entries_)) {
   other.handle_ = VK_NULL_HANDLE;
 }
 

backends/vulkan/runtime/vk_api/Pipeline.h

@@ -156,7 +156,6 @@ class ComputePipeline final {
     VkPipelineLayout pipeline_layout;
     VkShaderModule shader_module;
     SpecVarList specialization_constants;
-    utils::WorkgroupSize local_wg_size;
   };
 
   explicit ComputePipeline(
@@ -175,6 +174,7 @@ class ComputePipeline final {
  private:
   VkDevice device_;
   VkPipeline handle_;
+  std::vector<VkSpecializationMapEntry> map_entries_;
 
  public:
   inline VkPipeline handle() const {
@@ -274,9 +274,6 @@ class ComputePipelineCache final {
         seed = utils::hash_combine(seed, new_seed);
       }
 
-      seed = utils::hash_combine(
-          seed, std::hash<uint32_t>()((uint32_t)descriptor.local_wg_size));
-
       return seed;
     }
   };