1 // 2 // Copyright © 2021 Arm Ltd and Contributors. All rights reserved. 3 // SPDX-License-Identifier: MIT 4 // 5 6 #include <backendsCommon/MemoryManager.hpp> 7 #include <armnn/utility/IgnoreUnused.hpp> 8 9 #include <doctest/doctest.h> 10 #include <numeric> 11 12 namespace armnn 13 { 14 /// @brief Class that implements a sample custom allocator. 15 class SampleCustomAllocator : public armnn::ICustomAllocator 16 { 17 public: 18 SampleCustomAllocator() = default; 19 allocate(size_t size,size_t alignment)20 void* allocate(size_t size, size_t alignment) override 21 { 22 IgnoreUnused(alignment); 23 CHECK(size == m_Values.size()); 24 m_CounterAllocate+=1; 25 return m_Values.data(); 26 } 27 free(void * ptr)28 void free(void* ptr) override 29 { 30 CHECK(ptr == m_Values.data()); 31 m_CounterFree+=1; 32 } 33 GetMemorySourceType()34 armnn::MemorySource GetMemorySourceType() override 35 { 36 return armnn::MemorySource::Malloc; 37 } 38 GetMemoryRegionAtOffset(void * buffer,size_t offset,size_t alignment=0)39 virtual void* GetMemoryRegionAtOffset(void* buffer, size_t offset, size_t alignment = 0 ) override 40 { 41 IgnoreUnused(alignment); 42 return (static_cast<char*>(buffer) + offset); 43 } 44 45 /// Holds the data in the tensors. Create for testing purposes. 46 std::vector<uint8_t> m_Values; 47 /// Counts the number of times the function allocate is called. 48 unsigned long m_CounterAllocate= 0; 49 /// Counts the number of times the function free is called. 50 unsigned long m_CounterFree = 0; 51 }; 52 53 TEST_SUITE("MemoryManagerTests") 54 { 55 /// Unit test Storing, Allocating and Deallocating with a custom allocator. 56 TEST_CASE("MemoryManagerTest") 57 { 58 using namespace armnn; 59 60 // Create mock up bufferStorageVector with 2 BufferStorage with the same TensorMemory 61 size_t numTensors = 5; 62 std::vector<std::shared_ptr<TensorMemory>> tensorMemoryPointerVector(numTensors); 63 std::vector<std::shared_ptr<TensorMemory>> tensorMemoryVector; 64 tensorMemoryVector.reserve(numTensors); 65 66 std::vector<size_t> offsets(numTensors); 67 std::iota(std::begin(offsets), std::end(offsets), 0); 68 69 for (uint32_t idx = 0; idx < tensorMemoryPointerVector.size(); ++idx) 70 { 71 tensorMemoryVector.emplace_back(std::make_shared<TensorMemory>(TensorMemory{offsets[idx], 0, nullptr})); 72 73 tensorMemoryPointerVector[idx] = tensorMemoryVector[idx]; 74 } 75 76 std::vector<BufferStorage> bufferStorageVector; 77 bufferStorageVector.emplace_back(BufferStorage{tensorMemoryPointerVector, numTensors}); 78 bufferStorageVector.emplace_back(BufferStorage{tensorMemoryPointerVector, numTensors}); 79 80 // Create an instance of the SampleCustomAllocator 81 std::shared_ptr<SampleCustomAllocator> customAllocator = 82 std::make_unique<SampleCustomAllocator>(SampleCustomAllocator()); 83 84 customAllocator->m_Values = {10, 11, 12, 13, 14}; 85 // Check that the test was set up correctly 86 CHECK(customAllocator->m_Values.size() == numTensors); 87 88 size_t bufferVecSize = bufferStorageVector.size(); 89 // Utilise 3 functions in the MemoryManager. Check the counters and the pointer to the values are correct. 90 MemoryManager memoryManager; 91 memoryManager.StoreMemToAllocate(bufferStorageVector, customAllocator); 92 93 memoryManager.Allocate(); 94 CHECK(customAllocator->m_CounterAllocate == bufferVecSize); 95 96 uint32_t idx = 0; 97 for (auto tensorMemory : tensorMemoryVector) 98 { 99 auto value = reinterpret_cast<uint8_t *>(tensorMemory->m_Data); 100 CHECK(customAllocator->m_Values[idx] == *value); 101 idx += 1; 102 } 103 104 memoryManager.Deallocate(); 105 CHECK(customAllocator->m_CounterFree == bufferStorageVector.size()); 106 } 107 } 108 109 } // namespace armnn