[AutoDiff] Fix SILGen JVP/VJP thunking bug. (#26448)

SILGen JVP/VJP thunks for methods now always perform self reordering.
Add leak checking test.

Previously, self reordering was not performed if actual JVP/VJP type
matched expected JVP/VJP type, which is incorrect.

Gardening included:
- Change `SILFunction *SILGenFunction::getOrCreateAutoDiffLinearMapThunk` to
  `ManagedValue SILGenFunction::getThunkedAutoDiffLinearMap`. The new helper
  is more ergonomic and directly returns thunked values.
- Use ad-hoc `_vtable_entry_thunk` suffix for JVP/VJP vtable entry thunks.
- Add todo comments for TF-685: principled AD thunk mangling.

Resolves TF-698.

Author: dan-zheng

lib/SIL/SILFunctionType.cpp

@@ -157,7 +157,6 @@ static CanGenericSignature getAutoDiffAssociatedFunctionGenericSignature(
     CanGenericSignature assocFnGenSig,
     ArrayRef<SILParameterInfo> originalParameters,
     AutoDiffIndexSubset *parameterIndices, SILModule &module) {
-  // If associated function has no
   if (!assocFnGenSig)
     return nullptr;
   auto &ctx = module.getASTContext();

lib/SILGen/SILGen.cpp

@@ -810,13 +810,26 @@ void SILGenModule::postEmitFunction(SILDeclRef constant,
           indices.parameters, indices.source, /*differentiationOrder*/ 1,
           AutoDiffAssociatedFunctionKind::VJP, M, lookUpConformance);
 
+      // Self reordering is necessary if wrt at least two parameters, including
+      // self.
+      auto shouldReorderSelf = [&]() {
+        if (!F->hasSelfParam())
+          return false;
+        auto selfParamIndex = origSilFnType->getNumParameters() - 1;
+        if (!indices.isWrtParameter(selfParamIndex))
+          return false;
+        return indices.parameters->getNumIndices() > 1;
+      };
+      bool reorderSelf = shouldReorderSelf();
+
       // Thunk JVP method, if it is defined.
       if (auto *jvpDecl = diffAttr->getJVPFunction()) {
         SILFunction *jvpThunk;
         auto *jvpFn = getFunction(SILDeclRef(jvpDecl), NotForDefinition);
-        if (jvpFn->getLoweredFunctionType() != expectedJVPType) {
+        if (reorderSelf || jvpFn->getLoweredFunctionType() != expectedJVPType) {
           jvpThunk = getOrCreateAutoDiffAssociatedFunctionThunk(
-              F, indices, jvpFn, AutoDiffAssociatedFunctionKind::JVP);
+              F, indices, jvpFn, AutoDiffAssociatedFunctionKind::JVP,
+              reorderSelf);
         } else {
           auto *id = AutoDiffAssociatedFunctionIdentifier::get(
               AutoDiffAssociatedFunctionKind::JVP, /*differentiationOrder*/ 1,
@@ -831,9 +844,10 @@ void SILGenModule::postEmitFunction(SILDeclRef constant,
       if (auto *vjpDecl = diffAttr->getVJPFunction()) {
         SILFunction *vjpThunk;
         auto *vjpFn = getFunction(SILDeclRef(vjpDecl), NotForDefinition);
-        if (vjpFn->getLoweredFunctionType() != expectedVJPType) {
+        if (reorderSelf || vjpFn->getLoweredFunctionType() != expectedVJPType) {
           vjpThunk = getOrCreateAutoDiffAssociatedFunctionThunk(
-              F, indices, vjpFn, AutoDiffAssociatedFunctionKind::VJP);
+              F, indices, vjpFn, AutoDiffAssociatedFunctionKind::VJP,
+              reorderSelf);
         } else {
           auto *id = AutoDiffAssociatedFunctionIdentifier::get(
               AutoDiffAssociatedFunctionKind::VJP, /*differentiationOrder*/ 1,

lib/SILGen/SILGen.h

@@ -187,7 +187,8 @@ class LLVM_LIBRARY_VISIBILITY SILGenModule : public ASTVisitor<SILGenModule> {
   /// - The last result in the returned pullback.
   SILFunction *getOrCreateAutoDiffAssociatedFunctionThunk(
       SILFunction *original, SILAutoDiffIndices &indices,
-      SILFunction *assocFn, AutoDiffAssociatedFunctionKind assocFnKind);
+      SILFunction *assocFn, AutoDiffAssociatedFunctionKind assocFnKind,
+      bool reorderSelf);
 
   /// Determine whether the given class has any instance variables that
   /// need to be destroyed.

lib/SILGen/SILGenFunction.h

@@ -1780,14 +1780,15 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
   // Differentiation thunks
   //===--------------------------------------------------------------------===//
 
-  /// Get or create a thunk for reabstracting differentials/pullbacks returned
-  /// by user-defined JVP/VJP functions.
+  /// Get or create a thunk for reabstracting and self-reordering
+  /// differentials/pullbacks returned by user-defined JVP/VJP functions, and
+  /// apply it to the given differential/pullback.
   ///
   /// If `reorderSelf` is true, reorder self so that it appears as:
   /// - The last parameter, for differentials.
   /// - The last result, for pullbacks.
-  SILFunction *getOrCreateAutoDiffLinearMapThunk(
-      AutoDiffAssociatedFunctionKind assocFnKind,
+  ManagedValue getThunkedAutoDiffLinearMap(
+      ManagedValue linearMap, AutoDiffAssociatedFunctionKind assocFnKind,
       CanSILFunctionType fromType, CanSILFunctionType toType,
       bool reorderSelf);
 

lib/SILGen/SILGenPoly.cpp

@@ -3362,7 +3362,6 @@ static CanSILFunctionType buildWithoutActuallyEscapingThunkType(
 /// tuple. Otherwise, add this value directly to `result`.
 static void extractAllElements(SILValue val, SILBuilder &builder,
                                SmallVectorImpl<SILValue> &result) {
-  // auto &fn = builder.getFunction();
   if (auto tupleType = val->getType().getAs<TupleType>())
     for (auto i : range(tupleType->getNumElements()))
       result.push_back(builder.createTupleExtract(val.getLoc(), val, i));
@@ -3383,10 +3382,11 @@ static SILValue joinElements(ArrayRef<SILValue> elements, SILBuilder &builder,
 
 // SWIFT_ENABLE_TENSORFLOW
 /// Adapted from `SILGenModule::getOrCreateReabstractionThunk`.
-SILFunction *
-SILGenFunction::getOrCreateAutoDiffLinearMapThunk(
-    AutoDiffAssociatedFunctionKind assocFnKind, CanSILFunctionType fromType,
-    CanSILFunctionType toType, bool reorderSelf) {
+ManagedValue
+SILGenFunction::getThunkedAutoDiffLinearMap(
+    ManagedValue linearMap, AutoDiffAssociatedFunctionKind assocFnKind,
+    CanSILFunctionType fromType, CanSILFunctionType toType,
+    bool reorderSelf) {
   // Compute the thunk type.
   SubstitutionMap interfaceSubs;
   GenericEnvironment *genericEnv = nullptr;
@@ -3407,7 +3407,7 @@ SILGenFunction::getOrCreateAutoDiffLinearMapThunk(
   std::string name = mangler.mangleReabstractionThunkHelper(
       thunkType, fromInterfaceType, toInterfaceType,
       Type(), getModule().getSwiftModule());
-  // TODO: Use principled mangling.
+  // TODO(TF-685): Use principled thunk mangling.
   if (reorderSelf) {
     switch (assocFnKind) {
     case AutoDiffAssociatedFunctionKind::JVP:
@@ -3426,8 +3426,20 @@ SILGenFunction::getOrCreateAutoDiffLinearMapThunk(
   auto *thunk = fb.getOrCreateSharedFunction(
       loc, name, thunkDeclType, IsBare, IsTransparent, IsSerialized,
       ProfileCounter(), IsReabstractionThunk, IsNotDynamic);
+
+  // Partially-apply the thunk to `linearMap` and return the thunked value.
+  auto getThunkedResult = [&]() {
+    auto thunkedFn = createPartialApplyOfThunk(
+        *this, loc, thunk, interfaceSubs, dynamicSelfType, toType, linearMap);
+    if (!toType->isNoEscape())
+      return thunkedFn;
+    // Handle escaping to noescape conversion.
+    return B.createConvertEscapeToNoEscape(
+        loc, thunkedFn, SILType::getPrimitiveObjectType(toType));
+  };
+
   if (!thunk->empty())
-    return thunk;
+    return getThunkedResult();
   thunk->setGenericEnvironment(genericEnv);
   thunk->setOwnershipEliminated();
 
@@ -3558,9 +3570,9 @@ SILGenFunction::getOrCreateAutoDiffLinearMapThunk(
     arguments.push_back(load);
   }
 
-  auto linearMap = thunk->getArgumentsWithoutIndirectResults().back();
+  auto *linearMapArg = thunk->getArgumentsWithoutIndirectResults().back();
   auto *apply = thunkSGF.B.createApply(
-      loc, linearMap, SubstitutionMap(), arguments, /*isNonThrowing*/ false);
+      loc, linearMapArg, SubstitutionMap(), arguments, /*isNonThrowing*/ false);
 
   // Get return elements.
   SmallVector<SILValue, 4> results;
@@ -3628,7 +3640,7 @@ SILGenFunction::getOrCreateAutoDiffLinearMapThunk(
 
   // Create return.
   thunkSGF.B.createReturn(loc, retVal);
-  return thunk;
+  return getThunkedResult();
 }
 
 /// Forward function arguments, converting ownership.
@@ -3685,9 +3697,12 @@ static void forwardFunctionArgumentsConvertingOwnership(
 SILFunction *
 SILGenModule::getOrCreateAutoDiffAssociatedFunctionThunk(
     SILFunction *original, SILAutoDiffIndices &indices,
-    SILFunction *assocFn, AutoDiffAssociatedFunctionKind assocFnKind) {
+    SILFunction *assocFn, AutoDiffAssociatedFunctionKind assocFnKind,
+    bool reorderSelf) {
   auto assocFnType = assocFn->getLoweredFunctionType();
 
+  // TODO(TF-685): Use principled thunk mangling.
+  // Do not simply reuse reabstraction thunk mangling.
   Mangle::ASTMangler mangler;
   auto name = getASTContext().getIdentifier(
       mangler.mangleAutoDiffAssociatedFunctionHelper(
@@ -3744,8 +3759,8 @@ SILGenModule::getOrCreateAutoDiffAssociatedFunctionThunk(
 
   SmallVector<SILValue, 8> directResults;
   extractAllElements(apply, thunkSGF.B, directResults);
-  auto linearMap = directResults.back();
-  auto linearMapFnType = linearMap->getType().castTo<SILFunctionType>();
+  auto linearMap = ManagedValue::forBorrowedObjectRValue(directResults.back());
+  auto linearMapFnType = linearMap.getType().castTo<SILFunctionType>();
   auto targetLinearMapFnType = thunk->mapTypeIntoContext(
       origAssocFnType->getResults().back().getSILStorageType())
           .castTo<SILFunctionType>();
@@ -3767,33 +3782,17 @@ SILGenModule::getOrCreateAutoDiffAssociatedFunctionThunk(
     thunkSGF.B.createReturn(loc, retValue);
   };
 
-  // If linear map types are unchanged, return the `apply` instruction.
-  if (linearMapFnType == targetLinearMapFnType) {
+  // If self ordering is not necessary and linear map types are unchanged,
+  // return the `apply` instruction.
+  if (!reorderSelf && linearMapFnType == targetLinearMapFnType) {
     createReturn(apply);
     return thunk;
   }
 
-  // Generate linear map thunk for reabstraction/self reordering.
-  auto shouldReorderSelf = [&]() {
-    if (!original->hasSelfParam())
-      return false;
-    auto selfParamIndex =
-        original->getArgumentsWithoutIndirectResults().size() - 1;
-    if (!indices.isWrtParameter(selfParamIndex))
-      return false;
-    return indices.parameters->getNumIndices() > 1;
-  };
-  bool reorderSelf = shouldReorderSelf();
-  auto *linearMapThunk = thunkSGF.getOrCreateAutoDiffLinearMapThunk(
-      assocFnKind, linearMapFnType, targetLinearMapFnType, reorderSelf);
-  auto linearMapThunkValue =
-      thunkSGF.B.createFunctionRefFor(loc, linearMapThunk);
-  SubstitutionMap linearMapSubs;
-  if (linearMapThunk->getLoweredFunctionType()->isPolymorphic())
-    linearMapSubs = thunk->getForwardingSubstitutionMap();
-  linearMap = thunkSGF.B.createPartialApply(
-      loc, linearMapThunkValue, linearMapSubs, {linearMap},
-      linearMapFnType->getCalleeConvention());
+  // Otherwise, apply reabstraction/self reordering thunk to linear map.
+  linearMap = thunkSGF.getThunkedAutoDiffLinearMap(
+      linearMap, assocFnKind, linearMapFnType, targetLinearMapFnType,
+      reorderSelf);
 
   // Return original results and thunked differential/pullback.
   if (directResults.size() > 1) {
@@ -3802,10 +3801,10 @@ SILGenModule::getOrCreateAutoDiffAssociatedFunctionThunk(
     auto originalDirectResult =
         joinElements(originalDirectResults, thunkSGF.B, apply.getLoc());
     auto thunkResult = joinElements(
-        {originalDirectResult, linearMap}, thunkSGF.B, loc);
+        {originalDirectResult, linearMap.getValue()}, thunkSGF.B, loc);
     createReturn(thunkResult);
   } else {
-    createReturn(linearMap);
+    createReturn(linearMap.getValue());
   }
   return thunk;
 }

lib/SILGen/SILGenThunk.cpp

@@ -84,28 +84,27 @@ SILGenModule::getOrCreateAutoDiffThunk(SILDeclRef assocFnDeclRef,
   auto linkage = autodiff::getAutoDiffAssociatedFunctionLinkage(
       originalLinkage, /*isAssocFnExported*/ true);
   auto name = assocFnDeclRef.mangle();
-  auto *F = builder.getOrCreateFunction(
+  auto *thunk = builder.getOrCreateFunction(
       assocFnDecl, name, linkage, assocFnTy, IsBare, IsTransparent,
       assocFnDeclRef.isSerialized(), IsNotDynamic, ProfileCounter(), IsThunk);
-
-  if (F->empty()) {
-    if (auto genSig = assocFnTy->getGenericSignature())
-      F->setGenericEnvironment(genSig->createGenericEnvironment());
-    SILGenFunction SGF(*this, *F, SwiftModule);
-    SmallVector<ManagedValue, 4> params;
-    auto loc = assocFnDeclRef.getAsRegularLocation();
-    SGF.collectThunkParams(loc, params);
-    auto assocFnRef = SGF.B.createFunctionRef(loc, assocFn);
-    auto autoDiffAssocFnSILTy = SILType::getPrimitiveObjectType(assocFnTy);
-    SmallVector<SILValue, 4> args(F->getArguments().begin(),
-                                  F->getArguments().end());
-    auto apply = SGF.emitApplyWithRethrow(
-        loc, assocFnRef, autoDiffAssocFnSILTy,
-        SGF.getForwardingSubstitutionMap(), args);
-    SGF.B.createReturn(loc, apply);
-  }
-
-  return F;
+  if (!thunk->empty())
+    return thunk;
+
+  if (auto genSig = assocFnTy->getGenericSignature())
+    thunk->setGenericEnvironment(genSig->createGenericEnvironment());
+  SILGenFunction SGF(*this, *thunk, SwiftModule);
+  SmallVector<ManagedValue, 4> params;
+  auto loc = assocFnDeclRef.getAsRegularLocation();
+  SGF.collectThunkParams(loc, params);
+  auto assocFnRef = SGF.B.createFunctionRef(loc, assocFn);
+  auto autoDiffAssocFnSILTy = SILType::getPrimitiveObjectType(assocFnTy);
+  SmallVector<SILValue, 4> args(thunk->getArguments().begin(),
+                                thunk->getArguments().end());
+  auto apply = SGF.emitApplyWithRethrow(
+      loc, assocFnRef, autoDiffAssocFnSILTy,
+      SGF.getForwardingSubstitutionMap(), args);
+  SGF.B.createReturn(loc, apply);
+  return thunk;
 }
 
 // SWIFT_ENABLE_TENSORFLOW
@@ -120,37 +119,38 @@ SILFunction *SILGenModule::getOrCreateAutoDiffClassMethodThunk(
   auto originalLinkage = originalFn.getLinkage(ForDefinition);
   auto linkage = autodiff::getAutoDiffAssociatedFunctionLinkage(
       originalLinkage, /*isAssocFnExported*/ true);
-  auto name = assocFnDeclRef.mangle() + "_thunk";
-  auto *F = builder.getOrCreateFunction(
+  // TODO(TF-685): Use principled thunk mangling.
+  // Do not simply reuse reabstraction thunk mangling.
+  auto name = assocFnDeclRef.mangle() + "_vtable_entry_thunk";
+  auto *thunk = builder.getOrCreateFunction(
       assocFnDecl, name, linkage, constantTy, IsBare, IsTransparent,
       assocFnDeclRef.isSerialized(), IsNotDynamic, ProfileCounter(), IsThunk);
-
-  if (F->empty()) {
-    if (auto genSig = constantTy->getGenericSignature())
-      F->setGenericEnvironment(genSig->createGenericEnvironment());
-    SILGenFunction SGF(*this, *F, SwiftModule);
-    SmallVector<ManagedValue, 4> params;
-    auto loc = assocFnDeclRef.getAsRegularLocation();
-    SGF.collectThunkParams(loc, params);
-    auto originalFnRef = SGF.emitGlobalFunctionRef(loc, originalFn);
-    auto loweredIndices = autoDiffFuncId->getParameterIndices()->getLowered(
-        SGF.getASTContext(),
-        assocFnDecl->getInterfaceType()->castTo<AnyFunctionType>());
-    auto autoDiffFn = SGF.B.createAutoDiffFunction(
-        loc, loweredIndices, /*differentiationOrder*/ 1, originalFnRef);
-    auto autoDiffAssocFn = SGF.B.createAutoDiffFunctionExtract(
-        loc, AutoDiffFunctionExtractInst::Extractee(autoDiffFuncId->getKind()),
-        /*differentiationOrder*/ 1, autoDiffFn);
-    auto autoDiffAssocFnSILTy = SILType::getPrimitiveObjectType(constantTy);
-    SmallVector<SILValue, 4> args(F->getArguments().begin(),
-                                  F->getArguments().end());
-    auto apply = SGF.emitApplyWithRethrow(
-        loc, autoDiffAssocFn, autoDiffAssocFnSILTy,
-        SGF.getForwardingSubstitutionMap(), args);
-    SGF.B.createReturn(loc, apply);
-  }
-
-  return F;
+  if (!thunk->empty())
+    return thunk;
+
+  if (auto genSig = constantTy->getGenericSignature())
+    thunk->setGenericEnvironment(genSig->createGenericEnvironment());
+  SILGenFunction SGF(*this, *thunk, SwiftModule);
+  SmallVector<ManagedValue, 4> params;
+  auto loc = assocFnDeclRef.getAsRegularLocation();
+  SGF.collectThunkParams(loc, params);
+  auto originalFnRef = SGF.emitGlobalFunctionRef(loc, originalFn);
+  auto loweredIndices = autoDiffFuncId->getParameterIndices()->getLowered(
+      SGF.getASTContext(),
+      assocFnDecl->getInterfaceType()->castTo<AnyFunctionType>());
+  auto autoDiffFn = SGF.B.createAutoDiffFunction(
+      loc, loweredIndices, /*differentiationOrder*/ 1, originalFnRef);
+  auto autoDiffAssocFn = SGF.B.createAutoDiffFunctionExtract(
+      loc, AutoDiffFunctionExtractInst::Extractee(autoDiffFuncId->getKind()),
+      /*differentiationOrder*/ 1, autoDiffFn);
+  auto autoDiffAssocFnSILTy = SILType::getPrimitiveObjectType(constantTy);
+  SmallVector<SILValue, 4> args(thunk->getArguments().begin(),
+                                thunk->getArguments().end());
+  auto apply = SGF.emitApplyWithRethrow(
+      loc, autoDiffAssocFn, autoDiffAssocFnSILTy,
+      SGF.getForwardingSubstitutionMap(), args);
+  SGF.B.createReturn(loc, apply);
+  return thunk;
 }
 
 ManagedValue