Repository navigation
[Thrust] Implement a single-pass OpenMP copy_if - #11746
SamMausberg wants to merge 4 commits into
Conversation
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 NVIDIA#1715
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
📝 Summary
Priority: ➖ Normal Change: Bug fix · Severity of issue fixed: Medium Merge Risk: 🔵 Low · up to
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository: NVIDIA/cccl/.coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: a1aa2a1e-da76-473f-99f5-249a46845715
📒 Files selected for processing (3)
thrust/testing/omp/copy_if.cuthrust/testing/unique_by_key.cuthrust/thrust/system/omp/detail/copy_if.h
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.
|
Hi @SamMausberg first and foremost thank you for the contribution. Unfortunately, however, thrust is more or less in maintenance-mode at this point and we are unlikely to accept large algorithmic additions or changes like this PR. I will let some other maintainers chime in to confirm, but unless someone objects we are unlikely to review or merge this PR. As a result the attached issue is also unlikely to be resolved. |
|
@Jacobfaib No worries at all, I genuinely don't mind. |
|
/ok to test 32942cb |
|
I think the strong improvements and broad impact at a moderate implementation size convince me to give this PR a chance to land. Most importantly, we must not introduce new bugs though. |
Thank you for giving this a chance. I'll try to spend some more time validating this ASAP (1-2 days?) so you guys don't have any redundant review, and this doesn't introduce bugs somewhere. Edit: taking a look at CI now. Found other issues as well; I'll fix these as I take note of them in this comment.
|
This comment has been minimized.
This comment has been minimized.
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.
|
Alright, I spent a couple of hours going over everything I could think of and pushed fixes for the above. |
There was a problem hiding this comment.
Actionable comments posted: 1
ℹ️ Review info
⚙️ Run configuration
- Configuration used: Repository: NVIDIA/cccl/.coderabbit.yaml
- Review profile: CHILL
- Plan: Enterprise
- Run ID:
f9f56b7b-a5a7-446d-8356-2d25bab16233
📒 Files selected for processing: 2
thrust/testing/omp/copy_if.cuthrust/thrust/system/omp/detail/copy_if.h
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.
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.
| // 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 |
There was a problem hiding this comment.
If these are just using inside the algorithm body, then let's move them there instead.
| } | ||
| else | ||
| { | ||
| bool selected[copy_if_detail::max_tile_size]; |
There was a problem hiding this comment.
Nit: suggest moving this block into a separate function. The current function is already quite large.
| if (team_of_one_result) | ||
| { | ||
| return *team_of_one_result; | ||
| } |
There was a problem hiding this comment.
We should just return the result immediately above instead of holding in an optional until the end of the function no? I don't think this is set when team_of_one is false.
There was a problem hiding this comment.
I'm not so sure here, but I don't think we can return from inside the parallel region. OpenMP, afaik, doesn't allow branching out of it. But agreed it was clunky. I'll fix these shortly, along with your other suggestions.
|
/ok to test 20c5f61 |
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.
|
/ok to test 81755da |
😬 CI Workflow Results🟥 Finished in 1h 38m: Pass: 97%/154 | Total: 1d 14h | Max: 1h 24m | Hits: 77%/312283See results here. AI failure analysis1. cuda.compute captured-bool transform strict xfail unexpectedly passes with numba-cuda-mlir 0.5.4 · 3 jobsExplanation: All three jobs installed numba-cuda-mlir 0.5.4, where the tested conversion now behaves correctly, so pytest converts the passing test into a failure because its xfail is strict. The same signature occurs on Linux, Windows, and the v2/HostJIT backend. Evidence: Copy this prompt into a coding agentJobs: 2. CUDASTF stream-context lifetime test reports CUDA error 719 at synchronization · 1 jobExplanation: The caller-stream synchronization reports an asynchronous failure from one of the previously queued `slow_set_kernel` launches. The log has no launch-site CUDA error check, compute-sanitizer output, or GPU Xid diagnostic, so it cannot distinguish a CUDASTF resource-lifetime defect from a stress-kernel or runner GPU failure; the PR diff does not modify CUDASTF. Evidence: Copy this prompt into a coding agentJobs: |
This branch has not been deployed
3 participants
Description
closes #1715
The OpenMP backend forwarded
copy_ifto the generic implementation, which allocates two temporary arrays ofnindices and runstransform,exclusive_scanandscatter_ifover them: 16 bytes per element with 64-bit indices. The two 4.4e9-elementunique_by_key_copytests disabled for #1715 peak at 65.6 GiB RSS on main, consistent with the OOM kill reported there.This PR gives the OpenMP backend a single-pass
copy_ifwith decoupled look-back. Threads take tiles of up to 8192 elements in order from a shared counter. For its tile, a thread evaluates the predicate 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 preceding tiles until it reaches a published prefix, publishes its own prefix, and copies the selected elements. Temporary storage is one 64-bit word per tile, from the policy's allocator. The words are aligned inside a byte array, because a custom allocator only has to align for its ownvalue_type. They usecuda::atomic_ref:cuda::std::atomic_refdoes not compile with host-only MSVC and NVHPC, because it reads anis_always_lock_freethat libcu++ only declares when_CCCL_ATOMIC_ALWAYS_LOCK_FREEis defined, which is under CUDA compilation and for GCC and Clang (#12120). #10968 as written would makecuda::atomic_refread the same member. Inputs below 1024 elements, and teams of a single thread (a nested region with nesting disabled), use the sequential loop.A simpler two-pass design (count per thread, scan, copy) needs only per-thread storage but evaluates the predicate twice. A two-pass version I measured first took 48 to 50 ms against main's 24 ms for a 200-round hash predicate on 2^23 elements; the look-back takes 24 ms.
A thread only waits for tiles handed out before its own, and their owners publish their counts before waiting on anything, so waits cannot form a cycle, whatever the team size. Publishers call
notify_all; libcu++ implements the hostatomic_ref::waitby polling with a back-off from spinning tosched_yieldand sleeping. The back-off also happens on an idle machine (about 38Ksched_yieldcalls percopy_ifcall on 2^28 elements) and lets more threads than cores make progress.In the cases I tested (see Validation), the new code accepts the same iterators and predicates as the generic implementation. It calls the predicate with const lvalues through
predicate_to_integral, and needs nothing of the output iterator beyondresult + n,*resultand++result. Data iterators whose reference is an rvalue reference, like amove_iterator, are rejected with astatic_assert. The generic implementation did not compile for them either, and moving the selected elements out would breakunique_copy,unique_by_key_copy,partition_copyandstable_partition_copy, which read their input again.#11847 has since fixed the OpenMP scans that main's
copy_ifuses (#11757), socopy_ifinside another parallel region returns correct results on main as well.thrust/testing/omp/copy_if.custays as a regression test for the new implementation.remove_copy_if,remove_copy,remove_if,remove,unique_copy,unique,unique_by_key_copy,unique_by_key,stable_partition,stable_partition_copyandpartition_copyare built oncopy_ifin the OpenMP backend and use the new code. This addresses thecopy_ifpart of #1709;count_ifis unchanged.Tests
TestUniqueCopyByKeyLargeInputandTestUniqueCopyByKeyLargeOutCountrun for OpenMP again. Under a 4 GiB address-space limit,TestUniqueCopyByKeyLargeInputfails withstd::bad_allocon main and passes here, at a sampled peak RSS (/proc/<pid>/smaps_rollup) of 13 MiB. It takes 13.2 s atOMP_NUM_THREADS=4, 3.3 s at 16 and 0.9 s at 64;TestUniqueCopyByKeyLargeOutCounttakes under 1 s. Both run in thecpp.ompandomp.ompconfigurations.thrust/testing/omp/copy_if.cucallscopy_iffrom inside a parallel region; it failed on main before [Thrust] Fix OpenMP scans with constrained thread teams #11847 and passes here.chartemporary allocator that misaligns the status words, an output iterator without comparison operators, a predicate takingauto&over azip_iterator, and ashortdifference type. On the first commit, 14 to 33 % ofcopy_ifcalls with the misaligning allocator returned wrong results, and the other three tests did not compile.Measurements
Idle 64-core Grace, g++ 11.4
-O2, 64 OpenMP threads (unpinned for large inputs, one per core for small ones), medians of 11 calls (2000 for small inputs) averaged over two runs. "Hash" runs 16 or 200 rounds of an LCG and xorshift per element. The main column was measured before #11847, which changed the OpenMP scan that main'scopy_ifuses.std::copy_if(par), libstdc++ + TBBint, 33.7 % selectedint, 1 % selectedint, 100 % selectedint, hash (16 rounds)int, hash (200 rounds)unique_by_key_copy, 2^27intpairsx % 3x % 3The last row is the cost of the serial threshold: a predicate costing about 180 ns per element is slower than main's parallel path below 1024 elements. Below about 18 ns per element, the serial path is not slower than main at any size under the threshold. clang-15 with libomp behaves the same (2^28 at 33.7 %: 48 ms on main, 7.4 ms here).
A single pass also schedules uneven predicate costs differently. A thread waiting behind a slow tile cannot take a new one, so expensive elements in different tiles are evaluated one after another, while main's static schedule spreads them over threads. With 16 elements of about 10 ms each at the ends of tiles 32 tiles apart (4M
ints, 16 threads on a 16-vCPU x86 VM), this PR takes 163 ms against main's 40 ms; the ratio grows towards the number of threads as the elements get more expensive. When the expensive elements fall into one thread's chunk of main's static schedule instead, this PR takes about a fifth of main's time.With
OMP_NUM_THREADS=128on the 64 cores (500 calls, 2^24 elements), g++/libgomp gives a 0.7 ms median and at most about 4 ms, against 4.5 ms and up to 163 ms on main. clang/libomp gives a 0.5 ms median; in three of eight runs single calls reached 16 to 21 ms (p99 up to 11 ms), below main's worst in the same runs (26 to 40 ms).<cuda/std/atomic>adds about 110 ms to parsingthrust/copy.hwith the OpenMP system (762 to 872 ms, measured on the first commit). libcu++'s internal__atomic/apiheaders would save only 15 ms of that and do not defineatomic_refitself.Validation
CUDA 12.8,
CMAKE_CUDA_ARCHITECTURES=90,THRUST_MULTICONFIG_WORKLOAD=MEDIUM:thrust-cpp17:thrust.{cpp.omp,omp.omp}.test.{copy,remove,unique,unique_by_key,partition,omp.copy_if,omp.reduce_intervals,omp.nvcc_independence}andthrust.{cpp.cpp,cpp.tbb,tbb.tbb}.test.{copy,remove,unique,unique_by_key,partition}pass; OMP header tests (host compiler and nvcc) build; no warnings.thrust-cpp20: the same OpenMP targets andthrust.cpp.tbb.test.{copy,unique_by_key}pass; OMP header tests build; no warnings.transform_output_iteratoroutputs; derived algorithms; a skewed-cost input; 5e9 elements; exactly one predicate call per element) passed 60 runs of 3 to 15 s: g++ and clang, default and 16 to 64-element tiles,OMP_NUM_THREADS1 to 128,OMP_DYNAMIC,OMP_THREAD_LIMIT=4,OMP_WAIT_POLICY=passive, nested regions with nesting off and on.pre-commitpasses.The second commit was validated on a 16-vCPU x86 VM (WSL2) with g++ 15.2 and clang 21:
move_iteratorinputs and stencils; misaligning allocators), throughcopy_ifand the derived algorithms, in C++17 and C++20: everything that compiles on main compiles here with the same results, and thestatic_assertfires only where main does not compile.THRUST_MULTICONFIG_WORKLOAD=MEDIUM, g++ and clang (with libomp), C++17 and C++20:thrust.{cpp.omp,omp.omp}.test.{copy,remove,unique,unique_by_key,partition,omp.copy_if,omp.reduce_intervals,omp.nvcc_independence}andthrust.cpp.cpp.test.{copy,remove,unique,unique_by_key,partition}pass with CI's warning flags and no warnings. Thecpp.ompandomp.ompheader test TUs build with the host compilers (not with nvcc), and clang-tidy andpre-commitpass._CCCL_ATOMIC_ALWAYS_LOCK_FREEundefined, which reproduces the MSVC and NVHPC failure of the first commit.omp.copy_iftests passed 1400 runs with 2 to 128 threads andOMP_WAIT_POLICY=passive. About 33K randomized stress iterations againststd::copy_if(including injected sleeps at the publish and look-back points, 1 to 128 threads and nested regions) and an ASan and UBSan build show no failures.Not tested: MSVC and NVHPC builds (beyond the emulation above), nvcc and TBB configurations for the second commit, the second commit on aarch64, CTK 13. On Windows, libcu++'s host wait falls back to
Sleep()once a wait exceeds about 40 us, which rounds up to the timer resolution (15.6 ms by default); I have not measured how oftencopy_ifhits this there.Checklist