Skip to content

Introduce int32 index_fill and index_copy indices#142090

Closed
rpsilva-aws wants to merge 1 commit intopytorch:mainfrom
rpsilva-aws:rpsilva_pt_int32
Closed

Introduce int32 index_fill and index_copy indices#142090
rpsilva-aws wants to merge 1 commit intopytorch:mainfrom
rpsilva-aws:rpsilva_pt_int32

Conversation

@rpsilva-aws
Copy link
Copy Markdown
Contributor

@rpsilva-aws rpsilva-aws commented Dec 4, 2024

This PR extends index_fill and index_copy operations to support int32 indices in addition to the existing int64 support:

  1. Memory Efficiency and potential performance for handling a large number of indices, particularly when needing to transfer to accelerator backends. In some cases, the compiler may not support 64-bit integers, and in some cases, may cause conflicting casts when used with TorchXLA. I do not see the immediate need to propagate these to the dim input, since it is tied to the APIs and is relatively negligible.
  2. Framework interoperability: As mentioned above, this gives more flexibility when working with TorchXLA, since some operations require the same type of physical raw representations for the tensors as the XLA tensors. In some cases, for Neuron, XLA generates S32 types which are not compatible with some operations (e.g. casts) when needing to convert across tensors.
  3. Consistency with other APIs, such as index_add and index_select.

cc @jgong5 @mingfeima @XiaobingSuper @sanchitintel @ashokei @jingxu10

@pytorch-bot pytorch-bot Bot added the module: cpu CPU specific problem (e.g., perf, algorithm) label Dec 4, 2024
@pytorch-bot
Copy link
Copy Markdown

pytorch-bot Bot commented Dec 4, 2024

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/142090

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit e1a4376 with merge base 0318589 (image):
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@linux-foundation-easycla
Copy link
Copy Markdown

linux-foundation-easycla Bot commented Dec 4, 2024

CLA Not Signed

@rpsilva-aws
Copy link
Copy Markdown
Contributor Author

@pytorchbot label "release notes: cpu"

@pytorch-bot
Copy link
Copy Markdown

pytorch-bot Bot commented Dec 5, 2024

Didn't find following labels among repository labels: release notes: cpu

@rpsilva-aws
Copy link
Copy Markdown
Contributor Author

@pytorchbot label "release notes: cuda"

@pytorch-bot pytorch-bot Bot added the release notes: cuda release notes category label Dec 5, 2024
@rpsilva-aws rpsilva-aws changed the title Extend IndexFill and IndexCopy for int32 Introduce int32 index_fill and index_copy indices Dec 5, 2024
@albanD albanD added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Dec 5, 2024
@rpsilva-aws
Copy link
Copy Markdown
Contributor Author

Recreated with signed CLA: #142160

@rpsilva-aws rpsilva-aws closed this Dec 5, 2024
@rpsilva-aws rpsilva-aws deleted the rpsilva_pt_int32 branch December 5, 2024 19:27
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

module: cpu CPU specific problem (e.g., perf, algorithm) open source release notes: cuda release notes category triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants