[SimplifyLibCalls] Fold memchr() with size 1

If the memchr() size is 1, then we can convert the call into a single-byte comparison. This works even if both the string and the character are unknown. Split off from https://reviews.llvm.org/D122836.
llvm · Apr 4, 2022 · d18991d · d18991d
1 parent 0f08875
commit d18991d
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 9 deletions.
diff --git a/llvm/lib/Transforms/Utils/SimplifyLibCalls.cpp b/llvm/lib/Transforms/Utils/SimplifyLibCalls.cpp
@@ -884,13 +884,25 @@ Value *LibCallSimplifier::optimizeMemChr(CallInst *CI, IRBuilderBase &B) {
   Value *SrcStr = CI->getArgOperand(0);
   Value *Size = CI->getArgOperand(2);
   annotateNonNullAndDereferenceable(CI, 0, Size, DL);
-  ConstantInt *CharC = dyn_cast<ConstantInt>(CI->getArgOperand(1));
+  Value *CharVal = CI->getArgOperand(1);
+  ConstantInt *CharC = dyn_cast<ConstantInt>(CharVal);
   ConstantInt *LenC = dyn_cast<ConstantInt>(Size);
 
   // memchr(x, y, 0) -> null
   if (LenC) {
     if (LenC->isZero())
       return Constant::getNullValue(CI->getType());
+
+    if (LenC->isOne()) {
+      // Fold memchr(x, y, 1) --> *x == y ? x : null for any x and y,
+      // constant or otherwise.
+      Value *Val = B.CreateLoad(B.getInt8Ty(), SrcStr, "memchr.char0");
+      // Slice off the character's high end bits.
+      CharVal = B.CreateTrunc(CharVal, B.getInt8Ty());
+      Value *Cmp = B.CreateICmpEQ(Val, CharVal, "memchr.char0cmp");
+      Value *NullPtr = Constant::getNullValue(CI->getType());
+      return B.CreateSelect(Cmp, SrcStr, NullPtr, "memchr.sel");
+    }
   } else {
     // From now on we need at least constant length and string.
     return nullptr;
@@ -939,7 +951,7 @@ Value *LibCallSimplifier::optimizeMemChr(CallInst *CI, IRBuilderBase &B) {
     Value *BitfieldC = B.getInt(Bitfield);
 
     // Adjust width of "C" to the bitfield width, then mask off the high bits.
-    Value *C = B.CreateZExtOrTrunc(CI->getArgOperand(1), BitfieldC->getType());
+    Value *C = B.CreateZExtOrTrunc(CharVal, BitfieldC->getType());
     C = B.CreateAnd(C, B.getIntN(Width, 0xFF));
 
     // First check that the bit field access is within bounds.

diff --git a/llvm/test/Transforms/InstCombine/memchr-3.ll b/llvm/test/Transforms/InstCombine/memchr-3.ll
@@ -41,8 +41,10 @@ define i8* @fold_memchr_a12345_2_1() {
 
 define i8* @fold_memchr_ax_257_1(i32 %chr, i64 %n) {
 ; CHECK-LABEL: @fold_memchr_ax_257_1(
-; CHECK-NEXT:    [[RES:%.*]] = call i8* @memchr(i8* noundef nonnull dereferenceable(1) getelementptr inbounds ([0 x i8], [0 x i8]* @ax, i64 0, i64 0), i32 257, i64 1)
-; CHECK-NEXT:    ret i8* [[RES]]
+; CHECK-NEXT:    [[MEMCHR_CHAR0:%.*]] = load i8, i8* getelementptr inbounds ([0 x i8], [0 x i8]* @ax, i64 0, i64 0), align 1
+; CHECK-NEXT:    [[MEMCHR_CHAR0CMP:%.*]] = icmp eq i8 [[MEMCHR_CHAR0]], 1
+; CHECK-NEXT:    [[MEMCHR_SEL:%.*]] = select i1 [[MEMCHR_CHAR0CMP]], i8* getelementptr inbounds ([0 x i8], [0 x i8]* @ax, i64 0, i64 0), i8* null
+; CHECK-NEXT:    ret i8* [[MEMCHR_SEL]]
 ;
 
   %ptr = getelementptr [0 x i8], [0 x i8]* @ax, i32 0, i32 0
@@ -55,8 +57,11 @@ define i8* @fold_memchr_ax_257_1(i32 %chr, i64 %n) {
 
 define i8* @fold_memchr_ax_c_1(i32 %chr, i64 %n) {
 ; CHECK-LABEL: @fold_memchr_ax_c_1(
-; CHECK-NEXT:    [[RES:%.*]] = call i8* @memchr(i8* noundef nonnull dereferenceable(1) getelementptr inbounds ([0 x i8], [0 x i8]* @ax, i64 0, i64 0), i32 [[CHR:%.*]], i64 1)
-; CHECK-NEXT:    ret i8* [[RES]]
+; CHECK-NEXT:    [[MEMCHR_CHAR0:%.*]] = load i8, i8* getelementptr inbounds ([0 x i8], [0 x i8]* @ax, i64 0, i64 0), align 1
+; CHECK-NEXT:    [[TMP1:%.*]] = trunc i32 [[CHR:%.*]] to i8
+; CHECK-NEXT:    [[MEMCHR_CHAR0CMP:%.*]] = icmp eq i8 [[MEMCHR_CHAR0]], [[TMP1]]
+; CHECK-NEXT:    [[MEMCHR_SEL:%.*]] = select i1 [[MEMCHR_CHAR0CMP]], i8* getelementptr inbounds ([0 x i8], [0 x i8]* @ax, i64 0, i64 0), i8* null
+; CHECK-NEXT:    ret i8* [[MEMCHR_SEL]]
 ;
 
   %ptr = getelementptr [0 x i8], [0 x i8]* @ax, i32 0, i32 0

diff --git a/llvm/test/Transforms/InstCombine/memchr.ll b/llvm/test/Transforms/InstCombine/memchr.ll
@@ -174,9 +174,9 @@ define i1 @test13(i32 %C) {
 
 define i1 @test14(i32 %C) {
 ; CHECK-LABEL: @test14(
-; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[C:%.*]], 255
-; CHECK-NEXT:    [[MEMCHR_BITS:%.*]] = icmp eq i32 [[TMP1]], 31
-; CHECK-NEXT:    ret i1 [[MEMCHR_BITS]]
+; CHECK-NEXT:    [[TMP1:%.*]] = trunc i32 [[C:%.*]] to i8
+; CHECK-NEXT:    [[MEMCHR_CHAR0CMP:%.*]] = icmp eq i8 [[TMP1]], 31
+; CHECK-NEXT:    ret i1 [[MEMCHR_CHAR0CMP]]
 ;
   %dst = call i8* @memchr(i8* getelementptr inbounds ([2 x i8], [2 x i8]* @single, i64 0, i64 0), i32 %C, i32 1)
   %cmp = icmp ne i8* %dst, null