From d981dda97f1a68c34cb7f80f6cf1307ccabbd25e Mon Sep 17 00:00:00 2001
From: He Wei <hewei@huawei.com>
Date: Fri, 31 Jul 2020 18:50:46 +0800
Subject: [PATCH] Decouple infer trace from LogWriter

---
 mindspore/ccsrc/debug/trace.cc      | 13 +++++++++++++
 mindspore/core/utils/log_adapter.cc | 10 +++-------
 mindspore/core/utils/log_adapter.h  |  3 +++
 3 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/mindspore/ccsrc/debug/trace.cc b/mindspore/ccsrc/debug/trace.cc
index 7ceb368290b..605d09ce7fc 100644
--- a/mindspore/ccsrc/debug/trace.cc
+++ b/mindspore/ccsrc/debug/trace.cc
@@ -33,6 +33,7 @@
 #include "ir/tensor.h"
 #include "debug/anf_ir_utils.h"
 #include "pipeline/jit/static_analysis/evaluator.h"
+#include "utils/log_adapter.h"
 
 namespace mindspore {
 // namespace to support debug trace infomation
@@ -495,5 +496,17 @@ void ClearTraceStack() {
   }
   cnode_debug_stack.clear();
 }
+
+// Register trace provider to LogWriter.
+struct TraceProviderRegister {
+  TraceProviderRegister() {
+    LogWriter::set_trace_provider([](std::ostringstream &oss) {
+      TraceGraphEval();
+      GetEvalStackInfo(oss);
+    });
+  }
+  ~TraceProviderRegister() = default;
+} trace_provider_regsiter;
+
 }  // namespace trace
 }  // namespace mindspore
diff --git a/mindspore/core/utils/log_adapter.cc b/mindspore/core/utils/log_adapter.cc
index 3a3f0e296ce..8d91a6f7768 100644
--- a/mindspore/core/utils/log_adapter.cc
+++ b/mindspore/core/utils/log_adapter.cc
@@ -19,9 +19,6 @@
 #include <unistd.h>
 #include <sys/time.h>
 #include <map>
-#ifndef USE_ANDROID_LOG
-#include "debug/trace.h"
-#endif
 
 // namespace to support utils module definition
 namespace mindspore {
@@ -244,10 +241,9 @@ void LogWriter::operator^(const LogStream &stream) const {
   }
   oss << msg.str();
 
-#ifndef USE_ANDROID_LOG
-  trace::TraceGraphEval();
-  trace::GetEvalStackInfo(oss);
-#endif
+  if (trace_provider_ != nullptr) {
+    trace_provider_(oss);
+  }
 
   if (exception_handler_ != nullptr) {
     exception_handler_(exception_type_, oss.str());
diff --git a/mindspore/core/utils/log_adapter.h b/mindspore/core/utils/log_adapter.h
index d354f96e68d..ce31ce9ab8e 100644
--- a/mindspore/core/utils/log_adapter.h
+++ b/mindspore/core/utils/log_adapter.h
@@ -142,6 +142,7 @@ extern int g_ms_submodule_log_levels[] __attribute__((visibility("default")));
 class LogWriter {
  public:
   using ExceptionHandler = std::function<void(ExceptionType, const std::string &msg)>;
+  using TraceProvider = std::function<void(std::ostringstream &oss)>;
 
   LogWriter(const LocationInfo &location, MsLogLevel log_level, SubModuleId submodule,
             ExceptionType excp_type = NoExceptionType)
@@ -152,6 +153,7 @@ class LogWriter {
   void operator^(const LogStream &stream) const __attribute__((noreturn, visibility("default")));
 
   static void set_exception_handler(ExceptionHandler exception_handler) { exception_handler_ = exception_handler; }
+  static void set_trace_provider(TraceProvider trace_provider) { trace_provider_ = trace_provider; }
 
  private:
   void OutputLog(const std::ostringstream &msg) const;
@@ -162,6 +164,7 @@ class LogWriter {
   ExceptionType exception_type_;
 
   inline static ExceptionHandler exception_handler_ = nullptr;
+  inline static TraceProvider trace_provider_ = nullptr;
 };
 
 #define MSLOG_IF(level, condition, excp_type)                                                                       \