[onert] Configure context_data_map using the backend info of lgraph

[jyoungyun]

This commit configures the context_data_map using the backend
information of lgraph.

ONE-DCO-1.0-Signed-off-by: Jiyoung Yun [email protected]
Co-authored-by: Hyeongseok Oh [email protected]

[zetwhite]

As i understood correctly, the difference is here :

Before

create backend from backend_manager.getAll()

After

create backend from CompilerOptions config

Is there a difference between backend_manager.getAll() and config lists?

[jyoungyun]

Yes, it can be different.
A scenario like this is as follows:

flowchart TD
   create_session --> id0["set_available_backends('train')"] -- "BACKENDS config: ['train']\nbackend_manager: ['train']" --> train_prepare --> train --> id1["set_available_backends('cpu')"] -- "BACKENDS config: ['cpu']\nbackend_manager: ['train', 'cpu']" --> prepare --> run --> close_session

Loading

[hseok-oh]

I think backend configuration has already been reflected in the lowered graph. So I suggest to use lgraph iteration instead of getting backend list

@@ -160,18 +160,17 @@ createBackendContexts(compiler::ILoweredGraph &lgraph, bool linear_executor,
                       std::shared_ptr<backend::custom::IKernelBuilder> custom_kernel_builder)
 {
   backend::BackendContexts contexts;
-  auto &backend_manager = compiler::BackendManager::get();
-
   std::unordered_map<const backend::Backend *, backend::ContextData> context_data_map;
 
-  // Generate partial graphs for each backend
-  for (auto &&backend : backend_manager.getAll())
-  {
+  auto init_context_data = [&](const backend::Backend *backend) {
     auto &data = context_data_map[backend];
     auto graph = std::make_unique<ir::Graph>();
     graph->setLayout(lgraph.graph().layout());
     data.graph = std::move(graph);
-  }
+  };
 
   auto &whole_graph = lgraph.graph();
   // Separate operands into partial graphs
@@ -182,6 +181,9 @@ createBackendContexts(compiler::ILoweredGraph &lgraph, bool linear_executor,
       return;
     const auto &def_factor = def_factors.getOnlyElement();
     const auto backend = def_factor.backend();
+    if (context_data_map.find(backend) == context_data_map.end())
+      init_context_data(backend);
+
     auto &partial_graph = *context_data_map[backend].graph;
     auto &operand_layouts = context_data_map[backend].operand_layouts;
     assert(operand_layouts.find(operand_ind) == operand_layouts.end());
@@ -200,6 +202,9 @@ createBackendContexts(compiler::ILoweredGraph &lgraph, bool linear_executor,
     [&](const ir::OperationIndex &op_ind, const ir::IOperation &operation) {
       auto &op_li = lgraph.lower_info().operation;
       auto backend = op_li.at(op_ind).backend();
+      if (context_data_map.find(backend) == context_data_map.end())
+        init_context_data(backend);
+
       auto &partial_graph = *context_data_map[backend].graph;
       auto &external_operands = context_data_map[backend].external_operands;
       auto &operand_layouts = context_data_map[backend].operand_layouts;

[YongseopKim]

I'll review after @jyoungyun responds to @hseok-oh 's comment.

[jyoungyun]

@hseok-oh Your suggestion looks good. The backend of each operation and operand has already been specified. So we can use the backend information of each operation and operand. One concern is that each operation and operand must find a maching backend each time. This is done in the preparation stage, so it is not time-sensitive. However, if possible, I'd like to define the context_data_map in advance. I will take a closer look. :)

[jyoungyun]

The context_data_map is an unordered_map. And the time complexity of finding key is just O(1) on average. So I don't think this will affect the performance! I will fix this PR. :)

[YongseopKim]

LGTM! It works well on my draft.

[hseok-oh]

LGTM

[jyoungyun]

/cc @zetwhite