[NeoMathEngine] The threshold in bytes for buffer pools allocation fo…

…r the thread (#1061)

Signed-off-by: Kirill Golikov <kirill.golikov@abbyy.com>

NeoML/test/src/DnnBlobTest.cpp

@@ -28,16 +28,16 @@ TEST( CDnnBlobTest, InitWindowBlob )
     CPtr<CDnnBlob> parent = CDnnBlob::CreateDataBlob( MathEngine(), CT_Float, 16, 1, 1 );
     CPtr<CDnnBlob> blob = CDnnBlob::CreateWindowBlob( parent );
 
-    ASSERT_FALSE( blob->GetData().IsNull() );
+    EXPECT_FALSE( blob->GetData().IsNull() );
 }
 
 TEST( CDnnBlobTest, BufferTest )
 {
     CPtr<CDnnBlob> blob = CDnnBlob::CreateDataBlob( MathEngine(), CT_Float, 16, 1, 1 );
-    ASSERT_FALSE( blob->GetData().IsNull() );
+    EXPECT_FALSE( blob->GetData().IsNull() );
 
     CDnnBlobBuffer<float> 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<CDnnBlob> blob1 = CDnnBlob::CreateDataBlob( MathEngine(), CT_Float, int(threshold / sizeof( float )), 1, 1 );
+            ASSERT_TRUE( blob1 != nullptr && !blob1->GetData().IsNull() );
+            CPtr<CDnnBlob> 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 {

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;

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<std::mutex> lock( mutex );
+	memoryPool->SetThreadBufferMemoryThreshold( threshold );
+}
+
 CMemoryHandle CCpuMathEngine::HeapAlloc( size_t size )
 {
 	std::lock_guard<std::mutex> lock( mutex );

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;

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<std::mutex> lock( mutex );
+	memoryPool->SetThreadBufferMemoryThreshold( threshold );
+}
+
 CMemoryHandle CCudaMathEngine::HeapAlloc( size_t size )
 {
 	std::lock_guard<std::mutex> lock( mutex );

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;

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 ); }

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<std::mutex> lock( *mutex );
+	memoryPool->SetThreadBufferMemoryThreshold( threshold );
+}
+
 CMemoryHandle CMetalMathEngine::HeapAlloc( size_t size )
 {
 	std::lock_guard<CMutex> lock( *mutex );

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<std::mutex> lock( mutex );
+	memoryPool->SetThreadBufferMemoryThreshold( threshold );
+}
+
 CMemoryHandle CVulkanMathEngine::HeapAlloc( size_t size )
 {
 	std::lock_guard<std::mutex> lock( mutex );

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 ); }

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 <typename T, int size>
 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() ) {

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<CMemoryBufferPool*, CrtAllocator<CMemoryBufferPool*>>;
 	// 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)