Skip to content

parallel_reduce_sum: drop caller scratch, keep default pool cached #10

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
shsahiti merged 1 commit into from
Sep 30, 2026
Merged

parallel_reduce_sum: drop caller scratch, keep default pool cached #10

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 2 additions & 22 deletions tests/reduce/cases.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,6 @@ struct nonconst_sum {
template <typename T, typename F>
concept supports_sum = requires(const xpu::range<2uz>& range, T* output, F contribution) {
xpu::parallel_reduce_sum(range, output, contribution);
xpu::parallel_reduce_sum_bytes<T>(range, contribution);
};

static_assert(supports_sum<int, sum_values<int>>);
Expand All @@ -54,16 +53,6 @@ inline auto run_sum_type() -> int {
const auto range = xpu::range<2uz>{
{2uz, 1uz}, {6uz, 1uz + 3uz * count}, {2uz, 3uz}
};
const auto required_bytes{xpu::parallel_reduce_sum_bytes<T>(range, contribution)};
const auto empty{count == 0uz};
const auto invalid_empty_query{empty && required_bytes != 0uz};

if (invalid_empty_query) {
const auto failure{test::fail("empty reduction requested scratch storage")};

return failure;
}

auto expected{T{}};

for (auto row{2uz}; row < 6uz; row += 2uz) {
Expand All @@ -72,16 +61,11 @@ inline auto run_sum_type() -> int {
}
}

const auto capacity{required_bytes + 64uz};
auto scratch{xpu::buffer<std::byte>{capacity}};
const auto initial{T{91}};
xpu::copy_n(output.data(), &initial, 1uz);

for (auto repeat{0uz}; repeat < 2uz; ++repeat) {
xpu::parallel_reduce_sum(
range, output.data(), contribution,
empty ? nullptr : scratch.data(), empty ? 0uz : capacity
);
xpu::parallel_reduce_sum(range, output.data(), contribution);

auto actual{T{}};
xpu::copy_n(&actual, output.data(), 1uz);
Expand Down Expand Up @@ -116,13 +100,9 @@ struct sum_coordinates {
template <std::size_t dims>
inline auto check_coordinate_sum(const xpu::range<dims>& range, int expected) -> int {
const auto contribution{sum_coordinates{}};
const auto scratch_bytes{xpu::parallel_reduce_sum_bytes<int>(range, contribution)};
auto scratch{xpu::buffer<std::byte>{scratch_bytes}};
auto output{xpu::buffer<int>{1uz}};

xpu::parallel_reduce_sum(
range, output.data(), contribution, scratch.data(), scratch_bytes
);
xpu::parallel_reduce_sum(range, output.data(), contribution);

auto actual{0};
xpu::copy_n(&actual, output.data(), 1uz);
Expand Down
94 changes: 26 additions & 68 deletions xpu/launch.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
#pragma once

#include <cstddef>
#include <cstdint>
#include <limits>
#include <type_traits>
#include <utility>
#include <xpu/config.hpp>
Expand Down Expand Up @@ -40,6 +42,24 @@ inline auto device_SMs() -> unsigned int {
return cached;
}

inline auto retain_default_pool() -> void {
static const auto once{[] {
auto device{0};
auto pool{cudaMemPool_t{}};
auto threshold{std::numeric_limits<std::uint64_t>::max()};
xpu::cu_check(cudaGetDevice(&device));
xpu::cu_check(cudaDeviceGetDefaultMemPool(&pool, device));
xpu::cu_check(cudaMemPoolSetAttribute(
pool, cudaMemPoolAttrReleaseThreshold, &threshold
));

return true;
}()
};

static_cast<void>(once);
}

template <typename Kernel> [[nodiscard]]
inline auto wave_blocks(Kernel kernel, dim3 threads, std::size_t smem = 0uz) -> unsigned int {
const auto thread_budget{xpu::detail::checked_mul(
Expand Down Expand Up @@ -249,17 +269,6 @@ inline auto reduction_input(
return input;
}

template <arithmetic T, typename Input> [[nodiscard]]
inline auto reduction_bytes(Input input, std::ptrdiff_t count) -> std::size_t {
auto required_bytes{0uz};

xpu::cu_check(cub::DeviceReduce::Sum(
nullptr, required_bytes, input, static_cast<T*>(nullptr), count
));

return required_bytes;
}

template <std::size_t dims, typename F>
struct parallel_for_function {
xpu::range<dims> range;
Expand Down Expand Up @@ -476,34 +485,14 @@ inline auto parallel_reduce_sum_impl(
const xpu::range<dims>& range,
std::size_t total,
T* output_ptr,
F&& contribution,
[[maybe_unused]] void* scratch = nullptr,
[[maybe_unused]] std::size_t scratch_bytes = 0uz
F&& contribution
) -> void {
#if defined(XPU_CUDA)
using function_t = std::decay_t<F>;

const auto transform{reduction_contribution<dims, T, function_t>{
range, function_t{std::forward<F>(contribution)}
}};

const auto indices{thrust::counting_iterator<std::size_t>{0uz}};
const auto input{thrust::make_transform_iterator(indices, transform)};

const auto input{reduction_input<T>(range, std::forward<F>(contribution))};
const auto count{xpu::detail::checked_cast<std::ptrdiff_t>(total)};
const auto required_bytes{reduction_bytes<T>(input, count)};

const auto insufficient_scratch{
scratch == nullptr || scratch_bytes < required_bytes
};

if (insufficient_scratch) {
xpu::detail::checked_error("insufficient reduction scratch storage");
}

xpu::cu_check(cub::DeviceReduce::Sum(
scratch, scratch_bytes, input, output_ptr, count
));
xpu::detail::retain_default_pool();
xpu::cu_check(cub::DeviceReduce::Sum(input, output_ptr, count));
#else
if (range.step[dims - 1uz] == 1uz) {
*output_ptr = parallel_reduce_cpu<true, T>(range, total, contribution);
Expand All @@ -525,43 +514,12 @@ inline auto parallel_for(
detail::parallel_for_impl(range, std::forward<F>(fcn));
}

template <arithmetic T, std::size_t dims, typename F>
requires detail::sum_contribution<F, T, dims>
[[nodiscard]]
inline auto parallel_reduce_sum_bytes(
[[maybe_unused]] const xpu::range<dims>& range,
[[maybe_unused]] F&& contribution
) -> std::size_t {
static_assert(dims > 0uz, "ERROR: Dimension must be greater than 0.");

auto required_bytes{0uz};

#if defined(XPU_CUDA)
const auto total{xpu::detail::num_itrs(range)};
const auto empty{total == 0uz};

if (empty) {

return required_bytes;
}

const auto input{detail::reduction_input<T>(range, std::forward<F>(contribution))};
const auto count{xpu::detail::checked_cast<std::ptrdiff_t>(total)};

required_bytes = detail::reduction_bytes<T>(input, count);
#endif

return required_bytes;
}

template <std::size_t dims, arithmetic T, typename F>
requires detail::sum_contribution<F, T, dims>
inline auto parallel_reduce_sum(
const xpu::range<dims>& range,
T* output_ptr,
F&& contribution,
void* scratch = nullptr,
std::size_t scratch_bytes = 0uz
F&& contribution
) -> void {
static_assert(dims > 0uz, "ERROR: Dimension must be greater than 0.");

Expand All @@ -579,7 +537,7 @@ inline auto parallel_reduce_sum(
}

detail::parallel_reduce_sum_impl(
range, total, output_ptr, std::forward<F>(contribution), scratch, scratch_bytes
range, total, output_ptr, std::forward<F>(contribution)
);
}

Expand Down