[EarlyCSE] Propagate conditions of AND and OR instructions

[mkazantsev]

Summary

This patches teaches EarlyCSE to figure out that if and i1 %x, %y is true then both
%x and %y are true in the taken branch, and if or i1 %x, %y is false then both
%x and %y are false in non-taken branch. Fix for PR37635.

[lebedev.ri]

lib/Transforms/Scalar/EarlyCSE.cpp

723 ↗

(On Diff #149244)

It would be less confusing to name it BOp

728 ↗

(On Diff #149244)

It would be less confusing to name it BOp

[• dberlin]

If we want to start expanding this, i would suggest just using predicateinfo, which wouldn't require worklisting/etc to achieve this.

[mkazantsev]

I'm not sure if it gives benefits, but the idea is interesting. Do you mind if it goes as a follow-up?

[• dberlin]

So, the goal would be to reduce the maintenance burden and significantly more complex logic you are starting to introduce here about what can be propagated where.
If it doesn't bring benefits, i'd like to understand why.

I'm fine doing it as a followup.

[mkazantsev]

I will do it as a follow-up then.

[reames]

Current patch LGTM.

As a follow up, I think it might make sense to refactor the interface of the helper a bit. In particular, you have the SimpleValue::canHandle logic both in the caller and the callee. Might make sense to sink it into the callee only and check when taking items out of the worklist.

Danny, I'm not quite clear why PredicateInfo is the right direction to take here? Can you explain? Probably best done on llvm-dev instead of an individual review thread just so we have a wider audience.

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

bool EarlyCSE::handleBranchCondition(Instruction *CondInst,
                                      const BranchInst *BI, const BasicBlock *BB,
                                      const BasicBlock *Pred) {
   assert(BI->isConditional() && "Should be a conditional branch!");
   assert(BI->getCondition() == CondInst && "Wrong condition?");
   assert(BI->getSuccessor(0) == BB || BI->getSuccessor(1) == BB);
   auto *TorF = (BI->getSuccessor(0) == BB)
                    ? ConstantInt::getTrue(BB->getContext())
                    : ConstantInt::getFalse(BB->getContext());
-  AvailableValues.insert(CondInst, TorF);
+  auto IsAnd = [](Instruction *I) {
+    if (BinaryOperator *BOp = dyn_cast<BinaryOperator>(I))
+      return BOp->getOpcode() == Instruction::And;
+    return false;
+  };
+  auto IsOr = [](Instruction *I) {
+    if (BinaryOperator *BOp = dyn_cast<BinaryOperator>(I))
+      return BOp->getOpcode() == Instruction::Or;
+    return false;
+  };
+  // If the condition is AND operation, we can propagate its operands into the
+  // true branch. If it is OR operation, we can propagate them into the false
+  // branch.
+  auto CanPropagateOperands = (BI->getSuccessor(0) == BB) ? IsAnd : IsOr;
+
+  bool MadeChanges = false;
+  SmallVector<Instruction *, 4> WorkList;
+  SmallPtrSet<Instruction *, 4> Visited;
+  WorkList.push_back(CondInst);
+  while (!WorkList.empty()) {
+    Instruction *Curr = WorkList.pop_back_val();
+
+    AvailableValues.insert(Curr, TorF);
     LLVM_DEBUG(dbgs() << "EarlyCSE CVP: Add conditional value for '"
-                    << CondInst->getName() << "' as " << *TorF << " in "
+                      << Curr->getName() << "' as " << *TorF << " in "
                       << BB->getName() << "\n");
     if (!DebugCounter::shouldExecute(CSECounter)) {
       LLVM_DEBUG(dbgs() << "Skipping due to debug counter\n");
     } else {
       // Replace all dominated uses with the known value.
-    if (unsigned Count = replaceDominatedUsesWith(CondInst, TorF, DT,
+      if (unsigned Count = replaceDominatedUsesWith(Curr, TorF, DT,
                                                     BasicBlockEdge(Pred, BB))) {
-
         NumCSECVP += Count;
-      return true;
+        MadeChanges = true;
       }
     }
-  return false;
+
+    if (CanPropagateOperands(Curr))
+      for (auto &Op : cast<BinaryOperator>(Curr)->operands())
+        if (Instruction *OPI = dyn_cast<Instruction>(Op))
+          if (SimpleValue::canHandle(OPI) && Visited.insert(OPI).second)
+            WorkList.push_back(OPI);
+  }
+
+  return MadeChanges;
 }
 
 bool EarlyCSE::processNode(DomTreeNode *Node) {
   bool Changed = false;
   BasicBlock *BB = Node->getBlock();
 
   // If this block has a single predecessor, then the predecessor is the parent
   // of the domtree node and all of the live out memory values are still current

llvm/trunk/test/Transforms/EarlyCSE/and_or.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -early-cse -S < %s | FileCheck %s
+; RUN: opt -basicaa -early-cse-memssa -S < %s | FileCheck %s
+
+define i32 @test_01(i32 %a, i32 %b) {
+; CHECK-LABEL: @test_01(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[COND:%.*]] = icmp slt i32 [[A:%.*]], [[B:%.*]]
+; CHECK-NEXT:    br i1 [[COND]], label [[IF_TRUE:%.*]], label [[IF_FALSE:%.*]]
+; CHECK:       if.true:
+; CHECK-NEXT:    ret i32 [[A]]
+; CHECK:       if.false:
+; CHECK-NEXT:    ret i32 [[B]]
+;
+entry:
+  %cond = icmp slt i32 %a, %b
+  br i1 %cond, label %if.true, label %if.false
+
+if.true:
+  %cond2 = icmp slt i32 %a, %b
+  %x = select i1 %cond2, i32 %a, i32 %b
+  ret i32 %x
+
+if.false:
+  %cond3 = icmp slt i32 %a, %b
+  %y = select i1 %cond3, i32 %a, i32 %b
+  ret i32 %y
+}
+
+define i32 @test_02(i32 %a, i32 %b, i1 %c) {
+; CHECK-LABEL: @test_02(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[COND:%.*]] = icmp slt i32 [[A:%.*]], [[B:%.*]]
+; CHECK-NEXT:    [[AND_COND:%.*]] = and i1 [[COND]], [[C:%.*]]
+; CHECK-NEXT:    br i1 [[AND_COND]], label [[IF_TRUE:%.*]], label [[IF_FALSE:%.*]]
+; CHECK:       if.true:
+; CHECK-NEXT:    ret i32 [[A]]
+; CHECK:       if.false:
+; CHECK-NEXT:    [[Y:%.*]] = select i1 [[COND]], i32 [[A]], i32 [[B]]
+; CHECK-NEXT:    ret i32 [[Y]]
+;
+entry:
+  %cond = icmp slt i32 %a, %b
+  %and.cond = and i1 %cond, %c
+  br i1 %and.cond, label %if.true, label %if.false
+
+if.true:
+  %cond2 = icmp slt i32 %a, %b
+  %x = select i1 %cond2, i32 %a, i32 %b
+  ret i32 %x
+
+if.false:
+  %cond3 = icmp slt i32 %a, %b
+  %y = select i1 %cond3, i32 %a, i32 %b
+  ret i32 %y
+}
+
+define i32 @test_03(i32 %a, i32 %b, i1 %c) {
+; CHECK-LABEL: @test_03(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[COND:%.*]] = icmp slt i32 [[A:%.*]], [[B:%.*]]
+; CHECK-NEXT:    [[OR_COND:%.*]] = or i1 [[COND]], [[C:%.*]]
+; CHECK-NEXT:    br i1 [[OR_COND]], label [[IF_TRUE:%.*]], label [[IF_FALSE:%.*]]
+; CHECK:       if.true:
+; CHECK-NEXT:    [[X:%.*]] = select i1 [[COND]], i32 [[A]], i32 [[B]]
+; CHECK-NEXT:    ret i32 [[X]]
+; CHECK:       if.false:
+; CHECK-NEXT:    ret i32 [[B]]
+;
+entry:
+  %cond = icmp slt i32 %a, %b
+  %or.cond = or i1 %cond, %c
+  br i1 %or.cond, label %if.true, label %if.false
+
+if.true:
+  %cond2 = icmp slt i32 %a, %b
+  %x = select i1 %cond2, i32 %a, i32 %b
+  ret i32 %x
+
+if.false:
+  %cond3 = icmp slt i32 %a, %b
+  %y = select i1 %cond3, i32 %a, i32 %b
+  ret i32 %y
+}
+
+define i32 @test_04(i32 %a, i32 %b, i1 %c1, i1 %c2) {
+; CHECK-LABEL: @test_04(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[COND:%.*]] = icmp slt i32 [[A:%.*]], [[B:%.*]]
+; CHECK-NEXT:    [[AND_COND1:%.*]] = and i1 [[COND]], [[C1:%.*]]
+; CHECK-NEXT:    [[AND_COND2:%.*]] = and i1 [[AND_COND1]], [[C2:%.*]]
+; CHECK-NEXT:    br i1 [[AND_COND2]], label [[IF_TRUE:%.*]], label [[IF_FALSE:%.*]]
+; CHECK:       if.true:
+; CHECK-NEXT:    ret i32 [[A]]
+; CHECK:       if.false:
+; CHECK-NEXT:    [[Y:%.*]] = select i1 [[COND]], i32 [[A]], i32 [[B]]
+; CHECK-NEXT:    ret i32 [[Y]]
+;
+entry:
+  %cond = icmp slt i32 %a, %b
+  %and.cond1 = and i1 %cond, %c1
+  %and.cond2 = and i1 %and.cond1, %c2
+  br i1 %and.cond2, label %if.true, label %if.false
+
+if.true:
+  %cond2 = icmp slt i32 %a, %b
+  %x = select i1 %cond2, i32 %a, i32 %b
+  ret i32 %x
+
+if.false:
+  %cond3 = icmp slt i32 %a, %b
+  %y = select i1 %cond3, i32 %a, i32 %b
+  ret i32 %y
+}
+
+define i32 @test_05(i32 %a, i32 %b, i1 %c1, i1 %c2) {
+; CHECK-LABEL: @test_05(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[COND:%.*]] = icmp slt i32 [[A:%.*]], [[B:%.*]]
+; CHECK-NEXT:    [[OR_COND1:%.*]] = or i1 [[COND]], [[C1:%.*]]
+; CHECK-NEXT:    [[OR_COND2:%.*]] = or i1 [[OR_COND1]], [[C2:%.*]]
+; CHECK-NEXT:    br i1 [[OR_COND2]], label [[IF_TRUE:%.*]], label [[IF_FALSE:%.*]]
+; CHECK:       if.true:
+; CHECK-NEXT:    [[X:%.*]] = select i1 [[COND]], i32 [[A]], i32 [[B]]
+; CHECK-NEXT:    ret i32 [[X]]
+; CHECK:       if.false:
+; CHECK-NEXT:    ret i32 [[B]]
+;
+entry:
+  %cond = icmp slt i32 %a, %b
+  %or.cond1 = or i1 %cond, %c1
+  %or.cond2 = or i1 %or.cond1, %c2
+  br i1 %or.cond2, label %if.true, label %if.false
+
+if.true:
+  %cond2 = icmp slt i32 %a, %b
+  %x = select i1 %cond2, i32 %a, i32 %b
+  ret i32 %x
+
+if.false:
+  %cond3 = icmp slt i32 %a, %b
+  %y = select i1 %cond3, i32 %a, i32 %b
+  ret i32 %y
+}