From 45e206020ad5a6bb36e8491a9fd0b21b53035545 Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Mon, 28 Dec 2020 16:12:48 +0800
Subject: [PATCH] feat(mgb/tensor): add proxy_to_comp_node for HostTensorND

GitOrigin-RevId: 8f03dbcbd110010008926019adca7a8854eb2607
---
 src/core/include/megbrain/tensor.h | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/src/core/include/megbrain/tensor.h b/src/core/include/megbrain/tensor.h
index aae45495..3a613382 100644
--- a/src/core/include/megbrain/tensor.h
+++ b/src/core/include/megbrain/tensor.h
@@ -548,6 +548,15 @@ class TensorND {
             ret.reset(storage().proxy_to_default_cpu(), layout());
             return ret;
         }
+
+        template<bool x = true, typename = std::enable_if_t<x && std::is_same<TensorStorage, HostTensorStorage>::value>>
+        HostTensorND proxy_to_comp_node(CompNode cn) const {
+            HostTensorStorage host_storage;
+            host_storage.reset(cn, m_storage.size(), m_storage.raw_storage());
+            HostTensorND ret;
+            ret.reset(host_storage, m_layout);
+            return ret;
+        }
 };
 
 /*!