Implement mutex-free spin lock for task queue (#14834)

Implemented "lock-free" spinlock to save CPU usage on context switching.
The change has been tested on queene service of Ads team, the lock-free
version of ort (40 threads) saves CPU usage on gen8 (128 logical
processors on 8 numa nodes) windows by nearly half, from 65% to 35%.

For 32 cores, the curve is flat:

Anubis, 32 vCPU, windows, hugging face models,
95 percentile E2E latency in ms:

model | mutex(ms) | mutex-free
--- | --- | ---
 alvert_base_v2 | 34.21 | 34.09
 bert_large_uncased | 116.27| 117.84
 bart_base | 72.06 | 71.99
 distilgpt2 | 25.43 | 25.02
 vit_base_patch16_224 | 37.33 | 37.76

Anubis, 32 vCPU win, Linux, 1st party models,
95 percentile E2E latency in ms:

model | mutex(ms) | mutex-free
--- | --- | ---
deepthink_v2 | 24.35 | 22.95
bing_feeds |  36.96 | 36.48
deep_writes |  14.46 | 14.32
keypoints |  9.34 | 7.69
model11 |  1.71 | 1.66
model12 |  1.82 | 1.44
model2 |  4.21 | 3.95
model6 |  1.08 | 1.05
agiencoder |  0.99 | 0.93
geminet_transformer |  5.32 | 5.24

---------

Co-authored-by: Randy Shuai <rashuai@microsoft.com>

cmake/CMakeLists.txt

@@ -230,6 +230,7 @@ option(onnxruntime_BUILD_CACHE "onnxruntime build with cache" OFF)
 cmake_dependent_option(MSVC_Z7_OVERRIDE "replacing /Zi and /ZI with /Z7 when using MSVC with CCache" ON "onnxruntime_BUILD_CACHE; MSVC" OFF)
 
 option(onnxruntime_USE_AZURE "Build with azure inferencing support" OFF)
+option(onnxruntime_USE_LOCK_FREE_QUEUE "Build with lock-free task queue for threadpool." OFF)
 
 # ENABLE_TRAINING includes all training functionality
 # The following 2 entry points
@@ -752,7 +753,9 @@ if (onnxruntime_USE_AZURE)
     list(APPEND ORT_PROVIDER_CMAKE_FLAGS -Donnxruntime_USE_AZURE=1)
     list(APPEND ONNXRUNTIME_PROVIDER_NAMES azure)
 endif()
-
+if (onnxruntime_USE_LOCK_FREE_QUEUE)
+    add_compile_definitions(USE_LOCK_FREE_QUEUE)
+endif()
 
 if (onnxruntime_ENABLE_LAZY_TENSOR)
   # To support LazyTensor, ORT needs to call Python function from C/C++.

include/onnxruntime/core/platform/EigenNonBlockingThreadPool.h

@@ -44,6 +44,7 @@
 #include "core/common/inlined_containers_fwd.h"
 #include "core/common/spin_pause.h"
 #include "core/platform/ort_mutex.h"
+#include "core/platform/ort_spin_lock.h"
 #include "core/platform/Barrier.h"
 
 // ORT thread pool overview
@@ -449,7 +450,11 @@ class RunQueue {
   // PushBack adds w at the end of the queue.
   // If queue is full returns w, otherwise returns default-constructed Work.
   Work PushBack(Work w) {
+#ifdef USE_LOCK_FREE_QUEUE
+    std::lock_guard<OrtSpinLock> mtx(spin_lock_);
+#else
     std::lock_guard<OrtMutex> lock(mutex_);
+#endif
     unsigned back = back_.load(std::memory_order_relaxed);
     Elem& e = array_[(back - 1) & kMask];
     ElemState s = e.state.load(std::memory_order_relaxed);
@@ -469,7 +474,11 @@ class RunQueue {
   // with w_idx.  Typically the tag will be a per-thread ID to distinguish work
   // submitted from different threads.
   PushResult PushBackWithTag(Work w, Tag tag, unsigned& w_idx) {
+#ifdef USE_LOCK_FREE_QUEUE
+    std::lock_guard<OrtSpinLock> mtx(spin_lock_);
+#else
     std::lock_guard<OrtMutex> lock(mutex_);
+#endif
     unsigned back = back_.load(std::memory_order_relaxed);
     w_idx = (back - 1) & kMask;
     Elem& e = array_[w_idx];
@@ -490,7 +499,11 @@ class RunQueue {
   Work PopBack() {
     if (Empty())
       return Work();
+#ifdef USE_LOCK_FREE_QUEUE
+    std::lock_guard<OrtSpinLock> mtx(spin_lock_);
+#else
     std::lock_guard<OrtMutex> lock(mutex_);
+#endif
     unsigned back;
     Elem* e;
     ElemState s;
@@ -532,7 +545,11 @@ class RunQueue {
 
   bool RevokeWithTag(Tag tag, unsigned w_idx) {
     bool revoked = false;
+#ifdef USE_LOCK_FREE_QUEUE
+    std::lock_guard<OrtSpinLock> mtx(spin_lock_);
+#else
     std::lock_guard<OrtMutex> lock(mutex_);
+#endif
     Elem& e = array_[w_idx];
     ElemState s = e.state.load(std::memory_order_relaxed);
 
@@ -604,7 +621,11 @@ class RunQueue {
     Work w;
   };
 
+#ifdef USE_LOCK_FREE_QUEUE
+  OrtSpinLock spin_lock_;
+#else
   OrtMutex mutex_;
+#endif
 
   // Low log(kSize) + 1 bits in front_ and back_ contain rolling index of
   // front/back, respectively. The remaining bits contain modification counters

include/onnxruntime/core/platform/ort_spin_lock.h

@@ -0,0 +1,35 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+#pragma once
+#include "core/common/spin_pause.h"
+#include <atomic>
+
+namespace onnxruntime {
+/*
+OrtSpinLock implemented mutex semantic "lock-freely",
+calling thread will not be put to sleep on blocked,
+which reduces cpu usage on context switching.
+*/
+struct OrtSpinLock {
+  using LockState = enum { Locked = 0,
+                           Unlocked };
+
+  void lock() noexcept {
+    LockState state = Unlocked;
+    while (!state_.compare_exchange_weak(state, Locked, std::memory_order_acq_rel, std::memory_order_relaxed)) {
+      state = Unlocked;
+      concurrency::SpinPause();  // pause and retry
+    }
+  }
+  bool try_lock() noexcept {
+    LockState state = Unlocked;
+    return state_.compare_exchange_weak(state, Locked, std::memory_order_acq_rel, std::memory_order_relaxed);
+  }
+  void unlock() noexcept {
+    state_.store(Unlocked, std::memory_order_release);
+  }
+
+ private:
+  std::atomic<LockState> state_{Unlocked};
+};
+}  // namespace onnxruntime

tools/ci_build/build.py

@@ -687,6 +687,7 @@ def parse_arguments():
     parser.add_argument("--use_cache", action="store_true", help="Use compiler cache in CI")
 
     parser.add_argument("--use_triton_kernel", action="store_true", help="Use triton compiled kernels")
+    parser.add_argument("--use_lock_free_queue", action="store_true", help="Use lock-free task queue for threadpool.")
 
     if not is_windows():
         parser.add_argument(
@@ -1325,6 +1326,9 @@ def generate_build_tree(
     if args.use_azure:
         add_default_definition(cmake_extra_defines, "onnxruntime_USE_AZURE", "ON")
 
+    if args.use_lock_free_queue:
+        add_default_definition(cmake_extra_defines, "onnxruntime_USE_LOCK_FREE_QUEUE", "ON")
+
     cmake_args += [f"-D{define}" for define in cmake_extra_defines]
 
     cmake_args += cmake_extra_args

tools/ci_build/github/azure-pipelines/win-ci-pipeline.yml

@@ -200,7 +200,7 @@ stages:
         BuildConfig: 'RelWithDebInfo'
         EnvSetupScript: setup_env_azure.bat
         buildArch: x64
-        additionalBuildFlags: --use_azure
+        additionalBuildFlags: --use_azure --use_lock_free_queue
         msbuildPlatform: x64
         isX86: false
         job_name_suffix: x64_release_azure