From 8b7d8d290bd8b60a5382ab0dd41e8e92eef6d264 Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Mon, 24 May 2021 19:39:02 +0800
Subject: [PATCH] fix(core): fix json dump when weight preprocess

GitOrigin-RevId: 6cd882b10ddc0405f4475d2c118e030715153e37
---
 src/core/impl/graph/var_node.cpp |  1 +
 src/core/test/graph/misc.cpp     | 14 ++++++++++++++
 2 files changed, 15 insertions(+)

diff --git a/src/core/impl/graph/var_node.cpp b/src/core/impl/graph/var_node.cpp
index ac25c1333..7a4272fc7 100644
--- a/src/core/impl/graph/var_node.cpp
+++ b/src/core/impl/graph/var_node.cpp
@@ -528,6 +528,7 @@ std::shared_ptr<json::Value> VarNode::to_json() const {
         CHK(PERSISTENT_DEVICE_VALUE);
         CHK(DISALLOW_RT_FORCE_DYNAMIC_MEM_ALLOC);
         CHK(DISALLOW_VAR_SANITY_CHECK);
+        CHK(MEMORY_NO_NEED);
 #undef CHK
 
         mgb_assert(flag_checked == static_cast<size_t>(m_flag));
diff --git a/src/core/test/graph/misc.cpp b/src/core/test/graph/misc.cpp
index 602a36c6a..961465768 100644
--- a/src/core/test/graph/misc.cpp
+++ b/src/core/test/graph/misc.cpp
@@ -25,6 +25,7 @@
 #include "megbrain/utils/timer.h"
 #include "megbrain/comp_node_env.h"
 #include "megbrain/gopt/inference.h"
+#include "megbrain/plugin/profiler.h"
 
 #include "megbrain/test/helper.h"
 #include "megdnn/oprs/base.h"
@@ -1993,6 +1994,12 @@ typename megdnn::ExecutionPolicy try_find_any_bias_preprocess_algo(
 void test_free_memory_in_weight_preprocess(int record_level, CompNode cn) {
     HostTensorGenerator<> gen;
     auto graph = ComputingGraph::make();
+#if MGB_ENABLE_JSON
+    std::unique_ptr<GraphProfiler> profiler;
+    if(!record_level){
+        profiler = std::make_unique<GraphProfiler>(graph.get());
+    }
+#endif
     graph->options().graph_opt.weight_preprocess = true;
     graph->options().comp_node_seq_record_level = record_level;
     auto mkvar = [&](const char* name, const TensorShape& shp) {
@@ -2055,6 +2062,13 @@ void test_free_memory_in_weight_preprocess(int record_level, CompNode cn) {
     if (wp2.val()) {
         check(w2);
     }
+#if MGB_ENABLE_JSON
+    if (profiler) {
+        func->wait();
+        profiler->to_json_full(func.get())
+                ->writeto_fpath(output_file("weight_preprocess.json"));
+    }
+#endif
 }
 } // anonymous namespace
 
-- 
GitLab