[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,10 +25,31 @@ 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 ) * 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 ) {}
+	// Any possible number of all mathEngines
+	static constexpr int MaxMathEngineEntities = 1024;
+
+	CMemoryHandle() : object( nullptr ), offset( 0 ), entity( MathEngineEntityInvalid ) {}
 	CMemoryHandle( CMemoryHandle&& other ) = default;
 	CMemoryHandle( const CMemoryHandle& other ) = default;
 
@@ -36,23 +58,37 @@ class NEOMATHENGINE_API CMemoryHandle {
 
 	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 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;
 
-	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( mathEngine, object, offset + shift ); }
+
+	CMemoryHandle CopyMemoryHandle( 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 )
+	{
+#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-2024 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

@@ -1233,7 +1233,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 ) :
+	IMemoryEngine( 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 ) :
+	IMemoryEngine( 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 ) :
+	IMemoryEngine( 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 ) :
+	IMemoryEngine( 0 ),
 	dllLoader( CDllLoader::VULKAN_DLL ),
 	device( std::move( _device ) ),
 	tmpImages( TVI_Count, nullptr )

NeoMathEngine/src/MemoryEngine.cpp

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

NeoMathEngine/src/MemoryEngine.h

@@ -59,6 +59,28 @@ class IMemoryEngine : 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 IMemoryEngine( 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