[AggressiveInstCombine] Inline strcmp/strncmp

Inline calls to strcmp(s1, s2) and strncmp(s1, s2, N),
where N and exactly one of s1 and s2 are constant.

Author: FLZ101

llvm/lib/Transforms/AggressiveInstCombine/AggressiveInstCombine.cpp

@@ -19,6 +19,7 @@
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/BasicAliasAnalysis.h"
 #include "llvm/Analysis/ConstantFolding.h"
+#include "llvm/Analysis/DomTreeUpdater.h"
 #include "llvm/Analysis/GlobalsModRef.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
@@ -28,6 +29,7 @@
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/PatternMatch.h"
+#include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/BuildLibCalls.h"
 #include "llvm/Transforms/Utils/Local.h"
 
@@ -922,6 +924,251 @@ static bool foldPatternedLoads(Instruction &I, const DataLayout &DL) {
   return true;
 }
 
+static cl::opt<unsigned> StrNCmpInlineThreshold(
+    "strncmp-inline-threshold", cl::init(3), cl::Hidden,
+    cl::desc("The maximum length of a constant string for a builtin string cmp "
+             "call eligible for inlining. The default value is 3."));
+
+namespace {
+class StrNCmpInliner {
+public:
+  StrNCmpInliner(CallInst *CI, LibFunc Func, Function::iterator &BBNext,
+                 DomTreeUpdater *DTU, const DataLayout &DL)
+      : CI(CI), Func(Func), BBNext(BBNext), DTU(DTU), DL(DL) {}
+
+  bool optimizeStrNCmp();
+
+private:
+  bool inlineCompare(Value *LHS, StringRef RHS, uint64_t N, bool Switched);
+
+  CallInst *CI;
+  LibFunc Func;
+  Function::iterator &BBNext;
+  DomTreeUpdater *DTU;
+  const DataLayout &DL;
+};
+
+} // namespace
+
+/// First we normalize calls to strncmp/strcmp to the form of
+/// compare(s1, s2, N), which means comparing first N bytes of s1 and s2
+/// (without considering '\0')
+///
+/// Examples:
+///
+/// \code
+///   strncmp(s, "a", 3) -> compare(s, "a", 2)
+///   strncmp(s, "abc", 3) -> compare(s, "abc", 3)
+///   strncmp(s, "a\0b", 3) -> compare(s, "a\0b", 2)
+///   strcmp(s, "a") -> compare(s, "a", 2)
+///
+///   char s2[] = {'a'}
+///   strncmp(s, s2, 3) -> compare(s, s2, 3)
+///
+///   char s2[] = {'a', 'b', 'c', 'd'}
+///   strncmp(s, s2, 3) -> compare(s, s2, 3)
+/// \endcode
+///
+/// We only handle cases that N and exactly one of s1 and s2 are constant. Cases
+/// that s1 and s2 are both constant are already handled by the instcombine
+/// pass.
+///
+/// We do not handle cases that N > StrNCmpInlineThreshold.
+///
+/// We also do not handles cases that N < 2, which are already
+/// handled by the instcombine pass.
+///
+bool StrNCmpInliner::optimizeStrNCmp() {
+  if (StrNCmpInlineThreshold < 2)
+    return false;
+
+  Value *Str1P = CI->getArgOperand(0);
+  Value *Str2P = CI->getArgOperand(1);
+  // should be handled elsewhere
+  if (Str1P == Str2P)
+    return false;
+
+  StringRef Str1, Str2;
+  bool HasStr1 = getConstantStringInfo(Str1P, Str1, false);
+  bool HasStr2 = getConstantStringInfo(Str2P, Str2, false);
+  if (!(HasStr1 ^ HasStr2))
+    return false;
+
+  // note that '\0' and characters after it are not trimmed
+  StringRef Str = HasStr1 ? Str1 : Str2;
+
+  size_t Idx = Str.find('\0');
+  uint64_t N = Idx == StringRef::npos ? UINT64_MAX : Idx + 1;
+  if (Func == LibFunc_strncmp) {
+    if (!isa<ConstantInt>(CI->getArgOperand(2)))
+      return false;
+    N = std::min(N, cast<ConstantInt>(CI->getArgOperand(2))->getZExtValue());
+  }
+  // now N means how many bytes we need to compare at most
+  if (N > Str.size() || N < 2 || N > StrNCmpInlineThreshold)
+    return false;
+
+  Value *StrP = HasStr1 ? Str2P : Str1P;
+
+  // cases that StrP has two or more dereferenceable bytes might be better
+  // optimized elsewhere
+  bool CanBeNull = false, CanBeFreed = false;
+  if (StrP->getPointerDereferenceableBytes(DL, CanBeNull, CanBeFreed) > 1)
+    return false;
+
+  return inlineCompare(StrP, Str, N, HasStr1);
+}
+
+/// Convert
+///
+/// \code
+///   ret = compare(s1, s2, N)
+/// \endcode
+///
+/// into
+///
+/// \code
+///   ret = (int)s1[0] - (int)s2[0]
+///   if (ret != 0)
+///     goto NE
+///   ...
+///   ret = (int)s1[N-2] - (int)s2[N-2]
+///   if (ret != 0)
+///     goto NE
+///   ret = (int)s1[N-1] - (int)s2[N-1]
+///   NE:
+/// \endcode
+///
+/// CFG before and after the transformation:
+///
+/// (before)
+/// BBCI
+///
+/// (after)
+/// BBBefore -> BBSubs[0] (sub,icmp) --NE-> BBNE -> BBCI
+///                      |                    ^
+///                      E                    |
+///                      |                    |
+///             BBSubs[1] (sub,icmp) --NE-----+
+///                     ...                   |
+///             BBSubs[N-1]    (sub) ---------+
+///
+bool StrNCmpInliner::inlineCompare(Value *LHS, StringRef RHS, uint64_t N,
+                                   bool Switched) {
+  IRBuilder<> B(CI->getContext());
+
+  BasicBlock *BBCI = CI->getParent();
+  bool IsEntry = BBCI->isEntryBlock();
+  BasicBlock *BBBefore = splitBlockBefore(BBCI, CI, DTU, nullptr, nullptr,
+                                          BBCI->getName() + ".before");
+
+  SmallVector<BasicBlock *> BBSubs;
+  for (uint64_t i = 0; i < N + 1; ++i)
+    BBSubs.push_back(
+        BasicBlock::Create(CI->getContext(), "sub", BBCI->getParent(), BBCI));
+  BasicBlock *BBNE = BBSubs[N];
+
+  cast<BranchInst>(BBBefore->getTerminator())->setSuccessor(0, BBSubs[0]);
+
+  B.SetInsertPoint(BBNE);
+  PHINode *Phi = B.CreatePHI(CI->getType(), N);
+  B.CreateBr(BBCI);
+
+  Value *Base = LHS;
+  for (uint64_t i = 0; i < N; ++i) {
+    B.SetInsertPoint(BBSubs[i]);
+    Value *VL = B.CreateZExt(
+        B.CreateLoad(B.getInt8Ty(),
+                     B.CreateInBoundsGEP(B.getInt8Ty(), Base, B.getInt64(i))),
+        CI->getType());
+    Value *VR = ConstantInt::get(CI->getType(), RHS[i]);
+    Value *Sub = Switched ? B.CreateSub(VR, VL) : B.CreateSub(VL, VR);
+    if (i < N - 1)
+      B.CreateCondBr(B.CreateICmpNE(Sub, ConstantInt::get(CI->getType(), 0)),
+                     BBNE, BBSubs[i + 1]);
+    else
+      B.CreateBr(BBNE);
+
+    Phi->addIncoming(Sub, BBSubs[i]);
+  }
+
+  CI->replaceAllUsesWith(Phi);
+  CI->eraseFromParent();
+
+  BBNext = BBCI->getIterator();
+
+  // Update DomTree
+  if (DTU) {
+    if (IsEntry) {
+      DTU->recalculate(*BBCI->getParent());
+    } else {
+      SmallVector<DominatorTree::UpdateType, 8> Updates;
+      Updates.push_back({DominatorTree::Delete, BBBefore, BBCI});
+      Updates.push_back({DominatorTree::Insert, BBBefore, BBSubs[0]});
+      for (uint64_t i = 0; i < N; ++i) {
+        if (i < N - 1)
+          Updates.push_back({DominatorTree::Insert, BBSubs[i], BBSubs[i + 1]});
+        Updates.push_back({DominatorTree::Insert, BBSubs[i], BBNE});
+      }
+      Updates.push_back({DominatorTree::Insert, BBNE, BBCI});
+      DTU->applyUpdates(Updates);
+    }
+  }
+  return true;
+}
+
+static bool inlineLibCalls(Function &F, TargetLibraryInfo &TLI,
+                           const TargetTransformInfo &TTI, DominatorTree &DT,
+                           bool &MadeCFGChange) {
+  MadeCFGChange = false;
+  DomTreeUpdater DTU(&DT, DomTreeUpdater::UpdateStrategy::Lazy);
+
+  bool MadeChange = false;
+
+  Function::iterator CurrBB;
+  for (Function::iterator BB = F.begin(), BE = F.end(); BB != BE;) {
+    CurrBB = BB++;
+
+    for (BasicBlock::iterator II = CurrBB->begin(), IE = CurrBB->end();
+         II != IE; ++II) {
+      CallInst *Call = dyn_cast<CallInst>(&*II);
+      Function *CalledFunc;
+
+      if (!Call || !(CalledFunc = Call->getCalledFunction()))
+        continue;
+
+      if (Call->isNoBuiltin())
+        continue;
+
+      // Skip if function either has local linkage or is not a known library
+      // function.
+      LibFunc LF;
+      if (CalledFunc->hasLocalLinkage() || !TLI.getLibFunc(*CalledFunc, LF) ||
+          !TLI.has(LF))
+        continue;
+
+      switch (LF) {
+      case LibFunc_strcmp:
+      case LibFunc_strncmp: {
+        auto &DL = F.getParent()->getDataLayout();
+        if (StrNCmpInliner(Call, LF, BB, &DTU, DL).optimizeStrNCmp()) {
+          MadeCFGChange = true;
+          break;
+        }
+        continue;
+      }
+      default:
+        continue;
+      }
+
+      MadeChange = true;
+      break;
+    }
+  }
+
+  return MadeChange;
+}
+
 /// 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.
@@ -969,11 +1216,12 @@ static bool foldUnusualPatterns(Function &F, DominatorTree &DT,
 /// handled in the callers of this function.
 static bool runImpl(Function &F, AssumptionCache &AC, TargetTransformInfo &TTI,
                     TargetLibraryInfo &TLI, DominatorTree &DT,
-                    AliasAnalysis &AA) {
+                    AliasAnalysis &AA, bool &MadeCFGChange) {
   bool MadeChange = false;
   const DataLayout &DL = F.getParent()->getDataLayout();
   TruncInstCombine TIC(AC, TLI, DL, DT);
   MadeChange |= TIC.run(F);
+  MadeChange |= inlineLibCalls(F, TLI, TTI, DT, MadeCFGChange);
   MadeChange |= foldUnusualPatterns(F, DT, TTI, TLI, AA, AC);
   return MadeChange;
 }
@@ -985,12 +1233,16 @@ PreservedAnalyses AggressiveInstCombinePass::run(Function &F,
   auto &DT = AM.getResult<DominatorTreeAnalysis>(F);
   auto &TTI = AM.getResult<TargetIRAnalysis>(F);
   auto &AA = AM.getResult<AAManager>(F);
-  if (!runImpl(F, AC, TTI, TLI, DT, AA)) {
+  bool MadeCFGChange = false;
+  if (!runImpl(F, AC, TTI, TLI, DT, AA, MadeCFGChange)) {
     // No changes, all analyses are preserved.
     return PreservedAnalyses::all();
   }
   // Mark all the analyses that instcombine updates as preserved.
   PreservedAnalyses PA;
-  PA.preserveSet<CFGAnalyses>();
+  if (MadeCFGChange)
+    PA.preserve<DominatorTreeAnalysis>();
+  else
+    PA.preserveSet<CFGAnalyses>();
   return PA;
 }