From ed7e765e6fd6df727884ddfad17e5ab54ec1bb1c Mon Sep 17 00:00:00 2001
From: Stephen Jia <ssjia@meta.com>
Date: Wed, 19 Mar 2025 11:51:07 -0700
Subject: [PATCH] [ET-VK][ez] Fix using a temporary variable when creating
 `ComputePipeline`

## Context

Fix changes introduced in https://github.com/pytorch/executorch/pull/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]
---
 backends/vulkan/runtime/api/Context.cpp     | 10 +++++++--
 backends/vulkan/runtime/vk_api/Pipeline.cpp | 24 +++++++--------------
 backends/vulkan/runtime/vk_api/Pipeline.h   |  5 +----
 3 files changed, 17 insertions(+), 22 deletions(-)

diff --git a/backends/vulkan/runtime/api/Context.cpp b/backends/vulkan/runtime/api/Context.cpp
index 42006bd4564..258afba7502 100644
--- a/backends/vulkan/runtime/api/Context.cpp
+++ b/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);
 
diff --git a/backends/vulkan/runtime/vk_api/Pipeline.cpp b/backends/vulkan/runtime/vk_api/Pipeline.cpp
index b5ee47cd2c7..5dcb00168b2 100644
--- a/backends/vulkan/runtime/vk_api/Pipeline.cpp
+++ b/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;
 }
 
diff --git a/backends/vulkan/runtime/vk_api/Pipeline.h b/backends/vulkan/runtime/vk_api/Pipeline.h
index 3248051d12a..1e0fc1e28aa 100644
--- a/backends/vulkan/runtime/vk_api/Pipeline.h
+++ b/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;
     }
   };