Simplify and refactor contractUncondEdges as a utility in TFCanonicalizeCFG.

lib/SILOptimizer/Mandatory/TFCanonicalizeCFG.cpp

@@ -118,6 +118,29 @@ void ConditionalSESERegion::print(llvm::raw_ostream &OS, unsigned indent) const
 // CFG Canonicalization Implementation
 //===----------------------------------------------------------------------===//
 
+// Our partitioning and other transformations can leave around lots of
+// unconditional branches between blocks that formerly had control edges.  Go
+// through and merge those to make later passes simpler.
+bool tf::contractUncondBranches(SILFunction *fn, DominanceInfo* DI, SILLoopInfo *LI) {
+  bool changed = false;
+  // Iterate carefully to avoid invalidating iterators: we mutate the block list
+  // while we walk it.
+  for (auto bbi = fn->begin(), e = fn->end(); bbi != e;) {
+    auto *bb = &*bbi;
+    if (mergeBasicBlockWithSuccessor(bb, DI, LI)) {
+      // The block was merged with this successor. Therefore, revisit this node:
+      // we have new successor(s) and may need to contract them as well.  Also,
+      // bbi may be invalidated at this point.
+      changed = true;
+      bbi = SILFunction::iterator(bb);
+    } else {
+      // Move to the next block if this was not merged.
+      ++bbi;
+    }
+  }
+  return changed;
+}
+
 namespace {
   class SESERegionBuilder {
     DominanceInfo DI;
@@ -1156,8 +1179,8 @@ namespace {
     /// The entry point to the transformation.
     void run() override {
       auto fn = getFunction();
+      contractUncondBranches(fn, /*DI*/nullptr, /*LI*/ nullptr);
       auto region = canonicalizeCFGForXLA(fn);
-
       llvm::outs() << "--- XLA CFG Canonicalize: " << fn->getName() << "\n";
       region->print(llvm::outs());
       llvm::outs() << "\n--- XLA CFG Canonicalize end\n";

lib/SILOptimizer/Mandatory/TFCanonicalizeCFG.h

@@ -20,6 +20,7 @@
 #define SWIFT_SILOPTIMIZER_TFCANONICALIZECFG_H
 
 #include "swift/SIL/SILBasicBlock.h"
+#include "swift/SIL/LoopInfo.h"
 
 namespace swift {
 namespace tf {
@@ -157,6 +158,11 @@ namespace tf {
     }
   };
 
+  // Our partitioning and other transformations can leave around lots of
+  // unconditional branches between blocks that formerly had control edges.  Go
+  // through and merge those to make later passes simpler.
+  bool contractUncondBranches(SILFunction *fn, DominanceInfo *DI,
+                              SILLoopInfo *LI);
 
   /// Transform the function into a properly nested series of
   /// single-entry-single-exit regions and return the data structure that

lib/SILOptimizer/Mandatory/TFPartition.cpp

@@ -16,6 +16,7 @@
 //===----------------------------------------------------------------------===//
 
 #define DEBUG_TYPE "tf-partition"
+#include "TFCanonicalizeCFG.h"
 #include "TFUtilities.h"
 #include "swift/AST/DiagnosticsSIL.h"
 #include "swift/AST/Expr.h"
@@ -4037,41 +4038,6 @@ bool TFFunctionPartition::partitionAndLowerGraph(bool isTest) {
   return partition(isTest) || lowerGraph(isTest);
 }
 
-// Our partitioning can leave around lots of unconditional branches between
-// blocks that formerly had control edges.  Go through and merge those to make
-// later passes simpler.
-static void contractUncondBranches(SILFunction *fn) {
-  // Iterate carefully to avoid invalidating iterators: we mutate the block list
-  // while we walk it.
-  for (auto bbi = fn->begin(), e = fn->end(); bbi != e;) {
-    auto *bb = &*bbi;
-    ++bbi; // Increment the iterator in case we do no transformation.
-
-    if (auto succ = bb->getSingleSuccessorBlock()) {
-      if (succ != bb && succ->getSinglePredecessorBlock()) {
-        if (auto *BI = dyn_cast<BranchInst>(bb->getTerminator())) {
-          // If there are any BB arguments in the destination, replace them with
-          // the branch operands, since they must dominate the dest block.
-          for (unsigned i = 0, e = BI->getArgs().size(); i != e; ++i) {
-            assert(succ->getArgument(i) != BI->getArg(i) &&
-                   "Cloned code regions are always reachable");
-            succ->getArgument(i)->replaceAllUsesWith(BI->getArg(i));
-          }
-
-          // Zap BI and move all of the instructions from DestBB into this one.
-          BI->eraseFromParent();
-          bb->spliceAtEnd(succ);
-          succ->eraseFromParent();
-
-          // Revisit this node: we have new successor(s) and may need to
-          // contract them as well.  Also, bbi may be invalidated at this point.
-          bbi = SILFunction::iterator(bb);
-        }
-      }
-    }
-  }
-}
-
 /// Return true if a user instruction is returning or forming a return value.
 static bool isReturning(SILInstruction *user) {
   if (isa<ReturnInst>(user)) return true;
@@ -4262,7 +4228,7 @@ bool TFFunctionPartition::partition(bool isTest) {
   // Our partitioning can leave around lots of unconditional branches between
   // blocks that formerly had control edges.  Go through and merge those to
   // make later passes simpler.
-  contractUncondBranches(acceleratorFn);
+  contractUncondBranches(acceleratorFn, /*DI*/ nullptr, /*LI*/ nullptr);
 
   if (auto outs = getTFDumpIntermediateStream()) {
     *outs << "--- TFPartition Accelerator Result: " << acceleratorFn->getName()

test/TensorFlow/contract_uncond_edges_test.sil

@@ -0,0 +1,113 @@
+// RUN: %target-sil-opt -tf-xla-cfg-canonicalize -tf-ensure-single-loop-exit -assume-parsing-unqualified-ownership-sil %s -o /dev/null | %FileCheck %s
+
+import Builtin
+import Swift
+import TensorFlow
+
+// Straight line blocks
+sil @$straightLine : $@convention(thin) (Builtin.Int32) -> Builtin.Int32 {
+bb0(%0 : $Builtin.Int32):
+  br bb1(%0 : $Builtin.Int32)
+
+bb1(%1: $Builtin.Int32):
+  br bb2(%1 : $Builtin.Int32)
+
+bb2(%2: $Builtin.Int32):
+  br bb3(%2 : $Builtin.Int32)
+
+bb3(%3 : $Builtin.Int32):
+	return %3 : $Builtin.Int32
+}
+// CHECK-LABEL: --- XLA CFG Canonicalize: $straightLine
+// CHECK:  block bb0
+// CHECK: --- XLA CFG Canonicalize end
+// CHECK: sil @$straightLine : $@convention(thin) (Builtin.Int32) -> Builtin.Int32 {
+// CHECK: bb0(%0 : $Builtin.Int32):
+// CHECK-NEXT:  return %0 : $Builtin.Int32
+// CHECK: }
+
+
+// Conditionals
+sil @$condRegion : $@convention(thin) (Builtin.Int32, Builtin.Int32) -> Builtin.Int32 {
+bb0(%0 : $Builtin.Int32, %1 : $Builtin.Int32):
+  %2 = builtin "cmp_slt_Int32"(%0 : $Builtin.Int32, %1 : $Builtin.Int32) : $Builtin.Int1
+  cond_br %2, bb1, bb4
+
+bb1:
+  br bb2(%0 : $Builtin.Int32)
+
+bb2(%3 : $Builtin.Int32):
+  br bb3(%3 : $Builtin.Int32)
+
+bb3(%4 : $Builtin.Int32):
+  br bb5(%4 : $Builtin.Int32)
+
+bb4:
+  br bb5(%1 : $Builtin.Int32)
+
+bb5(%5 : $Builtin.Int32):
+	return %5 : $Builtin.Int32
+}
+// CHECK-LABEL: --- XLA CFG Canonicalize: $condRegion
+// CHECK:[sequence
+// CHECK:  {condition Header: bb0
+// CHECK:    block bb1
+// CHECK:    block bb2}
+// CHECK:  block bb3]
+// CHECK: --- XLA CFG Canonicalize end
+// CHECK: sil @$condRegion : $@convention(thin) {{.*}} {
+// CHECK: bb0(%0 : $Builtin.Int32, %1 : $Builtin.Int32):
+// CHECK:   cond_br %2, bb1, bb2
+// CHECK: bb1:
+// CHECK:   br bb3(%0 : $Builtin.Int32)
+// CHECK: bb2:
+// CHECK:   br bb3(%1 : $Builtin.Int32)
+// CHECK: bb3([[A:%.*]] : $Builtin.Int32):
+// CHECK:   return [[A]] : $Builtin.Int32
+// CHECK: }
+
+// Simple loop
+sil @$simpleLoop : $@convention(thin) (Builtin.Int32, Builtin.Int32) -> Builtin.Int32 {
+bb0(%0 : $Builtin.Int32, %1 : $Builtin.Int32):
+  br bb1(%0 : $Builtin.Int32)
+
+bb1(%2 : $Builtin.Int32):
+  %3 = builtin "cmp_slt_Int32"(%2 : $Builtin.Int32, %1 : $Builtin.Int32) : $Builtin.Int1
+  cond_br %3, bb3, bb2
+
+bb2:
+  br bb6(%2 : $Builtin.Int32)
+
+bb3:
+  br bb4(%2 : $Builtin.Int32)
+
+bb4(%4 : $Builtin.Int32):
+  br bb5(%4 : $Builtin.Int32)
+
+bb5(%5 : $Builtin.Int32):
+  br bb1(%5 : $Builtin.Int32)
+
+bb6(%6 : $Builtin.Int32):
+  return %6 : $Builtin.Int32
+}
+// CHECK-LABEL:--- XLA CFG Canonicalize: $simpleLoop
+// CHECK:[sequence
+// CHECK:  <while Preheader: bb0, Header: bb1, exit: bb2
+// CHECK:    block bb3>
+// CHECK:  block bb2]
+// CHECK:--- XLA CFG Canonicalize end
+// CHECK:sil @$simpleLoop : $@convention(thin) (Builtin.Int32, Builtin.Int32) -> Builtin.Int32 {
+// CHECK:bb0(%0 : $Builtin.Int32, %1 : $Builtin.Int32):
+// CHECK:  br bb1(%0 : $Builtin.Int32)
+// CHECK:bb1(%3 : $Builtin.Int32):
+// CHECK:  %4 = builtin "cmp_slt_Int32"(%3 : $Builtin.Int32, %1 : $Builtin.Int32) : $Builtin.Int1
+// CHECK:  cond_br %4, bb3, bb2
+// CHECK:bb2:
+// CHECK:  return %3 : $Builtin.Int32
+// CHECK:bb3:
+// CHECK:  br bb1(%3 : $Builtin.Int32)
+// CHECK:}
+
+
+
+