diff --git a/cpp/include/cuml/fil/detail/bitset.hpp b/cpp/include/cuml/fil/detail/bitset.hpp index 6981c29d1a..edbe5140ed 100644 --- a/cpp/include/cuml/fil/detail/bitset.hpp +++ b/cpp/include/cuml/fil/detail/bitset.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -22,6 +22,9 @@ struct bitset { using storage_type = storage_t; using index_type = index_t; + // Ensrue 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} {} @@ -39,12 +42,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/cuml/fil/detail/decision_forest_builder.hpp b/cpp/include/cuml/fil/detail/decision_forest_builder.hpp index 0fa8ac029c..488e1be893 100644 --- a/cpp/include/cuml/fil/detail/decision_forest_builder.hpp +++ b/cpp/include/cuml/fil/detail/decision_forest_builder.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -18,26 +18,58 @@ #include #include #include +#include #include #include +#include +#include +#include +#include #include namespace ML { namespace fil { namespace detail { -/* - * Exception indicating that FIL 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() {} + floating_point_truncation_error(std::string msg) : msg_{msg} {} + floating_point_truncation_error(char const* msg) : msg_{msg} {} + virtual char const* what() const noexcept { return msg_.c_str(); } private: - char const* msg_; + std::string msg_; }; +template +To safe_cast_floating_point(From 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(To) >= sizeof(From)) { + // Widening cast + return static_cast(x); + } else { + // Narrowing cast: should be checked + 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 FIL forests */ @@ -57,20 +89,36 @@ 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; + + // Check invariants for data types + using cat_t = typename std::iterator_traits::value_type; + using index_t = typename node_type::index_type; + static_assert(std::is_same_v, "Category value must be uint32_t"); + static_assert(std::is_same_v || std::is_same_v, + "Index type in tree node must be either uint32_t or uint64_t"); + + // Ensure that (max_cat + 1) can be represented as index_t to prevent integer overflow. + auto max_cat = (vec_begin != vec_end) ? *std::max_element(vec_begin, vec_end) : cat_t{0}; + if constexpr (std::is_same_v) { + if (max_cat == std::numeric_limits::max()) { + throw model_import_error{std::string{"Category index must be at most "} + + std::to_string(std::numeric_limits::max() - 1)}; + } + } + auto max_cat_plus_one = static_cast(max_cat) + index_t{1}; + 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( @@ -153,7 +201,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 unusable_model_exception("Inconsistent leaf vector size"); } output_size_ = val; } @@ -182,11 +230,78 @@ struct decision_forest_builder { int device = 0, raft_proto::cuda_stream stream = raft_proto::cuda_stream{}) { - // 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. + + // tree_index arithmetic in the kernel uses index_type, so the tree count + // must fit without narrowing. + 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()) + ")"}; + } + + // forest::get_tree_root(tree_index) dereferences nodes_ + root_index. + // Ensure each root index points into the nodes buffer. + 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(); + + // evaluate_tree_impl() reads categorical_storage[offset] as the number + // of categories for this node; offset must be in-range. + 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); + + // evaluate_tree_impl() reconstructs a bitset from + // [offset + 1, offset + 1 + bins_required). Compute this range using + // size_t to keep the arithmetic explicit and overflow-safe. + 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"}; + } + } + } + + // Safely cast average_factor_ and postproc_constant_ to node_type::threshold_type + 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_); + // We can't use cuda::narrow here, because it throws for imprecise conversion, i.e. casting + // double{3.1} to float. + } catch (const floating_point_truncation_error& e) { + throw unusable_model_exception{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 unusable_model_exception{ + 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}, @@ -219,9 +334,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/cuml/fil/detail/degenerate_trees.hpp b/cpp/include/cuml/fil/detail/degenerate_trees.hpp index a0d20714f7..5dcf19d208 100644 --- a/cpp/include/cuml/fil/detail/degenerate_trees.hpp +++ b/cpp/include/cuml/fil/detail/degenerate_trees.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -17,7 +17,7 @@ namespace ML::fil::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 = ML::forest::tree_accumulate(tl_model, false, [](auto&& contains, auto&& tree) { diff --git a/cpp/include/cuml/fil/detail/node.hpp b/cpp/include/cuml/fil/detail/node.hpp index 6a0a3765d4..0a5bf88e2c 100644 --- a/cpp/include/cuml/fil/detail/node.hpp +++ b/cpp/include/cuml/fil/detail/node.hpp @@ -1,11 +1,12 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once #include #include +#include #include #include @@ -104,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/cuml/fil/detail/raft_proto/ceildiv.hpp b/cpp/include/cuml/fil/detail/raft_proto/ceildiv.hpp index 61887f6ffb..429b8af64a 100644 --- a/cpp/include/cuml/fil/detail/raft_proto/ceildiv.hpp +++ b/cpp/include/cuml/fil/detail/raft_proto/ceildiv.hpp @@ -1,14 +1,17 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #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/cuml/fil/exceptions.hpp b/cpp/include/cuml/fil/exceptions.hpp index 8a8c2a00fd..39c00a43e3 100644 --- a/cpp/include/cuml/fil/exceptions.hpp +++ b/cpp/include/cuml/fil/exceptions.hpp @@ -23,11 +23,12 @@ struct unusable_model_exception : std::exception { /** 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} {} model_import_error(char const* msg) : msg_{msg} {} - virtual char const* what() const noexcept { return msg_; } + virtual char const* what() const noexcept { return msg_.c_str(); } private: - char const* msg_; + std::string msg_; }; /** diff --git a/cpp/include/cuml/fil/treelite_importer.hpp b/cpp/include/cuml/fil/treelite_importer.hpp index be3fdb217c..853fc5d39b 100644 --- a/cpp/include/cuml/fil/treelite_importer.hpp +++ b/cpp/include/cuml/fil/treelite_importer.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -186,7 +186,7 @@ struct treelite_importer { result.constant = tl_model.sigmoid_alpha; result.element = element_op::sigmoid; } else { - throw model_import_error{"Unrecognized Treelite pred_transform string"}; + throw unusable_model_exception{"Unrecognized Treelite pred_transform string"}; } return result; } @@ -197,7 +197,7 @@ struct treelite_importer { switch (tl_model.GetThresholdType()) { case treelite::TypeInfo::kFloat64: result = true; break; case treelite::TypeInfo::kFloat32: result = false; break; - default: throw model_import_error("Unrecognized Treelite threshold type"); + default: throw unusable_model_exception("Unrecognized Treelite threshold type"); } return result; } @@ -209,7 +209,7 @@ struct treelite_importer { case treelite::TypeInfo::kFloat64: result = true; break; case treelite::TypeInfo::kFloat32: result = false; break; case treelite::TypeInfo::kUInt32: result = false; break; - default: throw model_import_error("Unrecognized Treelite threshold type"); + default: throw unusable_model_exception("Unrecognized Treelite threshold type"); } return result; } @@ -221,7 +221,7 @@ struct treelite_importer { case treelite::TypeInfo::kFloat64: result = false; break; case treelite::TypeInfo::kFloat32: result = false; break; case treelite::TypeInfo::kUInt32: result = true; break; - default: throw model_import_error("Unrecognized Treelite threshold type"); + default: throw unusable_model_exception("Unrecognized Treelite threshold type"); } return result; } @@ -258,37 +258,44 @@ 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(); + 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; }); @@ -429,13 +436,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) { @@ -479,13 +487,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 25e6109cf9..89061b14b9 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -109,7 +109,12 @@ if(all_algo OR fil_algo) ConfigureTest(PREFIX SG NAME DEVICE_BUFFER_TEST sg/fil/raft_proto/buffer.cu ML_INCLUDE) ConfigureTest(PREFIX SG NAME FOREST_TRAVERSAL_TEST sg/forest/traversal_forest.cpp ML_INCLUDE) ConfigureTest(PREFIX SG NAME TREELITE_TRAVERSAL_TEST sg/forest/treelite_traversal.cpp ML_INCLUDE) - ConfigureTest(PREFIX SG NAME TREELITE_IMPORTER_TEST sg/fil/treelite_importer.cpp ML_INCLUDE) + ConfigureTest( + PREFIX SG + NAME TREELITE_IMPORTER_TEST + sg/fil/treelite_importer.cpp sg/fil/treelite_importer_invalid_inputs.cpp + sg/fil/decision_forest_builder_invalid_inputs.cpp ML_INCLUDE + ) endif() # todo: organize linear models better diff --git a/cpp/tests/sg/fil/decision_forest_builder_invalid_inputs.cpp b/cpp/tests/sg/fil/decision_forest_builder_invalid_inputs.cpp new file mode 100644 index 0000000000..5e508573a9 --- /dev/null +++ b/cpp/tests/sg/fil/decision_forest_builder_invalid_inputs.cpp @@ -0,0 +1,82 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include + +#include +#include + +#include +#include + +namespace ML { +namespace fil { +namespace 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}); + + // Construct a malformed categorical node that references non-local category + // storage at an out-of-range offset. This should be rejected by the invariant + // checks in get_decision_forest(). + 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}); + + // Create a valid categorical node first, which allocates non-local storage: + // categorical_storage_ = [num_categories, packed_bin_data] + 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); + + // Construct another categorical node that points at offset=1, i.e. the first + // packed bin entry rather than the metadata entry. The value at offset=1 is + // interpreted as stored_num_cats, and the resulting bins_required exceeds the + // available headroom. + 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 detail +} // namespace fil +} // namespace ML diff --git a/cpp/tests/sg/fil/treelite_importer_invalid_inputs.cpp b/cpp/tests/sg/fil/treelite_importer_invalid_inputs.cpp new file mode 100644 index 0000000000..aa0bae0661 --- /dev/null +++ b/cpp/tests/sg/fil/treelite_importer_invalid_inputs.cpp @@ -0,0 +1,248 @@ +/* + * 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 ML { +namespace fil { + +TEST(TreeliteImporter, large_category_value) +{ + // For tree models with 32-bit data storage, + // attempting to use UINT32_MAX in a categorical test node + // must throw an exception. + + 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, {std::numeric_limits::max()}, 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(); + + auto tl_model = model_builder->CommitModel(); + + 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))); +} + +TEST(TreeliteImporter, large_category_value2) +{ + // For tree models with 32-bit data storage, + // it should be possible to use (UINT32_MAX - 1) in a categorical test node + + 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, {std::numeric_limits::max() - 1}, 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(); + + auto tl_model = model_builder->CommitModel(); + ASSERT_NO_THROW(import_from_treelite_model(*tl_model, tree_layout::breadth_first)); +} + +TEST(TreeliteImporter, large_feature_id) +{ + // Tree models with 16-bit storage for node metadata should throw + // an exception for feature IDs larger than 0x1FFF. + + 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); + // Use a "large" feature ID here + 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(); + + // Normally, treelite_importer::import() would choose the right size + // for the metadata storage, sufficient to hold all given feature IDs. + // For this example, it chooses 32-bit storage type (due to the use of feature ID 8999). + ASSERT_NO_THROW(import_from_treelite_model(*tl_model, tree_layout::breadth_first)); + + // Trick the importer to pick 16-bit storage type for metadata storage. + auto variant_index = get_forest_variant_index(false, 2, 1); + auto importer = treelite_importer{}; + + // The importer should throw an informative error message rather than silently + // truncating the feature ID. + 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, safe_cast_floating_point) +{ + /* Valid casts */ + ASSERT_NO_THROW( + detail::safe_cast_floating_point(double{3.1})); // Some loss of precision, but o.k. + ASSERT_NO_THROW(detail::safe_cast_floating_point(std::numeric_limits::max())); + + // INFs and NANs are allowed for widening cast + 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())); + + // Invalid casts + 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, invalid_postproc_constant) +{ + 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 fil +} // namespace ML