[atomics] Rewrite the fallback path of llvm::call_once to use std::atomic instead of the home-grown LLVM atomic type.

[chandlerc]

Summary

This seems strictly more likely to be correct and is also probably more
efficient on some platforms. I think I've gotten the memory orderings
correct, but a second pair of eyes is very welcome here, notably to make
sure that in all cases the Done synchronizes with works correctly and it
is not possible to release threads prior to the effects of calling the
function being visible to all threads.

If this looks good, and it sticks, and ...., I'll be able to remove
llvm::cas_flag entirely.

include/llvm/Support/Threading.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_SUPPORT_THREADING_H
 #define LLVM_SUPPORT_THREADING_H
 
 #include "llvm/Config/llvm-config.h" // for LLVM_ON_UNIX
 #include "llvm/Support/Compiler.h"
+#include <cassert>
 #include <ciso646> // So we can check the C++ standard lib macros.
 #include <functional>
 
 // We use std::call_once on all Unix platforms except for NetBSD with
 // libstdc++. That platform has a bug they are working to fix, and they'll
 // remove the NetBSD checks once fixed.
 #if defined(LLVM_ON_UNIX) &&                                                   \
     !(defined(__NetBSD__) && !defined(_LIBCPP_VERSION)) && !defined(__ppc__)
 #define LLVM_THREADING_USE_STD_CALL_ONCE 1
 #else
 #define LLVM_THREADING_USE_STD_CALL_ONCE 0
 #endif
 
 #if LLVM_THREADING_USE_STD_CALL_ONCE
 #include <mutex>
 #else
-#include "llvm/Support/Atomic.h"
+#include <atomic>
+#include <thread>
 #endif
 
 namespace llvm {
   /// Returns true if LLVM is compiled with support for multi-threading, and
   /// false otherwise.
   bool llvm_is_multithreaded();
 
   /// llvm_execute_on_thread - Execute the given \p UserFn on a separate
#if LLVM_THREADING_USE_STD_CALL_ONCE
 
   /// This macro is the only way you should define your once flag for LLVM's
   /// call_once.
 #define LLVM_DEFINE_ONCE_FLAG(flag) static once_flag flag
 
 #else
 
   enum InitStatus { Uninitialized = 0, Wait = 1, Done = 2 };
-  typedef volatile sys::cas_flag once_flag;
+  typedef std::atomic<InitStatus> once_flag;
 
   /// This macro is the only way you should define your once flag for LLVM's
   /// call_once.
-#define LLVM_DEFINE_ONCE_FLAG(flag) static once_flag flag = Uninitialized
+#define LLVM_DEFINE_ONCE_FLAG(flag) static once_flag flag = {Uninitialized}
 
 #endif
 
   /// \brief Execute the function specified as a parameter once.
   ///
   /// Typical usage:
   /// \code
   ///   void foo() {...};
   ///   ...
   ///   LLVM_DEFINE_ONCE_FLAG(flag);
   ///   call_once(flag, foo);
   /// \endcode
   ///
   /// \param flag Flag used for tracking whether or not this has run.
   /// \param F Function to call once.
   template <typename Function, typename... Args>
-  void call_once(once_flag &flag, Function &&F, Args &&... ArgList) {
+  void call_once(once_flag &Flag, Function &&F, Args &&... ArgList) {
 #if LLVM_THREADING_USE_STD_CALL_ONCE
-    std::call_once(flag, std::forward<Function>(F),
+    std::call_once(Flag, std::forward<Function>(F),
                    std::forward<Args>(ArgList)...);
 #else
     // For other platforms we use a generic (if brittle) version based on our
     // atomics.
-    sys::cas_flag old_val = sys::CompareAndSwap(&flag, Wait, Uninitialized);
-    if (old_val == Uninitialized) {
+    InitStatus OldFlag = Uninitialized;
+    if (Flag.compare_exchange_strong(OldFlag, Wait, std::memory_order_acq_rel)) {
+      // We've moved the flag to the Wait state which means we have exclusive
+      // access to call the function.
       std::forward<Function>(F)(std::forward<Args>(ArgList)...);
-      sys::MemoryFence();
-      TsanIgnoreWritesBegin();
-      TsanHappensBefore(&flag);
-      flag = Done;
-      TsanIgnoreWritesEnd();
-    } else {
-      // Wait until any thread doing the call has finished.
-      sys::cas_flag tmp = flag;
-      sys::MemoryFence();
-      while (tmp != Done) {
-        tmp = flag;
-        sys::MemoryFence();
-      }
-    }
-    TsanHappensAfter(&flag);
+      // Now release any other threads spinning and mark that no other threads need to acquire
+      Flag.store(Done, std::memory_order_release);
+      return;
+    assert(OldFlag != Uninitialized &&
+           "compare_exchange_strong failed to update!");
+
+    // If we already acquired the done state, no need to spin (and do more
+    // atomic operations).
+    if (OldFlag == Done)
+      return;
+
+    // Otherwise we're in the wait state, so spin until we acquire the done
+    // state.
+    assert(OldFlag == Wait && "Corrupted flag value!");
+    while (Flag.load(std::memory_order_acquire) != Done)
+      std::this_thread::yield();
 #endif
   }
 }
 
 #endif