[VULKAN] Add support for specialization constants (#499)

This commit introduces support for specialization constants in the
Vulkan backend.

Key changes:
- Added struct to to represent a specialization constant with its ID,
type, and value.
- Updated YAML mapping in to parse specialization constants from the
test configuration.
- Modified to create and use when creating the compute pipeline,
allowing specialization constants to be passed to the shader.
- Added a new test case in to verify the functionality of specialization
constants with various data types (bool, int, uint, float).

Fixes llvm/llvm-project#142992

Author: s-perron

include/Support/Pipeline.h

@@ -282,11 +282,18 @@ struct IOBindings {
   }
 };
 
+struct SpecializationConstant {
+  uint32_t ConstantID;
+  DataFormat Type;
+  std::string Value;
+};
+
 struct Shader {
   Stages Stage;
   std::string Entry;
   std::unique_ptr<llvm::MemoryBuffer> Shader;
   int DispatchSize[3];
+  llvm::SmallVector<SpecializationConstant> SpecializationConstants;
 };
 
 struct Pipeline {
@@ -335,6 +342,7 @@ LLVM_YAML_IS_SEQUENCE_VECTOR(offloadtest::Shader)
 LLVM_YAML_IS_SEQUENCE_VECTOR(offloadtest::dx::RootParameter)
 LLVM_YAML_IS_SEQUENCE_VECTOR(offloadtest::Result)
 LLVM_YAML_IS_SEQUENCE_VECTOR(offloadtest::VertexAttribute)
+LLVM_YAML_IS_SEQUENCE_VECTOR(offloadtest::SpecializationConstant)
 
 namespace llvm {
 namespace yaml {
@@ -399,6 +407,10 @@ template <> struct MappingTraits<offloadtest::RuntimeSettings> {
   static void mapping(IO &I, offloadtest::RuntimeSettings &S);
 };
 
+template <> struct MappingTraits<offloadtest::SpecializationConstant> {
+  static void mapping(IO &I, offloadtest::SpecializationConstant &C);
+};
+
 template <> struct ScalarEnumerationTraits<offloadtest::Rule> {
   static void enumeration(IO &I, offloadtest::Rule &V) {
 #define ENUM_CASE(Val) I.enumCase(V, #Val, offloadtest::Rule::Val)

lib/API/VK/Device.cpp

@@ -11,6 +11,7 @@
 
 #include "API/Device.h"
 #include "Support/Pipeline.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/Support/Error.h"
 
 #include <memory>
@@ -20,22 +21,34 @@
 
 using namespace offloadtest;
 
-#define VKFormats(FMT)                                                         \
+#define VKFormats(FMT, BITS)                                                   \
   if (Channels == 1)                                                           \
-    return VK_FORMAT_R32_##FMT;                                                \
+    return VK_FORMAT_R##BITS##_##FMT;                                          \
   if (Channels == 2)                                                           \
-    return VK_FORMAT_R32G32_##FMT;                                             \
+    return VK_FORMAT_R##BITS##G##BITS##_##FMT;                                 \
   if (Channels == 3)                                                           \
-    return VK_FORMAT_R32G32B32_##FMT;                                          \
+    return VK_FORMAT_R##BITS##G##BITS##B##BITS##_##FMT;                        \
   if (Channels == 4)                                                           \
-    return VK_FORMAT_R32G32B32A32_##FMT;
+    return VK_FORMAT_R##BITS##G##BITS##B##BITS##A##BITS##_##FMT;
 
 static VkFormat getVKFormat(DataFormat Format, int Channels) {
   switch (Format) {
+  case DataFormat::Int16:
+    VKFormats(SINT, 16) break;
+  case DataFormat::UInt16:
+    VKFormats(UINT, 16) break;
   case DataFormat::Int32:
-    VKFormats(SINT) break;
+    VKFormats(SINT, 32) break;
+  case DataFormat::UInt32:
+    VKFormats(UINT, 32) break;
   case DataFormat::Float32:
-    VKFormats(SFLOAT) break;
+    VKFormats(SFLOAT, 32) break;
+  case DataFormat::Int64:
+    VKFormats(SINT, 64) break;
+  case DataFormat::UInt64:
+    VKFormats(UINT, 64) break;
+  case DataFormat::Float64:
+    VKFormats(SFLOAT, 64) break;
   default:
     llvm_unreachable("Unsupported Resource format specified");
   }
@@ -1273,6 +1286,105 @@ class VKDevice : public offloadtest::Device {
     return llvm::Error::success();
   }
 
+  static llvm::Error
+  parseSpecializationConstant(const SpecializationConstant &SpecConst,
+                              VkSpecializationMapEntry &Entry,
+                              llvm::SmallVector<char> &SpecData) {
+    Entry.constantID = SpecConst.ConstantID;
+    Entry.offset = SpecData.size();
+    switch (SpecConst.Type) {
+    case DataFormat::Float32: {
+      float Value = 0.0f;
+      double Tmp = 0.0;
+      if (llvm::StringRef(SpecConst.Value).getAsDouble(Tmp))
+        return llvm::createStringError(
+            std::errc::invalid_argument,
+            "Invalid float value for specialization constant '%s'",
+            SpecConst.Value.c_str());
+      Value = static_cast<float>(Tmp);
+      Entry.size = sizeof(float);
+      SpecData.resize(SpecData.size() + sizeof(float));
+      memcpy(SpecData.data() + Entry.offset, &Value, sizeof(float));
+      break;
+    }
+    case DataFormat::Float64: {
+      double Value = 0.0;
+      if (llvm::StringRef(SpecConst.Value).getAsDouble(Value))
+        return llvm::createStringError(
+            std::errc::invalid_argument,
+            "Invalid double value for specialization constant '%s'",
+            SpecConst.Value.c_str());
+      Entry.size = sizeof(double);
+      SpecData.resize(SpecData.size() + sizeof(double));
+      memcpy(SpecData.data() + Entry.offset, &Value, sizeof(double));
+      break;
+    }
+    case DataFormat::Int16: {
+      int16_t Value = 0;
+      if (llvm::StringRef(SpecConst.Value).getAsInteger(0, Value))
+        return llvm::createStringError(
+            std::errc::invalid_argument,
+            "Invalid int16 value for specialization constant '%s'",
+            SpecConst.Value.c_str());
+      Entry.size = sizeof(int16_t);
+      SpecData.resize(SpecData.size() + sizeof(int16_t));
+      memcpy(SpecData.data() + Entry.offset, &Value, sizeof(int16_t));
+      break;
+    }
+    case DataFormat::UInt16: {
+      uint16_t Value = 0;
+      if (llvm::StringRef(SpecConst.Value).getAsInteger(0, Value))
+        return llvm::createStringError(
+            std::errc::invalid_argument,
+            "Invalid uint16 value for specialization constant '%s'",
+            SpecConst.Value.c_str());
+      Entry.size = sizeof(uint16_t);
+      SpecData.resize(SpecData.size() + sizeof(uint16_t));
+      memcpy(SpecData.data() + Entry.offset, &Value, sizeof(uint16_t));
+      break;
+    }
+    case DataFormat::Int32: {
+      int32_t Value = 0;
+      if (llvm::StringRef(SpecConst.Value).getAsInteger(0, Value))
+        return llvm::createStringError(
+            std::errc::invalid_argument,
+            "Invalid int32 value for specialization constant '%s'",
+            SpecConst.Value.c_str());
+      Entry.size = sizeof(int32_t);
+      SpecData.resize(SpecData.size() + sizeof(int32_t));
+      memcpy(SpecData.data() + Entry.offset, &Value, sizeof(int32_t));
+      break;
+    }
+    case DataFormat::UInt32: {
+      uint32_t Value = 0;
+      if (llvm::StringRef(SpecConst.Value).getAsInteger(0, Value))
+        return llvm::createStringError(
+            std::errc::invalid_argument,
+            "Invalid uint32 value for specialization constant '%s'",
+            SpecConst.Value.c_str());
+      Entry.size = sizeof(uint32_t);
+      SpecData.resize(SpecData.size() + sizeof(uint32_t));
+      memcpy(SpecData.data() + Entry.offset, &Value, sizeof(uint32_t));
+      break;
+    }
+    case DataFormat::Bool: {
+      bool Value = false;
+      if (llvm::StringRef(SpecConst.Value).getAsInteger(0, Value))
+        return llvm::createStringError(
+            std::errc::invalid_argument,
+            "Invalid bool value for specialization constant '%s'",
+            SpecConst.Value.c_str());
+      Entry.size = sizeof(bool);
+      SpecData.resize(SpecData.size() + sizeof(bool));
+      memcpy(SpecData.data() + Entry.offset, &Value, sizeof(bool));
+      break;
+    }
+    default:
+      llvm_unreachable("Unsupported specialization constant type");
+    }
+    return llvm::Error::success();
+  }
+
   llvm::Error createPipeline(Pipeline &P, InvocationState &IS) {
     VkPipelineCacheCreateInfo CacheCreateInfo = {};
     CacheCreateInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_CACHE_CREATE_INFO;
@@ -1282,15 +1394,43 @@ class VKDevice : public offloadtest::Device {
                                      "Failed to create pipeline cache.");
 
     if (P.isCompute()) {
-      const CompiledShader &S = IS.Shaders[0];
+      const offloadtest::Shader &Shader = P.Shaders[0];
       assert(IS.Shaders.size() == 1 &&
              "Currently only support one compute shader");
+      const CompiledShader &S = IS.Shaders[0];
       VkPipelineShaderStageCreateInfo StageInfo = {};
       StageInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
       StageInfo.stage = VK_SHADER_STAGE_COMPUTE_BIT;
       StageInfo.module = S.Shader;
       StageInfo.pName = S.Entry.c_str();
 
+      llvm::SmallVector<VkSpecializationMapEntry> SpecEntries;
+      llvm::SmallVector<char> SpecData;
+      VkSpecializationInfo SpecInfo = {};
+      if (!Shader.SpecializationConstants.empty()) {
+        llvm::DenseSet<uint32_t> SeenConstantIDs;
+        for (const auto &SpecConst : Shader.SpecializationConstants) {
+          if (!SeenConstantIDs.insert(SpecConst.ConstantID).second)
+            return llvm::createStringError(
+                std::errc::invalid_argument,
+                "Test configuration contains multiple entries for "
+                "specialization constant ID %u.",
+                SpecConst.ConstantID);
+
+          VkSpecializationMapEntry Entry;
+          if (auto Err =
+                  parseSpecializationConstant(SpecConst, Entry, SpecData))
+            return Err;
+          SpecEntries.push_back(Entry);
+        }
+
+        SpecInfo.mapEntryCount = SpecEntries.size();
+        SpecInfo.pMapEntries = SpecEntries.data();
+        SpecInfo.dataSize = SpecData.size();
+        SpecInfo.pData = SpecData.data();
+        StageInfo.pSpecializationInfo = &SpecInfo;
+      }
+
       VkComputePipelineCreateInfo PipelineCreateInfo = {};
       PipelineCreateInfo.sType = VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO;
       PipelineCreateInfo.stage = StageInfo;

lib/Support/Pipeline.cpp

@@ -372,6 +372,7 @@ void MappingTraits<offloadtest::Shader>::mapping(IO &I,
                                                  offloadtest::Shader &S) {
   I.mapRequired("Stage", S.Stage);
   I.mapRequired("Entry", S.Entry);
+  I.mapOptional("SpecializationConstants", S.SpecializationConstants);
 
   if (S.Stage == Stages::Compute) {
     // Stage-specific data, not sure if this should be optional
@@ -380,6 +381,7 @@ void MappingTraits<offloadtest::Shader>::mapping(IO &I,
     I.mapRequired("DispatchSize", MutableDispatchSize);
   }
 }
+
 void MappingTraits<offloadtest::Result>::mapping(IO &I,
                                                  offloadtest::Result &R) {
   I.mapRequired("Result", R.Name);
@@ -402,5 +404,13 @@ void MappingTraits<offloadtest::Result>::mapping(IO &I,
     break;
   }
 }
+
+void MappingTraits<offloadtest::SpecializationConstant>::mapping(
+    IO &I, offloadtest::SpecializationConstant &C) {
+  I.mapRequired("ConstantID", C.ConstantID);
+  I.mapRequired("Type", C.Type);
+  I.mapRequired("Value", C.Value);
+}
+
 } // namespace yaml
 } // namespace llvm

test/Feature/SpecializationConstant/duplicate_spec_id.test

@@ -0,0 +1,39 @@
+#--- duplicate_spec_id.hlsl
+[[vk::constant_id(0)]]
+const int spec_int_A = 0;
+[[vk::constant_id(0)]]
+const int spec_int_B = 0;
+
+RWStructuredBuffer<int> Out : register(u0, space0);
+
+[numthreads(1,1,1)]
+void main(uint GI : SV_GroupIndex) {
+  Out[0] = spec_int_A;
+  Out[1] = spec_int_B;
+}
+#--- duplicate_spec_id.yaml
+---
+Shaders:
+  - Stage: Compute
+    Entry: main
+    DispatchSize: [1, 1, 1]
+    SpecializationConstants:
+      - { ConstantID: 0, Value: 123, Type: Int32 }
+Buffers:
+  - { Name: Out, Format: Int32, Stride: 4, FillSize: 8 }
+DescriptorSets:
+  - Resources:
+    - Name: Out
+      Kind: RWStructuredBuffer
+      DirectXBinding: { Register: 0, Space: 0 }
+      VulkanBinding: { Binding: 0 }
+...
+#--- end
+
+# REQUIRES: Vulkan
+
+# RUN: split-file %s %t
+# RUN: %dxc_target -T cs_6_2 -Fo %t.o %t/duplicate_spec_id.hlsl
+# RUN: %offloader %t/duplicate_spec_id.yaml %t.o | FileCheck %s
+
+# CHECK: Data: [ 123, 123 ]

test/Feature/SpecializationConstant/duplicate_spec_id_in_config.test

@@ -0,0 +1,37 @@
+#--- duplicate_spec_id_in_config.hlsl
+[[vk::constant_id(0)]]
+const int spec_int = 0;
+
+RWStructuredBuffer<int> Out : register(u0, space0);
+
+[numthreads(1,1,1)]
+void main(uint GI : SV_GroupIndex) {
+  Out[GI] = spec_int;
+}
+#--- duplicate_spec_id_in_config.yaml
+---
+Shaders:
+  - Stage: Compute
+    Entry: main
+    DispatchSize: [1, 1, 1]
+    SpecializationConstants:
+      - { ConstantID: 0, Value: 123, Type: Int32 }
+      - { ConstantID: 0, Value: 456, Type: Int32 }
+Buffers:
+  - { Name: Out, Format: Int32, Stride: 4, FillSize: 4 }
+DescriptorSets:
+  - Resources:
+    - Name: Out
+      Kind: RWStructuredBuffer
+      DirectXBinding: { Register: 0, Space: 0 }
+      VulkanBinding: { Binding: 0 }
+...
+#--- end
+
+# REQUIRES: Vulkan
+
+# RUN: split-file %s %t
+# RUN: %dxc_target -T cs_6_2 -Fo %t.o %t/duplicate_spec_id_in_config.hlsl
+# RUN: not %offloader %t/duplicate_spec_id_in_config.yaml %t.o 2>&1 | FileCheck %s
+
+# CHECK: gpu-exec: error: Test configuration contains multiple entries for specialization constant ID 0.

test/Feature/SpecializationConstant/invalid_bool.test

@@ -0,0 +1,35 @@
+#--- invalid_bool.hlsl
+[[vk::constant_id(0)]]
+const bool spec_bool = false;
+RWStructuredBuffer<uint> OutBool : register(u0, space0);
+
+[numthreads(1,1,1)]
+void main(uint GI : SV_GroupIndex) {
+  OutBool[GI] = (uint)spec_bool;
+}
+#--- invalid_bool.yaml
+---
+Shaders:
+  - Stage: Compute
+    Entry: main
+    DispatchSize: [1, 1, 1]
+    SpecializationConstants:
+      - { ConstantID: 0, Value: "not a number", Type: Bool }
+Buffers:
+  - { Name: OutBool, Format: Int32, Stride: 4, FillSize: 4 }
+DescriptorSets:
+  - Resources:
+    - Name: OutBool
+      Kind: RWStructuredBuffer
+      DirectXBinding: { Register: 0, Space: 0 }
+      VulkanBinding: { Binding: 0 }
+...
+#--- end
+
+# REQUIRES: Vulkan
+
+# RUN: split-file %s %t
+# RUN: %dxc_target -T cs_6_2 -Fo %t.o %t/invalid_bool.hlsl
+# RUN: not %offloader %t/invalid_bool.yaml %t.o 2>&1 | FileCheck %s
+
+# CHECK: Invalid bool value for specialization constant 'not a number'