[SimplifyLibCalls] Transform memchr(STR, C, N) to chain of ORs

[xbolva00]

Summary

Motivation:

#include <string_view>

size_t findFirst_ABCDEF(std::string_view sv) {
  return sv.find_first_of("ABCDEF");
}

memchr("ABCDEF", C, 6) != NULL -> (C == 'A' || C == 'B' || C == 'C' || C == 'D' || C == 'E' || C == 'F') != 0

[nikic]

This looks quite nice for the examples -- but can't this create a very large chain of comparisons in some cases? Like if you check every second character you'd end up with 128 comparisons? (Or maybe you need every third, I can imagine that we manage to fold every second character to am mask).

[xbolva00]

Added PhaseOrdering test to ensure instcombine + simplifycfg produces optimal switch code.

[xbolva00]

In D128011#3591797, @nikic wrote:

This looks quite nice for the examples -- but can't this create a very large chain of comparisons in some cases? Like if you check every second character you'd end up with 128 comparisons? (Or maybe you need every third, I can imagine that we manage to fold every second character to am mask).

Maybe we could introduce some limit (feel free to suggest some), but on the other hand, we often end up with nice switch or very simple code.

int g(char a)
{
   if (__builtin_memchr ("abcdefghijklmnopqrstuvwxyz", a, 27) != 0)
        return foo();
   else return -1;
}

define dso_local noundef i32 @_Z1gc(i8 noundef signext %a) local_unnamed_addr #1 {
entry:
  %0 = icmp ne i8 %a, 0
  %1 = add i8 %a, -123
  %2 = icmp ult i8 %1, -26
  %cmp.not = and i1 %0, %2
  br i1 %cmp.not, label %return, label %if.then

if.then:                                          ; preds = %entry
  %call1 = tail call noundef i32 @_Z3foov()
  br label %return

return:                                           ; preds = %entry, %if.then
  %retval.0 = phi i32 [ %call1, %if.then ], [ -1, %entry ]
  ret i32 %retval.0
}

or

int g(char a)
{
   if (__builtin_memchr ("acegikmoqsuwyz13456789", a, 22) != 0)
        return foo();
   else return -1;
}

define dso_local noundef i32 @_Z1gc(i8 noundef signext %a) local_unnamed_addr #1 {
entry:
  switch i8 %a, label %return [
    i8 122, label %if.then
    i8 121, label %if.then
    i8 119, label %if.then
    i8 117, label %if.then
    i8 115, label %if.then
    i8 113, label %if.then
    i8 111, label %if.then
    i8 109, label %if.then
    i8 107, label %if.then
    i8 105, label %if.then
    i8 103, label %if.then
    i8 101, label %if.then
    i8 99, label %if.then
    i8 97, label %if.then
    i8 57, label %if.then
    i8 56, label %if.then
    i8 55, label %if.then
    i8 54, label %if.then
    i8 53, label %if.then
    i8 52, label %if.then
    i8 51, label %if.then
    i8 49, label %if.then
  ]

...

X86: https://pastebin.com/m5gMEwGZ

[msebor]

An analogous optimization is applicable to other string functions besides memchr (and strchr), including for example strlen or even memcmp and strcmp (with one short constant array/string). With this one accepted as a precedent the other handlers can be enhanced as follow ups. It would be nice if the algorithm could be made general enough to let them all share it. It does seem that the expansion in all of them should be capped at some small number.

For what it's worth, a prototype of the corresponding patch for GCC is attached to PR 103798. It looks like it has a hardwired limit of 8 bytes.

[xbolva00]

strchr

strchr is transformed to memchr, if possible. See test strchr_to_memchr_n_equals_len.

It looks like it has a hardwired limit of 8 bytes.

Very small, imho, We usually end up with one switch instruction, not a big deal. But I am fine with 32 I think.

An analogous optimization is applicable to other string functions ... strlen or even memcmp and strcmp

I would not overcomplicate things for know, currently just 5 lines of code. Any possible followup could create create a new helper function.

For what it's worth, a prototype of the corresponding patch for GCC is attached to PR 103798.

Yes, I saw it. Oh, that expanded implementation in match.pd = D

[xbolva00]

Update memchr.ll

[xbolva00]

Ping

[msebor]

It seems important to put a limit on the length of the chain to avoid pathological expansion (that might come up in generated code, but not only there). I expect the optimal maximum will vary across different targets but a suitably low minimum should be target-independent. I don't have a good sense of what the ideal lower bound might be but I would hesitate to go over the proposed GCC limit of 8 unless the ORs could be assumed to collapse to something much simpler (like in memchr-7.ll). Instrumenting the compiler and building a few projects might help dial it in.

Having said that, I don't think it's a guaranteed win to set the limit to zero (or to disable the subsequent folding) when optimizing for size. But whatever the effect of -Os might be here, it should be exercised by the test suite if only to make the decision explicit.

[nikic]

I think for the limit, the important thing to consider is comparisons that collapse into range checks. We don't want to limit the total number of characters being checked, but rather the number of of non-contiguous ranges. The motivating cases reduce down to just one or two range checks, which is certainly profitable. The cases that would require a switch for decent lowering are less clearly profitable, especially when we take into account that such a switch is not being reliably produced (as your phase ordering test shows).

[xbolva00]

We don't want to limit the total number of characters being checked, but rather the number of of non-contiguous ranges.

I believe there is no easy StringRef API, right? Not sure if I want to work on such new API.

For now, I would go with GCC's limit of 8.

[msebor]

Something like this seems to do the right thing for a sorted list (s is a container like StringRef):

int last = 0;
int n = std::count_if (s.begin(), s.end(),
                       [&last](char c) {
                         bool ret = (unsigned char)c > (unsigned char)last + 1;
                         last = c;
                         return ret; });

(It would be a bit more involved if it had to handle different charsets between the host and the target.)

[xbolva00]

Restrict transformation based on the number of of non-contiguous ranges.

[xbolva00]

Added PhaseOrdering test.

[xbolva00]

Rebased.

Ping.

[nikic]

Why does this need a phase ordering test? What is the output after just InstCombine for it? Does InstCombine not reliably reduce the comparisons and require SimplifyCFG switch optimization? What if there is no actual branch involved and SimplifyCFG doesn't reduce it?

[xbolva00]

In D128011#4314930, @nikic wrote:

Why does this need a phase ordering test? What is the output after just InstCombine for it? Does InstCombine not reliably reduce the comparisons and require SimplifyCFG switch optimization? What if there is no actual branch involved and SimplifyCFG doesn't reduce it?

Indeed, we dont need it now. I will remove it.

[nikic]

llvm/lib/Transforms/Utils/SimplifyLibCalls.cpp

1405

So, the reason why this checks > rather than != is to handle the case where the same char appears in the string multiple times, right? I don't think there is a test for this case right now, could you please add one?

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

2

Can you please add -instcombine-infinite-loop-threshold=2 to this test, to make sure this doesn't require multiple iteration to fold.

[xbolva00]

Addressed review feedback. Added new test with char duplication.

[xbolva00]

Remove useless PhaseOrdering test.

[xbolva00]

any futher comments?

[nikic]

LGTM

llvm/lib/Transforms/Utils/SimplifyLibCalls.cpp

1395

Drop "logical". The or is not logical in the sense we use the word.

1408

clang-format. The indentation is wrong here.

[xbolva00]

Addressed review feedback, thanks.

llvm/lib/Transforms/Utils/SimplifyLibCalls.cpp

if (isOnlyUsedInEqualityComparison(CI, SrcStr))
       //   memchr(S, C, N) == S to N && *S == C for any C and N.
       // TODO: This is safe even even for nonconstant S.
       return memChrToCharCompare(CI, Size, B, DL);
 
     // From now on we need a constant length and constant array.
     return nullptr;
   }
 
+  bool OptForSize = CI->getFunction()->hasOptSize() ||
+                    llvm::shouldOptimizeForSize(CI->getParent(), PSI, BFI,
+                                                PGSOQueryType::IRPass);
+
   // 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 (Str.empty() || !isOnlyUsedInZeroEqualityComparison(CI))
+  if (OptForSize || Str.empty() || !isOnlyUsedInZeroEqualityComparison(CI))
     return nullptr;
 
   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;
+  if (!DL.fitsInLegalInteger(Max + 1)) {
+    // Build chain of logical ORs
+    // Transform:
+    //    memchr("abcd", C, 4) != nullptr
+    // to:
+    //    (C == 'a' || C == 'b' || C == 'c' || C == 'd') != 0
+    std::string SortedStr = Str.str();
+    llvm::sort(SortedStr);
+    // Compute the number of of non-contiguous ranges.
+    unsigned NonContRanges = 1;
+    for (size_t i = 1; i < SortedStr.size(); ++i) {
+        if (SortedStr[i] > SortedStr[i - 1] + 1) {
+            NonContRanges++;
+        }
+    }
+
+    // Restrict this optimization to profitable cases with one or two range checks.
+    if (NonContRanges > 2)
+      return nullptr;
+
+    SmallVector<Value *> CharCompares;
+    for (char C : SortedStr)
+      CharCompares.push_back(
+          B.CreateICmpEQ(CharVal, ConstantInt::get(CharVal->getType(), C)));
+
+    return B.CreateIntToPtr(B.CreateOr(CharCompares), CI->getType());
+  }
 
   // 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)

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

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -passes=instcombine -S | FileCheck %s
 
 @.str = private unnamed_addr constant [27 x i8] c"abcdefghijklmnopqrstuvwxyz\00", align 1
 @.str.1 = private unnamed_addr constant [2 x i8] c"\0D\0A", align 1
 
 declare ptr @strchr(ptr, i32)
 declare ptr @memchr(ptr, i32, i64)
 
 define zeroext i1 @strchr_to_memchr_n_equals_len(i32 %c) {
 ; CHECK-LABEL: @strchr_to_memchr_n_equals_len(
-; CHECK-NEXT:    [[MEMCHR:%.*]] = tail call ptr @memchr(ptr noundef nonnull dereferenceable(1) @.str, i32 [[C:%.*]], i64 27)
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ne ptr [[MEMCHR]], null
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i32 [[C:%.*]], 0
+; CHECK-NEXT:    [[TMP2:%.*]] = add i32 [[C]], -97
+; CHECK-NEXT:    [[TMP3:%.*]] = icmp ult i32 [[TMP2]], 26
+; CHECK-NEXT:    [[TMP4:%.*]] = or i1 [[TMP1]], [[TMP3]]
+; CHECK-NEXT:    ret i1 [[TMP4]]
 ;
   %call = tail call ptr @strchr(ptr nonnull dereferenceable(27) @.str, i32 %c)
   %cmp = icmp ne ptr %call, null
   ret i1 %cmp
 }
 
 define zeroext i1 @memchr_n_equals_len(i32 %c) {
 ; CHECK-LABEL: @memchr_n_equals_len(
 ; CHECK-NEXT:    [[TMP1:%.*]] = trunc i32 [[C:%.*]] to i8
 ; CHECK-NEXT:    [[TMP2:%.*]] = icmp ne i8 [[TMP1]], 10
 ; CHECK-NEXT:    [[TMP3:%.*]] = icmp ne i8 [[TMP1]], 13
 ; CHECK-NEXT:    [[CMP:%.*]] = and i1 [[TMP3]], [[TMP2]]
 ; CHECK-NEXT:    ret i1 [[CMP]]
 ;
   %call = tail call ptr @memchr(ptr nonnull dereferenceable(3) @.str.1, i32 %c, i64 2)
   %cmp = icmp eq ptr %call, null
   ret i1 %cmp
 }
 
 define zeroext i1 @memchr_n_less_than_len(i32 %c) {
 ; CHECK-LABEL: @memchr_n_less_than_len(
-; CHECK-NEXT:    [[CALL:%.*]] = tail call ptr @memchr(ptr noundef nonnull dereferenceable(1) @.str, i32 [[C:%.*]], i64 15)
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ne ptr [[CALL]], null
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    [[TMP1:%.*]] = add i32 [[C:%.*]], -97
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ult i32 [[TMP1]], 15
+; CHECK-NEXT:    ret i1 [[TMP2]]
 ;
   %call = tail call ptr @memchr(ptr @.str, i32 %c, i64 15)
   %cmp = icmp ne ptr %call, null
   ret i1 %cmp
 }
 
 
 define zeroext i1 @memchr_n_more_than_len(i32 %c) {
 ; CHECK-LABEL: @memchr_n_more_than_len(
-; CHECK-NEXT:    [[CALL:%.*]] = tail call ptr @memchr(ptr noundef nonnull dereferenceable(1) @.str, i32 [[C:%.*]], i64 30)
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ne ptr [[CALL]], null
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i32 [[C:%.*]], 0
+; CHECK-NEXT:    [[TMP2:%.*]] = add i32 [[C]], -97
+; CHECK-NEXT:    [[TMP3:%.*]] = icmp ult i32 [[TMP2]], 26
+; CHECK-NEXT:    [[TMP4:%.*]] = or i1 [[TMP1]], [[TMP3]]
+; CHECK-NEXT:    ret i1 [[TMP4]]
 ;
   %call = tail call ptr @memchr(ptr @.str, i32 %c, i64 30)
   %cmp = icmp ne ptr %call, null
   ret i1 %cmp
 }
 
 ; Negative test - no comparison with zero
 

llvm/test/Transforms/InstCombine/memchr.ll

 ;
   %dst = call ptr @memchr(ptr @single, i32 %C, i32 1)
   %cmp = icmp ne ptr %dst, null
   ret i1 %cmp
 }
 
 define i1 @test15(i32 %C) {
 ; CHECK-LABEL: @test15(
-; CHECK-NEXT:    [[DST:%.*]] = call ptr @memchr(ptr noundef nonnull dereferenceable(1) @negative, i32 [[C:%.*]], i32 3)
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ne ptr [[DST]], null
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    [[TMP1:%.*]] = add i32 [[C:%.*]], 2
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ult i32 [[TMP1]], 3
+; CHECK-NEXT:    ret i1 [[TMP2]]
 ;
   %dst = call ptr @memchr(ptr @negative, i32 %C, i32 3)
   %cmp = icmp ne ptr %dst, null
   ret i1 %cmp
 }
 
 @s = internal constant [1 x i8] [i8 0], align 1
 define ptr @pr32124() {

llvm/test/Transforms/PhaseOrdering/memchr-expansion.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -passes='instcombine,simplifycfg' -S < %s  | FileCheck %s
 
 @.str = private constant [7 x i8] c"abcdef\00", align 1
 
 declare ptr @memchr(ptr noundef, i32 noundef, i64 noundef)
 declare i32 @foo()
 
-define i32 @memchr_to_switch_br(i32 %a) {
-; CHECK-LABEL: @memchr_to_switch_br(
+define i32 @memchr_to_rangecheck_br(i32 %a) {
+; CHECK-LABEL: @memchr_to_rangecheck_br(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[CALL:%.*]] = call ptr @memchr(ptr noundef nonnull dereferenceable(1) @.str, i32 noundef [[A:%.*]], i64 noundef 6)
-; CHECK-NEXT:    [[CMP_NOT:%.*]] = icmp eq ptr [[CALL]], null
-; CHECK-NEXT:    br i1 [[CMP_NOT]], label [[RETURN:%.*]], label [[IF_THEN:%.*]]
+; CHECK-NEXT:    [[TMP0:%.*]] = add i32 [[A:%.*]], -103
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ult i32 [[TMP0]], -6
+; CHECK-NEXT:    br i1 [[TMP1]], label [[RETURN:%.*]], label [[IF_THEN:%.*]]
 ; CHECK:       if.then:
 ; CHECK-NEXT:    [[CALL1:%.*]] = call i32 @foo()
 ; CHECK-NEXT:    br label [[RETURN]]
 ; CHECK:       return:
 ; CHECK-NEXT:    [[RETVAL_0:%.*]] = phi i32 [ [[CALL1]], [[IF_THEN]] ], [ -1, [[ENTRY:%.*]] ]
 ; CHECK-NEXT:    ret i32 [[RETVAL_0]]
 ;
 entry:
   %call = call ptr @memchr(ptr noundef nonnull @.str, i32 noundef %a, i64 noundef 6) #3
   %cmp.not = icmp eq ptr %call, null
   br i1 %cmp.not, label %return, label %if.then
 
 if.then:                                          ; preds = %entry
   %call1 = call i32 @foo() #3
   br label %return
 
 return:                                           ; preds = %entry, %if.then
   %retval.0 = phi i32 [ %call1, %if.then ], [ -1, %entry ]
   ret i32 %retval.0
 }
 
-define i1 @memchr_to_switch_eq_ret(i32 %c) {
-; CHECK-LABEL: @memchr_to_switch_eq_ret(
-; CHECK-NEXT:    [[P:%.*]] = tail call ptr @memchr(ptr noundef nonnull dereferenceable(1) @.str, i32 noundef [[C:%.*]], i64 noundef 6)
-; CHECK-NEXT:    [[RET:%.*]] = icmp eq ptr [[P]], null
-; CHECK-NEXT:    ret i1 [[RET]]
+define i1 @memchr_to_rangecheck_eq_ret(i32 %c) {
+; CHECK-LABEL: @memchr_to_rangecheck_eq_ret(
+; CHECK-NEXT:    [[TMP1:%.*]] = add i32 [[C:%.*]], -103
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ult i32 [[TMP1]], -6
+; CHECK-NEXT:    ret i1 [[TMP2]]
 ;
   %p = tail call ptr @memchr(ptr noundef nonnull @.str, i32 noundef %c, i64 noundef 6)
   %ret = icmp eq ptr %p, null
   ret i1 %ret
 }
 
-define i1 @memchr_to_switch_ne_ret(i32 %c) {
-; CHECK-LABEL: @memchr_to_switch_ne_ret(
-; CHECK-NEXT:    [[P:%.*]] = tail call ptr @memchr(ptr noundef nonnull dereferenceable(1) @.str, i32 noundef [[C:%.*]], i64 noundef 6)
-; CHECK-NEXT:    [[RET:%.*]] = icmp ne ptr [[P]], null
-; CHECK-NEXT:    ret i1 [[RET]]
+define i1 @memchr_to_rangecheck_ne_ret(i32 %c) {
+; CHECK-LABEL: @memchr_to_rangecheck_ne_ret(
+; CHECK-NEXT:    [[TMP1:%.*]] = add i32 [[C:%.*]], -97
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ult i32 [[TMP1]], 6
+; CHECK-NEXT:    ret i1 [[TMP2]]
 ;
   %p = tail call ptr @memchr(ptr noundef nonnull @.str, i32 noundef %c, i64 noundef 6)
   %ret = icmp ne ptr %p, null
   ret i1 %ret
 }