You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
Add a dispatch to count size for int32 when the number of elements fits within (int32 / 2). Speeds up large arrays of int8 elements by around 5%, close to thrust's performance. Does not dispatch for int64 elements because a performance regression was measured for that case.
Benchmarks were done on RTX 3060.
With guarantee:
n=2^24 threshold=-128 reset=0 thrust_partition min=293.740 median=298.340 us type=I8
n=2^24 threshold=-128 reset=0 pstl_partition min=287.040 median=293.671 us type=I8 guarantee=1
n=2^24 threshold=-128 reset=0 thrust_partition min=458.820 median=467.310 us type=I16
n=2^24 threshold=-128 reset=0 pstl_partition min=443.750 median=452.831 us type=I16 guarantee=1
n=2^24 threshold=-128 reset=0 thrust_partition min=835.321 median=841.091 us type=I32
n=2^24 threshold=-128 reset=0 pstl_partition min=827.601 median=835.821 us type=I32 guarantee=1
n=2^24 threshold=-128 reset=0 thrust_partition min=1638.382 median=1645.371 us type=I64
n=2^24 threshold=-128 reset=0 pstl_partition min=1631.601 median=1639.571 us type=I64 guarantee=1
Without guarantee:
n=2^24 threshold=-128 reset=0 thrust_partition min=277.470 median=281.200 us type=I8
n=2^24 threshold=-128 reset=0 pstl_partition min=290.080 median=297.370 us type=I8 guarantee=0
n=2^24 threshold=-128 reset=0 thrust_partition min=454.730 median=463.260 us type=I16
n=2^24 threshold=-128 reset=0 pstl_partition min=450.720 median=461.190 us type=I16 guarantee=0
n=2^24 threshold=-128 reset=0 thrust_partition min=831.331 median=837.151 us type=I32
n=2^24 threshold=-128 reset=0 pstl_partition min=831.321 median=838.091 us type=I32 guarantee=0
n=2^24 threshold=-128 reset=0 thrust_partition min=1637.621 median=1642.771 us type=I64
n=2^24 threshold=-128 reset=0 pstl_partition min=1630.171 median=1636.321 us type=I64 guarantee=0
Benchmark script
#include <algorithm>
#include <chrono>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cuda/execution>
#include <cuda/memory_resource>
#include <cuda/std/execution>
#include <cuda/std/iterator>
#include <cuda/stream>
#include <map>
#include <string>
#include <thrust/device_vector.h>
#include <thrust/execution_policy.h>
#include <thrust/partition.h>
#include <vector>
void check(cudaError_t e) {
if (e != cudaSuccess) {
std::fprintf(stderr, "%s\n", cudaGetErrorString(e));
std::exit(1);
}
}
struct cache {
using value_type = char;
std::multimap<size_t, void *> free;
std::vector<void *> all;
char *allocate(std::ptrdiff_t bytes) {
auto it = free.find(bytes);
if (it != free.end()) {
auto p = it->second;
free.erase(it);
return static_cast<char *>(p);
}
void *p;
check(cudaMalloc(&p, bytes));
all.push_back(p);
return static_cast<char *>(p);
}
void deallocate(char *p, size_t bytes) { free.emplace(bytes, p); }
void *allocate_sync(size_t b, size_t) { return allocate(b); }
void deallocate_sync(void *p, size_t b, size_t) {
deallocate(static_cast<char *>(p), b);
}
void *allocate(cuda::stream_ref, size_t b, size_t) { return allocate(b); }
void deallocate(cuda::stream_ref, void *p, size_t b, size_t) {
deallocate(static_cast<char *>(p), b);
}
friend bool operator==(const cache &a, const cache &b) { return &a == &b; }
friend bool operator!=(const cache &a, const cache &b) { return !(a == b); }
friend void get_property(const cache &,
cuda::mr::device_accessible) noexcept {}
~cache() {
for (auto p : all)
cudaFree(p);
}
};
// Keep values in [-128, 127] for every element type. Threshold=128 selects all.
template <class T> struct pred {
int threshold;
__host__ __device__ bool operator()(T x) const { return x < threshold; }
};
template <class T> __global__ void initialize(T *data, size_t n) {
for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < n;
i += size_t(gridDim.x) * blockDim.x)
data[i] = static_cast<T>(int(i % 256) - 128);
}
template <class T> void validate(const T *data, size_t n, pred<T> predicate) {
std::vector<T> host(std::min(n, size_t(1) << 20));
size_t counts[256] = {};
bool rejected = false;
for (size_t offset = 0; offset < n; offset += host.size()) {
size_t count = std::min(host.size(), n - offset);
check(cudaMemcpy(host.data(), data + offset, count * sizeof(T),
cudaMemcpyDeviceToHost));
for (size_t i = 0; i < count; ++i) {
T value = host[i];
if (value < T(-128) || value > T(127)) {
std::fprintf(stderr, "invalid value\n");
std::exit(1);
}
++counts[int(value) + 128];
if (!predicate(value))
rejected = true;
else if (rejected) {
std::fprintf(stderr, "invalid partition\n");
std::exit(1);
}
}
}
for (auto count : counts)
if (count != n / 256) {
std::fprintf(stderr, "invalid multiset\n");
std::exit(1);
}
}
constexpr int guarantee_power = 29;
template <class T> int run(int power, int threshold, int reset, bool use_guarantee) {
const size_t n = size_t(1) << power;
thrust::device_vector<T> source(n), thrust_input(n), pstl_input(n);
auto src = thrust::raw_pointer_cast(source.data());
initialize<<<512, 256>>>(src, n);
check(cudaGetLastError());
check(cudaDeviceSynchronize());
check(cudaMemcpy(thrust::raw_pointer_cast(thrust_input.data()), src,
n * sizeof(T), cudaMemcpyDeviceToDevice));
check(cudaMemcpy(thrust::raw_pointer_cast(pstl_input.data()), src,
n * sizeof(T), cudaMemcpyDeviceToDevice));
cache alloc;
auto tp = thrust::cuda::par(alloc);
auto cp = cuda::execution::gpu.with(cuda::mr::get_memory_resource, alloc);
auto bounded_cp = cp.with(cuda::execution::guarantee(
cuda::execution::max_num_items<(size_t{1} << guarantee_power)>{}));
pred<T> predicate{threshold};
std::vector<double> times[2];
// Separate, identically initialized inputs. Five warmups and 40 measured
// calls.
for (int round = 0; round < 45; ++round)
for (int j = 0; j < 2; ++j) {
int mode = (j + round) % 2;
auto &input = mode == 0 ? thrust_input : pstl_input;
if (reset)
check(cudaMemcpy(thrust::raw_pointer_cast(input.data()), src,
n * sizeof(T), cudaMemcpyDeviceToDevice));
auto start = std::chrono::steady_clock::now();
auto end =
mode == 0
? thrust::partition(tp, input.begin(), input.end(), predicate)
: use_guarantee
? cuda::std::partition(bounded_cp, input.begin(), input.end(), predicate)
: cuda::std::partition(cp, input.begin(), input.end(), predicate);
auto finish = std::chrono::steady_clock::now();
if (end - input.begin() != std::ptrdiff_t(n / 256 * (threshold + 128))) {
std::fprintf(stderr, "wrong selected count\n");
return 1;
}
if (round >= 5)
times[mode].push_back(
std::chrono::duration<double, std::micro>(finish - start).count());
}
validate(thrust::raw_pointer_cast(thrust_input.data()), n, predicate);
validate(thrust::raw_pointer_cast(pstl_input.data()), n, predicate);
const char *names[] = {"thrust_partition", "pstl_partition"};
for (int mode = 0; mode < 2; ++mode) {
std::sort(times[mode].begin(), times[mode].end());
std::printf(
"n=2^%d threshold=%d reset=%d %s min=%.3f median=%.3f us type=I%zu%s\n",
power, threshold, reset, names[mode], times[mode].front(),
times[mode][times[mode].size() / 2], sizeof(T) * 8,
mode == 0 ? "" : (use_guarantee ? " guarantee=1" : " guarantee=0"));
}
return 0;
}
int main(int argc, char **argv) {
int power = argc > 1 ? std::atoi(argv[1]) : 24;
int threshold = argc > 2 ? std::atoi(argv[2]) : -128;
int reset = argc > 3 ? std::atoi(argv[3]) : 0;
const std::string type = argc > 4 ? argv[4] : "i8";
const std::string guarantee = argc > 5 ? argv[5] : "0";
int bits = 0;
if (type == "i8" || type == "I8" || type == "8")
bits = 8;
if (type == "i16" || type == "I16" || type == "16")
bits = 16;
if (type == "i32" || type == "I32" || type == "32")
bits = 32;
if (type == "i64" || type == "I64" || type == "64")
bits = 64;
if (argc > 6 || power < 8 || power > 30 || threshold < -128 ||
threshold > 128 || (reset != 0 && reset != 1) || bits == 0 ||
(guarantee != "0" && guarantee != "1")) {
std::fprintf(stderr,
"usage: %s [power=8..30] [threshold=-128..128] [reset=0|1] "
"[type=i8|i16|i32|i64] [guarantee=0|1]\n",
argv[0]);
return 2;
}
const bool use_guarantee = guarantee == "1";
if (use_guarantee && power > guarantee_power) {
std::fprintf(stderr, "guarantee=1 requires power <= %d\n", guarantee_power);
return 2;
}
switch (bits) {
case 8:
return run<std::int8_t>(power, threshold, reset, use_guarantee);
case 16:
return run<std::int16_t>(power, threshold, reset, use_guarantee);
case 32:
return run<std::int32_t>(power, threshold, reset, use_guarantee);
case 64:
return run<std::int64_t>(power, threshold, reset, use_guarantee);
}
return 2;
}
Checklist
New or existing tests cover these changes.
The documentation is up to date with these changes.
It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.
No actionable comments were generated in the recent review.
Summary
Summary by CodeRabbit
New Features
Added a CUDA execution guarantee for specifying the maximum number of items an algorithm will process. Algorithms can use this bound when handling item counts.
Bug Fixes
Improved range-size handling in CUDA partition operations, including support for counts beyond int limits while retaining existing handling for 64-bit integral values.
Tests
Added coverage for partition results around integer count boundaries across signed and unsigned integer types and floating-point values, including cases with different partition outcomes.
CUDA partition dispatch now passes an element count to __par_impl and selects int or size_t offsets using the value type and policy item-count bound. The change adds the max_num_items execution guarantee and tests partitioning around the int offset threshold.
Changes: CUDA partition dispatch
Layer / File(s)
Summary
Maximum item-count guarantee libcudacxx/include/cuda/__execution/max_num_items.h, libcudacxx/include/cuda/execution, libcudacxx/include/cuda/execution.max_num_items.h
max_num_items<_MaxItems> exposes the bound as an integral constant. The execution headers export the guarantee.
__par_impl accepts a count. The dispatcher checks the policy bound and selects int or size_t offsets based on the value type and bound.
Partition count boundary test libcudacxx/test/libcudacxx/std/algorithms/alg.modifying/alg.partitions/pstl_partition.cu
Tests apply the item-count guarantee and cover counts immediately below, at, and above half of int’s maximum. They check the returned boundary and both resulting groups.
Assessment against linked issues
Objective
Addressed
Explanation
Investigate and avoid the reported performance regressions for large arrays [#8112]
❓
The dispatch logic changes, but the provided context has no benchmark results confirming whether the reported regressions are avoided.
Suggested reviewers:miscco
Change: Bug fix · Severity of issue fixed: Low
Merge Risk: Minimal · up to c305a
This change lets CUDA partition use 32-bit counting when the caller guarantees a small enough input, which targets the reported int8 slowdown. Policies without the guarantee keep the previous behavior. The earlier naming concern in the test is resolved, and no actionable merge risk remains.
Security architecture risk: Minimal · up to c305a
The optimization is opt-in: existing callers retain wide offsets, while calls satisfying the new size promise fit the narrower representation. No introduced security issue or materially worsened ownership or failure-containment behavior was established.
Retained concerns
No architecture-level concerns identified.
Security review details
Security Blast Radius
inferred — The changed behavior is scoped to partition invocations whose application-supplied policy carries the new guarantee. Those invocations continue to operate on the caller’s iterator range using the policy’s existing CUDA context, stream, and memory resource; the guarantee itself grants no additional authority.
Trust Boundaries and Controls
observed — The size promise is not an unconditional input-validation boundary: _CCCL_ASSERT can compile to a no-op when host assertions are disabled. Callers must establish the documented bound themselves. For valid calls, count is no greater than the declared bound, which is below INT_MAX/2 on the narrow path, so the int conversion is representable.
- [ ] Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 10 remain after this review.
The reason will be displayed to describe this comment to others. Learn more.
Caution
Some comments are outside the diff and can’t be posted inline due to GitHub limitations.
⚠️ Outside diff range comments (1)
Major · important: Make the padded partition scan safe at the 32-bit limit. · partition.h:149-160
libcudacxx/include/cuda/std/__pstl/cuda/partition.h:149-160 Major · Functional Correctness · Heavy lift
important: Make the padded partition scan safe at the 32-bit limit.
On SM100, an INT_MAX-element int16_t range with 3,000 rejected elements overflows CUB’s signed scan with the new dispatch. This is undefined behavior and may corrupt the partition or returned iterator. Falling back to size_t avoids this example, but not all predicates: both dispatches use a 32-bit per-partition scan.
130-153: Trivial · Maintainability & Code Quality · Heavy lift
suggestion: Test the count type at the dispatcher.
If the dispatcher passes size_t again, these result checks can still pass while the performance regression returns. Add a dispatcher-level test that records the type passed to __par_impl. Check INT_MAX - 1, INT_MAX, INT_MAX + 1, and the 64-bit integer exception. Test-only specializations can avoid large allocations.
suggestion: Test a mixed partition at the count boundary.
A large-count wrapper regression that returns the input length instead of the selected count would pass this fixture. Add a case with rejected values before selected values, then check the returned boundary and both output ranges. The CUB large-input test does not exercise the libcudacxx wrapper.
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.
The reason will be displayed to describe this comment to others. Learn more.
Thanks for the PR.
Unfortunately, this is something we will need to put on hold for now, because it has nontrivial side effect.
Having a runtime dispatch for both kernels will instantiate both of those kernels, which unfortunately greatly increases binary size and compile time.
Thrust uses this approach but we have users that have a lot of issues with this.
I believe the proper approach is to handle this via the guarantees framework, where a user can give us the guarantee that an input sequence does not exceed a given length.
@gevtushenko for visibility because we discussed this at some point but it is not on the roadmap
@miscco Thanks for the explanation. For the guarantee framework, would adding a max_num_items execution::guarantee or execution::require be in the right direction?
auto policy = cuda::execution::gpu.with(
cuda::execution::guarantee(cuda::execution::max_num_items<(1u << 24)>{}));
auto middle = cuda::std::partition(policy, first, last, pred);
Or should I leave the PR on hold or close it for now, since it isn't on the roadmap?
suggestion: Exercise narrow offsets near the cutoff.
This policy has no max_num_items guarantee. All three counts therefore use wide offsets. The guarantee-bearing section covers only 1000 items.
Add cases with compile-time bounds of cutoff - 1, cutoff, and cutoff + 1, with matching counts. This covers the narrow path near its limit and the transition to wide offsets. As per path instructions, “Cover relevant edge cases.” (raw.githubusercontent.com)
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
miscco
left a comment
Description
closes
Fixes #8112
Add a dispatch to count size for int32 when the number of elements fits within (int32 / 2). Speeds up large arrays of int8 elements by around 5%, close to thrust's performance. Does not dispatch for int64 elements because a performance regression was measured for that case.
Benchmarks were done on RTX 3060.
With guarantee:
n=2^24 threshold=-128 reset=0 thrust_partition min=293.740 median=298.340 us type=I8
n=2^24 threshold=-128 reset=0 pstl_partition min=287.040 median=293.671 us type=I8 guarantee=1
n=2^24 threshold=-128 reset=0 thrust_partition min=458.820 median=467.310 us type=I16
n=2^24 threshold=-128 reset=0 pstl_partition min=443.750 median=452.831 us type=I16 guarantee=1
n=2^24 threshold=-128 reset=0 thrust_partition min=835.321 median=841.091 us type=I32
n=2^24 threshold=-128 reset=0 pstl_partition min=827.601 median=835.821 us type=I32 guarantee=1
n=2^24 threshold=-128 reset=0 thrust_partition min=1638.382 median=1645.371 us type=I64
n=2^24 threshold=-128 reset=0 pstl_partition min=1631.601 median=1639.571 us type=I64 guarantee=1
Without guarantee:
n=2^24 threshold=-128 reset=0 thrust_partition min=277.470 median=281.200 us type=I8
n=2^24 threshold=-128 reset=0 pstl_partition min=290.080 median=297.370 us type=I8 guarantee=0
n=2^24 threshold=-128 reset=0 thrust_partition min=454.730 median=463.260 us type=I16
n=2^24 threshold=-128 reset=0 pstl_partition min=450.720 median=461.190 us type=I16 guarantee=0
n=2^24 threshold=-128 reset=0 thrust_partition min=831.331 median=837.151 us type=I32
n=2^24 threshold=-128 reset=0 pstl_partition min=831.321 median=838.091 us type=I32 guarantee=0
n=2^24 threshold=-128 reset=0 thrust_partition min=1637.621 median=1642.771 us type=I64
n=2^24 threshold=-128 reset=0 pstl_partition min=1630.171 median=1636.321 us type=I64 guarantee=0
Benchmark script
Checklist