[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));
+      }
 
-  // Check whether inner loops are uniform. At this point, we only support
-  // simple outer loops scenarios with uniform nested loops.
-  if (!isUniformLoopNest(TheLoop /*loop nest*/,
-                         TheLoop /*context outer loop*/)) {
-    reportVectorizationFailure(
-        "Outer loop contains divergent loops",
-        "loop control flow is not understood by vectorizer", "CFGNotUnderstood",
-        ORE, TheLoop);
-    if (DoExtraAnalysis)
-      Result = false;
-    else
-      return false;
+      // If the condition is not uniform, report a failure. We currently require
+      // uniform conditions to avoid the complexity of vectorizing divergent
+      // control flow in the outer loop.
+      if (!isUniformCondBr) {
+        reportVectorizationFailure(
+            "Outer loop contains divergent conditional branch",
+            "loop control flow is not understood by vectorizer",
+            "CFGNotUnderstood", ORE, TheLoop);
----------------
NexMing wrote:

I removed the previous overly specialized checks. I am not quite sure what more specific information you would like to see here.

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