[AutoDiff] Fix crash during generic curry thunk cloning. (#26404)

Create simple `BasicTypeSubstCloner` inheriting from `TypeSubstCloner`.
Use `BasicTypeSubstCloner` to clone generic curry thunks, remapping types.

Resolves TF-688.

Author: dan-zheng

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -78,30 +78,6 @@ template <typename T> static inline void debugDump(T &v) {
                           << v << "\n==== END DEBUG DUMP ====\n");
 }
 
-/// Creates arguments in the entry block based on the function type.
-static void createEntryArguments(SILFunction *f) {
-  auto *entry = f->getEntryBlock();
-  auto conv = f->getConventions();
-  auto &ctx = f->getASTContext();
-  auto moduleDecl = f->getModule().getSwiftModule();
-  assert((entry->getNumArguments() == 0 || conv.getNumSILArguments() == 0) &&
-         "Entry already has arguments?!");
-  auto createFunctionArgument = [&](SILType type) {
-    // Create a dummy parameter declaration.
-    // Necessary to prevent crash during argument explosion optimization.
-    auto loc = f->getLocation().getSourceLoc();
-    auto *decl = new (ctx)
-        ParamDecl(VarDecl::Specifier::Default, loc, loc, Identifier(), loc,
-                  Identifier(), moduleDecl);
-    decl->setType(type.getASTType());
-    entry->createFunctionArgument(type, decl);
-  };
-  for (auto indResTy : conv.getIndirectSILResultTypes())
-    createFunctionArgument(f->mapTypeIntoContext(indResTy).getAddressType());
-  for (auto paramTy : conv.getParameterSILTypes())
-    createFunctionArgument(f->mapTypeIntoContext(paramTy));
-}
-
 static bool isWithoutDerivative(SILValue v) {
   if (auto *fnRef = dyn_cast<FunctionRefInst>(v))
     return fnRef->getReferencedFunctionOrNull()->hasSemanticsAttr(
@@ -3003,9 +2979,8 @@ class VJPEmitter final
         original->isBare(), IsNotTransparent, original->isSerialized(),
         original->isDynamicallyReplaceable());
     pullback->setOwnershipEliminated();
-    pullback->setDebugScope(new (module)
-                                SILDebugScope(original->getLocation(),
-                                              pullback));
+    pullback->setDebugScope(
+        new (module) SILDebugScope(original->getLocation(), pullback));
     return pullback;
   }
 
@@ -3677,7 +3652,7 @@ class JVPEmitter final
   void visitAutoDiffFunctionInst(AutoDiffFunctionInst *adfi) {
     // Clone `autodiff_function` from original to JVP, then add the cloned
     // instruction to the `autodiff_function` worklist.
-    SILClonerWithScopes::visitAutoDiffFunctionInst(adfi);
+    TypeSubstCloner::visitAutoDiffFunctionInst(adfi);
     auto *newADFI = cast<AutoDiffFunctionInst>(getOpValue(adfi));
     context.getAutoDiffFunctionInsts().push_back(newADFI);
   }
@@ -6588,10 +6563,12 @@ SILValue ADContext::promoteToDifferentiableFunction(
         // returned function value with an `autodiff_function` instruction,
         // and process the `autodiff_function` instruction.
         if (newThunk->empty()) {
+          if (auto newThunkGenSig = thunkType->getGenericSignature())
+            newThunk->setGenericEnvironment(
+                newThunkGenSig->createGenericEnvironment());
           newThunk->setOwnershipEliminated();
-          SILFunctionCloner cloner(newThunk);
-          cloner.cloneFunction(thunk);
-
+          BasicTypeSubstCloner cloner(thunk, newThunk);
+          cloner.run();
           auto *retInst =
               cast<ReturnInst>(newThunk->findReturnBB()->getTerminator());
           SILBuilder thunkBuilder(retInst);

lib/SILOptimizer/Mandatory/Differentiation.h

@@ -24,6 +24,7 @@
 #ifndef SWIFT_SILOPTIMIZER_MANDATORY_DIFFERENTIATION_H
 #define SWIFT_SILOPTIMIZER_MANDATORY_DIFFERENTIATION_H
 
+#include "swift/SIL/TypeSubstCloner.h"
 #include "swift/SILOptimizer/Analysis/DominanceAnalysis.h"
 #include "swift/SILOptimizer/Utils/Local.h"
 
@@ -88,6 +89,58 @@ class PostOrderPostDominanceOrder {
   }
 };
 
+/// Creates arguments in the entry block based on the function type.
+void createEntryArguments(SILFunction *f) {
+  auto *entry = f->getEntryBlock();
+  auto conv = f->getConventions();
+  auto &ctx = f->getASTContext();
+  auto moduleDecl = f->getModule().getSwiftModule();
+  assert((entry->getNumArguments() == 0 || conv.getNumSILArguments() == 0) &&
+         "Entry already has arguments?!");
+  auto createFunctionArgument = [&](SILType type) {
+    // Create a dummy parameter declaration.
+    // Necessary to prevent crash during argument explosion optimization.
+    auto loc = f->getLocation().getSourceLoc();
+    auto *decl = new (ctx)
+        ParamDecl(VarDecl::Specifier::Default, loc, loc, Identifier(), loc,
+                  Identifier(), moduleDecl);
+    decl->setType(type.getASTType());
+    entry->createFunctionArgument(type, decl);
+  };
+  for (auto indResTy : conv.getIndirectSILResultTypes())
+    createFunctionArgument(f->mapTypeIntoContext(indResTy).getAddressType());
+  for (auto paramTy : conv.getParameterSILTypes())
+    createFunctionArgument(f->mapTypeIntoContext(paramTy));
+}
+
+/// Cloner that remaps types using the target function's generic environment.
+class BasicTypeSubstCloner final
+    : public TypeSubstCloner<BasicTypeSubstCloner, SILOptFunctionBuilder> {
+
+  static SubstitutionMap getSubstitutionMap(SILFunction *target) {
+    if (auto *targetGenEnv = target->getGenericEnvironment())
+      return targetGenEnv->getForwardingSubstitutionMap();
+    return SubstitutionMap();
+  }
+
+public:
+  explicit BasicTypeSubstCloner(SILFunction *original, SILFunction *target)
+      : TypeSubstCloner(*target, *original, getSubstitutionMap(target)) {}
+
+  void postProcess(SILInstruction *orig, SILInstruction *cloned) {
+    SILClonerWithScopes::postProcess(orig, cloned);
+  }
+
+  void run() {
+    auto &target = Builder.getFunction();
+    auto *entry = target.createBasicBlock();
+    createEntryArguments(&target);
+    SmallVector<SILValue, 8> entryArguments(target.getArguments().begin(),
+                                            target.getArguments().end());
+    cloneFunctionBody(&Original, entry, entryArguments);
+  }
+};
+
 } // end namespace swift
 
 #endif // SWIFT_SILOPTIMIZER_MANDATORY_DIFFERENTIATION_H

test/AutoDiff/generics.swift

@@ -229,4 +229,22 @@ extension TF_682_Proto where Self : Differentiable,
   }
 }
 
+// TF-688: Test generic curry thunk cloning.
+public struct TF_688_Struct<Scalar> {
+  var x: Scalar
+}
+extension TF_688_Struct: Differentiable where Scalar: Differentiable {
+  @differentiable
+  public static func id(x: Self) -> Self {
+    return x
+  }
+}
+@differentiable(wrt: x)
+public func TF_688<Scalar: Differentiable>(
+  _ x: TF_688_Struct<Scalar>,
+  reduction: @differentiable (TF_688_Struct<Scalar>) -> TF_688_Struct<Scalar> = TF_688_Struct.id
+) -> TF_688_Struct<Scalar> {
+  reduction(x)
+}
+
 // TODO: add more tests.