[llvm-branch-commits] [llvm] [HLSL] Refactor Offset calculation while writing Root Signatures. (PR #128577)

[via llvm-branch-commits]

================
@@ -7,53 +7,109 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/MC/DXContainerRootSignature.h"
-#include "llvm/Support/EndianStream.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/Twine.h"
+#include "llvm/Support/BinaryStreamWriter.h"
+#include <cstdint>
+#include <sys/types.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 setRewrite(BinaryStreamWriter &Stream, uint32_t &Offset) {
+  const uint32_t DummyValue = std::numeric_limits<uint32_t>::max();
+
+  Offset = Stream.getOffset();
+
+  if (Error Err = Stream.writeInteger(DummyValue))
+    return Err;
+
+  return Error::success();
+}
+
+Error rewriteOffset(BinaryStreamWriter &Stream, uint32_t Offset) {
+  uint64_t Value = Stream.getOffset();
+  Stream.setOffset(Offset);
+  if (Error Err = Stream.writeInteger((uint32_t)Value))
+    return Err;
+
+  Stream.setOffset(Value);
+
+  return Error::success();
+}
+
+Error RootSignatureDesc::write(raw_ostream &OS) const {
+  std::vector<uint8_t> Buffer(getSizeInBytes());
+  BinaryStreamWriter Writer(Buffer, llvm::endianness::little);
+
+  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;
+
+  uint32_t HeaderPoint;
+  if (Error Err = setRewrite(Writer, HeaderPoint))
+    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;
+
+  if (Error Err = rewriteOffset(Writer, HeaderPoint))
+    return Err;
+
+  SmallVector<uint32_t> ParamsOffset;
+  for (const auto &P : Parameters) {
 
-  support::endian::write(OS, Header.Flags, llvm::endianness::little);
+    if (Error Err = Writer.writeEnum(P.ParameterType))
+      return Err;
 
-  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);
+    if (Error Err = Writer.writeEnum(P.ShaderVisibility))
+      return Err;
 
-    // Size of root parameter, removing the ParameterType and ShaderVisibility.
-    ParamsOffset += sizeof(dxbc::RootParameter) - 2 * sizeof(uint32_t);
+    uint32_t Offset;
+    if (Error Err = setRewrite(Writer, Offset))
+      return Err;
+    ParamsOffset.push_back(Offset);
   }
 
-  for (const dxbc::RootParameter &P : Parameters) {
+  size_t It = 0;
+  for (const auto &P : Parameters) {
+
+    auto Offset = ParamsOffset[It];
+    if (Error Err = rewriteOffset(Writer, Offset))
+      return Err;
+    It++;
+
     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");
     }
   }
+
+  assert(It == NumParameters);
+
+  llvm::ArrayRef<char> BufferRef(reinterpret_cast<char *>(Buffer.data()),
+                                 Buffer.size());
+  OS.write(BufferRef.data(), BufferRef.size());
----------------
joaosaffran wrote:

This is not currently possible to be removed, `OS` expects `char *` and `BinaryStreamWriter` expects `uint_8t`

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