initiate zero-copy output.

src/plugins/intel_cpu/src/cpu_memory.cpp

@@ -34,10 +34,26 @@ namespace {
     }
 }   // namespace
 
-Memory::Memory(const dnnl::engine& eng) :
-    eng(eng), mgrHandle(std::make_shared<DnnlMemoryMngr>(std::unique_ptr<MemoryMngrWithReuse>(new MemoryMngrWithReuse())), this), dnnlMemHandle(this) {}
-Memory::Memory(const dnnl::engine& eng, std::unique_ptr<IMemoryMngr> mngr) :
-    eng(eng), mgrHandle(std::make_shared<DnnlMemoryMngr>(std::move(mngr)), this), dnnlMemHandle(this) {}
+Memory::Memory(const dnnl::engine& eng, MemoryDescPtr desc, const void* data, bool pads_zeroing) :
+    pMemDesc(desc),
+    eng(eng),
+    mgrHandle(std::make_shared<DnnlMemoryMngr>(std::unique_ptr<MemoryMngrWithReuse>(new MemoryMngrWithReuse())), this),
+    dnnlMemHandle(this) {
+        Create(pMemDesc, data, pads_zeroing);
+    }
+
+Memory::Memory(const dnnl::engine& eng, const MemoryDesc& desc, const void* data, bool pads_zeroing) :
+    Memory::Memory(eng, desc.clone(), data, pads_zeroing) {}
+
+Memory::Memory(const dnnl::engine& eng, MemoryDescPtr desc, MemoryMngrPtr mngr) :
+    pMemDesc(desc), eng(eng), mgrHandle(mngr, this), dnnlMemHandle(this) {
+        bool memAllocated = mgrHandle->getRawPtr();
+
+        Create(desc, nullptr, !memAllocated);
+    }
+
+Memory::Memory(const dnnl::engine& eng, const MemoryDesc& desc, MemoryMngrPtr mngr) :
+    Memory::Memory(eng, desc.clone(), mngr) {}
 
 size_t Memory::GetSize() const {
     auto size = getDesc().getCurrentMemSize();
@@ -68,7 +84,7 @@ void Memory::Create(MemoryDescPtr desc, const void* data, bool pads_zeroing) {
     }
 }
 
-void Memory::SetData(const Memory& src, bool ftz) const {
+void Memory::SetData(const IMemory& src, bool ftz) const {
     node::Reorder::reorderData(src, *this);
 
     auto localPrim = GetPrimitive();
@@ -94,12 +110,6 @@ void Memory::FillZero() {
         memset(dataPtr, 0, getDesc().getCurrentMemSize());
 }
 
-void *Memory::GetPtr() const  {
-    auto ptr = static_cast<uint8_t*>(GetData());
-    ptr += pMemDesc->getOffsetPadding() * pMemDesc->getPrecision().size();
-    return ptr;
-}
-
 void Memory::redefineDesc(MemoryDescPtr desc) {
     if (!desc->hasDefinedMaxSize()) {
         IE_THROW() << "Can not reset descriptor, memory upper bound is unknown.";
@@ -109,23 +119,13 @@ void Memory::redefineDesc(MemoryDescPtr desc) {
 }
 
 template<>
-DnnlMemoryDescPtr Memory::GetDescWithType<DnnlMemoryDesc, 0, 0>() const {
-    return MemoryDescUtils::convertToDnnlMemoryDesc(pMemDesc);
+DnnlMemoryDescPtr IMemory::GetDescWithType<DnnlMemoryDesc, 0, 0>() const {
+    return MemoryDescUtils::convertToDnnlMemoryDesc(getDescPtr());
 }
 
-void Memory::setDataHandle(void *data) {
-    if (!mgrHandle->hasExtBuffer()) {
-        mgrHandle = DnnlMemMngrHandle(
-            std::make_shared<DnnlMemoryMngr>(std::unique_ptr<MemoryMngrWithReuse>(new MemoryMngrWithReuse())),
-            this);
-    }
-
-    size_t maxMemSize = pMemDesc->isDefined() ?  pMemDesc->getCurrentMemSize() : 0;
-    mgrHandle->setExtBuff(data, maxMemSize);
-    if (dnnlMemHandle.isInit()) {
-        auto prim = dnnlMemHandle.getPrim();
-        prim.set_data_handle(mgrHandle->getRawPtr()); // for pads zeroing, to preserve dnnl::memory::set_data_handle behaviour
-    }
+template<>
+BlockedMemoryDescPtr IMemory::GetDescWithType<BlockedMemoryDesc, 0, 0>() const {
+    return MemoryDescUtils::convertToBlockedMemoryDesc(getDescPtr());
 }
 
 void Memory::update() {
@@ -135,22 +135,6 @@ void Memory::update() {
     }
 }
 
-void Memory::Create(const MemoryDesc &desc, DnnlMemoryMngrPtr memMgr) {
-    Create(desc.clone(), memMgr);
-}
-
-void Memory::Create(MemoryDescPtr desc, DnnlMemoryMngrPtr memMgr) {
-    mgrHandle = DnnlMemMngrHandle(memMgr, this);
-    bool memAllocated = mgrHandle->getRawPtr();
-
-    Create(desc, nullptr, !memAllocated);
-}
-
-template<>
-BlockedMemoryDescPtr Memory::GetDescWithType<BlockedMemoryDesc, 0, 0>() const {
-    return MemoryDescUtils::convertToBlockedMemoryDesc(pMemDesc);
-}
-
 dnnl::memory Memory::GetPrimitive() const {
     return dnnlMemHandle.getPrim();
 }
@@ -190,6 +174,31 @@ dnnl::memory Memory::DnnlMemPrimHandle::getPrim() const {
     return m_prim;
 }
 
+bool Memory::isAllocated() const noexcept {
+    if (mgrHandle->getRawPtr()) {
+        return true;
+    }
+    if (!pMemDesc) {
+        return false;
+    }
+    if (!(pMemDesc->isDefined())) {
+        return true;
+    }
+    if (pMemDesc->getCurrentMemSize() == 0) {
+        return true;
+    }
+    return false;
+}
+
+void* Memory::GetData() const {
+    void* data = getDataNoThrow();
+    if (data == nullptr &&
+        pMemDesc->getShape().isStatic() &&
+        pMemDesc->getShape().getElementsCount() != 0)
+        IE_THROW() << "Memory has not been allocated";
+    return data;
+}
+
 void* MemoryMngrWithReuse::getRawPtr() const noexcept {
     return _data.get();
 }

src/plugins/intel_cpu/src/cpu_memory.h

@@ -89,18 +89,24 @@ class MemoryMngrWithReuse : public IMemoryMngr {
     static void destroy(void *ptr);
 };
 
+class IMemoryMngrObserver : public IMemoryMngr {
+public:
+    virtual void registerMemory(Memory* memPtr) = 0;
+    virtual void unregisterMemory(Memory* memPtr) = 0;
+};
+
 /**
  * @brief A proxy object that additionally implements observer pattern
  */
-class DnnlMemoryMngr : public IMemoryMngr {
+class DnnlMemoryMngr : public IMemoryMngrObserver {
 public:
     explicit DnnlMemoryMngr(std::unique_ptr<IMemoryMngr> mngr) : _pMemMngr(std::move(mngr)) {}
     void* getRawPtr() const noexcept override;
     void setExtBuff(void* ptr, size_t size) override;
     bool resize(size_t size) override;
     bool hasExtBuffer() const noexcept override;
-    void registerMemory(Memory* memPtr);
-    void unregisterMemory(Memory* memPtr);
+    void registerMemory(Memory* memPtr) override;
+    void unregisterMemory(Memory* memPtr) override;
 
 private:
     void notifyUpdate();
@@ -110,12 +116,12 @@ class DnnlMemoryMngr : public IMemoryMngr {
     std::unique_ptr<IMemoryMngr> _pMemMngr;
 };
 
-using DnnlMemoryMngrPtr = std::shared_ptr<DnnlMemoryMngr>;
-using DnnlMemoryMngrCPtr = std::shared_ptr<const DnnlMemoryMngr>;
+using MemoryMngrPtr = std::shared_ptr<IMemoryMngrObserver>;
+using MemoryMngrCPtr = std::shared_ptr<const IMemoryMngrObserver>;
 
 class DnnlMemMngrHandle {
 public:
-    DnnlMemMngrHandle(DnnlMemoryMngrPtr pMgr, Memory* pMem) : _pMgr(pMgr), _pMem(pMem) {
+    DnnlMemMngrHandle(MemoryMngrPtr pMgr, Memory* pMem) : _pMgr(pMgr), _pMem(pMem) {
         if (_pMgr) {
             _pMgr->registerMemory(_pMem);
         }
@@ -140,123 +146,109 @@ class DnnlMemMngrHandle {
         }
     }
 
-    DnnlMemoryMngrPtr get() const {
+    MemoryMngrPtr get() const {
         return _pMgr;
     }
 
-    DnnlMemoryMngrPtr::element_type* operator->() const noexcept {
+    MemoryMngrPtr::element_type* operator->() const noexcept {
         return _pMgr.get();
     }
 
 private:
-    DnnlMemoryMngrPtr _pMgr = nullptr;
+    MemoryMngrPtr _pMgr = nullptr;
     Memory* _pMem = nullptr;
 };
 
-class Memory {
+class IMemory {
 public:
-    explicit Memory(const dnnl::engine& eng);
-    Memory(const dnnl::engine& eng, std::unique_ptr<IMemoryMngr> mngr);
+    virtual ~IMemory() = default;
 
-    Memory(const Memory&) = delete;
-    Memory& operator= (const Memory&) = delete;
+    virtual bool isAllocated() const noexcept = 0;
 
-    Memory(Memory&&) = delete;
-    Memory& operator= (Memory&&) = delete;
+    virtual const MemoryDesc& getDesc() const = 0;
+    virtual MemoryDescPtr getDescPtr() const = 0;
 
-    dnnl::memory GetPrimitive() const;
+    virtual void* GetData() const = 0; // pointer to the actual memory
 
-    bool isAllocated() const noexcept {
-        if (mgrHandle->getRawPtr()) {
-            return true;
-        }
-        if (!pMemDesc) {
-            return false;
-        }
-        if (!(pMemDesc->isDefined())) {
-            return true;
-        }
-        if (pMemDesc->getCurrentMemSize() == 0) {
-            return true;
-        }
-        return false;
-    }
+    virtual size_t GetSize() const = 0; // in bytes
+    virtual const Shape& GetShape() const = 0;
+    virtual const VectorDims& getStaticDims() const = 0;
 
-    /**
-     * @brief Resets the memory manager to a new one created with the provided raw memory
-     */
-    void setDataHandle(void* data);
+    // Redefines descriptor. The memory descriptor will be replaced with the new one.
+    // Memory will not be reallocated if the new tensor size is less or equal the upper bound.
+    // Caution!!! This action invalidates the previous data layout. The old data may become unreachable.
+    virtual void redefineDesc(MemoryDescPtr desc) = 0;
 
-    const MemoryDesc& getDesc() const {
-        return *pMemDesc;
-    }
+    virtual void SetData(const IMemory& memory, bool ftz = true) const = 0;
 
-    MemoryDescPtr getDescPtr() const {
-        return pMemDesc;
-    }
+    virtual MemoryMngrPtr getMemoryMngr() const = 0;
+
+    //oneDNN specifics for backward compatibility
+    virtual dnnl::memory GetPrimitive() const = 0;
+    virtual dnnl::memory::data_type GetDataType() const = 0;
+
+    virtual void FillZero() = 0;
 
     template <typename T,
             typename std::enable_if<!std::is_pointer<T>::value && !std::is_reference<T>::value, int>::type = 0,
             typename std::enable_if<std::is_base_of<MemoryDesc, T>::value, int>::type = 0>
     std::shared_ptr<T> GetDescWithType() const;
+};
 
-    /**
-     * Return handler of buffer. Real data may starts from some other offset
-     * @return
-     */
-    void* GetData() const {
-        void* data = getDataNoThrow();
-        if (data == nullptr &&
-            pMemDesc->getShape().isStatic() &&
-            pMemDesc->getShape().getElementsCount() != 0)
-            IE_THROW() << "Memory has not been allocated";
-        return data;
+class Memory : public IMemory {
+public:
+    Memory(const dnnl::engine& eng, MemoryDescPtr desc, const void* data = nullptr, bool pads_zeroing = true);
+    Memory(const dnnl::engine& eng, const MemoryDesc& desc, const void* data = nullptr, bool pads_zeroing = true);
+    Memory(const dnnl::engine& eng, MemoryDescPtr desc, MemoryMngrPtr mngr);
+    Memory(const dnnl::engine& eng, const MemoryDesc& desc, MemoryMngrPtr mbgr);
+
+    Memory(const Memory&) = delete;
+    Memory& operator= (const Memory&) = delete;
+
+    Memory(Memory&&) = delete;
+    Memory& operator= (Memory&&) = delete;
+
+    dnnl::memory GetPrimitive() const override;
+
+    bool isAllocated() const noexcept override;
+
+    const MemoryDesc& getDesc() const override {
+        return *pMemDesc;
     }
 
-    /**
-     * Return raw pointer on first element
-     * Like a GetData() but offset is applied.
-     * @return
-     */
-    void* GetPtr() const;
+    MemoryDescPtr getDescPtr() const override {
+        return pMemDesc;
+    }
+
+    void* GetData() const override;
 
-    dnnl::memory::data_type GetDataType() const {
+    dnnl::memory::data_type GetDataType() const override {
         return DnnlExtensionUtils::IEPrecisionToDataType(getDesc().getPrecision());
     }
 
-    size_t GetSize() const;
+    size_t GetSize() const override;
 
-    const Shape& GetShape() const {
+    const Shape& GetShape() const override {
         return getDesc().getShape();
     }
 
-    void Create(const MemoryDesc& desc, const void* data = nullptr, bool pads_zeroing = true);
-    void Create(MemoryDescPtr desc, const void* data = nullptr, bool pads_zeroing = true);
-
-    void Create(const MemoryDesc& desc, DnnlMemoryMngrPtr memMgr);
-    void Create(MemoryDescPtr desc, DnnlMemoryMngrPtr memMgr);
+    const VectorDims& getStaticDims() const override {
+        return getDesc().getShape().getStaticDims();
+    }
 
     // Redefines descriptor. The memory descriptor will be replaced with the new one.
     // Memory will not be reallocated if the new tensor size is less or equal the upper bound.
     // Caution!!! This action invalidates the previous data layout. The old data may become unreachable.
-    void redefineDesc(MemoryDescPtr desc);
+    void redefineDesc(MemoryDescPtr desc) override;
 
-    void SetData(const Memory& memory, bool ftz = true) const;
-    void FillZero();
-
-    const VectorDims& getStaticDims() const {
-        return getDesc().getShape().getStaticDims();
-    }
+    void SetData(const IMemory& memory, bool ftz = true) const override;
+    void FillZero() override;
 
     dnnl::engine getEngine() const {
         return eng;
     }
 
-    bool isUsedExternalStorage() const {
-        return mgrHandle->hasExtBuffer();
-    }
-
-    DnnlMemoryMngrPtr getDnnlMemoryMngr() const {
+    MemoryMngrPtr getMemoryMngr() const override {
         return mgrHandle.get();
     }
 
@@ -266,6 +258,9 @@ class Memory {
 private:
     void update();
 
+    void Create(const MemoryDesc& desc, const void* data = nullptr, bool pads_zeroing = true);
+    void Create(MemoryDescPtr desc, const void* data = nullptr, bool pads_zeroing = true);
+
 private:
     MemoryDescPtr pMemDesc;
     dnnl::engine eng;
@@ -291,8 +286,8 @@ class Memory {
     }
 };
 
-using MemoryPtr = std::shared_ptr<Memory>;
-using MemoryCPtr = std::shared_ptr<const Memory>;
+using MemoryPtr = std::shared_ptr<IMemory>;
+using MemoryCPtr = std::shared_ptr<const IMemory>;
 
 }   // namespace intel_cpu
 }   // namespace ov