1 //===- DXILResourceAccess.cpp - Resource access via load/store ------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8
9 #include "DXILResourceAccess.h"
10 #include "DirectX.h"
11 #include "llvm/Analysis/DXILResource.h"
12 #include "llvm/IR/Dominators.h"
13 #include "llvm/IR/IRBuilder.h"
14 #include "llvm/IR/Instructions.h"
15 #include "llvm/IR/IntrinsicInst.h"
16 #include "llvm/IR/Intrinsics.h"
17 #include "llvm/IR/IntrinsicsDirectX.h"
18 #include "llvm/InitializePasses.h"
19
20 #define DEBUG_TYPE "dxil-resource-access"
21
22 using namespace llvm;
23
calculateGEPOffset(GetElementPtrInst * GEP,Value * PrevOffset,dxil::ResourceTypeInfo & RTI)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
createTypedBufferStore(IntrinsicInst * II,StoreInst * SI,Value * Offset,dxil::ResourceTypeInfo & RTI)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
createRawStore(IntrinsicInst * II,StoreInst * SI,Value * Offset)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
createStoreIntrinsic(IntrinsicInst * II,StoreInst * SI,Value * Offset,dxil::ResourceTypeInfo & RTI)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 reportFatalUsageError("DXIL Load not implemented yet");
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
createTypedBufferLoad(IntrinsicInst * II,LoadInst * LI,Value * Offset,dxil::ResourceTypeInfo & RTI)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());
136
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});
141
142 if (Offset)
143 V = Builder.CreateExtractElement(V, Offset);
144
145 // If we loaded a <1 x ...> instead of a scalar (presumably to feed a
146 // shufflevector), then make sure we're maintaining the resulting type.
147 if (auto *VT = dyn_cast<FixedVectorType>(LI->getType()))
148 if (VT->getNumElements() == 1 && !isa<FixedVectorType>(V->getType()))
149 V = Builder.CreateInsertElement(PoisonValue::get(VT), V,
150 Builder.getInt32(0));
151
152 LI->replaceAllUsesWith(V);
153 }
154
createRawLoad(IntrinsicInst * II,LoadInst * LI,Value * Offset)155 static void createRawLoad(IntrinsicInst *II, LoadInst *LI, Value *Offset) {
156 IRBuilder<> Builder(LI);
157 // TODO: break up larger types
158 Type *LoadType = StructType::get(LI->getType(), Builder.getInt1Ty());
159 if (!Offset)
160 Offset = ConstantInt::get(Builder.getInt32Ty(), 0);
161 Value *V =
162 Builder.CreateIntrinsic(LoadType, Intrinsic::dx_resource_load_rawbuffer,
163 {II->getOperand(0), II->getOperand(1), Offset});
164 V = Builder.CreateExtractValue(V, {0});
165
166 LI->replaceAllUsesWith(V);
167 }
168
createLoadIntrinsic(IntrinsicInst * II,LoadInst * LI,Value * Offset,dxil::ResourceTypeInfo & RTI)169 static void createLoadIntrinsic(IntrinsicInst *II, LoadInst *LI, Value *Offset,
170 dxil::ResourceTypeInfo &RTI) {
171 switch (RTI.getResourceKind()) {
172 case dxil::ResourceKind::TypedBuffer:
173 return createTypedBufferLoad(II, LI, Offset, RTI);
174 case dxil::ResourceKind::RawBuffer:
175 case dxil::ResourceKind::StructuredBuffer:
176 return createRawLoad(II, LI, Offset);
177 case dxil::ResourceKind::Texture1D:
178 case dxil::ResourceKind::Texture2D:
179 case dxil::ResourceKind::Texture2DMS:
180 case dxil::ResourceKind::Texture3D:
181 case dxil::ResourceKind::TextureCube:
182 case dxil::ResourceKind::Texture1DArray:
183 case dxil::ResourceKind::Texture2DArray:
184 case dxil::ResourceKind::Texture2DMSArray:
185 case dxil::ResourceKind::TextureCubeArray:
186 case dxil::ResourceKind::FeedbackTexture2D:
187 case dxil::ResourceKind::FeedbackTexture2DArray:
188 case dxil::ResourceKind::CBuffer:
189 case dxil::ResourceKind::TBuffer:
190 // TODO: handle these
191 return;
192 case dxil::ResourceKind::Sampler:
193 case dxil::ResourceKind::RTAccelerationStructure:
194 case dxil::ResourceKind::Invalid:
195 case dxil::ResourceKind::NumEntries:
196 llvm_unreachable("Invalid resource kind for load");
197 }
198 llvm_unreachable("Unhandled case in switch");
199 }
200
replaceAccess(IntrinsicInst * II,dxil::ResourceTypeInfo & RTI)201 static void replaceAccess(IntrinsicInst *II, dxil::ResourceTypeInfo &RTI) {
202 // Process users keeping track of indexing accumulated from GEPs.
203 struct AccessAndOffset {
204 User *Access;
205 Value *Offset;
206 };
207 SmallVector<AccessAndOffset> Worklist;
208 for (User *U : II->users())
209 Worklist.push_back({U, nullptr});
210
211 SmallVector<Instruction *> DeadInsts;
212 while (!Worklist.empty()) {
213 AccessAndOffset Current = Worklist.back();
214 Worklist.pop_back();
215
216 if (auto *GEP = dyn_cast<GetElementPtrInst>(Current.Access)) {
217 IRBuilder<> Builder(GEP);
218
219 Value *Offset = calculateGEPOffset(GEP, Current.Offset, RTI);
220 for (User *U : GEP->users())
221 Worklist.push_back({U, Offset});
222 DeadInsts.push_back(GEP);
223
224 } else if (auto *SI = dyn_cast<StoreInst>(Current.Access)) {
225 assert(SI->getValueOperand() != II && "Pointer escaped!");
226 createStoreIntrinsic(II, SI, Current.Offset, RTI);
227 DeadInsts.push_back(SI);
228
229 } else if (auto *LI = dyn_cast<LoadInst>(Current.Access)) {
230 createLoadIntrinsic(II, LI, Current.Offset, RTI);
231 DeadInsts.push_back(LI);
232
233 } else
234 llvm_unreachable("Unhandled instruction - pointer escaped?");
235 }
236
237 // Traverse the now-dead instructions in RPO and remove them.
238 for (Instruction *Dead : llvm::reverse(DeadInsts))
239 Dead->eraseFromParent();
240 II->eraseFromParent();
241 }
242
transformResourcePointers(Function & F,DXILResourceTypeMap & DRTM)243 static bool transformResourcePointers(Function &F, DXILResourceTypeMap &DRTM) {
244 bool Changed = false;
245 SmallVector<std::pair<IntrinsicInst *, dxil::ResourceTypeInfo>> Resources;
246 for (BasicBlock &BB : F)
247 for (Instruction &I : BB)
248 if (auto *II = dyn_cast<IntrinsicInst>(&I))
249 if (II->getIntrinsicID() == Intrinsic::dx_resource_getpointer) {
250 auto *HandleTy = cast<TargetExtType>(II->getArgOperand(0)->getType());
251 Resources.emplace_back(II, DRTM[HandleTy]);
252 }
253
254 for (auto &[II, RI] : Resources)
255 replaceAccess(II, RI);
256
257 return Changed;
258 }
259
run(Function & F,FunctionAnalysisManager & FAM)260 PreservedAnalyses DXILResourceAccess::run(Function &F,
261 FunctionAnalysisManager &FAM) {
262 auto &MAMProxy = FAM.getResult<ModuleAnalysisManagerFunctionProxy>(F);
263 DXILResourceTypeMap *DRTM =
264 MAMProxy.getCachedResult<DXILResourceTypeAnalysis>(*F.getParent());
265 assert(DRTM && "DXILResourceTypeAnalysis must be available");
266
267 bool MadeChanges = transformResourcePointers(F, *DRTM);
268 if (!MadeChanges)
269 return PreservedAnalyses::all();
270
271 PreservedAnalyses PA;
272 PA.preserve<DXILResourceTypeAnalysis>();
273 PA.preserve<DominatorTreeAnalysis>();
274 return PA;
275 }
276
277 namespace {
278 class DXILResourceAccessLegacy : public FunctionPass {
279 public:
runOnFunction(Function & F)280 bool runOnFunction(Function &F) override {
281 DXILResourceTypeMap &DRTM =
282 getAnalysis<DXILResourceTypeWrapperPass>().getResourceTypeMap();
283
284 return transformResourcePointers(F, DRTM);
285 }
getPassName() const286 StringRef getPassName() const override { return "DXIL Resource Access"; }
DXILResourceAccessLegacy()287 DXILResourceAccessLegacy() : FunctionPass(ID) {}
288
289 static char ID; // Pass identification.
getAnalysisUsage(llvm::AnalysisUsage & AU) const290 void getAnalysisUsage(llvm::AnalysisUsage &AU) const override {
291 AU.addRequired<DXILResourceTypeWrapperPass>();
292 AU.addPreserved<DominatorTreeWrapperPass>();
293 }
294 };
295 char DXILResourceAccessLegacy::ID = 0;
296 } // end anonymous namespace
297
298 INITIALIZE_PASS_BEGIN(DXILResourceAccessLegacy, DEBUG_TYPE,
299 "DXIL Resource Access", false, false)
INITIALIZE_PASS_DEPENDENCY(DXILResourceTypeWrapperPass)300 INITIALIZE_PASS_DEPENDENCY(DXILResourceTypeWrapperPass)
301 INITIALIZE_PASS_END(DXILResourceAccessLegacy, DEBUG_TYPE,
302 "DXIL Resource Access", false, false)
303
304 FunctionPass *llvm::createDXILResourceAccessLegacyPass() {
305 return new DXILResourceAccessLegacy();
306 }
307