[llvm] 5ccfd5f - [SimplifyLibCalls] Optimize memchr() with known char+str and unknown length

[Nikita Popov via llvm-commits]

Author: Martin Sebor
Date: 2022-04-04T11:01:33+02:00
New Revision: 5ccfd5f6d43069a5bc152b335bc704c6d2584ddc

URL: https://github.com/llvm/llvm-project/commit/5ccfd5f6d43069a5bc152b335bc704c6d2584ddc
DIFF: https://github.com/llvm/llvm-project/commit/5ccfd5f6d43069a5bc152b335bc704c6d2584ddc.diff

LOG: [SimplifyLibCalls] Optimize memchr() with known char+str and unknown length

If both the character and string are known, but the length
potentially isn't, we can optimize the memchr() call to a select
of either the known position of the character or null.

Split off from https://reviews.llvm.org/D122836.

Added: 
    

Modified: 
    llvm/lib/Transforms/Utils/SimplifyLibCalls.cpp
    llvm/test/Transforms/InstCombine/memchr-2.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Utils/SimplifyLibCalls.cpp b/llvm/lib/Transforms/Utils/SimplifyLibCalls.cpp
index 60c4d4a4ab933..329ffcdd48f2b 100644
--- a/llvm/lib/Transforms/Utils/SimplifyLibCalls.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyLibCalls.cpp
@@ -903,30 +903,37 @@ Value *LibCallSimplifier::optimizeMemChr(CallInst *CI, IRBuilderBase &B) {
       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;
   }
 
   StringRef Str;
   if (!getConstantStringInfo(SrcStr, Str, 0, /*TrimAtNul=*/false))
     return nullptr;
 
-  // Truncate the string to LenC. If Str is smaller than LenC we will still only
-  // scan the string, as reading past the end of it is undefined and we can just
-  // return null if we don't find the char.
-  Str = Str.substr(0, LenC->getZExtValue());
-
   if (CharC) {
-    // Compute the offset.
-    size_t I = Str.find(CharC->getSExtValue() & 0xFF);
-    if (I == StringRef::npos) // Didn't find the char.  memchr returns null.
+    size_t Pos = Str.find(CharC->getZExtValue());
+    if (Pos == StringRef::npos)
+      // When the character is not in the source array fold the result
+      // to null regardless of Size.
       return Constant::getNullValue(CI->getType());
 
-    // memchr(s+n,c,l) -> gep(s+n+i,c)
-    return B.CreateGEP(B.getInt8Ty(), SrcStr, B.getInt64(I), "memchr");
+    // Fold memchr(s, c, n) -> n <= Pos ? null : s + Pos
+    // When the constant Size is less than or equal to the character
+    // position also fold the result to null.
+    Value *Cmp = B.CreateICmpULE(Size, ConstantInt::get(Size->getType(), Pos),
+                                 "memchr.cmp");
+    Value *NullPtr = Constant::getNullValue(CI->getType());
+    Value *SrcPlus =
+        B.CreateGEP(B.getInt8Ty(), SrcStr, B.getInt64(Pos), "memchr.ptr");
+    return B.CreateSelect(Cmp, NullPtr, SrcPlus);
   }
 
+  if (!LenC)
+    // From now on we need a constant length and constant array.
+    return nullptr;
+
+  // Truncate the string to LenC.
+  Str = Str.substr(0, LenC->getZExtValue());
+
   // If the char is variable but the input str and length are not we can turn
   // this memchr call into a simple bit field test. Of course this only works
   // when the return value is only checked against null.

diff  --git a/llvm/test/Transforms/InstCombine/memchr-2.ll b/llvm/test/Transforms/InstCombine/memchr-2.ll
index d35ba52b276a1..c826753d7898c 100644
--- a/llvm/test/Transforms/InstCombine/memchr-2.ll
+++ b/llvm/test/Transforms/InstCombine/memchr-2.ll
@@ -15,8 +15,7 @@ declare i8* @memchr(i8*, i32, i64)
 
 define i8* @fold_memchr_a12345_6_n(i64 %n) {
 ; CHECK-LABEL: @fold_memchr_a12345_6_n(
-; CHECK-NEXT:    [[RES:%.*]] = call i8* @memchr(i8* getelementptr inbounds ([5 x i8], [5 x i8]* @a12345, i64 0, i64 0), i32 6, i64 [[N:%.*]])
-; CHECK-NEXT:    ret i8* [[RES]]
+; CHECK-NEXT:    ret i8* null
 ;
 
   %ptr = getelementptr [5 x i8], [5 x i8]* @a12345, i32 0, i32 0
@@ -95,8 +94,9 @@ define i8* @fold_memchr_a123f45_500_9() {
 
 define i8* @fold_a12345_3_n(i64 %n) {
 ; CHECK-LABEL: @fold_a12345_3_n(
-; CHECK-NEXT:    [[RES:%.*]] = call i8* @memchr(i8* getelementptr inbounds ([5 x i8], [5 x i8]* @a12345, i64 0, i64 0), i32 3, i64 [[N:%.*]])
-; CHECK-NEXT:    ret i8* [[RES]]
+; CHECK-NEXT:    [[MEMCHR_CMP:%.*]] = icmp ult i64 [[N:%.*]], 3
+; CHECK-NEXT:    [[TMP1:%.*]] = select i1 [[MEMCHR_CMP]], i8* null, i8* getelementptr inbounds ([5 x i8], [5 x i8]* @a12345, i64 0, i64 2)
+; CHECK-NEXT:    ret i8* [[TMP1]]
 ;
 
   %ptr = getelementptr [5 x i8], [5 x i8]* @a12345, i32 0, i32 0
@@ -110,8 +110,9 @@ define i8* @fold_a12345_3_n(i64 %n) {
 
 define i8* @fold_a12345_259_n(i64 %n) {
 ; CHECK-LABEL: @fold_a12345_259_n(
-; CHECK-NEXT:    [[RES:%.*]] = call i8* @memchr(i8* getelementptr inbounds ([5 x i8], [5 x i8]* @a12345, i64 0, i64 0), i32 259, i64 [[N:%.*]])
-; CHECK-NEXT:    ret i8* [[RES]]
+; CHECK-NEXT:    [[MEMCHR_CMP:%.*]] = icmp ult i64 [[N:%.*]], 3
+; CHECK-NEXT:    [[TMP1:%.*]] = select i1 [[MEMCHR_CMP]], i8* null, i8* getelementptr inbounds ([5 x i8], [5 x i8]* @a12345, i64 0, i64 2)
+; CHECK-NEXT:    ret i8* [[TMP1]]
 ;
 
   %ptr = getelementptr [5 x i8], [5 x i8]* @a12345, i32 0, i32 0