[NeoML] CDnnBlob less memory consumption

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

Author: favorart

NeoML/Python/src/PyCustomLossLayer.cpp

@@ -1,4 +1,4 @@
-/* Copyright © 2017-2021 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.
@@ -60,16 +60,14 @@ static CArchive& operator >>( CArchive& archive, py::object& obj )
 
 //------------------------------------------------------------------------------------------------------------
 
-class CTempBlob : public CDnnBlob {
+// CPyTempDnnBlob does not own the handler data
+class CPyTempDnnBlob : public CDnnBlobView {
 public:
-	CTempBlob( IMathEngine& mathEngine, const CConstFloatHandle& data, const CBlobDesc& dataDesc );
+	CPyTempDnnBlob( IMathEngine& mathEngine, const CConstFloatHandle& data, const CBlobDesc& dataDesc ) :
+		CDnnBlobView( mathEngine, dataDesc, data )
+	{}
 };
 
-CTempBlob::CTempBlob( IMathEngine& mathEngine, const CConstFloatHandle& data, const CBlobDesc& dataDesc ) :
-	CDnnBlob( mathEngine, dataDesc, data, false )
-{
-}
-
 //------------------------------------------------------------------------------------------------------------
 
 class CPythonLossLayer : public CLossLayer {
@@ -89,11 +87,11 @@ class CPythonLossLayer : public CLossLayer {
 
 		CPtr<CPyMathEngineOwner> mathEngineOwner = new CPyMathEngineOwner( &MathEngine(), false );
 
-		CPtr<const CDnnBlob> dataBlob = new CTempBlob( mathEngineOwner->MathEngine(), data, inputBlobs[0]->GetDesc() );
+		CPtr<const CDnnBlob> dataBlob = new CPyTempDnnBlob( mathEngineOwner->MathEngine(), data, inputBlobs[0]->GetDesc() );
 		CPtr<const CDnnBlob> var = tape.Variable( *dataBlob.Ptr() );
 		CPyBlob dataPyBlob( *mathEngineOwner, const_cast<CDnnBlob*>(var.Ptr()) );
 
-		CPtr<CDnnBlob> labelBlob( new CTempBlob( mathEngineOwner->MathEngine(), label, inputBlobs[1]->GetDesc() ) );
+		CPtr<CDnnBlob> labelBlob( new CPyTempDnnBlob( mathEngineOwner->MathEngine(), label, inputBlobs[1]->GetDesc() ) );
 		CPyBlob labelPyBlob( *mathEngineOwner, labelBlob );
 
 		CPtr<CDnnBlob> value;
@@ -121,7 +119,7 @@ class CPythonLossLayer : public CLossLayer {
 		CheckArchitecture( false, GetName(), "The custom loss layer doesn't support int labels!" );
 	}
 
-	void Serialize( CArchive& archive )
+	void Serialize( CArchive& archive ) override
 	{
 		archive.SerializeVersion( PythonLossLayerVersion, 1 );
 		CLossLayer::Serialize( archive );

NeoML/Python/src/PyDnnBlob.cpp

@@ -1,4 +1,4 @@
-/* Copyright © 2017-2021 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.
@@ -115,7 +115,10 @@ static CBlobDesc createBlobDesc( TBlobType type, std::initializer_list<int> dime
 	return desc;
 }
 
-class CPyDnnBlob : public CDnnBlob {
+//------------------------------------------------------------------------------------------------------------
+
+// CPyDnnBlob does not own the handler data
+class CPyDnnBlob : public CDnnBlobView {
 public:
 	CPyDnnBlob( IMathEngine& mathEngine, TBlobType type, std::initializer_list<int> dimension, py::buffer_info&& _info );
 	virtual ~CPyDnnBlob();
@@ -128,7 +131,7 @@ class CPyDnnBlob : public CDnnBlob {
 };
 
 CPyDnnBlob::CPyDnnBlob( IMathEngine& mathEngine, TBlobType type, std::initializer_list<int> dimension, py::buffer_info&& _info ) :
-	CDnnBlob( mathEngine, createBlobDesc( type, dimension ), CPyMemoryHandle( &mathEngine, _info.ptr ), false ),
+	CDnnBlobView( mathEngine, createBlobDesc( type, dimension ), CPyMemoryHandle( &mathEngine, _info.ptr ) ),
 	info( new py::buffer_info( std::move( _info ) ) )
 {
 }

NeoML/include/NeoML/Dnn/DnnBlob.h

@@ -34,9 +34,9 @@ namespace NeoML {
 
 class NEOML_API CDnnBlob : public IObject {
 public:
-	explicit CDnnBlob( IMathEngine& mathEngine );
+	explicit CDnnBlob( IMathEngine& mathEngine ) : mathEngine( mathEngine ) {}
 
-	// Move other's Blob state to this Blob and transfer its data (if dataOwned) to this thread
+	// Move other's Blob state to this Blob and transfer its data to this thread
 	CDnnBlob( CDnnBlob&& other );
 	CDnnBlob& operator=( CDnnBlob&& other );
 
@@ -65,7 +65,7 @@ class NEOML_API CDnnBlob : public IObject {
 	static CDnnBlob* CreateBlob(IMathEngine& mathEngine, TBlobType type, const CBlobDesc& pattern);
 
 	// Checks if the dimensions of another blob are the same
-	bool HasEqualDimensions(const CDnnBlob* other) const;
+	bool HasEqualDimensions( const CDnnBlob* other ) const { return desc.HasEqualDimensions( other->desc ); }
 
 	// Gets the blob size along the specified dimension
 	int DimSize(int d) const { return desc.DimSize(d); }
@@ -143,7 +143,7 @@ class NEOML_API CDnnBlob : public IObject {
 	// Transfers CDnnBlob data from other thread owner to this thread.
 	// By default memory underneath each blob is associated with the thread on which its allocation has occurred.
 	// This method switches this association to the calling thread.
-	void TransferDataToThisThread();
+	virtual void TransferDataToThisThread();
 
 	// Elementwise adds a blob of the same dimensions
 	void Add(const CDnnBlob* other);
@@ -165,7 +165,7 @@ class NEOML_API CDnnBlob : public IObject {
 	// Changes the blob dimensions "names" without moving the data
 	// In effect, only the blob description is changed
 	// As the data is unaffected, the total blob size specified by the new descriptor should be the same
-	void ReinterpretDimensions( const CBlobDesc& newDesc );
+	virtual void ReinterpretDimensions( const CBlobDesc& newDesc );
 
 	// Merges blobs along the given dimension
 	static void MergeByDim( IMathEngine& mathEngine, TBlobDim d, const CObjectArray<CDnnBlob>& from, const CPtr<CDnnBlob>& to );
@@ -193,31 +193,31 @@ class NEOML_API CDnnBlob : public IObject {
 
 	// Gets the pointer to the MathEngine on which the blob was created
 	IMathEngine& GetMathEngine() const { return mathEngine; }
-
 	// Gets the blob descriptor
 	const CBlobDesc& GetDesc() const { return desc; }
 	// Gets the type of data in the blob
 	TBlobType GetDataType() const { return desc.GetDataType(); }
 
+	// All methods below are just the interface for the CDnnWindowBlob representation
+
 	// Gets the parent blob
-	CDnnBlob* GetParent() { return parent; }
-	const CDnnBlob* GetParent() const { return parent; }
+	virtual CDnnBlob* GetParent() { return nullptr; }
+	const CDnnBlob* GetParent() const { return const_cast<CDnnBlob*>( this )->GetParent(); }
 	// Gets the blob that owns the data (and has no parent)
-	CDnnBlob* GetOwner();
-	const CDnnBlob* GetOwner() const { return const_cast<CDnnBlob*>(this)->GetOwner(); }
-
+	virtual CDnnBlob* GetOwner() { return this; }
+	const CDnnBlob* GetOwner() const { return const_cast<CDnnBlob*>( this )->GetOwner(); }
 	// Gets the shift in data relative to the parent blob
 	// The position in the parent blob is calculated along the BatchLength dimension
 	// The position equal to N would correspond to a N*BatchWidth*ListSize*Height*Width*Depth*Channels shift in the one-dimensional array
-	int GetParentPos() const;
-	void SetParentPos( int pos );
-	void ShiftParentPos( int shift );
+	virtual int GetParentPos() const { return 0; }
+	virtual void SetParentPos( int /*pos*/ ) { NeoAssert( false ); }
+	virtual void ShiftParentPos( int /*shift*/ ) { NeoAssert( false ); }
 
 protected:
 	~CDnnBlob() override;
 
-	CDnnBlob( IMathEngine& _mathEngine, const CBlobDesc& _desc, CMemoryHandle _data, bool _dataOwned ) :
-		mathEngine( _mathEngine ), desc( _desc ), data( _data ), dataOwned( _dataOwned ), parentPos( 0 )
+	CDnnBlob( IMathEngine& _mathEngine, const CBlobDesc& _desc, CMemoryHandle _data ) :
+		mathEngine( _mathEngine ), desc( _desc ), data( _data )
 	{
 		NeoAssert( desc.GetDataType() != CT_Invalid );
 		NeoAssert( &mathEngine == data.GetMathEngine() );
@@ -227,20 +227,79 @@ class NEOML_API CDnnBlob : public IObject {
 	IMathEngine& mathEngine;
 	CBlobDesc desc;
 	CMemoryHandle data;
-	bool dataOwned;
-
-	CPtr<CDnnBlob> parent;	// parent blob
-	int parentPos;
 
 	void initializeBlob(TBlobType _type, int batchLength, int batchWidth, int listSize, int height, int width,
 		int depth, int channels);
 	void initializeTensor(TBlobType _type, std::initializer_list<int> dimensions);
-	void initializeWindow(const CPtr<CDnnBlob>& _parent, int windowSize);
 	void initializeByPattern(TBlobType type, const CBlobDesc& pattern);
 
 	friend class CDnnBlobClassRegistrar;
+	friend class CDnnWindowBlob;
+	friend class CDnnBlobView;
 };
 
+//---------------------------------------------------------------------------------------------------------------------
+
+// The kind of CDnnBlob does not own the data
+// CDnnBlobView does not clear data handler
+// Used in python wrappers
+class NEOML_API CDnnBlobView : public CDnnBlob {
+protected:
+	CDnnBlobView( IMathEngine& mathEngine ) : CDnnBlob( mathEngine ) {}
+	CDnnBlobView( IMathEngine& mathEngine, const CBlobDesc& desc, CMemoryHandle data ) :
+		CDnnBlob( mathEngine, desc, data )
+	{}
+
+	~CDnnBlobView() { if( !data.IsNull() ) { data = CMemoryHandle{}; } } // no need to free
+
+	// Prohibited methods
+	//void ReinterpretDimensions( const CBlobDesc& ) override { NeoAssert( false ); } // impossible !!! USED in Python !!!
+	//void Serialize( CArchive& ) override // !!! PICKLED in Python !!!
+	//{ NeoAssert( false ); } // a blob that links to another may not be serialized
+	void TransferDataToThisThread() override {} // !!! MOVED in Python !!!
+};
+
+//---------------------------------------------------------------------------------------------------------------------
+
+// The kind of CDnnBlob that is view of the some parent Blob as sequence (BatchLength > 1).
+// This CDnnWindowBlob do not owner of its memory, technical CDnnBlob representation.
+// This CDnnWindowBlob represents 1 element (BatchLength == 1) of the sequence for the recursive networks.
+class NEOML_API CDnnWindowBlob : public CDnnBlob {
+public:
+	// Creates a "window" blob to represent a subsequence of objects from the parent blob
+	static CDnnBlob* CreateWindowBlob( const CPtr<CDnnBlob>& parent, int windowSize = 1 );
+
+	// Prohibited methods
+	void ReinterpretDimensions( const CBlobDesc& ) override { NeoAssert( false ); } // impossible
+	void Serialize( CArchive& ) override { NeoAssert( false ); } // a blob that links to another may not be serialized
+	void TransferDataToThisThread() override { NeoAssert( false ); }
+
+	// Interface of communication
+	CDnnBlob* GetParent() override { return parent; }
+	CDnnBlob* GetOwner() override;
+	int GetParentPos() const override;
+	void SetParentPos( int pos ) override;
+	void ShiftParentPos( int shift ) override;
+
+protected:
+	CDnnWindowBlob( IMathEngine& mathEngine ) : CDnnBlob( mathEngine ) {}
+	CDnnWindowBlob( IMathEngine& mathEngine, const CBlobDesc& desc, CMemoryHandle data ) :
+		CDnnBlob( mathEngine, desc, data )
+	{}
+	CDnnWindowBlob( CDnnWindowBlob&& other ) = delete;
+	CDnnWindowBlob& operator=( CDnnWindowBlob&& other ) = delete;
+
+	~CDnnWindowBlob() { if( parent != nullptr ) { data = CMemoryHandle{}; } } // no need to free
+
+private:
+	CPtr<CDnnBlob> parent; // parent blob
+	int parentPos = 0;
+
+	void initializeWindow( const CPtr<CDnnBlob>& parent, int windowSize );
+};
+
+//---------------------------------------------------------------------------------------------------------------------
+
 inline void SerializeBlob( IMathEngine& mathEngine, CArchive& archive, CPtr<CDnnBlob>& blob )
 {
 	if( archive.IsStoring() ) {
@@ -286,6 +345,8 @@ enum class TDnnBlobBufferAccess {
 	ReadWrite
 };
 
+//---------------------------------------------------------------------------------------------------------------------
+
 // RAII-helper to safely work with `CDnnBlob::GetBuffer`/`CDnnBlob::ReleaseBuffer`.
 template<typename TBufferType = float>
 class CDnnBlobBuffer {
@@ -326,6 +387,8 @@ class CDnnBlobBuffer {
 	TBufferType* ptr;
 };
 
+//---------------------------------------------------------------------------------------------------------------------
+
 inline CDnnBlob* CDnnBlob::CreateBlob( IMathEngine& mathEngine, const CBlobDesc& pattern )
 {
 	return CreateBlob(mathEngine, CT_Float, pattern);
@@ -449,13 +512,15 @@ inline T* CDnnBlob::GetBuffer( int pos, int size, bool exchange )
 	return static_cast<T*>( mathEngine.GetBuffer( data, pos * dataSize, size * dataSize, exchange ) );
 }
 
-inline int CDnnBlob::GetParentPos() const
+//---------------------------------------------------------------------------------------------------------------------
+
+inline int CDnnWindowBlob::GetParentPos() const
 {
 	NeoAssert(parent != 0);
 	return parentPos;
 }
 
-inline void CDnnBlob::SetParentPos(int pos)
+inline void CDnnWindowBlob::SetParentPos(int pos)
 {
 	int arrayPos = pos * (desc.BlobSize() / desc.BatchLength());
 	NeoAssert(parent != 0);
@@ -473,25 +538,24 @@ inline void CDnnBlob::SetParentPos(int pos)
 	}
 }
 
-inline void CDnnBlob::ShiftParentPos(int shift)
+inline void CDnnWindowBlob::ShiftParentPos(int shift)
 {
 	SetParentPos(parentPos + shift);
 }
 
-inline bool CDnnBlob::HasEqualDimensions(const CDnnBlob* other) const
-{
-	return desc.HasEqualDimensions(other->desc);
-}
-
-inline CDnnBlob* CDnnBlob::GetOwner()
+inline CDnnBlob* CDnnWindowBlob::GetOwner()
 {
 	CDnnBlob* result = this;
-	while( result->parent != 0 ) {
-		result = result->parent;
+	CDnnWindowBlob* window = this;
+	while( window != nullptr && window->parent != nullptr ) {
+		result = window->parent;
+		window = dynamic_cast<CDnnWindowBlob*>( result );
 	}
 	return result;
 }
 
+//---------------------------------------------------------------------------------------------------------------------
+
 template<typename TBufferType>
 inline CDnnBlobBuffer<TBufferType>::~CDnnBlobBuffer()
 {

NeoML/src/Dnn/AutoDiff.cpp

@@ -1,4 +1,4 @@
-/* Copyright © 2017-2021 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.
@@ -21,14 +21,14 @@ limitations under the License.
 namespace NeoML {
 
 CTapeBlob::CTapeBlob( IGradientTape* _tape, const CDnnBlob& blob ) :
-	CDnnBlob( blob.GetMathEngine(), blob.GetDesc(), blob.GetMathEngine().HeapAlloc( blob.GetDataSize() * sizeof(float) ), true ),
+	CDnnBlob( blob.GetMathEngine(), blob.GetDesc(), blob.GetMathEngine().HeapAlloc( blob.GetDataSize() * sizeof(float) ) ),
 	tape( _tape )
 {
 	blob.GetMathEngine().VectorCopy( GetData(), blob.GetData(), blob.GetDataSize() );
 }
 
 CTapeBlob::CTapeBlob( IGradientTape* _tape, IMathEngine& mathEngine, const CBlobDesc& desc ) :
-	CDnnBlob( mathEngine, desc, mathEngine.HeapAlloc( desc.BlobSize() * sizeof(float) ), true ),
+	CDnnBlob( mathEngine, desc, mathEngine.HeapAlloc( desc.BlobSize() * sizeof(float) ) ),
 	tape( _tape )
 {
 }