Lightning-AI / Lightning-AI/lightning-thunder
Handling inplace through SSA
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
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
No contributing guide indexed for this repository
Research direction
No files or tests are named. Start by reviewing the interpretation phase and the existing tensor-proxy handling, then trace how inplace updates, views, and aliases are represented. Done means an agreed SSA-based approach that tracks affected tensors and aliasing without ad-hoc pass-specific workarounds.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- compilers
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100