[llvm] [LV] Vectorize early exit loops with stores using masking (PR #178454)

[Florian Hahn via llvm-commits]

================
@@ -2017,12 +2021,23 @@ bool LoopVectorizationLegality::canVectorize(bool UseVPlanNativePath) {
       return false;
   }
 
-  // Bail out for ReadWrite loops with uncountable exits for now.
-  if (UncountableExitType == UncountableExitTrait::ReadWrite) {
-    reportVectorizationFailure(
-        "Writes to memory unsupported in early exit loops",
-        "Cannot vectorize early exit loop with writes to memory",
-        "WritesInEarlyExitLoop", ORE, TheLoop);
+  // TODO: Remove this restriction once we're sure it's safe to do so.
+  //       Handling stores to invariant addresses will be slightly different
+  //       based on the vectorization style chosen. If we bail out to a scalar
+  //       tail before executing any lane that would take the uncountable exit,
+  //       then the store that occurs in the scalar loop would suffice.
+  //
+  //       If we instead handle the lane taking the uncountable exit within the
+  //       vectorized loop, then we will have to ensure that we extract the
+  //       last active lane at that point in the loop instead of the last lane
+  //       of the vector before performing a scalar store.
+  if (UncountableExitType == UncountableExitTrait::ReadWrite &&
----------------
fhahn wrote:

Yep that would probably be good

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