diff options
author | Manjunath Kudlur <keveman@gmail.com> | 2015-11-06 16:27:58 -0800 |
---|---|---|
committer | Manjunath Kudlur <keveman@gmail.com> | 2015-11-06 16:27:58 -0800 |
commit | f41959ccb2d9d4c722fe8fc3351401d53bcf4900 (patch) | |
tree | ef0ca22cb2a5ac4bdec9d080d8e0788a53ed496d /tensorflow/core/framework/tracking_allocator_test.cc |
TensorFlow: Initial commit of TensorFlow library.
TensorFlow is an open source software library for numerical computation
using data flow graphs.
Base CL: 107276108
Diffstat (limited to 'tensorflow/core/framework/tracking_allocator_test.cc')
-rw-r--r-- | tensorflow/core/framework/tracking_allocator_test.cc | 115 |
1 files changed, 115 insertions, 0 deletions
diff --git a/tensorflow/core/framework/tracking_allocator_test.cc b/tensorflow/core/framework/tracking_allocator_test.cc new file mode 100644 index 0000000000..90ce851775 --- /dev/null +++ b/tensorflow/core/framework/tracking_allocator_test.cc @@ -0,0 +1,115 @@ +#include "tensorflow/core/framework/tracking_allocator.h" + +#include <unordered_map> + +#include "tensorflow/core/framework/allocator.h" +#include "tensorflow/core/platform/logging.h" +#include <gtest/gtest.h> + +namespace tensorflow { + +class TestableSizeTrackingAllocator : public Allocator { + public: + string Name() override { return "test"; } + void* AllocateRaw(size_t /*alignment*/, size_t num_bytes) override { + void* ptr = malloc(num_bytes); + size_map_[ptr] = num_bytes; + return ptr; + } + void DeallocateRaw(void* ptr) override { + const auto& iter = size_map_.find(ptr); + EXPECT_NE(size_map_.end(), iter); + size_map_.erase(iter); + free(ptr); + } + bool TracksAllocationSizes() override { return true; } + size_t RequestedSize(void* ptr) override { + const auto& iter = size_map_.find(ptr); + EXPECT_NE(size_map_.end(), iter); + return iter->second; + } + + private: + std::unordered_map<void*, size_t> size_map_; +}; + +class NoMemoryAllocator : public Allocator { + public: + string Name() override { return "test"; } + void* AllocateRaw(size_t /*alignment*/, size_t num_bytes) override { + return nullptr; + } + void DeallocateRaw(void* ptr) override {} + bool TracksAllocationSizes() override { return true; } +}; + +TEST(TrackingAllocatorTest, SimpleNoTracking) { + Allocator* a = cpu_allocator(); + + EXPECT_FALSE(a->TracksAllocationSizes()); + + TrackingAllocator* ta = new TrackingAllocator(a); + + void* p1 = ta->AllocateRaw(4, 4); + ta->Deallocate(p1); + void* p2 = ta->AllocateRaw(4, 12); + + std::pair<size_t, size_t> sizes = ta->GetSizesAndUnRef(); + + EXPECT_EQ(16, sizes.first); + EXPECT_EQ(0, sizes.second); + + ta->Deallocate(p2); +} + +TEST(TrackingAllocatorTest, SimpleTracking) { + TestableSizeTrackingAllocator a = TestableSizeTrackingAllocator(); + + EXPECT_TRUE(a.TracksAllocationSizes()); + + TrackingAllocator* ta = new TrackingAllocator(&a); + + void* p1 = ta->AllocateRaw(4, 12); + ta->Deallocate(p1); + void* p2 = ta->AllocateRaw(4, 4); + + std::pair<size_t, size_t> sizes = ta->GetSizesAndUnRef(); + + EXPECT_EQ(16, sizes.first); + EXPECT_EQ(12, sizes.second); + + ta->Deallocate(p2); +} + +TEST(TrackingAllocatorTest, OutOfMemory) { + NoMemoryAllocator a; + + EXPECT_TRUE(a.TracksAllocationSizes()); + + TrackingAllocator* ta = new TrackingAllocator(&a); + + void* p1 = ta->AllocateRaw(4, 12); + EXPECT_EQ(nullptr, p1); + + std::pair<size_t, size_t> sizes = ta->GetSizesAndUnRef(); + + EXPECT_EQ(0, sizes.first); + EXPECT_EQ(0, sizes.second); +} + +TEST(TrackingAllocatorTest, FreeNullPtr) { + NoMemoryAllocator a; + + EXPECT_TRUE(a.TracksAllocationSizes()); + + TrackingAllocator* ta = new TrackingAllocator(&a); + + ta->DeallocateRaw(nullptr); + + std::pair<size_t, size_t> sizes = ta->GetSizesAndUnRef(); + + EXPECT_EQ(0, sizes.first); + EXPECT_EQ(0, sizes.second); +} + +} // namespace tensorflow |