Lightning-AI / Lightning-AI/lightning-thunder

Too many `update_aliases` after in-place op

Open
#2,768 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

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

Description

```py
import torch, thunder

@thunder.jit
def fn(a):
a.tanh_()
return a * a * a * a

fn(torch.randn(5, 5, device='cuda'))
print(*thunder.last_traces(fn), sep='\n\n')
```

Ideally this function should be compiled into a single nvfuser region, but Thunder inserts `update_aliases` before every multiplication, causing fusion breaks.

```py
# ...

# Constructed by Update aliases for in-place ops
import thunder
import thunder.core.prims as prims
import thunder.torch as ltorch
import torch
from thunder.executors.torchex import no_autocast

@torch.no_grad()
@no_autocast
def computation(a):
# a: "cuda:0 f32[5, 5]"
(t5,) = prims.update_aliases((a,))

# /opt/pytorch/lightning-thunder/tmp/main.py:5: a.tanh_()
t1 = ltorch.tanh_(t5) # t1: "cuda:0 f32[5, 5]"
# t0 = ltorch.tanh(t5) # t0: "cuda:0 f32[5, 5]"
# t0 = prims.tanh(t5) # t0: "cuda:0 f32[5, 5]"
# t1 = prims.copy_(t0, t5, grad_enabled=True) # t1: "cuda:0 f32[5, 5]"
(t6,) = prims.update_aliases((t1,))

# /opt/pytorch/lightning-thunder/tmp/main.py:6: return a * a * a * a
t2 = ltorch.mul(t6, t6) # t2: "cuda:0 f32[5, 5]"
# t2 = prims.mul(t6, t6) # t2: "cuda:0 f32[5, 5]"
(t7,) = prims.update_aliases((t6,))

# /opt/pytorch/lightning-thunder/tmp/main.py:6: return a * a * a * a
t3 = ltorch.mul(t2, t7) # t3: "cuda:0 f32[5, 5]"
# t3 = prims.mul(t2, t7) # t3: "cuda:0 f32[5, 5]"
(t8,) = prims.update_aliases((t7,))

# /opt/pytorch/lightning-thunder/tmp/main.py:6: return a * a * a * a
t4 = ltorch.mul(t3, t8) # t4: "cuda:0 f32[5, 5]"
# t4 = prims.mul(t3, t8) # t4: "cuda:0 f32[5, 5]"
return {'output': (t4,), 'flat_args': [t8]}

# ...

# Constructed by Unwrap the actual return value
import torch
from thunder.executors.torchex import no_autocast

@torch.no_grad()
@no_autocast
def computation(a):
# a: "cuda:0 f32[5, 5]"
(t13,) = update_aliases((a,))
del a
[t1] = nvFusion0(t13)
# t0 = prims.tanh(t13) # t0: "cuda:0 f32[5, 5]"
# t1 = prims.copy_(t0, t13, grad_enabled=True) # t1: "cuda:0 f32[5, 5]"
del t13
(t14,) = update_aliases((t1,))
del t1
(t15,) = update_aliases((t14,))
[t3] = nvFusion2(t14, t15)
# t2 = prims.mul(t14, t14) # t2: "cuda:0 f32[5, 5]"
# t3 = prims.mul(t2, t15) # t3: "cuda:0 f32[5, 5]"
del t14
(t16,) = update_aliases((t15,))
del t15
[t4] = nvFusion3(t3, t16)
# t4 = prims.mul(t3, t16) # t4: "cuda:0 f32[5, 5]"
del t3
return (t4,)
```

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 provided Python reproducer using thunder.jit and thunder.last_traces, then trace where prims.update_aliases is inserted around the in-place tanh_ and multiplication operations. Done means the repeated multiplications can remain in a single nvfuser region without unnecessary update_aliases calls.

Written by the indexing model from the issue text.

Assessment

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.