treelite
builder.h
1 #ifndef TREELITE_COMPILER_AST_BUILDER_H_
2 #define TREELITE_COMPILER_AST_BUILDER_H_
3 
4 #include <treelite/common.h>
5 #include <treelite/tree.h>
6 #include "./ast.h"
7 
8 namespace treelite {
9 namespace compiler {
10 
11 // forward declaration
12 class ASTBuilder;
13 bool breakup(ASTNode* node, int num_descendant_limit, int* num_tu,
14  ASTBuilder* builder);
15 
16 class ASTBuilder {
17  public:
18  ASTBuilder() : output_vector_flag(false), main_node(nullptr),
19  quantize_threshold_flag(false) {}
20 
21  void Build(const Model& model);
22  void Split(int parallel_comp);
23  void QuantizeThresholds();
24  void CountDescendant();
25  void BreakUpLargeUnits(int num_descendant_limit);
26  void AnnotateBranches(const std::vector<std::vector<size_t>>& counts);
27  void Dump();
28 
29  inline const ASTNode* GetRootNode() {
30  return main_node;
31  }
32 
33  private:
34  friend bool treelite::compiler::breakup(ASTNode*, int, int*, ASTBuilder*);
35 
36  template <typename NodeType, typename ...Args>
37  NodeType* AddNode(ASTNode* parent, Args&& ...args) {
38  std::unique_ptr<NodeType> node
39  = common::make_unique<NodeType>(std::forward<Args>(args)...);
40  NodeType* ref = node.get();
41  ref->parent = parent;
42  nodes.push_back(std::move(node));
43  return ref;
44  }
45  ASTNode* WalkTree(const Tree& tree, ASTNode* parent);
46  ASTNode* WalkTree(const Tree& tree, int nid, ASTNode* parent);
47 
48  // keep tract of all nodes built so far, to prevent memory leak
49  std::vector<std::unique_ptr<ASTNode>> nodes;
50  bool output_vector_flag;
51  bool quantize_threshold_flag;
52  int num_feature;
53  ASTNode* main_node;
54 };
55 
56 } // namespace compiler
57 } // namespace treelite
58 
59 #endif // TREELITE_COMPILER_AST_BUILDER_H_
thin wrapper for tree ensemble model
Definition: tree.h:351
model structure for tree
in-memory representation of a decision tree
Definition: tree.h:19
Some useful utilities.