[LogicCombine 5/?] Fold constant in logical expression

[bcl5980]

rebase

llvm/include/llvm/Analysis/LogicCombine.h

public:
   bool worthToCombine(unsigned InstCnt) const {
     return (OneUseWeight + InstCnt) < Weight;
   }
   void print(raw_ostream &OS) const;
 };
 
 class LogicCombiner {
 public:
-  LogicCombiner() : LogicalOpNodes(), LeafValues(), LeafsMayPoison() {}
+  LogicCombiner()
+      : LogicalOpNodes(), LeafValues(), LeafsMayPoison(), ConstantLeafs() {}
   ~LogicCombiner() { clear(); }
 
   Value *simplify(Value *Root);
 
 private:
   friend class LogicalOpNode;
 
   SpecificBumpPtrAllocator<LogicalOpNode> Alloc;
   SmallDenseMap<Value *, LogicalOpNode *, 16> LogicalOpNodes;
   SmallSetVector<Value *, 8> LeafValues;
   uint64_t LeafsMayPoison;
+  uint64_t ConstantLeafs;
+  Constant *ConstAllOne;
+  Constant *ConstZero;
 
   void clear();
 
   LogicalOpNode *visitLeafNode(Value *Val, unsigned Depth);
   LogicalOpNode *visitBinOp(BinaryOperator *BO, unsigned Depth);
   LogicalOpNode *visitSelect(SelectInst *SI, unsigned Depth);
   LogicalOpNode *getLogicalOpNode(Value *Val, unsigned Depth = 0);
+  void foldConstForExpr(LogicalExpr &Expr);
+
   Value *logicalOpToValue(LogicalOpNode *Node);
   Value *buildAndChain(IRBuilder<> &Builder, Type *Ty, uint64_t LeafBits);
 };
 
 inline raw_ostream &operator<<(raw_ostream &OS, const LogicalOpNode &I) {
   I.print(OS);
   return OS;
 }
 
 } // namespace llvm

llvm/lib/Analysis/LogicCombine.cpp

void LogicalOpNode::print(raw_ostream &OS) const {
   OS << "\nWeight: " << Weight << "; OneUseWeight: " << OneUseWeight << ";\n";
   OS << "PoisonSrc: 0x"
      << Twine::utohexstr(Expr.getLeafMask() & Helper->LeafsMayPoison)
      << "; PoisonMaskSI: 0x" << Twine::utohexstr(PoisonMaskSI) << ";\n\n";
 }
 
 void LogicCombiner::clear() {
   LeafsMayPoison = 0;
+  ConstantLeafs = 0;
   LogicalOpNodes.clear();
   LeafValues.clear();
 }
 
 LogicalOpNode *LogicCombiner::visitLeafNode(Value *Val, unsigned Depth) {
   // Depth is 0 means the root is not logical operation. We can't
   // do anything for that.
   if (Depth == 0 || LeafValues.size() >= MaxLogicOpLeafsToScan)
     return nullptr;
 
   uint64_t ExprVal = 1ULL << LeafValues.size();
   // Constant Zero,AllOne are special leaf nodes. They involve
   // LogicalExpr's calculation so we must detect them at first.
   if (auto ConstVal = dyn_cast<ConstantInt>(Val)) {
     if (ConstVal->isZero())
       ExprVal = 0;
     else if (ConstVal->isAllOnesValue())
       ExprVal = LogicalExpr::ExprAllOne;
   }
   if (ExprVal != LogicalExpr::ExprAllOne && ExprVal != 0) {
     if (!isGuaranteedNotToBeUndefOrPoison(Val))
       LeafsMayPoison |= ExprVal;
+    if (isa<Constant>(Val))
+      ConstantLeafs |= ExprVal;
     LeafValues.insert(Val);
   }
   LogicalOpNode *Node = new (Alloc.Allocate())
       LogicalOpNode(this, Val, LogicalExpr(ExprVal), 0, 0, 0);
   LogicalOpNodes[Val] = Node;
   return Node;
 }
 
LogicalOpNode *LogicCombiner::visitBinOp(BinaryOperator *BO, unsigned Depth) {
 
   LogicalExpr NewExpr;
   if (BO->getOpcode() == Instruction::And)
     NewExpr = LHS->getExpr() & RHS->getExpr();
   else if (BO->getOpcode() == Instruction::Or)
     NewExpr = LHS->getExpr() | RHS->getExpr();
   else
     NewExpr = LHS->getExpr() ^ RHS->getExpr();
+  foldConstForExpr(NewExpr);
 
   uint64_t PoisonMaskSI = 0;
   uint64_t LPI = LHS->getExpr().getLeafMask() ^ LHS->getPoisonMaskSI();
   uint64_t RPI = RHS->getExpr().getLeafMask() ^ RHS->getPoisonMaskSI();
   PoisonMaskSI =
       ((~LPI) & RHS->getPoisonMaskSI()) | ((~RPI) & LHS->getPoisonMaskSI());
   PoisonMaskSI &= NewExpr.getLeafMask() & LeafsMayPoison;
 
if (TrueVal == nullptr)
     return nullptr;
 
   LogicalOpNode *FalseVal = getLogicalOpNode(SI->getFalseValue(), Depth + 1);
   if (FalseVal == nullptr)
     return nullptr;
 
   LogicalExpr NewExpr = (Cond->getExpr() & TrueVal->getExpr()) ^
                         ((~Cond->getExpr()) & FalseVal->getExpr());
+  foldConstForExpr(NewExpr);
+
   // TODO: We can reduce the weight if th node can be simplified even if
   // it is not the root node.
   unsigned Weight =
       Cond->getWeight() + TrueVal->getWeight() + FalseVal->getWeight() + 1;
   unsigned OneUseWeight = Weight;
   if (SI->hasOneUse())
     OneUseWeight = Cond->getOneUseWeight() + TrueVal->getOneUseWeight() +
                    FalseVal->getOneUseWeight();
if (LogicalOpNodes.find(Val) == LogicalOpNodes.end()) {
 
     if (!Node)
       return nullptr;
     LLVM_DEBUG(dbgs() << *Node);
   }
   return LogicalOpNodes[Val];
 }
 
+void LogicCombiner::foldConstForExpr(LogicalExpr &Expr) {
+  uint64_t ConstLeafBits = Expr.getLeafMask() & ConstantLeafs;
+  if (popcount(ConstLeafBits) <= 1)
+    return;
+
+  bool Changed = false;
+  ExprAddChain NewAddChain;
+  Constant *ConstXor = nullptr;
+  for (auto BitSet : Expr) {
+    ConstLeafBits = BitSet & ConstantLeafs;
+    int ConstLeafCnt = popcount(ConstLeafBits);
+    if (ConstLeafCnt > 1) {
+      Constant *AndChain = ConstAllOne;
+      unsigned LeafIdx;
+      uint64_t LeafIterBits = ConstLeafBits;
+      for (int I = 0; I < ConstLeafCnt; I++) {
+        LeafIdx = countr_zero(LeafIterBits);
+        Value *LeafVal = LeafValues[LeafIdx];
+        if (auto *ConstVal = dyn_cast<Constant>(LeafVal))
+          AndChain = ConstantExpr::getAnd(ConstVal, AndChain);
+        LeafIterBits -= (1ULL << LeafIdx);
+      }
+
+      if (auto *NewConstNode = getLogicalOpNode(AndChain, 1)) {
+        Changed = true;
+        if (NewConstNode->getExpr().size() == 0) {
+          BitSet = 0;
+          continue;
+        } else {
+          BitSet ^= ConstLeafBits;
+          BitSet |= *NewConstNode->getExpr().begin();
+        }
+      }
+    }
+
+    ConstLeafBits = BitSet & ConstantLeafs;
+    if (ConstLeafBits == BitSet && popcount(BitSet) == 1) {
+      unsigned LeafIdx = countr_zero(BitSet);
+      Constant *ConstVal = cast<Constant>(LeafValues[LeafIdx]);
+      if (!ConstXor)
+        ConstXor = ConstVal;
+      else
+        ConstXor = ConstantExpr::getXor(ConstXor, ConstVal);
+    } else
+      NewAddChain.insert(BitSet);
+  }
+
+  if (ConstXor && !ConstXor->isNullValue()) {
+    if (auto *NewConstNode = getLogicalOpNode(ConstXor, 1)) {
+      Changed = true;
+      NewAddChain.insert(*NewConstNode->getExpr().begin());
+    }
+  }
+
+  if (Changed)
+    Expr = LogicalExpr(NewAddChain);
+}
+
 Value *LogicCombiner::logicalOpToValue(LogicalOpNode *Node) {
   const LogicalExpr &Expr = Node->getExpr();
   // Empty when all leaf bits are erased from the set because a ^ a = 0.
   if (Expr.size() == 0)
     return Constant::getNullValue(Node->getValue()->getType());
 
   if (Node->getPoisonMaskSI() != 0)
     return nullptr;
Value *LogicCombiner::logicalOpToValue(LogicalOpNode *Node) {
   // 2 and chains.
 
   return nullptr;
 }
 
 Value *LogicCombiner::buildAndChain(IRBuilder<> &Builder, Type *Ty,
                                     uint64_t LeafBits) {
   if (LeafBits == 0)
-    return Constant::getNullValue(Ty);
+    return ConstZero;
 
   // ExprAllOne is not in the LeafValues
   if (LeafBits == LogicalExpr::ExprAllOne)
-    return Constant::getAllOnesValue(Ty);
+    return ConstAllOne;
 
   unsigned LeafCnt = popcount(LeafBits);
   if (LeafCnt == 1)
     return LeafValues[Log2_64(LeafBits)];
 
   unsigned LeafIdx = countr_zero(LeafBits);
   Value *AndChain = LeafValues[LeafIdx];
   LeafBits -= (1ULL << LeafIdx);
   for (unsigned I = 1; I < LeafCnt; I++) {
     LeafIdx = countr_zero(LeafBits);
     AndChain = Builder.CreateAnd(AndChain, LeafValues[LeafIdx]);
     LeafBits -= (1ULL << LeafIdx);
   }
   return AndChain;
 }
 
 Value *LogicCombiner::simplify(Value *Root) {
   assert(MaxLogicOpLeafsToScan <= 63 &&
          "Logical leaf node can't be larger than 63.");
+
+  ConstAllOne = Constant::getAllOnesValue(Root->getType());
+  ConstZero = Constant::getNullValue(Root->getType());
+
   LogicalOpNode *RootNode = getLogicalOpNode(Root);
   if (RootNode == nullptr)
     return nullptr;
 
   Value *NewRoot = logicalOpToValue(RootNode);
   if (NewRoot == nullptr || NewRoot == Root)
     return nullptr;
 
   LogicalOpNodes.erase(Root);
   NumLogicalOpsSimplified++;
   return NewRoot;
 }

llvm/test/Transforms/AggressiveInstCombine/logic-combine.ll

 ; CHECK-LABEL: @leaf1_or_not(
 ; CHECK-NEXT:    ret i8 -1
 ;
   %not.a = xor i8 %a, -1
   %or = or i8 %a, %not.a
   ret i8 %or
 }
 
+define i8 @leaf1_and_const_fold(i8 %a)  {
+; CHECK-LABEL: @leaf1_and_const_fold(
+; CHECK-NEXT:    [[TMP1:%.*]] = and i8 [[A:%.*]], 15
+; CHECK-NEXT:    ret i8 [[TMP1]]
+;
+  %and1 = and i8 %a, 111
+  %and2 = and i8 %and1, 31
+  ret i8 %and2
+}
+
+define i8 @leaf1_xor_const_fold(i8 %a)  {
+; CHECK-LABEL: @leaf1_xor_const_fold(
+; CHECK-NEXT:    [[TMP1:%.*]] = xor i8 112, [[A:%.*]]
+; CHECK-NEXT:    ret i8 [[TMP1]]
+;
+  %xor1 = xor i8 %a, 111
+  %xor2 = xor i8 %xor1, 31
+  ret i8 %xor2
+}
+
 define i8 @leaf2_xor(i8 %a, i8 %b)  {
 ; CHECK-LABEL: @leaf2_xor(
 ; CHECK-NEXT:    ret i8 [[B:%.*]]
 ;
   %ab = xor i8 %a, %b
   %xor.ab.a = xor i8 %ab, %a
   ret i8 %xor.ab.a
 }

 ; CHECK-NEXT:    ret i16 [[TMP1]]
 ;
   %or.ab = or i16 %a, %b
   %and.ab = and i16 %a, %b
   %xor = xor i16 %or.ab, %and.ab
   ret i16 %xor
 }
 
+define i8 @leaf2_ret_const_fold(i8 %a, i8 %b)  {
+; CHECK-LABEL: @leaf2_ret_const_fold(
+; CHECK-NEXT:    [[TMP1:%.*]] = xor i8 [[B:%.*]], -1
+; CHECK-NEXT:    [[TMP2:%.*]] = and i8 [[A:%.*]], 18
+; CHECK-NEXT:    [[TMP3:%.*]] = and i8 [[TMP1]], [[TMP2]]
+; CHECK-NEXT:    ret i8 [[TMP3]]
+;
+  %and1 = and i8 %a, 122
+  %or = or i8 %and1, %b
+  %and2 = and i8 %or, 23
+  %and3 = and i8 %b, 23
+  %r = xor i8 %and2, %and3
+  ret i8 %r
+}
+
 define i8 @leaf3_complex_ret_const_false(i8 %a, i8 %b, i8 %c)  {
 ; CHECK-LABEL: @leaf3_complex_ret_const_false(
 ; CHECK-NEXT:    ret i8 0
 ;
   %ab = or i8 %a, %b
   %abc = or i8 %ab, %c
   %not.abc = xor i8 %abc, -1
   %r = and i8 %not.abc, %a