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

Adds [[vk::resource_heap_stride_constant_id(id)]] and [[vk::sampler_heap_stride_constant_id(id)]] attributes that emit the descriptor heap ArrayStride as a SPIR-V specialization constant using ArrayStrideIdEXT, allowing applications to override the stride at pipeline creation time via VkSpecializationInfo without recompiling the shader.

The attribute initializer supplies the default stride value and must be a power of two in [8, 256]. The CLI flags -fvk-resource-heap-stride and -fvk-sampler-heap-stride take higher precedence and suppress these attributes with a warning when both are specified.

Author: jzakharovnv

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

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 {

tools/clang/lib/SPIRV/DeclResultIdMapper.cpp

@@ -1889,6 +1889,22 @@ void DeclResultIdMapper::registerSpecConstant(const VarDecl *decl,
                                               SpirvInstruction *specConstant) {
   specConstant->setRValue();
   registerVariableForDecl(decl, createDeclSpirvInfo(specConstant));
+  if (decl->hasAttr<VKResourceHeapStrideConstantIdAttr>()) {
+    resourceHeapStride = HeapStrideSpecConst{
+        specConstant,
+        static_cast<uint32_t>(
+            decl->getAttr<VKResourceHeapStrideConstantIdAttr>()
+                ->getSpecConstId()),
+        decl};
+  } else if (decl->hasAttr<VKSamplerHeapStrideConstantIdAttr>()) {
+    samplerHeapStride = HeapStrideSpecConst{
+        specConstant,
+        static_cast<uint32_t>(decl->getAttr<VKSamplerHeapStrideConstantIdAttr>()
+                                  ->getSpecConstId()),
+        decl};
+  } else if (const auto *attr = decl->getAttr<VKConstantIdAttr>()) {
+    userSpecConstIdMap[attr->getSpecConstId()] = decl;
+  }
 }
 
 void DeclResultIdMapper::createCounterVar(

tools/clang/lib/SPIRV/DeclResultIdMapper.h

@@ -296,6 +296,29 @@ class DeclResultIdMapper {
   /// |type|.
   SpirvVariableLike *createResourceHeap(const VarDecl *var, QualType type);
 
+  /// Records the [[vk::*_heap_stride_constant_id]] spec constant for one heap.
+  /// Present iff such an attribute was declared in the translation unit.
+  struct HeapStrideSpecConst {
+    SpirvInstruction *specConst; // the OpSpecConstant
+    uint32_t specId;             // its SpecId
+    const VarDecl *decl;         // declaring var (diagnostics)
+  };
+
+  /// Returns the resource/sampler heap stride spec constant, or None if no
+  /// [[vk::*_heap_stride_constant_id]] was declared.
+  const llvm::Optional<HeapStrideSpecConst> &getResourceHeapStride() const {
+    return resourceHeapStride;
+  }
+  const llvm::Optional<HeapStrideSpecConst> &getSamplerHeapStride() const {
+    return samplerHeapStride;
+  }
+  /// Returns the user [[vk::constant_id]] VarDecl that owns \p specId, or
+  /// nullptr if no user spec-const has claimed that ID.
+  const VarDecl *getUserSpecConstForId(uint32_t specId) const {
+    auto it = userSpecConstIdMap.find(specId);
+    return it != userSpecConstIdMap.end() ? it->second : nullptr;
+  }
+
   /// \brief Creates an external-visible variable and returns its instruction.
   SpirvVariable *createExternVar(const VarDecl *var);
 
@@ -1058,6 +1081,11 @@ class DeclResultIdMapper {
 
   SpirvUntypedVariableKHR *ResourceHeapVar = nullptr;
   SpirvUntypedVariableKHR *SamplerHeapVar = nullptr;
+  llvm::Optional<HeapStrideSpecConst> resourceHeapStride;
+  llvm::Optional<HeapStrideSpecConst> samplerHeapStride;
+  /// Maps SpecId -> VarDecl for user [[vk::constant_id]] declarations, used to
+  /// detect SpecId collisions with heap-stride attributes.
+  llvm::DenseMap<uint32_t, const VarDecl *> userSpecConstIdMap;
 
   /// Mapping from {RW|Append|Consume}StructuredBuffers to their
   /// counter variables' (instr-ptr, is-alias-or-not) pairs

tools/clang/lib/SPIRV/EmitVisitor.cpp

@@ -2695,9 +2695,15 @@ uint32_t EmitTypeHandler::emitType(const SpirvType *type) {
     curTypeInst.push_back(elemTypeId);
     finalizeTypeInstruction();
 
-    auto stride = raType->getStride();
-    if (stride.hasValue())
-      emitDecoration(id, spv::Decoration::ArrayStride, {stride.getValue()});
+    if (auto *sc = raType->getStrideSpecConst()) {
+      const uint32_t scId = getOrAssignResultId<SpirvInstruction>(sc);
+      emitDecoration(id, spv::Decoration::ArrayStrideIdEXT, {scId}, llvm::None,
+                     /*usesIdParams=*/true);
+    } else {
+      auto stride = raType->getStride();
+      if (stride.hasValue())
+        emitDecoration(id, spv::Decoration::ArrayStride, {stride.getValue()});
+    }
   }
   // NodePayloadArray types
   else if (const auto *npaType = dyn_cast<NodePayloadArrayType>(type)) {

tools/clang/lib/SPIRV/SpirvContext.cpp

@@ -268,14 +268,15 @@ SpirvContext::getArrayType(const SpirvType *elemType, uint32_t elemCount,
 
 const RuntimeArrayType *
 SpirvContext::getRuntimeArrayType(const SpirvType *elemType,
-                                  llvm::Optional<uint32_t> arrayStride) {
-  RuntimeArrayType type(elemType, arrayStride);
+                                  llvm::Optional<uint32_t> arrayStride,
+                                  SpirvInstruction *strideSpecConst) {
+  RuntimeArrayType type(elemType, arrayStride, strideSpecConst);
   auto found = runtimeArrayTypes.find(&type);
   if (found != runtimeArrayTypes.end())
     return *found;
 
   auto inserted = runtimeArrayTypes.insert(
-      new (this) RuntimeArrayType(elemType, arrayStride));
+      new (this) RuntimeArrayType(elemType, arrayStride, strideSpecConst));
   return *(inserted.first);
 }