Add 'strlen' formation to LoopIdiomRecognize

[nlewycky]

Summary

This recognizer finds loops of the form:

for (const char *ptr = p; *ptr; ++ptr, ++len)
  /*empty*/ ;

and converts them into calls to strlen().

[Prazek]

I don't think I have enough power to say LGTM, but looks legit to me.

lib/Transforms/Scalar/LoopIdiomRecognize.cpp

1105

I think it would look better if You would put this in the same line. Now it somehow look like
for (... );

IndVar;
IndVar = ...

[hfinkel]

lib/Transforms/Scalar/LoopIdiomRecognize.cpp

1081

Do you also need to skip debug intrinsics, etc.?

1119

Can you please add a test case for this variant with vectors?

[sanjoy]

lib/Transforms/Scalar/LoopIdiomRecognize.cpp

64

I'd open the namespace in just the one function that uses it.

205

Why this->?

1043

Why not declare Ptr as a LoadInst *? Then you won't need this second cast<Instruction>(.

1045

Why not isIntegerTy(8)?

1047

auto * might be cleaner.

1052

auto *?

1057

Why not use a range for here?

1072

How can Start ever not be loop invariant?

My reasoning went like this: it'll have to dominate the entry edge, and therefore the preheader, but if it is defined inside the only BB, then the only BB == the header dominates the preheader, making it part of the loop.

1083

What if these have uses outside the loop? Also why these specifically? Why not also handle, say, mul?

1087

Why not use a range for loop?

1089

Why not

if (auto *OpI = dyn_cast<..)
  if (OpI->getParent() != ...
    ...

Alternatively, why not use Loop::hasLoopInvariantOperands ?

[sanjoy]

lib/Transforms/Scalar/LoopIdiomRecognize.cpp

205

Never mind this one -- I missed the LPM on both sides initially.

[sanjoy]

lib/Transforms/Scalar/LoopIdiomRecognize.cpp

1083

What if these have uses outside the loop?

I'll redact this bit -- I had initially missed what you were doing with the splicing logic below.

[nlewycky]

lib/Transforms/Scalar/LoopIdiomRecognize.cpp

1043

Because then it can't be used in m_Value(Load).

Do you think I should I add an m_Load instead?

1045

'i8*' not 'i8'. I could do "cast<PointerType>(Ptr->getType())->getElementType()->isIntegerTy(8)" if you prefer?

1057

I wanted to exclude the terminator.

I could add a simple "if (isa<TerminatorInst>(Inst)) continue;".

I don't think we can use "E = llvm::prior(LoopBody->end())" because it may walk off the beginning of the block. At this point, we haven't yet shown that LoopCond, PtrGEP or PtrPHI are actually members instructions inside LoopBody.

1072

I concur. This is now an assertion.

1081

Also, I'm not updating DebugLoc anywhere in this code. If there's debug intrinsics they'll move along with the splice, which is probably close enough to right, as much as any tracking for optimized debug info.

1083

Gave this some thought, decided we can safely extend it.

1087

Not applicable.

[sanjoy]

lib/Transforms/Scalar/LoopIdiomRecognize.cpp

1043

Do you think I should I add an m_Load instead?

I think that's a good idea, it will be generally useful.

1045

I think what you have is fine. Sorry for misreading the code originally.

1057

I could add a simple "if (isa<TerminatorInst>(Inst)) continue;".

I mildly prefer that.

1087

I may be missing something here, but can't you do for (Value *Op : Inst->operands())?

[alexr]

This makes me very nervous. We have fielded a lot of user complaints about canonicalizing loops into memcpy calls. strlen feels like a more dubious choice.

My past experience writing SIMD versions of string functions tells me that the typical string length that strlen is tuned for may be very different from the amount of data an open-coded loop is intended to handle.

I would want to see very tuned lowering code to go with this. I'd also very much like to see the lowering code pay attention to target-specific cost information such as the possible cost to cross a dynamic library boundary when weighing the lowering of a call vs expanding inline.

lib/Transforms/Scalar/LoopIdiomRecognize.cpp

 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Module.h"
+#include "llvm/IR/PatternMatch.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/Transforms/Utils/BuildLibCalls.h"
 #include "llvm/Transforms/Utils/Local.h"
 using namespace llvm;
+using namespace llvm::PatternMatch;
 
 #define DEBUG_TYPE "loop-idiom"
 
 STATISTIC(NumMemSet, "Number of memset's formed from loop stores");
 STATISTIC(NumMemCpy, "Number of memcpy's formed from loop load+stores");
 
 namespace {
 
 class LoopIdiomRecognize : public LoopPass {
   Loop *CurLoop;
   AliasAnalysis *AA;
   DominatorTree *DT;
   LoopInfo *LI;
   ScalarEvolution *SE;
   TargetLibraryInfo *TLI;
   const TargetTransformInfo *TTI;
+  LPPassManager *LPM;
 
 public:
   static char ID;
   explicit LoopIdiomRecognize() : LoopPass(ID) {
     initializeLoopIdiomRecognizePass(*PassRegistry::getPassRegistry());
   }
 
   bool runOnLoop(Loop *L, LPPassManager &LPM) override;
private:
   /// @{
 
   bool runOnNoncountableLoop();
 
   bool recognizePopcount();
   void transformLoopToPopcount(BasicBlock *PreCondBB, Instruction *CntInst,
                                PHINode *CntPhi, Value *Var);
 
+  bool recognizeStrlen();
+
   /// @}
 };
 
 } // End anonymous namespace.
 
 char LoopIdiomRecognize::ID = 0;
 INITIALIZE_PASS_BEGIN(LoopIdiomRecognize, "loop-idiom", "Recognize loop idioms",
                       false, false)
bool LoopIdiomRecognize::runOnLoop(Loop *L, LPPassManager &LPM) {
   CurLoop = L;
   // If the loop could not be converted to canonical form, it must have an
   // indirectbr in it, just give up.
   if (!L->getLoopPreheader())
     return false;
 
   // Disable loop idiom recognition if the function's name is a common idiom.
   StringRef Name = L->getHeader()->getParent()->getName();
-  if (Name == "memset" || Name == "memcpy")
+  if (Name == "memset" || Name == "memcpy" || Name == "strlen")
     return false;
 
   AA = &getAnalysis<AliasAnalysis>();
   DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
   LI = &getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
   SE = &getAnalysis<ScalarEvolutionWrapperPass>().getSE();
   TLI = &getAnalysis<TargetLibraryInfoWrapperPass>().getTLI();
   TTI = &getAnalysis<TargetTransformInfoWrapperPass>().getTTI(
       *CurLoop->getHeader()->getParent());
+  this->LPM = &LPM;
 
   if (SE->hasLoopInvariantBackedgeTakenCount(L))
     return runOnCountableLoop();
 
   return runOnNoncountableLoop();
 }
 
 bool LoopIdiomRecognize::runOnCountableLoop() {
bool LoopIdiomRecognize::processLoopStoreOfLoopLoad(
   ++NumMemCpy;
   return true;
 }
 
 bool LoopIdiomRecognize::runOnNoncountableLoop() {
   if (recognizePopcount())
     return true;
 
+  if (recognizeStrlen())
+    return true;
+
   return false;
 }
 
 /// Check if the given conditional branch is based on the comparison between
 /// a variable and zero, and if the variable is non-zero, the control yields to
 /// the loop entry. If the branch matches the behavior, the variable involved
 /// in the comparion is returned. This function will be called to see if the
 /// precondition and postcondition of the loop are in desirable form.
void LoopIdiomRecognize::transformLoopToPopcount(BasicBlock *PreCondBB,
   //  the loop are replaced with the NewCount -- the value returned from
   //  __builtin_ctpop().
   CntInst->replaceUsesOutsideBlock(NewCount, Body);
 
   // step 5: Forget the "non-computable" trip-count SCEV associated with the
   //   loop. The loop would otherwise not be deleted even if it becomes empty.
   SE->forgetLoop(CurLoop);
 }
+
+bool LoopIdiomRecognize::recognizeStrlen() {
+  // If we're not allowed to introduce a strlen, don't try.
+  if (!TLI->has(LibFunc::strlen))
+    return false;
+
+  // Give up if the loop has multiple blocks or multiple backedges.
+  if (CurLoop->getNumBackEdges() != 1 || CurLoop->getNumBlocks() != 1)
+    return false;
+
+  BasicBlock *LoopBody = *(CurLoop->block_begin());
+  if (LoopBody->size() >= 20) {
+    // The loop is too big, bail out.
+    return false;
+  }
+
+  // If successful, there won't be a loop to take the address of.
+  if (LoopBody->hasAddressTaken())
+    return false;
+
+  // It should have a preheader containing nothing but an unconditional branch.
+  BasicBlock *PH = CurLoop->getLoopPreheader();
+  if (!PH)
+    return false;
+  if (&PH->front() != PH->getTerminator())
+    return false;
+  auto *EntryBI = dyn_cast<BranchInst>(PH->getTerminator());
+  if (!EntryBI || EntryBI->isConditional())
+    return false;
+
+  // Check that this loop has a pointer that increments by one byte each
+  // iteration and exits the loop when a load of our pointer returns i8 0.
+  auto *Backedge = dyn_cast<BranchInst>(LoopBody->getTerminator());
+  if (!Backedge || Backedge->isUnconditional())
+    return false;
+  auto *LoopCond = Backedge->getCondition();
+  Value *Load, *Ptr;
+  CmpInst::Predicate Pred;
+  if (!(match(LoopCond, m_ICmp(Pred, m_Value(Load), m_Zero())) &&
+        Pred == CmpInst::ICMP_EQ && isa<LoadInst>(Load) &&
+        (Ptr = cast<LoadInst>(Load)->getPointerOperand()) &&
+        cast<Instruction>(Ptr)->getParent() == LoopBody))
+    return false;
+  if (Ptr->getType() != Type::getInt8PtrTy(Ptr->getContext()))
+    return false;
+  GetElementPtrInst *PtrGEP = dyn_cast<GetElementPtrInst>(Ptr);
+  auto &DL = CurLoop->getHeader()->getModule()->getDataLayout();
+  APInt Offset(DL.getPointerTypeSizeInBits(Ptr->getType()), 0);
+  if (!PtrGEP || !PtrGEP->accumulateConstantOffset(DL, Offset) || Offset != 1)
+    return false;
+  PHINode *PtrPHI = dyn_cast<PHINode>(PtrGEP->getPointerOperand());
+  if (!PtrPHI || PtrPHI->getIncomingValueForBlock(LoopBody) != PtrGEP)
+    return false;
+
+  // Check that this loop does nothing else each iteration.
+  for (Instruction *Inst = LoopBody->begin(); !isa<TerminatorInst>(Inst);
+       Inst = Inst->getNextNode()) {
+    if (Inst == PtrPHI || Inst == Load || Inst == LoopCond) {
+      if (!Inst->hasOneUse())
+        return false;
+      continue;
+    }
+    if (!Inst->getType()->isIntOrIntVectorTy() &&
+        !Inst->getType()->isPointerTy())
+      return false;
+    if (Inst->mayHaveSideEffects())
+      return false;
+    if (PHINode *IndVar = dyn_cast<PHINode>(Inst)) {
+      Value *Start = IndVar->getIncomingValueForBlock(PH);
+      if (!CurLoop->isLoopInvariant(Start))
+        return false;
+      Value *Step;
+      if (!match(IndVar->getIncomingValueForBlock(LoopBody),
+                 m_Add(m_Specific(IndVar), m_Value(Step))))
+        return false;
+      if (!CurLoop->isLoopInvariant(Step))
+        return false;
+      continue;
+    }
+    switch (Inst->getOpcode()) {
+    case Instruction::Add:
+    case Instruction::GetElementPtr:
+    case Instruction::ZExt:
+    case Instruction::SExt:
+    case Instruction::Trunc:
+      for (unsigned i = 0, e = Inst->getNumOperands(); i != e; ++i) {
+        Instruction *Op = dyn_cast<Instruction>(Inst->getOperand(i));
+        if (!Op) continue;
+        if (Op->getParent() != LoopBody)
+          return false;
+      }
+      break;
+    default:
+      return false;
+    }
+  }
+
+  // Rewrite the integer counters in terms of strlen().
+  IRBuilder<> Builder(PH->getTerminator());
+  Value *StrLen = llvm::EmitStrLen(PtrPHI->getIncomingValueForBlock(PH),
+                                   Builder, DL, TLI);
+  StrLen = Builder.CreateSub(StrLen, ConstantInt::get(StrLen->getType(), 1));
+  for (PHINode *IndVar = dyn_cast<PHINode>(LoopBody->begin());
+       IndVar;
+       IndVar = dyn_cast<PHINode>(IndVar->getNextNode())) {
+    if (IndVar == PtrPHI)
+      continue;
+
+    Value *Start = IndVar->getIncomingValueForBlock(PH);
+    Value *Step;
+    if (!match(IndVar->getIncomingValueForBlock(LoopBody),
+               m_Add(m_Specific(IndVar), m_Value(Step))))
+      llvm_unreachable("Failed to find loop step that we already found");
+
+    Value *StrLenV =
+        Builder.CreateZExtOrTrunc(StrLen, IndVar->getType()->getScalarType());
+    if (VectorType *VTy = dyn_cast<VectorType>(IndVar->getType()))
+      StrLenV = Builder.CreateVectorSplat(VTy->getNumElements(), StrLenV);
+    Value *Result = Builder.CreateAdd(Start, Builder.CreateMul(Step, StrLenV));
+
+    IndVar->replaceAllUsesWith(Result);
+  }
+
+  BasicBlock *ExitBlock = CurLoop->getExitBlock();
+  assert(ExitBlock && "Loop ends with branch but has multiple exit blocks?");
+
+  SE->forgetLoop(CurLoop);
+  LPM->deleteLoopFromQueue(CurLoop);
+
+  // We know these instructions are certainly dead. So is the pointer increment,
+  // but we didn't keep a Value* around for that.
+  LoopCond->replaceAllUsesWith(UndefValue::get(LoopCond->getType()));
+  Load->replaceAllUsesWith(UndefValue::get(Load->getType()));
+  PtrPHI->replaceAllUsesWith(UndefValue::get(PtrPHI->getType()));
+  cast<Instruction>(LoopCond)->eraseFromParent();
+  cast<Instruction>(Load)->eraseFromParent();
+  PtrPHI->eraseFromParent();
+
+  PH->getInstList().splice(Builder.GetInsertPoint(), LoopBody->getInstList(),
+                           LoopBody->getFirstNonPHI(),
+                           LoopBody->getTerminator());
+
+  // At this stage, LoopBody contains only use_empty PHI nodes and a terminator.
+
+  LoopBody->replaceSuccessorsPhiUsesWith(PH);
+  LoopBody->replaceAllUsesWith(ExitBlock);
+  LoopBody->eraseFromParent();
+  
+  return true;
+}

test/Transforms/LoopIdiom/strlen.ll

+; RUN: opt -loop-idiom -S < %s | FileCheck %s
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+define i32 @test1(i8* %p) {
+; CHECK-LABEL: @test1
+entry:
+  %0 = load i8, i8* %p, align 1
+  %tobool.5 = icmp eq i8 %0, 0
+  br i1 %tobool.5, label %for.cond.cleanup, label %for.inc.lr.ph
+
+for.inc.lr.ph:                                    ; preds = %entry
+; CHECK: for.inc.lr.ph:
+; CHECK-NEXT: @strlen
+; CHECK-NOT: icmp
+  br label %for.inc
+
+for.cond.for.cond.cleanup_crit_edge:              ; preds = %for.inc
+  %inc.lcssa = phi i32 [ %inc, %for.inc ]
+  br label %for.cond.cleanup
+
+for.cond.cleanup:                                 ; preds = %for.cond.for.cond.cleanup_crit_edge, %entry
+  %len.0.lcssa = phi i32 [ %inc.lcssa, %for.cond.for.cond.cleanup_crit_edge ], [ 0, %entry ]
+  ret i32 %len.0.lcssa
+
+for.inc:                                          ; preds = %for.inc.lr.ph, %for.inc
+  %ptr.07 = phi i8* [ %p, %for.inc.lr.ph ], [ %incdec.ptr, %for.inc ]
+  %len.06 = phi i32 [ 0, %for.inc.lr.ph ], [ %inc, %for.inc ]
+  %incdec.ptr = getelementptr inbounds i8, i8* %ptr.07, i64 1
+  %inc = add nuw nsw i32 %len.06, 1
+  %1 = load i8, i8* %incdec.ptr, align 1
+  %tobool = icmp eq i8 %1, 0
+  br i1 %tobool, label %for.cond.for.cond.cleanup_crit_edge, label %for.inc
+}