Merge pull request #14745 from sneaxiy/fix_eigen_deallocate

Fix eigen deallocate bug
revert-14786-revert-14782-revert-14398-imperative
Zeng Jinle 6 years ago committed by GitHub
commit add98c9e7d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -120,15 +120,24 @@ 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 { void deallocate(void* buffer) const override {
allocations_.erase(allocations_.find(buffer)); if (LIKELY(buffer)) {
std::lock_guard<std::mutex> lock(mtx_);
allocations_.erase(buffer);
}
} }
void* scratchpad() const override { void* scratchpad() const override {
@ -155,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_;
}; };

Loading…
Cancel
Save