[NeoMathEngine] CMemoryHandle less memory consumption

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

Author: favorart

NeoMathEngine/include/NeoMathEngine/MemoryHandle.h

@@ -16,6 +16,7 @@ limitations under the License.
 #pragma once
 
 #include <NeoMathEngine/NeoMathEngineDefs.h>
+#include <NeoMathEngine/NeoMathEngineException.h>
 #include <cstddef>
 #include <type_traits>
 
@@ -24,31 +25,66 @@ namespace NeoML {
 class IMathEngine;
 class CMemoryHandleInternal;
 
+// Get pointer to IMathEngine by the given current entity
+NEOMATHENGINE_API IMathEngine* GetMathEngineByIndex( size_t currentEntity );
+// Get current entity from the given pointer to IMathEngine
+NEOMATHENGINE_API 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 ) * CHAR_BIT ) - mathEngineCountWidth;
+	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 ) * CHAR_BIT;
+	static constexpr int mathEngineCountShift = sizeof( size_t ) * CHAR_BIT;
+#endif // FINE_32_BIT
+
 public:
-	constexpr CMemoryHandle() = default;
+	// Any possible number of all mathEngines
+	static constexpr int MaxMathEngineEntities = 1024;
+	static constexpr size_t MathEngineEntityInvalid = size_t( -1 );
+
+	CMemoryHandle() : CMemoryHandle( nullptr, 0, MathEngineEntityInvalid ) {}
 	// Be copied and moved by default
-	
+
 	bool operator!=( const CMemoryHandle& other ) const { return !operator==( other ); }
 	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 IsNull() const { return *this == CMemoryHandle{}; }
 
-	IMathEngine* GetMathEngine() const { return MathEngine; }
+	IMathEngine* GetMathEngine() const { return GetMathEngineByIndex( Entity ); }
 
 protected:
-	IMathEngine* MathEngine = nullptr; // the math engine owner
+	// struct of (16 bytes size for x64 and arm-x64) and (12 bytes size for x86 and arm-x32)
 	const void* Object = nullptr; // the memory allocated base pointer
-	std::ptrdiff_t Offset = 0; // the offset in the memory allocated volume, in bytes
+	// The offset in the memory allocated volume, in bytes
+	size_t Offset : mathEngineCountShift; // (x64) the less significant bits of size_t stores offset in the base object, in bytes
+	// The math engine owner
+	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 ) {}
+		CMemoryHandle( object , offset, GetIndexOfMathEngine( mathEngine ) ) {}
 
-	CMemoryHandle Copy( ptrdiff_t shift ) const { return CMemoryHandle( MathEngine, Object, Offset + shift ); }
+	CMemoryHandle Copy( ptrdiff_t shift ) const { return CMemoryHandle( Object, Offset + shift, Entity ); }
+
+private:
+	explicit CMemoryHandle( const void* object, ptrdiff_t offset, size_t entity ) :
+		Object( object ), Offset( offset ), Entity( entity & ( MaxMathEngineEntities - 1 ) )
+	{
+#if FINE_PLATFORM( FINE_64_BIT )
+		static_assert( MaxMathEngineEntities == ( 1 << mathEngineCountWidth ), "Invalid max MathEngine entities" );
+		// 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 == MathEngineEntityInvalid || entity < ( MaxMathEngineEntities - 1/*Invalid*/ ) );
+	}
 };
 
 //---------------------------------------------------------------------------------------------------------------------

NeoMathEngine/include/NeoMathEngine/MemoryHandle.inl

@@ -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

@@ -1234,7 +1234,7 @@ 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;
 
-	// For Distributed only
+	// 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;

NeoMathEngine/src/CPU/CpuMathEngine.cpp

@@ -51,6 +51,7 @@ int NEOMATHENGINE_API FloatAlignment = CCPUInfo::DefineFloatAlignment();
 CCpuMathEngine::CCpuMathEngine( size_t _memoryLimit,
 		std::shared_ptr<CMultiThreadDistributedCommunicator> communicator,
 		const CMathEngineDistributedInfo& distributedInfo ) :
+	CMemoryEngineMixin( 0 ),
 	floatAlignment( FloatAlignment ),
 	communicator( communicator ),
 	distributedInfo( distributedInfo ),

NeoMathEngine/src/GPU/CUDA/CudaMathEngine.cu

@@ -40,6 +40,7 @@ const int CudaMemoryAlignment = 4;
 
 CCudaMathEngine::CCudaMathEngine( const CCusparse* _cusparse, const CCublas* _cublas,
 		std::unique_ptr<CCudaDevice>& _device, int flags ) :
+	CMemoryEngineMixin( 0 ),
 	loader( CDllLoader::CUDA_DLL ),
 	cusparse( _cusparse ),
 	cublas( _cublas ),

NeoMathEngine/src/GPU/Metal/MetalMathEngine.mm

@@ -66,6 +66,7 @@ bool LoadMetalEngineInfo( CMathEngineInfo& info )
 const int MetalMemoryAlignment = 16;
 
 CMetalMathEngine::CMetalMathEngine( size_t memoryLimit ) :
+	CMemoryEngineMixin( 0 ),
 	queue( new CMetalCommandQueue() )
 {
 	ASSERT_EXPR( queue->Create() );

NeoMathEngine/src/GPU/Vulkan/VulkanMathEngine.cpp

@@ -69,6 +69,7 @@ bool LoadVulkanEngineInfo( const CVulkanDll& dll, std::vector< CMathEngineInfo,
 constexpr int VulkanMemoryAlignment = 16;
 
 CVulkanMathEngine::CVulkanMathEngine( std::unique_ptr<const CVulkanDevice>& _device, size_t memoryLimit ) :
+	CMemoryEngineMixin( 0 ),
 	dllLoader( CDllLoader::VULKAN_DLL ),
 	device( std::move( _device ) ),
 	tmpImages( TVI_Count, nullptr )

NeoMathEngine/src/MemoryEngineMixin.cpp

@@ -24,6 +24,25 @@ limitations under the License.
 
 namespace NeoML {
 
+size_t CMemoryEngineMixin::MathEngineEntitiesNumerator = 0;
+IMathEngine* CMemoryEngineMixin::MathEngineEntitiesArray[CMemoryHandle::MaxMathEngineEntities]{};
+
+// Get pointer to IMathEngine by the given current entity
+IMathEngine* GetMathEngineByIndex( size_t currentEntity )
+{
+	return CMemoryEngineMixin::MathEngineEntitiesArray[currentEntity];
+}
+
+// Get current entity from the given pointer to IMathEngine
+size_t GetIndexOfMathEngine( const IMathEngine* mathEngine )
+{
+	return ( mathEngine == nullptr )
+		? CMemoryHandle::MathEngineEntityInvalid
+		: static_cast<const CMemoryEngineMixin*>( mathEngine )->CurrentEntity;
+}
+
+//---------------------------------------------------------------------------------------------------------------------
+
 void CMemoryEngineMixin::InitializeMemory( IRawMemoryManager* _rawManager, size_t _memoryLimit, int _memoryAlignment,
 		bool _reuse, bool _hostStack )
 {

NeoMathEngine/src/MemoryEngineMixin.h

@@ -59,6 +59,28 @@ class CMemoryEngineMixin : public IMathEngine {
 	std::unique_ptr<IStackAllocator, CStackAllocatorDeleter> HostStackAllocator; // stack allocator for regular memory
 
 	void CleanUpSpecial() override {}
+
+	// All below is need to avoid excess (8 bytes) field in each CMemoryHandler
+
+	// Special constructor
+	explicit CMemoryEngineMixin( 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 );
+	friend size_t GetIndexOfMathEngine( const IMathEngine* mathEngine );
 };
 
 } // namespace NeoML