Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
26 changes: 26 additions & 0 deletions cpp/include/cuml/ensemble/randomforest_mg_utils.hpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,26 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*/

#pragma once

#include <raft/core/comms.hpp>
#include <raft/core/handle.hpp>
#include <raft/core/resource/comms.hpp>
#include <raft/core/resource/cuda_stream.hpp>

namespace ML::detail {

inline void cuml_rf_allreduce_validation_status(const raft::handle_t& handle,
const int* local_status,
int* global_status)
Comment on lines +15 to +17

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

I'm putting this in a separate header. If this function had gone into cuml/ensemble/randomforest.hpp, we'd need to include raft/core/handle.hpp and increate the compilation time.

{
auto const& comm = raft::resource::get_comms(handle);
cudaStream_t stream = raft::resource::get_cuda_stream(handle);
comm.allreduce(local_status, global_status, 1, raft::comms::op_t::MAX, stream);
RAFT_EXPECTS(comm.sync_stream(stream) == raft::comms::status_t::SUCCESS,
"Input validation status all-reduce failed");
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.

} // namespace ML::detail
10 changes: 6 additions & 4 deletions cpp/src/randomforest/randomforest.cuh
Original file line number Diff line number Diff line change
Expand Up @@ -330,10 +330,12 @@ class RandomForest {
// 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<std::size_t>(n_streams) <= handle.get_stream_pool_size(),
"effective RF n_streams (=%d) should be <= raft::handle_t.n_streams (=%lu)",
n_streams,
handle.get_stream_pool_size());
auto stream_pool_size = handle.get_stream_pool_size();
if (static_cast<std::size_t>(n_streams) > stream_pool_size) {
CUML_LOG_WARN("Resizing n_streams to fit the available stream pool size (%lu)",
stream_pool_size);
n_streams = ML::narrow_cast<int>(stream_pool_size);
}

auto quantile_result = DT::computeQuantiles(handle,
input,
Expand Down
Loading
Loading