feat(core): add zero copy input output from out of megengine

GitOrigin-RevId: 388d9d9551916b24ab28c206ecc3906965154ac2

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

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();