mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
74dc4460eb
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/12029 In order to remove New() function in StaticContext(to remove StaticContext) and converge to the Allocator design, we'll first change the return type of New to at::DataPtr. Reviewed By: ezyang Differential Revision: D9889990 fbshipit-source-id: 3257c763530b987025f428741bdd2e089d11bad4
50 lines
1.3 KiB
C++
50 lines
1.3 KiB
C++
#include "caffe2/core/context.h"
|
|
#include "caffe2/core/logging.h"
|
|
#include "caffe2/core/typeid.h"
|
|
|
|
CAFFE2_DEFINE_bool(
|
|
caffe2_report_cpu_memory_usage,
|
|
false,
|
|
"If set, print out detailed memory usage");
|
|
|
|
CAFFE2_DEFINE_bool(
|
|
caffe2_cpu_allocator_do_zero_fill,
|
|
true,
|
|
"If set, do memory zerofilling when allocating on CPU");
|
|
|
|
namespace caffe2 {
|
|
|
|
void NoDelete(void*) {}
|
|
|
|
static std::unique_ptr<at::Allocator> g_cpu_allocator(
|
|
new DefaultCPUAllocator());
|
|
at::Allocator* GetCPUAllocator() {
|
|
return g_cpu_allocator.get();
|
|
}
|
|
|
|
void SetCPUAllocator(at::Allocator* alloc) {
|
|
g_cpu_allocator.reset(alloc);
|
|
}
|
|
|
|
MemoryAllocationReporter DefaultCPUAllocator::reporter_;
|
|
|
|
void MemoryAllocationReporter::New(void* ptr, size_t nbytes) {
|
|
std::lock_guard<std::mutex> guard(mutex_);
|
|
size_table_[ptr] = nbytes;
|
|
allocated_ += nbytes;
|
|
LOG(INFO) << "Caffe2 alloc " << nbytes << " bytes, total alloc " << allocated_
|
|
<< " bytes.";
|
|
}
|
|
|
|
void MemoryAllocationReporter::Delete(void* ptr) {
|
|
std::lock_guard<std::mutex> guard(mutex_);
|
|
auto it = size_table_.find(ptr);
|
|
CHECK(it != size_table_.end());
|
|
allocated_ -= it->second;
|
|
LOG(INFO) << "Caffe2 deleted " << it->second << " bytes, total alloc "
|
|
<< allocated_ << " bytes.";
|
|
size_table_.erase(it);
|
|
}
|
|
|
|
} // namespace caffe2
|