Train distributed RF over global rows

cpp/CMakeLists.txt

@@ -219,6 +219,12 @@ if(BUILD_CUML_MG_TESTS AND NOT SINGLEGPU)
   set(BUILD_CUML_MPI_COMMS ON)
 endif()
 
+if(BUILD_CUML_MPI_COMMS)
+  find_package(MPI COMPONENTS CXX)
+  find_package(ucx REQUIRED)
+  find_package(ucxx REQUIRED)
+endif()
+
 if(USE_CCACHE)
   set(CMAKE_C_COMPILER_LAUNCHER ccache)
   set(CMAKE_CXX_COMPILER_LAUNCHER ccache)

cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh

@@ -17,6 +17,8 @@
 
 #include <cub/cub.cuh>
 
+#include <algorithm>
+
 namespace ML {
 namespace DT {
 
@@ -39,9 +41,8 @@ struct NodeWorkItem {
  */
 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 nodeid;          // Node in the batch on which the threadblock needs to work
+  IdxT large_nodeid;    // legacy field; histogram offsets now use nodeid for all nodes
   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
@@ -73,17 +74,27 @@ void launchNodeSplitKernel(const IdxT min_samples_leaf,
                            const NodeWorkItem* work_items,
                            const size_t work_items_size,
                            const Split<DataT, IdxT>* splits,
+                           Split<DataT, IdxT>* local_splits,
                            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.
@@ -371,25 +382,91 @@ template <typename DataT,
           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* colids,
-                              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 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* colids,
+                                       ObjectiveT& objective,
+                                       IdxT treeid,
+                                       const WorkloadInfo<IdxT>* workload_info,
+                                       uint64_t seed,
+                                       dim3 grid,
+                                       size_t smem_size,
+                                       cudaStream_t builder_stream);
+
+template <typename DataT,
+          typename LabelT,
+          typename IdxT,
+          int TPB,
+          typename ObjectiveT,
+          typename BinT>
+void launchEvaluateSplitKernel(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* colids,
+                               int* mutex,
+                               volatile Split<DataT, IdxT>* splits,
+                               ObjectiveT& objective,
+                               IdxT treeid,
+                               dim3 grid,
+                               size_t smem_size,
+                               cudaStream_t builder_stream);
+
+template <typename InT, typename OutT, typename OpT>
+static __global__ void transformHistogramKernel(const InT* in, OutT* out, std::size_t len, OpT op)
+{
+  std::size_t tid = blockIdx.x * blockDim.x + threadIdx.x;
+  for (std::size_t i = tid; i < len; i += std::size_t(blockDim.x) * gridDim.x) {
+    op(in, out, i);
+  }
+}
+
+inline std::size_t histogramTransformBlocks(std::size_t len)
+{
+  return std::max<std::size_t>(std::size_t{1}, raft::ceildiv<std::size_t>(len, 256));
+}
+
+inline void packHistograms(const CountBin* in, int* out, std::size_t len, cudaStream_t stream)
+{
+  auto op = [] __device__(const CountBin* in, int* out, std::size_t i) { out[i] = in[i].x; };
+  transformHistogramKernel<<<histogramTransformBlocks(len), 256, 0, stream>>>(in, out, len, op);
+}
+
+inline void unpackHistograms(const int* in, CountBin* out, std::size_t len, cudaStream_t stream)
+{
+  auto op = [] __device__(const int* in, CountBin* out, std::size_t i) { out[i].x = in[i]; };
+  transformHistogramKernel<<<histogramTransformBlocks(len), 256, 0, stream>>>(in, out, len, op);
+}
+
+inline void packHistograms(const AggregateBin* in,
+                           double* out,
+                           std::size_t len,
+                           cudaStream_t stream)
+{
+  auto op = [] __device__(const AggregateBin* in, double* out, std::size_t i) {
+    out[2 * i]     = in[i].label_sum;
+    out[2 * i + 1] = static_cast<double>(in[i].count);
+  };
+  transformHistogramKernel<<<histogramTransformBlocks(len), 256, 0, stream>>>(in, out, len, op);
+}
+
+inline void unpackHistograms(const double* in,
+                             AggregateBin* out,
+                             std::size_t len,
+                             cudaStream_t stream)
+{
+  auto op = [] __device__(const double* in, AggregateBin* out, std::size_t i) {
+    out[i].label_sum = in[2 * i];
+    out[i].count     = static_cast<int>(in[2 * i + 1]);
+  };
+  transformHistogramKernel<<<histogramTransformBlocks(len), 256, 0, stream>>>(in, out, len, op);
+}
 
 }  // namespace DT
 }  // namespace ML