Add vectorized_sort function

[Berkant03]

Due Diligence

• General:
  • title of the PR is suitable to appear in the Release Notes
• Implementation:
  • unit tests: all split configurations tested
  • unit tests: multiple dtypes tested
  • NEW unit tests: MPS tested (1 MPI process, 1 GPU)
  • benchmarks: created for new functionality
  • benchmarks: performance improved or maintained
  • documentation updated where needed

Description

Issue/s resolved: #363

Changes proposed:

Introduce the vectorized sorting mentioned in #363, which is also a lexicographical sort
Allows to finalize #2332.

Type of change

• New feature (non-breaking change which adds functionality)

Memory requirements

A complete alltoallv requires a total of 3x the memory here: the original buffer, contiguous send buffer and the receive buffer.

One could in theory implement an own chunking algorithm to reduce this, if needed. That might make the performance worse or better. Testing is the only way to know for sure.

Custom datatypes should also be considered, but would need extensive analysis on the MPI implementation itself, since the sending rank might want to send multiple packets and might need to pack them together, creating again a copy, unless the network protocol can completely handle the non-contiguous data.

Performance

Does this change modify the behaviour of other functions? If so, which?

no

[Berkant03]

I have written some pytest tests, but would wait until #777 is merged, to write them into the TestSorting class.

[brownbaerchen]

I have written some pytest tests, but would wait until #777 is merged, to write them into the TestSorting class.

#2386 was actually already merged.
For future things like this, see the very cool new GitHub stack feature. Using this feature, you can branch out of the development branch of #2386, implement something on top, start a PR and stack it in top of #2386. Every PR only shows it's own diff and not the one from PRs below in the stack, making it easy to review. PRs below in the stack can be merged independently, or the entire stack can be merged at once. #2422 is an example where we stacked PRs.

Don't worry about this now. I just love this new feature and want to advertise it ;)

[brownbaerchen]

It's a bit too hot for me to understand everything. I marked a bunch of small things that could be a bit cleaner. I suggest you add some tests and then I'll have another look.
Since we want to extract the side effect of sorting from unique, a comparison with unique would be a useful test.

[brownbaerchen]

I failed at creating a stack of this and #2332, which depends on this. Getting the commit histories in order after the fact is not so easy :D I'll take care of fixing the merge conflicts that will arise when we merge #2332 after merging this. But on the bright side, the tests seem to pass :) Sorry about the mess.

[codecov]

Codecov Report

❌ Patch coverage is 92.72727% with 8 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.79%. Comparing base (6611beb) to head (82f0132).
⚠️ Report is 31 commits behind head on main.
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
heat/core/manipulations.py	92.72%	8 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2431      +/-   ##
==========================================
+ Coverage   91.74%   91.79%   +0.04%     
==========================================
  Files          87       87              
  Lines       14168    14278     +110     
==========================================
+ Hits        12998    13106     +108     
- Misses       1170     1172       +2     

Flag	Coverage Δ
unit	91.79% <92.72%> (+0.04%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

[Berkant03]

There was the Idea to also implement a ht.lexsort using vectorized_sort to mimic similar behaviour as to numpy's np.lexsort. Unfortunetely, this is not possible for arrays with more than two dimensions, since numpy does not reduce the keys a two dimensional array like vectorized_sort.

[Berkant03]

The other Idea was to use vectorized_sort for sorting complex types, which also would not work for arrays with more than one dimensions, since the sorting should happen on the specified axis, while ignoring other axis, which is not possible using vectorized_sort. One could use it for one dimensional data, but performance might suffer.