diff --git a/include/treelite/compiler.h b/include/treelite/compiler.h index 87004021..b3fe6810 100644 --- a/include/treelite/compiler.h +++ b/include/treelite/compiler.h @@ -7,7 +7,6 @@ #ifndef TREELITE_COMPILER_H_ #define TREELITE_COMPILER_H_ -#include #include #include #include @@ -75,31 +74,6 @@ class Compiler { const char* param_json_str); }; -/*! - * \brief Registry entry for compiler - */ -struct CompilerReg - : public dmlc::FunctionRegEntryBase > { -}; - -/*! - * \brief Macro to register compiler. - * - * \code - * // example of registering the simple compiler - * TREELITE_REGISTER_COMPILER(SimpleCompiler, "simple") - * .describe("Bare-bones simple compiler") - * .set_body([]() { - * return new SimpleCompiler(); - * }); - * \endcode - */ -#define TREELITE_REGISTER_COMPILER(UniqueId, Name) \ - static DMLC_ATTRIBUTE_UNUSED ::treelite::CompilerReg & \ - __make_ ## CompilerReg ## _ ## UniqueId ## __ = \ - ::dmlc::Registry< ::treelite::CompilerReg>::Get()->__REGISTER__(Name) - } // namespace treelite #endif // TREELITE_COMPILER_H_ diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index b56222d5..e2dda71d 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -84,8 +84,10 @@ target_sources(objtreelite compiler/native/pred_transform.h compiler/native/qnode_template.h compiler/native/typeinfo_ctypes.h + compiler/ast_native.h compiler/ast_native.cc compiler/compiler.cc + compiler/failsafe.h compiler/failsafe.cc compiler/pred_transform.cc compiler/pred_transform.h diff --git a/src/compiler/ast_native.cc b/src/compiler/ast_native.cc index cbd8933e..62902d1c 100644 --- a/src/compiler/ast_native.cc +++ b/src/compiler/ast_native.cc @@ -1,8 +1,8 @@ /*! - * Copyright (c) 2017-2020 by Contributors + * Copyright (c) 2017-2021 by Contributors * \file ast_native.cc - * \author Hyunsu Cho * \brief C code generator + * \author Hyunsu Cho */ #include #include @@ -18,6 +18,7 @@ #include #include #include +#include "./ast_native.h" #include "./pred_transform.h" #include "./ast/builder.h" #include "./native/main_template.h" @@ -39,20 +40,9 @@ using namespace fmt::literals; namespace treelite { namespace compiler { -DMLC_REGISTRY_FILE_TAG(ast_native); - -class ASTNativeCompiler : public Compiler { +class ASTNativeCompilerImpl { public: - explicit ASTNativeCompiler(const CompilerParam& param) - : param(param) { - if (param.verbose > 0) { - LOG(INFO) << "Using ASTNativeCompiler"; - } - if (param.dump_array_as_elf > 0) { - LOG(INFO) << "Warning: 'dump_array_as_elf' parameter is not applicable " - "for ASTNativeCompiler"; - } - } + explicit ASTNativeCompilerImpl(const CompilerParam& param) : param_(param) {} template CompiledModel CompileImpl(const ModelImpl& model) { @@ -74,22 +64,22 @@ class ASTNativeCompiler : public Compiler { ASTBuilder builder; builder.BuildAST(model); - if (builder.FoldCode(param.code_folding_req) || param.quantize > 0) { + if (builder.FoldCode(param_.code_folding_req) || param_.quantize > 0) { // is_categorical[i] : is i-th feature categorical? array_is_categorical_ = RenderIsCategoricalArray(builder.GenerateIsCategoricalArray()); } - if (param.annotate_in != "NULL") { + if (param_.annotate_in != "NULL") { BranchAnnotator annotator; - std::ifstream fi(param.annotate_in.c_str()); + std::ifstream fi(param_.annotate_in.c_str()); annotator.Load(fi); const auto annotation = annotator.Get(); builder.LoadDataCounts(annotation); LOG(INFO) << "Loading node frequencies from `" - << param.annotate_in << "'"; + << param_.annotate_in << "'"; } - builder.Split(param.parallel_comp); - if (param.quantize > 0) { + builder.Split(param_.parallel_comp); + if (param_.quantize > 0) { builder.QuantizeThresholds(); } @@ -113,7 +103,7 @@ class ASTNativeCompiler : public Compiler { writer.StartObject(); writer.Key("target"); - writer.String(param.native_lib_name.data(), param.native_lib_name.size()); + writer.String(param_.native_lib_name.data(), param_.native_lib_name.size()); writer.Key("sources"); writer.StartArray(); for (const auto& kv : files_) { @@ -138,7 +128,7 @@ class ASTNativeCompiler : public Compiler { return cm; } - CompiledModel Compile(const Model& model) override { + CompiledModel Compile(const Model& model) { CHECK(model.GetLeafOutputType() != TypeInfo::kUInt32) << "Integer leaf outputs not yet supported"; this->pred_tranform_func_ = PredTransformFunction("native", model); @@ -147,12 +137,12 @@ class ASTNativeCompiler : public Compiler { }); } - CompilerParam QueryParam() const override { - return param; + CompilerParam QueryParam() const { + return param_; } private: - CompilerParam param; + CompilerParam param_; int num_feature_; TaskType task_type_; TaskParam task_param_; @@ -255,7 +245,7 @@ class ASTNativeCompiler : public Compiler { "predict_function_signature"_a = predict_function_signature, "query_functions_prototype"_a = query_functions_prototype, "threshold_type"_a = threshold_type, - "threshold_type_Node"_a = (param.quantize > 0 ? std::string("int") : threshold_type)), + "threshold_type_Node"_a = (param_.quantize > 0 ? std::string("int") : threshold_type)), indent); CHECK_EQ(node->children.size(), 1); @@ -522,11 +512,11 @@ class ASTNativeCompiler : public Compiler { std::string output_switch_statement; Operator common_comp_op; - common_util::RenderCodeFolderArrays(node, param.quantize, false, - "{{ {default_left}, {split_index}, {threshold}, {left_child}, {right_child} }}", + common_util::RenderCodeFolderArrays(node, param_.quantize, + false, "{{ {default_left}, {split_index}, {threshold}, {left_child}, {right_child} }}", [this](const OutputNode* node) { return RenderOutputStatement(node); }, - &array_nodes, &array_cat_bitmap, &array_cat_begin, - &output_switch_statement, &common_comp_op); + &array_nodes, &array_cat_bitmap, &array_cat_begin, &output_switch_statement, + &common_comp_op); if (!array_nodes.empty()) { AppendToBuffer("header.h", fmt::format("extern const struct Node {node_array_name}[];\n", @@ -576,7 +566,7 @@ class ASTNativeCompiler : public Compiler { "node_array_name"_a = node_array_name, "cat_bitmap_name"_a = cat_bitmap_name, "cat_begin_name"_a = cat_begin_name, - "data_field"_a = (param.quantize > 0 ? "qvalue" : "fvalue"), + "data_field"_a = (param_.quantize > 0 ? "qvalue" : "fvalue"), "comp_op"_a = OpName(common_comp_op), "output_switch_statement"_a = output_switch_statement), indent); @@ -584,7 +574,7 @@ class ASTNativeCompiler : public Compiler { AppendToBuffer(dest, fmt::format(native::eval_loop_template_without_categorical_feature, "node_array_name"_a = node_array_name, - "data_field"_a = (param.quantize > 0 ? "qvalue" : "fvalue"), + "data_field"_a = (param_.quantize > 0 ? "qvalue" : "fvalue"), "comp_op"_a = OpName(common_comp_op), "output_switch_statement"_a = output_switch_statement), indent); @@ -708,10 +698,28 @@ class ASTNativeCompiler : public Compiler { } }; -TREELITE_REGISTER_COMPILER(ASTNativeCompiler, "ast_native") -.describe("AST-based compiler that produces C code") -.set_body([](const CompilerParam& param) -> Compiler* { - return new ASTNativeCompiler(param); - }); +ASTNativeCompiler::ASTNativeCompiler(const CompilerParam& param) + : pimpl_(std::make_unique(param)) { + if (param.verbose > 0) { + LOG(INFO) << "Using ASTNativeCompiler"; + } + if (param.dump_array_as_elf > 0) { + LOG(INFO) << "Warning: 'dump_array_as_elf' parameter is not applicable " + "for ASTNativeCompiler"; + } +} + +ASTNativeCompiler::~ASTNativeCompiler() = default; + +CompiledModel +ASTNativeCompiler::Compile(const Model &model) { + return pimpl_->Compile(model); +} + +CompilerParam +ASTNativeCompiler::QueryParam() const { + return pimpl_->QueryParam(); +} + } // namespace compiler } // namespace treelite diff --git a/src/compiler/ast_native.h b/src/compiler/ast_native.h new file mode 100644 index 00000000..117b91c4 --- /dev/null +++ b/src/compiler/ast_native.h @@ -0,0 +1,33 @@ +/*! + * Copyright (c) 2021 by Contributors + * \file ast_native.h + * \brief C code generator + * \author Hyunsu Cho + */ + +#ifndef TREELITE_COMPILER_AST_NATIVE_H_ +#define TREELITE_COMPILER_AST_NATIVE_H_ + +#include +#include +#include + +namespace treelite { +namespace compiler { + +class ASTNativeCompilerImpl; + +class ASTNativeCompiler : public Compiler { + public: + explicit ASTNativeCompiler(const CompilerParam& param); + virtual ~ASTNativeCompiler(); + CompiledModel Compile(const Model& model) override; + CompilerParam QueryParam() const override; + private: + std::unique_ptr pimpl_; +}; + +} // namespace compiler +} // namespace treelite + +#endif // TREELITE_COMPILER_AST_NATIVE_H_ diff --git a/src/compiler/compiler.cc b/src/compiler/compiler.cc index 8dd074d5..3af2ae08 100644 --- a/src/compiler/compiler.cc +++ b/src/compiler/compiler.cc @@ -5,22 +5,23 @@ */ #include #include -#include +#include #include #include - -namespace dmlc { -DMLC_REGISTRY_ENABLE(::treelite::CompilerReg); -} // namespace dmlc +#include "./ast_native.h" +#include "./failsafe.h" namespace treelite { Compiler* Compiler::Create(const std::string& name, const char* param_json_str) { compiler::CompilerParam param = compiler::CompilerParam::ParseFromJSON(param_json_str); - auto *e = ::dmlc::Registry< ::treelite::CompilerReg>::Get()->Find(name); - if (e == nullptr) { - LOG(FATAL) << "Unknown compiler type " << name; + if (name == "ast_native") { + return new compiler::ASTNativeCompiler(param); + } else if (name == "failsafe") { + return new compiler::FailSafeCompiler(param); + } else { + LOG(FATAL) << "Unrecognized compiler '" << name << "'"; + return nullptr; } - return (e->body)(param); } namespace compiler { @@ -75,9 +76,6 @@ CompilerParam::ParseFromJSON(const char* param_json_str) { return param; } -// List of files that will be force linked in static links. -DMLC_REGISTRY_LINK_TAG(ast_native); -DMLC_REGISTRY_LINK_TAG(failsafe); } // namespace compiler } // namespace treelite diff --git a/src/compiler/failsafe.cc b/src/compiler/failsafe.cc index a31540b7..ef196cbb 100644 --- a/src/compiler/failsafe.cc +++ b/src/compiler/failsafe.cc @@ -1,14 +1,15 @@ /*! - * Copyright (c) 2019-2020 by Contributors + * Copyright (c) 2019-2021 by Contributors * \file failsafe.cc - * \author Hyunsu Cho * \brief C code generator (fail-safe). The generated code will mimic prediction logic found in * XGBoost + * \author Hyunsu Cho */ #include #include #include +#include #include #include #include @@ -17,6 +18,7 @@ #include #include #include +#include "./failsafe.h" #include "./pred_transform.h" #include "./common/format_util.h" #include "./elf/elf_formatter.h" @@ -234,34 +236,11 @@ inline bool EndsWith(const std::string& str, const std::string& suffix) { namespace treelite { namespace compiler { -DMLC_REGISTRY_FILE_TAG(failsafe); - -class FailSafeCompiler : public Compiler { +class FailSafeCompilerImpl { public: - explicit FailSafeCompiler(const CompilerParam& param) - : param(param) { - if (param.verbose > 0) { - LOG(INFO) << "Using FailSafeCompiler"; - } - if (param.annotate_in != "NULL") { - LOG(INFO) << "Warning: 'annotate_in' parameter is not applicable for " - "FailSafeCompiler"; - } - if (param.quantize > 0) { - LOG(INFO) << "Warning: 'quantize' parameter is not applicable for " - "FailSafeCompiler"; - } - if (param.parallel_comp > 0) { - LOG(INFO) << "Warning: 'parallel_comp' parameter is not applicable for " - "FailSafeCompiler"; - } - if (std::isfinite(param.code_folding_req)) { - LOG(INFO) << "Warning: 'code_folding_req' parameter is not applicable " - "for FailSafeCompiler"; - } - } + explicit FailSafeCompilerImpl(const CompilerParam& param) : param_(param) {} - CompiledModel Compile(const Model& model_ptr) override { + CompiledModel Compile(const Model& model_ptr) { CHECK(model_ptr.GetThresholdType() == TypeInfo::kFloat32 && model_ptr.GetLeafOutputType() == TypeInfo::kFloat32) << "Failsafe compiler only supports models with float32 thresholds and float32 leaf outputs"; @@ -312,8 +291,8 @@ class FailSafeCompiler : public Compiler { std::string nodes, nodes_row_ptr; std::vector nodes_elf; - if (param.dump_array_as_elf > 0) { - if (param.verbose > 0) { + if (param_.dump_array_as_elf > 0) { + if (param_.verbose > 0) { LOG(INFO) << "Dumping arrays as an ELF relocatable object..."; } std::tie(nodes_elf, nodes_row_ptr) = FormatNodesArrayELF(model); @@ -345,7 +324,7 @@ class FailSafeCompiler : public Compiler { files_["main.c"] = CompiledModel::FileEntry(main_program.str()); - if (param.dump_array_as_elf > 0) { + if (param_.dump_array_as_elf > 0) { files_["arrays.o"] = CompiledModel::FileEntry(std::move(nodes_elf)); } else { files_["arrays.c"] = CompiledModel::FileEntry(fmt::format(arrays_template, @@ -367,7 +346,7 @@ class FailSafeCompiler : public Compiler { writer.StartObject(); writer.Key("target"); - writer.String(param.native_lib_name.data(), param.native_lib_name.size()); + writer.String(param_.native_lib_name.data(), param_.native_lib_name.size()); writer.Key("sources"); writer.StartArray(); std::vector extra_file_list; @@ -403,22 +382,52 @@ class FailSafeCompiler : public Compiler { return cm; } - CompilerParam QueryParam() const override { - return param; + CompilerParam QueryParam() const { + return param_; } private: - CompilerParam param; + CompilerParam param_; int num_feature_; unsigned int num_class_; std::string pred_tranform_func_; std::unordered_map files_; }; -TREELITE_REGISTER_COMPILER(FailSafeCompiler, "failsafe") -.describe("Simple compiler to express trees as a tight for-loop") -.set_body([](const CompilerParam& param) -> Compiler* { - return new FailSafeCompiler(param); - }); +FailSafeCompiler::FailSafeCompiler(const CompilerParam& param) + : pimpl_(std::make_unique(param)) { + if (param.verbose > 0) { + LOG(INFO) << "Using FailSafeCompiler"; + } + if (param.annotate_in != "NULL") { + LOG(INFO) << "Warning: 'annotate_in' parameter is not applicable for " + "FailSafeCompiler"; + } + if (param.quantize > 0) { + LOG(INFO) << "Warning: 'quantize' parameter is not applicable for " + "FailSafeCompiler"; + } + if (param.parallel_comp > 0) { + LOG(INFO) << "Warning: 'parallel_comp' parameter is not applicable for " + "FailSafeCompiler"; + } + if (std::isfinite(param.code_folding_req)) { + LOG(INFO) << "Warning: 'code_folding_req' parameter is not applicable " + "for FailSafeCompiler"; + } +} + +FailSafeCompiler::~FailSafeCompiler() = default; + +CompiledModel +FailSafeCompiler::Compile(const Model& model) { + return pimpl_->Compile(model); +} + +CompilerParam +FailSafeCompiler::QueryParam() const { + return pimpl_->QueryParam(); +} + } // namespace compiler } // namespace treelite diff --git a/src/compiler/failsafe.h b/src/compiler/failsafe.h new file mode 100644 index 00000000..09aa86f8 --- /dev/null +++ b/src/compiler/failsafe.h @@ -0,0 +1,34 @@ +/*! + * Copyright (c) 2021 by Contributors + * \file failsafe.h + * \brief C code generator (fail-safe). The generated code will mimic prediction logic found in + * XGBoost + * \author Hyunsu Cho + */ + +#ifndef TREELITE_COMPILER_FAILSAFE_H_ +#define TREELITE_COMPILER_FAILSAFE_H_ + +#include +#include +#include + +namespace treelite { +namespace compiler { + +class FailSafeCompilerImpl; + +class FailSafeCompiler : public Compiler { + public: + explicit FailSafeCompiler(const CompilerParam& param); + virtual ~FailSafeCompiler(); + CompiledModel Compile(const Model& model) override; + CompilerParam QueryParam() const override; + private: + std::unique_ptr pimpl_; +}; + +} // namespace compiler +} // namespace treelite + +#endif // TREELITE_COMPILER_FAILSAFE_H_