---

yaml
---
r: 221183
b: refs/heads/tensorflow-merge
c: 175eb5e
h: refs/heads/master
i:
  221181: 3c86bb8
  221179: 36eb32f
  221175: 0845776
  221167: 91b1b6d
  221151: 30fea0d
  221119: 197e1cd
  221055: b3fa5f1
  220927: fdf1979
  220671: 807e010
  220159: 82ba960
  219135: 94af4e4
  217087: ea6baf1
  212991: 1aeddf8

Author: bgogul

[refs]

@@ -906,7 +906,7 @@ refs/tags/swift-DEVELOPMENT-SNAPSHOT-2018-06-14-a: 888f12b456d2b9b0e35518c1972bb
 refs/tags/swift-DEVELOPMENT-SNAPSHOT-2018-06-15-a: 683450174579052b46e32f3b8f7fcc6fe8127db9
 refs/tags/swift-DEVELOPMENT-SNAPSHOT-2018-06-16-a: 5d81d51e075adbc6128d5612ac1d9a43186d0396
 refs/heads/move-debugging-executables-into-its-own-section: 968835a1ac02f692a0f5d8ebd60949ab6b0f14bc
-refs/heads/tensorflow-merge: fa69bf24de77e0f654d3345eda6ad934b363ad93
+refs/heads/tensorflow-merge: 175eb5e2d1374eb43bd8e83d5162b88862da6156
 refs/tags/swift-4.2-DEVELOPMENT-SNAPSHOT-2018-06-17-a: b947b41f2e23855c11bd483879813857b41e63ec
 refs/tags/swift-4.2-DEVELOPMENT-SNAPSHOT-2018-06-18-a: 76370219d53fcd9b9d7440166a992c9fd21909f7
 refs/tags/swift-4.2-DEVELOPMENT-SNAPSHOT-2018-06-19-a: 171b02e41ec9c3c8b592cecf7da70ab560d06cc5

branches/tensorflow-merge/lib/SILOptimizer/Mandatory/TFCanonicalizeCFG.cpp

@@ -316,9 +316,13 @@ class SingleExitLoopTransformer {
   /// Create a new latch block and clone all the arguments from new header.
   SILBasicBlock *createNewLatch(SILBasicBlock *newHeader);
 
-  // Patch the edges that go to the header and exit blocks to the new header or
-  // latch block as appropriate. Also, return the map consisting of indices
-  // assigned to the exit blocks.
+  /// Replace the preheader->header edge with preheader->newHeader edge
+  /// and update the arguments of preheader to match that of newheader.
+  void patchPreheader(SILBasicBlock *newHeader);
+
+  /// Patch the edges that go to the header and exit blocks to the new header or
+  /// latch block as appropriate. Also, return the map consisting of indices
+  /// assigned to the exit blocks.
   llvm::DenseMap<SILBasicBlock *, intmax_t>
   patchEdges(SILBasicBlock *newHeader, SILBasicBlock *latchBlock);
 
@@ -368,9 +372,18 @@ void SingleExitLoopTransformer::initialize() {
   loop->getExitBlocks(exitBlocks);
   // All the exiting edges need to be rewired.
   loop->getExitEdges(edgesToFix);
-  // All the edges to the header need to be rewired.
-  for (SILBasicBlock *headerPred : header->getPredecessorBlocks()) {
-    edgesToFix.emplace_back(headerPred, header);
+  edgesToFix.emplace_back(latch, header);
+  // Split critical edges in edgesToFix before we do any transformations.
+  for (auto &edge : edgesToFix) {
+    SILBasicBlock *src = const_cast<SILBasicBlock *>(edge.first);
+    SILBasicBlock *tgt = const_cast<SILBasicBlock *>(edge.second);
+    SILBasicBlock *splitBlock = splitIfCriticalEdge(src, tgt, /*DT*/ nullptr, LI);
+    if (splitBlock != nullptr) {
+      // If the edge was critical then splitBlock would have been inserted
+      // between src and tgt as follows: src -> splitBlock -> tgt.  Therefore,
+      // update src as the new edge to patch would be splitBlock -> tgt.
+      edge.first = splitBlock;
+    }
   }
 
   for (const SILBasicBlock *bb : loop->getBlocks()) {
@@ -502,6 +515,30 @@ SingleExitLoopTransformer::createNewLatch(SILBasicBlock *newHeader) {
   return latchBlock;
 }
 
+void SingleExitLoopTransformer::patchPreheader(SILBasicBlock *newHeader) {
+  // Update edge
+  replaceBranchTarget(preheader->getTerminator(), header, newHeader,
+                      /*preserveArgs*/ true);
+  SILBuilder builder(preheader->getTerminator());
+  SILLocation location(
+    getUserSourceLocation(preheader->getTerminator()->getDebugLocation()));
+  // Add arguments corresponding to escaping arguments.
+  // State from within the loop is not available in the preheader.
+  // Simply pass in an undef. This will never be accessed at runtime.
+  SmallVector<SILValue, 8> newArgs;
+  for (const SILValue &escapingValue : escapingValues) {
+    newArgs.push_back(getUndef(escapingValue->getType()));
+  }
+  // `exitIndex` to identify the block to which we exit from the loop.
+  newArgs.push_back(createTFIntegerConst(*deviceInfo, builder, location,
+                                         /*bitwidth*/ 32,
+                                         /*exitIndex*/ 0));
+  // `stayInLoop` flag
+  newArgs.push_back(createTFIntegerConst(*deviceInfo, builder, location,
+                                         /*bitwidth*/ 1, true));
+  appendArguments(preheader->getTerminator(), newHeader, newArgs);
+}
+
 llvm::DenseMap<SILBasicBlock *, intmax_t>
 SingleExitLoopTransformer::patchEdges(SILBasicBlock *newHeader,
                                       SILBasicBlock *latchBlock) {
@@ -537,14 +574,7 @@ SingleExitLoopTransformer::patchEdges(SILBasicBlock *newHeader,
   for (const auto &edge : edgesToFix) {
     SILBasicBlock *src = const_cast<SILBasicBlock *>(edge.first);
     SILBasicBlock *tgt = const_cast<SILBasicBlock *>(edge.second);
-    SILBasicBlock *newTgt = (src == preheader) ? newHeader : latchBlock;
-    SILBasicBlock *splitBlock = splitIfCriticalEdge(src, tgt, /*DT*/ nullptr, LI);
-    if (splitBlock != nullptr) {
-      // If the edge was critical then splitBlock would have been inserted
-      // between src and tgt as follows: src -> splitBlock -> tgt.  Therefore,
-      // update src as the new edge to patch would be splitBlock -> tgt.
-      src = splitBlock;
-    }
+    SILBasicBlock *newTgt = latchBlock;
     bool stayInLoop = loop->contains(tgt);
     replaceBranchTarget(src->getTerminator(), tgt, newTgt,
                         /*preserveArgs=*/stayInLoop);
@@ -564,13 +594,8 @@ SingleExitLoopTransformer::patchEdges(SILBasicBlock *newHeader,
     for (const SILValue &escapingValue : escapingValues) {
       if (DI->properlyDominates(escapingValue, src->getTerminator())) {
         newArgs.push_back(escapingValue);
-      } else if (src != preheader) {
-        // newHeader arguments are available if src is not the preheader.
-        newArgs.push_back(newHeader->getArgument(argIndex));
       } else {
-        // State from within the loop is not available in the preheader.
-        // Simply pass in an undef. This will never be accessed at runtime.
-        newArgs.push_back(getUndef(escapingValue->getType()));
+        newArgs.push_back(newHeader->getArgument(argIndex));
       }
       ++argIndex;
     }
@@ -690,6 +715,8 @@ bool SingleExitLoopTransformer::transform() {
   // Create a new latch block.
   SILBasicBlock *latchBlock = createNewLatch(newHeader);
 
+  patchPreheader(newHeader);
+
   // Patch the edges and return the map consisting of indices assigned
   // to the exit blocks.
   llvm::DenseMap<SILBasicBlock *, intmax_t> exitIndices =