Skip to content

[shamalgs] Add a tiled merge sort for sort_by_keys - #2233

Draft
tdavidcl wants to merge 2 commits into
Shamrock-code:mainfrom
tdavidcl:claude/tiled-merge-sort-gpu-r8cw0d
Draft

tdavidcl wants to merge 2 commits into
Shamrock-code:mainfrom
tdavidcl:claude/tiled-merge-sort-gpu-r8cw0d

Conversation

@tdavidcl

Copy link
Copy Markdown
Member

No description provided.

sort_by_keys only had host side implementations, so every call round
tripped the buffers through host memory, and the only device sort in the
tree, sort_by_key_pow2_len, requires a power of two length.

Add a device tiled merge sort: every work item first sorts a tile of Vt
consecutive elements with a fully unrolled odd-even transpose network
held in thread private arrays, then adjacent sorted runs are merged
pairwise, doubling the run length each round. A merge round is one
kernel where a work item owns a fixed size chunk of the output and finds
its own input sub-ranges with a Merge Path co-rank search. The total
work is O(n log n), and the length does not need to be a power of two.

A trailing partial tile is padded up to Vt with shambase::get_max. Since
the network only swaps on a strictly smaller key, the padding never
overtakes a real key comparing equal to the sentinel.

 - shambase: odd_even_transpose_sort_by_key, the key/value twin of the
   existing keys only OddEvenTransposeSortT
 - shamalgs: co_rank, a device and host compatible Merge Path partition
   search, alongside lower_bound and upper_bound
 - shamalgs: tiled_merge_sort.hpp for the device implementation, and
   tiled_merge_sort_host.hpp for the host serial reference, kept free of
   any SYCL dependency so the algorithm can be checked without a device
 - sort_by_keys: new tiled_merge_sort and tiled_merge_host_serial
   implementation variants, the default stays std_sort

Assisted-by: Claude Code
@coderabbitai

coderabbitai Bot commented Aug 30, 2026

Copy link
Copy Markdown

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

  • Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@github-actions

Copy link
Copy Markdown
Contributor

Thanks @tdavidcl for opening this PR!

You can do multiple things directly here:
1 - Comment pre-commit.ci run to run pre-commit checks.
2 - Comment pre-commit.ci autofix to apply fixes.
3 - Add label autofix.ci to fix authorship & pre-commit for every commit made.
4 - Add label full-ci to run the full test suite (default is light CI; full CI also runs on Mergify merge-queue branches).
5 - Add label profile-build to run the compile-time build profile job even in light CI.
6 - Add label trigger-ci to create an empty commit to trigger the CI.

Once the workflow completes a message will appear displaying informations related to the run.

Also the PR gets automatically reviewed by gemini, you can:
1 - Comment /gemini review to trigger a review
2 - Comment /gemini summary for a summary
3 - Tag it using @gemini-code-assist either in the PR or in review comments on files

@github-actions

Copy link
Copy Markdown
Contributor

Workflow report

workflow report corresponding to commit 2bc3669
Commiter email is timothee.davidcleris@proton.me

Light CI is enabled (the default for pull requests). This will only run the basic tests and not the full tests.
Full CI runs if the full-ci label is set, or automatically on Mergify merge-queue branches (mergify/merge-queue/*).
The merge gate job "on PR / all" is skipped in this case. Queue entry uses "on PR / all_light"; full CI runs in the merge queue.

Pre-commit check report

Pre-commit check: ✅

trim trailing whitespace.................................................Passed
fix end of files.........................................................Passed
check for merge conflicts................................................Passed
check that executables have shebangs.....................................Passed
check that scripts with shebangs are executable..........................Passed
check for added large files..............................................Passed
check for case conflicts.................................................Passed
check for broken symlinks................................................Passed
check yaml...............................................................Passed
detect private key.......................................................Passed
No-tabs checker..........................................................Passed
Tabs remover.............................................................Passed
cmake-format.............................................................Passed
Validate GitHub Workflows................................................Passed
clang-format.............................................................Passed
ruff check...............................................................Passed
ruff format..............................................................Passed
Check doxygen headers....................................................Passed
Check license headers....................................................Passed
Check #pragma once.......................................................Passed
Check SYCL #include......................................................Passed
No ssh in git submodules remote..........................................Passed
No UTF-8 in files (except for authors)...................................Passed

Test pipeline can run.

Clang-tidy diff report


455 warnings generated.
Suppressed 456 warnings (455 in non-user code, 1 NOLINT).
Use -header-filter=.* to display errors from all non-system headers. Use -system-headers to display errors from system headers as well.

Doxygen diff with main

Removed warnings : 8
New warnings : 14
Warnings count : 8244 → 8250 (0.1%)

Detailed changes :
+ src/shamalgs/src/primitives/sort_by_keys.cpp:126: warning: Member variant_type_name (variable) of struct shamalgs::primitives::impl::BatcherOddEvenHostSerial is not documented.
+ src/shamalgs/src/primitives/sort_by_keys.cpp:131: warning: Member variant_type_name (variable) of struct shamalgs::primitives::impl::BatcherOddEven is not documented.
+ src/shamalgs/src/primitives/sort_by_keys.cpp:136: warning: Member variant_type_name (variable) of struct shamalgs::primitives::impl::TiledMergeSort is not documented.
+ src/shamalgs/src/primitives/sort_by_keys.cpp:142: warning: Member variant_type_name (variable) of struct shamalgs::primitives::impl::TiledMergeHostSerial is not documented.
+ src/shamalgs/src/primitives/sort_by_keys.cpp:151: warning: Member sort_by_keys_impl (variable) of namespace shamalgs::primitives::impl is not documented.
- src/shamalgs/src/primitives/sort_by_keys.cpp:155: warning: Member sort_by_keys(sham::DeviceBuffer< u32 > &buf_key, sham::DeviceBuffer< u32 > &buf_values, u32 len) (function) of namespace shamalgs::primitives is not documented.
- src/shamalgs/src/primitives/sort_by_keys.cpp:158: warning: Member sort_by_keys(sham::DeviceBuffer< u64 > &buf_key, sham::DeviceBuffer< u32 > &buf_values, u32 len) (function) of namespace shamalgs::primitives is not documented.
- src/shamalgs/src/primitives/sort_by_keys.cpp:161: warning: Member sort_by_keys(sham::DeviceBuffer< f64 > &buf_key, sham::DeviceBuffer< f64 > &buf_values, u32 len) (function) of namespace shamalgs::primitives is not documented.
- src/shamalgs/src/primitives/sort_by_keys.cpp:164: warning: Member sort_by_keys(sham::DeviceBuffer< f32 > &buf_key, sham::DeviceBuffer< f32 > &buf_values, u32 len) (function) of namespace shamalgs::primitives is not documented.
+ src/shamalgs/src/primitives/sort_by_keys.cpp:213: warning: Member sort_by_keys(sham::DeviceBuffer< u32 > &buf_key, sham::DeviceBuffer< u32 > &buf_values, u32 len) (function) of namespace shamalgs::primitives is not documented.
+ src/shamalgs/src/primitives/sort_by_keys.cpp:216: warning: Member sort_by_keys(sham::DeviceBuffer< u64 > &buf_key, sham::DeviceBuffer< u32 > &buf_values, u32 len) (function) of namespace shamalgs::primitives is not documented.
+ src/shamalgs/src/primitives/sort_by_keys.cpp:219: warning: Member sort_by_keys(sham::DeviceBuffer< f64 > &buf_key, sham::DeviceBuffer< f64 > &buf_values, u32 len) (function) of namespace shamalgs::primitives is not documented.
+ src/shamalgs/src/primitives/sort_by_keys.cpp:222: warning: Member sort_by_keys(sham::DeviceBuffer< f32 > &buf_key, sham::DeviceBuffer< f32 > &buf_values, u32 len) (function) of namespace shamalgs::primitives is not documented.
- src/shamalgs/src/primitives/sort_by_keys.cpp:90: warning: Member variant_type_name (variable) of struct shamalgs::primitives::impl::BatcherOddEvenHostSerial is not documented.
- src/shamalgs/src/primitives/sort_by_keys.cpp:95: warning: Member variant_type_name (variable) of struct shamalgs::primitives::impl::BatcherOddEven is not documented.
- src/shamalgs/src/primitives/sort_by_keys.cpp:99: warning: Member sort_by_keys_impl (variable) of namespace shamalgs::primitives::impl is not documented.
+ src/shambase/include/shambase/alg_primitives.hpp:67: warning: Compound shambase::OddEvenTransposeSortByKeyT is not documented.
+ src/shambase/include/shambase/alg_primitives.hpp:69: warning: Member Sort(Tkey *keys, Tval *vals, Comp comp) (function) of struct shambase::OddEvenTransposeSortByKeyT is not documented.
+ src/shambase/include/shambase/alg_primitives.hpp:81: warning: Compound shambase::OddEvenTransposeSortByKeyT< I, I > is not documented.
+ src/shambase/include/shambase/alg_primitives.hpp:83: warning: Member Sort(Tkey *keys, Tval *vals, Comp comp) (function) of struct shambase::OddEvenTransposeSortByKeyT< I, I > is not documented.

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