[spirv] Add support for [[vk::push_constant]] (#791)

[[vk::push_constant]] can be attached to a global variable of
struct type to put that variable in the PushConstant storage
class.

PushConstant should be of OpTypeStruct type with std430 layout.

docs/SPIR-V.rst

@@ -152,13 +152,35 @@ Shader Stage Input Output
 
 More details in the `gl_PerVertex`_ section.
 
+Vulkan specific features
+------------------------
+
+We try to implement Vulkan specific features using the most intuitive and
+non-intrusive ways in HLSL, which means we will prefer native language
+constructs when possible. If that is inadequate, we then consider attaching
+`Vulkan specific attributes`_ to them, or introducing new syntax.
+
+Descriptor sets
+~~~~~~~~~~~~~~~
+
+To specify which Vulkan descriptor a particular resource binds to, use the
+``[[vk::binding(X[, Y])]]`` attribute.
+
+Push constants
+~~~~~~~~~~~~~~
+
+Vulkan push constant blocks are represented using normal global variables of
+struct types in HLSL. The variables (not the underlying struct types) should be
+annotated with the ``[[vk::push_constant]]`` attribute.
+
+Please note as per the requirements of Vulkan, "there must be no more than one
+push constant block statically used per shader entry point."
+
 Vulkan specific attributes
 --------------------------
 
-To provide additional information required by Vulkan in HLSL, we need to extend
-the syntax of HLSL.
 `C++ attribute specifier sequence <http://en.cppreference.com/w/cpp/language/attributes>`_
-is a non-intrusive way of achieving such purpose.
+is a non-intrusive way of providing Vulkan specific information in HLSL.
 
 The namespace ``vk`` will be used for all Vulkan attributes:
 
@@ -171,6 +193,9 @@ The namespace ``vk`` will be used for all Vulkan attributes:
 - ``counter_binding(X)``: For specifying the binding number (``X``) for the
   associated counter for RW/Append/Consume structured buffer. The descriptor
   set number for the associated counter is always the same as the main resource.
+- ``push_constant``: For marking a variable as the push constant block. Allowed
+  on global variables of struct type. At most one variable can be marked as
+  ``push_constant`` in a shader.
 
 Only ``vk::`` attributes in the above list are supported. Other attributes will
 result in warnings and be ignored by the compiler. All C++11 attributes will

tools/clang/include/clang/Basic/Attr.td

@@ -879,7 +879,7 @@ def VKBinding : InheritableAttr {
   let Documentation = [Undocumented];
 }
 
-// StructuredBuffer types that can have associated counters.
+// StructuredBuffer types that can have associated counters
 def CounterStructuredBuffer : SubsetSubject<
     Var,
     [{S->hasGlobalStorage() && S->getType()->getAs<RecordType>() &&
@@ -895,6 +895,17 @@ def VKCounterBinding : InheritableAttr {
   let Documentation = [Undocumented];
 }
 
+// Global variables that are of struct type
+def StructGlobalVar : SubsetSubject<Var, [{S->hasGlobalStorage() && S->getType()->isStructureType()}]>;
+
+def VKPushConstant : InheritableAttr {
+  let Spellings = [CXX11<"vk", "push_constant">];
+  let Subjects = SubjectList<[StructGlobalVar], ErrorDiag, "ExpectedStructGlobalVar">;
+  let Args = [];
+  let LangOpts = [SPIRV];
+  let Documentation = [Undocumented];
+}
+
 // SPIRV Change Ends
 
 def C11NoReturn : InheritableAttr {

tools/clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -2326,6 +2326,7 @@ def warn_attribute_wrong_decl_type : Warning<
   "functions and global variables|structs, unions, and typedefs|structs and typedefs|"
   "interface or protocol declarations|kernel functions|"
   // SPIRV Change Starts
+  "global variables of struct type|"
   "global variables, cbuffers, and tbuffers|"
   "RWStructuredBuffers, AppendStructuredBuffers, and ConsumeStructuredBuffers|"
   // SPIRV Change Ends

tools/clang/include/clang/Sema/AttributeList.h

@@ -856,6 +856,7 @@ enum AttributeDeclKind {
   ExpectedObjectiveCInterfaceOrProtocol,
   ExpectedKernelFunction
   // SPIRV Change Begins
+  ,ExpectedStructGlobalVar
   ,ExpectedGlobalVarOrCTBuffer
   ,ExpectedCounterStructuredBuffer
   // SPIRV Change Ends

tools/clang/lib/SPIRV/DeclResultIdMapper.cpp

@@ -289,20 +289,24 @@ uint32_t DeclResultIdMapper::createExternVar(const VarDecl *var) {
 }
 
 uint32_t DeclResultIdMapper::createVarOfExplicitLayoutStruct(
-    const DeclContext *decl, llvm::StringRef typeName, llvm::StringRef varName,
-    bool isCBuffer) {
-  LayoutRule layoutRule =
-      isCBuffer ? LayoutRule::GLSLStd140 : LayoutRule::GLSLStd430;
-  const auto *blockDec =
-      isCBuffer ? Decoration::getBlock(*theBuilder.getSPIRVContext())
-                : Decoration::getBufferBlock(*theBuilder.getSPIRVContext());
-
-  // Get the type for the whole buffer
+    const DeclContext *decl, const ContextUsageKind usageKind,
+    llvm::StringRef typeName, llvm::StringRef varName) {
   // cbuffers are translated into OpTypeStruct with Block decoration.
   // tbuffers are translated into OpTypeStruct with BufferBlock decoration.
+  // PushConstants are translated into OpTypeStruct with Block decoration.
+  //
   // Both cbuffers and tbuffers have the SPIR-V Uniform storage class. cbuffers
-  // satisfy GLSL std140 layout rules, and tbuffers satisfy GLSL std430 layout
-  // rules.
+  // follow GLSL std140 layout rules, and tbuffers follow GLSL std430 layout
+  // rules. PushConstants follow GLSL std430 layout rules.
+
+  auto &context = *theBuilder.getSPIRVContext();
+  const LayoutRule layoutRule = usageKind == ContextUsageKind::CBuffer
+                                    ? LayoutRule::GLSLStd140
+                                    : LayoutRule::GLSLStd430;
+  const auto *blockDec = usageKind == ContextUsageKind::TBuffer
+                             ? Decoration::getBufferBlock(context)
+                             : Decoration::getBlock(context);
+
   auto decorations = typeTranslator.getLayoutDecorations(decl, layoutRule);
   decorations.push_back(blockDec);
 
@@ -330,26 +334,32 @@ uint32_t DeclResultIdMapper::createVarOfExplicitLayoutStruct(
 
     // tbuffer/TextureBuffers are non-writable SSBOs. OpMemberDecorate
     // NonWritable must be applied to all fields.
-    if (!isCBuffer) {
+    if (usageKind == ContextUsageKind::TBuffer) {
       decorations.push_back(Decoration::getNonWritable(
           *theBuilder.getSPIRVContext(), fieldIndex));
     }
     ++fieldIndex;
   }
 
+  // Get the type for the whole struct
   const uint32_t structType =
       theBuilder.getStructType(fieldTypes, typeName, fieldNames, decorations);
 
-  // Create the variable for the whole buffer
-  return theBuilder.addModuleVar(structType, spv::StorageClass::Uniform,
-                                 varName);
+  const auto sc = usageKind == ContextUsageKind::PushConstant
+                      ? spv::StorageClass::PushConstant
+                      : spv::StorageClass::Uniform;
+
+  // Create the variable for the whole struct
+  return theBuilder.addModuleVar(structType, sc, varName);
 }
 
 uint32_t DeclResultIdMapper::createCTBuffer(const HLSLBufferDecl *decl) {
+  const auto usageKind =
+      decl->isCBuffer() ? ContextUsageKind::CBuffer : ContextUsageKind::TBuffer;
   const std::string structName = "type." + decl->getName().str();
   const std::string varName = "var." + decl->getName().str();
-  const uint32_t bufferVar = createVarOfExplicitLayoutStruct(
-      decl, structName, varName, decl->isCBuffer());
+  const uint32_t bufferVar =
+      createVarOfExplicitLayoutStruct(decl, usageKind, structName, varName);
 
   // We still register all VarDecls seperately here. All the VarDecls are
   // mapped to the <result-id> of the buffer object, which means when querying
@@ -374,18 +384,20 @@ uint32_t DeclResultIdMapper::createCTBuffer(const VarDecl *decl) {
   const auto *recordType = decl->getType()->getAs<RecordType>();
   assert(recordType);
   const auto *context = cast<HLSLBufferDecl>(decl->getDeclContext());
-  const bool isCBuffer = context->isCBuffer();
+  const auto usageKind = context->isCBuffer() ? ContextUsageKind::CBuffer
+                                              : ContextUsageKind::TBuffer;
 
   const std::string structName =
-      "type." + std::string(isCBuffer ? "ConstantBuffer." : "TextureBuffer.") +
+      "type." +
+      std::string(context->isCBuffer() ? "ConstantBuffer." : "TextureBuffer.") +
       recordType->getDecl()->getName().str();
   const uint32_t bufferVar = createVarOfExplicitLayoutStruct(
-      recordType->getDecl(), structName, decl->getName(), isCBuffer);
+      recordType->getDecl(), usageKind, structName, decl->getName());
 
   // We register the VarDecl here.
   astDecls[decl] = {bufferVar, spv::StorageClass::Uniform,
-                    isCBuffer ? LayoutRule::GLSLStd140
-                              : LayoutRule::GLSLStd430};
+                    context->isCBuffer() ? LayoutRule::GLSLStd140
+                                         : LayoutRule::GLSLStd430};
   resourceVars.emplace_back(
       bufferVar, ResourceVar::Category::Other, getResourceBinding(context),
       decl->getAttr<VKBindingAttr>(), decl->getAttr<VKCounterBindingAttr>());
@@ -393,6 +405,25 @@ uint32_t DeclResultIdMapper::createCTBuffer(const VarDecl *decl) {
   return bufferVar;
 }
 
+uint32_t DeclResultIdMapper::createPushConstant(const VarDecl *decl) {
+  const auto *recordType = decl->getType()->getAs<RecordType>();
+  assert(recordType);
+
+  const std::string structName =
+      "type.PushConstant." + recordType->getDecl()->getName().str();
+  const uint32_t var = createVarOfExplicitLayoutStruct(
+      recordType->getDecl(), ContextUsageKind::PushConstant, structName,
+      decl->getName());
+
+  // Register the VarDecl
+  astDecls[decl] = {var, spv::StorageClass::PushConstant,
+                    LayoutRule::GLSLStd430};
+  // Do not push this variable into resourceVars since it does not need
+  // descriptor set.
+
+  return var;
+}
+
 uint32_t DeclResultIdMapper::getOrRegisterFnResultId(const FunctionDecl *fn) {
   if (const auto *info = getDeclSpirvInfo(fn))
     return info->resultId;

tools/clang/lib/SPIRV/DeclResultIdMapper.h

@@ -203,6 +203,9 @@ public:
   /// VarDecl does not need an extra OpAccessChain.
   uint32_t createCTBuffer(const VarDecl *decl);
 
+  /// \brief Creates a PushConstant block from the given decl.
+  uint32_t createPushConstant(const VarDecl *decl);
+
   /// \brief Sets the <result-id> of the entry function.
   void setEntryFunctionId(uint32_t id) { entryFunctionId = id; }
 
@@ -282,8 +285,7 @@ private:
   /// \brief Wrapper method to create an error message and report it
   /// in the diagnostic engine associated with this consumer.
   template <unsigned N>
-  DiagnosticBuilder emitError(const char (&message)[N],
-                              SourceLocation loc) {
+  DiagnosticBuilder emitError(const char (&message)[N], SourceLocation loc) {
     const auto diagId =
         diags.getCustomDiagID(clang::DiagnosticsEngine::Error, message);
     return diags.Report(loc, diagId);
@@ -301,20 +303,27 @@ private:
   /// construction.
   bool finalizeStageIOLocations(bool forInput);
 
+  /// \brief An enum class for representing what the DeclContext is used for
+  enum class ContextUsageKind {
+    CBuffer,
+    TBuffer,
+    PushConstant,
+  };
+
   /// Creates a variable of struct type with explicit layout decorations.
   /// The sub-Decls in the given DeclContext will be treated as the struct
   /// fields. The struct type will be named as typeName, and the variable
   /// will be named as varName.
   ///
-  /// This method should only be used for cbuffers/ContantBuffers and
-  /// tbuffers/TextureBuffers. isCBuffer must be set appropriately based on the
-  /// type of the buffer.
+  /// This method should only be used for cbuffers/ContantBuffers, tbuffers/
+  /// TextureBuffers, and PushConstants. usageKind must be set properly
+  /// depending on the usage kind.
   ///
   /// Panics if the DeclContext is neither HLSLBufferDecl or RecordDecl.
   uint32_t createVarOfExplicitLayoutStruct(const DeclContext *decl,
+                                           ContextUsageKind usageKind,
                                            llvm::StringRef typeName,
-                                           llvm::StringRef varName,
-                                           bool isCBuffer);
+                                           llvm::StringRef varName);
 
   /// Creates all the stage variables mapped from semantics on the given decl.
   /// Returns true on sucess.

tools/clang/lib/SPIRV/SPIRVEmitter.cpp

@@ -295,12 +295,7 @@ void SPIRVEmitter::HandleTranslationUnit(ASTContext &context) {
         patchConstFunc = funcDecl;
       }
     } else if (auto *varDecl = dyn_cast<VarDecl>(decl)) {
-      if (isa<HLSLBufferDecl>(varDecl->getDeclContext())) {
-        // This is a VarDecl of a ConstantBuffer/TextureBuffer type.
-        (void)declIdMapper.createCTBuffer(varDecl);
-      } else {
-        doVarDecl(varDecl);
-      }
+      doVarDecl(varDecl);
     } else if (auto *bufferDecl = dyn_cast<HLSLBufferDecl>(decl)) {
       // This is a cbuffer/tbuffer decl.
 
@@ -637,6 +632,18 @@ void SPIRVEmitter::doFunctionDecl(const FunctionDecl *decl) {
 }
 
 void SPIRVEmitter::doVarDecl(const VarDecl *decl) {
+  if (decl->hasAttr<VKPushConstantAttr>()) {
+    // This is a VarDecl for PushConstant block.
+    (void)declIdMapper.createPushConstant(decl);
+    return;
+  }
+
+  if (isa<HLSLBufferDecl>(decl->getDeclContext())) {
+    // This is a VarDecl of a ConstantBuffer/TextureBuffer type.
+    (void)declIdMapper.createCTBuffer(decl);
+    return;
+  }
+
   uint32_t varId = 0;
 
   // The contents in externally visible variables can be updated via the

tools/clang/lib/Sema/SemaHLSL.cpp

@@ -10267,6 +10267,10 @@ void hlsl::HandleDeclAttributeForHLSL(Sema &S, Decl *D, const AttributeList &A,
     declAttr = ::new (S.Context) VKCounterBindingAttr(A.getRange(), S.Context,
       ValidateAttributeIntArg(S, A), A.getAttributeSpellingListIndex());
     break;
+  case AttributeList::AT_VKPushConstant:
+    declAttr = ::new (S.Context) VKPushConstantAttr(A.getRange(), S.Context,
+      A.getAttributeSpellingListIndex());
+    break;
   default:
     Handled = false;
     return;

tools/clang/test/CodeGenSPIRV/vk.attribute.error.hlsl

@@ -28,6 +28,20 @@ float4 main([[vk::binding(15)]] float4 a: A // error
  return 1.0;
 }
 
+struct T {
+    [[vk::push_constant]] // error
+    float4 f;
+};
+
+[[vk::push_constant]] // error
+float foo([[vk::push_constant]] int param) // error
+{
+    return param;
+}
+
+[[vk::push_constant(5)]]
+T pcs;
+
 // CHECK:   :4:7: error: 'binding' attribute only applies to global variables, cbuffers, and tbuffers
 // CHECK:   :6:7: error: 'counter_binding' attribute only applies to RWStructuredBuffers, AppendStructuredBuffers, and ConsumeStructuredBuffers
 // CHECK:  :10:3: error: 'counter_binding' attribute only applies to RWStructuredBuffers, AppendStructuredBuffers, and ConsumeStructuredBuffers
@@ -37,3 +51,7 @@ float4 main([[vk::binding(15)]] float4 a: A // error
 // CHECK:  :20:3: error: 'location' attribute only applies to functions, parameters, and fields
 // CHECK: :26:15: error: 'binding' attribute only applies to global variables, cbuffers, and tbuffers
 // CHECK:  :25:3: error: 'binding' attribute only applies to global variables, cbuffers, and tbuffers
+// CHECK:  :32:7: error: 'push_constant' attribute only applies to global variables of struct type
+// CHECK: :37:13: error: 'push_constant' attribute only applies to global variables of struct type
+// CHECK:  :36:3: error: 'push_constant' attribute only applies to global variables of struct type
+// CHECK:  :42:3: error: 'push_constant' attribute takes no arguments

tools/clang/test/CodeGenSPIRV/vk.layout.push-constant.std430.hlsl

@@ -0,0 +1,35 @@
+// Run: %dxc -T vs_6_0 -E main
+
+// CHECK: OpDecorate %_arr_v2float_uint_3 ArrayStride 8
+// CHECK: OpDecorate %_arr_mat3v2float_uint_2 ArrayStride 32
+
+// CHECK: OpMemberDecorate %T 0 Offset 0
+// CHECK: OpMemberDecorate %T 1 Offset 32
+// CHECK: OpMemberDecorate %T 1 MatrixStride 16
+// CHECK: OpMemberDecorate %T 1 RowMajor
+struct T {
+                 float2   f1[3];
+    column_major float3x2 f2[2];
+};
+
+// CHECK: OpMemberDecorate %type_PushConstant_S 0 Offset 0
+// CHECK: OpMemberDecorate %type_PushConstant_S 1 Offset 16
+// CHECK: OpMemberDecorate %type_PushConstant_S 2 Offset 32
+// CHECK: OpMemberDecorate %type_PushConstant_S 3 Offset 128
+// CHECK: OpMemberDecorate %type_PushConstant_S 3 MatrixStride 16
+// CHECK: OpMemberDecorate %type_PushConstant_S 3 ColMajor
+
+// CHECK: OpDecorate %type_PushConstant_S Block
+struct S {
+              float    f1;
+              float3   f2;
+              T        f4;
+    row_major float2x3 f3;
+};
+
+[[vk::push_constant]]
+S pcs;
+
+float main() : A {
+    return pcs.f1;
+}

tools/clang/test/CodeGenSPIRV/vk.push-constant.hlsl

@@ -0,0 +1,38 @@
+// Run: %dxc -T vs_6_0 -E main
+
+struct T {
+    float2 val[3];
+};
+
+// CHECK: OpName %type_PushConstant_S "type.PushConstant.S"
+// CHECK: OpMemberName %type_PushConstant_S 0 "f1"
+// CHECK: OpMemberName %type_PushConstant_S 1 "f2"
+// CHECK: OpMemberName %type_PushConstant_S 2 "f3"
+// CHECK: OpMemberName %type_PushConstant_S 3 "f4"
+
+// CHECK: %type_PushConstant_S = OpTypeStruct %float %v3float %mat2v3float %T
+struct S {
+    float    f1;
+    float3   f2;
+    float2x3 f3;
+    T        f4;
+};
+// CHECK: %_ptr_PushConstant_type_PushConstant_S = OpTypePointer PushConstant %type_PushConstant_S
+
+// CHECK: %pcs = OpVariable %_ptr_PushConstant_type_PushConstant_S PushConstant
+[[vk::push_constant]]
+S pcs;
+
+float main() : A {
+    return
+// CHECK:     {{%\d+}} = OpAccessChain %_ptr_PushConstant_float %pcs %int_0
+        pcs.f1 +
+// CHECK: [[ptr:%\d+]] = OpAccessChain %_ptr_PushConstant_v3float %pcs %int_1
+// CHECK:     {{%\d+}} = OpAccessChain %_ptr_PushConstant_float [[ptr]] %int_2
+        pcs.f2.z +
+// CHECK:     {{%\d+}} = OpAccessChain %_ptr_PushConstant_float %pcs %int_2 %uint_1 %uint_2
+        pcs.f3[1][2] +
+// CHECK: [[ptr:%\d+]] = OpAccessChain %_ptr_PushConstant_v2float %pcs %int_3 %int_0 %int_2
+// CHECK:     {{%\d+}} = OpAccessChain %_ptr_PushConstant_float [[ptr]] %int_1
+        pcs.f4.val[2].y;
+}

tools/clang/unittests/SPIRV/CodeGenSPIRVTest.cpp

@@ -913,6 +913,8 @@ TEST_F(FileTest, VulkanStructuredBufferCounter) {
   runFileTest("vk.binding.counter.hlsl");
 }
 
+TEST_F(FileTest, VulkanPushConstant) { runFileTest("vk.push-constant.hlsl"); }
+
 TEST_F(FileTest, VulkanLayoutCBufferStd140) {
   runFileTest("vk.layout.cbuffer.std140.hlsl");
 }
@@ -941,6 +943,10 @@ TEST_F(FileTest, VulkanLayoutTextureBufferStd430) {
   runFileTest("vk.layout.texture-buffer.std430.hlsl");
 }
 
+TEST_F(FileTest, VulkanLayoutPushConstantStd430) {
+  runFileTest("vk.layout.push-constant.std430.hlsl");
+}
+
 // HS: for different Patch Constant Functions
 TEST_F(FileTest, HullShaderPCFVoid) { runFileTest("hs.pcf.void.hlsl"); }
 TEST_F(FileTest, HullShaderPCFTakesInputPatch) {