Skip to content

Commit 685b9c3

Browse files
bognerbogner
committed
[DirectX] Handle dx.RawBuffer in DXILResourceAccess
This adds handling for raw and structured buffers when lowering resource access via `llvm.dx.resource.getpointer`. Fixes #121714
1 parent acbd822 commit 685b9c3

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+
// TODO: handle these
119+
return;
120+
case dxil::ResourceKind::CBuffer:
121+
case dxil::ResourceKind::Sampler:
122+
case dxil::ResourceKind::TBuffer:
123+
case dxil::ResourceKind::RTAccelerationStructure:
124+
case dxil::ResourceKind::Invalid:
125+
case dxil::ResourceKind::NumEntries:
126+
llvm_unreachable("Invalid resource kind for store");
127+
}
128+
llvm_unreachable("Unhandled case in switch");
129+
}
130+
131+
static void createTypedBufferLoad(IntrinsicInst *II, LoadInst *LI,
132+
Value *Offset, dxil::ResourceTypeInfo &RTI) {
133+
IRBuilder<> Builder(LI);
134+
Type *ContainedType = RTI.getHandleTy()->getTypeParameter(0);
135+
Type *LoadType = StructType::get(ContainedType, Builder.getInt1Ty());
27136

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);
137+
Value *V =
138+
Builder.CreateIntrinsic(LoadType, Intrinsic::dx_resource_load_typedbuffer,
139+
{II->getOperand(0), II->getOperand(1)});
140+
V = Builder.CreateExtractValue(V, {0});
32141

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

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;
145+
LI->replaceAllUsesWith(V);
146+
}
40147

148+
static void createRawLoad(IntrinsicInst *II, LoadInst *LI, Value *Offset) {
149+
IRBuilder<> Builder(LI);
150+
// TODO: break up larger types
151+
Type *LoadType = StructType::get(LI->getType(), Builder.getInt1Ty());
152+
if (!Offset)
153+
Offset = ConstantInt::get(Builder.getInt32Ty(), 0);
154+
Value *V =
155+
Builder.CreateIntrinsic(LoadType, Intrinsic::dx_resource_load_rawbuffer,
156+
{II->getOperand(0), II->getOperand(1), Offset});
157+
V = Builder.CreateExtractValue(V, {0});
158+
159+
LI->replaceAllUsesWith(V);
160+
}
161+
162+
static void createLoadIntrinsic(IntrinsicInst *II, LoadInst *LI, Value *Offset,
163+
dxil::ResourceTypeInfo &RTI) {
164+
switch (RTI.getResourceKind()) {
165+
case dxil::ResourceKind::TypedBuffer:
166+
return createTypedBufferLoad(II, LI, Offset, RTI);
167+
case dxil::ResourceKind::RawBuffer:
168+
case dxil::ResourceKind::StructuredBuffer:
169+
return createRawLoad(II, LI, Offset);
170+
case dxil::ResourceKind::Texture1D:
171+
case dxil::ResourceKind::Texture2D:
172+
case dxil::ResourceKind::Texture2DMS:
173+
case dxil::ResourceKind::Texture3D:
174+
case dxil::ResourceKind::TextureCube:
175+
case dxil::ResourceKind::Texture1DArray:
176+
case dxil::ResourceKind::Texture2DArray:
177+
case dxil::ResourceKind::Texture2DMSArray:
178+
case dxil::ResourceKind::TextureCubeArray:
179+
case dxil::ResourceKind::FeedbackTexture2D:
180+
case dxil::ResourceKind::FeedbackTexture2DArray:
181+
case dxil::ResourceKind::CBuffer:
182+
case dxil::ResourceKind::TBuffer:
183+
// TODO: handle these
184+
return;
185+
case dxil::ResourceKind::Sampler:
186+
case dxil::ResourceKind::RTAccelerationStructure:
187+
case dxil::ResourceKind::Invalid:
188+
case dxil::ResourceKind::NumEntries:
189+
llvm_unreachable("Invalid resource kind for load");
190+
}
191+
llvm_unreachable("Unhandled case in switch");
192+
}
193+
194+
static void
195+
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)