From bdb853ee6f88a806ae0c7988becb8992807b009b Mon Sep 17 00:00:00 2001 From: Megvii Engine Team Date: Tue, 8 Feb 2022 19:03:58 +0800 Subject: [PATCH] fix(mgb): fix extra device malloc when load MultipleDeviceTensorWithFormatHolder GitOrigin-RevId: adf4a7f77a7b08c5e9a29ff126ece64a74f263e7 --- src/opr/impl/io.sereg.h | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/opr/impl/io.sereg.h b/src/opr/impl/io.sereg.h index 00865c48e..9260a1b64 100644 --- a/src/opr/impl/io.sereg.h +++ b/src/opr/impl/io.sereg.h @@ -218,11 +218,8 @@ struct OprLoadDumpImpl { auto handle = MegDNNHandle::get(CompNodeEnv::from_comp_node(i->comp_node())) .handle(); auto format = TensorFormat::deserialize(ctx.load_buf_with_len(), handle); - DeviceTensorStorage storage(i->comp_node()); TensorLayout layout_with_format{i->layout(), i->layout().dtype, format}; - auto size = layout_with_format.span().dist_byte(); - storage.ensure_size(size); if (i->storage().comp_node().mem_node() == CompNode::default_cpu().mem_node()) { mgb_assert( @@ -233,6 +230,13 @@ struct OprLoadDumpImpl { src.copy_from_fixlayout(*i).sync(); *i = DeviceTensorND::make_proxy(src); } else { + //! actually only layout of this tensor will be used later, see + //! src/serialization/impl/batched_device_value_loader.cpp:49. But we + //! have no way to reset layout only, so just construct a invalid + //! storage instead + auto size = layout_with_format.span().dist_byte(); + DeviceTensorStorage storage; + storage.reset(i->comp_node(), size, nullptr); i->reset(storage, layout_with_format); } } -- GitLab