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

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.