[SPIR-V] Add descriptor heap RaytracingAccelerationStructure support

docs/SPIR-V.rst

@@ -336,6 +336,8 @@ Supported extensions
 * SPV_KHR_float_controls
 * SPV_NV_shader_subgroup_partitioned
 * SPV_KHR_quad_control
+* SPV_KHR_untyped_pointers
+* SPV_EXT_descriptor_heap
 
 Vulkan specific attributes
 --------------------------
@@ -1993,10 +1995,14 @@ responsibility to provide proper numbers and avoid binding overlaps.
 ResourceDescriptorHeaps & SamplerDescriptorHeaps
 ------------------------------------------------
 
-The SPIR-V backend supported SM6.6 resource heaps, using 2 extensions:
+By default, the SPIR-V backend supports SM6.6 resource heaps by emulating the
+heaps with descriptor-indexing runtime arrays, using 2 extensions:
+
 - `SPV_EXT_descriptor_indexing`
 - `VK_EXT_mutable_descriptor_type`
 
+This is also the behavior selected by ``-fspv-use-emulated-heap``.
+
 Each type loaded from a heap is considered to be an unbounded RuntimeArray
 bound to the descriptor set 0.
 
@@ -2074,6 +2080,85 @@ Bindings & sets associated with each heap can be explicitly set using:
 - `-fvk-bind-counter-heap <binding> <set>`: Specify Vulkan binding number
   and set number for the counter heap.
 
+Native descriptor heap extension lowering
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+When ``-fspv-use-descriptor-heap`` is specified, DXC lowers
+``ResourceDescriptorHeap`` and ``SamplerDescriptorHeap`` through
+``SPV_EXT_descriptor_heap`` instead of the default emulated heap path. This
+also requires ``SPV_KHR_untyped_pointers`` and ``-fspv-target-env=vulkan1.3``
+(targeting a lower environment is an error), and a SPIRV-Headers / SPIRV-Tools
+build that defines these extensions. The emitted module declares the heap
+objects as untyped variables in ``UniformConstant`` storage class:
+
+.. code:: spirv
+
+  %uptr_uc = OpTypeUntypedPointerKHR UniformConstant
+  %resource_heap = OpUntypedVariableKHR %uptr_uc UniformConstant
+  %sampler_heap  = OpUntypedVariableKHR %uptr_uc UniformConstant
+  OpDecorate %resource_heap BuiltIn ResourceHeapEXT
+  OpDecorate %sampler_heap BuiltIn SamplerHeapEXT
+
+The concrete descriptor type is selected at each heap access. For image,
+sampler, and texel buffer resources, DXC forms a runtime array of that
+descriptor type and decorates the array with a byte ``ArrayStride``, then uses
+``OpUntypedAccessChainKHR`` followed by ``OpLoad``:
+
+.. code:: spirv
+
+  %image_type = OpTypeImage %float 2D 2 0 0 1 Unknown
+  %image_array = OpTypeRuntimeArray %image_type
+  OpDecorate %image_array ArrayStride 64
+  %descriptor = OpUntypedAccessChainKHR %uptr_uc %image_array %resource_heap %index
+  %image = OpLoad %image_type %descriptor
+
+For buffer-like resources, DXC uses ``OpTypeBufferEXT`` as the descriptor type
+and ``OpBufferPointerEXT`` to recover the pointer to the buffer data. The
+descriptor storage class matches the recovered buffer pointer storage class; for
+example, ``ConstantBuffer<T>`` uses ``Uniform`` and ``TextureBuffer<T>`` uses
+``StorageBuffer``:
+
+.. code:: spirv
+
+  %buffer_type = OpTypeBufferEXT Uniform
+  %buffer_array = OpTypeRuntimeArray %buffer_type
+  OpDecorate %buffer_array ArrayStride 64
+  %descriptor = OpUntypedAccessChainKHR %uptr_uc %buffer_array %resource_heap %index
+  %buffer_ptr = OpBufferPointerEXT %_ptr_Uniform_type_BufferData %descriptor
+
+For ``RWTexture`` resources loaded from ``ResourceDescriptorHeap``, interlocked
+operations that need a texel pointer use ``OpUntypedImageTexelPointerEXT``.
+The image descriptor pointer produced by ``OpUntypedAccessChainKHR`` is passed
+directly to the texel-pointer instruction instead of first storing the image
+handle into a function-scope image variable:
+
+.. code:: spirv
+
+  %image_type = OpTypeImage %uint 2D 2 0 0 2 R32ui
+  %image_array = OpTypeRuntimeArray %image_type
+  %descriptor = OpUntypedAccessChainKHR %uptr_uc %image_array %resource_heap %index
+  %uptr_image = OpTypeUntypedPointerKHR Image
+  %texel_ptr = OpUntypedImageTexelPointerEXT %uptr_image %image_type %descriptor %coord %sample
+  %old = OpAtomicIAdd %uint %texel_ptr %scope %semantics %value
+
+This path supports texture, RWTexture, sampler, Buffer/RWBuffer,
+StructuredBuffer/RWStructuredBuffer without associated counter operations,
+ByteAddressBuffer/RWByteAddressBuffer, ConstantBuffer, and TextureBuffer heap
+loads, including direct field and array-element accesses for
+``ConstantBuffer<T>`` and ``TextureBuffer<T>``. ``NonUniformResourceIndex`` is
+accepted but the ``NonUniform`` decoration is not emitted on
+``OpUntypedAccessChainKHR`` or the loaded value; ``SPV_EXT_descriptor_heap``
+deprecates the ``NonUniform`` decoration for heap accesses.
+
+Append/consume structured buffers and UAV counter heap lowering are not
+supported by the native descriptor heap path yet. Those forms should continue
+to use the default emulated heap lowering, or DXC will emit a diagnostic for
+unsupported append/consume structured-buffer heap loads. Heap-loaded
+``RWStructuredBuffer`` resources are supported for ordinary data access, but
+associated counter operations such as ``IncrementCounter`` and
+``DecrementCounter`` emit a diagnostic because the native descriptor heap path
+does not recover an associated counter descriptor.
+
 HLSL Expressions
 ================
 

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

@@ -250,6 +250,10 @@ bool isBuffer(QualType type);
 /// \brief Returns true if the given type is the HLSL RWBuffer type.
 bool isRWBuffer(QualType type);
 
+/// \brief Returns true if the given type is the HLSL
+/// RaytracingAccelerationStructure type.
+bool isRaytracingAccelerationStructure(QualType type);
+
 /// \brief Returns true if the given type is an HLSL Texture type.
 bool isTexture(QualType);
 

tools/clang/include/clang/SPIRV/SpirvBuilder.h

@@ -277,7 +277,7 @@ class SpirvBuilder {
   /// \brief Creates an OpUntypedImageTexelPointerEXT SPIR-V instruction with
   /// the given parameters.
   SpirvUntypedImageTexelPointerEXT *createUntypedImageTexelPointerEXT(
-      QualType resultType, SpirvInstruction *image,
+      QualType resultType, const SpirvType *imageType, SpirvInstruction *image,
       SpirvInstruction *coordinate, SpirvInstruction *sample, SourceLocation);
 
   /// \brief Creates an OpConverPtrToU SPIR-V instruction with the given

tools/clang/include/clang/SPIRV/SpirvContext.h

@@ -298,6 +298,7 @@ class SpirvContext {
                                          spv::StorageClass);
 
   const UntypedPointerKHRType *getUntypedPointerKHRType(spv::StorageClass sc);
+  const BufferEXTType *getBufferEXTType(spv::StorageClass sc);
 
   FunctionType *getFunctionType(const SpirvType *ret,
                                 llvm::ArrayRef<const SpirvType *> param);
@@ -536,6 +537,9 @@ class SpirvContext {
   llvm::DenseMap<spv::StorageClass, const UntypedPointerKHRType *,
                  StorageClassDenseMapInfo>
       untypedPointerKHRTypes;
+  llvm::DenseMap<spv::StorageClass, const BufferEXTType *,
+                 StorageClassDenseMapInfo>
+      bufferEXTTypes;
   llvm::MapVector<QualType, const ForwardPointerType *> forwardPointerTypes;
   llvm::MapVector<QualType, const SpirvPointerType *> forwardReferences;
   llvm::DenseSet<FunctionType *, FunctionTypeMapInfo> functionTypes;

tools/clang/include/clang/SPIRV/SpirvInstruction.h

@@ -2065,6 +2065,7 @@ class SpirvImageTexelPointer : public SpirvInstruction {
 class SpirvUntypedImageTexelPointerEXT : public SpirvInstruction {
 public:
   SpirvUntypedImageTexelPointerEXT(QualType resultType, SourceLocation loc,
+                                   const SpirvType *imageType,
                                    SpirvInstruction *image,
                                    SpirvInstruction *coordinate,
                                    SpirvInstruction *sample);
@@ -2078,11 +2079,22 @@ class SpirvUntypedImageTexelPointerEXT : public SpirvInstruction {
 
   bool invokeVisitor(Visitor *v) override;
 
+  const SpirvType *getImageType() const { return imageType; }
   SpirvInstruction *getImage() const { return image; }
   SpirvInstruction *getCoordinate() const { return coordinate; }
   SpirvInstruction *getSample() const { return sample; }
 
+  void replaceOperand(
+      llvm::function_ref<SpirvInstruction *(SpirvInstruction *)> remapOp,
+      bool inEntryFunctionWrapper) override {
+    // imageType is a compile-time SpirvType, not an SSA operand.
+    image = remapOp(image);
+    coordinate = remapOp(coordinate);
+    sample = remapOp(sample);
+  }
+
 private:
+  const SpirvType *imageType;
   SpirvInstruction *image;
   SpirvInstruction *coordinate;
   SpirvInstruction *sample;

tools/clang/lib/SPIRV/AstTypeProbe.cpp

@@ -917,6 +917,13 @@ bool isBuffer(QualType type) {
   return false;
 }
 
+bool isRaytracingAccelerationStructure(QualType type) {
+  if (const auto *rt = type->getAs<RecordType>()) {
+    return rt->getDecl()->getName() == "RaytracingAccelerationStructure";
+  }
+  return false;
+}
+
 bool isRWTexture(QualType type) {
   if (const auto *rt = type->getAs<RecordType>()) {
     const auto name = rt->getDecl()->getName();

tools/clang/lib/SPIRV/CapabilityVisitor.cpp

@@ -955,8 +955,10 @@ bool CapabilityVisitor::visit(SpirvModule *, Visitor::Phase phase) {
                                        {spv::Capability::QuadControlKHR});
 
   if (spvOptions.useDescriptorHeap) {
-    addExtension(Extension::EXT_descriptor_heap, "DescriptorHeap", {});
-    addExtension(Extension::KHR_untyped_pointers, "DescriptorHeap", {});
+    const llvm::StringRef feature = "DescriptorHeap";
+    featureManager.requestTargetEnv(SPV_ENV_VULKAN_1_3, feature, {});
+    addExtension(Extension::EXT_descriptor_heap, feature, {});
+    addExtension(Extension::KHR_untyped_pointers, feature, {});
     addCapability(spv::Capability::DescriptorHeapEXT);
   }
 

tools/clang/lib/SPIRV/DeclResultIdMapper.cpp

@@ -1136,6 +1136,12 @@ DeclResultIdMapper::createFnVar(const VarDecl *var,
   return varInstr;
 }
 
+void DeclResultIdMapper::registerFnVarAlias(const VarDecl *var,
+                                            SpirvInstruction *varInstr) {
+  if (varInstr)
+    registerVariableForDecl(var, createDeclSpirvInfo(varInstr));
+}
+
 SpirvDebugGlobalVariable *DeclResultIdMapper::createDebugGlobalVariable(
     SpirvVariable *var, const QualType &type, const SourceLocation &loc,
     const StringRef &name) {

tools/clang/lib/SPIRV/DeclResultIdMapper.h

@@ -10,7 +10,6 @@
 #ifndef LLVM_CLANG_LIB_SPIRV_DECLRESULTIDMAPPER_H
 #define LLVM_CLANG_LIB_SPIRV_DECLRESULTIDMAPPER_H
 
-#include <tuple>
 #include <vector>
 
 #include "dxc/Support/SPIRVOptions.h"
@@ -286,6 +285,9 @@ class DeclResultIdMapper {
   SpirvVariable *createFnVar(const VarDecl *var,
                              llvm::Optional<SpirvInstruction *> init);
 
+  /// \brief Registers a function-scope alias to an existing instruction.
+  void registerFnVarAlias(const VarDecl *var, SpirvInstruction *varInstr);
+
   /// \brief Creates a file-scope variable and returns its instruction.
   SpirvVariable *createFileVar(const VarDecl *var,
                                llvm::Optional<SpirvInstruction *> init);

tools/clang/lib/SPIRV/EmitVisitor.cpp

@@ -820,6 +820,7 @@ bool EmitVisitor::visit(SpirvUntypedImageTexelPointerEXT *inst) {
   initInstruction(inst);
   curInst.push_back(inst->getResultTypeId());
   curInst.push_back(getOrAssignResultId<SpirvInstruction>(inst));
+  curInst.push_back(typeHandler.emitType(inst->getImageType()));
   curInst.push_back(getOrAssignResultId<SpirvInstruction>(inst->getImage()));
   curInst.push_back(
       getOrAssignResultId<SpirvInstruction>(inst->getCoordinate()));

tools/clang/lib/SPIRV/LowerTypeVisitor.cpp

@@ -224,6 +224,12 @@ bool LowerTypeVisitor::visitInstruction(SpirvInstruction *instr) {
     instr->setResultType(pointerType);
     break;
   }
+  case spv::Op::OpUntypedImageTexelPointerEXT: {
+    instr->setResultType(
+        spvContext.getUntypedPointerKHRType(spv::StorageClass::Image));
+    instr->setStorageClass(spv::StorageClass::Image);
+    break;
+  }
   // Sparse image operations return a sparse residency struct.
   case spv::Op::OpImageSparseSampleImplicitLod:
   case spv::Op::OpImageSparseSampleExplicitLod:

tools/clang/lib/SPIRV/SpirvBuilder.cpp

@@ -533,13 +533,14 @@ SpirvImageTexelPointer *SpirvBuilder::createImageTexelPointer(
 
 SpirvUntypedImageTexelPointerEXT *
 SpirvBuilder::createUntypedImageTexelPointerEXT(QualType resultType,
+                                                const SpirvType *imageType,
                                                 SpirvInstruction *image,
                                                 SpirvInstruction *coordinate,
                                                 SpirvInstruction *sample,
                                                 SourceLocation loc) {
   assert(insertPoint && "null insert point");
   auto *instruction = new (context) SpirvUntypedImageTexelPointerEXT(
-      resultType, loc, image, coordinate, sample);
+      resultType, loc, imageType, image, coordinate, sample);
   insertPoint->addInstruction(instruction);
   return instruction;
 }

tools/clang/lib/SPIRV/SpirvContext.cpp

@@ -65,6 +65,9 @@ SpirvContext::~SpirvContext() {
   for (auto *npaType : nodePayloadArrayTypes)
     npaType->~NodePayloadArrayType();
 
+  for (auto &pair : bufferEXTTypes)
+    pair.second->~BufferEXTType();
+
   for (auto *fnType : functionTypes)
     fnType->~FunctionType();
 
@@ -402,6 +405,14 @@ const StructType *SpirvContext::getByteAddressBufferType(bool isWritable) {
                        !isWritable, StructInterfaceType::StorageBuffer);
 }
 
+const BufferEXTType *SpirvContext::getBufferEXTType(spv::StorageClass sc) {
+  auto found = bufferEXTTypes.find(sc);
+  if (found != bufferEXTTypes.end())
+    return found->second;
+
+  return bufferEXTTypes[sc] = new (this) BufferEXTType(sc);
+}
+
 const StructType *SpirvContext::getACSBufferCounterType() {
   // Create int32.
   const auto *int32Type = getSIntType(32);