[LV] Enforce addrec in current loop for uncountable exit load address check

Addresses post-commit review raised for #145663

Author: huntergr-arm

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

@@ -1878,11 +1878,11 @@ bool LoopVectorizationLegality::canUncountableExitConditionLoadBeMoved(
 
   // Make sure that the load address is not loop invariant; we want an
   // address calculation that we can rotate to the next vector iteration.
-  const SCEV *PtrScev = PSE.getSE()->getSCEV(Ptr);
-  if (!isa<SCEVAddRecExpr>(PtrScev)) {
+  const auto *AR = dyn_cast<SCEVAddRecExpr>(PSE.getSE()->getSCEV(Ptr));
+  if (!AR || AR->getLoop() != TheLoop || !AR->isAffine()) {
     reportVectorizationFailure(
         "Uncountable exit condition depends on load with an address that is "
-        "not an add recurrence",
+        "not an add recurrence in the loop",
         "EarlyExitLoadInvariantAddress", ORE, TheLoop);
     return false;
   }

llvm/test/Transforms/LoopVectorize/early_exit_store_legality.ll

@@ -353,7 +353,7 @@ exit:
 
 define void @loop_contains_store_condition_load_is_chained(ptr dereferenceable(40) noalias %array, ptr align 8 dereferenceable(160) readonly %offsets, ptr align 2 dereferenceable(40) readonly %pred) {
 ; CHECK-LABEL: LV: Checking a loop in 'loop_contains_store_condition_load_is_chained'
-; CHECK:       LV: Not vectorizing: Uncountable exit condition depends on load with an address that is not an add recurrence.
+; CHECK:       LV: Not vectorizing: Uncountable exit condition depends on load with an address that is not an add recurrence in the loop.
 entry:
   br label %for.body
 
@@ -407,7 +407,7 @@ exit:
 
 define void @loop_contains_store_condition_load_requires_gather(ptr dereferenceable(40) noalias %array, ptr align 2 dereferenceable(512) readonly %pred, ptr align 1 dereferenceable(20) readonly %offsets) {
 ; CHECK-LABEL: LV: Checking a loop in 'loop_contains_store_condition_load_requires_gather'
-; CHECK:       LV: Not vectorizing: Uncountable exit condition depends on load with an address that is not an add recurrence.
+; CHECK:       LV: Not vectorizing: Uncountable exit condition depends on load with an address that is not an add recurrence in the loop.
 entry:
   br label %for.body
 
@@ -544,7 +544,7 @@ exit:
 
 define void @uncountable_exit_condition_address_is_invariant(ptr dereferenceable(40) noalias %array, ptr align 2 dereferenceable(2) readonly %pred) {
 ; CHECK-LABEL: LV: Checking a loop in 'uncountable_exit_condition_address_is_invariant'
-; CHECK:       LV: Not vectorizing: Uncountable exit condition depends on load with an address that is not an add recurrence.
+; CHECK:       LV: Not vectorizing: Uncountable exit condition depends on load with an address that is not an add recurrence in the loop.
 entry:
   br label %for.body
 
@@ -567,5 +567,40 @@ exit:
   ret void
 }
 
+define void @uncountable_exit_condition_address_is_addrec_in_outer_loop(ptr dereferenceable(40) noalias %array, ptr align 2 dereferenceable(2) readonly %pred) {
+; CHECK-LABEL: LV: Checking a loop in 'uncountable_exit_condition_address_is_addrec_in_outer_loop'
+; CHECK:       LV: Not vectorizing: Uncountable exit condition depends on load with an address that is not an add recurrence in the loop.
+entry:
+  br label %outer.body
+
+outer.body:
+  %outer.iv = phi i64 [ 0, %entry ], [ %outer.iv.next, %outer.inc ]
+  %ee.addr = getelementptr inbounds nuw i16, ptr %pred, i64 %outer.iv
+  br label %for.body
+
+for.body:
+  %iv = phi i64 [ 0, %outer.body ], [ %iv.next, %for.inc ]
+  %st.addr = getelementptr inbounds nuw i16, ptr %array, i64 %iv
+  %data = load i16, ptr %st.addr, align 2
+  %inc = add nsw i16 %data, 1
+  store i16 %inc, ptr %st.addr, align 2
+  %ee.val = load i16, ptr %ee.addr, align 2
+  %ee.cond = icmp sgt i16 %ee.val, 500
+  br i1 %ee.cond, label %exit, label %for.inc
+
+for.inc:
+  %iv.next = add nuw nsw i64 %iv, 1
+  %counted.cond = icmp eq i64 %iv.next, 20
+  br i1 %counted.cond, label %outer.inc, label %for.body
+
+outer.inc:
+  %outer.iv.next = add nuw nsw i64 %outer.iv, 1
+  %outer.cond = icmp eq i64 %outer.iv.next, 2
+  br i1 %outer.cond, label %exit, label %outer.body
+
+exit:
+  ret void
+}
+
 declare void @init_mem(ptr, i64);
 declare i64 @get_an_unknown_offset();