[HLSL][RootSignature] Metadata generation of StaticSampler

[inbelic]

Implements metadata generation of a Root Signature from its in-memory representation. It follows the same style as: #139633.

This pr handles StaticSamplers. It also handles converting the else-if chain into a std::visit to allow for future compiler warnings when adding additional RootElement variants.

The metadata follows the format described here.

• Implement BuildStaticSampler into HLSLRootSignature.h
• Add sample testcases demonstrating functionality

Note: there is no validation of metadata nodes as the llvm::hlsl::rootsig::RootElement that generates it will have already been validated.

Resolves #126586

[llvmbot]

@llvm/pr-subscribers-hlsl

@llvm/pr-subscribers-clang

Author: Finn Plummer (inbelic)

Changes

Implements metadata generation of a Root Signature from its in-memory representation. It follows the same style as: #139633.

This pr handles StaticSamplers. It also handles converting the else-if chain into a std::visit to allow for future compiler warnings when adding additional RootElement variants.

The metadata follows the format described here.

• Implement BuildStaticSampler into HLSLRootSignature.h
• Add sample testcases demonstrating functionality

Note: there is no validation of metadata nodes as the llvm::hlsl::rootsig::RootElement that generates it will have already been validated.

Resolves #126586

---

Full diff: https://github.com/llvm/llvm-project/pull/142642.diff

3 Files Affected:

• (modified) clang/test/CodeGenHLSL/RootSignature.hlsl (+35-1)
• (modified) llvm/include/llvm/Frontend/HLSL/HLSLRootSignatureUtils.h (+1)
• (modified) llvm/lib/Frontend/HLSL/HLSLRootSignatureUtils.cpp (+66-16)

diff --git a/clang/test/CodeGenHLSL/RootSignature.hlsl b/clang/test/CodeGenHLSL/RootSignature.hlsl
index 34caa3eb6b714..ca843ffbb1ced 100644
--- a/clang/test/CodeGenHLSL/RootSignature.hlsl
+++ b/clang/test/CodeGenHLSL/RootSignature.hlsl
@@ -1,7 +1,7 @@
 // RUN: %clang_cc1 -triple dxil-pc-shadermodel6.3-library -emit-llvm -o - %s | FileCheck %s
 
 // CHECK: !dx.rootsignatures = !{![[#EMPTY_ENTRY:]], ![[#DT_ENTRY:]],
-// CHECK-SAME: ![[#RF_ENTRY:]], ![[#RC_ENTRY:]], ![[#RD_ENTRY:]]}
+// CHECK-SAME: ![[#RF_ENTRY:]], ![[#RC_ENTRY:]], ![[#RD_ENTRY:]], ![[#SS_ENTRY:]]}
 
 // CHECK: ![[#EMPTY_ENTRY]] = !{ptr @EmptyEntry, ![[#EMPTY:]]}
 // CHECK: ![[#EMPTY]] = !{}
@@ -66,6 +66,40 @@ void RootConstantsEntry() {}
 [numthreads(1,1,1)]
 void RootDescriptorsEntry() {}
 
+// CHECK: ![[#SS_ENTRY]] = !{ptr @StaticSamplerEntry, ![[#SS_RS:]]}
+// CHECK: ![[#SS_RS]] = !{![[#STATIC_SAMPLER:]]}
+
+// checking filter = 0x4
+// CHECK: ![[#STATIC_SAMPLER]] = !{!"StaticSampler", i32 4,
+
+// checking texture address[U|V|W]
+// CHECK-SAME: i32 2, i32 3, i32 5,
+
+// checking mipLODBias, maxAnisotropy, comparisonFunc, borderColor
+// CHECK-SAME: float 0x40403999A0000000, i32 9, i32 3, i32 2,
+
+// checking minLOD, maxLOD
+// CHECK-SAME: float -1.280000e+02, float 1.280000e+02,
+
+// checking register, space and visibility
+// CHECK-SAME: i32 42, i32 0, i32 0}
+
+#define SampleStaticSampler \
+  "StaticSampler(s42, " \
+  " filter = FILTER_MIN_POINT_MAG_LINEAR_MIP_POINT, " \
+  " addressU = TEXTURE_ADDRESS_MIRROR, " \
+  " addressV = TEXTURE_ADDRESS_CLAMP, " \
+  " addressW = TEXTURE_ADDRESS_MIRRORONCE, " \
+  " mipLODBias = 32.45f, maxAnisotropy = 9, " \
+  " comparisonFunc = COMPARISON_EQUAL, " \
+  " borderColor = STATIC_BORDER_COLOR_OPAQUE_WHITE, " \
+  " minLOD = -128.f, maxLOD = 128.f, " \
+  " space = 0, visibility = SHADER_VISIBILITY_ALL, " \
+  ")"
+[shader("compute"), RootSignature(SampleStaticSampler)]
+[numthreads(1,1,1)]
+void StaticSamplerEntry() {}
+
 // Sanity test to ensure no root is added for this function as there is only
 // two entries in !dx.roosignatures
 [shader("compute")]
diff --git a/llvm/include/llvm/Frontend/HLSL/HLSLRootSignatureUtils.h b/llvm/include/llvm/Frontend/HLSL/HLSLRootSignatureUtils.h
index 365197a4dfdb5..6d959ad5bdc7f 100644
--- a/llvm/include/llvm/Frontend/HLSL/HLSLRootSignatureUtils.h
+++ b/llvm/include/llvm/Frontend/HLSL/HLSLRootSignatureUtils.h
@@ -52,6 +52,7 @@ class MetadataBuilder {
   MDNode *BuildRootDescriptor(const RootDescriptor &Descriptor);
   MDNode *BuildDescriptorTable(const DescriptorTable &Table);
   MDNode *BuildDescriptorTableClause(const DescriptorTableClause &Clause);
+  MDNode *BuildStaticSampler(const StaticSampler &Sampler);
 
   llvm::LLVMContext &Ctx;
   ArrayRef<RootElement> Elements;
diff --git a/llvm/lib/Frontend/HLSL/HLSLRootSignatureUtils.cpp b/llvm/lib/Frontend/HLSL/HLSLRootSignatureUtils.cpp
index aa1f1957d9cbe..0a866374bf5dc 100644
--- a/llvm/lib/Frontend/HLSL/HLSLRootSignatureUtils.cpp
+++ b/llvm/lib/Frontend/HLSL/HLSLRootSignatureUtils.cpp
@@ -169,25 +169,44 @@ void dumpRootElements(raw_ostream &OS, ArrayRef<RootElement> Elements) {
   OS << "}";
 }
 
+namespace {
+
+// We use the OverloadBuild with std::visit to ensure the compiler catches if a
+// new RootElement variant type is added but it's metadata generation isn't
+// handled.
+template <class... Ts> struct OverloadedBuild : Ts... {
+  using Ts::operator()...;
+};
+template <class... Ts> OverloadedBuild(Ts...) -> OverloadedBuild<Ts...>;
+
+} // namespace
+
 MDNode *MetadataBuilder::BuildRootSignature() {
+  const auto Visitor = OverloadedBuild{
+      [this](const RootFlags &Flags) -> MDNode * {
+        return BuildRootFlags(Flags);
+      },
+      [this](const RootConstants &Constants) -> MDNode * {
+        return BuildRootConstants(Constants);
+      },
+      [this](const RootDescriptor &Descriptor) -> MDNode * {
+        return BuildRootDescriptor(Descriptor);
+      },
+      [this](const DescriptorTableClause &Clause) -> MDNode * {
+        return BuildDescriptorTableClause(Clause);
+      },
+      [this](const DescriptorTable &Table) -> MDNode * {
+        return BuildDescriptorTable(Table);
+      },
+      [this](const StaticSampler &Sampler) -> MDNode * {
+        return BuildStaticSampler(Sampler);
+      },
+  };
+
   for (const RootElement &Element : Elements) {
-    MDNode *ElementMD = nullptr;
-    if (const auto &Flags = std::get_if<RootFlags>(&Element))
-      ElementMD = BuildRootFlags(*Flags);
-    else if (const auto &Constants = std::get_if<RootConstants>(&Element))
-      ElementMD = BuildRootConstants(*Constants);
-    else if (const auto &Descriptor = std::get_if<RootDescriptor>(&Element))
-      ElementMD = BuildRootDescriptor(*Descriptor);
-    else if (const auto &Clause = std::get_if<DescriptorTableClause>(&Element))
-      ElementMD = BuildDescriptorTableClause(*Clause);
-    else if (const auto &Table = std::get_if<DescriptorTable>(&Element))
-      ElementMD = BuildDescriptorTable(*Table);
-
-    // FIXME(#126586): remove once all RootElemnt variants are handled in a
-    // visit or otherwise
+    MDNode *ElementMD = std::visit(Visitor, Element);
     assert(ElementMD != nullptr &&
-           "Constructed an unhandled root element type.");
-
+           "Root Element must be initialized and validated");
     GeneratedMetadata.push_back(ElementMD);
   }
 
@@ -274,6 +293,37 @@ MDNode *MetadataBuilder::BuildDescriptorTableClause(
            });
 }
 
+MDNode *MetadataBuilder::BuildStaticSampler(const StaticSampler &Sampler) {
+  IRBuilder<> Builder(Ctx);
+  Metadata *Operands[] = {
+      MDString::get(Ctx, "StaticSampler"),
+      ConstantAsMetadata::get(
+          Builder.getInt32(llvm::to_underlying(Sampler.Filter))),
+      ConstantAsMetadata::get(
+          Builder.getInt32(llvm::to_underlying(Sampler.AddressU))),
+      ConstantAsMetadata::get(
+          Builder.getInt32(llvm::to_underlying(Sampler.AddressV))),
+      ConstantAsMetadata::get(
+          Builder.getInt32(llvm::to_underlying(Sampler.AddressW))),
+      ConstantAsMetadata::get(llvm::ConstantFP::get(llvm::Type::getFloatTy(Ctx),
+                                                    Sampler.MipLODBias)),
+      ConstantAsMetadata::get(Builder.getInt32(Sampler.MaxAnisotropy)),
+      ConstantAsMetadata::get(
+          Builder.getInt32(llvm::to_underlying(Sampler.CompFunc))),
+      ConstantAsMetadata::get(
+          Builder.getInt32(llvm::to_underlying(Sampler.BorderColor))),
+      ConstantAsMetadata::get(
+          llvm::ConstantFP::get(llvm::Type::getFloatTy(Ctx), Sampler.MinLOD)),
+      ConstantAsMetadata::get(
+          llvm::ConstantFP::get(llvm::Type::getFloatTy(Ctx), Sampler.MaxLOD)),
+      ConstantAsMetadata::get(Builder.getInt32(Sampler.Reg.Number)),
+      ConstantAsMetadata::get(Builder.getInt32(Sampler.Space)),
+      ConstantAsMetadata::get(
+          Builder.getInt32(llvm::to_underlying(Sampler.Visibility))),
+  };
+  return MDNode::get(Ctx, Operands);
+}
+
 } // namespace rootsig
 } // namespace hlsl
 } // namespace llvm