swift
542c2368 - [AutoDiff] Differentiate partially-applied reabstraction thunks. (#28570)

Commit
5 years ago
[AutoDiff] Differentiate partially-applied reabstraction thunks. (#28570) Support reabstraction thunk differentiation. Specifically, support this reabstraction thunk `partial_apply` pattern: ``` %generic_fn = function_ref @generic : $<τ_0_0> (@in_guaranteed τ_0_0) -> @out τ_0_0 %specialized_fn = partial_apply %generic_fn<Float>() %reabs_thunk = function_ref @reabs_thunk : $(Float, @guaranteed @callee_guaranteed (@in_guaranteed Float) -> @out Float) -> Float %thunked_fn = partial_apply %reabs_thunk(%specialized_fn) %diff_fn = differentiable_function %thunked_fn ``` - SIL: add special reabstraction thunk case to `SILFunctionType::getAutoDiffDerivativeFunctionType`. - Reabstraction thunks have a function-typed last argument, representing the function to reabstract. - That argument is transformed into a `@differentiable` function-typed argument for reabstraction thunk JVPs/VJPs. This enables differentiation transform support: reabstraction thunk JVP/VJP callers are responsible for passing a `@differentiable` function. Otherwise, the function-typed argument is opaque and cannot be differentiated. - Differentiation transform: - Handle reabstraction thunk `partial_apply` reapplications in `reapplyFunctionConversion`. Form a `differentiable_function` of the function-typed argument and `partial_apply` it to the reabstraction thunk JVP/VJP. This enables differentiation of direct references to generic functions, which generate reabstraction thunk `partial_apply`s in SIL. Resolves TF-201. Exposes TF-1033: differential generation ownership error. TF-1036 tracks future reabstraction thunk derivative generation optimizations.
Author
Parents
Loading