Fixes a race condition in DeviceTopK

[elstehle]

Description

Fixes a potential race condition, where thread 0 signals a CTA's completion of merge_histograms while some other threads have not concluded contributing their counts to the global yet.

[coderabbitai]

📝 Walkthrough

Summary by CodeRabbit

• Bug Fixes
  • Improved synchronization during Top-K processing to ensure consistent results and more reliable execution across GPU threads.

Walkthrough

finalize_pass now issues __threadfence() before __syncthreads(). This ensures that histogram merges complete before block completion is published.

Changes

TopK synchronization

Layer / File(s)	Summary
Synchronize histogram publication cub/cub/agent/agent_topk.cuh	finalize_pass fences memory before synchronizing block threads. The synchronization ensures that all histogram merges complete before completion is published.

Suggested reviewers: bernhardmgruber, jacobfaib

---

_{Comment @coderabbitai help to get the list of available commands.}

[pauleonix]

I hope this is not another false alarm from my side, but is the order of the two correct? Don't we need all results flushed to L2 before synchronizing the CTA so that the atomic certainly arrives after stores from all threads in the CTA?

[coderabbitai]

🧹 Nitpick comments (1)

cub/cub/agent/agent_topk.cuh (1)

650-655: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

suggestion: Add a deterministic multi-CTA regression test for this ordering. The supplied cub/test/catch2_test_device_topk_keys.cu:73-112 test validates final keys, but its random size selection does not guarantee multiple CTAs or repeated execution of this publication path. Add a fixed case that launches multiple CTAs, reaches finalize_pass more than once, and repeats the operation.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 87d7e43c-158a-42c5-aede-a3d2a356bc8d

📥 Commits

Reviewing files that changed from the base of the PR and between 5080f54 and 9434918.

📒 Files selected for processing (1)

• cub/cub/agent/agent_topk.cuh

[elstehle]

I hope this is not another false alarm from my side, but is the order of the two correct? Don't we need all results flushed to L2 before synchronizing the CTA so that the atomic certainly arrives after stores from all threads in the CTA?

I see. I changed the order, to be on the safe side.

Under the plain documented semantics, __syncthreads(); __threadfence(); if (tid==0) atomicInc(...) leaves a gap: nothing in the docs promises the other threads' contributions reached the device coherence point before the retirement atomic.

However, based on the AI assessment, the current order apparently would still hold up, but for less obvious reasons:

Formally it is also sound, but only via a non-obvious property. In the PTX memory model, bar.sync "has the same effect as release-acquire synchronization performed at .cta scope", and base causality order is transitive and cumulative — the ASPLOS'19 PTX memory model paper states the recursion "ensures that synchronization composes transitively... It also enforces the cumulativity property."

Applied here: thread 100's atomicAdd precedes its barrier arrival in program order; the barrier synchronizes thread 100 with thread 0 (CTA scope covers both, so they're morally strong); thread 0's fence and atomicInc follow its barrier in program order. Cumulativity therefore drags thread 100's write into causality order before thread 0's atomicInc, and the device-scope fence makes that a release at .gpu scope. So the fence does carry the whole CTA's prior writes — but you need the cumulativity clause to say so, which is not something a reader should have to reconstruct.

There's also a practical cushion specific to this code: merge_histograms publishes via device-scope atomicAdd, which is performed at L2 by definition, so "visible to the block" after __syncthreads() already implies "landed at L2". That cushion would evaporate if those ever became plain stores.

[github-actions]

⏱️ CCCL compile-time benchmark comparison: Public headers compile-time bench

Result: 0 regression row(s), 7 improvement row(s) above threshold.

Run	Value
Config	public-headers-gcc13
Baseline	origin/main
Preset	all-dev
Targets	cub.headers.base, thrust.cpp.cuda.headers.base, libcudacxx.test.public_headers
GPU / launch args	rtx2080 / --cuda 13.3 --host gcc13

Artifacts: reports and traces

Direct file processing

-f file-processing exclusive --sort total

🟢 Direct file processing — Improvements

Rank	Improvement impact	Selected Δ	Baseline	Current	Event	Matched traces
1	1.595311	-1.595311	20.507159	18.911848	Processing Header File: libcudacxx/include/cuda/__device/physical_device.h	94
2	1.340935	-1.340935	9.648752	8.307817	Processing Header File: libcudacxx/include/cuda/std/__cccl/prologue.h	548
3	0.491711	-0.491711	3.432253	2.940542	Processing Header File: libcudacxx/include/cuda/std/__cccl/epilogue.h	548
4	0.426785	-0.426785	11.039697	10.612912	Processing Header File: libcudacxx/include/cuda/std/__iterator/concepts.h	436
5	0.273014	-0.273014	2.152568	1.879554	Processing Header File: libcudacxx/include/cuda/std/string_view	98
6	0.242502	-0.242502	5.671503	5.429001	Processing Header File: libcudacxx/include/cuda/std/__iterator/iterator_traits.h	463
7	0.216030	-0.216030	11.297008	11.080978	Processing Header File: libcudacxx/include/cuda/std/__type_traits/type_list.h	475

[github-actions]

🥳 CI Workflow Results

🟩 Finished in 1h 25m: Pass: 100%/270 | Total: 2d 01h | Max: 56m 05s | Hits: 87%/222326

See results here.

[miscco]

Looks reasonable

[elstehle]

/backport

[github-actions]

Successfully created backport PR for branch/3.4.x:

• [Backport branch/3.4.x] Fixes a race condition in DeviceTopK #10682

[github-actions]

Successfully created backport PR for branch/3.5.x:

• [Backport branch/3.5.x] Fixes a race condition in DeviceTopK #10683