treelite
split.cc
1 #include "./builder.h"
2 
3 namespace treelite {
4 namespace compiler {
5 
6 DMLC_REGISTRY_FILE_TAG(split);
7 
8 void ASTBuilder::Split(int parallel_comp) {
9  if (parallel_comp <= 0) {
10  LOG(INFO) << "Parallel compilation disabled; all member trees will be "
11  << "dumped to a single source file. This may increase "
12  << "compilation time and memory usage.";
13  return;
14  }
15  LOG(INFO) << "Parallel compilation enabled; member trees will be "
16  << "divided into " << parallel_comp << " translation units.";
17  CHECK_EQ(this->main_node->children.size(), 1);
18  ASTNode* top_ac_node = this->main_node->children[0];
19  CHECK(dynamic_cast<AccumulatorContextNode*>(top_ac_node));
20 
21  /* tree_head[i] stores reference to head of tree i */
22  std::vector<ConditionNode*> tree_head;
23  for (ASTNode* node : top_ac_node->children) {
24  ConditionNode* tree_head_node = dynamic_cast<ConditionNode*>(node);
25  CHECK(tree_head_node);
26  tree_head.push_back(tree_head_node);
27  }
28  /* dynamic_cast<> is used here to check node types. This is to ensure
29  that we don't accidentally call Split() twice. */
30 
31  const int ntree = static_cast<int>(tree_head.size());
32  const int nunit = parallel_comp;
33  const int unit_size = (ntree + nunit - 1) / nunit;
34  std::vector<ASTNode*> tu_list; // list of translation units
35  for (int unit_id = 0; unit_id < nunit; ++unit_id) {
36  const int tree_begin = unit_id * unit_size;
37  const int tree_end = std::min((unit_id + 1) * unit_size, ntree);
38  if (tree_begin < tree_end) {
39  TranslationUnitNode* tu = AddNode<TranslationUnitNode>(top_ac_node,
40  unit_id);
41  tu_list.push_back(tu);
42  AccumulatorContextNode* ac = AddNode<AccumulatorContextNode>(tu);
43  tu->children.push_back(ac);
44  for (int tree_id = tree_begin; tree_id < tree_end; ++tree_id) {
45  ConditionNode* tree_head_node = tree_head[tree_id];
46  tree_head_node->parent = ac;
47  ac->children.push_back(tree_head_node);
48  }
49  }
50  }
51  top_ac_node->children = tu_list;
52 }
53 
54 } // namespace compiler
55 } // namespace treelite