7
7
// ===----------------------------------------------------------------------===//
8
8
9
9
#include " llvm/MC/DXContainerRootSignature.h"
10
+ #include " llvm/ADT/Twine.h"
10
11
#include " llvm/Support/BinaryStreamWriter.h"
11
- #include < vector>
12
12
13
13
using namespace llvm ;
14
14
using namespace llvm ::mcdxbc;
15
15
16
+ Error StreamOffsetHelper::addOffset (std::string Key) {
17
+ const uint32_t DummyValue = std::numeric_limits<uint32_t >::max ();
18
+
19
+ uint32_t Offset = Stream.getOffset ();
20
+ auto Value = std::make_pair (Offset, DummyValue);
21
+
22
+ OffsetsMaping.insert_or_assign (Key, Value);
23
+
24
+ if (Error Err = Stream.writeInteger (DummyValue))
25
+ return Err;
26
+
27
+ return Error::success ();
28
+ }
29
+
30
+ void StreamOffsetHelper::addRewriteValue (std::string Key) {
31
+ auto It = OffsetsMaping.find (Key);
32
+ assert (It != OffsetsMaping.end () && " Offset address was not found." );
33
+ auto [Offset, _] = It->second ;
34
+
35
+ uint32_t Value = Stream.getOffset ();
36
+
37
+ std::pair<uint32_t , uint32_t > NewValue = std::make_pair (Offset, Value);
38
+ OffsetsMaping.insert_or_assign (Key, NewValue);
39
+ }
40
+
41
+ Error StreamOffsetHelper::rewrite () {
42
+ for (auto &[Key, RewriteInfo] : OffsetsMaping) {
43
+ auto [Position, Value] = RewriteInfo;
44
+ assert (Value != std::numeric_limits<uint32_t >::max ());
45
+
46
+ Stream.setOffset (Position);
47
+ if (Error Err = Stream.writeInteger (Value))
48
+ return Err;
49
+ }
50
+
51
+ return Error::success ();
52
+ }
53
+
16
54
Error RootSignatureDesc::write (raw_ostream &OS) const {
17
- // Header Size + accounting for parameter offset + parameters size
18
- std::vector<uint8_t > Buffer (24 + (Parameters.size () * 4 ) +
19
- Parameters.size_in_bytes ());
55
+ std::vector<uint8_t > Buffer (getSizeInBytes ());
20
56
BinaryStreamWriter Writer (Buffer, llvm::endianness::little);
21
57
22
- SmallVector<uint64_t > OffsetsToReplace;
23
- SmallVector<uint32_t > ValuesToReplaceOffsetsWith;
24
- const uint32_t Dummy = std::numeric_limits<uint32_t >::max ();
58
+ StreamOffsetHelper OffsetMap (Writer);
25
59
26
60
const uint32_t NumParameters = Parameters.size ();
27
61
const uint32_t Zero = 0 ;
@@ -32,8 +66,7 @@ Error RootSignatureDesc::write(raw_ostream &OS) const {
32
66
if (Error Err = Writer.writeInteger (NumParameters))
33
67
return Err;
34
68
35
- OffsetsToReplace.push_back (Writer.getOffset ());
36
- if (Error Err = Writer.writeInteger (Dummy))
69
+ if (Error Err = OffsetMap.addOffset (" header" ))
37
70
return Err;
38
71
39
72
// Static samplers still not implemented
@@ -46,21 +79,28 @@ Error RootSignatureDesc::write(raw_ostream &OS) const {
46
79
if (Error Err = Writer.writeInteger (Header.Flags ))
47
80
return Err;
48
81
49
- ValuesToReplaceOffsetsWith.push_back (Writer.getOffset ());
82
+ OffsetMap.addRewriteValue (" header" );
83
+
84
+ for (size_t It = 0 ; It < Parameters.size (); It++) {
85
+ const auto &P = Parameters[It];
50
86
51
- for (const dxbc::RootParameter &P : Parameters) {
52
87
if (Error Err = Writer.writeEnum (P.ParameterType ))
53
88
return Err;
89
+
54
90
if (Error Err = Writer.writeEnum (P.ShaderVisibility ))
55
91
return Err;
56
92
57
- OffsetsToReplace. push_back (Writer. getOffset () );
58
- if (Error Err = Writer. writeInteger (Dummy ))
93
+ std::string Key = ( " parameters " + Twine (It)). str ( );
94
+ if (Error Err = OffsetMap. addOffset (Key ))
59
95
return Err;
60
96
}
61
97
62
- for (const dxbc::RootParameter &P : Parameters) {
63
- ValuesToReplaceOffsetsWith.push_back (Writer.getOffset ());
98
+ for (size_t It = 0 ; It < Parameters.size (); It++) {
99
+ const auto &P = Parameters[It];
100
+
101
+ std::string Key = (" parameters" + Twine (It)).str ();
102
+ OffsetMap.addRewriteValue (Key);
103
+
64
104
switch (P.ParameterType ) {
65
105
case dxbc::RootParameterType::Constants32Bit: {
66
106
if (Error Err = Writer.writeInteger (P.Constants .ShaderRegister ))
@@ -75,17 +115,8 @@ Error RootSignatureDesc::write(raw_ostream &OS) const {
75
115
}
76
116
}
77
117
78
- assert (ValuesToReplaceOffsetsWith.size () == OffsetsToReplace.size () &&
79
- " Offset missing value to replace with." );
80
-
81
- for (size_t It = 0 ; It < ValuesToReplaceOffsetsWith.size (); It++) {
82
- uint32_t Position = OffsetsToReplace[It];
83
- uint32_t Value = ValuesToReplaceOffsetsWith[It];
84
-
85
- Writer.setOffset (Position);
86
- if (Error Err = Writer.writeInteger (Value))
87
- return Err;
88
- }
118
+ if (Error Err = OffsetMap.rewrite ())
119
+ return Err;
89
120
90
121
llvm::ArrayRef<char > BufferRef (reinterpret_cast <char *>(Buffer.data ()),
91
122
Buffer.size ());
0 commit comments