pytorch / pytorch/executorch

Replace remaining uses of linearize_access_indexes for broadcasting with BroadcastIndexesRange

Open
#8,965 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

module: kernels triaged
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_index_put.cpp#L145.(I think?)

https://github.com/pytorch/executorch/blob/95f779ae2120d94e20bb95ae6af45da76ce3ff52/kernels/portable/cpu/op_split_with_sizes_copy.cpp#L129

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

Open the contributing guide

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.

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.