[HLSL] Support packoffset attribute in AST (#89836)

Add HLSLPackOffsetAttr to save packoffset in AST.

Since we have to parse the attribute manually in ParseHLSLAnnotations,
we could create the ParsedAttribute with a integer offset parameter
instead of string. This approach avoids parsing the string if the offset
is saved as a string in HLSLPackOffsetAttr.

For #57914.

Author: python3kgae

clang/include/clang/Basic/Attr.td

@@ -4415,6 +4415,18 @@ def HLSLResourceBinding: InheritableAttr {
   let Documentation = [HLSLResourceBindingDocs];
 }
 
+def HLSLPackOffset: HLSLAnnotationAttr {
+  let Spellings = [HLSLAnnotation<"packoffset">];
+  let LangOpts = [HLSL];
+  let Args = [IntArgument<"Subcomponent">, IntArgument<"Component">];
+  let Documentation = [HLSLPackOffsetDocs];
+  let AdditionalMembers = [{
+      unsigned getOffset() {
+        return subcomponent * 4 + component;
+      }
+  }];
+}
+
 def HLSLSV_DispatchThreadID: HLSLAnnotationAttr {
   let Spellings = [HLSLAnnotation<"SV_DispatchThreadID">];
   let Subjects = SubjectList<[ParmVar, Field]>;

clang/include/clang/Basic/AttrDocs.td

@@ -7408,6 +7408,26 @@ The full documentation is available here: https://docs.microsoft.com/en-us/windo
   }];
 }
 
+def HLSLPackOffsetDocs : Documentation {
+  let Category = DocCatFunction;
+  let Content = [{
+The packoffset attribute is used to change the layout of a cbuffer.
+Attribute spelling in HLSL is: ``packoffset( c[Subcomponent][.component] )``.
+A subcomponent is a register number, which is an integer. A component is in the form of [.xyzw].
+
+Examples:
+
+.. code-block:: c++
+
+  cbuffer A {
+    float3 a : packoffset(c0.y);
+    float4 b : packoffset(c4);
+  }
+
+The full documentation is available here: https://learn.microsoft.com/en-us/windows/win32/direct3dhlsl/dx-graphics-hlsl-variable-packoffset
+  }];
+}
+
 def HLSLSV_DispatchThreadIDDocs : Documentation {
   let Category = DocCatFunction;
   let Content = [{

clang/include/clang/Basic/DiagnosticGroups.td

@@ -1507,6 +1507,9 @@ def BranchProtection : DiagGroup<"branch-protection">;
 // Warnings for HLSL Clang extensions
 def HLSLExtension : DiagGroup<"hlsl-extensions">;
 
+// Warning for mix packoffset and non-packoffset.
+def HLSLMixPackOffset : DiagGroup<"mix-packoffset">;
+
 // Warnings for DXIL validation
 def DXILValidation : DiagGroup<"dxil-validation">;
 

clang/include/clang/Basic/DiagnosticParseKinds.td

@@ -1754,5 +1754,7 @@ def err_hlsl_separate_attr_arg_and_number : Error<"wrong argument format for hls
 def ext_hlsl_access_specifiers : ExtWarn<
   "access specifiers are a clang HLSL extension">,
   InGroup<HLSLExtension>;
+def err_hlsl_unsupported_component : Error<"invalid component '%0' used; expected 'x', 'y', 'z', or 'w'">;
+def err_hlsl_packoffset_invalid_reg : Error<"invalid resource class specifier '%0' for packoffset, expected 'c'">;
 
 } // end of Parser diagnostics

clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -12184,6 +12184,11 @@ def err_hlsl_init_priority_unsupported : Error<
 def err_hlsl_unsupported_register_type : Error<"invalid resource class specifier '%0' used; expected 'b', 's', 't', or 'u'">;
 def err_hlsl_unsupported_register_number : Error<"register number should be an integer">;
 def err_hlsl_expected_space : Error<"invalid space specifier '%0' used; expected 'space' followed by an integer, like space1">;
+def warn_hlsl_packoffset_mix : Warning<"cannot mix packoffset elements with nonpackoffset elements in a cbuffer">,
+    InGroup<HLSLMixPackOffset>;
+def err_hlsl_packoffset_overlap : Error<"packoffset overlap between %0, %1">;
+def err_hlsl_packoffset_cross_reg_boundary : Error<"packoffset cannot cross register boundary">;
+def err_hlsl_packoffset_alignment_mismatch : Error<"packoffset at 'y' not match alignment %0 required by %1">;
 def err_hlsl_pointers_unsupported : Error<
   "%select{pointers|references}0 are unsupported in HLSL">;
 

clang/lib/Parse/ParseHLSL.cpp

@@ -183,6 +183,94 @@ void Parser::ParseHLSLAnnotations(ParsedAttributes &Attrs,
       return;
     }
   } break;
+  case ParsedAttr::AT_HLSLPackOffset: {
+    // Parse 'packoffset( c[Subcomponent][.component] )'.
+    // Check '('.
+    if (ExpectAndConsume(tok::l_paren, diag::err_expected_lparen_after)) {
+      SkipUntil(tok::r_paren, StopAtSemi); // skip through )
+      return;
+    }
+    // Check c[Subcomponent] as an identifier.
+    if (!Tok.is(tok::identifier)) {
+      Diag(Tok.getLocation(), diag::err_expected) << tok::identifier;
+      SkipUntil(tok::r_paren, StopAtSemi); // skip through )
+      return;
+    }
+    StringRef OffsetStr = Tok.getIdentifierInfo()->getName();
+    SourceLocation SubComponentLoc = Tok.getLocation();
+    if (OffsetStr[0] != 'c') {
+      Diag(Tok.getLocation(), diag::err_hlsl_packoffset_invalid_reg)
+          << OffsetStr;
+      SkipUntil(tok::r_paren, StopAtSemi); // skip through )
+      return;
+    }
+    OffsetStr = OffsetStr.substr(1);
+    unsigned SubComponent = 0;
+    if (!OffsetStr.empty()) {
+      // Make sure SubComponent is a number.
+      if (OffsetStr.getAsInteger(10, SubComponent)) {
+        Diag(SubComponentLoc.getLocWithOffset(1),
+             diag::err_hlsl_unsupported_register_number);
+        SkipUntil(tok::r_paren, StopAtSemi); // skip through )
+        return;
+      }
+    }
+    unsigned Component = 0;
+    ConsumeToken(); // consume identifier.
+    SourceLocation ComponentLoc;
+    if (Tok.is(tok::period)) {
+      ConsumeToken(); // consume period.
+      if (!Tok.is(tok::identifier)) {
+        Diag(Tok.getLocation(), diag::err_expected) << tok::identifier;
+        SkipUntil(tok::r_paren, StopAtSemi); // skip through )
+        return;
+      }
+      StringRef ComponentStr = Tok.getIdentifierInfo()->getName();
+      ComponentLoc = Tok.getLocation();
+      ConsumeToken(); // consume identifier.
+      // Make sure Component is a single character.
+      if (ComponentStr.size() != 1) {
+        Diag(ComponentLoc, diag::err_hlsl_unsupported_component)
+            << ComponentStr;
+        SkipUntil(tok::r_paren, StopAtSemi); // skip through )
+        return;
+      }
+      switch (ComponentStr[0]) {
+      case 'x':
+      case 'r':
+        Component = 0;
+        break;
+      case 'y':
+      case 'g':
+        Component = 1;
+        break;
+      case 'z':
+      case 'b':
+        Component = 2;
+        break;
+      case 'w':
+      case 'a':
+        Component = 3;
+        break;
+      default:
+        Diag(ComponentLoc, diag::err_hlsl_unsupported_component)
+            << ComponentStr;
+        SkipUntil(tok::r_paren, StopAtSemi); // skip through )
+        return;
+      }
+    }
+    ASTContext &Ctx = Actions.getASTContext();
+    QualType SizeTy = Ctx.getSizeType();
+    uint64_t SizeTySize = Ctx.getTypeSize(SizeTy);
+    ArgExprs.push_back(IntegerLiteral::Create(
+        Ctx, llvm::APInt(SizeTySize, SubComponent), SizeTy, SubComponentLoc));
+    ArgExprs.push_back(IntegerLiteral::Create(
+        Ctx, llvm::APInt(SizeTySize, Component), SizeTy, ComponentLoc));
+    if (ExpectAndConsume(tok::r_paren, diag::err_expected)) {
+      SkipUntil(tok::r_paren, StopAtSemi); // skip through )
+      return;
+    }
+  } break;
   case ParsedAttr::UnknownAttribute:
     Diag(Loc, diag::err_unknown_hlsl_semantic) << II;
     return;

clang/lib/Sema/SemaDeclAttr.cpp

@@ -7309,6 +7309,55 @@ static void handleHLSLSV_DispatchThreadIDAttr(Sema &S, Decl *D,
   D->addAttr(::new (S.Context) HLSLSV_DispatchThreadIDAttr(S.Context, AL));
 }
 
+static void handleHLSLPackOffsetAttr(Sema &S, Decl *D, const ParsedAttr &AL) {
+  if (!isa<VarDecl>(D) || !isa<HLSLBufferDecl>(D->getDeclContext())) {
+    S.Diag(AL.getLoc(), diag::err_hlsl_attr_invalid_ast_node)
+        << AL << "shader constant in a constant buffer";
+    return;
+  }
+
+  uint32_t SubComponent;
+  if (!checkUInt32Argument(S, AL, AL.getArgAsExpr(0), SubComponent))
+    return;
+  uint32_t Component;
+  if (!checkUInt32Argument(S, AL, AL.getArgAsExpr(1), Component))
+    return;
+
+  QualType T = cast<VarDecl>(D)->getType().getCanonicalType();
+  // Check if T is an array or struct type.
+  // TODO: mark matrix type as aggregate type.
+  bool IsAggregateTy = (T->isArrayType() || T->isStructureType());
+
+  // Check Component is valid for T.
+  if (Component) {
+    unsigned Size = S.getASTContext().getTypeSize(T);
+    if (IsAggregateTy || Size > 128) {
+      S.Diag(AL.getLoc(), diag::err_hlsl_packoffset_cross_reg_boundary);
+      return;
+    } else {
+      // Make sure Component + sizeof(T) <= 4.
+      if ((Component * 32 + Size) > 128) {
+        S.Diag(AL.getLoc(), diag::err_hlsl_packoffset_cross_reg_boundary);
+        return;
+      }
+      QualType EltTy = T;
+      if (const auto *VT = T->getAs<VectorType>())
+        EltTy = VT->getElementType();
+      unsigned Align = S.getASTContext().getTypeAlign(EltTy);
+      if (Align > 32 && Component == 1) {
+        // NOTE: Component 3 will hit err_hlsl_packoffset_cross_reg_boundary.
+        // So we only need to check Component 1 here.
+        S.Diag(AL.getLoc(), diag::err_hlsl_packoffset_alignment_mismatch)
+            << Align << EltTy;
+        return;
+      }
+    }
+  }
+
+  D->addAttr(::new (S.Context)
+                 HLSLPackOffsetAttr(S.Context, AL, SubComponent, Component));
+}
+
 static void handleHLSLShaderAttr(Sema &S, Decl *D, const ParsedAttr &AL) {
   StringRef Str;
   SourceLocation ArgLoc;
@@ -9730,6 +9779,9 @@ ProcessDeclAttribute(Sema &S, Scope *scope, Decl *D, const ParsedAttr &AL,
   case ParsedAttr::AT_HLSLSV_DispatchThreadID:
     handleHLSLSV_DispatchThreadIDAttr(S, D, AL);
     break;
+  case ParsedAttr::AT_HLSLPackOffset:
+    handleHLSLPackOffsetAttr(S, D, AL);
+    break;
   case ParsedAttr::AT_HLSLShader:
     handleHLSLShaderAttr(S, D, AL);
     break;

clang/lib/Sema/SemaHLSL.cpp

@@ -39,9 +39,89 @@ Decl *SemaHLSL::ActOnStartBuffer(Scope *BufferScope, bool CBuffer,
   return Result;
 }
 
+// Calculate the size of a legacy cbuffer type based on
+// https://learn.microsoft.com/en-us/windows/win32/direct3dhlsl/dx-graphics-hlsl-packing-rules
+static unsigned calculateLegacyCbufferSize(const ASTContext &Context,
+                                           QualType T) {
+  unsigned Size = 0;
+  constexpr unsigned CBufferAlign = 128;
+  if (const RecordType *RT = T->getAs<RecordType>()) {
+    const RecordDecl *RD = RT->getDecl();
+    for (const FieldDecl *Field : RD->fields()) {
+      QualType Ty = Field->getType();
+      unsigned FieldSize = calculateLegacyCbufferSize(Context, Ty);
+      unsigned FieldAlign = 32;
+      if (Ty->isAggregateType())
+        FieldAlign = CBufferAlign;
+      Size = llvm::alignTo(Size, FieldAlign);
+      Size += FieldSize;
+    }
+  } else if (const ConstantArrayType *AT = Context.getAsConstantArrayType(T)) {
+    if (unsigned ElementCount = AT->getSize().getZExtValue()) {
+      unsigned ElementSize =
+          calculateLegacyCbufferSize(Context, AT->getElementType());
+      unsigned AlignedElementSize = llvm::alignTo(ElementSize, CBufferAlign);
+      Size = AlignedElementSize * (ElementCount - 1) + ElementSize;
+    }
+  } else if (const VectorType *VT = T->getAs<VectorType>()) {
+    unsigned ElementCount = VT->getNumElements();
+    unsigned ElementSize =
+        calculateLegacyCbufferSize(Context, VT->getElementType());
+    Size = ElementSize * ElementCount;
+  } else {
+    Size = Context.getTypeSize(T);
+  }
+  return Size;
+}
+
 void SemaHLSL::ActOnFinishBuffer(Decl *Dcl, SourceLocation RBrace) {
   auto *BufDecl = cast<HLSLBufferDecl>(Dcl);
   BufDecl->setRBraceLoc(RBrace);
+
+  // Validate packoffset.
+  llvm::SmallVector<std::pair<VarDecl *, HLSLPackOffsetAttr *>> PackOffsetVec;
+  bool HasPackOffset = false;
+  bool HasNonPackOffset = false;
+  for (auto *Field : BufDecl->decls()) {
+    VarDecl *Var = dyn_cast<VarDecl>(Field);
+    if (!Var)
+      continue;
+    if (Field->hasAttr<HLSLPackOffsetAttr>()) {
+      PackOffsetVec.emplace_back(Var, Field->getAttr<HLSLPackOffsetAttr>());
+      HasPackOffset = true;
+    } else {
+      HasNonPackOffset = true;
+    }
+  }
+
+  if (HasPackOffset && HasNonPackOffset)
+    Diag(BufDecl->getLocation(), diag::warn_hlsl_packoffset_mix);
+
+  if (HasPackOffset) {
+    ASTContext &Context = getASTContext();
+    // Make sure no overlap in packoffset.
+    // Sort PackOffsetVec by offset.
+    std::sort(PackOffsetVec.begin(), PackOffsetVec.end(),
+              [](const std::pair<VarDecl *, HLSLPackOffsetAttr *> &LHS,
+                 const std::pair<VarDecl *, HLSLPackOffsetAttr *> &RHS) {
+                return LHS.second->getOffset() < RHS.second->getOffset();
+              });
+
+    for (unsigned i = 0; i < PackOffsetVec.size() - 1; i++) {
+      VarDecl *Var = PackOffsetVec[i].first;
+      HLSLPackOffsetAttr *Attr = PackOffsetVec[i].second;
+      unsigned Size = calculateLegacyCbufferSize(Context, Var->getType());
+      unsigned Begin = Attr->getOffset() * 32;
+      unsigned End = Begin + Size;
+      unsigned NextBegin = PackOffsetVec[i + 1].second->getOffset() * 32;
+      if (End > NextBegin) {
+        VarDecl *NextVar = PackOffsetVec[i + 1].first;
+        Diag(NextVar->getLocation(), diag::err_hlsl_packoffset_overlap)
+            << NextVar << Var;
+      }
+    }
+  }
+
   SemaRef.PopDeclContext();
 }