Lightning-AI / Lightning-AI/lightning-thunder

Modeling of shape queries

Open
#1,133 9 comments 0 reactions 1 assignee View on GitHub

@jjsjann123 is already working on this.

Since Sep 10, 2024.

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

  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.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.