[AutoDiff] Add control flow AD leak checking tests.

stdlib/private/DifferentiationUnittest/GenericLifetimeTracked.swift

@@ -20,20 +20,35 @@ public enum _GlobalLeakCount {
 /// automatic differentiation.
 public struct Tracked<T> {
   fileprivate class Box {
-    fileprivate let value : T
+    fileprivate var value : T
     init(_ value: T) {
       self.value = value
-        _GlobalLeakCount.count += 1
+      _GlobalLeakCount.count += 1
     }
     deinit {
       _GlobalLeakCount.count -= 1
     }
   }
-  private let handle: Box
+  private var handle: Box
+
+  @differentiable(
+    vjp: _vjpInit
+    where T : Differentiable, T == T.AllDifferentiableVariables,
+          T == T.TangentVector
+  )
   public init(_ value: T) {
     self.handle = Box(value)
   }
-  public var value: T { return handle.value }
+
+  @differentiable(
+    vjp: _vjpValue
+    where T : Differentiable, T == T.AllDifferentiableVariables,
+          T == T.TangentVector
+  )
+  public var value: T {
+    get { handle.value }
+    set { handle.value = newValue }
+  }
 }
 
 extension Tracked : ExpressibleByFloatLiteral where T : ExpressibleByFloatLiteral {
@@ -123,17 +138,37 @@ extension Tracked : Differentiable
   public typealias TangentVector = Tracked<T.TangentVector>
 }
 
-@differentiable(vjp: _vjpAdd)
-public func + (_ a: Tracked<Float>, _ b: Tracked<Float>) -> Tracked<Float> {
-  return Tracked<Float>(a.value + b.value)
+extension Tracked where T : Differentiable, T == T.AllDifferentiableVariables,
+                        T == T.TangentVector
+{
+  @usableFromInline
+  internal static func _vjpInit(_ value: T)
+      -> (value: Self, pullback: (Self.TangentVector) -> (T.TangentVector)) {
+    return (Tracked(value), { v in v.value })
+  }
+
+  @usableFromInline
+  internal func _vjpValue() -> (T, (T.TangentVector) -> Self.TangentVector) {
+    return (value, { v in Tracked(v) })
+  }
 }
 
-@usableFromInline
-func _vjpAdd(_ a: Tracked<Float>, _ b: Tracked<Float>)
-    -> (Tracked<Float>, (Tracked<Float>) -> (Tracked<Float>, Tracked<Float>)) {
-  return (Tracked<Float>(a.value + b.value), { v in
-    return (v, v)
-  })
+extension Tracked where T : Differentiable, T == T.AllDifferentiableVariables,
+                        T == T.TangentVector
+{
+  @usableFromInline
+  @differentiating(+)
+  internal static func _vjpAdd(lhs: Self, rhs: Self)
+      -> (value: Self, pullback: (Self) -> (Self, Self)) {
+    return (lhs + rhs, { v in (v, v) })
+  }
+
+  @usableFromInline
+  @differentiating(-)
+  internal static func _vjpSubtract(lhs: Self, rhs: Self)
+      -> (value: Self, pullback: (Self) -> (Self, Self)) {
+    return (lhs - rhs, { v in (v, .zero - v) })
+  }
 }
 
 // Differential operators for `Tracked<Float>`.
@@ -151,4 +186,20 @@ public extension Differentiable {
   ) -> (TangentVector, T.TangentVector) {
     return self.pullback(at: x, in: f)(1)
   }
+
+  @inlinable
+  func valueWithGradient(
+    in f: @differentiable (Self) -> Tracked<Float>
+  ) -> (value: Tracked<Float>, gradient: TangentVector) {
+    let (y, pb) = self.valueWithPullback(in: f)
+    return (y, pb(1))
+  }
+
+  @inlinable
+  func valueWithGradient<T : Differentiable>(
+    at x: T, in f: @differentiable (Self, T) -> Tracked<Float>
+  ) -> (value: Tracked<Float>, gradient: (TangentVector, T.TangentVector)) {
+    let (y, pb) = self.valueWithPullback(at: x, in: f)
+    return (y, pb(1))
+  }
 }

test/AutoDiff/leakchecking.swift

@@ -1,17 +1,22 @@
 // RUN: %target-run-simple-swift-control-flow-differentiation
 // REQUIRES: executable_test
 
-// A test that we can properly differentiate types that require refcounting.
+// Test differentiation-related memory leaks.
 
 import StdlibUnittest
 import DifferentiationUnittest
 
 var LeakCheckingTests = TestSuite("LeakChecking")
 
 /// Execute body, check expected leak count, and reset global leak count.
-func testWithLeakChecking(expectedLeakCount: Int = 0, _ body: () -> Void) {
+func testWithLeakChecking(
+  expectedLeakCount: Int = 0, file: String = #file, line: UInt = #line,
+  _ body: () -> Void
+) {
   body()
-  expectEqual(expectedLeakCount, _GlobalLeakCount.count, "Leak detected.")
+  expectEqual(
+    expectedLeakCount, _GlobalLeakCount.count, "Leak detected.",
+    file: file, line: line)
   _GlobalLeakCount.count = 0
 }
 
@@ -23,27 +28,122 @@ struct ExampleLeakModel : Differentiable {
   }
 }
 
+struct FloatPair : Differentiable & AdditiveArithmetic {
+  var first, second: Tracked<Float>
+  init(_ first: Tracked<Float>, _ second: Tracked<Float>) {
+    self.first = first
+    self.second = second
+  }
+}
+
+struct Pair<T : Differentiable, U : Differentiable> : Differentiable
+  where T == T.AllDifferentiableVariables, T == T.TangentVector,
+        U == U.AllDifferentiableVariables, U == U.TangentVector
+{
+  var first: Tracked<T>
+  var second: Tracked<U>
+  init(_ first: Tracked<T>, _ second: Tracked<U>) {
+    self.first = first
+    self.second = second
+  }
+}
+
 LeakCheckingTests.test("BasicVarLeakChecking") {
-  do {
+  testWithLeakChecking {
+    var model = ExampleLeakModel()
+    let x: Tracked<Float> = 1.0
+    _ = model.gradient(at: x) { m, x in m.applied(to: x) }
+  }
+
+  testWithLeakChecking {
+    var model = ExampleLeakModel()
+    let x: Tracked<Float> = 1.0
+
+    _ = model.gradient { m in m.applied(to: x) }
+    for _ in 0..<10 {
+      _ = model.gradient { m in m.applied(to: x) }
+    }
+  }
+
+  testWithLeakChecking {
+    var model = ExampleLeakModel()
+    var x: Tracked<Float> = 1.0
+    _ = model.gradient { m in
+      x = x + x
+      var y = x + Tracked<Float>(x.value)
+      return m.applied(to: y)
+    }
+  }
+
+  // TODO: Fix memory leak.
+  testWithLeakChecking(expectedLeakCount: 1) {
     var model = ExampleLeakModel()
     let x: Tracked<Float> = 1.0
-    let _ = model.gradient(at: x) { m, x in m.applied(to: x) }
+    _ = model.gradient { m in
+      var model = m
+      // Next line causes leak.
+      model.bias = x
+      return model.applied(to: x)
+    }
   }
-  expectEqual(0, _GlobalLeakCount.count, "Leak detected.")
 }
 
 LeakCheckingTests.test("ControlFlow") {
-  // TODO: Add more `var` + control flow tests.
-  // Porting tests from test/AutoDiff/control_flow.swift requires more support
-  // for `Tracked<Float>`.
+  // FIXME: Fix control flow AD memory leaks.
+  // See related FIXME comments in adjoint value/buffer propagation in
+  // lib/SILOptimizer/Mandatory/Differentiation.cpp.
+  testWithLeakChecking(expectedLeakCount: 105) {
+    func cond_nestedtuple_var(_ x: Tracked<Float>) -> Tracked<Float> {
+      // Convoluted function returning `x + x`.
+      var y = (x + x, x - x)
+      var z = (y, x)
+      if x > 0 {
+        var w = (x, x)
+        y.0 = w.1
+        y.1 = w.0
+        z.0.0 = z.0.0 - y.0
+        z.0.1 = z.0.1 + y.0
+      } else {
+        z = ((y.0 - x, y.1 + x), x)
+      }
+      return y.0 + y.1 - z.0.0 + z.0.1
+    }
+    expectEqual((8, 2), Tracked<Float>(4).valueWithGradient(in: cond_nestedtuple_var))
+    expectEqual((-20, 2), Tracked<Float>(-10).valueWithGradient(in: cond_nestedtuple_var))
+    expectEqual((-2674, 2), Tracked<Float>(-1337).valueWithGradient(in: cond_nestedtuple_var))
+  }
+
+  // FIXME: Fix control flow AD memory leaks.
+  // See related FIXME comments in adjoint value/buffer propagation in
+  // lib/SILOptimizer/Mandatory/Differentiation.cpp.
+  testWithLeakChecking(expectedLeakCount: 379) {
+    func cond_nestedstruct_var(_ x: Tracked<Float>) -> Tracked<Float> {
+      // Convoluted function returning `x + x`.
+      var y = FloatPair(x + x, x - x)
+      var z = Pair(Tracked(y), x)
+      if x > 0 {
+        var w = FloatPair(x, x)
+        y.first = w.second
+        y.second = w.first
+        z.first = Tracked(FloatPair(z.first.value.first - y.first,
+                                    z.first.value.second + y.first))
+      } else {
+        z = Pair(Tracked(FloatPair(y.first - x, y.second + x)), x)
+      }
+      return y.first + y.second - z.first.value.first + z.first.value.second
+    }
+    expectEqual((8, 2), Tracked<Float>(4).valueWithGradient(in: cond_nestedstruct_var))
+    expectEqual((-20, 2), Tracked<Float>(-10).valueWithGradient(in: cond_nestedstruct_var))
+    expectEqual((-2674, 2), Tracked<Float>(-1337).valueWithGradient(in: cond_nestedstruct_var))
+  }
 
   // FIXME: Fix control flow AD memory leaks.
   // See related FIXME comments in adjoint value/buffer propagation in
   // lib/SILOptimizer/Mandatory/Differentiation.cpp.
   testWithLeakChecking(expectedLeakCount: 9) {
     var model = ExampleLeakModel()
     let x: Tracked<Float> = 1.0
-    let _ = model.gradient(at: x) { m, x in
+    _ = model.gradient(at: x) { m, x in
       let result: Tracked<Float>
       if x > 0 {
         result = m.applied(to: x)
@@ -60,7 +160,7 @@ LeakCheckingTests.test("ControlFlow") {
   testWithLeakChecking(expectedLeakCount: 14) {
     var model = ExampleLeakModel()
     let x: Tracked<Float> = 1.0
-    let _ = model.gradient(at: x) { m, x in
+    _ = model.gradient(at: x) { m, x in
       var result: Tracked<Float> = x
       if x > 0 {
         result = result + m.applied(to: x)