[DirectX] Add static sampler support to root signature (#143422)

Implements static samplers parsing from root signature metadata
representation. This is required to support Root Signatures in HLSL.
Closes: #[126641](#126641)

---------

Co-authored-by: joaosaffran <[email protected]>

Author: joaosaffran

llvm/lib/MC/DXContainerRootSignature.cpp

@@ -71,12 +71,16 @@ void RootSignatureDesc::write(raw_ostream &OS) const {
   BOS.reserveExtraSpace(getSize());
 
   const uint32_t NumParameters = ParametersContainer.size();
-
+  const uint32_t NumSamplers = StaticSamplers.size();
   support::endian::write(BOS, Version, llvm::endianness::little);
   support::endian::write(BOS, NumParameters, llvm::endianness::little);
   support::endian::write(BOS, RootParameterOffset, llvm::endianness::little);
-  support::endian::write(BOS, NumStaticSamplers, llvm::endianness::little);
-  support::endian::write(BOS, StaticSamplersOffset, llvm::endianness::little);
+  support::endian::write(BOS, NumSamplers, llvm::endianness::little);
+  uint32_t SSO = StaticSamplersOffset;
+  if (NumSamplers > 0)
+    SSO = writePlaceholder(BOS);
+  else
+    support::endian::write(BOS, SSO, llvm::endianness::little);
   support::endian::write(BOS, Flags, llvm::endianness::little);
 
   SmallVector<uint32_t> ParamsOffsets;
@@ -142,20 +146,23 @@ void RootSignatureDesc::write(raw_ostream &OS) const {
     }
     }
   }
-  for (const auto &S : StaticSamplers) {
-    support::endian::write(BOS, S.Filter, llvm::endianness::little);
-    support::endian::write(BOS, S.AddressU, llvm::endianness::little);
-    support::endian::write(BOS, S.AddressV, llvm::endianness::little);
-    support::endian::write(BOS, S.AddressW, llvm::endianness::little);
-    support::endian::write(BOS, S.MipLODBias, llvm::endianness::little);
-    support::endian::write(BOS, S.MaxAnisotropy, llvm::endianness::little);
-    support::endian::write(BOS, S.ComparisonFunc, llvm::endianness::little);
-    support::endian::write(BOS, S.BorderColor, llvm::endianness::little);
-    support::endian::write(BOS, S.MinLOD, llvm::endianness::little);
-    support::endian::write(BOS, S.MaxLOD, llvm::endianness::little);
-    support::endian::write(BOS, S.ShaderRegister, llvm::endianness::little);
-    support::endian::write(BOS, S.RegisterSpace, llvm::endianness::little);
-    support::endian::write(BOS, S.ShaderVisibility, llvm::endianness::little);
+  if (NumSamplers > 0) {
+    rewriteOffsetToCurrentByte(BOS, SSO);
+    for (const auto &S : StaticSamplers) {
+      support::endian::write(BOS, S.Filter, llvm::endianness::little);
+      support::endian::write(BOS, S.AddressU, llvm::endianness::little);
+      support::endian::write(BOS, S.AddressV, llvm::endianness::little);
+      support::endian::write(BOS, S.AddressW, llvm::endianness::little);
+      support::endian::write(BOS, S.MipLODBias, llvm::endianness::little);
+      support::endian::write(BOS, S.MaxAnisotropy, llvm::endianness::little);
+      support::endian::write(BOS, S.ComparisonFunc, llvm::endianness::little);
+      support::endian::write(BOS, S.BorderColor, llvm::endianness::little);
+      support::endian::write(BOS, S.MinLOD, llvm::endianness::little);
+      support::endian::write(BOS, S.MaxLOD, llvm::endianness::little);
+      support::endian::write(BOS, S.ShaderRegister, llvm::endianness::little);
+      support::endian::write(BOS, S.RegisterSpace, llvm::endianness::little);
+      support::endian::write(BOS, S.ShaderVisibility, llvm::endianness::little);
+    }
   }
   assert(Storage.size() == getSize());
   OS.write(Storage.data(), Storage.size());

llvm/lib/Target/DirectX/DXILRootSignature.cpp

@@ -27,6 +27,7 @@
 #include "llvm/Support/Error.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/raw_ostream.h"
+#include <cmath>
 #include <cstdint>
 #include <optional>
 #include <utility>
@@ -55,6 +56,13 @@ static std::optional<uint32_t> extractMdIntValue(MDNode *Node,
   return std::nullopt;
 }
 
+static std::optional<float> extractMdFloatValue(MDNode *Node,
+                                                unsigned int OpId) {
+  if (auto *CI = mdconst::dyn_extract<ConstantFP>(Node->getOperand(OpId).get()))
+    return CI->getValueAPF().convertToFloat();
+  return std::nullopt;
+}
+
 static std::optional<StringRef> extractMdStringValue(MDNode *Node,
                                                      unsigned int OpId) {
   MDString *NodeText = dyn_cast<MDString>(Node->getOperand(OpId));
@@ -261,6 +269,81 @@ static bool parseDescriptorTable(LLVMContext *Ctx,
   return false;
 }
 
+static bool parseStaticSampler(LLVMContext *Ctx, mcdxbc::RootSignatureDesc &RSD,
+                               MDNode *StaticSamplerNode) {
+  if (StaticSamplerNode->getNumOperands() != 14)
+    return reportError(Ctx, "Invalid format for Static Sampler");
+
+  dxbc::RTS0::v1::StaticSampler Sampler;
+  if (std::optional<uint32_t> Val = extractMdIntValue(StaticSamplerNode, 1))
+    Sampler.Filter = *Val;
+  else
+    return reportError(Ctx, "Invalid value for Filter");
+
+  if (std::optional<uint32_t> Val = extractMdIntValue(StaticSamplerNode, 2))
+    Sampler.AddressU = *Val;
+  else
+    return reportError(Ctx, "Invalid value for AddressU");
+
+  if (std::optional<uint32_t> Val = extractMdIntValue(StaticSamplerNode, 3))
+    Sampler.AddressV = *Val;
+  else
+    return reportError(Ctx, "Invalid value for AddressV");
+
+  if (std::optional<uint32_t> Val = extractMdIntValue(StaticSamplerNode, 4))
+    Sampler.AddressW = *Val;
+  else
+    return reportError(Ctx, "Invalid value for AddressW");
+
+  if (std::optional<float> Val = extractMdFloatValue(StaticSamplerNode, 5))
+    Sampler.MipLODBias = *Val;
+  else
+    return reportError(Ctx, "Invalid value for MipLODBias");
+
+  if (std::optional<uint32_t> Val = extractMdIntValue(StaticSamplerNode, 6))
+    Sampler.MaxAnisotropy = *Val;
+  else
+    return reportError(Ctx, "Invalid value for MaxAnisotropy");
+
+  if (std::optional<uint32_t> Val = extractMdIntValue(StaticSamplerNode, 7))
+    Sampler.ComparisonFunc = *Val;
+  else
+    return reportError(Ctx, "Invalid value for ComparisonFunc ");
+
+  if (std::optional<uint32_t> Val = extractMdIntValue(StaticSamplerNode, 8))
+    Sampler.BorderColor = *Val;
+  else
+    return reportError(Ctx, "Invalid value for ComparisonFunc ");
+
+  if (std::optional<float> Val = extractMdFloatValue(StaticSamplerNode, 9))
+    Sampler.MinLOD = *Val;
+  else
+    return reportError(Ctx, "Invalid value for MinLOD");
+
+  if (std::optional<float> Val = extractMdFloatValue(StaticSamplerNode, 10))
+    Sampler.MaxLOD = *Val;
+  else
+    return reportError(Ctx, "Invalid value for MaxLOD");
+
+  if (std::optional<uint32_t> Val = extractMdIntValue(StaticSamplerNode, 11))
+    Sampler.ShaderRegister = *Val;
+  else
+    return reportError(Ctx, "Invalid value for ShaderRegister");
+
+  if (std::optional<uint32_t> Val = extractMdIntValue(StaticSamplerNode, 12))
+    Sampler.RegisterSpace = *Val;
+  else
+    return reportError(Ctx, "Invalid value for RegisterSpace");
+
+  if (std::optional<uint32_t> Val = extractMdIntValue(StaticSamplerNode, 13))
+    Sampler.ShaderVisibility = *Val;
+  else
+    return reportError(Ctx, "Invalid value for ShaderVisibility");
+
+  RSD.StaticSamplers.push_back(Sampler);
+  return false;
+}
+
 static bool parseRootSignatureElement(LLVMContext *Ctx,
                                       mcdxbc::RootSignatureDesc &RSD,
                                       MDNode *Element) {
@@ -276,6 +359,7 @@ static bool parseRootSignatureElement(LLVMContext *Ctx,
           .Case("RootSRV", RootSignatureElementKind::SRV)
           .Case("RootUAV", RootSignatureElementKind::UAV)
           .Case("DescriptorTable", RootSignatureElementKind::DescriptorTable)
+          .Case("StaticSampler", RootSignatureElementKind::StaticSamplers)
           .Default(RootSignatureElementKind::Error);
 
   switch (ElementKind) {
@@ -290,6 +374,8 @@ static bool parseRootSignatureElement(LLVMContext *Ctx,
     return parseRootDescriptors(Ctx, RSD, Element, ElementKind);
   case RootSignatureElementKind::DescriptorTable:
     return parseDescriptorTable(Ctx, RSD, Element);
+  case RootSignatureElementKind::StaticSamplers:
+    return parseStaticSampler(Ctx, RSD, Element);
   case RootSignatureElementKind::Error:
     return reportError(Ctx, "Invalid Root Signature Element: " + *ElementText);
   }
@@ -406,6 +492,58 @@ static bool verifyDescriptorRangeFlag(uint32_t Version, uint32_t Type,
   return (Flags & ~Mask) == FlagT::NONE;
 }
 
+static bool verifySamplerFilter(uint32_t Value) {
+  switch (Value) {
+#define STATIC_SAMPLER_FILTER(Num, Val)                                        \
+  case llvm::to_underlying(dxbc::StaticSamplerFilter::Val):
+#include "llvm/BinaryFormat/DXContainerConstants.def"
+    return true;
+  }
+  return false;
+}
+
+// Values allowed here:
+// https://learn.microsoft.com/en-us/windows/win32/api/d3d12/ne-d3d12-d3d12_texture_address_mode#syntax
+static bool verifyAddress(uint32_t Address) {
+  switch (Address) {
+#define TEXTURE_ADDRESS_MODE(Num, Val)                                         \
+  case llvm::to_underlying(dxbc::TextureAddressMode::Val):
+#include "llvm/BinaryFormat/DXContainerConstants.def"
+    return true;
+  }
+  return false;
+}
+
+static bool verifyMipLODBias(float MipLODBias) {
+  return MipLODBias >= -16.f && MipLODBias <= 15.99f;
+}
+
+static bool verifyMaxAnisotropy(uint32_t MaxAnisotropy) {
+  return MaxAnisotropy <= 16u;
+}
+
+static bool verifyComparisonFunc(uint32_t ComparisonFunc) {
+  switch (ComparisonFunc) {
+#define COMPARISON_FUNCTION(Num, Val)                                          \
+  case llvm::to_underlying(dxbc::SamplersComparisonFunction::Val):
+#include "llvm/BinaryFormat/DXContainerConstants.def"
+    return true;
+  }
+  return false;
+}
+
+static bool verifyBorderColor(uint32_t BorderColor) {
+  switch (BorderColor) {
+#define STATIC_BORDER_COLOR(Num, Val)                                          \
+  case llvm::to_underlying(dxbc::SamplersBorderColor::Val):
+#include "llvm/BinaryFormat/DXContainerConstants.def"
+    return true;
+  }
+  return false;
+}
+
+static bool verifyLOD(float LOD) { return !std::isnan(LOD); }
+
 static bool validate(LLVMContext *Ctx, const mcdxbc::RootSignatureDesc &RSD) {
 
   if (!verifyVersion(RSD.Version)) {
@@ -463,6 +601,48 @@ static bool validate(LLVMContext *Ctx, const mcdxbc::RootSignatureDesc &RSD) {
     }
   }
 
+  for (const dxbc::RTS0::v1::StaticSampler &Sampler : RSD.StaticSamplers) {
+    if (!verifySamplerFilter(Sampler.Filter))
+      return reportValueError(Ctx, "Filter", Sampler.Filter);
+
+    if (!verifyAddress(Sampler.AddressU))
+      return reportValueError(Ctx, "AddressU", Sampler.AddressU);
+
+    if (!verifyAddress(Sampler.AddressV))
+      return reportValueError(Ctx, "AddressV", Sampler.AddressV);
+
+    if (!verifyAddress(Sampler.AddressW))
+      return reportValueError(Ctx, "AddressW", Sampler.AddressW);
+
+    if (!verifyMipLODBias(Sampler.MipLODBias))
+      return reportValueError(Ctx, "MipLODBias", Sampler.MipLODBias);
+
+    if (!verifyMaxAnisotropy(Sampler.MaxAnisotropy))
+      return reportValueError(Ctx, "MaxAnisotropy", Sampler.MaxAnisotropy);
+
+    if (!verifyComparisonFunc(Sampler.ComparisonFunc))
+      return reportValueError(Ctx, "ComparisonFunc", Sampler.ComparisonFunc);
+
+    if (!verifyBorderColor(Sampler.BorderColor))
+      return reportValueError(Ctx, "BorderColor", Sampler.BorderColor);
+
+    if (!verifyLOD(Sampler.MinLOD))
+      return reportValueError(Ctx, "MinLOD", Sampler.MinLOD);
+
+    if (!verifyLOD(Sampler.MaxLOD))
+      return reportValueError(Ctx, "MaxLOD", Sampler.MaxLOD);
+
+    if (!verifyRegisterValue(Sampler.ShaderRegister))
+      return reportValueError(Ctx, "ShaderRegister", Sampler.ShaderRegister);
+
+    if (!verifyRegisterSpace(Sampler.RegisterSpace))
+      return reportValueError(Ctx, "RegisterSpace", Sampler.RegisterSpace);
+
+    if (!dxbc::isValidShaderVisibility(Sampler.ShaderVisibility))
+      return reportValueError(Ctx, "ShaderVisibility",
+                              Sampler.ShaderVisibility);
+  }
+
   return false;
 }
 
@@ -542,6 +722,9 @@ analyzeModule(Module &M) {
     // offset will always equal to the header size.
     RSD.RootParameterOffset = sizeof(dxbc::RTS0::v1::RootSignatureHeader);
 
+    // static sampler offset is calculated when writting dxcontainer.
+    RSD.StaticSamplersOffset = 0u;
+
     if (parse(Ctx, RSD, RootElementListNode) || validate(Ctx, RSD)) {
       return RSDMap;
     }

llvm/lib/Target/DirectX/DXILRootSignature.h

@@ -32,6 +32,7 @@ enum class RootSignatureElementKind {
   UAV = 4,
   CBV = 5,
   DescriptorTable = 6,
+  StaticSamplers = 7
 };
 class RootSignatureAnalysis : public AnalysisInfoMixin<RootSignatureAnalysis> {
   friend AnalysisInfoMixin<RootSignatureAnalysis>;

llvm/test/CodeGen/DirectX/ContainerData/RootSignature-StaticSamplers-Invalid-AddressU.ll

@@ -0,0 +1,19 @@
+; RUN: not opt -passes='print<dxil-root-signature>' %s -S -o - 2>&1 | FileCheck %s
+
+
+target triple = "dxil-unknown-shadermodel6.0-compute"
+
+; CHECK: error: Invalid value for AddressU: 666
+; CHECK-NOT: Root Signature Definitions
+
+define void @main() #0 {
+entry:
+  ret void
+}
+attributes #0 = { "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" }
+
+
+!dx.rootsignatures = !{!2} ; list of function/root signature pairs
+!2 = !{ ptr @main, !3, i32 2 } ; function, root signature
+!3 = !{ !5 } ; list of root signature elements
+!5 = !{ !"StaticSampler", i32 4, i32 666, i32 3, i32 5, float 0x3FF6CCCCC0000000, i32 9, i32 3, i32 2, float -1.280000e+02, float 1.280000e+02, i32 42, i32 0, i32 0 }

llvm/test/CodeGen/DirectX/ContainerData/RootSignature-StaticSamplers-Invalid-AddressV.ll

@@ -0,0 +1,19 @@
+; RUN: not opt -passes='print<dxil-root-signature>' %s -S -o - 2>&1 | FileCheck %s
+
+
+target triple = "dxil-unknown-shadermodel6.0-compute"
+
+; CHECK: error: Invalid value for AddressV: 666
+; CHECK-NOT: Root Signature Definitions
+
+define void @main() #0 {
+entry:
+  ret void
+}
+attributes #0 = { "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" }
+
+
+!dx.rootsignatures = !{!2} ; list of function/root signature pairs
+!2 = !{ ptr @main, !3, i32 2 } ; function, root signature
+!3 = !{ !5 } ; list of root signature elements
+!5 = !{ !"StaticSampler", i32 4, i32 2, i32 666, i32 5, float 0x3FF6CCCCC0000000, i32 9, i32 3, i32 2, float -1.280000e+02, float 1.280000e+02, i32 42, i32 0, i32 0 }

llvm/test/CodeGen/DirectX/ContainerData/RootSignature-StaticSamplers-Invalid-AddressW.ll

@@ -0,0 +1,19 @@
+; RUN: not opt -passes='print<dxil-root-signature>' %s -S -o - 2>&1 | FileCheck %s
+
+
+target triple = "dxil-unknown-shadermodel6.0-compute"
+
+; CHECK: error: Invalid value for AddressW: 666
+; CHECK-NOT: Root Signature Definitions
+
+define void @main() #0 {
+entry:
+  ret void
+}
+attributes #0 = { "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" }
+
+
+!dx.rootsignatures = !{!2} ; list of function/root signature pairs
+!2 = !{ ptr @main, !3, i32 2 } ; function, root signature
+!3 = !{ !5 } ; list of root signature elements
+!5 = !{ !"StaticSampler", i32 4, i32 2, i32 3, i32 666, float 0x3FF6CCCCC0000000, i32 9, i32 3, i32 2, float -1.280000e+02, float 1.280000e+02, i32 42, i32 0, i32 0 }

llvm/test/CodeGen/DirectX/ContainerData/RootSignature-StaticSamplers-Invalid-BorderColor.ll

@@ -0,0 +1,19 @@
+; RUN: not opt -passes='print<dxil-root-signature>' %s -S -o - 2>&1 | FileCheck %s
+
+
+target triple = "dxil-unknown-shadermodel6.0-compute"
+
+; CHECK: error: Invalid value for BorderColor: 666
+; CHECK-NOT: Root Signature Definitions
+
+define void @main() #0 {
+entry:
+  ret void
+}
+attributes #0 = { "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" }
+
+
+!dx.rootsignatures = !{!2} ; list of function/root signature pairs
+!2 = !{ ptr @main, !3, i32 2 } ; function, root signature
+!3 = !{ !5 } ; list of root signature elements
+!5 = !{ !"StaticSampler", i32 4, i32 2, i32 3, i32 5, float 0x3FF6CCCCC0000000, i32 9, i32 3, i32 666, float -1.280000e+02, float 1.280000e+02, i32 42, i32 0, i32 0 }