diff --git a/paddle/fluid/memory/allocation/npu_pinned_allocator.cc b/paddle/fluid/memory/allocation/npu_pinned_allocator.cc index 507a8589d94ddd1adf925aa5e01c787439624c62..9178825efa9e1da0ad43fee13f5f64d992f91cd9 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 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 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 lock(mtx_); + // Empty implementation return static_cast(0); } void NPUPinnedAllocator::RecordEvent(Allocation *allocation, aclrtStream stream) { + std::lock_guard 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 4c856b931ee2cf5b5734d90636b4bfd3dad138da..b330b6e352ce42140132d64994f2d76d3ff71e41 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 npu_events_; + mutable std::mutex mtx_; }; } // namespace allocation