Update on "[ET-VK] Enable additional specialization constants in compute shaders"

## Context

Building on top of the previous changeset in the stack, this changeset modifies shader dispatch APIs to accept additional specialization constants for a shader.

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

[ghstack-poisoned]

Author: SS-JIA

backends/vulkan/runtime/api/Pipeline.cpp

@@ -157,34 +157,38 @@ bool operator==(const SpecVar& lhs, const SpecVar& rhs) {
   return false;
 }
 
+SpecVarList::SpecVarList() {
+  vars.reserve(8);
+}
+
 SpecVarList::SpecVarList(std::initializer_list<SpecVar> init_list) {
-  VK_CHECK_COND(init_list.size() <= SPECVAR_LIST_LIMIT);
-  arr_size = init_list.size();
-  std::copy(init_list.begin(), init_list.end(), arr);
-  uint32_t cur_offset = 0u;
-  for (uint32_t i = 0; i < arr_size; ++i) {
-    map_entries[i] = {i, cur_offset + arr[i].val_offset(), arr[i].val_size()};
-    cur_offset += sizeof(SpecVar);
-  }
+  vars.resize(init_list.size());
+  std::copy(init_list.begin(), init_list.end(), vars.begin());
 }
 
 void SpecVarList::append(const SpecVarList& other) {
-  VK_CHECK_COND(arr_size + other.size() <= SPECVAR_LIST_LIMIT);
-  std::copy(other.arr, other.arr + other.size(), arr + arr_size);
-  uint32_t cur_offset = arr_size * sizeof(SpecVar);
-  for (uint32_t i = arr_size; i < arr_size + other.size(); ++i) {
-    map_entries[i] = {i, cur_offset + arr[i].val_offset(), arr[i].val_size()};
+  vars.insert(vars.end(), other.vars.begin(), other.vars.end());
+}
+
+std::vector<VkSpecializationMapEntry> SpecVarList::generate_map_entries()
+    const {
+  std::vector<VkSpecializationMapEntry> map_entries;
+  map_entries.resize(vars.size());
+  uint32_t cur_offset = 0u;
+  for (uint32_t i = 0; i < vars.size(); ++i) {
+    map_entries[i] = {
+        i, cur_offset + vars.at(i).val_offset(), vars.at(i).val_size()};
     cur_offset += sizeof(SpecVar);
   }
-  arr_size += other.size();
+  return map_entries;
 }
 
 bool operator==(const SpecVarList& lhs, const SpecVarList& rhs) {
   if (lhs.size() != rhs.size()) {
     return false;
   }
   for (uint32_t i = 0; i < lhs.size(); ++i) {
-    if (lhs.var_data()[i] != rhs.var_data()[i]) {
+    if (lhs.vars.at(i) != rhs.vars.at(i)) {
       return false;
     }
   }
@@ -247,10 +251,13 @@ ComputePipeline::ComputePipeline(
     const ComputePipeline::Descriptor& descriptor,
     VkPipelineCache pipeline_cache)
     : device_(device), handle_{VK_NULL_HANDLE} {
+  std::vector<VkSpecializationMapEntry> map_entries =
+      descriptor.specialization_constants.generate_map_entries();
+
   const VkSpecializationInfo specialization_info{
       descriptor.specialization_constants.size(), // mapEntryCount
-      descriptor.specialization_constants.map_entries_data(), // pMapEntries
-      descriptor.specialization_constants.map_entries_data_size(), // dataSize
+      map_entries.data(), // pMapEntries
+      descriptor.specialization_constants.data_nbytes(), // dataSize
       descriptor.specialization_constants.data(), // pData
   };
 

backends/vulkan/runtime/api/Pipeline.h

@@ -18,8 +18,6 @@
 #include <mutex>
 #include <unordered_map>
 
-#define SPECVAR_LIST_LIMIT 8
-
 #define SV(x) ::vkcompute::api::SpecVar(x)
 
 namespace vkcompute {
@@ -55,38 +53,29 @@ struct SpecVar final {
 
 bool operator==(const SpecVar& lhs, const SpecVar& rhs);
 
-// using SpecVarList = std::vector<SpecVar>;
-
-class SpecVarList final {
-  SpecVar arr[SPECVAR_LIST_LIMIT];
-  VkSpecializationMapEntry map_entries[SPECVAR_LIST_LIMIT];
-  uint32_t arr_size;
+struct SpecVarList final {
+  std::vector<SpecVar> vars;
 
- public:
-  SpecVarList() : arr_size(0) {}
+  SpecVarList();
   SpecVarList(std::initializer_list<SpecVar> init_list);
 
-  inline const SpecVar* var_data() const {
-    return &(arr[0]);
-  }
-
-  inline const void* data() const {
-    return &(arr[0]);
+  inline const SpecVar* data() const {
+    return vars.data();
   }
 
   inline uint32_t size() const {
-    return arr_size;
+    return api::utils::safe_downcast<uint32_t>(vars.size());
   }
 
-  inline const VkSpecializationMapEntry* map_entries_data() const {
-    return &(map_entries[0]);
-  }
-
-  inline size_t map_entries_data_size() const {
-    return arr_size * sizeof(VkSpecializationMapEntry);
+  inline uint32_t data_nbytes() const {
+    return vars.size() * sizeof(SpecVar);
   }
 
   void append(const SpecVarList& other);
+
+  std::vector<VkSpecializationMapEntry> generate_map_entries() const;
+
+  friend bool operator==(const SpecVarList& lhs, const SpecVarList& rhs);
 };
 
 bool operator==(const SpecVarList& lhs, const SpecVarList& rhs);
@@ -246,7 +235,7 @@ class ComputePipelineCache final {
       seed = utils::hash_combine(seed, std::hash<uint32_t>()(spec_vars.size()));
 
       for (int i = 0; i < spec_vars.size(); ++i) {
-        const SpecVar& spec_var = spec_vars.var_data()[i];
+        const SpecVar& spec_var = spec_vars.vars.at(i);
         size_t new_seed = 0;
         switch (spec_var.type) {
           case SpecVar::Type::FLOAT:

backends/vulkan/test/vulkan_compute_api_test.cpp

@@ -67,13 +67,14 @@ TEST_F(VulkanComputeAPITest, spec_var_classes_test) {
   ASSERT_TRUE(spec_vars.size() == 7);
 
   // Check validity of the data
-  const api::SpecVar* data = spec_vars.var_data();
+  const api::SpecVar* data = spec_vars.data();
   ASSERT_TRUE(*(reinterpret_cast<const float*>(data + 3)) == 2.6f);
   ASSERT_TRUE(*(reinterpret_cast<const int32_t*>(data + 1)) == 32);
   ASSERT_TRUE(*(reinterpret_cast<const int32_t*>(data + 5)) == 78u);
 
   // Check validity of the map entries
-  const VkSpecializationMapEntry* entries = spec_vars.map_entries_data();
+  std::vector<VkSpecializationMapEntry> entries =
+      spec_vars.generate_map_entries();
 
   for (size_t i = 0; i < spec_vars.size(); ++i) {
     ASSERT_TRUE(entries[i].constantID == i);
@@ -86,11 +87,11 @@ TEST_F(VulkanComputeAPITest, spec_var_classes_test) {
   }
 
   // Check copy
-  api::SpecVarList spec_vars_2(spec_vars);
-  ASSERT_TRUE(spec_vars.size() == 7);
+  api::SpecVarList spec_vars_copy(spec_vars);
+  ASSERT_TRUE(spec_vars_copy.size() == 7);
 
   // Check validity of the copied data
-  const api::SpecVar* copy_data = spec_vars.var_data();
+  const api::SpecVar* copy_data = spec_vars_copy.data();
   ASSERT_TRUE(*(reinterpret_cast<const bool*>(copy_data + 4)) == true);
   ASSERT_TRUE(*(reinterpret_cast<const int32_t*>(copy_data + 2)) == 45);
   ASSERT_TRUE(*(reinterpret_cast<const float*>(copy_data + 6)) == 5.5f);