[HLSL] Add resource binding attribute for HLSL.

The resource binding attribute is to set the virtual registers and logical register spaces resources in HLSL are bound to.
Format is ''register(ID,  space)'' like register(t3,  space1).
ID must be start with
t – for shader resource views (SRV),
s – for samplers,
u – for unordered access views (UAV),
b – for constant buffer views (CBV).

Register space is default to space0.

The full documentation is available here: https://docs.microsoft.com/en-us/windows/win32/direct3d12/resource-binding-in-hlsl

Reviewed By: aaron.ballman

Differential Revision: https://reviews.llvm.org/D130033

Author: python3kgae

clang/include/clang/Basic/Attr.td

@@ -145,6 +145,9 @@ def HLSLEntry
     : SubsetSubject<Function,
                     [{S->isExternallyVisible() && !isa<CXXMethodDecl>(S)}],
                     "global functions">;
+def HLSLBufferObj : SubsetSubject<HLSLBuffer,
+                    [{isa<HLSLBufferDecl>(S)}],
+                    "cbuffer/tbuffer">;
 
 def ClassTmpl : SubsetSubject<CXXRecord, [{S->getDescribedClassTemplate()}],
                               "class templates">;
@@ -4041,6 +4044,14 @@ def HLSLSV_GroupIndex: HLSLAnnotationAttr {
   let Documentation = [HLSLSV_GroupIndexDocs];
 }
 
+def HLSLResourceBinding: InheritableAttr {
+  let Spellings = [HLSLSemantic<"register">];
+  let Subjects = SubjectList<[HLSLBufferObj, GlobalVar]>;
+  let LangOpts = [HLSL];
+  let Args = [StringArgument<"Slot">, StringArgument<"Space", 1>];
+  let Documentation = [HLSLResourceBindingDocs];
+}
+
 def HLSLShader : InheritableAttr {
   let Spellings = [Microsoft<"shader">];
   let Subjects = SubjectList<[HLSLEntry]>;

clang/include/clang/Basic/AttrDocs.td

@@ -6590,6 +6590,31 @@ The full documentation is available here: https://docs.microsoft.com/en-us/windo
   }];
 }
 
+def HLSLResourceBindingDocs : Documentation {
+  let Category = DocCatFunction;
+  let Content = [{
+The resource binding attribute sets the virtual register and logical register space for a resource.
+Attribute spelling in HLSL is: ``register(slot [, space])``.
+``slot`` takes the format ``[type][number]``,
+where ``type`` is a single character specifying the resource type and ``number`` is the virtual register number.
+
+Register types are:
+t for shader resource views (SRV),
+s for samplers,
+u for unordered access views (UAV),
+b for constant buffer views (CBV).
+
+Register space is specified in the format ``space[number]`` and defaults to ``space0`` if omitted.
+Here're resource binding examples with and without space:
+.. code-block:: c++
+
+  RWBuffer<float> Uav : register(u3, space1);
+  Buffer<float> Buf : register(t1);
+
+The full documentation is available here: https://docs.microsoft.com/en-us/windows/win32/direct3d12/resource-binding-in-hlsl
+  }];
+}
+
 def AnnotateTypeDocs : Documentation {
   let Category = DocCatType;
   let Heading = "annotate_type";

clang/include/clang/Basic/DiagnosticParseKinds.td

@@ -1619,6 +1619,7 @@ def err_expected_semantic_identifier : Error<
 def err_invalid_declaration_in_hlsl_buffer : Error<
   "invalid declaration inside %select{tbuffer|cbuffer}0">;
 def err_unknown_hlsl_semantic : Error<"unknown HLSL semantic %0">;
+def err_hlsl_separate_attr_arg_and_number : Error<"wrong argument format for hlsl attribute, use %0 instead">;
 def ext_hlsl_access_specifiers : ExtWarn<
   "access specifiers are a clang HLSL extension">,
   InGroup<HLSLExtension>;

clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -11672,6 +11672,9 @@ def err_hlsl_missing_semantic_annotation : Error<
 def err_hlsl_init_priority_unsupported : Error<
   "initializer priorities are not supported in HLSL">;
 
+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 err_hlsl_pointers_unsupported : Error<
   "%select{pointers|references}0 are unsupported in HLSL">;
 

clang/lib/Parse/ParseHLSL.cpp

@@ -10,6 +10,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "clang/AST/Attr.h"
 #include "clang/Basic/AttributeCommonInfo.h"
 #include "clang/Parse/ParseDiagnostic.h"
 #include "clang/Parse/Parser.h"
@@ -60,6 +61,9 @@ Decl *Parser::ParseHLSLBuffer(SourceLocation &DeclEnd) {
   IdentifierInfo *Identifier = Tok.getIdentifierInfo();
   SourceLocation IdentifierLoc = ConsumeToken();
 
+  ParsedAttributes Attrs(AttrFactory);
+  MaybeParseHLSLSemantics(Attrs, nullptr);
+
   ParseScope BufferScope(this, Scope::DeclScope);
   BalancedDelimiterTracker T(*this, tok::l_brace);
   if (T.consumeOpen()) {
@@ -71,7 +75,6 @@ Decl *Parser::ParseHLSLBuffer(SourceLocation &DeclEnd) {
                                          Identifier, IdentifierLoc,
                                          T.getOpenLocation());
 
-  // FIXME: support attribute on cbuffer/tbuffer.
   while (Tok.isNot(tok::r_brace) && Tok.isNot(tok::eof)) {
     // FIXME: support attribute on constants inside cbuffer/tbuffer.
     ParsedAttributes Attrs(AttrFactory);
@@ -92,30 +95,103 @@ Decl *Parser::ParseHLSLBuffer(SourceLocation &DeclEnd) {
   BufferScope.Exit();
   Actions.ActOnFinishHLSLBuffer(D, DeclEnd);
 
+  Actions.ProcessDeclAttributeList(Actions.CurScope, D, Attrs);
   return D;
 }
 
+static void fixSeparateAttrArgAndNumber(StringRef ArgStr, SourceLocation ArgLoc,
+                                        Token Tok, ArgsVector &ArgExprs,
+                                        Parser &P, ASTContext &Ctx,
+                                        Preprocessor &PP) {
+  StringRef Num = StringRef(Tok.getLiteralData(), Tok.getLength());
+  SourceLocation EndNumLoc = Tok.getEndLoc();
+
+  P.ConsumeToken(); // consume constant.
+  std::string FixedArg = ArgStr.str() + Num.str();
+  P.Diag(ArgLoc, diag::err_hlsl_separate_attr_arg_and_number)
+      << FixedArg
+      << FixItHint::CreateReplacement(SourceRange(ArgLoc, EndNumLoc), FixedArg);
+  ArgsUnion &Slot = ArgExprs.back();
+  Slot = IdentifierLoc::create(Ctx, ArgLoc, PP.getIdentifierInfo(FixedArg));
+}
+
 void Parser::ParseHLSLSemantics(ParsedAttributes &Attrs,
                                 SourceLocation *EndLoc) {
+  // FIXME: HLSLSemantic is shared for Semantic and resource binding which is
+  // confusing. Need a better name to avoid misunderstanding. Issue
+  // https://github.com/llvm/llvm-project/issues/57882
   assert(Tok.is(tok::colon) && "Not a HLSL Semantic");
   ConsumeToken();
 
-  if (!Tok.is(tok::identifier)) {
+  IdentifierInfo *II = nullptr;
+  if (Tok.is(tok::kw_register))
+    II = PP.getIdentifierInfo("register");
+  else if (Tok.is(tok::identifier))
+    II = Tok.getIdentifierInfo();
+
+  if (!II) {
     Diag(Tok.getLocation(), diag::err_expected_semantic_identifier);
     return;
   }
 
-  IdentifierInfo *II = Tok.getIdentifierInfo();
   SourceLocation Loc = ConsumeToken();
   if (EndLoc)
     *EndLoc = Tok.getLocation();
   ParsedAttr::Kind AttrKind =
       ParsedAttr::getParsedKind(II, nullptr, ParsedAttr::AS_HLSLSemantic);
 
-  if (AttrKind == ParsedAttr::UnknownAttribute) {
+  ArgsVector ArgExprs;
+  switch (AttrKind) {
+  case ParsedAttr::AT_HLSLResourceBinding: {
+    if (ExpectAndConsume(tok::l_paren, diag::err_expected_lparen_after)) {
+      SkipUntil(tok::r_paren, StopAtSemi); // skip through )
+      return;
+    }
+    if (!Tok.is(tok::identifier)) {
+      Diag(Tok.getLocation(), diag::err_expected) << tok::identifier;
+      SkipUntil(tok::r_paren, StopAtSemi); // skip through )
+      return;
+    }
+    StringRef SlotStr = Tok.getIdentifierInfo()->getName();
+    SourceLocation SlotLoc = Tok.getLocation();
+    ArgExprs.push_back(ParseIdentifierLoc());
+
+    // Add numeric_constant for fix-it.
+    if (SlotStr.size() == 1 && Tok.is(tok::numeric_constant))
+      fixSeparateAttrArgAndNumber(SlotStr, SlotLoc, Tok, ArgExprs, *this,
+                                  Actions.Context, PP);
+
+    if (Tok.is(tok::comma)) {
+      ConsumeToken(); // consume comma
+      if (!Tok.is(tok::identifier)) {
+        Diag(Tok.getLocation(), diag::err_expected) << tok::identifier;
+        SkipUntil(tok::r_paren, StopAtSemi); // skip through )
+        return;
+      }
+      StringRef SpaceStr = Tok.getIdentifierInfo()->getName();
+      SourceLocation SpaceLoc = Tok.getLocation();
+      ArgExprs.push_back(ParseIdentifierLoc());
+
+      // Add numeric_constant for fix-it.
+      if (SpaceStr.equals("space") && Tok.is(tok::numeric_constant))
+        fixSeparateAttrArgAndNumber(SpaceStr, SpaceLoc, Tok, ArgExprs, *this,
+                                    Actions.Context, PP);
+    }
+    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;
+  case ParsedAttr::AT_HLSLSV_GroupIndex:
+    break;
+  default:
+    llvm_unreachable("invalid HLSL Semantic");
+    break;
   }
-  Attrs.addNew(II, Loc, nullptr, SourceLocation(), nullptr, 0,
-               ParsedAttr::AS_HLSLSemantic);
+
+  Attrs.addNew(II, Loc, nullptr, SourceLocation(), ArgExprs.data(),
+               ArgExprs.size(), ParsedAttr::AS_HLSLSemantic);
 }

clang/lib/Sema/SemaDeclAttr.cpp

@@ -6947,6 +6947,78 @@ Sema::mergeHLSLShaderAttr(Decl *D, const AttributeCommonInfo &AL,
   return HLSLShaderAttr::Create(Context, ShaderType, AL);
 }
 
+static void handleHLSLResourceBindingAttr(Sema &S, Decl *D,
+                                          const ParsedAttr &AL) {
+  StringRef Space = "space0";
+  StringRef Slot = "";
+
+  if (!AL.isArgIdent(0)) {
+    S.Diag(AL.getLoc(), diag::err_attribute_argument_type)
+        << AL << AANT_ArgumentIdentifier;
+    return;
+  }
+
+  IdentifierLoc *Loc = AL.getArgAsIdent(0);
+  StringRef Str = Loc->Ident->getName();
+  SourceLocation ArgLoc = Loc->Loc;
+
+  SourceLocation SpaceArgLoc;
+  if (AL.getNumArgs() == 2) {
+    Slot = Str;
+    if (!AL.isArgIdent(1)) {
+      S.Diag(AL.getLoc(), diag::err_attribute_argument_type)
+          << AL << AANT_ArgumentIdentifier;
+      return;
+    }
+
+    IdentifierLoc *Loc = AL.getArgAsIdent(1);
+    Space = Loc->Ident->getName();
+    SpaceArgLoc = Loc->Loc;
+  } else {
+    Slot = Str;
+  }
+
+  // Validate.
+  if (!Slot.empty()) {
+    switch (Slot[0]) {
+    case 'u':
+    case 'b':
+    case 's':
+    case 't':
+      break;
+    default:
+      S.Diag(ArgLoc, diag::err_hlsl_unsupported_register_type)
+          << Slot.substr(0, 1);
+      return;
+    }
+
+    StringRef SlotNum = Slot.substr(1);
+    unsigned Num = 0;
+    if (SlotNum.getAsInteger(10, Num)) {
+      S.Diag(ArgLoc, diag::err_hlsl_unsupported_register_number);
+      return;
+    }
+  }
+
+  if (!Space.startswith("space")) {
+    S.Diag(SpaceArgLoc, diag::err_hlsl_expected_space) << Space;
+    return;
+  }
+  StringRef SpaceNum = Space.substr(5);
+  unsigned Num = 0;
+  if (SpaceNum.getAsInteger(10, Num)) {
+    S.Diag(SpaceArgLoc, diag::err_hlsl_expected_space) << Space;
+    return;
+  }
+
+  // FIXME: check reg type match decl. Issue
+  // https://github.com/llvm/llvm-project/issues/57886.
+  HLSLResourceBindingAttr *NewAttr =
+      HLSLResourceBindingAttr::Create(S.getASTContext(), Slot, Space, AL);
+  if (NewAttr)
+    D->addAttr(NewAttr);
+}
+
 static void handleMSInheritanceAttr(Sema &S, Decl *D, const ParsedAttr &AL) {
   if (!S.LangOpts.CPlusPlus) {
     S.Diag(AL.getLoc(), diag::err_attribute_not_supported_in_lang)
@@ -8926,6 +8998,9 @@ ProcessDeclAttribute(Sema &S, Scope *scope, Decl *D, const ParsedAttr &AL,
   case ParsedAttr::AT_HLSLShader:
     handleHLSLShaderAttr(S, D, AL);
     break;
+  case ParsedAttr::AT_HLSLResourceBinding:
+    handleHLSLResourceBindingAttr(S, D, AL);
+    break;
 
   case ParsedAttr::AT_AbiTag:
     handleAbiTagAttr(S, D, AL);

clang/test/AST/HLSL/resource_binding_attr.hlsl

@@ -0,0 +1,24 @@
+// RUN: %clang_cc1 -triple dxil-pc-shadermodel6.3-library -x hlsl -ast-dump -o - %s | FileCheck %s
+
+// CHECK:HLSLBufferDecl 0x[[CB:[0-9a-f]+]] {{.*}} line:6:9 cbuffer CB
+// CHECK-NEXT:HLSLResourceBindingAttr 0x{{[0-9a-f]+}} <col:14> "b3" "space2"
+// CHECK-NEXT:VarDecl 0x[[A:[0-9a-f]+]] {{.*}} col:9 used a 'float'
+cbuffer CB : register(b3, space2) {
+  float a;
+}
+
+// CHECK:HLSLBufferDecl 0x[[TB:[0-9a-f]+]] {{.*}} line:13:9 tbuffer TB
+// CHECK-NEXT:HLSLResourceBindingAttr 0x{{[0-9a-f]+}} <col:14> "t2" "space1"
+// CHECK-NEXT:VarDecl 0x[[B:[0-9a-f]+]] {{.*}} col:9 used b 'float'
+tbuffer TB : register(t2, space1) {
+  float b;
+}
+
+float foo() {
+// CHECK: BinaryOperator 0x{{[0-9a-f]+}} <col:10, col:14> 'float' '+'
+// CHECK-NEXT: ImplicitCastExpr 0x{{[0-9a-f]+}} <col:10> 'float' <LValueToRValue>
+// CHECK-NEXT: DeclRefExpr 0x{{[0-9a-f]+}} <col:10> 'float' lvalue Var 0x[[A]] 'a' 'float'
+// CHECK-NEXT: ImplicitCastExpr 0x{{[0-9a-f]+}} <col:14> 'float' <LValueToRValue>
+// CHECK-NEXT: DeclRefExpr 0x{{[0-9a-f]+}} <col:14> 'float' lvalue Var 0x[[B]] 'b' 'float'
+  return a + b;
+}

clang/test/SemaHLSL/resource_binding_attr_error.hlsl

@@ -0,0 +1,38 @@
+// RUN: %clang_cc1 -triple dxil-pc-shadermodel6.3-library -x hlsl -o - -fsyntax-only %s -verify
+
+// expected-error@+5 {{expected ';' after top level declarator}}
+// expected-error@+4 {{expected ')'}}
+// expected-note@+3 {{to match this '('}}
+// expected-error@+2 {{a type specifier is required for all declarations}}
+// expected-error@+1 {{illegal storage class on file-scoped variable}}
+float a : register(c0, space1);
+
+// expected-error@+1 {{invalid resource class specifier 'i' used; expected 'b', 's', 't', or 'u'}}
+cbuffer b : register(i0) {
+
+}
+// expected-error@+1 {{invalid space specifier 's2' used; expected 'space' followed by an integer, like space1}}
+cbuffer c : register(b0, s2) {
+
+}
+// expected-error@+1 {{register number should be an integer}}
+cbuffer d : register(bf, s2) {
+
+}
+// expected-error@+1 {{invalid space specifier 'spaces' used; expected 'space' followed by an integer, like space1}}
+cbuffer e : register(b2, spaces) {
+
+}
+
+// expected-error@+1 {{expected identifier}}
+cbuffer A : register() {}
+
+// expected-error@+1 {{register number should be an integer}}
+cbuffer B : register(space1) {}
+
+// expected-error@+1 {{wrong argument format for hlsl attribute, use b2 instead}}
+cbuffer C : register(b 2) {}
+
+// expected-error@+2 {{wrong argument format for hlsl attribute, use b2 instead}}
+// expected-error@+1 {{wrong argument format for hlsl attribute, use space3 instead}}
+cbuffer D : register(b 2, space 3) {}