[SYCLLowerIR][GlobalOffset] Fix collection order of load's defs

llvm/lib/SYCLLowerIR/GlobalOffset.cpp

@@ -60,16 +60,52 @@ ModulePass *llvm::createGlobalOffsetPassLegacy() {
   return new GlobalOffsetLegacy();
 }
 
-// Recursive helper function to collect Loads from GEPs in a BFS fashion.
-static void getLoads(Instruction *P, SmallVectorImpl<Instruction *> &Traversed,
-                     SmallVectorImpl<LoadInst *> &Loads) {
-  Traversed.push_back(P);
-  if (auto *L = dyn_cast<LoadInst>(P)) // Base case for recursion
-    Loads.push_back(L);
-  else {
-    assert(isa<GetElementPtrInst>(*P));
-    for (Value *V : P->users())
-      getLoads(cast<Instruction>(V), Traversed, Loads);
+// Helper function to collect all GEPs, PHIs and Loads in post-order.
+static void getLoads(Function *ImplicitOffsetIntrinsic,
+                     SmallVectorImpl<Instruction *> &Traversed,
+                     SmallVectorImpl<Instruction *> &Loads) {
+  SmallVector<Instruction *, 4> WorkList;
+  SmallPtrSet<Value *, 4> Visited;
+
+  // Find load instructions.
+  for (auto *U : ImplicitOffsetIntrinsic->users()) {
+    for (auto *U2 : cast<CallInst>(U)->users()) {
+      auto *I = cast<Instruction>(U2);
+      WorkList.push_back(I);
+      Visited.insert(I);
+    }
+  }
+  while (!WorkList.empty()) {
+    Instruction *I = WorkList.pop_back_val();
+    if (isa<PHINode>(I) || isa<GetElementPtrInst>(I)) {
+      for (User *U : I->users())
+        if (Visited.insert(U).second)
+          WorkList.push_back(cast<Instruction>(U));
+    }
+    if (isa<LoadInst>(I))
+      Loads.push_back(I);
+  }
+
+  // For each load, find its defs by post-order walking operand use.
+  Visited.clear();
+  for (auto *LI : Loads) {
+    Use *OpUse0 = &LI->getOperandUse(0);
+    auto PostOrderTraveral = [&](auto &Self, Use &U) -> void {
+      auto *I = cast<Instruction>(U.get());
+      Visited.insert(I);
+      for (auto &Op : I->operands()) {
+        auto *OpI = dyn_cast<Instruction>(Op.get());
+        if (!OpI || isa<CallInst>(OpI))
+          continue;
+        if (!Visited.contains(OpI))
+          Self(Self, Op);
+      }
+      if (!isa<CallInst>(I))
+        Traversed.push_back(I);
+    };
+    Visited.insert(LI);
+    if (!Visited.contains(OpUse0->get()))
+      PostOrderTraveral(PostOrderTraveral, *OpUse0);
   }
 }
 
@@ -199,20 +235,16 @@ PreservedAnalyses GlobalOffsetPass::run(Module &M, ModuleAnalysisManager &) {
     // Add implicit parameters to all direct and indirect users of the offset
     addImplicitParameterToCallers(M, ImplicitOffsetIntrinsic, nullptr, KCache);
   }
-  SmallVector<CallInst *, 4> Worklist;
-  SmallVector<LoadInst *, 4> Loads;
+  SmallVector<Instruction *, 4> Loads;
   SmallVector<Instruction *, 4> PtrUses;
 
-  // Collect all GEPs and Loads from the intrinsic's CallInsts
-  for (Value *V : ImplicitOffsetIntrinsic->users()) {
-    Worklist.push_back(cast<CallInst>(V));
-    for (Value *V2 : V->users())
-      getLoads(cast<Instruction>(V2), PtrUses, Loads);
-  }
+  getLoads(ImplicitOffsetIntrinsic, PtrUses, Loads);
 
   // Replace each use of a collected Load with a Constant 0
-  for (LoadInst *L : Loads)
+  for (Instruction *L : Loads) {
     L->replaceAllUsesWith(ConstantInt::get(L->getType(), 0));
+    L->eraseFromParent();
+  }
 
   // Remove all collected Loads and GEPs from the kernel.
   // PtrUses is returned by `getLoads` in topological order.
@@ -221,10 +253,8 @@ PreservedAnalyses GlobalOffsetPass::run(Module &M, ModuleAnalysisManager &) {
     I->eraseFromParent();
 
   // Remove all collected CallInsts from the kernel.
-  for (CallInst *CI : Worklist) {
-    auto *I = cast<Instruction>(CI);
-    I->eraseFromParent();
-  }
+  for (auto *U : make_early_inc_range(ImplicitOffsetIntrinsic->users()))
+    cast<Instruction>(U)->eraseFromParent();
 
   // Assert that all uses of `ImplicitOffsetIntrinsic` are removed and delete
   // it.

llvm/test/CodeGen/AMDGPU/global-offset-phi.ll

@@ -0,0 +1,45 @@
+; RUN: opt -bugpoint-enable-legacy-pm -globaloffset %s -S -o - | FileCheck %s
+
+; Check that phi is correctly handled in load's defs collection.
+
+target datalayout = "e-p:64:64-p1:64:64-p2:32:32-p3:32:32-p4:64:64-p5:32:32-p6:32:32-p7:160:256:256:32-p8:128:128:128:48-p9:192:256:256:32-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-v2048:2048-n32:64-S32-A5-G1-ni:7:8:9"
+target triple = "amdgcn-amd-amdhsa"
+
+define internal i64 @_Z33__spirv_BuiltInGlobalInvocationIdi(i32 %0) {
+; CHECK-NOT: call ptr addrspace(5) @llvm.amdgcn.implicit.offset()
+
+  switch i32 %0, label %14 [
+    i32 0, label %2
+    i32 1, label %4
+    i32 2, label %7
+  ]
+
+2:                                                ; preds = %1
+  %3 = tail call ptr addrspace(5) @llvm.amdgcn.implicit.offset()
+  br label %10
+
+4:                                                ; preds = %1
+  %5 = tail call ptr addrspace(5) @llvm.amdgcn.implicit.offset()
+  %6 = getelementptr inbounds nuw i8, ptr addrspace(5) %5, i32 4
+  br label %10
+
+7:                                                ; preds = %1
+  %8 = tail call ptr addrspace(5) @llvm.amdgcn.implicit.offset()
+  %9 = getelementptr inbounds nuw i8, ptr addrspace(5) %8, i32 8
+  br label %10
+
+10:                                               ; preds = %7, %4, %2
+  %11 = phi ptr addrspace(5) [ %3, %2 ], [ %6, %4 ], [ %9, %7 ]
+  %12 = load i32, ptr addrspace(5) %11, align 4
+  %13 = load i32, ptr addrspace(5) %11, align 4
+  ret i64 0
+
+14:                                               ; preds = %1
+  unreachable
+}
+
+declare ptr addrspace(5) @llvm.amdgcn.implicit.offset()
+
+!llvm.module.flags = !{!0}
+
+!0 = !{i32 1, !"sycl-device", i32 1}

llvm/test/CodeGen/NVPTX/global-offset-two-loads.ll

@@ -0,0 +1,21 @@
+; RUN: opt -bugpoint-enable-legacy-pm -globaloffset %s -S -o - | FileCheck %s
+
+target datalayout = "e-p6:32:32-i64:64-i128:128-i256:256-v16:16-v32:32-n16:32:64"
+target triple = "nvptx64-nvidia-cuda"
+
+declare ptr @llvm.nvvm.implicit.offset()
+
+define weak_odr dso_local ptx_kernel void @test() {
+entry:
+; CHECK-NOT: call ptr @llvm.nvvm.implicit.offset()
+
+  %0 = tail call ptr @llvm.nvvm.implicit.offset()
+  %1 = getelementptr inbounds nuw i8, ptr %0, i64 4
+  %2 = load i32, ptr %1, align 4
+  %3 = load i32, ptr %0, align 4
+  ret void
+}
+
+!llvm.module.flags = !{!0}
+
+!0 = !{i32 1, !"sycl-device", i32 1}