differentiate struct_extract instructions using VJP

[marcrasi]

Differentiates struct_extract instructions using the VJP of the corresponding getter.

Float._value has no VJP defined so AD can't differentiate public static prefix func - (x: ${Self}). This could be fixed by defining a VJP for Float._value or by adding a custom adjoint for -. The latter sounded less invasive so I did that.

[marcrasi]

@swift-ci please test tensorflow

[marcrasi]

@swift-ci please test tensorflow

[marcrasi]

@swift-ci please test tensorflow

[rxwei]

Great work!

[rxwei]

Instead of always allocating stack and emitting a load when the adjoint is direct, could you call materializeAdjointDirect directly? It'd simplify both your logic and possibly generated code.

[marcrasi]

Oh cool, in fact, I should always be able to materializeAdjointDirect because struct_extract only works on object types. This simplifies everything.

[rxwei]

The helper materializeAdjointIndirectHelper is not supposed to be used outside of materializeAdjointIndirect. It expects the buffer argument to be the result of begin_access. Please switch to materializeAdjointIndirect if that does what you need.

[rxwei]

When we have derived conformances to Differentiable, it'd be much more efficient to use the original struct_extract differentiation logic when we know that the cotangent type equals the original type and that the property vjp is compiler-synthesized (this can be checked via Decl::isImplicit(), I think).

[marcrasi]

In fact, this gives me an idea. The original struct_extract differentiation logic does not need getter VJPs. So if I can get the original struct_extract differentiation logic working on structs with derived conformance to Differentiable (even those with cotangent type not equal to self), then I don't need to write any code that synthesizes getter VJPs.

I think that the original struct_extract logic can be made to work even on derived conformances to Differentiable with cotangent type not equal to self, because it should be possible to determine which fields correspond.

I will try out this approach.

[rxwei]

Sounds reasonable for internal structs. Public structs need property VJPs in any case because of resilience requirements.

[rxwei]

Nit

Suggested change

	// Get the VJP results (original results and pullback)
	// Get the VJP results (original results and pullback).

[rxwei]

When will a getter not exist?

[marcrasi]

One example is the case that I added to test/AutoDiff/autodiff_diagnostics.swift.

I don't know why the getter doesn't exist in that case, and if I put the same code in a file an compile it with a plain call to swiftc, the getter does exist and the code ends up triggering the case where the getter exists but doesn't have a VJP.

I haven't investigated any farther than that, so I don't really know what's going on but the logic seems to handle all the cases I can think of correctly.

[marcrasi]

reminder to self: must add a retain here

    getBuilder().createRetainValue(loc, originalDirectResult,
                                   getBuilder().getDefaultAtomicity());

[rxwei]

I don’t think inlinable is necessary.

[marcrasi]

I get a SILVerifier error about referencing a non-ABI-public function from a fragile function if I remove inlinable. (simple_math.swift triggers it)

I think what's happening is that _transparent doesn't make the function ABI public, and mandatory inlining isn't inlining the function because we're partially applying it.

[rxwei]

Gotcha. If you just need it to be ABI-public without needing it serializable, just use @usableFromInline.

[marcrasi]

@swift-ci please test tensorflow

[marcrasi]

@swift-ci please test tensorflow

[marcrasi]

@swift-ci please test tensorflow