[llvm-branch-commits] [llvm] [HLSL] Refactor Offset calculation while writing Root Signatures. (PR #128577)
via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Mon Feb 24 13:06:39 PST 2025
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-llvm-binary-utilities
Author: None (joaosaffran)
<details>
<summary>Changes</summary>
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
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,
``````````
</details>
https://github.com/llvm/llvm-project/pull/128577
More information about the llvm-branch-commits
mailing list