[HLSL] Refactor Offset calculation while writing Root Signatures.

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

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();
 }

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));
 

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;

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);

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

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

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

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,