Lightning-AI / Lightning-AI/lightning-thunder

Using `symbolic values` leads to unfused program.

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

Nobody has claimed this yet.

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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.