Lightning-AI / Lightning-AI/lightning-thunder
executor should be able to put checks into prologue trace
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
## 🚀 Feature
dynamic constraints would want to insert executor specific checks into prologue trace, given that backends might have specific dynamic constraints.
A quick example as below:
Given a program to be compiled, where we would expect the reduction axis to be a direct input to the program vvv
```
def foo(a, reducedim):
return a.sum(reducedim)
```
The compute trace looks like below vvv.
```
import thunder
import thunder.torch as ltorch
import torch
from thunder.executors.torchex import no_autocast
@torch.no_grad()
@no_autocast()
def computation(a, i0, i1):
# a: "cuda:0 f32[8, 16, 32]"
# i0: "int 0"
# i1: "int 1"
t2 = ltorch.sum(a, (i0, i1), False, dtype=None) # t2: "cuda:0 f32[32]"
# t2 = prims.sum(a, (i0, i1)) # t2: "cuda:0 f32[32]"
return t2
```
Depends on which backend claims the `ltorch.sum`. i.e. for torchex, since aten can handle arbitrary `reducedim` given at runtime, we can re-use the cache and there's no need to insert any check on `arg[1]`.
On the contrary, nvfuserex would require the program to bake in reduction axis as compile time constant. So we'd want to insert that as part of the prologue trace checks.
### Alternatives
Alternative 0: we can converge on the *most conservative* backends and apply a simpler caching strategy at the primitive level. In the example above, we'll just require reduction axis to stay as a compile time constant thing across the board, even though it could be re-used for some executors. This would unfortunately gives us some negative cache hit but would be easier to plumb through.
Alternative 1: thunder as a system can establish a caching strategy. When a backend sees a cache requirement on a certain op that it cannot fulfill, the backend could just reject the operation.
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with the prologue trace and backend handling for the `ltorch.sum` example, comparing the torchex and nvfuserex cases described in the issue. Determine how executor-specific dynamic constraints would be represented and how they affect caching; done should allow each backend to enforce its required checks without imposing the most conservative strategy on all executors.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- backend, compilers
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100