[AMDGPU] Fixed constexpr expansion to handle multiple uses

[rampitec]

Summary

Recently added convertConstantExprsToInstructions() does not handle
a case when a same ConstantExpr used multiple times in the same
instruction. A first use is replaced and the rest of the uses in the
instruction are replaced as well with the replaceUsesOfWith(). Then
function attempts to replace a constant already destroyed.

So far this interface is only used by the AMDGPU BE.

[tra]

I've tested it on the full code which triggered the issue and this patch fixes the issue. Thank you for fixing it promptly.

[tra]

llvm/lib/IR/ReplaceConstant.cpp

87

Do we need/want to bump the initial set size? I'd assume that it would need to be larger now that we keep track of all operands.

[rampitec]

llvm/lib/IR/ReplaceConstant.cpp

87

Actually I think even 8 is too much :) Usually expressions are way shorter.

llvm/lib/IR/ReplaceConstant.cpp

void convertConstantExprsToInstructions(Instruction *I, ConstantExpr *CE,
   // CEPaths.
   convertConstantExprsToInstructions(I, CEPaths, Insts);
 }
 
 void convertConstantExprsToInstructions(
     Instruction *I,
     std::map<Use *, std::vector<std::vector<ConstantExpr *>>> &CEPaths,
     SmallPtrSetImpl<Instruction *> *Insts) {
+  SmallPtrSet<ConstantExpr *, 8> Visited;
   for (Use &U : I->operands()) {
     // The operand U is either not a constant expression operand or the
     // constant expression paths do not belong to U, ignore U.
     if (!CEPaths.count(&U))
       continue;
 
     // If the instruction I is a PHI instruction, then fix the instruction
     // insertion point to the entry of the incoming basic block for operand U.
     auto *BI = I;
     if (auto *Phi = dyn_cast<PHINode>(I)) {
       BasicBlock *BB = Phi->getIncomingBlock(U);
       BI = &(*(BB->getFirstInsertionPt()));
     }
 
     // Go through the paths associated with operand U, and convert all the
     // constant expressions along all paths to corresponding instructions.
     auto *II = I;
     auto &Paths = CEPaths[&U];
-    SmallPtrSet<ConstantExpr *, 8> Visited;
     for (auto &Path : Paths) {
       for (auto *CE : Path) {
         if (!Visited.insert(CE).second)
           continue;
         auto *NI = CE->getAsInstruction();
         NI->insertBefore(BI);
         II->replaceUsesOfWith(CE, NI);
         CE->removeDeadConstantUsers();

llvm/test/CodeGen/AMDGPU/lower-kernel-lds-constexpr.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: -p --check-globals
 ; RUN: opt -S -mtriple=amdgcn-- -amdgpu-lower-module-lds < %s | FileCheck %s
 ; RUN: opt -S -mtriple=amdgcn-- -passes=amdgpu-lower-module-lds < %s | FileCheck %s
 
 @lds.1 = internal unnamed_addr addrspace(3) global [2 x i8] undef, align 1
 
 ; CHECK: %llvm.amdgcn.kernel.k0.lds.t = type { [2 x i8] }
 ; CHECK: %llvm.amdgcn.kernel.k1.lds.t = type { [2 x i8] }
 ; CHECK: %llvm.amdgcn.kernel.k2.lds.t = type { i32 }
 ; CHECK: %llvm.amdgcn.kernel.k3.lds.t = type { [32 x i8] }
 ; CHECK: %llvm.amdgcn.kernel.k4.lds.t = type { [2 x i8] }
+; CHECK: %llvm.amdgcn.kernel.k5.lds.t = type { [505 x i32] }
 
 ; Use constant from different kernels
 ;.
 ; CHECK: @llvm.amdgcn.kernel.k0.lds = internal addrspace(3) global %llvm.amdgcn.kernel.k0.lds.t undef, align 2
 ; CHECK: @llvm.amdgcn.kernel.k1.lds = internal addrspace(3) global %llvm.amdgcn.kernel.k1.lds.t undef, align 2
 ; CHECK: @llvm.amdgcn.kernel.k2.lds = internal addrspace(3) global %llvm.amdgcn.kernel.k2.lds.t undef, align 4
 ; CHECK: @llvm.amdgcn.kernel.k3.lds = internal addrspace(3) global %llvm.amdgcn.kernel.k3.lds.t undef, align 16
 ; CHECK: @llvm.amdgcn.kernel.k4.lds = internal addrspace(3) global %llvm.amdgcn.kernel.k4.lds.t undef, align 2
+; CHECK: @llvm.amdgcn.kernel.k5.lds = internal addrspace(3) global %llvm.amdgcn.kernel.k5.lds.t undef, align 16
 ;.
 define amdgpu_kernel void @k0(i64 %x) {
 ; CHECK-LABEL: @k0(
 ; CHECK-NEXT:    %1 = getelementptr inbounds [2 x i8], [2 x i8] addrspace(3)* getelementptr inbounds (%llvm.amdgcn.kernel.k0.lds.t, %llvm.amdgcn.kernel.k0.lds.t addrspace(3)* @llvm.amdgcn.kernel.k0.lds, i32 0, i32 0), i32 0, i32 0
 ; CHECK-NEXT:    %2 = addrspacecast i8 addrspace(3)* %1 to i8*
 ; CHECK-NEXT:    %ptr = getelementptr inbounds i8, i8* %2, i64 %x
 ; CHECK-NEXT:    store i8 1, i8* %ptr, align 1
 ; CHECK-NEXT:    ret void

 ; CHECK-NEXT:    %ptr = getelementptr inbounds i8, i8* %2, i64 %x
 ; CHECK-NEXT:    store i8 1, i8* %ptr, align 1
 ; CHECK-NEXT:    ret void
 ;
   %ptr = getelementptr inbounds i8, i8* addrspacecast ([2 x i8] addrspace(3)* @lds.1 to i8*), i64 %x
   store i8 1, i8 addrspace(0)* %ptr, align 1
   ret void
 }
+
+@lds.4 = internal unnamed_addr addrspace(3) global [505 x i32] undef, align 4
+
+; Multiple constexpr use in a same instruction.
+define amdgpu_kernel void @k5() {
+; CHECK-LABEL: @k5(
+; CHECK-NEXT: %1 = addrspacecast [505 x i32] addrspace(3)* getelementptr inbounds (%llvm.amdgcn.kernel.k5.lds.t, %llvm.amdgcn.kernel.k5.lds.t addrspace(3)* @llvm.amdgcn.kernel.k5.lds, i32 0, i32 0) to [505 x i32]*
+; CHECK-NEXT: %2 = getelementptr inbounds [505 x i32], [505 x i32]* %1, i64 0, i64 0
+; CHECK-NEXT: call void undef(i32* %2, i32* %2)
+;
+  call void undef(i32* getelementptr inbounds ([505 x i32], [505 x i32]* addrspacecast ([505 x i32] addrspace(3)* @lds.4 to [505 x i32]*), i64 0, i64 0), i32* getelementptr inbounds ([505 x i32], [505 x i32]* addrspacecast ([505 x i32] addrspace(3)* @lds.4 to [505 x i32]*), i64 0, i64 0))
+  ret void
+}