Update comments and tests for better readability.

Author: bgogul

lib/SILOptimizer/Mandatory/TFCanonicalizeCFG.cpp

@@ -420,11 +420,13 @@ class BasicBlockCloner : public SILClonerWithScopes<BasicBlockCloner> {
                                    "times during SESE cloning.");
   }
 
-  /// Clone the body of `loop` starting from `startBlock` and nest the cloned
-  /// fragment into the parent loop. If `startBlock` is the same as the header
-  /// of `loop`, we clone the entire loop including the back edge.  Otherwise,
-  /// we clone one iteration of the loop body without the back edge.
-  SILLoop *cloneLoop(SILLoopInfo *LI, SILLoop *loop, SILBasicBlock *startBlock) {
+  /// Utility to unroll one iteration of the loop or to clone the entire loop.
+  ///   - If `startBlock` is the same as the header of `loop`, we clone the
+  ///     entire loop including the back edge.
+  ///   - Otherwise, we unroll one iteration of the loop body starting from
+  ///     `startBlock' to the latch.
+  /// The unrolled or cloned version is nested into the parent loop.
+  SILLoop *cloneOrUnrollLoop(SILLoopInfo *LI, SILLoop *loop, SILBasicBlock *startBlock) {
     llvm::DenseMap<SILLoop*, SILLoop*> loopClones;
     // This is for convenience as top-level loops have nullptr for parent loop.
     loopClones[nullptr] = nullptr;
@@ -515,6 +517,8 @@ class BasicBlockCloner : public SILClonerWithScopes<BasicBlockCloner> {
 // A helper class to transform a loop to have a single exit from the header.
 class SingleExitLoopTransformer {
   public:
+    // Convert the given loop into a SESE form. Returns false if the loop was
+    // already in SESE form. Otherwise, returns true.
     static bool doIt(GraphFunctionDeviceInfo *deviceInfo, SILLoopInfo *LI,
                      DominanceInfo *DI, SILLoop *loop, PostDominanceInfo *PDI) {
       SingleExitLoopTransformer transformer(deviceInfo, LI, DI, loop, PDI);
@@ -528,7 +532,7 @@ class SingleExitLoopTransformer {
 #ifndef NDEBUG
       {
         // Verify that the loop is OK after all the transformations.
-        llvm::DenseSet<const SILLoop*> nestedLoops;
+        llvm::DenseSet<const SILLoop *> nestedLoops;
         loop->verifyLoopNest(&nestedLoops);
       }
 #endif
@@ -742,9 +746,9 @@ void SingleExitLoopTransformer::ensureSingleExitBlock() {
   //       break;
   //     }
   //   }
+  // The unrelated loops are those that are not contained within each other.
   SmallPtrSet<SILBasicBlock *, 32> unrelatedPreheaders;
   for (auto *otherLoop : *LI) {
-    // If loop and otherLoop are not nested into either, remember the preheader.
     if (!otherLoop->contains(loop) && !loop->contains(otherLoop)) {
       unrelatedPreheaders.insert(otherLoop->getLoopPreheader());
     }
@@ -808,8 +812,8 @@ void SingleExitLoopTransformer::ensureSingleExitBlock() {
         // If `succ` is a preheader of an unrelated loop, we will have to clone
         // the entire loop now so that we can also incrementally update LoopInfo.
         if (succBlockLoop) {
-          SILLoop *clonedLoop =
-              cloner.cloneLoop(LI, succBlockLoop, succBlockLoop->getHeader());
+          SILLoop *clonedLoop = cloner.cloneOrUnrollLoop(
+              LI, succBlockLoop, succBlockLoop->getHeader());
           changeBranchTarget(clonedSucc->getTerminator(), 0,
                              clonedLoop->getHeader(), /*preserveArgs*/ true);
           // Note that all the nodes of `clonedLoop` should be moved into the
@@ -1501,7 +1505,7 @@ void SingleExitLoopTransformer::unrollLoopBodyOnce() {
   }
 
   // Clone everything starting from the old header.
-  cloner.cloneLoop(LI, loop, header);
+  cloner.cloneOrUnrollLoop(LI, loop, header);
 
   // Get the clone for old header.
   SILBasicBlock *clonedOldHeader = cloner.remapBasicBlock(header);

test/TensorFlow/sese_loop_canonicalization.sil

@@ -294,17 +294,17 @@ bb3 (%9 : $Builtin.Int32):
 // }
 // The CFG is shown below:
 //
-//  [0]-----+
+// [0]-------+
 //  |         \
 // [1]       [2]
-//  |          |
-//  |   +--- [4]
-//  |   |    /   \
-//  |   | [6]    [5]
-//  |   +/ \      |
-// [3]----[8]     |
-//  |             |
-//  +-->[9]<------+
+//  |         |
+//  |   +--->[4]
+//  |  [7]   /  \
+//  |   |  [6]  [5]
+//  V   +-+ \    |
+// [3]<-----[8]  |
+//  |            |
+//  +-->[9]<-----+
 //
 // This is an example where the nodes that lead to the common post dominator of all exit nodes
 // is also reachable from nodes outside of the loop. The key issue is that bb9 is the common
@@ -380,8 +380,20 @@ bb9(%18 : $Builtin.Int32):
 // reachable from nodes outside of the loop. The only difference in this example
 // is that some of these nodes also belong to a loop. Therefore, it is required
 // for us to update the LoopInfo in addition to cloning. The loop in question
-// consists of the single node bb3.
-
+// consists of the single node bb3. The CFG is shown below:
+//
+//     [0]-------+
+//      |         \
+//     [1]       [2]
+//      |         |
+//      |   +--->[4]
+//      |  [7]   /  \
+//      |   |  [6]  [5]
+//      V   +--+ \   |
+// +-->[3]<-----[8]  |
+// | /  |            |
+// +    +-->[9]<-----+
+//
 // CHECK-LABEL:--- XLA CFG Loops Before Canonicalize: $loopThatRequiresNodeAndLoopCloning
 //----The following loop will be cloned in the second loop.
 // CHECK:Loop at depth 1 containing: {{bb[0-9]+}}<header><latch><exiting>

test/TensorFlow/sese_loop_canonicalization.swift

@@ -5,19 +5,21 @@
 
 import TensorFlow
 
-//expected-warning @+1 {{value implicitly copied to the host}}
+// expected-warning @+1 {{value implicitly copied to the host}}
 public func testLoopMovementFromOutside(_ breakIndex: Int32) -> Tensor<Int32> {
   var i: Int32 = 1
   var sum = Tensor<Int32>(0)
   let maxCount: Int32 = 100
-  while i <= maxCount {
+  while i <= maxCount { // expected-note {{value used here}}
     sum += i
-    if (i == breakIndex) {
+    if i == breakIndex {
       var k: Int32 = 0
-      while (k < i) {
+      while k < i {
         sum += i
         k += 1
       }
+      // expected-warning @+1 {{value implicitly copied to the host}}
+      sum += 1
       break // expected-note {{value used here}}
     }
     i += 1
@@ -27,10 +29,12 @@ public func testLoopMovementFromOutside(_ breakIndex: Int32) -> Tensor<Int32> {
 
 // CHECK-LABEL: --- XLA CFG Loops Before Canonicalize: {{.*}}testLoopMovementFromOutside{{.*}}
 // CHECK: Loop at depth 1 containing: {{.*}}
+//--- The following loop is the inner while loop, but is unnested in the CFG due to the `break`
 // CHECK: Loop at depth 1 containing: {{.*}}
 
 // CHECK-LABEL: --- XLA CFG Loops After Canonicalize: {{.*}}testLoopMovementFromOutside{{.*}}
 // CHECK: Loop at depth 1 containing: {{.*}}
+//--- The inner while loop gets nested into the outer while loop by SESE canonicalzation.
 // CHECK:     Loop at depth 2 containing: {{.*}}
 
 // CHECK: --- XLA CFG Canonicalize:  {{.*}}testLoopMovementFromOutside{{.*}}