Lightning-AI / Lightning-AI/lightning-thunder

executor should be able to put checks into prologue trace

Open
#263 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

dynamic constraints enhancement
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

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.