[EM] Have one partitioner for each batch. (dmlc#10760)

- Initialize one partitioner for each batch.
- Collect partition size during initialization.
- Support base ridx in the finalization.

Author: trivialfis

src/common/device_helpers.cuh

@@ -387,11 +387,6 @@ void CopyTo(Src const &src, Dst *dst) {
                                 src.size() * sizeof(SVT), cudaMemcpyDefault));
 }
 
-template <class HContainer, class DContainer>
-void CopyToD(HContainer const &h, DContainer *d) {
-  CopyTo(h, d);
-}
-
 // Keep track of pinned memory allocation
 struct PinnedMemory {
   void *temp_storage{nullptr};

src/common/threading_utils.cc

@@ -124,15 +124,15 @@ void NameThread(std::thread* t, StringView name) {
   char old[16];
   auto ret = pthread_getname_np(handle, old, 16);
   if (ret != 0) {
-    LOG(WARNING) << "Failed to get the name from thread";
+    LOG(DEBUG) << "Failed to get the name from thread";
   }
   auto new_name = std::string{old} + ">" + name.c_str();  // NOLINT
   if (new_name.size() > 15) {
     new_name = new_name.substr(new_name.size() - 15);
   }
   ret = pthread_setname_np(handle, new_name.c_str());
   if (ret != 0) {
-    LOG(WARNING) << "Failed to name thread:" << ret << " :" << new_name;
+    LOG(DEBUG) << "Failed to name thread:" << ret << " :" << new_name;
   }
 #else
   (void)name;

src/tree/gpu_hist/gradient_based_sampler.cu

@@ -152,7 +152,7 @@ NoSampling::NoSampling(BatchParam batch_param) : batch_param_(std::move(batch_pa
 
 GradientBasedSample NoSampling::Sample(Context const*, common::Span<GradientPair> gpair,
                                        DMatrix* dmat) {
-  return {dmat->Info().num_row_, dmat, gpair};
+  return {dmat, gpair};
 }
 
 ExternalMemoryNoSampling::ExternalMemoryNoSampling(BatchParam batch_param)
@@ -179,7 +179,7 @@ GradientBasedSample ExternalMemoryNoSampling::Sample(Context const* ctx,
     this->p_fmat_new_ =
         std::make_unique<data::IterativeDMatrix>(new_page, p_fmat->Info(), batch_param_);
   }
-  return {p_fmat->Info().num_row_, this->p_fmat_new_.get(), gpair};
+  return {this->p_fmat_new_.get(), gpair};
 }
 
 UniformSampling::UniformSampling(BatchParam batch_param, float subsample)
@@ -192,7 +192,7 @@ GradientBasedSample UniformSampling::Sample(Context const* ctx, common::Span<Gra
   thrust::replace_if(cuctx->CTP(), dh::tbegin(gpair), dh::tend(gpair),
                      thrust::counting_iterator<std::size_t>(0),
                      BernoulliTrial(common::GlobalRandom()(), subsample_), GradientPair());
-  return {p_fmat->Info().num_row_, p_fmat, gpair};
+  return {p_fmat, gpair};
 }
 
 ExternalMemoryUniformSampling::ExternalMemoryUniformSampling(size_t n_rows,
@@ -252,7 +252,8 @@ GradientBasedSample ExternalMemoryUniformSampling::Sample(Context const* ctx,
   // Create the new DMatrix
   this->p_fmat_new_ = std::make_unique<data::IterativeDMatrix>(
       new_page, dmat->Info().Slice(ctx, dh::ToSpan(compact_row_index_), nnz), batch_param_);
-  return {sample_rows, this->p_fmat_new_.get(), dh::ToSpan(gpair_)};
+  CHECK_EQ(sample_rows, this->p_fmat_new_->Info().num_row_);
+  return {this->p_fmat_new_.get(), dh::ToSpan(gpair_)};
 }
 
 GradientBasedSampling::GradientBasedSampling(std::size_t n_rows, BatchParam batch_param,
@@ -274,7 +275,7 @@ GradientBasedSample GradientBasedSampling::Sample(Context const* ctx,
                     thrust::counting_iterator<size_t>(0), dh::tbegin(gpair),
                     PoissonSampling(dh::ToSpan(threshold_), threshold_index,
                                     RandomWeight(common::GlobalRandom()())));
-  return {n_rows, dmat, gpair};
+  return {dmat, gpair};
 }
 
 ExternalMemoryGradientBasedSampling::ExternalMemoryGradientBasedSampling(size_t n_rows,
@@ -334,7 +335,8 @@ GradientBasedSample ExternalMemoryGradientBasedSampling::Sample(Context const* c
   // Create the new DMatrix
   this->p_fmat_new_ = std::make_unique<data::IterativeDMatrix>(
       new_page, dmat->Info().Slice(ctx, dh::ToSpan(compact_row_index_), nnz), batch_param_);
-  return {sample_rows, this->p_fmat_new_.get(), dh::ToSpan(gpair_)};
+  CHECK_EQ(sample_rows, this->p_fmat_new_->Info().num_row_);
+  return {this->p_fmat_new_.get(), dh::ToSpan(gpair_)};
 }
 
 GradientBasedSampler::GradientBasedSampler(Context const* /*ctx*/, size_t n_rows,

src/tree/gpu_hist/gradient_based_sampler.cuh

@@ -12,11 +12,9 @@
 
 namespace xgboost::tree {
 struct GradientBasedSample {
-  /*!\brief Number of sampled rows. */
-  bst_idx_t sample_rows;
-  /*!\brief Sampled rows in ELLPACK format. */
+  /** @brief Sampled rows in ELLPACK format. */
   DMatrix* p_fmat;
-  /*!\brief Gradient pairs for the sampled rows. */
+  /** @brief Gradient pairs for the sampled rows. */
   common::Span<GradientPair const> gpair;
 };
 

src/tree/gpu_hist/row_partitioner.cu

@@ -31,7 +31,7 @@ common::Span<const RowPartitioner::RowIndexT> RowPartitioner::GetRows(bst_node_t
   return dh::ToSpan(ridx_).subspan(segment.begin, segment.Size());
 }
 
-common::Span<const RowPartitioner::RowIndexT> RowPartitioner::GetRows() {
+common::Span<const RowPartitioner::RowIndexT> RowPartitioner::GetRows() const {
   return dh::ToSpan(ridx_);
 }
 

src/tree/gpu_hist/row_partitioner.cuh

@@ -200,11 +200,11 @@ XGBOOST_DEV_INLINE int GetPositionFromSegments(std::size_t idx,
 
 template <int kBlockSize, typename RowIndexT, typename OpT>
 __global__ __launch_bounds__(kBlockSize) void FinalisePositionKernel(
-    const common::Span<const NodePositionInfo> d_node_info,
+    const common::Span<const NodePositionInfo> d_node_info, bst_idx_t base_ridx,
     const common::Span<const RowIndexT> d_ridx, common::Span<bst_node_t> d_out_position, OpT op) {
   for (auto idx : dh::GridStrideRange<std::size_t>(0, d_ridx.size())) {
     auto position = GetPositionFromSegments(idx, d_node_info.data());
-    RowIndexT ridx = d_ridx[idx];
+    RowIndexT ridx = d_ridx[idx] - base_ridx;
     bst_node_t new_position = op(ridx, position);
     d_out_position[ridx] = new_position;
   }
@@ -264,7 +264,12 @@ class RowPartitioner {
   /**
    * \brief Gets all training rows in the set.
    */
-  common::Span<const RowIndexT> GetRows();
+  common::Span<const RowIndexT> GetRows() const;
+  /**
+   * @brief Get the number of rows in this partitioner.
+   */
+  std::size_t Size() const { return this->GetRows().size(); }
+
   [[nodiscard]] bst_node_t GetNumNodes() const { return n_nodes_; }
 
   /**
@@ -351,7 +356,8 @@ class RowPartitioner {
    *           argument and return the new position for this training instance.
    */
   template <typename FinalisePositionOpT>
-  void FinalisePosition(common::Span<bst_node_t> d_out_position, FinalisePositionOpT op) const {
+  void FinalisePosition(common::Span<bst_node_t> d_out_position, bst_idx_t base_ridx,
+                        FinalisePositionOpT op) const {
     dh::TemporaryArray<NodePositionInfo> d_node_info_storage(ridx_segments_.size());
     dh::safe_cuda(cudaMemcpyAsync(d_node_info_storage.data().get(), ridx_segments_.data(),
                                   sizeof(NodePositionInfo) * ridx_segments_.size(),
@@ -361,8 +367,8 @@ class RowPartitioner {
     const int kItemsThread = 8;
     const int grid_size = xgboost::common::DivRoundUp(ridx_.size(), kBlockSize * kItemsThread);
     common::Span<RowIndexT const> d_ridx{ridx_.data(), ridx_.size()};
-    FinalisePositionKernel<kBlockSize>
-        <<<grid_size, kBlockSize, 0>>>(dh::ToSpan(d_node_info_storage), d_ridx, d_out_position, op);
+    FinalisePositionKernel<kBlockSize><<<grid_size, kBlockSize, 0>>>(
+        dh::ToSpan(d_node_info_storage), base_ridx, d_ridx, d_out_position, op);
   }
 };
 };  // namespace xgboost::tree