[AutoDiff] Lookup for custom derivatives in non-primary source files #58965

asl · 2022-05-18T17:26:33Z

This is done after typecheck is finished for the primary source.

This registers all custom derivatives before autodiff transformations and makes them available to them.

Fully resolves #55170

…heck is finished for the primary source. This registers all custom derivatives before autodiff transformations and makes them available to them. Fully resolves swiftlang#55170

asl · 2022-05-18T17:27:15Z

Note that this PR effectively makes large part of a1e138b obsolete and, more importantly, removes cycle-breaking hacks as they are not needed anymore

asl · 2022-05-18T17:27:23Z

@swift-ci please test

asl · 2022-05-18T17:27:37Z

Tagging @rxwei @BradLarson

test/AutoDiff/Sema/DerivativeRegistrationCrossFile/main.swift

Co-authored-by: Philip Turner <[email protected]>

asl · 2022-05-18T20:27:22Z

@swift-ci please test

dan-zheng

Looks good, thanks!

dan-zheng · 2022-05-18T23:05:36Z

lib/Frontend/Frontend.cpp

+  forEachSourceFile([](SourceFile &SF) {
+    loadDerivativeConfigurations(SF);
+    return false;
+  });


Do you have any sense of the performance implications (empirical or theoretical) of this derivative function resolution approach, compared with the old approach?

A high-level performance summary seems to be:

No overhead if import _Differentiation is missing.

Otherwise, every declaration in every source file is visited.

Right. And any custom derivative declaration is typechecked. Regardless whether the derivative is used or not. This is quite a big hammer, but for now I do not see other way to "expose" them.

Old approach had its own pros and cons:

Pros: other source files were visited only when derivative was needed. And only those required were typechecked

Cons: other source files were visited on every derivative lookup. And we had typechecking cycles.

dan-zheng · 2022-05-19T00:16:16Z

test/AutoDiff/Sema/DerivativeRegistrationCrossFile/main.swift

+@differentiable(reverse)
+func clamp(_ value: Double, _ lowerBound: Double, _ upperBound: Double) -> Double {
+    // No error expected
+    return max(min(value, upperBound), lowerBound)


Could you please clarify the functionality added by this PR on top of #58644?

My understanding:

[AutoDiff] Implement cross-file lookup of derivatives #58644 adds an option for non-primary-file derivative resolution to AbstractFunctionDecl::getDerivativeFunctionConfigurations, as well as original declaration serialization for @derivative attribute.

This PR moves the non-primary-file derivative resolution logic to CompilerInstance::finishTypeChecking instead, visiting each source file.

This resolves the "cross-file derivative not found" issue in [SR-14235] [AutoDiff] Cross-file differentiation error: derivative configuration not registered for @derivative attribute #55170.

Does this also avoid circular requests?

Are there other implications to this new resolution-triggering condition? Ostensibly, to fix [SR-14235] [AutoDiff] Cross-file differentiation error: derivative configuration not registered for @derivative attribute #55170, this PR is causing derivative resolution to happen in more cases than previously in [AutoDiff] Implement cross-file lookup of derivatives #58644.

Yes, as I mentioned, this avoids circular requests. As custom derivative resolution is decoupled from obtaining derivatives configurations. Also, after this PR the differentiation will not trigger typechecking

After this PR indeed, we lookup and register all explicit derivates regardless of their original function and regardless whether this function is used. The nice side effect is that we will error out in case of @derivative(of: foo) with foo not defined anywhere.

Effectively this resolves the following cross-file derivate requests:

File A.swift is using (wants a derivative of) function from B.swift with derivatives defined in C.swift

The important special case is when B.swift is effectively a part of stdlib (like in the testcase). Here C.swift is a non-primary source but we have no way to resolve the derivative in other way as there is no mechanism to "pull" anything from C.swift – there are simply no forward dependency edges in the dependency graph, only backward.

slavapestov · 2022-05-19T19:56:27Z

lib/Sema/TypeChecker.cpp

@@ -375,6 +375,44 @@ void swift::performWholeModuleTypeChecking(SourceFile &SF) {
  }
 }

+void swift::loadDerivativeConfigurations(SourceFile &SF) {
+  if (!isDifferentiableProgrammingEnabled(SF))


When is this true? The recursive walk over the AST in secondary files is going to trigger delayed parsing in all nominal types and extensions, which will impact performance in non-WMO builds.

isDifferentiableProgrammingEnabled (implementation) returns true only for files that have import _Differentiation, so I hope there would be negligible impact for other files.

Ah ok. Can you add a comment to that effect? Something that mentions that this defeats delayed parsing, but only triggers with 'import _Differentiation'.

@slavapestov This is true when there is an import _Differentiation in the given SourceFile (so, corresponding ImportDecl).

What is the viable alternative? Essentially the task is: for function f find all functions that has @derivative(of : f) attribute specified. Both in primary and secondary sources.

What is the viable alternative? Essentially the task is: for function f find all functions that has @Derivative(of : f) attribute specified. Both in primary and secondary sources.

The only real alternative would be to change the design of this attribute so that this kind of search does not have to be performed. By definition, it requires parsing all source files ahead of time which is incompatible with our strategy of skipping bodies of nominal types and extensions in secondary files.

Yeah. Though the alternative design would not easily allow "pluggable" derivatives.... (see #58644 (comment) for the design explanation, etc. by @dan-zheng )

One alternative is to handle this like we do operators and AnyObject lookup. When skipping over the body of a nominal type or extension, make a note if the lexer sees '@_derivative', then eagerly parse only those delayed bodies when performing this global lookup. That's still not great but better than disabling delayed parsing entirely when 'import _Differentiation' appears in the program.

@slavapestov Yes, I though about this as well, thanks! Though I was a bit hesitant to introduce such derivative-specific logic into lexer and parser :) But if you're thinking it's acceptable, I can try to do something around this in the subsequent PRs.

philipturner · 2022-06-18T01:03:31Z

I wanted to make the #59467 regression test as a "suggestion" in a GitHub review, but it's too large to merge a contribution that way. Please check out asl#1, which will put the test into asl:55170-fix. For more context, see #59467 (comment).

asl · 2022-06-18T09:50:21Z

@rxwei Do you happen to have any further comments on this? Or we can land this PR?

Thanks!

rxwei

LGTM but I would also get a ✅ from @slavapestov

philipturner · 2022-06-29T04:37:07Z

Could we get this merged? The longer we wait, the more nightly snapshots are produced that are unusable with my AutoDiff client library. If general consensus is that it’s acceptable, let’s merge it (and the uncommitted regression test) sooner rather than later.

asl · 2022-06-29T10:15:53Z

We are currently waiting for @slavapestov's opinion per @rxwei request

philipturner · 2022-07-08T12:48:28Z

We are currently waiting for slavapestov's opinion per rxwei request

Would it be possible to temporarily revert #58644 until #58965 is approved? That would remove the crash it introduced, letting other modifications to the compiler (such as RQM, SIL patches) be tested against AutoDiff. Furthermore, #58644 is like an incomplete version of #58965. We should wait to incorporate cross-file lookup the Swift compiler until its implementation is complete. This solution would also provide reviewers with as much time as necessary to decide whether to merge this PR.

There is also a precedent for this, as it is how the merging of #42189 played out. The initial version caused a regression on Windows, much like the how #58644 introduced #59467. After being temporarily reverted, the PR was reapplied with fixes for Windows.

asl · 2022-07-18T18:07:45Z

@slavapestov gently ping :)

slavapestov · 2022-07-18T18:31:55Z

I would still like to see you redo this to work like the global operator and AnyObject lookup, however, avoiding paging in all delayed parsed bodies in all source files.

slavapestov · 2022-07-18T18:32:21Z

lib/Frontend/Frontend.cpp

 void CompilerInstance::finishTypeChecking() {
  forEachFileToTypeCheck([](SourceFile &SF) {
    performWholeModuleTypeChecking(SF);
    return false;
  });
+
+  forEachSourceFile([](SourceFile &SF) {


Is there a way to avoid the walk over all source files in the case where none of them have differentiable programming enabled?

asl · 2022-07-18T18:51:44Z

@slavapestov Thanks! I created #58965 to track it. Will post follow-up patch

philipturner · 2022-07-18T20:10:55Z

@asl Thank you so much! But what happened to the regression test from asl#1? Should I make a new PR to merge it into apple/swift, or did you decide that it's not a good idea?

asl · 2022-07-18T20:11:40Z

@philipturner I will take care of it separately, do not worry :)

…wiftlang#58965) * Lookup for custom derivatives in non-primary source files after typecheck is finished for the primary source. This registers all custom derivatives before autodiff transformations and makes them available to them. Fully resolves swiftlang#55170

In swiftlang#58965, lookup for custom derivatives in non-primary source files was introduced. It required traversing all delayed parsed function bodies of a file if the file was compiled with differential programming enabled (even for functions with no `@derivative` attribute). This patch introduces `CustomDerivativesLookupRequest` to address the issue. Resolves swiftlang#60102

In swiftlang#58965, lookup for custom derivatives in non-primary source files was introduced. It required traversing all delayed parsed function bodies of a file if the file was compiled with differential programming enabled (even for functions with no `@derivative` attribute). This patch introduces `CustomDerivativesRequest` to address the issue. Resolves swiftlang#60102

In #58965, lookup for custom derivatives in non-primary source files was introduced. It required triggering delayed members parsing of nominal types in a file if the file was compiled with differential programming enabled. This patch introduces `CustomDerivativesRequest` to address the issue. We only parse delayed members if tokens `@` and `derivative` appear together inside skipped nominal type body (similar to how member operators are handled). Resolves #60102

Lookup for custom derivatives in non-primary source files after typec…

0b06fb4

…heck is finished for the primary source. This registers all custom derivatives before autodiff transformations and makes them available to them. Fully resolves swiftlang#55170

asl requested a review from dan-zheng May 18, 2022 17:26

philipturner suggested changes May 18, 2022

View reviewed changes

test/AutoDiff/Sema/DerivativeRegistrationCrossFile/main.swift Outdated Show resolved Hide resolved

Add newline to end of file

c7b6df7

Co-authored-by: Philip Turner <[email protected]>

dan-zheng approved these changes May 19, 2022

View reviewed changes

rxwei self-requested a review May 19, 2022 06:00

slavapestov requested changes May 19, 2022

View reviewed changes

This was referenced Jun 15, 2022

[AutoDiff][IRGen] Can't compile Swift for TensorFlow on top-of-tree Swift #59467

Closed

Add S4TF crasher test asl/swift#1

Closed

rxwei approved these changes Jun 19, 2022

View reviewed changes

asl requested a review from slavapestov June 21, 2022 15:30

slavapestov approved these changes Jul 18, 2022

View reviewed changes

slavapestov reviewed Jul 18, 2022

View reviewed changes

asl mentioned this pull request Jul 18, 2022

I would still like to see you redo this to work like the global operator and AnyObject lookup, however, avoiding paging in all delayed parsed bodies in all source files. #60102

Closed

asl merged commit 6e2c4fa into swiftlang:main Jul 18, 2022

asl mentioned this pull request Jul 20, 2022

Add regression test for #59467 #60158

Merged

kovdan01 mentioned this pull request Oct 10, 2024

[AutoDiff] Enhance performance of custom derivatives lookup #76951

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[AutoDiff] Lookup for custom derivatives in non-primary source files #58965

[AutoDiff] Lookup for custom derivatives in non-primary source files #58965

asl commented May 18, 2022

asl commented May 18, 2022

asl commented May 18, 2022

asl commented May 18, 2022

asl commented May 18, 2022

dan-zheng left a comment

dan-zheng May 18, 2022

asl May 19, 2022

dan-zheng May 19, 2022

asl May 19, 2022

slavapestov May 19, 2022

dan-zheng May 19, 2022 •

edited

Loading

slavapestov May 19, 2022

asl May 19, 2022

slavapestov May 19, 2022

asl May 19, 2022

slavapestov Jun 18, 2022

asl Jun 18, 2022

philipturner commented Jun 18, 2022 •

edited

Loading

asl commented Jun 18, 2022

rxwei left a comment

philipturner commented Jun 29, 2022

asl commented Jun 29, 2022

philipturner commented Jul 8, 2022

asl commented Jul 18, 2022

slavapestov commented Jul 18, 2022

slavapestov Jul 18, 2022

asl commented Jul 18, 2022

philipturner commented Jul 18, 2022 •

edited

Loading

asl commented Jul 18, 2022 •

edited

Loading

[AutoDiff] Lookup for custom derivatives in non-primary source files #58965

[AutoDiff] Lookup for custom derivatives in non-primary source files #58965

Conversation

asl commented May 18, 2022

asl commented May 18, 2022

asl commented May 18, 2022

asl commented May 18, 2022

asl commented May 18, 2022

dan-zheng left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

dan-zheng May 19, 2022 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

philipturner commented Jun 18, 2022 • edited Loading

asl commented Jun 18, 2022

rxwei left a comment

Choose a reason for hiding this comment

philipturner commented Jun 29, 2022

asl commented Jun 29, 2022

philipturner commented Jul 8, 2022

asl commented Jul 18, 2022

slavapestov commented Jul 18, 2022

Choose a reason for hiding this comment

asl commented Jul 18, 2022

philipturner commented Jul 18, 2022 • edited Loading

asl commented Jul 18, 2022 • edited Loading

dan-zheng May 19, 2022 •

edited

Loading

philipturner commented Jun 18, 2022 •

edited

Loading

philipturner commented Jul 18, 2022 •

edited

Loading

asl commented Jul 18, 2022 •

edited

Loading