[ET-VK][ez] Fix using a temporary variable when creating ComputePipeline (pytorch#9464)

This PR was created by the merge bot to help merge the original PR into
the main branch.
ghstack PR number: pytorch#9405 by
@SS-JIA
^ Please use this as the source of truth for the PR details, comments,
and reviews
ghstack PR base:
https://github.com/pytorch/executorch/tree/gh/SS-JIA/197/base
ghstack PR head:
https://github.com/pytorch/executorch/tree/gh/SS-JIA/197/head
Merge bot PR base: https://github.com/pytorch/executorch/tree/main
Merge bot PR head:
https://github.com/pytorch/executorch/tree/gh/SS-JIA/197/orig
@diff-train-skip-merge

Co-authored-by: Stephen Jia <[email protected]>

Author: 2 people

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;
     }
   };