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

[Ashutosh Nema via llvm-branch-commits]

================
@@ -478,30 +609,190 @@ void VPPredicator::run() {
     // Mask all VPInstructions in the block.
     for (VPRecipeBase &R : *VPBB) {
       if (auto *VPI = dyn_cast<VPInstruction>(&R))
-        VPI->addMask(BlockMask);
+        if (VPI->getOpcode() != VPInstruction::BranchOnCond)
+          VPI->addMask(BlockMask);
     }
   }
 
-  for (VPBlockBase *VPBB : reverse(BlocksInCompactRPOTOrder))
-    if (VPBB != Header)
-      convertPhisToBlends(cast<VPBasicBlock>(VPBB));
+  // Cache blend terms before linearization. Computing them requires the
+  // original phi predecessor mappings and CFG successor relation, both of
+  // which are rewritten below. Skip the header which is the first block.
+  for (auto [BlockIdx, VPBB] : drop_begin(enumerate(
+           VPBlockUtils::blocksOnly<VPBasicBlock>(BlocksInCompactRPOTOrder)))) {
+    auto &Terms = BlendTerms[BlockIdx];
+    for (VPRecipeBase &R : VPBB->phis()) {
+      auto *Phi = cast<VPPhi>(&R);
+      Terms[Phi] = computeBlendTerms(Phi);
+      LLVM_DEBUG({
+        dbgs() << "Computed blend terms for ";
+        Phi->dump();
+        dbgs() << "\n";
+        for (BlendTermTy &Term : Terms[Phi]) {
+          dbgs() << "  " << Term.second->getName() << ": ";
+          Term.first->dump();
+        }
+      });
+    }
+  }
+
+  // The following implements "Partial Control-Flow Linearization" by Simon Moll
+  // and Sebastian Hack.
+
+  using DeferredSuccessorsTy = SmallPtrSet<VPBlockBase *, 4>;
+  DenseMap<VPBlockBase *, DeferredSuccessorsTy> DeferredMap;
+  auto PopDeferred = [&](VPBlockBase *VPBB) -> DeferredSuccessorsTy {
+    auto It = DeferredMap.find(VPBB);
+    if (It != DeferredMap.end()) {
+      DeferredSuccessorsTy Res = std::move(It->second);
+      DeferredMap.erase(It);
+      return Res;
+    }
+    return {};
+  };
 
-  // Linearize the blocks of the loop into one serial chain.
-  VPBlockBase *PrevVPBB = nullptr;
   for (VPBasicBlock *VPBB :
        VPBlockUtils::blocksOnly<VPBasicBlock>(BlocksInCompactRPOTOrder)) {
+    LLVM_DEBUG(dbgs() << "Setting successors for " << VPBB->getName() << "\n");
     auto Successors = to_vector(VPBB->getSuccessors());
-    if (Successors.size() > 1)
-      VPBB->getTerminator()->eraseFromParent();
 
-    // Flatten the CFG in the loop. To do so, first disconnect VPBB from its
-    // successors. Then connect VPBB to the previously visited VPBB.
-    for (auto *Succ : Successors)
-      VPBlockUtils::disconnectBlocks(VPBB, Succ);
-    if (PrevVPBB)
-      VPBlockUtils::connectBlocks(PrevVPBB, VPBB);
+    if (shouldPreserveTerminator(VPBB)) {
----------------
nema-ashutosh wrote:

It keeps the uniform br but never drops then/else masks ?
due to this mask loads appears for following case:

  for (unsigned i = 0; i < len; i++)
    if (flag) A[i] = B[i] + C[i];
    else       X[i] = Y[i] * Z[i];



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