treelite
breakup.cc
1 #include "./builder.h"
2 
3 namespace treelite {
4 namespace compiler {
5 
6 DMLC_REGISTRY_FILE_TAG(breakup);
7 
8 static int count_tu(ASTNode* node) {
9  int accum = (dynamic_cast<TranslationUnitNode*>(node)) ? 1 : 0;
10  for (ASTNode* child : node->children) {
11  accum += count_tu(child);
12  }
13  return accum;
14 }
15 
16 bool breakup(ASTNode* node, int num_descendant_limit, int* num_tu,
17  ASTBuilder* builder) {
18  bool flag = false;
19  if (dynamic_cast<ConditionNode*>(node)
20  && node->num_descendant > num_descendant_limit) {
21  bool break_here = true;
22  for (ASTNode* child : node->children) {
23  if (child->num_descendant > num_descendant_limit) {
24  break_here = false; // don't break this node; break the child instead
25  }
26  }
27  if (break_here) {
28  ASTNode* parent = node->parent;
29 
30  int node_idx = -1;
31  for (size_t i = 0; i < parent->children.size(); ++i) {
32  if (parent->children[i] == node) {
33  node_idx = static_cast<int>(i);
34  break;
35  }
36  }
37  CHECK_GE(node_idx, 0);
38 
39  const int unit_id = (*num_tu)++;
40  TranslationUnitNode* tu
41  = builder->AddNode<TranslationUnitNode>(parent, unit_id);
42  AccumulatorContextNode* ac
43  = builder->AddNode<AccumulatorContextNode>(tu);
44  parent->children[node_idx] = tu;
45  tu->children.push_back(ac);
46  ac->children.push_back(node);
47  node->parent = ac;
48  tu->num_descendant = 0;
49  ASTNode* n = tu->parent;
50  while (n) {
51  n->num_descendant -= node->num_descendant;
52  CHECK_GE(n->num_descendant, 0);
53  n = n->parent;
54  }
55  flag = true;
56  }
57  }
58  for (ASTNode* child : node->children) {
59  if (child->num_descendant > 0) {
60  flag |= breakup(child, num_descendant_limit, num_tu, builder);
61  }
62  }
63  return flag;
64 }
65 
66 void ASTBuilder::BreakUpLargeUnits(int num_descendant_limit) {
67  CHECK_GT(num_descendant_limit, 0);
68  int num_tu = count_tu(this->main_node);
69  while (breakup(this->main_node, num_descendant_limit, &num_tu, this)) {}
70 }
71 
72 } // namespace compiler
73 } // namespace treelite