diff --git a/src/core/impl/comp_node/comp_node.cpp b/src/core/impl/comp_node/comp_node.cpp index 46ef3fbda0c9d650901911e2298c18084613bd63..094c8c33c487bf1f2e1a66c6a2cf86ca67f04d38 100644 --- a/src/core/impl/comp_node/comp_node.cpp +++ b/src/core/impl/comp_node/comp_node.cpp @@ -659,6 +659,14 @@ void* CompNode::get_logical_addr_by_host_ptr(void* ptr, size_t size) { return m_impl->get_logical_addr_by_host_ptr(ptr, size); } +void* CompNode::register_external_device_ptr(void* ptr, size_t size) { + return m_impl->register_external_device_ptr(ptr, size); +} + +void* CompNode::unregister_external_device_ptr(void* ptr, size_t size) { + return m_impl->unregister_external_device_ptr(ptr, size); +} + void CompNode::ImplBase::map_to_cpu(void* ptr, size_t size, bool blocking) { MGB_MARK_USED_VAR(ptr); MGB_MARK_USED_VAR(size); @@ -679,4 +687,14 @@ void* CompNode::ImplBase::get_logical_addr_by_host_ptr(void* ptr, size_t size) { return nullptr; } +void* CompNode::ImplBase::register_external_device_ptr(void* ptr, size_t size) { + MGB_MARK_USED_VAR(size); + return ptr; +} + +void* CompNode::ImplBase::unregister_external_device_ptr(void* ptr, size_t size) { + MGB_MARK_USED_VAR(size); + return ptr; +} + // vim: syntax=cpp.doxygen foldmethod=marker foldmarker=f{{{,f}}} diff --git a/src/core/include/megbrain/comp_node.h b/src/core/include/megbrain/comp_node.h index 72fc680eb9539cf6875151e67c1e900fa9f527a8..6fcc3651f75aad55f4c5e389cbd7da15b0fd853a 100644 --- a/src/core/include/megbrain/comp_node.h +++ b/src/core/include/megbrain/comp_node.h @@ -324,23 +324,22 @@ public: /*! * \brief get logical address by host ptr */ - //! api for register memory info, in order compat old user code, do not change the - //! api name, this api have different behavior for with different OpenCL memory - //! type. - //! Map/UnMap: use to get logical_addr of cl_mem alloc by MegEngine, which can - //! be reset to TensorND storage raw_ptr, NOTICE: Map/UnMap do not support cl_mem - //! not alloc from MegEngine. - //! SVM: in fact do nothing, just return ptr same with args. - //! ION: this case support two case: - //! ION case one: IonHostPtr info is alloc by MegEngine, this case, args ptr is - //! tensor raw_ptr, and args size is tensor size in bytes, api will return - //! IonHostPtr ptr, then user can use it out of EegEngine. - //! ION case two: use to register user IonHostPtr info, args ptr is IonHostPtr, size - //! is IonHostPtr->ion_hostptr size_in_bytes, MegEngine will register it into DNN, - //! then just return IonHostPtr->ion_hostptr, which can be reset to TensorND storage - //! raw_ptr. MGE_WIN_DECLSPEC_FUC void* get_logical_addr_by_host_ptr(void* ptr, size_t size); + /*! + * \brief register user external device ptr, which means not malloc by MegEngine + * case 1: cpu and cuda compnode will do nothing, just return args ptr + * case 2: OpenCL(ION) compnode will do real register, OpenCL(map/svm) compnode will + * trigger assert, caused by OpenCL only can use extern ION ptr, can not use map/svm + * with different OpenCL context. + */ + MGE_WIN_DECLSPEC_FUC void* register_external_device_ptr(void* ptr, size_t size); + + /*! + * \brief unregister user external device ptr, which means not malloc by MegEngine + */ + MGE_WIN_DECLSPEC_FUC void* unregister_external_device_ptr(void* ptr, size_t size); + /* =================== synchronization ======================== */ class Event; @@ -598,6 +597,10 @@ protected: virtual void* get_logical_addr_by_host_ptr(void* ptr, size_t size); + virtual void* register_external_device_ptr(void* ptr, size_t size); + + virtual void* unregister_external_device_ptr(void* ptr, size_t size); + virtual size_t get_mem_addr_alignment() = 0; virtual size_t get_mem_padding();