---

yaml
---
r: 294867
b: refs/heads/tensorflow
c: 5ac83e2
h: refs/heads/master
i:
  294865: a99532e
  294863: 4d69ce8

Author: rxwei

[refs]

@@ -816,7 +816,7 @@ refs/tags/swift-DEVELOPMENT-SNAPSHOT-2018-04-25-a: 22f738a831d43aff2b9c9773bcb65
 refs/tags/swift-DEVELOPMENT-SNAPSHOT-2018-05-08-a: 7d98cc16689baba5c8a3b90a9329bdcc1a12b4e9
 refs/heads/cherr42: a566ad54b073c2c56ac0a705d0a5bed9743135a5
 "refs/heads/codable_test_comment_fix": fc8f6824f7f347e1e8db55bff62db385c5728b5a
-refs/heads/tensorflow: 1c1305ca1059a45e4f5a90bedbdb008ea0a10ab2
+refs/heads/tensorflow: 5ac83e25b48a79ed6a093a7e32a98a52fc138836
 refs/tags/swift-4.1-DEVELOPMENT-SNAPSHOT-2018-05-11-a: 8126fd7a652e2f70ad6d76505239e34fb2ef3e1a
 refs/tags/swift-4.1-DEVELOPMENT-SNAPSHOT-2018-05-12-a: b3fd3dd84df6717f2e2e9df58c6d7e99fed57086
 refs/tags/swift-4.1-DEVELOPMENT-SNAPSHOT-2018-05-13-a: 71135119579039dc321c5f65d870050fe36efda2

branches/tensorflow/lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -4131,6 +4131,10 @@ class AdjointEmitter final : public SILInstructionVisitor<AdjointEmitter> {
         : nullptr;
     Lowering::GenericContextScope genericContextScope(
         getContext().getTypeConverter(), adjGenSig);
+    auto origExitIt = original.findReturnBB();
+    assert(origExitIt != original.end() &&
+           "Functions without returns must have been diagnosed");
+    auto *origExit = &*origExitIt;
 
     // Get dominated active values in original blocks.
     // Adjoint values of dominated active values are passed as adjoint block
@@ -4175,12 +4179,22 @@ class AdjointEmitter final : public SILInstructionVisitor<AdjointEmitter> {
     for (auto *origBB : postOrderInfo->getPostOrder()) {
       auto *adjointBB = adjoint.createBasicBlock();
       adjointBBMap.insert({origBB, adjointBB});
-      // If adjoint block is the adjoint entry, continue.
-      if (adjointBB->isEntry())
-        continue;
-      // Otherwise, add a pullback struct argument to the adjoint block.
       auto pbStructLoweredType =
           remapType(getPullbackInfo().getPullbackStructLoweredType(origBB));
+      // If the BB is the original exit, then the adjoint block that we just
+      // created must be the adjoint function's entry. For the adjoint entry,
+      // create entry arguments and continue to the next block.
+      if (origBB == origExit) {
+        assert(adjointBB->isEntry());
+        createEntryArguments(&getAdjoint());
+        auto *lastArg = adjointBB->getArguments().back();
+        assert(lastArg->getType() == pbStructLoweredType);
+        adjointPullbackStructArguments[origBB] = lastArg;
+        continue;
+      }
+      
+      // Otherwise, we create a phi argument for the corresponding pullback
+      // struct, and handle dominated active values/buffers.
       auto *pbStructArg = adjointBB->createPhiArgument(
           pbStructLoweredType, ValueOwnershipKind::Guaranteed);
       adjointPullbackStructArguments[origBB] = pbStructArg;
@@ -4222,15 +4236,11 @@ class AdjointEmitter final : public SILInstructionVisitor<AdjointEmitter> {
       }
     }
 
-    auto *origEntry = original.getEntryBlock();
-    auto *origExit = &*original.findReturnBB();
     auto *adjointEntry = adjoint.getEntryBlock();
-    createEntryArguments(&adjoint);
     // The adjoint function has type (seed, exit_pbs) -> ([arg0], ..., [argn]).
     auto adjParamArgs = adjoint.getArgumentsWithoutIndirectResults();
     assert(adjParamArgs.size() == 2);
     seed = adjParamArgs[0];
-    adjointPullbackStructArguments[origExit] = adjParamArgs[1];
 
     // Assign adjoint for original result.
     SmallVector<SILValue, 8> origFormalResults;
@@ -4434,6 +4444,7 @@ class AdjointEmitter final : public SILInstructionVisitor<AdjointEmitter> {
     // Place the builder at the adjoint exit, i.e. the adjoint block
     // corresponding to the original entry. Return the adjoints wrt parameters
     // in the adjoint exit.
+    auto *origEntry = getOriginal().getEntryBlock();
     builder.setInsertionPoint(getAdjointBlock(origEntry));
 
     // This vector will contain all the materialized return elements.