[SPIR-V] Add descriptor heap stride specialization constant attribute

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,111 @@ 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, decorates the array with a byte ``ArrayStride`` (the stride is
+configurable; see `Descriptor heap array stride`_ below), and 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
+
+Descriptor heap array stride
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+The ``ArrayStride`` of each heap runtime array defaults to 64 bytes for the
+resource heap and 32 bytes for the sampler heap. The stride
+can be overridden, in increasing order of precedence:
+
+- ``[[vk::resource_heap_stride_constant_id(id)]]`` and
+  ``[[vk::sampler_heap_stride_constant_id(id)]]`` on a ``uint`` global emit the
+  stride as a specialization constant decorated ``ArrayStrideIdEXT`` (an ``<id>``)
+  instead of a literal, letting the application override it at pipeline creation
+  through ``VkSpecializationInfo``. The attribute initializer supplies the default
+  value and must be a power of two in [8, 256]. This attribute is mutually
+  exclusive with ``[[vk::constant_id]]`` on the same declaration.
+- ``-fvk-resource-heap-stride <N>`` and ``-fvk-sampler-heap-stride <N>``
+  emit a fixed literal ``OpDecorate ... ArrayStride N`` on the resource and
+  sampler heap arrays respectively. ``N`` must be a power of two in the inclusive
+  range [8, 256]. The command-line override has the **highest** precedence: when
+  set, the matching ``[[vk::*_heap_stride_constant_id]]`` attribute is ignored
+  (DXC emits a warning at the attribute and no ``ArrayStrideIdEXT`` is emitted for
+  that heap) and the literal stride is used.
+
+So the command-line literal takes precedence over the spec-constant attribute,
+which in turn takes precedence over the built-in defaults.
+
+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
 ================
 

include/dxc/Support/HLSLOptions.td

@@ -454,6 +454,20 @@ def fvk_bind_sampler_heap : MultiArg<["-"], "fvk-bind-sampler-heap", 2>, MetaVar
   HelpText<"Specify Vulkan binding number and set number for the sampler heap.">;
 def fvk_bind_counter_heap : MultiArg<["-"], "fvk-bind-counter-heap", 2>, MetaVarName<"<binding> <set>">, Group<spirv_Group>, Flags<[CoreOption, DriverOption]>,
   HelpText<"Specify Vulkan binding number and set number for the counter heap.">;
+def fvk_resource_heap_stride
+    : Separate<["-"], "fvk-resource-heap-stride">,
+      MetaVarName<"<stride>">,
+      Group<spirv_Group>,
+      Flags<[CoreOption, DriverOption]>,
+      HelpText<"Override the byte ArrayStride of the resource descriptor heap "
+               "runtime array. Must be a power of 2 in [8, 256].">;
+def fvk_sampler_heap_stride
+    : Separate<["-"], "fvk-sampler-heap-stride">,
+      MetaVarName<"<stride>">,
+      Group<spirv_Group>,
+      Flags<[CoreOption, DriverOption]>,
+      HelpText<"Override the byte ArrayStride of the sampler descriptor heap "
+               "runtime array. Must be a power of 2 in [8, 256].">;
 // SPIRV Change Ends
 
 //////////////////////////////////////////////////////////////////////////////

include/dxc/Support/SPIRVOptions.h

@@ -110,6 +110,14 @@ struct SpirvCodeGenOptions {
   std::optional<BindingInfo> samplerHeapBinding;
   std::optional<BindingInfo> counterHeapBinding;
 
+  // User-defined byte ArrayStride overrides for the resource/sampler descriptor
+  // heap runtime arrays (-fvk-resource-heap-stride / -fvk-sampler-heap-stride).
+  // When set, the value is a literal power of 2 in [8, 256] and takes the
+  // highest precedence: it overrides any [[vk::*_heap_stride_constant_id]]
+  // spec-constant attribute (which is suppressed with a warning).
+  std::optional<uint32_t> resourceHeapStride;
+  std::optional<uint32_t> samplerHeapStride;
+
   bool signaturePacking =
       false; ///< Whether signature packing is enabled or not
 

lib/DxcSupport/HLSLOptions.cpp

@@ -357,6 +357,43 @@ handleFixedBinding(const InputArgList &args, OptSpecifier id,
   return true;
 }
 
+// Parses the single-integer descriptor-heap stride flag |id| in |args|. If
+// present, validates that the value is a power of 2 in [8, 256] and stores it
+// in |stride|. Returns true on success (including when the flag is absent).
+// Returns false and writes to |errors| when the value is malformed or invalid,
+// using |name| as the pretty flag name.
+static bool handleHeapStride(const InputArgList &args, OptSpecifier id,
+                             std::optional<uint32_t> *stride,
+                             llvm::StringRef name, llvm::raw_ostream &errors) {
+  Arg *arg = args.getLastArg(id);
+  if (!arg) {
+    *stride = std::nullopt;
+    return true;
+  }
+
+  if (!args.hasArg(OPT_spirv)) {
+    errors << name << " requires -spirv";
+    return false;
+  }
+
+  llvm::StringRef value = arg->getValue();
+  uint32_t number = 0;
+  if (value.getAsInteger(10, number)) {
+    errors << "invalid " << name << " argument: '" << value << "'";
+    return false;
+  }
+  // Power of 2 in [8, 256] inclusive.
+  if (number < 8 || number > 256 || (number & (number - 1)) != 0) {
+    errors << name
+           << " must be a power of 2 between 8 and 256 (inclusive); got "
+           << value;
+    return false;
+  }
+
+  *stride = number;
+  return true;
+}
+
 // Check if any options that are unsupported with SPIR-V are used.
 static bool hasUnsupportedSpirvOption(const InputArgList &args,
                                       llvm::raw_ostream &errors) {
@@ -1175,6 +1212,15 @@ int ReadDxcOpts(const OptTable *optionTable, unsigned flagsToInclude,
     return 1;
   }
 
+  if (!handleHeapStride(Args, OPT_fvk_resource_heap_stride,
+                        &opts.SpirvOptions.resourceHeapStride,
+                        "-fvk-resource-heap-stride", errors) ||
+      !handleHeapStride(Args, OPT_fvk_sampler_heap_stride,
+                        &opts.SpirvOptions.samplerHeapStride,
+                        "-fvk-sampler-heap-stride", errors)) {
+    return 1;
+  }
+
   for (const Arg *A : Args.filtered(OPT_fspv_extension_EQ)) {
     opts.SpirvOptions.allowedExtensions.push_back(A->getValue());
   }
@@ -1316,7 +1362,9 @@ int ReadDxcOpts(const OptTable *optionTable, unsigned flagsToInclude,
       !Args.getLastArgValue(OPT_fvk_u_shift).empty() ||
       !Args.getLastArgValue(OPT_fvk_bind_resource_heap).empty() ||
       !Args.getLastArgValue(OPT_fvk_bind_sampler_heap).empty() ||
-      !Args.getLastArgValue(OPT_fvk_bind_counter_heap).empty()) {
+      !Args.getLastArgValue(OPT_fvk_bind_counter_heap).empty() ||
+      !Args.getLastArgValue(OPT_fvk_resource_heap_stride).empty() ||
+      !Args.getLastArgValue(OPT_fvk_sampler_heap_stride).empty()) {
     errors << "SPIR-V CodeGen not available. "
               "Please recompile with -DENABLE_SPIRV_CODEGEN=ON.";
     return 1;

tools/clang/include/clang/Basic/Attr.td

@@ -1611,6 +1611,24 @@ def VKConstantId : InheritableAttr {
   let Documentation = [Undocumented];
 }
 
+def VKResourceHeapStrideConstantId : InheritableAttr {
+  let Spellings = [CXX11<"vk", "resource_heap_stride_constant_id">];
+  let Subjects =
+      SubjectList<[ScalarGlobalVar], ErrorDiag, "ExpectedScalarGlobalVar">;
+  let Args = [IntArgument<"SpecConstId">];
+  let LangOpts = [SPIRV];
+  let Documentation = [Undocumented];
+}
+
+def VKSamplerHeapStrideConstantId : InheritableAttr {
+  let Spellings = [CXX11<"vk", "sampler_heap_stride_constant_id">];
+  let Subjects =
+      SubjectList<[ScalarGlobalVar], ErrorDiag, "ExpectedScalarGlobalVar">;
+  let Args = [IntArgument<"SpecConstId">];
+  let LangOpts = [SPIRV];
+  let Documentation = [Undocumented];
+}
+
 def VKPostDepthCoverage : InheritableAttr {
   let Spellings = [CXX11<"vk", "post_depth_coverage">];
   let Subjects = SubjectList<[Function], ErrorDiag>;

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

@@ -89,7 +89,8 @@ struct RuntimeArrayTypeMapInfo {
   static inline RuntimeArrayType *getTombstoneKey() { return nullptr; }
   static unsigned getHashValue(const RuntimeArrayType *Val) {
     return llvm::hash_combine(Val->getElementType(),
-                              Val->getStride().hasValue());
+                              Val->getStride().hasValue(),
+                              Val->getStrideSpecConst());
   }
   static bool isEqual(const RuntimeArrayType *LHS,
                       const RuntimeArrayType *RHS) {
@@ -284,7 +285,8 @@ class SpirvContext {
                                 llvm::Optional<uint32_t> arrayStride);
   const RuntimeArrayType *
   getRuntimeArrayType(const SpirvType *elemType,
-                      llvm::Optional<uint32_t> arrayStride);
+                      llvm::Optional<uint32_t> arrayStride,
+                      SpirvInstruction *strideSpecConst = nullptr);
   const NodePayloadArrayType *
   getNodePayloadArrayType(const SpirvType *elemType,
                           const ParmVarDecl *nodeDecl);
@@ -298,6 +300,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 +539,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/include/clang/SPIRV/SpirvType.h

@@ -23,6 +23,7 @@ namespace clang {
 namespace spirv {
 
 class HybridType;
+class SpirvInstruction;
 
 enum class StructInterfaceType : uint32_t {
   InternalStorage = 0,
@@ -274,9 +275,10 @@ class ArrayType : public SpirvType {
 class RuntimeArrayType : public SpirvType {
 public:
   RuntimeArrayType(const SpirvType *elemType,
-                   llvm::Optional<uint32_t> arrayStride)
-      : SpirvType(TK_RuntimeArray), elementType(elemType), stride(arrayStride) {
-  }
+                   llvm::Optional<uint32_t> arrayStride,
+                   SpirvInstruction *strideSpecConst = nullptr)
+      : SpirvType(TK_RuntimeArray), elementType(elemType), stride(arrayStride),
+        strideSpecConst(strideSpecConst) {}
 
   static bool classof(const SpirvType *t) {
     return t->getKind() == TK_RuntimeArray;
@@ -286,12 +288,16 @@ class RuntimeArrayType : public SpirvType {
 
   const SpirvType *getElementType() const { return elementType; }
   llvm::Optional<uint32_t> getStride() const { return stride; }
+  SpirvInstruction *getStrideSpecConst() const { return strideSpecConst; }
 
 private:
   const SpirvType *elementType;
   // Two runtime arrays with different ArrayStride decorations, are in fact two
   // different types. If no layout information is needed, use llvm::None.
+  // Ignored when strideSpecConst is non-null (the spec-const wins).
   llvm::Optional<uint32_t> stride;
+  // When non-null, ArrayStrideIdEXT %id is emitted instead of ArrayStride N.
+  SpirvInstruction *strideSpecConst;
 };
 
 class NodePayloadArrayType : public SpirvType {