提交 0f96c2e8 编写于 作者: S sneaxiy

fix thread-safety bug

test=develop
上级 90076522
...@@ -120,14 +120,25 @@ class EigenCudaStreamDevice : public Eigen::StreamInterface { ...@@ -120,14 +120,25 @@ class EigenCudaStreamDevice : public Eigen::StreamInterface {
} }
void* allocate(size_t num_bytes) const override { void* allocate(size_t num_bytes) const override {
if (UNLIKELY(num_bytes == 0)) {
return nullptr;
}
auto buf = paddle::memory::Alloc(place_, num_bytes, auto buf = paddle::memory::Alloc(place_, num_bytes,
memory::Allocator::kScratchpad); memory::Allocator::kScratchpad);
void* retv = buf->ptr(); void* retv = buf->ptr();
allocations_[buf->ptr()] = std::move(buf); {
std::lock_guard<std::mutex> lock(mtx_);
allocations_.emplace(retv, std::move(buf));
}
return retv; return retv;
} }
void deallocate(void* buffer) const override { allocations_.erase(buffer); } void deallocate(void* buffer) const override {
if (LIKELY(buffer)) {
std::lock_guard<std::mutex> lock(mtx_);
allocations_.erase(buffer);
}
}
void* scratchpad() const override { void* scratchpad() const override {
if (scratch_ == NULL) { if (scratch_ == NULL) {
...@@ -153,6 +164,7 @@ class EigenCudaStreamDevice : public Eigen::StreamInterface { ...@@ -153,6 +164,7 @@ class EigenCudaStreamDevice : public Eigen::StreamInterface {
const cudaDeviceProp* device_prop_; // not owned; const cudaDeviceProp* device_prop_; // not owned;
mutable void* scratch_; mutable void* scratch_;
mutable unsigned int* semaphore_; mutable unsigned int* semaphore_;
mutable std::mutex mtx_; // to protect allocations_
mutable std::unordered_map<void*, memory::AllocationPtr> allocations_; mutable std::unordered_map<void*, memory::AllocationPtr> allocations_;
}; };
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册