[llvm] [LV] Vectorize early exit loops with multiple exits. (PR #174864)

[Florian Hahn via llvm-commits]

================
@@ -241,7 +241,64 @@ exit:
 }
 
 define i64 @two_early_exits_same_exit_with_constant_live_outs() {
-; CHECK: LV: Not vectorizing:  Auto-vectorization of loops with multiple uncountable early exits is not yet supported.
+; CHECK:      VPlan 'Initial VPlan for VF={4},UF>=1' {
+; CHECK-NEXT: Live-in vp<[[VF:%.+]]> = VF
+; CHECK-NEXT: Live-in vp<[[VFxUF:%.+]]> = VF * UF
+; CHECK-NEXT: Live-in vp<[[VTC:%.+]]> = vector-trip-count
+; CHECK-NEXT: Live-in ir<67> = original trip-count
+; CHECK-EMPTY:
+; CHECK-NEXT: ir-bb<entry>:
+; CHECK-NEXT:   IR   %A = alloca [1024 x i8], align 1
+; CHECK-NEXT:   IR   %B = alloca [1024 x i8], align 1
+; CHECK-NEXT:   IR   call void @init(ptr %A, i64 1024)
+; CHECK-NEXT:   IR   call void @init(ptr %B, i64 1024)
+; CHECK-NEXT: Successor(s): scalar.ph, vector.ph
+; CHECK-EMPTY:
+; CHECK-NEXT: vector.ph:
+; CHECK-NEXT: Successor(s): vector loop
+; CHECK-EMPTY:
+; CHECK-NEXT: <x1> vector loop: {
+; CHECK-NEXT:   vector.body:
+; CHECK-NEXT:     EMIT vp<[[CAN_IV:%.+]]> = CANONICAL-INDUCTION ir<0>, vp<%index.next>
+; CHECK-NEXT:     vp<[[SCALAR_STEPS:%.+]]> = SCALAR-STEPS vp<[[CAN_IV]]>, ir<1>, vp<[[VF]]>
+; CHECK-NEXT:     CLONE ir<%gep.A> = getelementptr inbounds ir<%A>, vp<[[SCALAR_STEPS]]>
+; CHECK-NEXT:     vp<[[PTRA:%.+]]> = vector-pointer inbounds ir<%gep.A>
+; CHECK-NEXT:     WIDEN ir<%ld.A> = load vp<[[PTRA]]>
+; CHECK-NEXT:     WIDEN ir<%cmp1> = icmp eq ir<%ld.A>, ir<42>
+; CHECK-NEXT:     CLONE ir<%gep.B> = getelementptr inbounds ir<%B>, vp<[[SCALAR_STEPS]]>
+; CHECK-NEXT:     vp<[[PTRB:%.+]]> = vector-pointer inbounds ir<%gep.B>
+; CHECK-NEXT:     WIDEN ir<%ld.B> = load vp<[[PTRB]]>
+; CHECK-NEXT:     WIDEN ir<%cmp2> = icmp eq ir<%ld.A>, ir<%ld.B>
+; CHECK-NEXT:     EMIT vp<%index.next> = add nuw vp<[[CAN_IV]]>, vp<[[VFxUF]]>
+; CHECK-NEXT:     EMIT vp<[[OR:%.+]]> = or ir<%cmp1>, ir<%cmp2>
+; CHECK-NEXT:     EMIT vp<[[ANY_OF:%.+]]> = any-of vp<[[OR]]>
+; CHECK-NEXT:     EMIT vp<[[CMP:%.+]]> = icmp eq vp<%index.next>, vp<[[VTC]]>
+; CHECK-NEXT:     EMIT branch-on-two-conds vp<[[ANY_OF]]>, vp<[[CMP]]>
+; CHECK-NEXT:   No successors
+; CHECK-NEXT: }
+; CHECK-NEXT: Successor(s): vector.early.exit, middle.block
+; CHECK-EMPTY:
+; CHECK-NEXT: middle.block:
+; CHECK-NEXT:   EMIT vp<%cmp.n> = icmp eq ir<67>, vp<[[VTC]]>
+; CHECK-NEXT:   EMIT branch-on-cond vp<%cmp.n>
+; CHECK-NEXT: Successor(s): ir-bb<exit>, scalar.ph
+; CHECK-EMPTY:
+; CHECK-NEXT: vector.early.exit:
+; CHECK-NEXT:   EMIT vp<%first.active.lane> = first-active-lane vp<[[OR]]>
+; CHECK-NEXT:   EMIT vp<%exit.cond.at.lane> = extract-lane vp<%first.active.lane>, ir<%cmp1>
----------------
fhahn wrote:

Yep that would be an option, although it would require a chain of horizontal reductions, which is quite expensive on AArch64 at least. The benefit of the current representation is that there's just a single any_of in the vector loop. It might be a good optimization on targets with fast any_of for loops where we otherwise run out of vector registers.

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