From 32942cb7b29b571276363c3f94cac8a9a56572a3 Mon Sep 17 00:00:00 2001 From: Sam Date: Tue, 29 Sep 2026 17:30:02 +0000 Subject: [PATCH 1/4] [Thrust] Implement a single-pass OpenMP copy_if The OpenMP backend forwarded copy_if to the generic implementation, which allocates two temporary arrays of n indices (predicate flags and scatter indices) and runs a transform, an exclusive_scan and a scatter_if over them. With 64-bit indices that is 16 bytes of temporary storage per element, about 70 GB for the 4.4e9-element unique_by_key_copy tests, which is why those tests were disabled for OpenMP. The new implementation is a single pass with decoupled look-back. Threads take tiles of up to 8192 elements in increasing order from a shared counter, evaluate the predicate once per element into a tile-local flag buffer, publish the tile's count, add up the counts of the preceding tiles until they reach a published prefix, publish their own prefix and copy the selected elements to their final position. The temporary storage is one 64-bit status word per tile. A thread only waits for tiles that were handed out before its own, and their owners publish their counts before waiting on anything, so the waits cannot form a cycle for any team size, including more threads than cores. Inputs with fewer than 1024 elements, and teams of a single thread (e.g. inside another parallel region with nested parallelism disabled), use the sequential loop instead. The generic implementation also returned wrong results when OpenMP provided fewer threads than omp_get_max_threads() reports, inside another parallel region or with OMP_THREAD_LIMIT, because the OpenMP exclusive_scan it relies on processes one block per thread id. The shared tile counter makes the new code independent of the team size. On the OpenMP backend, remove_copy_if, remove_copy, remove_if, remove, unique_copy, unique, unique_by_key_copy, unique_by_key, stable_partition, stable_partition_copy and partition_copy are built on copy_if and use the new code as well. Re-enable TestUniqueCopyByKeyLargeInput and TestUniqueCopyByKeyLargeOutCount for OpenMP, and add a test that calls copy_if from inside a parallel region. Fixes #1715 --- thrust/testing/omp/copy_if.cu | 52 +++++++ thrust/testing/unique_by_key.cu | 9 +- thrust/thrust/system/omp/detail/copy_if.h | 181 +++++++++++++++++++++- 3 files changed, 232 insertions(+), 10 deletions(-) create mode 100644 thrust/testing/omp/copy_if.cu diff --git a/thrust/testing/omp/copy_if.cu b/thrust/testing/omp/copy_if.cu new file mode 100644 index 000000000000..dabba946fb2a --- /dev/null +++ b/thrust/testing/omp/copy_if.cu @@ -0,0 +1,52 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#include +#include +#include + +#include + +#include + +namespace +{ +struct is_even +{ + bool operator()(int x) const + { + return x % 2 == 0; + } +}; +} // namespace + +TEST_CASE("TestOmpCopyIfInsideParallelRegion", "[copy_if]") +{ + // Large enough that copy_if splits the work between threads when it is called outside of a parallel region + const int n = 1 << 20; + + const thrust::host_vector h_input = unittest::random_integers(n); + const thrust::device_vector d_input = h_input; + + thrust::host_vector h_result(n, thrust::no_init); + h_result.resize(thrust::copy_if(h_input.begin(), h_input.end(), h_result.begin(), is_even{}) - h_result.begin()); + + thrust::device_vector d_results[2] = { + thrust::device_vector(n, thrust::no_init), thrust::device_vector(n, thrust::no_init)}; + std::ptrdiff_t d_result_sizes[2] = {}; + + // Nested parallelism is disabled by default, so the parallel regions started by copy_if get a single thread each, + // although omp_get_max_threads() still reports the full number of threads +#pragma omp parallel for num_threads(2) + for (int i = 0; i < 2; ++i) + { + d_result_sizes[i] = + thrust::copy_if(d_input.begin(), d_input.end(), d_results[i].begin(), is_even{}) - d_results[i].begin(); + } + + for (int i = 0; i < 2; ++i) + { + d_results[i].resize(d_result_sizes[i]); + REQUIRE(d_results[i] == h_result); + } +} diff --git a/thrust/testing/unique_by_key.cu b/thrust/testing/unique_by_key.cu index b5397a11b571..62bcc5e67315 100644 --- a/thrust/testing/unique_by_key.cu +++ b/thrust/testing/unique_by_key.cu @@ -391,10 +391,7 @@ struct TestUniqueCopyByKeyToDiscardIterator }; DECLARE_GENERIC_SIZED_UNITTEST_WITH_TYPES(TestUniqueCopyByKeyToDiscardIterator, IntegralTypes); -// OpenMP has issues with these tests, NVIDIA/cccl#1715 -#if THRUST_DEVICE_SYSTEM != THRUST_DEVICE_SYSTEM_OMP - -# ifndef THRUST_FORCE_32_BIT_OFFSET_TYPE +#ifndef THRUST_FORCE_32_BIT_OFFSET_TYPE template struct TestUniqueCopyByKeyLargeInput @@ -453,9 +450,7 @@ struct TestUniqueCopyByKeyLargeOutCount }; DECLARE_GENERIC_UNITTEST_WITH_TYPES(TestUniqueCopyByKeyLargeOutCount, IntegralTypes); -# endif // THRUST_FORCE_32_BIT_OFFSET_TYPE - -#endif // non-OpenMP backend +#endif // THRUST_FORCE_32_BIT_OFFSET_TYPE // This test fails only on GCC 6 #if !defined(__GNUC__) || __GNUC__ != 6 diff --git a/thrust/thrust/system/omp/detail/copy_if.h b/thrust/thrust/system/omp/detail/copy_if.h index 82b93bfe4a4c..61b2a4350db9 100644 --- a/thrust/thrust/system/omp/detail/copy_if.h +++ b/thrust/thrust/system/omp/detail/copy_if.h @@ -13,12 +13,49 @@ # pragma system_header #endif // no system header -#include +#include +#include +#include // for depend_on_instantiation +#include +#include +#include #include +#include + +#include +#include +#include +#include +#include +#include +#include + +// don't attempt to #include this file without omp support +#if (THRUST_DEVICE_COMPILER_IS_OMP_CAPABLE == THRUST_TRUE) +# include +#endif // omp support THRUST_NAMESPACE_BEGIN namespace system::omp::detail { +// Below this many elements, copy_if runs on the calling thread +inline constexpr int parallel_copy_if_threshold = 1024; + +namespace copy_if_detail +{ +// Measured on a 64-core Grace with 1- to 32-byte elements: small tiles spread a costly predicate over all threads just +// above the threshold, while more and smaller tiles make the shared tile counter a bottleneck for large inputs, and +// larger tiles lose cache locality for large elements. +inline constexpr int min_tile_size = 64; +inline constexpr int max_tile_size = 8192; + +using status_word = ::cuda::std::uint64_t; + +// A tile's status word holds a number of selected elements, shifted left by two, and one of these flags +inline constexpr status_word status_partial = 1; // selected elements in this tile +inline constexpr status_word status_prefix = 2; // selected elements in this tile and all earlier ones +} // namespace copy_if_detail + template ::value, + "OpenMP compiler support is not enabled"); + + using Size = thrust::detail::it_difference_t; + using copy_if_detail::status_word; + + const Size n = ::cuda::std::distance(first, last); + +#if (THRUST_DEVICE_COMPILER_IS_OMP_CAPABLE == THRUST_TRUE) + const int num_threads = omp_get_max_threads(); +#else + const int num_threads = 1; +#endif // THRUST_DEVICE_COMPILER_IS_OMP_CAPABLE + + if (n < parallel_copy_if_threshold || num_threads <= 1) + { + return thrust::system::detail::sequential::copy_if(exec, first, last, stencil, result, pred); + } + + // Single pass with decoupled look-back. Threads take tiles in increasing order from a shared counter. For its tile, a + // thread evaluates pred once per element into a local flag buffer and publishes the tile's count in the tile's status + // word. It then adds up the counts of the preceding tiles, walking back until it finds a published prefix, publishes + // its own prefix, and copies the selected elements of the tile to their final position. + // + // Progress: a thread only waits for the status of an earlier tile. That tile was handed out earlier, so a thread is + // working on it, and that thread publishes the tile's count before it waits for anything. The waits cannot form a + // cycle and only require the owners of earlier tiles to get scheduled, whatever the team size. Publishers call + // atomic_ref::notify_all for the waiters. libcu++ implements the host wait by polling with a back-off from spinning + // to yielding and sleeping, which it does on an idle machine too, and which lets more threads than cores progress. + // + // Memory order: a status word packs the flag with the count, so a reader never sees one without the other, and the + // counts are the only data a thread reads from another one. The output ranges of the tiles are disjoint and only read + // after the barrier that ends the parallel region, so relaxed accesses are sufficient. + const Size tile_size = ::cuda::std::clamp( + ::cuda::ceil_div(n, static_cast(num_threads)), + static_cast(copy_if_detail::min_tile_size), + static_cast(copy_if_detail::max_tile_size)); + const Size num_tiles = ::cuda::ceil_div(n, tile_size); + + thrust::detail::temporary_array status_storage(exec, num_tiles); + status_word* const tile_status = thrust::raw_pointer_cast(status_storage.data()); + ::cuda::std::fill_n(tile_status, num_tiles, status_word{0}); + + ::cuda::std::atomic next_tile{0}; + + // Set by the thread that handles the last tile, or by a team of one thread + Size num_selected = 0; + + const thrust::detail::wrapped_function wrapped_pred{pred}; + + THRUST_PRAGMA_OMP(parallel) + { +#if (THRUST_DEVICE_COMPILER_IS_OMP_CAPABLE == THRUST_TRUE) + const bool team_of_one = omp_get_num_threads() == 1; +#else + const bool team_of_one = true; +#endif // THRUST_DEVICE_COMPILER_IS_OMP_CAPABLE + + if (team_of_one) + { + // E.g. inside another parallel region with nested parallelism disabled, where tiles only add overhead + num_selected = static_cast(::cuda::std::distance( + result, thrust::system::detail::sequential::copy_if(exec, first, last, stencil, result, pred))); + } + else + { + bool selected[copy_if_detail::max_tile_size]; + + for (;;) + { + const Size tile = next_tile.fetch_add(1, ::cuda::std::memory_order_relaxed); + if (tile >= num_tiles) + { + break; + } + + const Size tile_begin = tile * tile_size; + const Size tile_items = (::cuda::std::min) (tile_size, n - tile_begin); + + Size count = 0; + InputIterator2 stencil_it = stencil + tile_begin; + for (Size i = 0; i < tile_items; ++i, ++stencil_it) + { + selected[i] = wrapped_pred(*stencil_it); + count += static_cast(selected[i]); + } + + // Number of selected elements in all earlier tiles + Size offset = 0; + ::cuda::std::atomic_ref own_status(tile_status[tile]); + if (tile > 0) + { + own_status.store((static_cast(count) << 2) | copy_if_detail::status_partial, + ::cuda::std::memory_order_relaxed); + own_status.notify_all(); + for (Size earlier = tile - 1;; --earlier) + { + ::cuda::std::atomic_ref earlier_status(tile_status[earlier]); + status_word word; + while ((word = earlier_status.load(::cuda::std::memory_order_relaxed)) == 0) + { + earlier_status.wait(status_word{0}, ::cuda::std::memory_order_relaxed); + } + offset += static_cast(word >> 2); + if (word & copy_if_detail::status_prefix) + { + break; + } + } + } + own_status.store((static_cast(offset + count) << 2) | copy_if_detail::status_prefix, + ::cuda::std::memory_order_relaxed); + own_status.notify_all(); + + if (tile == num_tiles - 1) + { + num_selected = offset + count; + } + + OutputIterator out = result + offset; + InputIterator1 in = first + tile_begin; + for (Size i = 0; i < tile_items; ++i, ++in) + { + if (selected[i]) + { + *out = *in; + ++out; + } + } + } + } + } + + return result + num_selected; } // end copy_if() } // end namespace system::omp::detail THRUST_NAMESPACE_END From 552b2dfccc619a7c342ebf2d666aacc85e460a27 Mon Sep 17 00:00:00 2001 From: Sam Date: Tue, 20 Oct 2026 21:59:38 -0600 Subject: [PATCH 2/4] [Thrust] Fix the MSVC and NVHPC build and regressions in OpenMP copy_if cuda::std::atomic_ref does not compile with host-only MSVC and NVHPC: the public wrapper reads __atomic_ref_impl::is_always_lock_free, which libcu++ only declares when _CCCL_ATOMIC_ALWAYS_LOCK_FREE is defined (CUDA compilation, GCC and Clang). Use cuda::atomic_ref, which defines is_always_lock_free itself on top of the same implementation. A custom temporary allocator only has to align memory for its own value_type. With a char allocator, a status word could cross a cache line and be read torn, which gave wrong results. Allocate the status words as bytes with room for one more word and align the start. Restore what the generic implementation accepted: - call pred with const lvalues through predicate_to_integral, so that predicates taking auto& compile for zip_iterator inputs again, - for a team of one, keep the end returned by sequential::copy_if in a cuda::std::optional instead of its distance to result, so that the output iterator needs no comparison or difference, - cast n - tile_begin to Size for difference types narrower than int. Reject data inputs whose reference is an rvalue reference, like a move_iterator, with a static_assert. The generic implementation did not compile for them, and moving the selected elements out made unique_copy, unique_by_key_copy, partition_copy and stable_partition_copy return wrong results, because they read their input again. Add tests for the alignment and the three compile regressions to thrust/testing/omp/copy_if.cu. --- thrust/testing/omp/copy_if.cu | 141 ++++++++++++++++++++++ thrust/thrust/system/omp/detail/copy_if.h | 61 ++++++++-- 2 files changed, 192 insertions(+), 10 deletions(-) diff --git a/thrust/testing/omp/copy_if.cu b/thrust/testing/omp/copy_if.cu index dabba946fb2a..bc9b131e4ece 100644 --- a/thrust/testing/omp/copy_if.cu +++ b/thrust/testing/omp/copy_if.cu @@ -3,9 +3,17 @@ #include #include +#include #include +#include +#include +#include +#include +#include #include +#include +#include #include @@ -18,6 +26,56 @@ struct is_even return x % 2 == 0; } }; + +// Returns memory 7 bytes past a cache line boundary, which is legal for a char allocator. Every eighth 8-byte word in +// it then crosses a cache line. +struct misaligning_allocator +{ + using value_type = char; + + char* allocate(std::size_t n) + { + return static_cast(::operator new(n + 7, std::align_val_t{64})) + 7; + } + + void deallocate(char* p, std::size_t) + { + ::operator delete(p - 7, std::align_val_t{64}); + } +}; + +// An output iterator without comparison operators or a difference +struct output_iterator_without_comparisons +{ + using iterator_category = std::output_iterator_tag; + using value_type = void; + using difference_type = std::ptrdiff_t; + using pointer = void; + using reference = int&; + + int* ptr; + + int& operator*() const + { + return *ptr; + } + + output_iterator_without_comparisons& operator++() + { + ++ptr; + return *this; + } + + output_iterator_without_comparisons operator++(int) + { + return {ptr++}; + } + + output_iterator_without_comparisons operator+(std::ptrdiff_t i) const + { + return {ptr + i}; + } +}; } // namespace TEST_CASE("TestOmpCopyIfInsideParallelRegion", "[copy_if]") @@ -50,3 +108,86 @@ TEST_CASE("TestOmpCopyIfInsideParallelRegion", "[copy_if]") REQUIRE(d_results[i] == h_result); } } + +TEST_CASE("TestOmpCopyIfMisalignedTemporaryStorage", "[copy_if]") +{ + const int n = 1 << 14; + + const thrust::host_vector input = unittest::random_integers(n); + + thrust::host_vector expected(n, thrust::no_init); + expected.resize(thrust::copy_if(input.begin(), input.end(), expected.begin(), is_even{}) - expected.begin()); + + misaligning_allocator alloc; + thrust::host_vector result(n, thrust::no_init); + + // A torn read of a status word only shows up in some runs + for (int run = 0; run < 100; ++run) + { + const int* const end = + thrust::copy_if(thrust::omp::par(alloc), input.data(), input.data() + n, result.data(), is_even{}); + REQUIRE(end - result.data() == static_cast(expected.size())); + REQUIRE(thrust::equal(expected.begin(), expected.end(), result.begin())); + } +} + +TEST_CASE("TestOmpCopyIfOutputIteratorWithoutComparisons", "[copy_if]") +{ + for (const int n : {100, 1 << 16}) + { + thrust::host_vector input(n); + thrust::sequence(input.begin(), input.end()); + thrust::host_vector result(n); + + const output_iterator_without_comparisons end = thrust::copy_if( + thrust::omp::par, input.data(), input.data() + n, output_iterator_without_comparisons{result.data()}, is_even{}); + + REQUIRE(end.ptr - result.data() == n / 2); + for (int i = 0; i < n / 2; ++i) + { + REQUIRE(result[i] == 2 * i); + } + } +} + +TEST_CASE("TestOmpCopyIfPredicateTakingLvalueReferenceToZipReference", "[copy_if]") +{ + // Like the generic implementation, copy_if passes elements to pred as const lvalues, also for a zip_iterator, whose + // reference is a prvalue + for (const int n : {100, 1 << 16}) + { + thrust::device_vector a(n); + thrust::device_vector b(n, n / 2); + thrust::sequence(a.begin(), a.end()); + thrust::device_vector out_a(n); + thrust::device_vector out_b(n); + + const auto first = thrust::make_zip_iterator(a.begin(), b.begin()); + const auto end = + thrust::copy_if(first, first + n, thrust::make_zip_iterator(out_a.begin(), out_b.begin()), [](auto& t) { + return thrust::get<0>(t) < thrust::get<1>(t); + }); + + REQUIRE(end - thrust::make_zip_iterator(out_a.begin(), out_b.begin()) == n / 2); + out_a.resize(n / 2); + a.resize(n / 2); + REQUIRE(out_a == a); + } +} + +TEST_CASE("TestOmpCopyIfNarrowDifferenceType", "[copy_if]") +{ + // n - tile_begin has type int for a short difference_type + const int n = 20000; + + const thrust::counting_iterator first(0); + thrust::device_vector result(n); + + const auto end = thrust::copy_if(first, first + static_cast(n), result.begin(), is_even{}); + + REQUIRE(end - result.begin() == n / 2); + for (int i = 0; i < n / 2; ++i) + { + REQUIRE(result[i] == 2 * i); + } +} diff --git a/thrust/thrust/system/omp/detail/copy_if.h b/thrust/thrust/system/omp/detail/copy_if.h index 61b2a4350db9..c72f03b6de0b 100644 --- a/thrust/thrust/system/omp/detail/copy_if.h +++ b/thrust/thrust/system/omp/detail/copy_if.h @@ -14,21 +14,28 @@ #endif // no system header #include +#include #include #include // for depend_on_instantiation #include +#include #include #include #include #include #include +#include +#include #include #include #include #include +#include #include #include +#include +#include // don't attempt to #include this file without omp support #if (THRUST_DEVICE_COMPILER_IS_OMP_CAPABLE == THRUST_TRUE) @@ -54,6 +61,18 @@ using status_word = ::cuda::std::uint64_t; // A tile's status word holds a number of selected elements, shifted left by two, and one of these flags inline constexpr status_word status_partial = 1; // selected elements in this tile inline constexpr status_word status_prefix = 2; // selected elements in this tile and all earlier ones + +// Whether an iterator's reference is an rvalue reference or a tuple holding one, like a zip_iterator over a +// move_iterator +template +inline constexpr bool has_rvalue_reference = ::cuda::std::is_rvalue_reference_v; + +template +inline constexpr bool has_rvalue_reference<::cuda::std::tuple> = (has_rvalue_reference || ...); + +template +inline constexpr bool has_rvalue_reference> = + (has_rvalue_reference || ...); } // namespace copy_if_detail template ::value, "OpenMP compiler support is not enabled"); + // Moving the selected elements out would break the algorithms built on copy_if that read their input again, like + // unique_copy and stable_partition_copy + static_assert(!copy_if_detail::has_rvalue_reference>, + "copy_if and the algorithms built on it do not support an input iterator with an rvalue reference, " + "like a move_iterator, with the OpenMP system"); + using Size = thrust::detail::it_difference_t; using copy_if_detail::status_word; @@ -89,9 +114,12 @@ OutputIterator copy_if( const int num_threads = 1; #endif // THRUST_DEVICE_COMPILER_IS_OMP_CAPABLE + // Like generic::copy_if, call pred with a const lvalue, also for prvalue references like zip_iterator's + using const_lvalue_predicate = thrust::detail::predicate_to_integral; + if (n < parallel_copy_if_threshold || num_threads <= 1) { - return thrust::system::detail::sequential::copy_if(exec, first, last, stencil, result, pred); + return thrust::system::detail::sequential::copy_if(exec, first, last, stencil, result, const_lvalue_predicate{pred}); } // Single pass with decoupled look-back. Threads take tiles in increasing order from a shared counter. For its tile, a @@ -114,16 +142,24 @@ OutputIterator copy_if( static_cast(copy_if_detail::max_tile_size)); const Size num_tiles = ::cuda::ceil_div(n, tile_size); - thrust::detail::temporary_array status_storage(exec, num_tiles); - status_word* const tile_status = thrust::raw_pointer_cast(status_storage.data()); + // A custom temporary allocator may return memory that is only aligned for its own value_type, e.g. char, while a + // misaligned status word can be read torn. Allocate bytes with room for one more word and align the start. + thrust::detail::temporary_array<::cuda::std::uint8_t, DerivedPolicy> status_storage( + exec, (num_tiles + 1) * static_cast(sizeof(status_word))); + status_word* const tile_status = reinterpret_cast(::cuda::align_up( + thrust::raw_pointer_cast(status_storage.data()), ::cuda::atomic_ref::required_alignment)); ::cuda::std::fill_n(tile_status, num_tiles, status_word{0}); ::cuda::std::atomic next_tile{0}; - // Set by the thread that handles the last tile, or by a team of one thread + // Set by the thread that handles the last tile Size num_selected = 0; - const thrust::detail::wrapped_function wrapped_pred{pred}; + // Set by a team of one thread. Keeping the end returned by sequential::copy_if, instead of its distance to result, + // requires no comparison or difference of OutputIterators. + ::cuda::std::optional team_of_one_result; + + const thrust::detail::wrapped_function wrapped_pred{const_lvalue_predicate{pred}}; THRUST_PRAGMA_OMP(parallel) { @@ -136,8 +172,8 @@ OutputIterator copy_if( if (team_of_one) { // E.g. inside another parallel region with nested parallelism disabled, where tiles only add overhead - num_selected = static_cast(::cuda::std::distance( - result, thrust::system::detail::sequential::copy_if(exec, first, last, stencil, result, pred))); + team_of_one_result.emplace( + thrust::system::detail::sequential::copy_if(exec, first, last, stencil, result, const_lvalue_predicate{pred})); } else { @@ -152,7 +188,8 @@ OutputIterator copy_if( } const Size tile_begin = tile * tile_size; - const Size tile_items = (::cuda::std::min) (tile_size, n - tile_begin); + // n - tile_begin is an int for a Size like short + const Size tile_items = (::cuda::std::min) (tile_size, static_cast(n - tile_begin)); Size count = 0; InputIterator2 stencil_it = stencil + tile_begin; @@ -164,7 +201,7 @@ OutputIterator copy_if( // Number of selected elements in all earlier tiles Size offset = 0; - ::cuda::std::atomic_ref own_status(tile_status[tile]); + const ::cuda::atomic_ref own_status(tile_status[tile]); if (tile > 0) { own_status.store((static_cast(count) << 2) | copy_if_detail::status_partial, @@ -172,7 +209,7 @@ OutputIterator copy_if( own_status.notify_all(); for (Size earlier = tile - 1;; --earlier) { - ::cuda::std::atomic_ref earlier_status(tile_status[earlier]); + const ::cuda::atomic_ref earlier_status(tile_status[earlier]); status_word word; while ((word = earlier_status.load(::cuda::std::memory_order_relaxed)) == 0) { @@ -208,6 +245,10 @@ OutputIterator copy_if( } } + if (team_of_one_result) + { + return *team_of_one_result; + } return result + num_selected; } // end copy_if() } // end namespace system::omp::detail From 20c5f61050b15983d7facb25fc40834b4e751053 Mon Sep 17 00:00:00 2001 From: Sam Date: Wed, 21 Oct 2026 04:59:19 -0600 Subject: [PATCH 3/4] [Thrust] Trim the comments in OpenMP copy_if Make the comments shorter. Also compute the size of the status storage in size_t instead of in the iterator's difference type, and use thrust::no_init for test vectors that are overwritten before they are read. No behavior changes. --- thrust/testing/omp/copy_if.cu | 18 ++++----- thrust/thrust/system/omp/detail/copy_if.h | 46 ++++++++++------------- 2 files changed, 28 insertions(+), 36 deletions(-) diff --git a/thrust/testing/omp/copy_if.cu b/thrust/testing/omp/copy_if.cu index bc9b131e4ece..fe3f71f49eef 100644 --- a/thrust/testing/omp/copy_if.cu +++ b/thrust/testing/omp/copy_if.cu @@ -27,8 +27,8 @@ struct is_even } }; -// Returns memory 7 bytes past a cache line boundary, which is legal for a char allocator. Every eighth 8-byte word in -// it then crosses a cache line. +// Returns memory 7 bytes past a cache line boundary, which is legal for a char allocator, so every eighth 8-byte word +// in it crosses a cache line struct misaligning_allocator { using value_type = char; @@ -135,9 +135,9 @@ TEST_CASE("TestOmpCopyIfOutputIteratorWithoutComparisons", "[copy_if]") { for (const int n : {100, 1 << 16}) { - thrust::host_vector input(n); + thrust::host_vector input(n, thrust::no_init); thrust::sequence(input.begin(), input.end()); - thrust::host_vector result(n); + thrust::host_vector result(n, thrust::no_init); const output_iterator_without_comparisons end = thrust::copy_if( thrust::omp::par, input.data(), input.data() + n, output_iterator_without_comparisons{result.data()}, is_even{}); @@ -152,15 +152,15 @@ TEST_CASE("TestOmpCopyIfOutputIteratorWithoutComparisons", "[copy_if]") TEST_CASE("TestOmpCopyIfPredicateTakingLvalueReferenceToZipReference", "[copy_if]") { - // Like the generic implementation, copy_if passes elements to pred as const lvalues, also for a zip_iterator, whose + // copy_if passes elements to pred as const lvalues, like the generic implementation, although zip_iterator's // reference is a prvalue for (const int n : {100, 1 << 16}) { - thrust::device_vector a(n); + thrust::device_vector a(n, thrust::no_init); thrust::device_vector b(n, n / 2); thrust::sequence(a.begin(), a.end()); - thrust::device_vector out_a(n); - thrust::device_vector out_b(n); + thrust::device_vector out_a(n, thrust::no_init); + thrust::device_vector out_b(n, thrust::no_init); const auto first = thrust::make_zip_iterator(a.begin(), b.begin()); const auto end = @@ -181,7 +181,7 @@ TEST_CASE("TestOmpCopyIfNarrowDifferenceType", "[copy_if]") const int n = 20000; const thrust::counting_iterator first(0); - thrust::device_vector result(n); + thrust::device_vector result(n, thrust::no_init); const auto end = thrust::copy_if(first, first + static_cast(n), result.begin(), is_even{}); diff --git a/thrust/thrust/system/omp/detail/copy_if.h b/thrust/thrust/system/omp/detail/copy_if.h index c72f03b6de0b..dbb2950f8795 100644 --- a/thrust/thrust/system/omp/detail/copy_if.h +++ b/thrust/thrust/system/omp/detail/copy_if.h @@ -33,6 +33,7 @@ #include #include #include +#include #include #include #include @@ -50,20 +51,18 @@ inline constexpr int parallel_copy_if_threshold = 1024; namespace copy_if_detail { -// Measured on a 64-core Grace with 1- to 32-byte elements: small tiles spread a costly predicate over all threads just -// above the threshold, while more and smaller tiles make the shared tile counter a bottleneck for large inputs, and -// larger tiles lose cache locality for large elements. +// Tuned on a 64-core Grace with 1- to 32-byte elements: the minimum lets a costly predicate use all threads just above +// the threshold, and the maximum balances contention on the tile counter against cache locality. inline constexpr int min_tile_size = 64; inline constexpr int max_tile_size = 8192; using status_word = ::cuda::std::uint64_t; -// A tile's status word holds a number of selected elements, shifted left by two, and one of these flags +// A status word holds a count of selected elements, shifted left by two, and one of these flags inline constexpr status_word status_partial = 1; // selected elements in this tile inline constexpr status_word status_prefix = 2; // selected elements in this tile and all earlier ones -// Whether an iterator's reference is an rvalue reference or a tuple holding one, like a zip_iterator over a -// move_iterator +// Whether a reference type is or holds an rvalue reference, like that of a move_iterator or a zip_iterator over one template inline constexpr bool has_rvalue_reference = ::cuda::std::is_rvalue_reference_v; @@ -97,8 +96,7 @@ OutputIterator copy_if( (THRUST_DEVICE_COMPILER_IS_OMP_CAPABLE == THRUST_TRUE)>::value, "OpenMP compiler support is not enabled"); - // Moving the selected elements out would break the algorithms built on copy_if that read their input again, like - // unique_copy and stable_partition_copy + // Moving elements out would break algorithms built on copy_if that read the input again, like unique_copy static_assert(!copy_if_detail::has_rvalue_reference>, "copy_if and the algorithms built on it do not support an input iterator with an rvalue reference, " "like a move_iterator, with the OpenMP system"); @@ -114,7 +112,7 @@ OutputIterator copy_if( const int num_threads = 1; #endif // THRUST_DEVICE_COMPILER_IS_OMP_CAPABLE - // Like generic::copy_if, call pred with a const lvalue, also for prvalue references like zip_iterator's + // Pass elements to pred as const lvalues, as generic::copy_if does using const_lvalue_predicate = thrust::detail::predicate_to_integral; if (n < parallel_copy_if_threshold || num_threads <= 1) @@ -122,30 +120,25 @@ OutputIterator copy_if( return thrust::system::detail::sequential::copy_if(exec, first, last, stencil, result, const_lvalue_predicate{pred}); } - // Single pass with decoupled look-back. Threads take tiles in increasing order from a shared counter. For its tile, a - // thread evaluates pred once per element into a local flag buffer and publishes the tile's count in the tile's status - // word. It then adds up the counts of the preceding tiles, walking back until it finds a published prefix, publishes - // its own prefix, and copies the selected elements of the tile to their final position. + // Single pass with decoupled look-back. Threads take tiles in order from a shared counter. For its tile, a thread + // evaluates pred once per element into a flag buffer and publishes the tile's count. It then adds up the counts of + // earlier tiles back to the first published prefix, publishes its own prefix and copies the selected elements. // - // Progress: a thread only waits for the status of an earlier tile. That tile was handed out earlier, so a thread is - // working on it, and that thread publishes the tile's count before it waits for anything. The waits cannot form a - // cycle and only require the owners of earlier tiles to get scheduled, whatever the team size. Publishers call - // atomic_ref::notify_all for the waiters. libcu++ implements the host wait by polling with a back-off from spinning - // to yielding and sleeping, which it does on an idle machine too, and which lets more threads than cores progress. + // Progress: a thread only waits for earlier tiles, whose owners publish their counts before waiting on anything, so + // the waits cannot form a cycle, whatever the team size. libcu++'s host wait polls with a back-off from spinning to + // yielding and sleeping, so more threads than cores also make progress. // - // Memory order: a status word packs the flag with the count, so a reader never sees one without the other, and the - // counts are the only data a thread reads from another one. The output ranges of the tiles are disjoint and only read - // after the barrier that ends the parallel region, so relaxed accesses are sufficient. + // Memory order: a status word packs the flag with the count, the counts are the only data threads share, and the + // disjoint output ranges are only read after the region's closing barrier, so relaxed accesses are sufficient. const Size tile_size = ::cuda::std::clamp( ::cuda::ceil_div(n, static_cast(num_threads)), static_cast(copy_if_detail::min_tile_size), static_cast(copy_if_detail::max_tile_size)); const Size num_tiles = ::cuda::ceil_div(n, tile_size); - // A custom temporary allocator may return memory that is only aligned for its own value_type, e.g. char, while a - // misaligned status word can be read torn. Allocate bytes with room for one more word and align the start. + // A custom allocator may only align for its own value_type, e.g. char, and a misaligned status word can be read torn thrust::detail::temporary_array<::cuda::std::uint8_t, DerivedPolicy> status_storage( - exec, (num_tiles + 1) * static_cast(sizeof(status_word))); + exec, (static_cast<::cuda::std::size_t>(num_tiles) + 1) * sizeof(status_word)); status_word* const tile_status = reinterpret_cast(::cuda::align_up( thrust::raw_pointer_cast(status_storage.data()), ::cuda::atomic_ref::required_alignment)); ::cuda::std::fill_n(tile_status, num_tiles, status_word{0}); @@ -155,8 +148,7 @@ OutputIterator copy_if( // Set by the thread that handles the last tile Size num_selected = 0; - // Set by a team of one thread. Keeping the end returned by sequential::copy_if, instead of its distance to result, - // requires no comparison or difference of OutputIterators. + // Set by a team of one thread. Keeping the returned end avoids comparing or subtracting OutputIterators. ::cuda::std::optional team_of_one_result; const thrust::detail::wrapped_function wrapped_pred{const_lvalue_predicate{pred}}; @@ -188,7 +180,7 @@ OutputIterator copy_if( } const Size tile_begin = tile * tile_size; - // n - tile_begin is an int for a Size like short + // n - tile_begin is an int for a narrow Size like short const Size tile_items = (::cuda::std::min) (tile_size, static_cast(n - tile_begin)); Size count = 0; From 81755daa0063e031217f618f25a5634289a141c2 Mon Sep 17 00:00:00 2001 From: Sam Date: Wed, 21 Oct 2026 10:08:35 -0600 Subject: [PATCH 4/4] [Thrust] Split the tiled path out of OpenMP copy_if Move the tiled path into copy_if_detail::tiled_copy_if, with its constants as locals there. For a team of one thread, only set a flag in the parallel region and call sequential::copy_if after it, instead of keeping its result in a cuda::std::optional. An exception from pred in that case now propagates instead of calling std::terminate. --- thrust/thrust/system/omp/detail/copy_if.h | 167 ++++++++++++---------- 1 file changed, 90 insertions(+), 77 deletions(-) diff --git a/thrust/thrust/system/omp/detail/copy_if.h b/thrust/thrust/system/omp/detail/copy_if.h index dbb2950f8795..219c6d38f838 100644 --- a/thrust/thrust/system/omp/detail/copy_if.h +++ b/thrust/thrust/system/omp/detail/copy_if.h @@ -35,7 +35,6 @@ #include #include #include -#include #include // don't attempt to #include this file without omp support @@ -46,22 +45,8 @@ THRUST_NAMESPACE_BEGIN namespace system::omp::detail { -// Below this many elements, copy_if runs on the calling thread -inline constexpr int parallel_copy_if_threshold = 1024; - namespace copy_if_detail { -// Tuned on a 64-core Grace with 1- to 32-byte elements: the minimum lets a costly predicate use all threads just above -// the threshold, and the maximum balances contention on the tile counter against cache locality. -inline constexpr int min_tile_size = 64; -inline constexpr int max_tile_size = 8192; - -using status_word = ::cuda::std::uint64_t; - -// A status word holds a count of selected elements, shifted left by two, and one of these flags -inline constexpr status_word status_partial = 1; // selected elements in this tile -inline constexpr status_word status_prefix = 2; // selected elements in this tile and all earlier ones - // Whether a reference type is or holds an rvalue reference, like that of a move_iterator or a zip_iterator over one template inline constexpr bool has_rvalue_reference = ::cuda::std::is_rvalue_reference_v; @@ -72,68 +57,48 @@ inline constexpr bool has_rvalue_reference<::cuda::std::tuple> = (has_rva template inline constexpr bool has_rvalue_reference> = (has_rvalue_reference || ...); -} // namespace copy_if_detail +// Single pass with decoupled look-back. Threads take tiles in order from a shared counter. For its tile, a thread +// evaluates pred once per element into a flag buffer and publishes the tile's count. It then adds up the counts of +// earlier tiles back to the first published prefix, publishes its own prefix and copies the selected elements. +// +// Progress: a thread only waits for earlier tiles, whose owners publish their counts before waiting on anything, so +// the waits cannot form a cycle, whatever the team size. libcu++'s host wait polls with a back-off from spinning to +// yielding and sleeping, so more threads than cores also make progress. +// +// Memory order: a status word packs the flag with the count, the counts are the only data threads share, and the +// disjoint output ranges are only read after the region's closing barrier, so relaxed accesses are sufficient. template -OutputIterator copy_if( + typename Predicate, + typename Size> +OutputIterator tiled_copy_if( execution_policy& exec, InputIterator1 first, InputIterator1 last, + Size n, InputIterator2 stencil, OutputIterator result, - Predicate pred) + Predicate pred, + int num_threads) { - // we're attempting to launch an omp kernel, assert we're compiling with omp support - // ======================================================================== - // X Note to the user: If you've found this line due to a compiler error, X - // X you need to enable OpenMP support in your compiler. X - // ======================================================================== - static_assert(thrust::detail::depend_on_instantiation::value, - "OpenMP compiler support is not enabled"); + // Tuned on a 64-core Grace with 1- to 32-byte elements: the minimum lets a costly predicate use all threads just + // above the threshold, and the maximum balances contention on the tile counter against cache locality. + constexpr int min_tile_size = 64; + constexpr int max_tile_size = 8192; - // Moving elements out would break algorithms built on copy_if that read the input again, like unique_copy - static_assert(!copy_if_detail::has_rvalue_reference>, - "copy_if and the algorithms built on it do not support an input iterator with an rvalue reference, " - "like a move_iterator, with the OpenMP system"); + using status_word = ::cuda::std::uint64_t; - using Size = thrust::detail::it_difference_t; - using copy_if_detail::status_word; + // A status word holds a count of selected elements, shifted left by two, and one of these flags + constexpr status_word status_partial = 1; // selected elements in this tile + constexpr status_word status_prefix = 2; // selected elements in this tile and all earlier ones - const Size n = ::cuda::std::distance(first, last); - -#if (THRUST_DEVICE_COMPILER_IS_OMP_CAPABLE == THRUST_TRUE) - const int num_threads = omp_get_max_threads(); -#else - const int num_threads = 1; -#endif // THRUST_DEVICE_COMPILER_IS_OMP_CAPABLE - - // Pass elements to pred as const lvalues, as generic::copy_if does - using const_lvalue_predicate = thrust::detail::predicate_to_integral; - - if (n < parallel_copy_if_threshold || num_threads <= 1) - { - return thrust::system::detail::sequential::copy_if(exec, first, last, stencil, result, const_lvalue_predicate{pred}); - } - - // Single pass with decoupled look-back. Threads take tiles in order from a shared counter. For its tile, a thread - // evaluates pred once per element into a flag buffer and publishes the tile's count. It then adds up the counts of - // earlier tiles back to the first published prefix, publishes its own prefix and copies the selected elements. - // - // Progress: a thread only waits for earlier tiles, whose owners publish their counts before waiting on anything, so - // the waits cannot form a cycle, whatever the team size. libcu++'s host wait polls with a back-off from spinning to - // yielding and sleeping, so more threads than cores also make progress. - // - // Memory order: a status word packs the flag with the count, the counts are the only data threads share, and the - // disjoint output ranges are only read after the region's closing barrier, so relaxed accesses are sufficient. const Size tile_size = ::cuda::std::clamp( ::cuda::ceil_div(n, static_cast(num_threads)), - static_cast(copy_if_detail::min_tile_size), - static_cast(copy_if_detail::max_tile_size)); + static_cast(min_tile_size), + static_cast(max_tile_size)); const Size num_tiles = ::cuda::ceil_div(n, tile_size); // A custom allocator may only align for its own value_type, e.g. char, and a misaligned status word can be read torn @@ -148,28 +113,26 @@ OutputIterator copy_if( // Set by the thread that handles the last tile Size num_selected = 0; - // Set by a team of one thread. Keeping the returned end avoids comparing or subtracting OutputIterators. - ::cuda::std::optional team_of_one_result; + // Set if the team has a single thread + bool team_of_one = false; - const thrust::detail::wrapped_function wrapped_pred{const_lvalue_predicate{pred}}; + const thrust::detail::wrapped_function wrapped_pred{pred}; THRUST_PRAGMA_OMP(parallel) { #if (THRUST_DEVICE_COMPILER_IS_OMP_CAPABLE == THRUST_TRUE) - const bool team_of_one = omp_get_num_threads() == 1; + const bool single_thread = omp_get_num_threads() == 1; #else - const bool team_of_one = true; + const bool single_thread = true; #endif // THRUST_DEVICE_COMPILER_IS_OMP_CAPABLE - if (team_of_one) + if (single_thread) { - // E.g. inside another parallel region with nested parallelism disabled, where tiles only add overhead - team_of_one_result.emplace( - thrust::system::detail::sequential::copy_if(exec, first, last, stencil, result, const_lvalue_predicate{pred})); + team_of_one = true; } else { - bool selected[copy_if_detail::max_tile_size]; + bool selected[max_tile_size]; for (;;) { @@ -196,8 +159,7 @@ OutputIterator copy_if( const ::cuda::atomic_ref own_status(tile_status[tile]); if (tile > 0) { - own_status.store((static_cast(count) << 2) | copy_if_detail::status_partial, - ::cuda::std::memory_order_relaxed); + own_status.store((static_cast(count) << 2) | status_partial, ::cuda::std::memory_order_relaxed); own_status.notify_all(); for (Size earlier = tile - 1;; --earlier) { @@ -208,13 +170,13 @@ OutputIterator copy_if( earlier_status.wait(status_word{0}, ::cuda::std::memory_order_relaxed); } offset += static_cast(word >> 2); - if (word & copy_if_detail::status_prefix) + if (word & status_prefix) { break; } } } - own_status.store((static_cast(offset + count) << 2) | copy_if_detail::status_prefix, + own_status.store((static_cast(offset + count) << 2) | status_prefix, ::cuda::std::memory_order_relaxed); own_status.notify_all(); @@ -237,11 +199,62 @@ OutputIterator copy_if( } } - if (team_of_one_result) + // E.g. inside another parallel region with nested parallelism disabled, where tiles only add overhead + if (team_of_one) { - return *team_of_one_result; + return thrust::system::detail::sequential::copy_if(exec, first, last, stencil, result, pred); } return result + num_selected; +} +} // namespace copy_if_detail + +template +OutputIterator copy_if( + execution_policy& exec, + InputIterator1 first, + InputIterator1 last, + InputIterator2 stencil, + OutputIterator result, + Predicate pred) +{ + // we're attempting to launch an omp kernel, assert we're compiling with omp support + // ======================================================================== + // X Note to the user: If you've found this line due to a compiler error, X + // X you need to enable OpenMP support in your compiler. X + // ======================================================================== + static_assert(thrust::detail::depend_on_instantiation::value, + "OpenMP compiler support is not enabled"); + + // Moving elements out would break algorithms built on copy_if that read the input again, like unique_copy + static_assert(!copy_if_detail::has_rvalue_reference>, + "copy_if and the algorithms built on it do not support an input iterator with an rvalue reference, " + "like a move_iterator, with the OpenMP system"); + + // Below this many elements, copy_if runs on the calling thread + constexpr int parallel_copy_if_threshold = 1024; + + using Size = thrust::detail::it_difference_t; + const Size n = ::cuda::std::distance(first, last); + +#if (THRUST_DEVICE_COMPILER_IS_OMP_CAPABLE == THRUST_TRUE) + const int num_threads = omp_get_max_threads(); +#else + const int num_threads = 1; +#endif // THRUST_DEVICE_COMPILER_IS_OMP_CAPABLE + + // Pass elements to pred as const lvalues, as generic::copy_if does + using const_lvalue_predicate = thrust::detail::predicate_to_integral; + + if (n < parallel_copy_if_threshold || num_threads <= 1) + { + return thrust::system::detail::sequential::copy_if(exec, first, last, stencil, result, const_lvalue_predicate{pred}); + } + return copy_if_detail::tiled_copy_if(exec, first, last, n, stencil, result, const_lvalue_predicate{pred}, num_threads); } // end copy_if() } // end namespace system::omp::detail THRUST_NAMESPACE_END