[InferAddressSpaces] Fix bad addrspacecast insertion for phinode

llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp

@@ -617,6 +617,41 @@ InferAddressSpacesImpl::collectFlatAddressExpressions(Function &F) const {
   return Postorder;
 }
 
+// Inserts an addrspacecast for a phi node operand, handling the proper
+// insertion position based on the operand type.
+static Value *phiNodeOperandWithNewAddressSpace(AddrSpaceCastInst *NewI,
+                                                Value *Operand) {
+  auto InsertBefore = [NewI](auto It) {
+    NewI->insertBefore(It);
+    NewI->setDebugLoc(It->getDebugLoc());
+    return NewI;
+  };
+
+  if (auto *Arg = dyn_cast<Argument>(Operand)) {
+    // For arguments, insert the cast at the beginning of entry block.
+    // Consider inserting at the dominating block for better placement.
+    Function *F = Arg->getParent();
+    auto InsertI = F->getEntryBlock().getFirstNonPHIIt();
+    return InsertBefore(InsertI);
+  }
+
+  // No check for Constant here, as constants are already handled.
+  assert(isa<Instruction>(Operand));
+
+  Instruction *OpInst = cast<Instruction>(Operand);
+  if (LLVM_UNLIKELY(OpInst->getOpcode() == Instruction::PHI)) {
+    // If the operand is defined by another PHI node, insert after the first
+    // non-PHI instruction at the corresponding basic block.
+    auto InsertI = OpInst->getParent()->getFirstNonPHIIt();
+    return InsertBefore(InsertI);
+  }
+
+  // Otherwise, insert immediately after the operand definition.
+  NewI->insertAfter(OpInst->getIterator());
+  NewI->setDebugLoc(OpInst->getDebugLoc());
+  return NewI;
+}
+
 // A helper function for cloneInstructionWithNewAddressSpace. Returns the clone
 // of OperandUse.get() in the new address space. If the clone is not ready yet,
 // returns poison in the new address space as a placeholder.
@@ -642,6 +677,10 @@ static Value *operandWithNewAddressSpaceOrCreatePoison(
     unsigned NewAS = I->second;
     Type *NewPtrTy = getPtrOrVecOfPtrsWithNewAS(Operand->getType(), NewAS);
     auto *NewI = new AddrSpaceCastInst(Operand, NewPtrTy);
+
+    if (LLVM_UNLIKELY(Inst->getOpcode() == Instruction::PHI))
+      return phiNodeOperandWithNewAddressSpace(NewI, Operand);
+
     NewI->insertBefore(Inst->getIterator());
     NewI->setDebugLoc(Inst->getDebugLoc());
     return NewI;

llvm/test/Transforms/InferAddressSpaces/NVPTX/phinode-address-infer.ll

@@ -0,0 +1,57 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -passes='require<domtree>,infer-address-spaces' %s | FileCheck %s
+
+;;; Handle write corner case for infer-address-spaces with phi-nodes. The
+;;; verifier will crash if we insert `addrspacecast` before phi-node.
+
+target triple = "nvptx64-nvidia-cuda"
+
+declare void @llvm.assume(i1 noundef)
+declare i1 @llvm.nvvm.isspacep.shared(ptr) readnone noinline
+declare i1 @llvm.nvvm.isspacep.global(ptr) readnone noinline
+
+define ptr @phinode_instr() {
+; CHECK-LABEL: @phinode_instr(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[PTR_1:%.*]] = load ptr, ptr null, align 8
+; CHECK-NEXT:    [[TMP0:%.*]] = addrspacecast ptr [[PTR_1]] to ptr addrspace(3)
+; CHECK-NEXT:    [[BOOL_1:%.*]] = tail call i1 @llvm.nvvm.isspacep.shared(ptr [[PTR_1]])
+; CHECK-NEXT:    tail call void @llvm.assume(i1 [[BOOL_1]])
+; CHECK-NEXT:    br label [[IF_SINK_SPLIT:%.*]]
+; CHECK:       if.sink.split:
+; CHECK-NEXT:    [[PTR_SINK:%.*]] = phi ptr addrspace(3) [ [[TMP0]], [[ENTRY:%.*]] ]
+; CHECK-NEXT:    [[TMP1:%.*]] = addrspacecast ptr addrspace(3) [[PTR_SINK]] to ptr
+; CHECK-NEXT:    ret ptr [[TMP1]]
+;
+entry:
+  %ptr.1 = load ptr, ptr null, align 8
+  %bool.1 = tail call i1 @llvm.nvvm.isspacep.shared(ptr %ptr.1)
+  tail call void @llvm.assume(i1 %bool.1)
+  br label %if.sink.split
+
+if.sink.split:                                    ; preds = %entry
+  %ptr.sink = phi ptr [ %ptr.1, %entry ]
+  ret ptr %ptr.sink
+}
+
+define ptr @phinode_argument(ptr %lhs_ptr) {
+; CHECK-LABEL: @phinode_argument(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = addrspacecast ptr [[LHS_PTR:%.*]] to ptr addrspace(1)
+; CHECK-NEXT:    [[BOOL_1:%.*]] = tail call i1 @llvm.nvvm.isspacep.global(ptr [[LHS_PTR]])
+; CHECK-NEXT:    tail call void @llvm.assume(i1 [[BOOL_1]])
+; CHECK-NEXT:    br label [[IF_SINK_SPLIT:%.*]]
+; CHECK:       if.sink.split:
+; CHECK-NEXT:    [[PTR_SINK:%.*]] = phi ptr addrspace(1) [ [[TMP0]], [[ENTRY:%.*]] ]
+; CHECK-NEXT:    [[TMP1:%.*]] = addrspacecast ptr addrspace(1) [[PTR_SINK]] to ptr
+; CHECK-NEXT:    ret ptr [[TMP1]]
+;
+entry:
+  %bool.1 = tail call i1 @llvm.nvvm.isspacep.global(ptr %lhs_ptr)
+  tail call void @llvm.assume(i1 %bool.1)
+  br label %if.sink.split
+
+if.sink.split:                                    ; preds = %entry
+  %ptr.sink = phi ptr [ %lhs_ptr, %entry ]
+  ret ptr %ptr.sink
+}