Lightning-AI / Lightning-AI/lightning-thunder

dce pass does not treat correctly `DONT_DCE` tag in sub-symbols

Open
#2,132 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

optimization passes
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

  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 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.