Method try_lock

system/lib/llvm-libc/shared/rpc.h:171–207

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.

  /// 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.
  RPC_ATTRS bool try_lock(uint64_t lane_mask, uint32_t index) {
    // On amdgpu, test and set to the nth lock bit and a sync_lane would suffice
    // On volta, 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)
      __scoped_atomic_thread_fence(__ATOMIC_ACQUIRE, __MEMORY_SCOPE_DEVICE);
    return holding_lock;
  }

  /// Unlock the lock at index. We need a lane sync to keep this function
  /// convergent, otherwise the compiler will sink the store and deadlock.