Lightning-AI / Lightning-AI/lightning-thunder
Modeling of shape queries
@jjsjann123 is already working on this.
Since Sep 10, 2024.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
## 🚀 Feature
Modeling of accessing shape attribute of Tensor/TensorProxy has raising up to discussion in separate PRs & discussions.
We are trying to debate on whether we would want to lift shape inference logic into prologue trace in Thunder in general.
i.e. For a program
```
def foo(...):
# assume we computed a tensor `t0`
t1 = t0.reshape(t0.size(0), -1)
```
If we leave all shape logic in the computation, it should be simplified as below:
```
# produce t0 from some earlier trace
(i0, i1) = prims.shape(t0)
i2 = prims.mul(1, i0)
i3 = prims.mul(i2, i1) # this is t0.numel
i4 = prims.mul(1, i0)
i5 = prims.div(i3, i4) # this is the simplified logic in clang.reshape with `-1` in the entry
t1 = prims.reshape(t0, (i0, i5))
```
One alternative is to lift all shape logic in the prologue trace, so we'll have
```
@prologue trace:
foo(...):
# all the shape logic to compute i0 & i1 from original input
i2 = prims.mul(1, i0)
i3 = prims.mul(i2, i1) # this is t0.numel
i4 = prims.mul(1, i0)
i5 = prims.div(i3, i4) # this is the simplified logic in clang.reshape with `-1` in the entry
return (..., i0, i5, ...) # NOTE: we are not necessarily passing i0 / i5, it could be any equivalent symbols.
@compute trace:
foo(..., i0, i5, ...):
# compute t0
t1 = prims.reshape(t0, (i0, i5)) # here we do not seeing that i0 is equivalent to t0.shape[0]
```
I think the first version where we see how the shape operation is defined in the operation would simplify generated kernel, since we do not need to resolve/validate reshape concretization.
- [ ] Follow up with codegen example on the impact of shape operation vs opaque scalar reshape.
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.