[llvm] [LV] Use SCEV loop-uniformity for outer-loop branch legality (PR #199632)

[Ming Yan via llvm-commits]

================
@@ -621,38 +541,39 @@ bool LoopVectorizationLegality::canVectorizeOuterLoop() {
     }
 
     // Check whether the branch is a supported one. Only unconditional
-    // branches, conditional branches with an outer loop invariant condition or
+    // branches, conditional branches with an outer loop uniform condition or
     // backedges are supported.
     // FIXME: We skip these checks when VPlan predication is enabled as we
     // want to allow divergent branches. This whole check will be removed
     // once VPlan predication is on by default.
     auto *Br = dyn_cast<CondBrInst>(Term);
-    if (Br && !TheLoop->isLoopInvariant(Br->getCondition()) &&
-        !LI->isLoopHeader(Br->getSuccessor(0)) &&
-        !LI->isLoopHeader(Br->getSuccessor(1))) {
-      reportVectorizationFailure(
-          "Unsupported conditional branch",
-          "loop control flow is not understood by vectorizer",
-          "CFGNotUnderstood", ORE, TheLoop);
-      if (DoExtraAnalysis)
-        Result = false;
-      else
-        return false;
-    }
-  }
+    if (Br && !TheLoop->isLoopLatch(BB)) {
+      bool isUniformCondBr = TheLoop->isLoopInvariant(Br->getCondition());
+
+      auto *Cmp = dyn_cast<CmpInst>(Br->getCondition());
+      auto *SE = PSE.getSE();
+      if (!isUniformCondBr && Cmp &&
+          SE->isSCEVable(Cmp->getOperand(0)->getType())) {
+        const SCEV *LhsExpr = PSE.getSCEV(Cmp->getOperand(0));
+        const SCEV *RhsExpr = PSE.getSCEV(Cmp->getOperand(1));
+        isUniformCondBr |= (SE->isLoopUniform(LhsExpr, TheLoop) &&
+                            SE->isLoopUniform(RhsExpr, TheLoop));
+      }
----------------
NexMing wrote:

Yes, I plan to drop the old behavior checks because they are too specialized. I do not really care about inner-loop IVs; as long as the control flow within the loop is uniform, the vectorization is legal. This patch is also intended to improve outer-loop vectorization support.

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