[CGP] In optimizeMemoryInst handle select similarly to phi

[john.brawn]

Summary

This lets us optimize away selects that perform the same address computation in two different ways and is also the first step towards being able to handle selects between two different, but compatible, address computations.

[davide]

shouldn't GVN handle this case?

[skatkov]

lib/CodeGen/CodeGenPrepare.cpp

4462 ↗

(On Diff #116568)

Please update at least this comment. Now the variable PhiSeen means not only Phi but also Select.

[john.brawn]

Rename PhiSeen to PhiOrSelectSeen and adjust a comment.

[john.brawn]

In D38242#880428, @davide wrote:

shouldn't GVN handle this case?

It's quite possible that we never actually see this kind of input, but this patch is just part of D38133 that's been split out to simplify things.

[skatkov]

LGTM

[skatkov]

lib/CodeGen/CodeGenPrepare.cpp

4396 ↗

(On Diff #116659)

Please consider separation PhiSeen and SelectSeen, we probably will need it in future.
From other side we can do it later if we really need it.
So it is up to you.

llvm/trunk/lib/CodeGen/CodeGenPrepare.cpp

bool CodeGenPrepare::optimizeMemoryInst(Instruction *MemoryInst, Value *Addr,
   Value *Repl = Addr;
 
   // Try to collapse single-value PHI nodes.  This is necessary to undo
   // unprofitable PRE transformations.
   SmallVector<Value*, 8> worklist;
   SmallPtrSet<Value*, 16> Visited;
   worklist.push_back(Addr);
 
-  // Use a worklist to iteratively look through PHI nodes, and ensure that
-  // the addressing mode obtained from the non-PHI roots of the graph
-  // are equivalent.
+  // Use a worklist to iteratively look through PHI and select nodes, and
+  // ensure that the addressing mode obtained from the non-PHI/select roots of
+  // the graph are equivalent.
   bool AddrModeFound = false;
-  bool PhiSeen = false;
+  bool PhiOrSelectSeen = false;
   SmallVector<Instruction*, 16> AddrModeInsts;
   ExtAddrMode AddrMode;
   TypePromotionTransaction TPT(RemovedInsts);
   TypePromotionTransaction::ConstRestorationPt LastKnownGood =
       TPT.getRestorationPoint();
   while (!worklist.empty()) {
     Value *V = worklist.back();
     worklist.pop_back();
while (!worklist.empty()) {
     // as Phi. So we can safely sink address computation to memory instruction.
     if (!Visited.insert(V).second)
       continue;
 
     // For a PHI node, push all of its incoming values.
     if (PHINode *P = dyn_cast<PHINode>(V)) {
       for (Value *IncValue : P->incoming_values())
         worklist.push_back(IncValue);
-      PhiSeen = true;
+      PhiOrSelectSeen = true;
+      continue;
+    }
+    // Similar for select.
+    if (SelectInst *SI = dyn_cast<SelectInst>(V)) {
+      worklist.push_back(SI->getFalseValue());
+      worklist.push_back(SI->getTrueValue());
+      PhiOrSelectSeen = true;
       continue;
     }
 
     // For non-PHIs, determine the addressing mode being computed.  Note that
     // the result may differ depending on what other uses our candidate
     // addressing instructions might have.
     AddrModeInsts.clear();
     ExtAddrMode NewAddrMode = AddressingModeMatcher::Match(
bool CodeGenPrepare::optimizeMemoryInst(Instruction *MemoryInst, Value *Addr,
   // ones were determined, bail out now.
   if (!AddrModeFound) {
     TPT.rollback(LastKnownGood);
     return false;
   }
   TPT.commit();
 
   // If all the instructions matched are already in this BB, don't do anything.
-  // If we saw Phi node then it is not local definitely.
-  if (!PhiSeen && none_of(AddrModeInsts, [&](Value *V) {
+  // If we saw a Phi node then it is not local definitely, and if we saw a select
+  // then we want to push the address calculation past it even if it's already
+  // in this BB.
+  if (!PhiOrSelectSeen && none_of(AddrModeInsts, [&](Value *V) {
         return IsNonLocalValue(V, MemoryInst->getParent());
                   })) {
     DEBUG(dbgs() << "CGP: Found      local addrmode: " << AddrMode << "\n");
     return false;
   }
 
   // Insert this computation right after this user.  Since our caller is
   // scanning from the top of the BB to the bottom, reuse of the expr are

llvm/trunk/test/Transforms/CodeGenPrepare/X86/sink-addrmode.ll

; CHECK: getelementptr i8, {{.+}} 40
   call void @foo(i32 %v)
   %iv.inc = add i32 %iv, 1
   %cmp = icmp slt i32 %iv.inc, 1000
   br i1 %cmp, label %header, label %exit
 
 exit:
   ret void
 }
+
+; Make sure we can eliminate a select when both arguments perform equivalent
+; address computation.
+define void @test10(i1 %cond, i64* %base) {
+; CHECK-LABEL: @test10
+; CHECK: getelementptr i8, {{.+}} 40
+; CHECK-NOT: select
+entry:
+  %gep1 = getelementptr inbounds i64, i64* %base, i64 5
+  %gep1.casted = bitcast i64* %gep1 to i32*
+  %base.casted = bitcast i64* %base to i32*
+  %gep2 = getelementptr inbounds i32, i32* %base.casted, i64 10
+  %casted.merged = select i1 %cond, i32* %gep1.casted, i32* %gep2
+  %v = load i32, i32* %casted.merged, align 4
+  call void @foo(i32 %v)
+  ret void
+}