[AutoDiff] Add cloned curry thunks to generated function list. (#27720)

Add cloned curry thunks to generated function list so that they will be
deleted upon clean up if any diagnostics are emitted.

Resolves [TF-918](https://bugs.swift.org/browse/TF-918): crash due to invalid cloned curry thunk that was not deleted
during clean up.

Author: dan-zheng

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -878,10 +878,9 @@ class ADContext {
   /// Saved for deletion during cleanup.
   SmallVector<SILFunction *, 32> generatedFunctions;
 
-  /// List of derivative function references, generated via
-  /// `emitDerivativeFunctionReference`.
+  /// List of references to generated functions.
   /// Saved for deletion during cleanup.
-  SmallVector<SILValue, 32> generatedDerivativeFunctionReferences;
+  SmallVector<SILValue, 32> generatedFunctionReferences;
 
   /// The AdditiveArithmetic protocol in the standard library.
   ProtocolDecl *additiveArithmeticProtocol =
@@ -933,8 +932,8 @@ class ADContext {
     return generatedFunctions;
   }
 
-  SmallVector<SILValue, 32> &getGeneratedDerivativeFunctionReferences() {
-    return generatedDerivativeFunctionReferences;
+  SmallVector<SILValue, 32> &getGeneratedFunctionReferences() {
+    return generatedFunctionReferences;
   }
 
   ProtocolDecl *getAdditiveArithmeticProtocol() const {
@@ -969,11 +968,11 @@ class ADContext {
       original->removeDifferentiableAttr(attr);
     }
     // Delete all references to generated functions.
-    for (auto derivativeFn : generatedDerivativeFunctionReferences) {
-      if (auto *fnRef =
-              peerThroughFunctionConversions<FunctionRefInst>(derivativeFn)) {
-        fnRef->replaceAllUsesWithUndef();
-        fnRef->eraseFromParent();
+    for (auto fnRef : generatedFunctionReferences) {
+      if (auto *fnRefInst =
+              peerThroughFunctionConversions<FunctionRefInst>(fnRef)) {
+        fnRefInst->replaceAllUsesWithUndef();
+        fnRefInst->eraseFromParent();
       }
     }
     // Delete all generated functions.
@@ -1226,6 +1225,10 @@ ADContext::emitNondifferentiabilityError(SILValue value,
     getADDebugStream() << "With invoker:\n" << invoker << '\n';
   });
   auto valueLoc = value.getLoc().getSourceLoc();
+  // If instruction does not have a valid location, use the function location
+  // as a fallback. Improves diagnostics in some cases.
+  if (valueLoc.isInvalid())
+    valueLoc = value->getFunction()->getLocation().getSourceLoc();
   return emitNondifferentiabilityError(valueLoc, invoker, diag,
                                        std::forward<U>(args)...);
 }
@@ -8566,13 +8569,15 @@ SILValue ADContext::promoteToDifferentiableFunction(
           thunkBuilder.createReturn(loc, dfi);
           retInst->eraseFromParent();
 
+          getGeneratedFunctions().push_back(newThunk);
           getDifferentiableFunctionInsts().push_back(dfi);
           if (processDifferentiableFunctionInst(dfi))
             return nullptr;
         }
 
         // Apply the new curry thunk.
         auto *newThunkRef = builder.createFunctionRef(loc, newThunk);
+        getGeneratedFunctionReferences().push_back(newThunkRef);
         SmallVector<SILValue, 8> newArgs;
         SmallVector<SILValue, 8> newArgsToDestroy;
         SmallVector<AllocStackInst *, 1> newBuffersToDealloc;
@@ -8608,7 +8613,7 @@ SILValue ADContext::promoteToDifferentiableFunction(
       return nullptr;
 
     auto derivativeFn = derivativeFnAndIndices->first;
-    getGeneratedDerivativeFunctionReferences().push_back(derivativeFn);
+    getGeneratedFunctionReferences().push_back(derivativeFn);
 
     // If desired indices are a subset of actual indices, create a "subset
     // indices thunk" and destroy the emitted derivative function reference.

test/AutoDiff/differentiation_transform_diagnostics.swift

@@ -332,6 +332,11 @@ struct TF_675 : Differentiable {
 // expected-error @+1 {{function is not differentiable}}
 let _: @differentiable (Float) -> Float = TF_675().method
 
+// TF-918: test parameter subset thunk + partially-applied original function.
+// expected-error @+2 {{function is not differentiable}}
+// expected-note @+1 {{cannot convert a direct method reference to a '@differentiable' function; use an explicit closure instead}}
+_ = gradient(at: Float(1), Float(2), in: (+) as @differentiable (Float, @nondiff Float) -> Float)
+
 //===----------------------------------------------------------------------===//
 // Conversion to `@differentiable(linear)` (not yet supported)
 //===----------------------------------------------------------------------===//