diff --git a/cpp/src/decisiontree/batched-levelalgo/builder.cuh b/cpp/src/decisiontree/batched-levelalgo/builder.cuh index cf168dec03..8d7b305880 100644 --- a/cpp/src/decisiontree/batched-levelalgo/builder.cuh +++ b/cpp/src/decisiontree/batched-levelalgo/builder.cuh @@ -15,13 +15,16 @@ #include #include +#include #include #include #include +#include #include #include +#include #include #include @@ -41,12 +44,17 @@ class NodeQueue { std::deque work_items_; public: - NodeQueue(DecisionTreeParams params, size_t max_nodes, size_t sampled_rows, int num_outputs) + NodeQueue(DecisionTreeParams params, + size_t max_nodes, + size_t sampled_rows, + int num_outputs, + size_t global_sampled_rows = 0) : params(params), tree(std::make_shared>()) { + if (global_sampled_rows == 0) { global_sampled_rows = sampled_rows; } tree->num_outputs = num_outputs; tree->sparsetree.reserve(max_nodes); - tree->sparsetree.emplace_back(NodeT::CreateLeafNode(sampled_rows)); + tree->sparsetree.emplace_back(NodeT::CreateLeafNode(global_sampled_rows)); tree->leaf_counter = 1; tree->depth_counter = 0; node_instances_.reserve(max_nodes); @@ -76,55 +84,50 @@ class NodeQueue { bool IsExpandable(const NodeT& n, int depth) { if (depth >= params.max_depth) return false; - if (int(n.InstanceCount()) < params.min_samples_split) return false; + if (std::size_t(n.InstanceCount()) < std::size_t(params.min_samples_split)) return false; if (params.max_leaves != -1 && tree->leaf_counter >= params.max_leaves) return false; return true; } - template - void Push(const std::vector& work_items, SplitT* h_splits) + void Push(const std::vector& work_items, const Split* h_splits) { // Update node queue based on splits for (std::size_t i = 0; i < work_items.size(); i++) { - auto split = h_splits[i]; - auto item = work_items[i]; - auto parent_range = node_instances_.at(item.idx); - if (SplitNotValid( - split, params.min_impurity_decrease, params.min_samples_leaf, parent_range.count)) { - continue; - } + auto global_split = h_splits[i]; + auto item = work_items[i]; + auto parent_range = node_instances_.at(item.idx); + auto parent_global_count = std::size_t(tree->sparsetree.at(item.idx).InstanceCount()); + if (global_split.best_metric_val <= params.min_impurity_decrease) { continue; } if (params.max_leaves != -1 && tree->leaf_counter >= params.max_leaves) break; - // parent - tree->sparsetree.at(item.idx) = NodeT::CreateSplitNode(split.colid, - split.quesval, - split.best_metric_val, + auto left_global_count = std::size_t(global_split.global_nLeft); + auto right_global_count = parent_global_count - left_global_count; + auto left_local_count = std::size_t(global_split.local_nLeft); + auto right_local_count = parent_range.count - left_local_count; + + tree->sparsetree.at(item.idx) = NodeT::CreateSplitNode(global_split.colid, + global_split.quesval, + global_split.best_metric_val, int64_t(tree->sparsetree.size()), - parent_range.count); + parent_global_count); tree->leaf_counter++; - // left - tree->sparsetree.emplace_back(NodeT::CreateLeafNode(split.nLeft)); - node_instances_.emplace_back(InstanceRange{parent_range.begin, std::size_t(split.nLeft)}); - // Do not add a work item if this child is definitely a leaf + tree->sparsetree.emplace_back(NodeT::CreateLeafNode(left_global_count)); + node_instances_.emplace_back(InstanceRange{parent_range.begin, left_local_count}); if (this->IsExpandable(tree->sparsetree.back(), item.depth + 1)) { work_items_.emplace_back( NodeWorkItem{tree->sparsetree.size() - 1, item.depth + 1, node_instances_.back()}); } - // right - tree->sparsetree.emplace_back(NodeT::CreateLeafNode(parent_range.count - split.nLeft)); + tree->sparsetree.emplace_back(NodeT::CreateLeafNode(right_global_count)); node_instances_.emplace_back( - InstanceRange{parent_range.begin + split.nLeft, parent_range.count - split.nLeft}); - - // Do not add a work item if this child is definitely a leaf + InstanceRange{parent_range.begin + left_local_count, right_local_count}); if (this->IsExpandable(tree->sparsetree.back(), item.depth + 1)) { work_items_.emplace_back( NodeWorkItem{tree->sparsetree.size() - 1, item.depth + 1, node_instances_.back()}); } - // update depth tree->depth_counter = max(tree->depth_counter, item.depth + 1); } } @@ -140,7 +143,7 @@ struct Builder { typedef typename ObjectiveT::IdxT IdxT; typedef typename ObjectiveT::BinT BinT; typedef SparseTreeNode NodeT; - typedef Split SplitT; + typedef Split SplitT; typedef Dataset DatasetT; typedef Quantiles QuantilesT; @@ -160,12 +163,8 @@ struct Builder { IdxT treeid; /** Seed used for randomization */ uint64_t seed; - /** number of nodes created in the current batch */ - IdxT* n_nodes; /** buffer of segmented histograms*/ BinT* histograms; - /** threadblock arrival count */ - int* done_count; /** mutex array used for atomically updating best split */ int* mutex; /** best splits for the current batch of nodes */ @@ -173,13 +172,15 @@ struct Builder { /** current batch of nodes */ NodeWorkItem* d_work_items; /** device AOS to map CTA blocks along dimx to nodes of a batch */ - WorkloadInfo* workload_info; + WorkloadInfo* workload_info; /** host AOS to map CTA blocks along dimx to nodes of a batch */ - WorkloadInfo* h_workload_info; + WorkloadInfo* h_workload_info; /** maximum CTA blocks along dimx */ int max_blocks_dimx = 0; /** host array of splits */ SplitT* h_splits; + /** packed histogram buffer used by distributed all-reduce */ + void* packed_histograms; /** number of blocks used to parallelize column-wise computations */ int n_blks_for_cols = 10; /** Memory alignment value */ @@ -191,6 +192,10 @@ struct Builder { rmm::device_uvector d_buff; /** pinned host buffer to store the trained nodes */ ML::pinned_host_vector h_buff; + /** true when a communicator with more than one rank is available */ + bool distributed; + /** global root sample count in distributed mode */ + std::size_t global_sampled_rows; Builder(const raft::handle_t& handle, cudaStream_t s, @@ -218,8 +223,25 @@ struct Builder { row_ids->data(), n_classes}, quantiles(q), - d_buff(0, builder_stream) + d_buff(0, builder_stream), + distributed(raft::resource::comms_initialized(handle) && handle.get_comms().get_size() > 1), + global_sampled_rows(row_ids->size()) { + if (distributed) { + // Each rank stores only its local row_ids, but tree metadata and split + // validity checks need the global bootstrap sample count. + auto local_count = static_cast(row_ids->size()); + rmm::device_uvector count_buffer(1, builder_stream); + raft::update_device(count_buffer.data(), &local_count, 1, builder_stream); + handle.get_comms().allreduce( + count_buffer.data(), count_buffer.data(), 1, raft::comms::op_t::SUM, builder_stream); + ASSERT(handle.get_comms().sync_stream(builder_stream) == raft::comms::status_t::SUCCESS, + "An error occurred in the distributed RF row-count all-reduce."); + auto global_count = std::uint64_t{0}; + raft::update_host(&global_count, count_buffer.data(), 1, builder_stream); + handle.sync_stream(builder_stream); + global_sampled_rows = global_count; + } max_blocks_dimx = 1 + params.max_batch_size + dataset.n_sampled_rows / TPB_DEFAULT; ASSERT(q.quantiles_array != nullptr && q.n_bins_array != nullptr, "Currently quantiles need to be computed before this call!"); @@ -262,6 +284,15 @@ struct Builder { * * @return a pair of device workspace and host workspace size requirements */ + template + size_t packedHistogramWorkspaceSize(size_t len) const + { + size_t size = calculateAlignedBytes(sizeof(std::uint64_t) * len); + if constexpr (has_label_sum_v) { size += calculateAlignedBytes(sizeof(double) * len); } + if constexpr (has_weight_v) { size += calculateAlignedBytes(sizeof(double) * len); } + return size; + } + auto workspaceSize() const { size_t d_wsize = 0, h_wsize = 0; @@ -270,22 +301,21 @@ struct Builder { size_t max_len_histograms = max_batch * params.max_n_bins * n_blks_for_cols * dataset.num_outputs; - d_wsize += calculateAlignedBytes(sizeof(IdxT)); // n_nodes - d_wsize += calculateAlignedBytes(sizeof(BinT) * max_len_histograms); // histograms - d_wsize += calculateAlignedBytes(sizeof(int) * max_batch * n_blks_for_cols); // done_count - d_wsize += calculateAlignedBytes(sizeof(int) * max_batch); // mutex - d_wsize += calculateAlignedBytes(sizeof(SplitT) * max_batch); // splits - d_wsize += calculateAlignedBytes(sizeof(NodeWorkItem) * max_batch); // d_work_Items - d_wsize += // workload_info - calculateAlignedBytes(sizeof(WorkloadInfo) * max_blocks_dimx); + d_wsize += calculateAlignedBytes(sizeof(BinT) * max_len_histograms); // histograms + d_wsize += calculateAlignedBytes(sizeof(int) * max_batch); // mutex + d_wsize += calculateAlignedBytes(sizeof(SplitT) * max_batch); // splits + d_wsize += calculateAlignedBytes(sizeof(NodeWorkItem) * max_batch); // d_work_Items + d_wsize += // workload_info + calculateAlignedBytes(sizeof(WorkloadInfo) * max_blocks_dimx); d_wsize += calculateAlignedBytes(sizeof(IdxT) * max_batch * dataset.n_sampled_cols); // column_samples d_wsize += calculateAlignedBytes(sizeof(IdxT) * dataset.n_sampled_rows); // partition row IDs // all nodes in the tree h_wsize += // h_workload_info - calculateAlignedBytes(sizeof(WorkloadInfo) * max_blocks_dimx); + calculateAlignedBytes(sizeof(WorkloadInfo) * max_blocks_dimx); h_wsize += calculateAlignedBytes(sizeof(SplitT) * max_batch); // splits + d_wsize += packedHistogramWorkspaceSize(max_len_histograms); return std::make_pair(d_wsize, h_wsize); } @@ -301,39 +331,33 @@ struct Builder { { raft::common::nvtx::range fun_scope( "Builder::assignWorkspace @builder.cuh [batched-levelalgo]"); - auto max_batch = params.max_batch_size; - auto n_col_blks = n_blks_for_cols; + auto max_batch = params.max_batch_size; size_t max_len_histograms = max_batch * (params.max_n_bins) * n_blks_for_cols * dataset.num_outputs; // device - n_nodes = reinterpret_cast(d_wspace); - d_wspace += calculateAlignedBytes(sizeof(IdxT)); histograms = reinterpret_cast(d_wspace); d_wspace += calculateAlignedBytes(sizeof(BinT) * max_len_histograms); - done_count = reinterpret_cast(d_wspace); - d_wspace += calculateAlignedBytes(sizeof(int) * max_batch * n_col_blks); mutex = reinterpret_cast(d_wspace); d_wspace += calculateAlignedBytes(sizeof(int) * max_batch); splits = reinterpret_cast(d_wspace); d_wspace += calculateAlignedBytes(sizeof(SplitT) * max_batch); d_work_items = reinterpret_cast(d_wspace); d_wspace += calculateAlignedBytes(sizeof(NodeWorkItem) * max_batch); - workload_info = reinterpret_cast*>(d_wspace); - d_wspace += calculateAlignedBytes(sizeof(WorkloadInfo) * max_blocks_dimx); + workload_info = reinterpret_cast(d_wspace); + d_wspace += calculateAlignedBytes(sizeof(WorkloadInfo) * max_blocks_dimx); column_samples = reinterpret_cast(d_wspace); d_wspace += calculateAlignedBytes(sizeof(IdxT) * max_batch * dataset.n_sampled_cols); partition_row_ids = reinterpret_cast(d_wspace); d_wspace += calculateAlignedBytes(sizeof(IdxT) * dataset.n_sampled_rows); - RAFT_CUDA_TRY( - cudaMemsetAsync(done_count, 0, sizeof(int) * max_batch * n_col_blks, builder_stream)); RAFT_CUDA_TRY(cudaMemsetAsync(mutex, 0, sizeof(int) * max_batch, builder_stream)); // host - h_workload_info = reinterpret_cast*>(h_wspace); - h_wspace += calculateAlignedBytes(sizeof(WorkloadInfo) * max_blocks_dimx); + h_workload_info = reinterpret_cast(h_wspace); + h_wspace += calculateAlignedBytes(sizeof(WorkloadInfo) * max_blocks_dimx); h_splits = reinterpret_cast(h_wspace); h_wspace += calculateAlignedBytes(sizeof(SplitT) * max_batch); + packed_histograms = reinterpret_cast(d_wspace); } /** @@ -346,11 +370,10 @@ struct Builder { raft::common::nvtx::range fun_scope("Builder::train @builder.cuh [batched-levelalgo]"); MLCommon::TimerCPU timer; NodeQueue queue( - params, this->maxNodes(), dataset.n_sampled_rows, dataset.num_outputs); + params, this->maxNodes(), dataset.n_sampled_rows, dataset.num_outputs, global_sampled_rows); while (queue.HasWork()) { - auto work_items = queue.Pop(); - auto [splits_host_ptr, splits_count] = doSplit(work_items); - queue.Push(work_items, splits_host_ptr); + auto work_items = queue.Pop(); + queue.Push(work_items, doSplit(work_items)); } auto tree = queue.GetTree(); this->SetLeafPredictions(tree, queue.GetInstanceRanges()); @@ -361,54 +384,33 @@ struct Builder { private: auto updateWorkloadInfo(const std::vector& work_items) { - int n_large_nodes = 0; // large nodes are nodes having training instances larger than block - // size, hence require global memory for histogram construction - int n_blocks_dimx = 0; // gridDim.x required for computeSplitKernel + int n_blocks_dimx = 0; // gridDim.x required for split histogram construction for (std::size_t i = 0; i < work_items.size(); i++) { - auto item = work_items[i]; - int n_blocks_per_node = - std::max(raft::ceildiv(item.instances.count, size_t(TPB_DEFAULT)), size_t(1)); - - if (n_blocks_per_node > 1) ++n_large_nodes; + auto item = work_items[i]; + const auto node_id = static_cast(i); + int n_blocks_per_node = static_cast( + std::max(raft::ceildiv(item.instances.count, size_t(TPB_DEFAULT)), size_t(1))); for (int b = 0; b < n_blocks_per_node; b++) { - h_workload_info[n_blocks_dimx + b] = {int(i), n_large_nodes - 1, b, n_blocks_per_node}; + h_workload_info[n_blocks_dimx + b] = {node_id, b, n_blocks_per_node}; } n_blocks_dimx += n_blocks_per_node; } raft::update_device(workload_info, h_workload_info, n_blocks_dimx, builder_stream); - return std::make_pair(n_blocks_dimx, n_large_nodes); + return n_blocks_dimx; } auto doSplit(const std::vector& work_items) { raft::common::nvtx::range fun_scope("Builder::doSplit @builder.cuh [batched-levelalgo]"); - // start fresh on the number of *new* nodes created in this batch - RAFT_CUDA_TRY(cudaMemsetAsync(n_nodes, 0, sizeof(IdxT), builder_stream)); - const IdxT original_n_sampled_cols = dataset.n_sampled_cols; ASSERT(original_n_sampled_cols > 0 && original_n_sampled_cols <= dataset.N, "n_sampled_cols must be in [1, n_cols]"); const std::size_t max_sampling_rounds = std::size_t((dataset.N + original_n_sampled_cols - 1) / original_n_sampled_cols); - struct HostSplit { - DataT quesval; - IdxT colid; - DataT best_metric_val; - int nLeft; - }; - static_assert(sizeof(HostSplit) == sizeof(SplitT)); - static_assert(alignof(HostSplit) == alignof(SplitT)); - - // The final split chosen for each original work item. Nodes that need - // additional feature samples are compacted in active_items, so successful - // splits must be copied back to their original batch position. - std::vector final_splits(work_items.size()); - // Current retry batch. It starts as the full batch and shrinks to only - // nodes whose sampled features did not produce a valid split. + + std::vector final_splits(work_items.size()); std::vector active_items(work_items); - // active_items[i] maps back to the corresponding index in the original - // work_items/final_splits arrays. std::vector active_to_original(work_items.size()); for (std::size_t i = 0; i < active_to_original.size(); ++i) { active_to_original[i] = i; @@ -425,8 +427,7 @@ struct Builder { std::vector retry_to_original; for (std::size_t i = 0; i < active_items.size(); ++i) { const auto original_idx = active_to_original[i]; - final_splits[original_idx] = HostSplit{ - h_splits[i].quesval, h_splits[i].colid, h_splits[i].best_metric_val, h_splits[i].nLeft}; + final_splits[original_idx] = h_splits[i]; if (SplitPartitionNotValid( h_splits[i], params.min_samples_leaf, active_items[i].instances.count)) { retry_items.push_back(active_items[i]); @@ -440,8 +441,6 @@ struct Builder { } dataset.n_sampled_cols = original_n_sampled_cols; - // Partition samples once, using the valid split found for each node. Nodes - // still without a valid split after all features have been visited remain leaves. RAFT_CUDA_TRY(cudaMemcpyAsync(splits, final_splits.data(), sizeof(SplitT) * work_items.size(), @@ -450,20 +449,19 @@ struct Builder { raft::update_device(d_work_items, work_items.data(), work_items.size(), builder_stream); const auto partition_workload = this->updateWorkloadInfo(work_items); raft::common::nvtx::push_range("nodeSplitKernel @builder.cuh [batched-levelalgo]"); - launchNodeSplitKernel(params.min_samples_leaf, - params.min_impurity_decrease, + launchNodeSplitKernel(params.min_impurity_decrease, dataset, d_work_items, splits, workload_info, - partition_workload.first, + partition_workload, partition_row_ids, builder_stream); RAFT_CUDA_TRY(cudaPeekAtLastError()); raft::common::nvtx::pop_range(); raft::update_host(h_splits, splits, work_items.size(), builder_stream); handle.sync_stream(builder_stream); - return std::make_tuple(h_splits, work_items.size()); + return h_splits; } void computeBestSplits(const std::vector& work_items, @@ -471,16 +469,14 @@ struct Builder { IdxT sample_offset) { initSplit(splits, work_items.size(), builder_stream); - RAFT_CUDA_TRY(cudaMemsetAsync( - done_count, 0, sizeof(int) * params.max_batch_size * n_blks_for_cols, builder_stream)); RAFT_CUDA_TRY(cudaMemsetAsync(mutex, 0, sizeof(int) * params.max_batch_size, builder_stream)); raft::update_device(d_work_items, work_items.data(), work_items.size(), builder_stream); - auto [n_blocks_dimx, n_large_nodes] = this->updateWorkloadInfo(work_items); + auto n_blocks_dimx = this->updateWorkloadInfo(work_items); sampleFeatures(work_items, sampling_seed, sample_offset); for (IdxT c = 0; c < dataset.n_sampled_cols; c += n_blks_for_cols) { - computeSplit(c, n_blocks_dimx, n_large_nodes); + computeSplit(c, n_blocks_dimx, work_items.size()); RAFT_CUDA_TRY(cudaPeekAtLastError()); } raft::update_host(h_splits, splits, work_items.size(), builder_stream); @@ -508,11 +504,9 @@ struct Builder { { size_t smem_size_1 = params.max_n_bins * dataset.num_outputs * sizeof(BinT) + // shared_histogram size - params.max_n_bins * sizeof(DataT) + // shared_quantiles size - sizeof(int); // shared_done size - // Extra room for alignment (see alignPointer in - // computeSplitKernel) - smem_size_1 += sizeof(DataT) + 3 * sizeof(int); + params.max_n_bins * sizeof(DataT); // shared_quantiles size + // Extra room for alignment (see alignPointer in the split kernels). + smem_size_1 += sizeof(DataT); // Calculate the shared memory needed for evalBestSplit size_t smem_size_2 = raft::ceildiv(TPB_DEFAULT, raft::WarpSize) * sizeof(SplitT); // Pick the max of two @@ -522,7 +516,57 @@ struct Builder { return smem_size; } - void computeSplit(IdxT col, size_t n_blocks_dimx, size_t n_large_nodes) + void allReduceHistograms(BinT* histograms_to_reduce, std::size_t len_histograms) + { + auto const& comm = handle.get_comms(); + auto* packed_base = reinterpret_cast(packed_histograms); + double* packed_label_sums = nullptr; + if constexpr (has_label_sum_v) { + packed_label_sums = reinterpret_cast(packed_base); + packed_base += calculateAlignedBytes(sizeof(double) * len_histograms); + } + auto* packed_counts = reinterpret_cast(packed_base); + packed_base += calculateAlignedBytes(sizeof(std::uint64_t) * len_histograms); + double* packed_weights = nullptr; + if constexpr (has_weight_v) { packed_weights = reinterpret_cast(packed_base); } + + packHistograms(histograms_to_reduce, + packed_label_sums, + packed_counts, + packed_weights, + len_histograms, + builder_stream); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + if constexpr (has_label_sum_v) { + comm.allreduce(packed_label_sums, + packed_label_sums, + len_histograms, + raft::comms::op_t::SUM, + builder_stream); + ASSERT(comm.sync_stream(builder_stream) == raft::comms::status_t::SUCCESS, + "An error occurred in the distributed RF label-sum histogram all-reduce."); + } + comm.allreduce( + packed_counts, packed_counts, len_histograms, raft::comms::op_t::SUM, builder_stream); + ASSERT(comm.sync_stream(builder_stream) == raft::comms::status_t::SUCCESS, + "An error occurred in the distributed RF count histogram all-reduce."); + if constexpr (has_weight_v) { + comm.allreduce( + packed_weights, packed_weights, len_histograms, raft::comms::op_t::SUM, builder_stream); + ASSERT(comm.sync_stream(builder_stream) == raft::comms::status_t::SUCCESS, + "An error occurred in the distributed RF weight histogram all-reduce."); + } + unpackHistograms(packed_label_sums, + packed_counts, + packed_weights, + histograms_to_reduce, + len_histograms, + builder_stream); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + + void computeSplit(IdxT col, size_t n_blocks_dimx, size_t work_items_size) { // if no instances to split, return if (n_blocks_dimx == 0) return; @@ -535,32 +579,37 @@ struct Builder { auto smem_size = computeSplitSmemSize(); dim3 grid(n_blocks_dimx, n_blocks_dimy, 1); // required total length (in bins) of the global segmented histograms over all - // classes, features and (large)nodes. - int len_histograms = n_bins * n_classes * n_blocks_dimy * n_large_nodes; + // classes, features and nodes. + int len_histograms = n_bins * n_classes * n_blocks_dimy * work_items_size; RAFT_CUDA_TRY(cudaMemsetAsync(histograms, 0, sizeof(BinT) * len_histograms, builder_stream)); - // create the objective function object + raft::common::nvtx::range kernel_scope("split kernels @builder.cuh [batched-levelalgo]"); + launchComputeSplitHistogramKernel(histograms, + params.max_n_bins, + dataset, + quantiles, + d_work_items, + col, + column_samples, + workload_info, + grid, + smem_size, + builder_stream); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + if (distributed) { allReduceHistograms(histograms, len_histograms); } ObjectiveT objective(dataset.num_outputs, params.min_samples_leaf, params.split_criterion); - // call the computeSplitKernel - raft::common::nvtx::range kernel_scope("computeSplitKernel @builder.cuh [batched-levelalgo]"); - launchComputeSplitKernel(histograms, - params.max_n_bins, - params.min_samples_split, - params.max_leaves, - dataset, - quantiles, - d_work_items, - col, - column_samples, - done_count, - mutex, - splits, - objective, - treeid, - workload_info, - seed, - grid, - smem_size, - builder_stream); + dim3 eval_grid(work_items_size, n_blocks_dimy, 1); + launchEvaluateSplitKernel(histograms, + params.max_n_bins, + dataset, + quantiles, + col, + column_samples, + mutex, + splits, + objective, + eval_grid, + smem_size, + builder_stream); } // Set the leaf value predictions in batch @@ -574,6 +623,8 @@ struct Builder { std::size_t max_batch_size = min(std::size_t(100000), tree->sparsetree.size()); rmm::device_uvector d_tree(max_batch_size, builder_stream); rmm::device_uvector d_instance_ranges(max_batch_size, builder_stream); + rmm::device_uvector d_leaf_histograms(max_batch_size * dataset.num_outputs, + builder_stream); rmm::device_uvector d_leaves(max_batch_size * dataset.num_outputs, builder_stream); ObjectiveT objective(dataset.num_outputs, params.min_samples_leaf, params.split_criterion); @@ -586,17 +637,31 @@ struct Builder { raft::update_device( d_instance_ranges.data(), instance_ranges.data() + batch_begin, batch_size, builder_stream); + RAFT_CUDA_TRY(cudaMemsetAsync( + d_leaf_histograms.data(), 0, sizeof(BinT) * d_leaf_histograms.size(), builder_stream)); RAFT_CUDA_TRY( cudaMemsetAsync(d_leaves.data(), 0, sizeof(DataT) * d_leaves.size(), builder_stream)); size_t smem_size = sizeof(BinT) * dataset.num_outputs; - launchLeafKernel(objective, - dataset, - d_tree.data(), - d_instance_ranges.data(), - d_leaves.data(), - batch_size, - smem_size, - builder_stream); + launchLeafHistogramKernel(objective, + dataset, + d_tree.data(), + d_instance_ranges.data(), + d_leaf_histograms.data(), + batch_size, + smem_size, + builder_stream); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + if (distributed) { + allReduceHistograms(d_leaf_histograms.data(), batch_size * dataset.num_outputs); + } + launchFinalizeLeafKernel(objective, + d_tree.data(), + d_leaf_histograms.data(), + d_leaves.data(), + batch_size, + dataset.num_outputs, + builder_stream); + RAFT_CUDA_TRY(cudaPeekAtLastError()); raft::update_host(tree->vector_leaf.data() + batch_begin * dataset.num_outputs, d_leaves.data(), batch_size * dataset.num_outputs, diff --git a/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh b/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh index 5251f3cc61..fcc783be57 100644 --- a/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh +++ b/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh @@ -12,12 +12,18 @@ #include +#include + #include #include #include #include #include +#include +#include +#include + namespace ML { namespace DT { @@ -38,31 +44,23 @@ struct NodeWorkItem { * This struct has information about workload of a single threadblock of * computeSplit kernels of classification and regression */ -template 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 -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 -HDI bool SplitNotValid(const SplitT& split, - DataT min_impurity_decrease, - IdxT min_samples_leaf, - std::size_t num_rows) +template +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; } /* Returns 'dataset' rounded up to a correctly-aligned pointer of type OutT* */ @@ -103,25 +101,33 @@ void sample_features(IdxT* column_samples, } template -void launchNodeSplitKernel(const IdxT min_samples_leaf, - const DataT min_impurity_decrease, +void launchNodeSplitKernel(const DataT min_impurity_decrease, const Dataset& dataset, const NodeWorkItem* work_items, - const Split* splits, - const WorkloadInfo* workload_info, + Split* splits, + const WorkloadInfo* workload_info, size_t n_blocks_dimx, IdxT* partition_row_ids, cudaStream_t builder_stream); -template -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 +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 +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. @@ -142,31 +148,128 @@ HDI IdxT lower_bound(DataT* array, IdxT len, DataT element) return start; } +template +void launchComputeSplitHistogramKernel(BinT* histograms, + IdxT max_n_bins, + const Dataset& dataset, + const Quantiles& 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 -void launchComputeSplitKernel(BinT* histograms, - IdxT n_bins, - IdxT min_samples_split, - IdxT max_leaves, - const Dataset& dataset, - const Quantiles& quantiles, - const NodeWorkItem* work_items, - IdxT colStart, - const IdxT* column_samples, - int* done_count, - int* mutex, - volatile Split* splits, - ObjectiveT& objective, - IdxT treeid, - const WorkloadInfo* 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& dataset, + const Quantiles& quantiles, + IdxT colStart, + const IdxT* column_samples, + int* mutex, + volatile Split* splits, + ObjectiveT& objective, + dim3 grid, + size_t smem_size, + cudaStream_t builder_stream); + +template +inline constexpr bool has_label_sum_v = + std::is_same_v || std::is_same_v; + +template +inline constexpr bool has_weight_v = + std::is_same_v || std::is_same_v; + +template +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) { + auto label_sum_op = [in] __device__(double* out, std::size_t i) { *out = in[i].LabelSum(); }; + raft::linalg::writeOnlyUnaryOp( + 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( + counts, len, count_op, stream); + + if constexpr (has_weight_v) { + auto weight_op = [in] __device__(double* out, std::size_t i) { *out = in[i].Weight(); }; + raft::linalg::writeOnlyUnaryOp( + 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( + 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( + 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( + 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( + out, len, op, stream); +} } // namespace DT } // namespace ML diff --git a/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels_impl.cuh b/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels_impl.cuh index f2c4035685..e4e2cf249b 100644 --- a/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels_impl.cuh +++ b/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels_impl.cuh @@ -22,6 +22,7 @@ #include #include +#include #include namespace ML { @@ -29,17 +30,17 @@ namespace DT { static constexpr int TPB_DEFAULT = 128; -template +template struct NodeSplitPartitionState { - IdxT left_count; + CountT left_count; bool valid_row; bool goes_left; }; -template +template struct NodeSplitPartitionScanOp { - __host__ __device__ NodeSplitPartitionState operator()( - const NodeSplitPartitionState& lhs, const NodeSplitPartitionState& rhs) const + __host__ __device__ NodeSplitPartitionState operator()( + const NodeSplitPartitionState& lhs, const NodeSplitPartitionState& rhs) const { return {lhs.left_count + rhs.left_count, rhs.valid_row, rhs.goes_left}; } @@ -51,14 +52,16 @@ struct NodeSplitPartitionScanOp { // partition buffer. template struct NodeSplitPartitionWriter { + using CountT = typename Split::CountT; + Dataset dataset; const NodeWorkItem* work_items; - const Split* splits; - const WorkloadInfo* workload_info; + const Split* splits; + const WorkloadInfo* workload_info; IdxT* partition_row_ids; __host__ __device__ void operator()(std::ptrdiff_t index, - NodeSplitPartitionState state) const + NodeSplitPartitionState state) const { if (!state.valid_row) { return; } @@ -71,32 +74,54 @@ struct NodeSplitPartitionWriter { const auto range_start = work_item.instances.begin; const auto range_pos = std::size_t(workload_info_cta.offset_blockid) * TPB + slot % TPB; - const auto row = dataset.row_ids[range_start + range_pos]; - const auto rank = - state.goes_left ? state.left_count - IdxT(1) : IdxT(range_pos) - state.left_count; - const auto out_idx = range_start + (state.goes_left ? rank : split.nLeft + rank); - partition_row_ids[out_idx] = row; + const auto row = dataset.row_ids[range_start + range_pos]; + const auto rank = state.goes_left ? std::size_t(state.left_count - CountT{1}) + : range_pos - std::size_t(state.left_count); + const auto local_left_count = std::size_t(split.local_nLeft); + const auto out_idx = range_start + (state.goes_left ? rank : local_left_count + rank); + partition_row_ids[out_idx] = row; } }; +template +static __global__ void nodeSplitLocalCountKernel(const DataT min_impurity_decrease, + const Dataset dataset, + const NodeWorkItem* work_items, + Split* splits, + const WorkloadInfo* workload_info) +{ + using CountT = typename Split::CountT; + + const auto workload_info_cta = workload_info[blockIdx.x]; + const auto nid = workload_info_cta.nodeid; + const auto work_item = work_items[nid]; + const auto split = splits[nid]; + if (SplitNotValid(split, min_impurity_decrease)) { return; } + + const auto range_pos = std::size_t(workload_info_cta.offset_blockid) * blockDim.x + threadIdx.x; + if (range_pos >= work_item.instances.count) { return; } + + const auto row = dataset.row_ids[work_item.instances.begin + range_pos]; + const auto col_idx = std::size_t(split.colid) * dataset.M + row; + const auto goes_left = dataset.data[col_idx] <= split.quesval; + if (goes_left) { atomicAdd(&splits[nid].local_nLeft, CountT{1}); } +} + // Copy back only ranges for nodes that actually split. Leaf/invalid nodes keep // their existing row-id order because the scan writer skips them too. template -static __global__ void nodeSplitCopyBackKernel(const IdxT min_samples_leaf, - const DataT min_impurity_decrease, +static __global__ void nodeSplitCopyBackKernel(const DataT min_impurity_decrease, const Dataset dataset, const NodeWorkItem* work_items, - const Split* splits, - const WorkloadInfo* workload_info, + const Split* splits, + const WorkloadInfo* workload_info, const IdxT* partition_row_ids) { const auto workload_info_cta = workload_info[blockIdx.x]; const auto nid = workload_info_cta.nodeid; const auto work_item = work_items[nid]; const auto split = splits[nid]; - if (SplitNotValid(split, min_impurity_decrease, min_samples_leaf, work_item.instances.count)) { - return; - } + if (SplitNotValid(split, min_impurity_decrease)) { return; } const auto range_start = work_item.instances.begin; const auto range_len = work_item.instances.count; @@ -108,18 +133,21 @@ static __global__ void nodeSplitCopyBackKernel(const IdxT min_samples_leaf, } template -void launchNodeSplitKernel(const IdxT min_samples_leaf, - const DataT min_impurity_decrease, +void launchNodeSplitKernel(const DataT min_impurity_decrease, const Dataset& dataset, const NodeWorkItem* work_items, - const Split* splits, - const WorkloadInfo* workload_info, + Split* splits, + const WorkloadInfo* workload_info, size_t n_blocks_dimx, IdxT* partition_row_ids, cudaStream_t builder_stream) { if (n_blocks_dimx == 0) return; + using CountT = typename Split::CountT; + nodeSplitLocalCountKernel<<>>( + min_impurity_decrease, dataset, work_items, splits, workload_info); + // Each slot corresponds to one thread lane in the tiled workload_info layout. // workload_info is grouped by node, so scan-by-key resets ranks at node boundaries. const auto n_slots = n_blocks_dimx * TPB; @@ -134,19 +162,19 @@ void launchNodeSplitKernel(const IdxT min_samples_leaf, const auto nid = workload_info_cta.nodeid; const auto work_item = work_items[nid]; const auto split = splits[nid]; - if (SplitNotValid(split, min_impurity_decrease, min_samples_leaf, work_item.instances.count)) { - return NodeSplitPartitionState{IdxT(0), false, false}; + if (SplitNotValid(split, min_impurity_decrease)) { + return NodeSplitPartitionState{CountT{0}, false, false}; } const auto range_pos = std::size_t(workload_info_cta.offset_blockid) * TPB + slot % TPB; if (range_pos >= work_item.instances.count) { - return NodeSplitPartitionState{IdxT(0), false, false}; + return NodeSplitPartitionState{CountT{0}, false, false}; } const auto row = dataset.row_ids[work_item.instances.begin + range_pos]; const auto col_idx = std::size_t(split.colid) * dataset.M + row; const auto goes_left = dataset.data[col_idx] <= split.quesval; - return NodeSplitPartitionState{goes_left ? IdxT(1) : IdxT(0), true, goes_left}; + return NodeSplitPartitionState{goes_left ? CountT{1} : CountT{0}, true, goes_left}; }; // The scan input is a stream of per-slot partition states keyed by node id. @@ -162,26 +190,20 @@ void launchNodeSplitKernel(const IdxT min_samples_leaf, node_keys + n_slots, partition_states, partition_writer, - thrust::equal_to{}, - NodeSplitPartitionScanOp{}); + thrust::equal_to{}, + NodeSplitPartitionScanOp{}); // The original row_ids buffer remains the source during the scan, so copy back after it finishes. - nodeSplitCopyBackKernel - <<>>(min_samples_leaf, - min_impurity_decrease, - dataset, - work_items, - splits, - workload_info, - partition_row_ids); + nodeSplitCopyBackKernel<<>>( + min_impurity_decrease, dataset, work_items, splits, workload_info, partition_row_ids); } -template +template static __global__ void leafKernel(ObjectiveT objective, DatasetT dataset, const NodeT* tree, const InstanceRange* instance_ranges, - DataT* leaves) + typename ObjectiveT::BinT* leaf_histograms) { using BinT = typename ObjectiveT::BinT; extern __shared__ char shared_memory[]; @@ -201,24 +223,59 @@ static __global__ void leafKernel(ObjectiveT objective, } __syncthreads(); if (tid == 0) { - ObjectiveT::SetLeafVector( - histogram, dataset.num_outputs, leaves + dataset.num_outputs * node_id); + auto leaf_histogram = leaf_histograms + dataset.num_outputs * node_id; + for (int i = 0; i < dataset.num_outputs; ++i) { + leaf_histogram[i] = histogram[i]; + } + } +} + +template +static __global__ void finalizeLeafKernel(ObjectiveT objective, + const NodeT* tree, + const typename ObjectiveT::BinT* leaf_histograms, + DataT* leaves, + int num_outputs) +{ + auto node_id = blockIdx.x; + auto leaf = leaves + num_outputs * node_id; + auto& node = tree[node_id]; + if (!node.IsLeaf()) { + for (int i = 0; i < num_outputs; ++i) { + leaf[i] = DataT(0); + } + return; } + auto leaf_histogram = leaf_histograms + num_outputs * node_id; + ObjectiveT::SetLeafVector(leaf_histogram, num_outputs, leaf); } -template -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 +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) +{ + finalizeLeafKernel<<>>( + objective, tree, leaf_histograms, leaves, num_outputs); +} + +template +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) { int num_blocks = batch_size; leafKernel<<>>( - objective, dataset, tree, instance_ranges, leaves); + objective, dataset, tree, instance_ranges, leaf_histograms); } /** @@ -251,216 +308,232 @@ DI BinT pdf_to_cdf(BinT* shared_histogram, IdxT n_bins) return total_aggregate; } -template -static __global__ void computeSplitKernel(BinT* histograms, - IdxT max_n_bins, - IdxT min_samples_split, - IdxT max_leaves, - const Dataset dataset, - const Quantiles quantiles, - const NodeWorkItem* work_items, - IdxT colStart, - const IdxT* column_samples, - int* done_count, - int* mutex, - volatile Split* splits, - ObjectiveT objective, - IdxT treeid, - const WorkloadInfo* workload_info, - uint64_t seed) +template +DI BinCountT bin_count(BinT const& bin) { - // dynamic shared memory - extern __shared__ char smem[]; - - // Read workload info for this block - WorkloadInfo workload_info_cta = workload_info[blockIdx.x]; - IdxT nid = workload_info_cta.nodeid; - IdxT large_nid = workload_info_cta.large_nodeid; - const auto work_item = work_items[nid]; - auto range_start = work_item.instances.begin; - auto range_len = work_item.instances.count; - - IdxT offset_blockid = workload_info_cta.offset_blockid; - IdxT num_blocks = workload_info_cta.num_blocks; + return bin.Count(); +} - // obtaining the feature to test split on - IdxT colIndex = colStart + blockIdx.y; - IdxT col = column_samples[nid * dataset.n_sampled_cols + colIndex]; +template +static __global__ void computeSplitHistogramKernel(BinT* histograms, + IdxT max_n_bins, + const Dataset dataset, + const Quantiles quantiles, + const NodeWorkItem* work_items, + IdxT colStart, + const IdxT* column_samples, + const WorkloadInfo* workload_info) +{ + extern __shared__ char smem[]; - // getting the n_bins for that feature - int n_bins = quantiles.n_bins_array[col]; + auto workload_info_cta = workload_info[blockIdx.x]; + IdxT nid = workload_info_cta.nodeid; + const auto work_item = work_items[nid]; + auto range_start = work_item.instances.begin; + auto range_len = work_item.instances.count; + IdxT offset_blockid = workload_info_cta.offset_blockid; + IdxT num_blocks = workload_info_cta.num_blocks; + + IdxT col; + if (dataset.n_sampled_cols == dataset.N) { + col = colStart + blockIdx.y; + } else { + IdxT colIndex = colStart + blockIdx.y; + col = column_samples[nid * dataset.n_sampled_cols + colIndex]; + } - auto end = range_start + range_len; - auto shared_histogram_len = n_bins * objective.NumClasses(); + int n_bins = quantiles.n_bins_array[col]; + auto shared_histogram_len = n_bins * dataset.num_outputs; auto* shared_histogram = alignPointer(smem); auto* shared_quantiles = alignPointer(shared_histogram + shared_histogram_len); - auto* shared_done = alignPointer(shared_quantiles + n_bins); IdxT stride = blockDim.x * num_blocks; IdxT tid = threadIdx.x + offset_blockid * blockDim.x; + auto histograms_offset = ((nid * gridDim.y) + blockIdx.y) * max_n_bins * dataset.num_outputs; - // populating shared memory with initial values - for (IdxT i = threadIdx.x; i < shared_histogram_len; i += blockDim.x) + for (IdxT i = threadIdx.x; i < shared_histogram_len; i += blockDim.x) { shared_histogram[i] = BinT(); - for (IdxT b = threadIdx.x; b < n_bins; b += blockDim.x) + } + for (IdxT b = threadIdx.x; b < n_bins; b += blockDim.x) { shared_quantiles[b] = quantiles.quantiles_array[max_n_bins * col + b]; - - // synchronizing above changes across block + } __syncthreads(); - // compute pdf shared histogram for all bins for all classes in shared mem - - // Must be 64 bit - can easily grow larger than a 32 bit int std::size_t col_offset = std::size_t(col) * dataset.M; - for (auto i = range_start + tid; i < end; i += stride) { - // each thread works over a data point and strides to the next + for (auto i = range_start + tid; i < range_start + range_len; i += stride) { auto row = dataset.row_ids[i]; auto data = dataset.data[row + col_offset]; auto label = dataset.labels[row]; - - // `start` is lowest index such that data <= shared_quantiles[start] IdxT start = lower_bound(shared_quantiles, n_bins, data); - // ++shared_histogram[start] BinT::IncrementHistogram(shared_histogram, n_bins, start, label); } - - // synchronizing above changes across block __syncthreads(); - if (num_blocks > 1) { - // update the corresponding global location - auto histograms_offset = - ((large_nid * gridDim.y) + blockIdx.y) * max_n_bins * objective.NumClasses(); - for (IdxT i = threadIdx.x; i < shared_histogram_len; i += blockDim.x) { - BinT::AtomicAdd(histograms + histograms_offset + i, shared_histogram[i]); - } - __threadfence(); // for commit guarantee - __syncthreads(); + for (IdxT i = threadIdx.x; i < shared_histogram_len; i += blockDim.x) { + BinT::AtomicAdd(histograms + histograms_offset + i, shared_histogram[i]); + } +} - // last threadblock will go ahead and compute the best split - bool last = MLCommon::signalDone( - done_count + nid * gridDim.y + blockIdx.y, num_blocks, offset_blockid == 0, shared_done); - // if not the last threadblock, exit - if (!last) return; +template +static __global__ void evaluateSplitKernel(BinT* histograms, + IdxT max_n_bins, + const Dataset dataset, + const Quantiles quantiles, + IdxT colStart, + const IdxT* column_samples, + int* mutex, + volatile Split* splits, + ObjectiveT objective) +{ + extern __shared__ char smem[]; - // store the complete global histogram in shared memory of last block - for (IdxT i = threadIdx.x; i < shared_histogram_len; i += blockDim.x) - shared_histogram[i] = histograms[histograms_offset + i]; + IdxT nid = blockIdx.x; + IdxT col; + if (dataset.n_sampled_cols == dataset.N) { + col = colStart + blockIdx.y; + } else { + IdxT colIndex = colStart + blockIdx.y; + col = column_samples[nid * dataset.n_sampled_cols + colIndex]; + } - __syncthreads(); + int n_bins = quantiles.n_bins_array[col]; + auto shared_histogram_len = n_bins * objective.NumClasses(); + auto* shared_histogram = alignPointer(smem); + auto* shared_quantiles = alignPointer(shared_histogram + shared_histogram_len); + auto histograms_offset = ((nid * gridDim.y) + blockIdx.y) * max_n_bins * objective.NumClasses(); + + for (IdxT i = threadIdx.x; i < shared_histogram_len; i += blockDim.x) { + shared_histogram[i] = histograms[histograms_offset + i]; + } + for (IdxT b = threadIdx.x; b < n_bins; b += blockDim.x) { + shared_quantiles[b] = quantiles.quantiles_array[max_n_bins * col + b]; } + __syncthreads(); - // PDF to CDF inplace in `shared_histogram` + typename ObjectiveT::CountT split_len = 0; for (IdxT c = 0; c < objective.NumClasses(); ++c) { - // left to right scan operation for scanning - // "lesser-than-or-equal" counts - BinT total_sum = pdf_to_cdf(shared_histogram + n_bins * c, n_bins); - // now, `shared_histogram[n_bins * c + i]` will have count of datapoints of class `c` - // that are less than or equal to `shared_quantiles[i]`. + auto total_sum = pdf_to_cdf(shared_histogram + n_bins * c, n_bins); + split_len += bin_count(total_sum); } - __syncthreads(); - // calculate the best candidate bins (one for each thread in the block) in current feature and - // corresponding information gain for splitting - Split sp = - objective.Gain(shared_histogram, shared_quantiles, col, range_len, n_bins); - + Split sp = + objective.Gain(shared_histogram, shared_quantiles, static_cast(col), split_len, n_bins); __syncthreads(); - - // calculate best bins among candidate bins per feature using warp reduce - // then atomically update across features to get best split per node - // (in split[nid]) sp.evalBestSplit(smem, splits + nid, mutex + nid); } +template +void launchComputeSplitHistogramKernel(BinT* histograms, + IdxT max_n_bins, + const Dataset& dataset, + const Quantiles& 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) +{ + computeSplitHistogramKernel + <<>>(histograms, + max_n_bins, + dataset, + quantiles, + work_items, + colStart, + column_samples, + workload_info); +} + template -void launchComputeSplitKernel(BinT* histograms, - IdxT max_n_bins, - IdxT min_samples_split, - IdxT max_leaves, - const Dataset& dataset, - const Quantiles& quantiles, - const NodeWorkItem* work_items, - IdxT colStart, - const IdxT* column_samples, - int* done_count, - int* mutex, - volatile Split* splits, - ObjectiveT& objective, - IdxT treeid, - const WorkloadInfo* 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& dataset, + const Quantiles& quantiles, + IdxT colStart, + const IdxT* column_samples, + int* mutex, + volatile Split* splits, + ObjectiveT& objective, + dim3 grid, + size_t smem_size, + cudaStream_t builder_stream) { - computeSplitKernel + evaluateSplitKernel <<>>(histograms, max_n_bins, - min_samples_split, - max_leaves, dataset, quantiles, - work_items, colStart, column_samples, - done_count, mutex, splits, - objective, - treeid, - workload_info, - seed); + objective); } template void launchNodeSplitKernel<_DataT, _LabelT, _IdxT, TPB_DEFAULT>( - const _IdxT min_samples_leaf, 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 void launchLeafKernel<_DatasetT, _NodeT, _ObjectiveT, _DataT>( +template void launchLeafHistogramKernel<_DatasetT, _NodeT, _ObjectiveT>( _ObjectiveT objective, _DatasetT& dataset, const _NodeT* tree, const InstanceRange* instance_ranges, - _DataT* leaves, + typename _ObjectiveT::BinT* leaf_histograms, int batch_size, size_t smem_size, cudaStream_t builder_stream); -template void launchComputeSplitKernel<_DataT, _LabelT, _IdxT, TPB_DEFAULT, _ObjectiveT, _BinT>( +template void launchFinalizeLeafKernel<_NodeT, _ObjectiveT, _DataT>( + _ObjectiveT objective, + const _NodeT* tree, + const typename _ObjectiveT::BinT* leaf_histograms, + _DataT* leaves, + int batch_size, + int num_outputs, + cudaStream_t builder_stream); + +template void launchComputeSplitHistogramKernel<_DataT, _LabelT, _IdxT, TPB_DEFAULT, _BinT>( _BinT* histograms, - _IdxT n_bins, - _IdxT min_samples_split, - _IdxT max_leaves, + _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, - int* done_count, + const WorkloadInfo* workload_info, + dim3 grid, + size_t smem_size, + cudaStream_t builder_stream); + +template void launchEvaluateSplitKernel<_DataT, _LabelT, _IdxT, TPB_DEFAULT, _ObjectiveT, _BinT>( + _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, _IdxT>* splits, + volatile Split<_DataT>* splits, _ObjectiveT& objective, - _IdxT treeid, - const WorkloadInfo<_IdxT>* workload_info, - uint64_t seed, dim3 grid, size_t smem_size, cudaStream_t builder_stream); diff --git a/cpp/src/decisiontree/batched-levelalgo/objectives.cuh b/cpp/src/decisiontree/batched-levelalgo/objectives.cuh index 5ec6d42cc7..bc54daede2 100644 --- a/cpp/src/decisiontree/batched-levelalgo/objectives.cuh +++ b/cpp/src/decisiontree/batched-levelalgo/objectives.cuh @@ -24,6 +24,7 @@ class ClassificationObjectiveFunction { using LabelT = LabelT_; using IdxT = IdxT_; using BinT = std::conditional_t; + using CountT = BinCountT; static constexpr bool weighted = weighted_; private: @@ -31,13 +32,13 @@ class ClassificationObjectiveFunction { IdxT min_samples_leaf; CRITERION criterion; - DI IdxT CountLeft(BinT const* hist, IdxT i, IdxT n_bins) const + DI CountT CountLeft(BinT const* hist, IdxT i, IdxT n_bins) const { - BinCountT nLeft = 0; + CountT nLeft = 0; for (IdxT j = 0; j < nclasses; ++j) { nLeft += hist[n_bins * j + i].Count(); } - return static_cast(nLeft); + return nLeft; } HDI double WeightAt(BinT const* hist, IdxT i, IdxT n_bins) const @@ -49,7 +50,7 @@ class ClassificationObjectiveFunction { return weight; } - HDI DataT GiniGain(BinT const* hist, IdxT i, IdxT n_bins, IdxT, IdxT, IdxT) const + HDI DataT GiniGain(BinT const* hist, IdxT i, IdxT n_bins, CountT, CountT, CountT) const { constexpr DataT One = DataT(1.0); auto total_weight = WeightAt(hist, n_bins - 1, n_bins); @@ -84,7 +85,7 @@ class ClassificationObjectiveFunction { return gain; } - HDI DataT EntropyGain(BinT const* hist, IdxT i, IdxT n_bins, IdxT, IdxT, IdxT) const + HDI DataT EntropyGain(BinT const* hist, IdxT i, IdxT n_bins, CountT, CountT, CountT) const { auto total_weight = WeightAt(hist, n_bins - 1, n_bins); auto left_weight = WeightAt(hist, i, n_bins); @@ -125,9 +126,9 @@ class ClassificationObjectiveFunction { public: HDI DataT - GainPerSplit(BinT const* hist, IdxT i, IdxT n_bins, IdxT len, IdxT nLeft, IdxT nRight) const + GainPerSplit(BinT const* hist, IdxT i, IdxT n_bins, CountT len, CountT nLeft, CountT nRight) const { - if (nLeft < min_samples_leaf || nRight < min_samples_leaf) + if (nLeft < CountT(min_samples_leaf) || nRight < CountT(min_samples_leaf)) return -std::numeric_limits::max(); switch (criterion) { @@ -144,18 +145,17 @@ class ClassificationObjectiveFunction { DI IdxT NumClasses() const { return nclasses; } - DI Split Gain( - BinT const* shist, DataT const* squantiles, IdxT col, IdxT len, IdxT n_bins) const + DI Split Gain( + BinT const* shist, DataT const* squantiles, int col, CountT len, IdxT n_bins) const { - Split sp; + Split sp; for (IdxT i = threadIdx.x; i < n_bins; i += blockDim.x) { auto nLeft = CountLeft(shist, i, n_bins); auto nRight = len - nLeft; - auto gain = -std::numeric_limits::max(); - if (nLeft >= min_samples_leaf && nRight >= min_samples_leaf) { - gain = GainPerSplit(shist, i, n_bins, len, nLeft, nRight); + if (nLeft >= CountT(min_samples_leaf) && nRight >= CountT(min_samples_leaf)) { + auto gain = GainPerSplit(shist, i, n_bins, len, nLeft, nRight); + sp.update({squantiles[i], col, gain, nLeft}); } - sp.update({squantiles[i], col, gain, nLeft}); } return sp; } @@ -186,6 +186,7 @@ class RegressionObjectiveFunction { using LabelT = LabelT_; using IdxT = IdxT_; using BinT = std::conditional_t; + using CountT = BinCountT; static constexpr bool weighted = weighted_; private: @@ -193,7 +194,7 @@ class RegressionObjectiveFunction { CRITERION criterion; static constexpr auto eps_ = 10 * std::numeric_limits::epsilon(); - HDI DataT MSEGain(BinT const* hist, IdxT i, IdxT n_bins, IdxT, IdxT, IdxT) const + HDI DataT MSEGain(BinT const* hist, IdxT i, IdxT n_bins, CountT, CountT, CountT) const { auto parent_weight = hist[n_bins - 1].Weight(); auto left_weight = hist[i].Weight(); @@ -215,7 +216,7 @@ class RegressionObjectiveFunction { return gain; } - HDI DataT PoissonGain(BinT const* hist, IdxT i, IdxT n_bins, IdxT, IdxT, IdxT) const + HDI DataT PoissonGain(BinT const* hist, IdxT i, IdxT n_bins, CountT, CountT, CountT) const { auto parent_weight = hist[n_bins - 1].Weight(); auto left_weight = hist[i].Weight(); @@ -242,7 +243,7 @@ class RegressionObjectiveFunction { return gain; } - HDI DataT GammaGain(BinT const* hist, IdxT i, IdxT n_bins, IdxT, IdxT, IdxT) const + HDI DataT GammaGain(BinT const* hist, IdxT i, IdxT n_bins, CountT, CountT, CountT) const { auto parent_weight = hist[n_bins - 1].Weight(); auto left_weight = hist[i].Weight(); @@ -269,7 +270,7 @@ class RegressionObjectiveFunction { return gain; } - HDI DataT InverseGaussianGain(BinT const* hist, IdxT i, IdxT n_bins, IdxT, IdxT, IdxT) const + HDI DataT InverseGaussianGain(BinT const* hist, IdxT i, IdxT n_bins, CountT, CountT, CountT) const { auto parent_weight = hist[n_bins - 1].Weight(); auto left_weight = hist[i].Weight(); @@ -297,9 +298,9 @@ class RegressionObjectiveFunction { public: HDI DataT - GainPerSplit(BinT const* hist, IdxT i, IdxT n_bins, IdxT len, IdxT nLeft, IdxT nRight) const + GainPerSplit(BinT const* hist, IdxT i, IdxT n_bins, CountT len, CountT nLeft, CountT nRight) const { - if (nLeft < min_samples_leaf || nRight < min_samples_leaf) + if (nLeft < CountT(min_samples_leaf) || nRight < CountT(min_samples_leaf)) return -std::numeric_limits::max(); switch (criterion) { @@ -319,18 +320,17 @@ class RegressionObjectiveFunction { DI IdxT NumClasses() const { return 1; } - DI Split Gain( - BinT const* shist, DataT const* squantiles, IdxT col, IdxT len, IdxT n_bins) const + DI Split Gain( + BinT const* shist, DataT const* squantiles, int col, CountT len, IdxT n_bins) const { - Split sp; + Split sp; for (IdxT i = threadIdx.x; i < n_bins; i += blockDim.x) { - auto nLeft = static_cast(shist[i].Count()); + auto nLeft = shist[i].Count(); auto nRight = len - nLeft; - auto gain = -std::numeric_limits::max(); - if (nLeft >= min_samples_leaf && nRight >= min_samples_leaf) { - gain = GainPerSplit(shist, i, n_bins, len, nLeft, nRight); + if (nLeft >= CountT(min_samples_leaf) && nRight >= CountT(min_samples_leaf)) { + auto gain = GainPerSplit(shist, i, n_bins, len, nLeft, nRight); + sp.update({squantiles[i], col, gain, nLeft}); } - sp.update({squantiles[i], col, gain, nLeft}); } return sp; } diff --git a/cpp/src/decisiontree/batched-levelalgo/quantiles.cuh b/cpp/src/decisiontree/batched-levelalgo/quantiles.cuh index f0d3fd000f..bfa1216922 100644 --- a/cpp/src/decisiontree/batched-levelalgo/quantiles.cuh +++ b/cpp/src/decisiontree/batched-levelalgo/quantiles.cuh @@ -6,6 +6,7 @@ #pragma once #include "quantiles.h" +#include "random_utils.cuh" #include #include @@ -127,7 +128,7 @@ struct QuantileResult { * @param handle RAFT handle used for stream and resource access. * @param data Column-major input matrix with shape `[n_cols, n_rows]`. * @param max_n_bins Maximum number of quantile candidates to retain per feature. - * @param n_rows Number of local rows in `data` for this rank. + * @param n_rows Number of local rows in `data` for this rank; may be zero in distributed mode. * @param n_cols Number of columns in `data`. * @param oversampling_factor Multiplier applied to `max_n_bins` to choose the * sampled row budget per feature before sorting and quantile extraction. The @@ -147,9 +148,9 @@ CUML_EXPORT QuantileResult computeQuantiles(const raft::handle_t& handle, uint64_t seed = uint64_t{0}) { raft::common::nvtx::push_range("computeQuantiles"); - RAFT_EXPECTS(data != nullptr, "data pointer must not be null"); + RAFT_EXPECTS(data != nullptr || n_rows == 0, "data pointer must not be null"); RAFT_EXPECTS(max_n_bins > 0, "max_n_bins must be positive"); - RAFT_EXPECTS(n_rows > 0, "n_rows must be positive"); + RAFT_EXPECTS(n_rows >= 0, "n_rows must be non-negative"); RAFT_EXPECTS(n_cols > 0, "n_cols must be positive"); RAFT_EXPECTS(oversampling_factor > 0, "oversampling_factor must be positive"); diff --git a/cpp/src/decisiontree/batched-levelalgo/split.cuh b/cpp/src/decisiontree/batched-levelalgo/split.cuh index b0790945f1..0091998c1f 100644 --- a/cpp/src/decisiontree/batched-levelalgo/split.cuh +++ b/cpp/src/decisiontree/batched-levelalgo/split.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2022, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -16,9 +16,11 @@ namespace DT { * * @tparam DataT input data type */ -template +template struct Split { - typedef Split SplitT; + typedef Split SplitT; + using CountT = unsigned long long int; + static_assert(sizeof(CountT) == 8, "RF split counts must be 64-bit."); /** start with this as the initial gain */ static constexpr DataT Min = -std::numeric_limits::max(); @@ -26,14 +28,24 @@ struct Split { /** threshold to compare in this node */ DataT quesval; /** feature index */ - IdxT colid; + int colid; /** best info gain on this node */ DataT best_metric_val; - /** number of samples in the left child */ - int nLeft; + /** global number of samples in the left child */ + CountT global_nLeft; + /** rank-local number of samples in the left child */ + CountT local_nLeft; - DI Split(DataT quesval, IdxT colid, DataT best_metric_val, IdxT nLeft) - : quesval(quesval), colid(colid), best_metric_val(best_metric_val), nLeft(nLeft) + DI Split(DataT quesval, + int colid, + DataT best_metric_val, + CountT global_nLeft, + CountT local_nLeft = CountT{0}) + : quesval(quesval), + colid(colid), + best_metric_val(best_metric_val), + global_nLeft(global_nLeft), + local_nLeft(local_nLeft) { } @@ -41,7 +53,8 @@ struct Split { { quesval = best_metric_val = Min; colid = -1; - nLeft = 0; + global_nLeft = 0; + local_nLeft = 0; } /** @@ -56,7 +69,8 @@ struct Split { quesval = other.quesval; colid = other.colid; best_metric_val = other.best_metric_val; - nLeft = other.nLeft; + global_nLeft = other.global_nLeft; + local_nLeft = other.local_nLeft; return *this; } @@ -87,12 +101,13 @@ struct Split { auto lane = raft::laneId(); #pragma unroll for (int i = raft::WarpSize / 2; i >= 1; i /= 2) { - auto id = lane + i; - auto qu = raft::shfl(quesval, id); - auto co = raft::shfl(colid, id); - auto be = raft::shfl(best_metric_val, id); - auto nl = raft::shfl(nLeft, id); - update({qu, co, be, nl}); + auto id = lane + i; + auto qu = raft::shfl(quesval, id); + auto co = raft::shfl(colid, id); + auto be = raft::shfl(best_metric_val, id); + auto gnl = raft::shfl(global_nLeft, id); + auto lnl = raft::shfl(local_nLeft, id); + update({qu, co, be, gnl, lnl}); } } @@ -130,14 +145,19 @@ struct Split { split_reg.quesval = split->quesval; split_reg.colid = split->colid; split_reg.best_metric_val = split->best_metric_val; - split_reg.nLeft = split->nLeft; - bool update_result = - split_reg.update({this->quesval, this->colid, this->best_metric_val, this->nLeft}); + split_reg.global_nLeft = split->global_nLeft; + split_reg.local_nLeft = split->local_nLeft; + bool update_result = split_reg.update({this->quesval, + this->colid, + this->best_metric_val, + this->global_nLeft, + this->local_nLeft}); if (update_result) { split->quesval = split_reg.quesval; split->colid = split_reg.colid; split->best_metric_val = split_reg.best_metric_val; - split->nLeft = split_reg.nLeft; + split->global_nLeft = split_reg.global_nLeft; + split->local_nLeft = split_reg.local_nLeft; } __threadfence(); atomicExch(mutex, 0); @@ -154,23 +174,25 @@ struct Split { * @param[in] s cuda stream where to schedule work */ template -void initSplit(Split* splits, IdxT len, cudaStream_t s) +void initSplit(Split* splits, IdxT len, cudaStream_t s) { - auto op = [] __device__(Split * ptr, IdxT idx) { *ptr = Split(); }; - raft::linalg::writeOnlyUnaryOp, decltype(op), IdxT, TPB>(splits, len, op, s); + auto op = [] __device__(Split * ptr, IdxT idx) { *ptr = Split(); }; + raft::linalg::writeOnlyUnaryOp, decltype(op), IdxT, TPB>(splits, len, op, s); } template -void printSplits(Split* splits, IdxT len, cudaStream_t s) +void printSplits(Split* splits, IdxT len, cudaStream_t s) { - auto op = [] __device__(Split * ptr, IdxT idx) { - printf("quesval = %e, colid = %d, best_metric_val = %e, nLeft = %d\n", - ptr->quesval, - ptr->colid, - ptr->best_metric_val, - ptr->nLeft); + auto op = [] __device__(Split * ptr, IdxT idx) { + printf( + "quesval = %e, colid = %d, best_metric_val = %e, global_nLeft = %llu, local_nLeft = %llu\n", + ptr->quesval, + ptr->colid, + ptr->best_metric_val, + ptr->global_nLeft, + ptr->local_nLeft); }; - raft::linalg::writeOnlyUnaryOp, decltype(op), IdxT, TPB>(splits, len, op, s); + raft::linalg::writeOnlyUnaryOp, decltype(op), IdxT, TPB>(splits, len, op, s); RAFT_CUDA_TRY(cudaDeviceSynchronize()); } diff --git a/cpp/src/randomforest/randomforest.cu b/cpp/src/randomforest/randomforest.cu index f77c4a3c7c..9baf2e5f1a 100644 --- a/cpp/src/randomforest/randomforest.cu +++ b/cpp/src/randomforest/randomforest.cu @@ -18,6 +18,7 @@ #include #include +#include #include #include #include @@ -738,10 +739,12 @@ void compute_feature_importances(const RandomForestMetaData* forest, T* im if (forest->n_features == 0) { return; } int n_cols = forest->n_features; - std::vector accumulated_importances(n_cols, T(0)); + std::vector accumulated_importances(n_cols, 0.0); for (const auto& tree : forest->trees) { - std::vector tree_importances(n_cols, T(0)); + std::vector finite_importances(n_cols, 0.0); + std::vector infinite_importances(n_cols, 0.0); + bool has_infinite_importance = false; if (tree->sparsetree.empty()) continue; int root_sample_count = tree->sparsetree[0].InstanceCount(); @@ -751,29 +754,36 @@ void compute_feature_importances(const RandomForestMetaData* forest, T* im for (const auto& node : tree->sparsetree) { if (!node.IsLeaf()) { int feature_id = node.ColumnId(); - tree_importances[feature_id] += node.BestMetric() * node.InstanceCount(); + double contribution = + static_cast(node.BestMetric()) * static_cast(node.InstanceCount()); + if (std::isfinite(contribution)) { + if (contribution > 0.0) { finite_importances[feature_id] += contribution; } + } else if (std::isinf(contribution) && contribution > 0.0) { + infinite_importances[feature_id] += 1.0; + has_infinite_importance = true; + } } } - T sum = T(0); + auto& tree_importances = has_infinite_importance ? infinite_importances : finite_importances; + double sum = 0.0; for (int i = 0; i < n_cols; i++) { sum += tree_importances[i]; } if (sum > 0) { for (int i = 0; i < n_cols; i++) { - tree_importances[i] /= sum; - accumulated_importances[i] += tree_importances[i]; + accumulated_importances[i] += tree_importances[i] / sum; } } } - T sum = T(0); + double sum = 0.0; for (auto i = 0; i < n_cols; i++) { sum += accumulated_importances[i]; } if (sum > 0) { for (auto i = 0; i < n_cols; i++) { - importances[i] = accumulated_importances[i] / sum; + importances[i] = T(accumulated_importances[i] / sum); } } else { for (auto i = 0; i < n_cols; i++) { diff --git a/cpp/src/randomforest/randomforest.cuh b/cpp/src/randomforest/randomforest.cuh index 7e0ebbb229..b1ec7fac54 100644 --- a/cpp/src/randomforest/randomforest.cuh +++ b/cpp/src/randomforest/randomforest.cuh @@ -9,6 +9,7 @@ #include #include +#include #include #include #include @@ -49,6 +50,7 @@ class RandomForest { const cudaStream_t stream) { raft::common::nvtx::range fun_scope("bootstrapping row IDs @randomforest.cuh"); + if (selected_rows->size() == 0) { return; } // Hash these together so they are uncorrelated auto rs = DT::fnv1a32_basis; @@ -65,14 +67,21 @@ class RandomForest { } } - void error_checking(const T* input, L* predictions, int n_rows, int n_cols, bool predict) const + void error_checking(const T* input, + L* predictions, + int n_rows, + int n_cols, + bool predict, + bool allow_empty_local_rows = false) const { if (predict) { ASSERT(predictions != nullptr, "Error! User has not allocated memory for predictions."); } - ASSERT((n_rows > 0), "Invalid n_rows %d", n_rows); + ASSERT(allow_empty_local_rows ? (n_rows >= 0) : (n_rows > 0), "Invalid n_rows %d", n_rows); ASSERT((n_cols > 0), "Invalid n_cols %d", n_cols); + if (n_rows == 0) { return; } + bool input_is_dev_ptr = DT::is_dev_ptr(input); bool preds_is_dev_ptr = DT::is_dev_ptr(predictions); @@ -120,9 +129,13 @@ class RandomForest { bool* bootstrap_masks = nullptr) { raft::common::nvtx::range fun_scope("RandomForest::fit @randomforest.cuh"); - this->error_checking(input, labels, n_rows, n_cols, false); const raft::handle_t& handle = user_handle; - int n_sampled_rows = 0; + bool distributed = + raft::resource::comms_initialized(handle) && handle.get_comms().get_size() > 1; + this->error_checking(input, labels, n_rows, n_cols, false, distributed); + int device = 0; + RAFT_CUDA_TRY(cudaGetDevice(&device)); + int n_sampled_rows = 0; if (this->rf_params.bootstrap) { n_sampled_rows = std::round(this->rf_params.max_samples * n_rows); } else { @@ -135,8 +148,11 @@ class RandomForest { n_sampled_rows = n_rows; } int n_streams = this->rf_params.n_streams; + // Distributed tree builders issue collectives independently, so train them serially until + // the forest-level scheduler can impose a global collective order across concurrent trees. + if (distributed) { n_streams = 1; } ASSERT(static_cast(n_streams) <= handle.get_stream_pool_size(), - "rf_params.n_streams (=%d) should be <= raft::handle_t.n_streams (=%lu)", + "effective RF n_streams (=%d) should be <= raft::handle_t.n_streams (=%lu)", n_streams, handle.get_stream_pool_size()); @@ -161,6 +177,7 @@ class RandomForest { #pragma omp parallel for num_threads(n_streams) for (int i = 0; i < this->rf_params.n_trees; i++) { + RAFT_CUDA_TRY(cudaSetDevice(device)); int stream_id = omp_get_thread_num(); auto s = handle.get_stream_from_stream_pool(stream_id); diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index e257adab18..2f8f4d73b4 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -210,6 +210,7 @@ if(BUILD_CUML_MG_TESTS) ConfigureTest( PREFIX MG NAME RF_QUANTILE_TEST mg/rf_quantile_test.cu MPI RAFT_DISTRIBUTED ML_INCLUDE ) + ConfigureTest(PREFIX MG NAME RF_TEST mg/rf_test.cu MPI RAFT_DISTRIBUTED ML_INCLUDE) else(MPI_CXX_FOUND) message("OpenMPI not found. Skipping MultiGPU tests '${CUML_MG_TEST_TARGET}'") endif() diff --git a/cpp/tests/mg/rf_quantile_test.cu b/cpp/tests/mg/rf_quantile_test.cu index 3a39651a24..f21096dbd9 100644 --- a/cpp/tests/mg/rf_quantile_test.cu +++ b/cpp/tests/mg/rf_quantile_test.cu @@ -206,3 +206,11 @@ TEST_F(RfMgQuantileTestD, SharedAcrossRanks) {} } // namespace opg } // namespace Test } // namespace ML + +int main(int argc, char** argv) +{ + ::testing::InitGoogleTest(&argc, argv); + ::testing::AddGlobalTestEnvironment(new MLCommon::Test::opg::MPIEnvironment()); + + return RUN_ALL_TESTS(); +} diff --git a/cpp/tests/mg/rf_test.cu b/cpp/tests/mg/rf_test.cu new file mode 100644 index 0000000000..c9f909db23 --- /dev/null +++ b/cpp/tests/mg/rf_test.cu @@ -0,0 +1,431 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "../prims/test_utils.h" +#include "test_opg_utils.h" + +#include + +#include +#include +#include + +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace ML { +namespace Test { +namespace opg { + +enum class PartitionKind { Contiguous, Strided, Imbalanced, EmptyNonRootRanks }; + +struct RfMgTestParams { + int n_rows; + int n_cols; + int n_trees; + float max_features; + int max_depth; + int max_leaves; + int max_n_bins; + int min_samples_leaf; + int min_samples_split; + float min_impurity_decrease; + int n_streams; + int handle_n_streams; + CRITERION split_criterion; + int seed; + int n_labels; + bool double_precision; + PartitionKind partition_kind; +}; + +template +void hash_value(uint64_t& hash, const T& value) +{ + const auto* bytes = reinterpret_cast(&value); + for (size_t i = 0; i < sizeof(T); ++i) { + hash ^= bytes[i]; + hash *= 1099511628211ULL; + } +} + +template +uint64_t hash_forest_structure(const RandomForestMetaData& forest) +{ + uint64_t hash = 1469598103934665603ULL; + hash_value(hash, forest.n_features); + hash_value(hash, forest.rf_params.n_trees); + hash_value(hash, forest.rf_params.bootstrap); + hash_value(hash, forest.rf_params.max_samples); + hash_value(hash, forest.rf_params.seed); + hash_value(hash, forest.rf_params.tree_params.max_depth); + hash_value(hash, forest.rf_params.tree_params.max_leaves); + hash_value(hash, forest.rf_params.tree_params.max_n_bins); + hash_value(hash, forest.rf_params.tree_params.min_samples_leaf); + hash_value(hash, forest.rf_params.tree_params.min_samples_split); + hash_value(hash, forest.rf_params.tree_params.min_impurity_decrease); + hash_value(hash, forest.rf_params.tree_params.split_criterion); + for (auto const& tree : forest.trees) { + hash_value(hash, tree->treeid); + hash_value(hash, tree->depth_counter); + hash_value(hash, tree->leaf_counter); + hash_value(hash, tree->num_outputs); + hash_value(hash, tree->sparsetree.size()); + for (auto const& node : tree->sparsetree) { + hash_value(hash, node.ColumnId()); + hash_value(hash, node.QueryValue()); + hash_value(hash, node.BestMetric()); + hash_value(hash, node.LeftChildId()); + hash_value(hash, node.InstanceCount()); + hash_value(hash, node.IsLeaf()); + } + } + return hash; +} + +template +uint64_t hash_forest_leaf_values(const RandomForestMetaData& forest) +{ + uint64_t hash = 1469598103934665603ULL; + for (auto const& tree : forest.trees) { + hash_value(hash, tree->vector_leaf.size()); + for (auto const& leaf : tree->vector_leaf) { + hash_value(hash, leaf); + } + } + return hash; +} + +template +uint64_t hash_host_vector(std::vector const& values) +{ + uint64_t hash = 1469598103934665603ULL; + hash_value(hash, values.size()); + for (auto const& value : values) { + hash_value(hash, value); + } + return hash; +} + +std::vector local_rows_for_rank(int n_rows, int rank, int size, PartitionKind kind) +{ + std::vector rows; + if (kind == PartitionKind::Strided) { + for (int row = rank; row < n_rows; row += size) { + rows.push_back(row); + } + return rows; + } + + std::vector counts(size, n_rows / size); + for (int i = 0; i < n_rows % size; ++i) { + counts[i]++; + } + if (kind == PartitionKind::Imbalanced && size > 1) { + counts.assign(size, 0); + counts[0] = std::max(1, (n_rows * 3) / 4); + int remaining = n_rows - counts[0]; + for (int i = 1; i < size; ++i) { + counts[i] = remaining / (size - 1); + } + for (int i = 1; i <= remaining % (size - 1); ++i) { + counts[i]++; + } + } else if (kind == PartitionKind::EmptyNonRootRanks && size > 1) { + counts.assign(size, 0); + counts[0] = n_rows; + } + + int begin = std::accumulate(counts.begin(), counts.begin() + rank, 0); + rows.resize(counts[rank]); + std::iota(rows.begin(), rows.end(), begin); + rows.erase(std::remove_if(rows.begin(), rows.end(), [=](int row) { return row >= n_rows; }), + rows.end()); + return rows; +} + +template +void make_local_dataset(RfMgTestParams const& params, + std::vector const& rows, + std::vector& X, + std::vector& y) +{ + X.resize(rows.size() * params.n_cols); + y.resize(rows.size()); + for (size_t i = 0; i < rows.size(); ++i) { + int global_row = rows[i]; + DataT signal = static_cast((global_row % 97) - 48); + for (int col = 0; col < params.n_cols; ++col) { + DataT feature = signal * static_cast(col + 1); + feature += static_cast(((global_row + 13 * col + params.seed) % 11) - 5) / + static_cast(10); + X[static_cast(col) * rows.size() + i] = feature; + } + if constexpr (std::is_integral_v) { + y[i] = (signal >= DataT(0)) ? 1 : 0; + if (params.n_labels > 2 && global_row % 17 == 0) { y[i] = 2; } + } else { + y[i] = signal * DataT(0.5) + static_cast((global_row % 7) - 3); + } + } +} + +template +void make_prediction_dataset(RfMgTestParams const& params, std::vector& X) +{ + X.resize(params.n_rows * params.n_cols); + for (int row = 0; row < params.n_rows; ++row) { + for (int col = 0; col < params.n_cols; ++col) { + X[row * params.n_cols + col] = + static_cast(((row * 5 + col * 11 + params.seed) % 101) - 50); + } + } +} + +template +void expect_global_tree_counts(RandomForestMetaData const& forest, int n_rows) +{ + for (auto const& tree : forest.trees) { + ASSERT_FALSE(tree->sparsetree.empty()); + EXPECT_EQ(tree->sparsetree.front().InstanceCount(), n_rows); + for (auto const& node : tree->sparsetree) { + if (!node.IsLeaf()) { + auto left_count = tree->sparsetree[node.LeftChildId()].InstanceCount(); + auto right_count = tree->sparsetree[node.RightChildId()].InstanceCount(); + EXPECT_EQ(left_count + right_count, node.InstanceCount()); + } + } + } +} + +template +void expect_tree_limits(RandomForestMetaData const& forest, RfMgTestParams const& params) +{ + for (auto const& tree : forest.trees) { + EXPECT_LE(tree->depth_counter, params.max_depth); + if (params.max_leaves > 0) { EXPECT_LE(tree->leaf_counter, params.max_leaves); } + for (auto const& node : tree->sparsetree) { + if (!node.IsLeaf()) { EXPECT_GT(node.BestMetric(), params.min_impurity_decrease); } + } + } +} + +void initialize_mpi_once() +{ + int mpi_initialized = 0; + MPI_Initialized(&mpi_initialized); + if (!mpi_initialized) { MPI_Init(nullptr, nullptr); } +} + +void get_mpi_local_rank_size(int& local_rank, int& local_size) +{ + MPI_Comm local_comm{}; + MPI_Comm_split_type(MPI_COMM_WORLD, MPI_COMM_TYPE_SHARED, 0, MPI_INFO_NULL, &local_comm); + MPI_Comm_rank(local_comm, &local_rank); + MPI_Comm_size(local_comm, &local_size); + MPI_Comm_free(&local_comm); +} + +template +class RfMgPropertyTestImpl { + public: + explicit RfMgPropertyTestImpl(RfMgTestParams const& params) : params(params) + { + initialize_mpi_once(); + int rank = 0; + int size = 1; + MPI_Comm_rank(MPI_COMM_WORLD, &rank); + MPI_Comm_size(MPI_COMM_WORLD, &size); + + int local_rank = 0; + int local_size = 1; + get_mpi_local_rank_size(local_rank, local_size); + + int n_gpus = 0; + RAFT_CUDA_TRY(cudaGetDeviceCount(&n_gpus)); + if (n_gpus < local_size) { + ADD_FAILURE() << "Number of GPUs is smaller than local MPI ranks: ngpus=" << n_gpus + << ", local_ranks=" << local_size; + return; + } + RAFT_CUDA_TRY(cudaSetDevice(local_rank)); + + auto stream_pool = std::make_shared(params.handle_n_streams); + raft::handle_t handle(rmm::cuda_stream_per_thread, stream_pool); + raft::comms::initialize_mpi_comms(&handle, MPI_COMM_WORLD); + + auto local_rows = local_rows_for_rank(params.n_rows, rank, size, params.partition_kind); + std::vector h_X; + std::vector h_y; + make_local_dataset(params, local_rows, h_X, h_y); + + rmm::device_uvector X(h_X.size(), handle.get_stream()); + rmm::device_uvector y(h_y.size(), handle.get_stream()); + raft::update_device(X.data(), h_X.data(), h_X.size(), handle.get_stream()); + raft::update_device(y.data(), h_y.data(), h_y.size(), handle.get_stream()); + + auto rf_params = set_rf_params(params.max_depth, + params.max_leaves, + params.max_features, + params.max_n_bins, + params.min_samples_leaf, + params.min_samples_split, + params.min_impurity_decrease, + false, + params.n_trees, + 1.0f, + params.seed, + params.split_criterion, + params.n_streams, + 128); + + RandomForestMetaData forest; + if constexpr (std::is_integral_v) { + fit(handle, + &forest, + X.data(), + static_cast(local_rows.size()), + params.n_cols, + y.data(), + params.n_labels, + rf_params); + } else { + fit(handle, + &forest, + X.data(), + static_cast(local_rows.size()), + params.n_cols, + y.data(), + rf_params); + } + + expect_global_tree_counts(forest, params.n_rows); + expect_tree_limits(forest, params); + expect_identical_across_ranks(handle, hash_forest_structure(forest), "tree structure"); + expect_identical_across_ranks(handle, hash_forest_leaf_values(forest), "leaf values"); + expect_identical_predictions_across_ranks(handle, &forest); + } + + private: + void expect_identical_across_ranks(raft::handle_t const& handle, + uint64_t local_hash, + char const* label) + { + auto const& comm = handle.get_comms(); + std::vector hashes(comm.get_size()); + MPI_Allgather(&local_hash, 1, MPI_UINT64_T, hashes.data(), 1, MPI_UINT64_T, MPI_COMM_WORLD); + for (auto hash : hashes) { + EXPECT_EQ(hash, hashes.front()) << "Mismatched distributed RF " << label; + } + } + + void expect_identical_predictions_across_ranks(raft::handle_t const& handle, + RandomForestMetaData* forest) + { + std::vector h_X; + make_prediction_dataset(params, h_X); + rmm::device_uvector X(h_X.size(), handle.get_stream()); + rmm::device_uvector predictions(params.n_rows, handle.get_stream()); + raft::update_device(X.data(), h_X.data(), h_X.size(), handle.get_stream()); + predict(handle, forest, X.data(), params.n_rows, params.n_cols, predictions.data()); + std::vector h_predictions(params.n_rows); + raft::update_host( + h_predictions.data(), predictions.data(), h_predictions.size(), handle.get_stream()); + handle.sync_stream(); + expect_identical_across_ranks(handle, hash_host_vector(h_predictions), "predictions"); + } + + RfMgTestParams params; +}; + +class RfMgPropertyTest : public ::testing::TestWithParam { + public: + void SetUp() override + { + auto params = GetParam(); + bool is_regression = params.split_criterion != GINI && params.split_criterion != ENTROPY; + if (params.double_precision) { + if (is_regression) { + RfMgPropertyTestImpl test(params); + } else { + RfMgPropertyTestImpl test(params); + } + } else { + if (is_regression) { + RfMgPropertyTestImpl test(params); + } else { + RfMgPropertyTestImpl test(params); + } + } + } +}; + +TEST_P(RfMgPropertyTest, DistributedProperties) {} + +std::vector inputs = { + {128, 4, 1, 1.0f, 3, -1, 16, 1, 2, 0.0f, 1, 1, GINI, 7, 2, false, PartitionKind::Contiguous}, + {128, 4, 3, 0.5f, 4, 16, 32, 1, 2, 0.0f, 4, 4, ENTROPY, 11, 2, false, PartitionKind::Strided}, + {192, 6, 1, 1.0f, 5, -1, 32, 2, 4, 0.0f, 1, 1, MSE, 13, 2, false, PartitionKind::Imbalanced}, + {96, 3, 2, 1.0f, 4, 8, 8, 1, 2, 0.0f, 1, 1, GINI, 17, 2, true, PartitionKind::Imbalanced}, + {144, 5, 2, 0.8f, 4, -1, 16, 1, 2, 0.0f, 1, 1, GINI, 31, 3, false, PartitionKind::Strided}, + {160, 5, 1, 0.8f, 4, -1, 16, 1, 2, 0.0f, 1, 1, MSE, 19, 2, true, PartitionKind::Contiguous}, + {64, + 4, + 2, + 1.0f, + 4, + -1, + 16, + 1, + 2, + 0.0f, + 3, + 3, + GINI, + 23, + 2, + false, + PartitionKind::EmptyNonRootRanks}, + {80, + 5, + 2, + 0.8f, + 4, + -1, + 16, + 1, + 2, + 0.0f, + 3, + 3, + MSE, + 29, + 2, + false, + PartitionKind::EmptyNonRootRanks}}; + +INSTANTIATE_TEST_CASE_P(RfTests, RfMgPropertyTest, ::testing::ValuesIn(inputs)); + +} // namespace opg +} // namespace Test +} // namespace ML + +int main(int argc, char** argv) +{ + ::testing::InitGoogleTest(&argc, argv); + ::testing::AddGlobalTestEnvironment(new MLCommon::Test::opg::MPIEnvironment()); + + return RUN_ALL_TESTS(); +} diff --git a/cpp/tests/sg/rf_test.cu b/cpp/tests/sg/rf_test.cu index 2d32b72580..dfe40e55f5 100644 --- a/cpp/tests/sg/rf_test.cu +++ b/cpp/tests/sg/rf_test.cu @@ -748,6 +748,21 @@ TEST(RfTests, IntegerOverflow) handle.sync_stream_pool(); } +TEST(RfTests, EmptyGlobalRowsRejected) +{ + thrust::device_vector X(1); + thrust::device_vector y(1); + auto forest = std::make_shared>(); + auto forest_ptr = forest.get(); + auto stream_pool = std::make_shared(1); + raft::handle_t handle(rmm::cuda_stream_per_thread, stream_pool); + RF_params rf_params = + set_rf_params(3, 100, 1.0, 16, 1, 2, 0.0, false, 1, 1.0, 0, CRITERION::MSE, 1, 128); + + EXPECT_THROW(fit(handle, forest_ptr, X.data().get(), 0, 1, y.data().get(), rf_params), + raft::exception); +} + //------------------------------------------------------------------------------------------------------------------------------------- struct QuantileTestParameters { int n_rows;