[HLSL] [SPIR-V] Add initial support for typed buffer counters

This is part 1 of implementing the typed buffer counters proposal:
https://github.com/llvm/wg-hlsl/blob/main/proposals/0023-typed-buffer-counters.md

This patch adds the initial plumbing for supporting counter variables
associated with structured buffers for the SPIR-V backend. It introduces
an `IsCounter` attribute to `HLSLAttributedResourceType` and threads it
through the AST, type printing, and mangling. It also adds a
`__counter_handle` member to the relevant buffer types in
`HLSLBuiltinTypeDeclBuilder`.

Author: s-perron

clang/include/clang/AST/TypeBase.h

@@ -6700,15 +6700,21 @@ class HLSLAttributedResourceType : public Type, public llvm::FoldingSetNode {
     LLVM_PREFERRED_TYPE(bool)
     uint8_t RawBuffer : 1;
 
+    LLVM_PREFERRED_TYPE(bool)
+    uint8_t IsCounter : 1;
+
     Attributes(llvm::dxil::ResourceClass ResourceClass, bool IsROV = false,
-               bool RawBuffer = false)
-        : ResourceClass(ResourceClass), IsROV(IsROV), RawBuffer(RawBuffer) {}
+               bool RawBuffer = false, bool IsCounter = false)
+        : ResourceClass(ResourceClass), IsROV(IsROV), RawBuffer(RawBuffer),
+          IsCounter(IsCounter) {}
 
-    Attributes() : Attributes(llvm::dxil::ResourceClass::UAV, false, false) {}
+    Attributes()
+        : Attributes(llvm::dxil::ResourceClass::UAV, false, false, false) {}
 
     friend bool operator==(const Attributes &LHS, const Attributes &RHS) {
-      return std::tie(LHS.ResourceClass, LHS.IsROV, LHS.RawBuffer) ==
-             std::tie(RHS.ResourceClass, RHS.IsROV, RHS.RawBuffer);
+      return std::tie(LHS.ResourceClass, LHS.IsROV, LHS.RawBuffer,
+                      LHS.IsCounter) == std::tie(RHS.ResourceClass, RHS.IsROV,
+                                                 RHS.RawBuffer, RHS.IsCounter);
     }
     friend bool operator!=(const Attributes &LHS, const Attributes &RHS) {
       return !(LHS == RHS);
@@ -6749,6 +6755,7 @@ class HLSLAttributedResourceType : public Type, public llvm::FoldingSetNode {
     ID.AddInteger(static_cast<uint32_t>(Attrs.ResourceClass));
     ID.AddBoolean(Attrs.IsROV);
     ID.AddBoolean(Attrs.RawBuffer);
+    ID.AddBoolean(Attrs.IsCounter);
   }
 
   static bool classof(const Type *T) {

clang/include/clang/AST/TypeProperties.td

@@ -662,14 +662,17 @@ let Class = HLSLAttributedResourceType in {
   def : Property<"rawBuffer", Bool> {
     let Read = [{ node->getAttrs().RawBuffer }];
   }
+  def : Property<"isCounter", Bool> {
+    let Read = [{ node->getAttrs().IsCounter }];
+  }
   def : Property<"wrappedTy", QualType> {
     let Read = [{ node->getWrappedType() }];
   }
   def : Property<"containedTy", QualType> {
     let Read = [{ node->getContainedType() }];
   }
   def : Creator<[{
-    HLSLAttributedResourceType::Attributes attrs(static_cast<llvm::dxil::ResourceClass>(resClass), isROV, rawBuffer);
+    HLSLAttributedResourceType::Attributes attrs(static_cast<llvm::dxil::ResourceClass>(resClass), isROV, rawBuffer, isCounter);
     return ctx.getHLSLAttributedResourceType(wrappedTy, containedTy, attrs);
   }]>;
 }

clang/include/clang/Basic/Attr.td

@@ -5059,6 +5059,12 @@ def HLSLRawBuffer : TypeAttr {
   let Documentation = [InternalOnly];
 }
 
+def HLSLIsCounter : TypeAttr {
+  let Spellings = [CXX11<"hlsl", "is_counter">];
+  let LangOpts = [HLSL];
+  let Documentation = [InternalOnly];
+}
+
 def HLSLGroupSharedAddressSpace : TypeAttr {
   let Spellings = [CustomKeyword<"groupshared">];
   let Subjects = SubjectList<[Var]>;

clang/lib/AST/ItaniumMangle.cpp

@@ -4624,6 +4624,8 @@ void CXXNameMangler::mangleType(const HLSLAttributedResourceType *T) {
     Str += "_ROV";
   if (Attrs.RawBuffer)
     Str += "_Raw";
+  if (Attrs.IsCounter)
+    Str += "_Counter";
   if (T->hasContainedType())
     Str += "_CT";
   mangleVendorQualifier(Str);

clang/lib/AST/TypePrinter.cpp

@@ -2033,6 +2033,7 @@ void TypePrinter::printAttributedAfter(const AttributedType *T,
   case attr::HLSLROV:
   case attr::HLSLRawBuffer:
   case attr::HLSLContainedType:
+  case attr::HLSLIsCounter:
     llvm_unreachable("HLSL resource type attributes handled separately");
 
   case attr::OpenCLPrivateAddressSpace:
@@ -2181,6 +2182,8 @@ void TypePrinter::printHLSLAttributedResourceAfter(
     OS << " [[hlsl::is_rov]]";
   if (Attrs.RawBuffer)
     OS << " [[hlsl::raw_buffer]]";
+  if (Attrs.IsCounter)
+    OS << " [[hlsl::is_counter]]";
 
   QualType ContainedTy = T->getContainedType();
   if (!ContainedTy.isNull()) {

clang/lib/CodeGen/Targets/SPIR.cpp

@@ -486,6 +486,12 @@ llvm::Type *CommonSPIRTargetCodeGenInfo::getHLSLType(
       return getSPIRVImageTypeFromHLSLResource(ResAttrs, ContainedTy, CGM);
     }
 
+    if (ResAttrs.IsCounter) {
+      llvm::Type *ElemType = llvm::Type::getInt32Ty(Ctx);
+      uint32_t StorageClass = /* StorageBuffer storage class */ 12;
+      return llvm::TargetExtType::get(Ctx, "spirv.VulkanBuffer", {ElemType},
+                                      {StorageClass, true});
+    }
     llvm::Type *ElemType = CGM.getTypes().ConvertTypeForMem(ContainedTy);
     llvm::ArrayType *RuntimeArrayType = llvm::ArrayType::get(ElemType, 0);
     uint32_t StorageClass = /* StorageBuffer storage class */ 12;

clang/lib/Sema/HLSLBuiltinTypeDeclBuilder.cpp

@@ -57,6 +57,15 @@ CXXConstructorDecl *lookupCopyConstructor(QualType ResTy) {
       return CD;
   return nullptr;
 }
+
+/// Set up common members and attributes for buffer types
+static bool resourceHasCounter(const CXXRecordDecl *Decl) {
+  StringRef Name = Decl->getName();
+  return Name == "RWStructuredBuffer" || Name == "AppendStructuredBuffer" ||
+         Name == "ConsumeStructuredBuffer" ||
+         Name == "RasterizerOrderedStructuredBuffer";
+}
+
 } // namespace
 
 // Builder for template arguments of builtin types. Used internally
@@ -138,7 +147,16 @@ struct BuiltinTypeMethodBuilder {
   //   LastStmt - refers to the last statement in the method body; referencing
   //              LastStmt will remove the statement from the method body since
   //              it will be linked from the new expression being constructed.
-  enum class PlaceHolder { _0, _1, _2, _3, _4, Handle = 128, LastStmt };
+  enum class PlaceHolder {
+    _0,
+    _1,
+    _2,
+    _3,
+    _4,
+    Handle = 128,
+    CounterHandle,
+    LastStmt
+  };
 
   Expr *convertPlaceholder(PlaceHolder PH);
   Expr *convertPlaceholder(LocalVar &Var);
@@ -178,10 +196,17 @@ struct BuiltinTypeMethodBuilder {
   template <typename ResourceT, typename ValueT>
   BuiltinTypeMethodBuilder &setHandleFieldOnResource(ResourceT ResourceRecord,
                                                      ValueT HandleValue);
+  template <typename T>
+  BuiltinTypeMethodBuilder &
+  accessCounterHandleFieldOnResource(T ResourceRecord);
+  template <typename ResourceT, typename ValueT>
+  BuiltinTypeMethodBuilder &
+  setCounterHandleFieldOnResource(ResourceT ResourceRecord, ValueT HandleValue);
   template <typename T> BuiltinTypeMethodBuilder &returnValue(T ReturnValue);
   BuiltinTypeMethodBuilder &returnThis();
   BuiltinTypeDeclBuilder &finalize();
   Expr *getResourceHandleExpr();
+  Expr *getResourceCounterHandleExpr();
 
 private:
   void createDecl();
@@ -346,6 +371,8 @@ TemplateParameterListBuilder::finalizeTemplateArgs(ConceptDecl *CD) {
 Expr *BuiltinTypeMethodBuilder::convertPlaceholder(PlaceHolder PH) {
   if (PH == PlaceHolder::Handle)
     return getResourceHandleExpr();
+  if (PH == PlaceHolder::CounterHandle)
+    return getResourceCounterHandleExpr();
 
   if (PH == PlaceHolder::LastStmt) {
     assert(!StmtsList.empty() && "no statements in the list");
@@ -467,6 +494,18 @@ Expr *BuiltinTypeMethodBuilder::getResourceHandleExpr() {
                                     OK_Ordinary);
 }
 
+Expr *BuiltinTypeMethodBuilder::getResourceCounterHandleExpr() {
+  ensureCompleteDecl();
+
+  ASTContext &AST = DeclBuilder.SemaRef.getASTContext();
+  CXXThisExpr *This = CXXThisExpr::Create(
+      AST, SourceLocation(), Method->getFunctionObjectParameterType(), true);
+  FieldDecl *HandleField = DeclBuilder.getResourceCounterHandleField();
+  return MemberExpr::CreateImplicit(AST, This, false, HandleField,
+                                    HandleField->getType(), VK_LValue,
+                                    OK_Ordinary);
+}
+
 BuiltinTypeMethodBuilder &
 BuiltinTypeMethodBuilder::declareLocalVar(LocalVar &Var) {
   ensureCompleteDecl();
@@ -583,6 +622,44 @@ BuiltinTypeMethodBuilder::setHandleFieldOnResource(ResourceT ResourceRecord,
   return *this;
 }
 
+template <typename T>
+BuiltinTypeMethodBuilder &
+BuiltinTypeMethodBuilder::accessCounterHandleFieldOnResource(T ResourceRecord) {
+  ensureCompleteDecl();
+
+  Expr *ResourceExpr = convertPlaceholder(ResourceRecord);
+
+  ASTContext &AST = DeclBuilder.SemaRef.getASTContext();
+  FieldDecl *HandleField = DeclBuilder.getResourceCounterHandleField();
+  MemberExpr *HandleExpr = MemberExpr::CreateImplicit(
+      AST, ResourceExpr, false, HandleField, HandleField->getType(), VK_LValue,
+      OK_Ordinary);
+  StmtsList.push_back(HandleExpr);
+  return *this;
+}
+
+template <typename ResourceT, typename ValueT>
+BuiltinTypeMethodBuilder &
+BuiltinTypeMethodBuilder::setCounterHandleFieldOnResource(
+    ResourceT ResourceRecord, ValueT HandleValue) {
+  ensureCompleteDecl();
+
+  Expr *ResourceExpr = convertPlaceholder(ResourceRecord);
+  Expr *HandleValueExpr = convertPlaceholder(HandleValue);
+
+  ASTContext &AST = DeclBuilder.SemaRef.getASTContext();
+  FieldDecl *HandleField = DeclBuilder.getResourceCounterHandleField();
+  MemberExpr *HandleMemberExpr = MemberExpr::CreateImplicit(
+      AST, ResourceExpr, false, HandleField, HandleField->getType(), VK_LValue,
+      OK_Ordinary);
+  Stmt *AssignStmt = BinaryOperator::Create(
+      DeclBuilder.SemaRef.getASTContext(), HandleMemberExpr, HandleValueExpr,
+      BO_Assign, HandleMemberExpr->getType(), ExprValueKind::VK_PRValue,
+      ExprObjectKind::OK_Ordinary, SourceLocation(), FPOptionsOverride());
+  StmtsList.push_back(AssignStmt);
+  return *this;
+}
+
 template <typename T>
 BuiltinTypeMethodBuilder &BuiltinTypeMethodBuilder::returnValue(T ReturnValue) {
   ensureCompleteDecl();
@@ -722,6 +799,15 @@ BuiltinTypeDeclBuilder::addMemberVariable(StringRef Name, QualType Type,
   return *this;
 }
 
+BuiltinTypeDeclBuilder &BuiltinTypeDeclBuilder::addHandleMembers(
+    ResourceClass RC, bool IsROV, bool RawBuffer, AccessSpecifier Access) {
+  addHandleMember(RC, IsROV, RawBuffer, Access);
+  if (resourceHasCounter(Record)) {
+    addCounterHandleMember(RC, IsROV, RawBuffer, Access);
+  }
+  return *this;
+}
+
 BuiltinTypeDeclBuilder &BuiltinTypeDeclBuilder::addHandleMember(
     ResourceClass RC, bool IsROV, bool RawBuffer, AccessSpecifier Access) {
   assert(!Record->isCompleteDefinition() && "record is already complete");
@@ -745,6 +831,30 @@ BuiltinTypeDeclBuilder &BuiltinTypeDeclBuilder::addHandleMember(
   return *this;
 }
 
+BuiltinTypeDeclBuilder &BuiltinTypeDeclBuilder::addCounterHandleMember(
+    ResourceClass RC, bool IsROV, bool RawBuffer, AccessSpecifier Access) {
+  assert(!Record->isCompleteDefinition() && "record is already complete");
+
+  ASTContext &Ctx = SemaRef.getASTContext();
+  TypeSourceInfo *ElementTypeInfo =
+      Ctx.getTrivialTypeSourceInfo(getHandleElementType(), SourceLocation());
+
+  // add handle member with resource type attributes
+  QualType AttributedResTy = QualType();
+  SmallVector<const Attr *> Attrs = {
+      HLSLResourceClassAttr::CreateImplicit(Ctx, RC),
+      IsROV ? HLSLROVAttr::CreateImplicit(Ctx) : nullptr,
+      RawBuffer ? HLSLRawBufferAttr::CreateImplicit(Ctx) : nullptr,
+      ElementTypeInfo
+          ? HLSLContainedTypeAttr::CreateImplicit(Ctx, ElementTypeInfo)
+          : nullptr,
+      HLSLIsCounterAttr::CreateImplicit(Ctx)};
+  if (CreateHLSLAttributedResourceType(SemaRef, Ctx.HLSLResourceTy, Attrs,
+                                       AttributedResTy))
+    addMemberVariable("__counter_handle", AttributedResTy, {}, Access);
+  return *this;
+}
+
 // Adds default constructor to the resource class:
 // Resource::Resource()
 BuiltinTypeDeclBuilder &BuiltinTypeDeclBuilder::addDefaultHandleConstructor() {
@@ -848,12 +958,18 @@ BuiltinTypeDeclBuilder &BuiltinTypeDeclBuilder::addCopyConstructor() {
 
   using PH = BuiltinTypeMethodBuilder::PlaceHolder;
 
-  return BuiltinTypeMethodBuilder(*this, /*Name=*/"", AST.VoidTy,
-                                  /*IsConst=*/false, /*IsCtor=*/true)
-      .addParam("other", ConstRecordRefType)
+  BuiltinTypeMethodBuilder MMB(*this, /*Name=*/"", AST.VoidTy,
+                               /*IsConst=*/false, /*IsCtor=*/true);
+  MMB.addParam("other", ConstRecordRefType)
       .accessHandleFieldOnResource(PH::_0)
-      .assign(PH::Handle, PH::LastStmt)
-      .finalize();
+      .assign(PH::Handle, PH::LastStmt);
+
+  if (getResourceCounterHandleField()) {
+    MMB.accessCounterHandleFieldOnResource(PH::_0).assign(PH::CounterHandle,
+                                                          PH::LastStmt);
+  }
+
+  return MMB.finalize();
 }
 
 BuiltinTypeDeclBuilder &BuiltinTypeDeclBuilder::addCopyAssignmentOperator() {
@@ -868,12 +984,17 @@ BuiltinTypeDeclBuilder &BuiltinTypeDeclBuilder::addCopyAssignmentOperator() {
 
   using PH = BuiltinTypeMethodBuilder::PlaceHolder;
   DeclarationName Name = AST.DeclarationNames.getCXXOperatorName(OO_Equal);
-  return BuiltinTypeMethodBuilder(*this, Name, RecordRefType)
-      .addParam("other", ConstRecordRefType)
+  BuiltinTypeMethodBuilder MMB(*this, Name, RecordRefType);
+  MMB.addParam("other", ConstRecordRefType)
       .accessHandleFieldOnResource(PH::_0)
-      .assign(PH::Handle, PH::LastStmt)
-      .returnThis()
-      .finalize();
+      .assign(PH::Handle, PH::LastStmt);
+
+  if (getResourceCounterHandleField()) {
+    MMB.accessCounterHandleFieldOnResource(PH::_0).assign(PH::CounterHandle,
+                                                          PH::LastStmt);
+  }
+
+  return MMB.returnThis().finalize();
 }
 
 BuiltinTypeDeclBuilder &BuiltinTypeDeclBuilder::addArraySubscriptOperators() {
@@ -909,6 +1030,13 @@ FieldDecl *BuiltinTypeDeclBuilder::getResourceHandleField() const {
   return I->second;
 }
 
+FieldDecl *BuiltinTypeDeclBuilder::getResourceCounterHandleField() const {
+  auto I = Fields.find("__counter_handle");
+  if (I == Fields.end())
+    return nullptr;
+  return I->second;
+}
+
 QualType BuiltinTypeDeclBuilder::getFirstTemplateTypeParam() {
   assert(Template && "record it not a template");
   if (const auto *TTD = dyn_cast<TemplateTypeParmDecl>(

clang/lib/Sema/HLSLBuiltinTypeDeclBuilder.h

@@ -72,8 +72,14 @@ class BuiltinTypeDeclBuilder {
                     AccessSpecifier Access = AccessSpecifier::AS_private);
 
   BuiltinTypeDeclBuilder &
+  addHandleMembers(ResourceClass RC, bool IsROV, bool RawBuffer,
+                   AccessSpecifier Access = AccessSpecifier::AS_private);
+  BuiltinTypeDeclBuilder &
   addHandleMember(ResourceClass RC, bool IsROV, bool RawBuffer,
                   AccessSpecifier Access = AccessSpecifier::AS_private);
+  BuiltinTypeDeclBuilder &
+  addCounterHandleMember(ResourceClass RC, bool IsROV, bool RawBuffer,
+                         AccessSpecifier Access = AccessSpecifier::AS_private);
   BuiltinTypeDeclBuilder &addArraySubscriptOperators();
 
   // Builtin types constructors
@@ -96,6 +102,7 @@ class BuiltinTypeDeclBuilder {
 
 private:
   FieldDecl *getResourceHandleField() const;
+  FieldDecl *getResourceCounterHandleField() const;
   QualType getFirstTemplateTypeParam();
   QualType getHandleElementType();
   Expr *getConstantIntExpr(int value);

clang/lib/Sema/HLSLExternalSemaSource.cpp

@@ -130,7 +130,7 @@ static BuiltinTypeDeclBuilder setupBufferType(CXXRecordDecl *Decl, Sema &S,
                                               ResourceClass RC, bool IsROV,
                                               bool RawBuffer) {
   return BuiltinTypeDeclBuilder(S, Decl)
-      .addHandleMember(RC, IsROV, RawBuffer)
+      .addHandleMembers(RC, IsROV, RawBuffer)
       .addDefaultHandleConstructor()
       .addCopyConstructor()
       .addCopyAssignmentOperator()

clang/lib/Sema/SemaHLSL.cpp

@@ -1811,6 +1811,13 @@ bool clang::CreateHLSLAttributedResourceType(
       }
       ResAttrs.RawBuffer = true;
       break;
+    case attr::HLSLIsCounter:
+      if (ResAttrs.IsCounter) {
+        S.Diag(A->getLocation(), diag::warn_duplicate_attribute_exact) << A;
+        return false;
+      }
+      ResAttrs.IsCounter = true;
+      break;
     case attr::HLSLContainedType: {
       const HLSLContainedTypeAttr *CTAttr = cast<HLSLContainedTypeAttr>(A);
       QualType Ty = CTAttr->getType();
@@ -1903,6 +1910,10 @@ bool SemaHLSL::handleResourceTypeAttr(QualType T, const ParsedAttr &AL) {
     A = HLSLRawBufferAttr::Create(getASTContext(), ACI);
     break;
 
+  case ParsedAttr::AT_HLSLIsCounter:
+    A = HLSLIsCounterAttr::Create(getASTContext(), ACI);
+    break;
+
   case ParsedAttr::AT_HLSLContainedType: {
     if (AL.getNumArgs() != 1 && !AL.hasParsedType()) {
       Diag(AL.getLoc(), diag::err_attribute_wrong_number_arguments) << AL << 1;