[AutoDiff] Implement cross-file lookup of derivatives

[asl]

Look-up for functions with @derivative attributes defined in non-primary source files

Fixes #55170

[asl]

@swift-ci Please test

[asl]

Tagging @BradLarson

[asl]

Ok, look like the patch exposed the lookup cycle we're having here:

1. When we typecheck @derivative(of:) attribute we need to find the original function decl
2. This might trigger various requirement checks if the original function decl is generic
3. One of these checks is differentiability check that also checks what are the derivative function configurations out there
4. However, typechecking of @derivative(of:) also registers derivative function configuration

This was the cycle before, however it only caused correctness issue (e.g. we might miss the derivative configuration and produce an error). Now the cycle is explicitly exposed as we try to look for derivative configurations in no-primary files as well.

@rxwei Any preferences how to break the cycle? I'm thinking of removing the lookup logic from getDerivativeFunctionConfigurations and instead add "retry" logic into findMinimalDerivativeConfiguration. Alternatively we can to hack around and add a dedicated flag to getDerivativeFunctionConfigurations not to perform cross-module lookup during sema stage.

[dan-zheng]

Drive-by question: does @derivative attribute finding have an impact on incremental compilation?

---

When a file is updated, it and all files that depend on it need to be recompiled. Incremental compilation is fast when there are few dependencies between files.

I wonder if @derivative finding can cause incremental compilation to become asymptotically worse. Thinking through an example:

1. a.swift defines func foo. b.swift exists in the same module, defining a func vjpFoo function with no attributes.
2. Many files in the same module ask for the derivative of foo, e.g. by calling gradient(of: foo).

Initially, foo has no derivative functions registered via @derivative, so gradient(of: foo) triggers the compiler to automatically generate derivative functions.

3. func vjpFoo inb.swift is now marked with @derivative(of: foo).

Since all files (a.swift, b.swift, gradient(of: foo) files) are in the same module, the addition of @derivative(of: foo) in b.swift means that foo now has a registered derivative, and all functions that call gradient(of: foo) need to be recompiled†.

Is the † recompilation behavior asymptotically worse than existing incremental compilation? It might not be, if the same amount of work (e.g. recompiling the entire module) was necessary previous to the implementation of correct @derivative attribute finding.

[dan-zheng]

(It may be that the impact of @derivative attribute finding on incremental compilation was introduced prior to this PR.)

[asl]

@dan-zheng The complete recompilation is certainly necessary when optimizations are enabled. As derivatives might be inlined and subsequently optimized. So, effectively we'd need to throw out all the code and start afresh.

I do not like the existing getDerivativeFunctionConfigurations implementation as it has lots of side effects (even before this PR). Though lots of things are computed lazily here and there... Maybe cross-module lookup should be as a last-resort thing as it might be quite expensive on large modules – maybe we'd just iterate once and explicitly register all explicit derivatives.

[CodaFi]

The complete recompilation is certainly necessary when optimizations are enabled.

When optimizations are enabled you're building in whole-module mode anyways. Swift will refuse to perform the incremental build if you ask for -wmo as well.

An even simpler failure mode occurs here since this analysis is not going through name lookup. From the incremental build's perspective, there is no link between the derivative and the original.

[asl]

When optimizations are enabled you're building in whole-module mode anyways.

I'm confused. Will you please elaborate? The original testcase from #55170 fails regardless whether -O is provided to swiftc.

[asl]

@swift-ci please test

[rxwei]

Sorry I don't have the bandwidth in the near term to review this. Maybe @dan-zheng can help?

[CodaFi]

This analysis needs to go through name lookup. If that's triggering cycles then we need to figure out why getDerivativeFunctionConfigurations is on the lookup path.

[asl]

It cannot by design. See the full description in #58644 (comment) There is no way you could discover custom derivatives via name lookup as name is not known and the original function does not know anything about custom derivatives.

[CodaFi]

Please refactor this code to go through name lookup. You'll get access control, shadowing, and incremental build metadata done for free as a result. If there are cycles that is evidence that autodiff has a layering problem. We should perform minimal validation on the lookup path.

[dan-zheng]

Please refactor this code to go through name lookup. You'll get access control, shadowing, and incremental build metadata done for free as a result. If there are cycles that is evidence that autodiff has a layering problem. We should perform minimal validation on the lookup path.

This current approach for stateful derivative function configuration serialization (#30672) was adapted from prior work for serializing and loading extensions and Obj-C methods.

From current main:

• ASTContext::loadDerivativeFunctionConfigurations
• ASTContext::loadExtensions
• ASTContext::loadObjCMethods

If there's a better approach for handling derivative function configuration lookup that avoids this statefulness, that would be great.

---

One design choice is that derivative registration is done ① "retroactively" from @derivative functions to the original function, not ② from the original function to @derivative functions – so I'm not sure how to straightforwardly implement "get all derivative function configurations of foo" using name lookup, since foo does not know which @derivative functions have been registered for it (potentially across modules) or what those derivative functions' names are.

// Previous approach ②: derivative registration is done from original function to derivative functions.
// This does not enable retroactive derivative registration, e.g. registering derivatives for tgmath functions.
@differentiable(vjp: vjpFoo)
func foo(_ x: Float) -> Float { x }
func vjpFoo(_ x: Float) -> (value: Float, pullback: (Float) -> Float) { (x, { $0 }) }

// Current approach ①: derivative registration is done from derivative functions to original function.
//
// This supports retroactive registration: derivatives can be registered for functions in other modules.
// For example, `vjpFoo` could be defined in module B, importing module A where `foo` is defined.
//
// However, determining "registered derivatives" for a function is less straightforward, hence the
// `ASTContext::loadDerivativeFunctionConfigurations` complications.
func foo(_ x: Float) -> Float { x }

@derivative(of: foo)
func vjpFoo(_ x: Float) -> (value: Float, pullback: (Float) -> Float) { (x, { $0 }) }

[dan-zheng]

This PR as an incremental fix/improvement looks good!

If tests pass, merging this sounds good to me unless others have concerns.

[asl]

@swift-ci please test