Strlen loop idiom recognition

[musman]

Summary

This patch adds strlen to LoopIdiomRecognize.cpp. It is the first part of 3 patches:

1. Strlen loop idiom recognition
2. strlen16 recognition and creation of new strlen16 intrinsic
3. Folding of strlen/strlen16 call if they are only used for zero equality comparison (replace with load of first char)

Example that this recognizes:

unsigned strlen(char *Str) {
  char *Src = Str;
  if (!Src)
    return 0;
  for (; *Src;)
    Src++;
  return Src - Str;
}

[efriedma]

I don't like the way the pattern-matching is written here; analyzing the exact instruction patterns is complicated, error-prone, and likely to miss cases.

I think I'd structure the analysis like this:

1. Is the latch condition of the form "%loadedval = load i8, i8* %p; icmp eq i8 %loadedval, 0"
2. Query SCEV to check %p is an AddRec with step 1.
3. Query SCEV to check all the LCSSA PHI nodes are AddRecs.
4. Check that the loop doesn't contain any operations with side-effects.

Then to apply the transform, you take the SCEV for the LCSSA PHI node, transform the AddRec to an Add (for example, {n,+,1} to n + strlen(p)), then use SCEVExpander to expand it.

Leveraging SCEV like this will make the transform a lot more flexible and easier to read.

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

1606

"BasicBlock::size()" should never be used as a threshold; it's sensitive to debug info.

[musman]

@efriedma Thanks for the review, I didn't think to use SCEV for this. I will modify the code as you suggested.

[shawnl]

This comment was removed by shawnl.

[musman]

Hi @efriedma, I am having some trouble understanding how to transform the AddRec to Add ({n,+,1} to n + strlen(p)). Should this be done with ScalarEvolution::getAddExpr, and then expanded? Thanks.

[efriedma]

I was thinking of something like that, yes. But thinking a bit more, maybe simpler to just use SCEVExpander on the base (e.g. "n") and emit the add directly in IR.

[xbolva00]

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

110

Can you add similar statistic for strlen?

[musman]

Change to use SCEV instead of pattern match for strlen recognition.

[efriedma]

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

1610

I don't understand this bit about the precondition block; why are we not inserting the strlen call in the preheader?

1628

You probably want to check that the load is in address-space zero.

1642

for (auto *I : *LoopBody) {

1649

This checks there's at least one PHI, but not exactly one PHI.

1658

I don't understand what you're doing here.

Probably the simplest thing to do here would be to directly expand LCSSAPHI based on the result of strlen, instead of trying to dig into the exit block and find a subtraction operation.

[musman]

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

1658

I was thinking the subtraction operation was needed in order to do replaceAllUsesWith on it, is there a way to do that without needing the subtract?
Also, to expand the LCSSAPHI would you take the incoming value and check if it's an AddRec, then expand that?
Thanks again.

[efriedma]

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

1658

The actual output of the loop reflects the true count; you shouldn't need to dig through instructions outside it. (We can simplify the subtraction, if there is one, as a separate transform. I assume instcombine already handles this.)

To expand the LCSSAPHI, yes, you'd check if it's an appropriate addrec, then replace it with the sum of the addrec's base and the result of the strlen call.

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

 #include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/PassManager.h"
+#include "llvm/IR/PatternMatch.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/User.h"
 #include "llvm/IR/Value.h"
 #include "llvm/IR/ValueHandle.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Scalar.h"
 #include "llvm/Transforms/Utils/BuildLibCalls.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Transforms/Utils/LoopUtils.h"
 #include "llvm/Transforms/Utils/ScalarEvolutionExpander.h"
 #include <algorithm>
 #include <cassert>
 #include <cstdint>
 #include <utility>
 #include <vector>
 
 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");
 
 static cl::opt<bool> UseLIRCodeSizeHeurs(
     "use-lir-code-size-heurs",
     cl::desc("Use loop idiom recognition code size heuristics when compiling"
              "with -Os/-Oz"),
     cl::init(true), cl::Hidden);
 
+static cl::opt<bool> DisableStrlenIdiom(
+    "disable-strlen-idiom",
+    cl::desc("Disable strlen loop idiom recognize and transform"),
+    cl::init(false), cl::Hidden);
+
 namespace {
 
 class LoopIdiomRecognize {
   Loop *CurLoop = nullptr;
   AliasAnalysis *AA;
   DominatorTree *DT;
   LoopInfo *LI;
   ScalarEvolution *SE;
private:
   void transformLoopToPopcount(BasicBlock *PreCondBB, Instruction *CntInst,
                                PHINode *CntPhi, Value *Var);
   bool recognizeAndInsertFFS();  /// Find First Set: ctlz or cttz
   void transformLoopToCountable(Intrinsic::ID IntrinID, BasicBlock *PreCondBB,
                                 Instruction *CntInst, PHINode *CntPhi,
                                 Value *Var, Instruction *DefX,
                                 const DebugLoc &DL, bool ZeroCheck,
                                 bool IsCntPhiUsedOutsideLoop);
+  bool recognizeAndInsertStrlen();
 
   /// @}
 };
 
 class LoopIdiomRecognizeLegacyPass : public LoopPass {
 public:
   static char ID;
 

 }
 
 bool LoopIdiomRecognize::runOnNoncountableLoop() {
   LLVM_DEBUG(dbgs() << DEBUG_TYPE " Scanning: F["
                     << CurLoop->getHeader()->getParent()->getName()
                     << "] Noncountable Loop %"
                     << CurLoop->getHeader()->getName() << "\n");
 
-  return recognizePopcount() || recognizeAndInsertFFS();
+  return recognizePopcount() || recognizeAndInsertFFS() ||
+         recognizeAndInsertStrlen();
 }
 
 /// Check if the given conditional branch is based on the comparison between
 /// a variable and zero, and if the variable is non-zero or zero (JmpOnZero is
 /// true), the control yields to the loop entry. If the branch matches the
 /// behavior, the variable involved in the comparison is returned. This function
 /// will be called to see if the precondition and postcondition of the loop are
 /// in desirable form.
if (HeaderSize != IdiomCanonicalSize &&
     return false;
 
   transformLoopToCountable(IntrinID, PH, CntInst, CntPhi, InitX, DefX,
                            DefX->getDebugLoc(), ZeroCheck,
                            IsCntPhiUsedOutsideLoop);
   return true;
 }
 
+/// Recognizes a strlen idiom by checking for loops that either keep
+/// a variable counter or increment a char pointer and then subtract
+/// with the base pointer.
+///
+/// If detected, transforms the relevant code to a strlen function
+/// call, and returns true; otherwise, returns false.
+bool LoopIdiomRecognize::recognizeAndInsertStrlen() {
+  if (DisableStrlenIdiom)
+    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());
+  // The strlen loop idiom can only have a maximum of 7 instructions:
+  // this arises when there is an increment variable in the loop.
+  const unsigned MaxLoopSize = 7;
+  if (LoopBody->size() > MaxLoopSize)
+    return false;
+
+  // It should have a preheader containing nothing but an unconditional branch.
+  BasicBlock *PH = CurLoop->getLoopPreheader();
+  if (!PH || &PH->front() != PH->getTerminator())
+    return false;
+  auto *EntryBI = dyn_cast<BranchInst>(PH->getTerminator());
+  if (!EntryBI || EntryBI->isConditional())
+    return false;
+
+  // It should have a precondition block where the generated strlen intrinsic
+  // function can be inserted.
+  auto *PreCondBB = PH->getSinglePredecessor();
+  if (!PreCondBB)
+    return false;
+  auto *PreCondBI = dyn_cast<BranchInst>(PreCondBB->getTerminator());
+  if (!PreCondBI || PreCondBI->isUnconditional())
+    return false;
+
+  // The loop exit must be conditioned on an icmp with 0.
+  // The icmp operand has to be a load on some SSA reg that increments
+  // by 1 in the loop.
+  auto *LoopTerm = dyn_cast<BranchInst>(LoopBody->getTerminator());
+  auto *LoopCond = matchCondition(LoopTerm, LoopBody);
+  if (!LoopCond || !LoopCond->hasOneUse())
+    return false;
+  auto *LoopLoad = dyn_cast<LoadInst>(LoopCond);
+  if (!LoopLoad)
+    return false;
+
+  // Check that the vregs (phi and gep) are of valid form.
+  auto *IncPtr = LoopLoad->getPointerOperand();
+  PHINode *LoopPHI = nullptr;
+  GetElementPtrInst *LoopGEP = nullptr;
+  if ((LoopGEP = dyn_cast<GetElementPtrInst>(IncPtr))) {
+    LoopPHI = dyn_cast<PHINode>(LoopGEP->getPointerOperand());
+    if (!LoopPHI || LoopPHI->getParent() != LoopBody)
+      return false;
+    // The phi node has to have incoming values from loop body and preheader
+    if (LoopPHI->getBasicBlockIndex(LoopBody) == -1 ||
+        LoopPHI->getBasicBlockIndex(PH) == -1)
+      return false;
+    // The IncPtr could be a phi node if the first char was already null checked
+  } else if ((LoopPHI = dyn_cast<PHINode>(IncPtr))) {
+    if (LoopPHI->getBasicBlockIndex(LoopBody) == -1 ||
+        LoopPHI->getBasicBlockIndex(PH) == -1)
+      return false;
+    LoopGEP = dyn_cast<GetElementPtrInst>(
+        LoopPHI->getIncomingValueForBlock(LoopBody));
+  } else
+    return false;
+  if (!LoopPHI || !LoopGEP || LoopGEP->getParent() != LoopBody)
+    return false;
+  if (LoopPHI->getNumIncomingValues() != 2 ||
+      LoopPHI->getIncomingValueForBlock(LoopBody) != LoopGEP)
+    return false;
+  
+  // The increment can only accumulate constant offset of 1
+  const auto &DL = CurLoop->getHeader()->getModule()->getDataLayout();
+  APInt Offset(DL.getPointerTypeSizeInBits(LoopGEP->getType()), 0);
+  if (!LoopGEP->accumulateConstantOffset(DL, Offset))
+    return false;
+
+  // The string base (source) must be of pointer type.
+  Value *StrSrc = LoopPHI->getIncomingValueForBlock(PH);
+  PointerType *PT = dyn_cast<PointerType>(StrSrc->getType());
+  if (!PT)
+    return false;
+  // Make sure the char width is 8
+  if (PT->getPointerElementType()->getTypeID() != Type::IntegerTyID)
+    return false;
+  unsigned PtBW = PT->getPointerElementType()->getIntegerBitWidth();
+  if (PtBW != 8)
+    return false;
+  if (Offset != 1)
+    return false;
+
+  Value *ResInstr = nullptr;
+  if (IncPtr->hasNUses(3)) {
+    // Check if the increment vreg computes a pointer difference
+    // with the string base. If so, the loop should do nothing but increment it.
+    PtrToIntInst *SubLHS = nullptr;
+    Value *SubRHS = nullptr;
+    for (auto *U : IncPtr->users()) {
+      if (U == LoopLoad || U == LoopPHI || U == LoopGEP)
+        continue;
+      auto *LHSPhi = dyn_cast<PHINode>(U);
+      if (!LHSPhi || LHSPhi->getNumIncomingValues() != 1 ||
+          !LHSPhi->hasOneUse())
+        return false;
+      SubLHS = dyn_cast<PtrToIntInst>(LHSPhi->user_back());
+      break;
+    }
+    // Patternmatching IR that subtracts pointers of this form:
+    // %inc.lcssa = phi i8* [ %inc, %loop ]
+    // %sub.ptr.lhs = ptrtoint i8* %inc.lcssa to i64
+    // %sub.ptr.rhs = ptrtoint i8* %str.base to i64
+    // %sub.ptr.sub = sub i64 %sub.ptr.lhs, %sub.ptr.rhs
+    if (!SubLHS || !SubLHS->hasOneUse())
+      return false;
+    ResInstr = SubLHS->user_back();
+    if (!match(ResInstr, m_Sub(m_Specific(SubLHS), m_Value(SubRHS))) ||
+        !isa<PtrToIntInst>(SubRHS) || !SubRHS->hasOneUse())
+      return false;
+    // Check that the RHS is the string base
+    if (cast<PtrToIntInst>(SubRHS)->getPointerOperand() != StrSrc)
+      return false;
+
+    // loop can only have phi, gep, load, icmp, and br
+    if (LoopBody->size() != 5)
+      return false;
+
+  } else if (IncPtr->hasNUses(2)) {
+    // If the increment vreg does not compute a pointer difference,
+    // check if there is an add instruction that is used as the strlen result.
+    // If so, the loop should do nothing else.
+    // Seeking form:
+    // %add.phi = phi i64 [ %inc, %loop ], [ 0, %loop.preheader ]
+    // %inc = add i64 %add.phi, 1
+    for (BasicBlock::iterator I = LoopBody->begin(), E = LoopBody->end();
+         I != E;) {
+      Instruction *Inst = &*I++;
+      BinaryOperator *AddOneOp;
+      if ((AddOneOp = dyn_cast<BinaryOperator>(Inst))) {
+        if (AddOneOp->getOpcode() != Instruction::Add)
+          return false;
+        ConstantInt *AddInc = dyn_cast<ConstantInt>(AddOneOp->getOperand(1));
+        if (!AddInc || !AddInc->isOne())
+          return false;
+        auto *AddPhi = dyn_cast<PHINode>(AddOneOp->getOperand(0));
+        if (!AddPhi || !AddPhi->hasOneUse() ||
+            AddPhi->getNumIncomingValues() != 2 ||
+            AddPhi->getBasicBlockIndex(LoopBody) == -1 ||
+            AddPhi->getBasicBlockIndex(PH) == -1 ||
+            AddPhi->getIncomingValueForBlock(LoopBody) != AddOneOp ||
+            AddPhi->getParent() != LoopBody)
+          return false;
+        ConstantInt *StartVal =
+            dyn_cast<ConstantInt>(AddPhi->getIncomingValueForBlock(PH));
+        if (!StartVal || !StartVal->isZero())
+          return false;
+        ResInstr = AddOneOp;
+        break;
+      }
+    }
+    // loop can only have ptr phi, ptr gep, inc phi, inc add,
+    // load, icmp, and br.
+    if (LoopBody->size() != 7)
+      return false;
+  } else
+    return false;
+
+  // Scan every instruction in the loop to ensure it is doing nothing else.
+  for (BasicBlock::iterator I = LoopBody->begin(), E = LoopBody->end();
+       I != E;) {
+    Instruction *Inst = &*I++;
+    if (Inst == LoopLoad || Inst == LoopCond)
+      if (!Inst->hasOneUse())
+        return false;
+    if (Inst->mayHaveSideEffects())
+      return false;
+    if (Inst != LoopTerm && !Inst->getType()->isIntegerTy() &&
+        !Inst->getType()->isPointerTy())
+      return false;
+  }
+
+  // Insert strlen and RAUW
+  IRBuilder<> Builder(PreCondBI);
+  Value *FuncCall = emitStrLen(StrSrc, Builder, DL, TLI);
+  Value *Strlen = Builder.CreateZExtOrTrunc(FuncCall, ResInstr->getType());
+  ResInstr->replaceAllUsesWith(Strlen);
+  RecursivelyDeleteTriviallyDeadInstructions(ResInstr, TLI);
+
+  // Remove the loop-exit branch
+  ConstantInt *NewLoopCond = LoopTerm->getSuccessor(0) == LoopBody
+                                 ? Builder.getFalse()
+                                 : Builder.getTrue();
+  LoopTerm->setCondition(NewLoopCond);
+
+  // Delete the dead loop instructions
+  deleteDeadInstruction(cast<Instruction>(LoopCond));
+  deleteDeadInstruction(cast<Instruction>(IncPtr));
+  SE->forgetLoop(CurLoop);
+  return true;
+}
+
 /// Recognizes a population count idiom in a non-countable loop.
 ///
 /// If detected, transforms the relevant code to issue the popcount intrinsic
 /// function call, and returns true; otherwise, returns false.
 bool LoopIdiomRecognize::recognizePopcount() {
   if (TTI->getPopcntSupport(32) != TargetTransformInfo::PSK_FastHardware)
     return false;
 

llvm/test/Transforms/LoopIdiom/recognize-strlen.ll

+; RUN: opt -loop-idiom < %s -S | FileCheck %s
+target datalayout = "e-m:e-i64:64-n32:64"
+target triple = "powerpc64le-unknown-linux-gnu"
+
+define i64 @valid_strlen1(i8* %Str) {
+entry:
+  %tobool = icmp eq i8* %Str, null
+  br i1 %tobool, label %cleanup, label %lor.lhs.false
+
+lor.lhs.false:                                    ; preds = %entry
+  %0 = load i8, i8* %Str, align 1
+  %cmp = icmp eq i8 %0, 0
+  br i1 %cmp, label %cleanup, label %for.inc
+
+; CHECK:  call i64 @strlen(i8* %Str)
+; CHECK:  br i1 true
+for.inc:                                          ; preds = %lor.lhs.false, %for.inc
+  %Src.09 = phi i8* [ %incdec.ptr, %for.inc ], [ %Str, %lor.lhs.false ]
+  %incdec.ptr = getelementptr inbounds i8, i8* %Src.09, i64 1
+  %.pr = load i8, i8* %incdec.ptr, align 1
+  %tobool2 = icmp eq i8 %.pr, 0
+  br i1 %tobool2, label %for.end, label %for.inc
+
+; CHECK-NOT:    sub
+for.end:                                          ; preds = %for.inc
+  %sub.ptr.lhs.cast = ptrtoint i8* %incdec.ptr to i64
+  %sub.ptr.rhs.cast = ptrtoint i8* %Str to i64
+  %sub.ptr.sub = sub i64 %sub.ptr.lhs.cast, %sub.ptr.rhs.cast
+  br label %cleanup
+
+cleanup:                                          ; preds = %lor.lhs.false, %entry, %for.end
+  %retval.0 = phi i64 [ %sub.ptr.sub, %for.end ], [ 0, %entry ], [ 0, %lor.lhs.false ]
+  ret i64 %retval.0
+}
+
+
+define i64 @valid_strlen2(i8* %Str) {
+entry:
+  %tobool = icmp eq i8* %Str, null
+  br i1 %tobool, label %cleanup, label %for.cond
+
+; CHECK:  call i64 @strlen(i8* %Str)
+; CHECK:  br i1 true
+for.cond:                                         ; preds = %entry, %for.cond
+  %Src.0 = phi i8* [ %incdec.ptr, %for.cond ], [ %Str, %entry ]
+  %0 = load i8, i8* %Src.0, align 1
+  %tobool1 = icmp eq i8 %0, 0
+  %incdec.ptr = getelementptr inbounds i8, i8* %Src.0, i64 1
+  br i1 %tobool1, label %for.end, label %for.cond
+
+; CHECK-NOT:    sub
+for.end:                                          ; preds = %for.cond
+  %sub.ptr.lhs.cast = ptrtoint i8* %Src.0 to i64
+  %sub.ptr.rhs.cast = ptrtoint i8* %Str to i64
+  %sub.ptr.sub = sub i64 %sub.ptr.lhs.cast, %sub.ptr.rhs.cast
+  br label %cleanup
+
+ cleanup:                                          ; preds = %entry, %for.end
+  %retval.0 = phi i64 [ %sub.ptr.sub, %for.end ], [ 0, %entry ]
+  ret i64 %retval.0
+}
+
+
+define i64 @valid_strlen3(i8* readonly %Str) {
+entry:
+  %tobool = icmp eq i8* %Str, null
+  br i1 %tobool, label %cleanup, label %while.cond.preheader
+
+while.cond.preheader:                             ; preds = %entry
+  %0 = load i8, i8* %Str, align 1
+  %tobool16 = icmp eq i8 %0, 0
+  br i1 %tobool16, label %cleanup, label %while.body
+
+; CHECK:  call i64 @strlen(i8* %Str)
+; CHECK:  br i1 true
+while.body:                                       ; preds = %while.cond.preheader, %while.body
+  %Len.08 = phi i64 [ %inc, %while.body ], [ 0, %while.cond.preheader ]
+  %Str.addr.07 = phi i8* [ %incdec.ptr, %while.body ], [ %Str, %while.cond.preheader ]
+  %incdec.ptr = getelementptr inbounds i8, i8* %Str.addr.07, i64 1
+  %inc = add i64 %Len.08, 1
+  %1 = load i8, i8* %incdec.ptr, align 1
+  %tobool1 = icmp eq i8 %1, 0
+  br i1 %tobool1, label %cleanup, label %while.body
+
+cleanup:                                          ; preds = %while.body, %while.cond.preheader, %entry
+  %retval.0 = phi i64 [ 0, %entry ], [ 0, %while.cond.preheader ], [ %inc, %while.body ]
+  ret i64 %retval.0
+}
+
+
+define zeroext i32 @valid_strlen4(i8* %Str) {
+entry:
+  %tobool = icmp eq i8* %Str, null
+  br i1 %tobool, label %cleanup, label %lor.lhs.false
+
+lor.lhs.false:                                    ; preds = %entry
+  %0 = load i8, i8* %Str, align 1
+  %cmp = icmp eq i8 %0, 0
+  br i1 %cmp, label %cleanup, label %for.inc
+
+; CHECK:  call i64 @strlen(i8* %Str)
+; CHECK:  br i1 true
+for.inc:                                          ; preds = %lor.lhs.false, %for.inc
+  %Src.010 = phi i8* [ %incdec.ptr, %for.inc ], [ %Str, %lor.lhs.false ]
+  %incdec.ptr = getelementptr inbounds i8, i8* %Src.010, i64 1
+  %.pr = load i8, i8* %incdec.ptr, align 1
+  %tobool2 = icmp eq i8 %.pr, 0
+  br i1 %tobool2, label %for.end, label %for.inc
+
+; CHECK:      trunc i64 %strlen to i32
+; CHECK-NOT:  sub
+for.end:                                          ; preds = %for.inc
+  %sub.ptr.lhs.cast = ptrtoint i8* %incdec.ptr to i64
+  %sub.ptr.rhs.cast = ptrtoint i8* %Str to i64
+  %sub.ptr.sub = sub i64 %sub.ptr.lhs.cast, %sub.ptr.rhs.cast
+  %conv3 = trunc i64 %sub.ptr.sub to i32
+  br label %cleanup
+
+cleanup:                                          ; preds = %lor.lhs.false, %entry, %for.end
+  %retval.0 = phi i32 [ %conv3, %for.end ], [ 0, %entry ], [ 0, %lor.lhs.false ]
+  ret i32 %retval.0
+}