diff options
Diffstat (limited to 'tensorflow/core/framework/allocator.cc')
-rw-r--r-- | tensorflow/core/framework/allocator.cc | 25 |
1 files changed, 25 insertions, 0 deletions
diff --git a/tensorflow/core/framework/allocator.cc b/tensorflow/core/framework/allocator.cc new file mode 100644 index 0000000000..93f68dcccb --- /dev/null +++ b/tensorflow/core/framework/allocator.cc @@ -0,0 +1,25 @@ +#include "tensorflow/core/framework/allocator.h" +#include "tensorflow/core/platform/port.h" + +namespace tensorflow { + +Allocator::~Allocator() {} + +class CPUAllocator : public Allocator { + public: + ~CPUAllocator() override {} + + string Name() override { return "cpu"; } + void* AllocateRaw(size_t alignment, size_t num_bytes) override { + return port::aligned_malloc(num_bytes, alignment); + } + + void DeallocateRaw(void* ptr) override { port::aligned_free(ptr); } +}; + +Allocator* cpu_allocator() { + static CPUAllocator* cpu_alloc = new CPUAllocator; + return cpu_alloc; +} + +} // namespace tensorflow |