Lightning-AI / Lightning-AI/lightning-thunder
Using `symbolic values` leads to unfused program.
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
While using the default option leads to a fused program
Repro
```python
import torch
import thunder
def clamp_position(seq_lens):
return torch.clamp((seq_lens - 1), min=0).to(torch.int64)
seq_lens = torch.randint(10, 100, size=(1024,), device="cuda")
clamp_position(seq_lens)
tjit = thunder.jit(clamp_position)
tjit(seq_lens)
print(thunder.last_traces(tjit)[-1])
# Constructed by Unwrap the actual return value
# import torch
# from thunder.executors.torchex import no_autocast
# @torch.no_grad()
# @no_autocast
# def computation(seq_lens):
# # seq_lens: "cuda:0 i64[1024]"
# [t5] = nvFusion0(seq_lens)
# # t0 = prims.sub(seq_lens, 1) # t0: "cuda:0 i64[1024]"
# # t1 = prims.ne(t0, t0) # t1: "cuda:0 b8[1024]"
# # t2 = prims.gt(t0, 0) # t2: "cuda:0 b8[1024]"
# # t3 = prims.where(t2, t0, 0) # t3: "cuda:0 i64[1024]"
# # t4 = prims.where(t1, t0, t3) # t4: "cuda:0 i64[1024]"
# # t5 = prims.shallow_copy(t4) # t5: "cuda:0 i64[1024]"
# return (t5,)
tjit = thunder.jit(clamp_position, cache="symbolic values")
tjit(seq_lens)
print(thunder.last_traces(tjit)[-1])
# Constructed by Unwrap the actual return value
# import torch
# from torch import Tensor
# from thunder.executors.torchex import no_autocast
# @torch.no_grad()
# @no_autocast
# def computation(seq_lens):
# # seq_lens: "cuda:0 i64[[IntegerProxy name=i0, value=1024, static=CONSTRAINT.CONSTRAINABLE]]"
# # /opt/pytorch/lightning-thunder/test.py:5: return torch.clamp((seq_lens - 1), min=0).to(torch.int64)
# t38 = torch.sub(seq_lens, 1, alpha=1) # t38: "cuda:0 i64[[IntegerProxy name=i0, value=1024, static=CONSTRAINT.CONSTRAINABLE]]"
# # t38 = ltorch.sub(seq_lens, 1, alpha=1) # t38: "cuda:0 i64[[IntegerProxy name=i0, value=1024, static=CONSTRAINT.CONSTRAINABLE]]"
# # (i0,) = prims.shape(seq_lens)
# # (i0,) = prims.shape(seq_lens)
# # i31 = prims.eq(i0, 1) # i31: "bool False"
# # i37 = prims.eq(i0, i0) # i37: "bool True"
# # (i0,) = prims.shape(seq_lens)
# # (i0,) = prims.shape(seq_lens)
# # t38 = prims.sub(seq_lens, 1) # t38: "cuda:0 i64[[IntegerProxy name=i0, value=1024, static=CONSTRAINT.CONSTRAINABLE]]"
# t76 = torch.clamp(t38, 0, None) # t76: "cuda:0 i64[[IntegerProxy name=i0, value=1024, static=CONSTRAINT.CONSTRAINABLE]]"
# # t76 = ltorch.clamp(t38, 0, None) # t76: "cuda:0 i64[[IntegerProxy name=i0, value=1024, static=CONSTRAINT.CONSTRAINABLE]]"
# # t58 = ltorch.ne(t38, t38) # t58: "cuda:0 b8[[IntegerProxy name=i0, value=1024, static=CONSTRAINT.CONSTRAINABLE]]"
# # (i0,) = prims.shape(t38)
# # (i0,) = prims.shape(t38)
# # (i0,) = prims.shape(t38)
# # (i0,) = prims.shape(t38)
# # i40 = prims.eq(i0, 1) # i40: "bool False"
# # i46 = prims.eq(i0, i0) # i46: "bool True"
# # i55 = prims.eq(i0, 1) # i55: "bool False"
# # i56 = prims.eq(i0, 1) # i56: "bool False"
# # i57 = prims.eq(i0, i0) # i57: "bool True"
# # (i0,) = prims.shape(t38)
# # (i0,) = prims.shape(t38)
# # (i0,) = prims.shape(t38)
# # (i0,) = prims.shape(t38)
# # t58 = prims.ne(t38, t38) # t58: "cuda:0 b8[[IntegerProxy name=i0, value=1024, static=CONSTRAINT.CONSTRAINABLE]]"
# # t61 = ltorch.gt(t38, 0) # t61: "cuda:0 b8[[IntegerProxy name=i0, value=1024, static=CONSTRAINT.CONSTRAINABLE]]"
# # (i0,) = prims.shape(t38)
# # (i0,) = prims.shape(t38)
# # i59 = prims.eq(i0, 1) # i59: "bool False"
# # i60 = prims.eq(i0, i0) # i60: "bool True"
# # (i0,) = prims.shape(t38)
# # (i0,) = prims.shape(t38)
# # t61 = prims.gt(t38, 0) # t61: "cuda:0 b8[[IntegerProxy name=i0, value=1024, static=CONSTRAINT.CONSTRAINABLE]]"
# # t67 = ltorch.where(t61, t38, 0) # t67: "cuda:0 i64[[IntegerProxy name=i0, value=1024, static=CONSTRAINT.CONSTRAINABLE]]"
# # (i0,) = prims.shape(t61)
# # (i0,) = prims.shape(t61)
# # (i0,) = prims.shape(t38)
# # (i0,) = prims.shape(t38)
# # i62 = prims.eq(i0, 1) # i62: "bool False"
# # i63 = prims.eq(i0, i0) # i63: "bool True"
# # i64 = prims.eq(i0, 1) # i64: "bool False"
# # i65 = prims.eq(i0, 1) # i65: "bool False"
# # i66 = prims.eq(i0, i0) # i66: "bool True"
# # (i0,) = prims.shape(t61)
# # (i0,) = prims.shape(t61)
# # (i0,) = prims.shape(t38)
# # (i0,) = prims.shape(t38)
# # t67 = prims.where(t61, t38, 0) # t67: "cuda:0 i64[[IntegerProxy name=i0, value=1024, static=CONSTRAINT.CONSTRAINABLE]]"
# # t76 = ltorch.where(t58, t38, t67) # t76: "cuda:0 i64[[IntegerProxy name=i0, value=1024, static=CONSTRAINT.CONSTRAINABLE]]"
# # (i0,) = prims.shape(t58)
# # (i0,) = prims.shape(t58)
# # (i0,) = prims.shape(t38)
# # (i0,) = prims.shape(t38)
# # (i0,) = prims.shape(t67)
# # (i0,) = prims.shape(t67)
# # i68 = prims.eq(i0, 1) # i68: "bool False"
# # i69 = prims.eq(i0, i0) # i69: "bool True"
# # i70 = prims.eq(i0, 1) # i70: "bool False"
# # i71 = prims.eq(i0, 1) # i71: "bool False"
# # i72 = prims.eq(i0, i0) # i72: "bool True"
# # i73 = prims.eq(i0, 1) # i73: "bool False"
# # i74 = prims.eq(i0, 1) # i74: "bool False"
# # i75 = prims.eq(i0, i0) # i75: "bool True"
# # (i0,) = prims.shape(t58)
# # (i0,) = prims.shape(t58)
# # (i0,) = prims.shape(t38)
# # (i0,) = prims.shape(t38)
# # (i0,) = prims.shape(t67)
# # (i0,) = prims.shape(t67)
# # t76 = prims.where(t58, t38, t67) # t76: "cuda:0 i64[[IntegerProxy name=i0, value=1024, static=CONSTRAINT.CONSTRAINABLE]]"
# del t38
# # /opt/pytorch/lightning-thunder/test.py:5: return torch.clamp((seq_lens - 1), min=0).to(torch.int64)
# t77 = Tensor.to(t76, copy=False, dtype=torch.int64) # t77: "cuda:0 i64[[IntegerProxy name=i0, value=1024, static=CONSTRAINT.CONSTRAINABLE]]"
# # t77 = ltorch.to(t76, None, None, device=None, dtype=torch.int64, copy=False, memory_format=None) # t77: "cuda:0 i64[[IntegerProxy name=i0, value=1024, static=CONSTRAINT.CONSTRAINABLE]]"
# # t77 = prims.shallow_copy(t76) # t77: "cuda:0 i64[[IntegerProxy name=i0, value=1024, static=CONSTRAINT.CONSTRAINABLE]]"
# del t76
# return (t77,)
```
cc: @beverlylytle
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 with the Python repro using clamp_position, thunder.jit, and cache="symbolic values", then compare the traces from thunder.last_traces(tjit)[-1]. Investigate why the symbolic-values trace remains unfused while the default trace uses nvFusion0; done means the symbolic-values case produces an equivalent fused program.
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
- 35/100