From e285258e189d4b0bb60519f2a27a68c88389ad80 Mon Sep 17 00:00:00 2001
From: Leo Chen <chenqiuliang@baidu.com>
Date: Mon, 9 Aug 2021 18:57:36 +0800
Subject: [PATCH] [NPU] add lock for npu_pinned_allocator (#34700)

* add lock

* fix typo
---
 paddle/fluid/memory/allocation/npu_pinned_allocator.cc | 5 +++++
 paddle/fluid/memory/allocation/npu_pinned_allocator.h  | 1 +
 2 files changed, 6 insertions(+)

diff --git a/paddle/fluid/memory/allocation/npu_pinned_allocator.cc b/paddle/fluid/memory/allocation/npu_pinned_allocator.cc
index 507a8589d94..9178825efa9 100644
--- a/paddle/fluid/memory/allocation/npu_pinned_allocator.cc
+++ b/paddle/fluid/memory/allocation/npu_pinned_allocator.cc
@@ -39,6 +39,7 @@ void NPUPinnedAllocator::ProcessEventsAndFree() {
 }
 
 Allocation *NPUPinnedAllocator::AllocateImpl(size_t size) {
+  std::lock_guard<std::mutex> lock(mtx_);
   ProcessEventsAndFree();
   void *ptr;
   int error = posix_memalign(&ptr, kAlignment, size);
@@ -50,6 +51,7 @@ Allocation *NPUPinnedAllocator::AllocateImpl(size_t size) {
 }
 
 void NPUPinnedAllocator::FreeImpl(Allocation *allocation) {
+  std::lock_guard<std::mutex> lock(mtx_);
   void *ptr = allocation->ptr();
   auto iter = npu_events_.find(allocation);
   aclrtEvent event = iter->second;
@@ -65,11 +67,14 @@ void NPUPinnedAllocator::FreeImpl(Allocation *allocation) {
 }
 
 uint64_t NPUPinnedAllocator::ReleaseImpl(const platform::Place &place) {
+  std::lock_guard<std::mutex> lock(mtx_);
+  // Empty implementation
   return static_cast<uint64_t>(0);
 }
 
 void NPUPinnedAllocator::RecordEvent(Allocation *allocation,
                                      aclrtStream stream) {
+  std::lock_guard<std::mutex> lock(mtx_);
   aclrtEvent event = nullptr;
   PADDLE_ENFORCE_NPU_SUCCESS(aclrtCreateEvent(&event));
   PADDLE_ENFORCE_NPU_SUCCESS(aclrtRecordEvent(event, stream));
diff --git a/paddle/fluid/memory/allocation/npu_pinned_allocator.h b/paddle/fluid/memory/allocation/npu_pinned_allocator.h
index 4c856b931ee..b330b6e352c 100644
--- a/paddle/fluid/memory/allocation/npu_pinned_allocator.h
+++ b/paddle/fluid/memory/allocation/npu_pinned_allocator.h
@@ -42,6 +42,7 @@ class NPUPinnedAllocator : public Allocator {
 
  private:
   std::unordered_map<Allocation *, aclrtEvent> npu_events_;
+  mutable std::mutex mtx_;
 };
 
 }  // namespace allocation
-- 
GitLab