[AutoDiff upstream] Add differentiability_witness_function instruction.

[dan-zheng]

The differentiability_witness_function instruction looks up a
differentiability witness function (JVP, VJP, or transpose) for a referenced
function via SIL differentiability witnesses.

// Example differentiability witnesses.
sil_differentiability_witness [parameters 0] [results 0] @foo : $@convention(thin) (Float, Float, Float) -> Float

sil_differentiability_witness [parameters 0 1] [results 0] <T where T : Differentiable, T == T.TangentVector> @generic : $@convention(thin) <T> (@in_guaranteed T, Float) -> @out T

// Example `differentiability_witness_function` instructions.
sil @test : $@convention(thin) () -> () {
bb0:
  %foo_jvp_wrt_0 = differentiability_witness_function [jvp] [parameters 0] [results 0] @foo : $@convention(thin) (Float, Float, Float) -> Float
  %generic_vjp_wrt_0_1 = differentiability_witness_function [vjp] [parameters 0 1] [results 0] <T where T: Differentiable, T == T.TangentVector> @generic : $@convention(thin) <T> (@in_guaranteed T, Float) -> @out T
  return undef : $()
}

Add round-trip parsing/serialization and IRGen tests.

Notes:

• Differentiability witnesses for linear functions require more support.
  differentiability_witness_function [transpose] instructions do not yet
  have IRGen.
• Nothing currently generates differentiability_witness_function instructions.
  The differentiation transform does, but it hasn't been upstreamed yet.

Resolves TF-1141.

[rxwei]

Very nice!

[dan-zheng]

@swift-ci Please test

[swift-ci]

Build failed
Swift Test OS X Platform
Git Sha - eaba367

[dan-zheng]

@swift-ci Please test

[swift-ci]

Build failed
Swift Test Linux Platform
Git Sha - eaba367

[swift-ci]

Build failed
Swift Test OS X Platform
Git Sha - eaba367

[dan-zheng]

FileCheck tests failed for the 32-bit iOS simulator due to LLVM IR differences. The test should be amended now.

[dan-zheng]

@swift-ci Please test

[swift-ci]

Build failed
Swift Test Linux Platform
Git Sha - 4ea1240

[swift-ci]

Build failed
Swift Test OS X Platform
Git Sha - 4ea1240

[compnerd]

This instruction is very confusing to me (not sure if its just me though). differentiability_witness_function makes me think that this returns the witness. If this is looking up the associated differentiation function associated with a function would you be open to renaming this to something like differential?

[rxwei]

This name is similar to ’witness_method’, where it returns a method from a witness. ‘differentiability_witness_function’ returns a function from a differentiability witness. “Differential” is not the technically right name and can cause confusion with the “differential” in our API. The returned function is not a differential, but a JVP/VJP (derivative) function.

I’m entirely open to renaming this to ‘derivative_function’, but I just wanted to clarify that the current name falls in line with the name ‘witness_method’.

[compnerd]

Oops, I thought Id looked through the rest and then came back and changed differential to derivative. Id be okay with derivative_function as well.

The thing is that function becomes a fuzzy term here. It fetches the function pointer from the witness, returning a function which implements the function which performs a derivative over a function. The way that I initially read it, it made me question whether it was getting a getter for the differentiability witness or an entry in the witness table.

[dan-zheng]

derivative_function is perhaps not wholly precise, because the instruction has a [transpose] option for returning a transpose function, in addition to [jvp] and [vjp] options for returning derivative functions. I'm open to derivative_function though.

Do we have consensus on derivative_function (or some other name) as an alternative name for differentiability_witness_function? cc @rxwei @marcrasi

If not, let's merge and defer renaming until later. I'll merge within a day if no one responds.

[rxwei]

In the future each witness will be either a linear witness or non-linear witness, so it’s possible to define both a ‘derivative_function’ and a ‘transpose_function’.

That said, I don’t like the idea of dropping the word “witness” from this instruction. It makes it sounds like it’s differentiating or transposing something, whereas all it does is retrieve the pointer to a function in the witness.

In the future when we unify JVP and VJP into a single derivative, a differentiability witness will be equivalent to a function forward declaration, so calling it ‘differentiability_witness_function’ is quite clear. I’m supporting keeping the existing name.

[dan-zheng]

In the future when we unify JVP and VJP into a single derivative, a differentiability witness will be equivalent to a function forward declaration, so calling it ‘differentiability_witness_function’ is quite clear.

Interesting point. I'll go ahead and merge this patch now.