[neurun] Apply SubgraphContext(1 subgraph 1 node) (#5014)

Apply SubgraphContext(1 subgraph 1 node) to Graph, Linear, Compiler and
{Dataflow|Linear}Executor

Signed-off-by: Yongseop Kim <yons.kim@samsung.com>

9 files changed, 207 insertions, 140 deletions

diff --git a/runtimes/neurun/core/include/graph/Graph.h b/runtimes/neurun/core/include/graph/Graph.h
index c4b4ffb6b..504303ea8 100644
--- a/runtimes/neurun/core/include/graph/Graph.h
+++ b/runtimes/neurun/core/include/graph/Graph.h
@@ -23,6 +23,7 @@
 #include "model/Model.h"
 #include "graph/LowerInfoMap.h"
 #include "model/operation/Subgraph.h"
+#include "model/operation/SubgraphContext.h"
 
 namespace neurun
 {
@@ -127,11 +128,9 @@ public:
   bool isBuildingPhase(void) const { return _phase == Phase::BUILDING; }
   std::shared_ptr<const model::Model> shareModel() { return _model; }
   std::unique_ptr<graph::LowerInfoMap> releaseLowerInfo() { return std::move(_lower_info_map); }
-  // TODO Change this to releaseSubgraphContext()
-  std::unique_ptr<std::vector<std::unique_ptr<model::operation::Subgraph>>>
-  releaseSubgraphSequence()
+  std::unique_ptr<model::operation::SubgraphContext> releaseSubgraphContext()
   {
-    return std::move(_subg_seq);
+    return std::move(_subg_ctx);
   }
 
 private:
@@ -158,8 +157,8 @@ private:
 
   // For LOWERED phase
 public:
-  const operation::LowerInfo *getLowerInfo(const model::OperationIndex &index) const;
-  void setLowerInfo(const model::OperationIndex &index,
+  const operation::LowerInfo *getLowerInfo(const model::OperationIndex &subg_index) const;
+  void setLowerInfo(const model::OperationIndex &subg_index,
                     std::unique_ptr<operation::LowerInfo> &&lower_info);
   const operand::LowerInfo *getLowerInfo(const model::operand::Index &index) const;
   operand::LowerInfo *getLowerInfo(const model::operand::Index &index);
@@ -171,11 +170,16 @@ private:
   std::unique_ptr<LowerInfoMap> _lower_info_map;
 
   // For Subgraph
-private:
-  void partition();
+public:
+  model::operation::SubgraphContext &subg_ctx()
+  {
+    assert(_subg_ctx);
+    return *_subg_ctx;
+  }
 
 private:
-  std::unique_ptr<std::vector<std::unique_ptr<model::operation::Subgraph>>> _subg_seq;
+  // Pass(for Perm) can accept only graph so that Graph has SubgraphContext as a member
+  std::unique_ptr<model::operation::SubgraphContext> _subg_ctx;
 };
 
 } // namespace graph
diff --git a/runtimes/neurun/core/src/compiler/Compiler.cc b/runtimes/neurun/core/src/compiler/Compiler.cc
index 6d3143a6c..8051e8ff2 100644
--- a/runtimes/neurun/core/src/compiler/Compiler.cc
+++ b/runtimes/neurun/core/src/compiler/Compiler.cc
@@ -151,13 +151,14 @@ std::shared_ptr<exec::IExecutor> Compiler::createDataflowExecutor(graph::Graph &
   auto operand_context = std::make_shared<OperandContext>();
   std::unordered_map<model::OperationIndex, std::unique_ptr<backend::IStage>> stages;
 
-  model.operations().iterate([&](const model::OperationIndex &index, const model::Operation &node) {
-    auto backend = model.getLowerInfo(index)->backend();
+  model.subg_ctx().iterate(
+      [&](const model::OperationIndex &subg_index, const model::operation::Subgraph &subg) {
+        auto backend = model.getLowerInfo(subg_index)->backend();
 
-    // Generate Stage
-    auto stage_gen = backend->stage_gen();
-    stages[index] = stage_gen->generate(node);
-  });
+        // Generate Stage
+        auto stage_gen = backend->stage_gen();
+        stages[subg_index] = stage_gen->generate(subg);
+      });
 
   backend::TensorBuilderSet tensor_builders;
 
@@ -229,9 +230,9 @@ std::shared_ptr<exec::IExecutor> Compiler::createDataflowExecutor(graph::Graph &
 
   ConstantInitializer{model, *operand_context, *lower_info}();
 
-  return std::make_shared<exec::DataflowExecutor>(model.shareModel(), operand_context,
-                                                  std::move(lower_info),
-                                                  std::move(execution_builder.releaseCodeMap()));
+  return std::make_shared<exec::DataflowExecutor>(
+      model.shareModel(), std::move(model.releaseSubgraphContext()), operand_context,
+      std::move(lower_info), std::move(execution_builder.releaseCodeMap()));
 }
 
 bool Compiler::checkCompilable()
diff --git a/runtimes/neurun/core/src/exec/DataflowExecutor.cc b/runtimes/neurun/core/src/exec/DataflowExecutor.cc
index 92d5a86ff..f78f1504a 100644
--- a/runtimes/neurun/core/src/exec/DataflowExecutor.cc
+++ b/runtimes/neurun/core/src/exec/DataflowExecutor.cc
@@ -47,20 +47,24 @@ void DataflowExecutor::notify(const model::operand::IndexSet &operands)
 }
 
 DataflowExecutor::DataflowExecutor(const std::shared_ptr<const model::Model> &model,
+                                   std::unique_ptr<model::operation::SubgraphContext> subg_ctx,
                                    const std::shared_ptr<compiler::OperandContext> &operand_context,
                                    std::unique_ptr<graph::LowerInfoMap> lower_info,
                                    CodeMap &&code_map)
     : ExecutorBase{model, operand_context, nullptr, std::move(lower_info)},
-      _code_map{std::move(code_map)}
+      _code_map{std::move(code_map)}, _subg_ctx(std::move(subg_ctx))
 {
   VERBOSE(DataflowExecutor) << "Constructing Dataflow Executor" << std::endl;
 
+  assert(_subg_ctx);
+
   // Create jobs
-  _model->operations.iterate([&](const model::OperationIndex &index, const model::Operation &node) {
-    VERBOSE(DataflowExecutor) << "Add a job #" << index.value() << std::endl;
-    _finished_jobs.emplace_back(nnfw::cpp14::make_unique<Job>(index, _code_map.at(index).get(),
-                                                              node.getInputs(), node.getOutputs()));
-  });
+  _subg_ctx->iterate(
+      [&](const model::OperationIndex &subg_index, const model::operation::Subgraph &subg) {
+        VERBOSE(DataflowExecutor) << "Add a job #" << subg_index.value() << std::endl;
+        _finished_jobs.emplace_back(nnfw::cpp14::make_unique<Job>(
+            subg_index, _code_map.at(subg_index).get(), subg.getInputs(), subg.getOutputs()));
+      });
 
   // Save operands that are initially ready
   {
diff --git a/runtimes/neurun/core/src/exec/DataflowExecutor.h b/runtimes/neurun/core/src/exec/DataflowExecutor.h
index e9530b70c..55c257285 100644
--- a/runtimes/neurun/core/src/exec/DataflowExecutor.h
+++ b/runtimes/neurun/core/src/exec/DataflowExecutor.h
@@ -28,6 +28,7 @@
 #include "model/Model.h"
 #include "cpp14/memory.h"
 #include "exec/ExecutorBase.h"
+#include "model/operation/SubgraphContext.h"
 
 namespace neurun
 {
@@ -52,6 +53,7 @@ public:
    * @param code_map Compiled code map
    */
   DataflowExecutor(const std::shared_ptr<const model::Model> &model,
+                   std::unique_ptr<model::operation::SubgraphContext> subg_ctx,
                    const std::shared_ptr<compiler::OperandContext> &operand_context,
                    std::unique_ptr<graph::LowerInfoMap> lower_info, CodeMap &&code_map);
 
@@ -59,6 +61,7 @@ public:
 
 private:
   CodeMap _code_map;
+  std::unique_ptr<model::operation::SubgraphContext> _subg_ctx;
   model::operand::IndexSet _initially_ready_operands;
   std::list<std::unique_ptr<Job>> _finished_jobs;
   std::list<std::unique_ptr<Job>> _waiting_jobs;
diff --git a/runtimes/neurun/core/src/graph/Graph.cc b/runtimes/neurun/core/src/graph/Graph.cc
index 18e3e9b42..3a5166382 100644
--- a/runtimes/neurun/core/src/graph/Graph.cc
+++ b/runtimes/neurun/core/src/graph/Graph.cc
@@ -94,7 +94,7 @@ void Graph::lower(void)
 {
   assert(_phase == Phase::MODEL);
 
-  partition();
+  _subg_ctx = nnfw::cpp14::make_unique<model::operation::SubgraphContext>();
 
   // Lower
   {
@@ -110,11 +110,25 @@ void Graph::lower(void)
     _backend_resolver = nnfw::cpp14::make_unique<compiler::BackendResolver>(_model->operands);
     _lower_info_map = nnfw::cpp14::make_unique<LowerInfoMap>();
 
-    _model->operations.iterate([&](const model::OperationIndex &index, model::Operation &node) {
+    const auto &make_subgraph = [&](const model::OperationIndex &node_index,
+                                    const model::Operation &node) {
+      auto subg_index = _subg_ctx->append(node_index, node);
+      auto &subg = _subg_ctx->at(subg_index);
+      subg.setOutputs(node.getOutputs());
+      subg.setInputs(node.getInputs());
+      return subg_index;
+    };
+
+    // TODO Add code for merging(exactly appending properly) nodes in subgraph
+    // 1 subgraph 1 node
+    _model->operations.iterate([&](const model::OperationIndex &node_index,
+                                   model::Operation &node) {
+      auto new_subg_index = make_subgraph(node_index, node);
+
       auto backend = _backend_resolver->getBackend(typeid(node));
 
       // Operation LowerInfo
-      setLowerInfo(index, nnfw::cpp14::make_unique<graph::operation::LowerInfo>(backend));
+      setLowerInfo(new_subg_index, nnfw::cpp14::make_unique<graph::operation::LowerInfo>(backend));
 
       // LowerInfo for in/output operands
       for (auto operand : node.getInputs())
@@ -128,6 +142,7 @@ void Graph::lower(void)
         lower_info->addDefBackend(backend);
       }
     });
+    _subg_ctx->dump("1 subgraph 1 node without permutation");
 
 // NOTE This is desired way to handle model input and outputs however getDefaultBackend() is
 // cpu backend dependent for now we cannot use it.
@@ -235,6 +250,9 @@ void Graph::lower(void)
     pi_pass.run();
     pass::PermutationEliminationPass pe_pass(*this);
     pe_pass.run();
+
+    // TODO merge perm subgraphs if possible
+    _subg_ctx->dump("merged and sorted operations with permutation");
   }
 
   // Graph verifications for the LOWERED phase
@@ -248,7 +266,7 @@ std::unique_ptr<linear::Linear> Graph::linearize(void)
 {
   assert(_phase == Phase::MODEL);
 
-  auto linear = nnfw::cpp14::make_unique<linear::Linear>(shareModel(), releaseSubgraphSequence(),
+  auto linear = nnfw::cpp14::make_unique<linear::Linear>(shareModel(), releaseSubgraphContext(),
                                                          releaseLowerInfo());
 
   // TODO Move the operations and operands to linear object
@@ -273,21 +291,21 @@ void Graph::initializeUseDef()
       });
 }
 
-const operation::LowerInfo *Graph::getLowerInfo(const model::OperationIndex &index) const
+const operation::LowerInfo *Graph::getLowerInfo(const model::OperationIndex &subg_index) const
 {
   if (!_lower_info_map)
     return nullptr;
-  auto itr = _lower_info_map->operation.find(index);
+  auto itr = _lower_info_map->operation.find(subg_index);
   if (itr == _lower_info_map->operation.end())
     return nullptr;
   return itr->second.get();
 }
 
-void Graph::setLowerInfo(const model::OperationIndex &index,
+void Graph::setLowerInfo(const model::OperationIndex &subg_index,
                          std::unique_ptr<operation::LowerInfo> &&lower_info)
 {
   assert(_lower_info_map);
-  _lower_info_map->operation.insert(std::make_pair(index, std::move(lower_info)));
+  _lower_info_map->operation.insert(std::make_pair(subg_index, std::move(lower_info)));
 }
 
 const operand::LowerInfo *Graph::getLowerInfo(const model::operand::Index &index) const
@@ -317,99 +335,6 @@ void Graph::setLowerInfo(const model::operand::Index &index,
   _lower_info_map->operand.insert(std::make_pair(index, std::move(lower_info)));
 }
 
-void Graph::partition()
-{
-  // Partition the graph into some subgraphs by topological sort while assuming that
-  // a subgraph has linear form
-  //
-  // algorithm
-  //   0. Create new subgraph
-  //   1. Add a node into current subgraph
-  //   2. Test two stuff for checking new subgraph is needed
-  //     - Current node has multiple inputs like concat?
-  //       - Does current node have two or more than previous operation?
-  //
-  //        [CONV] [CONV] [CONV]  [MAX_POOL]
-  //         |      |      |       |
-  //        [0]    [1]    [2]     [3]
-  //         \      |      |      /
-  //          [    C O N C A T   ]  # current node
-  //
-  //     - Current node is on the separated branch at the beginning?
-  //       - Does current node's input operand's uses have two or more than?
-  //
-  //       [CONV]
-  //         |
-  //        [0]----.
-  //         |     |
-  //       [CONV] [CONV]  # current node
-  //         |      |
-  //        [1]    [2]
-  //         \      /
-  //         [CONCAT]
-  //
-
-  _subg_seq = nnfw::cpp14::make_unique<std::vector<std::unique_ptr<model::operation::Subgraph>>>();
-
-  {
-    std::unique_ptr<model::operation::Subgraph> subg = nullptr;
-    Graph::PostDfsConstIterator().iterate(
-        *this, [&](const model::OperationIndex &index, const model::Operation &node) {
-
-          if (!subg)
-            subg = nnfw::cpp14::make_unique<model::operation::Subgraph>();
-
-          subg->appendOperation(index, node);
-
-          bool finish_subg = false;
-          size_t prev_op_cnt = 0;
-          for (auto input : node.getInputs())
-          {
-            const auto &operand = this->operands().at(input);
-            if (operand.getDef().list().size() > 0)
-              ++prev_op_cnt;
-
-            if (prev_op_cnt > 1 || operand.getUses().list().size() > 1)
-            {
-              finish_subg = true;
-              break;
-            }
-          }
-
-          if (finish_subg)
-          {
-            _subg_seq->emplace_back(std::move(subg));
-            subg = nullptr;
-          }
-        });
-
-    // If the last subgraph leaves, append it to the subgraph set
-    if (subg && subg->operations().size() > 0)
-      _subg_seq->emplace_back(std::move(subg));
-
-    // NOTE. Now these subgraph are on the reverse order
-  }
-
-  // Set input/output of each subgraph while reversing
-  std::reverse(_subg_seq->begin(), _subg_seq->end());
-  for (auto &subg : *_subg_seq)
-  {
-    // output
-    auto it = std::begin(subg->operations());
-    subg->setOutputs((*it).node->getOutputs());
-
-    std::reverse(std::begin(subg->operations()), std::end(subg->operations()));
-
-    // input
-    it = std::begin(subg->operations());
-    subg->setInputs((*it).node->getInputs());
-  }
-
-  VERBOSE(Subgraph) << "Subgraphs" << std::endl;
-  for (const auto &subg : *_subg_seq)
-    VERBOSE(Subgraph) << subg->getStr() << std::endl;
-}
-
 } // namespace graph
 } // namespace neurun
 
diff --git a/runtimes/neurun/core/src/graph/pass/PermutationEliminationPass.cc b/runtimes/neurun/core/src/graph/pass/PermutationEliminationPass.cc
index 9173b7a18..ba91e971e 100644
--- a/runtimes/neurun/core/src/graph/pass/PermutationEliminationPass.cc
+++ b/runtimes/neurun/core/src/graph/pass/PermutationEliminationPass.cc
@@ -79,6 +79,9 @@ void PermutationEliminationPass::eliminateInput(const model::operand::Index &inp
     _graph.removeOperand(inp_index);
 
     // remove permutation operation
+    assert(_graph.subg_ctx().hasNode(input_use));
+    auto subg_idx = _graph.subg_ctx().findNode(input_use);
+    _graph.subg_ctx().remove(subg_idx);
     _graph.operations().remove(input_use);
 
     VERBOSE(PermutationEliminationPass::EliminateInput)
@@ -131,6 +134,9 @@ void PermutationEliminationPass::eliminateOutput(const model::operand::Index &ou
     _graph.removeOperand(out_index);
 
     // remove permutation operation
+    assert(_graph.subg_ctx().hasNode(output_def));
+    auto subg_idx = _graph.subg_ctx().findNode(output_def);
+    _graph.subg_ctx().remove(subg_idx);
     _graph.operations().remove(output_def);
 
     VERBOSE(PermutationEliminationPass::EliminateOutput)
diff --git a/runtimes/neurun/core/src/graph/pass/PermutationInsertionPass.cc b/runtimes/neurun/core/src/graph/pass/PermutationInsertionPass.cc
index 25f4a28a8..47a730bbb 100644
--- a/runtimes/neurun/core/src/graph/pass/PermutationInsertionPass.cc
+++ b/runtimes/neurun/core/src/graph/pass/PermutationInsertionPass.cc
@@ -86,16 +86,21 @@ void PermutationInsertionPass::callback(const model::operand::Index &index,
         continue;
 
       auto &operation = _graph.operations().at(use);
-      auto operation_li = _graph.getLowerInfo(use);
-      assert(operation_li);
-      auto backend = operation_li->backend();
-
+      assert(_graph.subg_ctx().hasNode(use));
+      auto subg_index = _graph.subg_ctx().findNode(use);
+      auto subg_li = _graph.getLowerInfo(subg_index);
+      assert(subg_li);
+      const backend::Backend *backend = subg_li->backend();
+      assert(backend);
       auto use_node_inputs = operation.getInputs();
       assert(use_node_inputs.contains(index));
 
       auto new_index = backend_to_index.at(backend);
       if (index != new_index)
       {
+        // Update from subgraph
+        _graph.subg_ctx().at(subg_index).replaceInput(index, new_index);
+
         // Update from operation
         operation.replaceInput(index, new_index);
 
@@ -176,8 +181,15 @@ PermutationInsertionPass::insertPermute(const model::operand::Index &operand_ind
   auto node_index = _graph.operations().append(std::move(insert_node));
   const auto &node = _graph.operations().at(node_index);
 
-  _graph.setLowerInfo(node_index, nnfw::cpp14::make_unique<graph::operation::LowerInfo>(
-                                      _graph.backend_resolver()->getDefaultBackend()));
+  // Subgraph
+  {
+    auto subg_index = _graph.subg_ctx().append(node_index, node);
+    auto &subg = _graph.subg_ctx().at(subg_index);
+    subg.setInputs(node.getInputs());
+    subg.setOutputs(node.getOutputs());
+    _graph.setLowerInfo(subg_index, nnfw::cpp14::make_unique<graph::operation::LowerInfo>(
+                                        _graph.backend_resolver()->getDefaultBackend()));
+  }
 
   // Update Use/Def info
   {
diff --git a/runtimes/neurun/core/src/linear/Linear.cc b/runtimes/neurun/core/src/linear/Linear.cc
index 2ca58a597..674808c9a 100644
--- a/runtimes/neurun/core/src/linear/Linear.cc
+++ b/runtimes/neurun/core/src/linear/Linear.cc
@@ -35,16 +35,128 @@ namespace linear
 {
 
 Linear::Linear(const std::shared_ptr<const model::Model> &model,
-               std::unique_ptr<std::vector<std::unique_ptr<model::operation::Subgraph>>> subg_seq,
+               std::unique_ptr<model::operation::SubgraphContext> subg_ctx,
                std::unique_ptr<graph::LowerInfoMap> lower_info_map)
-    : _model(model), _subg_seq(std::move(subg_seq)), _lower_info_map(std::move(lower_info_map))
+    : _model(model), _lower_info_map(std::move(lower_info_map))
 {
-  assert(_model && _subg_seq && _lower_info_map);
-  for (const auto &subg : *_subg_seq)
+  assert(_model && subg_ctx && _lower_info_map);
+
+  _subg_seq = nnfw::cpp14::make_unique<std::vector<std::unique_ptr<model::operation::Subgraph>>>();
+  std::vector<model::OperationIndex> index_list;
+
+  // Get SubgraphSequence by topological sorting
   {
-    // Assume that the lower_infos of all nodes on a subgraph are identified on the subgraph
-    const auto &first_ind = subg->operations()[0].index;
-    _elements.emplace_back(subg.get(), getLowerInfo(first_ind));
+    // subg_ctx can't access a subgraph by an operand so that input_to_subgs can offer it
+    std::unordered_map<model::operand::Index, model::operation::IndexList> input_to_subgs;
+
+    // Get the relations between input/subgraph to be used for dfs-post-iter
+    //
+    //      [0]             # input -> _input_to_subgraphes[0] = {SUBG0}
+    //       |
+    //     [SUBG0]
+    //       |
+    //      [1]-----.       # input -> _input_to_subgraphes[1] = {SUBG1, SUBG2}
+    //       |      |
+    //     [SUBG1] [SUBG2]
+    //       |      |
+    //      [2]    [3]      # input -> _input_to_subgraphes[2] = {SUBG3}
+    //       \      /       # input -> _input_to_subgraphes[3] = {SUBG3}
+    //       [SUBG3]
+    //         |
+    //        [4]
+    subg_ctx->iterate([&](const model::OperationIndex &subg_idx, model::operation::Subgraph &subg) {
+      for (auto input : subg.getInputs())
+      {
+        // only valid_inputs
+        const auto &operand = _model->operands.at(input);
+        if (operand.usage() == model::operand::Usage::CONSTANT)
+          continue;
+
+        auto it = input_to_subgs.find(input);
+        if (it == input_to_subgs.end())
+        {
+          model::operation::IndexList list{subg_idx};
+          input_to_subgs[input] = list;
+        }
+        else
+        {
+          it->second.append(subg_idx);
+        }
+      }
+    });
+
+    std::unordered_map<model::OperationIndex, bool> visited;
+    subg_ctx->iterate([&](const model::OperationIndex &index, const model::operation::Subgraph &) {
+      visited[index] = false;
+    });
+
+    std::function<void(const model::OperationIndex &, model::operation::Subgraph &)> dfs_recursive =
+        [&](const model::OperationIndex &index, model::operation::Subgraph &subg) -> void {
+      if (visited[index])
+        return;
+      visited[index] = true;
+
+      // The outputs should be not constants
+      for (auto output : subg.getOutputs())
+      {
+        const auto it = input_to_subgs.find(output);
+        if (it != input_to_subgs.end())
+        {
+          const auto &subg_index_list = it->second;
+          for (const auto &index : subg_index_list.list())
+          {
+            auto &subg = subg_ctx->at(index);
+            dfs_recursive(index, subg);
+          }
+        }
+      }
+
+      _subg_seq->emplace_back(std::move(subg_ctx->releaseAt(index)));
+      index_list.emplace_back(index);
+    };
+
+    subg_ctx->iterate(dfs_recursive);
+
+    assert(_subg_seq->size() == index_list.size());
+
+    // All of the nodes must have been visited.
+    assert(std::all_of(
+        visited.begin(), visited.end(),
+        [](const std::pair<const model::OperationIndex, bool> &v) { return v.second; }));
+
+    // NOTE. Now these subgraph are on the reverse order
+    std::reverse(_subg_seq->begin(), _subg_seq->end());
+    std::reverse(index_list.begin(), index_list.end());
+  }
+
+  {
+    const auto &backendToString = [](const neurun::backend::Backend *backend) {
+      assert(backend);
+      std::string str;
+      str += backend->config()->id();
+      str += " ";
+      return "{ " + str + "}";
+    };
+
+    VERBOSE(Linear) << "Final SubgraphSequence" << std::endl;
+    for (size_t i = 0; i < _subg_seq->size(); ++i)
+    {
+      const auto &subg = (*_subg_seq)[i];
+      const auto &index = index_list[i];
+      const auto lower_info = getLowerInfo(index);
+      VERBOSE(Linear) << "* SUBG#" << index.value() << " " << backendToString(lower_info->backend())
+                      << " " << subg->getStr() << std::endl;
+    }
+  }
+
+  for (size_t i = 0; i < _subg_seq->size(); ++i)
+  {
+    {
+      // Assume that the lower_infos of all nodes on a subgraph are identified on the subgraph
+      const auto &subg = (*_subg_seq)[i];
+      const auto &index = index_list[i];
+      _elements.emplace_back(subg.get(), getLowerInfo(index));
+    }
   }
 }
 
diff --git a/runtimes/neurun/core/src/linear/Linear.h b/runtimes/neurun/core/src/linear/Linear.h
index e00bf72fb..f12dd269c 100644
--- a/runtimes/neurun/core/src/linear/Linear.h
+++ b/runtimes/neurun/core/src/linear/Linear.h
@@ -21,7 +21,7 @@
 #include <memory>
 
 #include "model/Model.h"
-#include "model/operation/Subgraph.h"
+#include "model/operation/SubgraphContext.h"
 #include "backend/ITensorBuilder.h"
 #include "graph/LowerInfoMap.h"
 
@@ -57,7 +57,7 @@ class Linear
 {
 public:
   Linear(const std::shared_ptr<const model::Model> &model,
-         std::unique_ptr<std::vector<std::unique_ptr<model::operation::Subgraph>>> subg_seq,
+         std::unique_ptr<model::operation::SubgraphContext> subg_ctx,
          std::unique_ptr<graph::LowerInfoMap> lower_info_map);
 
 public: