[SimplifyLibCalls] Fold more memchr calls

[msebor]

Summary

This patch extends the LibCallSimplifier::optimizeMemChr folder to perform two more simplifications:

1. folding the case where the size is equal to 1 to a byte comparison, analogously to other such cases already handled by other folders
2. simplifying to a size comparison the case where the sought character with constant value is known either to appear in the constant array or not to appear, regardless of the size argument

I have tested the patch by bootstrapping LLVM + Clang on x86_64-linux and running check-all.

[nikic]

I've committed your tests with baseline checks in https://github.com/llvm/llvm-project/commit/371d2ed3f3d7f67602ff9956510f3e2d9df3d5b0, can you please rebase on top of that?

llvm/lib/Transforms/Utils/SimplifyLibCalls.cpp

907

Why does this zext Val and CharVal? Isn't the trunc of CharVal to i8 and a comparison in i8 type sufficient?

925

Should be n <= Pos in the comment to match the code.

934

Should be MaxLen <= Pos?

941

Not sure why we need to fall through to the code below here -- or maybe more generally, can we just drop the MaxLen == UINT64_MAX condition above and always emit the icmp+select sequence? Will get folded for the case where the size is known.

1003

Is the code below here rendered redundant by your new code above?

[xbolva00]

llvm/lib/Transforms/Utils/SimplifyLibCalls.cpp

924

!LenC ?

942

If possible better to check for LenC than comparing with UINT64_MAX..

[msebor]

Thanks again for committing the baseline tests for me and (to both of you) for your comments. Let me upload a revised patch.

llvm/lib/Transforms/Utils/SimplifyLibCalls.cpp

907

I believe it is. It just needs an adjustment to the existing memchr.ll test.

924

Sure.

925

Done.

934

Either works thanks to what follows but sure.

941

Good idea! It simplifies the rest of the function as well.

942

Done.

1003

Yes, I have removed it.

[msebor]

Revised patch with changes suggested in review.

[xbolva00]

llvm/test/Transforms/InstCombine/memchr-2.ll

108

But IR says N < 3

And gep is 2, not 3.

Confusing.

[nikic]

LGTM

llvm/lib/Transforms/Utils/SimplifyLibCalls.cpp

900

I think MaxLen has lost most of it's purpose now, and you can just do a LenC->isOne() check in the preceding branch?

[msebor]

llvm/lib/Transforms/Utils/SimplifyLibCalls.cpp

900

There's another use in Str = Str.substr(0, MaxLen); on line 939.

The latest diff is harder to follow because of the new early return from the function. Sorry about that. I debated making that change and maybe I shouldn't have...

llvm/test/Transforms/InstCombine/memchr-2.ll

108

Sadly there are a few copy and paste typos in the comments and names of the functions (should be fold_xxx not call_xxx). Let me fix them.

[msebor]

Correct typos in test comments and test function names in memchr-2.ll

@nikic if the final version looks good can you please commit the final patch for me when you have a chance?

[xbolva00]

I think you could ask for commit access now if you plan to contribute more things than just this patch.

Then you can precommit tests with fixes for names and typos. After than, rebase, regenerate tests, and commit .

[xbolva00]

llvm/lib/Transforms/Utils/SimplifyLibCalls.cpp

900

LenC->isOne() seems like as a good suggestion.

Then you can use..

Str = Str.substr(0, LenC->getZExtValue());

One less variable to keep in mind when reading.

[efriedma]

llvm/lib/Transforms/Utils/SimplifyLibCalls.cpp

973

This implicitly truncates MaxLen to 32 bits on 32-bit targets, which is probably not what you want.

[msebor]

llvm/lib/Transforms/Utils/SimplifyLibCalls.cpp

973

You're right, silent truncation indeed isn't what I want here (or what I would have expected). Thanks for pointing out this gotcha!

It seems to be a pre-existing bug in the code made more apparent by the introduction of the MaxLen variable. Let me fix it and post another update.

[nikic]

I've taken the liberty of splitting this up into the following four commits while landing the change:
https://github.com/llvm/llvm-project/commit/0f088757449d63182fb424dce7856213739c563e
https://github.com/llvm/llvm-project/commit/d18991debfde6be335a14af3c7d2dcdecbcd758d
https://github.com/llvm/llvm-project/commit/5197d2791f908815134ad48d7b966de2d8c47eeb
https://github.com/llvm/llvm-project/commit/5ccfd5f6d43069a5bc152b335bc704c6d2584ddc

The length truncation can be addressed separately, as it's a pre-existing issue (and it looks like there are a few more substr calls in the same file that should be handled at the same time).

[xbolva00]

In D122836#3425739, @nikic wrote:

I've taken the liberty of splitting this up into the following four commits while landing the change:
https://github.com/llvm/llvm-project/commit/0f088757449d63182fb424dce7856213739c563e
https://github.com/llvm/llvm-project/commit/d18991debfde6be335a14af3c7d2dcdecbcd758d
https://github.com/llvm/llvm-project/commit/5197d2791f908815134ad48d7b966de2d8c47eeb
https://github.com/llvm/llvm-project/commit/5ccfd5f6d43069a5bc152b335bc704c6d2584ddc

The length truncation can be addressed separately, as it's a pre-existing issue (and it looks like there are a few more substr calls in the same file that should be handled at the same time).

Great, nice small patches!

[nickdesaulniers]

Thanks for the patches! :)

llvm/lib/Transforms/Utils/SimplifyLibCalls.cpp

if (isKnownNonZero(CI->getOperand(2), DL))
     annotateNonNullNoUndefBasedOnAccess(CI, 0);
   return nullptr;
 }
 
 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);
 
+  // Set to the Size argument value if it's constant or left at maximum.
+  uint64_t MaxLen = UINT64_MAX;
   // memchr(x, y, 0) -> null
   if (LenC) {
     if (LenC->isZero())
       return Constant::getNullValue(CI->getType());
-  } else {
-    // From now on we need at least constant length and string.
-    return nullptr;
+    MaxLen = LenC->getZExtValue();
+  }
+
+  if (MaxLen == 1) {
+    // 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");
+    Val = B.CreateZExt(Val, CharVal->getType());
+    // Slice off the character's high end bits.
+    CharVal = B.CreateTrunc(CharVal, B.getInt8Ty());
+    CharVal = B.CreateZExt(CharVal, Val->getType());
+    Value *Cmp = B.CreateICmpEQ(Val, CharVal, "memchr.char0cmp");
+    Value *NullPtr = Constant::getNullValue(CI->getType());
+    return B.CreateSelect(Cmp, SrcStr, NullPtr, "memchr.sel");
   }
 
   StringRef Str;
   if (!getConstantStringInfo(SrcStr, Str, 0, /*TrimAtNul=*/false))
     return nullptr;
 
+  if (CharC) {
+    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());
+
+    if (MaxLen == UINT64_MAX) {
+      // Fold memchr(s, c, n) -> n < Pos ? null : s + Pos
+      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 (MaxLen < Pos)
+      // When the constant Size is less than the character position
+      // fold the result to null.
+      return Constant::getNullValue(CI->getType());
+
+    // Set MaxLen so that the handling below folds the result either
+    // to a pointer to the array or to null.
+    MaxLen = Pos < MaxLen ? Pos + 1 : Pos;
+  } else if (MaxLen == UINT64_MAX)
+    // From now on we need at least constant length and constant array.
+    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());
+  Str = Str.substr(0, MaxLen);
 
   // 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.
   //
   // It would be really nice to reuse switch lowering here but we can't change
   // the CFG at this point.
   //
   // memchr("\r\n", C, 2) != nullptr -> (1 << C & ((1 << '\r') | (1 << '\n')))
   // != 0
   //   after bounds check.
   if (!CharC && !Str.empty() && isOnlyUsedInZeroEqualityComparison(CI)) {
     unsigned char Max =
         *std::max_element(reinterpret_cast<const unsigned char *>(Str.begin()),
                           reinterpret_cast<const unsigned char *>(Str.end()));
 
     // Make sure the bit field we're about to create fits in a register on the
     // target.
     // FIXME: On a 64 bit architecture this prevents us from using the
     // interesting range of alpha ascii chars. We could do better by emitting
     // two bitfields or shifting the range by 64 if no lower chars are used.
     if (!DL.fitsInLegalInteger(Max + 1))
       return nullptr;
 
     // For the bit field use a power-of-2 type with at least 8 bits to avoid
     // creating unnecessary illegal types.
     unsigned char Width = NextPowerOf2(std::max((unsigned char)7, Max));
 
     // Now build the bit field.
     APInt Bitfield(Width, 0);
     for (char C : Str)
       Bitfield.setBit((unsigned char)C);
     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.
     Value *Bounds = B.CreateICmp(ICmpInst::ICMP_ULT, C, B.getIntN(Width, Width),
                                  "memchr.bounds");
 
     // Create code that checks if the given bit is set in the field.
     Value *Shl = B.CreateShl(B.getIntN(Width, 1ULL), C);
     Value *Bits = B.CreateIsNotNull(B.CreateAnd(Shl, BitfieldC), "memchr.bits");
 
     // Finally merge both checks and cast to pointer type. The inttoptr
     // implicitly zexts the i1 to intptr type.
     return B.CreateIntToPtr(B.CreateLogicalAnd(Bounds, Bits, "memchr"),
                             CI->getType());
   }
 
   // Check if all arguments are constants.  If so, we can constant fold.
   if (!CharC)
     return nullptr;
 
   // Compute the offset.
   size_t I = Str.find(CharC->getSExtValue() & 0xFF);
   if (I == StringRef::npos) // Didn't find the char.  memchr returns null.
     return Constant::getNullValue(CI->getType());
 
   // memchr(s+n,c,l) -> gep(s+n+i,c)

llvm/test/Transforms/InstCombine/memchr-2.ll

+; Verify that memchr calls with constant arrays, or constant characters,
+; or constant bounds are folded (or not) as expected.
+;
+; RUN: opt < %s -passes=instcombine -S | FileCheck %s
+
+declare i8* @memchr(i8*, i32, i64)
+
+@ax = external global [0 x i8]
+@a12345 = constant [5 x i8] c"\01\02\03\04\05"
+
+
+; Fold memchr(a12345, '\06', n) to null.
+
+define i8* @fold_memchr_a12345_6_n(i64 %n) {
+; CHECK-LABEL: @fold_memchr_a12345_6_n(
+; CHECK-NEXT:    ret i8* null
+
+  %ptr = getelementptr [5 x i8], [5 x i8]* @a12345, i32 0, i32 0
+  %res = call i8* @memchr(i8* %ptr, i32 6, i64 %n)
+  ret i8* %res
+}
+
+
+; Fold memchr(a12345, '\04', 2) to null.
+
+define i8* @fold_memchr_a12345_4_2() {
+; CHECK-LABEL: @fold_memchr_a12345_4_2(
+; CHECK-NEXT:    ret i8* null
+
+  %ptr = getelementptr [5 x i8], [5 x i8]* @a12345, i32 0, i32 0
+  %res = call i8* @memchr(i8* %ptr, i32 4, i64 2)
+  ret i8* %res
+}
+
+
+; Fold memchr(a12345, '\04', 3) to null.
+
+define i8* @fold_memchr_a12345_4_3() {
+; CHECK-LABEL: @fold_memchr_a12345_4_3(
+; CHECK-NEXT:    ret i8* null
+
+  %ptr = getelementptr [5 x i8], [5 x i8]* @a12345, i32 0, i32 0
+  %res = call i8* @memchr(i8* %ptr, i32 4, i64 3)
+  ret i8* %res
+}
+
+
+; Fold memchr(a12345, '\03', 3) to a12345 + 2.
+
+define i8* @fold_memchr_a12345_3_3() {
+; CHECK-LABEL: @fold_memchr_a12345_3_3(
+; CHECK-NEXT:    ret i8* getelementptr inbounds ([5 x i8], [5 x i8]* @a12345, i64 0, i64 2)
+
+  %ptr = getelementptr [5 x i8], [5 x i8]* @a12345, i32 0, i32 0
+  %res = call i8* @memchr(i8* %ptr, i32 3, i64 3)
+  ret i8* %res
+}
+
+
+; Fold memchr(a12345, '\03', 9) to a12345 + 2.
+
+define i8* @fold_memchr_a12345_3_9() {
+; CHECK-LABEL: @fold_memchr_a12345_3_9(
+; CHECK-NEXT:    ret i8* getelementptr inbounds ([5 x i8], [5 x i8]* @a12345, i64 0, i64 2)
+
+  %ptr = getelementptr [5 x i8], [5 x i8]* @a12345, i32 0, i32 0
+  %res = call i8* @memchr(i8* %ptr, i32 3, i64 9)
+  ret i8* %res
+}
+
+
+; Fold memchr(a12345, '\03', n) to n < 3 ? null : a12345 + 3.
+
+define i8* @call_a12345_3_n(i64 %n) {
+; CHECK-LABEL: @call_a12345_3_n(
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i64 %n, 3
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8* null, i8* getelementptr inbounds ([5 x i8], [5 x i8]* @a12345, i64 0, i64 2)
+
+  %ptr = getelementptr [5 x i8], [5 x i8]* @a12345, i32 0, i32 0
+  %res = call i8* @memchr(i8* %ptr, i32 3, i64 %n)
+  ret i8* %res
+}
+
+
+; Fold memchr(a12345, 259, n) to n < 4 ? null : a12345 + 3
+; to verify the constant 259 is converted to unsigned char (yielding 3).
+
+define i8* @call_a12345_259_n(i64 %n) {
+; CHECK-LABEL: @call_a12345_259_n(
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i64 %n, 3
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8* null, i8* getelementptr inbounds ([5 x i8], [5 x i8]* @a12345, i64 0, i64 2)
+
+  %ptr = getelementptr [5 x i8], [5 x i8]* @a12345, i32 0, i32 0
+  %res = call i8* @memchr(i8* %ptr, i32 259, i64 %n)
+  ret i8* %res
+}
+
+
+; Do no fold memchr(ax, 1, n).
+
+define i8* @call_ax_1_n(i64 %n) {
+; CHECK-LABEL: @call_ax_1_n(
+; CHECK-NEXT:    [[RES:%.*]] = call i8* @memchr
+; CHECK-NEXT:    ret i8* [[RES]]
+
+  %ptr = getelementptr [0 x i8], [0 x i8]* @ax, i32 0, i32 0
+  %res = call i8* @memchr(i8* %ptr, i32 1, i64 %n)
+  ret i8* %res
+}

llvm/test/Transforms/InstCombine/memchr-3.ll

+; Verify that the special case of memchr calls with the size of 1 are
+; folded as expected.
+;
+; RUN: opt < %s -passes=instcombine -S | FileCheck %s
+
+declare i8* @memchr(i8*, i32, i64)
+
+@ax = external global [0 x i8]
+@a12345 = constant [5 x i8] c"\01\02\03\04\05"
+
+
+; Fold memchr(a12345, 1, 1) to a12345.
+
+define i8* @fold_memchr_a12345_1_1() {
+; CHECK-LABEL: @fold_memchr_a12345_1_1(
+; CHECK-NEXT:    ret i8* getelementptr inbounds ([5 x i8], [5 x i8]* @a12345, i64 0, i64 0)
+
+  %ptr = getelementptr [5 x i8], [5 x i8]* @a12345, i32 0, i32 0
+  %res = call i8* @memchr(i8* %ptr, i32 1, i64 1)
+  ret i8* %res
+}
+
+
+; Fold memchr(a12345, 2, 1) to null.
+
+define i8* @fold_memchr_a12345_2_1() {
+; CHECK-LABEL: @fold_memchr_a12345_2_1(
+; CHECK-NEXT:    ret i8* null
+
+  %ptr = getelementptr [5 x i8], [5 x i8]* @a12345, i32 0, i32 0
+  %res = call i8* @memchr(i8* %ptr, i32 2, i64 1)
+  ret i8* %res
+}
+
+
+; Fold memchr(ax, 257, 1) to (unsigned char)*ax == 1 ? ax : null
+; to verify the constant 257 is converted to unsigned char (yielding 1).
+
+define i8* @fold_memchr_ax_257_1(i32 %chr, i64 %n) {
+; CHECK-LABEL: @fold_memchr_ax_257_1(
+; CHECK-NEXT:    [[AX_0:%.*]] = load i8, i8* getelementptr inbounds ([0 x i8], [0 x i8]* @ax, i64 0, i64 0)
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[AX_0]], 1
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8* getelementptr inbounds ([0 x i8], [0 x i8]* @ax, i64 0, i64 0), i8* null
+; CHECK-NEXT:    ret i8* [[SEL]]
+
+  %ptr = getelementptr [0 x i8], [0 x i8]* @ax, i32 0, i32 0
+  %res = call i8* @memchr(i8* %ptr, i32 257, i64 1)
+  ret i8* %res
+}
+
+
+; Fold memchr(ax, c, 1) to (unsigned char)*ax == (unsigned char)c ? ax : null.
+
+define i8* @fold_memchr_ax_c_1(i32 %chr, i64 %n) {
+; CHECK-LABEL: @fold_memchr_ax_c_1(
+; CHECK-NEXT:    [[AX_0:%.*]] = load i8, i8* getelementptr inbounds ([0 x i8], [0 x i8]* @ax, i64 0, i64 0)
+; CHECK-NEXT:    [[CHR_TRUNC:%.*]] = trunc i32 %chr to i8
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[AX_0]], [[CHR_TRUNC]]
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8* getelementptr inbounds ([0 x i8], [0 x i8]* @ax, i64 0, i64 0), i8* null
+; CHECK-NEXT:    ret i8* [[SEL]]
+
+  %ptr = getelementptr [0 x i8], [0 x i8]* @ax, i32 0, i32 0
+  %res = call i8* @memchr(i8* %ptr, i32 %chr, i64 1)
+  ret i8* %res
+}