From 28ed74e3f6127be585049518de99bf6039e73cee Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Timoth=C3=A9e=20David--Cl=C3=A9ris?= Date: Sun, 30 Aug 2026 13:05:53 +0200 Subject: [PATCH 1/7] prepare impl test --- .../benchmarks/run_sort_by_keys_tmp_dev.py | 141 +++++++++++++++++ .../device/details/modern_gpu_merge_sort.hpp | 146 ++++++++++++++++++ .../workitem/odd_even_transpose_sort.hpp | 52 +++++++ src/shamalgs/src/primitives/sort_by_keys.cpp | 15 +- 4 files changed, 352 insertions(+), 2 deletions(-) create mode 100644 examples/benchmarks/run_sort_by_keys_tmp_dev.py create mode 100644 src/shamalgs/include/shamalgs/primitives/device/details/modern_gpu_merge_sort.hpp create mode 100644 src/shamalgs/include/shamalgs/primitives/workitem/odd_even_transpose_sort.hpp diff --git a/examples/benchmarks/run_sort_by_keys_tmp_dev.py b/examples/benchmarks/run_sort_by_keys_tmp_dev.py new file mode 100644 index 0000000000..80e090a29f --- /dev/null +++ b/examples/benchmarks/run_sort_by_keys_tmp_dev.py @@ -0,0 +1,141 @@ +""" +sort by keys performance benchmarks +==================================== + +This example benchmarks the sort by keys (general, non power-of-2 length) performance for +the different algorithms available in Shamrock, as well as the sort_by_key_pow2_len +(power-of-2 length only) performance in a second figure. +""" + +# sphinx_gallery_multi_image = "single" + +import json +import random +import time + +import matplotlib.pyplot as plt +import numpy as np + +import shamrock + +# If we use the shamrock executable to run this script instead of the python interpreter, +# we should not initialize the system as the shamrock executable needs to handle specific MPI logic +if not shamrock.sys.is_initialized(): + shamrock.change_loglevel(1) + shamrock.sys.init("0:0") + + +# %% +# List current implementation +if not shamrock.algs.is_impl_set_sort_by_keys(): + shamrock.algs.autoselect_impl_sort_by_keys() + +current_impl = shamrock.algs.get_current_impl_sort_by_keys() + +print(current_impl) + +# %% +# List all implementations available +all_default_impls = shamrock.algs.get_default_impl_list_sort_by_keys() + +print(all_default_impls) + +# %% +# Pretty printer + + +_GREEN = "\033[32m" +_RED = "\033[31m" +_RESET = "\033[0m" + + +def _sorted_run_mask(values): + """Flag indices that belong to a non-decreasing run of length >= 2.""" + mask = [False for i in range(len(values))] + + for i,v in enumerate(values): + if i > 0: + if values[i-1] < values[i]: + mask[i] = True + else: + mask[i] = True + return mask + + +def print_key_val_table(keys, vals, labels=("i", "key", "val")): + """Pretty-print indices/keys/values as an aligned table, coloring any + sorted (non-decreasing) run of >= 2 consecutive keys in green. + + i | 0 1 2 ... + ---------------------- + key | ? ? ? ... + val | ? ? ? ... + """ + idx_label, key_label, val_label = labels + rows = [idx_label, key_label, val_label] + columns = [list(range(len(keys))), keys, vals] + + label_width = max(len(row) for row in rows) + col_width = max(len(str(v)) for col in columns for v in col) + + def cell(v, highlight=None): + text = f"{str(v):>{col_width}}" + if highlight is None: + return text + return f"{_GREEN}{text}{_RESET}" if highlight else f"{_RED}{text}{_RESET}" + + def fmt_row(label, values, mask=None): + cells = " ".join(cell(v, mask[i] if mask else None) for i, v in enumerate(values)) + return f"{label:<{label_width}} | {cells}" + + idx_line = fmt_row(idx_label, columns[0]) + key_line = fmt_row(key_label, columns[1], _sorted_run_mask(keys)) + val_line = fmt_row(val_label, columns[2]) + + print("-" * len(idx_line)) + print(idx_line) + print("-" * len(idx_line)) + print(key_line) + print(val_line) + +def to_buf(lst): + buf = shamrock.backends.DeviceBuffer_u32() + buf.resize(len(lst)) + buf.copy_from_stdvec(lst) + return buf + + +#%% +# Dataset +N = 25 +key_init = [i for i in range(N)][::-1] +val_init = [i for i in range(N)] + +print("Initial state:") +print_key_val_table(key_init, val_init) + + +#%% +# Switch impl +shamrock.algs.set_impl_sort_by_keys('{"implementation":"modern_gpu_mergesort","parameters":{}}') + +#%% +# End state +buffer_key = to_buf(key_init) +buffer_val = to_buf(val_init) + +shamrock.algs.sort_by_keys(buffer_key, buffer_val, N) + +print("End state:") +print_key_val_table(buffer_key.copy_to_stdvec(), buffer_val.copy_to_stdvec()) + +#%% +# Expected state +buffer_key = to_buf(key_init) +buffer_val = to_buf(val_init) + +shamrock.algs.autoselect_impl_sort_by_keys() +shamrock.algs.sort_by_keys(buffer_key, buffer_val, N) + +print("Expected state:") +print_key_val_table(buffer_key.copy_to_stdvec(), buffer_val.copy_to_stdvec()) diff --git a/src/shamalgs/include/shamalgs/primitives/device/details/modern_gpu_merge_sort.hpp b/src/shamalgs/include/shamalgs/primitives/device/details/modern_gpu_merge_sort.hpp new file mode 100644 index 0000000000..b33e6c518c --- /dev/null +++ b/src/shamalgs/include/shamalgs/primitives/device/details/modern_gpu_merge_sort.hpp @@ -0,0 +1,146 @@ +// -------------------------------------------------------// +// +// SHAMROCK code for hydrodynamics +// Copyright (c) 2021-2026 Timothée David--Cléris +// SPDX-License-Identifier: CeCILL Free Software License Agreement v2.1 +// Shamrock is licensed under the CeCILL 2.1 License, see LICENSE for more information +// +// -------------------------------------------------------// + +#pragma once + +/** + * @file modern_gpu_merge_sort.hpp + * @author Timothée David--Cléris (tim.shamrock@proton.me) + * @brief + */ + +#include "shambackends/DeviceBuffer.hpp" +#include +#include +#include +#include +#include +#include + +namespace shamalgs::primitives::device::details { + + namespace debug { + + inline constexpr const char *ansi_green = "\033[32m"; + inline constexpr const char *ansi_red = "\033[31m"; + inline constexpr const char *ansi_reset = "\033[0m"; + + /// Flags indices that continue an ascending run from the previous element (mirrors + /// `_sorted_run_mask` in examples/benchmarks/run_sort_by_keys_tmp_dev.py). + template + inline std::vector sorted_run_mask(const std::vector &values) { + std::vector mask(values.size(), false); + for (size_t i = 0; i < values.size(); ++i) { + mask[i] = (i == 0) || (values[i - 1] < values[i]); + } + return mask; + } + + template + inline std::string to_str(const T &v) { + std::ostringstream oss; + oss << v; + return oss.str(); + } + + } // namespace debug + + /** + * @brief Debug pretty-printer for key/value arrays, mirroring `print_key_val_table` in + * examples/benchmarks/run_sort_by_keys_tmp_dev.py. + * + * Prints the indices, the keys (highlighted green if ascending from the previous key, red + * otherwise), and the values as an aligned table. + * + * @tparam Tkey Type of the keys. + * @tparam Tval Type of the values. + * @param keys Host-side keys to print. + * @param vals Host-side values to print. + * @param idx_label Label of the index row. + * @param key_label Label of the key row. + * @param val_label Label of the value row. + */ + template + inline void print_key_val_table( + const std::vector &keys, + const std::vector &vals, + const std::string &idx_label = "i", + const std::string &key_label = "key", + const std::string &val_label = "val") { + + using namespace debug; + + size_t n = keys.size(); + + std::vector idx_str(n), key_str(n), val_str(n); + size_t col_width = 0; + for (size_t i = 0; i < n; ++i) { + idx_str[i] = to_str(i); + key_str[i] = to_str(keys[i]); + val_str[i] = to_str(vals[i]); + col_width + = std::max({col_width, idx_str[i].size(), key_str[i].size(), val_str[i].size()}); + } + + size_t label_width = std::max({idx_label.size(), key_label.size(), val_label.size()}); + + auto pad_label = [&](const std::string &label) { + std::string s = label; + s.resize(label_width, ' '); + return s; + }; + + auto pad_cell = [&](const std::string &s) { + return std::string(col_width - s.size(), ' ') + s; + }; + + auto fmt_row = [&](const std::string &label, + const std::vector &values, + const std::vector *mask) { + std::string line = pad_label(label) + " | "; + for (size_t i = 0; i < values.size(); ++i) { + if (i > 0) { + line += " "; + } + std::string cell = pad_cell(values[i]); + if (mask != nullptr) { + cell = std::string((*mask)[i] ? ansi_green : ansi_red) + cell + ansi_reset; + } + line += cell; + } + return line; + }; + + std::vector mask = sorted_run_mask(keys); + + std::string idx_line = fmt_row(idx_label, idx_str, nullptr); + std::string key_line = fmt_row(key_label, key_str, &mask); + std::string val_line = fmt_row(val_label, val_str, nullptr); + + std::cout << std::string(idx_line.size(), '-') << "\n"; + std::cout << idx_line << "\n"; + std::cout << std::string(idx_line.size(), '-') << "\n"; + std::cout << key_line << "\n"; + std::cout << val_line << "\n"; + } + + template + inline void sort_by_keys_modern_gpu_mergesort( + sham::DeviceBuffer &buf_key, sham::DeviceBuffer &buf_values, u32 len) { + std::cout << "-------------------------------------------------" << std::endl; + std::cout << "------- sort_by_keys_modern_gpu_mergesort -------" << std::endl; + std::cout << "-------------------------------------------------" << std::endl; + std::cout << "init state:" << std::endl; + print_key_val_table(buf_key.copy_to_stdvec(), buf_values.copy_to_stdvec()); + std::cout << "-------------------------------------------------" << std::endl; + std::cout << "------- sort_by_keys_modern_gpu_mergesort end -------" << std::endl; + std::cout << "-------------------------------------------------" << std::endl; + } + +} // namespace shamalgs::primitives::device::details diff --git a/src/shamalgs/include/shamalgs/primitives/workitem/odd_even_transpose_sort.hpp b/src/shamalgs/include/shamalgs/primitives/workitem/odd_even_transpose_sort.hpp new file mode 100644 index 0000000000..e82270db7a --- /dev/null +++ b/src/shamalgs/include/shamalgs/primitives/workitem/odd_even_transpose_sort.hpp @@ -0,0 +1,52 @@ +// -------------------------------------------------------// +// +// SHAMROCK code for hydrodynamics +// Copyright (c) 2021-2026 Timothée David--Cléris +// SPDX-License-Identifier: CeCILL Free Software License Agreement v2.1 +// Shamrock is licensed under the CeCILL 2.1 License, see LICENSE for more information +// +// -------------------------------------------------------// + +#pragma once + +/** + * @file odd_even_transpose_sort.hpp + * @author Timothée David--Cléris (tim.shamrock@proton.me) + * @brief Work-item local odd-even transposition sort + */ + +#include + +namespace shamalgs::primitives::workitem { + + /** + * @brief Sorts keys (and their associated values) held in per-work-item registers using an + * odd-even transposition sort. + * + * Adapted from moderngpu's OddEvenTransposeSort, see + * https://moderngpu.github.io/mergesort.html + * + * @tparam VT Number of elements per work-item. + * @tparam T Type of the keys. + * @tparam V Type of the values. + * @tparam Comp Type of the comparator. + * @param keys Pointer to the local array of VT keys to sort. + * @param values Pointer to the local array of VT values, permuted alongside the keys. + * @param comp Comparator used to order the keys. + */ + template + inline void odd_even_transpose_sort(T *keys, V *values, Comp comp) { +#pragma unroll + for (int level = 0; level < VT; ++level) { + +#pragma unroll + for (int i = 1 & level; i < VT - 1; i += 2) { + if (comp(keys[i + 1], keys[i])) { + std::swap(keys[i], keys[i + 1]); + std::swap(values[i], values[i + 1]); + } + } + } + } + +} // namespace shamalgs::primitives::workitem diff --git a/src/shamalgs/src/primitives/sort_by_keys.cpp b/src/shamalgs/src/primitives/sort_by_keys.cpp index c3907440b5..77b8c97a89 100644 --- a/src/shamalgs/src/primitives/sort_by_keys.cpp +++ b/src/shamalgs/src/primitives/sort_by_keys.cpp @@ -18,6 +18,7 @@ #include "shambase/overloaded.hpp" #include "shamalgs/ImplVariant.hpp" #include "shamalgs/details/algorithm/batcherOddEvenSort.hpp" +#include "shamalgs/primitives/device/details/modern_gpu_merge_sort.hpp" #include "shamalgs/primitives/sort_by_keys.hpp" #include "shamcomm/logs.hpp" #include @@ -95,8 +96,14 @@ namespace shamalgs::primitives { static constexpr std::string_view variant_type_name = "batcher_odd_even"; }; - shamalgs::ImplVariantGlobal - sort_by_keys_impl; + /// Copy the buffers to host, sort with Batcher's odd-even merge sort, and copy back + struct ModernGPUMergeSort { + static constexpr std::string_view variant_type_name = "modern_gpu_mergesort"; + }; + + shamalgs:: + ImplVariantGlobal + sort_by_keys_impl; /// Get list of available sort by keys implementations std::vector get_default_impl_list_sort_by_keys() { @@ -148,6 +155,10 @@ namespace shamalgs::primitives { algorithm::details::sort_by_key_batcher_odd_even( buf_key.get_dev_scheduler_ptr(), buf_key, buf_values, len); }, + [&](impl::ModernGPUMergeSort) { + primitives::device::details::sort_by_keys_modern_gpu_mergesort( + buf_key, buf_values, len); + }, }, impl::sort_by_keys_impl.get()); } From 4568bf124bb9ab5dcdaeda0377493d252ed17cb6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Timoth=C3=A9e=20David--Cl=C3=A9ris?= Date: Sun, 30 Aug 2026 18:06:10 +0200 Subject: [PATCH 2/7] make the CTA block sort work --- doc/sphinx/source/dev_doc/kernel_call.md | 88 ++++ .../benchmarks/run_sort_by_keys_tmp_dev.py | 17 +- .../device/details/modern_gpu_merge_sort.hpp | 443 +++++++++++++++++- .../workitem/odd_even_transpose_sort.hpp | 2 +- .../include/shambackends/kernel_call.hpp | 51 ++ 5 files changed, 584 insertions(+), 17 deletions(-) diff --git a/doc/sphinx/source/dev_doc/kernel_call.md b/doc/sphinx/source/dev_doc/kernel_call.md index cc611e3bf3..ea723941b8 100644 --- a/doc/sphinx/source/dev_doc/kernel_call.md +++ b/doc/sphinx/source/dev_doc/kernel_call.md @@ -39,6 +39,94 @@ sham::kernel_call( - Read buffers → `const T*`; write buffers → `T*` - Prefer `const sham::DeviceBuffer&` for pure inputs +## `sham::kernel_call_hndl` + +Same wrapper, but the functor doesn't get called per-thread directly: it takes the +thread count and the buffer pointers, and must **return** a +`[=](sycl::handler &cgh) { ... }` lambda — that's what actually gets submitted to the +queue. Use this variant when the kernel needs the `sycl::handler` itself, e.g. +`sycl::local_accessor`s, `cgh.depends_on(...)`, etc. + +```cpp +sham::kernel_call_hndl( + queue, + sham::MultiRef{/* inputs */}, + sham::MultiRef{/* outputs (in-out) */}, + n, // thread count + [](u32 n, /* input ptrs */, /* output ptrs */) { + return [=](sycl::handler &cgh) { + cgh.parallel_for(sycl::range<1>{n}, [=](sycl::item<1> item) { + u32 i = item.get_linear_id(); + // ... + }); + }; + }); +``` + +Minimal example: + +```cpp +sham::kernel_call_hndl( + q, + sham::MultiRef{buf_in}, + sham::MultiRef{buf_out}, + n, + [](u32 n, const T *in, T *out) { + return [=](sycl::handler &cgh) { + cgh.parallel_for(sycl::range<1>{n}, [=](sycl::item<1> item) { + u32 i = item.get_linear_id(); + out[i] = in[i]; + }); + }; + }); +``` + +- Functor first argument is the thread count (`u32 n`), not the index — the index + only shows up inside the `parallel_for` you write +- The functor must return the `[=](sycl::handler &cgh) { ... }` lambda; nothing is + submitted for you +- Pointers still follow in `MultiRef` order (inputs, then outputs), same as `kernel_call` + +## Side by side + +`kernel_call` is `kernel_call_hndl` with the `sycl::handler` boilerplate filled in +for you — compare the minimal examples above: + +::::{grid} 2 + +:::{grid-item-card} `kernel_call` +```cpp +sham::kernel_call( + q, + sham::MultiRef{buf_in}, + sham::MultiRef{buf_out}, + n, + [](u32 i, const T *in, T *out) { + out[i] = in[i]; + }); +``` +::: + +:::{grid-item-card} `kernel_call_hndl` +```cpp +sham::kernel_call_hndl( + q, + sham::MultiRef{buf_in}, + sham::MultiRef{buf_out}, + n, + [](u32 n, const T *in, T *out) { + return [=](sycl::handler &cgh) { + cgh.parallel_for(sycl::range<1>{n}, [=](sycl::item<1> item) { + u32 i = item.get_linear_id(); + out[i] = in[i]; + }); + }; + }); +``` +::: + +:::: + ## `MultiRef` `MultiRef` holds references to buffer-like objects passed to `kernel_call`. diff --git a/examples/benchmarks/run_sort_by_keys_tmp_dev.py b/examples/benchmarks/run_sort_by_keys_tmp_dev.py index 80e090a29f..7b2cd85b4d 100644 --- a/examples/benchmarks/run_sort_by_keys_tmp_dev.py +++ b/examples/benchmarks/run_sort_by_keys_tmp_dev.py @@ -53,9 +53,9 @@ def _sorted_run_mask(values): """Flag indices that belong to a non-decreasing run of length >= 2.""" mask = [False for i in range(len(values))] - for i,v in enumerate(values): + for i, v in enumerate(values): if i > 0: - if values[i-1] < values[i]: + if values[i - 1] < v: mask[i] = True else: mask[i] = True @@ -79,7 +79,7 @@ def print_key_val_table(keys, vals, labels=("i", "key", "val")): col_width = max(len(str(v)) for col in columns for v in col) def cell(v, highlight=None): - text = f"{str(v):>{col_width}}" + text = f"{v!s:>{col_width}}" if highlight is None: return text return f"{_GREEN}{text}{_RESET}" if highlight else f"{_RED}{text}{_RESET}" @@ -98,6 +98,7 @@ def fmt_row(label, values, mask=None): print(key_line) print(val_line) + def to_buf(lst): buf = shamrock.backends.DeviceBuffer_u32() buf.resize(len(lst)) @@ -105,9 +106,9 @@ def to_buf(lst): return buf -#%% +# %% # Dataset -N = 25 +N = 100 key_init = [i for i in range(N)][::-1] val_init = [i for i in range(N)] @@ -115,11 +116,11 @@ def to_buf(lst): print_key_val_table(key_init, val_init) -#%% +# %% # Switch impl shamrock.algs.set_impl_sort_by_keys('{"implementation":"modern_gpu_mergesort","parameters":{}}') -#%% +# %% # End state buffer_key = to_buf(key_init) buffer_val = to_buf(val_init) @@ -129,7 +130,7 @@ def to_buf(lst): print("End state:") print_key_val_table(buffer_key.copy_to_stdvec(), buffer_val.copy_to_stdvec()) -#%% +# %% # Expected state buffer_key = to_buf(key_init) buffer_val = to_buf(val_init) diff --git a/src/shamalgs/include/shamalgs/primitives/device/details/modern_gpu_merge_sort.hpp b/src/shamalgs/include/shamalgs/primitives/device/details/modern_gpu_merge_sort.hpp index b33e6c518c..03bc2f5c4a 100644 --- a/src/shamalgs/include/shamalgs/primitives/device/details/modern_gpu_merge_sort.hpp +++ b/src/shamalgs/include/shamalgs/primitives/device/details/modern_gpu_merge_sort.hpp @@ -15,7 +15,10 @@ * @brief */ +#include "shambase/integer.hpp" +#include "shamalgs/primitives/workitem/odd_even_transpose_sort.hpp" #include "shambackends/DeviceBuffer.hpp" +#include "shambackends/kernel_call.hpp" #include #include #include @@ -130,17 +133,441 @@ namespace shamalgs::primitives::device::details { std::cout << val_line << "\n"; } + enum MgpuBounds { MgpuBoundsLower, MgpuBoundsUpper }; + + template + inline int MergePath(It1 a, int aCount, It2 b, int bCount, int diag, Comp comp) { + + typedef typename std::iterator_traits::value_type T; + int begin = std::max(0, diag - bCount); + int end = std::min(diag, aCount); + + while (begin < end) { + int mid = (begin + end) >> 1; + T aKey = a[mid]; + T bKey = b[diag - 1 - mid]; + bool pred = (MgpuBoundsUpper == Bounds) ? comp(aKey, bKey) : !comp(bKey, aKey); + if (pred) + begin = mid + 1; + else + end = mid; + } + return begin; + } + + template + inline void SerialMerge( + sycl::nd_item<1> &item, + const T *keys_shared, + int aBegin, + int aEnd, + int bBegin, + int bEnd, + T *results, + int *indices, + Comp comp) { + + T aKey = keys_shared[aBegin]; + T bKey = keys_shared[bBegin]; + +#pragma unroll + for (int i = 0; i < VT; ++i) { + bool p; + if (RangeCheck) + p = (bBegin >= bEnd) || ((aBegin < aEnd) && !comp(bKey, aKey)); + else + p = !comp(bKey, aKey); + + results[i] = p ? aKey : bKey; + indices[i] = p ? aBegin : bBegin - !RangeCheck; + + if (p) + aKey = keys_shared[++aBegin]; + else + bKey = keys_shared[++bBegin]; + } + + item.barrier(sycl::access::fence_space::local_space); + } + + template + inline void CTABlocksortPass( + sycl::nd_item<1> &item, + T *keys_shared, + int tid, + int count, + int coop, + T *keys, + int *indices, + Comp comp) { + + int list = ~(coop - 1) & tid; + int diag = std::min(count, VT * ((coop - 1) & tid)); + int start = VT * list; + int a0 = std::min(count, start); + int b0 = std::min(count, start + VT * (coop / 2)); + int b1 = std::min(count, start + VT * coop); + + int p = MergePath( + keys_shared + a0, b0 - a0, keys_shared + b0, b1 - b0, diag, comp); + + SerialMerge( + item, keys_shared, a0 + p, b0, b0 + diag - p, b1, keys, indices, comp); + } + + template + inline void DeviceThreadToShared( + sycl::nd_item<1> &item, const T *threadReg, int tid, T *shared, bool sync = true) { + +// Odd grain size. Store as type T. +#pragma unroll + for (int i = 0; i < VT; ++i) + shared[VT * tid + i] = threadReg[i]; + + // In modern GPU there is 8-byte branch to exploit some instruction to store element 2 by 2 + + if (sync) + item.barrier(sycl::access::fence_space::local_space); + } + + template + inline void DeviceRegToShared( + sycl::nd_item<1> &item, const T *reg, int tid, OutputIt dest, bool sync) { + + typedef typename std::iterator_traits::value_type T2; +#pragma unroll + for (int i = 0; i < VT; ++i) + dest[NT * i + tid] = (T2) reg[i]; + + if (sync) + item.barrier(sycl::access::fence_space::local_space); + } + + template + inline void DeviceRegToGlobal( + sycl::nd_item<1> &item, int count, const T *reg, int tid, OutputIt dest, bool sync) { + +#pragma unroll + for (int i = 0; i < VT; ++i) { + int index = NT * i + tid; + if (index < count) + dest[index] = reg[i]; + } + if (sync) + item.barrier(sycl::access::fence_space::local_space); + } + + template + inline void DeviceSharedToThread( + sycl::nd_item<1> &item, const T *shared, int tid, T *threadReg, bool sync = true) { + +#pragma unroll + for (int i = 0; i < VT; ++i) + threadReg[i] = shared[VT * tid + i]; + + // In modern GPU there is 8-byte branch to exploit some instruction to store element 2 by 2 + + if (sync) + item.barrier(sycl::access::fence_space::local_space); + } + + template + inline void DeviceGlobalToRegPred( + sycl::nd_item<1> &item, int count, InputIt data, int tid, T *reg, bool sync) { + +// TODO: Attempt to issue 4 loads at a time. +#pragma unroll + for (int i = 0; i < VT; ++i) { + int index = NT * i + tid; + if (index < count) + reg[i] = data[index]; + } + if (sync) + item.barrier(sycl::access::fence_space::local_space); + } + + template + inline void DeviceGlobalToReg( + sycl::nd_item<1> &item, int count, InputIt data, int tid, T *reg, bool sync) { + + if (count >= NT * VT) { +#pragma unroll + for (int i = 0; i < VT; ++i) + reg[i] = data[NT * i + tid]; + } else + DeviceGlobalToRegPred(item, count, data, tid, reg, false); + + if (sync) + item.barrier(sycl::access::fence_space::local_space); + } + + template + inline void DeviceSharedToGlobal( + sycl::nd_item<1> &item, + int count, + const T *source, + int tid, + OutputIt dest, + bool sync = true) { + + typedef typename std::iterator_traits::value_type T2; +#pragma unroll + for (int i = 0; i < VT; ++i) { + int index = NT * i + tid; + if (index < count) + dest[index] = (T2) source[index]; + } + if (sync) + item.barrier(sycl::access::fence_space::local_space); + } + + template + inline void DeviceGlobalToShared( + sycl::nd_item<1> &item, int count, InputIt source, int tid, T *dest, bool sync = true) { + + T reg[VT]; + DeviceGlobalToReg(item, count, source, tid, reg, false); + DeviceRegToShared(item, reg, tid, dest, sync); + } + + template + inline void DeviceGather( + sycl::nd_item<1> &item, + int count, + InputIt data, + int indices[VT], + int tid, + T *reg, + bool sync = true) { + + if (count >= NT * VT) { +#pragma unroll + for (int i = 0; i < VT; ++i) + reg[i] = data[indices[i]]; + } else { +#pragma unroll + for (int i = 0; i < VT; ++i) { + int index = NT * i + tid; + if (index < count) + reg[i] = data[indices[i]]; + } + } + + if (sync) + item.barrier(sycl::access::fence_space::local_space); + } + + template + inline void CTABlocksortLoop( + sycl::nd_item<1> &item, + ValType threadValues[VT], + KeyType *keys_shared, + ValType *values_shared, + int tid, + int count, + Comp comp) { + +#pragma unroll + for (int coop = 2; coop <= NT; coop *= 2) { + int indices[VT]; + KeyType keys[VT]; + CTABlocksortPass(item, keys_shared, tid, count, coop, keys, indices, comp); + + if (HasValues) { + // Exchange the values through shared memory. + DeviceThreadToShared(item, threadValues, tid, values_shared); + DeviceGather(item, NT * VT, values_shared, indices, tid, threadValues); + } + + // Store results in shared memory in sorted order. + DeviceThreadToShared(item, keys, tid, keys_shared); + } + } + + template + inline void CTAMergesort( + sycl::nd_item<1> &item, + KeyType threadKeys[VT], + ValType threadValues[VT], + KeyType *keys_shared, + ValType *values_shared, + int count, + int tid, + Comp comp) { + + // Stable sort the keys in the thread. + if (VT * tid < count) { + workitem::odd_even_transpose_sort(threadKeys, threadValues, comp); + } + + // Store the locally sorted keys into shared memory. + DeviceThreadToShared(item, threadKeys, tid, keys_shared); + + // Recursively merge lists until the entire CTA is sorted. + CTABlocksortLoop( + item, threadValues, keys_shared, values_shared, tid, count, comp); + } + + template + union Shared { + static constexpr int NV = NT * VT; + + Tkey keys[NT * (VT + 1)]; + Tval values[NV]; + }; + + template + inline void KernelBlocksort( + sycl::nd_item<1> &item, + int tid, + int block, + Shared &shared, + Tkey *keysSource_global, + Tval *valsSource_global, + int count, + Tkey *keysDest_global, + Tval *valsDest_global, + Comp comp) { + + static constexpr int NV = NT * VT; + + int gid = NV * block; + int count2 = std::min(NV, count - gid); + + // Load the values into thread order. + Tval threadValues[VT]; + if (HasValues) { + DeviceGlobalToShared(item, count2, valsSource_global + gid, tid, shared.values); + DeviceSharedToThread(item, shared.values, tid, threadValues); + } + + // Load keys into shared memory and transpose into register in thread order. + Tkey threadKeys[VT]; + DeviceGlobalToShared(item, count2, keysSource_global + gid, tid, shared.keys); + DeviceSharedToThread(item, shared.keys, tid, threadKeys); + + // If we're in the last tile, set the uninitialized keys for the thread with + // a partial number of keys. + int first = VT * tid; + if (first + VT > count2 && first < count2) { + Tkey maxKey = threadKeys[0]; +#pragma unroll + for (int i = 1; i < VT; ++i) + if (first + i < count2) + maxKey = comp(maxKey, threadKeys[i]) ? threadKeys[i] : maxKey; + +// Fill in the uninitialized elements with max key. +#pragma unroll + for (int i = 0; i < VT; ++i) + if (first + i >= count2) + threadKeys[i] = maxKey; + } + + CTAMergesort( + item, threadKeys, threadValues, shared.keys, shared.values, count2, tid, comp); + + // Store the sorted keys to global. + DeviceSharedToGlobal(item, count2, shared.keys, tid, keysDest_global + gid); + + if (HasValues) { + DeviceThreadToShared(item, threadValues, tid, shared.values); + DeviceSharedToGlobal(item, count2, shared.values, tid, valsDest_global + gid); + } + } + + ///// Fine ////// + + /// Builds an `nd_range` with work-group size `wg_size`, rounding `nthread` up to the next + /// multiple of `wg_size` so every work-group launched is full. + inline sycl::nd_range<1> ndrange(u32 wg_size, u32 nthread) { + u32 corrected_len = shambase::group_count(nthread, wg_size) * wg_size; + return sycl::nd_range<1>{corrected_len, wg_size}; + } + template inline void sort_by_keys_modern_gpu_mergesort( sham::DeviceBuffer &buf_key, sham::DeviceBuffer &buf_values, u32 len) { - std::cout << "-------------------------------------------------" << std::endl; - std::cout << "------- sort_by_keys_modern_gpu_mergesort -------" << std::endl; - std::cout << "-------------------------------------------------" << std::endl; - std::cout << "init state:" << std::endl; - print_key_val_table(buf_key.copy_to_stdvec(), buf_values.copy_to_stdvec()); - std::cout << "-------------------------------------------------" << std::endl; - std::cout << "------- sort_by_keys_modern_gpu_mergesort end -------" << std::endl; - std::cout << "-------------------------------------------------" << std::endl; + + auto dev_sched = buf_key.get_dev_scheduler_ptr(); + + bool do_print = false; + + if (do_print) { + std::cout << "-------------------------------------------------" << std::endl; + std::cout << "------- sort_by_keys_modern_gpu_mergesort -------" << std::endl; + std::cout << "-------------------------------------------------" << std::endl; + + std::cout << "init state:" << std::endl; + print_key_val_table(buf_key.copy_to_stdvec(), buf_values.copy_to_stdvec()); + } + + static constexpr int VT = 7; + static constexpr int NT = 4; + static constexpr int NV = NT * VT; + + u32 nthreads = len / VT + 1; + + sham::kernel_call_hndl( + dev_sched->get_queue(), + sham::MultiRef{}, + sham::MultiRef{buf_key, buf_values}, + [nthreads, len](Tkey *__restrict keys, Tval *__restrict vals) { + return [=](sycl::handler &cgh) { + using Shared = Shared; + + sycl::local_accessor shared_mem(1, cgh); + + cgh.parallel_for(ndrange(NT, nthreads), [=](sycl::nd_item<1> item) { + u32 gid = item.get_global_linear_id(); + u32 lid = item.get_local_linear_id(); + u32 block = item.get_group_linear_id(); + + Tkey loc_key[VT]; + Tval loc_val[VT]; + + for (int i = 0; i < VT; i++) { + u32 idx = gid * VT + i; + loc_key[i] = (idx < len) ? keys[idx] : shambase::get_max(); + loc_val[i] = (idx < len) ? vals[idx] : Tval{}; + } + + workitem::odd_even_transpose_sort(loc_key, loc_val, [](Tkey a, Tkey b) { + return a < b; + }); + + for (int i = 0; i < VT; i++) { + u32 idx = gid * VT + i; + if (idx < len) { + keys[idx] = loc_key[i]; + vals[idx] = loc_val[i]; + } + } + + KernelBlocksort( + item, + lid, + block, + shared_mem[0], + keys, + vals, + len, + keys, + vals, + [](Tkey a, Tkey b) { + return a < b; + }); + }); + }; + }); + + if (do_print) { + std::cout << "after local sort state:" << std::endl; + print_key_val_table(buf_key.copy_to_stdvec(), buf_values.copy_to_stdvec()); + + std::cout << "-------------------------------------------------" << std::endl; + std::cout << "------- sort_by_keys_modern_gpu_mergesort end -------" << std::endl; + std::cout << "-------------------------------------------------" << std::endl; + } } } // namespace shamalgs::primitives::device::details diff --git a/src/shamalgs/include/shamalgs/primitives/workitem/odd_even_transpose_sort.hpp b/src/shamalgs/include/shamalgs/primitives/workitem/odd_even_transpose_sort.hpp index e82270db7a..2ab337e732 100644 --- a/src/shamalgs/include/shamalgs/primitives/workitem/odd_even_transpose_sort.hpp +++ b/src/shamalgs/include/shamalgs/primitives/workitem/odd_even_transpose_sort.hpp @@ -35,7 +35,7 @@ namespace shamalgs::primitives::workitem { * @param comp Comparator used to order the keys. */ template - inline void odd_even_transpose_sort(T *keys, V *values, Comp comp) { + inline void odd_even_transpose_sort(T keys[VT], V values[VT], Comp comp) { #pragma unroll for (int level = 0; level < VT; ++level) { diff --git a/src/shambackends/include/shambackends/kernel_call.hpp b/src/shambackends/include/shambackends/kernel_call.hpp index 7961a4d4d1..92fff1a638 100644 --- a/src/shambackends/include/shambackends/kernel_call.hpp +++ b/src/shambackends/include/shambackends/kernel_call.hpp @@ -68,6 +68,41 @@ namespace sham { in_out.complete_event_state(e); } + /// internal implementation of kernel_call_hndl without a thread count, for kernel + /// generators that build their own range/nd_range internally and have no use for `n` + template + void kernel_call_lambda( + sham::DeviceQueue &q, + RefIn in, + RefOut in_out, + Functor &&kernel_gen, + SourceLocation &&callsite = SourceLocation{}) { + + __shamrock_stack_entry_with_callsite(callsite); + + sham::EventList depends_list; + + auto acc_in = in.get_read_access(depends_list); + auto acc_in_out = in_out.get_write_access(depends_list); + + sycl::event e; + + // unpack the tuples of accessors + std::apply( + [&](auto &...__acc_in) { + std::apply( + [&](auto &...__acc_in_out) { + // submit the kernel generated by the functor + e = q.submit(depends_list, kernel_gen(__acc_in..., __acc_in_out...)); + }, + acc_in_out); + }, + acc_in); + + in.complete_event_state(e); + in_out.complete_event_state(e); + } + ////////////////////////////////////////////////////////////////////////////// // Helper types for kernel signature checking (and clean error messages) ////////////////////////////////////////////////////////////////////////////// @@ -387,6 +422,22 @@ namespace sham { q, in, in_out, n, std::forward(kernel_gen)); } + /// version of kernel_call_hndl without a thread count, for kernel generators that build + /// their own range/nd_range internally (e.g. from a size that isn't a simple element count) + template + void kernel_call_hndl( + sham::DeviceQueue &q, + RefIn in, + RefOut in_out, + Functor &&kernel_gen, + SourceLocation &&callsite = SourceLocation{}) { + + __shamrock_log_callsite(callsite); + + details::kernel_call_lambda( + q, in, in_out, std::forward(kernel_gen)); + } + /// u64 indexed variant of kernel_call_hndl template void kernel_call_hndl_u64( From 4451b05a12e90a9e296a924d2befcbef6b513007 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Timoth=C3=A9e=20David--Cl=C3=A9ris?= Date: Mon, 31 Aug 2026 08:08:03 +0200 Subject: [PATCH 3/7] Update modern_gpu_merge_sort.hpp --- .../device/details/modern_gpu_merge_sort.hpp | 445 +++++++++++++++++- 1 file changed, 426 insertions(+), 19 deletions(-) diff --git a/src/shamalgs/include/shamalgs/primitives/device/details/modern_gpu_merge_sort.hpp b/src/shamalgs/include/shamalgs/primitives/device/details/modern_gpu_merge_sort.hpp index 03bc2f5c4a..5b54fa6dd5 100644 --- a/src/shamalgs/include/shamalgs/primitives/device/details/modern_gpu_merge_sort.hpp +++ b/src/shamalgs/include/shamalgs/primitives/device/details/modern_gpu_merge_sort.hpp @@ -484,31 +484,19 @@ namespace shamalgs::primitives::device::details { return sycl::nd_range<1>{corrected_len, wg_size}; } - template - inline void sort_by_keys_modern_gpu_mergesort( - sham::DeviceBuffer &buf_key, sham::DeviceBuffer &buf_values, u32 len) { - - auto dev_sched = buf_key.get_dev_scheduler_ptr(); - - bool do_print = false; - - if (do_print) { - std::cout << "-------------------------------------------------" << std::endl; - std::cout << "------- sort_by_keys_modern_gpu_mergesort -------" << std::endl; - std::cout << "-------------------------------------------------" << std::endl; - - std::cout << "init state:" << std::endl; - print_key_val_table(buf_key.copy_to_stdvec(), buf_values.copy_to_stdvec()); - } + template + inline void kernel_blocksort( + sham::DeviceQueue &q, + sham::DeviceBuffer &buf_key, + sham::DeviceBuffer &buf_values, + u32 len) { - static constexpr int VT = 7; - static constexpr int NT = 4; static constexpr int NV = NT * VT; u32 nthreads = len / VT + 1; sham::kernel_call_hndl( - dev_sched->get_queue(), + q, sham::MultiRef{}, sham::MultiRef{buf_key, buf_values}, [nthreads, len](Tkey *__restrict keys, Tval *__restrict vals) { @@ -559,6 +547,361 @@ namespace shamalgs::primitives::device::details { }); }; }); + } + +#define MGPU_DIV_UP(x, y) (((x) + (y) - 1) / (y)) +#define MGPU_IS_POW_2(x) (0 == ((x) & ((x) - 1))) + + // Find log2(x) and optionally round up to the next integer logarithm. + inline int FindLog2(int x, bool roundUp = false) { + int a = 31 - sycl::clz(x); + if (roundUp) + a += !MGPU_IS_POW_2(x); + return a; + } + + inline sycl::vec FindMergesortFrame(int coop, int block, int nv) { + // coop is the number of CTAs or threads cooperating to merge two lists into + // one. We round block down to the first CTA's ID that is working on this + // merge. + int start = ~(coop - 1) & block; + int size = nv * (coop >> 1); + return {nv * start, nv * start + size, size}; + } + + template + inline void KernelMergePartition( + int block, + int tid, + It1 a_global, + int aCount, + It2 b_global, + int bCount, + int nv, + int coop, + int *mp_global, + int numSearches, + Comp comp) { + + int partition = NT * block + tid; + if (partition < numSearches) { + int a0 = 0, b0 = 0; + int gid = nv * partition; + if (coop) { + auto frame = FindMergesortFrame(coop, partition, nv); + a0 = frame.x(); + b0 = std::min(aCount, frame.y()); + bCount = std::min(aCount, frame.y() + frame.z()) - b0; + aCount = std::min(aCount, frame.x() + frame.z()) - a0; + + // Put the cross-diagonal into the coordinate system of the input + // lists. + gid -= a0; + } + int mp = MergePath( + a_global + a0, aCount, b_global + b0, bCount, std::min(gid, aCount + bCount), comp); + mp_global[partition] = mp; + } + } + + template + inline void MergePathPartitions( + sham::DeviceQueue &q, + sham::DeviceBuffer &a_global, + int aCount, + sham::DeviceBuffer &b_global, + int bCount, + int nv, + int coop, + Comp comp, + sham::DeviceBuffer &partitionsDevice) { + + const int NT = 64; + int numPartitions = MGPU_DIV_UP(aCount + bCount, nv); + int numSearches = numPartitions + 1; + int numPartitionBlocks = MGPU_DIV_UP(numSearches, NT); + partitionsDevice.resize(numSearches); + + sham::kernel_call_hndl( + q, sham::MultiRef{}, sham::MultiRef{a_global, b_global,partitionsDevice}, [=](auto a, auto b, int *__restrict part_dev) { + return [=](sycl::handler &cgh) { + cgh.parallel_for( + ndrange(NT, numPartitionBlocks * NT), [=](sycl::nd_item<1> item) { + int block = static_cast(item.get_group_linear_id()); + int tid = static_cast(item.get_local_linear_id()); + + KernelMergePartition( + block, + tid, + a, + aCount, + b, + bCount, + nv, + coop, + part_dev, + numSearches, + comp); + }); + }; + }); + } + + // Returns (a0, a1, b0, b1) into mergesort input lists between mp0 and mp1. + inline sycl::vec FindMergesortInterval( + sycl::vec frame, int coop, int block, int nv, int count, int mp0, int mp1) { + + // Locate diag from the start of the A sublist. + int diag = nv * block - frame.x(); + int a0 = frame.x() + mp0; + int a1 = std::min(count, frame.x() + mp1); + int b0 = std::min(count, frame.y() + diag - mp0); + int b1 = std::min(count, frame.y() + diag + nv - mp1); + + // The end partition of the last block for each merge operation is computed + // and stored as the begin partition for the subsequent merge. i.e. it is + // the same partition but in the wrong coordinate system, so its 0 when it + // should be listSize. Correct that by checking if this is the last block + // in this merge operation. + if (coop - 1 == ((coop - 1) & block)) { + a1 = std::min(count, frame.x() + frame.z()); + b1 = std::min(count, frame.y() + frame.z()); + } + return {a0, a1, b0, b1}; + } + + inline sycl::vec ComputeMergeRange( + int aCount, int bCount, int block, int coop, int NV, const int *mp_global) { + + // Load the merge paths computed by the partitioning kernel. + int mp0 = mp_global[block]; + int mp1 = mp_global[block + 1]; + int gid = NV * block; + + // Compute the ranges of the sources in global memory. + sycl::vec range; + if (coop) { + sycl::vec frame = FindMergesortFrame(coop, block, NV); + range = FindMergesortInterval(frame, coop, block, NV, aCount, mp0, mp1); + } else { + range.x() = mp0; // a0 + range.y() = mp1; // a1 + range.z() = gid - range.x(); // b0 + range.w() = std::min(aCount + bCount, gid + NV) - range.y(); // b1 + } + return range; + } + + + template + MGPU_DEVICE void DeviceMergeKeysIndices(It1 a_global, int aCount, It2 b_global, + int bCount, int4 range, int tid, T* keys_shared, T* results, int* indices, + Comp comp) { + + int a0 = range.x; + int a1 = range.y; + int b0 = range.z; + int b1 = range.w; + + if(LoadExtended) { + bool extended = (a1 < aCount) && (b1 < bCount); + aCount = a1 - a0; + bCount = b1 - b0; + int aCount2 = aCount + (int)extended; + int bCount2 = bCount + (int)extended; + + // Load one element past the end of each input to avoid having to use + // range checking in the merge loop. + DeviceLoad2ToShared(a_global + a0, aCount2, + b_global + b0, bCount2, tid, keys_shared); + + // Run a Merge Path search for each thread's starting point. + int diag = VT * tid; + int mp = MergePath(keys_shared, aCount, + keys_shared + aCount2, bCount, diag, comp); + + // Compute the ranges of the sources in shared memory. + int a0tid = mp; + int b0tid = aCount2 + diag - mp; + if(extended) { + SerialMerge(keys_shared, a0tid, 0, b0tid, 0, results, + indices, comp); + } else { + int a1tid = aCount; + int b1tid = aCount2 + bCount; + SerialMerge(keys_shared, a0tid, a1tid, b0tid, b1tid, + results, indices, comp); + } + } else { + // Use the input intervals from the ranges between the merge path + // intersections. + aCount = a1 - a0; + bCount = b1 - b0; + + // Load the data into shared memory. + DeviceLoad2ToShared(a_global + a0, aCount, b_global + b0, + bCount, tid, keys_shared); + + // Run a merge path to find the start of the serial merge for each + // thread. + int diag = VT * tid; + int mp = MergePath(keys_shared, aCount, + keys_shared + aCount, bCount, diag, comp); + + // Compute the ranges of the sources in shared memory. + int a0tid = mp; + int a1tid = aCount; + int b0tid = aCount + diag - mp; + int b1tid = aCount + bCount; + + // Serial merge into register. + SerialMerge(keys_shared, a0tid, a1tid, b0tid, b1tid, results, + indices, comp); + } + } + + template< + int NT, + int VT, + bool HasValues, + bool LoadExtended, + typename KeysIt1, + typename KeysIt2, + typename KeysIt3, + typename ValsIt1, + typename ValsIt2, + typename KeyType, + typename ValsIt3, + typename Comp> + inline void DeviceMerge( + sycl::nd_item<1> &item, + KeysIt1 aKeys_global, + ValsIt1 aVals_global, + int aCount, + KeysIt2 bKeys_global, + ValsIt2 bVals_global, + int bCount, + int tid, + int block, + sycl::vec range, + KeyType *keys_shared, + int *indices_shared, + KeysIt3 keys_global, + ValsIt3 vals_global, + Comp comp) { + + KeyType results[VT]; + int indices[VT]; + DeviceMergeKeysIndices( + aKeys_global, + aCount, + bKeys_global, + bCount, + range, + tid, + keys_shared, + results, + indices, + comp); + + // Store merge results back to shared memory. + DeviceThreadToShared(results, tid, keys_shared); + + // Store merged keys to global memory. + aCount = range.y() - range.x(); + bCount = range.w() - range.z(); + DeviceSharedToGlobal( + aCount + bCount, keys_shared, tid, keys_global + NT * VT * block); + + // Copy the values. + if (HasValues) { + DeviceThreadToShared(item,indices, tid, indices_shared); + + DeviceTransferMergeValuesShared( + aCount + bCount, + aVals_global + range.x(), + bVals_global + range.z(), + aCount, + indices_shared, + tid, + vals_global + NT * VT * block); + } + } + + template< + int NT, + int VT, + bool HasValues, + bool LoadExtended, + typename Tkey, + typename Tval, + typename Comp> + inline void KernelMerge( + sycl::nd_item<1> &item, + int tid, + int block, + Shared &shared, + Tkey *aKeys_global, + Tval *aVals_global, + int aCount, + Tkey *bKeys_global, + Tval *bVals_global, + int bCount, + const int *mp_global, + int coop, + Tkey *keys_global, + Tval *vals_global, + Comp comp) { + + static constexpr int NV = NT * VT; + + sycl::vec range + = ComputeMergeRange(aCount, bCount, block, coop, NT * VT, mp_global); + + DeviceMerge( + item, + aKeys_global, + aVals_global, + aCount, + bKeys_global, + bVals_global, + bCount, + tid, + block, + range, + shared.keys, + shared.values, + keys_global, + vals_global, + comp); + } + + template + inline void sort_by_keys_modern_gpu_mergesort( + sham::DeviceBuffer &buf_key, sham::DeviceBuffer &buf_values, u32 len) { + + auto dev_sched = buf_key.get_dev_scheduler_ptr(); + + bool do_print = false; + + static constexpr int VT = 7; + static constexpr int NT = 4; // 256 + static constexpr int NV = NT * VT; + + int numBlocks = MGPU_DIV_UP(len, NV); + int numPasses = FindLog2(numBlocks, true); + + if (do_print) { + std::cout << "-------------------------------------------------" << std::endl; + std::cout << "------- sort_by_keys_modern_gpu_mergesort -------" << std::endl; + std::cout << "-------------------------------------------------" << std::endl; + + std::cout << "init state:" << std::endl; + print_key_val_table(buf_key.copy_to_stdvec(), buf_values.copy_to_stdvec()); + } + + kernel_blocksort(dev_sched->get_queue(), buf_key, buf_values, len); if (do_print) { std::cout << "after local sort state:" << std::endl; @@ -568,6 +911,70 @@ namespace shamalgs::primitives::device::details { std::cout << "------- sort_by_keys_modern_gpu_mergesort end -------" << std::endl; std::cout << "-------------------------------------------------" << std::endl; } + + auto buf_key2 = sham::DeviceBuffer(len, dev_sched); + auto buf_values2 = sham::DeviceBuffer(len, dev_sched); + + if (1 & numPasses) { + std::swap(buf_key, buf_key2); + std::swap(buf_values, buf_values2); + } + + auto partitionsDevice = sham::DeviceBuffer(0, dev_sched); + + for (int pass = 0; pass < numPasses; ++pass) { + int coop = 2 << pass; + + MergePathPartitions( + dev_sched->get_queue(), + buf_key, + len, + buf_key, + 0, + NV, + coop, + [](Tkey a, Tkey b) { + return a < b; + }, + partitionsDevice); + + sham::kernel_call_hndl( + dev_sched->get_queue(), + sham::MultiRef{}, + sham::MultiRef{buf_key, buf_values, partitionsDevice, buf_key2, buf_values2}, + [=](auto key, auto values, auto part_dev, auto key2, auto values2) { + return [=](sycl::handler &cgh) { + + using Shared = Shared; + + sycl::local_accessor shared_mem(1, cgh); + + cgh.parallel_for(ndrange(NT, numBlocks * NT), [=](sycl::nd_item<1> item) { + KernelMerge( + item, + item.get_local_linear_id(), + item.get_group_linear_id(), + shared_mem[0], + key, + values, + len, + key, + values, + 0, + part_dev, + coop, + key2, + values2, + [](Tkey a, Tkey b) { + return a < b; + }); + }); + }; + }); + + std::swap(buf_key, buf_key2); + std::swap(buf_values, buf_values2); + } } } // namespace shamalgs::primitives::device::details From 726cdbdda5d79fa35d4bdaeddaba2025d1780fdd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Timoth=C3=A9e=20David--Cl=C3=A9ris?= Date: Mon, 31 Aug 2026 08:28:52 +0200 Subject: [PATCH 4/7] progress --- .../device/details/modern_gpu_merge_sort.hpp | 278 +++++++++--------- .../include/shambackends/make_ndrange.hpp | 31 ++ src/tests/shambackends/make_ndrangeTests.cpp | 55 ++++ 3 files changed, 232 insertions(+), 132 deletions(-) create mode 100644 src/shambackends/include/shambackends/make_ndrange.hpp create mode 100644 src/tests/shambackends/make_ndrangeTests.cpp diff --git a/src/shamalgs/include/shamalgs/primitives/device/details/modern_gpu_merge_sort.hpp b/src/shamalgs/include/shamalgs/primitives/device/details/modern_gpu_merge_sort.hpp index 5b54fa6dd5..2aebfb5299 100644 --- a/src/shamalgs/include/shamalgs/primitives/device/details/modern_gpu_merge_sort.hpp +++ b/src/shamalgs/include/shamalgs/primitives/device/details/modern_gpu_merge_sort.hpp @@ -19,6 +19,7 @@ #include "shamalgs/primitives/workitem/odd_even_transpose_sort.hpp" #include "shambackends/DeviceBuffer.hpp" #include "shambackends/kernel_call.hpp" +#include "shambackends/make_ndrange.hpp" #include #include #include @@ -409,81 +410,78 @@ namespace shamalgs::primitives::device::details { } template - union Shared { - static constexpr int NV = NT * VT; - - Tkey keys[NT * (VT + 1)]; - Tval values[NV]; - }; - - template - inline void KernelBlocksort( - sycl::nd_item<1> &item, - int tid, - int block, - Shared &shared, - Tkey *keysSource_global, - Tval *valsSource_global, - int count, - Tkey *keysDest_global, - Tval *valsDest_global, - Comp comp) { + struct KernelBlocksort { static constexpr int NV = NT * VT; - int gid = NV * block; - int count2 = std::min(NV, count - gid); + union Shared { - // Load the values into thread order. - Tval threadValues[VT]; - if (HasValues) { - DeviceGlobalToShared(item, count2, valsSource_global + gid, tid, shared.values); - DeviceSharedToThread(item, shared.values, tid, threadValues); - } + Tkey keys[NT * (VT + 1)]; + Tval values[NV]; + }; - // Load keys into shared memory and transpose into register in thread order. - Tkey threadKeys[VT]; - DeviceGlobalToShared(item, count2, keysSource_global + gid, tid, shared.keys); - DeviceSharedToThread(item, shared.keys, tid, threadKeys); + template + inline void Kernel( + sycl::nd_item<1> &item, + int tid, + int block, + Shared &shared, + Tkey *keysSource_global, + Tval *valsSource_global, + int count, + Tkey *keysDest_global, + Tval *valsDest_global, + Comp comp) { + + int gid = NV * block; + int count2 = std::min(NV, count - gid); + + // Load the values into thread order. + Tval threadValues[VT]; + if (HasValues) { + DeviceGlobalToShared( + item, count2, valsSource_global + gid, tid, shared.values); + DeviceSharedToThread(item, shared.values, tid, threadValues); + } + + // Load keys into shared memory and transpose into register in thread order. + Tkey threadKeys[VT]; + DeviceGlobalToShared(item, count2, keysSource_global + gid, tid, shared.keys); + DeviceSharedToThread(item, shared.keys, tid, threadKeys); - // If we're in the last tile, set the uninitialized keys for the thread with - // a partial number of keys. - int first = VT * tid; - if (first + VT > count2 && first < count2) { - Tkey maxKey = threadKeys[0]; + // If we're in the last tile, set the uninitialized keys for the thread with + // a partial number of keys. + int first = VT * tid; + if (first + VT > count2 && first < count2) { + Tkey maxKey = threadKeys[0]; #pragma unroll - for (int i = 1; i < VT; ++i) - if (first + i < count2) - maxKey = comp(maxKey, threadKeys[i]) ? threadKeys[i] : maxKey; + for (int i = 1; i < VT; ++i) + if (first + i < count2) + maxKey = comp(maxKey, threadKeys[i]) ? threadKeys[i] : maxKey; // Fill in the uninitialized elements with max key. #pragma unroll - for (int i = 0; i < VT; ++i) - if (first + i >= count2) - threadKeys[i] = maxKey; - } + for (int i = 0; i < VT; ++i) + if (first + i >= count2) + threadKeys[i] = maxKey; + } - CTAMergesort( - item, threadKeys, threadValues, shared.keys, shared.values, count2, tid, comp); + CTAMergesort( + item, threadKeys, threadValues, shared.keys, shared.values, count2, tid, comp); - // Store the sorted keys to global. - DeviceSharedToGlobal(item, count2, shared.keys, tid, keysDest_global + gid); + // Store the sorted keys to global. + DeviceSharedToGlobal(item, count2, shared.keys, tid, keysDest_global + gid); - if (HasValues) { - DeviceThreadToShared(item, threadValues, tid, shared.values); - DeviceSharedToGlobal(item, count2, shared.values, tid, valsDest_global + gid); + if (HasValues) { + DeviceThreadToShared(item, threadValues, tid, shared.values); + DeviceSharedToGlobal( + item, count2, shared.values, tid, valsDest_global + gid); + } } - } + }; ///// Fine ////// - /// Builds an `nd_range` with work-group size `wg_size`, rounding `nthread` up to the next - /// multiple of `wg_size` so every work-group launched is full. - inline sycl::nd_range<1> ndrange(u32 wg_size, u32 nthread) { - u32 corrected_len = shambase::group_count(nthread, wg_size) * wg_size; - return sycl::nd_range<1>{corrected_len, wg_size}; - } - template inline void kernel_blocksort( sham::DeviceQueue &q, @@ -495,17 +493,19 @@ namespace shamalgs::primitives::device::details { u32 nthreads = len / VT + 1; + using Kernel = KernelBlocksort; + sham::kernel_call_hndl( q, sham::MultiRef{}, sham::MultiRef{buf_key, buf_values}, [nthreads, len](Tkey *__restrict keys, Tval *__restrict vals) { return [=](sycl::handler &cgh) { - using Shared = Shared; + using Shared = Kernel::Shared; sycl::local_accessor shared_mem(1, cgh); - cgh.parallel_for(ndrange(NT, nthreads), [=](sycl::nd_item<1> item) { + cgh.parallel_for(sham::make_ndrange(NT, nthreads), [=](sycl::nd_item<1> item) { u32 gid = item.get_global_linear_id(); u32 lid = item.get_local_linear_id(); u32 block = item.get_group_linear_id(); @@ -531,7 +531,7 @@ namespace shamalgs::primitives::device::details { } } - KernelBlocksort( + Kernel::Kernel( item, lid, block, @@ -623,10 +623,14 @@ namespace shamalgs::primitives::device::details { partitionsDevice.resize(numSearches); sham::kernel_call_hndl( - q, sham::MultiRef{}, sham::MultiRef{a_global, b_global,partitionsDevice}, [=](auto a, auto b, int *__restrict part_dev) { + q, + sham::MultiRef{}, + sham::MultiRef{a_global, b_global, partitionsDevice}, + [=](auto a, auto b, int *__restrict part_dev) { return [=](sycl::handler &cgh) { cgh.parallel_for( - ndrange(NT, numPartitionBlocks * NT), [=](sycl::nd_item<1> item) { + sham::make_ndrange(NT, numPartitionBlocks * NT), + [=](sycl::nd_item<1> item) { int block = static_cast(item.get_group_linear_id()); int tid = static_cast(item.get_local_linear_id()); @@ -692,73 +696,84 @@ namespace shamalgs::primitives::device::details { return range; } + template< + int NT, + int VT, + bool LoadExtended, + typename It1, + typename It2, + typename T, + typename Comp> + MGPU_DEVICE void DeviceMergeKeysIndices( + It1 a_global, + int aCount, + It2 b_global, + int bCount, + int4 range, + int tid, + T *keys_shared, + T *results, + int *indices, + Comp comp) { - template - MGPU_DEVICE void DeviceMergeKeysIndices(It1 a_global, int aCount, It2 b_global, - int bCount, int4 range, int tid, T* keys_shared, T* results, int* indices, - Comp comp) { - - int a0 = range.x; - int a1 = range.y; - int b0 = range.z; - int b1 = range.w; - - if(LoadExtended) { - bool extended = (a1 < aCount) && (b1 < bCount); - aCount = a1 - a0; - bCount = b1 - b0; - int aCount2 = aCount + (int)extended; - int bCount2 = bCount + (int)extended; - - // Load one element past the end of each input to avoid having to use - // range checking in the merge loop. - DeviceLoad2ToShared(a_global + a0, aCount2, - b_global + b0, bCount2, tid, keys_shared); - - // Run a Merge Path search for each thread's starting point. - int diag = VT * tid; - int mp = MergePath(keys_shared, aCount, - keys_shared + aCount2, bCount, diag, comp); - - // Compute the ranges of the sources in shared memory. - int a0tid = mp; - int b0tid = aCount2 + diag - mp; - if(extended) { - SerialMerge(keys_shared, a0tid, 0, b0tid, 0, results, - indices, comp); - } else { - int a1tid = aCount; - int b1tid = aCount2 + bCount; - SerialMerge(keys_shared, a0tid, a1tid, b0tid, b1tid, - results, indices, comp); - } - } else { - // Use the input intervals from the ranges between the merge path - // intersections. - aCount = a1 - a0; - bCount = b1 - b0; - - // Load the data into shared memory. - DeviceLoad2ToShared(a_global + a0, aCount, b_global + b0, - bCount, tid, keys_shared); - - // Run a merge path to find the start of the serial merge for each - // thread. - int diag = VT * tid; - int mp = MergePath(keys_shared, aCount, - keys_shared + aCount, bCount, diag, comp); - - // Compute the ranges of the sources in shared memory. - int a0tid = mp; - int a1tid = aCount; - int b0tid = aCount + diag - mp; - int b1tid = aCount + bCount; - - // Serial merge into register. - SerialMerge(keys_shared, a0tid, a1tid, b0tid, b1tid, results, - indices, comp); - } + int a0 = range.x; + int a1 = range.y; + int b0 = range.z; + int b1 = range.w; + + if (LoadExtended) { + bool extended = (a1 < aCount) && (b1 < bCount); + aCount = a1 - a0; + bCount = b1 - b0; + int aCount2 = aCount + (int) extended; + int bCount2 = bCount + (int) extended; + + // Load one element past the end of each input to avoid having to use + // range checking in the merge loop. + DeviceLoad2ToShared( + a_global + a0, aCount2, b_global + b0, bCount2, tid, keys_shared); + + // Run a Merge Path search for each thread's starting point. + int diag = VT * tid; + int mp = MergePath( + keys_shared, aCount, keys_shared + aCount2, bCount, diag, comp); + + // Compute the ranges of the sources in shared memory. + int a0tid = mp; + int b0tid = aCount2 + diag - mp; + if (extended) { + SerialMerge(keys_shared, a0tid, 0, b0tid, 0, results, indices, comp); + } else { + int a1tid = aCount; + int b1tid = aCount2 + bCount; + SerialMerge( + keys_shared, a0tid, a1tid, b0tid, b1tid, results, indices, comp); + } + } else { + // Use the input intervals from the ranges between the merge path + // intersections. + aCount = a1 - a0; + bCount = b1 - b0; + + // Load the data into shared memory. + DeviceLoad2ToShared( + a_global + a0, aCount, b_global + b0, bCount, tid, keys_shared); + + // Run a merge path to find the start of the serial merge for each + // thread. + int diag = VT * tid; + int mp = MergePath( + keys_shared, aCount, keys_shared + aCount, bCount, diag, comp); + + // Compute the ranges of the sources in shared memory. + int a0tid = mp; + int a1tid = aCount; + int b0tid = aCount + diag - mp; + int b1tid = aCount + bCount; + + // Serial merge into register. + SerialMerge(keys_shared, a0tid, a1tid, b0tid, b1tid, results, indices, comp); + } } template< @@ -816,7 +831,7 @@ namespace shamalgs::primitives::device::details { // Copy the values. if (HasValues) { - DeviceThreadToShared(item,indices, tid, indices_shared); + DeviceThreadToShared(item, indices, tid, indices_shared); DeviceTransferMergeValuesShared( aCount + bCount, @@ -944,13 +959,12 @@ namespace shamalgs::primitives::device::details { sham::MultiRef{buf_key, buf_values, partitionsDevice, buf_key2, buf_values2}, [=](auto key, auto values, auto part_dev, auto key2, auto values2) { return [=](sycl::handler &cgh) { - using Shared = Shared; sycl::local_accessor shared_mem(1, cgh); cgh.parallel_for(ndrange(NT, numBlocks * NT), [=](sycl::nd_item<1> item) { - KernelMerge( + KernelMerge( item, item.get_local_linear_id(), item.get_group_linear_id(), diff --git a/src/shambackends/include/shambackends/make_ndrange.hpp b/src/shambackends/include/shambackends/make_ndrange.hpp new file mode 100644 index 0000000000..c4f0c57fa6 --- /dev/null +++ b/src/shambackends/include/shambackends/make_ndrange.hpp @@ -0,0 +1,31 @@ +// -------------------------------------------------------// +// +// SHAMROCK code for hydrodynamics +// Copyright (c) 2021-2026 Timothée David--Cléris +// SPDX-License-Identifier: CeCILL Free Software License Agreement v2.1 +// Shamrock is licensed under the CeCILL 2.1 License, see LICENSE for more information +// +// -------------------------------------------------------// + +#pragma once + +/** + * @file make_ndrange.hpp + * @author Timothée David--Cléris (tim.shamrock@proton.me) + * @brief + * + */ + +#include "shambase/aliases_int.hpp" +#include "shambackends/sycl.hpp" + +namespace sham { + + /// Builds an `nd_range` with work-group size `wg_size`, rounding `nthread` up to the next + /// multiple of `wg_size` so every work-group launched is full. + inline sycl::nd_range<1> make_ndrange(u32 wg_size, u32 nthread) { + u32 nthread_rounded = ((nthread + wg_size - 1) / wg_size) * wg_size; + return sycl::nd_range<1>(sycl::range<1>(nthread_rounded), sycl::range<1>(wg_size)); + } + +} // namespace sham diff --git a/src/tests/shambackends/make_ndrangeTests.cpp b/src/tests/shambackends/make_ndrangeTests.cpp new file mode 100644 index 0000000000..4d56306a00 --- /dev/null +++ b/src/tests/shambackends/make_ndrangeTests.cpp @@ -0,0 +1,55 @@ +// -------------------------------------------------------// +// +// SHAMROCK code for hydrodynamics +// Copyright (c) 2021-2026 Timothée David--Cléris +// SPDX-License-Identifier: CeCILL Free Software License Agreement v2.1 +// Shamrock is licensed under the CeCILL 2.1 License, see LICENSE for more information +// +// -------------------------------------------------------// + +#include "shambackends/make_ndrange.hpp" +#include "shamtest/details/TestResult.hpp" +#include "shamtest/shamtest.hpp" + +namespace { + /// Checks that the ndrange returned by `sham::make_ndrange` has the expected global and + /// local sizes for the given `wg_size` / `nthread` inputs. + void check_ndrange(u32 wg_size, u32 nthread, u32 expected_global) { + sycl::nd_range<1> ndr = sham::make_ndrange(wg_size, nthread); + + REQUIRE_EQUAL(ndr.get_global_range().size(), expected_global); + REQUIRE_EQUAL(ndr.get_local_range().size(), wg_size); + REQUIRE_EQUAL(ndr.get_group_range().size(), expected_global / wg_size); + } +} // namespace + +NEW_TEST(Unittest, "shambackends/make_ndrange.hpp:make_ndrange", 1) { + + // nthread already a multiple of wg_size -> no rounding needed + check_ndrange(32, 64, 64); + check_ndrange(32, 32, 32); + check_ndrange(16, 160, 160); + + // nthread one more than a multiple of wg_size -> rounds up to the next block + check_ndrange(32, 65, 96); + check_ndrange(16, 17, 32); + + // nthread one less than a multiple of wg_size -> rounds up to that multiple + check_ndrange(32, 63, 64); + check_ndrange(16, 15, 16); + + // nthread smaller than wg_size -> rounds up to a single full work-group + check_ndrange(32, 1, 32); + check_ndrange(1024, 5, 1024); + + // nthread == 1 with wg_size == 1 -> trivial single-thread launch + check_ndrange(1, 1, 1); + + // wg_size == 1 -> no rounding ever occurs, regardless of nthread + check_ndrange(1, 7, 7); + check_ndrange(1, 1000, 1000); + + // large values, well within u32 range, no overflow in the rounding arithmetic + check_ndrange(256, 1'000'000, 1'000'192); + check_ndrange(128, 1'048'576, 1'048'576); +} From 8ca3665e03498fe3ccbae18ff7f0096660461491 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Timoth=C3=A9e=20David--Cl=C3=A9ris?= Date: Mon, 31 Aug 2026 09:22:29 +0200 Subject: [PATCH 5/7] yolo --- .../include/shambackends/Device.hpp | 3 + .../include/shambackends/make_ndrange.hpp | 31 +++++++++++ src/shambackends/src/Device.cpp | 1 + src/tests/shambackends/make_ndrangeTests.cpp | 55 +++++++++++++++++++ 4 files changed, 90 insertions(+) create mode 100644 src/shambackends/include/shambackends/make_ndrange.hpp create mode 100644 src/tests/shambackends/make_ndrangeTests.cpp diff --git a/src/shambackends/include/shambackends/Device.hpp b/src/shambackends/include/shambackends/Device.hpp index 5259f747f6..eb75a1f38a 100644 --- a/src/shambackends/include/shambackends/Device.hpp +++ b/src/shambackends/include/shambackends/Device.hpp @@ -112,6 +112,9 @@ namespace sham { /// The number of compute units on the device uint32_t max_compute_units; + /// The maximum work-group size supported by the device + size_t max_work_group_size; + /// The maximum size of memory that can be allocated on the device in bytes uint64_t max_mem_alloc_size_dev; diff --git a/src/shambackends/include/shambackends/make_ndrange.hpp b/src/shambackends/include/shambackends/make_ndrange.hpp new file mode 100644 index 0000000000..c4f0c57fa6 --- /dev/null +++ b/src/shambackends/include/shambackends/make_ndrange.hpp @@ -0,0 +1,31 @@ +// -------------------------------------------------------// +// +// SHAMROCK code for hydrodynamics +// Copyright (c) 2021-2026 Timothée David--Cléris +// SPDX-License-Identifier: CeCILL Free Software License Agreement v2.1 +// Shamrock is licensed under the CeCILL 2.1 License, see LICENSE for more information +// +// -------------------------------------------------------// + +#pragma once + +/** + * @file make_ndrange.hpp + * @author Timothée David--Cléris (tim.shamrock@proton.me) + * @brief + * + */ + +#include "shambase/aliases_int.hpp" +#include "shambackends/sycl.hpp" + +namespace sham { + + /// Builds an `nd_range` with work-group size `wg_size`, rounding `nthread` up to the next + /// multiple of `wg_size` so every work-group launched is full. + inline sycl::nd_range<1> make_ndrange(u32 wg_size, u32 nthread) { + u32 nthread_rounded = ((nthread + wg_size - 1) / wg_size) * wg_size; + return sycl::nd_range<1>(sycl::range<1>(nthread_rounded), sycl::range<1>(wg_size)); + } + +} // namespace sham diff --git a/src/shambackends/src/Device.cpp b/src/shambackends/src/Device.cpp index a6112f067e..01aa853ded 100644 --- a/src/shambackends/src/Device.cpp +++ b/src/shambackends/src/Device.cpp @@ -361,6 +361,7 @@ namespace sham { .global_mem_cache_size = shambase::get_check_ref(global_mem_cache_size), .local_mem_size = shambase::get_check_ref(local_mem_size), .max_compute_units = shambase::get_check_ref(max_compute_units), + .max_work_group_size = shambase::get_check_ref(max_work_group_size), .max_mem_alloc_size_dev = max_alloc_dev, .max_mem_alloc_size_host = max_alloc_host, // the SYCL standard returns the alignment in bits, we convert to bytes for convenience diff --git a/src/tests/shambackends/make_ndrangeTests.cpp b/src/tests/shambackends/make_ndrangeTests.cpp new file mode 100644 index 0000000000..4d56306a00 --- /dev/null +++ b/src/tests/shambackends/make_ndrangeTests.cpp @@ -0,0 +1,55 @@ +// -------------------------------------------------------// +// +// SHAMROCK code for hydrodynamics +// Copyright (c) 2021-2026 Timothée David--Cléris +// SPDX-License-Identifier: CeCILL Free Software License Agreement v2.1 +// Shamrock is licensed under the CeCILL 2.1 License, see LICENSE for more information +// +// -------------------------------------------------------// + +#include "shambackends/make_ndrange.hpp" +#include "shamtest/details/TestResult.hpp" +#include "shamtest/shamtest.hpp" + +namespace { + /// Checks that the ndrange returned by `sham::make_ndrange` has the expected global and + /// local sizes for the given `wg_size` / `nthread` inputs. + void check_ndrange(u32 wg_size, u32 nthread, u32 expected_global) { + sycl::nd_range<1> ndr = sham::make_ndrange(wg_size, nthread); + + REQUIRE_EQUAL(ndr.get_global_range().size(), expected_global); + REQUIRE_EQUAL(ndr.get_local_range().size(), wg_size); + REQUIRE_EQUAL(ndr.get_group_range().size(), expected_global / wg_size); + } +} // namespace + +NEW_TEST(Unittest, "shambackends/make_ndrange.hpp:make_ndrange", 1) { + + // nthread already a multiple of wg_size -> no rounding needed + check_ndrange(32, 64, 64); + check_ndrange(32, 32, 32); + check_ndrange(16, 160, 160); + + // nthread one more than a multiple of wg_size -> rounds up to the next block + check_ndrange(32, 65, 96); + check_ndrange(16, 17, 32); + + // nthread one less than a multiple of wg_size -> rounds up to that multiple + check_ndrange(32, 63, 64); + check_ndrange(16, 15, 16); + + // nthread smaller than wg_size -> rounds up to a single full work-group + check_ndrange(32, 1, 32); + check_ndrange(1024, 5, 1024); + + // nthread == 1 with wg_size == 1 -> trivial single-thread launch + check_ndrange(1, 1, 1); + + // wg_size == 1 -> no rounding ever occurs, regardless of nthread + check_ndrange(1, 7, 7); + check_ndrange(1, 1000, 1000); + + // large values, well within u32 range, no overflow in the rounding arithmetic + check_ndrange(256, 1'000'000, 1'000'192); + check_ndrange(128, 1'048'576, 1'048'576); +} From 96f55d6e4ef5fe0afca1643953261f3c6f754e6e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Timoth=C3=A9e=20David--Cl=C3=A9ris?= Date: Mon, 31 Aug 2026 09:27:20 +0200 Subject: [PATCH 6/7] yolo --- .../device/details/modern_gpu_merge_sort.hpp | 4 +- .../include/shambackends/make_ndrange.hpp | 82 +++++++++++++++++-- 2 files changed, 80 insertions(+), 6 deletions(-) diff --git a/src/shamalgs/include/shamalgs/primitives/device/details/modern_gpu_merge_sort.hpp b/src/shamalgs/include/shamalgs/primitives/device/details/modern_gpu_merge_sort.hpp index 2aebfb5299..3f8b012d9d 100644 --- a/src/shamalgs/include/shamalgs/primitives/device/details/modern_gpu_merge_sort.hpp +++ b/src/shamalgs/include/shamalgs/primitives/device/details/modern_gpu_merge_sort.hpp @@ -495,6 +495,8 @@ namespace shamalgs::primitives::device::details { using Kernel = KernelBlocksort; + auto range = sham::make_check_ndrange(NT, nthreads, q); + sham::kernel_call_hndl( q, sham::MultiRef{}, @@ -505,7 +507,7 @@ namespace shamalgs::primitives::device::details { sycl::local_accessor shared_mem(1, cgh); - cgh.parallel_for(sham::make_ndrange(NT, nthreads), [=](sycl::nd_item<1> item) { + cgh.parallel_for(range, [=](sycl::nd_item<1> item) { u32 gid = item.get_global_linear_id(); u32 lid = item.get_local_linear_id(); u32 block = item.get_group_linear_id(); diff --git a/src/shambackends/include/shambackends/make_ndrange.hpp b/src/shambackends/include/shambackends/make_ndrange.hpp index c4f0c57fa6..7b30eabd09 100644 --- a/src/shambackends/include/shambackends/make_ndrange.hpp +++ b/src/shambackends/include/shambackends/make_ndrange.hpp @@ -12,20 +12,92 @@ /** * @file make_ndrange.hpp * @author Timothée David--Cléris (tim.shamrock@proton.me) - * @brief + * @brief Helper to build a `sycl::nd_range<1>` from a work-group size and a thread count * */ #include "shambase/aliases_int.hpp" +#include "shambase/assert.hpp" +#include "shambase/exception.hpp" +#include "sham/format/format.hpp" +#include "shambackends/Device.hpp" +#include "shambackends/DeviceQueue.hpp" #include "shambackends/sycl.hpp" namespace sham { - /// Builds an `nd_range` with work-group size `wg_size`, rounding `nthread` up to the next - /// multiple of `wg_size` so every work-group launched is full. - inline sycl::nd_range<1> make_ndrange(u32 wg_size, u32 nthread) { - u32 nthread_rounded = ((nthread + wg_size - 1) / wg_size) * wg_size; + /** + * @brief Builds an `nd_range` with work-group size `wg_size`, rounding `nthread` up to the + * next multiple of `wg_size` so every work-group launched is full + * + * @param wg_size the work-group size, must be > 0 + * @param nthread the number of threads to cover, must be > 0 + * @return sycl::nd_range<1> the resulting nd range + */ + inline sycl::nd_range<1> make_ndrange(size_t wg_size, size_t nthread) { + if (nthread == 0) { + throw shambase::make_except_with_loc( + sham::format("make_ndrange: nthread must be > 0 (nthread = {})", nthread)); + } + if (wg_size == 0) { + throw shambase::make_except_with_loc( + sham::format("make_ndrange: wg_size must be > 0 (wg_size = {})", wg_size)); + } + size_t nthread_rounded = ((nthread + wg_size - 1) / wg_size) * wg_size; + SHAM_ASSERT(nthread_rounded % wg_size == 0); + SHAM_ASSERT(nthread_rounded > 0); return sycl::nd_range<1>(sycl::range<1>(nthread_rounded), sycl::range<1>(wg_size)); } + /** + * @brief Checks that an `nd_range` can be launched on the given device + * + * @param in the nd range to check + * @param d the device the nd range is meant to be launched on + * @throws std::invalid_argument if the local (work-group) range of `in` exceeds the + * device's max work-group size + */ + inline void check_ndrange(const sycl::nd_range<1> &in, DeviceProperties &d) { + if (in.get_local_range().size() > d.max_work_group_size) { + throw shambase::make_except_with_loc(sham::format( + "nd_range local size ({}) exceeds the device's max work-group size ({}) on " + "device {}", + in.get_local_range().size(), + d.max_work_group_size, + d.name)); + } + } + + /** + * @brief Builds an `nd_range` with `make_ndrange` and checks that it can be launched on the + * given device + * + * @param wg_size the work-group size, must be > 0 + * @param nthread the number of threads to cover, must be > 0 + * @param d the device the nd range is meant to be launched on + * @return sycl::nd_range<1> the resulting nd range + * @throws std::invalid_argument if the resulting local (work-group) range exceeds the + * device's max work-group size + */ + inline sycl::nd_range<1> make_check_ndrange(size_t wg_size, size_t nthread, Device &d) { + auto ret = sham::make_ndrange(wg_size, nthread); + check_ndrange(ret, d.prop); + return ret; + } + + /** + * @brief Builds an `nd_range` with `make_ndrange` and checks that it can be launched on the + * device backing the given queue + * + * @param wg_size the work-group size, must be > 0 + * @param nthread the number of threads to cover, must be > 0 + * @param d the queue whose device the nd range is meant to be launched on + * @return sycl::nd_range<1> the resulting nd range + * @throws std::invalid_argument if the resulting local (work-group) range exceeds the + * device's max work-group size + */ + inline sycl::nd_range<1> make_check_ndrange(size_t wg_size, size_t nthread, DeviceQueue &d) { + return make_check_ndrange(wg_size, nthread, *d.ctx->device); + } + } // namespace sham From 82816cae296fcfdd20b369c3767b0225f3f9996e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Timoth=C3=A9e=20David--Cl=C3=A9ris?= Date: Mon, 31 Aug 2026 22:22:52 +0200 Subject: [PATCH 7/7] tini tiny progress --- .../device/details/modern_gpu_merge_sort.hpp | 54 ++++++++++++++++--- 1 file changed, 48 insertions(+), 6 deletions(-) diff --git a/src/shamalgs/include/shamalgs/primitives/device/details/modern_gpu_merge_sort.hpp b/src/shamalgs/include/shamalgs/primitives/device/details/modern_gpu_merge_sort.hpp index 3f8b012d9d..3f18758e30 100644 --- a/src/shamalgs/include/shamalgs/primitives/device/details/modern_gpu_merge_sort.hpp +++ b/src/shamalgs/include/shamalgs/primitives/device/details/modern_gpu_merge_sort.hpp @@ -698,6 +698,48 @@ namespace shamalgs::primitives::device::details { return range; } + + template + inline void DeviceLoad2ToReg(InputIt1 a_global, int aCount, + InputIt2 b_global, int bCount, int tid, T* reg, bool sync=true) { + + b_global -= aCount; + int total = aCount + bCount; + if(total >= NT * VT0) { + #pragma unroll + for(int i = 0; i < VT0; ++i) { + int index = NT * i + tid; + if(index < aCount) reg[i] = a_global[index]; + else reg[i] = b_global[index]; + } + } else { + #pragma unroll + for(int i = 0; i < VT0; ++i) { + int index = NT * i + tid; + if(index < aCount) reg[i] = a_global[index]; + else if(index < total) reg[i] = b_global[index]; + } + } + #pragma unroll + for(int i = VT0; i < VT1; ++i) { + int index = NT * i + tid; + if(index < aCount) reg[i] = a_global[index]; + else if(index < total) reg[i] = b_global[index]; + } + } + + template + inline void DeviceLoad2ToShared( sycl::nd_item<1> &item,InputIt1 a_global, int aCount, + InputIt2 b_global, int bCount, int tid, T* shared, bool sync) { + + T reg[VT1]; + DeviceLoad2ToReg(a_global, aCount, b_global, bCount, tid, + reg, false); + DeviceRegToShared(item ,reg, tid, shared, sync); + } + template< int NT, int VT, @@ -706,22 +748,22 @@ namespace shamalgs::primitives::device::details { typename It2, typename T, typename Comp> - MGPU_DEVICE void DeviceMergeKeysIndices( + inline void DeviceMergeKeysIndices( It1 a_global, int aCount, It2 b_global, int bCount, - int4 range, + sycl::vec range, int tid, T *keys_shared, T *results, int *indices, Comp comp) { - int a0 = range.x; - int a1 = range.y; - int b0 = range.z; - int b1 = range.w; + int a0 = range.x(); + int a1 = range.y(); + int b0 = range.z(); + int b1 = range.w(); if (LoadExtended) { bool extended = (a1 < aCount) && (b1 < bCount);