Skip to content

Commit 2f39d13

Browse files
bognerbogner
authored
[DirectX] Handle dx.RawBuffer in DXILResourceAccess (#121725)
This adds handling for raw and structured buffers when lowering resource access via `llvm.dx.resource.getpointer`. Fixes #121714
1 parent 6db73fa commit 2f39d13

File tree

3 files changed

+469
-77
lines changed

3 files changed

+469
-77
lines changed

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

Lines changed: 178 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -21,99 +21,207 @@
2121

2222
using namespace llvm;
2323

24-
static void replaceTypedBufferAccess(IntrinsicInst *II,
25-
dxil::ResourceTypeInfo &RTI) {
26-
const DataLayout &DL = II->getDataLayout();
24+
static Value *calculateGEPOffset(GetElementPtrInst *GEP, Value *PrevOffset,
25+
dxil::ResourceTypeInfo &RTI) {
26+
assert(!PrevOffset && "Non-constant GEP chains not handled yet");
27+
28+
const DataLayout &DL = GEP->getDataLayout();
29+
30+
uint64_t ScalarSize = 1;
31+
if (RTI.isTyped()) {
32+
Type *ContainedType = RTI.getHandleTy()->getTypeParameter(0);
33+
// We need the size of an element in bytes so that we can calculate the
34+
// offset in elements given a total offset in bytes.
35+
Type *ScalarType = ContainedType->getScalarType();
36+
ScalarSize = DL.getTypeSizeInBits(ScalarType) / 8;
37+
}
38+
39+
APInt ConstantOffset(DL.getIndexTypeSizeInBits(GEP->getType()), 0);
40+
if (GEP->accumulateConstantOffset(DL, ConstantOffset)) {
41+
APInt Scaled = ConstantOffset.udiv(ScalarSize);
42+
return ConstantInt::get(Type::getInt32Ty(GEP->getContext()), Scaled);
43+
}
44+
45+
auto IndexIt = GEP->idx_begin();
46+
assert(cast<ConstantInt>(IndexIt)->getZExtValue() == 0 &&
47+
"GEP is not indexing through pointer");
48+
++IndexIt;
49+
Value *Offset = *IndexIt;
50+
assert(++IndexIt == GEP->idx_end() && "Too many indices in GEP");
51+
return Offset;
52+
}
53+
54+
static void createTypedBufferStore(IntrinsicInst *II, StoreInst *SI,
55+
Value *Offset, dxil::ResourceTypeInfo &RTI) {
56+
IRBuilder<> Builder(SI);
57+
Type *ContainedType = RTI.getHandleTy()->getTypeParameter(0);
58+
Type *LoadType = StructType::get(ContainedType, Builder.getInt1Ty());
59+
60+
Value *V = SI->getValueOperand();
61+
if (V->getType() == ContainedType) {
62+
// V is already the right type.
63+
assert(!Offset && "store of whole element has offset?");
64+
} else if (V->getType() == ContainedType->getScalarType()) {
65+
// We're storing a scalar, so we need to load the current value and only
66+
// replace the relevant part.
67+
auto *Load = Builder.CreateIntrinsic(
68+
LoadType, Intrinsic::dx_resource_load_typedbuffer,
69+
{II->getOperand(0), II->getOperand(1)});
70+
auto *Struct = Builder.CreateExtractValue(Load, {0});
71+
72+
// If we have an offset from seeing a GEP earlier, use that. Otherwise, 0.
73+
if (!Offset)
74+
Offset = ConstantInt::get(Builder.getInt32Ty(), 0);
75+
V = Builder.CreateInsertElement(Struct, V, Offset);
76+
} else {
77+
llvm_unreachable("Store to typed resource has invalid type");
78+
}
79+
80+
auto *Inst = Builder.CreateIntrinsic(
81+
Builder.getVoidTy(), Intrinsic::dx_resource_store_typedbuffer,
82+
{II->getOperand(0), II->getOperand(1), V});
83+
SI->replaceAllUsesWith(Inst);
84+
}
85+
86+
static void createRawStore(IntrinsicInst *II, StoreInst *SI, Value *Offset) {
87+
IRBuilder<> Builder(SI);
88+
89+
if (!Offset)
90+
Offset = ConstantInt::get(Builder.getInt32Ty(), 0);
91+
Value *V = SI->getValueOperand();
92+
// TODO: break up larger types
93+
auto *Inst = Builder.CreateIntrinsic(
94+
Builder.getVoidTy(), Intrinsic::dx_resource_store_rawbuffer,
95+
{II->getOperand(0), II->getOperand(1), Offset, V});
96+
SI->replaceAllUsesWith(Inst);
97+
}
98+
99+
static void createStoreIntrinsic(IntrinsicInst *II, StoreInst *SI,
100+
Value *Offset, dxil::ResourceTypeInfo &RTI) {
101+
switch (RTI.getResourceKind()) {
102+
case dxil::ResourceKind::TypedBuffer:
103+
return createTypedBufferStore(II, SI, Offset, RTI);
104+
case dxil::ResourceKind::RawBuffer:
105+
case dxil::ResourceKind::StructuredBuffer:
106+
return createRawStore(II, SI, Offset);
107+
case dxil::ResourceKind::Texture1D:
108+
case dxil::ResourceKind::Texture2D:
109+
case dxil::ResourceKind::Texture2DMS:
110+
case dxil::ResourceKind::Texture3D:
111+
case dxil::ResourceKind::TextureCube:
112+
case dxil::ResourceKind::Texture1DArray:
113+
case dxil::ResourceKind::Texture2DArray:
114+
case dxil::ResourceKind::Texture2DMSArray:
115+
case dxil::ResourceKind::TextureCubeArray:
116+
case dxil::ResourceKind::FeedbackTexture2D:
117+
case dxil::ResourceKind::FeedbackTexture2DArray:
118+
report_fatal_error("DXIL Load not implemented yet",
119+
/*gen_crash_diag=*/false);
120+
return;
121+
case dxil::ResourceKind::CBuffer:
122+
case dxil::ResourceKind::Sampler:
123+
case dxil::ResourceKind::TBuffer:
124+
case dxil::ResourceKind::RTAccelerationStructure:
125+
case dxil::ResourceKind::Invalid:
126+
case dxil::ResourceKind::NumEntries:
127+
llvm_unreachable("Invalid resource kind for store");
128+
}
129+
llvm_unreachable("Unhandled case in switch");
130+
}
131+
132+
static void createTypedBufferLoad(IntrinsicInst *II, LoadInst *LI,
133+
Value *Offset, dxil::ResourceTypeInfo &RTI) {
134+
IRBuilder<> Builder(LI);
135+
Type *ContainedType = RTI.getHandleTy()->getTypeParameter(0);
136+
Type *LoadType = StructType::get(ContainedType, Builder.getInt1Ty());
27137

28-
auto *HandleType = cast<TargetExtType>(II->getOperand(0)->getType());
29-
assert(HandleType->getName() == "dx.TypedBuffer" &&
30-
"Unexpected typed buffer type");
31-
Type *ContainedType = HandleType->getTypeParameter(0);
138+
Value *V =
139+
Builder.CreateIntrinsic(LoadType, Intrinsic::dx_resource_load_typedbuffer,
140+
{II->getOperand(0), II->getOperand(1)});
141+
V = Builder.CreateExtractValue(V, {0});
32142

33-
Type *LoadType =
34-
StructType::get(ContainedType, Type::getInt1Ty(II->getContext()));
143+
if (Offset)
144+
V = Builder.CreateExtractElement(V, Offset);
35145

36-
// We need the size of an element in bytes so that we can calculate the offset
37-
// in elements given a total offset in bytes later.
38-
Type *ScalarType = ContainedType->getScalarType();
39-
uint64_t ScalarSize = DL.getTypeSizeInBits(ScalarType) / 8;
146+
LI->replaceAllUsesWith(V);
147+
}
40148

149+
static void createRawLoad(IntrinsicInst *II, LoadInst *LI, Value *Offset) {
150+
IRBuilder<> Builder(LI);
151+
// TODO: break up larger types
152+
Type *LoadType = StructType::get(LI->getType(), Builder.getInt1Ty());
153+
if (!Offset)
154+
Offset = ConstantInt::get(Builder.getInt32Ty(), 0);
155+
Value *V =
156+
Builder.CreateIntrinsic(LoadType, Intrinsic::dx_resource_load_rawbuffer,
157+
{II->getOperand(0), II->getOperand(1), Offset});
158+
V = Builder.CreateExtractValue(V, {0});
159+
160+
LI->replaceAllUsesWith(V);
161+
}
162+
163+
static void createLoadIntrinsic(IntrinsicInst *II, LoadInst *LI, Value *Offset,
164+
dxil::ResourceTypeInfo &RTI) {
165+
switch (RTI.getResourceKind()) {
166+
case dxil::ResourceKind::TypedBuffer:
167+
return createTypedBufferLoad(II, LI, Offset, RTI);
168+
case dxil::ResourceKind::RawBuffer:
169+
case dxil::ResourceKind::StructuredBuffer:
170+
return createRawLoad(II, LI, Offset);
171+
case dxil::ResourceKind::Texture1D:
172+
case dxil::ResourceKind::Texture2D:
173+
case dxil::ResourceKind::Texture2DMS:
174+
case dxil::ResourceKind::Texture3D:
175+
case dxil::ResourceKind::TextureCube:
176+
case dxil::ResourceKind::Texture1DArray:
177+
case dxil::ResourceKind::Texture2DArray:
178+
case dxil::ResourceKind::Texture2DMSArray:
179+
case dxil::ResourceKind::TextureCubeArray:
180+
case dxil::ResourceKind::FeedbackTexture2D:
181+
case dxil::ResourceKind::FeedbackTexture2DArray:
182+
case dxil::ResourceKind::CBuffer:
183+
case dxil::ResourceKind::TBuffer:
184+
// TODO: handle these
185+
return;
186+
case dxil::ResourceKind::Sampler:
187+
case dxil::ResourceKind::RTAccelerationStructure:
188+
case dxil::ResourceKind::Invalid:
189+
case dxil::ResourceKind::NumEntries:
190+
llvm_unreachable("Invalid resource kind for load");
191+
}
192+
llvm_unreachable("Unhandled case in switch");
193+
}
194+
195+
static void replaceAccess(IntrinsicInst *II, dxil::ResourceTypeInfo &RTI) {
41196
// Process users keeping track of indexing accumulated from GEPs.
42-
struct AccessAndIndex {
197+
struct AccessAndOffset {
43198
User *Access;
44-
Value *Index;
199+
Value *Offset;
45200
};
46-
SmallVector<AccessAndIndex> Worklist;
201+
SmallVector<AccessAndOffset> Worklist;
47202
for (User *U : II->users())
48203
Worklist.push_back({U, nullptr});
49204

50205
SmallVector<Instruction *> DeadInsts;
51206
while (!Worklist.empty()) {
52-
AccessAndIndex Current = Worklist.back();
207+
AccessAndOffset Current = Worklist.back();
53208
Worklist.pop_back();
54209

55210
if (auto *GEP = dyn_cast<GetElementPtrInst>(Current.Access)) {
56211
IRBuilder<> Builder(GEP);
57212

58-
Value *Index;
59-
APInt ConstantOffset(DL.getIndexTypeSizeInBits(GEP->getType()), 0);
60-
if (GEP->accumulateConstantOffset(DL, ConstantOffset)) {
61-
APInt Scaled = ConstantOffset.udiv(ScalarSize);
62-
Index = ConstantInt::get(Builder.getInt32Ty(), Scaled);
63-
} else {
64-
auto IndexIt = GEP->idx_begin();
65-
assert(cast<ConstantInt>(IndexIt)->getZExtValue() == 0 &&
66-
"GEP is not indexing through pointer");
67-
++IndexIt;
68-
Index = *IndexIt;
69-
assert(++IndexIt == GEP->idx_end() && "Too many indices in GEP");
70-
}
71-
213+
Value *Offset = calculateGEPOffset(GEP, Current.Offset, RTI);
72214
for (User *U : GEP->users())
73-
Worklist.push_back({U, Index});
215+
Worklist.push_back({U, Offset});
74216
DeadInsts.push_back(GEP);
75217

76218
} else if (auto *SI = dyn_cast<StoreInst>(Current.Access)) {
77219
assert(SI->getValueOperand() != II && "Pointer escaped!");
78-
IRBuilder<> Builder(SI);
79-
80-
Value *V = SI->getValueOperand();
81-
if (V->getType() == ContainedType) {
82-
// V is already the right type.
83-
} else if (V->getType() == ScalarType) {
84-
// We're storing a scalar, so we need to load the current value and only
85-
// replace the relevant part.
86-
auto *Load = Builder.CreateIntrinsic(
87-
LoadType, Intrinsic::dx_resource_load_typedbuffer,
88-
{II->getOperand(0), II->getOperand(1)});
89-
auto *Struct = Builder.CreateExtractValue(Load, {0});
90-
91-
// If we have an offset from seeing a GEP earlier, use it.
92-
Value *IndexOp = Current.Index
93-
? Current.Index
94-
: ConstantInt::get(Builder.getInt32Ty(), 0);
95-
V = Builder.CreateInsertElement(Struct, V, IndexOp);
96-
} else {
97-
llvm_unreachable("Store to typed resource has invalid type");
98-
}
99-
100-
auto *Inst = Builder.CreateIntrinsic(
101-
Builder.getVoidTy(), Intrinsic::dx_resource_store_typedbuffer,
102-
{II->getOperand(0), II->getOperand(1), V});
103-
SI->replaceAllUsesWith(Inst);
220+
createStoreIntrinsic(II, SI, Current.Offset, RTI);
104221
DeadInsts.push_back(SI);
105222

106223
} else if (auto *LI = dyn_cast<LoadInst>(Current.Access)) {
107-
IRBuilder<> Builder(LI);
108-
Value *V = Builder.CreateIntrinsic(
109-
LoadType, Intrinsic::dx_resource_load_typedbuffer,
110-
{II->getOperand(0), II->getOperand(1)});
111-
V = Builder.CreateExtractValue(V, {0});
112-
113-
if (Current.Index)
114-
V = Builder.CreateExtractElement(V, Current.Index);
115-
116-
LI->replaceAllUsesWith(V);
224+
createLoadIntrinsic(II, LI, Current.Offset, RTI);
117225
DeadInsts.push_back(LI);
118226

119227
} else
@@ -137,15 +245,8 @@ static bool transformResourcePointers(Function &F, DXILResourceTypeMap &DRTM) {
137245
Resources.emplace_back(II, DRTM[HandleTy]);
138246
}
139247

140-
for (auto &[II, RI] : Resources) {
141-
if (RI.isTyped()) {
142-
Changed = true;
143-
replaceTypedBufferAccess(II, RI);
144-
}
145-
146-
// TODO: handle other resource types. We should probably have an
147-
// `unreachable` here once we've added support for all of them.
148-
}
248+
for (auto &[II, RI] : Resources)
249+
replaceAccess(II, RI);
149250

150251
return Changed;
151252
}

0 commit comments

Comments
 (0)