From 0efda9d970d6b8cca86392cb193de6a96eaef3bb Mon Sep 17 00:00:00 2001
From: wanghuancoder <wanghuancoder@163.com>
Date: Thu, 26 Aug 2021 09:28:28 +0800
Subject: [PATCH] use spinlock in auto growth (#35139)

* use spinlock in auto growth, test=develop

* refine,test=develop
---
 .../auto_growth_best_fit_allocator.cc         |  4 +-
 .../auto_growth_best_fit_allocator.h          |  3 +-
 paddle/fluid/memory/allocation/spin_lock.h    | 54 +++++++++++++++++++
 3 files changed, 58 insertions(+), 3 deletions(-)
 create mode 100644 paddle/fluid/memory/allocation/spin_lock.h

diff --git a/paddle/fluid/memory/allocation/auto_growth_best_fit_allocator.cc b/paddle/fluid/memory/allocation/auto_growth_best_fit_allocator.cc
index d29a33b4701..cca29797bb6 100644
--- a/paddle/fluid/memory/allocation/auto_growth_best_fit_allocator.cc
+++ b/paddle/fluid/memory/allocation/auto_growth_best_fit_allocator.cc
@@ -45,7 +45,7 @@ AutoGrowthBestFitAllocator::AutoGrowthBestFitAllocator(
 Allocation *AutoGrowthBestFitAllocator::AllocateImpl(size_t size) {
   size = AlignedSize(size, alignment_);
 
-  std::lock_guard<std::mutex> guard(mtx_);
+  std::lock_guard<SpinLock> guard(spinlock_);
   auto iter = free_blocks_.lower_bound(std::make_pair(size, nullptr));
   BlockIt block_it;
   if (iter != free_blocks_.end()) {
@@ -98,7 +98,7 @@ Allocation *AutoGrowthBestFitAllocator::AllocateImpl(size_t size) {
 }
 
 void AutoGrowthBestFitAllocator::FreeImpl(Allocation *allocation) {
-  std::lock_guard<std::mutex> guard(mtx_);
+  std::lock_guard<SpinLock> guard(spinlock_);
   auto block_it = static_cast<BlockAllocation *>(allocation)->block_it_;
   auto &blocks = block_it->chunk_->blocks_;
 
diff --git a/paddle/fluid/memory/allocation/auto_growth_best_fit_allocator.h b/paddle/fluid/memory/allocation/auto_growth_best_fit_allocator.h
index eb52cab2594..5ed6eb94f15 100644
--- a/paddle/fluid/memory/allocation/auto_growth_best_fit_allocator.h
+++ b/paddle/fluid/memory/allocation/auto_growth_best_fit_allocator.h
@@ -21,6 +21,7 @@
 #include <utility>
 
 #include "paddle/fluid/memory/allocation/allocator.h"
+#include "paddle/fluid/memory/allocation/spin_lock.h"
 
 namespace paddle {
 namespace memory {
@@ -86,7 +87,7 @@ class AutoGrowthBestFitAllocator : public Allocator {
   size_t alignment_;
   size_t chunk_size_;
 
-  mutable std::mutex mtx_;
+  SpinLock spinlock_;
 };
 
 }  // namespace allocation
diff --git a/paddle/fluid/memory/allocation/spin_lock.h b/paddle/fluid/memory/allocation/spin_lock.h
new file mode 100644
index 00000000000..bc46fa2b370
--- /dev/null
+++ b/paddle/fluid/memory/allocation/spin_lock.h
@@ -0,0 +1,54 @@
+// Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <atomic>
+#if !defined(_WIN32)
+#include <sched.h>
+#else
+#include <windows.h>
+#endif  // !_WIN32
+
+namespace paddle {
+namespace memory {
+
+class SpinLock {
+ public:
+  SpinLock() : mlock_(false) {}
+
+  void lock() {
+    bool expect = false;
+    uint64_t spin_cnt = 0;
+    while (!mlock_.compare_exchange_weak(expect, true)) {
+      expect = false;
+      if ((++spin_cnt & 0xFF) == 0) {
+#if defined(_WIN32)
+        SleepEx(50, FALSE);
+#else
+        sched_yield();
+#endif
+      }
+    }
+  }
+
+  void unlock() { mlock_.store(false); }
+  DISABLE_COPY_AND_ASSIGN(SpinLock);
+
+ private:
+  std::atomic<bool> mlock_;
+};
+
+}  // namespace memory
+}  // namespace paddle
-- 
GitLab