From 86c41dbfbf6a33d136983978d8ac62acf4b90c18 Mon Sep 17 00:00:00 2001 From: charan-003 Date: Wed, 19 Nov 2025 18:37:32 -0700 Subject: [PATCH] fix omp scan bug (#6560) Co-authored-by: Bernhard Manfred Gruber (cherry picked from commit bb8a5b9e26010f932042a302bddbccea8621819a) --- thrust/testing/scan.cu | 188 +++++++++++++++++++++++-- thrust/testing/transform_scan.cu | 140 ++++++++++++++++++ thrust/thrust/system/omp/detail/scan.h | 35 ++++- 3 files changed, 346 insertions(+), 17 deletions(-) diff --git a/thrust/testing/scan.cu b/thrust/testing/scan.cu index 3602f6214306..f692e41149d8 100644 --- a/thrust/testing/scan.cu +++ b/thrust/testing/scan.cu @@ -7,10 +7,13 @@ #include #include #include +#include #include #include +#include + #include template @@ -663,24 +666,40 @@ struct Int _CCCL_HOST_DEVICE explicit Int(int num) : i(num) {} - _CCCL_HOST_DEVICE Int() - : i{} - {} - _CCCL_HOST_DEVICE Int operator+(Int const& o) const + + // TODO(bgruber): I think we should not need this constructors, but CUB fails to compile without it + Int() = default; + + _CCCL_HOST_DEVICE friend Int operator+(Int const& a, Int const& b) { - return Int{this->i + o.i}; + return Int{a.i + b.i}; } -}; -void TestInclusiveScanWithUserDefinedType() -{ - thrust::device_vector vec(5, Int{1}); + _CCCL_HOST_DEVICE friend bool operator==(Int const& a, Int const& b) + { + return a.i == b.i; + } - thrust::inclusive_scan(thrust::device, vec.cbegin(), vec.cend(), vec.begin()); + friend std::ostream& operator<<(std::ostream& os, Int const& a) + { + return os << "Int{" << a.i << "}"; + } +}; - ASSERT_EQUAL(static_cast(vec.back()).i, 5); +void TestScanWithUserDefinedTypeAndInit() +{ + { + thrust::device_vector vec(5, Int{1}); + thrust::exclusive_scan(thrust::device, vec.cbegin(), vec.cend(), vec.begin(), Int{100}, ::cuda::std::plus()); + ASSERT_EQUAL(vec, (thrust::device_vector{Int{100}, Int{101}, Int{102}, Int{103}, Int{104}})); + } + { + thrust::device_vector vec(5, Int{1}); + thrust::inclusive_scan(thrust::device, vec.cbegin(), vec.cend(), vec.begin(), Int{100}, ::cuda::std::plus()); + ASSERT_EQUAL(vec, (thrust::device_vector{Int{101}, Int{102}, Int{103}, Int{104}, Int{105}})); + } } -DECLARE_UNITTEST(TestInclusiveScanWithUserDefinedType); +DECLARE_UNITTEST(TestScanWithUserDefinedTypeAndInit); // Represents a permutation as a tuple of integers, see also: https://en.wikipedia.org/wiki/Permutation // We need a distinct type (instead of an alias) for operator<< to be found via ADL @@ -768,3 +787,148 @@ void TestInclusiveScanWithNonCommutativeOp() {4, 0, 2, 1, 3}})); } DECLARE_UNITTEST(TestInclusiveScanWithNonCommutativeOp); + +// Adapted from issue: https://github.com/NVIDIA/cccl/issues/6317 +void TestScanBug6317() +{ + using T = unsigned int; + const int n = 1729; + { + thrust::device_vector s = unittest::random_integers(n); + thrust::device_vector d(n); + const auto r = thrust::inclusive_scan(s.cbegin(), s.cend(), d.begin(), ::cuda::std::multiplies<>{}); + ASSERT_EQUAL((d.end() == r), true); + std::partial_sum(s.cbegin(), s.cend(), s.begin(), std::multiplies<>{}); + ASSERT_EQUAL(s, d); + } + { + thrust::device_vector s = unittest::random_integers(n); + thrust::device_vector d(n); + const auto r = thrust::exclusive_scan(s.cbegin(), s.cend(), d.begin(), 42, ::cuda::std::multiplies<>{}); + ASSERT_EQUAL((d.end() == r), true); + thrust::exclusive_scan(s.cbegin(), s.cend(), s.begin(), 42, ::cuda::std::multiplies<>{}); + ASSERT_EQUAL(s, d); + } +} +DECLARE_UNITTEST(TestScanBug6317); + +// Test edge cases for parallel scan with non-additive operations +void TestScanEdgeCases() +{ + // Test 1: Large array with inclusive_scan, multiplies, and init + // Tests Step 1 fix with many threads + { + const int n = 10000; + thrust::device_vector d_input(n); + using thrust::placeholders::_1; + thrust::tabulate(d_input.begin(), d_input.end(), (_1 % 5) + 1); + + thrust::device_vector d_output(n); + auto r = thrust::inclusive_scan(d_input.begin(), d_input.end(), d_output.begin(), 2, ::cuda::std::multiplies<>{}); + ASSERT_EQUAL((d_output.end() == r), true); + + // Verify with host scan + thrust::host_vector h_input = d_input; + thrust::host_vector h_output(n); + thrust::inclusive_scan(h_input.begin(), h_input.end(), h_output.begin(), 2, ::cuda::std::multiplies<>{}); + + ASSERT_EQUAL(d_output, h_output); + } + + // Test 2: Boundary exactly at parallel_scan_threshold (1024) + { + const int n = 1024; + thrust::device_vector d_input(n); + using thrust::placeholders::_1; + thrust::tabulate(d_input.begin(), d_input.end(), _1 + 1); + + thrust::device_vector d_output(n); + auto r = thrust::inclusive_scan(d_input.begin(), d_input.end(), d_output.begin(), 10, ::cuda::std::multiplies<>{}); + ASSERT_EQUAL((d_output.end() == r), true); + + thrust::host_vector h_input = d_input; + thrust::host_vector h_output(n); + thrust::inclusive_scan(h_input.begin(), h_input.end(), h_output.begin(), 10, ::cuda::std::multiplies<>{}); + + ASSERT_EQUAL(d_output, h_output); + } + + // Test 3: Below threshold (1023) should use serial path + { + const int n = 1023; + thrust::device_vector d_input(n); + using thrust::placeholders::_1; + thrust::tabulate(d_input.begin(), d_input.end(), (_1 % 3) + 1); + + thrust::device_vector d_output(n); + auto r = thrust::inclusive_scan(d_input.begin(), d_input.end(), d_output.begin(), 5, ::cuda::std::multiplies<>{}); + ASSERT_EQUAL((d_output.end() == r), true); + + thrust::host_vector h_input = d_input; + thrust::host_vector h_output(n); + thrust::inclusive_scan(h_input.begin(), h_input.end(), h_output.begin(), 5, ::cuda::std::multiplies<>{}); + + ASSERT_EQUAL(d_output, h_output); + } + + // Test 4: Very small array (edge case for block distribution) + { + thrust::device_vector d_input = {3, 7}; + thrust::device_vector d_output(2); + + auto r = thrust::inclusive_scan(d_input.begin(), d_input.end(), d_output.begin(), 2, ::cuda::std::multiplies<>{}); + ASSERT_EQUAL((d_output.end() == r), true); + + thrust::device_vector expected = {6, 42}; + ASSERT_EQUAL(d_output, expected); + } + + // Test 5: exclusive_scan with large array and multiplies + { + const int n = 10000; + thrust::device_vector d_input(n); + using thrust::placeholders::_1; + thrust::tabulate(d_input.begin(), d_input.end(), (_1 % 3) + 1); + + thrust::device_vector d_output(n); + auto r = thrust::exclusive_scan(d_input.begin(), d_input.end(), d_output.begin(), 5, ::cuda::std::multiplies<>{}); + ASSERT_EQUAL((d_output.end() == r), true); + + thrust::host_vector h_input = d_input; + thrust::host_vector h_output(n); + thrust::exclusive_scan(h_input.begin(), h_input.end(), h_output.begin(), 5, ::cuda::std::multiplies<>{}); + + ASSERT_EQUAL(d_output, h_output); + } + + // Test 6: exclusive_scan at boundary (1024 elements) + { + const int n = 1024; + thrust::device_vector d_input(n); + using thrust::placeholders::_1; + thrust::tabulate(d_input.begin(), d_input.end(), (_1 % 5) + 1); + + thrust::device_vector d_output(n); + auto r = thrust::exclusive_scan(d_input.begin(), d_input.end(), d_output.begin(), 3, ::cuda::std::multiplies<>{}); + ASSERT_EQUAL((d_output.end() == r), true); + + thrust::host_vector h_input = d_input; + thrust::host_vector h_output(n); + thrust::exclusive_scan(h_input.begin(), h_input.end(), h_output.begin(), 3, ::cuda::std::multiplies<>{}); + + ASSERT_EQUAL(d_output, h_output); + } + + // Test 7: exclusive_scan with very small array + { + thrust::device_vector d_input = {2, 4}; + thrust::device_vector d_output(2); + + auto r = thrust::exclusive_scan(d_input.begin(), d_input.end(), d_output.begin(), 3, ::cuda::std::multiplies<>{}); + ASSERT_EQUAL((d_output.end() == r), true); + + thrust::device_vector expected = {3, 6}; + ASSERT_EQUAL(d_output, expected); + } +} +DECLARE_UNITTEST(TestScanEdgeCases); diff --git a/thrust/testing/transform_scan.cu b/thrust/testing/transform_scan.cu index c70cc5acfed2..0d0d1c56be8d 100644 --- a/thrust/testing/transform_scan.cu +++ b/thrust/testing/transform_scan.cu @@ -3,8 +3,12 @@ #include #include #include +#include #include +#include +#include + #include template @@ -400,3 +404,139 @@ void TestValueCategoryDeduction() ASSERT_EQUAL((thrust::device_vector{0, 5, 5, 5, 8, 8, 8, 8, 8, 8}), vec); } DECLARE_GENERIC_UNITTEST(TestValueCategoryDeduction); + +// User-defined Int type with poison default constructor +struct TransformInt +{ + int value; + + _CCCL_HOST_DEVICE TransformInt() + : value(999) + {} + _CCCL_HOST_DEVICE TransformInt(int v) + : value(v) + {} + + _CCCL_HOST_DEVICE friend TransformInt operator+(TransformInt a, TransformInt b) + { + return TransformInt(a.value + b.value); + } + _CCCL_HOST_DEVICE friend TransformInt operator*(TransformInt a, TransformInt b) + { + return TransformInt(a.value * b.value); + } + _CCCL_HOST_DEVICE friend bool operator==(TransformInt a, TransformInt b) + { + return a.value == b.value; + } + friend std::ostream& operator<<(std::ostream& os, TransformInt i) + { + return os << i.value; + } +}; + +// Test edge cases for parallel transform_scan with non-additive operations +void TestTransformScanEdgeCases() +{ + // Test 1: Large array with transform_inclusive_scan, multiplies, and init + { + const int n = 10000; + thrust::device_vector d_input(n); + using thrust::placeholders::_1; + thrust::tabulate(d_input.begin(), d_input.end(), (_1 % 5) + 1); + + thrust::device_vector d_output(n); + auto r = thrust::transform_inclusive_scan( + d_input.begin(), d_input.end(), d_output.begin(), ::cuda::std::negate(), 2, ::cuda::std::multiplies<>{}); + ASSERT_EQUAL((d_output.end() == r), true); + + // Verify with host scan + thrust::host_vector h_input = d_input; + thrust::host_vector h_output(n); + thrust::transform_inclusive_scan( + h_input.begin(), h_input.end(), h_output.begin(), ::cuda::std::negate(), 2, ::cuda::std::multiplies<>{}); + + ASSERT_EQUAL(d_output, h_output); + } + + // Test 2: transform_exclusive_scan with multiplies on large array + { + const int n = 10000; + thrust::device_vector d_input(n); + using thrust::placeholders::_1; + thrust::tabulate(d_input.begin(), d_input.end(), (_1 % 3) + 1); + + thrust::device_vector d_output(n); + auto r = thrust::transform_exclusive_scan( + d_input.begin(), d_input.end(), d_output.begin(), ::cuda::std::negate(), 5, ::cuda::std::multiplies<>{}); + ASSERT_EQUAL((d_output.end() == r), true); + + thrust::host_vector h_input = d_input; + thrust::host_vector h_output(n); + thrust::transform_exclusive_scan( + h_input.begin(), h_input.end(), h_output.begin(), ::cuda::std::negate(), 5, ::cuda::std::multiplies<>{}); + + ASSERT_EQUAL(d_output, h_output); + } + + // Test 3: User-defined type with poison default constructor + { + thrust::device_vector vec(5, TransformInt{1}); + thrust::transform_exclusive_scan( + vec.begin(), + vec.end(), + vec.begin(), + ::cuda::std::identity{}, + TransformInt{100}, + ::cuda::std::plus()); + + ASSERT_EQUAL(vec, + (thrust::device_vector{ + TransformInt{100}, TransformInt{101}, TransformInt{102}, TransformInt{103}, TransformInt{104}})); + } + + // Test 4: In-place transform_exclusive_scan with multiplies + { + thrust::device_vector vec = {2, 3, 4, 5}; + auto r = thrust::transform_exclusive_scan( + vec.begin(), vec.end(), vec.begin(), ::cuda::std::negate(), 10, ::cuda::std::multiplies<>{}); + ASSERT_EQUAL((vec.end() == r), true); + + thrust::device_vector expected = {10, -20, 60, -240}; + ASSERT_EQUAL(vec, expected); + } + + // Test 5: Boundary case at threshold (1024 elements) + { + const int n = 1024; + thrust::device_vector d_input(n); + using thrust::placeholders::_1; + thrust::tabulate(d_input.begin(), d_input.end(), _1 + 1); + + thrust::device_vector d_output(n); + auto r = thrust::transform_inclusive_scan( + d_input.begin(), d_input.end(), d_output.begin(), thrust::square(), 1, ::cuda::std::multiplies<>{}); + ASSERT_EQUAL((d_output.end() == r), true); + + thrust::host_vector h_input = d_input; + thrust::host_vector h_output(n); + thrust::transform_inclusive_scan( + h_input.begin(), h_input.end(), h_output.begin(), thrust::square(), 1, ::cuda::std::multiplies<>{}); + + ASSERT_EQUAL(d_output, h_output); + } + + // Test 6: Very small array (2 elements) + { + thrust::device_vector d_input = {3, 7}; + thrust::device_vector d_output(2); + + auto r = thrust::transform_inclusive_scan( + d_input.begin(), d_input.end(), d_output.begin(), ::cuda::std::negate(), 2, ::cuda::std::multiplies<>{}); + ASSERT_EQUAL((d_output.end() == r), true); + + thrust::device_vector expected = {-6, 42}; + ASSERT_EQUAL(d_output, expected); + } +} +DECLARE_UNITTEST(TestTransformScanEdgeCases); diff --git a/thrust/thrust/system/omp/detail/scan.h b/thrust/thrust/system/omp/detail/scan.h index d4e5c9db81f7..c3b9f9421416 100644 --- a/thrust/thrust/system/omp/detail/scan.h +++ b/thrust/thrust/system/omp/detail/scan.h @@ -75,7 +75,8 @@ OutputIterator scan_impl( const int num_threads = omp_get_max_threads(); // Use serial scan for small arrays where parallel overhead dominates - if (static_cast(n) < parallel_scan_threshold || num_threads <= 1) + if (static_cast(n) < ::cuda::std::max(parallel_scan_threshold, static_cast(num_threads)) + || num_threads <= 1) { if constexpr (IsInclusive) { @@ -94,6 +95,8 @@ OutputIterator scan_impl( } } + _CCCL_ASSERT(num_threads > 1, "Parallel scan requires multiple threads"); + temporary_array block_sums(exec, num_threads); // Step 1: Reduce each block (N reads) @@ -106,7 +109,9 @@ OutputIterator scan_impl( if (start < n) { - block_sums[tid] = ::cuda::std::reduce(first + start, first + end, accum_t{}, wrapped_binary_op); + // For both has_init and no-init cases: reduce each block using first element as init + accum_t first_elem = *(first + start); + block_sums[tid] = ::cuda::std::reduce(first + start + 1, first + end, first_elem, wrapped_binary_op); } } @@ -117,7 +122,10 @@ OutputIterator scan_impl( } else { - ::cuda::std::exclusive_scan(block_sums.begin(), block_sums.end(), block_sums.begin(), accum_t{}, wrapped_binary_op); + // For no init inclusive scan: exclusive_scan starting at second element with first element as init + accum_t first_block_sum = block_sums[0]; + ::cuda::std::exclusive_scan( + block_sums.begin() + 1, block_sums.end(), block_sums.begin() + 1, first_block_sum, wrapped_binary_op); } // Step 3: Scan each block with offset (N reads/writes) @@ -130,13 +138,30 @@ OutputIterator scan_impl( if (start < n) { - const accum_t prefix = block_sums[tid]; if constexpr (IsInclusive) { - ::cuda::std::inclusive_scan(first + start, first + end, result + start, wrapped_binary_op, prefix); + if constexpr (has_init) + { + const accum_t prefix = block_sums[tid]; + ::cuda::std::inclusive_scan(first + start, first + end, result + start, wrapped_binary_op, prefix); + } + else + { + // For no init: thread 0 has no prefix, others use block_sums + if (tid == 0) + { + ::cuda::std::inclusive_scan(first + start, first + end, result + start, wrapped_binary_op); + } + else + { + const accum_t prefix = block_sums[tid]; + ::cuda::std::inclusive_scan(first + start, first + end, result + start, wrapped_binary_op, prefix); + } + } } else { + const accum_t prefix = block_sums[tid]; ::cuda::std::exclusive_scan(first + start, first + end, result + start, prefix, wrapped_binary_op); } }