6 #ifndef TREELITE_COMPILER_AST_BUILDER_H_ 7 #define TREELITE_COMPILER_AST_BUILDER_H_ 15 #include <treelite/common.h> 24 struct CodeFoldingContext;
25 bool fold_code(ASTNode*, CodeFoldingContext*, ASTBuilder*);
26 bool breakup(ASTNode*,
int,
int*, ASTBuilder*);
30 ASTBuilder() : output_vector_flag(
false), main_node(
nullptr),
31 quantize_threshold_flag(
false) {}
34 void BuildAST(
const Model& model);
37 std::vector<bool> GenerateIsCategoricalArray();
52 bool FoldCode(
double magnitude_req,
bool create_new_translation_unit =
false);
57 void Split(
int parallel_comp);
59 void QuantizeThresholds();
61 void LoadDataCounts(
const std::vector<std::vector<size_t>>& counts);
65 std::string GetDump()
const;
67 inline const ASTNode* GetRootNode() {
75 template <
typename NodeType,
typename ...Args>
76 NodeType* AddNode(
ASTNode* parent, Args&& ...args) {
77 std::unique_ptr<NodeType> node
78 = common::make_unique<NodeType>(std::forward<Args>(args)...);
79 NodeType* ref = node.get();
81 nodes.push_back(std::move(node));
85 ASTNode* BuildASTFromTree(
const Tree& tree,
int tree_id,
int nid,
89 std::vector<std::unique_ptr<ASTNode>> nodes;
90 bool output_vector_flag;
91 bool quantize_threshold_flag;
94 bool random_forest_flag;
96 std::vector<bool> is_categorical;
97 std::map<std::string, std::string> model_param;
103 #endif // TREELITE_COMPILER_AST_BUILDER_H_ thin wrapper for tree ensemble model
Definition for AST classes.
in-memory representation of a decision tree