1 files changed, 81 insertions, 32 deletions
diff --git a/llvm/lib/Transforms/AggressiveInstCombine/AggressiveInstCombine.cpp b/llvm/lib/Transforms/AggressiveInstCombine/AggressiveInstCombine.cpp
index b61c0b89c00..16e50c2d203 100644
--- a/llvm/lib/Transforms/AggressiveInstCombine/AggressiveInstCombine.cpp
+++ b/llvm/lib/Transforms/AggressiveInstCombine/AggressiveInstCombine.cpp
@@ -57,18 +57,45 @@ public:
 };
 } // namespace
 
-/// This is a recursive helper for 'and X, 1' that walks through a chain of 'or'
-/// instructions looking for shift ops of a common source value (first member of
-/// the pair). The second member of the pair is a mask constant for all of the
-/// bits that are being compared. So this:
-/// or (or (or X, (X >> 3)), (X >> 5)), (X >> 8)
-/// returns {X, 0x129} and those are the operands of an 'and' that is compared
-/// to zero.
-static bool matchMaskedCmpOp(Value *V, std::pair<Value *, APInt> &Result) {
-  // Recurse through a chain of 'or' operands.
+/// This is used by foldAnyOrAllBitsSet() to capture a source value (Root) and
+/// the bit indexes (Mask) needed by a masked compare. If we're matching a chain
+/// of 'and' ops, then we also need to capture the fact that we saw an
+/// "and X, 1", so that's an extra return value for that case.
+struct MaskOps {
+  Value *Root;
+  APInt Mask;
+  bool MatchAndChain;
+  bool FoundAnd1;
+
+  MaskOps(unsigned BitWidth, bool MatchAnds) :
+      Root(nullptr), Mask(APInt::getNullValue(BitWidth)),
+      MatchAndChain(MatchAnds), FoundAnd1(false) {}
+};
+
+/// This is a recursive helper for foldAnyOrAllBitsSet() that walks through a
+/// chain of 'and' or 'or' instructions looking for shift ops of a common source
+/// value. Examples:
+///   or (or (or X, (X >> 3)), (X >> 5)), (X >> 8)
+/// returns { X, 0x129 }
+///   and (and (X >> 1), 1), (X >> 4)
+/// returns { X, 0x12 }
+static bool matchAndOrChain(Value *V, MaskOps &MOps) {
   Value *Op0, *Op1;
-  if (match(V, m_Or(m_Value(Op0), m_Value(Op1))))
-    return matchMaskedCmpOp(Op0, Result) && matchMaskedCmpOp(Op1, Result);
+  if (MOps.MatchAndChain) {
+    // Recurse through a chain of 'and' operands. This requires an extra check
+    // vs. the 'or' matcher: we must find an "and X, 1" instruction somewhere
+    // in the chain to know that all of the high bits are cleared.
+    if (match(V, m_And(m_Value(Op0), m_One()))) {
+      MOps.FoundAnd1 = true;
+      return matchAndOrChain(Op0, MOps);
+    }
+    if (match(V, m_And(m_Value(Op0), m_Value(Op1))))
+      return matchAndOrChain(Op0, MOps) && matchAndOrChain(Op1, MOps);
+  } else {
+    // Recurse through a chain of 'or' operands.
+    if (match(V, m_Or(m_Value(Op0), m_Value(Op1))))
+      return matchAndOrChain(Op0, MOps) && matchAndOrChain(Op1, MOps);
+  }
 
   // We need a shift-right or a bare value representing a compare of bit 0 of
   // the original source operand.
@@ -78,33 +105,50 @@ static bool matchMaskedCmpOp(Value *V, std::pair<Value *, APInt> &Result) {
     Candidate = V;
 
   // Initialize result source operand.
-  if (!Result.first)
-    Result.first = Candidate;
+  if (!MOps.Root)
+    MOps.Root = Candidate;
 
   // Fill in the mask bit derived from the shift constant.
-  Result.second.setBit(BitIndex);
-  return Result.first == Candidate;
+  MOps.Mask.setBit(BitIndex);
+  return MOps.Root == Candidate;
 }
 
-/// Match an 'and' of a chain of or-shifted bits from a common source value into
-/// a masked compare:
-/// and (or (lshr X, C), ...), 1 --> (X & C') != 0
-static bool foldToMaskedCmp(Instruction &I) {
-  // TODO: This is only looking for 'any-bits-set' and 'all-bits-clear'.
-  // We should also match 'all-bits-set' and 'any-bits-clear' by looking for a
-  // a chain of 'and'.
-  if (!match(&I, m_And(m_OneUse(m_Or(m_Value(), m_Value())), m_One())))
+/// Match patterns that correspond to "any-bits-set" and "all-bits-set".
+/// These will include a chain of 'or' or 'and'-shifted bits from a
+/// common source value:
+/// and (or  (lshr X, C), ...), 1 --> (X & CMask) != 0
+/// and (and (lshr X, C), ...), 1 --> (X & CMask) == CMask
+/// Note: "any-bits-clear" and "all-bits-clear" are variations of these patterns
+/// that differ only with a final 'not' of the result. We expect that final
+/// 'not' to be folded with the compare that we create here (invert predicate).
+static bool foldAnyOrAllBitsSet(Instruction &I) {
+  // The 'any-bits-set' ('or' chain) pattern is simpler to match because the
+  // final "and X, 1" instruction must be the final op in the sequence.
+  bool MatchAllBitsSet;
+  if (match(&I, m_c_And(m_OneUse(m_And(m_Value(), m_Value())), m_Value())))
+    MatchAllBitsSet = true;
+  else if (match(&I, m_And(m_OneUse(m_Or(m_Value(), m_Value())), m_One())))
+    MatchAllBitsSet = false;
+  else
     return false;
 
-  std::pair<Value *, APInt>
-  MaskOps(nullptr, APInt::getNullValue(I.getType()->getScalarSizeInBits()));
-  if (!matchMaskedCmpOp(cast<BinaryOperator>(&I)->getOperand(0), MaskOps))
-    return false;
+  MaskOps MOps(I.getType()->getScalarSizeInBits(), MatchAllBitsSet);
+  if (MatchAllBitsSet) {
+    if (!matchAndOrChain(cast<BinaryOperator>(&I), MOps) || !MOps.FoundAnd1)
+      return false;
+  } else {
+    if (!matchAndOrChain(cast<BinaryOperator>(&I)->getOperand(0), MOps))
+      return false;
+  }
 
+  // The pattern was found. Create a masked compare that replaces all of the
+  // shift and logic ops.
   IRBuilder<> Builder(&I);
-  Value *Mask = Builder.CreateAnd(MaskOps.first, MaskOps.second);
-  Value *CmpZero = Builder.CreateIsNotNull(Mask);
-  Value *Zext = Builder.CreateZExt(CmpZero, I.getType());
+  Constant *Mask = ConstantInt::get(I.getType(), MOps.Mask);
+  Value *And = Builder.CreateAnd(MOps.Root, Mask);
+  Value *Cmp = MatchAllBitsSet ? Builder.CreateICmpEQ(And, Mask) :
+                                 Builder.CreateIsNotNull(And);
+  Value *Zext = Builder.CreateZExt(Cmp, I.getType());
   I.replaceAllUsesWith(Zext);
   return true;
 }
@@ -119,8 +163,13 @@ static bool foldUnusualPatterns(Function &F, DominatorTree &DT) {
     if (!DT.isReachableFromEntry(&BB))
       continue;
     // Do not delete instructions under here and invalidate the iterator.
-    for (Instruction &I : BB)
-      MadeChange |= foldToMaskedCmp(I);
+    // Walk the block backwards for efficiency. We're matching a chain of
+    // use->defs, so we're more likely to succeed by starting from the bottom.
+    // Also, we want to avoid matching partial patterns.
+    // TODO: It would be more efficient if we removed dead instructions
+    // iteratively in this loop rather than waiting until the end.
+    for (Instruction &I : make_range(BB.rbegin(), BB.rend()))
+      MadeChange |= foldAnyOrAllBitsSet(I);
   }
 
   // We're done with transforms, so remove dead instructions.