[AutoDiff] Revamp differentiation transform. (#24845)

- Directly generate primal code in VJP functions.
  - Rename `PrimalGenCloner` to `VJPEmitter`.
  - This unblocks control-flow support.
    Primal data structures can be generated based on activity analysis.

- Refactor differentiation transform to use iterative `autodiff_function`
  instruction worklist instead of a fixed `DifferentiationTask` worklist.
  - `VJPEmitter::visitApply` now creates `autodiff_function` and
    `autodiff_function_extract` instructions instead of directly emitting
    associated function references.
  - An iterative loop canonicalizes `autodiff_function` instructions,
    promoting them to `@differentiable` function-typed values.

- Add `autodiff_function_extract` folding optimization.
  - Fold `autodiff_function_extract` users of `autodiff_function` instructions,
    directly replacing them with operands of the `autodiff_function`
    instruction.
  - If the `autodiff_function` instruction has only `autodiff_function_extract`
    users, delete the instruction itself after folding.

- Remove unnecessary auxiliary data structures.
  - `DifferentiationTask`: replace with map from `[differentiable]` attributes
    to `DifferentiationInvoker`.
  - `PrimalGen`, `AdjointGen`: replace with `autodiff_function` iterative loop.

Author: dan-zheng

include/swift/AST/AutoDiff.h

@@ -578,12 +578,12 @@ bool getBuiltinAutoDiffApplyConfig(StringRef operationName,
                                    unsigned &arity, unsigned &order,
                                    bool &rethrows);
 
-/// Computes the correct linkage for associated functions given the linkage of
+/// Computes the correct linkage for an associated function given the linkage of
 /// the original function. If the original linkage is not external and
 /// `isAssocFnExported` is true, use the original function's linkage. Otherwise,
 /// return hidden linkage.
-SILLinkage getAutoDiffFunctionLinkage(SILLinkage originalLinkage,
-                                      bool isAssocFnExported);
+SILLinkage getAutoDiffAssociatedFunctionLinkage(SILLinkage originalLinkage,
+                                                bool isAssocFnExported);
 
 } // end namespace autodiff
 

include/swift/AST/Types.h

@@ -4153,7 +4153,7 @@ class SILFunctionType final : public TypeBase, public llvm::FoldingSetNode,
       AutoDiffIndexSubset *parameterIndices, unsigned resultIndex,
       unsigned differentiationOrder, AutoDiffAssociatedFunctionKind kind,
       SILModule &module, LookupConformanceFn lookupConformance,
-      GenericSignature *whereClauseGenericSignature = nullptr);
+      CanGenericSignature whereClauseGenericSignature = nullptr);
 
   /// Returns a bit vector that specifices which parameters you can
   /// differentiate with respect to for this differentiable function type. (e.g.

lib/AST/AutoDiff.cpp

@@ -89,8 +89,8 @@ bool autodiff::getBuiltinAutoDiffApplyConfig(
   return operationName.empty();
 }
 
-SILLinkage autodiff::getAutoDiffFunctionLinkage(SILLinkage originalLinkage,
-                                                bool isAssocFnExported) {
+SILLinkage autodiff::getAutoDiffAssociatedFunctionLinkage(
+    SILLinkage originalLinkage, bool isAssocFnExported) {
   // If the original is defined externally, then the AD pass is just generating
   // associated functions for use in the current module and therefore these
   // associated functions should not be visible outside the module.

lib/SIL/SILFunctionType.cpp

@@ -150,18 +150,18 @@ CanSILFunctionType SILFunctionType::getAutoDiffAssociatedFunctionType(
     AutoDiffIndexSubset *parameterIndices, unsigned resultIndex,
     unsigned differentiationOrder, AutoDiffAssociatedFunctionKind kind,
     SILModule &module, LookupConformanceFn lookupConformance,
-    GenericSignature *whereClauseGenSig) {
+    CanGenericSignature whereClauseGenSig) {
   // JVP: (T...) -> ((R...),
   //                 (T.TangentVector...) -> (R.TangentVector...))
   // VJP: (T...) -> ((R...),
   //                 (R.TangentVector...) -> (T.TangentVector...))
 
   auto &ctx = getASTContext();
   auto &typeConverter = module.Types;
-  Lowering::GenericContextScope
-      genericContextScope(module.Types, getGenericSignature());
   if (!whereClauseGenSig)
     whereClauseGenSig = getGenericSignature();
+  Lowering::GenericContextScope genericContextScope(
+      module.Types, whereClauseGenSig);
 
   // Given a type, returns its formal SIL parameter info.
   auto getTangentParameterInfoForOriginalResult = [&](

lib/SILGen/SILGenPoly.cpp

@@ -3530,7 +3530,7 @@ SILGenModule::getOrCreateAutoDiffAssociatedFunctionReorderingThunk(
 
   auto loc = assocFn->getLocation();
   SILGenFunctionBuilder fb(*this);
-  auto linkage = autodiff::getAutoDiffFunctionLinkage(
+  auto linkage = autodiff::getAutoDiffAssociatedFunctionLinkage(
       original->getLinkage(), /*isAssocFnExported*/ true);
   auto *thunk = fb.getOrCreateFunction(
       loc, name, linkage, targetType, IsBare, IsNotTransparent,

lib/SILOptimizer/IPO/CapturePropagation.cpp

@@ -461,13 +461,20 @@ bool CapturePropagation::optimizePartialApply(PartialApplyInst *PAI) {
   SILOptFunctionBuilder FuncBuilder(*this);
   if (auto *NewFunc = getSpecializedWithDeadParams(FuncBuilder,
           PAI, SubstF, PAI->getNumArguments(), GenericSpecialized)) {
-    rewritePartialApply(PAI, NewFunc);
-    if (GenericSpecialized.first) {
-      // Notify the pass manager about the new function.
-      addFunctionToPassManagerWorklist(GenericSpecialized.first,
-                                       GenericSpecialized.second);
+    // SWIFT_ENABLE_TENSORFLOW
+    // Add a previously unexercised check to prevent AD crash. Rewrite
+    // `partial_apply` only if the specialized function is `@convention(thin)`.
+    // Revert check when `VJPEmitter::visitApplyInst` no longer produces
+    // argumentless `partial_apply` instructions.
+    if (NewFunc->getRepresentation() == SILFunctionTypeRepresentation::Thin) {
+      rewritePartialApply(PAI, NewFunc);
+      if (GenericSpecialized.first) {
+        // Notify the pass manager about the new function.
+        addFunctionToPassManagerWorklist(GenericSpecialized.first,
+                                         GenericSpecialized.second);
+      }
+      return true;
     }
-    return true;
   }
 
   // Second possibility: Are all partially applied arguments constant?