Skip to content

Arm backend: Support high rank index tensor#17951

Open
YufengShi-dudu wants to merge 3 commits intopytorch:mainfrom
YufengShi-dudu:support-high-rank-index-tensor
Open

Arm backend: Support high rank index tensor#17951
YufengShi-dudu wants to merge 3 commits intopytorch:mainfrom
YufengShi-dudu:support-high-rank-index-tensor

Conversation

@YufengShi-dudu
Copy link
Collaborator

@YufengShi-dudu YufengShi-dudu commented Mar 6, 2026

  • Drop the rank>=4 rejection from index.Tensor TOSA support checks
  • Add rank-4 and rank-5 index tensor test cases
  • Note: rank>=4 support has been covered by the index.Tensor
    refactor, removing the need for special handling in
    ToTosaMemoryFormatPass

Change-Id: Ief40942a94040c02e54c7f276eecd660d571e46d

Arm backend: Refactor aten.index.tensor suppport

  • Decompose edge.aten.index.Tensor via backend tosa.GATHER + shape ops
  • Remove the index.Tensor node visitor
  • Reorder AccumulateIndexPutPass (generates index.Tensor) and
    DecomposeSliceScatterPass (may generate index_put) before
    DecomposeIndexTensorToGatherPass
  • Correct index_put tests to not test int inputs under FP-only profile

Change-Id: I5cfc7c110f0074463043ef1cb61165cc784a4db2

cc @digantdesai @SS-JIA @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell

YufengShi-dudu and others added 2 commits March 5, 2026 14:21
- Decompose edge.aten.index.Tensor via backend tosa.GATHER + shape ops

- Remove the index.Tensor node visitor

- Reorder AccumulateIndexPutPass (generates index.Tensor) and
  DecomposeSliceScatterPass (may generate index_put) before
  DecomposeIndexTensorToGatherPass

- Correct index_put tests to not test int inputs under FP-only profile

Change-Id: I5cfc7c110f0074463043ef1cb61165cc784a4db2
Signed-off-by: Yufeng Shi <yufeng.shi@arm.com>
Co-authored-by: Erik Lundell <erik.lundell@arm.com>
  - Drop the rank>=4 rejection from index.Tensor TOSA support checks
  - Add rank-4 and rank-5 index tensor test cases
  - Note: rank>=4 support has been covered by the index.Tensor
    refactor, removing the need for special handling in
    ToTosaMemoryFormatPass

Change-Id: Ief40942a94040c02e54c7f276eecd660d571e46d
Signed-off-by: Yufeng Shi <yufeng.shi@arm.com>
@YufengShi-dudu YufengShi-dudu requested review from SS-JIA and zingo March 6, 2026 11:18
@YufengShi-dudu YufengShi-dudu added partner: arm For backend delegation, kernels, demo, etc. from the 3rd-party partner, Arm ciflow/trunk release notes: arm Changes to the ARM backend delegate labels Mar 6, 2026
@pytorch-bot
Copy link

pytorch-bot bot commented Mar 6, 2026

🔗 Helpful Links

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

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

❌ 8 New Failures, 1 Cancelled Job, 3 Unrelated Failures

As of commit fc32245 with merge base 6b5283c (image):

NEW FAILURES - The following jobs have failed:

CANCELLED JOB - The following job was cancelled. Please retry:

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

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

@meta-cla meta-cla bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Mar 6, 2026
@YufengShi-dudu
Copy link
Collaborator Author

Hi @SS-JIA this PR add and delete some files :)

@YufengShi-dudu
Copy link
Collaborator Author

cortex-m failures are unrelated

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ciflow/trunk CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. partner: arm For backend delegation, kernels, demo, etc. from the 3rd-party partner, Arm release notes: arm Changes to the ARM backend delegate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant