diff --git a/NeoML/test/src/DnnBlobTest.cpp b/NeoML/test/src/DnnBlobTest.cpp index 44678a1a4..e99be5c96 100644 --- a/NeoML/test/src/DnnBlobTest.cpp +++ b/NeoML/test/src/DnnBlobTest.cpp @@ -28,16 +28,16 @@ TEST( CDnnBlobTest, InitWindowBlob ) CPtr parent = CDnnBlob::CreateDataBlob( MathEngine(), CT_Float, 16, 1, 1 ); CPtr blob = CDnnBlob::CreateWindowBlob( parent ); - ASSERT_FALSE( blob->GetData().IsNull() ); + EXPECT_FALSE( blob->GetData().IsNull() ); } TEST( CDnnBlobTest, BufferTest ) { CPtr blob = CDnnBlob::CreateDataBlob( MathEngine(), CT_Float, 16, 1, 1 ); - ASSERT_FALSE( blob->GetData().IsNull() ); + EXPECT_FALSE( blob->GetData().IsNull() ); CDnnBlobBuffer buffer( *blob, TDnnBlobBufferAccess::Write ); - ASSERT_NE( nullptr, buffer.Ptr() ); + EXPECT_NE( nullptr, buffer.Ptr() ); ::memset( buffer, 0, buffer.Size() * sizeof( float ) ); EXPECT_FALSE( buffer.IsClosed() ); @@ -47,6 +47,48 @@ TEST( CDnnBlobTest, BufferTest ) //--------------------------------------------------------------------------------------------------------------------- +TEST( CDnnBlobTest, BufferMemoryThresholdTest ) +{ + auto testMethod = []( size_t threshold, bool init, size_t &sumMemoryInPools ) + { + if( init ) { + MathEngine().CleanUp(); + MathEngine().SetReuseMemoryMode( true ); + MathEngine().SetThreadBufferMemoryThreshold( threshold ); + } + + MathEngine().ResetPeakMemoryUsage(); + const size_t peakMemory = MathEngine().GetPeakMemoryUsage(); + const size_t reusedMemory = ( init ? 0 : threshold ); + { + CPtr blob1 = CDnnBlob::CreateDataBlob( MathEngine(), CT_Float, int(threshold / sizeof( float )), 1, 1 ); + ASSERT_TRUE( blob1 != nullptr && !blob1->GetData().IsNull() ); + CPtr blob2 = CDnnBlob::CreateDataBlob( MathEngine(), CT_Float, int( threshold / sizeof( float ) + 1), 1, 1 ); + ASSERT_TRUE( blob2 != nullptr && !blob2->GetData().IsNull() ); + EXPECT_EQ( MathEngine().GetPeakMemoryUsage(), peakMemory + threshold - reusedMemory + threshold + sizeof( float ) ); + } + const size_t memoryInPools = MathEngine().GetMemoryInPools() - reusedMemory; + EXPECT_EQ( memoryInPools, threshold - reusedMemory ); + + MathEngine().ResetPeakMemoryUsage(); + EXPECT_EQ( MathEngine().GetPeakMemoryUsage() - memoryInPools, peakMemory ); + sumMemoryInPools += memoryInPools; + }; + + size_t sumMemoryInPools = 0; + { + testMethod( 256, /*init*/true, sumMemoryInPools ); + + std::thread thread( testMethod, 512, /*init*/true, std::ref( sumMemoryInPools ) ); + thread.join(); + + testMethod( 256, /*init*/false, sumMemoryInPools ); + } + EXPECT_EQ( sumMemoryInPools, 256 + 512 ); +} + +//--------------------------------------------------------------------------------------------------------------------- + #if FINE_PLATFORM( FINE_WINDOWS ) || !defined( NEOML_USE_FINEOBJ ) namespace NeoMLTest { diff --git a/NeoMathEngine/include/NeoMathEngine/NeoMathEngine.h b/NeoMathEngine/include/NeoMathEngine/NeoMathEngine.h index 7ad8ac46c..fb5922ad6 100644 --- a/NeoMathEngine/include/NeoMathEngine/NeoMathEngine.h +++ b/NeoMathEngine/include/NeoMathEngine/NeoMathEngine.h @@ -1,4 +1,4 @@ -/* Copyright © 2017-2023 ABBYY +/* Copyright © 2017-2024 ABBYY Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. @@ -1160,11 +1160,18 @@ class NEOMATHENGINE_API IMathEngine : public IDnnEngine { virtual void GetMathEngineInfo( CMathEngineInfo& info ) const = 0; // Memory management + // Turns on and off the memory reuse mode // In this mode, the allocated memory blocks will not be deleted on HeapFree() and may be used until CleanUp() virtual void SetReuseMemoryMode( bool enable ) = 0; + // Specialize the size threshold in bytes for the current thread, so + // memory blocks of a size <= this threshold would be allocated in buffers if 'reuse' mode enabled + // memory blocks of a size > this threshold would be allocated in raw RAM memory (malloc/free) + virtual void SetThreadBufferMemoryThreshold( size_t threshold ) = 0; + virtual CMemoryHandle HeapAlloc( size_t count ) = 0; virtual void HeapFree( const CMemoryHandle& handle ) = 0; + // Transfers memory handle from other thread owner to this thread. // Caution! Do not use this method directly, only through the method CDnnBlob::TransferDataToThisThread() virtual void TransferHandleToThisThread( const CMemoryHandle& handle, size_t size ) = 0; diff --git a/NeoMathEngine/src/CPU/CpuMathEngine.cpp b/NeoMathEngine/src/CPU/CpuMathEngine.cpp index 15933d1be..7ba2c7ca9 100644 --- a/NeoMathEngine/src/CPU/CpuMathEngine.cpp +++ b/NeoMathEngine/src/CPU/CpuMathEngine.cpp @@ -1,4 +1,4 @@ -/* Copyright © 2017-2023 ABBYY +/* Copyright © 2017-2024 ABBYY Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. @@ -97,6 +97,12 @@ void CCpuMathEngine::SetReuseMemoryMode( bool enable ) memoryPool->SetReuseMemoryMode( enable ); } +void CCpuMathEngine::SetThreadBufferMemoryThreshold( size_t threshold ) +{ + std::lock_guard lock( mutex ); + memoryPool->SetThreadBufferMemoryThreshold( threshold ); +} + CMemoryHandle CCpuMathEngine::HeapAlloc( size_t size ) { std::lock_guard lock( mutex ); diff --git a/NeoMathEngine/src/CPU/CpuMathEngine.h b/NeoMathEngine/src/CPU/CpuMathEngine.h index e1d5e69dc..3a084c5a1 100644 --- a/NeoMathEngine/src/CPU/CpuMathEngine.h +++ b/NeoMathEngine/src/CPU/CpuMathEngine.h @@ -1,4 +1,4 @@ -/* Copyright © 2017-2023 ABBYY +/* Copyright © 2017-2024 ABBYY Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. @@ -45,6 +45,7 @@ class CCpuMathEngine : public IMathEngine, public IRawMemoryManager { // IMathEngine interface methods TMathEngineType GetType() const override { return MET_Cpu; } void SetReuseMemoryMode( bool enabled ) override; + void SetThreadBufferMemoryThreshold( size_t threshold ) override; CMemoryHandle HeapAlloc( size_t count ) override; void HeapFree( const CMemoryHandle& handle ) override; void TransferHandleToThisThread( const CMemoryHandle& handle, size_t size ) override; diff --git a/NeoMathEngine/src/GPU/CUDA/CudaMathEngine.cpp b/NeoMathEngine/src/GPU/CUDA/CudaMathEngine.cpp index 5cbb85bb0..0764959c8 100644 --- a/NeoMathEngine/src/GPU/CUDA/CudaMathEngine.cpp +++ b/NeoMathEngine/src/GPU/CUDA/CudaMathEngine.cpp @@ -1,4 +1,4 @@ -/* Copyright © 2017-2023 ABBYY +/* Copyright © 2017-2024 ABBYY Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. @@ -74,6 +74,12 @@ void CCudaMathEngine::SetReuseMemoryMode( bool ) // Always true, because allocation is sync } +void CCudaMathEngine::SetThreadBufferMemoryThreshold( size_t threshold ) +{ + std::lock_guard lock( mutex ); + memoryPool->SetThreadBufferMemoryThreshold( threshold ); +} + CMemoryHandle CCudaMathEngine::HeapAlloc( size_t size ) { std::lock_guard lock( mutex ); diff --git a/NeoMathEngine/src/GPU/CUDA/CudaMathEngine.h b/NeoMathEngine/src/GPU/CUDA/CudaMathEngine.h index 6793e74c3..18625bb06 100644 --- a/NeoMathEngine/src/GPU/CUDA/CudaMathEngine.h +++ b/NeoMathEngine/src/GPU/CUDA/CudaMathEngine.h @@ -1,4 +1,4 @@ -/* Copyright © 2017-2023 ABBYY +/* Copyright © 2017-2024 ABBYY Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. @@ -52,6 +52,7 @@ class CCudaMathEngine : public IMathEngine, public IRawMemoryManager { TMathEngineType GetType() const override { return MET_Cuda; } void GetMathEngineInfo( CMathEngineInfo& info ) const override; void SetReuseMemoryMode( bool enable ) override; + void SetThreadBufferMemoryThreshold( size_t threshold ) override; CMemoryHandle HeapAlloc( size_t count ) override; void HeapFree( const CMemoryHandle& handle ) override; void TransferHandleToThisThread( const CMemoryHandle& handle, size_t size ) override; diff --git a/NeoMathEngine/src/GPU/Metal/MetalMathEngine.h b/NeoMathEngine/src/GPU/Metal/MetalMathEngine.h index 9df1789b6..77317372a 100644 --- a/NeoMathEngine/src/GPU/Metal/MetalMathEngine.h +++ b/NeoMathEngine/src/GPU/Metal/MetalMathEngine.h @@ -1,4 +1,4 @@ -/* Copyright © 2017-2023 ABBYY +/* Copyright © 2017-2024 ABBYY Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. @@ -47,6 +47,7 @@ class CMetalMathEngine : public IMathEngine, public IRawMemoryManager { // IMathEngine interface methods TMathEngineType GetType() const override { return MET_Metal; } void SetReuseMemoryMode( bool enable ) override; + void SetThreadBufferMemoryThreshold( size_t threshold ) override; CMemoryHandle HeapAlloc( size_t count ) override; void HeapFree( const CMemoryHandle& handle ) override; void TransferHandleToThisThread( const CMemoryHandle& /*handle*/, size_t /*size*/ ) override { ASSERT_EXPR( false ); } diff --git a/NeoMathEngine/src/GPU/Metal/MetalMathEngine.mm b/NeoMathEngine/src/GPU/Metal/MetalMathEngine.mm index 60ad0b984..3ccad4ac9 100644 --- a/NeoMathEngine/src/GPU/Metal/MetalMathEngine.mm +++ b/NeoMathEngine/src/GPU/Metal/MetalMathEngine.mm @@ -1,4 +1,4 @@ -/* Copyright © 2017-2020 ABBYY Production LLC +/* Copyright © 2017-2024 ABBYY Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. @@ -91,6 +91,12 @@ bool LoadMetalEngineInfo( CMathEngineInfo& info ) memoryPool->SetReuseMemoryMode( enable ); } +void CVulkanMathEngine::SetThreadBufferMemoryThreshold( size_t threshold ) +{ + std::lock_guard lock( *mutex ); + memoryPool->SetThreadBufferMemoryThreshold( threshold ); +} + CMemoryHandle CMetalMathEngine::HeapAlloc( size_t size ) { std::lock_guard lock( *mutex ); diff --git a/NeoMathEngine/src/GPU/Vulkan/VulkanMathEngine.cpp b/NeoMathEngine/src/GPU/Vulkan/VulkanMathEngine.cpp index 2ac3d6d52..05a9e95db 100644 --- a/NeoMathEngine/src/GPU/Vulkan/VulkanMathEngine.cpp +++ b/NeoMathEngine/src/GPU/Vulkan/VulkanMathEngine.cpp @@ -1,4 +1,4 @@ -/* Copyright © 2017-2020 ABBYY Production LLC +/* Copyright © 2017-2024 ABBYY Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. @@ -97,6 +97,12 @@ void CVulkanMathEngine::SetReuseMemoryMode( bool enable ) memoryPool->SetReuseMemoryMode( enable ); } +void CVulkanMathEngine::SetThreadBufferMemoryThreshold( size_t threshold ) +{ + std::lock_guard lock( mutex ); + memoryPool->SetThreadBufferMemoryThreshold( threshold ); +} + CMemoryHandle CVulkanMathEngine::HeapAlloc( size_t size ) { std::lock_guard lock( mutex ); diff --git a/NeoMathEngine/src/GPU/Vulkan/VulkanMathEngine.h b/NeoMathEngine/src/GPU/Vulkan/VulkanMathEngine.h index 4c7a2d06a..a4e4c47b5 100644 --- a/NeoMathEngine/src/GPU/Vulkan/VulkanMathEngine.h +++ b/NeoMathEngine/src/GPU/Vulkan/VulkanMathEngine.h @@ -1,4 +1,4 @@ -/* Copyright © 2017-2023 ABBYY +/* Copyright © 2017-2024 ABBYY Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. @@ -59,6 +59,7 @@ class CVulkanMathEngine : public IMathEngine, public IRawMemoryManager { // IMathEngine interface methods TMathEngineType GetType() const override { return MET_Vulkan; } void SetReuseMemoryMode( bool enable ) override; + void SetThreadBufferMemoryThreshold( size_t threshold ) override; CMemoryHandle HeapAlloc( size_t count ) override; void HeapFree( const CMemoryHandle& handle ) override; void TransferHandleToThisThread( const CMemoryHandle& /*handle*/, size_t /*size*/ ) override { ASSERT_EXPR( false ); } diff --git a/NeoMathEngine/src/MemoryPool.cpp b/NeoMathEngine/src/MemoryPool.cpp index 97f113bb2..47c207318 100644 --- a/NeoMathEngine/src/MemoryPool.cpp +++ b/NeoMathEngine/src/MemoryPool.cpp @@ -1,4 +1,4 @@ -/* Copyright © 2017-2023 ABBYY +/* Copyright © 2017-2024 ABBYY Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. @@ -90,6 +90,10 @@ static const unsigned int BufferSizes[] = { template inline constexpr int lengthof( T(&)[size] ) { return size; } +const size_t CMemoryPool::CThreadData::DefaultBufferMemoryThreshold = BufferSizes[lengthof( BufferSizes ) - 1]; + +//------------------------------------------------------------------------------------------------------------ + CMemoryPool::CMemoryPool( size_t _memoryLimit, IRawMemoryManager* _rawMemoryManager, bool reuseMemoryMode ) : memoryLimit( _memoryLimit ), rawMemoryManager( _rawMemoryManager ), @@ -116,6 +120,12 @@ void CMemoryPool::SetReuseMemoryMode( bool enable ) getThreadData( id )->Enabled = enable; } +void CMemoryPool::SetThreadBufferMemoryThreshold( size_t threshold ) +{ + const std::thread::id id = std::this_thread::get_id(); + getThreadData( id )->BufferMemoryThreshold = threshold; +} + CMemoryHandle CMemoryPool::Alloc( size_t size ) { const std::thread::id id = std::this_thread::get_id(); @@ -197,7 +207,7 @@ void CMemoryPool::TransferHandleToThisThread( const CMemoryHandle& handle, size_ info.buffer->OwnerPool = thisThreadBufferPool; } } else { // Large buffers don't use the pools - const size_t validSize = *std::lower_bound( std::begin(BufferSizes), std::end( BufferSizes ), size ); + const size_t validSize = *std::lower_bound( std::begin( BufferSizes ), std::end( BufferSizes ), size ); ASSERT_EXPR( size == info.size || validSize == info.size ); // No need to transfer, because // it wouldn't be cleaned-up for that thread after mathEngine.CleanUp(). @@ -253,7 +263,7 @@ inline static bool poolsCompare( const CMemoryBufferPool* a, const size_t& b ) // Tries to allocate memory CMemoryHandle CMemoryPool::tryAlloc( size_t size, CThreadData& data ) { - if( !data.Enabled || size > BufferSizes[lengthof(BufferSizes) - 1] ) { + if( !data.Enabled || size > data.BufferMemoryThreshold ) { // Allocate without using the buffers pool CMemoryHandle result = alloc( size ); if( !result.IsNull() ) { diff --git a/NeoMathEngine/src/MemoryPool.h b/NeoMathEngine/src/MemoryPool.h index b9f6bb9a0..2cd53961e 100644 --- a/NeoMathEngine/src/MemoryPool.h +++ b/NeoMathEngine/src/MemoryPool.h @@ -1,4 +1,4 @@ -/* Copyright © 2017-2023 ABBYY +/* Copyright © 2017-2024 ABBYY Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. @@ -35,6 +35,8 @@ class CMemoryPool : public CCrtAllocatedObject { // Turns on and off the memory reuse mode for the current thread void SetReuseMemoryMode( bool enable ); + // Change the memory blocks' sizes threshold for this thread from 1GB to the user size in bytes + void SetThreadBufferMemoryThreshold( size_t threshold ); // Allocates the specified amount of memory CMemoryHandle Alloc( size_t size ); @@ -63,8 +65,10 @@ class CMemoryPool : public CCrtAllocatedObject { using TMemoryBufferPoolVector = std::vector>; // The information about all of memory buffers pools of unused non-cleared blocks struct CThreadData final { + static const size_t DefaultBufferMemoryThreshold; TMemoryBufferPoolVector Pool; - bool Enabled{}; + bool Enabled = false; // default 'reuse' mode is disabled + size_t BufferMemoryThreshold = DefaultBufferMemoryThreshold; // default max = 1GB }; using TThreadDataMap = std::unordered_map< std::thread::id, CThreadData, // (key, value)