1c8b59c046895fa5b6d79f73e0b5817330fcfbfc1A. Unique TensorFlower/* Copyright 2015 The TensorFlow Authors. All Rights Reserved. 29c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath Kudlur 39c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath KudlurLicensed under the Apache License, Version 2.0 (the "License"); 49c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath Kudluryou may not use this file except in compliance with the License. 59c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath KudlurYou may obtain a copy of the License at 69c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath Kudlur 79c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath Kudlur http://www.apache.org/licenses/LICENSE-2.0 89c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath Kudlur 99c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath KudlurUnless required by applicable law or agreed to in writing, software 109c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath Kudlurdistributed under the License is distributed on an "AS IS" BASIS, 119c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath KudlurWITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 129c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath KudlurSee the License for the specific language governing permissions and 139c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath Kudlurlimitations under the License. 149c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath Kudlur==============================================================================*/ 159c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath Kudlur 16f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur#include "tensorflow/core/framework/tracking_allocator.h" 17f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 18f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur#include <unordered_map> 19f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 20f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur#include "tensorflow/core/framework/allocator.h" 21f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur#include "tensorflow/core/platform/logging.h" 2283c6e0c63acdcab2c58c4ed7220bfa58879b1d57Jonathan Hseu#include "tensorflow/core/platform/mem.h" 23c8eaac926c929e07ac8db69f67803a2223ff2d93Josh Levenberg#include "tensorflow/core/platform/test.h" 24f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 25f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlurnamespace tensorflow { 26f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 27f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlurclass TestableSizeTrackingAllocator : public Allocator { 28f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur public: 29f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur string Name() override { return "test"; } 30f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur void* AllocateRaw(size_t /*alignment*/, size_t num_bytes) override { 3183c6e0c63acdcab2c58c4ed7220bfa58879b1d57Jonathan Hseu void* ptr = port::Malloc(num_bytes); 32f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur size_map_[ptr] = num_bytes; 33f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur return ptr; 34f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur } 35f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur void DeallocateRaw(void* ptr) override { 36f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur const auto& iter = size_map_.find(ptr); 37f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur EXPECT_NE(size_map_.end(), iter); 38f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur size_map_.erase(iter); 3983c6e0c63acdcab2c58c4ed7220bfa58879b1d57Jonathan Hseu port::Free(ptr); 40f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur } 41f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur bool TracksAllocationSizes() override { return true; } 429126444b41b243ca9bc2359d8e91a05fc0039e71Yuefeng Zhou size_t RequestedSize(const void* ptr) override { 43f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur const auto& iter = size_map_.find(ptr); 44f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur EXPECT_NE(size_map_.end(), iter); 45f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur return iter->second; 46f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur } 47648e2654e1c14ee965493749f1d3620498c44ddbA. Unique TensorFlower void GetStats(AllocatorStats* stats) override { stats->Clear(); } 48f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 49f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur private: 509126444b41b243ca9bc2359d8e91a05fc0039e71Yuefeng Zhou std::unordered_map<const void*, size_t> size_map_; 51f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur}; 52f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 53f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlurclass NoMemoryAllocator : public Allocator { 54f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur public: 55f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur string Name() override { return "test"; } 56f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur void* AllocateRaw(size_t /*alignment*/, size_t num_bytes) override { 57f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur return nullptr; 58f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur } 59f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur void DeallocateRaw(void* ptr) override {} 60f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur bool TracksAllocationSizes() override { return true; } 61648e2654e1c14ee965493749f1d3620498c44ddbA. Unique TensorFlower void GetStats(AllocatorStats* stats) override { stats->Clear(); } 62f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur}; 63f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 64f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath KudlurTEST(TrackingAllocatorTest, SimpleNoTracking) { 65f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur Allocator* a = cpu_allocator(); 66f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 67f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur EXPECT_FALSE(a->TracksAllocationSizes()); 68f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 698f8f7f11b9e49ed8fb297ec5b5cf5be30f93359dBenoit Steiner // Don't enable the tracking inside the tracking allocator. Since 708f8f7f11b9e49ed8fb297ec5b5cf5be30f93359dBenoit Steiner // the cpu_allocator doesn't track allocations itself the tracking 718f8f7f11b9e49ed8fb297ec5b5cf5be30f93359dBenoit Steiner // will be partial 728f8f7f11b9e49ed8fb297ec5b5cf5be30f93359dBenoit Steiner TrackingAllocator* ta = new TrackingAllocator(a, false); 73f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 74f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur void* p1 = ta->AllocateRaw(4, 4); 75861f8f01334d20e998eae9a759c8f5a1e07721caA. Unique TensorFlower ta->DeallocateRaw(p1); 76f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur void* p2 = ta->AllocateRaw(4, 12); 77f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 7886238e8d09efce59de038b062a230030aa8bdd3aA. Unique TensorFlower std::tuple<size_t, size_t, size_t> sizes = ta->GetSizes(); 79f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 80f1289358d6664d16a50b7f7de1023740b434474eYuefeng Zhou EXPECT_EQ(16, std::get<0>(sizes)); 81f1289358d6664d16a50b7f7de1023740b434474eYuefeng Zhou EXPECT_EQ(0, std::get<1>(sizes)); 82f1289358d6664d16a50b7f7de1023740b434474eYuefeng Zhou EXPECT_EQ(0, std::get<2>(sizes)); 83f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 84861f8f01334d20e998eae9a759c8f5a1e07721caA. Unique TensorFlower ta->DeallocateRaw(p2); 8586238e8d09efce59de038b062a230030aa8bdd3aA. Unique TensorFlower auto records = ta->GetRecordsAndUnRef(); 8686238e8d09efce59de038b062a230030aa8bdd3aA. Unique TensorFlower EXPECT_EQ(4, records[0].alloc_bytes); 8786238e8d09efce59de038b062a230030aa8bdd3aA. Unique TensorFlower EXPECT_EQ(12, records[1].alloc_bytes); 888f8f7f11b9e49ed8fb297ec5b5cf5be30f93359dBenoit Steiner 898f8f7f11b9e49ed8fb297ec5b5cf5be30f93359dBenoit Steiner // This time enable the tracking inside the tracking allocator 908f8f7f11b9e49ed8fb297ec5b5cf5be30f93359dBenoit Steiner ta = new TrackingAllocator(a, true); 918f8f7f11b9e49ed8fb297ec5b5cf5be30f93359dBenoit Steiner p1 = ta->AllocateRaw(4, 4); 928f8f7f11b9e49ed8fb297ec5b5cf5be30f93359dBenoit Steiner EXPECT_EQ(4, ta->RequestedSize(p1)); 938f8f7f11b9e49ed8fb297ec5b5cf5be30f93359dBenoit Steiner EXPECT_LE(4, ta->AllocatedSize(p1)); 948f8f7f11b9e49ed8fb297ec5b5cf5be30f93359dBenoit Steiner EXPECT_EQ(1, ta->AllocationId(p1)); 958f8f7f11b9e49ed8fb297ec5b5cf5be30f93359dBenoit Steiner 968f8f7f11b9e49ed8fb297ec5b5cf5be30f93359dBenoit Steiner ta->DeallocateRaw(p1); 978f8f7f11b9e49ed8fb297ec5b5cf5be30f93359dBenoit Steiner p2 = ta->AllocateRaw(4, 12); 988f8f7f11b9e49ed8fb297ec5b5cf5be30f93359dBenoit Steiner EXPECT_EQ(12, ta->RequestedSize(p2)); 998f8f7f11b9e49ed8fb297ec5b5cf5be30f93359dBenoit Steiner EXPECT_LE(12, ta->AllocatedSize(p2)); 1008f8f7f11b9e49ed8fb297ec5b5cf5be30f93359dBenoit Steiner EXPECT_EQ(2, ta->AllocationId(p2)); 1018f8f7f11b9e49ed8fb297ec5b5cf5be30f93359dBenoit Steiner 10286238e8d09efce59de038b062a230030aa8bdd3aA. Unique TensorFlower sizes = ta->GetSizes(); 1038f8f7f11b9e49ed8fb297ec5b5cf5be30f93359dBenoit Steiner 104f1289358d6664d16a50b7f7de1023740b434474eYuefeng Zhou EXPECT_LE(16, std::get<0>(sizes)); 105f1289358d6664d16a50b7f7de1023740b434474eYuefeng Zhou EXPECT_LE(12, std::get<1>(sizes)); 106f1289358d6664d16a50b7f7de1023740b434474eYuefeng Zhou EXPECT_LE(12, std::get<2>(sizes)); 1078f8f7f11b9e49ed8fb297ec5b5cf5be30f93359dBenoit Steiner 1088f8f7f11b9e49ed8fb297ec5b5cf5be30f93359dBenoit Steiner ta->DeallocateRaw(p2); 10986238e8d09efce59de038b062a230030aa8bdd3aA. Unique TensorFlower records = ta->GetRecordsAndUnRef(); 11086238e8d09efce59de038b062a230030aa8bdd3aA. Unique TensorFlower EXPECT_LE(4, records[0].alloc_bytes); 11186238e8d09efce59de038b062a230030aa8bdd3aA. Unique TensorFlower EXPECT_GE(-4, records[1].alloc_bytes); 11286238e8d09efce59de038b062a230030aa8bdd3aA. Unique TensorFlower EXPECT_LE(12, records[2].alloc_bytes); 11386238e8d09efce59de038b062a230030aa8bdd3aA. Unique TensorFlower EXPECT_GE(-12, records[3].alloc_bytes); 114f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur} 115f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 116f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath KudlurTEST(TrackingAllocatorTest, SimpleTracking) { 117f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur TestableSizeTrackingAllocator a = TestableSizeTrackingAllocator(); 118f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 119f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur EXPECT_TRUE(a.TracksAllocationSizes()); 120f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 1218f8f7f11b9e49ed8fb297ec5b5cf5be30f93359dBenoit Steiner TrackingAllocator* ta = new TrackingAllocator(&a, false); 122f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 123f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur void* p1 = ta->AllocateRaw(4, 12); 124861f8f01334d20e998eae9a759c8f5a1e07721caA. Unique TensorFlower ta->DeallocateRaw(p1); 125f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur void* p2 = ta->AllocateRaw(4, 4); 126f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 12786238e8d09efce59de038b062a230030aa8bdd3aA. Unique TensorFlower std::tuple<size_t, size_t, size_t> sizes = ta->GetSizes(); 128f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 129f1289358d6664d16a50b7f7de1023740b434474eYuefeng Zhou EXPECT_EQ(16, std::get<0>(sizes)); 130f1289358d6664d16a50b7f7de1023740b434474eYuefeng Zhou EXPECT_EQ(12, std::get<1>(sizes)); 131f1289358d6664d16a50b7f7de1023740b434474eYuefeng Zhou EXPECT_EQ(4, std::get<2>(sizes)); 132f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 133861f8f01334d20e998eae9a759c8f5a1e07721caA. Unique TensorFlower ta->DeallocateRaw(p2); 13486238e8d09efce59de038b062a230030aa8bdd3aA. Unique TensorFlower 13586238e8d09efce59de038b062a230030aa8bdd3aA. Unique TensorFlower auto records = ta->GetRecordsAndUnRef(); 13686238e8d09efce59de038b062a230030aa8bdd3aA. Unique TensorFlower EXPECT_EQ(12, records[0].alloc_bytes); 13786238e8d09efce59de038b062a230030aa8bdd3aA. Unique TensorFlower EXPECT_EQ(-12, records[1].alloc_bytes); 13886238e8d09efce59de038b062a230030aa8bdd3aA. Unique TensorFlower EXPECT_EQ(4, records[2].alloc_bytes); 13986238e8d09efce59de038b062a230030aa8bdd3aA. Unique TensorFlower EXPECT_EQ(-4, records[3].alloc_bytes); 140f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur} 141f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 142f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath KudlurTEST(TrackingAllocatorTest, OutOfMemory) { 143f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur NoMemoryAllocator a; 144f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 145f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur EXPECT_TRUE(a.TracksAllocationSizes()); 146f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 1478f8f7f11b9e49ed8fb297ec5b5cf5be30f93359dBenoit Steiner TrackingAllocator* ta = new TrackingAllocator(&a, false); 148f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 149f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur void* p1 = ta->AllocateRaw(4, 12); 150f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur EXPECT_EQ(nullptr, p1); 151f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 15286238e8d09efce59de038b062a230030aa8bdd3aA. Unique TensorFlower std::tuple<size_t, size_t, size_t> sizes = ta->GetSizes(); 153f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 154f1289358d6664d16a50b7f7de1023740b434474eYuefeng Zhou EXPECT_EQ(0, std::get<0>(sizes)); 155f1289358d6664d16a50b7f7de1023740b434474eYuefeng Zhou EXPECT_EQ(0, std::get<1>(sizes)); 156f1289358d6664d16a50b7f7de1023740b434474eYuefeng Zhou EXPECT_EQ(0, std::get<2>(sizes)); 15786238e8d09efce59de038b062a230030aa8bdd3aA. Unique TensorFlower 15886238e8d09efce59de038b062a230030aa8bdd3aA. Unique TensorFlower EXPECT_EQ(0, ta->GetRecordsAndUnRef().size()); 159f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur} 160f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 161f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath KudlurTEST(TrackingAllocatorTest, FreeNullPtr) { 162f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur NoMemoryAllocator a; 163f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 164f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur EXPECT_TRUE(a.TracksAllocationSizes()); 165f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 1668f8f7f11b9e49ed8fb297ec5b5cf5be30f93359dBenoit Steiner TrackingAllocator* ta = new TrackingAllocator(&a, false); 167f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 168f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur ta->DeallocateRaw(nullptr); 169f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 17086238e8d09efce59de038b062a230030aa8bdd3aA. Unique TensorFlower std::tuple<size_t, size_t, size_t> sizes = ta->GetSizes(); 171f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 172f1289358d6664d16a50b7f7de1023740b434474eYuefeng Zhou EXPECT_EQ(0, std::get<0>(sizes)); 173f1289358d6664d16a50b7f7de1023740b434474eYuefeng Zhou EXPECT_EQ(0, std::get<1>(sizes)); 174f1289358d6664d16a50b7f7de1023740b434474eYuefeng Zhou EXPECT_EQ(0, std::get<2>(sizes)); 17586238e8d09efce59de038b062a230030aa8bdd3aA. Unique TensorFlower 17686238e8d09efce59de038b062a230030aa8bdd3aA. Unique TensorFlower EXPECT_EQ(0, ta->GetRecordsAndUnRef().size()); 177f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur} 178f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 179f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur} // namespace tensorflow 180