[AutoDiff] Support curry thunks differentiation in fragile funcs (#77615)

Inside fragile functions, we expect function derivatives to be public, which could be achieved by either explicitly marking the functions as differentiable or having a public explicit derivative defined for them. This is obviously not
possible for single and double curry thunks which are a special case of `AutoClosureExpr`.

Instead of looking at the thunk itself, we unwrap it and look at the function being wrapped. While the thunk itself and its differentiability witness will not have public visibility, it's not an issue for the case where the function being wrapped (and its witness) have public visibility.

Fixes #54819
Fixes #75776

Author: kovdan01

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -47,6 +47,7 @@
 #include "swift/SILOptimizer/Utils/SILOptFunctionBuilder.h"
 #include "llvm/ADT/APSInt.h"
 #include "llvm/ADT/BreadthFirstIterator.h"
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/Support/CommandLine.h"
@@ -84,6 +85,9 @@ class DifferentiationTransformer {
   /// Context necessary for performing the transformations.
   ADContext context;
 
+  /// Cache used in getUnwrappedCurryThunkFunction.
+  llvm::DenseMap<AbstractFunctionDecl *, SILFunction *> afdToSILFn;
+
   /// Promotes the given `differentiable_function` instruction to a valid
   /// `@differentiable` function-typed value.
   SILValue promoteToDifferentiableFunction(DifferentiableFunctionInst *inst,
@@ -96,6 +100,25 @@ class DifferentiationTransformer {
                                    SILBuilder &builder, SILLocation loc,
                                    DifferentiationInvoker invoker);
 
+  /// Emits a reference to a derivative function of `original`, differentiated
+  /// with respect to a superset of `desiredIndices`. Returns the `SILValue` for
+  /// the derivative function and the actual indices that the derivative
+  /// function is with respect to.
+  ///
+  /// Returns `None` on failure, signifying that a diagnostic has been emitted
+  /// using `invoker`.
+  std::optional<std::pair<SILValue, AutoDiffConfig>>
+  emitDerivativeFunctionReference(
+      SILBuilder &builder, const AutoDiffConfig &desiredConfig,
+      AutoDiffDerivativeFunctionKind kind, SILValue original,
+      DifferentiationInvoker invoker,
+      SmallVectorImpl<AllocStackInst *> &newBuffersToDealloc);
+
+  /// If the given function corresponds to AutoClosureExpr with either
+  /// SingleCurryThunk or DoubleCurryThunk kind, get the SILFunction
+  /// corresponding to the function being wrapped in the thunk.
+  SILFunction *getUnwrappedCurryThunkFunction(SILFunction *originalFn);
+
 public:
   /// Construct an `DifferentiationTransformer` for the given module.
   explicit DifferentiationTransformer(SILModuleTransform &transform)
@@ -453,21 +476,63 @@ static SILValue reapplyFunctionConversion(
   llvm_unreachable("Unhandled function conversion instruction");
 }
 
-/// Emits a reference to a derivative function of `original`, differentiated
-/// with respect to a superset of `desiredIndices`. Returns the `SILValue` for
-/// the derivative function and the actual indices that the derivative function
-/// is with respect to.
-///
-/// Returns `None` on failure, signifying that a diagnostic has been emitted
-/// using `invoker`.
-static std::optional<std::pair<SILValue, AutoDiffConfig>>
-emitDerivativeFunctionReference(
-    DifferentiationTransformer &transformer, SILBuilder &builder,
-    const AutoDiffConfig &desiredConfig, AutoDiffDerivativeFunctionKind kind,
-    SILValue original, DifferentiationInvoker invoker,
-    SmallVectorImpl<AllocStackInst *> &newBuffersToDealloc) {
-  ADContext &context = transformer.getContext();
+SILFunction *DifferentiationTransformer::getUnwrappedCurryThunkFunction(
+    SILFunction *originalFn) {
+  auto *autoCE = dyn_cast_or_null<AutoClosureExpr>(
+      originalFn->getDeclRef().getAbstractClosureExpr());
+  if (autoCE == nullptr)
+    return nullptr;
+
+  auto *ae = dyn_cast_or_null<ApplyExpr>(autoCE->getUnwrappedCurryThunkExpr());
+  if (ae == nullptr)
+    return nullptr;
 
+  AbstractFunctionDecl *afd = cast<AbstractFunctionDecl>(ae->getCalledValue(
+      /*skipFunctionConversions=*/true));
+  auto silFnIt = afdToSILFn.find(afd);
+  if (silFnIt == afdToSILFn.end()) {
+    assert(afdToSILFn.empty() && "Expect all 'afdToSILFn' cache entries to be "
+                                 "filled at once on the first access attempt");
+
+    SILModule *module = getTransform().getModule();
+    for (SILFunction &currentFunc : module->getFunctions()) {
+      if (auto *currentAFD =
+              currentFunc.getDeclRef().getAbstractFunctionDecl()) {
+        // Update cache only with AFDs which might be potentially wrapped by a
+        // curry thunk. This includes member function references and references
+        // to functions having external property wrapper parameters (see
+        // ExprRewriter::buildDeclRef). If new use cases of curry thunks appear
+        // in future, the assertion after the loop will be a trigger for such
+        // cases being unhandled here.
+        //
+        // FIXME: References to functions having external property wrapper
+        // parameters are not handled since we can't now construct a test case
+        // for that due to the crash
+        // https://github.com/swiftlang/swift/issues/77613
+        if (currentAFD->hasCurriedSelf()) {
+          auto [_, wasEmplace] =
+              afdToSILFn.try_emplace(currentAFD, &currentFunc);
+          assert(wasEmplace && "Expect all 'afdToSILFn' cache entries to be "
+                               "filled at once on the first access attempt");
+        }
+      }
+    }
+
+    silFnIt = afdToSILFn.find(afd);
+    assert(silFnIt != afdToSILFn.end() &&
+           "Expect present curry thunk to SIL function mapping after "
+           "'afdToSILFn' cache fill");
+  }
+
+  return silFnIt->second;
+}
+
+std::optional<std::pair<SILValue, AutoDiffConfig>>
+DifferentiationTransformer::emitDerivativeFunctionReference(
+    SILBuilder &builder, const AutoDiffConfig &desiredConfig,
+    AutoDiffDerivativeFunctionKind kind, SILValue original,
+    DifferentiationInvoker invoker,
+    SmallVectorImpl<AllocStackInst *> &newBuffersToDealloc) {
   // If `original` is itself an `DifferentiableFunctionExtractInst` whose kind
   // matches the given kind and desired differentiation parameter indices,
   // simply extract the derivative function of its function operand, retain the
@@ -610,26 +675,36 @@ emitDerivativeFunctionReference(
           DifferentiabilityKind::Reverse, desiredParameterIndices,
           desiredResultIndices, derivativeConstrainedGenSig, /*jvp*/ nullptr,
           /*vjp*/ nullptr, /*isSerialized*/ false);
-      if (transformer.canonicalizeDifferentiabilityWitness(
-              minimalWitness, invoker, IsNotSerialized))
+      if (canonicalizeDifferentiabilityWitness(minimalWitness, invoker,
+                                               IsNotSerialized))
         return std::nullopt;
     }
     assert(minimalWitness);
-    if (original->getFunction()->isSerialized() &&
-        !hasPublicVisibility(minimalWitness->getLinkage())) {
-      enum { Inlinable = 0, DefaultArgument = 1 };
-      unsigned fragileKind = Inlinable;
-      // FIXME: This is not a very robust way of determining if the function is
-      // a default argument. Also, we have not exhaustively listed all the kinds
-      // of fragility.
-      if (original->getFunction()->getLinkage() == SILLinkage::PublicNonABI)
-        fragileKind = DefaultArgument;
-      context.emitNondifferentiabilityError(
-          original, invoker, diag::autodiff_private_derivative_from_fragile,
-          fragileKind,
-          isa_and_nonnull<AbstractClosureExpr>(
-              originalFRI->getLoc().getAsASTNode<Expr>()));
-      return std::nullopt;
+    if (original->getFunction()->isSerialized()) {
+      // When dealing with curry thunk, look at the function being wrapped
+      // inside implicit closure. If it has public visibility, the corresponding
+      // differentiability witness also has public visibility. It should be OK
+      // for implicit wrapper closure and its witness to have private linkage.
+      SILFunction *unwrappedFn = getUnwrappedCurryThunkFunction(originalFn);
+      bool isWitnessPublic =
+          unwrappedFn == nullptr
+              ? hasPublicVisibility(minimalWitness->getLinkage())
+              : hasPublicVisibility(unwrappedFn->getLinkage());
+      if (!isWitnessPublic) {
+        enum { Inlinable = 0, DefaultArgument = 1 };
+        unsigned fragileKind = Inlinable;
+        // FIXME: This is not a very robust way of determining if the function
+        // is a default argument. Also, we have not exhaustively listed all the
+        // kinds of fragility.
+        if (original->getFunction()->getLinkage() == SILLinkage::PublicNonABI)
+          fragileKind = DefaultArgument;
+        context.emitNondifferentiabilityError(
+            original, invoker, diag::autodiff_private_derivative_from_fragile,
+            fragileKind,
+            isa_and_nonnull<AbstractClosureExpr>(
+                originalFRI->getLoc().getAsASTNode<Expr>()));
+        return std::nullopt;
+      }
     }
     // TODO(TF-482): Move generic requirement checking logic to
     // `getExactDifferentiabilityWitness` and
@@ -1121,8 +1196,8 @@ SILValue DifferentiationTransformer::promoteToDifferentiableFunction(
   for (auto derivativeFnKind : {AutoDiffDerivativeFunctionKind::JVP,
                                 AutoDiffDerivativeFunctionKind::VJP}) {
     auto derivativeFnAndIndices = emitDerivativeFunctionReference(
-        *this, builder, desiredConfig, derivativeFnKind, origFnOperand,
-        invoker, newBuffersToDealloc);
+        builder, desiredConfig, derivativeFnKind, origFnOperand, invoker,
+        newBuffersToDealloc);
     // Show an error at the operator, highlight the argument, and show a note
     // at the definition site of the argument.
     if (!derivativeFnAndIndices)

test/AutoDiff/SILOptimizer/differentiation_diagnostics.swift

@@ -771,32 +771,6 @@ public func fragileDifferentiable(_ x: Float) -> Float {
   implicitlyDifferentiableFromFragile(x)
 }
 
-
-// FIXME: Differentiable curry thunk RequirementMachine error (rdar://87429620, https://github.com/apple/swift/issues/54819).
-#if false
-// TF-1208: Test curry thunk differentiation regression.
-public struct Struct_54819<Scalar> {
-  var x: Scalar
-}
-extension Struct_54819: Differentiable where Scalar: Differentiable {
-  @differentiable(reverse)
-  public static func id(x: Self) -> Self {
-    return x
-  }
-}
-@differentiable(reverse, wrt: x)
-public func f_54819<Scalar: Differentiable>(
-  _ x: Struct_54819<Scalar>,
-  // NOTE(TF-1208): This diagnostic is unexpected because `Struct_54819.id` is marked `@differentiable`.
-  // xpected-error @+3 2 {{function is not differentiable}}
-  // xpected-note @+2 {{differentiated functions in '@inlinable' functions must be marked '@differentiable' or have a public '@derivative'; this is not possible with a closure, make a top-level function instead}}
-  // xpected-note @+1 {{opaque non-'@differentiable' function is not differentiable}}
-  reduction: @differentiable(reverse) (Struct_54819<Scalar>) -> Struct_54819<Scalar> = Struct_54819.id
-) -> Struct_54819<Scalar> {
-  reduction(x)
-}
-#endif
-
 //===----------------------------------------------------------------------===//
 // Coroutines (SIL function yields, `begin_apply`) (not yet supported)
 //===----------------------------------------------------------------------===//

test/AutoDiff/SILOptimizer/fragile_curry_thunk.swift

@@ -0,0 +1,73 @@
+// RUN: %target-swift-frontend -emit-sil -verify -primary-file %s -o /dev/null
+
+import _Differentiation
+
+/// Minimal reproducer for both single and double curry thunk
+
+@inlinable
+func caller<Thing: Differentiable & FloatingPoint>(
+  of f: @differentiable(reverse) (_: Thing) -> Thing
+) -> Int  where Thing.TangentVector == Thing {
+   return 42
+}
+
+public struct Struct<Thing: Differentiable & FloatingPoint>: Differentiable where Thing.TangentVector == Thing {
+    @inlinable
+    static func foo_single() -> Int {
+        return caller(of: callee_single) // No error expected
+    }
+
+    @inlinable
+    @differentiable(reverse)
+    static func callee_single(input: Thing) -> Thing {
+        return input
+    }
+
+    @inlinable
+    func foo_double() -> Int {
+        return caller(of: callee_double) // No error expected
+    }
+
+    @inlinable
+    @differentiable(reverse)
+    func callee_double(input: Thing) -> Thing {
+        return input
+    }
+}
+
+/// Reproducer from https://github.com/swiftlang/swift/issues/75776
+
+public struct Solution2<Thing: Differentiable & FloatingPoint>: Differentiable where Thing.TangentVector == Thing {
+    @inlinable
+    public static func optimization() -> Thing {
+        var initial = Thing.zero
+        let (_, delta) = valueWithGradient(at: initial, of: simulationWithLoss) // No error expected
+        initial.move(by: delta)
+        return initial
+    }
+
+    @inlinable
+    @differentiable(reverse)
+    static func simulationWithLoss(input: Thing) -> Thing {
+        return input // implementation
+    }
+}
+
+/// Reproducer from https://github.com/swiftlang/swift/issues/54819
+
+public struct TF_688_Struct<Scalar> {
+  var x: Scalar
+}
+extension TF_688_Struct: Differentiable where Scalar: Differentiable {
+  @differentiable(reverse)
+  public static func id(x: Self) -> Self {
+    return x
+  }
+}
+@differentiable(reverse, wrt: x)
+public func TF_688<Scalar: Differentiable>(
+  _ x: TF_688_Struct<Scalar>,
+  reduction: @differentiable(reverse) (TF_688_Struct<Scalar>) -> TF_688_Struct<Scalar> = TF_688_Struct.id // No error expected
+) -> TF_688_Struct<Scalar> {
+  reduction(x)
+}

test/AutoDiff/SILOptimizer/generics.swift

@@ -250,27 +250,6 @@ extension TF_682_Proto where Self : Differentiable,
   }
 }
 
-// NOTE(TF-1208): Differentiation regression due to changes in curry thunk generation.
-/*
-// 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(reverse)
-  public static func id(x: Self) -> Self {
-    return x
-  }
-}
-@differentiable(reverse, wrt: x)
-public func TF_688<Scalar: Differentiable>(
-  _ x: TF_688_Struct<Scalar>,
-  reduction: @differentiable(reverse) (TF_688_Struct<Scalar>) -> TF_688_Struct<Scalar> = TF_688_Struct.id
-) -> TF_688_Struct<Scalar> {
-  reduction(x)
-}
-*/
-
 // TF-697: Test generic requirements of generated derivative function.
 protocol TF_697_Module: Differentiable {
     associatedtype Input

test/AutoDiff/compiler_crashers/rdar87429620-differentiable-curry-thunk-reqmachine.swift renamed to test/AutoDiff/compiler_crashers_fixed/rdar87429620-differentiable-curry-thunk-reqmachine.swift

@@ -1,5 +1,4 @@
 // RUN: %target-swift-frontend -emit-sil -verify %s
-// XFAIL: *
 
 // rdar://87429620
 // https://github.com/apple/swift/issues/54819