Fix remaining test failures.

Remaining test failures were related to `differentiable_function_extract`
cloning during generic specialization.

Add explanatory comment to
`TypeSubstCloner::visitDifferentiableFunctionExtractInst`:

    +----------------+        remap       +-------------------------+
    | orig.  fn type |  -------(A)------> | remapped orig.  fn type |
    +----------------+                    +-------------------------+
            |                                                |
       (B, SILGen)   getAutoDiffDerivativeFunctionType   (D, here)
            V                                                V
    +----------------+        remap       +-------------------------+
    | deriv. fn type |  -------(C)------> | remapped deriv. fn type |
    +----------------+                    +-------------------------+

(AD) does not always commute with (BC):
- (AD) is the result of remapping, then computing the derivative type.
  This is the default cloning behavior, but may break invariants in the
  initial SIL generated by SILGen.
- (BC) is the result of computing the derivative type (SILGen), then
  remapping. This is the expected type, preserving invariants from
  earlier transforms.

If (AD) is not equal to (BC), (BC) must be used as the explicit type.

Done with @marcrasi.

Author: dan-zheng

include/swift/SIL/SILCloner.h

@@ -2859,10 +2859,13 @@ template<typename ImplClass>
 void SILCloner<ImplClass>::
 visitDifferentiableFunctionExtractInst(DifferentiableFunctionExtractInst *Inst) {
   getBuilder().setCurrentDebugScope(getOpScope(Inst->getDebugScope()));
+  Optional<SILType> explicitExtracteeType = None;
+  if (Inst->hasExplicitExtracteeType())
+    explicitExtracteeType = Inst->getType();
   recordClonedInstruction(
       Inst, getBuilder().createDifferentiableFunctionExtract(
                 getOpLocation(Inst->getLoc()), Inst->getExtractee(),
-                getOpValue(Inst->getFunctionOperand())));
+                getOpValue(Inst->getFunctionOperand()), explicitExtracteeType));
 }
 
 template<typename ImplClass>

include/swift/SIL/TypeSubstCloner.h

@@ -315,6 +315,65 @@ class TypeSubstCloner : public SILClonerWithScopes<ImplClass> {
     super::visitDestroyValueInst(Destroy);
   }
 
+  // SWIFT_ENABLE_TENSORFLOW
+  void visitDifferentiableFunctionExtractInst(
+      DifferentiableFunctionExtractInst *dfei) {
+    // If the extractee is the original function, do regular cloning.
+    if (dfei->getExtractee() ==
+        NormalDifferentiableFunctionTypeComponent::Original) {
+      super::visitDifferentiableFunctionExtractInst(dfei);
+      return;
+    }
+    // If the extractee is a derivative function, check whether the *remapped
+    // derivative function type* (BC) is equal to the *derivative remapped
+    // function type* (AD).
+    //
+    // +----------------+        remap       +-------------------------+
+    // | orig.  fn type |  -------(A)------> | remapped orig.  fn type |
+    // +----------------+                    +-------------------------+
+    //         |                                                |
+    //    (B, SILGen)   getAutoDiffDerivativeFunctionType   (D, here)
+    //         V                                                V
+    // +----------------+        remap       +-------------------------+
+    // | deriv. fn type |  -------(C)------> | remapped deriv. fn type |
+    // +----------------+                    +-------------------------+
+    //
+    // (AD) does not always commute with (BC):
+    // - (AD) is the result of remapping, then computing the derivative type.
+    //   This is the default cloning behavior, but may break invariants in the
+    //   initial SIL generated by SILGen.
+    // - (BC) is the result of computing the derivative type (SILGen), then
+    //   remapping. This is the expected type, preserving invariants from
+    //   earlier transforms.
+    //
+    // If (AD) is not equal to (BC), use (BC) as the explicit type.
+    SILType remappedOrigType = getOpType(dfei->getFunctionOperand()->getType());
+    auto remappedOrigFnType = remappedOrigType.castTo<SILFunctionType>();
+    auto derivativeRemappedFnType =
+        remappedOrigFnType
+            ->getAutoDiffDerivativeFunctionType(
+                remappedOrigFnType->getDifferentiabilityParameterIndices(),
+                /*resultIndex*/ 0, dfei->getDerivativeFunctionKind(),
+                getBuilder().getModule().Types,
+                LookUpConformanceInModule(SwiftMod))
+            ->getWithoutDifferentiability();
+    SILType remappedDerivativeFnType = getOpType(dfei->getType());
+    // If remapped derivative type and derivative remapped type are equal, do
+    // regular cloning.
+    if (SILType::getPrimitiveObjectType(derivativeRemappedFnType) ==
+        remappedDerivativeFnType) {
+      super::visitDifferentiableFunctionExtractInst(dfei);
+      return;
+    }
+    // Otherwise, explicitly use the remapped derivative type.
+    recordClonedInstruction(
+        dfei,
+        getBuilder().createDifferentiableFunctionExtract(
+            getOpLocation(dfei->getLoc()), dfei->getExtractee(),
+            getOpValue(dfei->getFunctionOperand()), remappedDerivativeFnType));
+  }
+  // SWIFT_ENABLE_TENSORFLOW END
+
   /// One abstract function in the debug info can only have one set of variables
   /// and types. This function determines whether applying the substitutions in
   /// \p SubsMap on the generic signature \p Sig will change the generic type

lib/SIL/SILInstructions.cpp

@@ -709,7 +709,6 @@ getExtracteeType(
   auto resultFnTy = originalFnTy->getAutoDiffDerivativeFunctionType(
       fnTy->getDifferentiabilityParameterIndices(), /*resultIndex*/ 0,
       *kindOpt, module.Types, LookUpConformanceInModule(module.getSwiftModule()));
-  llvm::dbgs() << "getExtracteeType for " << function << ":\n" << SILType::getPrimitiveObjectType(resultFnTy) << "\n";
   return SILType::getPrimitiveObjectType(resultFnTy);
 }
 
@@ -725,8 +724,14 @@ DifferentiableFunctionExtractInst::DifferentiableFunctionExtractInst(
       HasExplicitExtracteeType(extracteeType.hasValue()) {
 #ifndef NDEBUG
   if (extracteeType.hasValue()) {
-    assert(module.getStage() == SILStage::Lowered &&
-           "Explicit type is valid only in lowered SIL");
+    // Note: explicit extractee type is used to avoid inconsistent typing in:
+    // - Canonical SIL, due to generic specialization.
+    // - Lowered SIL, due to LoadableByAddress.
+    // See `TypeSubstCloner::visitDifferentiableFunctionExtractInst` for an
+    // explanation of how explicit extractee type is used.
+    assert((module.getStage() == SILStage::Canonical ||
+            module.getStage() == SILStage::Lowered) &&
+           "Explicit type is valid only in canonical or lowered SIL");
   }
 #endif
 }

lib/SILOptimizer/Utils/Differentiation/JVPEmitter.cpp

@@ -1431,8 +1431,12 @@ void JVPEmitter::visitReturnInst(ReturnInst *ri) {
       loc, differentialRef, jvpSubstMap, {diffStructVal},
       ParameterConvention::Direct_Guaranteed);
 
-  auto differentialType = jvp->getLoweredFunctionType()->getResults().back().getSILStorageInterfaceType();
-  differentialType = differentialType.substGenericArgs(getModule(), jvpSubstMap, TypeExpansionContext::minimal());
+  auto differentialType = jvp->getLoweredFunctionType()
+                              ->getResults()
+                              .back()
+                              .getSILStorageInterfaceType();
+  differentialType = differentialType.substGenericArgs(
+      getModule(), jvpSubstMap, TypeExpansionContext::minimal());
   differentialType = differentialType.subst(getModule(), jvpSubstMap);
   auto differentialFnType = differentialType.castTo<SILFunctionType>();
 
@@ -1452,8 +1456,7 @@ void JVPEmitter::visitReturnInst(ReturnInst *ri) {
     // applies, the return type may be ABI-incomaptible with the type of the
     // partially applied differential. In these cases, produce an undef and rely
     // on other code to emit a diagnostic.
-    differentialValue = SILUndef::get(differentialType,
-                                      *differentialPartialApply->getFunction());
+    differentialValue = SILUndef::get(differentialType, *jvp);
   }
 
   // Return a tuple of the original result and differential.

lib/SILOptimizer/Utils/Differentiation/VJPEmitter.cpp

@@ -363,7 +363,7 @@ void VJPEmitter::visitReturnInst(ReturnInst *ri) {
     // applies, the return type may be ABI-incomaptible with the type of the
     // partially applied pullback. In these cases, produce an undef and rely on
     // other code to emit a diagnostic.
-    pullbackValue = SILUndef::get(pullbackType, *pullbackPartialApply->getFunction());
+    pullbackValue = SILUndef::get(pullbackType, *vjp);
   }
 
   // Return a tuple of the original result and pullback.