[AutoDiff] Handle materializing adjoints with non-differentiable fields #67319

jkshtj · 2023-07-14T23:01:44Z

This PR fixes a compiler crasher in AutoDiff.

The compiler used to crash while generating a pullback for differentiable functions that take a differentiable input, whose tangent vector contains non-differentiable fields.

Changes in this PR fix the issue by specializing adjoint materialization logic for user-defined tangent vectors containing non-differentiable fields.

Fixes #66522

jkshtj · 2023-07-14T23:02:45Z

@asl @BradLarson could one of you invoke swift-ci tests?

...6522-pullback-generation-when-tangentvector-of-input-contains-nondifferentiable-fields.swift

asl · 2023-07-14T23:21:42Z

I am confused. I thought the discussion was that we need to fix adjoint generation for struct_extract instruction. Are there other usecases besides struct_extract that should be covered?

What is the test coverage for newly added code?

asl · 2023-07-14T23:22:01Z

@swift-ci please test

jkshtj · 2023-07-14T23:26:09Z

I am confused. I thought the discussion was that we need to fix adjoint generation for struct_extract instruction.

I think I may have misunderstood how to go about fixing the issue. Are you suggesting that we never add non-differentiable fields to the adjoint for struct_extract?

jkshtj · 2023-07-14T23:29:32Z

Are there other usecases besides struct_extract that should be covered?

If we do make the fix in the adjoint generation code, then I believe we'd have to at least do the same for tuple_extract. There may be other instructions as well, but I'd need to take a closer look.

jkshtj · 2023-07-14T23:32:29Z

What is the test coverage for newly added code?

In unit tests, none.

In the end to end tests, I added a tests for the path where we need to materialize an adjoint with non-differentiable fields. I should perhaps also have added a test for materializing normal adjoints with all differentiable fields.

asl · 2023-07-14T23:36:15Z

I think I may have misunderstood how to go about fixing the issue. Are you suggesting that we never add non-differentiable fields to the adjoint for struct_extract?

There are multiple questions that need to be answered:

Are there any other instructions that are affected here. Besides struct_extract? For example, you've added handling of adjoint buffers in addition to values. Are there any testcases for this code?
Next, the current adjoint generation for struct_extract is essentially:

  ///   Original: y = struct_extract x, #field
  ///    Adjoint: adj[x] += struct (0, ..., #field': adj[y], ..., 0)
  ///                                       ^~~~~~~
  ///                     field in tangent space corresponding to #field

So, instead of materalizing that bunch of zeros, why we simply cannot do the following (in struct_extract adjoint generation):

adj[x].field += adj[y]

After all, we know that things should conform to AdditiveArithmetic and therefore a += 0 should be equal to a.

asl · 2023-07-14T23:37:24Z

In unit tests, none.

Please ensure that all codepaths are covered by tests.

jkshtj · 2023-07-15T00:06:38Z

1. Are there any other instructions that are affected here. Besides `struct_extract`? For example, you've added handling of adjoint buffers in addition to values. Are there any testcases for this code?

Since the adjoint materialization code is shared by all instructions, we shouldn't need to handle things separately for different instructions. As for the handling of adjoint buffer handling, it looks like that code isn't actually called anywhere.

jkshtj · 2023-07-15T00:14:20Z

2. Next, the current adjoint generation for `struct_extract` is essentially:
  ///   Original: y = struct_extract x, #field
  ///    Adjoint: adj[x] += struct (0, ..., #field': adj[y], ..., 0)
  ///                                       ^~~~~~~
  ///                     field in tangent space corresponding to #field
So, instead of materalizing that bunch of zeros, why we simply cannot do the following:
adj[x].field += adj[y]
After all, we know that things should conform to AdditiveArithmetic and therefore a += 0 should be equal to a.

Ah, indeed this is what I tried to do first and I ran into an issue with TangentBuilder's emitInPlaceAdd method, which takes the lhs as a buffer. This was causing SIL verification failures for '+='-ing trivial types such as Double, with the following error message -

SIL verification failed: operand of 'apply' doesn't match function input type
  $Double
  $*Double
Verifying instruction:
   %0 = argument of bb0 : $Double                 // user: %8
     %5 = struct_element_addr %1 : $*P<Double>, #P.value // user: %8
     %6 = witness_method $Double, #AdditiveArithmetic."+=" : <Self where Self : AdditiveArithmetic> (Self.Type) -> (inout Self, Self) -> () : $@convention(witness_method: AdditiveArithmetic) <τ_0_0 where τ_0_0 : AdditiveArithmetic> (@inout τ_0_0, @in_guaranteed τ_0_0, @thick τ_0_0.Type) -> () // user: %8
     %7 = metatype $@thick Double.Type            // user: %8
->   %8 = apply %6<Double>(%5, %0, %7) : $@convention(witness_method: AdditiveArithmetic) <τ_0_0 where τ_0_0 : AdditiveArithmetic> (@inout τ_0_0, @in_guaranteed τ_0_0, @thick τ_0_0.Type) -> ()

It looks like I may have misinterpreted the meaning of the error, however. And I think I simply needed to change the type of the input arguments.

I can definitely try this out.

jkshtj · 2023-07-15T00:16:04Z

In unit tests, none.

Please ensure that all codepaths are covered by tests.

Would you know of an existing unit test that tests auto diff code? I wasn't quite sure how I can setup PullbackCloner for unit testing.

asl · 2023-07-15T00:16:58Z

1. Are there any other instructions that are affected here. Besides `struct_extract`? For example, you've added handling of adjoint buffers in addition to values. Are there any testcases for this code?
Since the adjoint materialization code is shared by all instructions, we shouldn't need to handle things separately for different instructions. As for the handling of adjoint buffer handling, it looks like that code isn't actually called anywhere.

Allright, but what are those "different instruction"? Besides tuple_extract? Looks like you're adding whole bunch of generic code without any tests just to fix one particular corner case. Differential paths are known to have less-than-optimal test coverage and this usually cause all kinds of issues when the code is actually called in some obscure case. So we try to improve the test coverage, not to reduce it.

asl · 2023-07-15T00:17:38Z

Would you know of an existing unit test that tests auto diff code? I wasn't quite sure how I can setup PullbackCloner for unit testing.

test/AutoDiff contain the tests for different aspects of autodiff code.

asl · 2023-07-15T00:18:19Z

This was causing SIL verification failures for '+='-ing trivial types such as Double, with the following error message -

The lhs is a struct somewhere. So you should be able to take the address of its field and pass to += method. However, in your particular case it was rhs that caused the assertion as the argument is passed indirect.

jkshtj · 2023-07-17T16:23:06Z

Would you know of an existing unit test that tests auto diff code? I wasn't quite sure how I can setup PullbackCloner for unit testing.

test/AutoDiff contain the tests for different aspects of autodiff code.

Oh, I think I was mistaking unit tests for the ones that occur in the unittests directory in the Swift repo.

jkshtj · 2023-07-17T16:30:00Z

This was causing SIL verification failures for '+='-ing trivial types such as Double, with the following error message -

The lhs is a struct somewhere. So you should be able to take the address of its field and pass to += method. However, in your particular case it was rhs that caused the assertion as the argument is passed indirect.

Yeah it was the right one. I was a little bit surprised because I was thinking that only the left one needs to be taken as a pointer (or by reference). But I'm guessing this more generalized SIL definition exists to accomodate for different types that may want to conform to AdditiveArithmetic?

jkshtj · 2023-07-17T16:37:21Z

1. Are there any other instructions that are affected here. Besides `struct_extract`? For example, you've added handling of adjoint buffers in addition to values. Are there any testcases for this code?
Since the adjoint materialization code is shared by all instructions, we shouldn't need to handle things separately for different instructions. As for the handling of adjoint buffer handling, it looks like that code isn't actually called anywhere.
Allright, but what are those "different instruction"? Besides tuple_extract? Looks like you're adding whole bunch of generic code without any tests just to fix one particular corner case. Differential paths are known to have less-than-optimal test coverage and this usually cause all kinds of issues when the code is actually called in some obscure case. So we try to improve the test coverage, not to reduce it.

I see the point. Indeed, I missed a number of tests that could be added with this change. I was thinking we can start by covering pullbacks for instructions listed here.

Can you think of any others that I should also test?

jkshtj · 2023-07-17T23:47:57Z

2. Next, the current adjoint generation for `struct_extract` is essentially:
  ///   Original: y = struct_extract x, #field
  ///    Adjoint: adj[x] += struct (0, ..., #field': adj[y], ..., 0)
  ///                                       ^~~~~~~
  ///                     field in tangent space corresponding to #field
So, instead of materalizing that bunch of zeros, why we simply cannot do the following:
adj[x].field += adj[y]
After all, we know that things should conform to AdditiveArithmetic and therefore a += 0 should be equal to a.

@asl I've incorporated this feedback about adjoint materialization and am creating a new PR revision for some early feedback.

I'm still working on adding test coverage but am struggling a bit to understand what exactly the tests should be doing. I'll post my questions about tests in another comment shortly.

jkshtj · 2023-07-18T00:18:03Z

@rxwei Could you also please take a look at the approach and comment whether it makes sense and whether I am making changes at the right place?

asl · 2023-07-18T00:20:40Z

@asl I've incorporated this feedback about adjoint materialization and am creating a new PR revision for some early feedback.

Have you forgotten to push your new changes?

jkshtj · 2023-07-18T00:24:53Z

@asl I've incorporated this feedback about adjoint materialization and am creating a new PR revision for some early feedback.

Have you forgotten to push your new changes?

I don't think so. I amended the existing commit, however, instead of creating a new one.

lib/SILOptimizer/Differentiation/PullbackCloner.cpp

test/AutoDiff/SILOptimizer/pullback_generation.swift

...6522-pullback-generation-when-tangentvector-of-input-contains-nondifferentiable-fields.swift

jkshtj · 2023-09-01T16:49:04Z

@rxwei @asl Thanks a lot for reviewing this PR so far!

I have fixed the nits. Could you please take a look?

asl · 2023-09-01T16:49:40Z

@rxwei @asl Thanks a lot for reviewing this PR so far!

I have fixed the nits. Could you please take a look?

See new comment. Also, will you please squash the commits?

jkshtj · 2023-09-01T16:57:37Z

See new comment.
Did you forget to publish the comment 😅?

Also, will you please squash the commits?
Sure.

asl · 2023-09-01T16:59:28Z

lib/SILOptimizer/Differentiation/PullbackCloner.cpp

@@ -430,6 +430,8 @@ class PullbackCloner::Implementation final
    }
    case AdjointValueKind::AddElement: {
      auto baseAdjoint = val;
+      assert(baseAdjoint.getType().is<TupleType>() ||


This is not what should be checked here. This assertion is enforced during AddElement construction. Here you need to ensure that all nested adjoints are of the same kind: either all tuples or all structs, but not the mixture.

I see. Sorry I sent out that last revision in haste. Fixed in the latest revision and squashed commits.

asl · 2023-09-02T02:03:48Z

lib/SILOptimizer/Differentiation/PullbackCloner.cpp

+        addEltAdjValues.push_back(addElementValue);
+        baseAdjoint = addElementValue->baseAdjoint;
+        assert(baseAdjointType == baseAdjoint.getType());
+      } while (baseAdjoint.getKind() == AdjointValueKind::AddElement);


Do we have tests for this loop, btw? Maybe I'm missing something, but all loops are doing single struct/tuple extract?

Not explicit tests but the other tests I have added do use this code path -- I needed to fix new test failures after making these changes. But I think it's a good idea to add explicit tests for this.

@asl I have added tests exercising the nested add element adjoint materialization path.

asl

Thanks!

asl · 2023-09-07T07:22:44Z

@swift-ci please test

jkshtj · 2023-09-10T05:19:28Z

@asl Had to make some changes after I discovered some small bugs while compiling and running internal test suite. I have added tests for these cases as well.

Fixes swiftlang#66522

During internal testing we discovered 2 more bugs - 1. The element adjoint of a struct_extract can itself be an AddElement. 2. Indirect concrete adjoint materialization was missing a copy operation. This commit fixes these bugs and adds relevant test cases.

asl

Makes sense, thanks!

asl · 2023-09-10T17:49:13Z

@swift-ci please test

jkshtj · 2023-09-10T21:16:28Z

@asl looks like the "Swift Test Linux Platform" failed due to some network issues. Can we restart them?

BradLarson · 2023-09-11T00:38:35Z

@swift-ci Please test Linux platform

asl · 2023-09-11T02:21:45Z

@swift-ci please test linux

jkshtj · 2023-09-11T14:42:07Z

Huh I see the following error due to which the Linux tests failed -

ERROR: Error fetching remote repo 'origin'

Seems like something may be up with the Linux build hosts?

asl · 2023-09-11T16:56:23Z

@shahmishal Is CI down?

asl · 2023-09-12T15:33:24Z

@swift-ci please test linux

jkshtj commented Jul 14, 2023

View reviewed changes

...6522-pullback-generation-when-tangentvector-of-input-contains-nondifferentiable-fields.swift Outdated Show resolved Hide resolved

jkshtj force-pushed the main branch from 005a208 to 63cfd7c Compare July 17, 2023 23:34

asl reviewed Jul 18, 2023

View reviewed changes

lib/SILOptimizer/Differentiation/PullbackCloner.cpp Outdated Show resolved Hide resolved

asl reviewed Jul 18, 2023

View reviewed changes

lib/SILOptimizer/Differentiation/PullbackCloner.cpp Outdated Show resolved Hide resolved

jkshtj force-pushed the main branch from 8b645c3 to 11ebece Compare July 19, 2023 15:58

jkshtj commented Jul 19, 2023

View reviewed changes

lib/SILOptimizer/Differentiation/PullbackCloner.cpp Outdated Show resolved Hide resolved

jkshtj requested a review from asl July 19, 2023 16:11

asl reviewed Jul 25, 2023

View reviewed changes

test/AutoDiff/SILOptimizer/pullback_generation.swift Outdated Show resolved Hide resolved

asl reviewed Jul 25, 2023

View reviewed changes

...6522-pullback-generation-when-tangentvector-of-input-contains-nondifferentiable-fields.swift Show resolved Hide resolved

jkshtj force-pushed the main branch from bfadeb1 to cfae072 Compare September 1, 2023 16:45

jkshtj requested a review from asl September 1, 2023 16:48

asl reviewed Sep 1, 2023

View reviewed changes

jkshtj force-pushed the main branch from cfae072 to f81411a Compare September 1, 2023 17:27

jkshtj requested a review from asl September 1, 2023 17:28

asl reviewed Sep 2, 2023

View reviewed changes

jkshtj force-pushed the main branch from f81411a to 2c7ed0b Compare September 5, 2023 17:47

jkshtj requested a review from asl September 5, 2023 17:48

asl approved these changes Sep 7, 2023

View reviewed changes

jkshtj force-pushed the main branch 3 times, most recently from 0cf3dbb to 9343748 Compare September 10, 2023 05:16

jkshtj requested a review from asl September 10, 2023 05:19

jkshtj added 2 commits September 10, 2023 09:26

[AutoDiff] Handle materializing adjoints with non-differentiable fields

b9fc40b

Fixes swiftlang#66522

jkshtj force-pushed the main branch from 9343748 to 111917e Compare September 10, 2023 17:03

asl approved these changes Sep 10, 2023

View reviewed changes

asl merged commit d971f12 into swiftlang:main Sep 12, 2023

[AutoDiff] Handle materializing adjoints with non-differentiable fields #67319

[AutoDiff] Handle materializing adjoints with non-differentiable fields #67319

Uh oh!

Conversation

jkshtj commented Jul 14, 2023

Uh oh!

jkshtj commented Jul 14, 2023

Uh oh!

Uh oh!

asl commented Jul 14, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

asl commented Jul 14, 2023

Uh oh!

jkshtj commented Jul 14, 2023

Uh oh!

jkshtj commented Jul 14, 2023

Uh oh!

jkshtj commented Jul 14, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

asl commented Jul 14, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

asl commented Jul 14, 2023

Uh oh!

jkshtj commented Jul 15, 2023

Uh oh!

jkshtj commented Jul 15, 2023

Uh oh!

jkshtj commented Jul 15, 2023

Uh oh!

asl commented Jul 15, 2023

Uh oh!

asl commented Jul 15, 2023

Uh oh!

asl commented Jul 15, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jkshtj commented Jul 17, 2023

Uh oh!

jkshtj commented Jul 17, 2023

Uh oh!

jkshtj commented Jul 17, 2023

Uh oh!

jkshtj commented Jul 17, 2023

Uh oh!

jkshtj commented Jul 18, 2023

Uh oh!

asl commented Jul 18, 2023

Uh oh!

jkshtj commented Jul 18, 2023

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jkshtj commented Sep 1, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

asl commented Sep 1, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jkshtj commented Sep 1, 2023

Uh oh!

asl Sep 1, 2023

Choose a reason for hiding this comment

Uh oh!

jkshtj Sep 1, 2023

Choose a reason for hiding this comment

Uh oh!

asl Sep 2, 2023

Choose a reason for hiding this comment

Uh oh!

jkshtj Sep 2, 2023

Choose a reason for hiding this comment

Uh oh!

jkshtj Sep 5, 2023

Choose a reason for hiding this comment

asl commented Jul 14, 2023 •

edited

Loading

jkshtj commented Jul 14, 2023 •

edited

Loading

asl commented Jul 14, 2023 •

edited

Loading

asl commented Jul 15, 2023 •

edited

Loading

jkshtj commented Sep 1, 2023 •

edited

Loading

asl commented Sep 1, 2023 •

edited

Loading