[SPIR-V] Add descriptor heap RaytracingAccelerationStructure support

docs/ReleaseNotes.md

@@ -25,6 +25,11 @@ The included licenses apply to the following files:
   buffer-like resources (`StructuredBuffer`, `ByteAddressBuffer`,
   `ConstantBuffer`, `TextureBuffer`, and RW variants) via
   `SPV_KHR_untyped_pointers`. Requires `-fspv-target-env=vulkan1.3`.
+- Extended `-fspv-use-descriptor-heap` to support `RaytracingAccelerationStructure`
+  loaded from `ResourceDescriptorHeap`. Outside ray-tracing stages, list
+  `-fspv-extension=SPV_KHR_ray_query`, `SPV_KHR_ray_tracing` or
+  `SPV_NV_ray_tracing` so the resource heap stride covers acceleration structure
+  descriptors; otherwise compilation fails.
 
 Place release notes for the upcoming release below this line and remove this
 line upon naming the release. Refer to previous for appropriate section names.

docs/SPIR-V.rst

@@ -2143,13 +2143,27 @@ handle into a function-scope image variable:
   %texel_ptr = OpUntypedImageTexelPointerEXT %uptr_image %image_type %descriptor %coord %sample
   %old = OpAtomicIAdd %uint %texel_ptr %scope %semantics %value
 
+``RaytracingAccelerationStructure`` resources loaded from
+``ResourceDescriptorHeap`` follow the same access-chain-then-load shape over a
+runtime array of ``OpTypeAccelerationStructureKHR``:
+
+.. code:: spirv
+
+  %accel_type = OpTypeAccelerationStructureKHR
+  %accel_array = OpTypeRuntimeArray %accel_type
+  OpDecorateId %accel_array ArrayStrideIdEXT %resource_stride
+  %descriptor = OpUntypedAccessChainKHR %uptr_uc %accel_array %resource_heap %index
+  %accel = OpLoad %accel_type %descriptor
+
 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 no ``NonUniform`` decoration is emitted on the
-``OpUntypedAccessChainKHR`` result or on the loaded descriptor;
+ByteAddressBuffer/RWByteAddressBuffer, ConstantBuffer, TextureBuffer, and
+``RaytracingAccelerationStructure`` heap loads, including direct field and
+array-element accesses for ``ConstantBuffer<T>`` and ``TextureBuffer<T>``.
+Acceleration structure loads are additionally subject to the stride requirement
+described in `Descriptor heap array stride`_ below.
+``NonUniformResourceIndex`` is accepted, but no ``NonUniform`` decoration is
+emitted on the ``OpUntypedAccessChainKHR`` result or on the loaded descriptor;
 ``SPV_EXT_descriptor_heap`` deprecates the decoration for heap accesses and
 drivers handle divergent heap indices natively. The index operand itself may
 still carry ``NonUniform`` from the surrounding expression.
@@ -2167,29 +2181,45 @@ Buffer-like local variables initialized from a heap access are tracked by their
 heap index, and later uses are re-lowered as fresh access chains. Image
 variables retain their loaded descriptor handle for ordinary reads and writes;
 the tracked heap index is used when an atomic operation requires
-``OpUntypedImageTexelPointerEXT``. Acceleration-structure aliases retain the
-initially loaded handle. Because index-tracked aliases do not model
-control-flow merges, DXC rejects a variable that holds both a bound resource
-and a heap descriptor, whether through a conditional assignment, a
-reassignment back to a bound resource, or an assignment inside a loop::
+``OpUntypedImageTexelPointerEXT``. A ``RaytracingAccelerationStructure``
+variable initialized from a heap access is a name for the handle loaded by its
+initializer.
+
+These aliases are compile-time state and do not model control-flow merges, so
+DXC rejects a variable that holds both a bound resource and a heap descriptor,
+whether through a conditional assignment, a reassignment back to a bound
+resource, or an assignment inside a loop:
+
+.. code:: text
 
   error: mixing bound and descriptor heap resources in the same variable is not
   supported with SPV_EXT_descriptor_heap
 
+A heap-initialized ``RaytracingAccelerationStructure`` variable cannot be
+reassigned at all, whether from a bound resource or from another heap access:
+
+.. code:: text
+
+  error: reassigning a RaytracingAccelerationStructure initialized from
+  ResourceDescriptorHeap is not supported with SPV_EXT_descriptor_heap
+
 Supporting these forms requires modelling the alias as a value with real
 control-flow merges instead of as compile-time state.
 
 ``StructuredBuffer``, ``ByteAddressBuffer``, and related buffer-type aliases
 cannot yet be passed to or returned from user functions; access the resource
 directly at each call or return site.
 
-Two further restrictions on the heap access expression itself produce
-diagnostics. The object being subscripted must be a direct reference to the
-builtin ``ResourceDescriptorHeap`` or ``SamplerDescriptorHeap`` variable, and
-the subscript result must be immediately converted to a concrete resource type
-so that DXC can select a descriptor type for the access. A subscript whose
-result is discarded, or used in a context that supplies no target resource
-type, is rejected.
+Three further restrictions on heap accesses produce diagnostics:
+
+- The object being subscripted must be a direct reference to the builtin
+  ``ResourceDescriptorHeap`` or ``SamplerDescriptorHeap`` variable.
+- The subscript result must be immediately converted to a concrete resource
+  type so that DXC can select a descriptor type for the access. A subscript
+  whose result is discarded, or used in a context that supplies no target
+  resource type, is rejected.
+- A local ``RaytracingAccelerationStructure`` must be initialized from a
+  loadable heap access, since the alias has no backing descriptor otherwise.
 
 Descriptor heap array stride
 ++++++++++++++++++++++++++++
@@ -2205,8 +2235,29 @@ The ``OpConstantSizeOfEXT`` operands are placeholder types chosen only for their
 descriptor class. All image types report the same descriptor size, so the
 placeholder is a plain sampled 2D float image and bears no relation to the image
 types the shader actually uses; a module will normally contain both the
-placeholder type and the distinct image types its heap accesses lower to. The
-stride value is built once and cached on first use.
+placeholder type and the distinct image types its heap accesses lower to.
+
+The stride is a property of the heap layout, not of the shader's loads. DXC
+decides it before translating any function, from the shader stage and the
+``-fspv-extension`` list. When acceleration structure descriptors may be on the
+resource heap, the formula expands to a three-way max
+``max(max(sizeof(image_descriptor), sizeof(buffer_descriptor)), sizeof(acceleration_structure))``.
+A ray-tracing shader that only heap-loads a texture still gets the three-way
+max. DXC adds the acceleration structure term when either
+
+- any entry point is a ray-tracing stage, or
+- the user lists ``-fspv-extension=SPV_KHR_ray_query``,
+  ``-fspv-extension=SPV_KHR_ray_tracing`` or
+  ``-fspv-extension=SPV_NV_ray_tracing``.
+
+Extensions allowed by default do not satisfy the second condition. The
+acceleration structure type added by widening is declared with the capability
+of the listed extension (``RayQueryKHR``, ``RayTracingKHR`` or
+``RayTracingNV``); if several are listed, ``SPV_KHR_ray_query`` is preferred,
+then ``SPV_KHR_ray_tracing``, then ``SPV_NV_ray_tracing``. A shader that is not
+a ray-tracing stage and heap-loads a ``RaytracingAccelerationStructure``
+without such an extension is rejected, because its stride would not cover the
+acceleration structure 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/FeatureManager.h

@@ -117,6 +117,15 @@ class FeatureManager {
   /// Returns true if the given extension is set in allowedExtensions
   bool isExtensionEnabled(Extension) const;
 
+  /// \brief Returns the allowed extension that provides the acceleration
+  /// structure type, or Extension::Unknown if there is none.
+  Extension getAccelerationStructureProvider() const;
+
+  /// \brief Returns true if the user listed, with -fspv-extension, an extension
+  /// that provides the acceleration structure type.
+  /// When no list is given, returns false.
+  bool hasExplicitAccelerationStructureProvider() const;
+
   /// Returns true if the target environment is Vulkan 1.1 or above.
   /// Returns false otherwise.
   bool isTargetEnvVulkan1p1OrAbove() const;
@@ -181,6 +190,7 @@ class FeatureManager {
   DiagnosticsEngine &diags;
 
   llvm::SmallBitVector allowedExtensions;
+  bool hasExplicitExtensions = false;
   spv_target_env targetEnv;
   std::string targetEnvStr;
 };

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

@@ -231,6 +231,8 @@ class SpirvBuilder {
                                 SpirvInstruction *lhs, SpirvInstruction *rhs,
                                 SourceLocation loc, SourceRange range = {});
 
+  /// \brief Creates a binary OpSpecConstantOp at module level, not in the
+  /// current basic block.
   SpirvSpecConstantBinaryOp *createSpecConstantBinaryOp(spv::Op op,
                                                         QualType resultType,
                                                         SpirvInstruction *lhs,
@@ -838,13 +840,37 @@ class SpirvBuilder {
   /// integer.
   SpirvConstant *getConstantSizeOfEXT(const SpirvType *operandType);
 
+  /// \brief Creates a ternary OpSpecConstantOp at module level, not in the
+  /// current basic block.
   SpirvSpecConstantTernaryOp *
   createSpecConstantTernaryOp(spv::Op op, QualType resultType,
                               SpirvInstruction *op1, SpirvInstruction *op2,
                               SpirvInstruction *op3, SourceLocation loc);
 
-  /// \brief Shared ArrayStrideIdEXT operand for resource-heap runtime arrays:
-  /// max(sizeof(image), sizeof(buffer)).
+  /// \brief Records that acceleration structures may be loaded from the
+  /// resource heap, so the resource-heap stride must cover their descriptor.
+  ///
+  /// Must be called before getResourceHeapArrayStride(), which builds and
+  /// caches the stride. A later call is ignored, so
+  /// resourceHeapStrideIncludesAccelStruct() stays false and code generation
+  /// must reject acceleration structure heap loads.
+  void registerResourceHeapAccelStruct() {
+    if (!resourceHeapArrayStride)
+      resourceHeapHasAccelStruct = true;
+  }
+
+  /// \brief Returns true if the resource-heap stride covers the acceleration
+  /// structure descriptor. Acceleration structure heap loads are only valid
+  /// when this is true.
+  bool resourceHeapStrideIncludesAccelStruct() const {
+    return resourceHeapHasAccelStruct;
+  }
+
+  /// \brief Returns the ArrayStrideIdEXT operand shared by every resource-heap
+  /// runtime array: the largest descriptor size among images and buffers, and
+  /// among acceleration structures if registerResourceHeapAccelStruct() was
+  /// called first. The value is built from specialization constants because
+  /// descriptor sizes are only known at pipeline creation.
   SpirvInstruction *getResourceHeapArrayStride();
 
   /// \brief Shared ArrayStrideIdEXT operand for sampler-heap runtime arrays:
@@ -975,6 +1001,11 @@ class SpirvBuilder {
   SpirvInstruction *resourceHeapArrayStride = nullptr;
   SpirvInstruction *samplerHeapArrayStride = nullptr;
 
+  /// Set by registerResourceHeapAccelStruct(). When true,
+  /// getResourceHeapArrayStride() includes the acceleration structure
+  /// descriptor size.
+  bool resourceHeapHasAccelStruct = false;
+
   SpirvDebugInfoNone *debugNone;
 
   /// DebugExpression that does not reference any DebugOperation

tools/clang/include/clang/SPIRV/SpirvModule.h

@@ -149,6 +149,10 @@ class SpirvModule {
   // Adds a constant to the module.
   void addConstant(SpirvConstant *);
 
+  // Adds an OpSpecConstantOp to the module. These are emitted after all
+  // constants, in insertion order.
+  void addSpecConstantOp(SpirvInstruction *);
+
   // Adds an Undef to the module.
   void addUndef(SpirvUndef *);
 
@@ -216,6 +220,7 @@ class SpirvModule {
       decorations;
 
   std::vector<SpirvConstant *> constants;
+  std::vector<SpirvInstruction *> specConstantOps;
   std::vector<SpirvUndef *> undefs;
   std::vector<SpirvVariableLike *> variables;
   // A vector of functions in the module in the order that they should be

tools/clang/lib/SPIRV/AstTypeProbe.cpp

@@ -895,6 +895,11 @@ bool isBuffer(QualType type) {
          hlsl::GetHLSLResourceClass(type) == hlsl::DXIL::ResourceClass::SRV;
 }
 
+bool isRaytracingAccelerationStructure(QualType type) {
+  return hlsl::GetHLSLResourceKind(type) ==
+         hlsl::DXIL::ResourceKind::RTAccelerationStructure;
+}
+
 bool isRWTexture(QualType type) {
   return hlsl::DXIL::IsAnyTexture(hlsl::GetHLSLResourceKind(type)) &&
          hlsl::GetHLSLResourceClass(type) == hlsl::DXIL::ResourceClass::UAV;

tools/clang/lib/SPIRV/CapabilityVisitor.cpp

@@ -208,6 +208,17 @@ void CapabilityVisitor::addCapabilityForType(const SpirvType *type,
   else if (const auto *sampledImageType = dyn_cast<SampledImageType>(type)) {
     addCapabilityForType(sampledImageType->getImageType(), loc, sc);
   }
+  // Acceleration structure type
+  else if (isa<AccelerationStructureTypeNV>(type)) {
+    // visit(SpirvModule *) declares SPV_KHR_ray_query when enabled, otherwise
+    // SPV_KHR_ray_tracing. NV is not declared there because that would add it
+    // to modules that never use the type, so declare it here, and only when no
+    // KHR extension provides the type.
+    if (!featureManager.isExtensionEnabled(Extension::KHR_ray_query) &&
+        !featureManager.isExtensionEnabled(Extension::KHR_ray_tracing))
+      addExtensionAndCapabilitiesIfEnabled(Extension::NV_ray_tracing,
+                                           {spv::Capability::RayTracingNV});
+  }
   // Pointer type
   else if (const auto *ptrType = dyn_cast<SpirvPointerType>(type)) {
     addCapabilityForType(ptrType->getPointeeType(), loc,
@@ -239,6 +250,14 @@ void CapabilityVisitor::addCapabilityForType(const SpirvType *type,
   }
 }
 
+bool CapabilityVisitor::visit(SpirvConstantSizeOfEXT *inst) {
+  // The operand type is not the result type, so the generic handling of
+  // instructions does not see it.
+  addCapabilityForType(inst->getOperandType(), inst->getSourceLocation(),
+                       spv::StorageClass::Max);
+  return visitInstruction(inst);
+}
+
 bool CapabilityVisitor::visit(SpirvDecoration *decor) {
   const auto loc = decor->getSourceLocation();
   switch (decor->getDecoration()) {

tools/clang/lib/SPIRV/CapabilityVisitor.h

@@ -40,6 +40,7 @@ class CapabilityVisitor : public Visitor {
   bool visit(SpirvDemoteToHelperInvocation *) override;
   bool visit(SpirvIsHelperInvocationEXT *) override;
   bool visit(SpirvReadClock *) override;
+  bool visit(SpirvConstantSizeOfEXT *) override;
 
   using Visitor::visit;
 

tools/clang/lib/SPIRV/DeclResultIdMapper.cpp

@@ -1144,6 +1144,19 @@ DeclResultIdMapper::createFnVar(const VarDecl *var,
   return varInstr;
 }
 
+void DeclResultIdMapper::registerFnVarAlias(const VarDecl *var,
+                                            SpirvInstruction *varInstr) {
+  assert(varInstr);
+  registerVariableForDecl(var, createDeclSpirvInfo(varInstr));
+}
+
+bool DeclResultIdMapper::hasFnVarAlias(const VarDecl *var) const {
+  const DeclSpirvInfo *info = getDeclSpirvInfo(var);
+  if (!info || !info->instr)
+    return false;
+  return !isa<SpirvVariable>(info->instr);
+}
+
 SpirvDebugGlobalVariable *DeclResultIdMapper::createDebugGlobalVariable(
     SpirvVariable *var, const QualType &type, const SourceLocation &loc,
     const StringRef &name) {

tools/clang/lib/SPIRV/DeclResultIdMapper.h

@@ -285,6 +285,14 @@ class DeclResultIdMapper {
   SpirvVariable *createFnVar(const VarDecl *var,
                              llvm::Optional<SpirvInstruction *> init);
 
+  /// \brief Maps a function-scope variable to an existing instruction instead
+  /// of creating a function variable for it.
+  void registerFnVarAlias(const VarDecl *var, SpirvInstruction *varInstr);
+
+  /// \brief Returns true if the decl is mapped to an instruction that is not a
+  /// variable, as done by registerFnVarAlias.
+  bool hasFnVarAlias(const VarDecl *var) const;
+
   /// \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

@@ -460,9 +460,16 @@ void EmitVisitor::initInstruction(SpirvInstruction *inst) {
   if (auto *var = dyn_cast<SpirvVariable>(inst))
     isGlobalVar = var->getStorageClass() != spv::StorageClass::Function;
   const auto op = inst->getopcode();
-  emitDebugLine(op, inst->getSourceLocation(), inst->getSourceRange(),
-                isGlobalVar ? &globalVarsBinary : &mainBinary,
-                isa<SpirvDebugScope>(inst));
+
+  // OpSpecConstantOp is module-scope: it is emitted with the types and
+  // constants, before the debug sources and any function.
+  const bool isSpecConstantOp = isa<SpirvSpecConstantUnaryOp>(inst) ||
+                                isa<SpirvSpecConstantBinaryOp>(inst) ||
+                                isa<SpirvSpecConstantTernaryOp>(inst);
+  if (!isSpecConstantOp)
+    emitDebugLine(op, inst->getSourceLocation(), inst->getSourceRange(),
+                  isGlobalVar ? &globalVarsBinary : &mainBinary,
+                  isa<SpirvDebugScope>(inst));
 
   // Initialize the current instruction for emitting.
   curInst.clear();
@@ -1430,6 +1437,8 @@ bool EmitVisitor::visit(SpirvSelect *inst) {
   return true;
 }
 
+// OpSpecConstantOp is a module-scope instruction. All three forms below are
+// emitted with the types and constants.
 bool EmitVisitor::visit(SpirvSpecConstantTernaryOp *inst) {
   initInstruction(inst);
   curInst.push_back(inst->getResultTypeId());
@@ -1463,7 +1472,7 @@ bool EmitVisitor::visit(SpirvSpecConstantUnaryOp *inst) {
   curInst.push_back(getOrAssignResultId<SpirvInstruction>(inst));
   curInst.push_back(static_cast<uint32_t>(inst->getSpecConstantopcode()));
   curInst.push_back(getOrAssignResultId<SpirvInstruction>(inst->getOperand()));
-  finalizeInstruction(&mainBinary);
+  finalizeInstruction(&typeConstantBinary);
   emitDebugNameForInstruction(getOrAssignResultId<SpirvInstruction>(inst),
                               inst->getDebugName());
   return true;

tools/clang/lib/SPIRV/EmitVisitor.h

@@ -463,7 +463,7 @@ class EmitVisitor : public Visitor {
   std::vector<uint32_t> annotationsBinary;
   // All forward pointer type declaration instructions
   std::vector<uint32_t> fwdDeclBinary;
-  // All other type and constant instructions
+  // All other type and constant instructions, including OpSpecConstantOp
   std::vector<uint32_t> typeConstantBinary;
   // All global variable declarations (all OpVariable instructions whose Storage
   // Class is not Function)