[VULKAN] Add support for specialization constants

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

@@ -20,22 +20,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 +1285,76 @@ class VKDevice : public offloadtest::Device {
     return llvm::Error::success();
   }
 
+  static void
+  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;
+      llvm::StringRef(SpecConst.Value).getAsDouble(Tmp);
+      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;
+      llvm::StringRef(SpecConst.Value).getAsDouble(Value);
+      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;
+      llvm::StringRef(SpecConst.Value).getAsInteger(0, Value);
+      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;
+      llvm::StringRef(SpecConst.Value).getAsInteger(0, Value);
+      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;
+      llvm::StringRef(SpecConst.Value).getAsInteger(0, Value);
+      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;
+      llvm::StringRef(SpecConst.Value).getAsInteger(0, Value);
+      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;
+      llvm::StringRef(SpecConst.Value).getAsInteger(0, Value);
+      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");
+    }
+  }
+
   llvm::Error createPipeline(Pipeline &P, InvocationState &IS) {
     VkPipelineCacheCreateInfo CacheCreateInfo = {};
     CacheCreateInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_CACHE_CREATE_INFO;
@@ -1282,15 +1364,33 @@ 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()) {
+        for (const auto &SpecConst : Shader.SpecializationConstants) {
+          VkSpecializationMapEntry Entry;
+          parseSpecializationConstant(SpecConst, Entry, SpecData);
+          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/spec_const_32_bits.test

@@ -0,0 +1,87 @@
+#--- simple.hlsl
+// bool
+[[vk::constant_id(0)]]
+const bool spec_bool = false;
+RWBuffer<uint> OutBool : register(u0, space0);
+
+// int
+[[vk::constant_id(1)]]
+const int spec_int = 0;
+RWBuffer<int> OutInt : register(u1, space0);
+
+// unsigned int
+[[vk::constant_id(2)]]
+const uint spec_uint = 0;
+RWBuffer<uint> OutUInt : register(u2, space0);
+
+// float
+[[vk::constant_id(3)]]
+const float spec_float = 0.0;
+RWBuffer<float> OutFloat : register(u3, space0);
+
+// int with default value
+[[vk::constant_id(4)]]
+const int spec_int_default = 1234;
+RWBuffer<int> OutIntDefault : register(u4, space0);
+
+[numthreads(1,1,1)]
+void main(uint GI : SV_GroupIndex) {
+  OutBool[GI] = (uint)spec_bool;
+  OutInt[GI] = spec_int;
+  OutUInt[GI] = spec_uint;
+  OutFloat[GI] = spec_float;
+  OutIntDefault[GI] = spec_int_default;
+}
+#--- simple.yaml
+---
+Shaders:
+  - Stage: Compute
+    Entry: main
+    DispatchSize: [1, 1, 1]
+    SpecializationConstants:
+      - { ConstantID: 0, Value: 1, Type: Bool }
+      - { ConstantID: 1, Value: 42, Type: Int32 }
+      - { ConstantID: 2, Value: 0xDEADBEEF, Type: UInt32 }
+      - { ConstantID: 3, Value: 3.14, Type: Float32 }
+Buffers:
+  - { Name: OutBool, Format: Int32, FillSize: 4 }
+  - { Name: OutInt,  Format: Int32, FillSize: 4 }
+  - { Name: OutUInt, Format: UInt32, FillSize: 4 }
+  - { Name: OutFloat, Format: Float32, FillSize: 4 }
+  - { Name: OutIntDefault, Format: Int32, FillSize: 4 }
+DescriptorSets:
+  - Resources:
+    - Name: OutBool
+      Kind: RWBuffer
+      DirectXBinding: { Register: 0, Space: 0 }
+      VulkanBinding: { Binding: 0 }
+    - Name: OutInt
+      Kind: RWBuffer
+      DirectXBinding: { Register: 1, Space: 0 }
+      VulkanBinding: { Binding: 1 }
+    - Name: OutUInt
+      Kind: RWBuffer
+      DirectXBinding: { Register: 2, Space: 0 }
+      VulkanBinding: { Binding: 2 }
+    - Name: OutFloat
+      Kind: RWBuffer
+      DirectXBinding: { Register: 3, Space: 0 }
+      VulkanBinding: { Binding: 3 }
+    - Name: OutIntDefault
+      Kind: RWBuffer
+      DirectXBinding: { Register: 4, Space: 0 }
+      VulkanBinding: { Binding: 4 }
+...
+#--- end
+
+# REQUIRES: Vulkan
+
+# RUN: split-file %s %t
+# RUN: %dxc_target -T cs_6_0 -Fo %t.o %t/simple.hlsl
+# RUN: %offloader %t/simple.yaml %t.o | FileCheck %s
+
+# CHECK: Data: [ 1 ]
+# CHECK: Data: [ 42 ]
+# CHECK: Data: [ 3735928559 ]
+# CHECK: Data: [ {{3.14.*}} ]
+# CHECK: Data: [ 1234 ]

test/Feature/SpecializationConstant/spec_const_other_sizes.test

@@ -0,0 +1,64 @@
+#--- simple_64bit.hlsl
+// double
+[[vk::constant_id(0)]]
+const double spec_double = 0.0;
+RWStructuredBuffer<double> OutDouble : register(u0, space0);
+
+// short
+[[vk::constant_id(1)]]
+const int16_t spec_short = 0;
+RWStructuredBuffer<int16_t> OutShort : register(u1, space0);
+
+// ushort
+[[vk::constant_id(2)]]
+const uint16_t spec_ushort = 0;
+RWStructuredBuffer<uint16_t> OutUShort : register(u2, space0);
+
+[numthreads(1,1,1)]
+void main(uint GI : SV_GroupIndex) {
+  OutDouble[GI] = spec_double;
+  OutShort[GI] = spec_short;
+  OutUShort[GI] = spec_ushort;
+}
+#--- simple_64bit.yaml
+---
+Shaders:
+  - Stage: Compute
+    Entry: main
+    DispatchSize: [1, 1, 1]
+    SpecializationConstants:
+      - { ConstantID: 0, Value: 2.718, Type: Float64 }
+      - { ConstantID: 1, Value: 123, Type: Int16 }
+      - { ConstantID: 2, Value: 456, Type: UInt16 }
+Buffers:
+  - { Name: OutDouble, Format: Float64, Stride: 8, FillSize: 8 }
+  - { Name: OutShort, Format: Int16, Stride: 2, FillSize: 2 }
+  - { Name: OutUShort, Format: UInt16, Stride: 2, FillSize: 2 }
+DescriptorSets:
+  - Resources:
+    - Name: OutDouble
+      Kind: RWStructuredBuffer
+      DirectXBinding: { Register: 0, Space: 0 }
+      VulkanBinding: { Binding: 0 }
+    - Name: OutShort
+      Kind: RWStructuredBuffer
+      DirectXBinding: { Register: 1, Space: 0 }
+      VulkanBinding: { Binding: 1 }
+    - Name: OutUShort
+      Kind: RWStructuredBuffer
+      DirectXBinding: { Register: 2, Space: 0 }
+      VulkanBinding: { Binding: 2 }
+...
+#--- end
+
+# REQUIRES: Vulkan
+
+# XFAIL: DXC
+
+# RUN: split-file %s %t
+# RUN: %dxc_target -T cs_6_2 -enable-16bit-types -Fo %t.o %t/simple_64bit.hlsl
+# RUN: %offloader %t/simple_64bit.yaml %t.o | FileCheck %s
+
+# CHECK: Data: [ {{2.718}} ]
+# CHECK: Data: [ 123 ]
+# CHECK: Data: [ 456 ]