[SPIR-V] Add support for D3D12 descriptor heaps (microsoft#6650)

This commit adds support for Sampler/Resource descriptor heaps in DXC.
Support for those heaps on the SPIR-V side requires no other extension
than SPV_EXT_descriptor_indexing.
On the Vulkan side, the VK_EXT_mutable_descriptor_type will be required
as multiple descriptor types must be allowed on the same binding.

When loading a type from a heap, DXC generates a new OpRuntimeArray of
the correct type, and binds it to `set=0,
binding=<BindingNumberOfTheHeap>`.
This means multiple OpRuntimeArrays will share the same binding. This is
why VK_EXT_mutable_descriptor_type is required.

This implementation uses at most 3 bindings:
 - N OpRuntimeArray as binding A for the ResourceDescriptorHeap
 - N OpRuntimeArray as binding B for the SamplerDescriptorHeap
- 1 OpRuntimeArray %counter_type for the ResourceDescriptorHeap
counters.

The bindings are only allocated if used. If only the
SamplerDescriptorHeap is used, a single binding is required.

The binding allocation logic is:
 1. allocate bindings for every resources, excluding heaps.
2. If ResourceDescriptorHeap is used, find the first unused binding in
set=0 and use it.
 3. Same for the SamplerDescriptorHeap
 4. Same for the counters.

UAV counters are not always created, only if used.
When used, they are stored in an OpRuntimeArray. The index of a counter
in that array
is equal to the index of the associated resource in its own
OpRuntimeArray.

```hlsl
        RWStructuredBuffer a = ResourceDescriptorHeap[2];
        a.IncrementCounter();
        // buffer in descriptorSet 0, binding 0, OpRuntimeArray[index=2]
        // counter in descriptorSet 0, binding 1, OpRuntimeArray[index=2]
```

As-is, this PR doesn't allow resource heaps to alias regular resources,
or to overlap.
A follow-up PR will add 3 flags to override each binding/set pairs:
 - 'fvk-bind-resource-heap <set> <binding>'
 - 'fvk-bind-sampler-heap <set> <binding>'
 - 'fvk-bind-counter-heap <set> <binding>'

---------

Signed-off-by: Nathan Gauër <[email protected]>

Author: Keenuts

docs/SPIR-V.rst

@@ -1721,6 +1721,22 @@ attached to a RW/append/consume structured buffers to specify the binding number
 for the associated counter to ``Z``. Note that the set number of the counter is
 always the same as the main buffer.
 
+.. warning::
+   When a RW/append/consume structured buffer is accessed through a resource
+   heap, its associated counter is in its own binding, but shares the same
+   index in the binding as its associated resource.
+
+   Example:
+    - ResourceDescriptorHeap -> binding 0, set 0
+    - No other resources are used.
+
+    - RWStructuredBuffer buff = ResourceDescriptorHeap[3]
+    - buff.IncrementCounter()
+
+    - buff will be at index 3 of the array at binding 0, set 0.
+      buff.counter will be at index 3 of the array at binding 1, set 0
+
+
 Implicit binding number assignment
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
@@ -1923,6 +1939,81 @@ Example 3: (compiled with ``-fvk-bind-globals 2 1``)
 Note that if the developer chooses to use this command line option, it is their
 responsibility to provide proper numbers and avoid binding overlaps.
 
+ResourceDescriptorHeaps & SamplerDescriptorHeaps
+------------------------------------------------
+
+The SPIR-V backend supported SM6.6 resource heaps, using 2 extensions:
+- `SPV_EXT_descriptor_indexing`
+- `VK_EXT_mutable_descriptor_type`
+
+Each type loaded from a heap is considered to be an unbounded RuntimeArray
+bound to the descriptor set 0.
+
+Each heap uses at most 1 binding in that set. Meaning if 2 types are loaded
+from the same heap, DXC will generate 2 RuntimeArray, one for each type, and
+will bind them to the same binding/set.
+(This requires `VK_EXT_mutable_descriptor_type`).
+
+For resources with counters, like RW/Append/Consume structured buffers,
+DXC generates another RuntimeArray of counters, and binds it to a new
+binding in the set 0.
+
+This means Resource/Sampler heaps can use at most 3 bindings:
+    - 1 for all RuntimeArrays associated with the ResourceDescriptorHeap.
+    - 1 for all RuntimeArrays associated with the SamplerDescriptorHeaps.
+    - 1 for UAV counters.
+
+The index of a counter in the counters RuntimeArray matches the index of the
+associated ResourceDescriptorHeap RuntimeArray.
+
+The selection of the binding indices for those RuntimeArrays is done once all
+other resources are bound to their respective bindings/sets.
+DXC takes the first 3 unused bindings in the set 0, and distributes them in
+that order:
+    1. Resource heap.
+    2. Sampler heap.
+    3. Resouce heap counters.
+
+Bindings are lazily allocated: if only the sampler heap is used,
+1 binding will be used.
+
+.. code:: hlsl
+   Texture2D tex = ResourceDescriptorHeap[10];
+   // tex is in the descriptor set 0, binding 0.
+
+.. code:: hlsl
+   [[vk::binding(0, 0)]]
+   Texture2D Texture;
+   // Texture is using set=0, binding=0
+
+   Texture2D tex = ResourceDescriptorHeap[0];
+   // tex is in the descriptor set 0, binding 1.
+
+.. code:: hlsl
+   [[vk::binding(0, 0)]]
+   RWStructuredBuffer<int> buffer;
+   // Texture is using set=0, binding=0
+
+   RWStructuredBuffer<int> tmp = ResourceDescriptorHeap[0];
+   tmp.IncrementCounter();
+   // tmp is in the descriptor set 0, binding 1.
+   // tmp.counter is in the descriptor set 0, binding 2
+
+.. code:: hlsl
+   [[vk::binding(1, 0)]]
+   RWStructuredBuffer<int> buffer;
+   // Texture is using set=0, binding=1
+
+   RWStructuredBuffer<int> tmp = ResourceDescriptorHeap[0];
+   tmp.IncrementCounter();
+   // tmp is in the descriptor set 0, binding 0.
+   // tmp.counter is in the descriptor set 0, binding 2
+
+.. code:: hlsl
+   RWStructuredBuffer buffer = ResourceDescriptorHeap[2];
+   // buffer is in the descriptor set 0, binding 0.
+   // Counter not generated, because unused.
+
 HLSL Expressions
 ================
 

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

@@ -221,6 +221,12 @@ bool isRWStructuredBuffer(QualType type);
 /// StructuredBuffer type.
 bool isRWAppendConsumeSBuffer(QualType type);
 
+/// \brief Returns true if the given type is a ResourceDescriptorHeap.
+bool isResourceDescriptorHeap(QualType type);
+
+/// \brief Returns true if the given type is a SamplerDescriptorHeap.
+bool isSamplerDescriptorHeap(QualType type);
+
 /// \brief Returns true if the given type is the HLSL ByteAddressBufferType.
 bool isByteAddressBuffer(QualType type);
 

tools/clang/lib/SPIRV/AstTypeProbe.cpp

@@ -995,6 +995,20 @@ 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 isSamplerDescriptorHeap(QualType type) {
+  if (const auto *rt = type->getAs<RecordType>()) {
+    return rt->getDecl()->getName() == ".Sampler";
+  }
+  return false;
+}
+
 bool isAKindOfStructuredOrByteBuffer(QualType type) {
   // Strip outer arrayness first
   while (type->isArrayType())
@@ -1007,7 +1021,8 @@ bool isAKindOfStructuredOrByteBuffer(QualType type) {
            name == "ByteAddressBuffer" || name == "RWByteAddressBuffer" ||
            name == "RasterizerOrderedByteAddressBuffer" ||
            name == "AppendStructuredBuffer" ||
-           name == "ConsumeStructuredBuffer";
+           name == "ConsumeStructuredBuffer" || name == ".Resource" ||
+           name == ".Sampler";
   }
   return false;
 }
@@ -1022,7 +1037,8 @@ bool isOrContainsAKindOfStructuredOrByteBuffer(QualType type) {
         name == "RasterizerOrderedStructuredBuffer" ||
         name == "ByteAddressBuffer" || name == "RWByteAddressBuffer" ||
         name == "RasterizerOrderedByteAddressBuffer" ||
-        name == "AppendStructuredBuffer" || name == "ConsumeStructuredBuffer")
+        name == "AppendStructuredBuffer" || name == "ConsumeStructuredBuffer" ||
+        name == ".Resource" || name == ".Sampler")
       return true;
 
     for (const auto *field : recordType->getDecl()->fields()) {

tools/clang/lib/SPIRV/DeclResultIdMapper.cpp

@@ -976,14 +976,6 @@ SpirvInstruction *DeclResultIdMapper::getDeclEvalInfo(const ValueDecl *decl,
                          decl->getType(), spv::StorageClass::Output, loc);
   }
 
-  if (hlsl::IsHLSLDynamicResourceType(decl->getType()) ||
-      hlsl::IsHLSLDynamicSamplerType(decl->getType())) {
-    emitError("HLSL object %0 not yet supported with -spirv",
-              decl->getLocation())
-        << decl->getName();
-    return nullptr;
-  }
-
   const DeclSpirvInfo *info = getDeclSpirvInfo(decl);
 
   // If DeclSpirvInfo is not found for this decl, it might be because it is an
@@ -1150,8 +1142,19 @@ DeclResultIdMapper::createFileVar(const VarDecl *var,
   return varInstr;
 }
 
+SpirvVariable *DeclResultIdMapper::createResourceHeap(const VarDecl *var,
+                                                      QualType ResourceType) {
+  QualType ResourceArrayType = astContext.getIncompleteArrayType(
+      ResourceType, clang::ArrayType::Normal, 0);
+  return createExternVar(var, ResourceArrayType);
+}
+
 SpirvVariable *DeclResultIdMapper::createExternVar(const VarDecl *var) {
-  const auto type = var->getType();
+  return createExternVar(var, var->getType());
+}
+
+SpirvVariable *DeclResultIdMapper::createExternVar(const VarDecl *var,
+                                                   QualType type) {
   const bool isGroupShared = var->hasAttr<HLSLGroupSharedAttr>();
   const bool isACSBuffer =
       isAppendStructuredBuffer(type) || isConsumeStructuredBuffer(type);
@@ -1262,12 +1265,9 @@ SpirvVariable *DeclResultIdMapper::createExternVar(const VarDecl *var) {
   if (storageClass == spv::StorageClass::Workgroup)
     return varInstr;
 
-  const auto *regAttr = getResourceBinding(var);
   const auto *bindingAttr = var->getAttr<VKBindingAttr>();
-  const auto *counterBindingAttr = var->getAttr<VKCounterBindingAttr>();
-
-  resourceVars.emplace_back(varInstr, var, loc, regAttr, bindingAttr,
-                            counterBindingAttr);
+  resourceVars.emplace_back(varInstr, var, loc, getResourceBinding(var),
+                            bindingAttr, var->getAttr<VKCounterBindingAttr>());
 
   if (const auto *inputAttachment = var->getAttr<VKInputAttachmentIndexAttr>())
     spvBuilder.decorateInputAttachmentIndex(varInstr,
@@ -1819,6 +1819,10 @@ void DeclResultIdMapper::createCounterVar(
       assert(declType->isIncompleteArrayType());
       counterType = spvContext.getRuntimeArrayType(counterType, arrayStride);
     }
+  } else if (isResourceDescriptorHeap(decl->getType()) ||
+             isSamplerDescriptorHeap(decl->getType())) {
+    counterType =
+        spvContext.getRuntimeArrayType(counterType, /* arrayStride= */ 4);
   }
 
   // {RW|Append|Consume}StructuredBuffer are all in Uniform storage class.
@@ -2002,9 +2006,11 @@ class LocationSet {
   uint32_t nextAvailableLocation[kMaxIndex];
 };
 
+} // namespace
+
 /// A class for managing resource bindings to avoid duplicate uses of the same
 /// set and binding number.
-class BindingSet {
+class DeclResultIdMapper::BindingSet {
 public:
   /// Uses the given set and binding number. Returns false if the binding number
   /// was already occupied in the set, and returns true otherwise.
@@ -2070,7 +2076,6 @@ class BindingSet {
   ///< set number -> set of used binding number
   llvm::DenseMap<uint32_t, std::set<uint32_t>> usedBindings;
 };
-} // namespace
 
 bool DeclResultIdMapper::checkSemanticDuplication(bool forInput) {
   // Mapping from entry points to the corresponding set of semantics.
@@ -2619,6 +2624,13 @@ bool DeclResultIdMapper::decorateResourceBindings() {
       }
     }
 
+    if (var.getDeclaration()) {
+      const VarDecl *decl = dyn_cast<VarDecl>(var.getDeclaration());
+      if (decl && (isResourceDescriptorHeap(decl->getType()) ||
+                   isSamplerDescriptorHeap(decl->getType())))
+        continue;
+    }
+
     if (var.isCounter()) {
 
       if (!var.getCounterBinding()) {
@@ -2673,9 +2685,61 @@ bool DeclResultIdMapper::decorateResourceBindings() {
     }
   }
 
+  decorateResourceHeapsBindings(bindingSet);
   return true;
 }
 
+void DeclResultIdMapper::decorateResourceHeapsBindings(BindingSet &bindingSet) {
+  bool hasResource = false;
+  bool hasSamplers = false;
+  bool hasCounters = false;
+
+  // Determine which type of heap resource is used to lazily allocation
+  // bindings.
+  for (const auto &var : resourceVars) {
+    if (!var.getDeclaration())
+      continue;
+    const VarDecl *decl = dyn_cast<VarDecl>(var.getDeclaration());
+    if (!decl)
+      continue;
+
+    const bool isResourceHeap = isResourceDescriptorHeap(decl->getType());
+    const bool isSamplerHeap = isSamplerDescriptorHeap(decl->getType());
+
+    assert(!(var.isCounter() && isSamplerHeap));
+
+    hasResource |= isResourceHeap;
+    hasSamplers |= isSamplerHeap;
+    hasCounters |= isResourceHeap && var.isCounter();
+  }
+
+  // Allocate bindings only for used resources. The order of this allocation is
+  // important:
+  //  - First resource heaps, then sampler heaps, and finally counter heaps.
+  const uint32_t resourceBinding =
+      hasResource ? bindingSet.useNextBinding(0) : 0;
+  const uint32_t samplersBinding =
+      hasSamplers ? bindingSet.useNextBinding(0) : 0;
+  const uint32_t countersBinding =
+      hasCounters ? bindingSet.useNextBinding(0) : 0;
+
+  for (const auto &var : resourceVars) {
+    if (!var.getDeclaration())
+      continue;
+    const VarDecl *decl = dyn_cast<VarDecl>(var.getDeclaration());
+    if (!decl)
+      continue;
+
+    if (isResourceDescriptorHeap(decl->getType()))
+      spvBuilder.decorateDSetBinding(var.getSpirvInstr(), /* set= */ 0,
+                                     var.isCounter() ? countersBinding
+                                                     : resourceBinding);
+    else if (isSamplerDescriptorHeap(decl->getType()))
+      spvBuilder.decorateDSetBinding(var.getSpirvInstr(), /* set= */ 0,
+                                     samplersBinding);
+  }
+}
+
 bool DeclResultIdMapper::decorateResourceCoherent() {
   for (const auto &var : resourceVars) {
     if (const auto *decl = var.getDeclaration()) {

tools/clang/lib/SPIRV/DeclResultIdMapper.h

@@ -189,6 +189,9 @@ class CounterVarFields {
 /// type, the fields with attached semantics will need to be translated into
 /// stage variables per Vulkan's requirements.
 class DeclResultIdMapper {
+  /// \brief An internal class to handle binding number allocation.
+  class BindingSet;
+
 public:
   inline DeclResultIdMapper(ASTContext &context, SpirvContext &spirvContext,
                             SpirvBuilder &spirvBuilder, SpirvEmitter &emitter,
@@ -277,9 +280,17 @@ class DeclResultIdMapper {
   SpirvVariable *createFileVar(const VarDecl *var,
                                llvm::Optional<SpirvInstruction *> init);
 
+  /// Creates a global variable for resource heaps containing elements of type
+  /// |type|.
+  SpirvVariable *createResourceHeap(const VarDecl *var, QualType type);
+
   /// \brief Creates an external-visible variable and returns its instruction.
   SpirvVariable *createExternVar(const VarDecl *var);
 
+  /// \brief Creates an external-visible variable of type |type| and returns its
+  /// instruction.
+  SpirvVariable *createExternVar(const VarDecl *var, QualType type);
+
   /// \brief Returns an OpString instruction that represents the given VarDecl.
   /// VarDecl must be a variable of string type.
   ///
@@ -637,6 +648,10 @@ class DeclResultIdMapper {
       llvm::DenseSet<StageVariableLocationInfo, StageVariableLocationInfo>
           *stageVariableLocationInfo);
 
+  /// \bried Decorates used Resource/Sampler descriptor heaps with the correct
+  /// binding/set decorations.
+  void decorateResourceHeapsBindings(BindingSet &bindingSet);
+
   /// \brief Returns a map that divides all of the shader stage variables into
   /// separate vectors for each entry point.
   llvm::DenseMap<const SpirvFunction *, SmallVector<StageVar, 8>>