[SPIR-V] Add SPV_EXT_descriptor_heap + SPV_KHR_untyped_pointers codegen

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/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
@@ -825,6 +825,25 @@ class SpirvBuilder {
                        bool specConst = false);
   SpirvConstant *getConstantNull(QualType);
   SpirvConstant *getConstantString(llvm::StringRef str, bool specConst = false);
+  /// \brief Returns the OpConstantSizeOfEXT (SPV_EXT_descriptor_heap) for the
+  /// given descriptor operandType, yielding its client-API size in bytes as
+  /// a 32-bit unsigned value. The result is cached per operand type, so each
+  /// descriptor type emits at most one instruction.
+  SpirvConstant *getConstantSizeOfEXT(const SpirvType *operandType);
+
+  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)) computed via OpSpecConstantOp.
+  /// Cached per module.
+  SpirvInstruction *getResourceHeapArrayStride();
+
+  /// \brief Shared ArrayStrideIdEXT operand for sampler-heap runtime arrays:
+  /// the sampler descriptor size. Cached per module.
+  SpirvInstruction *getSamplerHeapArrayStride();
   SpirvUndef *getUndef(QualType);
 
   SpirvString *createString(llvm::StringRef str);
@@ -941,6 +960,16 @@ class SpirvBuilder {
   /// Used as caches for all created builtin variables to avoid duplication.
   llvm::SmallVector<BuiltInVarInfo, 16> builtinVars;
 
+  /// Cache of OpConstantSizeOfEXT instructions keyed on the descriptor operand
+  /// type, so each distinct descriptor type emits at most one instruction.
+  llvm::DenseMap<const SpirvType *, SpirvConstant *> constantSizeOfEXTMap;
+
+  /// Cached shared descriptor-heap array strides (SPV_EXT_descriptor_heap), so
+  /// each is emitted once per module (see
+  /// get{Resource,Sampler}HeapArrayStride).
+  SpirvInstruction *resourceHeapArrayStride = nullptr;
+  SpirvInstruction *samplerHeapArrayStride = nullptr;
+
   SpirvDebugInfoNone *debugNone;
 
   /// DebugExpression that does not reference any DebugOperation

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->getArrayStrideId());
   }
   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 *arrayStrideId = 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

@@ -69,6 +69,7 @@ class SpirvInstruction {
     IK_ConstantFloat,
     IK_ConstantComposite,
     IK_ConstantString,
+    IK_ConstantSizeOfEXT,
     IK_ConstantNull,
 
     // Pointer <-> uint conversions.
@@ -137,6 +138,7 @@ class SpirvInstruction {
     IK_ReadClock,                   // OpReadClock
     IK_SampledImage,                // OpSampledImage
     IK_Select,                      // OpSelect
+    IK_SpecConstantTernaryOp,       // SpecConstant ternary operations
     IK_SpecConstantBinaryOp,        // SpecConstant binary operations
     IK_SpecConstantUnaryOp,         // SpecConstant unary operations
     IK_Store,                       // OpStore
@@ -1537,6 +1539,33 @@ class SpirvConstantNull : public SpirvConstant {
   bool operator==(const SpirvConstantNull &that) const;
 };
 
+/// \brief Represents OpConstantSizeOfEXT (SPV_EXT_descriptor_heap).
+///
+/// Yields the client-API-defined size of a descriptor type in bytes. Unlike
+/// other constants its operand is a SPIR-V type (a descriptor type such as
+/// OpTypeBufferEXT/OpTypeImage/OpTypeSampler), not a value. Used as the operand
+/// of an ArrayStrideIdEXT decoration on descriptor-heap runtime arrays.
+class SpirvConstantSizeOfEXT : public SpirvConstant {
+public:
+  SpirvConstantSizeOfEXT(QualType resultType, const SpirvType *operandType);
+
+  DEFINE_RELEASE_MEMORY_FOR_CLASS(SpirvConstantSizeOfEXT)
+
+  // For LLVM-style RTTI
+  static bool classof(const SpirvInstruction *inst) {
+    return inst->getKind() == IK_ConstantSizeOfEXT;
+  }
+
+  bool invokeVisitor(Visitor *v) override;
+
+  bool operator==(const SpirvConstantSizeOfEXT &that) const;
+
+  const SpirvType *getOperandType() const { return operandType; }
+
+private:
+  const SpirvType *operandType;
+};
+
 class SpirvConstantString : public SpirvConstant {
 public:
   SpirvConstantString(llvm::StringRef stringLiteral, bool isSpecConst = false);
@@ -2065,6 +2094,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 +2108,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;
@@ -2215,6 +2256,35 @@ class SpirvSelect : public SpirvInstruction {
   SpirvInstruction *falseObject;
 };
 
+/// \brief OpSpecConstantOp instruction where the operation is ternary.
+class SpirvSpecConstantTernaryOp : public SpirvInstruction {
+public:
+  SpirvSpecConstantTernaryOp(spv::Op specConstantOp, QualType resultType,
+                             SourceLocation loc, SpirvInstruction *operand1,
+                             SpirvInstruction *operand2,
+                             SpirvInstruction *operand3);
+
+  DEFINE_RELEASE_MEMORY_FOR_CLASS(SpirvSpecConstantTernaryOp)
+
+  // For LLVM-style RTTI
+  static bool classof(const SpirvInstruction *inst) {
+    return inst->getKind() == IK_SpecConstantTernaryOp;
+  }
+
+  bool invokeVisitor(Visitor *v) override;
+
+  spv::Op getSpecConstantopcode() const { return specOp; }
+  SpirvInstruction *getOperand1() const { return operand1; }
+  SpirvInstruction *getOperand2() const { return operand2; }
+  SpirvInstruction *getOperand3() const { return operand3; }
+
+private:
+  spv::Op specOp;
+  SpirvInstruction *operand1;
+  SpirvInstruction *operand2;
+  SpirvInstruction *operand3;
+};
+
 /// \brief OpSpecConstantOp instruction where the operation is binary.
 class SpirvSpecConstantBinaryOp : public SpirvInstruction {
 public:

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,8 +275,12 @@ 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 *strideId = nullptr)
+      : SpirvType(TK_RuntimeArray), elementType(elemType), stride(arrayStride),
+        arrayStrideId(strideId) {
+    assert(!(arrayStride.hasValue() && strideId) &&
+           "literal stride and stride-id are mutually exclusive");
   }
 
   static bool classof(const SpirvType *t) {
@@ -286,12 +291,17 @@ class RuntimeArrayType : public SpirvType {
 
   const SpirvType *getElementType() const { return elementType; }
   llvm::Optional<uint32_t> getStride() const { return stride; }
+  // When non-null, the array is decorated with ArrayStrideIdEXT referencing
+  // this constant <id> (SPV_EXT_descriptor_heap) instead of a literal
+  // ArrayStride. Mutually exclusive with a literal stride.
+  SpirvInstruction *getArrayStrideId() const { return arrayStrideId; }
 
 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.
   llvm::Optional<uint32_t> stride;
+  SpirvInstruction *arrayStrideId = nullptr;
 };
 
 class NodePayloadArrayType : public SpirvType {

tools/clang/include/clang/SPIRV/SpirvVisitor.h

@@ -97,6 +97,7 @@ class Visitor {
   DEFINE_VISIT_METHOD(SpirvConstantFloat)
   DEFINE_VISIT_METHOD(SpirvConstantComposite)
   DEFINE_VISIT_METHOD(SpirvConstantString)
+  DEFINE_VISIT_METHOD(SpirvConstantSizeOfEXT)
   DEFINE_VISIT_METHOD(SpirvConstantNull)
   DEFINE_VISIT_METHOD(SpirvConvertPtrToU)
   DEFINE_VISIT_METHOD(SpirvConvertUToPtr)
@@ -118,6 +119,7 @@ class Visitor {
   DEFINE_VISIT_METHOD(SpirvCopyObject)
   DEFINE_VISIT_METHOD(SpirvSampledImage)
   DEFINE_VISIT_METHOD(SpirvSelect)
+  DEFINE_VISIT_METHOD(SpirvSpecConstantTernaryOp)
   DEFINE_VISIT_METHOD(SpirvSpecConstantBinaryOp)
   DEFINE_VISIT_METHOD(SpirvSpecConstantUnaryOp)
   DEFINE_VISIT_METHOD(SpirvStore)

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);
   }