xref: /freebsd/contrib/llvm-project/llvm/lib/Target/DirectX/DXILResourceAccess.cpp (revision 770cf0a5f02dc8983a89c6568d741fbc25baa999) 
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 
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     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 
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 
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 
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 
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 
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 
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:
280   bool runOnFunction(Function &F) override {
281     DXILResourceTypeMap &DRTM =
282         getAnalysis<DXILResourceTypeWrapperPass>().getResourceTypeMap();
283 
284     return transformResourcePointers(F, DRTM);
285   }
286   StringRef getPassName() const override { return "DXIL Resource Access"; }
287   DXILResourceAccessLegacy() : FunctionPass(ID) {}
288 
289   static char ID; // Pass identification.
290   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)
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