Repository navigation
perf(nx): Copy a host gather's dense tail by runs, not elements - #237
Closed
ghennequin wants to merge 1 commit into
Closed
ghennequin wants to merge 1 commit into
ghennequin wants to merge 1 commit into
Conversation
nx.cpu's gather copied one element at a time through a generic ND loop — an unravel and two dot products per element — for every axis and rank but axis 0 of a rank 2 tensor, which has a row path of its own. A take along axis 1 of a rank 3 tensor is the shape a minibatch of the lorenz example has, [8; 1023; 256], and the loop moved its 2.1 M elements at about 28 ns each. When the dims after the axis are one dense run in both the data and the output, and the index is constant along it — a broadcast [take ~axis] index is — one run is now copied whole, at one index load for the run, and the outer coordinates are walked by an odometer, so a run costs no unravel. take_along_axis, whose indices vary within the run, keeps the generic path. The lorenz example's minibatch falls from 58 ms to 2.4 ms; an isolated take along axis 1 of [8; 1023; 256] falls from 58 ms to 2.3 ms, and of [8; 1024; 4] from 0.96 ms to 0.085 ms. The nx suites pass, the indexing property suite (98 tests, every dtype and axis) across seeds among them.
Contributor
|
Merged manually, thanks! :) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
nx.cpu's gather copied one element at a time through a generic ND loop — an unravel and two dot products per element — for every axis and rank but axis 0 of a rank 2 tensor, which has a row path of its own. A take along axis 1 of a rank 3 tensor is the shape a minibatch of the lorenz example has, [8; 1023; 256], and the loop moved its 2.1 M elements at about 28 ns each.
When the dims after the axis are one dense run in both the data and the output, and the index is constant along it — a broadcast [take ~axis] index is — one run is now copied whole, at one index load for the run, and the outer coordinates are walked by an odometer, so a run costs no unravel. take_along_axis, whose indices vary within the run, keeps the generic path.
The lorenz example's minibatch falls from 58 ms to 2.4 ms; an isolated take along axis 1 of [8; 1023; 256] falls from 58 ms to 2.3 ms, and of [8; 1024; 4] from 0.96 ms to 0.085 ms. The nx suites pass, the indexing property suite (98 tests, every dtype and axis) across seeds among them.