Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Implement GatherND operator #155

Merged
merged 3 commits into from
May 6, 2024
Merged

Implement GatherND operator #155

merged 3 commits into from
May 6, 2024

Conversation

robertknight
Copy link
Owner

@robertknight robertknight commented May 6, 2024

Add an unoptimized GatherND implementation, along with some new TensorBase methods to support the implementation.

This is one of three changes needed to run Piper TTS models, the others being support for 1D ConvTranspose with padding, and bug fixes for ScatterND.

TODO:

  • Review code for obvious easy performance improvements

Add a variant of the `inner_iter` iterators for the case where the number of
inner dimensions is determined at runtime.
By ensuring the output and `indices` are contiguous, we can reduce dynamic-rank
`TensorView` creation in the inner loop and operate directly on slices. This
reduces time in GatherND ops in the Piper TTS model by ~50%, but this is only
doing fairly small gather ops.
@robertknight robertknight merged commit 90c7cf4 into main May 6, 2024
2 checks passed
@robertknight robertknight deleted the gathernd-op branch May 6, 2024 21:03
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant