Skip to content

Add distributed RF C++ tree builder core #8255

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
RAMitchell wants to merge 6 commits into from
Closed

Add distributed RF C++ tree builder core #8255

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
b0024b0
Add distributed RF C++ tree builder core
RAMitchell Jun 10, 2026
13d7a95
Fix distributed RF multi-GPU tests
RAMitchell Jun 11, 2026
3f06f4e
Address distributed RF PR feedback
RAMitchell Jun 11, 2026
e21d08e
Merge remote-tracking branch 'upstream/main' into codex/distributed-r…
RAMitchell Jun 22, 2026
e409597
Merge branch 'main' into codex/distributed-rf-cpp-core
RAMitchell Jun 23, 2026
365871c
Merge remote-tracking branch 'upstream/main' into codex/distributed-r…
RAMitchell Jun 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
345 changes: 205 additions & 140 deletions cpp/src/decisiontree/batched-levelalgo/builder.cuh
View file
Open in desktop

Large diffs are not rendered by default.

201 changes: 152 additions & 49 deletions cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,12 +12,18 @@

#include <cuml/common/utils.hpp>

#include <raft/linalg/unary_op.cuh>

#include <cuda/iterator>
#include <cuda/std/random>
#include <thrust/execution_policy.h>
#include <thrust/for_each.h>
#include <thrust/iterator/counting_iterator.h>

#include <cstddef>
#include <cstdint>
#include <type_traits>

namespace ML {
namespace DT {

Expand All @@ -38,31 +44,23 @@ struct NodeWorkItem {
* This struct has information about workload of a single threadblock of
* computeSplit kernels of classification and regression
*/
template <typename IdxT>
struct WorkloadInfo {
IdxT nodeid; // Node in the batch on which the threadblock needs to work
IdxT large_nodeid; // counts only large nodes (nodes that require more than one block along x-dim
// for histogram calculation)
IdxT offset_blockid; // Offset threadblock id among all the blocks that are
// working on this node
IdxT num_blocks; // Total number of blocks that are working on the node
int nodeid; // Node in the batch on which the threadblock needs to work
int offset_blockid; // Offset threadblock id among all the blocks that are
// working on this node
int num_blocks; // Total number of blocks that are working on the node
};

template <typename SplitT, typename IdxT>
HDI bool SplitPartitionNotValid(const SplitT& split, IdxT min_samples_leaf, std::size_t num_rows)
HDI bool SplitPartitionNotValid(const SplitT& split, IdxT, std::size_t)
{
return split.colid == IdxT(-1) || split.nLeft < min_samples_leaf ||
(IdxT(num_rows) - split.nLeft) < min_samples_leaf;
return split.colid == -1;
}

template <typename SplitT, typename DataT, typename IdxT>
HDI bool SplitNotValid(const SplitT& split,
DataT min_impurity_decrease,
IdxT min_samples_leaf,
std::size_t num_rows)
template <typename SplitT, typename DataT>
HDI bool SplitNotValid(const SplitT& split, DataT min_impurity_decrease)
{
return split.best_metric_val <= min_impurity_decrease ||
SplitPartitionNotValid(split, min_samples_leaf, num_rows);
return split.colid == -1 || split.best_metric_val <= min_impurity_decrease;
}
Comment on lines +60 to 64

@coderabbitai coderabbitai Bot Jun 25, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟠 Major

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '\n== builder_kernels.cuh (around SplitNotValid) ==\n'
nl -ba cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh | sed -n '1,180p'

printf '\n== usages of SplitNotValid ==\n'
rg -n "SplitNotValid|best_metric_val" cpp/src/decisiontree -g '!**/build/**'

Repository: rapidsai/cuml

Length of output: 241

🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '\n== builder_kernels.cuh (around SplitNotValid) ==\n'
nl -ba cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh | sed -n '1,180p'

printf '\n== usages of SplitNotValid ==\n'
rg -n "SplitNotValid|best_metric_val" cpp/src/decisiontree -g '!**/build/**'

Repository: rapidsai/cuml

Length of output: 241

🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '\n== builder_kernels.cuh (around SplitNotValid) ==\n'
sed -n '1,180p' cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh | cat -n

printf '\n== usages of SplitNotValid ==\n'
rg -n "SplitNotValid|best_metric_val" cpp/src/decisiontree -g '!**/build/**'

Repository: rapidsai/cuml

Length of output: 10841

🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '\n== builder.cuh around line 100 ==\n'
sed -n '80,130p' cpp/src/decisiontree/batched-levelalgo/builder.cuh | cat -n

printf '\n== split.cuh around metric initialization/comparison ==\n'
sed -n '1,220p' cpp/src/decisiontree/batched-levelalgo/split.cuh | cat -n

Repository: rapidsai/cuml

Length of output: 10767

🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '\n== builder.cuh around line 100 ==\n'
sed -n '80,130p' cpp/src/decisiontree/batched-levelalgo/builder.cuh | cat -n

printf '\n== split.cuh around metric initialization/comparison ==\n'
sed -n '1,220p' cpp/src/decisiontree/batched-levelalgo/split.cuh | cat -n

Repository: rapidsai/cuml

Length of output: 10767

Handle NaN split gains in the builder path

cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh:61-63 treats NaN as valid, and cpp/src/decisiontree/batched-levelalgo/builder.cuh:21 has the same direct compare. That lets invalid gains slip through and be written into the tree; reject non-finite gains here or use a NaN-safe comparison in both places.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh` around
lines 60 - 64, The split validation in SplitNotValid currently only checks colid
and a direct best_metric_val comparison, so NaN gains can still pass as valid.
Update SplitNotValid in builder_kernels.cuh and the matching builder path check
in builder.cuh to reject non-finite split gains (or use a NaN-safe comparison)
before a split is accepted or written into the tree. Keep the fix centered on
SplitNotValid and the builder validation logic so both paths behave
consistently.

Source: Coding guidelines


/* Returns 'dataset' rounded up to a correctly-aligned pointer of type OutT* */
Expand Down Expand Up @@ -103,25 +101,33 @@ void sample_features(IdxT* column_samples,
}

template <typename DataT, typename LabelT, typename IdxT, int TPB>
void launchNodeSplitKernel(const IdxT min_samples_leaf,
const DataT min_impurity_decrease,
void launchNodeSplitKernel(const DataT min_impurity_decrease,
const Dataset<DataT, LabelT, IdxT>& dataset,
const NodeWorkItem* work_items,
const Split<DataT, IdxT>* splits,
const WorkloadInfo<IdxT>* workload_info,
Split<DataT>* splits,
const WorkloadInfo* workload_info,
size_t n_blocks_dimx,
IdxT* partition_row_ids,
cudaStream_t builder_stream);

template <typename DatasetT, typename NodeT, typename ObjectiveT, typename DataT>
void launchLeafKernel(ObjectiveT objective,
DatasetT& dataset,
const NodeT* tree,
const InstanceRange* instance_ranges,
DataT* leaves,
int batch_size,
size_t smem_size,
cudaStream_t builder_stream);
template <typename DatasetT, typename NodeT, typename ObjectiveT>
void launchLeafHistogramKernel(ObjectiveT objective,
DatasetT& dataset,
const NodeT* tree,
const InstanceRange* instance_ranges,
typename ObjectiveT::BinT* leaf_histograms,
int batch_size,
size_t smem_size,
cudaStream_t builder_stream);

template <typename NodeT, typename ObjectiveT, typename DataT>
void launchFinalizeLeafKernel(ObjectiveT objective,
const NodeT* tree,
const typename ObjectiveT::BinT* leaf_histograms,
DataT* leaves,
int batch_size,
int num_outputs,
cudaStream_t builder_stream);
// Returns the lowest index in `array` whose value is greater or equal to `element`.
// Values outside the quantile range are clamped to the edge bins: values below the
// first quantile return 0, and values above the last quantile return len - 1.
Expand All @@ -142,31 +148,128 @@ HDI IdxT lower_bound(DataT* array, IdxT len, DataT element)
return start;
}

template <typename DataT, typename LabelT, typename IdxT, int TPB, typename BinT>
void launchComputeSplitHistogramKernel(BinT* histograms,
IdxT max_n_bins,
const Dataset<DataT, LabelT, IdxT>& dataset,
const Quantiles<DataT, IdxT>& quantiles,
const NodeWorkItem* work_items,
IdxT colStart,
const IdxT* column_samples,
const WorkloadInfo* workload_info,
dim3 grid,
size_t smem_size,
cudaStream_t builder_stream);

template <typename DataT,
typename LabelT,
typename IdxT,
int TPB,
typename ObjectiveT,
typename BinT>
void launchComputeSplitKernel(BinT* histograms,
IdxT n_bins,
IdxT min_samples_split,
IdxT max_leaves,
const Dataset<DataT, LabelT, IdxT>& dataset,
const Quantiles<DataT, IdxT>& quantiles,
const NodeWorkItem* work_items,
IdxT colStart,
const IdxT* column_samples,
int* done_count,
int* mutex,
volatile Split<DataT, IdxT>* splits,
ObjectiveT& objective,
IdxT treeid,
const WorkloadInfo<IdxT>* workload_info,
uint64_t seed,
dim3 grid,
size_t smem_size,
cudaStream_t builder_stream);
void launchEvaluateSplitKernel(BinT* histograms,
IdxT max_n_bins,
const Dataset<DataT, LabelT, IdxT>& dataset,
const Quantiles<DataT, IdxT>& quantiles,
IdxT colStart,
const IdxT* column_samples,
int* mutex,
volatile Split<DataT>* splits,
ObjectiveT& objective,
dim3 grid,
size_t smem_size,
cudaStream_t builder_stream);

template <typename BinT>
inline constexpr bool has_label_sum_v =
std::is_same_v<BinT, RegressionBin> || std::is_same_v<BinT, WeightedRegressionBin>;

template <typename BinT>
inline constexpr bool has_weight_v =
std::is_same_v<BinT, WeightedClassificationBin> || std::is_same_v<BinT, WeightedRegressionBin>;

template <typename BinT>
inline void packHistograms(const BinT* in,
double* label_sums,
std::uint64_t* counts,
double* weights,
std::size_t len,
cudaStream_t stream)
{
if constexpr (has_label_sum_v<BinT>) {
auto label_sum_op = [in] __device__(double* out, std::size_t i) { *out = in[i].LabelSum(); };
raft::linalg::writeOnlyUnaryOp<double, decltype(label_sum_op), std::size_t, 256>(
label_sums, len, label_sum_op, stream);
}

auto count_op = [in] __device__(std::uint64_t* out, std::size_t i) { *out = in[i].Count(); };
raft::linalg::writeOnlyUnaryOp<std::uint64_t, decltype(count_op), std::size_t, 256>(
counts, len, count_op, stream);

if constexpr (has_weight_v<BinT>) {
auto weight_op = [in] __device__(double* out, std::size_t i) { *out = in[i].Weight(); };
raft::linalg::writeOnlyUnaryOp<double, decltype(weight_op), std::size_t, 256>(
weights, len, weight_op, stream);
}
}

inline void unpackHistograms(const double*,
const std::uint64_t* counts,
const double*,
ClassificationBin* out,
std::size_t len,
cudaStream_t stream)
{
auto op = [counts] __device__(ClassificationBin * out, std::size_t i) { out->count = counts[i]; };
raft::linalg::writeOnlyUnaryOp<ClassificationBin, decltype(op), std::size_t, 256>(
out, len, op, stream);
}

inline void unpackHistograms(const double*,
const std::uint64_t* counts,
const double* weights,
WeightedClassificationBin* out,
std::size_t len,
cudaStream_t stream)
{
auto op = [counts, weights] __device__(WeightedClassificationBin * out, std::size_t i) {
out->count = counts[i];
out->weight = weights[i];
};
raft::linalg::writeOnlyUnaryOp<WeightedClassificationBin, decltype(op), std::size_t, 256>(
out, len, op, stream);
}

inline void unpackHistograms(const double* label_sums,
const std::uint64_t* counts,
const double*,
RegressionBin* out,
std::size_t len,
cudaStream_t stream)
{
auto op = [label_sums, counts] __device__(RegressionBin * out, std::size_t i) {
out->label_sum = label_sums[i];
out->count = counts[i];
};
raft::linalg::writeOnlyUnaryOp<RegressionBin, decltype(op), std::size_t, 256>(
out, len, op, stream);
}

inline void unpackHistograms(const double* label_sums,
const std::uint64_t* counts,
const double* weights,
WeightedRegressionBin* out,
std::size_t len,
cudaStream_t stream)
{
auto op = [label_sums, counts, weights] __device__(WeightedRegressionBin * out, std::size_t i) {
out->label_sum = label_sums[i];
out->count = counts[i];
out->weight = weights[i];
};
raft::linalg::writeOnlyUnaryOp<WeightedRegressionBin, decltype(op), std::size_t, 256>(
out, len, op, stream);
}

} // namespace DT
} // namespace ML
Loading
Loading