Replace remaining uses of linearize_access_indexes for broadcasting with BroadcastIndexesRange
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 5k
- Forks
- 1.2k
- Avg merge
- 2d 10h
- Merged PRs (30d)
- 581
Description
BroadcastIndexesRange is more efficient than linearize_access_indexes-based loops. We have several remaining places that use linearize_access_indexes, and what they have in common is that they're not looping over an existing Tensor's .sizes() and .strides(). To make these work nicely, we need a way for BroadcastIndexesRange to play nicely with .sizes() and .strides() ArrayRefs that aren't attached to a Tensor.
Specific usage sites that can be improved:
https://github.com/pytorch/executorch/blob/95f779ae2120d94e20bb95ae6af45da76ce3ff52/kernels/portable/cpu/op_cdist_forward.cpp#L79
https://github.com/pytorch/executorch/blob/95f779ae2120d94e20bb95ae6af45da76ce3ff52/kernels/portable/cpu/op_masked_select.cpp#L125 (might be straightforwardly fixable right away with no extra work)
cc @larryliu0820 @manuelcandales
Contributor guide
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.
Research direction
Start by reading BroadcastIndexesRange and the cited usage sites: kernels/portable/cpu/op_cdist_forward.cpp, op_index_put.cpp, op_split_with_sizes_copy.cpp, and op_masked_select.cpp. Determine how its iteration can accept standalone sizes and strides ArrayRefs, then replace the remaining linearize_access_indexes-based broadcasting loops and validate the affected kernels.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- machine-learning, performance
- Issue type
- Refactor
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100