[llvm] [SROA] When unfolding gep of phi, insert gep into predecessor (PR #127652)

[Nikita Popov via llvm-commits]

https://github.com/nikic created https://github.com/llvm/llvm-project/pull/127652

Instead of inserting all the geps into the entry block, insert them into the corresponding phi predecessor.

This is motivated by https://github.com/llvm/llvm-project/issues/127365, though it doesn't actually fix that issue, because the GEPs get hoisted by LICM later. If that didn't happen, having the GEPs in the predecessors would enable SimplifyCFG to undo the transform via sinking.

>From eedc55aa2b5f9feeb77a1836adbb8cfaf5da6f62 Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov at redhat.com>
Date: Tue, 18 Feb 2025 16:25:51 +0100
Subject: [PATCH] [SROA] When unfolding gep of phi, insert gep into predecessor

Instead of inserting all the geps into the entry block, insert
them into the corresponding phi predecessor.
---
 llvm/lib/Transforms/Scalar/SROA.cpp         | 2 +-
 llvm/test/Transforms/SROA/phi-and-select.ll | 2 +-
 llvm/test/Transforms/SROA/phi-gep.ll        | 4 ++--
 3 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/SROA.cpp b/llvm/lib/Transforms/Scalar/SROA.cpp
index e88c130cccf20..3dc125a0b6391 100644
--- a/llvm/lib/Transforms/Scalar/SROA.cpp
+++ b/llvm/lib/Transforms/Scalar/SROA.cpp
@@ -4123,7 +4123,6 @@ class AggLoadStoreRewriter : public InstVisitor<AggLoadStoreRewriter, bool> {
     Type *SourceTy = GEPI.getSourceElementType();
     // We only handle arguments, constants, and static allocas here, so we can
     // insert GEPs at the end of the entry block.
-    IRB.SetInsertPoint(GEPI.getFunction()->getEntryBlock().getTerminator());
     for (unsigned I = 0, E = Phi->getNumIncomingValues(); I != E; ++I) {
       Value *Op = Phi->getIncomingValue(I);
       BasicBlock *BB = Phi->getIncomingBlock(I);
@@ -4132,6 +4131,7 @@ class AggLoadStoreRewriter : public InstVisitor<AggLoadStoreRewriter, bool> {
         NewGEP = NewPhi->getIncomingValue(NI);
       } else {
         SmallVector<Value *> NewOps = GetNewOps(Op);
+        IRB.SetInsertPoint(BB->getTerminator());
         NewGEP =
             IRB.CreateGEP(SourceTy, NewOps[0], ArrayRef(NewOps).drop_front(),
                           Phi->getName() + ".sroa.gep", GEPI.getNoWrapFlags());
diff --git a/llvm/test/Transforms/SROA/phi-and-select.ll b/llvm/test/Transforms/SROA/phi-and-select.ll
index 616617b079828..4c87f090da94e 100644
--- a/llvm/test/Transforms/SROA/phi-and-select.ll
+++ b/llvm/test/Transforms/SROA/phi-and-select.ll
@@ -767,7 +767,6 @@ define void @PR20822(i1 %c1, i1 %c2, ptr %ptr) {
 ; CHECK-LABEL: @PR20822(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[F_SROA_0:%.*]] = alloca i32, align 4
-; CHECK-NEXT:    [[F1_SROA_GEP:%.*]] = getelementptr inbounds [[STRUCT_S:%.*]], ptr [[PTR:%.*]], i32 0, i32 0
 ; CHECK-NEXT:    br i1 [[C1:%.*]], label [[IF_END:%.*]], label [[FOR_COND:%.*]]
 ; CHECK:       for.cond:
 ; CHECK-NEXT:    br label [[IF_END]]
@@ -775,6 +774,7 @@ define void @PR20822(i1 %c1, i1 %c2, ptr %ptr) {
 ; CHECK-NEXT:    [[TMP0:%.*]] = phi i32 [ poison, [[ENTRY:%.*]] ], [ poison, [[FOR_COND]] ]
 ; CHECK-NEXT:    br i1 [[C2:%.*]], label [[IF_THEN5:%.*]], label [[IF_THEN2:%.*]]
 ; CHECK:       if.then2:
+; CHECK-NEXT:    [[F1_SROA_GEP:%.*]] = getelementptr inbounds [[STRUCT_S:%.*]], ptr [[PTR:%.*]], i32 0, i32 0
 ; CHECK-NEXT:    br label [[IF_THEN5]]
 ; CHECK:       if.then5:
 ; CHECK-NEXT:    [[F1_SROA_PHI:%.*]] = phi ptr [ [[F1_SROA_GEP]], [[IF_THEN2]] ], [ [[F_SROA_0]], [[IF_END]] ]
diff --git a/llvm/test/Transforms/SROA/phi-gep.ll b/llvm/test/Transforms/SROA/phi-gep.ll
index 776624c0798cf..c4a633b42b6c0 100644
--- a/llvm/test/Transforms/SROA/phi-gep.ll
+++ b/llvm/test/Transforms/SROA/phi-gep.ll
@@ -644,9 +644,9 @@ define i1 @test_phi_mem2reg_entry_block_alloca_not_at_beginning(i1 %arg) {
 ; CHECK-NEXT:    call void @f()
 ; CHECK-NEXT:    [[ALLOCA:%.*]] = alloca i64, align 8
 ; CHECK-NEXT:    [[PHI_SROA_GEP:%.*]] = getelementptr i64, ptr [[ALLOCA]], i64 1
-; CHECK-NEXT:    [[PHI_SROA_GEP1:%.*]] = getelementptr i64, ptr [[ALLOCA]], i64 2
 ; CHECK-NEXT:    br i1 [[ARG:%.*]], label [[BB2:%.*]], label [[BB3:%.*]]
 ; CHECK:       bb2:
+; CHECK-NEXT:    [[PHI_SROA_GEP1:%.*]] = getelementptr i64, ptr [[ALLOCA]], i64 2
 ; CHECK-NEXT:    br label [[BB3]]
 ; CHECK:       bb3:
 ; CHECK-NEXT:    [[PHI_SROA_PHI:%.*]] = phi ptr [ [[PHI_SROA_GEP]], [[BB:%.*]] ], [ [[PHI_SROA_GEP1]], [[BB2]] ]
@@ -706,9 +706,9 @@ define i64 @test_unfold_phi_duplicate_phi_entry(ptr %arg, i8 %arg1, i1 %arg2) {
 ; CHECK-LABEL: @test_unfold_phi_duplicate_phi_entry(
 ; CHECK-NEXT:  bb:
 ; CHECK-NEXT:    [[ALLOCA_SROA_0:%.*]] = alloca i64, align 8
-; CHECK-NEXT:    [[PHI_SROA_GEP:%.*]] = getelementptr i64, ptr [[ARG:%.*]], i64 1
 ; CHECK-NEXT:    br i1 [[ARG2:%.*]], label [[BB5:%.*]], label [[BB3:%.*]]
 ; CHECK:       bb3:
+; CHECK-NEXT:    [[PHI_SROA_GEP:%.*]] = getelementptr i64, ptr [[ARG:%.*]], i64 1
 ; CHECK-NEXT:    switch i8 [[ARG1:%.*]], label [[BB4:%.*]] [
 ; CHECK-NEXT:      i8 0, label [[BB5]]
 ; CHECK-NEXT:      i8 1, label [[BB5]]