blob: 93f68dcccb8a730ba9889e8240284cdde75f8dec (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
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
|
