@@ -60,16 +60,52 @@ ModulePass *llvm::createGlobalOffsetPassLegacy() {
6060 return new GlobalOffsetLegacy ();
6161}
6262
63- // Recursive helper function to collect Loads from GEPs in a BFS fashion.
64- static void getLoads (Instruction *P, SmallVectorImpl<Instruction *> &Traversed,
65- SmallVectorImpl<LoadInst *> &Loads) {
66- Traversed.push_back (P);
67- if (auto *L = dyn_cast<LoadInst>(P)) // Base case for recursion
68- Loads.push_back (L);
69- else {
70- assert (isa<GetElementPtrInst>(*P));
71- for (Value *V : P->users ())
72- getLoads (cast<Instruction>(V), Traversed, Loads);
63+ // Helper function to collect all GEPs, PHIs and Loads in post-order.
64+ static void collectGlobalOffsetUses (Function *ImplicitOffsetIntrinsic,
65+ SmallVectorImpl<Instruction *> &LoadPtrUses,
66+ SmallVectorImpl<Instruction *> &Loads) {
67+ SmallVector<Instruction *, 4 > WorkList;
68+ SmallPtrSet<Value *, 4 > Visited;
69+
70+ // Find load instructions.
71+ for (auto *U : ImplicitOffsetIntrinsic->users ()) {
72+ for (auto *U2 : cast<CallInst>(U)->users ()) {
73+ auto *I = cast<Instruction>(U2);
74+ WorkList.push_back (I);
75+ Visited.insert (I);
76+ }
77+ }
78+ while (!WorkList.empty ()) {
79+ Instruction *I = WorkList.pop_back_val ();
80+ if (isa<PHINode>(I) || isa<GetElementPtrInst>(I)) {
81+ for (User *U : I->users ())
82+ if (Visited.insert (U).second )
83+ WorkList.push_back (cast<Instruction>(U));
84+ }
85+ if (isa<LoadInst>(I))
86+ Loads.push_back (I);
87+ }
88+
89+ // For each load, find its defs by post-order walking operand use.
90+ Visited.clear ();
91+ for (auto *LI : Loads) {
92+ Use *OpUse0 = &LI->getOperandUse (0 );
93+ auto PostOrderTraveral = [&](auto &Self, Use &U) -> void {
94+ auto *I = cast<Instruction>(U.get ());
95+ Visited.insert (I);
96+ for (auto &Op : I->operands ()) {
97+ auto *OpI = dyn_cast<Instruction>(Op.get ());
98+ if (!OpI || isa<CallInst>(OpI))
99+ continue ;
100+ if (!Visited.contains (OpI))
101+ Self (Self, Op);
102+ }
103+ if (!isa<CallInst>(I))
104+ LoadPtrUses.push_back (I);
105+ };
106+ Visited.insert (LI);
107+ if (!Visited.contains (OpUse0->get ()))
108+ PostOrderTraveral (PostOrderTraveral, *OpUse0);
73109 }
74110}
75111
@@ -199,32 +235,26 @@ PreservedAnalyses GlobalOffsetPass::run(Module &M, ModuleAnalysisManager &) {
199235 // Add implicit parameters to all direct and indirect users of the offset
200236 addImplicitParameterToCallers (M, ImplicitOffsetIntrinsic, nullptr , KCache);
201237 }
202- SmallVector<CallInst *, 4 > Worklist;
203- SmallVector<LoadInst *, 4 > Loads;
238+ SmallVector<Instruction *, 4 > Loads;
204239 SmallVector<Instruction *, 4 > PtrUses;
205240
206- // Collect all GEPs and Loads from the intrinsic's CallInsts
207- for (Value *V : ImplicitOffsetIntrinsic->users ()) {
208- Worklist.push_back (cast<CallInst>(V));
209- for (Value *V2 : V->users ())
210- getLoads (cast<Instruction>(V2), PtrUses, Loads);
211- }
241+ collectGlobalOffsetUses (ImplicitOffsetIntrinsic, PtrUses, Loads);
212242
213243 // Replace each use of a collected Load with a Constant 0
214- for (LoadInst *L : Loads)
244+ for (Instruction *L : Loads) {
215245 L->replaceAllUsesWith (ConstantInt::get (L->getType (), 0 ));
246+ L->eraseFromParent ();
247+ }
216248
217249 // Remove all collected Loads and GEPs from the kernel.
218- // PtrUses is returned by `getLoads ` in topological order.
250+ // PtrUses is returned by `collectGlobalOffsetUses ` in topological order.
219251 // Walk it backwards so we don't violate users.
220252 for (auto *I : reverse (PtrUses))
221253 I->eraseFromParent ();
222254
223255 // Remove all collected CallInsts from the kernel.
224- for (CallInst *CI : Worklist) {
225- auto *I = cast<Instruction>(CI);
226- I->eraseFromParent ();
227- }
256+ for (auto *U : make_early_inc_range (ImplicitOffsetIntrinsic->users ()))
257+ cast<Instruction>(U)->eraseFromParent ();
228258
229259 // Assert that all uses of `ImplicitOffsetIntrinsic` are removed and delete
230260 // it.
0 commit comments