Skip to content

Commit c24411f

Browse files
IcohedronIcohedron
committed
Refactor DXIL CBuffer Access pass
Refactor to consolidate logic to be reused for implementing support for more cbuffer users.
1 parent f1988f4 commit c24411f

File tree

1 file changed

+128
-82
lines changed

1 file changed

+128
-82
lines changed

‎llvm/lib/Target/DirectX/DXILCBufferAccess.cpp‎

Lines changed: 128 additions & 82 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
#include "llvm/InitializePasses.h"
1818
#include "llvm/Pass.h"
1919
#include "llvm/Support/ErrorHandling.h"
20+
#include "llvm/Support/FormatVariadic.h"
2021
#include "llvm/Transforms/Utils/Local.h"
2122

2223
#define DEBUG_TYPE "dxil-cbuffer-access"
@@ -57,109 +58,153 @@ struct CBufferRowIntrin {
5758
}
5859
}
5960
};
60-
} // namespace
6161

62-
static size_t getOffsetForCBufferGEP(GEPOperator *GEP, GlobalVariable *Global,
63-
const DataLayout &DL) {
64-
// Since we should always have a constant offset, we should only ever have a
65-
// single GEP of indirection from the Global.
66-
assert(GEP->getPointerOperand() == Global &&
67-
"Indirect access to resource handle");
62+
// Helper for creating CBuffer handles and loading data from them
63+
struct CBufferResource {
64+
GlobalVariable *GVHandle;
65+
GlobalVariable *Member;
66+
size_t MemberOffset;
6867

69-
APInt ConstantOffset(DL.getIndexTypeSizeInBits(GEP->getType()), 0);
70-
bool Success = GEP->accumulateConstantOffset(DL, ConstantOffset);
71-
(void)Success;
72-
assert(Success && "Offsets into cbuffer globals must be constant");
68+
LoadInst *Handle;
7369

74-
if (auto *ATy = dyn_cast<ArrayType>(Global->getValueType()))
75-
ConstantOffset = hlsl::translateCBufArrayOffset(DL, ConstantOffset, ATy);
70+
CBufferResource(GlobalVariable *GVHandle, GlobalVariable *Member,
71+
size_t MemberOffset)
72+
: GVHandle(GVHandle), Member(Member), MemberOffset(MemberOffset) {}
7673

77-
return ConstantOffset.getZExtValue();
78-
}
74+
const DataLayout &getDataLayout() { return GVHandle->getDataLayout(); }
75+
Type *getValueType() { return Member->getValueType(); }
76+
iterator_range<ConstantDataSequential::user_iterator> users() {
77+
return Member->users();
78+
}
7979

80-
/// Replace access via cbuffer global with a load from the cbuffer handle
81-
/// itself.
82-
static void replaceAccess(LoadInst *LI, GlobalVariable *Global,
83-
GlobalVariable *HandleGV, size_t BaseOffset,
84-
SmallVectorImpl<WeakTrackingVH> &DeadInsts) {
85-
const DataLayout &DL = HandleGV->getDataLayout();
80+
/// Get the byte offset of a Pointer-typed Value * `Val` relative to Member.
81+
/// `Val` can either be Member itself, or a GEP of a constant offset from
82+
/// Member
83+
size_t getOffsetForCBufferGEP(Value *Val) {
84+
assert(isa<PointerType>(Val->getType()) &&
85+
"Expected a pointer-typed value");
86+
87+
if (Val == Member)
88+
return 0;
89+
90+
if (auto *GEP = dyn_cast<GEPOperator>(Val)) {
91+
// Since we should always have a constant offset, we should only ever have
92+
// a single GEP of indirection from the Global.
93+
assert(GEP->getPointerOperand() == Member &&
94+
"Indirect access to resource handle");
95+
96+
const DataLayout &DL = getDataLayout();
97+
APInt ConstantOffset(DL.getIndexTypeSizeInBits(GEP->getType()), 0);
98+
bool Success = GEP->accumulateConstantOffset(DL, ConstantOffset);
99+
(void)Success;
100+
assert(Success && "Offsets into cbuffer globals must be constant");
101+
102+
if (auto *ATy = dyn_cast<ArrayType>(Member->getValueType()))
103+
ConstantOffset =
104+
hlsl::translateCBufArrayOffset(DL, ConstantOffset, ATy);
105+
106+
return ConstantOffset.getZExtValue();
107+
}
86108

87-
size_t Offset = BaseOffset;
88-
if (auto *GEP = dyn_cast<GEPOperator>(LI->getPointerOperand()))
89-
Offset += getOffsetForCBufferGEP(GEP, Global, DL);
90-
else if (LI->getPointerOperand() != Global)
91-
llvm_unreachable("Load instruction doesn't reference cbuffer global");
109+
llvm_unreachable("Expected Val to be a GlobalVariable or GEP");
110+
}
92111

93-
IRBuilder<> Builder(LI);
94-
auto *Handle = Builder.CreateLoad(HandleGV->getValueType(), HandleGV,
95-
HandleGV->getName());
96-
97-
Type *Ty = LI->getType();
98-
CBufferRowIntrin Intrin(DL, Ty->getScalarType());
99-
// The cbuffer consists of some number of 16-byte rows.
100-
unsigned int CurrentRow = Offset / hlsl::CBufferRowSizeInBytes;
101-
unsigned int CurrentIndex =
102-
(Offset % hlsl::CBufferRowSizeInBytes) / Intrin.EltSize;
103-
104-
auto *CBufLoad = Builder.CreateIntrinsic(
105-
Intrin.RetTy, Intrin.IID,
106-
{Handle, ConstantInt::get(Builder.getInt32Ty(), CurrentRow)}, nullptr,
107-
LI->getName());
108-
auto *Elt =
109-
Builder.CreateExtractValue(CBufLoad, {CurrentIndex++}, LI->getName());
110-
111-
Value *Result = nullptr;
112-
unsigned int Remaining =
113-
((DL.getTypeSizeInBits(Ty) / 8) / Intrin.EltSize) - 1;
114-
if (Remaining == 0) {
115-
// We only have a single element, so we're done.
116-
Result = Elt;
117-
118-
// However, if we loaded a <1 x T>, then we need to adjust the type here.
119-
if (auto *VT = dyn_cast<FixedVectorType>(LI->getType())) {
120-
assert(VT->getNumElements() == 1 && "Can't have multiple elements here");
121-
Result = Builder.CreateInsertElement(PoisonValue::get(VT), Result,
122-
Builder.getInt32(0));
123-
}
124-
} else {
125-
// Walk each element and extract it, wrapping to new rows as needed.
126-
SmallVector<Value *> Extracts{Elt};
127-
while (Remaining--) {
128-
CurrentIndex %= Intrin.NumElts;
129-
130-
if (CurrentIndex == 0)
131-
CBufLoad = Builder.CreateIntrinsic(
132-
Intrin.RetTy, Intrin.IID,
133-
{Handle, ConstantInt::get(Builder.getInt32Ty(), ++CurrentRow)},
134-
nullptr, LI->getName());
135-
136-
Extracts.push_back(Builder.CreateExtractValue(CBufLoad, {CurrentIndex++},
137-
LI->getName()));
112+
/// Create a handle for this cbuffer resource using the IRBuilder `Builder`
113+
/// and sets the handle as the current one to use for subsequent calls to
114+
/// `loadValue`
115+
void createAndSetCurrentHandle(IRBuilder<> &Builder) {
116+
Handle = Builder.CreateLoad(GVHandle->getValueType(), GVHandle,
117+
GVHandle->getName());
118+
}
119+
120+
/// Load a value of type `Ty` at offset `Offset` using the handle from the
121+
/// last call to `createAndSetCurrentHandle`
122+
Value *loadValue(IRBuilder<> &Builder, Type *Ty, size_t Offset,
123+
const Twine &Name = "") {
124+
assert(Handle &&
125+
"Expected a handle for this cbuffer global resource to be created "
126+
"before loading a value from it");
127+
const DataLayout &DL = getDataLayout();
128+
129+
size_t TargetOffset = MemberOffset + Offset;
130+
CBufferRowIntrin Intrin(DL, Ty->getScalarType());
131+
// The cbuffer consists of some number of 16-byte rows.
132+
unsigned int CurrentRow = TargetOffset / hlsl::CBufferRowSizeInBytes;
133+
unsigned int CurrentIndex =
134+
(TargetOffset % hlsl::CBufferRowSizeInBytes) / Intrin.EltSize;
135+
136+
auto *CBufLoad = Builder.CreateIntrinsic(
137+
Intrin.RetTy, Intrin.IID,
138+
{Handle, ConstantInt::get(Builder.getInt32Ty(), CurrentRow)}, nullptr,
139+
Name + ".load");
140+
auto *Elt = Builder.CreateExtractValue(CBufLoad, {CurrentIndex++},
141+
Name + ".extract");
142+
143+
Value *Result = nullptr;
144+
unsigned int Remaining =
145+
((DL.getTypeSizeInBits(Ty) / 8) / Intrin.EltSize) - 1;
146+
if (Remaining == 0) {
147+
// We only have a single element, so we're done.
148+
Result = Elt;
149+
150+
// However, if we loaded a <1 x T>, then we need to adjust the type here.
151+
if (auto *VT = dyn_cast<FixedVectorType>(Ty)) {
152+
assert(VT->getNumElements() == 1 &&
153+
"Can't have multiple elements here");
154+
Result = Builder.CreateInsertElement(PoisonValue::get(VT), Result,
155+
Builder.getInt32(0), Name);
156+
}
157+
} else {
158+
// Walk each element and extract it, wrapping to new rows as needed.
159+
SmallVector<Value *> Extracts{Elt};
160+
while (Remaining--) {
161+
CurrentIndex %= Intrin.NumElts;
162+
163+
if (CurrentIndex == 0)
164+
CBufLoad = Builder.CreateIntrinsic(
165+
Intrin.RetTy, Intrin.IID,
166+
{Handle, ConstantInt::get(Builder.getInt32Ty(), ++CurrentRow)},
167+
nullptr, Name + ".load");
168+
169+
Extracts.push_back(Builder.CreateExtractValue(
170+
CBufLoad, {CurrentIndex++}, Name + ".extract"));
171+
}
172+
173+
// Finally, we build up the original loaded value.
174+
Result = PoisonValue::get(Ty);
175+
for (int I = 0, E = Extracts.size(); I < E; ++I)
176+
Result = Builder.CreateInsertElement(Result, Extracts[I],
177+
Builder.getInt32(I),
178+
Name + formatv(".upto{}", I));
138179
}
139180

140-
// Finally, we build up the original loaded value.
141-
Result = PoisonValue::get(Ty);
142-
for (int I = 0, E = Extracts.size(); I < E; ++I)
143-
Result =
144-
Builder.CreateInsertElement(Result, Extracts[I], Builder.getInt32(I));
181+
return Result;
145182
}
183+
};
146184

185+
} // namespace
186+
187+
/// Replace load via cbuffer global with a load from the cbuffer handle itself.
188+
static void replaceLoad(LoadInst *LI, CBufferResource &CBR,
189+
SmallVectorImpl<WeakTrackingVH> &DeadInsts) {
190+
size_t Offset = CBR.getOffsetForCBufferGEP(LI->getPointerOperand());
191+
IRBuilder<> Builder(LI);
192+
CBR.createAndSetCurrentHandle(Builder);
193+
Value *Result = CBR.loadValue(Builder, LI->getType(), Offset, LI->getName());
147194
LI->replaceAllUsesWith(Result);
148195
DeadInsts.push_back(LI);
149196
}
150197

151-
static void replaceAccessesWithHandle(GlobalVariable *Global,
152-
GlobalVariable *HandleGV,
153-
size_t BaseOffset) {
198+
static void replaceAccessesWithHandle(CBufferResource &CBR) {
154199
SmallVector<WeakTrackingVH> DeadInsts;
155200

156-
SmallVector<User *> ToProcess{Global->users()};
201+
SmallVector<User *> ToProcess{CBR.users()};
157202
while (!ToProcess.empty()) {
158203
User *Cur = ToProcess.pop_back_val();
159204

160205
// If we have a load instruction, replace the access.
161206
if (auto *LI = dyn_cast<LoadInst>(Cur)) {
162-
replaceAccess(LI, Global, HandleGV, BaseOffset, DeadInsts);
207+
replaceLoad(LI, CBR, DeadInsts);
163208
continue;
164209
}
165210

@@ -181,7 +226,8 @@ static bool replaceCBufferAccesses(Module &M) {
181226

182227
for (const hlsl::CBufferMapping &Mapping : *CBufMD)
183228
for (const hlsl::CBufferMember &Member : Mapping.Members) {
184-
replaceAccessesWithHandle(Member.GV, Mapping.Handle, Member.Offset);
229+
CBufferResource CBR(Mapping.Handle, Member.GV, Member.Offset);
230+
replaceAccessesWithHandle(CBR);
185231
Member.GV->removeFromParent();
186232
}
187233

0 commit comments

Comments
 (0)