[ET-VK] Introduce SpecVarList to represent specialization constants

Pull Request resolved: #3078

## 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.
ghstack-source-id: 222863753
@exported-using-ghexport

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

Author: SS-JIA

backends/vulkan/runtime/api/Context.cpp

@@ -59,17 +59,25 @@ Context::~Context() {
 
 DescriptorSet Context::get_descriptor_set(
     const ShaderInfo& shader_descriptor,
-    const utils::uvec3& local_workgroup_size) {
+    const utils::uvec3& local_workgroup_size,
+    const SpecVarList& additional_constants) {
   VkDescriptorSetLayout shader_layout =
       shader_layout_cache().retrieve(shader_descriptor.kernel_layout);
 
   VkPipelineLayout pipeline_layout =
       pipeline_layout_cache().retrieve(shader_layout);
 
+  SpecVarList spec_constants = {
+      SV(local_workgroup_size.data[0u]),
+      SV(local_workgroup_size.data[1u]),
+      SV(local_workgroup_size.data[2u])};
+
+  spec_constants.append(additional_constants);
+
   VkPipeline pipeline = pipeline_cache().retrieve(
       {pipeline_layout_cache().retrieve(shader_layout),
        shader_cache().retrieve(shader_descriptor),
-       local_workgroup_size});
+       spec_constants});
 
   cmd_.bind_pipeline(pipeline, pipeline_layout, local_workgroup_size);
 

backends/vulkan/runtime/api/Context.h

@@ -172,7 +172,16 @@ class Context final {
     }
   }
 
-  DescriptorSet get_descriptor_set(const ShaderInfo&, const utils::uvec3&);
+  DescriptorSet get_descriptor_set(
+      const ShaderInfo&,
+      const utils::uvec3&,
+      const SpecVarList&);
+
+  inline DescriptorSet get_descriptor_set(
+      const ShaderInfo& shader_descriptor,
+      const utils::uvec3& local_work_group_size) {
+    return get_descriptor_set(shader_descriptor, local_work_group_size, {});
+  }
 
   void register_shader_dispatch(
       const DescriptorSet&,

backends/vulkan/runtime/api/Pipeline.cpp

@@ -98,6 +98,103 @@ VkImageLayout vk_layout(
   return VK_IMAGE_LAYOUT_UNDEFINED;
 }
 
+//
+// SpecVar
+//
+
+SpecVar::SpecVar() : type(SpecVar::Type::INT) {
+  value.as_int32 = 0;
+}
+
+SpecVar::SpecVar(const float val) : type(SpecVar::Type::FLOAT) {
+  value.as_float = val;
+}
+
+SpecVar::SpecVar(const int32_t val) : type(SpecVar::Type::INT) {
+  value.as_int32 = val;
+}
+
+SpecVar::SpecVar(const uint32_t val) : type(SpecVar::Type::UINT) {
+  value.as_uint32 = val;
+}
+
+SpecVar::SpecVar(const bool val) : type(SpecVar::Type::BOOL) {
+  value.as_bool = val;
+}
+
+uint32_t SpecVar::val_size() const {
+  switch (type) {
+    case SpecVar::Type::FLOAT:
+      return sizeof(float);
+    case SpecVar::Type::INT:
+      return sizeof(int32_t);
+    case SpecVar::Type::UINT:
+      return sizeof(uint32_t);
+    case SpecVar::Type::BOOL:
+      return sizeof(bool);
+  }
+  return 4;
+}
+
+uint32_t SpecVar::val_offset() const {
+  return api::utils::safe_downcast<uint32_t>(offsetof(SpecVar, value));
+}
+
+bool operator==(const SpecVar& lhs, const SpecVar& rhs) {
+  if (lhs.type != rhs.type) {
+    return false;
+  }
+  switch (lhs.type) {
+    case SpecVar::Type::FLOAT:
+      return lhs.value.as_float == rhs.value.as_float;
+    case SpecVar::Type::INT:
+      return lhs.value.as_int32 == rhs.value.as_int32;
+    case SpecVar::Type::UINT:
+      return lhs.value.as_uint32 == rhs.value.as_uint32;
+    case SpecVar::Type::BOOL:
+      return lhs.value.as_bool == rhs.value.as_bool;
+  }
+  return false;
+}
+
+SpecVarList::SpecVarList() {
+  vars.reserve(8);
+}
+
+SpecVarList::SpecVarList(std::initializer_list<SpecVar> init_list) {
+  vars.resize(init_list.size());
+  std::copy(init_list.begin(), init_list.end(), vars.begin());
+}
+
+void SpecVarList::append(const SpecVarList& other) {
+  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);
+  }
+  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.vars.at(i) != rhs.vars.at(i)) {
+      return false;
+    }
+  }
+  return true;
+}
+
 //
 // PipelineLayout
 //
@@ -154,33 +251,14 @@ ComputePipeline::ComputePipeline(
     const ComputePipeline::Descriptor& descriptor,
     VkPipelineCache pipeline_cache)
     : device_(device), handle_{VK_NULL_HANDLE} {
-  // NOLINTNEXTLINE
-  constexpr VkSpecializationMapEntry specialization_map_entries[3]{
-      // X
-      {
-          0u,
-          offsetof(utils::uvec3, data[0u]),
-          sizeof(utils::uvec3::data[0u]),
-      },
-      // Y
-      {
-          1u,
-          offsetof(utils::uvec3, data[1u]),
-          sizeof(utils::uvec3::data[1u]),
-      },
-      // Z
-      {
-          2u,
-          offsetof(utils::uvec3, data[2u]),
-          sizeof(utils::uvec3::data[2u]),
-      },
-  };
+  std::vector<VkSpecializationMapEntry> map_entries =
+      descriptor.specialization_constants.generate_map_entries();
 
   const VkSpecializationInfo specialization_info{
-      3u, // mapEntryCount
-      specialization_map_entries, // pMapEntries
-      sizeof(descriptor.local_work_group), // dataSize
-      &descriptor.local_work_group, // 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{
@@ -242,7 +320,7 @@ bool operator==(
   return (
       _1.pipeline_layout == _2.pipeline_layout &&
       _1.shader_module == _2.shader_module &&
-      _1.local_work_group == _2.local_work_group);
+      _1.specialization_constants == _2.specialization_constants);
 }
 
 //

backends/vulkan/runtime/api/Pipeline.h

@@ -18,9 +18,68 @@
 #include <mutex>
 #include <unordered_map>
 
+#define SV(x) ::vkcompute::api::SpecVar(x)
+
 namespace vkcompute {
 namespace api {
 
+struct SpecVar final {
+  enum class Type : uint8_t {
+    FLOAT,
+    INT,
+    UINT,
+    BOOL,
+  };
+
+  union Value {
+    int32_t as_int32;
+    uint32_t as_uint32;
+    float as_float;
+    bool as_bool;
+  };
+
+  Value value;
+  Type type;
+
+  SpecVar();
+  SpecVar(const float val);
+  SpecVar(const int32_t val);
+  SpecVar(const uint32_t val);
+  SpecVar(const bool val);
+
+  uint32_t val_size() const;
+  uint32_t val_offset() const;
+};
+
+bool operator==(const SpecVar& lhs, const SpecVar& rhs);
+
+struct SpecVarList final {
+  std::vector<SpecVar> vars;
+
+  SpecVarList();
+  SpecVarList(std::initializer_list<SpecVar> init_list);
+
+  inline const SpecVar* data() const {
+    return vars.data();
+  }
+
+  inline uint32_t size() const {
+    return api::utils::safe_downcast<uint32_t>(vars.size());
+  }
+
+  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);
+
 struct PipelineBarrier final {
   struct Stages final {
     VkPipelineStageFlags src;
@@ -83,7 +142,7 @@ class ComputePipeline final {
   struct Descriptor final {
     VkPipelineLayout pipeline_layout;
     VkShaderModule shader_module;
-    utils::uvec3 local_work_group;
+    SpecVarList specialization_constants;
   };
 
   explicit ComputePipeline(
@@ -171,12 +230,29 @@ class ComputePipelineCache final {
           seed, std::hash<VkPipelineLayout>()(descriptor.pipeline_layout));
       seed = utils::hash_combine(
           seed, std::hash<VkShaderModule>()(descriptor.shader_module));
-      seed = utils::hash_combine(
-          seed, std::hash<uint32_t>()(descriptor.local_work_group.data[0u]));
-      seed = utils::hash_combine(
-          seed, std::hash<uint32_t>()(descriptor.local_work_group.data[1u]));
-      seed = utils::hash_combine(
-          seed, std::hash<uint32_t>()(descriptor.local_work_group.data[2u]));
+
+      const SpecVarList& spec_vars = descriptor.specialization_constants;
+      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.vars.at(i);
+        size_t new_seed = 0;
+        switch (spec_var.type) {
+          case SpecVar::Type::FLOAT:
+            new_seed = std::hash<float>()(spec_var.value.as_float);
+            break;
+          case SpecVar::Type::INT:
+            new_seed = std::hash<int32_t>()(spec_var.value.as_int32);
+            break;
+          case SpecVar::Type::UINT:
+            new_seed = std::hash<uint32_t>()(spec_var.value.as_uint32);
+            break;
+          case SpecVar::Type::BOOL:
+            new_seed = std::hash<bool>()(spec_var.value.as_bool);
+            break;
+        }
+        seed = utils::hash_combine(seed, new_seed);
+      }
 
       return seed;
     }

backends/vulkan/test/vulkan_compute_api_test.cpp

@@ -50,6 +50,53 @@ TEST_F(VulkanComputeAPITest, retrieve_custom_shader_test) {
   ASSERT_TRUE(kernel.kernel_name == "test_shader");
 }
 
+TEST_F(VulkanComputeAPITest, spec_var_classes_test) {
+  // Check equality operator
+  ASSERT_TRUE(SV(1.5f) == SV(1.5f));
+  ASSERT_FALSE(SV(15.0f) == SV(15));
+  ASSERT_FALSE(SV(1u) == SV(true));
+
+  size_t sv_size = sizeof(api::SpecVar);
+
+  api::SpecVarList spec_vars = {};
+  ASSERT_TRUE(spec_vars.size() == 0);
+  spec_vars = {SV(1.1f), SV(32), SV(45)};
+  ASSERT_TRUE(spec_vars.size() == 3);
+  api::SpecVarList spec_vars_other = {SV(2.6f), SV(true), SV(78u), SV(5.5f)};
+  spec_vars.append(spec_vars_other);
+  ASSERT_TRUE(spec_vars.size() == 7);
+
+  // Check validity of the 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
+  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);
+    ASSERT_TRUE(entries[i].offset == sv_size * i);
+    if (i != 4) {
+      ASSERT_TRUE(entries[i].size == 4);
+    } else {
+      ASSERT_TRUE(entries[i].size == 1);
+    }
+  }
+
+  // Check copy
+  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_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);
+}
+
 TEST_F(VulkanComputeAPITest, update_params_between_submit) {
   api::context()->set_cmd(/*reusable = */ true);
   std::vector<int64_t> sizes = {4, 4, 2};