[AutoDiff] Fix adjoints for loop-local active values #78374
Conversation
|
Tagging @asl @JaapWijnen @rxwei |
|
@swift-ci please test |
kovdan01
force-pushed
the
issue78264-new
branch
2 times, most recently
from
6960602 to
f9bb221
Compare
|
@swift-ci please test |
kovdan01
force-pushed
the
issue78264-new
branch
from
f9bb221 to
ffe5c3c
Compare
|
Would be glad to see feedback from everyone interested |
|
@swift-ci please test |
|
Would be glad to see feedback from everyone interested |
1 similar comment
|
Would be glad to see feedback from everyone interested |
|
Would be glad to see feedback from everyone interested |
1 similar comment
|
Would be glad to see feedback from everyone interested |
|
@kovdan01 I cherrypicked this change onto |
|
The change looks ok to me. The explanation of the problem is at #78264 Long story short: SSA values in the loop body should be considered as "new values" at each loop iteration. We used to reuse the adjoints for them before this patch (though we should create distinct adjoints on each loop iteration). This used to work before due to a particular CFG generated for for-style loops. As a result we've been lucky and simply zeroed the adjoints in the loop header on each iteration (loop header BB was the first BB we encountered the adjoint, so it was initialized there and this zero-initialization happened on each loop iteration effectively producing new adjoint value). However, this stopped to work for while-style loops as the CFG is different and while we produced adjoint initialization, it happened outside the loop and therefore the adjoint values got reused from the previous loop iteration. @rxwei will you please take a look as a second pair of eyes? :) |
To be a bit more precise: it not stopped working, it just never worked for repeat-while loops (if not never, at least, for several latest years). For regular while loops, we had the same situation that we observed with for loops, when the behavior looked correct because of a particular CFG. |
@clackary Great, thanks for letting me know! BTW, does "everything builds successfully on macOS" imply that your internal tests are also passing? |
@kovdan01 No, just the compile-time checks for now. Further testing is more of a manual process at the moment, but I can get to testing a few projects this week. |
|
@rxwei Could you please take a look at this PR? Thanks! |
|
@kovdan01 I was able to finish testing a bunch of our internal projects. The issue we discussed earlier was caused by a flaw in my invocation, so nothing to worry about there. This looks good from my end. |
|
@rxwei Could you please take a look at this PR? Thanks! |
|
@swift-ci please smoke test |
| assert(!arrayValue && "Array value already found"); | ||
| // The first `destructure_tuple` result is the `Array` value. | ||
| arrayValue = dti->getResult(0); | ||
| #ifdef NDEBUG |
There was a problem hiding this comment.
Consider using ASSERT / DEBUG_ASSERT from swift/Basic/Assertions.h instead of conditional asserts, etc. Overall, the intention is to get rid of conditional asserts in new code (and avoid #ifdef NDEBUG).
There was a problem hiding this comment.
I've put this break; under DEBUG_ASSERT_enabled condition since having the check present increases iterations count and this should probably not be done in non-debug builds. Please let me know if this looks reasonable for you.
| // invariants hold and then skip. | ||
| if (auto *ai = getAllocateUninitializedArrayIntrinsicElementAddress( | ||
| bbActiveValue)) { | ||
| #ifndef NDEBUG |
There was a problem hiding this comment.
Deleted this preprocessor condition and switched to ASSERT. These checks should be OK in terms of negative performance impact, so ASSERT seems preferrable over DEBUG_ASSERT
| continue; | ||
| } | ||
|
|
||
| assert(getTangentValueCategory(bbActiveValue) == |
There was a problem hiding this comment.
@asl As far as I understood, there is an intention to enable most checks in all builds including release ones. Given that, should I switch these and other remaining regular assert to ASSERT?
|
@swift-ci please test |
1 similar comment
|
@swift-ci please test |
|
@swift-ci please smoke test |
Fixes #78264