CVP: Improve handling of Selects used as incoming PHI values

[dotdash]

Summary

If the branch that leads to the PHI node and the Select instruction
depend on correlated conditions, we might be able to directly use the
corresponding value from the Select instruction as the incoming value
for the PHI node, allowing later removal of the select instruction.

[dotdash]

Updated the patch to only apply the optimization to selects with a scalar condition.

[dotdash]

FWIW, this helps with an optimization problem we hit with the rust compiler:

https://github.com/rust-lang/rust/issues/24489

[reames]

Seems like a reasonable approach. A couple of small comments, but once those are addressed, should be easy to sign off on.

lib/Transforms/Scalar/CorrelatedValuePropagation.cpp

115 ↗

(On Diff #23899)

I don't understand why we need the vector type check? Is this a limitation of LVI? If so, it should be pushed inside the LVI function being called. (i.e. returning Unknown for a vector type is perfectly reasonable)

118 ↗

(On Diff #23899)

I wonder whether a wrapper function on LVI which just take an ICmpInst would be more clear here. The comparison of against true is slightly opaque at first.

test/Transforms/CorrelatedValuePropagation/select.ll

8 ↗

(On Diff #23899)

A few more tests would be good. In particular:

• a negative test for the vector case
• a simple case without the loop to express the idea
• a negative test for when the select condition *isn't* implied by the branch test.

[dotdash]

lib/Transforms/Scalar/CorrelatedValuePropagation.cpp

118 ↗

(On Diff #23899)

Actually, I can just use getConstantOnEdge, which is easier to read and better at handling the vector types. Not sure what made me think that I could not use that. With the predicate approach, the false case matches all vectors except for the one with all ones, which is why I had excluded the vector types. I'm currently struggling at coming up with a test for the vector types though.

[dotdash]

Updated the patch to use getConstantOnEdge to improve readability. Added a
simpler test case without a loop and a negative test case. LVI doesn't support
vector types yet, so there's no test for that, but I left a note that support
for vector types may be added later.

[dotdash]

Could somebody have another look at this?

Thanks!

[reames]

LGTM.

llvm/trunk/lib/Transforms/Scalar/CorrelatedValuePropagation.cpp

bool CorrelatedValuePropagation::processPHI(PHINode *P) {
 
   BasicBlock *BB = P->getParent();
   for (unsigned i = 0, e = P->getNumIncomingValues(); i < e; ++i) {
     Value *Incoming = P->getIncomingValue(i);
     if (isa<Constant>(Incoming)) continue;
 
     Value *V = LVI->getConstantOnEdge(Incoming, P->getIncomingBlock(i), BB, P);
 
-    // Look if the incoming value is a select with a constant but LVI tells us
-    // that the incoming value can never be that constant. In that case replace
-    // the incoming value with the other value of the select. This often allows
-    // us to remove the select later.
+    // Look if the incoming value is a select with a scalar condition for which
+    // LVI can tells us the value. In that case replace the incoming value with
+    // the appropriate value of the select. This often allows us to remove the
+    // select later.
     if (!V) {
       SelectInst *SI = dyn_cast<SelectInst>(Incoming);
       if (!SI) continue;
 
+      Value *Condition = SI->getCondition();
+      if (!Condition->getType()->isVectorTy()) {
+        if (Constant *C = LVI->getConstantOnEdge(Condition, P->getIncomingBlock(i), BB, P)) {
+          if (C == ConstantInt::getTrue(Condition->getType())) {
+            V = SI->getTrueValue();
+          } else {
+            V = SI->getFalseValue();
+          }
+          // Once LVI learns to handle vector types, we could also add support
+          // for vector type constants that are not all zeroes or all ones.
+        }
+      }
+
+      // Look if the select has a constant but LVI tells us that the incoming
+      // value can never be that constant. In that case replace the incoming
+      // value with the other value of the select. This often allows us to
+      // remove the select later.
+      if (!V) {
         Constant *C = dyn_cast<Constant>(SI->getFalseValue());
         if (!C) continue;
 
         if (LVI->getPredicateOnEdge(ICmpInst::ICMP_EQ, SI, C,
-                                  P->getIncomingBlock(i), BB, P) !=
             LazyValueInfo::False)
           continue;
+        V = SI->getTrueValue();
+      }
 
       DEBUG(dbgs() << "CVP: Threading PHI over " << *SI << '\n');
-      V = SI->getTrueValue();
     }
 
     P->setIncomingValue(i, V);
     Changed = true;
   }
 
   // FIXME: Provide TLI, DT, AT to SimplifyInstruction.
   const DataLayout &DL = BB->getModule()->getDataLayout();

llvm/trunk/test/Transforms/CorrelatedValuePropagation/select.ll

+; RUN: opt < %s -correlated-propagation -S | FileCheck %s
+
+; CHECK-LABEL: @simple(
+define i8 @simple(i1) {
+entry:
+  %s = select i1 %0, i8 0, i8 1
+  br i1 %0, label %then, label %else
+
+then:
+; CHECK: ret i8 0
+  %a = phi i8 [ %s, %entry ]
+  ret i8 %a
+
+else:
+; CHECK: ret i8 1
+  %b = phi i8 [ %s, %entry ]
+  ret i8 %b
+}
+
+; CHECK-LABEL: @loop(
+define void @loop(i32) {
+entry:
+  br label %loop
+
+loop:
+  %idx = phi i32 [ %0, %entry ], [ %sel, %loop ]
+; CHECK: %idx = phi i32 [ %0, %entry ], [ %2, %loop ]
+  %1 = icmp eq i32 %idx, 0
+  %2 = add i32 %idx, -1
+  %sel = select i1 %1, i32 0, i32 %2
+  br i1 %1, label %out, label %loop
+
+out:
+  ret void
+}
+
+; CHECK-LABEL: @not_correlated(
+define i8 @not_correlated(i1, i1) {
+entry:
+  %s = select i1 %0, i8 0, i8 1
+  br i1 %1, label %then, label %else
+
+then:
+; CHECK: ret i8 %s
+  %a = phi i8 [ %s, %entry ]
+  ret i8 %a
+
+else:
+; CHECK: ret i8 %s
+  %b = phi i8 [ %s, %entry ]
+  ret i8 %b
+}
+