Faster general take

[blchu]

Description

Speeding up the take operator for the general case (when axis != 0). Speeds should now be comparable to when axis = 0 on GPU, as opposed to being multiple times slower (depending on the shape of the array). CPU also receives a speedup, but not necessarily as fast as axis = 0.

Checklist

Essentials

• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage:
• Unit tests are added for small changes to verify correctness (e.g. adding a new operator)
• Nightly tests are added for complicated/long-running ones (e.g. changing distributed kvstore)
• Build tests will be added for build configuration changes (e.g. adding a new build option with NCCL)
• To the my best knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change

Changes

• Modified Take kernel to run faster

Comments

• Interesting edge cases to note: The Take kernel now assumes that when the take operator is called with axis=0, TakeGPU/TakeCPU will be used instead, and thus is now incompatible with axis=0
• Requesting @ptrendx to review

[ptrendx]

Few questions:

• Can we rename the TakeCPU/TakeGPU to something like TakeZeroAxisCPU/TakeZeroAxisGPU and Take to TakeNonzeroAxis?
• Are there any other functions that use Take that could be affected by it not supporting axis=0 anymore?
• Also, could you explain what makes your version faster? Is it just replacing taking values from array with actual values? If so, I wonder whether the problem does not just come from the fact that the Shape<10> is passed to the function by value instead of by const reference const Shape<10>&?

[blchu]

I've done a check and from what I've found only the take op calls that exact function, so I believe that isn't a problem. In any case TakeCPU/TakeGPU gives the same results as Take when axis=0, so if now or in the future there are other functions that require Take along axis=0, that can be replaced with TakeCPU/TakeGPU.

The speedup when using actual values is significantly more than with const reference, nearly the same as opposed to several times slower still. Regardless, the function itself only uses exactly 3 of the 20 elements of in_strides and out_strides, and so replacing the arrays with those three specific values may be better anyway. Would there be a reason to keep the arrays?

Additionally, while removal of the ternary operators that check for the special case axis==0 does also give some speedup, it is very small (+1-2% speed increase) compared to the other changes. So if there is a reason to keep that option, the function should still see nearly the same speedup.

[ptrendx]

Ok, then please rename Take to TakeNonzeroAxis to show the limitations and we will be good to go.

[blchu]

Change has been made.