Handle RF split histogram shared-memory pressure

cpp/src/decisiontree/batched-levelalgo/builder.cuh

@@ -149,6 +149,11 @@ struct Builder {
 
   /** default threads per block for most kernels in here */
   static constexpr int TPB_DEFAULT = 128;
+  // Tunable performance heuristic for the shared-memory histogram path. Large per-block
+  // histograms, usually from large n_classes, can reduce occupancy enough that global memory is
+  // faster even when the histogram fits in shared memory. 16 KiB keeps small/default histograms in
+  // shared memory while avoiding the large-class shared-memory slowdown measured locally.
+  static constexpr size_t tunable_split_histogram_dynamic_smem_limit_bytes = 16 * 1024;
   /** handle to get device properties */
   const raft::handle_t& handle;
   /** stream to launch kernels */
@@ -487,11 +492,12 @@ struct Builder {
     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 split_smem_config              = computeSplitSharedMemoryConfig();
 
     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, n_large_nodes, work_items.size(), split_smem_config);
       RAFT_CUDA_TRY(cudaPeekAtLastError());
     }
     raft::update_host(h_splits, splits, work_items.size(), builder_stream);
@@ -515,72 +521,100 @@ struct Builder {
     RAFT_CUDA_TRY(cudaPeekAtLastError());
   }
 
-  auto computeSplitSmemSize()
+  struct SplitSharedMemoryConfig {
+    bool use_global_memory_histogram;
+    size_t dynamic_smem_size;
+  };
+
+  SplitSharedMemoryConfig computeSplitSharedMemoryConfig() const
   {
+    // Dynamic shared memory for the fast path: histogram, copied quantiles, and
+    // alignment padding for the kernel's shared-memory layout.
     auto shared_histogram_size =
       ML::checked_mul<std::size_t>(params.max_n_bins, dataset.num_outputs, sizeof(BinT));
     auto shared_quantiles_size = ML::checked_mul<std::size_t>(params.max_n_bins, sizeof(DataT));
-    auto dynamic_smem_size =
+    auto shared_dynamic_smem_size =
       ML::checked_add<std::size_t>(shared_histogram_size, shared_quantiles_size, sizeof(int));
-
-    // Extra room for alignment (see alignPointer in
-    // computeSplitKernel)
     auto alignment_smem_size =
       ML::checked_add<std::size_t>(sizeof(DataT), ML::checked_mul<std::size_t>(3, sizeof(int)));
-    dynamic_smem_size = ML::checked_add<std::size_t>(dynamic_smem_size, alignment_smem_size);
+    shared_dynamic_smem_size =
+      ML::checked_add<std::size_t>(shared_dynamic_smem_size, alignment_smem_size);
+
+    // Dynamic shared memory for the fallback path only needs the per-block done
+    // flag used by the inter-block completion handshake.
+    auto global_dynamic_smem_size = ML::checked_add<std::size_t>(sizeof(int), sizeof(int));
 
-    // computeSplitKernel also reserves static shared memory for CUB's scan temp
-    // storage and the per-warp split reduction scratch.
+    // Static shared memory is reserved by the kernel regardless of where the
+    // histogram lives.
     auto cdf_scan_smem_size = sizeof(typename cub::BlockScan<BinT, TPB_DEFAULT>::TempStorage);
     auto split_scratch_smem_size =
       ML::checked_mul<std::size_t>(raft::ceildiv(TPB_DEFAULT, raft::WarpSize), sizeof(SplitT));
-    auto total_smem_size =
-      ML::checked_add<std::size_t>(dynamic_smem_size, cdf_scan_smem_size, split_scratch_smem_size);
-    auto available_smem = handle.get_device_properties().sharedMemPerBlock;
-    ASSERT(available_smem >= total_smem_size,
-           "Not enough shared memory. Consider reducing max_n_bins.");
-    return dynamic_smem_size;
+    auto static_smem_size =
+      ML::checked_add<std::size_t>(cdf_scan_smem_size, split_scratch_smem_size);
+
+    auto available_smem = size_t(handle.get_device_properties().sharedMemPerBlock);
+    auto global_total_smem_size =
+      ML::checked_add<std::size_t>(global_dynamic_smem_size, static_smem_size);
+    ASSERT(available_smem >= global_total_smem_size,
+           "Not enough shared memory for RF split bookkeeping.");
+
+    // Prefer shared memory when it fits and stays small enough for good occupancy;
+    // otherwise use the global histogram path to avoid launch failure or slowdown.
+    auto shared_total_smem_size =
+      ML::checked_add<std::size_t>(shared_dynamic_smem_size, static_smem_size);
+    bool use_global_memory_histogram =
+      shared_total_smem_size > available_smem ||
+      shared_dynamic_smem_size > tunable_split_histogram_dynamic_smem_limit_bytes;
+
+    return {use_global_memory_histogram,
+            use_global_memory_histogram ? global_dynamic_smem_size : shared_dynamic_smem_size};
   }
 
-  void computeSplit(IdxT col, size_t n_blocks_dimx, size_t n_large_nodes)
+  void computeSplit(IdxT col,
+                    size_t n_blocks_dimx,
+                    size_t n_large_nodes,
+                    size_t n_work_items,
+                    const SplitSharedMemoryConfig& split_smem_config)
   {
     // if no instances to split, return
     if (n_blocks_dimx == 0) return;
     raft::common::nvtx::range fun_scope("Builder::computeSplit @builder.cuh [batched-levelalgo]");
-    auto n_bins    = params.max_n_bins;
-    auto n_classes = dataset.num_outputs;
+    auto n_bins                      = params.max_n_bins;
+    auto n_classes                   = dataset.num_outputs;
+    auto use_global_memory_histogram = split_smem_config.use_global_memory_histogram;
     // if columns left to be processed lesser than `n_blks_for_cols`, shrink the blocks along dimy
     auto n_blocks_dimy = std::min(n_blks_for_cols, dataset.n_sampled_cols - col);
-    // compute required dynamic shared memory
-    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;
-    RAFT_CUDA_TRY(cudaMemsetAsync(histograms, 0, sizeof(BinT) * len_histograms, builder_stream));
+    auto histogram_node_count = use_global_memory_histogram ? n_work_items : n_large_nodes;
+    auto len_histograms =
+      ML::checked_mul<std::size_t>(n_bins, n_classes, n_blocks_dimy, histogram_node_count);
+    auto histograms_bytes = ML::checked_mul<std::size_t>(sizeof(BinT), len_histograms);
+    RAFT_CUDA_TRY(cudaMemsetAsync(histograms, 0, histograms_bytes, builder_stream));
     // create the objective function object
     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<DataT, LabelT, IdxT, TPB_DEFAULT, ObjectiveT>(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);
+    launchComputeSplitKernel<DataT, LabelT, IdxT, TPB_DEFAULT, ObjectiveT>(
+      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,
+      use_global_memory_histogram,
+      grid,
+      split_smem_config.dynamic_smem_size,
+      builder_stream);
   }
 
   // Set the leaf value predictions in batch

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

@@ -126,7 +126,7 @@ void launchLeafKernel(ObjectiveT objective,
 // 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.
 template <typename DataT, typename IdxT>
-HDI IdxT lower_bound(DataT* array, IdxT len, DataT element)
+HDI IdxT lower_bound(DataT const* array, IdxT len, DataT element)
 {
   IdxT start = 0;
   IdxT end   = len - 1;
@@ -159,6 +159,7 @@ void launchComputeSplitKernel(typename ObjectiveT::BinT* histograms,
                               IdxT treeid,
                               const WorkloadInfo<IdxT>* workload_info,
                               uint64_t seed,
+                              bool use_global_memory_histogram,
                               dim3 grid,
                               size_t smem_size,
                               cudaStream_t builder_stream);

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

@@ -223,14 +223,14 @@ void launchLeafKernel(ObjectiveT objective,
 }
 
 /**
- * @brief For every threadblock, converts the smem pdf-histogram to
+ * @brief For every threadblock, converts a pdf-histogram to a
  *        cdf-histogram inplace using inclusive block-sum-scan and returns
  *        the total_sum
  * @return The total sum aggregated over the sumscan,
  *         as well as the modified cdf-histogram pointer
  */
 template <typename BinT, typename IdxT, int TPB>
-DI BinT pdf_to_cdf(BinT* shared_histogram, IdxT n_bins)
+DI BinT pdf_to_cdf(BinT* histogram, IdxT n_bins)
 {
   // Blockscan instance preparation
   typedef cub::BlockScan<BinT, TPB> BlockScan;
@@ -242,10 +242,10 @@ DI BinT pdf_to_cdf(BinT* shared_histogram, IdxT n_bins)
   for (IdxT tix = threadIdx.x; tix < raft::ceildiv(n_bins, TPB) * TPB; tix += blockDim.x) {
     BinT result;
     BinT block_aggregate;
-    BinT element = tix < n_bins ? shared_histogram[tix] : BinT();
+    BinT element = tix < n_bins ? histogram[tix] : BinT();
     BlockScan(temp_storage).InclusiveSum(element, result, block_aggregate);
     __syncthreads();
-    if (tix < n_bins) { shared_histogram[tix] = result + total_aggregate; }
+    if (tix < n_bins) { histogram[tix] = result + total_aggregate; }
     total_aggregate += block_aggregate;
   }
   // return the total sum
@@ -268,7 +268,8 @@ static __global__ void computeSplitKernel(typename ObjectiveT::BinT* histograms,
                                           ObjectiveT objective,
                                           IdxT treeid,
                                           const WorkloadInfo<IdxT>* workload_info,
-                                          uint64_t seed)
+                                          uint64_t seed,
+                                          bool use_global_memory_histogram)
 {
   using BinT = typename ObjectiveT::BinT;
   // dynamic shared memory
@@ -296,24 +297,35 @@ static __global__ void computeSplitKernel(typename ObjectiveT::BinT* histograms,
   // getting the n_bins for that feature
   int n_bins = quantiles.n_bins_array[col];
 
+  auto n_classes            = objective.NumClasses();
   auto end                  = range_start + range_len;
-  auto shared_histogram_len = n_bins * objective.NumClasses();
-  auto* shared_histogram    = alignPointer<BinT>(smem);
-  auto* shared_quantiles    = alignPointer<DataT>(shared_histogram + shared_histogram_len);
-  auto* shared_done         = alignPointer<int>(shared_quantiles + n_bins);
+  auto histogram_len        = n_bins * n_classes;
+  auto* histogram           = static_cast<BinT*>(nullptr);
+  auto* shared_done         = static_cast<int*>(nullptr);
+  auto* quantiles_for_split = quantiles.quantiles_array + std::size_t(max_n_bins) * col;
   IdxT stride               = blockDim.x * num_blocks;
   IdxT tid                  = threadIdx.x + offset_blockid * blockDim.x;
 
-  // populating shared memory with initial values
-  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)
-    shared_quantiles[b] = quantiles.quantiles_array[max_n_bins * col + b];
+  if (use_global_memory_histogram) {
+    auto histograms_offset = (std::size_t(nid) * gridDim.y + blockIdx.y) * max_n_bins * n_classes;
+    histogram              = histograms + histograms_offset;
+    shared_done            = alignPointer<int>(smem);
+  } else {
+    histogram              = alignPointer<BinT>(smem);
+    auto* shared_quantiles = alignPointer<DataT>(histogram + histogram_len);
+    shared_done            = alignPointer<int>(shared_quantiles + n_bins);
+    quantiles_for_split    = shared_quantiles;
+    for (IdxT i = threadIdx.x; i < histogram_len; i += blockDim.x) {
+      histogram[i] = BinT();
+    }
+    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
+  // compute pdf histogram for all bins for all classes
 
   // Must be 64 bit - can easily grow larger than a 32 bit int
   std::size_t col_offset = std::size_t(col) * dataset.M;
@@ -323,20 +335,28 @@ static __global__ void computeSplitKernel(typename ObjectiveT::BinT* histograms,
     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]
-    objective.IncrementHistogram(shared_histogram, n_bins, start, label, dataset, row);
+    // `start` is lowest index such that data <= quantiles_for_split[start]
+    IdxT start = lower_bound(quantiles_for_split, n_bins, data);
+    // ++histogram[start]
+    objective.IncrementHistogram(histogram, n_bins, start, label, dataset, row);
   }
 
-  // synchronizing above changes across block
   __syncthreads();
-  if (num_blocks > 1) {
+  if (use_global_memory_histogram) {
+    __threadfence();  // for commit guarantee before the last block scores the split
+    __syncthreads();
+
+    bool last = MLCommon::signalDone(
+      done_count + nid * gridDim.y + blockIdx.y, num_blocks, offset_blockid == 0, shared_done);
+    if (!last) return;
+  } else if (num_blocks > 1) {
+    // Shared-memory histogram path: each block built a partial histogram, so unify those
+    // partial histograms in global memory before scoring the split.
     // 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]);
+      (std::size_t(large_nid) * gridDim.y + blockIdx.y) * max_n_bins * n_classes;
+    for (IdxT i = threadIdx.x; i < histogram_len; i += blockDim.x) {
+      BinT::AtomicAdd(histograms + histograms_offset + i, histogram[i]);
     }
 
     __threadfence();  // for commit guarantee
@@ -349,34 +369,34 @@ static __global__ void computeSplitKernel(typename ObjectiveT::BinT* histograms,
     if (!last) return;
 
     // 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];
+    for (IdxT i = threadIdx.x; i < histogram_len; i += blockDim.x) {
+      histogram[i] = histograms[histograms_offset + i];
+    }
 
     __syncthreads();
   }
 
-  // PDF to CDF inplace in `shared_histogram`
-  for (IdxT c = 0; c < objective.NumClasses(); ++c) {
+  // PDF to CDF inplace in `histogram`
+  for (IdxT c = 0; c < n_classes; ++c) {
     // left to right scan operation for scanning
     // "lesser-than-or-equal" counts
-    BinT total_sum = pdf_to_cdf<BinT, IdxT, TPB>(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]`.
+    BinT total_sum = pdf_to_cdf<BinT, IdxT, TPB>(histogram + n_bins * c, n_bins);
+    // now, `histogram[n_bins * c + i]` will have count of datapoints of class `c`
+    // that are less than or equal to `quantiles_for_split[i]`.
   }
 
   __syncthreads();
 
   // calculate the best candidate bins (one for each thread in the block) in current feature and
   // corresponding information gain for splitting
-  Split<DataT, IdxT> sp =
-    objective.Gain(shared_histogram, shared_quantiles, col, range_len, n_bins);
+  Split<DataT, IdxT> sp = objective.Gain(histogram, quantiles_for_split, col, range_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(split_scratch, splits + nid, mutex + nid, shared_quantiles, n_bins);
+  sp.evalBestSplit(split_scratch, splits + nid, mutex + nid, quantiles_for_split, n_bins);
 }
 
 template <typename DataT, typename LabelT, typename IdxT, int TPB, typename ObjectiveT>
@@ -396,6 +416,7 @@ void launchComputeSplitKernel(typename ObjectiveT::BinT* histograms,
                               IdxT treeid,
                               const WorkloadInfo<IdxT>* workload_info,
                               uint64_t seed,
+                              bool use_global_memory_histogram,
                               dim3 grid,
                               size_t smem_size,
                               cudaStream_t builder_stream)
@@ -416,7 +437,8 @@ void launchComputeSplitKernel(typename ObjectiveT::BinT* histograms,
                                                objective,
                                                treeid,
                                                workload_info,
-                                               seed);
+                                               seed,
+                                               use_global_memory_histogram);
 }
 
 }  // namespace DT

cpp/src/decisiontree/batched-levelalgo/kernels/classification-double.cu

@@ -46,6 +46,7 @@ template void launchComputeSplitKernel<DataT, LabelT, IdxT, TPB_DEFAULT, Objecti
   IdxT treeid,
   const WorkloadInfo<IdxT>* workload_info,
   uint64_t seed,
+  bool use_global_memory_histogram,
   dim3 grid,
   size_t smem_size,
   cudaStream_t builder_stream);

cpp/src/decisiontree/batched-levelalgo/kernels/classification-float.cu

@@ -46,6 +46,7 @@ template void launchComputeSplitKernel<DataT, LabelT, IdxT, TPB_DEFAULT, Objecti
   IdxT treeid,
   const WorkloadInfo<IdxT>* workload_info,
   uint64_t seed,
+  bool use_global_memory_histogram,
   dim3 grid,
   size_t smem_size,
   cudaStream_t builder_stream);

cpp/src/decisiontree/batched-levelalgo/kernels/regression-double.cu

@@ -46,6 +46,7 @@ template void launchComputeSplitKernel<DataT, LabelT, IdxT, TPB_DEFAULT, Objecti
   IdxT treeid,
   const WorkloadInfo<IdxT>* workload_info,
   uint64_t seed,
+  bool use_global_memory_histogram,
   dim3 grid,
   size_t smem_size,
   cudaStream_t builder_stream);