Merge pull request swiftlang#33113 from atrick/silopt-critedge

SILCloner: fix critical edge handling; remove previous workarounds

Author: atrick

include/swift/SIL/SILCloner.h

@@ -18,9 +18,11 @@
 #define SWIFT_SIL_SILCLONER_H
 
 #include "swift/AST/ProtocolConformance.h"
-#include "swift/SIL/SILOpenedArchetypesTracker.h"
+#include "swift/SIL/BasicBlockUtils.h"
+#include "swift/SIL/Dominance.h"
 #include "swift/SIL/SILBuilder.h"
 #include "swift/SIL/SILDebugScope.h"
+#include "swift/SIL/SILOpenedArchetypesTracker.h"
 #include "swift/SIL/SILVisitor.h"
 
 namespace swift {
@@ -44,6 +46,7 @@ class SILCloner : protected SILInstructionVisitor<ImplClass> {
   /// MARK: Context shared with CRTP extensions.
 
   SILBuilder Builder;
+  DominanceInfo *DomTree = nullptr;
   TypeSubstitutionMap OpenedExistentialSubs;
   SILOpenedArchetypesTracker OpenedArchetypesTracker;
 
@@ -75,12 +78,15 @@ class SILCloner : protected SILInstructionVisitor<ImplClass> {
   using SILInstructionVisitor<ImplClass>::asImpl;
 
   explicit SILCloner(SILFunction &F,
-                     SILOpenedArchetypesTracker &OpenedArchetypesTracker)
-      : Builder(F), OpenedArchetypesTracker(OpenedArchetypesTracker) {
+                     SILOpenedArchetypesTracker &OpenedArchetypesTracker,
+                     DominanceInfo *DT = nullptr)
+      : Builder(F), DomTree(DT),
+        OpenedArchetypesTracker(OpenedArchetypesTracker) {
     Builder.setOpenedArchetypesTracker(&OpenedArchetypesTracker);
   }
 
-  explicit SILCloner(SILFunction &F) : Builder(F), OpenedArchetypesTracker(&F) {
+  explicit SILCloner(SILFunction &F, DominanceInfo *DT = nullptr)
+      : Builder(F), DomTree(DT), OpenedArchetypesTracker(&F) {
     Builder.setOpenedArchetypesTracker(&OpenedArchetypesTracker);
   }
 
@@ -442,8 +448,9 @@ class SILClonerWithScopes : public SILCloner<ImplClass> {
 public:
   SILClonerWithScopes(SILFunction &To,
                       SILOpenedArchetypesTracker &OpenedArchetypesTracker,
+                      DominanceInfo *DT = nullptr,
                       bool Disable = false)
-      : SILCloner<ImplClass>(To, OpenedArchetypesTracker) {
+      : SILCloner<ImplClass>(To, OpenedArchetypesTracker, DT) {
 
     // We only want to do this when we generate cloned functions, not
     // when we inline.
@@ -694,17 +701,36 @@ void SILCloner<ImplClass>::visitBlocksDepthFirst(SILBasicBlock *startBB) {
     asImpl().visitInstructionsInBlock(BB);
 
     unsigned dfsSuccStartIdx = dfsWorklist.size();
-    for (auto &succ : BB->getSuccessors()) {
+
+    // splitEdge may rewrite BB's successors during this loop.
+    for (unsigned succIdx = 0, numSucc = BB->getSuccessors().size();
+         succIdx != numSucc; ++succIdx) {
+
       // Only visit a successor that has not already been visited and was not
       // premapped by the client.
-      if (BBMap.count(succ))
-        continue;
-
+      if (BBMap.count(BB->getSuccessors()[succIdx])) {
+        // After cloning BB, this successor may be a new CFG merge. If it is
+        // valid to branch directly from the BB to its clone do nothing; if not,
+        // split the edge from BB->succ and clone the new block.
+        //
+        // A CFG merge may require new block arguments, so check for both a
+        // critical edge and the ability to add branch arguments (BranchInst).
+        if (BB->getSingleSuccessorBlock()
+            && isa<BranchInst>(BB->getTerminator())) {
+          continue;
+        }
+        // This predecessor has multiple successors, so cloning it without
+        // cloning its successors would create a critical edge.
+        splitEdge(BB->getTerminator(), succIdx, DomTree);
+        assert(!BBMap.count(BB->getSuccessors()[succIdx]));
+      }
       // Map the successor to a new BB. Layout the cloned blocks in the order
       // they are visited and cloned.
       lastClonedBB =
           getBuilder().getFunction().createBasicBlockAfter(lastClonedBB);
 
+      // After splitting, BB has stable successors.
+      auto &succ = BB->getSuccessors()[succIdx];
       BBMap.insert(std::make_pair(succ.getBB(), lastClonedBB));
 
       dfsWorklist.push_back(succ);

include/swift/SIL/TypeSubstCloner.h

@@ -144,8 +144,9 @@ class TypeSubstCloner : public SILClonerWithScopes<ImplClass> {
                   SILFunction &From,
                   SubstitutionMap ApplySubs,
                   SILOpenedArchetypesTracker &OpenedArchetypesTracker,
+                  DominanceInfo *DT = nullptr,
                   bool Inlining = false)
-    : SILClonerWithScopes<ImplClass>(To, OpenedArchetypesTracker, Inlining),
+    : SILClonerWithScopes<ImplClass>(To, OpenedArchetypesTracker, DT, Inlining),
       SwiftMod(From.getModule().getSwiftModule()),
       SubsMap(ApplySubs),
       Original(From),

include/swift/SILOptimizer/Utils/BasicBlockOptUtils.h

@@ -136,9 +136,9 @@ class SinkAddressProjections {
 /// block's branch to jump to the newly cloned block, call cloneBranchTarget
 /// instead.
 ///
-/// After cloning, call splitCriticalEdges, then updateSSAAfterCloning. This is
-/// decoupled from cloning becaused some clients perform CFG edges updates after
-/// cloning but before splitting CFG edges.
+/// After cloning, call updateSSAAfterCloning. This is decoupled from cloning
+/// becaused some clients perform CFG edges updates after cloning but before
+/// splitting CFG edges.
 class BasicBlockCloner : public SILCloner<BasicBlockCloner> {
   using SuperTy = SILCloner<BasicBlockCloner>;
   friend class SILCloner<BasicBlockCloner>;
@@ -225,12 +225,7 @@ class BasicBlockCloner : public SILCloner<BasicBlockCloner> {
 
   bool wasCloned() { return isBlockCloned(origBB); }
 
-  /// Call this after processing all instructions to fix the control flow
-  /// graph. The branch cloner may have left critical edges.
-  bool splitCriticalEdges(DominanceInfo *domInfo, SILLoopInfo *loopInfo);
-
-  /// Helper function to perform SSA updates after calling both
-  /// cloneBranchTarget and splitCriticalEdges.
+  /// Helper function to perform SSA updates after calling cloneBranchTarget.
   void updateSSAAfterCloning();
 
 protected:

lib/SILOptimizer/LoopTransforms/ArrayPropertyOpt.cpp

@@ -430,37 +430,24 @@ namespace {
 /// in a set of basic blocks. Updates the dominator tree with the cloned blocks.
 /// However, the client needs to update the dominator of the exit blocks.
 ///
-/// FIXME: SILCloner is used to cloned CFG regions by multiple clients. All
-/// functionality for generating valid SIL (including the DomTree) should be
-/// handled by the common SILCloner.
+/// FIXME: All functionality for generating valid SIL (including the DomTree)
+/// should be handled by the common SILCloner. Currently, SILCloner only updates
+/// the DomTree for original (non-cloned) blocks when splitting edges. The
+/// cloned blocks won't be mapped to dominator nodes until fixDomTree()
+/// runs. However, since SILCloner always handles single-entry regions,
+/// fixDomTree() could be part of SILCloner itself.
 class RegionCloner : public SILCloner<RegionCloner> {
-  DominanceInfo &DomTree;
   SILBasicBlock *StartBB;
 
   friend class SILInstructionVisitor<RegionCloner>;
   friend class SILCloner<RegionCloner>;
 
 public:
   RegionCloner(SILBasicBlock *EntryBB, DominanceInfo &DT)
-      : SILCloner<RegionCloner>(*EntryBB->getParent()), DomTree(DT),
-        StartBB(EntryBB) {}
+      : SILCloner<RegionCloner>(*EntryBB->getParent(), &DT), StartBB(EntryBB) {}
 
   SILBasicBlock *cloneRegion(ArrayRef<SILBasicBlock *> exitBBs) {
-    assert (DomTree.getNode(StartBB) != nullptr && "Can't cloned dead code");
-
-    // We need to split any edge from a non cond_br basic block leading to a
-    // exit block. After cloning this edge will become critical if it came from
-    // inside the cloned region. The SSAUpdater can't handle critical non
-    // cond_br edges.
-    //
-    // FIXME: remove this in the next commit. The SILCloner will always do it.
-    for (auto *BB : exitBBs) {
-      SmallVector<SILBasicBlock *, 8> Preds(BB->getPredecessorBlocks());
-      for (auto *Pred : Preds)
-        if (!isa<CondBranchInst>(Pred->getTerminator()) &&
-            !isa<BranchInst>(Pred->getTerminator()))
-          splitEdgesFromTo(Pred, BB, &DomTree, nullptr);
-    }
+    assert(DomTree->getNode(StartBB) != nullptr && "Can't cloned dead code");
 
     cloneReachableBlocks(StartBB, exitBBs);
 
@@ -477,25 +464,25 @@ class RegionCloner : public SILCloner<RegionCloner> {
   void fixDomTree() {
     for (auto *BB : originalPreorderBlocks()) {
       auto *ClonedBB = getOpBasicBlock(BB);
-      auto *OrigDomBB = DomTree.getNode(BB)->getIDom()->getBlock();
+      auto *OrigDomBB = DomTree->getNode(BB)->getIDom()->getBlock();
       if (BB == StartBB) {
         // The cloned start node shares the same dominator as the original node.
-        auto *ClonedNode = DomTree.addNewBlock(ClonedBB, OrigDomBB);
+        auto *ClonedNode = DomTree->addNewBlock(ClonedBB, OrigDomBB);
         (void)ClonedNode;
         assert(ClonedNode);
         continue;
       }
       // Otherwise, map the dominator structure using the mapped block.
-      DomTree.addNewBlock(ClonedBB, getOpBasicBlock(OrigDomBB));
+      DomTree->addNewBlock(ClonedBB, getOpBasicBlock(OrigDomBB));
     }
   }
 
   SILValue getMappedValue(SILValue V) {
     if (auto *BB = V->getParentBlock()) {
-      if (!DomTree.dominates(StartBB, BB)) {
+      if (!DomTree->dominates(StartBB, BB)) {
         // Must be a value that dominates the start basic block.
-        assert(DomTree.dominates(BB, StartBB) &&
-               "Must dominated the start of the cloned region");
+        assert(DomTree->dominates(BB, StartBB)
+               && "Must dominated the start of the cloned region");
         return V;
       }
     }

lib/SILOptimizer/LoopTransforms/LoopRotate.cpp

@@ -409,10 +409,13 @@ bool swift::rotateLoop(SILLoop *loop, DominanceInfo *domInfo,
   header->moveAfter(latch);
 
   // Merge the old latch with the old header if possible.
-  mergeBasicBlockWithSuccessor(latch, domInfo, loopInfo);
+  if (mergeBasicBlockWithSuccessor(latch, domInfo, loopInfo))
+    newHeader = latch; // The old Header is gone. Latch is now the Header.
 
-  // Create a new preheader.
-  splitIfCriticalEdge(preheader, newHeader, domInfo, loopInfo);
+  // Cloning the header into the preheader created critical edges from the
+  // preheader and original header to both the new header and loop exit.
+  splitCriticalEdgesFrom(preheader, domInfo, loopInfo);
+  splitCriticalEdgesFrom(newHeader, domInfo, loopInfo);
 
   if (shouldVerify) {
     domInfo->verify();

lib/SILOptimizer/LoopTransforms/LoopUnroll.cpp

@@ -74,7 +74,8 @@ void LoopCloner::cloneLoop() {
   Loop->getExitBlocks(ExitBlocks);
 
   // Clone the entire loop.
-  cloneReachableBlocks(Loop->getHeader(), ExitBlocks);
+  cloneReachableBlocks(Loop->getHeader(), ExitBlocks,
+                       /*insertAfter*/Loop->getLoopLatch());
 }
 
 /// Determine the number of iterations the loop is at most executed. The loop
@@ -93,7 +94,7 @@ static Optional<uint64_t> getMaxLoopTripCount(SILLoop *Loop,
   if (!Loop->isLoopExiting(Latch))
     return None;
 
-  // Get the loop exit condition.
+ // Get the loop exit condition.
   auto *CondBr = dyn_cast<CondBranchInst>(Latch->getTerminator());
   if (!CondBr)
     return None;