[llvm] [LV] Initial support for stores in early exit loops (PR #137774)

[Graham Hunter via llvm-commits]

================
@@ -654,6 +661,16 @@ class LoopVectorizationLegality {
   /// Keep track of the loop edge to an uncountable exit, comprising a pair
   /// of (Exiting, Exit) blocks, if there is exactly one early exit.
   std::optional<std::pair<BasicBlock *, BasicBlock *>> UncountableEdge;
+
+  /// Indicates that we will need to copy the early exit condition into
+  /// the vector preheader, as we will need to mask some operations in
+  /// the loop (e.g. stores).
+  bool RequiresEarlyExitConditionCopy = false;
+
+  /// 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;
----------------
huntergr-arm wrote:

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:

1.  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)

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

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

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

5. Tail folding of EE loops with a store.

6. 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) ])



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