[spirv] Validate vk::push_constant usage (#807)

Attaching vk::push_constant and vk::binding to the same global
variable is disallowed.

tools/clang/lib/SPIRV/SPIRVEmitter.cpp

@@ -273,7 +273,8 @@ SPIRVEmitter::SPIRVEmitter(CompilerInstance &ci,
       theContext(), theBuilder(&theContext),
       declIdMapper(shaderModel, astContext, theBuilder, spirvOptions),
       typeTranslator(astContext, theBuilder, diags), entryFunctionId(0),
-      curFunction(nullptr), curThis(0), needsLegalization(false) {
+      curFunction(nullptr), curThis(0), seenPushConstantAt(),
+      needsLegalization(false) {
   if (shaderModel.GetKind() == hlsl::ShaderModel::Kind::Invalid)
     emitError("unknown shader module: %0", {}) << shaderModel.GetName();
 }
@@ -310,6 +311,8 @@ void SPIRVEmitter::HandleTranslationUnit(ASTContext &context) {
                 << bufferDecl->isCBuffer() << init->getSourceRange();
       }
 
+      validateVKAttributes(decl);
+
       (void)declIdMapper.createCTBuffer(bufferDecl);
     }
   }
@@ -631,7 +634,42 @@ void SPIRVEmitter::doFunctionDecl(const FunctionDecl *decl) {
   curFunction = nullptr;
 }
 
+void SPIRVEmitter::validateVKAttributes(const Decl *decl) {
+  // The frontend will make sure that
+  // * vk::push_constant applies to global variables of struct type
+  // * vk::binding applies to global variables or cbuffers/tbuffers
+  // * vk::counter_binding applies to global variables of RW/Append/Consume
+  //   StructuredBuffer
+  // * vk::location applies to function parameters/returns and struct fields
+  // So the only case we need to check co-existence is vk::push_constant and
+  // vk::binding.
+
+  if (const auto *pcAttr = decl->getAttr<VKPushConstantAttr>()) {
+    if (seenPushConstantAt.isInvalid()) {
+      seenPushConstantAt = pcAttr->getLocation();
+    } else {
+      // TODO: Actually this is slightly incorrect. The Vulkan spec says:
+      //   There must be no more than one push constant block statically used
+      //   per shader entry point.
+      // But we are checking whether there are more than one push constant
+      // blocks defined. Tracking usage requires more work.
+      emitError("cannot have more than one push constant block",
+                pcAttr->getLocation());
+      emitNote("push constant block previously defined here",
+               seenPushConstantAt);
+    }
+
+    if (decl->hasAttr<VKBindingAttr>()) {
+      emitError("'push_constant' attribute cannot be used together with "
+                "'binding' attribute",
+                pcAttr->getLocation());
+    }
+  }
+}
+
 void SPIRVEmitter::doVarDecl(const VarDecl *decl) {
+  validateVKAttributes(decl);
+
   if (decl->hasAttr<VKPushConstantAttr>()) {
     // This is a VarDecl for PushConstant block.
     (void)declIdMapper.createPushConstant(decl);

tools/clang/lib/SPIRV/SPIRVEmitter.h

@@ -238,6 +238,10 @@ private:
   collectArrayStructIndices(const Expr *expr,
                             llvm::SmallVectorImpl<uint32_t> *indices);
 
+private:
+  /// Validates that vk::* attributes are used correctly.
+  void validateVKAttributes(const Decl *decl);
+
 private:
   /// Processes the given expr, casts the result into the given bool (vector)
   /// type and returns the <result-id> of the casted value.
@@ -675,6 +679,15 @@ private:
     return diags.Report(loc, diagId);
   }
 
+  /// \brief Wrapper method to create a note message and report it
+  /// in the diagnostic engine associated with this consumer
+  template <unsigned N>
+  DiagnosticBuilder emitNote(const char (&message)[N], SourceLocation loc) {
+    const auto diagId =
+        diags.getCustomDiagID(clang::DiagnosticsEngine::Note, message);
+    return diags.Report(loc, diagId);
+  }
+
 private:
   CompilerInstance &theCompilerInstance;
   ASTContext &astContext;
@@ -706,6 +719,10 @@ private:
   /// The SPIR-V function parameter for the current this object.
   uint32_t curThis;
 
+  /// The source location of a push constant block we have previously seen.
+  /// Invalid means no push constant blocks defined thus far.
+  SourceLocation seenPushConstantAt;
+
   /// Whether the translated SPIR-V binary needs legalization.
   ///
   /// The following cases will require legalization:

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

@@ -0,0 +1,14 @@
+// Run: %dxc -T ps_6_0 -E main
+
+struct S {
+    float f;
+};
+
+[[vk::push_constant, vk::binding(5)]]
+S pcs;
+
+float main() : A {
+    return 1.0;
+}
+
+// CHECK: :7:3: error: 'push_constant' attribute cannot be used together with 'binding' attribute

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

@@ -0,0 +1,24 @@
+// Run: %dxc -T vs_6_0 -E main
+
+struct S {
+    float4 f;
+};
+
+[[vk::push_constant]]
+S pcs1;
+
+[[vk::push_constant]] // error
+S pcs2;
+
+[[vk::push_constant]] // error
+S pcs3;
+
+float main() : A {
+    return 1.0;
+}
+
+// CHECK: :10:3: error: cannot have more than one push constant block
+// CHECK: :7:3: note: push constant block previously defined here
+
+// CHECK: :13:3: error: cannot have more than one push constant block
+// CHECK: :7:3: note: push constant block previously defined here

tools/clang/unittests/SPIRV/CodeGenSPIRVTest.cpp

@@ -872,6 +872,9 @@ TEST_F(FileTest, SpirvInterpolationError) {
 TEST_F(FileTest, VulkanAttributeErrors) {
   runFileTest("vk.attribute.error.hlsl", FileTest::Expect::Failure);
 }
+TEST_F(FileTest, VulkanAttributeInvalidUsages) {
+  runFileTest("vk.attribute.invalid.hlsl", FileTest::Expect::Failure);
+}
 
 // Vulkan specific
 TEST_F(FileTest, VulkanLocation) { runFileTest("vk.location.hlsl"); }
@@ -924,6 +927,9 @@ TEST_F(FileTest, VulkanStructuredBufferCounter) {
 }
 
 TEST_F(FileTest, VulkanPushConstant) { runFileTest("vk.push-constant.hlsl"); }
+TEST_F(FileTest, VulkanMultiplePushConstant) {
+  runFileTest("vk.push-constant.multiple.hlsl", FileTest::Expect::Failure);
+}
 
 TEST_F(FileTest, VulkanLayoutCBufferStd140) {
   runFileTest("vk.layout.cbuffer.std140.hlsl");