[llvm] [AggressiveInstCombine] Expand strchr/memchr with small constant strings (PR #98501)

via llvm-commits llvm-commits at lists.llvm.org
Thu Jul 11 08:58:19 PDT 2024


llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Yingwei Zheng (dtcxzyw)

<details>
<summary>Changes</summary>

This patch converts strchr/memchr with a small constant string into a switch. It will reduce overhead of libcall and enable more folds (e.g., comparing the result with null).

References:
https://en.cppreference.com/w/c/string/byte/memchr
https://en.cppreference.com/w/c/string/byte/strchr


---
Full diff: https://github.com/llvm/llvm-project/pull/98501.diff


2 Files Affected:

- (modified) llvm/lib/Transforms/AggressiveInstCombine/AggressiveInstCombine.cpp (+87) 
- (added) llvm/test/Transforms/AggressiveInstCombine/strchr.ll (+207) 


``````````diff
diff --git a/llvm/lib/Transforms/AggressiveInstCombine/AggressiveInstCombine.cpp b/llvm/lib/Transforms/AggressiveInstCombine/AggressiveInstCombine.cpp
index 1e0b8d448b9d1..c907fef1379e5 100644
--- a/llvm/lib/Transforms/AggressiveInstCombine/AggressiveInstCombine.cpp
+++ b/llvm/lib/Transforms/AggressiveInstCombine/AggressiveInstCombine.cpp
@@ -54,6 +54,11 @@ static cl::opt<unsigned> StrNCmpInlineThreshold(
     cl::desc("The maximum length of a constant string for a builtin string cmp "
              "call eligible for inlining. The default value is 3."));
 
+static cl::opt<unsigned>
+    StrChrInlineThreshold("strchr-inline-threshold", cl::init(3), cl::Hidden,
+                          cl::desc("The maximum length of a constant string to "
+                                   "inline a memchr/strchr call."));
+
 /// Match a pattern for a bitwise funnel/rotate operation that partially guards
 /// against undefined behavior by branching around the funnel-shift/rotation
 /// when the shift amount is 0.
@@ -1103,6 +1108,81 @@ void StrNCmpInliner::inlineCompare(Value *LHS, StringRef RHS, uint64_t N,
   }
 }
 
+/// Convert strchr/memchr with a small constant string into a switch
+static bool foldStrChr(CallInst *Call, LibFunc Func, DomTreeUpdater *DTU,
+                       const DataLayout &DL) {
+  assert((Func == LibFunc_strchr || Func == LibFunc_memchr) &&
+         "Unexpected LibFunc");
+  if (isa<Constant>(Call->getArgOperand(1)))
+    return false;
+
+  StringRef Str;
+  Value *Base = Call->getArgOperand(0);
+  if (!getConstantStringInfo(Base, Str, /*TrimAtNul=*/Func == LibFunc_strchr))
+    return false;
+
+  uint64_t N = Str.size();
+  if (Func == LibFunc_memchr) {
+    if (auto *ConstInt = dyn_cast<ConstantInt>(Call->getArgOperand(2)))
+      N = std::min(N, ConstInt->getZExtValue());
+    else
+      return false;
+  }
+
+  if (N > StrChrInlineThreshold)
+    return false;
+
+  BasicBlock *BB = Call->getParent();
+  BasicBlock *BBNext = SplitBlock(BB, Call, DTU);
+  IRBuilder<> IRB(BB);
+  IntegerType *ByteTy = IRB.getInt8Ty();
+  BB->getTerminator()->eraseFromParent();
+  SwitchInst *SI = IRB.CreateSwitch(
+      IRB.CreateTrunc(Call->getArgOperand(1), ByteTy), BBNext, N);
+  Type *IndexTy = DL.getIndexType(Call->getType());
+
+  PHINode *PHI = PHINode::Create(Call->getType(), 2, "", BBNext->begin());
+  PHI->addIncoming(Constant::getNullValue(Call->getType()), BB);
+
+  SmallVector<DominatorTree::UpdateType, 8> Updates;
+
+  BasicBlock *BBSuccess =
+      BasicBlock::Create(Call->getContext(), "", BB->getParent(), BBSuccess);
+  IRB.SetInsertPoint(BBSuccess);
+  PHINode *IndexPHI = IRB.CreatePHI(IndexTy, N);
+  Value *FirstOccursLocation = IRB.CreateInBoundsPtrAdd(Base, IndexPHI);
+  PHI->addIncoming(FirstOccursLocation, BBSuccess);
+  IRB.CreateBr(BBNext);
+  if (DTU)
+    Updates.push_back({DominatorTree::Insert, BBSuccess, BBNext});
+
+  SmallPtrSet<ConstantInt *, 4> Cases;
+  for (uint64_t I = 0; I < N; ++I) {
+    ConstantInt *CaseVal = ConstantInt::get(ByteTy, Str[I]);
+    if (!Cases.insert(CaseVal).second)
+      continue;
+
+    BasicBlock *BBCase =
+        BasicBlock::Create(Call->getContext(), "", BB->getParent(), BBNext);
+    SI->addCase(CaseVal, BBCase);
+    IRB.SetInsertPoint(BBCase);
+    IndexPHI->addIncoming(ConstantInt::get(IndexTy, I), BBCase);
+    IRB.CreateBr(BBSuccess);
+    if (DTU) {
+      Updates.push_back({DominatorTree::Insert, BB, BBCase});
+      Updates.push_back({DominatorTree::Insert, BBCase, BBSuccess});
+    }
+  }
+
+  Call->replaceAllUsesWith(PHI);
+  Call->eraseFromParent();
+
+  if (DTU)
+    DTU->applyUpdates(Updates);
+
+  return true;
+}
+
 static bool foldLibCalls(Instruction &I, TargetTransformInfo &TTI,
                          TargetLibraryInfo &TLI, AssumptionCache &AC,
                          DominatorTree &DT, const DataLayout &DL,
@@ -1135,6 +1215,13 @@ static bool foldLibCalls(Instruction &I, TargetTransformInfo &TTI,
       return true;
     }
     break;
+  case LibFunc_strchr:
+  case LibFunc_memchr:
+    if (foldStrChr(CI, LF, &DTU, DL)) {
+      MadeCFGChange = true;
+      return true;
+    }
+    break;
   default:;
   }
   return false;
diff --git a/llvm/test/Transforms/AggressiveInstCombine/strchr.ll b/llvm/test/Transforms/AggressiveInstCombine/strchr.ll
new file mode 100644
index 0000000000000..1692e5892c8ed
--- /dev/null
+++ b/llvm/test/Transforms/AggressiveInstCombine/strchr.ll
@@ -0,0 +1,207 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -passes=aggressive-instcombine --strchr-inline-threshold=5 < %s | FileCheck %s
+
+ at str = constant [5 x i8] c"01\002\00", align 1
+ at str_long = constant [8 x i8] c"0123456\00", align 1
+
+declare ptr @memchr(ptr, i32, i64)
+declare ptr @strchr(ptr, i32)
+
+define ptr @test_strchr(i32 %x) {
+; CHECK-LABEL: define ptr @test_strchr(
+; CHECK-SAME: i32 [[X:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*]]:
+; CHECK-NEXT:    [[TMP0:%.*]] = trunc i32 [[X]] to i8
+; CHECK-NEXT:    switch i8 [[TMP0]], label %[[ENTRY_SPLIT:.*]] [
+; CHECK-NEXT:      i8 48, label %[[BB1:.*]]
+; CHECK-NEXT:      i8 49, label %[[BB2:.*]]
+; CHECK-NEXT:    ]
+; CHECK:       [[BB1]]:
+; CHECK-NEXT:    br label %[[BB4:.*]]
+; CHECK:       [[BB2]]:
+; CHECK-NEXT:    br label %[[BB4]]
+; CHECK:       [[ENTRY_SPLIT]]:
+; CHECK-NEXT:    [[MEMCHR:%.*]] = phi ptr [ null, %[[ENTRY]] ], [ [[TMP6:%.*]], %[[BB4]] ]
+; CHECK-NEXT:    ret ptr [[MEMCHR]]
+; CHECK:       [[BB4]]:
+; CHECK-NEXT:    [[TMP5:%.*]] = phi i64 [ 0, %[[BB1]] ], [ 1, %[[BB2]] ]
+; CHECK-NEXT:    [[TMP6]] = getelementptr inbounds i8, ptr @str, i64 [[TMP5]]
+; CHECK-NEXT:    br label %[[ENTRY_SPLIT]]
+;
+entry:
+  %memchr = call ptr @strchr(ptr @str, i32 %x)
+  ret ptr %memchr
+}
+
+define i1 @test_strchr_null(i32 %x) {
+; CHECK-LABEL: define i1 @test_strchr_null(
+; CHECK-SAME: i32 [[X:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*]]:
+; CHECK-NEXT:    [[TMP0:%.*]] = trunc i32 [[X]] to i8
+; CHECK-NEXT:    switch i8 [[TMP0]], label %[[ENTRY_SPLIT:.*]] [
+; CHECK-NEXT:      i8 48, label %[[BB1:.*]]
+; CHECK-NEXT:      i8 49, label %[[BB2:.*]]
+; CHECK-NEXT:    ]
+; CHECK:       [[BB1]]:
+; CHECK-NEXT:    br label %[[BB4:.*]]
+; CHECK:       [[BB2]]:
+; CHECK-NEXT:    br label %[[BB4]]
+; CHECK:       [[ENTRY_SPLIT]]:
+; CHECK-NEXT:    [[MEMCHR:%.*]] = phi ptr [ null, %[[ENTRY]] ], [ [[TMP6:%.*]], %[[BB4]] ]
+; CHECK-NEXT:    [[ISNULL:%.*]] = icmp eq ptr [[MEMCHR]], null
+; CHECK-NEXT:    ret i1 [[ISNULL]]
+; CHECK:       [[BB4]]:
+; CHECK-NEXT:    [[TMP5:%.*]] = phi i64 [ 0, %[[BB1]] ], [ 1, %[[BB2]] ]
+; CHECK-NEXT:    [[TMP6]] = getelementptr inbounds i8, ptr @str, i64 [[TMP5]]
+; CHECK-NEXT:    br label %[[ENTRY_SPLIT]]
+;
+entry:
+  %memchr = call ptr @strchr(ptr @str, i32 %x)
+  %isnull = icmp eq ptr %memchr, null
+  ret i1 %isnull
+}
+
+define ptr @test_memchr(i32 %x) {
+; CHECK-LABEL: define ptr @test_memchr(
+; CHECK-SAME: i32 [[X:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*]]:
+; CHECK-NEXT:    [[TMP0:%.*]] = trunc i32 [[X]] to i8
+; CHECK-NEXT:    switch i8 [[TMP0]], label %[[ENTRY_SPLIT:.*]] [
+; CHECK-NEXT:      i8 48, label %[[BB1:.*]]
+; CHECK-NEXT:      i8 49, label %[[BB2:.*]]
+; CHECK-NEXT:      i8 0, label %[[BB3:.*]]
+; CHECK-NEXT:      i8 50, label %[[BB4:.*]]
+; CHECK-NEXT:    ]
+; CHECK:       [[BB1]]:
+; CHECK-NEXT:    br label %[[BB6:.*]]
+; CHECK:       [[BB2]]:
+; CHECK-NEXT:    br label %[[BB6]]
+; CHECK:       [[BB3]]:
+; CHECK-NEXT:    br label %[[BB6]]
+; CHECK:       [[BB4]]:
+; CHECK-NEXT:    br label %[[BB6]]
+; CHECK:       [[ENTRY_SPLIT]]:
+; CHECK-NEXT:    [[MEMCHR:%.*]] = phi ptr [ null, %[[ENTRY]] ], [ [[TMP8:%.*]], %[[BB6]] ]
+; CHECK-NEXT:    ret ptr [[MEMCHR]]
+; CHECK:       [[BB6]]:
+; CHECK-NEXT:    [[TMP7:%.*]] = phi i64 [ 0, %[[BB1]] ], [ 1, %[[BB2]] ], [ 2, %[[BB3]] ], [ 3, %[[BB4]] ]
+; CHECK-NEXT:    [[TMP8]] = getelementptr inbounds i8, ptr @str, i64 [[TMP7]]
+; CHECK-NEXT:    br label %[[ENTRY_SPLIT]]
+;
+entry:
+  %memchr = call ptr @memchr(ptr @str, i32 %x, i64 5)
+  ret ptr %memchr
+}
+
+define ptr @test_memchr_smaller_n(i32 %x) {
+; CHECK-LABEL: define ptr @test_memchr_smaller_n(
+; CHECK-SAME: i32 [[X:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*]]:
+; CHECK-NEXT:    [[TMP0:%.*]] = trunc i32 [[X]] to i8
+; CHECK-NEXT:    switch i8 [[TMP0]], label %[[ENTRY_SPLIT:.*]] [
+; CHECK-NEXT:      i8 48, label %[[BB1:.*]]
+; CHECK-NEXT:      i8 49, label %[[BB2:.*]]
+; CHECK-NEXT:      i8 0, label %[[BB3:.*]]
+; CHECK-NEXT:      i8 50, label %[[BB4:.*]]
+; CHECK-NEXT:    ]
+; CHECK:       [[BB1]]:
+; CHECK-NEXT:    br label %[[BB6:.*]]
+; CHECK:       [[BB2]]:
+; CHECK-NEXT:    br label %[[BB6]]
+; CHECK:       [[BB3]]:
+; CHECK-NEXT:    br label %[[BB6]]
+; CHECK:       [[BB4]]:
+; CHECK-NEXT:    br label %[[BB6]]
+; CHECK:       [[ENTRY_SPLIT]]:
+; CHECK-NEXT:    [[MEMCHR:%.*]] = phi ptr [ null, %[[ENTRY]] ], [ [[TMP8:%.*]], %[[BB6]] ]
+; CHECK-NEXT:    ret ptr [[MEMCHR]]
+; CHECK:       [[BB6]]:
+; CHECK-NEXT:    [[TMP7:%.*]] = phi i64 [ 0, %[[BB1]] ], [ 1, %[[BB2]] ], [ 2, %[[BB3]] ], [ 3, %[[BB4]] ]
+; CHECK-NEXT:    [[TMP8]] = getelementptr inbounds i8, ptr @str, i64 [[TMP7]]
+; CHECK-NEXT:    br label %[[ENTRY_SPLIT]]
+;
+entry:
+  %memchr = call ptr @memchr(ptr @str, i32 %x, i64 4)
+  ret ptr %memchr
+}
+
+define ptr @test_memchr_larger_n(i32 %x) {
+; CHECK-LABEL: define ptr @test_memchr_larger_n(
+; CHECK-SAME: i32 [[X:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*]]:
+; CHECK-NEXT:    [[TMP0:%.*]] = trunc i32 [[X]] to i8
+; CHECK-NEXT:    switch i8 [[TMP0]], label %[[ENTRY_SPLIT:.*]] [
+; CHECK-NEXT:      i8 48, label %[[BB1:.*]]
+; CHECK-NEXT:      i8 49, label %[[BB2:.*]]
+; CHECK-NEXT:      i8 0, label %[[BB3:.*]]
+; CHECK-NEXT:      i8 50, label %[[BB4:.*]]
+; CHECK-NEXT:    ]
+; CHECK:       [[BB1]]:
+; CHECK-NEXT:    br label %[[BB6:.*]]
+; CHECK:       [[BB2]]:
+; CHECK-NEXT:    br label %[[BB6]]
+; CHECK:       [[BB3]]:
+; CHECK-NEXT:    br label %[[BB6]]
+; CHECK:       [[BB4]]:
+; CHECK-NEXT:    br label %[[BB6]]
+; CHECK:       [[ENTRY_SPLIT]]:
+; CHECK-NEXT:    [[MEMCHR:%.*]] = phi ptr [ null, %[[ENTRY]] ], [ [[TMP8:%.*]], %[[BB6]] ]
+; CHECK-NEXT:    ret ptr [[MEMCHR]]
+; CHECK:       [[BB6]]:
+; CHECK-NEXT:    [[TMP7:%.*]] = phi i64 [ 0, %[[BB1]] ], [ 1, %[[BB2]] ], [ 2, %[[BB3]] ], [ 3, %[[BB4]] ]
+; CHECK-NEXT:    [[TMP8]] = getelementptr inbounds i8, ptr @str, i64 [[TMP7]]
+; CHECK-NEXT:    br label %[[ENTRY_SPLIT]]
+;
+entry:
+  %memchr = call ptr @memchr(ptr @str, i32 %x, i64 6)
+  ret ptr %memchr
+}
+
+; negative tests
+
+define ptr @test_strchr_non_constant(i32 %x, ptr %str) {
+; CHECK-LABEL: define ptr @test_strchr_non_constant(
+; CHECK-SAME: i32 [[X:%.*]], ptr [[STR:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[MEMCHR:%.*]] = call ptr @strchr(ptr [[STR]], i32 [[X]])
+; CHECK-NEXT:    ret ptr [[MEMCHR]]
+;
+entry:
+  %memchr = call ptr @strchr(ptr %str, i32 %x)
+  ret ptr %memchr
+}
+
+define ptr @test_strchr_constant_ch() {
+; CHECK-LABEL: define ptr @test_strchr_constant_ch() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[MEMCHR:%.*]] = call ptr @strchr(ptr @str, i32 49)
+; CHECK-NEXT:    ret ptr [[MEMCHR]]
+;
+entry:
+  %memchr = call ptr @strchr(ptr @str, i32 49)
+  ret ptr %memchr
+}
+
+define ptr @test_memchr_dynamic_n(i32 %x, i32 %y) {
+; CHECK-LABEL: define ptr @test_memchr_dynamic_n(
+; CHECK-SAME: i32 [[X:%.*]], i32 [[Y:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[MEMCHR:%.*]] = call ptr @memchr(ptr @str, i32 [[X]], i32 [[Y]])
+; CHECK-NEXT:    ret ptr [[MEMCHR]]
+;
+entry:
+  %memchr = call ptr @memchr(ptr @str, i32 %x, i32 %y)
+  ret ptr %memchr
+}
+
+define ptr @test_strchr_long(i32 %x) {
+; CHECK-LABEL: define ptr @test_strchr_long(
+; CHECK-SAME: i32 [[X:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[MEMCHR:%.*]] = call ptr @strchr(ptr @str_long, i32 [[X]])
+; CHECK-NEXT:    ret ptr [[MEMCHR]]
+;
+entry:
+  %memchr = call ptr @strchr(ptr @str_long, i32 %x)
+  ret ptr %memchr
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/98501


More information about the llvm-commits mailing list