[InstCombine] fold zexts and constants into a phi (PR24766)

[spatel]

Summary

This is one step towards solving PR24766:
https://llvm.org/bugs/show_bug.cgi?id=24766

We were not producing the same IR for these two C functions because the store to the temp bool causes extra zexts:

#include <stdbool.h>

bool switchy(char x1, char x2, char condition) {
   bool conditionMet = false;
   switch (condition) {
   case 0: conditionMet = (x1 == x2); break;
   case 1: conditionMet = (x1 <= x2); break;
   }
   return conditionMet;
}

bool switchy2(char x1, char x2, char condition) {
   switch (condition) {
   case 0: return (x1 == x2);
   case 1: return (x1 <= x2);
   }
  return false;
}

As noted in the code comments, this test case manages to avoid the more general existing phi optimizations where there are only 2 phi inputs or where there are no constant phi args mixed in with the casts ops. It seems like a corner case, but if we don't catch it, then I don't think we can get SimplifyCFG to further optimize towards the canonical form for this function shown in the bug report.

[sanjoy]

lib/Transforms/InstCombine/InstCombinePHI.cpp

399

When will Cast->getOpcode() != FirstCast->getOpcode()? Doesn't Cast->getOpcode() always have to be Instruction::ZExt?

416

I'd structure this slightly different -- I'd have

SmallVector<Value *, 4> NewIncoming;
Type *NarrowType = nullptr;  // to keep track of what you're zext'ing from is same for all inputs

for (Value *V : Phi.incoming_value()) {
   // Check V here, and either append to NewIncoming or
   // return nullptr if you see something you cannot handle
   // (including constants)
}

// Create new PHI node with all the incoming values in NewIncoming

That'll keep most of the interesting logic in this function in one place.

[spatel]

lib/Transforms/InstCombine/InstCombinePHI.cpp

399

Hi Sanjoy - thanks for looking at this!
Yes, this check is unnecessary. This and the code structure (and the variable names and comments) are more general than they need to be because I started off hoping to handle any cast type. Let me make everything zext-specific for this patch, and then we can generalize it later if needed.

[spatel]

Patch updated based on Sanjoy's feedback:

1. Reorganize code so all interesting logic is in one loop.
2. Remove unnecessary opcode check because we're only dealing with zexts.
3. Fix variable names and comments to match the fact that we're only dealing with zexts.

[hfinkel]

lib/Transforms/InstCombine/InstCombinePHI.cpp

475

Do we already have a regression test that covers these cases?

test/Transforms/InstCombine/phi.ll

662

Could you please actually check that the PHI itself is formed correctly, and add a test case that covers the handling of the truncated-constant incoming value?

[spatel]

lib/Transforms/InstCombine/InstCombinePHI.cpp

475

There are other tests that cover the basics here:
a. test6 ("Suck casts into phi")
b. test13 has a case with a constant in a phi

...but I think we should have tests for phis that include more than 2 operands, so I've added those to show that the existing functions are working as expected on the constructs that this new code doesn't cover.

test/Transforms/InstCombine/phi.ll

662

Yes - that was just lazy.
Let me know if the new test cases with 0/1 -> false/true cover the constant scenario you're thinking of.

[spatel]

Patch updated:

1. Add tests to show that existing optimizations are working on cases not handled by this new code.
2. Add test with 2 constants and 2 variables to show that constant shrinking is working in the new code.

[hfinkel]

LGTM.

lib/Transforms/InstCombine/InstCombineInternal.h

private:
   Instruction *FoldOpIntoPhi(Instruction &I);
 
   /// \brief Try to rotate an operation below a PHI node, using PHI nodes for
   /// its operands.
   Instruction *FoldPHIArgOpIntoPHI(PHINode &PN);
   Instruction *FoldPHIArgBinOpIntoPHI(PHINode &PN);
   Instruction *FoldPHIArgGEPIntoPHI(PHINode &PN);
   Instruction *FoldPHIArgLoadIntoPHI(PHINode &PN);
+  Instruction *FoldPHIArgZextsIntoPHI(PHINode &PN);
 
   Instruction *OptAndOp(Instruction *Op, ConstantInt *OpRHS,
                         ConstantInt *AndRHS, BinaryOperator &TheAnd);
 
   Value *FoldLogicalPlusAnd(Value *LHS, Value *RHS, ConstantInt *Mask,
                             bool isSub, Instruction &I);
   Value *InsertRangeTest(Value *V, Constant *Lo, Constant *Hi, bool isSigned,
                          bool Inside);

lib/Transforms/InstCombine/InstCombinePHI.cpp

if (isVolatile)
     for (Value *IncValue : PN.incoming_values())
       cast<LoadInst>(IncValue)->setVolatile(false);
 
   LoadInst *NewLI = new LoadInst(PhiVal, "", isVolatile, LoadAlignment);
   NewLI->setDebugLoc(FirstLI->getDebugLoc());
   return NewLI;
 }
 
+/// TODO: This function could handle other cast types, but then it might
+/// require special-casing a cast from the 'i1' type. See the comment in
+/// FoldPHIArgOpIntoPHI() about pessimizing illegal integer types.
+Instruction *InstCombiner::FoldPHIArgZextsIntoPHI(PHINode &Phi) {
+  SmallVector<Constant*, 4> Consts;
+  ZExtInst *FirstCast = nullptr;
+  unsigned NumIncomingValues = Phi.getNumIncomingValues();
+
+  // Walk the phi operands looking for casts and constants.
+  for (Value *V : Phi.incoming_values()) {
+    if (auto *Cast = dyn_cast<ZExtInst>(V)) {
+      if (FirstCast == nullptr)
+        // Initialize cast instruction for matching subsequent casts.
+        FirstCast = Cast;
+      else if (Cast->getSrcTy() != FirstCast->getSrcTy() ||
+               Cast->getOpcode() != FirstCast->getOpcode() ||
+               !Cast->hasOneUse())
+        // All casts must be identical and have one use.
+        return nullptr;
+    } else if (auto *C = dyn_cast<Constant>(V)) {
+      // Make sure that constants can fit in the new type in the next loop.
+      Consts.push_back(C);
+    } else {
+      // If it's not a cast or a constant, bail out.
+      return nullptr;
+    }
+  }
 
+  // The more common cases of a phi with no constant operands or just one
+  // variable operand are handled by FoldPHIArgOpIntoPHI() and FoldOpIntoPhi()
+  // respectively. FoldOpIntoPhi() wants to do the opposite transform that is
+  // performed here. It tries to replicate a cast in the phi operand's basic
+  // block to expose other folding opportunities. Thus, InstCombine will
+  // infinite loop without this check.
+  if (Consts.empty() || NumIncomingValues - Consts.size() < 2)
+    return nullptr;
+
+  // Compute the constants that would happen if we truncated to SrcTy then
+  // re-extended to DstTy.
+  for (Constant *C : Consts) {
+    Constant *Trunc = ConstantExpr::getTrunc(C, FirstCast->getSrcTy());
+    Constant *Ext = ConstantExpr::getZExt(Trunc, C->getType());
+
+    // If the re-extended constant changed, then we can't shrink this phi
+    // operand for free.
+    if (Ext != C)
+      return nullptr;
+  }
+
+  // All incoming values are the same cast operation or constants that are safe
+  // to truncate. Create a new phi node of the correct type, and phi together
+  // all of the cast operands and constants.
+  PHINode *NewPhi = PHINode::Create(FirstCast->getSrcTy(), NumIncomingValues,
+                                    Phi.getName() + ".shrunk");
+
+  // Add all operands to the new phi.
+  for (Use &U : Phi.incoming_values()) {
+    Value *NewInVal = nullptr;
+    if (auto *Cast = dyn_cast<ZExtInst>(U))
+      NewInVal = Cast->getOperand(0);
+    else if (auto *C = dyn_cast<Constant>(U))
+      NewInVal = ConstantExpr::getTrunc(C, FirstCast->getSrcTy());
+
+    NewPhi->addIncoming(NewInVal, Phi.getIncomingBlock(U));
+  }
+
+  InsertNewInstBefore(NewPhi, Phi);
+  return CastInst::Create(FirstCast->getOpcode(), NewPhi, Phi.getType());
+}
 
 /// If all operands to a PHI node are the same "unary" operator and they all are
 /// only used by the PHI, PHI together their inputs, and do the operation once,
 /// to the result of the PHI.
 Instruction *InstCombiner::FoldPHIArgOpIntoPHI(PHINode &PN) {
   Instruction *FirstInst = cast<Instruction>(PN.getIncomingValue(0));
 
   if (isa<GetElementPtrInst>(FirstInst))
     return FoldPHIArgGEPIntoPHI(PN);
   if (isa<LoadInst>(FirstInst))
     return FoldPHIArgLoadIntoPHI(PN);
 
   // Scan the instruction, looking for input operations that can be folded away.
   // If all input operands to the phi are the same instruction (e.g. a cast from
   // the same type or "+42") we can pull the operation through the PHI, reducing
   // code size and simplifying code.
   Constant *ConstantOp = nullptr;
   Type *CastSrcTy = nullptr;
   bool isNUW = false, isNSW = false, isExact = false;
 
   if (isa<CastInst>(FirstInst)) {
     CastSrcTy = FirstInst->getOperand(0)->getType();
 
     // Be careful about transforming integer PHIs.  We don't want to pessimize
     // the code by turning an i32 into an i1293.
     if (PN.getType()->isIntegerTy() && CastSrcTy->isIntegerTy()) {
       if (!ShouldChangeType(PN.getType(), CastSrcTy))
         return nullptr;
     }
   } else if (isa<BinaryOperator>(FirstInst) || isa<CmpInst>(FirstInst)) {
     // Can fold binop, compare or shift here if the RHS is a constant,
for (unsigned i = 1, e = PN.getNumIncomingValues(); i != e; ++i) {
     if (isExact)
       isExact = cast<PossiblyExactOperator>(I)->isExact();
   }
 
   // Okay, they are all the same operation.  Create a new PHI node of the
   // correct type, and PHI together all of the LHS's of the instructions.
   PHINode *NewPN = PHINode::Create(FirstInst->getOperand(0)->getType(),
                                    PN.getNumIncomingValues(),
-                                   PN.getName()+".in");
+                                   PN.getName() + ".in");
 
   Value *InVal = FirstInst->getOperand(0);
   NewPN->addIncoming(InVal, PN.getIncomingBlock(0));
 
   // Add all operands to the new PHI.
   for (unsigned i = 1, e = PN.getNumIncomingValues(); i != e; ++i) {
     Value *NewInVal = cast<Instruction>(PN.getIncomingValue(i))->getOperand(0);
     if (NewInVal != InVal)

 }
 
 // PHINode simplification
 //
 Instruction *InstCombiner::visitPHINode(PHINode &PN) {
   if (Value *V = SimplifyInstruction(&PN, DL, TLI, DT, AC))
     return ReplaceInstUsesWith(PN, V);
 
+  if (Instruction *Result = FoldPHIArgZextsIntoPHI(PN))
+    return Result;
+  
   // If all PHI operands are the same operation, pull them through the PHI,
   // reducing code size.
   if (isa<Instruction>(PN.getIncomingValue(0)) &&
       isa<Instruction>(PN.getIncomingValue(1)) &&
       cast<Instruction>(PN.getIncomingValue(0))->getOpcode() ==
       cast<Instruction>(PN.getIncomingValue(1))->getOpcode() &&
       // FIXME: The hasOneUse check will fail for PHIs that use the value more
       // than themselves more than once.

test/Transforms/InstCombine/phi.ll

 ; CHECK: ret i32 undef
 define i32 @test27(i1 %b) {
 entry:
   br label %done
 done:
   %y = phi i32 [ undef, %entry ]
   ret i32 %y
 }
+
+; We should be able to fold the zexts to the other side of the phi
+; even though there's a constant value input to the phi because we
+; can shrink that constant to the smaller phi type.
+
+define i1 @PR24766(i8 %x1, i8 %x2, i8 %condition) {
+entry:
+  %conv = sext i8 %condition to i32
+  switch i32 %conv, label %epilog [
+    i32 0, label %sw1
+    i32 1, label %sw2
+  ]
+
+sw1:
+  %cmp1 = icmp eq i8 %x1, %x2
+  %frombool1 = zext i1 %cmp1 to i8
+  br label %epilog
+
+sw2:
+  %cmp2 = icmp sle i8 %x1, %x2
+  %frombool2 = zext i1 %cmp2 to i8
+  br label %epilog
+
+epilog:
+  %conditionMet = phi i8 [ 0, %entry ], [ %frombool2, %sw2 ], [ %frombool1, %sw1 ]
+  %tobool = icmp ne i8 %conditionMet, 0
+  ret i1 %tobool
+
+; CHECK-LABEL: @PR24766(
+; CHECK: phi i1
+}
+