[InstCombine] extend "simplifyUsingControlFlow" supporting zext/sext/trunc for different sizes

[inclyc]

Summary

Fixes: https://github.com/llvm/llvm-project/issues/54561

[inclyc]

• Update comments

[inclyc]

ping :) @nikic can you please review this revision to make sure it is what you want in https://github.com/llvm/llvm-project/issues/54561#issuecomment-1079460876 ?

[khei4]

(I'm new to around this, so those comments are just my opinions.)

llvm/lib/Transforms/InstCombine/InstCombinePHI.cpp

1307–1308

IMHO, duplicating SuccForValue map for each casts seems a little redundant. I like collecting incoming phi values in the front and checking correspondences to unify those maps. Then we can also early return if the values don't correspond.

1384–1386

Inverting can't be determined on these lines. If you do so, these comments look obsolete, and better to be modified

1387–1408

IMHO, this checking, which is done when creating maps, seems a bit redundant, it's better to hold conditions in some way.

[inclyc]

llvm/lib/Transforms/InstCombine/InstCombinePHI.cpp

1307–1308

Thanks for the feedback!

I like collecting incoming phi values in the front and checking correspondences to unify those maps.

Can you elaborate your solution? Here are my worries:

Here we are going to find a operation that maps branch condition in to corresponding phi values. Not an operation that maps phi values back to branch condition.

I think I have to left "sext"/"zext"ed values here because we must know the result of the branch condition after "s/zext". That's different from truncating from PHI values. i.e. s/zext information is not recoverable. Truncating phi values to branch condition does not imply we can z/sext the branch condition back to "phi values".

1384–1386

Nice catch! I'd like to fixup this after all "functional change" reviews.

[khei4]

llvm/lib/Transforms/InstCombine/InstCombinePHI.cpp

1307–1308

Sorry for my vague comments. The idea was unifying later inverting checking and these casting checking. Finding SExt or ZExt is similar to decide an inversion

I think I have to left "sext"/"zext"ed values here because we must know the result of the branch condition after "s/zext".

I just thought it'll be neat to reverse the roll between Succ/Cond and Pred/Phi. Then we don't need to hold candidate values. But I feel this might be too detailed and optional, this requires the dominant condition reversed. (I just tried to write actually, but ugly wip https://github.com/llvm/llvm-project/compare/main...khei4:llvm-project:wip/khei4-sext

[inclyc]

ping :)

[inclyc]

ping :)

llvm/lib/Transforms/InstCombine/InstCombinePHI.cpp

static Value *simplifyUsingControlFlow(InstCombiner &Self, PHINode &PN,
   // case v1: /       \ case v2:
   //         ...      ...
   //          \       /
   //       phi [v1] [v2]
   // Make sure all inputs are constants.
   if (!all_of(PN.operands(), [](Value *V) { return isa<ConstantInt>(V); }))
     return nullptr;
 
+  unsigned PNSize = PN.getType()->getPrimitiveSizeInBits();
   BasicBlock *BB = PN.getParent();
   // Do not bother with unreachable instructions.
   if (!DT.isReachableFromEntry(BB))
     return nullptr;
 
   // Determine which value the condition of the idom has for which successor.
   LLVMContext &Context = PN.getContext();
   auto *IDom = DT.getNode(BB)->getIDom()->getBlock();
   Value *Cond;
-  SmallDenseMap<ConstantInt *, BasicBlock *, 8> SuccForValue;
+  SmallDenseMap<ConstantInt *, BasicBlock *, 8> SuccValue, SuccValueTrunc,
+      SuccValueZExt, SuccValueSExt;
   SmallDenseMap<BasicBlock *, unsigned, 8> SuccCount;
-  auto AddSucc = [&](ConstantInt *C, BasicBlock *Succ) {
-    SuccForValue[C] = Succ;
+  auto AddSucc = [&](Value *Cond, ConstantInt *C, BasicBlock *Succ) {
+    unsigned CondSize = Cond->getType()->getPrimitiveSizeInBits();
+    if (CondSize < PNSize) {
+      // The cond needs to be extended
+      SuccValueSExt[cast<ConstantInt>(ConstantExpr::getSExt(C, PN.getType()))] =
+          Succ;
+      SuccValueZExt[cast<ConstantInt>(ConstantExpr::getZExt(C, PN.getType()))] =
+          Succ;
+    } else if (CondSize == PNSize) {
+      SuccValue[C] = Succ;
+    } else {
+      SuccValueTrunc[cast<ConstantInt>(
+          ConstantExpr::getTrunc(C, PN.getType()))] = Succ;
+    }
     ++SuccCount[Succ];
   };
   if (auto *BI = dyn_cast<BranchInst>(IDom->getTerminator())) {
     if (BI->isUnconditional())
       return nullptr;
 
     Cond = BI->getCondition();
-    AddSucc(ConstantInt::getTrue(Context), BI->getSuccessor(0));
-    AddSucc(ConstantInt::getFalse(Context), BI->getSuccessor(1));
+    AddSucc(Cond, ConstantInt::getTrue(Context), BI->getSuccessor(0));
+    AddSucc(Cond, ConstantInt::getFalse(Context), BI->getSuccessor(1));
   } else if (auto *SI = dyn_cast<SwitchInst>(IDom->getTerminator())) {
     Cond = SI->getCondition();
     ++SuccCount[SI->getDefaultDest()];
     for (auto Case : SI->cases())
-      AddSucc(Case.getCaseValue(), Case.getCaseSuccessor());
+      AddSucc(Cond, Case.getCaseValue(), Case.getCaseSuccessor());
   } else {
     return nullptr;
   }
 
-  if (Cond->getType() != PN.getType())
-    return nullptr;
-
   // Check that edges outgoing from the idom's terminators dominate respective
   // inputs of the Phi.
+  auto CheckSuccValue =
+      [&](SmallDenseMap<ConstantInt *, BasicBlock *, 8> SuccForValue)
+      -> std::tuple<bool, std::optional<bool>> {
     std::optional<bool> Invert;
     for (auto Pair : zip(PN.incoming_values(), PN.blocks())) {
       auto *Input = cast<ConstantInt>(std::get<0>(Pair));
       BasicBlock *Pred = std::get<1>(Pair);
       auto IsCorrectInput = [&](ConstantInt *Input) {
-      // The input needs to be dominated by the corresponding edge of the idom.
-      // This edge cannot be a multi-edge, as that would imply that multiple
-      // different condition values follow the same edge.
+        // The input needs to be dominated by the corresponding edge of the
+        // idom. This edge cannot be a multi-edge, as that would imply that
+        // multiple different condition values follow the same edge.
         auto It = SuccForValue.find(Input);
         return It != SuccForValue.end() && SuccCount[It->second] == 1 &&
                DT.dominates(BasicBlockEdge(IDom, It->second),
                             BasicBlockEdge(Pred, BB));
       };
 
       // Depending on the constant, the condition may need to be inverted.
       bool NeedsInvert;
       if (IsCorrectInput(Input))
         NeedsInvert = false;
       else if (IsCorrectInput(cast<ConstantInt>(ConstantExpr::getNot(Input))))
         NeedsInvert = true;
       else
-      return nullptr;
+        return {false, std::nullopt};
 
       // Make sure the inversion requirement is always the same.
       if (Invert && *Invert != NeedsInvert)
-      return nullptr;
+        return {false, std::nullopt};
 
       Invert = NeedsInvert;
     }
+    return {true, Invert};
+  };
 
-  if (!*Invert)
-    return Cond;
-
+  unsigned CondSize = Cond->getType()->getPrimitiveSizeInBits();
+  // TODO(lyc): change the comment here
   // This Phi is actually opposite to branching condition of IDom. We invert
   // the condition that will potentially open up some opportunities for
   // sinking.
   auto InsertPt = BB->getFirstInsertionPt();
   if (InsertPt != BB->end()) {
     Self.Builder.SetInsertPoint(&*InsertPt);
-    return Self.Builder.CreateNot(Cond);
+    if (CondSize < PNSize) {
+      auto [ZExtOk, ZExtInvert] = CheckSuccValue(SuccValueZExt);
+      if (ZExtOk) {
+        Cond = Self.Builder.CreateZExt(Cond, PN.getType());
+        return *ZExtInvert ? Self.Builder.CreateNot(Cond) : Cond;
+      }
+      auto [SExtOk, SExtInvert] = CheckSuccValue(SuccValueSExt);
+      if (SExtOk) {
+        Cond = Self.Builder.CreateSExt(Cond, PN.getType());
+        return *SExtInvert ? Self.Builder.CreateNot(Cond) : Cond;
+      }
+    } else if (CondSize == PNSize) {
+      auto [Ok, Invert] = CheckSuccValue(SuccValue);
+      if (Ok)
+        return *Invert ? Self.Builder.CreateNot(Cond) : Cond;
+    } else {
+      auto [Ok, Invert] = CheckSuccValue(SuccValueTrunc);
+      if (Ok) {
+        Cond = Self.Builder.CreateTrunc(Cond, PN.getType());
+        return *Invert ? Self.Builder.CreateNot(Cond) : Cond;
+      }
+    }
   }
 
   return nullptr;
 }
 
 // PHINode simplification
 //
 Instruction *InstCombinerImpl::visitPHINode(PHINode &PN) {

llvm/test/Transforms/InstCombine/simple_phi_condition.ll

 ; CHECK-NEXT:    unreachable
 ; CHECK:       bb3:
 ; CHECK-NEXT:    br label [[BB5:%.*]]
 ; CHECK:       bb4:
 ; CHECK-NEXT:    br label [[BB5]]
 ; CHECK:       bb1:
 ; CHECK-NEXT:    br label [[BB5]]
 ; CHECK:       bb5:
-; CHECK-NEXT:    [[DOT0:%.*]] = phi i32 [ 1, [[BB1]] ], [ 0, [[BB4]] ], [ -1, [[BB3]] ]
+; CHECK-NEXT:    [[DOT0:%.*]] = sext i8 [[TMP0]] to i32
 ; CHECK-NEXT:    ret i32 [[DOT0]]
 ;
   switch i8 %0, label %bb2 [
   i8 -1, label %bb3
   i8 0, label %bb4
   i8 1, label %bb1
   ]
 

 ; CHECK-NEXT:    unreachable
 ; CHECK:       bb3:
 ; CHECK-NEXT:    br label [[BB5:%.*]]
 ; CHECK:       bb4:
 ; CHECK-NEXT:    br label [[BB5]]
 ; CHECK:       bb1:
 ; CHECK-NEXT:    br label [[BB5]]
 ; CHECK:       bb5:
-; CHECK-NEXT:    [[DOT0:%.*]] = phi i32 [ 1, [[BB1]] ], [ 0, [[BB4]] ], [ 255, [[BB3]] ]
+; CHECK-NEXT:    [[DOT0:%.*]] = zext i8 [[TMP0]] to i32
 ; CHECK-NEXT:    ret i32 [[DOT0]]
 ;
   switch i8 %0, label %bb2 [
   i8 -1, label %bb3
   i8 0, label %bb4
   i8 1, label %bb1
   ]
 

 ; CHECK-NEXT:    unreachable
 ; CHECK:       bb3:
 ; CHECK-NEXT:    br label [[BB5:%.*]]
 ; CHECK:       bb4:
 ; CHECK-NEXT:    br label [[BB5]]
 ; CHECK:       bb1:
 ; CHECK-NEXT:    br label [[BB5]]
 ; CHECK:       bb5:
-; CHECK-NEXT:    [[DOT0:%.*]] = phi i8 [ 1, [[BB1]] ], [ 0, [[BB4]] ], [ -1, [[BB3]] ]
+; CHECK-NEXT:    [[DOT0:%.*]] = trunc i32 [[TMP0]] to i8
 ; CHECK-NEXT:    ret i8 [[DOT0]]
 ;
   switch i32 %0, label %bb2 [
   i32 -1, label %bb3
   i32 0, label %bb4
   i32 1, label %bb1
   ]
 

 ; CHECK-NEXT:    unreachable
 ; CHECK:       bb3:
 ; CHECK-NEXT:    br label [[BB5:%.*]]
 ; CHECK:       bb4:
 ; CHECK-NEXT:    br label [[BB5]]
 ; CHECK:       bb1:
 ; CHECK-NEXT:    br label [[BB5]]
 ; CHECK:       bb5:
-; CHECK-NEXT:    [[DOT0:%.*]] = phi i32 [ -2, [[BB1]] ], [ -1, [[BB4]] ], [ 0, [[BB3]] ]
+; CHECK-NEXT:    [[TMP2:%.*]] = xor i8 [[TMP0]], -1
+; CHECK-NEXT:    [[DOT0:%.*]] = sext i8 [[TMP2]] to i32
 ; CHECK-NEXT:    ret i32 [[DOT0]]
 ;
   switch i8 %0, label %bb2 [
   i8 -1, label %bb3
   i8 0, label %bb4
   i8 1, label %bb1
   ]
 

 ; CHECK-NEXT:    unreachable
 ; CHECK:       bb3:
 ; CHECK-NEXT:    br label [[BB5:%.*]]
 ; CHECK:       bb4:
 ; CHECK-NEXT:    br label [[BB5]]
 ; CHECK:       bb1:
 ; CHECK-NEXT:    br label [[BB5]]
 ; CHECK:       bb5:
-; CHECK-NEXT:    [[DOT0:%.*]] = phi i32 [ -2, [[BB1]] ], [ -1, [[BB4]] ], [ -256, [[BB3]] ]
+; CHECK-NEXT:    [[TMP2:%.*]] = zext i8 [[TMP0]] to i32
+; CHECK-NEXT:    [[DOT0:%.*]] = xor i32 [[TMP2]], -1
 ; CHECK-NEXT:    ret i32 [[DOT0]]
 ;
   switch i8 %0, label %bb2 [
   i8 -1, label %bb3
   i8 0, label %bb4
   i8 1, label %bb1
   ]
 

 ; CHECK-NEXT:    unreachable
 ; CHECK:       bb3:
 ; CHECK-NEXT:    br label [[BB5:%.*]]
 ; CHECK:       bb4:
 ; CHECK-NEXT:    br label [[BB5]]
 ; CHECK:       bb1:
 ; CHECK-NEXT:    br label [[BB5]]
 ; CHECK:       bb5:
-; CHECK-NEXT:    [[DOT0:%.*]] = phi i8 [ -2, [[BB1]] ], [ -1, [[BB4]] ], [ 0, [[BB3]] ]
+; CHECK-NEXT:    [[TMP2:%.*]] = trunc i32 [[TMP0]] to i8
+; CHECK-NEXT:    [[DOT0:%.*]] = xor i8 [[TMP2]], -1
 ; CHECK-NEXT:    ret i8 [[DOT0]]
 ;
   switch i32 %0, label %bb2 [
   i32 -1, label %bb3
   i32 0, label %bb4
   i32 1, label %bb1
   ]