[HLSL][RootSignature] Define and integrate HLSLRootSignatureAttr

clang/include/clang/AST/Attr.h

@@ -26,6 +26,7 @@
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Support/Compiler.h"
 #include "llvm/Frontend/HLSL/HLSLResource.h"
+#include "llvm/Frontend/HLSL/HLSLRootSignature.h"
 #include "llvm/Support/CodeGen.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/VersionTuple.h"

clang/include/clang/Basic/Attr.td

@@ -4709,6 +4709,25 @@ def Error : InheritableAttr {
   let Documentation = [ErrorAttrDocs];
 }
 
+def HLSLRootSignature : Attr {
+  /// [RootSignature(Signature)]
+  let Spellings = [Microsoft<"RootSignature">];
+  let Args = [StringArgument<"Signature">];
+  let Subjects = SubjectList<[Function],
+                             ErrorDiag, "'function'">;
+  let LangOpts = [HLSL];
+  let Documentation = [HLSLRootSignatureDocs];
+  let AdditionalMembers = [{
+private:
+  ArrayRef<llvm::hlsl::rootsig::RootElement> RootElements;
+public:
+  void setElements(ArrayRef<llvm::hlsl::rootsig::RootElement> Elements) {
+    RootElements = Elements;
+  }
+  auto getElements() const { return RootElements; }
+}];
+}
+
 def HLSLNumThreads: InheritableAttr {
   let Spellings = [Microsoft<"numthreads">];
   let Args = [IntArgument<"X">, IntArgument<"Y">, IntArgument<"Z">];

clang/include/clang/Basic/AttrDocs.td

@@ -8143,6 +8143,17 @@ and https://microsoft.github.io/hlsl-specs/proposals/0013-wave-size-range.html
   }];
 }
 
+def HLSLRootSignatureDocs : Documentation {
+  let Category = DocCatFunction;
+  let Content = [{
+The ``RootSignature`` attribute applies to HLSL entry functions to define what
+types of resources are bound to the graphics pipeline.
+
+For details about the use and specification of Root Signatures please see here:
+https://learn.microsoft.com/en-us/windows/win32/direct3d12/root-signatures
+  }];
+}
+
 def NumThreadsDocs : Documentation {
   let Category = DocCatFunction;
   let Content = [{

clang/include/clang/Sema/SemaHLSL.h

@@ -118,6 +118,7 @@ class SemaHLSL : public SemaBase {
                                        bool IsCompAssign);
   void emitLogicalOperatorFixIt(Expr *LHS, Expr *RHS, BinaryOperatorKind Opc);
 
+  void handleRootSignatureAttr(Decl *D, const ParsedAttr &AL);
   void handleNumThreadsAttr(Decl *D, const ParsedAttr &AL);
   void handleWaveSizeAttr(Decl *D, const ParsedAttr &AL);
   void handleSV_DispatchThreadIDAttr(Decl *D, const ParsedAttr &AL);

clang/lib/Sema/SemaDeclAttr.cpp

@@ -7480,6 +7480,9 @@ ProcessDeclAttribute(Sema &S, Scope *scope, Decl *D, const ParsedAttr &AL,
     break;
 
   // HLSL attributes:
+  case ParsedAttr::AT_HLSLRootSignature:
+    S.HLSL().handleRootSignatureAttr(D, AL);
+    break;
   case ParsedAttr::AT_HLSLNumThreads:
     S.HLSL().handleNumThreadsAttr(D, AL);
     break;

clang/lib/Sema/SemaHLSL.cpp

@@ -28,6 +28,7 @@
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/Specifiers.h"
 #include "clang/Basic/TargetInfo.h"
+#include "clang/Parse/ParseHLSLRootSignature.h"
 #include "clang/Sema/Initialization.h"
 #include "clang/Sema/ParsedAttr.h"
 #include "clang/Sema/Sema.h"
@@ -941,6 +942,40 @@ void SemaHLSL::emitLogicalOperatorFixIt(Expr *LHS, Expr *RHS,
       << NewFnName << FixItHint::CreateReplacement(FullRange, OS.str());
 }
 
+void SemaHLSL::handleRootSignatureAttr(Decl *D, const ParsedAttr &AL) {
+  if (AL.getNumArgs() != 1) {
+    Diag(AL.getLoc(), diag::err_attribute_wrong_number_arguments) << AL << 1;
+    return;
+  }
+
+  StringRef Signature;
+  if (!SemaRef.checkStringLiteralArgumentAttr(AL, 0, Signature))
+    return;
+
+  SourceLocation Loc = AL.getArgAsExpr(0)->getExprLoc();
+  // TODO(#126565): pass down below to lexer when fp is supported
+  // llvm::RoundingMode RM = SemaRef.CurFPFeatures.getRoundingMode();
+  hlsl::RootSignatureLexer Lexer(Signature, Loc);
+  SmallVector<llvm::hlsl::rootsig::RootElement> Elements;
+  hlsl::RootSignatureParser Parser(Elements, Lexer, SemaRef.getPreprocessor());
+
+  if (Parser.parse())
+    return;
+
+  // Allocate elements onto AST context
+  unsigned N = Elements.size();
+  auto RootElements = MutableArrayRef<llvm::hlsl::rootsig::RootElement>(
+      ::new (getASTContext()) llvm::hlsl::rootsig::RootElement[N], N);
+  for (unsigned I = 0; I < N; ++I)
+    RootElements[I] = Elements[I];
+
+  // Set elements
+  auto *Result = ::new (getASTContext())
+      HLSLRootSignatureAttr(getASTContext(), AL, Signature);
+  Result->setElements(ArrayRef<llvm::hlsl::rootsig::RootElement>(RootElements));
+  D->addAttr(Result);
+}
+
 void SemaHLSL::handleNumThreadsAttr(Decl *D, const ParsedAttr &AL) {
   llvm::VersionTuple SMVersion =
       getASTContext().getTargetInfo().getTriple().getOSVersion();

clang/test/AST/HLSL/RootSignatures-AST.hlsl

@@ -0,0 +1,24 @@
+// RUN: %clang_cc1 -triple dxil-pc-shadermodel6.0-library -ast-dump \
+// RUN:  -disable-llvm-passes -o - %s | FileCheck %s
+
+// This test ensures that the sample root signature is parsed without error and
+// the Attr AST Node is created succesfully. If an invalid root signature was
+// passed in then we would exit out of Sema before the Attr is created.
+
+#define SampleRS \
+  "DescriptorTable( " \
+  "  CBV(), " \
+  "  SRV(), " \
+  "  UAV()" \
+  "), " \
+  "DescriptorTable(Sampler())"
+
+// CHECK:      HLSLRootSignatureAttr
+// CHECK-SAME: "DescriptorTable(
+// CHECK-SAME:   CBV(),
+// CHECK-SAME:   SRV(),
+// CHECK-SAME:   UAV()
+// CHECK-SAME: ),
+// CHECK-SAME: DescriptorTable(Sampler())"
+[RootSignature(SampleRS)]
+void main() {}

clang/test/SemaHLSL/RootSignature-err.hlsl

@@ -0,0 +1,9 @@
+// RUN: %clang_cc1 -triple dxil-pc-shadermodel6.3-library -x hlsl -o - %s -verify
+
+// Attr test
+
+[RootSignature()] // expected-error {{'RootSignature' attribute takes one argument}}
+void bad_root_signature_0() {}
+
+[RootSignature("Arg1", "Arg2")] // expected-error {{'RootSignature' attribute takes one argument}}
+void bad_root_signature_1() {}