Lightning-AI/lightning-thunder

Handling inplace through SSA

Open

#145 opened on Apr 9, 2024

View on GitHub
 (16 comments) (1 reaction) (0 assignees)Python (114 forks)github user discovery
enhancementhelp wanted

Repository metrics

Stars
 (1,460 stars)
PR merge metrics
 (PR metrics pending)

Description

This issue is to facilitate discussion of inplace handling, namely the "big" solution of having a static single assignment (SSA) representation.

For any handling of inplace, we want to make certain that two things are achieved:

  • we don't want to take shortcuts that complicate passes by introducing the need to detect obstacles to optimizations, because it would harm usability and extensibility of Thunder.
  • we don't want to create ad-hoc band-aids to get things working that we would need to regress on later to introduce more proper handling because developing in the open more or less means no regressions.

Some thoughts from video/chat discussions:

About the problem:

  • The key difficulty in SSA is that we would need to keep track of which tensors get modified by an inplace update (i.e. which have memory that is to be updated), so we would need to know about views (the fancy term is alias analysis),
  • this is difficult for some things in PyTorch (i.e. reshape),
  • "assuming the worst" works to some extend.

Solution considerations:

  • Likely we would want inplace updates to have all affected tensors as outputs.
  • on inputs we would need to check for aliases as part of the prologue (maybe with a separate "assume aliasing is the OK" cache mode or sorts later),
  • operations need to know if their output is a view of their inputs (difficult for reshape, easy for most others),
  • initially, we would only check if tensors share storage,
  • likely the translation could be done in the interpretation phase,
  • we would need to have versioning / disambiguation of versions for tensor proxies during this, but not when we have the SSA.

Later versions could refine the alias analysis as needed.

@tfogal @mruberry @IvanYashchuk

Contributor guide