Lightning-AI / Lightning-AI/lightning-thunder

Handling inplace through SSA

Open
#145 16 comments 1 reaction 0 assignees View on GitHub
enhancement help wanted
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.