[AutoDiff] Unify adjoint entry block generation with the general logic. (#25255)

In control flow AD support, the logic that handles entry adjoint block generation should not be different than the general logic that generates intermediate blocks and processes active values. The only special thing about the entry adjoint block is that it need no phi args because its corresponding original block (the original exit) has no successors.

This PR unifies the block generation logic, moving us a bit closer to the memory leak fix.

Author: rxwei

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.