[AggressiveInstCombine] convert a chain of 'or-shift' bits into masked compare

[spatel]

Summary

and (or (lshr X, C), ...), 1 --> (X & C') != 0

I initially thought about implementing the minimal pattern in instcombine as mentioned here:
https://bugs.llvm.org/show_bug.cgi?id=37098#c6

...but we need to do better to catch the more general sequence from the motivating test (more than 2 bits in the compare). And a test-suite run with statistics showed that this pattern only happened 2 times currently. It would potentially happen more often if reassociation worked better (D45842), but it's probably still not too frequent?

This is small enough that I didn't see a need to create a whole new class/file within AggressiveInstCombine. There are likely other relatively small matchers like what was discussed in D44266 that would slide under foldUnusualPatterns() (name suggestions welcome). We could potentially also consolidate matchers for ctpop, bswap, etc under here.

[spatel]

Ping.

[kparzysz]

This recognizes something like "bitmask-any", where the code checks if any of the bits of V indicated by the mask C' is set. It looks like by replacing or with and you would get the corresponding "bitmask-all". Would it be worthwhile to add it? I don't have a use case, but the symmetry seems somewhat compelling.

[spatel]

In D45986#1083922, @kparzysz wrote:

This recognizes something like "bitmask-any", where the code checks if any of the bits of V indicated by the mask C' is set. It looks like by replacing or with and you would get the corresponding "bitmask-all". Would it be worthwhile to add it? I don't have a use case, but the symmetry seems somewhat compelling.

Yes, that's correct, and I agree. There are actually 4 similar patterns: any-set, all-set, any-clear, all-clear. We should recognize all of them. I didn't want to go overboard on this patch though because I wasn't sure if this implementation would receive support. At the least, I will add a 'TODO' comment to this patch. Then, we can follow-up to catch the other patterns.

[kparzysz]

I think we should recognize as much as we can out of the common idioms, and this is one of them. Looks good to me.

lib/Transforms/AggressiveInstCombine/AggressiveInstCombine.cpp

76

You could also match arithmetic shifts-right as long as the shift amount does not move the sign bit into the 0th position. It's not necessary for this patch though.

[spatel]

lib/Transforms/AggressiveInstCombine/AggressiveInstCombine.cpp

76

That's a good point. I am assuming this pass runs after instcombine, and instcombine would transform to lshr as a canonicalization step. This is as it happens currently, but might change as you've proposed.

I'll look at adding a phase-ordering test, so we'll know that we don't lose this optimization if the pass order changes.

lib/Transforms/AggressiveInstCombine/AggressiveInstCombine.cpp

 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/AggressiveInstCombine/AggressiveInstCombine.h"
 #include "AggressiveInstCombineInternal.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/BasicAliasAnalysis.h"
 #include "llvm/Analysis/GlobalsModRef.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
+#include "llvm/Analysis/Utils/Local.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Dominators.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/PatternMatch.h"
 #include "llvm/Pass.h"
 #include "llvm/Transforms/Scalar.h"
 using namespace llvm;
-
+using namespace PatternMatch;
 #define DEBUG_TYPE "aggressive-instcombine"
 
 namespace {
 /// Contains expression pattern combiner logic.
 /// This class provides both the logic to combine expression patterns and
 /// combine them. It differs from InstCombiner class in that each pattern
 /// combiner runs only once as opposed to InstCombine's multi-iteration,
 /// which allows pattern combiner to have higher complexity than the O(1)
public:
   /// Run all expression pattern optimizations on the given /p F function.
   ///
   /// \param F function to optimize.
   /// \returns true if the IR is changed.
   bool runOnFunction(Function &F) override;
 };
 } // 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.
+  Value *Op0, *Op1;
+  if (match(V, m_Or(m_Value(Op0), m_Value(Op1))))
+    return matchMaskedCmpOp(Op0, Result) && matchMaskedCmpOp(Op1, Result);
+
+  // We need a shift-right or a bare value representing a compare of bit 0 of
+  // the original source operand.
+  Value *Candidate;
+  uint64_t BitIndex = 0;
+  if (!match(V, m_LShr(m_Value(Candidate), m_ConstantInt(BitIndex))))
+    Candidate = V;
+
+  // Initialize result source operand.
+  if (!Result.first)
+    Result.first = Candidate;
+
+  // Fill in the mask bit derived from the shift constant.
+  Result.second |= (1 << BitIndex);
+  return Result.first == 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) {
+  if (!match(&I, m_And(m_OneUse(m_Or(m_Value(), m_Value())), m_One())))
+    return false;
+
+  std::pair<Value *, APInt>
+  MaskOps(nullptr, APInt::getNullValue(I.getType()->getScalarSizeInBits()));
+  if (!matchMaskedCmpOp(cast<BinaryOperator>(&I)->getOperand(0), MaskOps))
+    return false;
+
+  IRBuilder<> Builder(&I);
+  Value *Mask = Builder.CreateAnd(MaskOps.first, MaskOps.second);
+  Value *CmpZero = Builder.CreateIsNotNull(Mask);
+  Value *Zext = Builder.CreateZExt(CmpZero, I.getType());
+  I.replaceAllUsesWith(Zext);
+  return true;
+}
+
+/// This is the entry point for folds that could be implemented in regular
+/// InstCombine, but they are separated because they are not expected to
+/// occur frequently and/or have more than a constant-length pattern match.
+static bool foldUnusualPatterns(Function &F, DominatorTree &DT) {
+  bool MadeChange = false;
+  for (BasicBlock &BB : F) {
+    // Ignore unreachable basic blocks.
+    if (!DT.isReachableFromEntry(&BB))
+      continue;
+    // Do not delete instructions under here and invalidate the iterator.
+    for (Instruction &I : BB)
+      MadeChange |= foldToMaskedCmp(I);
+  }
+
+  // We're done with transforms, so remove dead instructions.
+  if (MadeChange)
+    for (BasicBlock &BB : F)
+      SimplifyInstructionsInBlock(&BB);
+
+  return MadeChange;
+}
+
+/// This is the entry point for all transforms. Pass manager differences are
+/// handled in the callers of this function.
+static bool runImpl(Function &F, TargetLibraryInfo &TLI, DominatorTree &DT) {
+  bool MadeChange = false;
+  const DataLayout &DL = F.getParent()->getDataLayout();
+  TruncInstCombine TIC(TLI, DL, DT);
+  MadeChange |= TIC.run(F);
+  MadeChange |= foldUnusualPatterns(F, DT);
+  return MadeChange;
+}
+
 void AggressiveInstCombinerLegacyPass::getAnalysisUsage(
     AnalysisUsage &AU) const {
   AU.setPreservesCFG();
   AU.addRequired<DominatorTreeWrapperPass>();
   AU.addRequired<TargetLibraryInfoWrapperPass>();
   AU.addPreserved<AAResultsWrapperPass>();
   AU.addPreserved<BasicAAWrapperPass>();
   AU.addPreserved<DominatorTreeWrapperPass>();
   AU.addPreserved<GlobalsAAWrapperPass>();
 }
 
 bool AggressiveInstCombinerLegacyPass::runOnFunction(Function &F) {
-  auto &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
   auto &TLI = getAnalysis<TargetLibraryInfoWrapperPass>().getTLI();
-  auto &DL = F.getParent()->getDataLayout();
-
-  bool MadeIRChange = false;
-
-  // Handle TruncInst patterns
-  TruncInstCombine TIC(TLI, DL, DT);
-  MadeIRChange |= TIC.run(F);
-
-  // TODO: add more patterns to handle...
-
-  return MadeIRChange;
+  auto &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
+  return runImpl(F, TLI, DT);
 }
 
 PreservedAnalyses AggressiveInstCombinePass::run(Function &F,
                                                  FunctionAnalysisManager &AM) {
-  auto &DT = AM.getResult<DominatorTreeAnalysis>(F);
   auto &TLI = AM.getResult<TargetLibraryAnalysis>(F);
-  auto &DL = F.getParent()->getDataLayout();
-  bool MadeIRChange = false;
-
-  // Handle TruncInst patterns
-  TruncInstCombine TIC(TLI, DL, DT);
-  MadeIRChange |= TIC.run(F);
-  if (!MadeIRChange)
+  auto &DT = AM.getResult<DominatorTreeAnalysis>(F);
+  if (!runImpl(F, TLI, DT)) {
     // No changes, all analyses are preserved.
     return PreservedAnalyses::all();
-
+  }
   // Mark all the analyses that instcombine updates as preserved.
   PreservedAnalyses PA;
   PA.preserveSet<CFGAnalyses>();
   PA.preserve<AAManager>();
   PA.preserve<GlobalsAA>();
   return PA;
 }
 

test/Transforms/AggressiveInstCombine/masked-cmp.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -aggressive-instcombine -S | FileCheck %s
 
 ; PR37098 - https://bugs.llvm.org/show_bug.cgi?id=37098
 
 define i32 @two_bit_mask(i32 %x) {
 ; CHECK-LABEL: @two_bit_mask(
-; CHECK-NEXT:    [[S:%.*]] = lshr i32 [[X:%.*]], 3
-; CHECK-NEXT:    [[O:%.*]] = or i32 [[S]], [[X]]
-; CHECK-NEXT:    [[R:%.*]] = and i32 [[O]], 1
-; CHECK-NEXT:    ret i32 [[R]]
+; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[X:%.*]], 9
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ne i32 [[TMP1]], 0
+; CHECK-NEXT:    [[TMP3:%.*]] = zext i1 [[TMP2]] to i32
+; CHECK-NEXT:    ret i32 [[TMP3]]
 ;
   %s = lshr i32 %x, 3
   %o = or i32 %s, %x
   %r = and i32 %o, 1
   ret i32 %r
 }
 
 define i32 @four_bit_mask(i32 %x) {
 ; CHECK-LABEL: @four_bit_mask(
-; CHECK-NEXT:    [[T1:%.*]] = lshr i32 [[X:%.*]], 3
-; CHECK-NEXT:    [[T2:%.*]] = lshr i32 [[X]], 5
-; CHECK-NEXT:    [[T3:%.*]] = lshr i32 [[X]], 8
-; CHECK-NEXT:    [[O1:%.*]] = or i32 [[T1]], [[X]]
-; CHECK-NEXT:    [[O2:%.*]] = or i32 [[T2]], [[T3]]
-; CHECK-NEXT:    [[O3:%.*]] = or i32 [[O1]], [[O2]]
-; CHECK-NEXT:    [[R:%.*]] = and i32 [[O3]], 1
-; CHECK-NEXT:    ret i32 [[R]]
+; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[X:%.*]], 297
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ne i32 [[TMP1]], 0
+; CHECK-NEXT:    [[TMP3:%.*]] = zext i1 [[TMP2]] to i32
+; CHECK-NEXT:    ret i32 [[TMP3]]
 ;
   %t1 = lshr i32 %x, 3
   %t2 = lshr i32 %x, 5
   %t3 = lshr i32 %x, 8
   %o1 = or i32 %t1, %x
   %o2 = or i32 %t2, %t3
   %o3 = or i32 %o1, %o2
   %r = and i32 %o3, 1
   ret i32 %r
 }