[FIL] Validate Treelite model input to prevent integer overflow and OOB memory access - #8016
Merged
Merged
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds runtime bounds checks and an unsigned-index static_assert to bitset; tightens categorical-node validation in decision_forest_builder (adds tree_id, runtime overflow check, model_builder_error string storage); marks several header functions inline; changes ceildiv to require integral types and adjusts logic; adds tests and a CMake update; updates header years. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
cpp/include/cuml/fil/detail/raft_proto/ceildiv.hpp (1)
12-12: Add type constraints toceildivto prevent accidental misuse with non-integral types.The unconstrained template accepts any type, allowing potential misuse even though the algorithm (division and modulo) only makes sense for integral types. All current callsites pass integral values (size_t, indices, counts), but explicit constraints via
static_assertwould prevent future errors.Suggested improvement
+#include <type_traits> namespace raft_proto { template <typename T, typename U> HOST DEVICE auto constexpr ceildiv(T dividend, U divisor) { + static_assert(std::is_integral_v<T> && std::is_integral_v<U>, + "ceildiv requires integral operands"); return dividend / divisor + (dividend % divisor != 0); } } // namespace raft_protoAdditionally, consider adding Doxygen documentation to this public header function documenting the requirement for positive divisors.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cpp/include/cuml/fil/detail/raft_proto/ceildiv.hpp` at line 12, Constrain the template for ceildiv to integral types and document the positive-divisor requirement: add a compile-time check (e.g., static_assert(std::is_integral_v<T>, "ceildiv requires an integral type") or use std::integral concept if C++20) inside the ceildiv<T> template to prevent non-integral instantiations, and add a Doxygen comment above the ceildiv function specifying that the divisor must be non-zero and positive and describing the function semantics; reference the ceildiv function name so reviewers can find the change.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@cpp/include/cuml/fil/detail/bitset.hpp`:
- Around line 42-48: The current clear (and the corresponding set-like) method
only checks the upper bound and can underflow with signed index types; update
the guard in clear and the sibling setter to also verify the lower bound by
ensuring index >= index_type(0) && index < num_bits_ before calling
bin_from_index() or mask_in_bin(), so negative indices are rejected and no
signed-index underflow or out-of-bounds access can occur.
In `@cpp/include/cuml/fil/detail/decision_forest_builder.hpp`:
- Around line 66-82: Replace all uses of typename iter_t::value_type with
std::iterator_traits<iter_t>::value_type (e.g., in is_negative lambda and when
defaulting max_cat). Compute max_cat using that traits type, then immediately
cast it to node_type::index_type (e.g., auto max_cat_index =
static_cast<typename node_type::index_type>(max_cat)); compare max_cat_index
against max_representable (auto const max_representable =
std::numeric_limits<typename node_type::index_type>::max()) and throw only if
max_cat_index >= max_representable; finally compute max_cat_plus_one as
max_cat_index + typename node_type::index_type{1}. This fixes portability for
raw pointers and avoids incorrectly rejecting values when the node index type is
wider.
---
Nitpick comments:
In `@cpp/include/cuml/fil/detail/raft_proto/ceildiv.hpp`:
- Line 12: Constrain the template for ceildiv to integral types and document the
positive-divisor requirement: add a compile-time check (e.g.,
static_assert(std::is_integral_v<T>, "ceildiv requires an integral type") or use
std::integral concept if C++20) inside the ceildiv<T> template to prevent
non-integral instantiations, and add a Doxygen comment above the ceildiv
function specifying that the divisor must be non-zero and positive and
describing the function semantics; reference the ceildiv function name so
reviewers can find the change.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 809d9678-862b-425a-a798-e6d385ad3f16
📒 Files selected for processing (3)
cpp/include/cuml/fil/detail/bitset.hppcpp/include/cuml/fil/detail/decision_forest_builder.hppcpp/include/cuml/fil/detail/raft_proto/ceildiv.hpp
There was a problem hiding this comment.
♻️ Duplicate comments (1)
cpp/include/cuml/fil/detail/decision_forest_builder.hpp (1)
75-82:⚠️ Potential issue | 🟠 MajorFix the off-by-one and pre-cast overflow in the new bounds check.
Line 77 rejects the largest valid category (
max_cat == max_representable - 1), even thoughmax_cat_plus_onewould still fit exactly innode_type::index_type. It also evaluatesmax_cat + 1incat_tbefore the cast, so signed iterators can still overflow here. Comparemax_catagainstmax_representablein a widened type, then compute+ 1only after converting tonode_type::index_type.Suggested adjustment
- auto const max_representable = std::numeric_limits<typename node_type::index_type>::max(); - if (max_cat >= max_representable || max_cat + 1 >= max_representable) { + using index_t = typename node_type::index_type; + auto const max_representable = std::numeric_limits<index_t>::max(); + if (static_cast<std::uintmax_t>(max_cat) >= + static_cast<std::uintmax_t>(max_representable)) { throw model_builder_error( "Category index exceeds maximum representable value for this model's index type"); } - auto max_cat_plus_one = - static_cast<typename node_type::index_type>(max_cat) + typename node_type::index_type{1}; + auto max_cat_plus_one = static_cast<index_t>(max_cat) + index_t{1};Based on learnings: Handle numerical edge cases explicitly (near-zero eigenvalues, degenerate matrices, zero-norm vectors, extreme values) and verify unsafe type casting between numeric types does not cause unintended precision loss.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cpp/include/cuml/fil/detail/decision_forest_builder.hpp` around lines 75 - 82, The bounds check in decision_forest_builder.hpp incorrectly rejects max_cat == max_representable - 1 and risks overflow by evaluating max_cat + 1 in cat_t; change the logic in the block around max_cat, max_representable and max_cat_plus_one so you first promote/compare max_cat to a widened integer type (e.g., std::common_type_t<cat_t, typename node_type::index_type, std::uintmax_t>) against max_representable to detect values >= max_representable, and only after that cast max_cat to typename node_type::index_type and add 1 to compute max_cat_plus_one; update the throw condition to reject only when max_cat >= max_representable (in the widened type) and ensure max_cat_plus_one is computed after the safe cast to node_type::index_type.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@cpp/include/cuml/fil/detail/decision_forest_builder.hpp`:
- Around line 75-82: The bounds check in decision_forest_builder.hpp incorrectly
rejects max_cat == max_representable - 1 and risks overflow by evaluating
max_cat + 1 in cat_t; change the logic in the block around max_cat,
max_representable and max_cat_plus_one so you first promote/compare max_cat to a
widened integer type (e.g., std::common_type_t<cat_t, typename
node_type::index_type, std::uintmax_t>) against max_representable to detect
values >= max_representable, and only after that cast max_cat to typename
node_type::index_type and add 1 to compute max_cat_plus_one; update the throw
condition to reject only when max_cat >= max_representable (in the widened type)
and ensure max_cat_plus_one is computed after the safe cast to
node_type::index_type.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 991dc041-df90-4549-9b8a-54153a533e47
📒 Files selected for processing (2)
cpp/include/cuml/fil/detail/bitset.hppcpp/include/cuml/fil/detail/decision_forest_builder.hpp
🚧 Files skipped from review as they are similar to previous changes (1)
- cpp/include/cuml/fil/detail/bitset.hpp
chyunsu3
added
improvement
There was a problem hiding this comment.
♻️ Duplicate comments (1)
cpp/include/cuml/fil/detail/decision_forest_builder.hpp (1)
77-84:⚠️ Potential issue | 🟠 MajorOverflow guard still has a boundary bug and potential signed-overflow UB.
The condition at Line 79 rejects a valid case (
max_cat == max_representable - 1) and can evaluatemax_cat + 1incat_t, which is UB for signed overflow.🔧 Suggested fix
- auto max_cat = (vec_begin != vec_end) ? *std::max_element(vec_begin, vec_end) : cat_t{0}; - auto const max_representable = std::numeric_limits<typename node_type::index_type>::max(); - if (max_cat >= max_representable || max_cat + 1 >= max_representable) { + auto const max_cat = (vec_begin != vec_end) ? *std::max_element(vec_begin, vec_end) : cat_t{0}; + using index_t = typename node_type::index_type; + using compare_t = + std::common_type_t<std::make_unsigned_t<cat_t>, std::make_unsigned_t<index_t>>; + auto const max_cat_u = static_cast<compare_t>(max_cat); + auto const max_representable = std::numeric_limits<index_t>::max(); + auto const max_rep_u = static_cast<compare_t>(max_representable); + if (max_cat_u >= max_rep_u) { throw model_builder_error( "Category index exceeds maximum representable value for this model's index type"); } - auto max_cat_plus_one = - static_cast<typename node_type::index_type>(max_cat) + typename node_type::index_type{1}; + auto max_cat_plus_one = static_cast<index_t>(max_cat_u + compare_t{1});#!/bin/bash set -euo pipefail f="cpp/include/cuml/fil/detail/decision_forest_builder.hpp" echo "== Current overflow guard ==" sed -n '75,86p' "$f" echo echo "== Boundary behavior demonstration ==" python3 - <<'PY' max_rep = 255 for max_cat in (254, 255): cond = (max_cat >= max_rep) or (max_cat + 1 >= max_rep) print(f"max_cat={max_cat}, condition={cond}") print("Expected: max_cat=max_rep-1 should be allowed, but condition is True.") PYBased on learnings: "Handle numerical edge cases explicitly (near-zero eigenvalues, degenerate matrices, zero-norm vectors, extreme values)".
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cpp/include/cuml/fil/detail/decision_forest_builder.hpp` around lines 77 - 84, The overflow check incorrectly rejects max_cat == max_representable-1 and risks signed overflow by evaluating max_cat + 1 in cat_t; fix by computing a safe threshold in the target index type and comparing after casting: obtain max_representable (typename node_type::index_type)::max(), compute max_representable_minus_one = max_representable - typename node_type::index_type{1}, then if static_cast<typename node_type::index_type>(max_cat) > max_representable_minus_one (or >= max_representable) throw the model_builder_error; finally compute max_cat_plus_one by casting max_cat to typename node_type::index_type before adding one. Use the existing symbols vec_begin/vec_end, max_cat, node_type::index_type, max_representable, and max_cat_plus_one to locate and update the checks and the addition.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@cpp/include/cuml/fil/detail/decision_forest_builder.hpp`:
- Around line 77-84: The overflow check incorrectly rejects max_cat ==
max_representable-1 and risks signed overflow by evaluating max_cat + 1 in
cat_t; fix by computing a safe threshold in the target index type and comparing
after casting: obtain max_representable (typename node_type::index_type)::max(),
compute max_representable_minus_one = max_representable - typename
node_type::index_type{1}, then if static_cast<typename
node_type::index_type>(max_cat) > max_representable_minus_one (or >=
max_representable) throw the model_builder_error; finally compute
max_cat_plus_one by casting max_cat to typename node_type::index_type before
adding one. Use the existing symbols vec_begin/vec_end, max_cat,
node_type::index_type, max_representable, and max_cat_plus_one to locate and
update the checks and the addition.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 626839f9-eb75-4801-a98c-56dddd363c40
📒 Files selected for processing (1)
cpp/include/cuml/fil/detail/decision_forest_builder.hpp
csadorf
self-requested a review
April 27, 2026 15:30
dantegd
requested changes
Apr 27, 2026
dantegd
left a comment
dantegd
left a comment
Contributor
There was a problem hiding this comment.
Mostly looks good, the direction is the right one, just had some comments and concerns.
csadorf
requested changes
Apr 27, 2026
Contributor
There was a problem hiding this comment.
There is additional out-of-bound access for the unused bit-wise boolean operations. I will remove those in a follow-up.
chyunsu3
commented
Apr 28, 2026
Comment on lines
+70
to
+88
| // Ensure that all category indices are non-negative. | ||
| using cat_t = typename std::iterator_traits<iter_t>::value_type; | ||
| static_assert(std::is_unsigned_v<cat_t>, "Category value must be an unsigned integer type"); | ||
|
|
||
| auto max_cat = (vec_begin != vec_end) ? *std::max_element(vec_begin, vec_end) : cat_t{0}; | ||
|
|
||
| // Ensure that (max_cat + 1) can be represented as node_type::index_type | ||
| // to prevent integer overflow. | ||
| using index_t = typename node_type::index_type; | ||
| auto const max_index = static_cast<std::uintmax_t>(std::numeric_limits<index_t>::max()); | ||
| auto const cat_unsigned = static_cast<std::uintmax_t>(max_cat); | ||
| if (cat_unsigned >= max_index) { | ||
| auto node_id_repr = | ||
| tl_node_id.has_value() ? std::to_string(tl_node_id.value()) : std::string{"n/a"}; | ||
| throw model_builder_error( | ||
| std::string{"Tree "} + std::to_string(tree_id) + ", Node " + node_id_repr + | ||
| ": Category index exceeds maximum representable value for this model's index type"); | ||
| } | ||
| auto max_cat_plus_one = static_cast<index_t>(cat_unsigned + 1u); |
Contributor
Author
There was a problem hiding this comment.
After a deeper look at the FIL and Treelite code base, I found the following invariants for the types cat_t and index_t:
cat_tis alwaysstd::uint32_t, since Treelite stores all category values using unsigned 32-bit integers.index_tis eitherstd::uint32_torstd::uint64_t, depending on whether the tree model stores 32-bit or 64-bit threshold values.
Using those invariants, we can massively simplify the overflow check.
static_assert(std::is_same_v<cat_t, std::uint32_t>, "Category value must be uint32_t");
static_assert(std::is_same_v<index_t, std::uint32_t> || std::is_same_v<index_t, std::uint64_t>,
"Index type in Tree node must be either uint32_t or uint64_t");
// Ensure that max_cat + 1 does not overflow
if constexpr (std::is_same_v<cat_t, index_t>) {
if (max_cat == std::numeric_limits<index_t>::max()) {
auto node_id_repr =
tl_node_id.has_value() ? std::to_string(tl_node_id.value()) : std::string{"n/a"};
throw model_builder_error(
std::string{"Tree "} + std::to_string(tree_id) + ", Node " + node_id_repr +
": Category index must be at most "
+ std::to_string(std::numeric_limits<index_t>::max() - 1));
}
}
auto max_cat_plus_one = static_cast<index_t>(max_cat) + index_t{1};
chyunsu3
changed the title
There was a problem hiding this comment.
🧹 Nitpick comments (1)
cpp/tests/sg/fil/treelite_importer_invalid_inputs.cpp (1)
29-62: Consider extracting shared Treelite model setup into a helper.The setup in Line 29-Line 62 and Line 77-Line 110 is duplicated; a small helper would make follow-up boundary cases easier to add and maintain.
Also applies to: 77-110
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cpp/tests/sg/fil/treelite_importer_invalid_inputs.cpp` around lines 29 - 62, Duplicate Treelite model construction occurs twice; extract it into a single helper (e.g., MakeTestTreeliteModel or BuildCategoricalTestModel) that encapsulates creating Metadata and TreeAnnotation, calling treelite::model_builder::GetModelBuilder, running StartTree/StartNode/CategoricalTest/LeafScalar/EndNode/EndTree and returning the committed model (the result of CommitModel or a shared_ptr to the builder); replace both in-test blocks with calls to that helper so tests reuse the same construction logic and make future boundary-case additions easier.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@cpp/tests/sg/fil/treelite_importer_invalid_inputs.cpp`:
- Around line 29-62: Duplicate Treelite model construction occurs twice; extract
it into a single helper (e.g., MakeTestTreeliteModel or
BuildCategoricalTestModel) that encapsulates creating Metadata and
TreeAnnotation, calling treelite::model_builder::GetModelBuilder, running
StartTree/StartNode/CategoricalTest/LeafScalar/EndNode/EndTree and returning the
committed model (the result of CommitModel or a shared_ptr to the builder);
replace both in-test blocks with calls to that helper so tests reuse the same
construction logic and make future boundary-case additions easier.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 94f159f8-8f91-47ae-aed8-6975a13dd693
📒 Files selected for processing (5)
cpp/include/cuml/fil/detail/decision_forest_builder.hppcpp/include/cuml/fil/detail/degenerate_trees.hppcpp/include/cuml/fil/treelite_importer.hppcpp/tests/CMakeLists.txtcpp/tests/sg/fil/treelite_importer_invalid_inputs.cpp
✅ Files skipped from review due to trivial changes (1)
- cpp/include/cuml/fil/detail/degenerate_trees.hpp
🚧 Files skipped from review as they are similar to previous changes (1)
- cpp/include/cuml/fil/treelite_importer.hpp
…est_builder methods - Added try-catch block in treelite_importer to provide detailed error messages for model import errors. - Removed unnecessary tree_id parameter from add_categorical_node and add_leaf_vector_node methods in decision_forest_builder. - Streamlined node addition logic in decision_forest_builder to improve clarity and maintainability.
This reverts commit 99d13f6.
- Introduced a new test file for decision forest builder to check for out-of-bounds errors in categorical storage offsets and bitset extents. - Updated CMakeLists.txt to include the new test file in the build configuration.
csadorf
requested changes
Apr 29, 2026
| result.element = element_op::sigmoid; | ||
| } else { | ||
| throw model_import_error{"Unrecognized Treelite pred_transform string"}; | ||
| throw unusable_model_exception{"Unrecognized Treelite pred_transform string"}; |
Contributor
There was a problem hiding this comment.
Why did you reclassify the exception type here?
(aside) I think in a follow-up we should consider merging these two exception classes. I don't really see the benefit in keeping both.
| void add_categorical_node( | ||
| iter_t vec_begin, | ||
| iter_t vec_end, | ||
| std::size_t tree_id, |
Contributor
There was a problem hiding this comment.
I don't think we should be passing the tree_id here. Instead we can throw the exception, catch it at the call-site (where tree_id is known), add the tree_id to the error message and then re-throw.
Comment on lines
+129
to
+130
| auto max_tree_index = (task_count - 1) / chunk_size; | ||
| if (max_tree_index < forest.tree_count()) { |
Contributor
There was a problem hiding this comment.
This check doesn't make any sense. The computation of max_tree_index is equivalent to forest.tree_count() - 1. We are thus checking that forest (forest.tree_count() - 1) < forest.tree_count() which is true for all finite .tree_count().
I think we should fully revert 99d13f6 since it does not add any benefit IMO.
Contributor
dantegd
approved these changes
Apr 29, 2026
| 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 |
Contributor
There was a problem hiding this comment.
Suggested change
| // Ensrue that index_t is unsigned. Bound checks below rely on index_t being unsigned | |
| // Ensure that index_t is unsigned. Bound checks below rely on index_t being unsigned |
csadorf
approved these changes
Apr 29, 2026
csadorf
requested review from
harrism
and removed request for
harrism and
robertmaynard
April 29, 2026 21:01
Contributor
|
/merge |
7 tasks
rapids-bot Bot
pushed a commit
to rapidsai/nvforest
that referenced
this pull request
May 21, 2026
Ports the Treelite input validation work from `NVIDIA/cuml#8016` to `nvforest`. Adds import-time checks for Treelite model values that would otherwise be packed into narrower `nvforest` node metadata or categorical storage without validation. The branch also consolidates model import failures under `model_import_error` and adds context to import errors so failures identify the source Treelite tree and node. ## Summary - Validates node feature IDs before packing them into node metadata. - Validates categorical split values and categorical storage bounds before constructing the final forest. - Checks floating-point downcasts for postprocessing constants instead of relying on unchecked narrowing. - Hardens shared helpers used by the import path, including `bitset` writes and `ceildiv`. - Adds invalid-input coverage for Treelite import and `decision_forest_builder`. - Renames local test variables from `fil_model` to `nvforest_model`. Follow-up to NVIDIA/cuml#8016 Authors: - Philip Hyunsu Cho (https://github.com/hcho3) - Simon Adorf (https://github.com/csadorf) Approvers: - Simon Adorf (https://github.com/csadorf) URL: #104
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Introduce bounds check to prevent integer overflow when parsing the Treelite input for consumption in FIL.