Skip to content

Commit d391727

Browse files
author
joaosaffran
committed
moving the offset logic to it's own class
1 parent 767b7d0 commit d391727

File tree

2 files changed

+81
-26
lines changed

2 files changed

+81
-26
lines changed

llvm/include/llvm/MC/DXContainerRootSignature.h

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,18 +7,42 @@
77
//===----------------------------------------------------------------------===//
88

99
#include "llvm/BinaryFormat/DXContainer.h"
10+
#include "llvm/Support/BinaryStreamWriter.h"
1011
#include "llvm/Support/raw_ostream.h"
12+
#include <map>
13+
#include <string>
1114

1215
namespace llvm {
1316

1417
class raw_ostream;
1518

1619
namespace mcdxbc {
20+
21+
class StreamOffsetHelper {
22+
private:
23+
std::map<std::string, std::pair<uint32_t, uint32_t>> OffsetsMaping;
24+
BinaryStreamWriter &Stream;
25+
26+
public:
27+
explicit StreamOffsetHelper(BinaryStreamWriter &Stream) : Stream(Stream) {}
28+
29+
Error addOffset(std::string Key);
30+
31+
void addRewriteValue(std::string Key);
32+
33+
Error rewrite();
34+
};
35+
1736
struct RootSignatureDesc {
1837
dxbc::RootSignatureHeader Header;
1938
SmallVector<dxbc::RootParameter> Parameters;
2039

2140
Error write(raw_ostream &OS) const;
41+
42+
uint32_t getSizeInBytes() const {
43+
// Header Size + accounting for parameter offset + parameters size
44+
return 24 + (Parameters.size() * 4) + Parameters.size_in_bytes();
45+
}
2246
};
2347
} // namespace mcdxbc
2448
} // namespace llvm

llvm/lib/MC/DXContainerRootSignature.cpp

Lines changed: 57 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -7,21 +7,55 @@
77
//===----------------------------------------------------------------------===//
88

99
#include "llvm/MC/DXContainerRootSignature.h"
10+
#include "llvm/ADT/Twine.h"
1011
#include "llvm/Support/BinaryStreamWriter.h"
11-
#include <vector>
1212

1313
using namespace llvm;
1414
using namespace llvm::mcdxbc;
1515

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+
1654
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());
2056
BinaryStreamWriter Writer(Buffer, llvm::endianness::little);
2157

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

2660
const uint32_t NumParameters = Parameters.size();
2761
const uint32_t Zero = 0;
@@ -32,8 +66,7 @@ Error RootSignatureDesc::write(raw_ostream &OS) const {
3266
if (Error Err = Writer.writeInteger(NumParameters))
3367
return Err;
3468

35-
OffsetsToReplace.push_back(Writer.getOffset());
36-
if (Error Err = Writer.writeInteger(Dummy))
69+
if (Error Err = OffsetMap.addOffset("header"))
3770
return Err;
3871

3972
// Static samplers still not implemented
@@ -46,21 +79,28 @@ Error RootSignatureDesc::write(raw_ostream &OS) const {
4679
if (Error Err = Writer.writeInteger(Header.Flags))
4780
return Err;
4881

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];
5086

51-
for (const dxbc::RootParameter &P : Parameters) {
5287
if (Error Err = Writer.writeEnum(P.ParameterType))
5388
return Err;
89+
5490
if (Error Err = Writer.writeEnum(P.ShaderVisibility))
5591
return Err;
5692

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))
5995
return Err;
6096
}
6197

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+
64104
switch (P.ParameterType) {
65105
case dxbc::RootParameterType::Constants32Bit: {
66106
if (Error Err = Writer.writeInteger(P.Constants.ShaderRegister))
@@ -75,17 +115,8 @@ Error RootSignatureDesc::write(raw_ostream &OS) const {
75115
}
76116
}
77117

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

90121
llvm::ArrayRef<char> BufferRef(reinterpret_cast<char *>(Buffer.data()),
91122
Buffer.size());

0 commit comments

Comments
 (0)