Update on "[ET-VK] Introduce SpecVarList to represent specialization constants"

## Context

Specialization constants are a useful tool to compile compute shaders with constants defined at runtime. The primary application of specialization constants is to define variables which may have an impact on how the code is compiled, for example:

* the number of elements of an array
* the range of a loop

Compared to the shader codegen system, which produces a complete copy of the shader and for which variants must be defined at build time, specialization constants can be defined at runtime when the compute pipeline is built.

Specialization constants are currently used to define local work group sizes in Vulkan, but the Compute API hard-codes the number of specialization constants accepted by the shader to 3.

This changeset introduces the `SpecVar` and `SpecVarList` classes to manage specialization constants and enable additional specialization constants to be specified.

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

[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);