[HLSL] Add [[hlsl::contained_type()]] attribute (#108456)

hekota · web-flow · commit 0a7a1ef2205c · 2024-09-16T13:39:13.000-07:00
Introducing a new HLSL resource type attribute `[[contained_type(T)]]` which describes the "contained type" of a buffer or resource type. Specifically, the attribute will be used on the resource handle in templated buffer types like so: template <typename T> struct RWBuffer { __hlsl_resource_t [[hlsl::contained_type(T)]] [[hlsl::resource_class(UAV)]] h; }; Fixes #104855
diff --git a/clang/include/clang/AST/TypeLoc.h b/clang/include/clang/AST/TypeLoc.h
@@ -942,6 +942,7 @@ class BTFTagAttributedTypeLoc
 
 struct HLSLAttributedResourceLocInfo {
   SourceRange Range;
+  TypeSourceInfo *ContainedTyInfo;
 };
 
 /// Type source information for HLSL attributed resource type.
@@ -952,8 +953,11 @@ class HLSLAttributedResourceTypeLoc
 public:
   TypeLoc getWrappedLoc() const { return getInnerTypeLoc(); }
 
-  TypeLoc getContainedLoc() const {
-    return TypeLoc(getTypePtr()->getContainedType(), getNonLocalData());
+  TypeSourceInfo *getContainedTypeSourceInfo() const {
+    return getLocalData()->ContainedTyInfo;
+  }
+  void setContainedTypeSourceInfo(TypeSourceInfo *TSI) const {
+    getLocalData()->ContainedTyInfo = TSI;
   }
 
   void setSourceRange(const SourceRange &R) { getLocalData()->Range = R; }
diff --git a/clang/include/clang/Basic/Attr.td b/clang/include/clang/Basic/Attr.td
@@ -4667,6 +4667,13 @@ def HLSLResourceClass : TypeAttr {
   let Documentation = [InternalOnly];
 }
 
+def HLSLContainedType : TypeAttr {
+  let Spellings = [CXX11<"hlsl", "contained_type">];
+  let LangOpts = [HLSL];
+  let Args = [TypeArgument<"Type", /*opt=*/0>];
+  let Documentation = [InternalOnly];
+}
+
 def HLSLGroupSharedAddressSpace : TypeAttr {
   let Spellings = [CustomKeyword<"groupshared">];
   let Subjects = SubjectList<[Var]>;
diff --git a/clang/include/clang/Sema/SemaHLSL.h b/clang/include/clang/Sema/SemaHLSL.h
@@ -16,6 +16,7 @@
 #include "clang/AST/ASTFwd.h"
 #include "clang/AST/Attr.h"
 #include "clang/AST/Type.h"
+#include "clang/AST/TypeLoc.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Sema/SemaBase.h"
 #include "llvm/ADT/SmallVector.h"
@@ -30,9 +31,9 @@ class Scope;
 
 // FIXME: This can be hidden (as static function in SemaHLSL.cpp) once we no
 // longer need to create builtin buffer types in HLSLExternalSemaSource.
-bool CreateHLSLAttributedResourceType(Sema &S, QualType Wrapped,
-                                      ArrayRef<const Attr *> AttrList,
-                                      QualType &ResType);
+bool CreateHLSLAttributedResourceType(
+    Sema &S, QualType Wrapped, ArrayRef<const Attr *> AttrList,
+    QualType &ResType, HLSLAttributedResourceLocInfo *LocInfo = nullptr);
 
 class SemaHLSL : public SemaBase {
 public:
@@ -73,7 +74,8 @@ class SemaHLSL : public SemaBase {
 
   bool CheckBuiltinFunctionCall(unsigned BuiltinID, CallExpr *TheCall);
   QualType ProcessResourceTypeAttributes(QualType Wrapped);
-  SourceLocation TakeLocForHLSLAttribute(const HLSLAttributedResourceType *RT);
+  HLSLAttributedResourceLocInfo
+  TakeLocForHLSLAttribute(const HLSLAttributedResourceType *RT);
 
   // HLSL Type trait implementations
   bool IsScalarizedLayoutCompatible(QualType T1, QualType T2) const;
@@ -90,9 +92,10 @@ class SemaHLSL : public SemaBase {
   // This is a list to collect them.
   llvm::SmallVector<const Attr *> HLSLResourcesTypeAttrs;
 
-  /// SourceLocation corresponding to HLSLAttributedResourceTypeLocs that we
+  /// TypeLoc data for HLSLAttributedResourceType instances that we
   /// have not yet populated.
-  llvm::DenseMap<const HLSLAttributedResourceType *, SourceLocation>
+  llvm::DenseMap<const HLSLAttributedResourceType *,
+                 HLSLAttributedResourceLocInfo>
       LocsForHLSLAttributedResources;
 };
 
diff --git a/clang/lib/AST/TypePrinter.cpp b/clang/lib/AST/TypePrinter.cpp
@@ -1945,6 +1945,7 @@ void TypePrinter::printAttributedAfter(const AttributedType *T,
 
   case attr::HLSLResourceClass:
   case attr::HLSLROV:
+  case attr::HLSLContainedType:
     llvm_unreachable("HLSL resource type attributes handled separately");
 
   case attr::OpenCLPrivateAddressSpace:
@@ -2078,6 +2079,14 @@ void TypePrinter::printHLSLAttributedResourceAfter(
      << ")]]";
   if (Attrs.IsROV)
     OS << " [[hlsl::is_rov]]";
+
+  QualType ContainedTy = T->getContainedType();
+  if (!ContainedTy.isNull()) {
+    OS << " [[hlsl::contained_type(";
+    printBefore(ContainedTy, OS);
+    printAfter(ContainedTy, OS);
+    OS << ")]]";
+  }
 }
 
 void TypePrinter::printObjCInterfaceBefore(const ObjCInterfaceType *T,
diff --git a/clang/lib/Sema/HLSLExternalSemaSource.cpp b/clang/lib/Sema/HLSLExternalSemaSource.cpp
@@ -16,6 +16,7 @@
 #include "clang/AST/Type.h"
 #include "clang/Basic/AttrKinds.h"
 #include "clang/Basic/HLSLRuntime.h"
+#include "clang/Basic/SourceLocation.h"
 #include "clang/Sema/Lookup.h"
 #include "clang/Sema/Sema.h"
 #include "clang/Sema/SemaHLSL.h"
@@ -114,19 +115,30 @@ struct BuiltinTypeDeclBuilder {
                   AccessSpecifier Access = AccessSpecifier::AS_private) {
     if (Record->isCompleteDefinition())
       return *this;
+
+    TypeSourceInfo *ElementTypeInfo = nullptr;
+
     QualType Ty = Record->getASTContext().VoidPtrTy;
     if (Template) {
       if (const auto *TTD = dyn_cast<TemplateTypeParmDecl>(
-              Template->getTemplateParameters()->getParam(0)))
+              Template->getTemplateParameters()->getParam(0))) {
         Ty = Record->getASTContext().getPointerType(
             QualType(TTD->getTypeForDecl(), 0));
+        QualType ElemType = QualType(TTD->getTypeForDecl(), 0);
+        ElementTypeInfo = S.getASTContext().getTrivialTypeSourceInfo(
+            ElemType, SourceLocation());
+      }
     }
 
     // add handle member with resource type attributes
     QualType AttributedResTy = QualType();
     SmallVector<const Attr *> Attrs = {
         HLSLResourceClassAttr::CreateImplicit(Record->getASTContext(), RC),
-        IsROV ? HLSLROVAttr::CreateImplicit(Record->getASTContext()) : nullptr};
+        IsROV ? HLSLROVAttr::CreateImplicit(Record->getASTContext()) : nullptr,
+        ElementTypeInfo ? HLSLContainedTypeAttr::CreateImplicit(
+                              Record->getASTContext(), ElementTypeInfo)
+                        : nullptr,
+    };
     Attr *ResourceAttr =
         HLSLResourceAttr::CreateImplicit(Record->getASTContext(), RK);
     if (CreateHLSLAttributedResourceType(S, Ty, Attrs, AttributedResTy))
diff --git a/clang/lib/Sema/SemaHLSL.cpp b/clang/lib/Sema/SemaHLSL.cpp
@@ -10,12 +10,15 @@
 
 #include "clang/Sema/SemaHLSL.h"
 #include "clang/AST/ASTContext.h"
+#include "clang/AST/Attr.h"
+#include "clang/AST/Attrs.inc"
 #include "clang/AST/Decl.h"
 #include "clang/AST/DeclBase.h"
 #include "clang/AST/DeclCXX.h"
 #include "clang/AST/Expr.h"
 #include "clang/AST/RecursiveASTVisitor.h"
 #include "clang/AST/Type.h"
+#include "clang/AST/TypeLoc.h"
 #include "clang/Basic/Builtins.h"
 #include "clang/Basic/DiagnosticSema.h"
 #include "clang/Basic/LLVM.h"
@@ -564,18 +567,23 @@ void SemaHLSL::handleShaderAttr(Decl *D, const ParsedAttr &AL) {
     D->addAttr(NewAttr);
 }
 
-bool clang::CreateHLSLAttributedResourceType(Sema &S, QualType Wrapped,
-                                             ArrayRef<const Attr *> AttrList,
-                                             QualType &ResType) {
+bool clang::CreateHLSLAttributedResourceType(
+    Sema &S, QualType Wrapped, ArrayRef<const Attr *> AttrList,
+    QualType &ResType, HLSLAttributedResourceLocInfo *LocInfo) {
   assert(AttrList.size() && "expected list of resource attributes");
 
-  QualType Contained = QualType();
+  QualType ContainedTy = QualType();
+  TypeSourceInfo *ContainedTyInfo;
+  SourceLocation LocBegin = AttrList[0]->getRange().getBegin();
+  SourceLocation LocEnd = AttrList[0]->getRange().getEnd();
+
   HLSLAttributedResourceType::Attributes ResAttrs = {};
 
   bool HasResourceClass = false;
   for (const Attr *A : AttrList) {
     if (!A)
       continue;
+    LocEnd = A->getRange().getEnd();
     switch (A->getKind()) {
     case attr::HLSLResourceClass: {
       llvm::dxil::ResourceClass RC =
@@ -598,6 +606,20 @@ bool clang::CreateHLSLAttributedResourceType(Sema &S, QualType Wrapped,
       }
       ResAttrs.IsROV = true;
       break;
+    case attr::HLSLContainedType: {
+      const HLSLContainedTypeAttr *CTAttr = cast<HLSLContainedTypeAttr>(A);
+      QualType Ty = CTAttr->getType();
+      if (!ContainedTy.isNull()) {
+        S.Diag(A->getLocation(), ContainedTy == Ty
+                                     ? diag::warn_duplicate_attribute_exact
+                                     : diag::warn_duplicate_attribute)
+            << A;
+        return false;
+      }
+      ContainedTy = Ty;
+      ContainedTyInfo = CTAttr->getTypeLoc();
+      break;
+    }
     default:
       llvm_unreachable("unhandled resource attribute type");
     }
@@ -609,8 +631,13 @@ bool clang::CreateHLSLAttributedResourceType(Sema &S, QualType Wrapped,
     return false;
   }
 
-  ResType = S.getASTContext().getHLSLAttributedResourceType(Wrapped, Contained,
-                                                            ResAttrs);
+  ResType = S.getASTContext().getHLSLAttributedResourceType(
+      Wrapped, ContainedTy, ResAttrs);
+
+  if (LocInfo) {
+    LocInfo->Range = SourceRange(LocBegin, LocEnd);
+    LocInfo->ContainedTyInfo = ContainedTyInfo;
+  }
   return true;
 }
 
@@ -647,9 +674,27 @@ bool SemaHLSL::handleResourceTypeAttr(const ParsedAttr &AL) {
     A = HLSLResourceClassAttr::Create(getASTContext(), RC, AL.getLoc());
     break;
   }
+
   case ParsedAttr::AT_HLSLROV:
     A = HLSLROVAttr::Create(getASTContext(), AL.getLoc());
     break;
+
+  case ParsedAttr::AT_HLSLContainedType: {
+    if (AL.getNumArgs() != 1 && !AL.hasParsedType()) {
+      Diag(AL.getLoc(), diag::err_attribute_wrong_number_arguments) << AL << 1;
+      return false;
+    }
+
+    TypeSourceInfo *TSI = nullptr;
+    QualType QT = SemaRef.GetTypeFromParser(AL.getTypeArg(), &TSI);
+    assert(TSI && "no type source info for attribute argument");
+    if (SemaRef.RequireCompleteType(TSI->getTypeLoc().getBeginLoc(), QT,
+                                    diag::err_incomplete_type))
+      return false;
+    A = HLSLContainedTypeAttr::Create(getASTContext(), TSI, AL.getLoc());
+    break;
+  }
+
   default:
     llvm_unreachable("unhandled HLSL attribute");
   }
@@ -664,30 +709,34 @@ QualType SemaHLSL::ProcessResourceTypeAttributes(QualType CurrentType) {
     return CurrentType;
 
   QualType QT = CurrentType;
+  HLSLAttributedResourceLocInfo LocInfo;
   if (CreateHLSLAttributedResourceType(SemaRef, CurrentType,
-                                       HLSLResourcesTypeAttrs, QT)) {
+                                       HLSLResourcesTypeAttrs, QT, &LocInfo)) {
     const HLSLAttributedResourceType *RT =
-        dyn_cast<HLSLAttributedResourceType>(QT.getTypePtr());
-    // Use the location of the first attribute as the location of the aggregated
-    // type. The attributes are stored in HLSLResourceTypeAttrs in the same
-    // order as they are parsed.
-    SourceLocation Loc = HLSLResourcesTypeAttrs[0]->getLoc();
-    LocsForHLSLAttributedResources.insert(std::pair(RT, Loc));
+        cast<HLSLAttributedResourceType>(QT.getTypePtr());
+
+    // Temporarily store TypeLoc information for the new type.
+    // It will be transferred to HLSLAttributesResourceTypeLoc
+    // shortly after the type is created by TypeSpecLocFiller which
+    // will call the TakeLocForHLSLAttribute method below.
+    LocsForHLSLAttributedResources.insert(std::pair(RT, LocInfo));
   }
   HLSLResourcesTypeAttrs.clear();
   return QT;
 }
 
 // Returns source location for the HLSLAttributedResourceType
-SourceLocation
+HLSLAttributedResourceLocInfo
 SemaHLSL::TakeLocForHLSLAttribute(const HLSLAttributedResourceType *RT) {
+  HLSLAttributedResourceLocInfo LocInfo = {};
   auto I = LocsForHLSLAttributedResources.find(RT);
   if (I != LocsForHLSLAttributedResources.end()) {
-    SourceLocation Loc = I->second;
+    LocInfo = I->second;
     LocsForHLSLAttributedResources.erase(I);
-    return Loc;
+    return LocInfo;
   }
-  return SourceLocation();
+  LocInfo.Range = SourceRange();
+  return LocInfo;
 }
 
 struct RegisterBindingFlags {
diff --git a/clang/lib/Sema/SemaType.cpp b/clang/lib/Sema/SemaType.cpp
@@ -5801,8 +5801,10 @@ static void fillAttributedTypeLoc(AttributedTypeLoc TL,
 
 static void fillHLSLAttributedResourceTypeLoc(HLSLAttributedResourceTypeLoc TL,
                                               TypeProcessingState &State) {
-  TL.setSourceRange(
-      State.getSema().HLSL().TakeLocForHLSLAttribute(TL.getTypePtr()));
+  HLSLAttributedResourceLocInfo LocInfo =
+      State.getSema().HLSL().TakeLocForHLSLAttribute(TL.getTypePtr());
+  TL.setSourceRange(LocInfo.Range);
+  TL.setContainedTypeSourceInfo(LocInfo.ContainedTyInfo);
 }
 
 static void fillMatrixTypeLoc(MatrixTypeLoc MTL,
@@ -8843,7 +8845,8 @@ static void processTypeAttrs(TypeProcessingState &state, QualType &type,
       break;
     }
     case ParsedAttr::AT_HLSLResourceClass:
-    case ParsedAttr::AT_HLSLROV: {
+    case ParsedAttr::AT_HLSLROV:
+    case ParsedAttr::AT_HLSLContainedType: {
       // Only collect HLSL resource type attributes that are in
       // decl-specifier-seq; do not collect attributes on declarations or those
       // that get to slide after declaration name.
diff --git a/clang/lib/Sema/TreeTransform.h b/clang/lib/Sema/TreeTransform.h
@@ -7477,8 +7477,17 @@ QualType TreeTransform<Derived>::TransformHLSLAttributedResourceType(
     return QualType();
 
   QualType ContainedTy = QualType();
-  if (!oldType->getContainedType().isNull())
-    ContainedTy = getDerived().TransformType(TLB, TL.getContainedLoc());
+  QualType OldContainedTy = oldType->getContainedType();
+  if (!OldContainedTy.isNull()) {
+    TypeSourceInfo *oldContainedTSI = TL.getContainedTypeSourceInfo();
+    if (!oldContainedTSI)
+      oldContainedTSI = getSema().getASTContext().getTrivialTypeSourceInfo(
+          OldContainedTy, SourceLocation());
+    TypeSourceInfo *ContainedTSI = getDerived().TransformType(oldContainedTSI);
+    if (!ContainedTSI)
+      return QualType();
+    ContainedTy = ContainedTSI->getType();
+  }
 
   QualType Result = TL.getType();
   if (getDerived().AlwaysRebuild() || WrappedTy != oldType->getWrappedType() ||
diff --git a/clang/test/AST/HLSL/RWBuffer-AST.hlsl b/clang/test/AST/HLSL/RWBuffer-AST.hlsl
@@ -1,7 +1,6 @@
 // RUN: %clang_cc1 -triple dxil-pc-shadermodel6.0-library -x hlsl -ast-dump -DEMPTY %s | FileCheck -check-prefix=EMPTY %s
 // RUN: %clang_cc1 -triple dxil-pc-shadermodel6.0-library -x hlsl -ast-dump %s | FileCheck %s
 
-
 // This test tests two different AST generations. The "EMPTY" test mode verifies
 // the AST generated by forward declaration of the HLSL types which happens on
 // initializing the HLSL external AST with an AST Context.
@@ -30,15 +29,15 @@ RWBuffer<float> Buffer;
 // CHECK-NEXT: CXXRecordDecl 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> <invalid sloc> implicit class RWBuffer definition
 
 // CHECK: FinalAttr 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> Implicit final
-// CHECK-NEXT: implicit h 'element_type * {{\[\[}}hlsl::resource_class(UAV)]]':'element_type *'
+// CHECK-NEXT: implicit h 'element_type * {{\[\[}}hlsl::resource_class(UAV)]] {{\[\[}}hlsl::contained_type(element_type)]]':'element_type *'
 // CHECK-NEXT: HLSLResourceAttr 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> Implicit TypedBuffer
 
 // CHECK: CXXMethodDecl 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> <invalid sloc> operator[] 'element_type &const (unsigned int) const'
 // CHECK-NEXT: ParmVarDecl 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> <invalid sloc> Idx 'unsigned int'
 // CHECK-NEXT: CompoundStmt 0x{{[0-9A-Fa-f]+}} <<invalid sloc>>
 // CHECK-NEXT: ReturnStmt 0x{{[0-9A-Fa-f]+}} <<invalid sloc>>
 // CHECK-NEXT: ArraySubscriptExpr 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> 'element_type' lvalue
-// CHECK-NEXT: MemberExpr 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> 'element_type * {{\[\[}}hlsl::resource_class(UAV)]]':'element_type *' lvalue .h 0x{{[0-9A-Fa-f]+}}
+// CHECK-NEXT: MemberExpr 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> 'element_type * {{\[\[}}hlsl::resource_class(UAV)]]  {{\[\[}}hlsl::contained_type(element_type)]]':'element_type *' lvalue .h 0x{{[0-9A-Fa-f]+}}
 // CHECK-NEXT: CXXThisExpr 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> 'const RWBuffer<element_type>' lvalue implicit this
 // CHECK-NEXT: DeclRefExpr 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> 'unsigned int' ParmVar 0x{{[0-9A-Fa-f]+}} 'Idx' 'unsigned int'
 // CHECK-NEXT: AlwaysInlineAttr 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> Implicit always_inline
@@ -48,7 +47,7 @@ RWBuffer<float> Buffer;
 // CHECK-NEXT: CompoundStmt 0x{{[0-9A-Fa-f]+}} <<invalid sloc>>
 // CHECK-NEXT: ReturnStmt 0x{{[0-9A-Fa-f]+}} <<invalid sloc>>
 // CHECK-NEXT: ArraySubscriptExpr 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> 'element_type' lvalue
-// CHECK-NEXT: MemberExpr 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> 'element_type * {{\[\[}}hlsl::resource_class(UAV)]]':'element_type *' lvalue .h 0x{{[0-9A-Fa-f]+}}
+// CHECK-NEXT: MemberExpr 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> 'element_type * {{\[\[}}hlsl::resource_class(UAV)]]  {{\[\[}}hlsl::contained_type(element_type)]]':'element_type *' lvalue .h 0x{{[0-9A-Fa-f]+}}
 // CHECK-NEXT: CXXThisExpr 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> 'RWBuffer<element_type>' lvalue implicit this
 // CHECK-NEXT: DeclRefExpr 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> 'unsigned int' ParmVar 0x{{[0-9A-Fa-f]+}} 'Idx' 'unsigned int'
 // CHECK-NEXT: AlwaysInlineAttr 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> Implicit always_inline
@@ -58,5 +57,5 @@ RWBuffer<float> Buffer;
 // CHECK: TemplateArgument type 'float'
 // CHECK-NEXT: BuiltinType 0x{{[0-9A-Fa-f]+}} 'float'
 // CHECK-NEXT: FinalAttr 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> Implicit final
-// CHECK-NEXT: FieldDecl 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> <invalid sloc> implicit referenced h 'float * {{\[\[}}hlsl::resource_class(UAV)]]':'float *'
+// CHECK-NEXT: FieldDecl 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> <invalid sloc> implicit referenced h 'float * {{\[\[}}hlsl::resource_class(UAV)]]  {{\[\[}}hlsl::contained_type(float)]]':'float *'
 // CHECK-NEXT: HLSLResourceAttr 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> Implicit TypedBuffer
diff --git a/clang/test/AST/HLSL/StructuredBuffer-AST.hlsl b/clang/test/AST/HLSL/StructuredBuffer-AST.hlsl
diff --git a/clang/test/ParserHLSL/hlsl_contained_type_attr.hlsl b/clang/test/ParserHLSL/hlsl_contained_type_attr.hlsl
diff --git a/clang/test/ParserHLSL/hlsl_contained_type_attr_error.hlsl b/clang/test/ParserHLSL/hlsl_contained_type_attr_error.hlsl
diff --git a/clang/test/ParserHLSL/hlsl_resource_handle_attrs.hlsl b/clang/test/ParserHLSL/hlsl_resource_handle_attrs.hlsl
