[SPIR-V] Fix resource heap & fvk-bind-register interactions (microsoft#7858)

The interaction between `-fvk-bind-register` and the resource/sampler
heaps were not tested, and broken.

This commit solves this by requiring the `-fvk-bind-sampler-heap` or
`-fvk-bind-resource-heap` flag to be used if the `-fvk-bind-register`
flag is present and a resource/sampler heap is present.

An alternative solution would be to allow implicit bindings on heaps
while other resources require the register attribute, but I think it
would be less useful as the goal of using this flag is to have full
control over the bindings.

This commit refactorizes a bit the helper functions
isResourceHeap/isSamplerHeap, but this part is NFC. The main bit is in
DeclResultIdMapper, where a special case is added the the heaps.

Fixes microsoft#7857

Author: Keenuts

tools/clang/include/clang/SPIRV/AstTypeProbe.h

@@ -222,10 +222,12 @@ bool isRWStructuredBuffer(QualType type);
 bool isRWAppendConsumeSBuffer(QualType type);
 
 /// \brief Returns true if the given type is a ResourceDescriptorHeap.
-bool isResourceDescriptorHeap(QualType type);
+bool isResourceDescriptorHeap(QualType T);
+bool isResourceDescriptorHeap(const Decl *D);
 
 /// \brief Returns true if the given type is a SamplerDescriptorHeap.
-bool isSamplerDescriptorHeap(QualType type);
+bool isSamplerDescriptorHeap(const Decl *D);
+bool isSamplerDescriptorHeap(QualType T);
 
 /// \brief Returns true if the given type is the HLSL ByteAddressBufferType.
 bool isByteAddressBuffer(QualType type);

tools/clang/lib/SPIRV/AstTypeProbe.cpp

@@ -995,18 +995,24 @@ bool isRWAppendConsumeSBuffer(QualType type) {
          isAppendStructuredBuffer(type);
 }
 
-bool isResourceDescriptorHeap(QualType type) {
-  if (const auto *rt = type->getAs<RecordType>()) {
-    return rt->getDecl()->getName() == ".Resource";
-  }
-  return false;
+bool isResourceDescriptorHeap(const Decl *D) {
+  const VarDecl *VD = dyn_cast<VarDecl>(D);
+  return VD && isResourceDescriptorHeap(VD->getType());
 }
 
-bool isSamplerDescriptorHeap(QualType type) {
-  if (const auto *rt = type->getAs<RecordType>()) {
-    return rt->getDecl()->getName() == ".Sampler";
-  }
-  return false;
+bool isResourceDescriptorHeap(QualType T) {
+  const RecordType *RT = T->getAs<RecordType>();
+  return RT && RT->getDecl()->getName() == ".Resource";
+}
+
+bool isSamplerDescriptorHeap(const Decl *D) {
+  const VarDecl *VD = dyn_cast<VarDecl>(D);
+  return VD && isSamplerDescriptorHeap(VD->getType());
+}
+
+bool isSamplerDescriptorHeap(QualType T) {
+  const RecordType *RT = T->getAs<RecordType>();
+  return RT && RT->getDecl()->getName() == ".Sampler";
 }
 
 bool isAKindOfStructuredOrByteBuffer(QualType type) {

tools/clang/lib/SPIRV/DeclResultIdMapper.cpp

@@ -1850,8 +1850,7 @@ void DeclResultIdMapper::createCounterVar(
       assert(declType->isIncompleteArrayType());
       counterType = spvContext.getRuntimeArrayType(counterType, noArrayStride);
     }
-  } else if (isResourceDescriptorHeap(decl->getType()) ||
-             isSamplerDescriptorHeap(decl->getType())) {
+  } else if (isResourceDescriptorHeap(decl) || isSamplerDescriptorHeap(decl)) {
     counterType = spvContext.getRuntimeArrayType(counterType, noArrayStride);
   }
 
@@ -2517,18 +2516,33 @@ bool DeclResultIdMapper::decorateResourceBindings() {
 
         spvBuilder.decorateDSetBinding(var.getSpirvInstr(), globalsSetNo,
                                        globalsBindNo);
+      } else if (var.isResourceDescriptorHeap()) {
+        if (!spirvOptions.resourceHeapBinding) {
+          emitError("-fvk-bind-resource-heap is required when using "
+                    "-fvk-bind-register",
+                    var.getSourceLocation());
+          return false;
+        }
+      } else if (var.isSamplerDescriptorHeap()) {
+        if (!spirvOptions.samplerHeapBinding) {
+          emitError("-fvk-bind-sampler-heap is required when using "
+                    "-fvk-bind-register",
+                    var.getSourceLocation());
+          return false;
+        }
       } else {
         emitError(
             "-fvk-bind-register requires register annotations on all resources",
             var.getSourceLocation());
         return false;
       }
 
+    BindingSet bindingSet;
+    decorateResourceHeapsBindings(bindingSet);
     return true;
   }
 
   BindingSet bindingSet;
-
   // If some bindings are reserved for heaps, mark those are used.
   if (spirvOptions.resourceHeapBinding)
     bindingSet.useBinding(spirvOptions.resourceHeapBinding->binding,
@@ -2669,8 +2683,8 @@ bool DeclResultIdMapper::decorateResourceBindings() {
 
     if (var.getDeclaration()) {
       const VarDecl *decl = dyn_cast<VarDecl>(var.getDeclaration());
-      if (decl && (isResourceDescriptorHeap(decl->getType()) ||
-                   isSamplerDescriptorHeap(decl->getType())))
+      if (decl &&
+          (isResourceDescriptorHeap(decl) || isSamplerDescriptorHeap(decl)))
         continue;
     }
 
@@ -2755,8 +2769,8 @@ void DeclResultIdMapper::decorateResourceHeapsBindings(BindingSet &bindingSet) {
     if (!decl)
       continue;
 
-    const bool isResourceHeap = isResourceDescriptorHeap(decl->getType());
-    const bool isSamplerHeap = isSamplerDescriptorHeap(decl->getType());
+    const bool isResourceHeap = isResourceDescriptorHeap(decl);
+    const bool isSamplerHeap = isSamplerDescriptorHeap(decl);
 
     assert(!(var.isCounter() && isSamplerHeap));
 
@@ -2791,8 +2805,8 @@ void DeclResultIdMapper::decorateResourceHeapsBindings(BindingSet &bindingSet) {
     if (!decl)
       continue;
 
-    const bool isResourceHeap = isResourceDescriptorHeap(decl->getType());
-    const bool isSamplerHeap = isSamplerDescriptorHeap(decl->getType());
+    const bool isResourceHeap = isResourceDescriptorHeap(decl);
+    const bool isSamplerHeap = isSamplerDescriptorHeap(decl);
     if (!isSamplerHeap && !isResourceHeap)
       continue;
     const SpirvCodeGenOptions::BindingInfo &info =

tools/clang/lib/SPIRV/DeclResultIdMapper.h

@@ -16,6 +16,7 @@
 #include "dxc/Support/SPIRVOptions.h"
 #include "spirv/unified1/spirv.hpp11"
 #include "clang/AST/Attr.h"
+#include "clang/SPIRV/AstTypeProbe.h"
 #include "clang/SPIRV/FeatureManager.h"
 #include "clang/SPIRV/SpirvBuilder.h"
 #include "llvm/ADT/ArrayRef.h"
@@ -33,7 +34,7 @@ class SpirvEmitter;
 
 class ResourceVar {
 public:
-  ResourceVar(SpirvVariable *var, const Decl *decl, SourceLocation loc,
+  ResourceVar(SpirvVariable *var, const NamedDecl *decl, SourceLocation loc,
               const hlsl::RegisterAssignment *r, const VKBindingAttr *b,
               const VKCounterBindingAttr *cb, bool counter = false,
               bool globalsBuffer = false)
@@ -52,9 +53,17 @@ class ResourceVar {
     return counterBinding;
   }
 
+  bool isResourceDescriptorHeap() const {
+    return spirv::isResourceDescriptorHeap(declaration);
+  }
+
+  bool isSamplerDescriptorHeap() const {
+    return spirv::isSamplerDescriptorHeap(declaration);
+  }
+
 private:
   SpirvVariable *variable;                    ///< The variable
-  const Decl *declaration;                    ///< The declaration
+  const NamedDecl *declaration;               ///< The declaration
   SourceLocation srcLoc;                      ///< Source location
   const hlsl::RegisterAssignment *reg;        ///< HLSL register assignment
   const VKBindingAttr *binding;               ///< Vulkan binding assignment

tools/clang/test/CodeGenSPIRV/vk.binding.cl.register.heap.hlsl

@@ -0,0 +1,13 @@
+// RUN: %dxc -T ps_6_8 -E main -fvk-bind-register s10 0 1 2 -fcgl  -fvk-bind-resource-heap 3 4 %s -spirv | FileCheck %s
+
+// CHECK: OpDecorate %Sampler DescriptorSet 2
+// CHECK: OpDecorate %Sampler Binding 1
+SamplerState Sampler : register(s10);
+
+// CHECK: OpDecorate %ResourceDescriptorHeap DescriptorSet 4
+// CHECK: OpDecorate %ResourceDescriptorHeap Binding 3
+
+float4 main() : SV_Target {
+  Texture2D Texture = ResourceDescriptorHeap[0];
+  return Texture.Sample(Sampler, float2(0.1, 0.2));
+}

tools/clang/test/CodeGenSPIRV/vk.binding.cl.register.missing-heap.hlsl

@@ -0,0 +1,10 @@
+// RUN: not %dxc -T ps_6_8 -E main -fvk-bind-register s10 0 10 0 -fcgl  %s -spirv  2>&1 | FileCheck %s
+
+SamplerState Sampler : register(s10);
+
+float4 main() : SV_Target {
+  Texture2D Texture = ResourceDescriptorHeap[0];
+  return Texture.Sample(Sampler, float2(0.1, 0.2));
+}
+
+// CHECK: error: -fvk-bind-resource-heap is required when using -fvk-bind-register

tools/clang/test/CodeGenSPIRV/vk.binding.cl.register.missing-sampler.hlsl

@@ -0,0 +1,10 @@
+// RUN: not %dxc -T ps_6_8 -E main -fvk-bind-register s10 0 10 0 -fcgl  %s -spirv  2>&1 | FileCheck %s
+
+Texture2D Texture : register(s10);
+
+float4 main() : SV_Target {
+  SamplerState Sampler = SamplerDescriptorHeap[0];
+  return Texture.Sample(Sampler, float2(0.1, 0.2));
+}
+
+// CHECK: error: -fvk-bind-sampler-heap is required when using -fvk-bind-register

tools/clang/test/CodeGenSPIRV/vk.binding.cl.register.sampler.hlsl

@@ -0,0 +1,13 @@
+// RUN: %dxc -T ps_6_8 -E main -fvk-bind-register s10 0 1 2 -fcgl  -fvk-bind-sampler-heap 3 4 %s -spirv | FileCheck %s
+
+// CHECK: OpDecorate %Texture DescriptorSet 2
+// CHECK: OpDecorate %Texture Binding 1
+Texture2D Texture : register(s10);
+
+// CHECK: OpDecorate %SamplerDescriptorHeap DescriptorSet 4
+// CHECK: OpDecorate %SamplerDescriptorHeap Binding 3
+
+float4 main() : SV_Target {
+  SamplerState Sampler = SamplerDescriptorHeap[0];
+  return Texture.Sample(Sampler, float2(0.1, 0.2));
+}