[libc-commits] [libc] [libc] Simplify RPC lock handling to be scalar on the GPU (PR #221805)
via libc-commits
libc-commits at lists.llvm.org
Mon Sep 7 12:16:50 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-libc
Author: Joseph Huber (jhuber6)
<details>
<summary>Changes</summary>
Summary:
This interface was originally written in a complicated way as a
pessimistic bet against NVIDIA's independent thread scheduling. The idea
was that it would work regardless of which set actually executed the
function. The problem is that this forces the operation to be an
expensive vector operation.
This should be safe, we already have other safeguards against ITS in the
main loop. When this function is called, we know that the index and land
mask are convergent in the set. ITS says the true threads executing this
could be greater than returned. But in this case the 'locked' value
would not be shuffled in, it would not see that it got the lock, and
loop again in a different mask set.
The complicated thing about ITS is mostly that any lanes that share a PC
will report convergent, but that should not bite us here as we already
guard against that before calling this.
---
Full diff: https://github.com/llvm/llvm-project/pull/221805.diff
1 Files Affected:
- (modified) libc/shared/rpc.h (+18-53)
``````````diff
diff --git a/libc/shared/rpc.h b/libc/shared/rpc.h
index bd55bc34870af..c4a96bf33891c 100644
--- a/libc/shared/rpc.h
+++ b/libc/shared/rpc.h
@@ -237,47 +237,18 @@ template <bool Invert> struct Process {
return Invert ? !cond : cond;
}
- /// Attempt to claim the lock at index. Return true on lock taken.
- /// lane_mask is a bitmap of the threads in the warp that would hold the
- /// single lock on success, e.g. the result of rpc::get_lane_mask()
- /// The lock is held when the n-th bit of the lock bitfield is set.
+ /// Attempts to claim the lock at this index under the given execution mask.
+ /// Returns true on success.
RPC_ATTRS bool try_lock(uint64_t lane_mask, uint32_t index) {
- // On AMDGPU, test and set to the n-th lock bit and a sync_lane would
- // suffice On NVIDIA with ITS we need to handle differences between the
- // threads running and the threads that were detected in the previous call
- // to get_lane_mask()
- //
- // All threads in lane_mask try to claim the lock. At most one can succeed.
- // There may be threads active which are not in lane mask which must not
- // succeed in taking the lock, as otherwise it will leak. This is handled
- // by making threads which are not in lane_mask or with 0, a no-op.
- uint32_t id = rpc::get_lane_id();
- bool id_in_lane_mask = lane_mask & (1ul << id);
-
- // All threads in the warp call fetch_or. Possibly at the same time.
- bool before = set_nth(lock, index, id_in_lane_mask);
- uint64_t packed = rpc::ballot(lane_mask, before);
-
- // If every bit set in lane_mask is also set in packed, every single thread
- // in the warp failed to get the lock. Ballot returns unset for threads not
- // in the lane mask.
- //
- // Cases, per thread:
- // mask==0 -> unspecified before, discarded by ballot -> 0
- // mask==1 and before==0 (success), set zero by ballot -> 0
- // mask==1 and before==1 (failure), set one by ballot -> 1
- //
- // mask != packed implies at least one of the threads got the lock
- // atomic semantics of fetch_or mean at most one of the threads for the lock
-
- // If holding the lock then the caller can load values knowing said loads
- // won't move past the lock. No such guarantee is needed if the lock acquire
- // failed. This conditional branch is expected to fold in the caller after
- // inlining the current function.
- bool holding_lock = lane_mask != packed;
- if (holding_lock)
+ bool claimed = false;
+ if (rpc::is_first_lane(lane_mask))
+ claimed = !set_nth(lock, index);
+ claimed = rpc::broadcast_value(lane_mask, claimed);
+
+ // Do not move any reads past the point we obtain the lock.
+ if (claimed)
__scoped_atomic_thread_fence(__ATOMIC_ACQUIRE, __MEMORY_SCOPE_DEVICE);
- return holding_lock;
+ return claimed;
}
/// Unlock the lock at index. We need a lane sync to keep this function
@@ -286,11 +257,8 @@ template <bool Invert> struct Process {
// Do not move any writes past the unlock.
__scoped_atomic_thread_fence(__ATOMIC_RELEASE, __MEMORY_SCOPE_DEVICE);
- // Use exactly one thread to clear the nth bit in the lock array Must
- // restrict to a single thread to avoid one thread dropping the lock, then
- // an unrelated warp claiming the lock, then a second thread in this warp
- // dropping the lock again.
- clear_nth(lock, index, rpc::is_first_lane(lane_mask));
+ if (rpc::is_first_lane(lane_mask))
+ clear_nth(lock, index);
rpc::sync_lane(lane_mask);
}
@@ -331,23 +299,20 @@ template <bool Invert> struct Process {
}
/// Conditionally set the n-th bit in the atomic bitfield.
- RPC_ATTRS static constexpr uint32_t set_nth(uint32_t *bits, uint32_t index,
- bool cond) {
+ RPC_ATTRS static constexpr uint32_t set_nth(uint32_t *bits, uint32_t index) {
uint32_t slot = index / NUM_BITS_IN_WORD;
uint32_t bit = index % NUM_BITS_IN_WORD;
- return __scoped_atomic_fetch_or(&bits[slot],
- static_cast<uint32_t>(cond) << bit,
- __ATOMIC_RELAXED, __MEMORY_SCOPE_DEVICE) &
+ return __scoped_atomic_fetch_or(&bits[slot], 1u << bit, __ATOMIC_RELAXED,
+ __MEMORY_SCOPE_DEVICE) &
(1u << bit);
}
/// Conditionally clear the n-th bit in the atomic bitfield.
- RPC_ATTRS static constexpr uint32_t clear_nth(uint32_t *bits, uint32_t index,
- bool cond) {
+ RPC_ATTRS static constexpr uint32_t clear_nth(uint32_t *bits,
+ uint32_t index) {
uint32_t slot = index / NUM_BITS_IN_WORD;
uint32_t bit = index % NUM_BITS_IN_WORD;
- return __scoped_atomic_fetch_and(&bits[slot],
- ~0u ^ (static_cast<uint32_t>(cond) << bit),
+ return __scoped_atomic_fetch_and(&bits[slot], ~0u ^ (1u << bit),
__ATOMIC_RELAXED, __MEMORY_SCOPE_DEVICE) &
(1u << bit);
}
``````````
</details>
https://github.com/llvm/llvm-project/pull/221805
More information about the libc-commits
mailing list