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

[Damyan Pepper via llvm-branch-commits]

================
@@ -12,44 +12,59 @@
 using namespace llvm;
 using namespace llvm::mcdxbc;
 
+void setRewrite(raw_ostream &Stream, uint32_t &Offset) {
+  const uint32_t DummyValue = std::numeric_limits<uint32_t>::max();
+  Offset = Stream.tell();
+  support::endian::write(Stream, DummyValue, llvm::endianness::little);
+}
+
+void rewriteOffset(buffer_ostream &Stream, uint32_t Offset) {
+  uint32_t Value = Stream.tell();
+  auto *InsertPoint = &Stream.buffer()[Offset];
+  support::endian::write(InsertPoint, Value, llvm::endianness::little);
+}
+
 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();
+  buffer_ostream Writer(OS);
+  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);
+  support::endian::write(Writer, Header.Version, llvm::endianness::little);
+  support::endian::write(Writer, NumParameters, llvm::endianness::little);
 
-  // Static samplers still not implemented
-  support::endian::write(OS, Zero, llvm::endianness::little);
-  support::endian::write(OS, ParameterByteSize + HeaderSize,
-                         llvm::endianness::little);
+  uint32_t HeaderPoint;
+  setRewrite(Writer, HeaderPoint);
 
-  support::endian::write(OS, Header.Flags, llvm::endianness::little);
+  support::endian::write(Writer, Zero, llvm::endianness::little);
+  support::endian::write(Writer, Zero, llvm::endianness::little);
+  support::endian::write(Writer, Header.Flags, llvm::endianness::little);
 
-  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);
+  rewriteOffset(Writer, HeaderPoint);
 
-    // Size of root parameter, removing the ParameterType and ShaderVisibility.
-    ParamsOffset += sizeof(dxbc::RootParameter) - 2 * sizeof(uint32_t);
+  SmallVector<uint32_t> ParamsOffset;
+  for (const auto &P : Parameters) {
+    support::endian::write(Writer, P.ParameterType, llvm::endianness::little);
+    support::endian::write(Writer, P.ShaderVisibility,
+                           llvm::endianness::little);
+
+    uint32_t Offset;
+    setRewrite(Writer, Offset);
+
+    ParamsOffset.push_back(Offset);
----------------
damyanp wrote:

Take this or leave it, but I think this can be made a little nicer:

```suggestion
    ParamsOffset.push_back(writePlaceholder(Writer));
```


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