[LV] Initial support for stores in early exit loops #137774

huntergr-arm · 2025-04-29T09:16:40Z

Adds some basic support for a simple early exit loop with a store.

This is vectorized such that when the next vector iteration would exit, we bail out to the scalar loop to handle the exit.

The main complication faced is transforming the vplan; it feels quite fragile at the moment, and might be better done up front at initial recipe creation time. Thoughts?

An alternative would be to use tail-folding, but that requires doing brka/brkb inside the loop, which could be costly. And would still likely require moving some IR, just not copying to the preheader.

(One advantage to this approach is that architectures without predication could potentially take advantage of ee autovec, though I suspect we'd want this driven by pragmas instead of just guessing about a possible trip count)

Adds some basic support for a simple early exit loop with a store. This is vectorized such that when the next vector iteration would exit, we bail out to the scalar loop to handle the exit.

david-arm · 2025-04-29T11:04:44Z

It's worth noting that brka/brkb is very much SVE-specific here, but there are other targets where it may be less of an issue.

david-arm

Thanks for this! I've only reviewed the legality bits, but thought I'd leave the comments I have so far.

david-arm · 2025-05-09T13:23:27Z

llvm/include/llvm/Transforms/Vectorize/LoopVectorizationLegality.h

+  bool isConditionCopyRequired() const { return RequiresEarlyExitConditionCopy; }
+
+  /// Returns the load instruction, if any, nearest to an uncountable early
+  /// exit.


Does nearest refer to nearest load prior to the early exit, after, or either?

The load closest in the IR graph to the exit (ideally, the load whose result is used directly in a comparison). I didn't spend much time on the naming ;)

david-arm · 2025-05-09T13:23:52Z

llvm/include/llvm/Transforms/Vectorize/LoopVectorizationLegality.h

@@ -407,6 +407,13 @@ class LoopVectorizationLegality {
    return hasUncountableEarlyExit() ? getUncountableEdge()->second : nullptr;
  }

+  /// Returns true if this is an early exit loop containing a store.
+  bool isConditionCopyRequired() const { return RequiresEarlyExitConditionCopy; }


Can the store be anywhere in the loop, i.e. including after the early exit? Might be worth clarifying.

If we do stick with this code in the final version I think it probably needs a better name because it's not obvious that it specifically refers to the condition driving an early exit. How would this scale with multiple exits?

The store can be anywhere in the loop, since we either execute a full vector iteration as normal or we bail out to the scalar tail if an exit would occur in the middle.

We need to do this for all exits, so a single transformation to copy and reorder the exit condition IR is sufficient.

david-arm · 2025-05-09T13:25:41Z

llvm/include/llvm/Transforms/Vectorize/LoopVectorizationLegality.h

+  /// The load used to determine an uncountable early-exit condition. This is
+  /// only used to allow further analysis in canVectorizeMemory if we found
+  /// what looks like a valid early exit loop with store beforehand.
+  std::optional<LoadInst *> EarlyExitLoad;


What happens in the case of two loads?, i.e.

%ld1 = load i8, ...
%ld2 = load i8, ...
%cmp = icmp eq i8 %ld1, %ld2
br i1 %cmp, label %early.exit, %loop.inc

This would fail the check that the second operand to the compare is loop invariant. There's a backup check in the vplan transform that would reject it at that point too.

In future, it would be nice to increase the number of cases handled, but I've tried to keep it simple for now and just come up with a list of extra cases to handle later.

Some of the possible future work:

Supporting a combined condition; e.g. no escaping value, but two exit conditions or'd together. Needs work in ScalarEvolution? (This actually applies to the sample loop I chose; for the IR test I manually edited the IR to create 2 exits again)

Multiple uses of loads or comparisons in the exit IR chain; requires introducing a PHI node for current vector iteration value.

Supporting a chain of conditional loads; e.g. early exit is itself in a conditional block so might not always execute.

Supporting a non-const second term for comparison, or more varied comparisons (e.g. bit test instead of icmp).

Tail folding of EE loops with a store.

Dynamic bounds for known-dereferenceable load (via e.g. call void @llvm.assume(i1 true) [ "align"(ptr %pred, i64 2), "dereferenceable"(ptr %pred, i32 %n_bytes) ])

david-arm · 2025-05-09T13:30:35Z