mars-project / mars-project/mars
Calculate chunk_size for all input tensors given a tensor graph
- Dominant language
- Python
- Stars
- 2.7k
- Forks
- 325
- PR merge metrics
- No merged PRs in 30d
Description
**Is your feature request related to a problem? Please describe.**
As we all known, `chunk_size` could be a key factor to effect the performance.
Currently, we can automatically decide the `chunk_size` by given a chunk's physical occupation. However, this is not enough. For example, if I want to calculate a tensor A with shape (5million, 256) dot the other tensor B with shape (256, 5million), if assume the input chunk's physical size is 128M, the A's `chunk_size` would be (65536, 256), and B's `chunk_size` would be (256, 65536). That's cool, but when a chunk of A dot chunk of B, the chunk's shape would be (65536, 65536) which would occupy 32G memory. That's wild.
Consider another situation, assume we have a tensor A with dtype int32, and B with dtype int64, A's `chunk_size` could be (5792, 5792), and B's `chunk_size` could be (4096, 4096), if we do `A + B`, obviously, because the chunk shapes are not aligned, we need to rechunk A or B to make sure both chunks have the same shape. But if at first we can conclude that A's `chunk_size` should be identical to B's `chunk_size`, things could be much simpler.
**Describe the solution you'd like**
A final status could be that, given the tensor graph, we can conclude ideal `chunk_size` for all input tensors so that the entire computation cost is lowest. But as a start, I think we can just achieve two goals just like the examples above:
1. We should make sure every chunk cannot exceed a given physical occupation.
2. We can align the `chunk_size` on some axes if they need to be identical due to the calculation.
Contributor guide
Assessment
This issue has not been assessed yet.