[AutoDiff] Add 'Builtin.applyTranspose*'. #28469

rxwei · 2019-11-25T03:36:20Z

Add support for applying the transpose function in a @differentiable(linear) function. The Builtin.applyTranspose* builtin takes a @differentiable(linear) function and a tangent vector, and returns the result of applying the transpose to the tangent vector.

Pseudo-declaration:

func applyTranspose_arity{arity}[_throws?{r}]<T...{arity}, R>(
    _: @differentiable (T...) {r}?throws -> R,
    _: R
) {r}?rethrows -> (T...)
    where T: Differentiable & AdditiveArithmetic, R: Differentiable & AdditiveArithmetic

This patch also renames Builtin.autodiffApply to Builtin.applyDerivative for clarity, and fixes a bug in LinearDifferentiableSILFunctionTypeLowering where it expected 3 component values in rebuildAggregate instead of 2.

Resolves SR-11844 and SR-11851.

@differentiable

…ntiable(linear)` function from component functons. * `Builtin.differentiableFunction_*` Takes an original function, a JVP function and a VJP function and returns a `@differentiable` function. Pseudo-declaration: ```swift func differentiableFunction_arity{arity}[_throws]?{throws}<T...{arity}, U>( _ original: __owned @escaping (T...{arity}) {throws}? -> U, _ jvp: __owned @escaping (T...{arity}) {throws}? -> (value: U, differential: (T.TangentVector...{arity}) -> U.TangentVector), _ vjp: __owned @escaping (T...{arity}) {throws}? -> (value: U, pullback: (U.TangentVector) -> (T.TangentVector...{arity})) ) -> @differentiable (T...{arity}) {throws}? -> U where T...{arity} : Differentiable, U : Differentiable ``` * `Builtin.linearFunction_*` Takes an original function and a transpose function and returns a `@differentiable` function. Pseudo-declaration: ```swift func linearFunction_arity{arity}[_throws]?{throws}<T...{arity}, U>( _ original: __owned @escaping (T...{arity}) {throws}? -> U, _ transpose: __owned @escaping (U.TangentVector) {throws}? -> (T.TangentVector...{arity}) ) -> @differentiable (T...{arity}) {throws}? -> U where T...{arity} : Differentiable & AdditiveArithmetic, U : Differentiable & AdditiveArithmetic ``` These builtins will be used to write unit tests for `@differentiable` and `@differentiable(linear)` function types that do not necessarily depend on the differentiation transform. TODO: - SR-11848: For robustness, we need SIL FileCheck tests for all AD builtins. These have not been added for `Builtin.autodiffApply*`, so I'm leaving this as a future task. - SR-11847: Update `differentiableFunction(from:)` to use `Builtin.differentiableFunction*` in its implementation. - SR-11849: Disallow non-top-level derivative registration. Resolves SR-11846.

@differentiable

Add support for applying the transpose function in a `@differentiable(linear)` function. The `Builtin.applyTranspose*` builtin takes a `@differentiable(linear)` function and a tangent vector, and returns the result of applying the transpose to the tangent vector. Pseudo-declaration: ```swift func applyTranspose_arity{arity}[_rethrowws?{r}]<T...{arity}, R>( _: @differentiable (T...) {r}?throws -> R, _: R ) {r}?rethrows -> (T...) where T: Differentiable & AdditiveArithmetic, R: Differentiable & AdditiveArithmetic ``` This patch also renames `Builtin.autodiffApply` to `Builtin.applyDerivative` for clarity. Resolves SR-11844 and SR-11851.

…nspose

dan-zheng

Nice!

Note for self: Builtin.apply{Derivative,Transpose}* and Builtin.{differentiable,linear}Function* are all higher-order functions.

apply{Derivative,Transpose} are rethrows because they may throw depending on whether the argument function throws.

{differentiable,linear}Function are not throwing because they merely form a bundle of functions.

lib/SILGen/SILGenBuiltin.cpp

lib/AST/AutoDiff.cpp

rxwei · 2019-11-25T04:10:28Z

apply{Derivative,Transpose} are rethrows because they may throw depending on whether the argument function throws.

I just unified the terminology to just "throws" for simplicity :) The _throws configuration refers to the throwing-ness of the @differentiable/@differentiable(linear) function.

rxwei · 2019-11-25T04:11:12Z