[llvm] [AArch64] Test all needles when vectorising find_first_of loops. (PR #179298)

[David Sherwood via llvm-commits]

================
@@ -52,19 +53,21 @@ define ptr @find_first_of_i8(ptr %search_start, ptr %search_end, ptr %needle_sta
 ; CHECK-NEXT:    [[NEEDLE0_SPLAT:%.*]] = shufflevector <vscale x 16 x i8> [[NEEDLE0_SPLATINSERT]], <vscale x 16 x i8> poison, <vscale x 16 x i32> zeroinitializer
 ; CHECK-NEXT:    [[NEEDLE_SPLAT:%.*]] = select <vscale x 16 x i1> [[NEEDLE_MASKED]], <vscale x 16 x i8> [[NEEDLE_LOAD_VEC]], <vscale x 16 x i8> [[NEEDLE0_SPLAT]]
 ; CHECK-NEXT:    [[NEEDLE_VEC:%.*]] = call <16 x i8> @llvm.vector.extract.v16i8.nxv16i8(<vscale x 16 x i8> [[NEEDLE_SPLAT]], i64 0)
-; CHECK-NEXT:    [[MATCH_PRED:%.*]] = call <vscale x 16 x i1> @llvm.experimental.vector.match.nxv16i8.v16i8(<vscale x 16 x i8> [[SEARCH_LOAD_VEC]], <16 x i8> [[NEEDLE_VEC]], <vscale x 16 x i1> [[SEARCH_MASKED]])
-; CHECK-NEXT:    [[TMP3:%.*]] = call i1 @llvm.vector.reduce.or.nxv16i1(<vscale x 16 x i1> [[MATCH_PRED]])
-; CHECK-NEXT:    br i1 [[TMP3]], label %[[CALCULATE_MATCH:.*]], label %[[NEEDLE_CHECK_VEC]]
+; CHECK-NEXT:    [[MATCH_SEGMENT:%.*]] = call <vscale x 16 x i1> @llvm.experimental.vector.match.nxv16i8.v16i8(<vscale x 16 x i8> [[SEARCH_LOAD_VEC]], <16 x i8> [[NEEDLE_VEC]], <vscale x 16 x i1> [[SEARCH_MASKED]])
+; CHECK-NEXT:    [[MATCH_ACCUMULATOR]] = or <vscale x 16 x i1> [[PMATCH]], [[MATCH_SEGMENT]]
+; CHECK-NEXT:    [[NEEDLE_NEXT_VEC]] = getelementptr i8, ptr [[PNEEDLE]], i64 16
+; CHECK-NEXT:    [[TMP3:%.*]] = icmp ult ptr [[NEEDLE_NEXT_VEC]], [[NEEDLE_END]]
----------------
david-arm wrote:

Unfortunately this isn't safe because the pointer could wrap. Suppose NEEDLE_END=UINT64_MAX (a perfectly valid address) and NEEDLE_START=UINT64_MAX - 5 (another perfectly valid address). Both lie on the same page in memory so we pass the runtime checks. At the end of this loop we increment NEEDLE_START by 16, which causes it to wrap around to 10. This will then appear to be less than NEEDLE_END and you'll be stuck in an infinite loop.

There are a few ways of solving this problem - you can look at how we handle tail-folding in the loop vectoriser for SVE to see this. Two solutions are:

1. Add more runtime checks to ensure that NEEDLE_END is < UINT64_MAX - 16, which means that NEEDLE_NEXT_VEC can never wrap, or
2. Modify the check to use an extra integer induction variable based on the trip count, i.e. TRIP_COUNT = NEEDLE_END - NEEDLE_START. We know the trip count can never be > the page size, and so long as the minimum page size is something sensible (i.e. < UINT64_MAX/2 or whatever the limit of a ptrint_t type is) then the induction variable can never overflow.

For approach 2 it would be something like:

```
  [[TC:%.*]] = sub i64 [[NEEDLE_END]], [[NEEDLE_START]]
NEEDLE_CHECK_VEC:
  [[IDX:%.*]] = phi i64 [ 0, %[[FIND_FIRST_VEC_HEADER]] ], [ [[IDX_NEXT:%.*]], %[[NEEDLE_CHECK_VEC]] ]
  [[NEEDLE_PRED:%.*]] = call <vscale x 16 x i1> @llvm.get.active.lane.mask.nxv16i1.i64(i64 [[IDX]], i64 [[%tc]])
  [[PNEEDLE:%.*]] = getelementptr i8, ptr [[NEEDLE_START]], i64 [[IDX]]
  ...
  [[IDX_NEXT]] = add i64 [[IDX]], 16
  [[TMP3:%.*]] = icmp ult ptr [[IDX_NEXT]], [[TC]]
```

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