[HLSL][RootSignature] Add parsing infastructure for StaticSampler (#140180)

- define StaticSampler in-memory representation
- implement the infastructure for parsing parameters of StaticSampler
- define and implement parsing of the `s` reg to demonstrate
functionality
- add unit tests

First part of #126574

Author: inbelic

clang/include/clang/Lex/HLSLRootSignatureTokenKinds.def

@@ -80,6 +80,7 @@ KEYWORD(RootSignature) // used only for diagnostic messaging
 KEYWORD(RootFlags)
 KEYWORD(DescriptorTable)
 KEYWORD(RootConstants)
+KEYWORD(StaticSampler)
 
 // RootConstants Keywords:
 KEYWORD(num32BitConstants)

clang/include/clang/Parse/ParseHLSLRootSignature.h

@@ -77,6 +77,7 @@ class RootSignatureParser {
   std::optional<llvm::hlsl::rootsig::DescriptorTable> parseDescriptorTable();
   std::optional<llvm::hlsl::rootsig::DescriptorTableClause>
   parseDescriptorTableClause();
+  std::optional<llvm::hlsl::rootsig::StaticSampler> parseStaticSampler();
 
   /// Parameter arguments (eg. `bReg`, `space`, ...) can be specified in any
   /// order and only exactly once. The following methods define a
@@ -108,6 +109,11 @@ class RootSignatureParser {
   std::optional<ParsedClauseParams>
   parseDescriptorTableClauseParams(RootSignatureToken::Kind RegType);
 
+  struct ParsedStaticSamplerParams {
+    std::optional<llvm::hlsl::rootsig::Register> Reg;
+  };
+  std::optional<ParsedStaticSamplerParams> parseStaticSamplerParams();
+
   // Common parsing methods
   std::optional<uint32_t> parseUIntParam();
   std::optional<llvm::hlsl::rootsig::Register> parseRegister();

clang/lib/Parse/ParseHLSLRootSignature.cpp

@@ -53,6 +53,13 @@ bool RootSignatureParser::parse() {
         return true;
       Elements.push_back(*Descriptor);
     }
+
+    if (tryConsumeExpectedToken(TokenKind::kw_StaticSampler)) {
+      auto Sampler = parseStaticSampler();
+      if (!Sampler.has_value())
+        return true;
+      Elements.push_back(*Sampler);
+    }
   } while (tryConsumeExpectedToken(TokenKind::pu_comma));
 
   return consumeExpectedToken(TokenKind::end_of_stream,
@@ -346,6 +353,37 @@ RootSignatureParser::parseDescriptorTableClause() {
   return Clause;
 }
 
+std::optional<StaticSampler> RootSignatureParser::parseStaticSampler() {
+  assert(CurToken.TokKind == TokenKind::kw_StaticSampler &&
+         "Expects to only be invoked starting at given keyword");
+
+  if (consumeExpectedToken(TokenKind::pu_l_paren, diag::err_expected_after,
+                           CurToken.TokKind))
+    return std::nullopt;
+
+  StaticSampler Sampler;
+
+  auto Params = parseStaticSamplerParams();
+  if (!Params.has_value())
+    return std::nullopt;
+
+  // Check mandatory parameters were provided
+  if (!Params->Reg.has_value()) {
+    getDiags().Report(CurToken.TokLoc, diag::err_hlsl_rootsig_missing_param)
+        << TokenKind::sReg;
+    return std::nullopt;
+  }
+
+  Sampler.Reg = Params->Reg.value();
+
+  if (consumeExpectedToken(TokenKind::pu_r_paren,
+                           diag::err_hlsl_unexpected_end_of_params,
+                           /*param of=*/TokenKind::kw_StaticSampler))
+    return std::nullopt;
+
+  return Sampler;
+}
+
 // Parameter arguments (eg. `bReg`, `space`, ...) can be specified in any
 // order and only exactly once. The following methods will parse through as
 // many arguments as possible reporting an error if a duplicate is seen.
@@ -604,6 +642,30 @@ RootSignatureParser::parseDescriptorTableClauseParams(TokenKind RegType) {
   return Params;
 }
 
+std::optional<RootSignatureParser::ParsedStaticSamplerParams>
+RootSignatureParser::parseStaticSamplerParams() {
+  assert(CurToken.TokKind == TokenKind::pu_l_paren &&
+         "Expects to only be invoked starting at given token");
+
+  ParsedStaticSamplerParams Params;
+  do {
+    // `s` POS_INT
+    if (tryConsumeExpectedToken(TokenKind::sReg)) {
+      if (Params.Reg.has_value()) {
+        getDiags().Report(CurToken.TokLoc, diag::err_hlsl_rootsig_repeat_param)
+            << CurToken.TokKind;
+        return std::nullopt;
+      }
+      auto Reg = parseRegister();
+      if (!Reg.has_value())
+        return std::nullopt;
+      Params.Reg = Reg;
+    }
+  } while (tryConsumeExpectedToken(TokenKind::pu_comma));
+
+  return Params;
+}
+
 std::optional<uint32_t> RootSignatureParser::parseUIntParam() {
   assert(CurToken.TokKind == TokenKind::pu_equal &&
          "Expects to only be invoked starting at given keyword");

clang/unittests/Lex/LexHLSLRootSignatureTest.cpp

@@ -128,7 +128,7 @@ TEST_F(LexHLSLRootSignatureTest, ValidLexAllTokensTest) {
 
     RootSignature
 
-    RootFlags DescriptorTable RootConstants
+    RootFlags DescriptorTable RootConstants StaticSampler
 
     num32BitConstants
 

clang/unittests/Parse/ParseHLSLRootSignatureTest.cpp

@@ -223,6 +223,34 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseDTClausesTest) {
   ASSERT_TRUE(Consumer->isSatisfied());
 }
 
+TEST_F(ParseHLSLRootSignatureTest, ValidParseStaticSamplerTest) {
+  const llvm::StringLiteral Source = R"cc(
+    StaticSampler(s0)
+  )cc";
+
+  TrivialModuleLoader ModLoader;
+  auto PP = createPP(Source, ModLoader);
+  auto TokLoc = SourceLocation();
+
+  hlsl::RootSignatureLexer Lexer(Source, TokLoc);
+  SmallVector<RootElement> Elements;
+  hlsl::RootSignatureParser Parser(Elements, Lexer, *PP);
+
+  // Test no diagnostics produced
+  Consumer->setNoDiag();
+
+  ASSERT_FALSE(Parser.parse());
+
+  ASSERT_EQ(Elements.size(), 1u);
+
+  RootElement Elem = Elements[0];
+  ASSERT_TRUE(std::holds_alternative<StaticSampler>(Elem));
+  ASSERT_EQ(std::get<StaticSampler>(Elem).Reg.ViewType, RegisterType::SReg);
+  ASSERT_EQ(std::get<StaticSampler>(Elem).Reg.Number, 0u);
+
+  ASSERT_TRUE(Consumer->isSatisfied());
+}
+
 TEST_F(ParseHLSLRootSignatureTest, ValidSamplerFlagsTest) {
   // This test will checks we can set the valid enum for Sampler descriptor
   // range flags

llvm/include/llvm/Frontend/HLSL/HLSLRootSignature.h

@@ -155,23 +155,29 @@ struct DescriptorTableClause {
 
 raw_ostream &operator<<(raw_ostream &OS, const DescriptorTableClause &Clause);
 
-/// Models RootElement : RootFlags | RootConstants | RootDescriptor
-///  | DescriptorTable | DescriptorTableClause
+struct StaticSampler {
+  Register Reg;
+};
+
+/// Models RootElement : RootFlags | RootConstants | RootParam
+///  | DescriptorTable | DescriptorTableClause | StaticSampler
 ///
 /// A Root Signature is modeled in-memory by an array of RootElements. These
 /// aim to map closely to their DSL grammar reprsentation defined in the spec.
 ///
 /// Each optional parameter has its default value defined in the struct, and,
 /// each mandatory parameter does not have a default initialization.
 ///
-/// For the variants RootFlags, RootConstants and DescriptorTableClause: each
-/// data member maps directly to a parameter in the grammar.
+/// For the variants RootFlags, RootConstants, RootParam, StaticSampler and
+/// DescriptorTableClause: each data member maps directly to a parameter in the
+/// grammar.
 ///
 /// The DescriptorTable is modelled by having its Clauses as the previous
 /// RootElements in the array, and it holds a data member for the Visibility
 /// parameter.
-using RootElement = std::variant<RootFlags, RootConstants, RootDescriptor,
-                                 DescriptorTable, DescriptorTableClause>;
+using RootElement =
+    std::variant<RootFlags, RootConstants, RootDescriptor, DescriptorTable,
+                 DescriptorTableClause, StaticSampler>;
 
 void dumpRootElements(raw_ostream &OS, ArrayRef<RootElement> Elements);