[NeoMathEngine] CMemoryHandle less memory consumption

Signed-off-by: Kirill Golikov <[email protected]>

Author: favorart

NeoMathEngine/include/NeoMathEngine/MemoryHandle.h

@@ -1,4 +1,4 @@
-/* Copyright © 2017-2020 ABBYY Production LLC
+/* Copyright © 2017-2023 ABBYY
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -16,6 +16,7 @@ limitations under the License.
 #pragma once
 
 #include <NeoMathEngine/NeoMathEngineDefs.h>
+#include <NeoMathEngine/NeoMathEngineException.h>
 #include <cstddef>
 #include <type_traits>
 
@@ -24,47 +25,77 @@ namespace NeoML {
 class IMathEngine;
 class CMemoryHandleInternal;
 
+// Get pointer to IMathEngine by the given CurrentEntity
+inline IMathEngine* GetMathEngineByIndex( size_t currentEntity );
+// Get current entity from the given pointer to IMathEngine
+inline size_t GetIndexOfMathEngine( const IMathEngine* mathEngine );
+
 // Wraps the pointer to memory allocated by a math engine
 class NEOMATHENGINE_API CMemoryHandle {
+private:
+#if FINE_PLATFORM( FINE_64_BIT )
+	static constexpr int MathEngineCountWidth = 10; // compress to bitfield
+	static constexpr int MathEngineCountShift = ( sizeof( size_t ) * 8 /*bits*/ ) - MathEngineCountWidth;
+	static constexpr size_t MathEngineEntityInvalid = ( ( size_t( 1 ) << MathEngineCountWidth ) - 1 );
+	static constexpr size_t MathEngineMaxOffset = size_t( 1 ) << MathEngineCountShift;
+#else  // FINE_32_BIT
+	// only for bitfield compiles correct. no compress
+	static constexpr int MathEngineCountWidth = sizeof( size_t ) * 8;
+	static constexpr int MathEngineCountShift = sizeof( size_t ) * 8;
+	static constexpr size_t MathEngineEntityInvalid = size_t( -1 );
+#endif // FINE_32_BIT
+
 public:
-	CMemoryHandle() : mathEngine( 0 ), object( 0 ), offset( 0 ) {}
-	CMemoryHandle( const CMemoryHandle& other ) : mathEngine( other.mathEngine ), object( other.object ), offset( other.offset ) {}
+	// Any possible number of all mathEngines
+	static constexpr int MaxMathEngineEntities = 1024;
+
+	CMemoryHandle() : object( nullptr ), offset( 0 ), entity( MathEngineEntityInvalid ) {}
+	CMemoryHandle( const CMemoryHandle& other ) : object( other.object ), offset( other.offset ), entity( other.entity ) {}
 
 	CMemoryHandle& operator=( const CMemoryHandle& other )
 	{
-		mathEngine = other.mathEngine;
-		object = other.object;
-		offset = other.offset;
+		if( this != &other ) {
+			object = other.object;
+			offset = other.offset;
+			entity = other.entity;
+		}
 		return *this;
 	}
 
 	bool operator==( const CMemoryHandle& other ) const
-	{
-		return mathEngine == other.mathEngine && object == other.object && offset == other.offset;
-	}
+	{ return object == other.object && offset == other.offset && entity == other.entity; }
 
-	bool operator!=( const CMemoryHandle& other ) const
-	{
-		return !operator==( other );
-	}
+	bool operator!=( const CMemoryHandle& other ) const { return !operator==( other ); }
 
-	bool IsNull() const
-	{
-		return mathEngine == 0 && object == 0 && offset == 0;
-	}
+	bool IsNull() const { return object == nullptr && offset == 0 && entity == MathEngineEntityInvalid; }
 
-	IMathEngine* GetMathEngine() const { return mathEngine; }
+	IMathEngine* GetMathEngine() const { return GetMathEngineByIndex( entity ); }
 
 protected:
-	IMathEngine* mathEngine; // the math engine
+	// struct of (16 bytes size for x64 and arm-x64) and (12 bytes size for x86 and arm-x32)
 	const void* object; // the base object
-	std::ptrdiff_t offset; // the offset in the base object, in bytes
+	size_t offset : MathEngineCountShift; // (x64) the less significant bits of size_t stores offset in the base object, in bytes
+	size_t entity : MathEngineCountWidth; // (x64) the most significant bits of size_t stores the number of IMathEngine entity
 
 	friend class CMemoryHandleInternal;
 
-	explicit CMemoryHandle( IMathEngine* _mathEngine, const void* _object, ptrdiff_t _offset ) : mathEngine( _mathEngine ), object( _object ), offset( _offset ) {}
+	explicit CMemoryHandle( IMathEngine* _mathEngine, const void* _object, ptrdiff_t _offset ) :
+		CMemoryHandle( _object , _offset, GetIndexOfMathEngine( _mathEngine ) )
+	{}
+
+	CMemoryHandle CopyMemoryHandle( ptrdiff_t shift ) const { return CMemoryHandle( object, offset + shift, entity ); }
 
-	CMemoryHandle CopyMemoryHandle( ptrdiff_t shift ) const { return CMemoryHandle( mathEngine, object, offset + shift ); }
+private:
+	explicit CMemoryHandle( const void* _object, ptrdiff_t _offset, size_t _entity ) :
+		object( _object ), offset( _offset ), entity( _entity )
+	{
+#if FINE_PLATFORM( FINE_64_BIT )
+		static_assert( MaxMathEngineEntities == ( 1 << MathEngineCountWidth ), "Invalid MaxMathEngineEntities" );
+		// Checks that the most significant bits do not interfere the result
+		ASSERT_EXPR( 0 <= _offset && size_t( _offset ) < MathEngineMaxOffset );
+#endif // FINE_64_BIT
+		ASSERT_EXPR( _entity < MaxMathEngineEntities );
+	}
 };
 
 //------------------------------------------------------------------------------------------------------------

NeoMathEngine/include/NeoMathEngine/MemoryHandle.inl

@@ -1,4 +1,4 @@
-/* Copyright © 2017-2020 ABBYY Production LLC
+/* Copyright © 2017-2023 ABBYY
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -23,30 +23,30 @@ template<class T>
 inline void CTypedMemoryHandle<T>::SetValueAt( int index, T value ) const
 {
 	CTypedMemoryHandle<T> result = *this + index;
-	mathEngine->DataExchangeRaw( result, &value, sizeof( T ) );
+	GetMathEngine()->DataExchangeRaw( result, &value, sizeof( T ) );
 }
 
 template<class T>
 inline T CTypedMemoryHandle<T>::GetValueAt( int index ) const
 {
 	char result[sizeof(T)];
 	CTypedMemoryHandle<T> source = *this + index;
-	mathEngine->DataExchangeRaw( result, source, sizeof( T ) );
+	GetMathEngine()->DataExchangeRaw( result, source, sizeof( T ) );
 	T* value = reinterpret_cast<T*>( &result );
 	return *value;
 }
 
 template<class T>
 inline void CTypedMemoryHandle<T>::SetValue( T value ) const
 {
-	mathEngine->DataExchangeRaw( *this, &value, sizeof( T ) );
+	GetMathEngine()->DataExchangeRaw( *this, &value, sizeof( T ) );
 }
 
 template<class T>
 inline T CTypedMemoryHandle<T>::GetValue() const
 {
 	char result[sizeof(T)];
-	mathEngine->DataExchangeRaw( result, *this, sizeof( T ) );
+	GetMathEngine()->DataExchangeRaw( result, *this, sizeof( T ) );
 	T* value = reinterpret_cast<T*>( &result );
 	return *value;
 }

NeoMathEngine/include/NeoMathEngine/NeoMathEngine.h

@@ -1230,15 +1230,46 @@ class NEOMATHENGINE_API IMathEngine : public IDnnEngine {
 	// This object should be destroyed using the standard delete operator after use.
 	virtual IPerformanceCounters* CreatePerformanceCounters( bool isTimeOnly = false ) const = 0;
 
+	// Methods group for the DnnDistributed execution only
 	virtual CMathEngineDistributedInfo GetDistributedInfo() { return CMathEngineDistributedInfo(); }
 	virtual void AllReduce( const CFloatHandle& handle, int size ) = 0;
 	virtual void Broadcast( const CFloatHandle& handle, int size, int root ) = 0;
 	virtual void AbortDistributed() {};
 	virtual bool IsDistributed() { return false; }
+
+	// Get CurrentEntity for IMathEngine
+	size_t GetCurrentEntity() const { return CurrentEntity; }
+
+protected:
+	// All below is need to avoid excess (8 bytes) field in each CMemoryHandler
+
+	// Special constructor
+	explicit IMathEngine( int/*cannot be no call*/ ) :
+		CurrentEntity( ++MathEngineEntitiesNumerator )
+	{
+		ASSERT_EXPR( CurrentEntity < CMemoryHandle::MaxMathEngineEntities );
+		MathEngineEntitiesArray[CurrentEntity] = this;
+	}
+
+	// Generation for indices of all IMathEngine entities,
+	// Incremets evey moment new MathEngine created to generate its CurrentEntity value.
+	static size_t MathEngineEntitiesNumerator;
+	// Array for pointers to all entities of IMathEngine
+	// No cache ping-pong, because pointers are created once and never changes
+	static IMathEngine* MathEngineEntitiesArray[CMemoryHandle::MaxMathEngineEntities];
+	// Index of the current IMathEngine entity
+	const size_t CurrentEntity;
+
+	friend IMathEngine* GetMathEngineByIndex( size_t currentEntity );
 };
 
 //------------------------------------------------------------------------------------------------------------
 
+// Get pointer to IMathEngine by the given CurrentEntity
+inline IMathEngine* GetMathEngineByIndex( size_t currentEntity ) { return IMathEngine::MathEngineEntitiesArray[currentEntity]; }
+// Get current entity from the given pointer to IMathEngine
+inline size_t GetIndexOfMathEngine( const IMathEngine* mathEngine ) { return mathEngine->GetCurrentEntity(); }
+
 // Creates a math engine that uses a CPU for calculations.
 // You should call SetMathEngineExceptionHandler() before this call.
 // the default value is 0, which means as many memory as the system has.

NeoMathEngine/src/CPU/CpuMathEngine.cpp

@@ -53,6 +53,7 @@ int NEOMATHENGINE_API FloatAlignment = CCPUInfo::DefineFloatAlignment();
 CCpuMathEngine::CCpuMathEngine( size_t _memoryLimit,
 		std::shared_ptr<CMultiThreadDistributedCommunicator> communicator,
 		const CMathEngineDistributedInfo& distributedInfo ) :
+	IMathEngine( 0 ),
 	floatAlignment( FloatAlignment ),
 	memoryAlignment( floatAlignment * sizeof(float) ),
 	communicator( communicator ),

NeoMathEngine/src/GPU/CUDA/CudaMathEngine.cu

@@ -1,4 +1,4 @@
-/* Copyright © 2017-2020 ABBYY Production LLC
+/* Copyright © 2017-2023 ABBYY
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -41,6 +41,7 @@ const int CudaMemoryAlignment = 4;
 //------------------------------------------------------------------------------------------------------------
 
 CCudaMathEngine::CCudaMathEngine( const CCusparse* _cusparse, const CCublas* _cublas, std::unique_ptr<CCudaDevice>& _device, int flags ) :
+	IMathEngine( 0 ),
 	loader( CDllLoader::CUDA_DLL ),
 	cusparse( _cusparse ),
 	cublas( _cublas ),

NeoMathEngine/src/GPU/Metal/MetalMathEngine.mm

@@ -1,4 +1,4 @@
-/* Copyright © 2017-2020 ABBYY Production LLC
+/* Copyright © 2017-2023 ABBYY
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -73,6 +73,7 @@ bool LoadMetalEngineInfo( CMathEngineInfo& info )
 const int MetalMemoryAlignment = 16;
 
 CMetalMathEngine::CMetalMathEngine( size_t memoryLimit ) :
+	IMathEngine( 0 ),
 	queue( new CMetalCommandQueue() ),
 	memoryPool( new CMemoryPool( MIN( memoryLimit == 0 ? SIZE_MAX : memoryLimit, defineMemoryLimit() ), this, false ) ),
 	deviceStackAllocator( new CDeviceStackAllocator( *memoryPool, MetalMemoryAlignment ) ),

NeoMathEngine/src/GPU/Vulkan/VulkanMathEngine.cpp

@@ -1,4 +1,4 @@
-/* Copyright © 2017-2020 ABBYY Production LLC
+/* Copyright © 2017-2023 ABBYY
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -71,6 +71,7 @@ bool LoadVulkanEngineInfo( const CVulkanDll& dll, std::vector< CMathEngineInfo,
 constexpr int VulkanMemoryAlignment = 16;
 
 CVulkanMathEngine::CVulkanMathEngine( std::unique_ptr<const CVulkanDevice>& _device, size_t memoryLimit ) :
+	IMathEngine( 0 ),
 	dllLoader( CDllLoader::VULKAN_DLL ),
 	device( std::move( _device ) ),
 	tmpImages( TVI_Count, nullptr )

NeoMathEngine/src/MathEngine.cpp

@@ -218,6 +218,9 @@ IMathEngineExceptionHandler* GetMathEngineExceptionHandler()
 
 //------------------------------------------------------------------------------------------------------------
 
+size_t IMathEngine::MathEngineEntitiesNumerator = 0;
+IMathEngine* IMathEngine::MathEngineEntitiesArray[CMemoryHandle::MaxMathEngineEntities]{};
+
 IMathEngine* CreateCpuMathEngine( size_t memoryLimit )
 {
 	return new CCpuMathEngine( memoryLimit );