[llvm-branch-commits] [llvm] [VPlan][Predicator] Preserve some uniform control flow (PR #217485)

[Ashutosh Nema via llvm-branch-commits]

================
@@ -209,6 +236,80 @@ VPValue *VPPredicator::createMaskOr(VPValue *LHS, VPValue *RHS, DebugLoc DL) {
       HeaderMask, Builder.createOr(LHSRemainder, RHSRemainder, DL), DL);
 }
 
+VPValue *VPPredicator::reconstructSSA(VPBasicBlock *UseBB, VPValue *V,
+                                      bool IsMask) {
+  auto *RecipeValue = dyn_cast<VPRecipeValue>(V);
+  if (!RecipeValue)
+    return V;
+
+  auto &SSADefs = SSAReconstructionDefsMap[RecipeValue];
+  VPBasicBlock *DefBB = RecipeValue->getDefiningRecipe()->getParent();
+  SSADefs[DefBB] = RecipeValue;
+  VPBasicBlock *Header = Plan.getVectorLoopRegion()->getEntryBasicBlock();
+  if (DefBB != Header)
+    SSADefs[Header] =
+        IsMask ? Plan.getFalse() : Plan.getPoison(RecipeValue->getScalarType());
+  return vputils::reconstructSSA(UseBB, SSADefs,
+                                 /*CreateWidenPhis=*/true);
+}
+
+void VPPredicator::fixSSA(VPRecipeBase *U, VPValue *V, bool IsMask) {
+  VPValue *Fixed = reconstructSSA(U->getParent(), V, IsMask);
+  if (Fixed == V)
+    return;
+
+  LLVM_DEBUG({
+    dbgs() << "SSA Fixup in  ";
+    U->dump();
+    dbgs() << "  replacing ";
+    V->dump();
+    dbgs() << "  with ";
+    Fixed->dump();
+  });
+  U->replaceUsesOfWith(V, Fixed);
+}
+
+bool VPPredicator::shouldPreserveTerminator(VPBasicBlock *VPBB) {
+  LLVM_DEBUG(dbgs() << "Checking if can preserve the branch at the end of "
+                    << VPBB->getName() << "\n");
+  auto False = []([[maybe_unused]] StringRef Reason = "") {
+    LLVM_DEBUG(dbgs() << "  can't be preserved"
+                      << (Reason.empty() ? Twine() : (": " + Reason)) << "\n");
+    return false;
+  };
+  if (DisablePartialLinearization)
+    return False("Disabled");
+
+  if (VPBB->getNumSuccessors() != 2)
----------------
nema-ashutosh wrote:

worth extending it to uniform-switch same as uniform-if ?

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