From c87d998e591d05612c729b40627ba7e94b999be9 Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Wed, 9 Mar 2022 15:46:14 +0800
Subject: [PATCH] feat(mgb): add interface to support opencl IO zero copy when
 inference

GitOrigin-RevId: a1d7021892f430378c4eebd46152495d450efc61
---
 src/core/impl/comp_node/comp_node.cpp | 32 +++++++++++++++++++++++++++
 src/core/include/megbrain/comp_node.h | 21 ++++++++++++++++++
 2 files changed, 53 insertions(+)

diff --git a/src/core/impl/comp_node/comp_node.cpp b/src/core/impl/comp_node/comp_node.cpp
index bff186d18..f973c343b 100644
--- a/src/core/impl/comp_node/comp_node.cpp
+++ b/src/core/impl/comp_node/comp_node.cpp
@@ -658,4 +658,36 @@ void CompNode::ImplBase::add_callback(megdnn::thin_function<void()>&&) {
             locator().to_string().c_str());
 }
 
+void CompNode::map_to_cpu(void* ptr, size_t size, bool blocking) {
+    m_impl->map_to_cpu(ptr, size, blocking);
+}
+
+void CompNode::unmap_to_gpu(void* ptr, size_t size) {
+    m_impl->unmap_to_gpu(ptr, size);
+}
+
+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::ImplBase::map_to_cpu(void* ptr, size_t size, bool blocking) {
+    MGB_MARK_USED_VAR(ptr);
+    MGB_MARK_USED_VAR(size);
+    MGB_MARK_USED_VAR(blocking);
+    mgb_assert(false, "No map_to_cpu Impl");
+}
+
+void CompNode::ImplBase::unmap_to_gpu(void* ptr, size_t size) {
+    MGB_MARK_USED_VAR(ptr);
+    MGB_MARK_USED_VAR(size);
+    mgb_assert(false, "No unmap_to_gpu Impl");
+}
+
+void* CompNode::ImplBase::get_logical_addr_by_host_ptr(void* ptr, size_t size) {
+    MGB_MARK_USED_VAR(ptr);
+    MGB_MARK_USED_VAR(size);
+    mgb_assert(false, "No get_logical_addr_by_host_ptr Impl");
+    return nullptr;
+}
+
 // 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 2f9c986e5..a068fd352 100644
--- a/src/core/include/megbrain/comp_node.h
+++ b/src/core/include/megbrain/comp_node.h
@@ -322,6 +322,21 @@ public:
     MGE_WIN_DECLSPEC_FUC static DeviceProperties get_device_prop(
             int dev, DeviceType device_type);
 
+    /*!
+     * \brief get control of host ptr to user
+     */
+    MGE_WIN_DECLSPEC_FUC void map_to_cpu(void* ptr, size_t size, bool blocking = false);
+
+    /*!
+     * \brief release control of host ptr to system
+     */
+    MGE_WIN_DECLSPEC_FUC void unmap_to_gpu(void* ptr, size_t size);
+
+    /*!
+     * \brief get logical address by host ptr
+     */
+    MGE_WIN_DECLSPEC_FUC void* get_logical_addr_by_host_ptr(void* ptr, size_t size);
+
     /* =================== synchronization ======================== */
 
     class Event;
@@ -573,6 +588,12 @@ protected:
             peer_copy_to(dest_impl, dest.get_ptr(), src.get_ptr(), size);
         }
 
+        virtual void map_to_cpu(void* ptr, size_t size, bool blocking = false);
+
+        virtual void unmap_to_gpu(void* ptr, size_t size);
+
+        virtual void* get_logical_addr_by_host_ptr(void* ptr, size_t size);
+
         virtual size_t get_mem_addr_alignment() = 0;
         virtual size_t get_mem_padding();
 
-- 
GitLab