Lightning-AI / Lightning-AI/lightning-thunder
Make prims.shape fusible with nvFuser executor
@beverlylytle is already working on this.
Since Nov 12, 2025.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
## 🚀 Feature
Currently, `prims.shape` is not supported by the nvFuser executor leading to simple torch.add operation to not be fused with nvFuser if it involved broadcasting and it's not a simple translation to `prims.add`.
### Motivation
We need `prims.shape` to occur in the nvFuser fusible regions as a lot of operations have broadcasting behavior and for dynamic shapes efficient execution shape retrieval must be in the fused region and not outside.
Test program:
```python
import torch, thunder
from functools import partial
from thunder import prims
a = torch.randn(1, 1024, device="cuda", dtype=torch.float32)
b = torch.randn(2, 1024, device="cuda", dtype=torch.float32)
@partial(thunder.jit, cache="symbolic values")
def f(a, b):
(i0, i1) = prims.shape(a)
(i2, i3) = prims.shape(b)
a = prims.broadcast_in_dim(a, (i2, i1), (0, 1))
return prims.add(a, b)
f(a, b)
print(f._lc_cs.last_traces[-1])
```
Prints:
```py
def computation(a, b):
# a: "cuda:0 f32[[IntegerProxy name=i0, value=1, static=CONSTRAINT.CONSTRAINABLE], [IntegerProxy name=i1, value=1024, static=CONSTRAINT.CONSTRAINABLE]]"
# b: "cuda:0 f32[[IntegerProxy name=i2, value=2, static=CONSTRAINT.CONSTRAINABLE], [IntegerProxy name=i3, value=1024, static=CONSTRAINT.CONSTRAINABLE]]"
# /tmp/ipython-input-1588267087.py:10: (i0, i1) = prims.shape(a)
(i0, i1) = shape(a)
del i0
# /tmp/ipython-input-1588267087.py:11: (i2, i3) = prims.shape(b)
(i2, i3) = shape(b)
del i3
[t9] = nvFusion0(a, i2, i1, b)
# t7 = prims.broadcast_in_dim(a, (i2, i1), (0, 1)) # t7: "cuda:0 f32[[IntegerProxy name=i2, value=2, static=CONSTRAINT.CONSTRAINABLE], [IntegerProxy name=i1, value=1024, static=CONSTRAINT.CONSTRAINABLE]]"
# t9 = prims.add(t7, b) # t9: "cuda:0 f32[[IntegerProxy name=i2, value=2, static=CONSTRAINT.CONSTRAINABLE], [IntegerProxy name=i1, value=1024, static=CONSTRAINT.CONSTRAINABLE]]"
del i2, i1
return (t9,)
```
cc @tfogal
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.
Assessment
This issue has not been assessed yet.