[InstCombine] Handle select inst when eliminating constant memcpy

Allow iterating through SelectInst use of the alloca when
checking if it is only ever overwritten from constant memory.
Recursively determine if the SelectInst is replacable and insert
it into the Worklist if so. Finally, define a new SelectInst to
replace the old one, with both of it's values replaced according
to the WorkMap.

Differential Revision: https://reviews.llvm.org/D136524

Author: gandhi56

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

@@ -72,14 +72,14 @@ isOnlyCopiedFromConstantMemory(AAResults *AA, AllocaInst *V,
         continue;
       }
 
-      if (isa<PHINode>(I)) {
+      if (isa<PHINode, SelectInst>(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 (isa<BitCastInst, AddrSpaceCastInst>(I)) {
         // If uses of the bitcast are ok, we are ok.
         Worklist.emplace_back(I, IsOffset);
         continue;
@@ -320,6 +320,19 @@ bool PointerReplacer::collectUsersRecursive(Instruction &I) {
       Worklist.insert(PHI);
       if (!collectUsersRecursive(*PHI))
         return false;
+    } else if (auto *SI = dyn_cast<SelectInst>(Inst)) {
+      if (!isa<Instruction>(SI->getTrueValue()) ||
+          !isa<Instruction>(SI->getFalseValue()))
+        return false;
+
+      if (!isAvailable(cast<Instruction>(SI->getTrueValue())) ||
+          !isAvailable(cast<Instruction>(SI->getFalseValue()))) {
+        ValuesToRevisit.insert(Inst);
+        continue;
+      }
+      Worklist.insert(SI);
+      if (!collectUsersRecursive(*SI))
+        return false;
     } else if (isa<GetElementPtrInst, BitCastInst>(Inst)) {
       Worklist.insert(Inst);
       if (!collectUsersRecursive(*Inst))
@@ -385,6 +398,13 @@ void PointerReplacer::replace(Instruction *I) {
     IC.InsertNewInstWith(NewI, *BC);
     NewI->takeName(BC);
     WorkMap[BC] = NewI;
+  } else if (auto *SI = dyn_cast<SelectInst>(I)) {
+    auto *NewSI = SelectInst::Create(
+        SI->getCondition(), getReplacement(SI->getTrueValue()),
+        getReplacement(SI->getFalseValue()), SI->getName(), nullptr, SI);
+    IC.InsertNewInstWith(NewSI, *SI);
+    NewSI->takeName(SI);
+    WorkMap[SI] = NewSI;
   } 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

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

@@ -366,9 +366,7 @@ exit:
 define i8 @select_same_addrspace_remove_alloca(i1 %cond, ptr %p) {
 ; CHECK-LABEL: @select_same_addrspace_remove_alloca(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[ALLOCA:%.*]] = alloca [32 x i8], align 1
-; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i64(ptr noundef nonnull align 1 dereferenceable(32) [[ALLOCA]], ptr noundef nonnull align 16 dereferenceable(32) @g1, i64 32, i1 false)
-; CHECK-NEXT:    [[PTR:%.*]] = select i1 [[COND:%.*]], ptr [[ALLOCA]], ptr [[P:%.*]]
+; CHECK-NEXT:    [[PTR:%.*]] = select i1 [[COND:%.*]], ptr @g1, ptr [[P:%.*]]
 ; CHECK-NEXT:    [[LOAD:%.*]] = load i8, ptr [[PTR]], align 1
 ; CHECK-NEXT:    ret i8 [[LOAD]]
 ;
@@ -417,13 +415,8 @@ entry:
 define i8 @select_diff_addrspace_remove_alloca(i1 %cond, ptr %p) {
 ; CHECK-LABEL: @select_diff_addrspace_remove_alloca(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[ALLOCA:%.*]] = alloca [32 x i8], align 1
-; CHECK-NEXT:    call void @llvm.memcpy.p0.p1.i64(ptr noundef nonnull align 1 dereferenceable(32) [[ALLOCA]], ptr addrspace(1) noundef align 16 dereferenceable(32) @g2, i64 32, i1 false)
-; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds [32 x i8], ptr [[ALLOCA]], i64 0, i64 2
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], ptr [[ALLOCA]], ptr [[GEP]]
-; CHECK-NEXT:    [[GEP2:%.*]] = getelementptr inbounds i8, ptr [[SEL]], i64 4
-; CHECK-NEXT:    [[LOAD:%.*]] = load i8, ptr [[GEP2]], align 1
-; CHECK-NEXT:    ret i8 [[LOAD]]
+; CHECK-NOT:     [[ALLOCA:%.*]] = alloca [32 x i8]
+; CHECK-NEXT:    ret i8 0
 ;
 entry:
   %alloca = alloca [32 x i8]