[AutoDiff] Fix a 'partial_apply' leak caused by subset parameters thunks.

In `ADContext::promoteToDifferentiableFunction` when we emit a subset parameters thunk, the closure produced by `emitAssociatedFunctionReference` is not being used or released. This patch fixes that by releasing the closure.

In a future design, we should determine the actual parameter indices and whether thunking is needed before emitting an unused closure. This shall be addressed in a future patch.

A value leak checking test was added for protocols because this bug was originally discovered in eaplatanios/swift-ale#1 by using the newly split `Layer` protocol. However, value leak checking tests do not test for closure (`partial_apply`) leaks. The proper setup for catching closure leaks is moving AD before ownership stripping (swiftlang#26157).

Author: rxwei

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -6607,7 +6607,7 @@ SILValue ADContext::promoteToDifferentiableFunction(
     getGeneratedAssociatedFunctionReferences().push_back(assocFn);
 
     // If desired indices are a subset of actual indices, create a "subset
-    // indices thunk".
+    // indices thunk" and destroy the emitted associated function reference.
     // - For JVPs: the thunked JVP returns a differential taking fewer
     //   parameters (using `.zero` for the dropped parameters).
     // - For VJPs: the thunked VJP returns a pullback that drops the unused
@@ -6621,6 +6621,9 @@ SILValue ADContext::promoteToDifferentiableFunction(
         getASTContext(), actualIndices.parameters->getCapacity());
     if (actualIndices.source != desiredIndices.source ||
         !actualIndices.parameters->equals(extendedDesiredIndices)) {
+      // Destroy the already emitted associated function reference because it
+      // is no longer used.
+      builder.emitReleaseValueAndFold(loc, assocFn);
       // Check if underlying original function reference has been partially
       // applied with arguments. If so, produce an error: parameter subset
       // thunks do not yet support this case because partially applied arguments

test/AutoDiff/leakchecking.swift

@@ -74,10 +74,14 @@ extension DummyLayer {
   func defaultImpl(_ input: Input) -> Output {
     return requirement(input)
   }
-  func vjpDefaultImpl(_ input: Input) -> (Output, (Self.Output.TangentVector) -> (Self.TangentVector, Self.Input.TangentVector)) {
-    return Swift.valueWithPullback(at: self, input) { (m, i) in m.requirement(i) }
+  func vjpDefaultImpl(_ input: Input)
+  -> (Output,
+    (Self.Output.TangentVector)
+    -> (Self.TangentVector, Self.Input.TangentVector)) {
+    return Swift.valueWithPullback(at: self, input) { $0.requirement($1) }
   }
 }
+
 LeakCheckingTests.testWithLeakChecking("TestProtocolDefaultDerivative") {
   struct Foo : DummyLayer {
     // NOTE: Make sure not to override `defaultImpl`.
@@ -98,6 +102,40 @@ LeakCheckingTests.testWithLeakChecking("TestProtocolDefaultDerivative") {
   }
 }
 
+protocol Module : Differentiable {
+  associatedtype Input
+  associatedtype Output : Differentiable
+  @differentiable(wrt: self)
+  func callAsFunction(_ input: Input) -> Output
+}
+protocol Layer : Module where Input : Differentiable {
+  @differentiable(wrt: (self, input))
+  func callAsFunction(_ input: Input) -> Output
+}
+
+LeakCheckingTests.testWithLeakChecking("ProtocolRequirements") {
+  struct Dense: Layer {
+    var w = Tracked<Float>(1)
+    @differentiable
+    func callAsFunction(_ input: Tracked<Float>) -> Tracked<Float> {
+      input * w
+    }
+  }
+  struct Model: Module {
+    var dense1 = Dense()
+    var dense2 = Dense()
+    @differentiable
+    func callAsFunction(_ input: Tracked<Int>) -> Tracked<Float> {
+      dense2(dense1(Tracked(Float(input.value))))
+    }
+  }
+  let x = Tracked<Int>(1)
+  let model = Model()
+  _ = model.valueWithGradient { model in
+    model(x)
+  }
+}
+
 LeakCheckingTests.testWithLeakChecking("LetStructs") {
   func structConstructionWithOwnedParams(_ x: Tracked<Float>) -> Tracked<Float> {
     let z = Tracked(x)

test/AutoDiff/subset_parameters_thunk.swift

@@ -0,0 +1,32 @@
+// RUN: %target-swift-frontend -emit-sil %s | %FileCheck %s
+
+@differentiable(where T: Differentiable)
+func foo<T: Numeric>(_ x: T, _ y: T) -> T { x * y }
+
+@differentiating(foo)
+func foo_vjp<T: Numeric & Differentiable>(_ x: T, _ y: T) -> (value: T, pullback: (T.TangentVector) -> (T.TangentVector, T.TangentVector)) {
+  (foo(x, y), { _ in (.zero, .zero) })
+}
+
+let x = Float(1)
+@differentiable
+func differentiate_foo_wrt_0(_ x: Float) -> Float {
+  foo(x, 1)
+}
+
+// CHECK-LABEL: @{{.*}}differentiate_foo_wrt_0{{.*}}__vjp
+// CHECK: bb0
+// CHECK:   [[FOO_ORIG:%.*]] = function_ref @{{.*}}foo{{.*}} : $@convention(thin) <τ_0_0 where τ_0_0 : Numeric> (@in_guaranteed τ_0_0, @in_guaranteed τ_0_0) -> @out τ_0_0
+// CHECK:   [[FOO_FLOAT:%.*]] = partial_apply [callee_guaranteed] [[FOO_ORIG]]<Float>() : $@convention(thin) <τ_0_0 where τ_0_0 : Numeric> (@in_guaranteed τ_0_0, @in_guaranteed τ_0_0) -> @out τ_0_0
+// CHECK:   [[FOO_JVP:%.*]] = function_ref @AD__{{.*}}foo{{.*}}__jvp_src_0_wrt_0_1 : $@convention(thin) <τ_0_0 where τ_0_0 : Differentiable, τ_0_0 : Numeric> (@in_guaranteed τ_0_0, @in_guaranteed τ_0_0) -> (@out τ_0_0, @owned @callee_guaranteed (@in_guaranteed τ_0_0.TangentVector, @in_guaranteed τ_0_0.TangentVector) -> @out τ_0_0.TangentVector)
+// CHECK:   [[FOO_JVP_FLOAT:%.*]] = partial_apply [callee_guaranteed] [[FOO_JVP]]<Float>() : $@convention(thin) <τ_0_0 where τ_0_0 : Differentiable, τ_0_0 : Numeric> (@in_guaranteed τ_0_0, @in_guaranteed τ_0_0) -> (@out τ_0_0, @owned @callee_guaranteed (@in_guaranteed τ_0_0.TangentVector, @in_guaranteed τ_0_0.TangentVector) -> @out τ_0_0.TangentVector)
+// CHECK:   release_value [[FOO_JVP_FLOAT]]
+// CHECK:   [[FOO_JVP_SUBSET_THUNK_THIN:%.*]] = function_ref @AD__orig_{{.*}}foo{{.*}}_src_0_wrt_0_jvp_subset_parameters_thunk : $@convention(thin) (@in_guaranteed Float, @in_guaranteed Float) -> (@out Float, @owned @callee_guaranteed (@in_guaranteed Float) -> @out Float)
+// CHECK:   [[FOO_JVP_SUBSET_THUNK:%.*]] = thin_to_thick_function [[FOO_JVP_SUBSET_THUNK_THIN]] : $@convention(thin) (@in_guaranteed Float, @in_guaranteed Float) -> (@out Float, @owned @callee_guaranteed (@in_guaranteed Float) -> @out Float) to $@callee_guaranteed (@in_guaranteed Float, @in_guaranteed Float) -> (@out Float, @owned @callee_guaranteed (@in_guaranteed Float) -> @out Float)
+// CHECK:   [[FOO_VJP:%.*]] = function_ref @{{.*}}foo_vjp{{.*}} : $@convention(thin) <τ_0_0 where τ_0_0 : Differentiable, τ_0_0 : Numeric> (@in_guaranteed τ_0_0, @in_guaranteed τ_0_0) -> (@out τ_0_0, @owned @callee_guaranteed (@in_guaranteed τ_0_0.TangentVector) -> (@out τ_0_0.TangentVector, @out τ_0_0.TangentVector))
+// CHECK:   [[FOO_VJP_FLOAT:%.*]] = partial_apply [callee_guaranteed] [[FOO_VJP]]<Float>() : $@convention(thin) <τ_0_0 where τ_0_0 : Differentiable, τ_0_0 : Numeric> (@in_guaranteed τ_0_0, @in_guaranteed τ_0_0) -> (@out τ_0_0, @owned @callee_guaranteed (@in_guaranteed τ_0_0.TangentVector) -> (@out τ_0_0.TangentVector, @out τ_0_0.TangentVector))
+// CHECK:   release_value [[FOO_VJP_FLOAT]]
+// CHECK:   [[FOO_VJP_SUBSET_THUNK_THIN:%.*]] = function_ref @AD__orig_{{.*}}foo{{.*}}_src_0_wrt_0_vjp_subset_parameters_thunk : $@convention(thin) (@in_guaranteed Float, @in_guaranteed Float) -> (@out Float, @owned @callee_guaranteed (@in_guaranteed Float) -> @out Float)
+// CHECK:   [[FOO_VJP_SUBSET_THUNK:%.*]] = thin_to_thick_function [[FOO_VJP_SUBSET_THUNK_THIN]] : $@convention(thin) (@in_guaranteed Float, @in_guaranteed Float) -> (@out Float, @owned @callee_guaranteed (@in_guaranteed Float) -> @out Float) to $@callee_guaranteed (@in_guaranteed Float, @in_guaranteed Float) -> (@out Float, @owned @callee_guaranteed (@in_guaranteed Float) -> @out Float)
+// CHECK:   [[FOO_DIFF:%.*]] = autodiff_function [wrt 0] [order 1] [[FOO_FLOAT]] : $@callee_guaranteed (@in_guaranteed Float, @in_guaranteed Float) -> @out Float with {[[FOO_JVP_SUBSET_THUNK]] : $@callee_guaranteed (@in_guaranteed Float, @in_guaranteed Float) -> (@out Float, @owned @callee_guaranteed (@in_guaranteed Float) -> @out Float), [[FOO_VJP_SUBSET_THUNK]] : $@callee_guaranteed (@in_guaranteed Float, @in_guaranteed Float) -> (@out Float, @owned @callee_guaranteed (@in_guaranteed Float) -> @out Float)}
+// CHECK: }