[HLSL] Refactor Offset calculation while writing Root Signatures. #128577

joaosaffran · 2025-02-24T21:05:50Z

This PR is refactoring Root Signatures serialization process to remove the need to manually calculate offsets. The tests being changed are not functional, since Static Samplers are still not implemented.

llvmbot · 2025-02-24T21:06:30Z

@llvm/pr-subscribers-objectyaml
@llvm/pr-subscribers-mc

@llvm/pr-subscribers-backend-directx

Author: None (joaosaffran)

Changes

This PR is refactoring Root Signatures serialization process to remove the need to manually calculate offsets.

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

9 Files Affected:

(modified) llvm/include/llvm/MC/DXContainerRootSignature.h (+26-3)
(modified) llvm/lib/MC/DXContainerRootSignature.cpp (+96-29)
(modified) llvm/lib/Object/DXContainer.cpp (+6-4)
(modified) llvm/lib/ObjectYAML/DXContainerEmitter.cpp (+3-1)
(modified) llvm/lib/Target/DirectX/DXContainerGlobals.cpp (+4-1)
(modified) llvm/test/CodeGen/DirectX/ContainerData/RootSignature-Flags.ll (+1-1)
(modified) llvm/test/ObjectYAML/DXContainer/RootSignature-Flags.yaml (+1-1)
(modified) llvm/test/ObjectYAML/DXContainer/RootSignature-MultipleParameters.yaml (+1-1)
(modified) llvm/unittests/ObjectYAML/DXContainerYAMLTest.cpp (+4-4)

diff --git a/llvm/include/llvm/MC/DXContainerRootSignature.h b/llvm/include/llvm/MC/DXContainerRootSignature.h
index b31b0da352038..ca958c0780cdd 100644
--- a/llvm/include/llvm/MC/DXContainerRootSignature.h
+++ b/llvm/include/llvm/MC/DXContainerRootSignature.h
@@ -7,19 +7,42 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/BinaryFormat/DXContainer.h"
-#include <cstdint>
-#include <limits>
+#include "llvm/Support/BinaryStreamWriter.h"
+#include "llvm/Support/raw_ostream.h"
+#include <map>
+#include <string>
 
 namespace llvm {
 
 class raw_ostream;
 
 namespace mcdxbc {
+
+class StreamOffsetHelper {
+private:
+  std::map<std::string, std::pair<uint32_t, uint32_t>> OffsetsMaping;
+  BinaryStreamWriter &Stream;
+
+public:
+  explicit StreamOffsetHelper(BinaryStreamWriter &Stream) : Stream(Stream) {}
+
+  Error addOffset(std::string Key);
+
+  void addRewriteValue(std::string Key);
+
+  Error rewrite();
+};
+
 struct RootSignatureDesc {
   dxbc::RootSignatureHeader Header;
   SmallVector<dxbc::RootParameter> Parameters;
 
-  void write(raw_ostream &OS) const;
+  Error write(raw_ostream &OS) const;
+
+  uint32_t getSizeInBytes() const {
+    // Header Size + accounting for parameter offset + parameters size
+    return 24 + (Parameters.size() * 4) + Parameters.size_in_bytes();
+  }
 };
 } // namespace mcdxbc
 } // namespace llvm
diff --git a/llvm/lib/MC/DXContainerRootSignature.cpp b/llvm/lib/MC/DXContainerRootSignature.cpp
index 35a4ef322d01e..828785980c15d 100644
--- a/llvm/lib/MC/DXContainerRootSignature.cpp
+++ b/llvm/lib/MC/DXContainerRootSignature.cpp
@@ -7,53 +7,120 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/MC/DXContainerRootSignature.h"
-#include "llvm/Support/EndianStream.h"
+#include "llvm/ADT/Twine.h"
+#include "llvm/Support/BinaryStreamWriter.h"
 
 using namespace llvm;
 using namespace llvm::mcdxbc;
 
-void RootSignatureDesc::write(raw_ostream &OS) const {
-  // Root signature header in dxcontainer has 6 uint_32t values.
-  const uint32_t HeaderSize = 24;
-  const uint32_t ParameterByteSize = Parameters.size_in_bytes();
-  const uint32_t NumParametes = Parameters.size();
+Error StreamOffsetHelper::addOffset(std::string Key) {
+  const uint32_t DummyValue = std::numeric_limits<uint32_t>::max();
+
+  uint32_t Offset = Stream.getOffset();
+  auto Value = std::make_pair(Offset, DummyValue);
+
+  OffsetsMaping.insert_or_assign(Key, Value);
+
+  if (Error Err = Stream.writeInteger(DummyValue))
+    return Err;
+
+  return Error::success();
+}
+
+void StreamOffsetHelper::addRewriteValue(std::string Key) {
+  auto It = OffsetsMaping.find(Key);
+  assert(It != OffsetsMaping.end() && "Offset address was not found.");
+  auto [Offset, _] = It->second;
+
+  uint32_t Value = Stream.getOffset();
+
+  std::pair<uint32_t, uint32_t> NewValue = std::make_pair(Offset, Value);
+  OffsetsMaping.insert_or_assign(Key, NewValue);
+}
+
+Error StreamOffsetHelper::rewrite() {
+  for (auto &[Key, RewriteInfo] : OffsetsMaping) {
+    auto [Position, Value] = RewriteInfo;
+    assert(Value != std::numeric_limits<uint32_t>::max());
+
+    Stream.setOffset(Position);
+    if (Error Err = Stream.writeInteger(Value))
+      return Err;
+  }
+
+  return Error::success();
+}
+
+Error RootSignatureDesc::write(raw_ostream &OS) const {
+  std::vector<uint8_t> Buffer(getSizeInBytes());
+  BinaryStreamWriter Writer(Buffer, llvm::endianness::little);
+
+  StreamOffsetHelper OffsetMap(Writer);
+
+  const uint32_t NumParameters = Parameters.size();
   const uint32_t Zero = 0;
 
-  // Writing header information
-  support::endian::write(OS, Header.Version, llvm::endianness::little);
-  support::endian::write(OS, NumParametes, llvm::endianness::little);
-  support::endian::write(OS, HeaderSize, llvm::endianness::little);
+  if (Error Err = Writer.writeInteger(Header.Version))
+    return Err;
+
+  if (Error Err = Writer.writeInteger(NumParameters))
+    return Err;
+
+  if (Error Err = OffsetMap.addOffset("header"))
+    return Err;
 
   // Static samplers still not implemented
-  support::endian::write(OS, Zero, llvm::endianness::little);
-  support::endian::write(OS, ParameterByteSize + HeaderSize,
-                         llvm::endianness::little);
+  if (Error Err = Writer.writeInteger(Zero))
+    return Err;
+
+  if (Error Err = Writer.writeInteger(Zero))
+    return Err;
+
+  if (Error Err = Writer.writeInteger(Header.Flags))
+    return Err;
 
-  support::endian::write(OS, Header.Flags, llvm::endianness::little);
+  OffsetMap.addRewriteValue("header");
 
-  uint32_t ParamsOffset =
-      HeaderSize + (3 * sizeof(uint32_t) * Parameters.size());
-  for (const dxbc::RootParameter &P : Parameters) {
-    support::endian::write(OS, P.ParameterType, llvm::endianness::little);
-    support::endian::write(OS, P.ShaderVisibility, llvm::endianness::little);
-    support::endian::write(OS, ParamsOffset, llvm::endianness::little);
+  for (size_t It = 0; It < Parameters.size(); It++) {
+    const auto &P = Parameters[It];
 
-    // Size of root parameter, removing the ParameterType and ShaderVisibility.
-    ParamsOffset += sizeof(dxbc::RootParameter) - 2 * sizeof(uint32_t);
+    if (Error Err = Writer.writeEnum(P.ParameterType))
+      return Err;
+
+    if (Error Err = Writer.writeEnum(P.ShaderVisibility))
+      return Err;
+
+    std::string Key = ("parameters" + Twine(It)).str();
+    if (Error Err = OffsetMap.addOffset(Key))
+      return Err;
   }
 
-  for (const dxbc::RootParameter &P : Parameters) {
+  for (size_t It = 0; It < Parameters.size(); It++) {
+    const auto &P = Parameters[It];
+
+    std::string Key = ("parameters" + Twine(It)).str();
+    OffsetMap.addRewriteValue(Key);
+
     switch (P.ParameterType) {
     case dxbc::RootParameterType::Constants32Bit: {
-      support::endian::write(OS, P.Constants.ShaderRegister,
-                             llvm::endianness::little);
-      support::endian::write(OS, P.Constants.RegisterSpace,
-                             llvm::endianness::little);
-      support::endian::write(OS, P.Constants.Num32BitValues,
-                             llvm::endianness::little);
+      if (Error Err = Writer.writeInteger(P.Constants.ShaderRegister))
+        return Err;
+      if (Error Err = Writer.writeInteger(P.Constants.RegisterSpace))
+        return Err;
+      if (Error Err = Writer.writeInteger(P.Constants.Num32BitValues))
+        return Err;
     } break;
     case dxbc::RootParameterType::Empty:
       llvm_unreachable("Invalid RootParameterType");
     }
   }
+
+  if (Error Err = OffsetMap.rewrite())
+    return Err;
+
+  llvm::ArrayRef<char> BufferRef(reinterpret_cast<char *>(Buffer.data()),
+                                 Buffer.size());
+  OS.write(BufferRef.data(), BufferRef.size());
+
+  return Error::success();
 }
diff --git a/llvm/lib/Object/DXContainer.cpp b/llvm/lib/Object/DXContainer.cpp
index 010f70a952ebf..35261b661cf2f 100644
--- a/llvm/lib/Object/DXContainer.cpp
+++ b/llvm/lib/Object/DXContainer.cpp
@@ -296,16 +296,18 @@ Error DirectX::RootSignature::parse(StringRef Data) {
     NewParam.ParameterType =
         support::endian::read<dxbc::RootParameterType,
                               llvm::endianness::little>(Current);
-    if (!dxbc::RootSignatureValidations::isValidParameterType(NewParam.ParameterType))
+    if (!dxbc::RootSignatureValidations::isValidParameterType(
+            NewParam.ParameterType))
       return validationFailed("unsupported parameter type value read: " +
                               llvm::Twine((uint32_t)NewParam.ParameterType));
 
     Current += sizeof(dxbc::RootParameterType);
 
     NewParam.ShaderVisibility =
-        support::endian::read<dxbc::ShaderVisibility,
-                              llvm::endianness::little>(Current);
-    if (!dxbc::RootSignatureValidations::isValidShaderVisibility(NewParam.ShaderVisibility))
+        support::endian::read<dxbc::ShaderVisibility, llvm::endianness::little>(
+            Current);
+    if (!dxbc::RootSignatureValidations::isValidShaderVisibility(
+            NewParam.ShaderVisibility))
       return validationFailed("unsupported shader visility flag value read: " +
                               llvm::Twine((uint32_t)NewParam.ShaderVisibility));
 
diff --git a/llvm/lib/ObjectYAML/DXContainerEmitter.cpp b/llvm/lib/ObjectYAML/DXContainerEmitter.cpp
index 87ba16fd40ba9..a5831a69f9bca 100644
--- a/llvm/lib/ObjectYAML/DXContainerEmitter.cpp
+++ b/llvm/lib/ObjectYAML/DXContainerEmitter.cpp
@@ -271,7 +271,9 @@ void DXContainerWriter::writeParts(raw_ostream &OS) {
       RS.Header.Version = P.RootSignature->Version;
       RS.Parameters = std::move(P.RootSignature->Parameters);
 
-      RS.write(OS);
+      if (Error Err = RS.write(OS))
+        handleAllErrors(std::move(Err));
+
       break;
     }
     uint64_t BytesWritten = OS.tell() - DataStart;
diff --git a/llvm/lib/Target/DirectX/DXContainerGlobals.cpp b/llvm/lib/Target/DirectX/DXContainerGlobals.cpp
index 5508af40663b1..5801046f83674 100644
--- a/llvm/lib/Target/DirectX/DXContainerGlobals.cpp
+++ b/llvm/lib/Target/DirectX/DXContainerGlobals.cpp
@@ -25,10 +25,12 @@
 #include "llvm/InitializePasses.h"
 #include "llvm/MC/DXContainerPSVInfo.h"
 #include "llvm/Pass.h"
+#include "llvm/Support/Error.h"
 #include "llvm/Support/MD5.h"
 #include "llvm/TargetParser/Triple.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
 #include <optional>
+#include <utility>
 
 using namespace llvm;
 using namespace llvm::dxil;
@@ -173,7 +175,8 @@ void DXContainerGlobals::addRootSignature(Module &M,
   SmallString<256> Data;
   raw_svector_ostream OS(Data);
 
-  RS.write(OS);
+  if (Error Err = RS.write(OS))
+    handleAllErrors(std::move(Err));
 
   Constant *Constant =
       ConstantDataArray::getString(M.getContext(), Data, /*AddNull*/ false);
diff --git a/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-Flags.ll b/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-Flags.ll
index 035ea373d30ea..1ca6ebb7ddab8 100644
--- a/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-Flags.ll
+++ b/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-Flags.ll
@@ -23,6 +23,6 @@ attributes #0 = { "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" }
 ; DXC-NEXT:    RootSignature:
 ; DXC-NEXT:      Version:         2
 ; DXC-NEXT:      NumStaticSamplers: 0
-; DXC-NEXT:      StaticSamplersOffset: 24
+; DXC-NEXT:      StaticSamplersOffset: 0
 ; DXC-NEXT:      Parameters: []
 ; DXC-NEXT:      AllowInputAssemblerInputLayout: true
diff --git a/llvm/test/ObjectYAML/DXContainer/RootSignature-Flags.yaml b/llvm/test/ObjectYAML/DXContainer/RootSignature-Flags.yaml
index 1b73b830f015f..0d7902afdaa66 100644
--- a/llvm/test/ObjectYAML/DXContainer/RootSignature-Flags.yaml
+++ b/llvm/test/ObjectYAML/DXContainer/RootSignature-Flags.yaml
@@ -25,7 +25,7 @@ Parts:
 # CHECK-NEXT:    RootSignature:
 # CHECK-NEXT:      Version: 2
 # CHECK-NEXT:      NumStaticSamplers: 0
-# CHECK-NEXT:      StaticSamplersOffset: 24
+# CHECK-NEXT:      StaticSamplersOffset: 0
 # CHECK-NEXT:      Parameters: []
 # CHECK-NEXT:      AllowInputAssemblerInputLayout: true
 # CHECK-NEXT:      DenyGeometryShaderRootAccess: true
diff --git a/llvm/test/ObjectYAML/DXContainer/RootSignature-MultipleParameters.yaml b/llvm/test/ObjectYAML/DXContainer/RootSignature-MultipleParameters.yaml
index bccfacb72819b..8d5ab5c1b0b23 100644
--- a/llvm/test/ObjectYAML/DXContainer/RootSignature-MultipleParameters.yaml
+++ b/llvm/test/ObjectYAML/DXContainer/RootSignature-MultipleParameters.yaml
@@ -37,7 +37,7 @@ Parts:
 # CHECK-NEXT:    RootSignature:
 # CHECK-NEXT:      Version:         2
 # CHECK-NEXT:      NumStaticSamplers: 0
-# CHECK-NEXT:      StaticSamplersOffset: 64
+# CHECK-NEXT:      StaticSamplersOffset: 0
 # CHECK-NEXT:      Parameters:
 # CHECK-NEXT:        - ParameterType:   Constants32Bit
 # CHECK-NEXT:          ShaderVisibility: Hull
diff --git a/llvm/unittests/ObjectYAML/DXContainerYAMLTest.cpp b/llvm/unittests/ObjectYAML/DXContainerYAMLTest.cpp
index b5fb8d143c0b9..fed941f685272 100644
--- a/llvm/unittests/ObjectYAML/DXContainerYAMLTest.cpp
+++ b/llvm/unittests/ObjectYAML/DXContainerYAMLTest.cpp
@@ -134,12 +134,12 @@ TEST(RootSignature, ParseRootFlags) {
     )"));
 
   uint8_t Buffer[] = {
-      0x44, 0x58, 0x42, 0x43, 0x32, 0x9A, 0x53, 0xD8, 0xEC, 0xBE, 0x35, 0x6F,
-      0x05, 0x39, 0xE1, 0xFE, 0x31, 0x20, 0xF0, 0xC1, 0x01, 0x00, 0x00, 0x00,
+      0x44, 0x58, 0x42, 0x43, 0x32, 0x9a, 0x53, 0xd8, 0xec, 0xbe, 0x35, 0x6f,
+      0x05, 0x39, 0xe1, 0xfe, 0x31, 0x20, 0xf0, 0xc1, 0x01, 0x00, 0x00, 0x00,
       0x44, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x24, 0x00, 0x00, 0x00,
       0x52, 0x54, 0x53, 0x30, 0x18, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00,
       0x00, 0x00, 0x00, 0x00, 0x18, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
-      0x18, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
+      0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
   };
 
   EXPECT_EQ(Storage.size(), 68u);
@@ -184,7 +184,7 @@ TEST(RootSignature, ParseRootConstants) {
       0x85, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x24, 0x00, 0x00, 0x00,
       0x52, 0x54, 0x53, 0x30, 0x59, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00,
       0x01, 0x00, 0x00, 0x00, 0x18, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
-      0x2c, 0x00, 0x00, 0x00, 0x11, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
+      0x00, 0x00, 0x00, 0x00, 0x11, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
       0x02, 0x00, 0x00, 0x00, 0x24, 0x00, 0x00, 0x00, 0x0f, 0x00, 0x00, 0x00,
       0x0e, 0x00, 0x00, 0x00, 0x10, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
       0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,

llvmbot · 2025-02-24T21:06:30Z

@llvm/pr-subscribers-llvm-binary-utilities

Author: None (joaosaffran)

Changes

This PR is refactoring Root Signatures serialization process to remove the need to manually calculate offsets.

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

9 Files Affected:

(modified) llvm/include/llvm/MC/DXContainerRootSignature.h (+26-3)
(modified) llvm/lib/MC/DXContainerRootSignature.cpp (+96-29)
(modified) llvm/lib/Object/DXContainer.cpp (+6-4)
(modified) llvm/lib/ObjectYAML/DXContainerEmitter.cpp (+3-1)
(modified) llvm/lib/Target/DirectX/DXContainerGlobals.cpp (+4-1)
(modified) llvm/test/CodeGen/DirectX/ContainerData/RootSignature-Flags.ll (+1-1)
(modified) llvm/test/ObjectYAML/DXContainer/RootSignature-Flags.yaml (+1-1)
(modified) llvm/test/ObjectYAML/DXContainer/RootSignature-MultipleParameters.yaml (+1-1)
(modified) llvm/unittests/ObjectYAML/DXContainerYAMLTest.cpp (+4-4)

diff --git a/llvm/include/llvm/MC/DXContainerRootSignature.h b/llvm/include/llvm/MC/DXContainerRootSignature.h
index b31b0da352038..ca958c0780cdd 100644
--- a/llvm/include/llvm/MC/DXContainerRootSignature.h
+++ b/llvm/include/llvm/MC/DXContainerRootSignature.h
@@ -7,19 +7,42 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/BinaryFormat/DXContainer.h"
-#include <cstdint>
-#include <limits>
+#include "llvm/Support/BinaryStreamWriter.h"
+#include "llvm/Support/raw_ostream.h"
+#include <map>
+#include <string>
 
 namespace llvm {
 
 class raw_ostream;
 
 namespace mcdxbc {
+
+class StreamOffsetHelper {
+private:
+  std::map<std::string, std::pair<uint32_t, uint32_t>> OffsetsMaping;
+  BinaryStreamWriter &Stream;
+
+public:
+  explicit StreamOffsetHelper(BinaryStreamWriter &Stream) : Stream(Stream) {}
+
+  Error addOffset(std::string Key);
+
+  void addRewriteValue(std::string Key);
+
+  Error rewrite();
+};
+
 struct RootSignatureDesc {
   dxbc::RootSignatureHeader Header;
   SmallVector<dxbc::RootParameter> Parameters;
 
-  void write(raw_ostream &OS) const;
+  Error write(raw_ostream &OS) const;
+
+  uint32_t getSizeInBytes() const {
+    // Header Size + accounting for parameter offset + parameters size
+    return 24 + (Parameters.size() * 4) + Parameters.size_in_bytes();
+  }
 };
 } // namespace mcdxbc
 } // namespace llvm
diff --git a/llvm/lib/MC/DXContainerRootSignature.cpp b/llvm/lib/MC/DXContainerRootSignature.cpp
index 35a4ef322d01e..828785980c15d 100644
--- a/llvm/lib/MC/DXContainerRootSignature.cpp
+++ b/llvm/lib/MC/DXContainerRootSignature.cpp
@@ -7,53 +7,120 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/MC/DXContainerRootSignature.h"
-#include "llvm/Support/EndianStream.h"
+#include "llvm/ADT/Twine.h"
+#include "llvm/Support/BinaryStreamWriter.h"
 
 using namespace llvm;
 using namespace llvm::mcdxbc;
 
-void RootSignatureDesc::write(raw_ostream &OS) const {
-  // Root signature header in dxcontainer has 6 uint_32t values.
-  const uint32_t HeaderSize = 24;
-  const uint32_t ParameterByteSize = Parameters.size_in_bytes();
-  const uint32_t NumParametes = Parameters.size();
+Error StreamOffsetHelper::addOffset(std::string Key) {
+  const uint32_t DummyValue = std::numeric_limits<uint32_t>::max();
+
+  uint32_t Offset = Stream.getOffset();
+  auto Value = std::make_pair(Offset, DummyValue);
+
+  OffsetsMaping.insert_or_assign(Key, Value);
+
+  if (Error Err = Stream.writeInteger(DummyValue))
+    return Err;
+
+  return Error::success();
+}
+
+void StreamOffsetHelper::addRewriteValue(std::string Key) {
+  auto It = OffsetsMaping.find(Key);
+  assert(It != OffsetsMaping.end() && "Offset address was not found.");
+  auto [Offset, _] = It->second;
+
+  uint32_t Value = Stream.getOffset();
+
+  std::pair<uint32_t, uint32_t> NewValue = std::make_pair(Offset, Value);
+  OffsetsMaping.insert_or_assign(Key, NewValue);
+}
+
+Error StreamOffsetHelper::rewrite() {
+  for (auto &[Key, RewriteInfo] : OffsetsMaping) {
+    auto [Position, Value] = RewriteInfo;
+    assert(Value != std::numeric_limits<uint32_t>::max());
+
+    Stream.setOffset(Position);
+    if (Error Err = Stream.writeInteger(Value))
+      return Err;
+  }
+
+  return Error::success();
+}
+
+Error RootSignatureDesc::write(raw_ostream &OS) const {
+  std::vector<uint8_t> Buffer(getSizeInBytes());
+  BinaryStreamWriter Writer(Buffer, llvm::endianness::little);
+
+  StreamOffsetHelper OffsetMap(Writer);
+
+  const uint32_t NumParameters = Parameters.size();
   const uint32_t Zero = 0;
 
-  // Writing header information
-  support::endian::write(OS, Header.Version, llvm::endianness::little);
-  support::endian::write(OS, NumParametes, llvm::endianness::little);
-  support::endian::write(OS, HeaderSize, llvm::endianness::little);
+  if (Error Err = Writer.writeInteger(Header.Version))
+    return Err;
+
+  if (Error Err = Writer.writeInteger(NumParameters))
+    return Err;
+
+  if (Error Err = OffsetMap.addOffset("header"))
+    return Err;
 
   // Static samplers still not implemented
-  support::endian::write(OS, Zero, llvm::endianness::little);
-  support::endian::write(OS, ParameterByteSize + HeaderSize,
-                         llvm::endianness::little);
+  if (Error Err = Writer.writeInteger(Zero))
+    return Err;
+
+  if (Error Err = Writer.writeInteger(Zero))
+    return Err;
+
+  if (Error Err = Writer.writeInteger(Header.Flags))
+    return Err;
 
-  support::endian::write(OS, Header.Flags, llvm::endianness::little);
+  OffsetMap.addRewriteValue("header");
 
-  uint32_t ParamsOffset =
-      HeaderSize + (3 * sizeof(uint32_t) * Parameters.size());
-  for (const dxbc::RootParameter &P : Parameters) {
-    support::endian::write(OS, P.ParameterType, llvm::endianness::little);
-    support::endian::write(OS, P.ShaderVisibility, llvm::endianness::little);
-    support::endian::write(OS, ParamsOffset, llvm::endianness::little);
+  for (size_t It = 0; It < Parameters.size(); It++) {
+    const auto &P = Parameters[It];
 
-    // Size of root parameter, removing the ParameterType and ShaderVisibility.
-    ParamsOffset += sizeof(dxbc::RootParameter) - 2 * sizeof(uint32_t);
+    if (Error Err = Writer.writeEnum(P.ParameterType))
+      return Err;
+
+    if (Error Err = Writer.writeEnum(P.ShaderVisibility))
+      return Err;
+
+    std::string Key = ("parameters" + Twine(It)).str();
+    if (Error Err = OffsetMap.addOffset(Key))
+      return Err;
   }
 
-  for (const dxbc::RootParameter &P : Parameters) {
+  for (size_t It = 0; It < Parameters.size(); It++) {
+    const auto &P = Parameters[It];
+
+    std::string Key = ("parameters" + Twine(It)).str();
+    OffsetMap.addRewriteValue(Key);
+
     switch (P.ParameterType) {
     case dxbc::RootParameterType::Constants32Bit: {
-      support::endian::write(OS, P.Constants.ShaderRegister,
-                             llvm::endianness::little);
-      support::endian::write(OS, P.Constants.RegisterSpace,
-                             llvm::endianness::little);
-      support::endian::write(OS, P.Constants.Num32BitValues,
-                             llvm::endianness::little);
+      if (Error Err = Writer.writeInteger(P.Constants.ShaderRegister))
+        return Err;
+      if (Error Err = Writer.writeInteger(P.Constants.RegisterSpace))
+        return Err;
+      if (Error Err = Writer.writeInteger(P.Constants.Num32BitValues))
+        return Err;
     } break;
     case dxbc::RootParameterType::Empty:
       llvm_unreachable("Invalid RootParameterType");
     }
   }
+
+  if (Error Err = OffsetMap.rewrite())
+    return Err;
+
+  llvm::ArrayRef<char> BufferRef(reinterpret_cast<char *>(Buffer.data()),
+                                 Buffer.size());
+  OS.write(BufferRef.data(), BufferRef.size());
+
+  return Error::success();
 }
diff --git a/llvm/lib/Object/DXContainer.cpp b/llvm/lib/Object/DXContainer.cpp
index 010f70a952ebf..35261b661cf2f 100644
--- a/llvm/lib/Object/DXContainer.cpp
+++ b/llvm/lib/Object/DXContainer.cpp
@@ -296,16 +296,18 @@ Error DirectX::RootSignature::parse(StringRef Data) {
     NewParam.ParameterType =
         support::endian::read<dxbc::RootParameterType,
                               llvm::endianness::little>(Current);
-    if (!dxbc::RootSignatureValidations::isValidParameterType(NewParam.ParameterType))
+    if (!dxbc::RootSignatureValidations::isValidParameterType(
+            NewParam.ParameterType))
       return validationFailed("unsupported parameter type value read: " +
                               llvm::Twine((uint32_t)NewParam.ParameterType));
 
     Current += sizeof(dxbc::RootParameterType);
 
     NewParam.ShaderVisibility =
-        support::endian::read<dxbc::ShaderVisibility,
-                              llvm::endianness::little>(Current);
-    if (!dxbc::RootSignatureValidations::isValidShaderVisibility(NewParam.ShaderVisibility))
+        support::endian::read<dxbc::ShaderVisibility, llvm::endianness::little>(
+            Current);
+    if (!dxbc::RootSignatureValidations::isValidShaderVisibility(
+            NewParam.ShaderVisibility))
       return validationFailed("unsupported shader visility flag value read: " +
                               llvm::Twine((uint32_t)NewParam.ShaderVisibility));
 
diff --git a/llvm/lib/ObjectYAML/DXContainerEmitter.cpp b/llvm/lib/ObjectYAML/DXContainerEmitter.cpp
index 87ba16fd40ba9..a5831a69f9bca 100644
--- a/llvm/lib/ObjectYAML/DXContainerEmitter.cpp
+++ b/llvm/lib/ObjectYAML/DXContainerEmitter.cpp
@@ -271,7 +271,9 @@ void DXContainerWriter::writeParts(raw_ostream &OS) {
       RS.Header.Version = P.RootSignature->Version;
       RS.Parameters = std::move(P.RootSignature->Parameters);
 
-      RS.write(OS);
+      if (Error Err = RS.write(OS))
+        handleAllErrors(std::move(Err));
+
       break;
     }
     uint64_t BytesWritten = OS.tell() - DataStart;
diff --git a/llvm/lib/Target/DirectX/DXContainerGlobals.cpp b/llvm/lib/Target/DirectX/DXContainerGlobals.cpp
index 5508af40663b1..5801046f83674 100644
--- a/llvm/lib/Target/DirectX/DXContainerGlobals.cpp
+++ b/llvm/lib/Target/DirectX/DXContainerGlobals.cpp
@@ -25,10 +25,12 @@
 #include "llvm/InitializePasses.h"
 #include "llvm/MC/DXContainerPSVInfo.h"
 #include "llvm/Pass.h"
+#include "llvm/Support/Error.h"
 #include "llvm/Support/MD5.h"
 #include "llvm/TargetParser/Triple.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
 #include <optional>
+#include <utility>
 
 using namespace llvm;
 using namespace llvm::dxil;
@@ -173,7 +175,8 @@ void DXContainerGlobals::addRootSignature(Module &M,
   SmallString<256> Data;
   raw_svector_ostream OS(Data);
 
-  RS.write(OS);
+  if (Error Err = RS.write(OS))
+    handleAllErrors(std::move(Err));
 
   Constant *Constant =
       ConstantDataArray::getString(M.getContext(), Data, /*AddNull*/ false);
diff --git a/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-Flags.ll b/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-Flags.ll
index 035ea373d30ea..1ca6ebb7ddab8 100644
--- a/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-Flags.ll
+++ b/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-Flags.ll
@@ -23,6 +23,6 @@ attributes #0 = { "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" }
 ; DXC-NEXT:    RootSignature:
 ; DXC-NEXT:      Version:         2
 ; DXC-NEXT:      NumStaticSamplers: 0
-; DXC-NEXT:      StaticSamplersOffset: 24
+; DXC-NEXT:      StaticSamplersOffset: 0
 ; DXC-NEXT:      Parameters: []
 ; DXC-NEXT:      AllowInputAssemblerInputLayout: true
diff --git a/llvm/test/ObjectYAML/DXContainer/RootSignature-Flags.yaml b/llvm/test/ObjectYAML/DXContainer/RootSignature-Flags.yaml
index 1b73b830f015f..0d7902afdaa66 100644
--- a/llvm/test/ObjectYAML/DXContainer/RootSignature-Flags.yaml
+++ b/llvm/test/ObjectYAML/DXContainer/RootSignature-Flags.yaml
@@ -25,7 +25,7 @@ Parts:
 # CHECK-NEXT:    RootSignature:
 # CHECK-NEXT:      Version: 2
 # CHECK-NEXT:      NumStaticSamplers: 0
-# CHECK-NEXT:      StaticSamplersOffset: 24
+# CHECK-NEXT:      StaticSamplersOffset: 0
 # CHECK-NEXT:      Parameters: []
 # CHECK-NEXT:      AllowInputAssemblerInputLayout: true
 # CHECK-NEXT:      DenyGeometryShaderRootAccess: true
diff --git a/llvm/test/ObjectYAML/DXContainer/RootSignature-MultipleParameters.yaml b/llvm/test/ObjectYAML/DXContainer/RootSignature-MultipleParameters.yaml
index bccfacb72819b..8d5ab5c1b0b23 100644
--- a/llvm/test/ObjectYAML/DXContainer/RootSignature-MultipleParameters.yaml
+++ b/llvm/test/ObjectYAML/DXContainer/RootSignature-MultipleParameters.yaml
@@ -37,7 +37,7 @@ Parts:
 # CHECK-NEXT:    RootSignature:
 # CHECK-NEXT:      Version:         2
 # CHECK-NEXT:      NumStaticSamplers: 0
-# CHECK-NEXT:      StaticSamplersOffset: 64
+# CHECK-NEXT:      StaticSamplersOffset: 0
 # CHECK-NEXT:      Parameters:
 # CHECK-NEXT:        - ParameterType:   Constants32Bit
 # CHECK-NEXT:          ShaderVisibility: Hull
diff --git a/llvm/unittests/ObjectYAML/DXContainerYAMLTest.cpp b/llvm/unittests/ObjectYAML/DXContainerYAMLTest.cpp
index b5fb8d143c0b9..fed941f685272 100644
--- a/llvm/unittests/ObjectYAML/DXContainerYAMLTest.cpp
+++ b/llvm/unittests/ObjectYAML/DXContainerYAMLTest.cpp
@@ -134,12 +134,12 @@ TEST(RootSignature, ParseRootFlags) {
     )"));
 
   uint8_t Buffer[] = {
-      0x44, 0x58, 0x42, 0x43, 0x32, 0x9A, 0x53, 0xD8, 0xEC, 0xBE, 0x35, 0x6F,
-      0x05, 0x39, 0xE1, 0xFE, 0x31, 0x20, 0xF0, 0xC1, 0x01, 0x00, 0x00, 0x00,
+      0x44, 0x58, 0x42, 0x43, 0x32, 0x9a, 0x53, 0xd8, 0xec, 0xbe, 0x35, 0x6f,
+      0x05, 0x39, 0xe1, 0xfe, 0x31, 0x20, 0xf0, 0xc1, 0x01, 0x00, 0x00, 0x00,
       0x44, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x24, 0x00, 0x00, 0x00,
       0x52, 0x54, 0x53, 0x30, 0x18, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00,
       0x00, 0x00, 0x00, 0x00, 0x18, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
-      0x18, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
+      0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
   };
 
   EXPECT_EQ(Storage.size(), 68u);
@@ -184,7 +184,7 @@ TEST(RootSignature, ParseRootConstants) {
       0x85, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x24, 0x00, 0x00, 0x00,
       0x52, 0x54, 0x53, 0x30, 0x59, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00,
       0x01, 0x00, 0x00, 0x00, 0x18, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
-      0x2c, 0x00, 0x00, 0x00, 0x11, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
+      0x00, 0x00, 0x00, 0x00, 0x11, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
       0x02, 0x00, 0x00, 0x00, 0x24, 0x00, 0x00, 0x00, 0x0f, 0x00, 0x00, 0x00,
       0x0e, 0x00, 0x00, 0x00, 0x10, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
       0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,

bogner · 2025-02-24T21:10:21Z

llvm/include/llvm/MC/DXContainerRootSignature.h

+#include <map>
+#include <string>


We probably don't want to use std::map and std::string here - see the section on picking-the-right-data-structure-for-a-task from the LLVM programmer's manual for some good guidelines.

Refactored to remove the use of those.

llvm/lib/MC/DXContainerRootSignature.cpp

llvm/test/ObjectYAML/DXContainer/RootSignature-Flags.yaml

damyanp

A few nits, and I want to make sure that existing clients of this binary format can cope with a value of 0 for StaticSamplersOffset.

llvm/lib/MC/DXContainerRootSignature.cpp

llvm/unittests/ObjectYAML/DXContainerYAMLTest.cpp

joaosaffran added 2 commits February 22, 2025 06:19

first working version

767b7d0

moving the offset logic to it's own class

d391727

joaosaffran requested review from bogner, damyanp and llvm-beanz February 24, 2025 21:05

llvmbot added llvm:mc Machine (object) code backend:DirectX objectyaml llvm:binary-utilities labels Feb 24, 2025

bogner reviewed Feb 24, 2025

View reviewed changes

refactoring to remove use of map and string

68c7513

damyanp reviewed Feb 25, 2025

View reviewed changes

joaosaffran added 7 commits February 25, 2025 04:25

addressing comments

23069ab

using buffer_ostream

d14471b

remove getsize

216341c

clean up

85f012c

clean up

1e2bcf5

clean up

0e277d9

clean up

5cd0044

joaosaffran requested review from bogner and damyanp February 25, 2025 19:26

damyanp requested changes Feb 25, 2025

View reviewed changes

addressing pr comments

7a7c34d

joaosaffran requested a review from damyanp February 25, 2025 22:36

damyanp approved these changes Feb 25, 2025

View reviewed changes

clean up

d3fafab

bogner approved these changes Feb 25, 2025

View reviewed changes

joaosaffran merged commit cfdcbe7 into llvm:users/joaosaffran/127840 Feb 26, 2025
8 checks passed

joaosaffran deleted the refactor/improve-offset-calculation branch February 26, 2025 22:46

damyanp moved this to Closed in HLSL Support Apr 25, 2025

damyanp removed this from HLSL Support Jun 25, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[HLSL] Refactor Offset calculation while writing Root Signatures. #128577

[HLSL] Refactor Offset calculation while writing Root Signatures. #128577

Uh oh!

joaosaffran commented Feb 24, 2025 •

edited

Loading

Uh oh!

llvmbot commented Feb 24, 2025 •

edited

Loading

Uh oh!

llvmbot commented Feb 24, 2025

Uh oh!

bogner Feb 24, 2025

Uh oh!

joaosaffran Feb 24, 2025 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

damyanp left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

[HLSL] Refactor Offset calculation while writing Root Signatures. #128577

[HLSL] Refactor Offset calculation while writing Root Signatures. #128577

Uh oh!

Conversation

joaosaffran commented Feb 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Feb 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Feb 24, 2025

Uh oh!

bogner Feb 24, 2025

Choose a reason for hiding this comment

Uh oh!

joaosaffran Feb 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

damyanp left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

joaosaffran commented Feb 24, 2025 •

edited

Loading

llvmbot commented Feb 24, 2025 •

edited

Loading

joaosaffran Feb 24, 2025 •

edited

Loading