swift: [Autodiff] Memory leaks found under certain conditions.

Description The following code was found to leak memory upon the process exiting.

Steps to reproduce First, compile the following code in Debug mode (in a single file):

import _Differentiation; import Foundation
public struct B: Differentiable & Z {public var e: Float = 0}
extension B {@differentiable(reverse) func uu(dt: Float, nt: A<SIMD8<Float>>, f: A<Float>, zs: [S]) -> N {let nt = nt; let f = f; return N(nt: nt, f: f)}}
extension B {@differentiable(reverse) mutating func a(_ r: C) {}}
struct S: Differentiable & Z {}
struct C: Differentiable {var wm: Array<SIMD8<Float>>}
struct W: Differentiable {var h: B}
struct N: Differentiable {var nt: A<SIMD8<Float>>; var f: A<Float>}
struct O: Differentiable {@noDerivative public var h: B; @differentiable(reverse) public init(h: B) {self.h = h}}
public struct A<T>: Differentiable where T: Differentiable, T: AdditiveArithmetic {
    public struct TangentVector: Differentiable, AdditiveArithmetic {
        public typealias TangentVector = A.TangentVector
        public var _b: [T.TangentVector]
        public var _a: T.TangentVector
        public init(_b: [T.TangentVector], _a: T.TangentVector) {self._b = _b; self._a = _a}
    }
    @usableFromInline var _v: [T]
    @inlinable @differentiable(reverse) public init(_ values: [T], _a: T = .zero) {self._v = values}
    @inlinable @differentiable(reverse) public var _r: [T] { return _v }
    @inlinable @derivative(of: init(_:_a:)) static func _vjpInit(_ values: [T], _a: T = .zero) -> (value: A, pullback: (TangentVector) -> (Array<T>.TangentVector, T.TangentVector)){return (A(values, _a: _a), {v in return (Array<T>.TangentVector(v._b), v._a)})}
    @inlinable @derivative(of: _r) func vjpArray() -> (value: [T], pullback: (Array<T>.TangentVector) -> TangentVector) {func pullback(v: Array<T>.TangentVector) -> TangentVector {return TangentVector(_b: v.base, _a: T.TangentVector.zero)}; return (_v, pullback)}
    public mutating func move(by offset: TangentVector) {}
}
public extension A.TangentVector { // Not mathematically correct, of course, but simplified to this to demonstrate the memory leak(s).
    @inlinable static func + (lhs: Self, rhs: Self) -> Self {return lhs}
    @inlinable static func - (lhs: Self, rhs: Self) -> Self {return lhs}
    @inlinable static var zero: Self { Self(_b: [], _a: .zero) }
}
public protocol Z: Differentiable {}
func g(h: B, dt: Float) -> C {
    let nt = A([SIMD8<Float>(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)])
    let N = h.uu(dt: 180.0, nt: nt, f: A([Float(0.0)]), zs: [S]())
    return C(wm: N.nt._r)
}
func o<T, R>(_ x: T, _ f: @differentiable(reverse) (T) -> R) -> R {f(x)}
func p<T, R>(_ f: @escaping @differentiable(reverse) (T) -> R) -> @differentiable(reverse) (T) -> R {{ x in o(x, f) }}
@differentiable(reverse) func j(h: B, r: S) -> O {
    @differentiable(reverse) func q(_ l: W) -> W {var h = l.h; d(h: &h, r: r); var n = l; n.h = h; return n}
    var W = W(h: h)
    for _ in 0 ..< 1 {W = p(q)(W)} // The single-iteration for-loop needs to be here, otherwise the leak(s) won't occur.
    return O(h: W.h)
}
func b(h: B, r: S) -> (value: O, pullback: (O.TangentVector) -> (B.TangentVector, S.TangentVector))
{
    let s = valueWithPullback(at: h, r, of: j)
    return (value: s.value, pullback: s.pullback)
}
@differentiable(reverse) func d(h: inout B, r: S) {h.a(g(h: h, dt: h.e))}
func main() throws{_ = b(h: B(), r: S())}
try! main()

Next, run leaks, or another memory leak-checking tool. The command line for leaks to generate the following output is: leaks --atExit -- ./executableName. Note that Xcode’s memory leak checker may give inconsistent results from run to run with this code snippet.

The stack trace of the first leak is as follows:

STACK OF 1 INSTANCE OF 'ROOT LEAK: <Swift closure context>':
10  dyld                                  0x182ddff28 start + 2236
9   memleak2                              0x104cfaea0 main + 28  main.swift:50
8   memleak2                              0x104cfea38 main() + 36  main.swift:49
7   memleak2                              0x104cfe790 b(h:r:) + 436  main.swift:45
6   libswift_Differentiation.dylib        0x21081f0dc valueWithPullback<A, B, C>(at:_:of:) + 160
5   memleak2                              0x104cfe950 thunk for @callee_guaranteed (@unowned B, @unowned S) -> (@unowned O, @owned @escaping @callee_guaranteed (@unowned O.TangentVector) -> (@unowned B.TangentVector, @unowned S.TangentVector)) + 48  <compiler-generated>:0
4   memleak2                              0x104d00550 reverse-mode derivative of j(h:r:) + 1072  main.swift:40
3   memleak2                              0x104cfe570 thunk for @escaping @callee_guaranteed (@in_guaranteed W) -> (@out W, @owned @escaping @callee_guaranteed (@in_guaranteed W.TangentVector) -> (@out W.TangentVector)) + 100  <compiler-generated>:0
2   libswiftCore.dylib                    0x1921e0790 swift_allocObject + 64
1   libswiftCore.dylib                    0x1921e0594 swift_slowAlloc + 64
0   libsystem_malloc.dylib                0x182f78d88 _malloc_zone_malloc_instrumented_or_legacy + 128 
====
    14 (864 bytes) ROOT LEAK: <Swift closure context 0x147708560> [48]
       13 (816 bytes)  + 8 --> <Swift closure context 0x147708510> [80]
          12 (736 bytes)  + 8 --> <Swift closure context 0x1477084c0> [80]
             11 (656 bytes)  + 8 --> <Swift closure context 0x147708470> [80]
                10 (576 bytes)  + 8 --> <Swift closure context 0x147707ed0> [80]
                   9 (496 bytes)  + 8 --> <Swift closure context 0x147708440> [48]
                      8 (448 bytes)  + 8 --> <Swift closure context 0x147708410> [48]
                         7 (400 bytes)  + 8 --> <Swift closure context 0x147708360> [48]
                            6 (352 bytes)  + 8 --> <Swift closure context 0x1477081e0> [64]
                               5 (288 bytes)  + 8 --> <malloc in reverse-mode derivative of g(h:dt:) 0x1477083c0> [80]
                                  3 (160 bytes) <malloc in thunk for @escaping @callee_guaranteed (@guaranteed A<SIMD8<Float>>) -> (@owned [SIMD8<Float>], @owned @escaping @callee_guaranteed @substituted <A, B> (@guaranteed A) -> (@out B) for <[SIMD8<Float>]<A>.DifferentiableViewA<SIMD8<Float>>.TangentVector>) 0x147708390> [48]
                                     2 (112 bytes) <malloc in reverse-mode derivative of A._r.getter 0x147708140> [64]
                                        1 (48 bytes) <Swift closure context 0x147708330> [48]
                                  1 (48 bytes) <Swift closure context 0x147708300> [48]

The second one looks like this:

STACK OF 1 INSTANCE OF 'ROOT LEAK: <Swift closure context>':
20  dyld                                  0x182ddff28 start + 2236
19  memleak2                              0x104cfaea0 main + 28  main.swift:50
18  memleak2                              0x104cfea38 main() + 36  main.swift:49
17  memleak2                              0x104cfe790 b(h:r:) + 436  main.swift:45
16  libswift_Differentiation.dylib        0x21081f0dc valueWithPullback<A, B, C>(at:_:of:) + 160
15  memleak2                              0x104cfe950 thunk for @callee_guaranteed (@unowned B, @unowned S) -> (@unowned O, @owned @escaping @callee_guaranteed (@unowned O.TangentVector) -> (@unowned B.TangentVector, @unowned S.TangentVector)) + 48  <compiler-generated>:0
14  memleak2                              0x104d00550 reverse-mode derivative of j(h:r:) + 1072  main.swift:40
13  memleak2                              0x104cfe544 thunk for @escaping @callee_guaranteed (@in_guaranteed W) -> (@out W, @owned @escaping @callee_guaranteed (@in_guaranteed W.TangentVector) -> (@out W.TangentVector)) + 56  <compiler-generated>:0
12  memleak2                              0x104cff2dc partial apply + 96  <compiler-generated>:0
11  memleak2                              0x104d02498 reverse-mode derivative of closure #1 in p<A, B>(_:) + 636  main.swift:36
10  memleak2                              0x104d02850 reverse-mode derivative of o<A, B>(_:_:) + 188  main.swift:35
9   memleak2                              0x104cfe440 thunk for @escaping @callee_guaranteed (@unowned W) -> (@unowned W, @owned @escaping @callee_guaranteed (@unowned W.TangentVector) -> (@unowned W.TangentVector)) + 48  <compiler-generated>:0
8   memleak2                              0x104cfff88 reverse-mode derivative of q #1 (_:) in j(h:r:) + 60  main.swift:38
7   memleak2                              0x104d020d8 autodiff subset parameters thunk for reverse-mode derivative from d(h:r:) + 44  <compiler-generated>:0
6   memleak2                              0x104d00bf4 reverse-mode derivative of d(h:r:) + 84  main.swift:48
5   memleak2                              0x104d01b68 reverse-mode derivative of g(h:dt:) + 532  main.swift:32
4   memleak2                              0x104d01fcc autodiff subset parameters thunk for reverse-mode derivative from B.uu(dt:nt:f:zs:) + 44  <compiler-generated>:0
3   memleak2                              0x104cff8d0 reverse-mode derivative of B.uu(dt:nt:f:zs:) + 224  main.swift:3
2   libswiftCore.dylib                    0x1921e0790 swift_allocObject + 64
1   libswiftCore.dylib                    0x1921e0594 swift_slowAlloc + 64
0   libsystem_malloc.dylib                0x182f78d88 _malloc_zone_malloc_instrumented_or_legacy + 128 
====
    1 (48 bytes) ROOT LEAK: <Swift closure context 0x1477082d0> [48]

Expected behavior No leaks should be detected, and the program should exit with an exit code of 0.

Environment

  • Swift compiler version info: Toolchain 2023-07-10a. Toolchains as far back as 2023-01-09a will also exhibit this issue.
  • Xcode version info: 14.2
  • Deployment target: M1

Additional context

  • Removing the for loop in line 40, and running its contents once, will cause the leaks to vanish.
  • Removing the variable f from lines 3, 8 and 32 will cause one, but not both, of the leaks to vanish.
  • Removing the conformance to AdditiveArithmetic in line 11 will crash the compiler, with the following assertion failing: Assertion failed: (!ActiveDiagnostic && "Already have an active diagnostic")

About this issue

  • Original URL
  • State: closed
  • Created a year ago
  • Comments: 19 (3 by maintainers)

Commits related to this issue

Most upvoted comments

So, here is the proposal. Let’s make everything typed. We are having 3 functions / builtins:

AutoDiffLinearMapContext *swift_autoDiffCreateLinearMapContext(size_t);
void *swift_autoDiffProjectTopLevelSubcontext(AutoDiffLinearMapContext *);
void *swift_autoDiffAllocateSubcontext(AutoDiffLinearMapContext *, size_t);

They correspond to the following builtins:

autoDiffCreateLinearMapContext: (Builtin.Word) -> Builtin.NativeObject
autoDiffProjectTopLevelSubcontext: (Builtin.NativeObject) -> Builtin.RawPointer
autoDiffAllocateSubcontext: (Builtin.NativeObject, Builtin.Word) -> Builtin.RawPointer

Instead of amount of memory to be allocated, let us pass the types there. So, the swift builtins will become:

autoDiffCreateLinearMapContext: (T.Type) -> Builtin.NativeObject
autoDiffProjectTopLevelSubcontext: (Builtin.NativeObject) -> Builtin.RawPointer
autoDiffAllocateSubcontext: (Builtin.NativeObject, T.Type) -> Builtin.RawPointer

and the corresponding runtime functions will be:

AutoDiffLinearMapContext *swift_autoDiffCreateLinearMapContext(const Metadata*);
void *swift_autoDiffProjectTopLevelSubcontext(AutoDiffLinearMapContext *);
void *swift_autoDiffAllocateSubcontext(AutoDiffLinearMapContext *, const Metadata*);

The semantics would be as follows:

  • Instead of sizes of pullback tuples, the type will be passed
  • swift_autoDiffCreateLinearMapContext and swift_autoDiffAllocateSubcontext would create new boxes via swift_allocBox and store the resulting HeapObjects on the bumpptr-allocated storage.
  • swift_autoDiffAllocateSubcontext would return the corresponding box projection
  • swift_autoDiffProjectTopLevelSubcontext would return the previously-created (top-level) box projection
  • In the destructor of AutoDiffLinearMapContext it will iterate over bump-ptr allocated space and release boxes via swift_deallocBox.

@jkshtj this is for you 😃

Here is unoptimized code:

  // function_ref autodiff subset parameters thunk for reverse-mode derivative from static Float.+ infix(_:_:)
  %63 = function_ref @$sSf1poiyS2f_SftFZS3fXMtS3fIegyd_IetMyyydo_TJSrSSUpSrSUUP : $@convention(method) (Float, Float, @thin Float.Type) -> (Float, @owned @callee_guaranteed (Float) -> Float) // user: %64
  %64 = differentiable_function [parameters 0] [results 0] %59 : $@convention(method) (Float, Float, @thin Float.Type) -> Float with_derivative {%61 : $@convention(method) (Float, Float, @thin Float.Type) -> (Float
  %65 = differentiable_function_extract [vjp] %64 : $@differentiable(reverse) @convention(method) (Float, @noDerivative Float, @noDerivative @thin Float.Type) -> Float // user: %66
  %66 = apply %65(%0, %5, %58) : $@convention(method) (Float, Float, @thin Float.Type) -> (Float, @owned @callee_guaranteed (Float) -> Float) // users: %68, %67
  %67 = tuple_extract %66 : $(Float, @callee_guaranteed (Float) -> Float), 0 // user: %75
  %68 = tuple_extract %66 : $(Float, @callee_guaranteed (Float) -> Float), 1 // user: %70
  // function_ref pullback of f(x:)
  %69 = function_ref @$s4leak1f1xS2f_tFTJpSpSr : $@convention(thin) (Float, @guaranteed Builtin.NativeObject) -> Float // user: %74
  %70 = tuple $(predecessor: _AD__$s4leak1f1xS2f_tF_bb3__Pred__src_0_wrt_0, @callee_guaranteed (Float) -> Float) (%56, %68) // user: %73
  %71 = builtin "autoDiffProjectTopLevelSubcontext"(%2 : $Builtin.NativeObject) : $Builtin.RawPointer // user: %72
  %72 = pointer_to_address %71 : $Builtin.RawPointer to [strict] $*(predecessor: _AD__$s4leak1f1xS2f_tF_bb3__Pred__src_0_wrt_0, @callee_guaranteed (Float) -> Float) // user: %73
  store %70 to %72 : $*(predecessor: _AD__$s4leak1f1xS2f_tF_bb3__Pred__src_0_wrt_0, @callee_guaranteed (Float) -> Float) // id: %73
  %74 = partial_apply [callee_guaranteed] %69(%2) : $@convention(thin) (Float, @guaranteed Builtin.NativeObject) -> Float // user: %75
  %75 = tuple (%67 : $Float, %74 : $@callee_guaranteed (Float) -> Float) // user: %76
  return %75 : $(Float, @callee_guaranteed (Float) -> Float) // id: %76
} // end sil function '$s4leak1f1xS2f_tFTJrSpSr'

We’re capturing %68 here

Note that in optimized code due to inlining, etc. the context finally got optimized out by LLVM passes.

Thanks @jkshtj for nailing this down. This is optimized code for m():

// m()
sil hidden [noinline] @$s4leak1myyF : $@convention(thin) () -> () {
bb0:
  // function_ref f(x:)
  %0 = function_ref @$s4leak1f1xS2f_tF : $@convention(thin) (Float) -> Float // users: %6, %1
  %1 = thin_to_thick_function %0 : $@convention(thin) (Float) -> Float to $@callee_guaranteed (Float) -> Float // users: %39, %10
  // function_ref forward-mode derivative of f(x:)
  %2 = function_ref @$s4leak1f1xS2f_tFTJfSpSUr : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) // users: %7, %3
  %3 = thin_to_thick_function %2 : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) to $@callee_guaranteed (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) // users: %40, %11
  // function_ref reverse-mode derivative of f(x:)
  %4 = function_ref @$s4leak1f1xS2f_tFTJrSpSUr : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) // users: %8, %5
  %5 = thin_to_thick_function %4 : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) to $@callee_guaranteed (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) // users: %41, %12
  %6 = thin_to_thick_function %0 : $@convention(thin) (Float) -> Float to $@noescape @callee_guaranteed (Float) -> Float // user: %42
  %7 = thin_to_thick_function %2 : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) to $@noescape @callee_guaranteed (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) // user: %43
  %8 = thin_to_thick_function %4 : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) to $@noescape @callee_guaranteed (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) // user: %44
  %9 = integer_literal $Builtin.Word, 24          // user: %13
  strong_retain %1 : $@callee_guaranteed (Float) -> Float // id: %10
  strong_retain %3 : $@callee_guaranteed (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) // id: %11
  strong_retain %5 : $@callee_guaranteed (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) // id: %12
  %13 = builtin "autoDiffCreateLinearMapContext"(%9 : $Builtin.Word) : $Builtin.NativeObject // users: %38, %35, %28, %23, %18
  %14 = tuple ()                                  // users: %45, %15
  %15 = enum $_AD__$s4leak1f1xS2f_tF_bb1__Pred__src_0_wrt_0, #_AD__$s4leak1f1xS2f_tF_bb1__Pred__src_0_wrt_0.bb0!enumelt, %14 : $() // user: %17
  %16 = integer_literal $Builtin.Word, 8          // users: %28, %23, %18
  %17 = tuple $(predecessor: _AD__$s4leak1f1xS2f_tF_bb1__Pred__src_0_wrt_0) (%15) // user: %20
  %18 = builtin "autoDiffAllocateSubcontext"(%13 : $Builtin.NativeObject, %16 : $Builtin.Word) : $Builtin.RawPointer // users: %21, %19
  %19 = pointer_to_address %18 : $Builtin.RawPointer to [strict] $*(predecessor: _AD__$s4leak1f1xS2f_tF_bb1__Pred__src_0_wrt_0) // user: %20
  store %17 to %19 : $*(predecessor: _AD__$s4leak1f1xS2f_tF_bb1__Pred__src_0_wrt_0) // id: %20
  %21 = enum $_AD__$s4leak1f1xS2f_tF_bb2__Pred__src_0_wrt_0, #_AD__$s4leak1f1xS2f_tF_bb2__Pred__src_0_wrt_0.bb1!enumelt, %18 : $Builtin.RawPointer // user: %22
  %22 = tuple $(predecessor: _AD__$s4leak1f1xS2f_tF_bb2__Pred__src_0_wrt_0) (%21) // user: %25
  %23 = builtin "autoDiffAllocateSubcontext"(%13 : $Builtin.NativeObject, %16 : $Builtin.Word) : $Builtin.RawPointer // users: %26, %24
  %24 = pointer_to_address %23 : $Builtin.RawPointer to [strict] $*(predecessor: _AD__$s4leak1f1xS2f_tF_bb2__Pred__src_0_wrt_0) // user: %25
  store %22 to %24 : $*(predecessor: _AD__$s4leak1f1xS2f_tF_bb2__Pred__src_0_wrt_0) // id: %25
  %26 = enum $_AD__$s4leak1f1xS2f_tF_bb1__Pred__src_0_wrt_0, #_AD__$s4leak1f1xS2f_tF_bb1__Pred__src_0_wrt_0.bb2!enumelt, %23 : $Builtin.RawPointer // user: %27
  %27 = tuple $(predecessor: _AD__$s4leak1f1xS2f_tF_bb1__Pred__src_0_wrt_0) (%26) // user: %30
  %28 = builtin "autoDiffAllocateSubcontext"(%13 : $Builtin.NativeObject, %16 : $Builtin.Word) : $Builtin.RawPointer // users: %31, %29
  %29 = pointer_to_address %28 : $Builtin.RawPointer to [strict] $*(predecessor: _AD__$s4leak1f1xS2f_tF_bb1__Pred__src_0_wrt_0) // user: %30
  store %27 to %29 : $*(predecessor: _AD__$s4leak1f1xS2f_tF_bb1__Pred__src_0_wrt_0) // id: %30
  %31 = enum $_AD__$s4leak1f1xS2f_tF_bb3__Pred__src_0_wrt_0, #_AD__$s4leak1f1xS2f_tF_bb3__Pred__src_0_wrt_0.bb1!enumelt, %28 : $Builtin.RawPointer // user: %34
  // function_ref specialized autodiff subset parameters thunk for pullback from @escaping @callee_guaranteed (@unowned Float) -> (@unowned Float, @unowned Float)
  %32 = function_ref @$sS3fIegydd_TJSpSSUpSUUUrSUUP067$sSf16_DifferentiationE7_vjpAdd3lhs3rhsSf5value_Sf_SftSfc8pullbackth1_i5FZSf_I6SfcfU_Tf3npf_n : $@convention(thin) (Float) -> Float // user: %33
  %33 = thin_to_thick_function %32 : $@convention(thin) (Float) -> Float to $@callee_guaranteed (Float) -> Float // user: %34
  %34 = tuple $(predecessor: _AD__$s4leak1f1xS2f_tF_bb3__Pred__src_0_wrt_0, @callee_guaranteed (Float) -> Float) (%31, %33) // user: %37
  %35 = builtin "autoDiffProjectTopLevelSubcontext"(%13 : $Builtin.NativeObject) : $Builtin.RawPointer // user: %36
  %36 = pointer_to_address %35 : $Builtin.RawPointer to [strict] $*(predecessor: _AD__$s4leak1f1xS2f_tF_bb3__Pred__src_0_wrt_0, @callee_guaranteed (Float) -> Float) // user: %37
  store %34 to %36 : $*(predecessor: _AD__$s4leak1f1xS2f_tF_bb3__Pred__src_0_wrt_0, @callee_guaranteed (Float) -> Float) // id: %37
  strong_release %13 : $Builtin.NativeObject      // id: %38
  strong_release %1 : $@callee_guaranteed (Float) -> Float // id: %39
  strong_release %3 : $@callee_guaranteed (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) // id: %40
  strong_release %5 : $@callee_guaranteed (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) // id: %41
  strong_release %6 : $@noescape @callee_guaranteed (Float) -> Float // id: %42
  strong_release %7 : $@noescape @callee_guaranteed (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) // id: %43
  strong_release %8 : $@noescape @callee_guaranteed (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) // id: %44
  return %14 : $()                                // id: %45
} // end sil function '$s4leak1myyF'

Note the following:

  %33 = thin_to_thick_function %32 : $@convention(thin) (Float) -> Float to $@callee_guaranteed (Float) -> Float // user: %34
  %34 = tuple $(predecessor: _AD__$s4leak1f1xS2f_tF_bb3__Pred__src_0_wrt_0, @callee_guaranteed (Float) -> Float) (%31, %33) // user: %37
  %35 = builtin "autoDiffProjectTopLevelSubcontext"(%13 : $Builtin.NativeObject) : $Builtin.RawPointer // user: %36
  %36 = pointer_to_address %35 : $Builtin.RawPointer to [strict] $*(predecessor: _AD__$s4leak1f1xS2f_tF_bb3__Pred__src_0_wrt_0, @callee_guaranteed (Float) -> Float) // user: %37
  store %34 to %36 : $*(predecessor: _AD__$s4leak1f1xS2f_tF_bb3__Pred__src_0_wrt_0, @callee_guaranteed (Float) -> Float) // id: %37

Here we’re essentially capturing %33 into a loop context. And it’s the function context that we’re leaking here.