Hi All,
We would like to propose a generalized strategy for vectorizing early exit loops with
multiple early exits. The idea is to check the exiting conditions first and execute the loop body later.
Previous Work
A masking-based approach for vectorizing early-exit loops with
stores was recently added in the loop vectorizer.
It derives a per-lane predicate from the exit condition, uses a
first-active-lane reduction to locate the exiting lane, and predicates stores
in the body under that mask. It’s functionally correct, but comes with a few constraints:
- Single uncountable exit only.
- The exit condition has to match a strict shape:
icmp(load, loop-invariant) - Not supported on architechtures without masked stores.
We propose to introduce a new strategy ‘Check-first’ where we try to generalize
vectorizing loops with n early exits. Through check-first strategy we increase the
scope to early exiting loops with:
- Multiple early exits.
- Flexible exit conditions. Instead of the rigid
icmp(load, invariant). - Nested exit conditions.
Additionally it can even work on machines without masked stores.
Implementation
The check-first transform is implemented as a VPlan transformation
within the loop vectorizer.
At a high level, the transform restructures the original vector
loop into three regions:
-
Check cascade: A sequence of check blocks (one per early exit) that
evaluate vectorized exit conditions. Each check block holds the
condition slice: the minimal set of instructions needed to evaluate
that exit. If any exit fires, control transfers to the scalar loop. -
Vector body: Executes only when no exit fires. Contains the stores
and remaining computation, unchanged from the original loop. -
Exit handling: When an exit fires, the iterations within the
exiting vector chunk that precede the exit lane still need to execute.
Two strategies exist:- Scalar replay (default): control transfers to the standard scalar
epilogue loop, which re-executes from the start of the current
vector chunk. No special replay infrastructure is needed. - Masked replay (experimental): a dedicated block clones the body’s stores
with lane masks derived from the first active exit lane,
avoiding the scalar loop entirely.
- Scalar replay (default): control transfers to the standard scalar
In VPlan terms, the transformation:
- Computes a condition slice for each exit and partitions the header’s
recipes between check blocks and the body block. - Wires the check cascade:
check.0 -> check.1 -> ... -> body, with each
check branching to the replay path on exit. - Attaches guard masks to stores under conditional predicates so only
guard-true lanes are written. - For masked replay, clones stores into the replay block with combined
guard and lane masks.
Example
// Source Loop:
for (i = 0; i < n; ++i) {
A[i] = B[i] + C[i];
if (X[i])
break; // early exit
}
// Vector Loop:
for (i = 0; i < n; i += VF) {
if (X[i, ... (i + VF - 1)]) // check region
branch to scalar_loop;
else
A[i, ... (i + VF - 1)] =
B[i, ... (i + VF - 1)] + C[i, ... (i + VF - 1)]; // vector region
}
// scalar_loop:
for (i < n; ++i) {
A[i] = B[i] + C[i];
if (X[i])
break;
}
IR before vectorization:
loop.header:
%iv = phi i64 [ 0, %entry ], [ %iv.next, %loop.latch ]
%b.ptr = getelementptr inbounds [1024 x i32], ptr @B, i64 0, i64 %iv
%b = load i32, ptr %b.ptr, align 4
%c.ptr = getelementptr inbounds [1024 x i32], ptr @C, i64 0, i64 %iv
%c = load i32, ptr %c.ptr, align 4
%sum = add i32 %b, %c
%a.ptr = getelementptr inbounds [1024 x i32], ptr @A, i64 0, i64 %iv
store i32 %sum, ptr %a.ptr, align 4
%x.ptr = getelementptr inbounds [1024 x i32], ptr @X, i64 0, i64 %iv
%x = load i32, ptr %x.ptr, align 4
%exit.cond = icmp ne i32 %x, 0
br i1 %exit.cond, label %found, label %loop.latch
loop.latch:
%iv.next = add nuw nsw i64 %iv, 1
%latch.cond = icmp eq i64 %iv.next, 1024
br i1 %latch.cond, label %done, label %loop.header
IR after vectorization (VF=4):
vector.check:
%index = phi i64 [ 0, %vector.ph ], [ %index.next, %vector.body ]
%x.ptr = getelementptr inbounds [1024 x i32], ptr @X, i64 0, i64 %index
%wide.load = load <4 x i32>, ptr %x.ptr, align 4
%exit.cond = icmp ne <4 x i32> %wide.load, zeroinitializer
%any.exit = call i1 @llvm.vector.reduce.or.v4i1(<4 x i1> %exit.cond)
br i1 %any.exit, label %scalar.ph, label %vector.body
vector.body:
%b.ptr = getelementptr inbounds [1024 x i32], ptr @B, i64 0, i64 %index
%b.vec = load <4 x i32>, ptr %b.ptr, align 4
%c.ptr = getelementptr inbounds [1024 x i32], ptr @C, i64 0, i64 %index
%c.vec = load <4 x i32>, ptr %c.ptr, align 4
%sum = add <4 x i32> %b.vec, %c.vec
%a.ptr = getelementptr inbounds [1024 x i32], ptr @A, i64 0, i64 %index
store <4 x i32> %sum, ptr %a.ptr, align 4
%index.next = add nuw i64 %index, 4
%cond = icmp eq i64 %index.next, 1024
br i1 %cond, label %middle.block, label %vector.check
scalar.ph:
br label %loop.header ; resume scalar loop at i = %index
Restrictions and Limitations
The current implementation targets a deliberately narrow set of loops as a
first step. Known limitations include:
- Fixed-width vector factors only.
- No dedicated cost model for the replay path overhead.
- Interleaving is not supported.
- Masked replay requires target support for masked stores instruction.
- Masked replay is supported only for loops with a single early exit, similar to the existing mask-based early-exit vectorization.
- Reductions are not yet supported.
Performance Impact
We demonstrate the impact of our check-first approach using the below microbenchmark kernels.
Each kernel runs 10 million iterations over N=1024 elements.
Loops run full length to measure performance. We used the below flags to enable
check-first implementation.
Configuration:
Check-First: -mllvm -enable-early-exit-vectorization
-mllvm -enable-check-first-early-exit-vectorization=true
Native: -mllvm -enable-early-exit-vectorization
-mllvm -enable-early-exit-vectorization-with-side-effects
-mllvm -enable-check-first-early-exit-vectorization=false
Test Cases
Test 1: Single exit, single store.
for (i = 0; i < N; i++) {
dst[i] = src[i] + 1;
if (pred[i] == 0)
return i;
}
Test 2: Multiple stores, two exits.
for (i = 0; i < N; i++) {
dst1[i] = src[i] + 1;
dst2[i] = src[i] + scale[i];
dst3[i] = src[i] ^ scale[i];
if (pred[i] == 0)
return i;
}
Test 3: Two exits with a store in between.
for (i = 0; i < N; i++) {
if (pred1[i] == 0)
return i;
dst[i] = src[i] + 1;
if (pred2[i] == 0)
return i;
}
Test 4: Nested exit with a store.
for (i = 0; i < N; i++) {
dst[i] = src[i] + 1;
if (guard[i]) {
if (pred[i] == 0)
return i;
}
}
Test 5: Three exits with multiple stores.
for (i = 0; i < N; i++) {
if (pred1[i] == 0)
return i;
dst1[i] = src[i] + 1;
if (pred2[i] == 0)
return i;
dst2[i] = src[i] * 6;
if (pred3[i] == 0)
return i;
}
Results
| Test | Native (ms) | Check-First (ms) | Speedup |
|---|---|---|---|
| test1_single_exit_single_store | 433.207 | 334.749 | 1.29x |
| test2_single_exit_multi_store | 676.677 | 533.841 | 1.26x |
| test3_two_exits_store_between | 5383.772 | 497.132 | 10.82x |
| test4_nested_exit_store | 5319.367 | 468.918 | 11.34x |
| test5_three_exits_multi_store | 7669.573 | 756.151 | 10.14x |
The following experiment is a small example to demonstrate the benefit of our approach at various exit positions.
int kernel(uint8_t dst[restrict static N],
const uint8_t src[restrict static N],
const uint8_t pred1[static N],
const uint8_t pred2[static N]) {
for (int i = 0; i < N; ++i) {
if (pred1[i] == 0)
return i;
dst[i] = src[i] + 1;
if (pred2[i] == 0)
return i;
}
return -1;
}
pred1 triggers the exit at various positions. pred2 is never triggered.
Results with N=128, VF=16:
| Test | Native (ms) | Check-First (ms) | Speedup |
|---|---|---|---|
| exit @ position 2 | 18.259 | 19.161 | .95x |
| exit @ position 30 | 144.300 | 81.144 | 1.77x |
| exit @ position 62 | 293.708 | 99.296 | 2.95x |
| exit @ position 80 | 376.788 | 42.437 | 8.87x |
| no exit | 662.149 | 61.709 | 10.73x |
Happy to iterate on this proposal based on review.
We have kept the first patch simple just focusing on trivial loops. If the direction looks right, we’ll follow up with the incremental updates.
Patch
Thanks!
Arjun