[HLSL] [SPIR-V] Add counter member for typed buffer

clang/include/clang/AST/TypeBase.h

@@ -6702,15 +6702,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);
@@ -6751,6 +6757,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

@@ -5074,6 +5074,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

@@ -2062,6 +2062,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:
@@ -2210,6 +2211,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

@@ -138,7 +138,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 +187,14 @@ struct BuiltinTypeMethodBuilder {
   template <typename ResourceT, typename ValueT>
   BuiltinTypeMethodBuilder &setHandleFieldOnResource(ResourceT ResourceRecord,
                                                      ValueT HandleValue);
+  template <typename T>
+  BuiltinTypeMethodBuilder &
+  accessCounterHandleFieldOnResource(T ResourceRecord);
   template <typename T> BuiltinTypeMethodBuilder &returnValue(T ReturnValue);
   BuiltinTypeMethodBuilder &returnThis();
   BuiltinTypeDeclBuilder &finalize();
   Expr *getResourceHandleExpr();
+  Expr *getResourceCounterHandleExpr();
 
 private:
   void createDecl();
@@ -346,6 +359,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 +482,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 +610,22 @@ 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 T>
 BuiltinTypeMethodBuilder &BuiltinTypeMethodBuilder::returnValue(T ReturnValue) {
   ensureCompleteDecl();
@@ -722,8 +765,31 @@ BuiltinTypeDeclBuilder::addMemberVariable(StringRef Name, QualType Type,
   return *this;
 }
 
+BuiltinTypeDeclBuilder &
+BuiltinTypeDeclBuilder::addBufferHandles(ResourceClass RC, bool IsROV,
+                                         bool RawBuffer, bool HasCounter,
+                                         AccessSpecifier Access) {
+  addHandleMember(RC, IsROV, RawBuffer, Access);
+  if (HasCounter)
+    addCounterHandleMember(RC, IsROV, RawBuffer, Access);
+  return *this;
+}
+
 BuiltinTypeDeclBuilder &BuiltinTypeDeclBuilder::addHandleMember(
     ResourceClass RC, bool IsROV, bool RawBuffer, AccessSpecifier Access) {
+  return addResourceMember("__handle", RC, IsROV, RawBuffer,
+                           /*IsCounter=*/false, Access);
+}
+
+BuiltinTypeDeclBuilder &BuiltinTypeDeclBuilder::addCounterHandleMember(
+    ResourceClass RC, bool IsROV, bool RawBuffer, AccessSpecifier Access) {
+  return addResourceMember("__counter_handle", RC, IsROV, RawBuffer,
+                           /*IsCounter=*/true, Access);
+}
+
+BuiltinTypeDeclBuilder &BuiltinTypeDeclBuilder::addResourceMember(
+    StringRef MemberName, ResourceClass RC, bool IsROV, bool RawBuffer,
+    bool IsCounter, AccessSpecifier Access) {
   assert(!Record->isCompleteDefinition() && "record is already complete");
 
   ASTContext &Ctx = SemaRef.getASTContext();
@@ -739,9 +805,12 @@ BuiltinTypeDeclBuilder &BuiltinTypeDeclBuilder::addHandleMember(
       ElementTypeInfo
           ? HLSLContainedTypeAttr::CreateImplicit(Ctx, ElementTypeInfo)
           : nullptr};
+  if (IsCounter)
+    Attrs.push_back(HLSLIsCounterAttr::CreateImplicit(Ctx));
+
   if (CreateHLSLAttributedResourceType(SemaRef, Ctx.HLSLResourceTy, Attrs,
                                        AttributedResTy))
-    addMemberVariable("__handle", AttributedResTy, {}, Access);
+    addMemberVariable(MemberName, AttributedResTy, {}, Access);
   return *this;
 }
 
@@ -844,12 +913,17 @@ 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() {
@@ -863,12 +937,16 @@ 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() {
@@ -903,6 +981,14 @@ FieldDecl *BuiltinTypeDeclBuilder::getResourceHandleField() const {
   return I->second;
 }
 
+FieldDecl *BuiltinTypeDeclBuilder::getResourceCounterHandleField() const {
+  auto I = Fields.find("__counter_handle");
+  if (I == Fields.end() ||
+      !I->second->getType()->isHLSLAttributedResourceType())
+    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,9 @@ class BuiltinTypeDeclBuilder {
                     AccessSpecifier Access = AccessSpecifier::AS_private);
 
   BuiltinTypeDeclBuilder &
-  addHandleMember(ResourceClass RC, bool IsROV, bool RawBuffer,
-                  AccessSpecifier Access = AccessSpecifier::AS_private);
+  addBufferHandles(ResourceClass RC, bool IsROV, bool RawBuffer,
+                   bool HasCounter,
+                   AccessSpecifier Access = AccessSpecifier::AS_private);
   BuiltinTypeDeclBuilder &addArraySubscriptOperators();
 
   // Builtin types constructors
@@ -95,7 +96,18 @@ class BuiltinTypeDeclBuilder {
   BuiltinTypeDeclBuilder &addConsumeMethod();
 
 private:
+  BuiltinTypeDeclBuilder &addResourceMember(StringRef MemberName,
+                                            ResourceClass RC, bool IsROV,
+                                            bool RawBuffer, bool IsCounter,
+                                            AccessSpecifier Access);
+  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);
   FieldDecl *getResourceHandleField() const;
+  FieldDecl *getResourceCounterHandleField() const;
   QualType getFirstTemplateTypeParam();
   QualType getHandleElementType();
   Expr *getConstantIntExpr(int value);