libc/src/__support/threads/spin_lock.h

   1 //===-- TTAS Spin Lock ----------------------------------------------------===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8
   9 #ifndef LLVM_LIBC_SRC___SUPPORT_THREADS_SPIN_LOCK_H
  10 #define LLVM_LIBC_SRC___SUPPORT_THREADS_SPIN_LOCK_H
  11
  12 #include "src/__support/CPP/atomic.h"
  13 #include "src/__support/macros/attributes.h"
  14 #include "src/__support/threads/sleep.h"
  15
  16 namespace LIBC_NAMESPACE_DECL {
  17
  18 class SpinLock {
  19   cpp::Atomic<unsigned char> flag;
  20
  21 public:
  22   LIBC_INLINE constexpr SpinLock() : flag{0} {}
  23   LIBC_INLINE bool try_lock() {
  24     return !flag.exchange(1u, cpp::MemoryOrder::ACQUIRE);
  25   }
  26   LIBC_INLINE void lock() {
  27     // clang-format off
  28     // For normal TTAS, this compiles to the following on armv9a and x86_64:
  29     //         mov     w8, #1            |          .LBB0_1:
  30     // .LBB0_1:                          |                  mov     al, 1
  31     //         swpab   w8, w9, [x0]      |                  xchg    byte ptr [rdi], al
  32     //         tbnz    w9, #0, .LBB0_3   |                  test    al, 1
  33     //         b       .LBB0_4           |                  jne     .LBB0_3
  34     // .LBB0_2:                          |                  jmp     .LBB0_4
  35     //         isb                       |         .LBB0_2:
  36     // .LBB0_3:                          |                  pause
  37     //         ldrb    w9, [x0]          |         .LBB0_3:
  38     //         tbnz    w9, #0, .LBB0_2   |                  movzx   eax, byte ptr [rdi]
  39     //         b       .LBB0_1           |                  test    al, 1
  40     // .LBB0_4:                          |                  jne     .LBB0_2
  41     //         ret                       |                  jmp     .LBB0_1
  42     //                                   |          .LBB0_4:
  43     //                                   |                  ret
  44     // clang-format on
  45     // Notice that inside the busy loop .LBB0_2 and .LBB0_3, only instructions
  46     // with load semantics are used. swpab/xchg is only issued in outer loop
  47     // .LBB0_1. This is useful to avoid extra write traffic. The cache
  48     // coherence guarantees "write propagation", so even if the inner loop only
  49     // reads with relaxed ordering, the thread will evetually see the write.
  50     while (!try_lock())
  51       while (flag.load(cpp::MemoryOrder::RELAXED))
  52         sleep_briefly();
  53   }
  54   LIBC_INLINE void unlock() { flag.store(0u, cpp::MemoryOrder::RELEASE); }
  55   LIBC_INLINE bool is_locked() { return flag.load(cpp::MemoryOrder::ACQUIRE); }
  56   LIBC_INLINE bool is_invalid() {
  57     return flag.load(cpp::MemoryOrder::ACQUIRE) > 1;
  58   }
  59   // poison the lock
  60   LIBC_INLINE ~SpinLock() { flag.store(0xffu, cpp::MemoryOrder::RELEASE); }
  61 };
  62
  63 } // namespace LIBC_NAMESPACE_DECL
  64
  65 #endif // LLVM_LIBC_SRC___SUPPORT_THREADS_SPIN_LOCK_H