未验证 提交 e285258e 编写于 作者: L Leo Chen 提交者: GitHub

[NPU] add lock for npu_pinned_allocator (#34700)

* add lock

* fix typo
上级 80092576
...@@ -39,6 +39,7 @@ void NPUPinnedAllocator::ProcessEventsAndFree() { ...@@ -39,6 +39,7 @@ void NPUPinnedAllocator::ProcessEventsAndFree() {
} }
Allocation *NPUPinnedAllocator::AllocateImpl(size_t size) { Allocation *NPUPinnedAllocator::AllocateImpl(size_t size) {
std::lock_guard<std::mutex> lock(mtx_);
ProcessEventsAndFree(); ProcessEventsAndFree();
void *ptr; void *ptr;
int error = posix_memalign(&ptr, kAlignment, size); int error = posix_memalign(&ptr, kAlignment, size);
...@@ -50,6 +51,7 @@ Allocation *NPUPinnedAllocator::AllocateImpl(size_t size) { ...@@ -50,6 +51,7 @@ Allocation *NPUPinnedAllocator::AllocateImpl(size_t size) {
} }
void NPUPinnedAllocator::FreeImpl(Allocation *allocation) { void NPUPinnedAllocator::FreeImpl(Allocation *allocation) {
std::lock_guard<std::mutex> lock(mtx_);
void *ptr = allocation->ptr(); void *ptr = allocation->ptr();
auto iter = npu_events_.find(allocation); auto iter = npu_events_.find(allocation);
aclrtEvent event = iter->second; aclrtEvent event = iter->second;
...@@ -65,11 +67,14 @@ void NPUPinnedAllocator::FreeImpl(Allocation *allocation) { ...@@ -65,11 +67,14 @@ void NPUPinnedAllocator::FreeImpl(Allocation *allocation) {
} }
uint64_t NPUPinnedAllocator::ReleaseImpl(const platform::Place &place) { uint64_t NPUPinnedAllocator::ReleaseImpl(const platform::Place &place) {
std::lock_guard<std::mutex> lock(mtx_);
// Empty implementation
return static_cast<uint64_t>(0); return static_cast<uint64_t>(0);
} }
void NPUPinnedAllocator::RecordEvent(Allocation *allocation, void NPUPinnedAllocator::RecordEvent(Allocation *allocation,
aclrtStream stream) { aclrtStream stream) {
std::lock_guard<std::mutex> lock(mtx_);
aclrtEvent event = nullptr; aclrtEvent event = nullptr;
PADDLE_ENFORCE_NPU_SUCCESS(aclrtCreateEvent(&event)); PADDLE_ENFORCE_NPU_SUCCESS(aclrtCreateEvent(&event));
PADDLE_ENFORCE_NPU_SUCCESS(aclrtRecordEvent(event, stream)); PADDLE_ENFORCE_NPU_SUCCESS(aclrtRecordEvent(event, stream));
......
...@@ -42,6 +42,7 @@ class NPUPinnedAllocator : public Allocator { ...@@ -42,6 +42,7 @@ class NPUPinnedAllocator : public Allocator {
private: private:
std::unordered_map<Allocation *, aclrtEvent> npu_events_; std::unordered_map<Allocation *, aclrtEvent> npu_events_;
mutable std::mutex mtx_;
}; };
} // namespace allocation } // namespace allocation
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册