[HLSL][RootSignature] Add parsing of address params in StaticSampler (#140293)

inbelic · web-flow · commit 2b5b8db9dcf1 · 2025-05-29T15:03:49.000-07:00
- defines in-memory reprsentation of `address[U|V|W]` - defines parsing of the `TextureAddressMode` enum - integrates parsing of these number parameters with their respective, `parseTextureAddressMode` - adds basic unit tests to demonstrate setting functionality Part 4 of #126574
diff --git a/clang/include/clang/Lex/HLSLRootSignatureTokenKinds.def b/clang/include/clang/Lex/HLSLRootSignatureTokenKinds.def
@@ -53,6 +53,9 @@
 #ifndef SHADER_VISIBILITY_ENUM
 #define SHADER_VISIBILITY_ENUM(NAME, LIT) ENUM(NAME, LIT)
 #endif
+#ifndef TEXTURE_ADDRESS_MODE_ENUM
+#define TEXTURE_ADDRESS_MODE_ENUM(NAME, LIT) ENUM(NAME, LIT)
+#endif
 
 // General Tokens:
 TOK(invalid, "invalid identifier")
@@ -102,6 +105,9 @@ KEYWORD(offset)
 
 // StaticSampler Keywords:
 KEYWORD(mipLODBias)
+KEYWORD(addressU)
+KEYWORD(addressV)
+KEYWORD(addressW)
 KEYWORD(maxAnisotropy)
 KEYWORD(minLOD)
 KEYWORD(maxLOD)
@@ -148,6 +154,14 @@ SHADER_VISIBILITY_ENUM(Pixel, "SHADER_VISIBILITY_PIXEL")
 SHADER_VISIBILITY_ENUM(Amplification, "SHADER_VISIBILITY_AMPLIFICATION")
 SHADER_VISIBILITY_ENUM(Mesh, "SHADER_VISIBILITY_MESH")
 
+// Texture Address Mode Enums:
+TEXTURE_ADDRESS_MODE_ENUM(Wrap, "TEXTURE_ADDRESS_WRAP")
+TEXTURE_ADDRESS_MODE_ENUM(Mirror, "TEXTURE_ADDRESS_MIRROR")
+TEXTURE_ADDRESS_MODE_ENUM(Clamp, "TEXTURE_ADDRESS_CLAMP")
+TEXTURE_ADDRESS_MODE_ENUM(Border, "TEXTURE_ADDRESS_BORDER")
+TEXTURE_ADDRESS_MODE_ENUM(MirrorOnce, "TEXTURE_ADDRESS_MIRRORONCE")
+
+#undef TEXTURE_ADDRESS_MODE_ENUM
 #undef SHADER_VISIBILITY_ENUM
 #undef DESCRIPTOR_RANGE_FLAG_ENUM
 #undef DESCRIPTOR_RANGE_FLAG_ENUM_OFF
diff --git a/clang/include/clang/Parse/ParseHLSLRootSignature.h b/clang/include/clang/Parse/ParseHLSLRootSignature.h
@@ -112,6 +112,9 @@ class RootSignatureParser {
   struct ParsedStaticSamplerParams {
     std::optional<llvm::hlsl::rootsig::Register> Reg;
     std::optional<float> MipLODBias;
+    std::optional<llvm::hlsl::rootsig::TextureAddressMode> AddressU;
+    std::optional<llvm::hlsl::rootsig::TextureAddressMode> AddressV;
+    std::optional<llvm::hlsl::rootsig::TextureAddressMode> AddressW;
     std::optional<uint32_t> MaxAnisotropy;
     std::optional<float> MinLOD;
     std::optional<float> MaxLOD;
@@ -125,6 +128,8 @@ class RootSignatureParser {
 
   /// Parsing methods of various enums
   std::optional<llvm::hlsl::rootsig::ShaderVisibility> parseShaderVisibility();
+  std::optional<llvm::hlsl::rootsig::TextureAddressMode>
+  parseTextureAddressMode();
   std::optional<llvm::hlsl::rootsig::RootDescriptorFlags>
   parseRootDescriptorFlags();
   std::optional<llvm::hlsl::rootsig::DescriptorRangeFlags>
diff --git a/clang/lib/Parse/ParseHLSLRootSignature.cpp b/clang/lib/Parse/ParseHLSLRootSignature.cpp
@@ -377,6 +377,15 @@ std::optional<StaticSampler> RootSignatureParser::parseStaticSampler() {
   Sampler.Reg = Params->Reg.value();
 
   // Fill in optional values
+  if (Params->AddressU.has_value())
+    Sampler.AddressU = Params->AddressU.value();
+
+  if (Params->AddressV.has_value())
+    Sampler.AddressV = Params->AddressV.value();
+
+  if (Params->AddressW.has_value())
+    Sampler.AddressW = Params->AddressW.value();
+
   if (Params->MipLODBias.has_value())
     Sampler.MipLODBias = Params->MipLODBias.value();
 
@@ -675,6 +684,57 @@ RootSignatureParser::parseStaticSamplerParams() {
       Params.Reg = Reg;
     }
 
+    // `addressU` `=` TEXTURE_ADDRESS
+    if (tryConsumeExpectedToken(TokenKind::kw_addressU)) {
+      if (Params.AddressU.has_value()) {
+        getDiags().Report(CurToken.TokLoc, diag::err_hlsl_rootsig_repeat_param)
+            << CurToken.TokKind;
+        return std::nullopt;
+      }
+
+      if (consumeExpectedToken(TokenKind::pu_equal))
+        return std::nullopt;
+
+      auto AddressU = parseTextureAddressMode();
+      if (!AddressU.has_value())
+        return std::nullopt;
+      Params.AddressU = AddressU;
+    }
+
+    // `addressV` `=` TEXTURE_ADDRESS
+    if (tryConsumeExpectedToken(TokenKind::kw_addressV)) {
+      if (Params.AddressV.has_value()) {
+        getDiags().Report(CurToken.TokLoc, diag::err_hlsl_rootsig_repeat_param)
+            << CurToken.TokKind;
+        return std::nullopt;
+      }
+
+      if (consumeExpectedToken(TokenKind::pu_equal))
+        return std::nullopt;
+
+      auto AddressV = parseTextureAddressMode();
+      if (!AddressV.has_value())
+        return std::nullopt;
+      Params.AddressV = AddressV;
+    }
+
+    // `addressW` `=` TEXTURE_ADDRESS
+    if (tryConsumeExpectedToken(TokenKind::kw_addressW)) {
+      if (Params.AddressW.has_value()) {
+        getDiags().Report(CurToken.TokLoc, diag::err_hlsl_rootsig_repeat_param)
+            << CurToken.TokKind;
+        return std::nullopt;
+      }
+
+      if (consumeExpectedToken(TokenKind::pu_equal))
+        return std::nullopt;
+
+      auto AddressW = parseTextureAddressMode();
+      if (!AddressW.has_value())
+        return std::nullopt;
+      Params.AddressW = AddressW;
+    }
+
     // `mipLODBias` `=` NUMBER
     if (tryConsumeExpectedToken(TokenKind::kw_mipLODBias)) {
       if (Params.MipLODBias.has_value()) {
@@ -849,6 +909,32 @@ RootSignatureParser::parseShaderVisibility() {
   return std::nullopt;
 }
 
+std::optional<llvm::hlsl::rootsig::TextureAddressMode>
+RootSignatureParser::parseTextureAddressMode() {
+  assert(CurToken.TokKind == TokenKind::pu_equal &&
+         "Expects to only be invoked starting at given keyword");
+
+  TokenKind Expected[] = {
+#define TEXTURE_ADDRESS_MODE_ENUM(NAME, LIT) TokenKind::en_##NAME,
+#include "clang/Lex/HLSLRootSignatureTokenKinds.def"
+  };
+
+  if (!tryConsumeExpectedToken(Expected))
+    return std::nullopt;
+
+  switch (CurToken.TokKind) {
+#define TEXTURE_ADDRESS_MODE_ENUM(NAME, LIT)                                   \
+  case TokenKind::en_##NAME:                                                   \
+    return TextureAddressMode::NAME;                                           \
+    break;
+#include "clang/Lex/HLSLRootSignatureTokenKinds.def"
+  default:
+    llvm_unreachable("Switch for consumed enum token was not provided");
+  }
+
+  return std::nullopt;
+}
+
 std::optional<llvm::hlsl::rootsig::RootDescriptorFlags>
 RootSignatureParser::parseRootDescriptorFlags() {
   assert(CurToken.TokKind == TokenKind::pu_equal &&
diff --git a/clang/unittests/Lex/LexHLSLRootSignatureTest.cpp b/clang/unittests/Lex/LexHLSLRootSignatureTest.cpp
@@ -136,7 +136,8 @@ TEST_F(LexHLSLRootSignatureTest, ValidLexAllTokensTest) {
     space visibility flags
     numDescriptors offset
 
-    mipLODBias maxAnisotropy minLOD maxLOD
+    mipLODBias addressU addressV addressW
+    maxAnisotropy minLOD maxLOD
 
     unbounded
     DESCRIPTOR_RANGE_OFFSET_APPEND
@@ -168,6 +169,12 @@ TEST_F(LexHLSLRootSignatureTest, ValidLexAllTokensTest) {
     shader_visibility_pixel
     shader_visibility_amplification
     shader_visibility_mesh
+
+    TEXTURE_ADDRESS_WRAP
+    TEXTURE_ADDRESS_MIRROR
+    TEXTURE_ADDRESS_CLAMP
+    TEXTURE_ADDRESS_BORDER
+    TEXTURE_ADDRESS_MIRRORONCE
   )cc";
   auto TokLoc = SourceLocation();
   hlsl::RootSignatureLexer Lexer(Source, TokLoc);
diff --git a/clang/unittests/Parse/ParseHLSLRootSignatureTest.cpp b/clang/unittests/Parse/ParseHLSLRootSignatureTest.cpp
@@ -228,7 +228,9 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseStaticSamplerTest) {
     StaticSampler(s0),
     StaticSampler(s0, maxAnisotropy = 3,
       minLOD = 4.2f, mipLODBias = 0.23e+3,
-      maxLOD = 9000,
+      addressW = TEXTURE_ADDRESS_CLAMP,
+      addressV = TEXTURE_ADDRESS_BORDER,
+      maxLOD = 9000, addressU = TEXTURE_ADDRESS_MIRROR
     )
   )cc";
 
@@ -252,6 +254,9 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseStaticSamplerTest) {
   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_EQ(std::get<StaticSampler>(Elem).AddressU, TextureAddressMode::Wrap);
+  ASSERT_EQ(std::get<StaticSampler>(Elem).AddressV, TextureAddressMode::Wrap);
+  ASSERT_EQ(std::get<StaticSampler>(Elem).AddressW, TextureAddressMode::Wrap);
   ASSERT_FLOAT_EQ(std::get<StaticSampler>(Elem).MipLODBias, 0.f);
   ASSERT_EQ(std::get<StaticSampler>(Elem).MaxAnisotropy, 16u);
   ASSERT_FLOAT_EQ(std::get<StaticSampler>(Elem).MinLOD, 0.f);
@@ -262,6 +267,9 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseStaticSamplerTest) {
   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_EQ(std::get<StaticSampler>(Elem).AddressU, TextureAddressMode::Mirror);
+  ASSERT_EQ(std::get<StaticSampler>(Elem).AddressV, TextureAddressMode::Border);
+  ASSERT_EQ(std::get<StaticSampler>(Elem).AddressW, TextureAddressMode::Clamp);
   ASSERT_FLOAT_EQ(std::get<StaticSampler>(Elem).MipLODBias, 230.f);
   ASSERT_EQ(std::get<StaticSampler>(Elem).MaxAnisotropy, 3u);
   ASSERT_FLOAT_EQ(std::get<StaticSampler>(Elem).MinLOD, 4.2f);
diff --git a/llvm/include/llvm/Frontend/HLSL/HLSLRootSignature.h b/llvm/include/llvm/Frontend/HLSL/HLSLRootSignature.h
@@ -76,6 +76,14 @@ enum class ShaderVisibility {
   Mesh = 7,
 };
 
+enum class TextureAddressMode {
+  Wrap = 1,
+  Mirror = 2,
+  Clamp = 3,
+  Border = 4,
+  MirrorOnce = 5
+};
+
 // Definitions of the in-memory data layout structures
 
 // Models the different registers: bReg | tReg | uReg | sReg
@@ -157,6 +165,9 @@ raw_ostream &operator<<(raw_ostream &OS, const DescriptorTableClause &Clause);
 
 struct StaticSampler {
   Register Reg;
+  TextureAddressMode AddressU = TextureAddressMode::Wrap;
+  TextureAddressMode AddressV = TextureAddressMode::Wrap;
+  TextureAddressMode AddressW = TextureAddressMode::Wrap;
   float MipLODBias = 0.f;
   uint32_t MaxAnisotropy = 16;
   float MinLOD = 0.f;
