Skip to content

Conversation

@asl
Copy link
Contributor

@asl asl commented Feb 2, 2024

Unreachable blocks possess some challenges to autodiff since in reverse pass (pullback generation) we need to execute the function backwards, pushing the values from return BB back to entry block. As a result, unreachable blocks might become reachable from the return BB and this might cause all kind of issues.

This PR aims to make autodiff a bit more robust with respect to unreachable blocks. I tried to find what are the cases where we might hit or miss unreachable blocks and here is the breakdown:

  1. When building the linear map types and branch tracing enums we traverse the function in reverse post order traversal. Therefore unreachable blocks are naturally excluded from RPOT and we do not have linear map tuples for them. Ensure that when we are building branch tracing enums we're excluding predecessor BBs without linear map tuples.
  2. VJP cloner does BFS from entry block, so it does not see / clone unreachable blocks. We are fine here
  3. Pulllback cloner does BFS starting from function return block and therefore might see unreachable blocks. Ensure they are:
    • Skipped for cloning / pullback codegen
    • Since we have not created branch tracing enum entries (due to 1.), ensure we skip them while forming pullback successor blocks

I think this covers all cases. At least I was unable to bring more testcases that fail :)

Fixes #71164

Complements #71339

@asl asl requested review from eeckstein and rxwei February 2, 2024 22:16
@asl
Copy link
Contributor Author

asl commented Feb 2, 2024

@swift-ci please test

@asl
Copy link
Contributor Author

asl commented Feb 3, 2024

@swift-ci please test

@asl asl merged commit f5cce47 into main Feb 3, 2024
@asl asl deleted the 71164-fix branch February 3, 2024 16:33
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[AutoDiff] Assertion failed: (linearMapStructTy && "must have linear map struct type for predecessor BB")

3 participants