diff --git a/cpp/include/nvforest/detail/bitset.hpp b/cpp/include/nvforest/detail/bitset.hpp index 3cb5963..c11cdee 100644 --- a/cpp/include/nvforest/detail/bitset.hpp +++ b/cpp/include/nvforest/detail/bitset.hpp @@ -20,6 +20,9 @@ struct bitset { using storage_type = storage_t; using index_type = index_t; + // Ensure that index_t is unsigned. Bound checks below rely on index_t being unsigned + static_assert(std::is_unsigned_v, "index_t must be unsigned"); + auto constexpr static const bin_width = index_type(sizeof(storage_type) * 8); HOST DEVICE bitset() : data_{nullptr}, num_bits_{0} {} @@ -37,12 +40,13 @@ struct bitset { // Standard bit-wise mutators and accessor HOST DEVICE auto& set(index_type index) { - data_[bin_from_index(index)] |= mask_in_bin(index); + // Guard against OOB writes; silently ignored to preserve memory safety + if (index < num_bits_) { data_[bin_from_index(index)] |= mask_in_bin(index); } return *this; } HOST DEVICE auto& clear(index_type index) { - data_[bin_from_index(index)] &= ~mask_in_bin(index); + if (index < num_bits_) { data_[bin_from_index(index)] &= ~mask_in_bin(index); } return *this; } HOST DEVICE auto test(index_type index) const diff --git a/cpp/include/nvforest/detail/decision_forest_builder.hpp b/cpp/include/nvforest/detail/decision_forest_builder.hpp index 4b5d31a..8d33ed7 100644 --- a/cpp/include/nvforest/detail/decision_forest_builder.hpp +++ b/cpp/include/nvforest/detail/decision_forest_builder.hpp @@ -18,24 +18,56 @@ #include #include #include +#include +#include +#include #include #include +#include +#include +#include +#include #include namespace nvforest::detail { -/* - * Exception indicating that nvForest model could not be built from given input - */ -struct model_builder_error : std::exception { - model_builder_error() : model_builder_error("Error while building model") {} - model_builder_error(char const* msg) : msg_{msg} {} - virtual char const* what() const noexcept { return msg_; } +struct floating_point_truncation_error : std::exception { + floating_point_truncation_error() = default; + floating_point_truncation_error(std::string msg) : msg_{std::move(msg)} {} + floating_point_truncation_error(char const* msg) : msg_{msg} {} + char const* what() const noexcept override { return msg_.c_str(); } private: - char const* msg_; + std::string msg_; }; +template +dest_t safe_cast_floating_point(src_t x) +{ + static_assert(std::is_floating_point_v && std::is_floating_point_v, + "Source and destination types must be both floating-point types."); + if constexpr (sizeof(dest_t) >= sizeof(src_t)) { + return static_cast(x); + } else { + if (!std::isfinite(x)) { + throw floating_point_truncation_error{"Cannot cast an INF or NaN value"}; + } + auto constexpr lower_limit = static_cast(std::numeric_limits::lowest()); + auto constexpr upper_limit = static_cast(std::numeric_limits::max()); + if (x < lower_limit) { + std::ostringstream ss; + ss << "Input must be at least " << lower_limit << "."; + throw floating_point_truncation_error{ss.str()}; + } + if (x > upper_limit) { + std::ostringstream ss; + ss << "Input must be at most " << upper_limit << "."; + throw floating_point_truncation_error{ss.str()}; + } + return static_cast(x); + } +} + /* * Struct used to build nvForest forests */ @@ -55,20 +87,40 @@ struct decision_forest_builder { typename node_type::metadata_storage_type feature = typename node_type::metadata_storage_type{}, typename node_type::offset_type offset = typename node_type::offset_type{}) { - auto constexpr const bin_width = index_type(sizeof(typename node_type::index_type) * 8); - auto node_value = typename node_type::index_type{}; - auto set_storage = &node_value; - auto max_node_categories = - (vec_begin != vec_end) ? *std::max_element(vec_begin, vec_end) + 1 : 1; + auto constexpr const bin_width = + typename node_type::index_type{sizeof(typename node_type::index_type) * 8}; + auto node_value = typename node_type::index_type{}; + auto set_storage = &node_value; + + using cat_t = typename std::iterator_traits::value_type; + using index_t = typename node_type::index_type; + static_assert(std::is_unsigned_v, "Category value must be an unsigned integer type"); + static_assert(std::is_same_v || std::is_same_v, + "Index type in tree node must be either uint32_t or uint64_t"); + + auto max_cat = (vec_begin != vec_end) ? *std::max_element(vec_begin, vec_end) : cat_t{0}; + auto const max_index = static_cast(std::numeric_limits::max()); + auto const max_bitset_size = + static_cast(std::numeric_limits::max()); + auto const cat_unsigned = static_cast(max_cat); + auto const max_representable = std::min(max_index, max_bitset_size); + if (cat_unsigned >= max_representable) { + throw model_import_error{std::string{"Category index must be at most "} + + std::to_string(max_representable - 1)}; + } + auto max_cat_plus_one = static_cast(cat_unsigned + std::uintmax_t{1}); + if (max_num_categories_ != index_type{} && static_cast(max_cat_plus_one) > + static_cast(max_num_categories_)) { + throw model_import_error{"Category index exceeds configured max_num_categories"}; + } if (max_num_categories_ > bin_width) { - // TODO(wphicks): Check for overflow here node_value = categorical_storage_.size(); - auto bins_required = raft_proto::ceildiv(max_node_categories, bin_width); - categorical_storage_.push_back(max_node_categories); + auto bins_required = raft_proto::ceildiv(max_cat_plus_one, bin_width); + categorical_storage_.push_back(max_cat_plus_one); categorical_storage_.resize(categorical_storage_.size() + bins_required); set_storage = &(categorical_storage_[node_value + 1]); } - auto set = bitset{set_storage, max_node_categories}; + auto set = bitset{set_storage, max_cat_plus_one}; std::for_each(vec_begin, vec_end, [&set](auto&& cat_index) { set.set(cat_index); }); add_node( @@ -151,7 +203,7 @@ struct decision_forest_builder { void set_output_size(index_type val) { if (output_size_ != index_type{1} && output_size_ != val) { - throw model_import_error("Inconsistent leaf vector size"); + throw model_import_error{"Inconsistent leaf vector size"}; } output_size_ = val; } @@ -183,11 +235,65 @@ struct decision_forest_builder { // Set device = -1 when loading the model onto CPU if (mem_type == raft_proto::device_type::cpu) { device = -1; } - // Allow narrowing for preprocessing constants. They are stored as doubles - // for consistency in the builder but must be converted to the proper types - // for the concrete forest model. -#pragma GCC diagnostic push -#pragma GCC diagnostic ignored "-Wnarrowing" + // Validate forest invariants the inference kernel relies on. After this + // function returns, the forest is treated as trusted by the kernel. + if (root_node_indexes_.size() > std::numeric_limits::max()) { + throw model_import_error{std::string{"Forest has "} + + std::to_string(root_node_indexes_.size()) + + " trees, which exceeds the maximum representable in index_type (" + + std::to_string(std::numeric_limits::max()) + ")"}; + } + + for (auto i = std::size_t{0}; i < root_node_indexes_.size(); ++i) { + if (root_node_indexes_[i] >= nodes_.size()) { + throw model_import_error{ + std::string{"Tree "} + std::to_string(i) + ": root node index out of bounds (" + + std::to_string(root_node_indexes_[i]) + " >= " + std::to_string(nodes_.size()) + ")"}; + } + } + + auto constexpr const cat_bin_width = + typename node_type::index_type{sizeof(typename node_type::index_type) * 8}; + if (max_num_categories_ > cat_bin_width) { + auto const storage_size = categorical_storage_.size(); + for (auto i = std::size_t{0}; i < nodes_.size(); ++i) { + auto const& n = nodes_[i]; + if (n.is_leaf() || !n.is_categorical()) { continue; } + auto const offset = n.index(); + + if (offset >= storage_size) { + throw model_import_error{std::string{"Categorical node "} + std::to_string(i) + + ": storage offset out of bounds (" + std::to_string(offset) + + " >= " + std::to_string(storage_size) + ")"}; + } + auto const stored_num_cats = categorical_storage_[offset]; + auto const bins_required = raft_proto::ceildiv(stored_num_cats, cat_bin_width); + auto const bits_begin = static_cast(offset) + std::size_t{1}; + auto const bits_end = bits_begin + static_cast(bins_required); + if (bits_end > storage_size) { + throw model_import_error{std::string{"Categorical node "} + std::to_string(i) + + ": bitset extends past categorical_storage end"}; + } + } + } + + auto average_factor_casted = typename node_type::threshold_type{}; + auto postproc_constant_casted = typename node_type::threshold_type{}; + try { + average_factor_casted = + safe_cast_floating_point(average_factor_); + } catch (const floating_point_truncation_error& e) { + throw model_import_error{std::string{"Found an invalid value for averaging factor: "} + + e.what()}; + } + try { + postproc_constant_casted = + safe_cast_floating_point(postproc_constant_); + } catch (const floating_point_truncation_error& e) { + throw model_import_error{std::string{"Found an invalid value for postprocessing constant: "} + + e.what()}; + } + return decision_forest_t{ raft_proto::buffer{ raft_proto::buffer{nodes_.data(), nodes_.size()}, mem_type, device, stream}, @@ -220,9 +326,8 @@ struct decision_forest_builder { output_size_, row_postproc_, element_postproc_, - static_cast(average_factor_), - static_cast(postproc_constant_)}; -#pragma GCC diagnostic pop + average_factor_casted, + postproc_constant_casted}; } private: diff --git a/cpp/include/nvforest/detail/degenerate_trees.hpp b/cpp/include/nvforest/detail/degenerate_trees.hpp index 861a580..f6ecf49 100644 --- a/cpp/include/nvforest/detail/degenerate_trees.hpp +++ b/cpp/include/nvforest/detail/degenerate_trees.hpp @@ -17,7 +17,7 @@ namespace nvforest::detail { // This function returns a modified copy of a given Treelite model if it contains // at least one degenerate tree (a single root node with no child). // If the model contains no degenerate tree, then the function returns nullptr. -std::unique_ptr convert_degenerate_trees(treelite::Model const& tl_model) +inline std::unique_ptr convert_degenerate_trees(treelite::Model const& tl_model) { bool contains_degenerate = nvforest::integration::tree_accumulate(tl_model, false, [](auto&& contains, auto&& tree) { diff --git a/cpp/include/nvforest/detail/node.hpp b/cpp/include/nvforest/detail/node.hpp index 985a246..5cdf73f 100644 --- a/cpp/include/nvforest/detail/node.hpp +++ b/cpp/include/nvforest/detail/node.hpp @@ -6,9 +6,11 @@ #include #include +#include #include #include +#include #include namespace nvforest { @@ -103,14 +105,15 @@ struct alignas(detail::get_node_alignment, "metadata storage must be unsigned"); + if (feature > FEATURE_MASK) { + throw model_import_error{std::string{"The 'feature' value in the node must be at most "} + + std::to_string(FEATURE_MASK) + "."}; + } + return metadata_storage_type( (is_leaf_node << LEAF_BIT) + (default_to_distant_child << DEFAULT_DISTANT_BIT) + (is_categorical_node << CATEGORICAL_BIT) + (feature & FEATURE_MASK)); diff --git a/cpp/include/nvforest/detail/raft_proto/ceildiv.hpp b/cpp/include/nvforest/detail/raft_proto/ceildiv.hpp index 2a1cf3e..8f4fb6c 100644 --- a/cpp/include/nvforest/detail/raft_proto/ceildiv.hpp +++ b/cpp/include/nvforest/detail/raft_proto/ceildiv.hpp @@ -5,10 +5,13 @@ #pragma once #include +#include + namespace raft_proto { template HOST DEVICE auto constexpr ceildiv(T dividend, U divisor) { - return (dividend + divisor - T{1}) / divisor; + static_assert(std::is_integral_v && std::is_integral_v, "Arguments must be integers"); + return dividend / divisor + (dividend % divisor != 0); } } // namespace raft_proto diff --git a/cpp/include/nvforest/exceptions.hpp b/cpp/include/nvforest/exceptions.hpp index 3074934..6feebba 100644 --- a/cpp/include/nvforest/exceptions.hpp +++ b/cpp/include/nvforest/exceptions.hpp @@ -5,25 +5,16 @@ #pragma once #include #include +#include namespace nvforest { -/** Exception indicating model is incompatible with nvForest */ -struct unusable_model_exception : std::exception { - unusable_model_exception() : msg_{"Model is not compatible with nvForest"} {} - unusable_model_exception(std::string msg) : msg_{msg} {} - unusable_model_exception(char const* msg) : msg_{msg} {} - virtual char const* what() const noexcept { return msg_.c_str(); } - - private: - std::string msg_; -}; - /** Exception indicating model import failed */ struct model_import_error : std::exception { model_import_error() : model_import_error("Error while importing model") {} - model_import_error(std::string msg) : msg_{msg} {} - virtual char const* what() const noexcept { return msg_.c_str(); } + model_import_error(std::string msg) : msg_{std::move(msg)} {} + model_import_error(char const* msg) : msg_{msg} {} + char const* what() const noexcept override { return msg_.c_str(); } private: std::string msg_; diff --git a/cpp/include/nvforest/treelite_importer.hpp b/cpp/include/nvforest/treelite_importer.hpp index ccbcc13..6a9da79 100644 --- a/cpp/include/nvforest/treelite_importer.hpp +++ b/cpp/include/nvforest/treelite_importer.hpp @@ -257,37 +257,47 @@ struct treelite_importer { tl_model, [&builder, &offsets, &node_index]( auto&& tree_id, auto&& node, auto&& depth, auto&& parent_index) { - if (node.is_leaf()) { - auto output = node.get_output(); - builder.set_output_size(output.size()); - if (output.size() > index_type{1}) { - builder.add_leaf_vector_node( - std::begin(output), std::end(output), node.get_treelite_id(), depth); + try { + if (node.is_leaf()) { + auto output = node.get_output(); + if (output.empty()) { + throw model_import_error{"Leaf node must contain at least one output value"}; + } + builder.set_output_size(output.size()); + if (output.size() > index_type{1}) { + builder.add_leaf_vector_node( + std::begin(output), std::end(output), node.get_treelite_id(), depth); + } else { + builder.add_node(typename forest_model_t::io_type(output[0]), + node.get_treelite_id(), + depth, + true); + } } else { - builder.add_node( - typename forest_model_t::io_type(output[0]), node.get_treelite_id(), depth, true); - } - } else { - if (node.is_categorical()) { - auto categories = node.get_categories(); - builder.add_categorical_node(std::begin(categories), - std::end(categories), - node.get_treelite_id(), - depth, - node.default_distant(), - node.get_feature(), - offsets[node_index]); - } else { - builder.add_node(typename forest_model_t::threshold_type(node.threshold()), - node.get_treelite_id(), - depth, - false, - node.default_distant(), - false, - node.get_feature(), - offsets[node_index], - node.is_inclusive()); + if (node.is_categorical()) { + auto categories = node.get_categories(); + builder.add_categorical_node(std::begin(categories), + std::end(categories), + node.get_treelite_id(), + depth, + node.default_distant(), + node.get_feature(), + offsets[node_index]); + } else { + builder.add_node(typename forest_model_t::threshold_type(node.threshold()), + node.get_treelite_id(), + depth, + false, + node.default_distant(), + false, + node.get_feature(), + offsets[node_index], + node.is_inclusive()); + } } + } catch (const model_import_error& e) { + throw model_import_error{std::string{"Tree "} + std::to_string(tree_id) + ", Node " + + std::to_string(node.get_treelite_id()) + ": " + e.what()}; } ++node_index; }); @@ -438,13 +448,14 @@ struct treelite_importer { * @param stream The CUDA stream to use for loading this model (can be * omitted for CPU). */ -auto import_from_treelite_model(treelite::Model const& tl_model, - tree_layout layout = preferred_tree_layout, - index_type align_bytes = index_type{}, - std::optional use_double_precision = std::nullopt, - raft_proto::device_type dev_type = raft_proto::device_type::cpu, - int device = 0, - raft_proto::cuda_stream stream = raft_proto::cuda_stream{}) +inline auto import_from_treelite_model( + treelite::Model const& tl_model, + tree_layout layout = preferred_tree_layout, + index_type align_bytes = index_type{}, + std::optional use_double_precision = std::nullopt, + raft_proto::device_type dev_type = raft_proto::device_type::cpu, + int device = 0, + raft_proto::cuda_stream stream = raft_proto::cuda_stream{}) { auto result = forest_model{}; switch (layout) { @@ -488,13 +499,14 @@ auto import_from_treelite_model(treelite::Model const& tl_model, * @param stream The CUDA stream to use for loading this model (can be * omitted for CPU). */ -auto import_from_treelite_handle(TreeliteModelHandle tl_handle, - tree_layout layout = preferred_tree_layout, - index_type align_bytes = index_type{}, - std::optional use_double_precision = std::nullopt, - raft_proto::device_type dev_type = raft_proto::device_type::cpu, - int device = 0, - raft_proto::cuda_stream stream = raft_proto::cuda_stream{}) +inline auto import_from_treelite_handle( + TreeliteModelHandle tl_handle, + tree_layout layout = preferred_tree_layout, + index_type align_bytes = index_type{}, + std::optional use_double_precision = std::nullopt, + raft_proto::device_type dev_type = raft_proto::device_type::cpu, + int device = 0, + raft_proto::cuda_stream stream = raft_proto::cuda_stream{}) { return import_from_treelite_model(*static_cast(tl_handle), layout, diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 18bef95..4430c46 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -84,6 +84,7 @@ ConfigureTest(NAME HOST_BUFFER_TEST raft_proto/buffer.cpp) ConfigureTest(NAME DEVICE_BUFFER_TEST raft_proto/buffer.cu) ConfigureTest(NAME FOREST_TRAVERSAL_TEST forest/traversal_forest.cpp) ConfigureTest(NAME TREELITE_TRAVERSAL_TEST forest/treelite_traversal.cpp) -ConfigureTest(NAME TREELITE_IMPORTER_TEST treelite_importer.cpp) +ConfigureTest(NAME TREELITE_IMPORTER_TEST treelite_importer.cpp treelite_importer_invalid_inputs.cpp + decision_forest_builder_invalid_inputs.cpp) rapids_test_install_relocatable(INSTALL_COMPONENT_SET testing DESTINATION bin/gtests/libnvforest) diff --git a/cpp/tests/decision_forest_builder_invalid_inputs.cpp b/cpp/tests/decision_forest_builder_invalid_inputs.cpp new file mode 100644 index 0000000..b6a6637 --- /dev/null +++ b/cpp/tests/decision_forest_builder_invalid_inputs.cpp @@ -0,0 +1,69 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include + +#include +#include + +#include +#include + +namespace nvforest::detail { + +using test_forest_t = + decision_forest; + +TEST(DecisionForestBuilder, CategoricalStorageOffsetOutOfBounds) +{ + auto builder = decision_forest_builder( + /*max_num_categories=*/std::uint32_t{33}, /*align_bytes=*/std::uint32_t{0}); + + builder.add_node(std::uint32_t{1234}, + /*tl_node_id=*/0, + /*depth=*/0, + /*is_leaf_node=*/false, + /*default_to_distant_child=*/false, + /*is_categorical_node=*/true, + /*feature=*/0, + /*offset=*/1); + + ASSERT_THAT( + [&] { builder.get_decision_forest(/*num_feature=*/1, /*num_class=*/1); }, + testing::ThrowsMessage(testing::HasSubstr("storage offset out of bounds"))); +} + +TEST(DecisionForestBuilder, CategoricalBitsetExtentOutOfBounds) +{ + auto builder = decision_forest_builder( + /*max_num_categories=*/std::uint32_t{33}, /*align_bytes=*/std::uint32_t{0}); + + std::array categories{0}; + builder.add_categorical_node(categories.begin(), + categories.end(), + /*tl_node_id=*/0, + /*depth=*/0, + /*default_to_distant_child=*/false, + /*feature=*/0, + /*offset=*/1); + + builder.add_node(std::uint32_t{1}, + /*tl_node_id=*/1, + /*depth=*/1, + /*is_leaf_node=*/false, + /*default_to_distant_child=*/false, + /*is_categorical_node=*/true, + /*feature=*/0, + /*offset=*/1); + + ASSERT_THAT([&] { builder.get_decision_forest(/*num_feature=*/1, /*num_class=*/1); }, + testing::ThrowsMessage( + testing::HasSubstr("bitset extends past categorical_storage end"))); +} + +} // namespace nvforest::detail diff --git a/cpp/tests/treelite_importer.cpp b/cpp/tests/treelite_importer.cpp index f29baad..642786c 100644 --- a/cpp/tests/treelite_importer.cpp +++ b/cpp/tests/treelite_importer.cpp @@ -248,45 +248,45 @@ auto static const SAMPLE_FOREST = []() { TEST(TreeliteImporter, depth_first) { - auto fil_model = import_from_treelite_model(*SAMPLE_FOREST, tree_layout::depth_first); - ASSERT_EQ(fil_model.num_features(), 7); - ASSERT_EQ(fil_model.num_outputs(), 1); - ASSERT_EQ(fil_model.num_trees(), 6); - ASSERT_FALSE(fil_model.has_vector_leaves()); - ASSERT_EQ(fil_model.row_postprocessing(), row_op::disable); - ASSERT_EQ(fil_model.elem_postprocessing(), element_op::disable); - ASSERT_EQ(fil_model.memory_type(), raft_proto::device_type::cpu); - ASSERT_EQ(fil_model.device_index(), -1); - ASSERT_FALSE(fil_model.is_double_precision()); + auto nvforest_model = import_from_treelite_model(*SAMPLE_FOREST, tree_layout::depth_first); + ASSERT_EQ(nvforest_model.num_features(), 7); + ASSERT_EQ(nvforest_model.num_outputs(), 1); + ASSERT_EQ(nvforest_model.num_trees(), 6); + ASSERT_FALSE(nvforest_model.has_vector_leaves()); + ASSERT_EQ(nvforest_model.row_postprocessing(), row_op::disable); + ASSERT_EQ(nvforest_model.elem_postprocessing(), element_op::disable); + ASSERT_EQ(nvforest_model.memory_type(), raft_proto::device_type::cpu); + ASSERT_EQ(nvforest_model.device_index(), -1); + ASSERT_FALSE(nvforest_model.is_double_precision()); } TEST(TreeliteImporter, breadth_first) { - auto fil_model = import_from_treelite_model(*SAMPLE_FOREST, tree_layout::breadth_first); - ASSERT_EQ(fil_model.num_features(), 7); - ASSERT_EQ(fil_model.num_outputs(), 1); - ASSERT_EQ(fil_model.num_trees(), 6); - ASSERT_FALSE(fil_model.has_vector_leaves()); - ASSERT_EQ(fil_model.row_postprocessing(), row_op::disable); - ASSERT_EQ(fil_model.elem_postprocessing(), element_op::disable); - ASSERT_EQ(fil_model.memory_type(), raft_proto::device_type::cpu); - ASSERT_EQ(fil_model.device_index(), -1); - ASSERT_FALSE(fil_model.is_double_precision()); + auto nvforest_model = import_from_treelite_model(*SAMPLE_FOREST, tree_layout::breadth_first); + ASSERT_EQ(nvforest_model.num_features(), 7); + ASSERT_EQ(nvforest_model.num_outputs(), 1); + ASSERT_EQ(nvforest_model.num_trees(), 6); + ASSERT_FALSE(nvforest_model.has_vector_leaves()); + ASSERT_EQ(nvforest_model.row_postprocessing(), row_op::disable); + ASSERT_EQ(nvforest_model.elem_postprocessing(), element_op::disable); + ASSERT_EQ(nvforest_model.memory_type(), raft_proto::device_type::cpu); + ASSERT_EQ(nvforest_model.device_index(), -1); + ASSERT_FALSE(nvforest_model.is_double_precision()); } TEST(TreeliteImporter, layered_children_together) { - auto fil_model = + auto nvforest_model = import_from_treelite_model(*SAMPLE_FOREST, tree_layout::layered_children_together); - ASSERT_EQ(fil_model.num_features(), 7); - ASSERT_EQ(fil_model.num_outputs(), 1); - ASSERT_EQ(fil_model.num_trees(), 6); - ASSERT_FALSE(fil_model.has_vector_leaves()); - ASSERT_EQ(fil_model.row_postprocessing(), row_op::disable); - ASSERT_EQ(fil_model.elem_postprocessing(), element_op::disable); - ASSERT_EQ(fil_model.memory_type(), raft_proto::device_type::cpu); - ASSERT_EQ(fil_model.device_index(), -1); - ASSERT_FALSE(fil_model.is_double_precision()); + ASSERT_EQ(nvforest_model.num_features(), 7); + ASSERT_EQ(nvforest_model.num_outputs(), 1); + ASSERT_EQ(nvforest_model.num_trees(), 6); + ASSERT_FALSE(nvforest_model.has_vector_leaves()); + ASSERT_EQ(nvforest_model.row_postprocessing(), row_op::disable); + ASSERT_EQ(nvforest_model.elem_postprocessing(), element_op::disable); + ASSERT_EQ(nvforest_model.memory_type(), raft_proto::device_type::cpu); + ASSERT_EQ(nvforest_model.device_index(), -1); + ASSERT_FALSE(nvforest_model.is_double_precision()); } template @@ -334,9 +334,9 @@ auto make_degenerate_tree(const leaf_t& leaf) TEST(TreeliteImporter, DegenerateTree) { - auto tl_model = make_degenerate_tree(1.0); - auto fil_model = import_from_treelite_model(*tl_model, tree_layout::breadth_first); - ASSERT_FALSE(fil_model.has_vector_leaves()); + auto tl_model = make_degenerate_tree(1.0); + auto nvforest_model = import_from_treelite_model(*tl_model, tree_layout::breadth_first); + ASSERT_FALSE(nvforest_model.has_vector_leaves()); #ifdef NVFOREST_ENABLE_GPU auto raft_handle = raft::handle_t{}; @@ -347,22 +347,22 @@ TEST(TreeliteImporter, DegenerateTree) auto X = std::vector{0.0}; auto preds = std::vector(1, 0.0); auto expected_preds = std::vector{1.0}; - fil_model.predict(handle, - preds.data(), - X.data(), - 1, - raft_proto::device_type::cpu, - raft_proto::device_type::cpu, - nvforest::infer_kind::default_kind, - 1); + nvforest_model.predict(handle, + preds.data(), + X.data(), + 1, + raft_proto::device_type::cpu, + raft_proto::device_type::cpu, + nvforest::infer_kind::default_kind, + 1); ASSERT_EQ(preds, expected_preds); } TEST(TreeliteImporter, DegenerateTreeWithVectorLeaf) { - auto tl_model = make_degenerate_tree(std::vector{0.5, 0.5}); - auto fil_model = import_from_treelite_model(*tl_model, tree_layout::breadth_first); - ASSERT_TRUE(fil_model.has_vector_leaves()); + auto tl_model = make_degenerate_tree(std::vector{0.5, 0.5}); + auto nvforest_model = import_from_treelite_model(*tl_model, tree_layout::breadth_first); + ASSERT_TRUE(nvforest_model.has_vector_leaves()); #ifdef NVFOREST_ENABLE_GPU auto raft_handle = raft::handle_t{}; @@ -373,14 +373,14 @@ TEST(TreeliteImporter, DegenerateTreeWithVectorLeaf) auto X = std::vector{0.0}; auto preds = std::vector(2, 0.0); auto expected_preds = std::vector{0.5, 0.5}; - fil_model.predict(handle, - preds.data(), - X.data(), - 1, - raft_proto::device_type::cpu, - raft_proto::device_type::cpu, - nvforest::infer_kind::default_kind, - 1); + nvforest_model.predict(handle, + preds.data(), + X.data(), + 1, + raft_proto::device_type::cpu, + raft_proto::device_type::cpu, + nvforest::infer_kind::default_kind, + 1); ASSERT_EQ(preds, expected_preds); } diff --git a/cpp/tests/treelite_importer_invalid_inputs.cpp b/cpp/tests/treelite_importer_invalid_inputs.cpp new file mode 100644 index 0000000..d2b9ccb --- /dev/null +++ b/cpp/tests/treelite_importer_invalid_inputs.cpp @@ -0,0 +1,203 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace nvforest { + +namespace { + +auto make_categorical_model(std::uint32_t category) +{ + auto metadata = treelite::model_builder::Metadata{ + 1, + treelite::TaskType::kRegressor, + false, + 1, + {1}, + {1, 1}, + }; + auto tree_annotation = treelite::model_builder::TreeAnnotation{1, {0}, {0}}; + auto model_builder = + treelite::model_builder::GetModelBuilder(treelite::TypeInfo::kFloat32, + treelite::TypeInfo::kFloat32, + metadata, + tree_annotation, + treelite::model_builder::PostProcessorFunc{"identity"}, + {0.0}); + + model_builder->StartTree(); + model_builder->StartNode(0); + model_builder->CategoricalTest(0, false, {category}, false, 1, 2); + model_builder->EndNode(); + + model_builder->StartNode(1); + model_builder->LeafScalar(1.0f); + model_builder->EndNode(); + + model_builder->StartNode(2); + model_builder->LeafScalar(-1.0f); + model_builder->EndNode(); + + model_builder->EndTree(); + return model_builder->CommitModel(); +} + +} // namespace + +TEST(TreeliteImporter, LargeCategoryValue) +{ + auto tl_model = make_categorical_model(std::numeric_limits::max()); + auto expected_error_msg = std::string{"Tree 0, Node 0: Category index must be at most "} + + std::to_string(std::numeric_limits::max() - 1); + + ASSERT_THAT([&] { import_from_treelite_model(*tl_model, tree_layout::breadth_first); }, + testing::ThrowsMessage(testing::HasSubstr(expected_error_msg))); + ASSERT_THAT( + [&] { + import_from_treelite_model( + *tl_model, tree_layout::breadth_first, 0, /*use_double_precision=*/true); + }, + testing::ThrowsMessage(testing::HasSubstr(expected_error_msg))); +} + +TEST(TreeliteImporter, LargeCategoryValueBelowLimit) +{ + auto tl_model = make_categorical_model(std::numeric_limits::max() - 1); + ASSERT_NO_THROW(import_from_treelite_model(*tl_model, tree_layout::breadth_first)); +} + +TEST(TreeliteImporter, LargeFeatureId) +{ + auto metadata = treelite::model_builder::Metadata{ + 9000, + treelite::TaskType::kRegressor, + false, + 1, + {1}, + {1, 1}, + }; + auto tree_annotation = treelite::model_builder::TreeAnnotation{1, {0}, {0}}; + auto model_builder = + treelite::model_builder::GetModelBuilder(treelite::TypeInfo::kFloat32, + treelite::TypeInfo::kFloat32, + metadata, + tree_annotation, + treelite::model_builder::PostProcessorFunc{"identity"}, + {0.0}); + + model_builder->StartTree(); + model_builder->StartNode(0); + model_builder->NumericalTest(8999, 0.0, false, treelite::Operator::kGT, 1, 2); + model_builder->EndNode(); + + model_builder->StartNode(1); + model_builder->LeafScalar(1.0f); + model_builder->EndNode(); + + model_builder->StartNode(2); + model_builder->LeafScalar(-1.0f); + model_builder->EndNode(); + + model_builder->EndTree(); + + auto tl_model = model_builder->CommitModel(); + ASSERT_NO_THROW(import_from_treelite_model(*tl_model, tree_layout::breadth_first)); + + auto variant_index = get_forest_variant_index(false, 2, 1); + auto importer = treelite_importer{}; + + auto expected_error_msg = + std::string{"Tree 0, Node 0: The 'feature' value in the node must be at most "} + + std::to_string(0x1FFF); + ASSERT_THAT( + [&] { + importer.import_to_specific_variant(variant_index, + *tl_model, + importer.get_num_class(*tl_model), + importer.get_num_feature(*tl_model), + importer.get_max_num_categories(*tl_model), + importer.get_offsets(*tl_model)); + }, + testing::ThrowsMessage(testing::HasSubstr(expected_error_msg))); +} + +TEST(TreeliteImporter, SafeCastFloatingPoint) +{ + ASSERT_NO_THROW(detail::safe_cast_floating_point(double{3.1})); + ASSERT_NO_THROW(detail::safe_cast_floating_point(std::numeric_limits::max())); + + ASSERT_NO_THROW(detail::safe_cast_floating_point(std::numeric_limits::infinity())); + ASSERT_NO_THROW( + detail::safe_cast_floating_point(std::numeric_limits::infinity())); + ASSERT_NO_THROW(detail::safe_cast_floating_point(std::numeric_limits::infinity())); + ASSERT_NO_THROW(detail::safe_cast_floating_point(std::numeric_limits::quiet_NaN())); + ASSERT_NO_THROW( + detail::safe_cast_floating_point(std::numeric_limits::quiet_NaN())); + ASSERT_NO_THROW( + detail::safe_cast_floating_point(std::numeric_limits::quiet_NaN())); + + auto inf_msg = std::string{"Cannot cast an INF or NaN value"}; + ASSERT_THAT( + [] { detail::safe_cast_floating_point(std::numeric_limits::infinity()); }, + testing::ThrowsMessage(testing::HasSubstr(inf_msg))); + ASSERT_THAT( + [] { detail::safe_cast_floating_point(std::numeric_limits::quiet_NaN()); }, + testing::ThrowsMessage(testing::HasSubstr(inf_msg))); + ASSERT_THAT([] { detail::safe_cast_floating_point(double{1e100}); }, + testing::ThrowsMessage( + testing::HasSubstr("Input must be at most"))); + ASSERT_THAT([] { detail::safe_cast_floating_point(double{-1e100}); }, + testing::ThrowsMessage( + testing::HasSubstr("Input must be at least"))); +} + +TEST(TreeliteImporter, InvalidPostprocConstant) +{ + auto metadata = treelite::model_builder::Metadata{ + 1, + treelite::TaskType::kRegressor, + false, + 1, + {1}, + {1, 1}, + }; + auto tree_annotation = treelite::model_builder::TreeAnnotation{1, {0}, {0}}; + auto model_builder = treelite::model_builder::GetModelBuilder( + treelite::TypeInfo::kFloat32, + treelite::TypeInfo::kFloat32, + metadata, + tree_annotation, + treelite::model_builder::PostProcessorFunc{ + "sigmoid", {{"sigmoid_alpha", std::numeric_limits::quiet_NaN()}}}, + {0.0}); + + model_builder->StartTree(); + model_builder->StartNode(0); + model_builder->LeafScalar(0.0f); + model_builder->EndNode(); + model_builder->EndTree(); + + auto tl_model = model_builder->CommitModel(); + + ASSERT_THAT([&] { import_from_treelite_model(*tl_model, tree_layout::breadth_first); }, + testing::ThrowsMessage( + testing::HasSubstr("Found an invalid value for postprocessing constant"))); +} + +} // namespace nvforest