[SYCLLowerIR][GlobalOffset] Avoid erasing live non-GEP/PHI defs

llvm/lib/SYCLLowerIR/GlobalOffset.cpp

@@ -60,7 +60,7 @@ ModulePass *llvm::createGlobalOffsetPassLegacy() {
   return new GlobalOffsetLegacy();
 }
 
-// Helper function to collect all GEPs, PHIs and Loads in post-order.
+// Helper function to collect all Uses of Load's pointer operand in post-order.
 static void collectGlobalOffsetUses(Function *ImplicitOffsetIntrinsic,
                                     SmallVectorImpl<Instruction *> &LoadPtrUses,
                                     SmallVectorImpl<Instruction *> &Loads) {
@@ -246,11 +246,14 @@ PreservedAnalyses GlobalOffsetPass::run(Module &M, ModuleAnalysisManager &) {
     L->eraseFromParent();
   }
 
-  // Remove all collected Loads and GEPs from the kernel.
+  // Try to remove all collected Loads and their Defs from the kernel.
   // PtrUses is returned by `collectGlobalOffsetUses` in topological order.
   // Walk it backwards so we don't violate users.
-  for (auto *I : reverse(PtrUses))
-    I->eraseFromParent();
+  for (auto *I : reverse(PtrUses)) {
+    // A Def might not be a GEP. Remove it if it has no use.
+    if (I->use_empty())
+      I->eraseFromParent();
+  }
 
   // Remove all collected CallInsts from the kernel.
   for (auto *U : make_early_inc_range(ImplicitOffsetIntrinsic->users()))

llvm/test/CodeGen/NVPTX/global-offset-non-gep-phi-use.ll

@@ -0,0 +1,49 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; 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()
+declare i32 @llvm.nvvm.read.ptx.sreg.ctaid.x()
+
+define i64 @test_non_gep_phi_use(i32 %x) {
+; CHECK-LABEL: define i64 @test_non_gep_phi_use(
+; CHECK-SAME: i32 [[X:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[TRUNC:%.*]] = trunc nuw i32 [[X]] to i1
+; CHECK-NEXT:    [[CTAID_X:%.*]] = tail call i32 @llvm.nvvm.read.ptx.sreg.ctaid.x()
+; CHECK-NEXT:    [[CTAID_Y:%.*]] = tail call i32 @llvm.nvvm.read.ptx.sreg.ctaid.y()
+; CHECK-NEXT:    [[CTAID_XY:%.*]] = select i1 [[TRUNC]], i32 [[CTAID_Y]], i32 [[CTAID_X]]
+; CHECK-NEXT:    [[RES:%.*]] = zext i32 0 to i64
+; CHECK-NEXT:    ret i64 [[RES]]
+;
+entry:
+  %trunc = trunc nuw i32 %x to i1
+  %ctaid_x = tail call i32 @llvm.nvvm.read.ptx.sreg.ctaid.x()
+  %ctaid_y = tail call i32 @llvm.nvvm.read.ptx.sreg.ctaid.y()
+  %ctaid_xy = select i1 %trunc, i32 %ctaid_y, i32 %ctaid_x
+  %offset = tail call ptr @llvm.nvvm.implicit.offset()
+  %idx = select i1 %trunc, i64 4, i64 0
+  %gep = getelementptr inbounds nuw i8, ptr %offset, i64 %idx
+  %load = load i32, ptr %gep, align 4
+  %res = zext i32 %load to i64
+  ret i64 %res
+}
+
+; CHECK-LABEL: define i64 @test_non_gep_phi_use_with_offset(
+; CHECK-SAME: i32 [[X:%.*]], ptr [[PTR:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[TRUNC:%.*]] = trunc nuw i32 [[X]] to i1
+; CHECK-NEXT:    [[CTAID_X:%.*]] = tail call i32 @llvm.nvvm.read.ptx.sreg.ctaid.x()
+; CHECK-NEXT:    [[CTAID_Y:%.*]] = tail call i32 @llvm.nvvm.read.ptx.sreg.ctaid.y()
+; CHECK-NEXT:    [[CTAID_XY:%.*]] = select i1 [[TRUNC]], i32 [[CTAID_Y]], i32 [[CTAID_X]]
+; CHECK-NEXT:    [[IDX:%.*]] = select i1 [[TRUNC]], i64 4, i64 0
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds nuw i8, ptr [[PTR]], i64 [[IDX]]
+; CHECK-NEXT:    [[LOAD:%.*]] = load i32, ptr [[GEP]], align 4
+; CHECK-NEXT:    [[RES:%.*]] = zext i32 [[LOAD]] to i64
+; CHECK-NEXT:    ret i64 [[RES]]
+
+!llvm.module.flags = !{!0}
+
+!0 = !{i32 1, !"sycl-device", i32 1}