From 96d186c8d26f2978fe02d825545ffd8898fd387d Mon Sep 17 00:00:00 2001
From: Kee <xuke537@hotmail.com>
Date: Thu, 11 Aug 2022 17:37:52 +0800
Subject: [PATCH] Set graph attributes when compile graph to binary

Keep the same graph attributes as compile graph

Signed-off-by: Kee <xuke537@hotmail.com>
---
 src/tim/vx/graph.cc        | 22 ++++++++--------------
 src/tim/vx/graph_private.h |  3 +++
 2 files changed, 11 insertions(+), 14 deletions(-)

diff --git a/src/tim/vx/graph.cc b/src/tim/vx/graph.cc
index afb8491..5b1080c 100644
--- a/src/tim/vx/graph.cc
+++ b/src/tim/vx/graph.cc
@@ -154,7 +154,7 @@ std::shared_ptr<Tensor> GraphImpl::CreateTensorPlaceHolder() {
   return tensor_placeholder_;
 }
 
-bool GraphImpl::Compile() {
+bool GraphImpl::Setup() {
   bool status = true;
 
   auto major = vsi_nn_GetVersionMajor();
@@ -180,7 +180,13 @@ bool GraphImpl::Compile() {
   std::call_once(setup_once_, [&status, this]() {
     status = (VSI_SUCCESS == vsi_nn_SetupGraph(this->graph_, true));
   });
+  return status;
+}
 
+bool GraphImpl::Compile() {
+  bool status = true;
+
+  status = Setup();
   std::call_once(verify_graph_once_, [&status, this]() {
     status = (VSI_SUCCESS == vsi_nn_VerifyGraph(this->graph_));
   });
@@ -189,19 +195,7 @@ bool GraphImpl::Compile() {
 }
 
 bool GraphImpl::CompileToBinary(void* buf, size_t* size) {
-  bool status = true;
-  std::call_once(setio_once_, [&status, this]() {
-    status = (vsi_nn_SetGraphInputs(this->graph_, this->inputs_.data(),
-                                    this->inputs_.size()) &&
-              vsi_nn_SetGraphOutputs(this->graph_, this->outputs_.data(),
-                                     this->outputs_.size()));
-  });
-
-  std::call_once(setup_once_, [&status, this]() {
-    status = (VSI_SUCCESS == vsi_nn_SetupGraph(this->graph_, true));
-  });
-
-  return ((status) && (VSI_SUCCESS == vsi_nn_GenerateNBG(graph_, buf, size)));
+  return ((Setup()) && (VSI_SUCCESS == vsi_nn_GenerateNBG(graph_, buf, size)));
 }
 
 bool GraphImpl::Run() {
diff --git a/src/tim/vx/graph_private.h b/src/tim/vx/graph_private.h
index 9f80e69..9a8ab19 100644
--- a/src/tim/vx/graph_private.h
+++ b/src/tim/vx/graph_private.h
@@ -93,6 +93,9 @@ class GraphImpl : public Graph {
   std::map<std::shared_ptr<Tensor>, std::shared_ptr<Operation>> tensor_producer_;
 
   CompileOption options_;
+ private:
+ /// Setup graph
+  bool Setup();
 };
 
 }  // namespace vx