From ebe8689267b22fb8f9bf1a4aea839d39886ecb27 Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Fri, 8 Jan 2021 16:11:57 +0800
Subject: [PATCH] feat(sdk/load_and_run): add flags --io-dump-stdout and
 --io-dump-stderr

GitOrigin-RevId: 64d572bdfc4e2008497bc1e62c44e852acfd46a0
---
 sdk/load-and-run/src/mgblar.cpp | 20 ++++++++++++++++++++
 src/plugin/test/opr_io_dump.cpp | 18 ++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/sdk/load-and-run/src/mgblar.cpp b/sdk/load-and-run/src/mgblar.cpp
index 506998ae..b7888bc5 100644
--- a/sdk/load-and-run/src/mgblar.cpp
+++ b/sdk/load-and-run/src/mgblar.cpp
@@ -94,6 +94,8 @@ R"__usage__(
     Dump input/output values of all internal variables to output file or
     directory, in text or binary format. The binary file can be parsed by
     `megbrain.plugin.load_tensor_binary`.
+  --io-dump-stdout | --io-dump-stderr
+    Dump input/output values of all internal variables to stdout or stderr in text format
   --bin-out-dump <output dir>
     Dump output tensor values in binary format to given directory.
   --iter <num>
@@ -1200,6 +1202,24 @@ Args Args::from_argv(int argc, char **argv) {
             ret.iodump = std::move(iodump);
             continue;
         }
+        if (!strcmp(argv[i], "--io-dump-stdout")) {
+            mgb_log_warn("enable opr io dump to stdout");
+            std::shared_ptr<FILE> sp(stdout, [](FILE*){});
+            auto iodump = std::make_unique<TextOprIODump>(
+                    ret.load_config.comp_graph.get(), sp);
+            iodump->print_addr(false);
+            ret.iodump = std::move(iodump);
+            continue;
+        }
+        if (!strcmp(argv[i], "--io-dump-stderr")) {
+            mgb_log_warn("enable opr io dump to stderr");
+            std::shared_ptr<FILE> sp(stderr, [](FILE*){});
+            auto iodump = std::make_unique<TextOprIODump>(
+                    ret.load_config.comp_graph.get(), sp);
+            iodump->print_addr(false);
+            ret.iodump = std::move(iodump);
+            continue;
+        }
         if (!strcmp(argv[i], "--bin-io-dump")) {
             mgb_log_warn("enable opr binary io dump");
             ++ i;
diff --git a/src/plugin/test/opr_io_dump.cpp b/src/plugin/test/opr_io_dump.cpp
index 4907940f..3b88a2cc 100644
--- a/src/plugin/test/opr_io_dump.cpp
+++ b/src/plugin/test/opr_io_dump.cpp
@@ -181,6 +181,24 @@ TEST(TestOprIODump, Text) {
     run_test(make_plugin, check_result);
 }
 
+TEST(TestOprIODump, StdErr) {
+    HostTensorGenerator<> gen;
+    auto host_x = gen({5});
+    auto host_y = gen({5});
+
+    auto graph = ComputingGraph::make();
+    std::shared_ptr<FILE> sp(stdout, [](FILE*){});
+    auto plugin = std::make_unique<TextOprIODump>(graph.get(), sp);
+
+    auto x = opr::Host2DeviceCopy::make(*graph, host_x);
+    auto y = opr::Host2DeviceCopy::make(*graph, host_y);
+    auto z = x + y;
+
+    HostTensorND host_z;
+    auto func = graph->compile({make_callback_copy(z, host_z)});
+    func->execute();
+}
+
 TEST(TestOprIODump, Binary) {
     auto fname = output_file("");
     auto make_plugin = [&](ComputingGraph* graph, int level) {