[AutoDiff] Support curry thunks differentiation in inlinable funcs

Inside inlinable functions, we expect functions to either be explicitly
marked as differentiable or have a public explicit derivative defined.
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 swiftlang#54819
Fixes swiftlang#75776

Author: kovdan01

include/swift/AST/Expr.h

@@ -4387,6 +4387,8 @@ class ClosureExpr : public AbstractClosureExpr {
 class AutoClosureExpr : public AbstractClosureExpr {
   BraceStmt *Body;
 
+  ApplyExpr *getUnwrappedCurryThunkImpl() const;
+
 public:
   enum class Kind : uint8_t {
     // An autoclosure with type () -> Result. Formed from type checking an
@@ -4448,6 +4450,10 @@ class AutoClosureExpr : public AbstractClosureExpr {
   /// - otherwise, returns nullptr for convenience.
   Expr *getUnwrappedCurryThunkExpr() const;
 
+  /// Same as getUnwrappedCurryThunkExpr, but get the called ValueDecl instead
+  /// of the expr.
+  ValueDecl *getUnwrappedCurryThunkCalledValue() const;
+
   // Implement isa/cast/dyncast/etc.
   static bool classof(const Expr *E) {
     return E->getKind() == ExprKind::AutoClosure;

lib/AST/Expr.cpp

@@ -2133,7 +2133,7 @@ Expr *AutoClosureExpr::getSingleExpressionBody() const {
   return cast<ReturnStmt>(Body->getLastElement().get<Stmt *>())->getResult();
 }
 
-Expr *AutoClosureExpr::getUnwrappedCurryThunkExpr() const {
+ApplyExpr *AutoClosureExpr::getUnwrappedCurryThunkImpl() const {
   auto maybeUnwrapOpenExistential = [](Expr *expr) {
     if (auto *openExistential = dyn_cast<OpenExistentialExpr>(expr)) {
       expr = openExistential->getSubExpr()->getSemanticsProvidingExpr();
@@ -2177,7 +2177,7 @@ Expr *AutoClosureExpr::getUnwrappedCurryThunkExpr() const {
     body = maybeUnwrapConversions(body);
 
     if (auto *outerCall = dyn_cast<ApplyExpr>(body)) {
-      return outerCall->getFn();
+      return outerCall;
     }
 
     assert(false && "Malformed curry thunk?");
@@ -2198,7 +2198,7 @@ Expr *AutoClosureExpr::getUnwrappedCurryThunkExpr() const {
       if (auto *outerCall = dyn_cast<ApplyExpr>(innerBody)) {
         auto outerFn = maybeUnwrapConversions(outerCall->getFn());
         if (auto *innerCall = dyn_cast<ApplyExpr>(outerFn)) {
-          return innerCall->getFn();
+          return innerCall;
         }
       }
     }
@@ -2211,6 +2211,20 @@ Expr *AutoClosureExpr::getUnwrappedCurryThunkExpr() const {
   return nullptr;
 }
 
+Expr *AutoClosureExpr::getUnwrappedCurryThunkExpr() const {
+  ApplyExpr *ae = getUnwrappedCurryThunkImpl();
+  if (ae == nullptr)
+    return nullptr;
+  return ae->getFn();
+}
+
+ValueDecl *AutoClosureExpr::getUnwrappedCurryThunkCalledValue() const {
+  ApplyExpr *ae = getUnwrappedCurryThunkImpl();
+  if (ae == nullptr)
+    return nullptr;
+  return ae->getCalledValue(/*skipFunctionConversions=*/true);
+}
+
 FORWARD_SOURCE_LOCS_TO(UnresolvedPatternExpr, subPattern)
 
 TypeExpr::TypeExpr(TypeRepr *Repr)

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -51,6 +51,8 @@
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/Support/CommandLine.h"
 
+#include <unordered_map>
+
 using namespace swift;
 using namespace swift::autodiff;
 using llvm::DenseMap;
@@ -84,6 +86,9 @@ class DifferentiationTransformer {
   /// Context necessary for performing the transformations.
   ADContext context;
 
+  /// Cache used in getUnwrappedCurryThunkFunction.
+  std::unordered_map<AbstractFunctionDecl *, SILFunction *> afdToSILFn;
+
   /// Promotes the given `differentiable_function` instruction to a valid
   /// `@differentiable` function-typed value.
   SILValue promoteToDifferentiableFunction(DifferentiableFunctionInst *inst,
@@ -96,6 +101,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 +477,40 @@ 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 *abstractCE = originalFn->getDeclRef().getAbstractClosureExpr();
+  if (abstractCE == nullptr)
+    return nullptr;
+  auto *autoCE = dyn_cast<AutoClosureExpr>(abstractCE);
+  if (autoCE == nullptr)
+    return nullptr;
+
+  auto *afd =
+      cast<AbstractFunctionDecl>(autoCE->getUnwrappedCurryThunkCalledValue());
+
+  auto silFnIt = afdToSILFn.find(afd);
+  if (silFnIt == afdToSILFn.end()) {
+    assert(afdToSILFn.empty());
 
+    SILModule *module = getTransform().getModule();
+    for (SILFunction &currentFunc : module->getFunctions())
+      if (auto *currentAFD = currentFunc.getDeclRef().getAbstractFunctionDecl())
+        afdToSILFn.emplace(currentAFD, &currentFunc);
+
+    silFnIt = afdToSILFn.find(afd);
+    assert(silFnIt != afdToSILFn.end());
+  }
+
+  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 +653,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 +1174,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

@@ -772,32 +772,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/inlinable_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 issue 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/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