Lightning-AI / Lightning-AI/lightning-thunder
dce pass does not treat correctly `DONT_DCE` tag in sub-symbols
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
## 🐛 Bug
This is a very specific and hard to reach bug, but a bug nonetheless. If one of the subsymbols has the DONT_DCE tag, but the output of the parent bound symbol is not used, running dce will remove the output proxy but not the bound symbol.
This is due to `is_needed` being `True` but the output proxy is not in `needed_proxies`
https://github.com/Lightning-AI/lightning-thunder/blob/d0f647479f73787389c4fa6ced3b0b9ada4083c0/thunder/core/transform_common.py#L182-L187
### To Reproduce
```python
import torch
import thunder
from thunder.core.transform_common import dce
def foo(a):
b = a.add_(5)
return a.add_(4)
a = torch.randn(2, 2)
jf = thunder.jit(foo)
jf(a)
trace_before_functionalization = thunder.last_traces(jf)[-7]
broken_trace = dce(trace_before_functionalization)
print(broken_trace)
```
Will print something like this:
```python
# Constructed by Dead Code Elimination (took 0 milliseconds)
import thunder
import thunder.torch as ltorch
import torch
from thunder.executors.torchex import no_autocast
@torch.no_grad()
@no_autocast
def computation(a):
# a: "cpu f32[2, 2]"
# :6: b = a.add_(5)
ltorch.add_(a, 5, alpha=1)
# t0 = ltorch.add(a, 5, alpha=1) # t0: "cpu f32[2, 2]"
# t0 = prims.add(a, 5.0) # t0: "cpu f32[2, 2]"
# b = prims.copy_(t0, a, grad_enabled=True) # b: "cpu f32[2, 2]"
# :7: return a.add_(4)
t3 = ltorch.add_(a, 4, alpha=1) # t3: "cpu f32[2, 2]"
# t2 = ltorch.add(a, 4, alpha=1) # t2: "cpu f32[2, 2]"
# t2 = prims.add(a, 4.0) # t2: "cpu f32[2, 2]"
# t3 = prims.copy_(t2, a, grad_enabled=True) # t3: "cpu f32[2, 2]"
return {'output': (t3,), 'flat_args': [a]}
```
Discovered while working on #1961
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 in thunder/core/transform_common.py around the linked is_needed logic and run the supplied Python reproduction to inspect the trace before and after dce. Trace how DONT_DCE on a subsymbol affects needed_proxies and the parent bound symbol. Done means dead-code elimination preserves the required bound symbol while removing only genuinely unused output proxies, with regression coverage for this case.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100