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

[Graham Hunter via llvm-commits]

================
@@ -8139,6 +8146,24 @@ void LoopVectorizationPlanner::buildVPlansWithVPRecipes(ElementCount MinVF,
                        CM.getMaxSafeElements());
         RUN_VPLAN_PASS(VPlanTransforms::optimizeEVLMasks, *Plan);
       }
+
+      // If we have side effects (e.g. stores) in a loop with uncountable
+      // exits, we must ensure the early exit condition is known before any
+      // state changing operations take place.
+      if (!Plan->hasScalarVFOnly() &&
+          Legal->hasUncountableExitWithSideEffects()) {
+        ArrayRef<BasicBlock *> UncountableExits =
+            Legal->getUncountableExitingBlocks();
+        // TODO: Support more than one.
+        if (UncountableExits.size() > 1)
+          break;
+        auto *UncountableBR =
+            cast<BranchInst>(UncountableExits.front()->getTerminator());
+        if (!RUN_VPLAN_PASS(
+                VPlanTransforms::handleUncountableExitsWithSideEffects, *Plan,
----------------
huntergr-arm wrote:

I've merged the transform into the existing EE logic; I was still relying a bit on the loop region before. Still seems to work ok.

I did add the operations to MaskedOp, but they still weren't considered to be predicated instructions. Looking at `LoopVectorizationCostModel::isPredicatedInst`, it seems I was right -- while we pass the initial check because `Legal->isMaskRequired()` returns true, the header doesn't need predication, and we then return false because we're not tail folding.

So I've added a check for whether we're in an ee-with-side-effects loop check there, too, and we now get scalarization of the non-condition memory ops. I guess we don't have a cli option to force support for masked loads/stores the way we do for forcing scalable VF support.

This patch _does_ support loads and stores after the early exit -- we just mask off the exiting lane, so we only touch addresses the scalar loop would have. I already had transform tests for that, but I've added a vplan print test for it as well.

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