[llvm] ddc484b - [InstCombine] Handle select inst when eliminating constant memcpy

[Anshil Gandhi via llvm-commits]

Author: Anshil Gandhi
Date: 2023-01-23T12:02:57-07:00
New Revision: ddc484b8eaa4aaf518cfe38c555495ce22cbf50a

URL: https://github.com/llvm/llvm-project/commit/ddc484b8eaa4aaf518cfe38c555495ce22cbf50a
DIFF: https://github.com/llvm/llvm-project/commit/ddc484b8eaa4aaf518cfe38c555495ce22cbf50a.diff

LOG: [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

Added: 
    

Modified: 
    llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp
    llvm/test/Transforms/InstCombine/ptr-replace-alloca.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp b/llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp
index 71efe0f3391f1..41bc65620ff62 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp
+++ b/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

diff  --git a/llvm/test/Transforms/InstCombine/ptr-replace-alloca.ll b/llvm/test/Transforms/InstCombine/ptr-replace-alloca.ll
index 7821eb1491074..de8fcbb3802ac 100644
--- a/llvm/test/Transforms/InstCombine/ptr-replace-alloca.ll
+++ b/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]