[InstCombine] Handle PHI nodes in PtrReplacer

[gandhi21299]

Summary

This patch adds on to the functionality implemented
in rG42ab5dc5a5dd6c79476104bdc921afa2a18559cf,
where PHI nodes are supported in the use-def traversal
algorithm to determine if an alloca ever overwritten
in addition to a memmove/memcpy. This patch implements
the support needed by the PointerReplacer to collect
all (indirect) users of the alloca in the case where a PHI
is involved. Finally, a new PHI is defined in the replace
method which takes in replaced incoming values and
updates the WorkMap accordingly.

[nikic]

llvm/test/Transforms/InstCombine/replace-alloca-phi.ll

1

Use update_test_checks.py.

69

This test is too complicated. Please reduce it to the minimum needed to show the new transform.

There also needs to be a negative test where there is an unsupported use of the phi.

[gandhi21299]

• fixed posion value returns
• fixed infinite loop cases by inserting a false check in PointerReplacer::collectUsers(Inst)
• simplified test case and added negative tests for non-simple load and no memcpy
• ran update_test_check.py

[gandhi21299]

• combine two if statements due to common blocks

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

301

Why do we need to add the phi incoming values to the worklist? Can't this use the same code as bitcast/GEP?

Also, I think it would be good to add a test involving a loop phi, because I think as written, this code might go into infinite recursion in that case.

[gandhi21299]

• removed #include "llvm/ADT/STLExtras.h"

[gandhi21299]

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

301

I suppose the incoming values would have already been inserted prior to collecting users of the PHI. Thanks!

A loop test sounds like a great idea!

[gandhi21299]

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

301

@nikic Never mind, the incoming values are missing in the WorkMap when PHIs are being replaced replace if I skip inserting them.

[gandhi21299]

• adding a loop test

[gandhi21299]

ping

[gandhi21299]

• insert an assertion in the replace() method

[gandhi21299]

ping

[nikic]

Just took another look at this. I believe the change to isOnlyCopiedFromConstantMemory() is fine, but the changes to the PointerReplacer are not. Consider these test cases:

@g1 = constant [32 x i8] zeroinitializer
@g2 = addrspace(1) constant [32 x i8] zeroinitializer

define i32 @test1(i1 %c, ptr %ptr) {
entry:
  %a = alloca [32 x i8]
  call void @llvm.memcpy.p0.p0.i64(ptr %a, ptr @g1, i64 32, i1 false)
  br i1 %c, label %if, label %join

if:
  br label %join

join:
  %phi = phi ptr [ %a, %if ], [ %ptr, %entry ]
  %v = load i32, ptr %phi
  ret i32 %v
}

define i32 @test2(i1 %c, ptr %ptr) {
entry:
  %a = alloca [32 x i8]
  call void @llvm.memcpy.p0.p1.i64(ptr %a, ptr addrspace(1) @g2, i64 32, i1 false)
  br i1 %c, label %if, label %join

if:
  br label %join

join:
  %phi = phi ptr [ %a, %if ], [ %ptr, %entry ]
  %v = load i32, ptr %phi
  ret i32 %v
}

declare void @llvm.memcpy.p0.p0.i64(ptr, ptr, i64, i1)
declare void @llvm.memcpy.p0.p1.i64(ptr, ptr addrspace(1), i64, i1)

The first one works, the second one (using PointerReplacer) will cause an assertion failure. The problem is that once we start to handle phis (or selects), additional operands become involved that might not be rooted in the alloca at all. To handle such cases, we need to make sure that all the phi argument are based on the alloca.

I'd personally recommend to split this patch in two parts. First do only the isOnlyCopiedFromConstantMemory() change (which is trivially correct), and then a separate one to deal with the more complex case of address space conversions.

While not strictly related to the patch, we probably need a limit on number of visited users in isOnlyCopiedFromConstantMemory() -- it's a pre-existing problem, but once we start looking through phis, it becomes more likely that we start looking through large use graphs.

llvm/test/Transforms/InstCombine/replace-alloca-phi.ll

2

Please use opaque pointers for new tests ([32 x i8] addrspace(4)* -> ptr addrspace(4) etc).

[nikic]

I believe the change to isOnlyCopiedFromConstantMemory() is fine

I think I spoke too soon ... I can't test this right now, but probably this minor variant is going to be miscompiled:

define i32 @test1(i1 %c, ptr %ptr) {
entry:
  %a = alloca [32 x i8]
  br i1 %c, label %if, label %join

if:
  br label %join

join:
  %phi = phi ptr [ %a, %if ], [ %ptr, %entry ]
  call void @llvm.memcpy.p0.p0.i64(ptr %phi, ptr @g1, i64 32, i1 false)
  %v = load i32, ptr %phi
  ret i32 %v
}

While we can replace uses of %a with @g1, dropping the memcpy would be incorrect, as it would also drop the potential write to %ptr, which is not covered by the replacement.

[gandhi21299]

This comment was removed by gandhi21299.

[gandhi21299]

The minor variant compiles correctly actually, no transformation applies so the code stays the same.

[nikic]

In D136201#3897529, @gandhi21299 wrote:

The minor variant compiles correctly actually, no transformation applies so the code stays the same.

I just tried this, and the memcpy does seem to get dropped:

@g1 = constant [32 x i8] zeroinitializer

define i32 @test1(i1 %c, ptr %ptr) {
entry:
  %a = alloca [32 x i8]
  br i1 %c, label %if, label %join

if:
  br label %join

join:
  %phi = phi ptr [ %a, %if ], [ %ptr, %entry ]
  call void @llvm.memcpy.p0.p0.i64(ptr %phi, ptr @g1, i64 32, i1 false)
  %v = load i32, ptr %phi
  ret i32 %v
}
 
declare void @llvm.memcpy.p0.p0.i64(ptr, ptr, i64, i1)

; RUN: opt -S -instcombine < %s

@g1 = constant [32 x i8] zeroinitializer

define i32 @test1(i1 %c, ptr %ptr) {
entry:
  br i1 %c, label %if, label %join

if:                                               ; preds = %entry
  br label %join

join:                                             ; preds = %if, %entry
  %phi = phi ptr [ @g1, %if ], [ %ptr, %entry ]
  %v = load i32, ptr %phi, align 4
  ret i32 %v
}

; Function Attrs: argmemonly nocallback nofree nounwind willreturn
declare void @llvm.memcpy.p0.p0.i64(ptr noalias nocapture writeonly, ptr noalias nocapture readonly, i64, i1 immarg) #0

[arsenm]

llvm/test/Transforms/InstCombine/replace-alloca-phi.ll

3

Drop the triple, this just needs target datalayout ="A5". Otherwise this will fail without amdgpu built

3

Probably should do "p5-32-A5"

[gandhi21299]

@nikic Since alloca instructions are iterated more than once, shouldn't the Worklist under PointerReplacer remove instructions that have already been replaced (new instruction inserted without removing the old instruction)? The first function goes through a strange transformation when opaque pointers are used, where the replace(Instruction&) function is called on a GEP instruction which was already replaced in a previous instruction. This somehow causes the removal of the new GEP instruction.

[nikic]

In D136201#3900195, @gandhi21299 wrote:

@nikic Since alloca instructions are iterated more than once, shouldn't the Worklist under PointerReplacer remove instructions that have already been replaced (new instruction inserted without removing the old instruction)? The first function goes through a strange transformation when opaque pointers are used, where the replace(Instruction&) function is called on a GEP instruction which was already replaced in a previous instruction. This somehow causes the removal of the new GEP instruction.

Not sure I follow, which IR function are you referring to here? Generally InstCombine will automatically clean up dead (side-effect-free) instructions, but it may be possible that this needs to be done explicitly once phis are involved (as they might no longer be trivially dead if there is a cycle).

[gandhi21299]

llvm/test/Transforms/InstCombine/replace-alloca-phi.ll

3

Is that a complete specification? I am getting the following error:

error: Missing size specification for pointer in datalayout string

[jmmartinez]

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

82

Cosmetic review: What do you think about using the varadic template version of isa ? You could collapse this into isa<BitCastInst, AddrSpaceCastInst, PHINode>(I)

[arsenm]

llvm/test/Transforms/InstCombine/replace-alloca-phi.ll

3

target datalayout = "p5:32:32-A5"

[gandhi21299]

• In case one of the incoming values being replaced belongs to a different address space with respect to other incoming values, cast them to the target address space.
• Included tests suggested by @nikic
• Minor nits suggested by reviewers

[gandhi21299]

• removed the test backup file I added accidentally

[gandhi21299]

llvm/test/Transforms/InstCombine/replace-alloca-phi.ll

3

isOnlyCopiedFromConstantMemory seems to depend on AMDGPUAliasAnalysis results. It returns false for the tests I added when I replace the triple string with the datalayout you suggested.

cc @bcahoon

[gandhi21299]

@nikic Would you still prefer for this patch to be split?

[gandhi21299]

• removed Copy argument from PointerReplacer::collectUsers(...)

[gandhi21299]

• removed included headers

[gandhi21299]

• Eliminated noalias and readonly attributes for parameters as they are irrelevant to this patch

[nikic]

I just tested the updated patch, and this case still gets miscompiled:

@g1 = constant [32 x i8] zeroinitializer

define i32 @test_memcpy_after_phi(i1 %c, ptr %ptr) {
entry:
  %a = alloca [32 x i8]
  br i1 %c, label %if, label %join

if:
  br label %join

join:
  %phi = phi ptr [ %a, %if ], [ %ptr, %entry ]
  call void @llvm.memcpy.p0.p0.i64(ptr %phi, ptr @g1, i64 32, i1 false)
  %v = load i32, ptr %phi
  ret i32 %v
}
 
declare void @llvm.memcpy.p0.p0.i64(ptr, ptr, i64, i1)

The memcpy gets removed, even though it might be writing to the %ptr argument rather than the %a alloca.

llvm/test/Transforms/InstCombine/replace-alloca-phi.ll

3

You should be able to avoid making this AMDGPU specific by copying from a constant global, instead of copying from a special AMDGPU address space.

[gandhi21299]

@nikic Great catch! The memcpy needs to be preserved in this case. I suppose isOnlyCopiedFromConstantMemory(..) should return false when this case is hit.

[gandhi21299]

• When walking over a memtransfer inst in isOnlyCopiedFromConstantMemory(..), check if the destination of the instruction is a PHI node and return false if so.
• Added the @nikic 's test and another loop test

[gandhi21299]

• replaced target triple with target datalayout for generalizing across targets.

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

131

This is not sufficient: E.g. consider doing the memcpy to an addrspacecast of the phi or similar. I think what we'd need instead is that set IsOffset=true for phi nodes, so we'll recursively reject memcpys into it. Of course we would need a different name than IsOffset in that case. It would more generally indicate that the pointer is not equivalent to the base.

[gandhi21299]

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

131

@nikic What do you mean by "need a different name"? I could simply emplace_back a pair of the instruction and a true to ValuesToInspect vector right?

[gandhi21299]

ping

[gandhi21299]

If an alloca has a PHI use which happens to be an ancestor of a memtransfer
instruction, then we cannot determine at compile time whether the result of the
PHI is the alloca or an alternative incoming value. As a result, either the
source or the destination is indeterminate. I added a third variable to
ValuesToInspect, to track whether we have already seen a PHI.

[arsenm]

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

65–66

Use the new C++17 tuple syntax,

381–382

In general InstCombine can't introduce a new address space cast when one didn't exist in the first place

390

Ditto

[gandhi21299]

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

381–382

If an incoming value is eligible to be replaced with a pointer from a different addrspace, other incoming values likely do not belong to the addrspace. This makes addrspacecast unavoidable.

[gandhi21299]

Using C++17 syntax when storing value returned from the ValuesToInspect vector.

[arsenm]

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

381–382

Right, so I'm saying that would just make this illegal and you need to check the address spaces match. Is there a situation where we would want this to happen, that also isn't possible after InferAddressSpaces?

[gandhi21299]

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

381–382

The function @addrspace_cast_remove_alloca in the test I have attached addresses this issue.

[gandhi21299]

• Eliminated addrspace casts in the PointerReplacer::replace() function
• If the PHI node under inspection resides in the default addrspace, do not allow replacing a value which copied from a non-default addrspace.

[gandhi21299]

• Applied clang-format

[gandhi21299]

ping

[gandhi21299]

ping

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

76

Why do we care about default address space here? I think what you probably want to check is PHI->getType()->getPointerAddressSpace() != TheCopy->getPointerAddressSpace()? But also, this code assumes that TheCopy is available before the PHI is seen, and I don't think anything guarantees that -- order of uses is unpredictable. I think you'd want to perform the bailout in PointerReplacer::collectUsers() instead? Maybe I'm misunderstanding the purpose here.

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

68

Why do we need the LastPHI variable? I'd have expected that just setting IsOffset=true is equivalent? (With a rename IsOffset -> ForbidCopy or something?)

[gandhi21299]

• Eliminated LastPHI as it had no valuable impact
• Removed the addrspace cast checks in isOnlyCopiedFromConstantMemory(..) as pointed out. Instead, determine if the alloca is directly used as an incoming value in a PHI while collecting users of an alloca in PointerReplacer::collectUsers(..). If the address spaces of the PHI and the destination of the Copy inst differ and there is a direct use of the alloca, an addrspace cast will be required to further allow the PHI for replacement. Since we do not allow for addrspace casts, reject the PHI at that moment.

[gandhi21299]

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

76

I implemented it for the function @addrspace_diff_keep_alloca in the test file. I improved upon the solution by moving the changes over to collectUsers() which does make a lot more sense wrt my previous solution. I hope it seems fine now.

[gandhi21299]

removed dbg statements

[arsenm]

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

131

I think I'm getting lost in the different revisions of the patch; not sure what this comment is referring to anymore. There isn't anything about phi handling right here?

301

Don't like auto for addrspace

[gandhi21299]

Use unsigned instead of auto

[gandhi21299]

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

131

I had a revision previously where I was tracking PHI nodes seperately, as opposed to IsOffset. I replaced it with my changes in PointerReplacer::collectUsers(..).

[gandhi21299]

Allow PHI nodes only if the number of incoming values is less than 5

[gandhi21299]

ping

[gandhi21299]

• PHI's need to be kept track of to avoid inserting them repeatedly into ValuesToInspect.

[gandhi21299]

• Applied clang-format and rebased

[gandhi21299]

• Insert PHI into ValuesToInspect only when it hasn't seen before. Otherwise, continue with the use-def traversal.

[gandhi21299]

ping

[gandhi21299]

@nikic Thanks for committing on my behalf, may I close this revision?

[nikic]

@gandhi21299 I only landed the basic isOnlyCopiedFromConstantMemory() support, but the PointerReplacer part is still open. I'm think the logic there is still not quite correct -- will take another look tomorrow.

[nikic]

Rebased patch over committed changes: https://gist.github.com/nikic/ce69a66fdc0985d1af7a56bf024bc860

With additional tests added (in particular @addrspace_diff_keep_alloca_extra_gep) the patch causes an assertion failure.

You are currently going over phi operands and checking whether one of them is the memcpy dest, but I'm not sure what this check is intended to achieve, semantically. I think it just works by accident on existing test cases.

The problem is that we need a closed graph to perform the pointer replacement (can't have any external inputs) and additionally the worklist also needs to be ordered correctly, so that inputs are always replaced before the using instruction. The former is a hard requirement for this to work, the latter is a requirement of the current algorithm.

If you don't want to change the current replacement approach, then I think a good way to enforce this would be to do the following: When processing a phi node, check whether all phi operands are already in the worklist. If not, don't visit the phi yet (wait until everything is in the worklist). Add the phi to some extra set, and after collectUsers(), check that all phis are in the worklist. This is to handle the case where an input was not part of the graph and the phi is never added to the worklist.

[gandhi21299]

• Revised PointerReplacer::collectUsers(..) as requested.
• Passes all tests under check-llvm-transforms, including replace-alloca-phi.ll

[gandhi21299]

• Rebased and removed an irrelevant comment.

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

321

I don't think you need this check (and as such this whole loop can be dropped).

339

I think this should only be done after the whole collectUsers() phase, currently you do it on each recursive call.

[gandhi21299]

• Eliminated the redundant loop as pointed out by reviewer
• Created a new method that drives the recursive algorithm for collecting users, this was necessary to avoid publicizing the member variables of PointerReplacer.

[gandhi21299]

• clang-format

[nikic]

Logic looks correct to me now. Looks like the test changes got lost though?

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

15

Shouldn't be needed (already used in file).

264

Copy argument no longer needed

287

Could make ValuesToRevisit use Instruction * rather than Value * and save this cast.

314

Can use any_of instead of for here and avoid the ValueInserted variable.

[gandhi21299]

In D136201#4052834, @nikic wrote:

Logic looks correct to me now. Looks like the test changes got lost though?

I rebased the changes and there weren't any test changes in the reduced version of this patch.

[gandhi21299]

• Changes as requested by reviewer, updated tests accordingly

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

301

Drop llvm:: prefix.

308

Missing the insertion into ValuesToRevisit here?

[gandhi21299]

• Replaced calls from collectUsers to collectUsersRecursive, which now generates the code we expect for @loop_phi_remove_alloca and @loop_phi_late_memtransfer_remove_alloca.

[gandhi21299]

• removed #include "STLExtras.h"

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

68

Spurious change.

308–314

The ValuesToRevisit update should happen outside the any_of.

458

Spurious change.

llvm/test/Transforms/InstCombine/replace-alloca-phi.ll

11–20

Broken check line (should be CHECK).

[gandhi21299]

• changes as suggested by reviewer

[nikic]

LGTM

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

299

Spurious change

308

Unneeded Inst capture.

[gandhi21299]

• requested updates

Thanks a bunch for such in-depth review, I will commit this patch.

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

 //===- InstCombineLoadStoreAlloca.cpp -------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the visit functions for load, store and alloca.
 //
 //===----------------------------------------------------------------------===//
 
 #include "InstCombineInternal.h"
 #include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/Loads.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/LLVMContext.h"
isOnlyCopiedFromConstantMemory(AAResults *AA, AllocaInst *V,
   SmallPtrSet<ValueAndIsOffset, 32> Visited;
   Worklist.emplace_back(V, false);
   while (!Worklist.empty()) {
     ValueAndIsOffset Elem = Worklist.pop_back_val();
     if (!Visited.insert(Elem).second)
       continue;
     if (Visited.size() > MaxCopiedFromConstantUsers)
       return false;
 
     const auto [Value, IsOffset] = Elem;
     for (auto &U : Value->uses()) {
       auto *I = cast<Instruction>(U.getUser());
-
       if (auto *LI = dyn_cast<LoadInst>(I)) {
         // Ignore non-volatile loads, they are always ok.
         if (!LI->isSimple()) return false;
         continue;
       }
 
       if (isa<PHINode>(I)) {
         // We set IsOffset=true, to forbid the memcpy from occurring after the
         // phi: If one of the phi operands is not based on the alloca, we
         // would incorrectly omit a write.
         Worklist.emplace_back(I, true);
         continue;
       }
       if (isa<BitCastInst>(I) || isa<AddrSpaceCastInst>(I)) {
         // If uses of the bitcast are ok, we are ok.
         Worklist.emplace_back(I, IsOffset);
         continue;
       }
       if (auto *GEP = dyn_cast<GetElementPtrInst>(I)) {
         // If the GEP has all zero indices, it doesn't offset the pointer. If it
         // doesn't, it does.
         Worklist.emplace_back(I, IsOffset || !GEP->hasAllZeroIndices());
for (auto &U : Value->uses()) {
         assert(I->use_empty() && "Lifetime markers have no result to use!");
         ToDelete.push_back(I);
         continue;
       }
 
       // If this is isn't our memcpy/memmove, reject it as something we can't
       // handle.
       MemTransferInst *MI = dyn_cast<MemTransferInst>(I);
       if (!MI)
         return false;
 
       // If the transfer is volatile, reject it.
       if (MI->isVolatile())
         return false;
 
       // If the transfer is using the alloca as a source of the transfer, then
       // ignore it since it is a load (unless the transfer is volatile).

 // instructions, then replaces the old pointer in the load instructions with
 // the new pointer. If during the chasing it sees bitcast or GEP, it will
 // create new bitcast or GEP with the new pointer and use them in the load
 // instruction.
 class PointerReplacer {
 public:
   PointerReplacer(InstCombinerImpl &IC) : IC(IC) {}
 
   bool collectUsers(Instruction &I);
   void replacePointer(Instruction &I, Value *V);
 
 private:
+  bool collectUsersRecursive(Instruction &I);
   void replace(Instruction *I);
   Value *getReplacement(Value *I);
 
+  SmallPtrSet<Instruction *, 32> ValuesToRevisit;
   SmallSetVector<Instruction *, 4> Worklist;
   MapVector<Value *, Value *> WorkMap;
   InstCombinerImpl &IC;
 };
 } // end anonymous namespace
 
 bool PointerReplacer::collectUsers(Instruction &I) {
+  if (!collectUsersRecursive(I))
+    return false;
+
+  // Ensure that all outstanding (indirect) users of I
+  // are inserted into the Worklist. Return false
+  // otherwise.
+  for (auto *Inst : ValuesToRevisit)
+    if (!Worklist.contains(Inst))
+      return false;
+  return true;
+}
+
+bool PointerReplacer::collectUsersRecursive(Instruction &I) {
   for (auto *U : I.users()) {
     auto *Inst = cast<Instruction>(&*U);
     if (auto *Load = dyn_cast<LoadInst>(Inst)) {
       if (Load->isVolatile())
         return false;
       Worklist.insert(Load);
-    } else if (isa<GetElementPtrInst>(Inst) || isa<BitCastInst>(Inst)) {
+    } else if (auto *PHI = dyn_cast<PHINode>(Inst)) {
+      // All incoming values must be instructions for replacability
+      if (any_of(PHI->incoming_values(),
+                 [](Value *V) { return !isa<Instruction>(V); }))
+        return false;
+
+      // If at least one incoming value of the PHI is not in Worklist,
+      // store the PHI for revisiting and skip this iteration of the
+      // loop.
+      if (any_of(PHI->incoming_values(), [this, Inst](Value *V) {
+            if (!Worklist.contains(cast<Instruction>(V))) {
+              ValuesToRevisit.insert(Inst);
+              return true;
+            }
+            return false;
+          }))
+        continue;
+
+      Worklist.insert(PHI);
+      if (!collectUsersRecursive(*PHI))
+        return false;
+    } else if (isa<GetElementPtrInst, BitCastInst>(Inst)) {
       Worklist.insert(Inst);
-      if (!collectUsers(*Inst))
+      if (!collectUsersRecursive(*Inst))
         return false;
     } else if (auto *MI = dyn_cast<MemTransferInst>(Inst)) {
       if (MI->isVolatile())
         return false;
       Worklist.insert(Inst);
     } else if (Inst->isLifetimeStartOrEnd()) {
       continue;
     } else {
       LLVM_DEBUG(dbgs() << "Cannot handle pointer user: " << *U << '\n');
       return false;
     }
   }
-
   return true;
 }
 
 Value *PointerReplacer::getReplacement(Value *V) { return WorkMap.lookup(V); }
 
 void PointerReplacer::replace(Instruction *I) {
   if (getReplacement(I))
     return;
 
   if (auto *LT = dyn_cast<LoadInst>(I)) {
     auto *V = getReplacement(LT->getPointerOperand());
     assert(V && "Operand not replaced");
     auto *NewI = new LoadInst(LT->getType(), V, "", LT->isVolatile(),
                               LT->getAlign(), LT->getOrdering(),
                               LT->getSyncScopeID());
     NewI->takeName(LT);
     copyMetadataForLoad(*NewI, *LT);
 
     IC.InsertNewInstWith(NewI, *LT);
     IC.replaceInstUsesWith(*LT, NewI);
     WorkMap[LT] = NewI;
+  } else if (auto *PHI = dyn_cast<PHINode>(I)) {
+    Type *NewTy = getReplacement(PHI->getIncomingValue(0))->getType();
+    auto *NewPHI = PHINode::Create(NewTy, PHI->getNumIncomingValues(),
+                                   PHI->getName(), PHI);
+    for (unsigned int I = 0; I < PHI->getNumIncomingValues(); ++I)
+      NewPHI->addIncoming(getReplacement(PHI->getIncomingValue(I)),
+                          PHI->getIncomingBlock(I));
+    WorkMap[PHI] = NewPHI;
   } else if (auto *GEP = dyn_cast<GetElementPtrInst>(I)) {
     auto *V = getReplacement(GEP->getPointerOperand());
     assert(V && "Operand not replaced");
     SmallVector<Value *, 8> Indices;
     Indices.append(GEP->idx_begin(), GEP->idx_end());
     auto *NewI =
         GetElementPtrInst::Create(GEP->getSourceElementType(), V, Indices);
     IC.InsertNewInstWith(NewI, *GEP);
     NewI->takeName(GEP);
     WorkMap[GEP] = NewI;
   } else if (auto *BC = dyn_cast<BitCastInst>(I)) {
     auto *V = getReplacement(BC->getOperand(0));
     assert(V && "Operand not replaced");
     auto *NewT = PointerType::getWithSamePointeeType(
         cast<PointerType>(BC->getType()),
         V->getType()->getPointerAddressSpace());
     auto *NewI = new BitCastInst(V, NewT);
     IC.InsertNewInstWith(NewI, *BC);
     NewI->takeName(BC);
     WorkMap[BC] = NewI;
   } else if (auto *MemCpy = dyn_cast<MemTransferInst>(I)) {
     auto *SrcV = getReplacement(MemCpy->getRawSource());
     // The pointer may appear in the destination of a copy, but we don't want to
     // replace it.
     if (!SrcV) {
       assert(getReplacement(MemCpy->getRawDest()) &&
              "destination not in replace list");
       return;
     }
 
     IC.Builder.SetInsertPoint(MemCpy);
     auto *NewI = IC.Builder.CreateMemTransferInst(
         MemCpy->getIntrinsicID(), MemCpy->getRawDest(), MemCpy->getDestAlign(),
         SrcV, MemCpy->getSourceAlign(), MemCpy->getLength(),
         MemCpy->isVolatile());
if (AllocaAlign <= SourceAlign &&
         ++NumGlobalCopies;
         return NewI;
       }
 
       PointerReplacer PtrReplacer(*this);
       if (PtrReplacer.collectUsers(AI)) {
         for (Instruction *Delete : ToDelete)
           eraseInstFromFunction(*Delete);
-
         Value *Cast = Builder.CreateBitCast(TheSrc, DestTy);
         PtrReplacer.replacePointer(AI, Cast);
         ++NumGlobalCopies;
       }
     }
   }
 
   // At last, use the generic allocation site handler to aggressively remove

llvm/test/Transforms/InstCombine/replace-alloca-phi.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -passes=instcombine -S -o - %s | FileCheck %s
 
 target datalayout="p5:32:32-A5"
 
 @g1 = constant [32 x i8] zeroinitializer
 @g2 = addrspace(1) constant [32 x i8] zeroinitializer
 
 define i8 @remove_alloca_use_arg(i1 %cond) {
 ; CHECK-LABEL: @remove_alloca_use_arg(
-; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[ALLOCA:%.*]] = alloca [32 x i8], align 4, addrspace(1)
-; CHECK-NEXT:    call void @llvm.memcpy.p1.p0.i64(ptr addrspace(1) noundef align 4 dereferenceable(256) [[ALLOCA]], ptr noundef nonnull align 16 dereferenceable(256) @g1, i64 256, i1 false)
-; CHECK-NEXT:    br i1 [[COND:%.*]], label [[IF:%.*]], label [[ELSE:%.*]]
-; CHECK:       if:
-; CHECK-NEXT:    [[VAL_IF:%.*]] = getelementptr inbounds [32 x i8], ptr addrspace(1) [[ALLOCA]], i64 0, i64 2
-; CHECK-NEXT:    br label [[SINK:%.*]]
-; CHECK:       else:
-; CHECK-NEXT:    [[VAL_ELSE:%.*]] = getelementptr inbounds [32 x i8], ptr addrspace(1) [[ALLOCA]], i64 0, i64 1
-; CHECK-NEXT:    br label [[SINK]]
-; CHECK:       sink:
-; CHECK-NEXT:    [[PTR:%.*]] = phi ptr addrspace(1) [ [[VAL_IF]], [[IF]] ], [ [[VAL_ELSE]], [[ELSE]] ]
-; CHECK-NEXT:    [[LOAD:%.*]] = load i8, ptr addrspace(1) [[PTR]], align 1
-; CHECK-NEXT:    ret i8 [[LOAD]]
+; CHECk-NEXT:  entry:
+; CHECk-NEXT:    br i1 [[COND:%.*]], label [[IF:%.*]], label [[ELSE:%.*]]
+; CHECk:       if:                                               ; preds = %entry
+; CHECk-NEXT:    br label [[SINK:%.*]]
+; CHECk:       else:                                             ; preds = %entry
+; CHECk-NEXT:    br label [[SINK:%.*]]
+; CHECk:       sink:                                             ; preds = %else, %if
+; CHECk-NEXT:    [[PTR:%.*]] = phi ptr [ getelementptr inbounds ([32 x i8], ptr [[G1:@.*]], i64 0, i64 2), [[IF]] ], [ getelementptr inbounds ([32 x i8], ptr [[G1]], i64 0, i64 1), [[ELSE]] ]
+; CHECk-NEXT:    [[LOAD:%.*]] = load i8, ptr [[PTR]], align 1
+; CHECk-NEXT:    ret i8 [[LOAD]]
 ;
 entry:
   %alloca = alloca [32 x i8], align 4, addrspace(1)
   call void @llvm.memcpy.p1.p0.i64(ptr addrspace(1) %alloca, ptr @g1, i64 256, i1 false)
   br i1 %cond, label %if, label %else
 
 if:
   %val.if = getelementptr inbounds [32 x i8], ptr addrspace(1) %alloca, i32 0, i32 2
entry:
   br i1 %cond, label %if, label %else
 
 if:
   %val.if = getelementptr inbounds [32 x i8], ptr addrspace(1) %alloca, i32 0, i32 1
   br label %sink
 
 else:
   %val.else = getelementptr inbounds [32 x i8], ptr addrspace(1) %alloca, i32 0, i32 2
   br label %sink
 
 sink:
   %ptr = phi ptr addrspace(1) [ %val.if, %if ], [ %val.else, %else ]
   %load = load volatile i8, ptr addrspace(1) %ptr
   ret i8 %load
 }
 
 
sink:
   %ptr = phi ptr addrspace(1) [ %val.if, %if ], [ %val.else, %else ]
   %load = load volatile i8, ptr addrspace(1) %ptr
   ret i8 %load
 }
 
 define i8 @loop_phi_remove_alloca(i1 %cond) {
 ; CHECK-LABEL: @loop_phi_remove_alloca(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[ALLOCA:%.*]] = alloca [32 x i8], align 4, addrspace(1)
-; CHECK-NEXT:    call void @llvm.memcpy.p1.p0.i64(ptr addrspace(1) noundef align 4 dereferenceable(256) [[ALLOCA]], ptr noundef nonnull align 16 dereferenceable(256) @g1, i64 256, i1 false)
-; CHECK-NEXT:    [[VAL1:%.*]] = getelementptr inbounds [32 x i8], ptr addrspace(1) [[ALLOCA]], i64 0, i64 1
 ; CHECK-NEXT:    br label [[BB_0:%.*]]
 ; CHECK:       bb.0:
-; CHECK-NEXT:    [[PTR:%.*]] = phi ptr addrspace(1) [ [[VAL1]], [[ENTRY:%.*]] ], [ [[VAL2:%.*]], [[BB_1:%.*]] ]
+; CHECK-NEXT:    [[PTR:%.*]] = phi ptr [ getelementptr inbounds ([32 x i8], ptr [[G1:@.*]], i64 0, i64 1), [[ENTRY:%.*]] ], [ getelementptr inbounds ([32 x i8], ptr [[G1]], i64 0, i64 2), [[BB_1:%.*]] ]
 ; CHECK-NEXT:    br i1 [[COND:%.*]], label [[BB_1]], label [[EXIT:%.*]]
 ; CHECK:       bb.1:
-; CHECK-NEXT:    [[VAL2]] = getelementptr inbounds [32 x i8], ptr addrspace(1) [[ALLOCA]], i64 0, i64 2
 ; CHECK-NEXT:    br label [[BB_0]]
 ; CHECK:       exit:
-; CHECK-NEXT:    [[LOAD:%.*]] = load i8, ptr addrspace(1) [[PTR]], align 1
+; CHECK-NEXT:    [[LOAD:%.*]] = load i8, ptr [[PTR]], align 1
 ; CHECK-NEXT:    ret i8 [[LOAD]]
 ;
 entry:
   %alloca = alloca [32 x i8], align 4, addrspace(1)
   call void @llvm.memcpy.p1.p0.i64(ptr addrspace(1) %alloca, ptr @g1, i64 256, i1 false)
   %val1 = getelementptr inbounds [32 x i8], ptr addrspace(1) %alloca, i32 0, i32 1
   br label %bb.0
 
if:
   br label %join
 
 join:
   %phi = phi ptr [ %alloca, %if ], [ %ptr, %entry ]
   %v = load i32, ptr %phi
   ret i32 %v
 }
 
-define i8 @loop_phi_late_memtransfer(i1 %cond) {
-; CHECK-LABEL: @loop_phi_late_memtransfer(
+define i8 @loop_phi_late_memtransfer_remove_alloca(i1 %cond) {
+; CHECK-LABEL: @loop_phi_late_memtransfer_remove_alloca(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[ALLOCA:%.*]] = alloca [32 x i8], align 4, addrspace(1)
-; CHECK-NEXT:    [[VAL1:%.*]] = getelementptr inbounds [32 x i8], ptr addrspace(1) [[ALLOCA]], i64 0, i64 1
 ; CHECK-NEXT:    br label [[BB_0:%.*]]
 ; CHECK:       bb.0:
-; CHECK-NEXT:    [[PTR:%.*]] = phi ptr addrspace(1) [ [[VAL1]], [[ENTRY:%.*]] ], [ [[VAL2:%.*]], [[BB_1:%.*]] ]
+; CHECK-NEXT:    [[PTR:%.*]] = phi ptr [ getelementptr inbounds ([32 x i8], ptr [[G1:@.*]], i64 0, i64 1), [[ENTRY:%.*]] ], [ getelementptr inbounds ([32 x i8], ptr [[G1]], i64 0, i64 2), [[BB_1:%.*]] ]
 ; CHECK-NEXT:    br i1 [[COND:%.*]], label [[BB_1]], label [[EXIT:%.*]]
 ; CHECK:       bb.1:
-; CHECK-NEXT:    [[VAL2]] = getelementptr inbounds [32 x i8], ptr addrspace(1) [[ALLOCA]], i64 0, i64 2
-; CHECK-NEXT:    call void @llvm.memcpy.p1.p0.i64(ptr addrspace(1) noundef align 4 dereferenceable(256) [[ALLOCA]], ptr noundef nonnull align 16 dereferenceable(256) @g1, i64 256, i1 false)
 ; CHECK-NEXT:    br label [[BB_0]]
 ; CHECK:       exit:
-; CHECK-NEXT:    [[LOAD:%.*]] = load i8, ptr addrspace(1) [[PTR]], align 1
+; CHECK-NEXT:    [[LOAD:%.*]] = load i8, ptr [[PTR]], align 1
 ; CHECK-NEXT:    ret i8 [[LOAD]]
 ;
 entry:
   %alloca = alloca [32 x i8], align 4, addrspace(1)
   %val1 = getelementptr inbounds [32 x i8], ptr addrspace(1) %alloca, i32 0, i32 1
   br label %bb.0
 
 bb.0: