Lightning-AI / Lightning-AI/lightning-thunder

Unhelpful variable renaming

Open
#739 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

jit
Dominant language
Python
Stars
1.5k
Forks
121
PR merge metrics
No merged PRs in 30d

Description

The commit 7c916c13675bb05b1a5522a9c797b33e997e4f19 tries to give identifiers observed in the source to the proxies of variables. However such identifiers sometimes refer to thunder's internal implementation of the language's construct, and this can be confusing.

### Code sample

```py
@thunder.jit
def f(xs, s):
for i, x in enumerate(xs):
s += x
return s

n = 6
xs = [torch.zeros(n) for _ in range(n)]
s = torch.zeros(n)
f(xs, s)
print(thunder.last_traces(f)[-1])
```
```py
@torch.no_grad()
@no_autocast
def computation(res, elem, x, b, t_0_4, t_0_5, s):
# res: "cpu f32[6]"
# elem: "cpu f32[6]"
# xs: "cpu f32[6]"
# b: "cpu f32[6]"
# t_0_4: "cpu f32[6]"
# t_0_5: "cpu f32[6]"
# s: "cpu f32[6]"
t0 = torch.add(s, res) # t0: "cpu f32[6]"
# ...
```
These renamings happen in [`thunder/core/jit_ext._maybe_update_proxy_name`](https://github.com/Lightning-AI/lightning-thunder/blob/main/thunder/core/jit_ext.py#L1108), which is called from [`thunder.core.interpreter._load_fast_handler`](https://github.com/Lightning-AI/lightning-thunder/blob/main/thunder/core/interpreter.py#L4572). `frame.code` reveals that the irrelevant variable names come from the lookasides implemented in `thunder.core.interpreter`. `res` is from [`SequenceIter.__next__`](https://github.com/Lightning-AI/lightning-thunder/blob/main/thunder/core/interpreter.py#L1997), `elem` is from [`_enumerate_lookaside`](https://github.com/Lightning-AI/lightning-thunder/blob/main/thunder/core/interpreter.py#L1382), `b` from [`_binary_op`](https://github.com/Lightning-AI/lightning-thunder/blob/main/thunder/core/interpreter.py#L3228).

### Ideal behavior

When deciding the identifiers we can just ignore those in thunder's interpreter, which will result in
```py
def computation(x, t_0_1, t_0_2, t_0_3, t_0_4, t_0_5, s):
```
Altenatively, we can perhaps label the identifier `x` by an index when `x` is bound to multiple proxies, as
```py
def computation(x_0, x_1, x_2, x_3, x_4, t_0_5, s):
```

cc @t-vi @nikitaved

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 at thunder/core/jit_ext.py:_maybe_update_proxy_name and trace its call from thunder/core/interpreter.py:_load_fast_handler. Inspect the interpreter lookasides, including SequenceIter.__next__, _enumerate_lookaside, and _binary_op, to distinguish implementation identifiers from source identifiers. Done means generated proxy names no longer expose irrelevant interpreter names, with the intended naming behavior resolved for the provided example.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
compilers
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.